Prevent runaway AI costs with pre-flight checks and token budgeting

import { gateway, type ModelMessage } from "ai";
 
/**
 * Estimate tokens for messages array using rough character-to-token ratio
 * Use this for pre-flight cost estimation
 *
 * Note: This is an approximation. Actual tokens depend on tokenizer.
 * Rule of thumb: 1 token ≈ 4 characters for English text.
 * In production, use a real tokenizer like OpenAI's tiktoken for more accurate estimates:
 * https://github.com/openai/tiktoken
 */
export function estimateTokens(messages: ModelMessage[]): number {
  return messages.reduce((sum, msg) => {
    const content = typeof msg.content === 'string'
      ? msg.content
      : JSON.stringify(msg.content);
    // Rough estimate: 1 token per 4 characters
    return sum + Math.ceil(content.length / 4);
  }, 0);
}
 
/**
 * Look up a model's price per token from AI Gateway's model catalog.
 * Prices change, so read them instead of hardcoding a table.
 * The gateway provider caches the catalog, so this isn't a network call per request.
 */
async function getModelPricing(model: string) {
  const { models } = await gateway.getAvailableModels();
  const pricing = models.find((m) => m.id === model)?.pricing;
 
  if (!pricing) {
    throw new Error(`AI Gateway has no pricing for ${model}`);
  }
 
  return { input: Number(pricing.input), output: Number(pricing.output) };
}
 
/**
 * Calculate request cost before making it
 * Use for budget decisions and request rejection
 */
export async function estimateRequestCost(
  inputTokens: number,
  outputTokens: number,
  model: string
): Promise<number> {
  const pricing = await getModelPricing(model);
  return inputTokens * pricing.input + outputTokens * pricing.output;
}
 
/**
 * Calculate actual cost from AI SDK usage response
 * Use to compare estimated vs actual and tune your estimates
 */
export async function calculateActualCost(
  inputTokens: number,
  outputTokens: number,
  model: string
): Promise<number> {
  return estimateRequestCost(inputTokens, outputTokens, model);
}
 
/**
 * Pre-flight check: Should we reject this request due to cost?
 * Reject obviously expensive requests BEFORE they burn money
 */
export function shouldRejectRequest(estimatedCost: number): {
  reject: boolean;
  reason?: string;
} {
  const MAX_REQUEST_COST = 0.02; // Example cap: $0.02 per request. Pick your own.
 
  if (estimatedCost > MAX_REQUEST_COST) {
    return {
      reject: true,
      reason: `Request would cost $${estimatedCost.toFixed(4)}, exceeds $${MAX_REQUEST_COST} limit`
    };
  }
 
  return { reject: false };
}
 
/**
 * Check if we should use cheap mode based on environment variable
 * In production, set this based on your own budget tracking
 */
export function shouldUseCheapModel(): boolean {
  return process.env.FORCE_CHEAP_MODEL === 'true';
}

import type { KnownEventFromType } from "@slack/bolt";
import { generateText, isStepCount, type ModelMessage } from "ai";
import { app } from "~/app";
import {
  calculateActualCost,
  estimateRequestCost,
  estimateTokens,
  shouldRejectRequest,
  shouldUseCheapModel,
} from "./cost-control";
import { simulateFailure, withRetry } from "./retry-wrapper";
// ... rest of imports ...
 
export const respondToMessage = async ({
  messages,
  event,
  channel,
  thread_ts,
  botId,
  correlation,
}: RespondToMessageOptions) => {
  // Pre-flight cost estimation - catch expensive requests BEFORE they cost money
  const inputTokens = estimateTokens(messages);
  const estimatedOutputTokens = 500; // Conservative estimate for response
  const primaryModel = "openai/gpt-5.4-mini";
  const estimatedCost = await estimateRequestCost(
    inputTokens,
    estimatedOutputTokens,
    primaryModel
  );
 
  app.logger.info("Pre-flight cost check", {
    ...correlation,
    inputTokens,
    estimatedOutputTokens,
    estimatedCost,
    model: primaryModel,
    status: estimatedCost > 0.02 ? "rejected" : "approved",
  });
 
  // Reject expensive requests BEFORE making them
  const rejectCheck = shouldRejectRequest(estimatedCost);
  if (rejectCheck.reject) {
    app.logger.warn("Request rejected - cost limit exceeded", {
      ...correlation,
      estimatedCost,
      limit: 0.02,
      inputTokens,
      estimatedOutputTokens,
      model: primaryModel,
    });
 
    // Return user-friendly error instead of making expensive request
    return `Your request is too large. Please break it into smaller questions or provide less context. (Estimated cost: $${estimatedCost.toFixed(4)})`;
  }
 
  // Model selection with budget awareness
  // In production, set FORCE_CHEAP_MODEL based on your own budget tracking
  const useCheapMode = shouldUseCheapModel();
  const models = useCheapMode
    ? ["openai/gpt-5.4-nano"] // Force cheap model only
    : [primaryModel, "anthropic/claude-haiku-4.5"];
 
  if (useCheapMode) {
    app.logger.warn("Budget threshold triggered - using cheap model", {
      ...correlation,
      reason: "FORCE_CHEAP_MODEL environment variable set",
      selectedModel: "openai/gpt-5.4-nano",
      normalModel: primaryModel,
    });
  }
 
  let lastError: unknown;
 
  for (const currentModel of models) {
    try {
      const result = await withRetry(
        async () => {
          app.logger.info("Attempting AI request", {
            ...correlation,
            model: currentModel,
            inputTokens,
            estimatedCost,
          });
 
          return await generateText({
            model: currentModel,
            instructions: `You are Slack Agent, a helpful assistant in Slack.
            // ... rest of system prompt ...
            `,
            messages,
            // ... rest of config ...
          });
        },
        {
          maxRetries: 3,
          initialDelayMs: 1000,
        }
      );
 
      // Log actual cost after request completes
      const actualCost = await calculateActualCost(
        result.usage.inputTokens ?? 0,
        result.usage.outputTokens ?? 0,
        currentModel
      );
 
      const estimationAccuracy = ((actualCost / estimatedCost) * 100).toFixed(0);
 
      app.logger.info("AI request completed", {
        ...correlation,
        model: currentModel,
        usage: {
          promptTokens: result.usage.inputTokens,
          completionTokens: result.usage.outputTokens,
          totalTokens: result.usage.totalTokens,
        },
        actualCost,
        estimatedCost,
        estimationAccuracy: `${estimationAccuracy}%`,
      });
 
      return result.text;
    } catch (error) {
      // ... existing error handling ...
      lastError = error;
    }
  }
 
  throw lastError;
};

import type { AllMiddlewareArgs, SlackCommandMiddlewareArgs } from "@slack/bolt";
import { gateway } from "ai";
 
export const gatewayStatusCallback = async ({
  ack,
  respond,
}: AllMiddlewareArgs & SlackCommandMiddlewareArgs) => {
  await ack();
 
  const today = new Date().toISOString().slice(0, 10);
  const report = await gateway.getSpendReport({
    startDate: today,
    endDate: today,
    groupBy: "model",
  });
 
  const lines = report.results.map(
    (row) => `• ${row.model}: $${row.totalCost.toFixed(4)}`
  );
 
  await respond(lines.length ? lines.join("\n") : "No AI spend yet today.");
};