testronaut 1.1.2 → 1.2.0

package/tools/tokenControl.js

@@ -1,3 +1,25 @@
+/**
+ * tokenControl.js
+ * ----------------
+ * Purpose:
+ *   Centralized token counting and throttling utilities for Testronaut.
+ *   Provides per-model token estimation, dynamic token-per-minute (TPM)
+ *   limits, adaptive backoff, and header-learned limit updates.
+ *
+ * Responsibilities:
+ *   - Estimate token counts (prefer model encodings; fallback to bytes/4 heuristic).
+ *   - Determine current TPM limit (defaults → ENV override → header-learned).
+ *   - Track rolling token usage and apply cooldowns to avoid rate limits.
+ *   - Accept provider headers (e.g., OpenAI) to update live TPM caps.
+ *
+ * Related tests:
+ *   Located in `tests/toolsTests/`
+ *
+ * Used by:
+ *   - core/turnLoop.js (rate limiting & accounting)
+ *   - tools that need token estimation before pushing large DOM payloads
+ */
+
 import { encoding_for_model, get_encoding } from '@dqbd/tiktoken';
 import { wait } from './turnLoopUtils.js';
 
@@ -7,23 +29,32 @@ import { wait } from './turnLoopUtils.js';
  * - ENV override: TESTRONAUT_TOKENS_PER_MIN forces a single limit for everything.
  * - At runtime, you can call `updateLimitsFromHeaders(model, headers)` after a 429
  *   to adopt server-advertised limits (if present).
+ *
+ * NOTE: These are conservative defaults meant for backoff heuristics, not hard truths.
+ *       Providers may change limits; env/header learning will supersede these.
  */
 const DEFAULT_LIMITS = [
-  // Newer families first (checked in order)
-  { test: /^gpt-5(-|$)/i,               tpm: 90000 },
+  // ── OpenAI (newer first) ────────────────────────────────────────────────
+  { test: /^gpt-5(-|$)/i,               tpm:  90000 },
   { test: /^gpt-5-mini(-|$)/i,          tpm: 240000 },
   { test: /^gpt-5-nano(-|$)/i,          tpm: 600000 },
 
   { test: /^gpt-4o(-|$)/i,              tpm: 450000 },
-  { test: /^gpt-4\.1(-|$)/i,            tpm: 1000000 },
+  { test: /^gpt-4\.1(-|$)/i,            tpm:1000000 },
 
   { test: /^o3(-|$)/i,                  tpm: 300000 },
   { test: /^o4-mini(-|$)/i,             tpm: 600000 },
 
-  // Older / fallback
+  // Older / fallback (OpenAI)
   { test: /^gpt-4(-|$)/i,               tpm: 150000 },
   { test: /^gpt-3\.5(-|$)/i,            tpm: 600000 },
 
+  // ── Gemini (approximate, conservative) ─────────────────────────────────
+  // These are heuristic defaults to guide local throttling only.
+  { test: /^gemini-2\.5-pro(-|$)/i,      tpm: 120000 },
+  { test: /^gemini-2\.5-flash-8b(-|$)/i, tpm: 300000 },
+  { test: /^gemini-2\.5-flash(-|$)/i,    tpm: 300000 },
+
   // Ultimate fallback for anything else
   { test: /.*/,                         tpm: 150000 },
 ];
@@ -34,14 +65,24 @@ const liveLimits = new Map(); // modelId -> { tpm, source: 'default'|'env'|'head
 // One-time warning tracking for tokenizer fallback
 const warnedModels = new Set();
 
-/* ---------------- Tokenizer helpers (same idea as before) ---------------- */
-
+/* ---------------- Tokenizer helpers ----------------
+ * We prefer tiktoken's per-model encoding when available.
+ * For unknown models, we pick a close base:
+ *  - o200k_base for modern, long-context families (OpenAI 4o/4.1/5, Gemini 2.5, O*)
+ *  - cl100k_base as a broad fallback
+ */
 function getTokenizer(model) {
   try {
     return encoding_for_model(model);
   } catch (_) {
-    const m = (model || '').toLowerCase();
+    const m = String(model || '').toLowerCase();
+
+    // Treat Gemini 2.5 like modern long-context models
+    const isGemini25 = /^gemini-2\.5/.test(m);
+
+    // OpenAI modern families also map well to o200k_base
     const useO200k =
+      isGemini25 ||
       m.startsWith('gpt-5') ||
       m.startsWith('gpt-4o') ||
       m.startsWith('gpt-4.1') ||
@@ -60,6 +101,14 @@ function getTokenizer(model) {
   }
 }
 
+/**
+ * Estimate tokens for a given text under a model family.
+ * Prefers a real tokenizer; falls back to bytes/4 heuristic.
+ *
+ * @param {string} model - Model identifier (e.g., 'gpt-4o', 'gemini-2.5-flash')
+ * @param {string|any} text - Text or JSON-like payload to estimate
+ * @returns {Promise<number>} estimated token count
+ */
 export const tokenEstimate = async (model, text) => {
   const str = typeof text === 'string' ? text : JSON.stringify(text ?? '');
   const enc = getTokenizer(model);
@@ -88,8 +137,14 @@ export const tokenEstimate = async (model, text) => {
 
 /* ---------------- Dynamic limit resolution ---------------- */
 
+/**
+ * Resolve a default TPM for a model:
+ * - ENV override takes precedence (TESTRONAUT_TOKENS_PER_MIN)
+ * - Otherwise match regex patterns in DEFAULT_LIMITS
+ * @param {string} model
+ * @returns {{tpm:number, source:'default'|'env'}}
+ */
 function resolveDefaultLimitForModel(model) {
-  // Highest priority: ENV override (forces a single limit for all models)
   const envTPM = process.env.TESTRONAUT_TOKENS_PER_MIN
     ? Number(process.env.TESTRONAUT_TOKENS_PER_MIN)
     : undefined;
@@ -97,14 +152,16 @@ function resolveDefaultLimitForModel(model) {
     return { tpm: envTPM, source: 'env' };
   }
 
-  // Otherwise find best-matching default
   const hit = DEFAULT_LIMITS.find(entry => entry.test.test(model || ''));
   return { tpm: hit?.tpm ?? 150000, source: 'default' };
 }
 
 /**
- * Public: get the current token-per-minute limit for a model.
- * Checks liveOverrides first (header-learned), then ENV, then defaults.
+ * Get the current token-per-minute limit for a model.
+ * Priority: header-learned (live) → ENV → defaults.
+ *
+ * @param {string} model
+ * @returns {{tpm:number, source:'default'|'env'|'header'}}
  */
 export function getCurrentTokenLimit(model) {
   const m = (model || '').trim() || 'unknown';
@@ -117,11 +174,11 @@ export function getCurrentTokenLimit(model) {
 }
 
 /**
- * Public: update limits from HTTP response headers (e.g., after 429).
- * Looks for typical headers like:
- *  - x-ratelimit-limit-tokens: "90000"
- *  - x-ratelimit-remaining-tokens, x-ratelimit-reset-tokens, etc.
- * If found, we store the new TPM for this model (source='header').
+ * Update TPM from HTTP response headers (e.g., after 429).
+ * Looks for common provider headers (OpenAI/Azure style). No-op if absent.
+ *
+ * @param {string} model
+ * @param {Record<string, string|number>} headers
  */
 export function updateLimitsFromHeaders(model, headers = {}) {
   if (!model) return;
@@ -130,7 +187,6 @@ export function updateLimitsFromHeaders(model, headers = {}) {
   const lower = {};
   for (const k of Object.keys(headers || {})) lower[k.toLowerCase()] = headers[k];
 
-  // Common header names OpenAI/Azure may set (varies!)
   const tokenCap =
     Number(lower['x-ratelimit-limit-tokens']) ||
     Number(lower['x-ratelimit-limit-tpm']) ||
@@ -148,6 +204,14 @@ export function updateLimitsFromHeaders(model, headers = {}) {
 
 /* ---------------- Cooloff / backoff logic ---------------- */
 
+/**
+ * If usage exceeds TPM, wait until safe and reset rolling counters.
+ *
+ * @param {number} totalTokensUsed - rolling 60s total
+ * @param {Array<[number,number]>} turnTimestamps - [[tsMs, tokens], ...]
+ * @param {string} model - model id for TPM lookup
+ * @returns {Promise<{shouldBackoff:boolean,totalTokensUsed:number,turnTimestamps:Array}>}
+ */
 export const tokenUseCoolOff = async (totalTokensUsed, turnTimestamps, model) => {
   const { tpm } = getCurrentTokenLimit(model);
   if (totalTokensUsed > tpm) {
@@ -160,11 +224,22 @@ export const tokenUseCoolOff = async (totalTokensUsed, turnTimestamps, model) =>
   return { shouldBackoff: false, totalTokensUsed, turnTimestamps };
 };
 
+/**
+ * Record tokens for the current turn in the rolling window.
+ * @param {Array<[number,number]>} turnTimestamps
+ * @param {number} tokensUsed
+ */
 export const recordTokenUsage = (turnTimestamps, tokensUsed) => {
   const now = Date.now();
   turnTimestamps.push([now, tokensUsed]);
 };
 
+/**
+ * Remove entries older than `windowMs` from the rolling window.
+ * @param {Array<[number,number]>} turnTimestamps
+ * @param {number} windowMs
+ * @returns {{turnTimestamps:Array<[number,number]>, totalTokensUsed:number}}
+ */
 export const pruneOldTokenUsage = (turnTimestamps, windowMs = 60000) => {
   const cutoff = Date.now() - windowMs;
   const recentEntries = turnTimestamps.filter(([timestamp]) => timestamp > cutoff);
@@ -172,6 +247,14 @@ export const pruneOldTokenUsage = (turnTimestamps, windowMs = 60000) => {
   return { turnTimestamps: recentEntries, totalTokensUsed };
 };
 
+/**
+ * Compute milliseconds to wait until under TPM again.
+ * Walks the sorted window and finds when usage first exceeds TPM.
+ *
+ * @param {Array<[number,number]>} turnTimestamps
+ * @param {number} tokenLimit
+ * @returns {Promise<number>} ms to wait (>= 0)
+ */
 const getDynamicBackoffMs = async (turnTimestamps, tokenLimit) => {
   const now = Date.now();
   let runningTotal = 0;
@@ -187,3 +270,10 @@ const getDynamicBackoffMs = async (turnTimestamps, tokenLimit) => {
   }
   return 0;
 };
+
+
+// Test-only helper to clear internal state between tests
+export function __resetTokenControlForTests() {
+  liveLimits.clear();
+  warnedModels.clear();
+}

package/vitest.config.js

@@ -0,0 +1,8 @@
+import { defineConfig } from 'vitest/config';
+export default defineConfig({
+  test: {
+    environment: 'node',
+    include: ['tests/**/*.{test,spec}.js'],
+    coverage: { reporter: ['text', 'lcov'] }
+  }
+});

package/openAI/modelResolver.js

@@ -1,33 +0,0 @@
-import path from 'path';
-import fs from 'fs';
-
-// Resolve model from config (or ENV) with safe fallback
-const root = process.cwd();
-const configPath = path.join(root, 'testronaut-config.json');
-
-export function resolveModel() {
-  // Highest priority: env override
-  if (process.env.TESTRONAUT_MODEL?.trim()) {
-    return process.env.TESTRONAUT_MODEL.trim();
-  }
-
-  try {
-    if (fs.existsSync(configPath)) {
-      const cfg = JSON.parse(fs.readFileSync(configPath, 'utf-8'));
-      // Only use the model if provider is openai (future-proofing for other providers)
-      if (cfg?.provider === 'openai' && typeof cfg?.model === 'string' && cfg.model.trim()) {
-        return cfg.model.trim();
-      }
-      // Back-compat: older configs may have stored 'openai' in model
-      if (!cfg?.provider && cfg?.model === 'openai') {
-        return 'gpt-4o';
-      }
-    }
-  } catch (e) {
-    console.warn('⚠️ Could not read testronaut-config.json; using default model. Error:', e.message);
-  }
-
-  return 'gpt-4o';
-}
-
-