testronaut 1.1.2 → 1.2.1

package/tests/toolsTests/tokenControl.test.js

@@ -0,0 +1,150 @@
+// tests/tokenControlTests/tokenControl.test.js
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+
+// Hoist env so we can safely mutate per test
+const ORIGINAL_ENV = { ...process.env };
+
+// Mock wait so tests are instant
+vi.mock('../../tools/turnLoopUtils.js', () => ({
+  wait: vi.fn(async () => {}),
+}));
+
+// Mock tiktoken in a controllable way
+const encMock = {
+  encode: vi.fn((s) => Array.from(String(s)).map(() => 1)), // 1 token per char
+  free: vi.fn(() => {}),
+};
+
+const tiktokenMocks = vi.hoisted(() => ({
+  encoding_for_model_impl: vi.fn(),
+  get_encoding_impl: vi.fn(() => encMock),
+}));
+
+vi.mock('@dqbd/tiktoken', () => ({
+  encoding_for_model: (model) => tiktokenMocks.encoding_for_model_impl(model),
+  get_encoding: (name) => tiktokenMocks.get_encoding_impl(name),
+}));
+
+// Import after mocks
+import {
+  tokenEstimate,
+  getCurrentTokenLimit,
+  updateLimitsFromHeaders,
+  tokenUseCoolOff,
+  recordTokenUsage,
+  pruneOldTokenUsage,
+  __resetTokenControlForTests,
+} from '../../tools/tokenControl.js';
+
+describe('tokenControl', () => {
+  beforeEach(() => {
+    // reset env and mocks
+    process.env = { ...ORIGINAL_ENV };
+    __resetTokenControlForTests(); 
+    tiktokenMocks.encoding_for_model_impl.mockReset();
+    tiktokenMocks.get_encoding_impl.mockReset().mockReturnValue(encMock);
+    encMock.encode.mockClear();
+    encMock.free.mockClear();
+  });
+
+  afterEach(() => {
+    process.env = { ...ORIGINAL_ENV };
+  });
+
+  describe('tokenEstimate', () => {
+    it('uses encoding_for_model when available', async () => {
+      tiktokenMocks.encoding_for_model_impl.mockReturnValue(encMock);
+
+      const n = await tokenEstimate('gpt-4o', 'abcd'); // 4 chars → 4 tokens via mock
+      expect(n).toBe(4);
+      expect(encMock.encode).toHaveBeenCalledOnce();
+      expect(encMock.free).toHaveBeenCalledOnce();
+    });
+
+    it('falls back to get_encoding when encoding_for_model throws', async () => {
+      tiktokenMocks.encoding_for_model_impl.mockImplementation(() => { throw new Error('no direct encoding'); });
+      tiktokenMocks.get_encoding_impl.mockReturnValue(encMock);
+
+      const n = await tokenEstimate('gemini-2.5-flash', 'abc'); // 3
+      expect(n).toBe(3);
+      expect(encMock.encode).toHaveBeenCalledOnce();
+    });
+
+    it('falls back to bytes/4 when both tokenizers fail', async () => {
+      tiktokenMocks.encoding_for_model_impl.mockImplementation(() => { throw new Error('no direct encoding'); });
+      tiktokenMocks.get_encoding_impl.mockImplementation(() => { throw new Error('no base encoding'); });
+
+      const n = await tokenEstimate('unknown-model', 'abcdefgh'); // 8 bytes → 8/4 = 2
+      expect(n).toBe(2);
+    });
+
+    it('stringifies non-string input', async () => {
+      tiktokenMocks.encoding_for_model_impl.mockReturnValue(encMock);
+      const n = await tokenEstimate('gpt-4o', { a: 1, b: 2 });
+      expect(encMock.encode).toHaveBeenCalled();
+      expect(typeof n).toBe('number');
+    });
+  });
+
+  describe('getCurrentTokenLimit / updateLimitsFromHeaders', () => {
+    it('uses defaults by model pattern', () => {
+      const openai = getCurrentTokenLimit('gpt-4o');
+      const gemini = getCurrentTokenLimit('gemini-2.5-flash');
+      expect(openai.tpm).toBeGreaterThan(0);
+      expect(gemini.tpm).toBeGreaterThan(0);
+      expect(['default', 'env', 'header']).toContain(openai.source);
+    });
+
+    it('honors ENV override TESTRONAUT_TOKENS_PER_MIN', () => {
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '999';
+      const anyModel = getCurrentTokenLimit('what-ever-model');
+      expect(anyModel.tpm).toBe(999);
+      expect(anyModel.source).toBe('env');
+    });
+
+    it('updates limits from headers (header wins)', () => {
+      const before = getCurrentTokenLimit('gpt-4o');
+      updateLimitsFromHeaders('gpt-4o', { 'x-ratelimit-limit-tokens': '1234' });
+      const after = getCurrentTokenLimit('gpt-4o');
+      expect(after.tpm).toBe(1234);
+      expect(after.source).toBe('header');
+      // sanity check it actually changed (unless defaults already 1234)
+      if (before.tpm !== 1234) expect(after.tpm).not.toBe(before.tpm);
+    });
+  });
+
+  describe('rolling window + cooldown', () => {
+    it('prunes old entries outside the window', () => {
+      const now = Date.now();
+      const entries = [
+        [now - 70000, 100], // old → pruned
+        [now - 10000, 50],  // kept
+        [now - 5000, 30],   // kept
+      ];
+      const { turnTimestamps, totalTokensUsed } = pruneOldTokenUsage(entries, 60000);
+      expect(turnTimestamps.length).toBe(2);
+      expect(totalTokensUsed).toBe(80);
+    });
+
+    it('triggers backoff when usage exceeds TPM', async () => {
+      // small ENV cap to force backoff
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '50';
+      const entries = [];
+      recordTokenUsage(entries, 60);
+
+      const result = await tokenUseCoolOff(60, entries, 'any-model');
+      expect(result.shouldBackoff).toBe(true);
+      expect(result.totalTokensUsed).toBe(0);
+      expect(Array.isArray(result.turnTimestamps)).toBe(true);
+    });
+
+    it('does not backoff when under TPM', async () => {
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '100';
+      const entries = [];
+      recordTokenUsage(entries, 60);
+      const result = await tokenUseCoolOff(60, entries, 'any-model');
+      expect(result.shouldBackoff).toBe(false);
+      expect(result.totalTokensUsed).toBe(60);
+    });
+  });
+});

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';
-}
-
-