tokenfirewall 1.0.1

package/dist/core/budgetManager.js

@@ -0,0 +1,116 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.BudgetManager = void 0;
+/**
+ * Budget manager - encapsulates all budget tracking state
+ */
+class BudgetManager {
+    constructor(options) {
+        this.totalSpent = 0;
+        this.trackingLock = Promise.resolve();
+        // Validate monthly limit
+        if (typeof options.monthlyLimit !== 'number' || isNaN(options.monthlyLimit) || !isFinite(options.monthlyLimit)) {
+            throw new Error('TokenFirewall: monthlyLimit must be a valid number');
+        }
+        if (options.monthlyLimit <= 0) {
+            throw new Error('TokenFirewall: monthlyLimit must be greater than 0');
+        }
+        // Validate mode
+        const mode = options.mode ?? "block";
+        if (mode !== "warn" && mode !== "block") {
+            throw new Error('TokenFirewall: mode must be either "warn" or "block"');
+        }
+        this.limit = options.monthlyLimit;
+        this.mode = mode;
+    }
+    /**
+     * Track a new cost and enforce budget limits
+     * Uses locking to prevent race conditions
+     * @throws Error if budget exceeded in block mode
+     */
+    async track(cost) {
+        // Validate cost parameter
+        if (typeof cost !== 'number' || isNaN(cost) || !isFinite(cost)) {
+            throw new Error('TokenFirewall: Cost must be a valid number');
+        }
+        if (cost < 0) {
+            throw new Error('TokenFirewall: Cost cannot be negative');
+        }
+        // Wait for any pending tracking to complete
+        await this.trackingLock;
+        // Create new lock for this tracking operation
+        let releaseLock;
+        this.trackingLock = new Promise((resolve) => {
+            releaseLock = resolve;
+        });
+        try {
+            // Check if this would exceed budget BEFORE adding
+            const projectedTotal = this.totalSpent + cost;
+            if (projectedTotal > this.limit) {
+                const message = `TokenFirewall: Budget exceeded! Would spend $${projectedTotal.toFixed(4)} of $${this.limit.toFixed(2)} limit`;
+                if (this.mode === "block") {
+                    // In block mode, DON'T track the cost and throw error
+                    throw new Error(message);
+                }
+                else {
+                    console.warn(message);
+                }
+            }
+            // Only commit the cost if we didn't throw
+            this.totalSpent += cost;
+        }
+        finally {
+            // Release the lock
+            releaseLock();
+        }
+    }
+    /**
+     * Get current budget status
+     */
+    getStatus() {
+        const remaining = Math.max(0, this.limit - this.totalSpent);
+        const percentageUsed = (this.totalSpent / this.limit) * 100;
+        return {
+            totalSpent: this.totalSpent,
+            limit: this.limit,
+            remaining,
+            percentageUsed,
+        };
+    }
+    /**
+     * Reset budget tracking (useful for monthly resets)
+     */
+    reset() {
+        this.totalSpent = 0;
+    }
+    /**
+     * Export budget state for persistence
+     */
+    exportState() {
+        return {
+            totalSpent: this.totalSpent,
+            limit: this.limit,
+            mode: this.mode,
+        };
+    }
+    /**
+     * Import budget state from persistence
+     * Validates state before importing
+     */
+    importState(state) {
+        // Validate totalSpent
+        if (typeof state.totalSpent !== 'number' || isNaN(state.totalSpent) || !isFinite(state.totalSpent)) {
+            throw new Error('TokenFirewall: Invalid budget state - totalSpent must be a valid number');
+        }
+        if (state.totalSpent < 0) {
+            throw new Error('TokenFirewall: Invalid budget state - totalSpent cannot be negative');
+        }
+        // Allow importing state that exceeds limit (user may have increased limit)
+        // Just warn if it's suspicious
+        if (state.totalSpent > this.limit * 10) {
+            console.warn(`TokenFirewall: Imported totalSpent (${state.totalSpent}) is much larger than current limit (${this.limit})`);
+        }
+        this.totalSpent = state.totalSpent;
+    }
+}
+exports.BudgetManager = BudgetManager;

package/dist/core/costEngine.d.ts

@@ -0,0 +1,6 @@
+import { NormalizedUsage, CostBreakdown } from "./types";
+/**
+ * Pure function to calculate cost from normalized usage
+ * Pricing is per 1M tokens
+ */
+export declare function calculateCost(usage: NormalizedUsage): CostBreakdown;

package/dist/core/costEngine.js

@@ -0,0 +1,30 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.calculateCost = calculateCost;
+const pricingRegistry_1 = require("./pricingRegistry");
+/**
+ * Pure function to calculate cost from normalized usage
+ * Pricing is per 1M tokens
+ */
+function calculateCost(usage) {
+    try {
+        const pricing = pricingRegistry_1.pricingRegistry.getPricing(usage.provider, usage.model);
+        const inputCost = (usage.inputTokens / 1000000) * pricing.input;
+        const outputCost = (usage.outputTokens / 1000000) * pricing.output;
+        const totalCost = inputCost + outputCost;
+        return {
+            inputCost,
+            outputCost,
+            totalCost,
+        };
+    }
+    catch (error) {
+        // If pricing not found, log warning and return zero cost
+        console.warn(`TokenFirewall: ${error instanceof Error ? error.message : 'Unknown pricing error'} - tracking with $0 cost`);
+        return {
+            inputCost: 0,
+            outputCost: 0,
+            totalCost: 0,
+        };
+    }
+}

package/dist/core/pricingRegistry.d.ts

@@ -0,0 +1,27 @@
+import { ModelPricing } from "./types";
+/**
+ * Centralized pricing registry for all LLM providers
+ */
+declare class PricingRegistry {
+    private pricing;
+    constructor();
+    /**
+     * Initialize default pricing for supported providers
+     */
+    private initializeDefaultPricing;
+    /**
+     * Register pricing for a provider and model
+     */
+    register(provider: string, model: string, pricing: ModelPricing): void;
+    /**
+     * Get pricing for a specific provider and model
+     * @throws Error if pricing not found
+     */
+    getPricing(provider: string, model: string): ModelPricing;
+    /**
+     * Check if pricing exists for a provider and model
+     */
+    hasPricing(provider: string, model: string): boolean;
+}
+export declare const pricingRegistry: PricingRegistry;
+export {};

package/dist/core/pricingRegistry.js

@@ -0,0 +1,75 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.pricingRegistry = void 0;
+/**
+ * Centralized pricing registry for all LLM providers
+ */
+class PricingRegistry {
+    constructor() {
+        this.pricing = new Map();
+        this.initializeDefaultPricing();
+    }
+    /**
+     * Initialize default pricing for supported providers
+     */
+    initializeDefaultPricing() {
+        // OpenAI pricing (per 1M tokens)
+        this.register("openai", "gpt-4o", { input: 2.5, output: 10.0 });
+        this.register("openai", "gpt-4o-mini", { input: 0.15, output: 0.6 });
+        this.register("openai", "gpt-4-turbo", { input: 10.0, output: 30.0 });
+        this.register("openai", "gpt-3.5-turbo", { input: 0.5, output: 1.5 });
+        // Anthropic pricing (per 1M tokens)
+        this.register("anthropic", "claude-3-5-sonnet-20241022", { input: 3.0, output: 15.0 });
+        this.register("anthropic", "claude-3-5-haiku-20241022", { input: 0.8, output: 4.0 });
+        this.register("anthropic", "claude-3-opus-20240229", { input: 15.0, output: 75.0 });
+        // Google Gemini pricing (per 1M tokens)
+        this.register("gemini", "gemini-2.0-flash-exp", { input: 0.0, output: 0.0 });
+        this.register("gemini", "gemini-1.5-pro", { input: 1.25, output: 5.0 });
+        this.register("gemini", "gemini-1.5-flash", { input: 0.075, output: 0.3 });
+        // Grok pricing (per 1M tokens)
+        this.register("grok", "grok-beta", { input: 5.0, output: 15.0 });
+        this.register("grok", "grok-vision-beta", { input: 5.0, output: 15.0 });
+        this.register("grok", "grok-2-1212", { input: 2.0, output: 10.0 });
+        this.register("grok", "grok-2-vision-1212", { input: 2.0, output: 10.0 });
+        // Llama models via Grok API (per 1M tokens)
+        this.register("grok", "llama-3.1-70b", { input: 0.5, output: 0.8 });
+        this.register("grok", "llama-3.1-8b", { input: 0.1, output: 0.1 });
+        this.register("grok", "llama-3.2-90b-vision", { input: 0.6, output: 0.9 });
+        this.register("grok", "llama-3.3-70b", { input: 0.5, output: 0.8 });
+        // Kimi pricing (per 1M tokens)
+        this.register("kimi", "moonshot-v1-8k", { input: 0.12, output: 0.12 });
+        this.register("kimi", "moonshot-v1-32k", { input: 0.24, output: 0.24 });
+        this.register("kimi", "moonshot-v1-128k", { input: 0.6, output: 0.6 });
+    }
+    /**
+     * Register pricing for a provider and model
+     */
+    register(provider, model, pricing) {
+        if (!this.pricing.has(provider)) {
+            this.pricing.set(provider, new Map());
+        }
+        this.pricing.get(provider).set(model, pricing);
+    }
+    /**
+     * Get pricing for a specific provider and model
+     * @throws Error if pricing not found
+     */
+    getPricing(provider, model) {
+        const providerPricing = this.pricing.get(provider);
+        if (!providerPricing) {
+            throw new Error(`TokenFirewall: No pricing found for provider "${provider}"`);
+        }
+        const modelPricing = providerPricing.get(model);
+        if (!modelPricing) {
+            throw new Error(`TokenFirewall: No pricing found for model "${model}" from provider "${provider}"`);
+        }
+        return modelPricing;
+    }
+    /**
+     * Check if pricing exists for a provider and model
+     */
+    hasPricing(provider, model) {
+        return this.pricing.get(provider)?.has(model) ?? false;
+    }
+}
+exports.pricingRegistry = new PricingRegistry();

package/dist/core/types.d.ts

@@ -0,0 +1,66 @@
+/**
+ * Unified usage contract that all providers must normalize to
+ */
+export interface NormalizedUsage {
+    provider: string;
+    model: string;
+    inputTokens: number;
+    outputTokens: number;
+    totalTokens: number;
+}
+/**
+ * Cost breakdown for a single LLM call
+ */
+export interface CostBreakdown {
+    inputCost: number;
+    outputCost: number;
+    totalCost: number;
+}
+/**
+ * Provider adapter interface - all adapters must implement this
+ */
+export interface ProviderAdapter {
+    name: string;
+    detect(response: unknown): boolean;
+    normalize(response: unknown, request?: unknown): NormalizedUsage;
+}
+/**
+ * Pricing configuration for a specific model
+ */
+export interface ModelPricing {
+    input: number;
+    output: number;
+}
+/**
+ * Budget guard configuration
+ */
+export interface BudgetGuardOptions {
+    monthlyLimit: number;
+    mode?: "warn" | "block";
+}
+/**
+ * Budget status information
+ */
+export interface BudgetStatus {
+    totalSpent: number;
+    limit: number;
+    remaining: number;
+    percentageUsed: number;
+}
+/**
+ * Model information with context limits and budget usage
+ */
+export interface ModelInfo {
+    model: string;
+    contextLimit?: number;
+    budgetUsagePercentage?: number;
+}
+/**
+ * Options for listing available models
+ */
+export interface ListModelsOptions {
+    provider: string;
+    apiKey: string;
+    baseURL?: string;
+    includeBudgetUsage?: boolean;
+}

package/dist/core/types.js

@@ -0,0 +1,2 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });

package/dist/index.d.ts

@@ -0,0 +1,71 @@
+import { BudgetGuardOptions, ProviderAdapter, ModelPricing, BudgetStatus } from "./core/types";
+import { BudgetManager } from "./core/budgetManager";
+import { patchGlobalFetch } from "./interceptors/fetchInterceptor";
+import { patchProvider } from "./interceptors/sdkInterceptor";
+import { listAvailableModels, ModelInfo, ListModelsOptions } from "./introspection/modelLister";
+/**
+ * Create and configure a budget guard
+ * @param options - Budget configuration options
+ * @returns Budget manager instance
+ */
+export declare function createBudgetGuard(options: BudgetGuardOptions): BudgetManager;
+/**
+ * Patch global fetch to intercept LLM API calls
+ */
+export { patchGlobalFetch };
+/**
+ * Patch a specific provider SDK
+ * @param providerName - Name of the provider to patch
+ */
+export { patchProvider };
+/**
+ * Register a custom provider adapter
+ * @param adapter - Provider adapter implementation
+ */
+export declare function registerAdapter(adapter: ProviderAdapter): void;
+/**
+ * Register custom pricing for a provider and model
+ * @param provider - Provider name
+ * @param model - Model name
+ * @param pricing - Pricing configuration (per 1M tokens)
+ */
+export declare function registerPricing(provider: string, model: string, pricing: ModelPricing): void;
+/**
+ * Register custom context limit for a provider and model
+ * @param provider - Provider name
+ * @param model - Model name
+ * @param contextLimit - Context window size in tokens
+ */
+export declare function registerContextLimit(provider: string, model: string, contextLimit: number): void;
+/**
+ * Get current budget status
+ * @returns Budget status or null if no budget guard is active
+ */
+export declare function getBudgetStatus(): BudgetStatus | null;
+/**
+ * Reset budget tracking
+ */
+export declare function resetBudget(): void;
+/**
+ * Export budget state for persistence
+ */
+export declare function exportBudgetState(): {
+    totalSpent: number;
+    limit: number;
+    mode: string;
+} | null;
+/**
+ * Import budget state from persistence
+ */
+export declare function importBudgetState(state: {
+    totalSpent: number;
+}): void;
+/**
+ * List available models for a provider with context limits and budget usage
+ * @param options - Provider configuration and options
+ * @returns Array of model information
+ */
+export declare function listModels(options: Omit<ListModelsOptions, 'budgetManager'>): Promise<ModelInfo[]>;
+export { listAvailableModels };
+export type { BudgetGuardOptions, ProviderAdapter, ModelPricing, NormalizedUsage, CostBreakdown, BudgetStatus, ModelInfo, ListModelsOptions, } from "./core/types";
+export type { ModelInfo as ModelInfoType, ListModelsOptions as ListModelsOptionsType } from "./introspection/modelLister";

package/dist/index.js

@@ -0,0 +1,134 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.listAvailableModels = exports.patchProvider = exports.patchGlobalFetch = void 0;
+exports.createBudgetGuard = createBudgetGuard;
+exports.registerAdapter = registerAdapter;
+exports.registerPricing = registerPricing;
+exports.registerContextLimit = registerContextLimit;
+exports.getBudgetStatus = getBudgetStatus;
+exports.resetBudget = resetBudget;
+exports.exportBudgetState = exportBudgetState;
+exports.importBudgetState = importBudgetState;
+exports.listModels = listModels;
+const budgetManager_1 = require("./core/budgetManager");
+const pricingRegistry_1 = require("./core/pricingRegistry");
+const registry_1 = require("./registry");
+const fetchInterceptor_1 = require("./interceptors/fetchInterceptor");
+Object.defineProperty(exports, "patchGlobalFetch", { enumerable: true, get: function () { return fetchInterceptor_1.patchGlobalFetch; } });
+const sdkInterceptor_1 = require("./interceptors/sdkInterceptor");
+Object.defineProperty(exports, "patchProvider", { enumerable: true, get: function () { return sdkInterceptor_1.patchProvider; } });
+const modelLister_1 = require("./introspection/modelLister");
+Object.defineProperty(exports, "listAvailableModels", { enumerable: true, get: function () { return modelLister_1.listAvailableModels; } });
+const contextRegistry_1 = require("./introspection/contextRegistry");
+let globalBudgetManager = null;
+/**
+ * Create and configure a budget guard
+ * @param options - Budget configuration options
+ * @returns Budget manager instance
+ */
+function createBudgetGuard(options) {
+    // Warn if overwriting existing budget guard
+    if (globalBudgetManager) {
+        console.warn('TokenFirewall: Creating new budget guard will replace existing one. Previous budget state will be lost.');
+    }
+    const manager = new budgetManager_1.BudgetManager(options);
+    globalBudgetManager = manager;
+    (0, fetchInterceptor_1.setBudgetManager)(manager);
+    return manager;
+}
+/**
+ * Register a custom provider adapter
+ * @param adapter - Provider adapter implementation
+ */
+function registerAdapter(adapter) {
+    registry_1.adapterRegistry.register(adapter);
+}
+/**
+ * Register custom pricing for a provider and model
+ * @param provider - Provider name
+ * @param model - Model name
+ * @param pricing - Pricing configuration (per 1M tokens)
+ */
+function registerPricing(provider, model, pricing) {
+    // Validate inputs
+    if (!provider || typeof provider !== 'string') {
+        throw new Error('TokenFirewall: Provider must be a non-empty string');
+    }
+    if (!model || typeof model !== 'string') {
+        throw new Error('TokenFirewall: Model must be a non-empty string');
+    }
+    if (!pricing || typeof pricing !== 'object') {
+        throw new Error('TokenFirewall: Pricing must be an object');
+    }
+    if (typeof pricing.input !== 'number' || pricing.input < 0 || !isFinite(pricing.input)) {
+        throw new Error('TokenFirewall: Pricing input must be a non-negative number');
+    }
+    if (typeof pricing.output !== 'number' || pricing.output < 0 || !isFinite(pricing.output)) {
+        throw new Error('TokenFirewall: Pricing output must be a non-negative number');
+    }
+    pricingRegistry_1.pricingRegistry.register(provider, model, pricing);
+}
+/**
+ * Register custom context limit for a provider and model
+ * @param provider - Provider name
+ * @param model - Model name
+ * @param contextLimit - Context window size in tokens
+ */
+function registerContextLimit(provider, model, contextLimit) {
+    // Validate inputs
+    if (!provider || typeof provider !== 'string') {
+        throw new Error('TokenFirewall: Provider must be a non-empty string');
+    }
+    if (!model || typeof model !== 'string') {
+        throw new Error('TokenFirewall: Model must be a non-empty string');
+    }
+    if (typeof contextLimit !== 'number' || contextLimit <= 0 || !isFinite(contextLimit)) {
+        throw new Error('TokenFirewall: Context limit must be a positive number');
+    }
+    contextRegistry_1.contextRegistry.register(provider, model, { tokens: contextLimit });
+}
+/**
+ * Get current budget status
+ * @returns Budget status or null if no budget guard is active
+ */
+function getBudgetStatus() {
+    return globalBudgetManager ? globalBudgetManager.getStatus() : null;
+}
+/**
+ * Reset budget tracking
+ */
+function resetBudget() {
+    if (globalBudgetManager) {
+        globalBudgetManager.reset();
+    }
+}
+/**
+ * Export budget state for persistence
+ */
+function exportBudgetState() {
+    if (globalBudgetManager) {
+        return globalBudgetManager.exportState();
+    }
+    return null;
+}
+/**
+ * Import budget state from persistence
+ */
+function importBudgetState(state) {
+    if (!globalBudgetManager) {
+        throw new Error('TokenFirewall: Cannot import budget state - no budget guard exists. Call createBudgetGuard() first.');
+    }
+    globalBudgetManager.importState(state);
+}
+/**
+ * List available models for a provider with context limits and budget usage
+ * @param options - Provider configuration and options
+ * @returns Array of model information
+ */
+async function listModels(options) {
+    // Pass the global budget manager to avoid circular dependency
+    return (0, modelLister_1.listAvailableModels)({
+        ...options,
+        budgetManager: globalBudgetManager
+    });
+}

package/dist/interceptors/fetchInterceptor.d.ts

@@ -0,0 +1,13 @@
+import { BudgetManager } from "../core/budgetManager";
+/**
+ * Set the budget manager for fetch interception
+ */
+export declare function setBudgetManager(manager: BudgetManager | null): void;
+/**
+ * Patch global fetch to intercept LLM API calls
+ */
+export declare function patchGlobalFetch(): void;
+/**
+ * Restore original fetch
+ */
+export declare function unpatchGlobalFetch(): void;

package/dist/interceptors/fetchInterceptor.js

@@ -0,0 +1,73 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.setBudgetManager = setBudgetManager;
+exports.patchGlobalFetch = patchGlobalFetch;
+exports.unpatchGlobalFetch = unpatchGlobalFetch;
+const registry_1 = require("../registry");
+const costEngine_1 = require("../core/costEngine");
+const logger_1 = require("../logger");
+let isPatched = false;
+let budgetManager = null;
+const originalFetch = globalThis.fetch;
+/**
+ * Set the budget manager for fetch interception
+ */
+function setBudgetManager(manager) {
+    budgetManager = manager;
+}
+/**
+ * Patch global fetch to intercept LLM API calls
+ */
+function patchGlobalFetch() {
+    if (isPatched) {
+        return;
+    }
+    const interceptedFetch = async function (input, init) {
+        const response = await originalFetch(input, init);
+        // Try to clone response for tracking (may fail for some responses)
+        let clonedResponse;
+        try {
+            clonedResponse = response.clone();
+        }
+        catch (error) {
+            // If cloning fails, just return original response without tracking
+            console.warn('TokenFirewall: Failed to clone response for tracking');
+            return response;
+        }
+        // Process response and track budget BEFORE returning
+        try {
+            const responseData = await clonedResponse.json();
+            // Try to process with adapter registry
+            const normalizedUsage = registry_1.adapterRegistry.process(responseData);
+            if (normalizedUsage) {
+                // Calculate cost
+                const cost = (0, costEngine_1.calculateCost)(normalizedUsage);
+                // Track budget if manager exists - MUST await to enforce blocking
+                if (budgetManager) {
+                    await budgetManager.track(cost.totalCost);
+                }
+                // Log usage
+                logger_1.logger.logUsage(normalizedUsage, cost);
+            }
+        }
+        catch (error) {
+            // If it's a budget error, re-throw it
+            if (error instanceof Error && error.message.includes('TokenFirewall: Budget exceeded')) {
+                throw error;
+            }
+            // Otherwise, not JSON or not an LLM response - ignore silently
+        }
+        return response;
+    };
+    globalThis.fetch = interceptedFetch;
+    isPatched = true;
+}
+/**
+ * Restore original fetch
+ */
+function unpatchGlobalFetch() {
+    if (isPatched) {
+        globalThis.fetch = originalFetch;
+        isPatched = false;
+    }
+}

package/dist/interceptors/sdkInterceptor.d.ts

@@ -0,0 +1,9 @@
+/**
+ * SDK-specific interceptors for providers that need direct patching
+ * Currently placeholder - most providers work via fetch interception
+ */
+/**
+ * Patch a specific provider SDK
+ * @param providerName - Name of the provider to patch
+ */
+export declare function patchProvider(providerName: string): void;

package/dist/interceptors/sdkInterceptor.js

@@ -0,0 +1,37 @@
+"use strict";
+/**
+ * SDK-specific interceptors for providers that need direct patching
+ * Currently placeholder - most providers work via fetch interception
+ */
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.patchProvider = patchProvider;
+const patchedProviders = new Set();
+/**
+ * Patch a specific provider SDK
+ * @param providerName - Name of the provider to patch
+ */
+function patchProvider(providerName) {
+    if (patchedProviders.has(providerName)) {
+        return;
+    }
+    switch (providerName.toLowerCase()) {
+        case "openai":
+            // OpenAI SDK uses fetch internally - covered by fetch interceptor
+            break;
+        case "anthropic":
+            // Anthropic SDK uses fetch internally - covered by fetch interceptor
+            break;
+        case "gemini":
+            // Gemini SDK uses fetch internally - covered by fetch interceptor
+            break;
+        case "grok":
+            // Grok uses fetch - covered by fetch interceptor
+            break;
+        case "kimi":
+            // Kimi uses fetch - covered by fetch interceptor
+            break;
+        default:
+            console.warn(`TokenFirewall: Provider "${providerName}" not recognized`);
+    }
+    patchedProviders.add(providerName);
+}