testronaut 1.7.1 → 1.8.0

package/README.md

@@ -143,6 +143,25 @@ testronaut --init
 TESTRONAUT_PROVIDER=gemini TESTRONAUT_MODEL=gemini-2.5-pro testronaut
 ```
 
+### Token-rate override
+
+`TESTRONAUT_TOKENS_PER_MIN` is an emergency/manual hard override. A positive
+numeric value takes priority over configured fallbacks and limits learned from
+provider response headers. Testronaut prints a warning whenever it is active.
+
+Clear it permanently from the current shell and remove the corresponding line
+from `.env` or your shell profile:
+
+```bash
+unset TESTRONAUT_TOKENS_PER_MIN
+```
+
+To bypass a numeric value inherited from `.env` for one run:
+
+```bash
+TESTRONAUT_TOKENS_PER_MIN=auto testronaut login.mission.js
+```
+
 Current supported providers:
 
 | Provider | Example Models |

package/bin/init.js

@@ -38,6 +38,7 @@ import {
 } from './initHelpers.js';
 import { DEFAULT_OPENAI_MODEL, OPENAI_MODELS } from '../llm/openAI/models.js';
 import { ANTHROPIC_MODELS, DEFAULT_ANTHROPIC_MODEL } from '../llm/anthropic/models.js';
+import { defaultRateLimitConfig } from '../tools/rateLimitDefaults.js';
 
 export async function initializeTestronautProject() {
   // ─────────────────────────────────────────────
@@ -159,6 +160,7 @@ export async function initializeTestronautProject() {
     config.dom = { ...defaults.dom, ...(config.dom || {}) };
     config.resourceGuard = { ...defaults.resourceGuard, ...(config.resourceGuard || {}) };
     config.humanInput = { ...defaults.humanInput, ...(config.humanInput || {}) };
+    config.rateLimits = config.rateLimits || defaultRateLimitConfig(config.provider, config.model);
 
     // ─────────────────────────────────────────────
     // STEP 4: Persist config and scaffold .env (first run only)

package/core/turnLoop.js

@@ -42,6 +42,8 @@ import {
   recordTokenUsage, 
   pruneOldTokenUsage,
   updateLimitsFromHeaders,
+  updateLimitsFromError,
+  configureTokenControl,
   warnIfContextNearLimit
 } from '../tools/tokenControl.js';
 import { resolveProviderModel } from '../llm/modelResolver.js';
@@ -62,6 +64,13 @@ import {
 const { provider: PROVIDER_ID, model: MODEL_ID } = resolveProviderModel();
 console.log(`🧠 Using LLM → provider: ${PROVIDER_ID}, model: ${MODEL_ID}`);
 
+let rateLimitConfig;
+try {
+  const config = JSON.parse(fs.readFileSync('testronaut-config.json', 'utf8'));
+  rateLimitConfig = config.rateLimits;
+} catch { /* startup fallbacks remain active when config is absent or invalid */ }
+configureTokenControl({ provider: PROVIDER_ID, rateLimits: rateLimitConfig });
+
 const llm = getLLM(PROVIDER_ID);
 
 // Track resource/download coverage across turns (guard against partial loops).
@@ -353,21 +362,36 @@ export const turnLoop = async (
       // ─────────────────────────────────────────────
       // STEP 2: Request next reasoning turn from model
       // ─────────────────────────────────────────────
-      await warnIfContextNearLimit(MODEL_ID, { messages, tools: activeToolsSchema });
+      const contextCheck = await warnIfContextNearLimit(MODEL_ID, { messages, tools: activeToolsSchema });
+      const projectedTokens = contextCheck.estimatedTokens
+        ?? await tokenEstimate(MODEL_ID, { messages, tools: activeToolsSchema });
+      ({ totalTokensUsed, turnTimestamps, shouldBackoff } = await tokenUseCoolOff(
+        totalTokensUsed,
+        turnTimestamps,
+        MODEL_ID,
+        projectedTokens
+      ));
+      if (shouldBackoff) {
+        recordStep(step);
+        turn -= 1;
+        continue;
+      }
       const { message, usage, headers } = await llm.chat({
         model: MODEL_ID,
         messages,
         tools: activeToolsSchema,
       });
       response = { message, usage, headers };
+      updateLimitsFromHeaders(MODEL_ID, headers);
 
     } catch (err) {
       // ─────────────────────────────────────────────
       // STEP 3: Error and rate-limit handling
       // ─────────────────────────────────────────────
       if (err.status === 429) {
-        try { updateLimitsFromHeaders(MODEL_ID, err.headers || err.response?.headers || {}); } catch {}
-        const delay = Math.min(60000, 2 ** retryCount * 2000);
+        let learned = {};
+        try { learned = updateLimitsFromError(MODEL_ID, err); } catch {}
+        const delay = Math.min(60000, learned.retryAfterMs ?? 2 ** retryCount * 2000);
         console.warn(`⚠️ Rate limited. Retrying in ${delay / 1000}s... (retry ${retryCount + 1})`);
         step.events.push(`⏳ Rate limit: waiting ${Math.round(delay/1000)}s (retry ${retryCount + 1})`);
         recordStep(step);         // ✅ persist this step before sleeping

package/llm/openAI/openaiProvider.js

@@ -48,7 +48,16 @@ export class OpenAIProvider {
       request.reasoning_effort = chatToolReasoningEffort;
     }
 
-    const res = await this.client.chat.completions.create(request);
+    const pending = this.client.chat.completions.create(request);
+    let res;
+    let headers;
+    if (typeof pending?.withResponse === 'function') {
+      const wrapped = await pending.withResponse();
+      res = wrapped.data;
+      headers = wrapped.response?.headers;
+    } else {
+      res = await pending;
+    }
 
     // OpenAI already returns an OpenAI-like message and usage structure.
     const message = res.choices?.[0]?.message ?? { role: 'assistant', content: '' };
@@ -58,7 +67,7 @@ export class OpenAIProvider {
     };
 
     // Some SDK versions surface headers on `res.headers`, others on `res.response.headers`.
-    const headers = res.headers || res.response?.headers;
+    headers = headers || res.headers || res.response?.headers;
 
     return { message, usage, headers };
   }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "testronaut",
-  "version": "1.7.1",
+  "version": "1.8.0",
   "main": "./runner/testronaut.js",
   "bin": {
     "testronaut": "bin/cli.js"

package/tests/binTests/init.test.js

@@ -84,6 +84,12 @@ describe('initializeTestronautProject', () => {
     expect(cfg.resourceGuard?.enabled).toBe(true);
     expect(cfg.humanInput?.enabled).toBe(true);
     expect(cfg.humanInput?.timeoutSeconds).toBe(60);
+    expect(cfg.rateLimits).toMatchObject({
+      tier: 'unknown',
+      region: 'global',
+      safetyMargin: 0.9,
+      models: { 'gpt-4o': { fallbackTPM: 450000 } },
+    });
 
     // .env
     const envPath = path.join(temp, '.env');
@@ -125,6 +131,9 @@ describe('initializeTestronautProject', () => {
     expect(cfg.resourceGuard?.enabled).toBe(true);
     expect(cfg.humanInput?.enabled).toBe(true);
     expect(cfg.humanInput?.timeoutSeconds).toBe(60);
+    expect(cfg.rateLimits.models).toEqual({
+      'gemini-2.5-flash': { fallbackTPM: 300000 },
+    });
 
     const envTxt = fs.readFileSync(path.join(temp, '.env'), 'utf8');
     expect(envTxt).toMatch(/GEMINI_API_KEY=AIza/);
@@ -149,6 +158,9 @@ describe('initializeTestronautProject', () => {
     const cfg = readJson(path.join(temp, 'testronaut-config.json'));
     expect(cfg.provider).toBe('anthropic');
     expect(cfg.model).toBe('claude-sonnet-5');
+    expect(cfg.rateLimits.models).toEqual({
+      'claude-sonnet-5': { fallbackTPM: 80000 },
+    });
     expect(fs.readFileSync(path.join(temp, '.env'), 'utf8')).toMatch(/ANTHROPIC_API_KEY=sk-ant-/);
     expect(shared.calls).toBe(2);
   });

package/tests/coreTests/turnLoop.test.js

@@ -67,6 +67,8 @@ vi.mock('../../tools/tokenControl.js', () => ({
     totalTokensUsed: turnTimestamps.reduce((a, [, t]) => a + t, 0),
   })),
   updateLimitsFromHeaders: vi.fn(() => {}),
+  updateLimitsFromError: vi.fn(() => ({})),
+  configureTokenControl: vi.fn(() => {}),
   warnIfContextNearLimit: vi.fn(async () => ({ warned: false })),
 }));
 

package/tests/llmTests/openAIProvider.test.js

@@ -18,7 +18,7 @@ const { shared } = vi.hoisted(() => ({
 // Mock the OpenAI SDK surface we use
 vi.mock('openai', () => {
   class FakeChatCompletions {
-    async create(args) {
+    create(args) {
       shared.lastCreateArgs = args;
       return shared.responseFactory();
     }
@@ -121,6 +121,25 @@ describe('OpenAIProvider', () => {
     expect(headers).toEqual({ 'x-ratelimit-limit-tokens': '12345' });
   });
 
+  it('uses the SDK raw response wrapper when withResponse is available', async () => {
+    const rawHeaders = new Headers({ 'x-ratelimit-limit-tokens': '54321' });
+    shared.responseFactory = () => ({
+      withResponse: async () => ({
+        data: {
+          choices: [{ message: { role: 'assistant', content: 'wrapped' } }],
+          usage: { total_tokens: 12 },
+        },
+        response: { headers: rawHeaders },
+      }),
+    });
+
+    const prov = new OpenAIProvider({ apiKey: 'sk-abc' });
+    const result = await prov.chat({ model: 'gpt-4o', messages: [] });
+    expect(result.message.content).toBe('wrapped');
+    expect(result.usage.total_tokens).toBe(12);
+    expect(result.headers).toBe(rawHeaders);
+  });
+
   it('provides a default assistant message when choices are empty', async () => {
     shared.responseFactory = () => ({
       choices: [],

package/tests/toolsTests/rateLimitDefaults.test.js

@@ -0,0 +1,34 @@
+import { describe, expect, it } from 'vitest';
+import { defaultRateLimitConfig, getFallbackTPM } from '../../tools/rateLimitDefaults.js';
+
+describe('rateLimitDefaults', () => {
+  it.each([
+    ['openai', 'gpt-5.6', 500000],
+    ['openai', 'gpt-5.6-luna', 500000],
+    ['openai', 'gpt-5.4-nano', 200000],
+    ['openai', 'gpt-5-mini', 240000],
+    ['gemini', 'gemini-2.5-pro', 120000],
+    ['gemini', 'gemini-2.5-flash', 300000],
+    ['anthropic', 'claude-sonnet-5', 80000],
+    ['claude', 'claude-haiku-4-5', 100000],
+  ])('resolves %s/%s to %i TPM', (provider, model, expected) => {
+    expect(getFallbackTPM(provider, model)).toBe(expected);
+  });
+
+  it('does not match a model family belonging to a different provider', () => {
+    expect(getFallbackTPM('gemini', 'gpt-4o')).toBe(150000);
+  });
+
+  it('uses a conservative fallback for unknown models', () => {
+    expect(getFallbackTPM('custom', 'future-model')).toBe(150000);
+  });
+
+  it('generates config for only the selected model', () => {
+    expect(defaultRateLimitConfig('openai', 'gpt-4o')).toEqual({
+      tier: 'unknown',
+      region: 'global',
+      safetyMargin: 0.9,
+      models: { 'gpt-4o': { fallbackTPM: 450000 } },
+    });
+  });
+});

package/tests/toolsTests/tokenControl.test.js

@@ -30,12 +30,15 @@ import {
   tokenEstimate,
   getCurrentTokenLimit,
   updateLimitsFromHeaders,
+  updateLimitsFromError,
+  configureTokenControl,
   tokenUseCoolOff,
   recordTokenUsage,
   pruneOldTokenUsage,
   warnIfContextNearLimit,
   __resetTokenControlForTests,
 } from '../../tools/tokenControl.js';
+import { wait } from '../../tools/turnLoopUtils.js';
 
 describe('tokenControl', () => {
   beforeEach(() => {
@@ -46,6 +49,7 @@ describe('tokenControl', () => {
     tiktokenMocks.get_encoding_impl.mockReset().mockReturnValue(encMock);
     encMock.encode.mockClear();
     encMock.free.mockClear();
+    wait.mockClear();
   });
 
   afterEach(() => {
@@ -103,6 +107,58 @@ describe('tokenControl', () => {
       expect(anyModel.source).toBe('env');
     });
 
+    it('warns once when a numeric environment override is configured', () => {
+      const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '999';
+      configureTokenControl({ provider: 'openai' });
+      configureTokenControl({ provider: 'openai' });
+      expect(warn).toHaveBeenCalledOnce();
+      expect(warn.mock.calls[0][0]).toContain('TPM override active');
+      expect(warn.mock.calls[0][0]).toContain('unset TESTRONAUT_TOKENS_PER_MIN');
+      expect(warn.mock.calls[0][0]).toContain('TESTRONAUT_TOKENS_PER_MIN=auto');
+      warn.mockRestore();
+    });
+
+    it('treats auto as a one-run bypass for an inherited override', () => {
+      const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
+      process.env.TESTRONAUT_TOKENS_PER_MIN = 'auto';
+      configureTokenControl({
+        provider: 'openai',
+        rateLimits: { models: { 'gpt-4o': { fallbackTPM: 777 } } },
+      });
+      expect(getCurrentTokenLimit('gpt-4o')).toEqual({ tpm: 777, source: 'config-fallback' });
+      expect(warn).not.toHaveBeenCalled();
+      warn.mockRestore();
+    });
+
+    it('gives the environment override priority over config and learned headers', () => {
+      configureTokenControl({
+        provider: 'openai',
+        rateLimits: { models: { 'gpt-4o': { tpm: 700 } } },
+      });
+      updateLimitsFromHeaders('gpt-4o', { 'x-ratelimit-limit-tokens': '800' });
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '900';
+      expect(getCurrentTokenLimit('gpt-4o')).toEqual({ tpm: 900, source: 'env' });
+    });
+
+    it('gives an explicit model config override priority over learned headers', () => {
+      configureTokenControl({
+        provider: 'openai',
+        rateLimits: { models: { 'gpt-4o': { tpm: 700, fallbackTPM: 600 } } },
+      });
+      updateLimitsFromHeaders('gpt-4o', { 'x-ratelimit-limit-tokens': '800' });
+      expect(getCurrentTokenLimit('gpt-4o')).toEqual({ tpm: 700, source: 'config' });
+    });
+
+    it('ignores invalid environment and config values', () => {
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '-4';
+      configureTokenControl({
+        provider: 'openai',
+        rateLimits: { models: { 'future-model': { tpm: 'nope', fallbackTPM: 0 } } },
+      });
+      expect(getCurrentTokenLimit('future-model')).toEqual({ tpm: 150000, source: 'default' });
+    });
+
     it('updates limits from headers (header wins)', () => {
       const before = getCurrentTokenLimit('gpt-4o');
       updateLimitsFromHeaders('gpt-4o', { 'x-ratelimit-limit-tokens': '1234' });
@@ -113,6 +169,71 @@ describe('tokenControl', () => {
       if (before.tpm !== 1234) expect(after.tpm).not.toBe(before.tpm);
     });
 
+    it('uses selected-model config fallback but lets learned headers supersede it', () => {
+      configureTokenControl({
+        provider: 'openai',
+        rateLimits: { safetyMargin: 0.9, models: { 'gpt-4o': { fallbackTPM: 777 } } },
+      });
+      expect(getCurrentTokenLimit('gpt-4o')).toMatchObject({ tpm: 777, source: 'config-fallback' });
+      updateLimitsFromHeaders('gpt-4o', { 'x-ratelimit-limit-tokens': '888' });
+      expect(getCurrentTokenLimit('gpt-4o')).toMatchObject({ tpm: 888, source: 'header' });
+    });
+
+    it('reads Headers objects and captures remaining/reset/request metadata', () => {
+      const headers = new Headers({
+        'x-ratelimit-limit-tokens': '1234',
+        'x-ratelimit-remaining-tokens': '1000',
+        'x-ratelimit-reset-tokens': '2s',
+        'x-ratelimit-limit-requests': '50',
+      });
+      updateLimitsFromHeaders('gpt-4o', headers);
+      expect(getCurrentTokenLimit('gpt-4o')).toMatchObject({
+        tpm: 1234,
+        remainingTokens: 1000,
+        resetTokens: '2s',
+        rpm: 50,
+      });
+    });
+
+    it('accepts alternate token-limit header names case-insensitively', () => {
+      updateLimitsFromHeaders('gpt-4o', { 'X-RateLimit-Limit-TPM': '4321' });
+      expect(getCurrentTokenLimit('gpt-4o')).toMatchObject({ tpm: 4321, source: 'header' });
+    });
+
+    it('ignores missing models and invalid learned limits', () => {
+      updateLimitsFromHeaders('', { 'x-ratelimit-limit-tokens': '999' });
+      updateLimitsFromHeaders('gpt-4o', { 'x-ratelimit-limit-tokens': 'not-a-number' });
+      expect(getCurrentTokenLimit('gpt-4o')).toMatchObject({ tpm: 450000, source: 'default' });
+    });
+
+    it('learns retry timing from a structured rate-limit error message', () => {
+      expect(updateLimitsFromError('gemini-2.5-pro', {
+        message: 'Quota exhausted. Please retry in 11.5s',
+      })).toEqual({ retryAfterMs: 11500 });
+    });
+
+    it('prefers Retry-After from a Headers object', () => {
+      expect(updateLimitsFromError('gpt-4o', {
+        headers: new Headers({
+          'retry-after': '3.25',
+          'x-ratelimit-limit-tokens': '5000',
+        }),
+        message: 'retry in 20s',
+      })).toEqual({ retryAfterMs: 3250 });
+      expect(getCurrentTokenLimit('gpt-4o')).toMatchObject({ tpm: 5000, source: 'header' });
+    });
+
+    it('returns no retry duration when an error provides none', () => {
+      expect(updateLimitsFromError('gpt-4o', { message: 'rate limited' }))
+        .toEqual({ retryAfterMs: undefined });
+    });
+
+    it('clears learned state when runtime context is reconfigured', () => {
+      updateLimitsFromHeaders('gpt-4o', { 'x-ratelimit-limit-tokens': '1234' });
+      configureTokenControl({ provider: 'openai' });
+      expect(getCurrentTokenLimit('gpt-4o')).toMatchObject({ tpm: 450000, source: 'default' });
+    });
+
     it('uses tier-1 defaults without letting generic GPT-5 rules shadow variants', () => {
       expect(getCurrentTokenLimit('gpt-5.6').tpm).toBe(500000);
       expect(getCurrentTokenLimit('gpt-5.6-luna').tpm).toBe(500000);
@@ -153,7 +274,7 @@ describe('tokenControl', () => {
 
       const result = await tokenUseCoolOff(60, entries, 'any-model');
       expect(result.shouldBackoff).toBe(true);
-      expect(result.totalTokensUsed).toBe(0);
+      expect(result.totalTokensUsed).toBeGreaterThanOrEqual(0);
       expect(Array.isArray(result.turnTimestamps)).toBe(true);
     });
 
@@ -165,5 +286,43 @@ describe('tokenControl', () => {
       expect(result.shouldBackoff).toBe(false);
       expect(result.totalTokensUsed).toBe(60);
     });
+
+    it('reserves the projected next request against the safety margin', async () => {
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '100';
+      const entries = [];
+      recordTokenUsage(entries, 80);
+      const result = await tokenUseCoolOff(80, entries, 'any-model', 15);
+      expect(result.shouldBackoff).toBe(true);
+    });
+
+    it('honors a configured safety margin', async () => {
+      configureTokenControl({
+        provider: 'openai',
+        rateLimits: {
+          safetyMargin: 0.5,
+          models: { 'gpt-4o': { tpm: 100 } },
+        },
+      });
+      const entries = [];
+      recordTokenUsage(entries, 45);
+      expect((await tokenUseCoolOff(45, entries, 'gpt-4o', 6)).shouldBackoff).toBe(true);
+    });
+
+    it('does not wait forever when one projected request alone exceeds the limit', async () => {
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '10';
+      const result = await tokenUseCoolOff(0, [], 'any-model', 100);
+      expect(result.shouldBackoff).toBe(false);
+      expect(wait).not.toHaveBeenCalled();
+    });
+
+    it('waits until the oldest necessary rolling entry expires', async () => {
+      process.env.TESTRONAUT_TOKENS_PER_MIN = '100';
+      const now = 1_000_000;
+      const nowSpy = vi.spyOn(Date, 'now').mockReturnValue(now);
+      const entries = [[now - 50000, 40], [now - 10000, 40]];
+      await tokenUseCoolOff(80, entries, 'any-model', 15);
+      expect(wait).toHaveBeenCalledWith(10000);
+      nowSpy.mockRestore();
+    });
   });
 });

package/todo.md

@@ -15,3 +15,11 @@
 - Learn effective OpenAI TPM/RPM limits from successful response headers and
   rate-limit errors. Prefer observed project/model limits over inferring a named
   account tier; retain the environment override and conservative startup values.
+- Persist provider-observed rate limits in a separate, expiring local cache keyed
+  by provider, API project/account, endpoint/region, model, and traffic type.
+  Never rewrite user config or infer a named usage tier from a TPM value; expose
+  configured, cached, and live sources distinctly and discard stale observations.
+- After dynamic rate limits are established, make DOM injection quota-aware.
+  Budget DOM tokens against both remaining TPM and the model context window,
+  preserve mission instructions and actionable controls, and fall back to
+  semantic sections plus on-demand DOM chunks instead of silent truncation.

package/tools/rateLimitDefaults.js

@@ -0,0 +1,40 @@
+const MODEL_LIMITS = [
+  { provider: 'openai', test: /^gpt-5\.6(-|$)/i, tpm: 500000 },
+  { provider: 'openai', test: /^gpt-5\.5(-|$)/i, tpm: 500000 },
+  { provider: 'openai', test: /^gpt-5\.4-nano(-|$)/i, tpm: 200000 },
+  { provider: 'openai', test: /^gpt-5\.4(-|$)/i, tpm: 500000 },
+  { provider: 'openai', test: /^gpt-5\.2(-|$)/i, tpm: 500000 },
+  { provider: 'openai', test: /^gpt-5\.1(-|$)/i, tpm: 120000 },
+  { provider: 'openai', test: /^gpt-5-mini(-|$)/i, tpm: 240000 },
+  { provider: 'openai', test: /^gpt-5-nano(-|$)/i, tpm: 600000 },
+  { provider: 'openai', test: /^gpt-5(-|$)/i, tpm: 90000 },
+  { provider: 'openai', test: /^gpt-4o(-|$)/i, tpm: 450000 },
+  { provider: 'openai', test: /^gpt-4\.1(-|$)/i, tpm: 1000000 },
+  { provider: 'openai', test: /^o3(-|$)/i, tpm: 300000 },
+  { provider: 'openai', test: /^o4-mini(-|$)/i, tpm: 600000 },
+  { provider: 'openai', test: /^gpt-4(-|$)/i, tpm: 150000 },
+  { provider: 'openai', test: /^gpt-3\.5(-|$)/i, tpm: 600000 },
+  { provider: 'gemini', test: /^gemini-2\.5-pro(-|$)/i, tpm: 120000 },
+  { provider: 'gemini', test: /^gemini-2\.5-flash(-|$)/i, tpm: 300000 },
+  { provider: 'anthropic', test: /^claude-(opus|sonnet)-/i, tpm: 80000 },
+  { provider: 'anthropic', test: /^claude-haiku-/i, tpm: 100000 },
+];
+
+export function getFallbackTPM(provider, model) {
+  const normalizedProvider = provider === 'claude' ? 'anthropic' : provider;
+  const hit = MODEL_LIMITS.find(entry =>
+    (!normalizedProvider || entry.provider === normalizedProvider) && entry.test.test(model || '')
+  );
+  return hit?.tpm ?? 150000;
+}
+
+export function defaultRateLimitConfig(provider, model) {
+  return {
+    tier: 'unknown',
+    region: 'global',
+    safetyMargin: 0.9,
+    models: {
+      [model]: { fallbackTPM: getFallbackTPM(provider, model) },
+    },
+  };
+}

package/tools/tokenControl.js

@@ -23,6 +23,7 @@
 import { encoding_for_model, get_encoding } from '@dqbd/tiktoken';
 import { wait } from './turnLoopUtils.js';
 import { getOpenAIModel } from '../llm/openAI/models.js';
+import { getFallbackTPM } from './rateLimitDefaults.js';
 
 /**
  * Dynamic token-per-minute limits by model family.
@@ -34,48 +35,36 @@ import { getOpenAIModel } from '../llm/openAI/models.js';
  * 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 = [
-  // ── OpenAI (newer first) ────────────────────────────────────────────────
-  { test: /^gpt-5\.6-luna(-|$)/i,        tpm: 500000 },
-  { test: /^gpt-5\.6(-|$)/i,             tpm: 500000 },
-  { test: /^gpt-5\.5(-|$)/i,             tpm: 500000 },
-  { test: /^gpt-5\.4-nano(-|$)/i,        tpm: 200000 },
-  { test: /^gpt-5\.4(-|$)/i,             tpm: 500000 },
-  { test: /^gpt-5\.2(-|$)/i,             tpm: 500000 },
-  { test: /^gpt-5.1(-|$)/i,             tpm: 120000 },
-  { test: /^gpt-5-mini(-|$)/i,          tpm: 240000 },
-  { test: /^gpt-5-nano(-|$)/i,          tpm: 600000 },
-  { test: /^gpt-5(-|$)/i,               tpm:  90000 },
-
-  { test: /^gpt-4o(-|$)/i,              tpm: 450000 },
-  { test: /^gpt-4\.1(-|$)/i,            tpm:1000000 },
-
-  { test: /^o3(-|$)/i,                  tpm: 300000 },
-  { test: /^o4-mini(-|$)/i,             tpm: 600000 },
-
-  // 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 },
-
-  // Anthropic limits vary by usage tier; this conservative fallback is only for local pacing.
-  { test: /^claude-(opus|sonnet)-/i,     tpm: 80000 },
-  { test: /^claude-haiku-/i,             tpm: 100000 },
-
-  // Ultimate fallback for anything else
-  { test: /.*/,                         tpm: 150000 },
-];
-
+/* Legacy model table moved to rateLimitDefaults.js so init and runtime share it. */
 // Live, mutable limits (can be updated by headers at runtime)
 const liveLimits = new Map(); // modelId -> { tpm, source: 'default'|'env'|'header' }
 
 // One-time warning tracking for tokenizer fallback
 const warnedModels = new Set();
+let runtimeContext = { provider: undefined, rateLimits: undefined };
+let overrideNoticeShown = false;
+
+function getEnvironmentTPMOverride() {
+  const raw = String(process.env.TESTRONAUT_TOKENS_PER_MIN ?? '').trim();
+  if (!raw || raw.toLowerCase() === 'auto') return null;
+  const value = Number(raw);
+  return Number.isFinite(value) && value > 0 ? value : null;
+}
+
+export function configureTokenControl({ provider, rateLimits } = {}) {
+  runtimeContext = { provider, rateLimits };
+  liveLimits.clear();
+  const envTPM = getEnvironmentTPMOverride();
+  if (envTPM && !overrideNoticeShown) {
+    console.warn(
+      `⚠️ TPM override active: TESTRONAUT_TOKENS_PER_MIN=${envTPM}. ` +
+      'This value overrides config and provider-advertised limits.\n' +
+      '   Clear it permanently with `unset TESTRONAUT_TOKENS_PER_MIN` and remove it from `.env` or your shell profile.\n' +
+      '   Bypass it for one run with `TESTRONAUT_TOKENS_PER_MIN=auto testronaut <mission>`.'
+    );
+    overrideNoticeShown = true;
+  }
+}
 
 /* ---------------- Tokenizer helpers ----------------
  * We prefer tiktoken's per-model encoding when available.
@@ -185,15 +174,33 @@ export async function warnIfContextNearLimit(model, payload, threshold = 0.9) {
  * @returns {{tpm:number, source:'default'|'env'}}
  */
 function resolveDefaultLimitForModel(model) {
-  const envTPM = process.env.TESTRONAUT_TOKENS_PER_MIN
-    ? Number(process.env.TESTRONAUT_TOKENS_PER_MIN)
-    : undefined;
-  if (envTPM && Number.isFinite(envTPM) && envTPM > 0) {
+  const envTPM = getEnvironmentTPMOverride();
+  if (envTPM) {
     return { tpm: envTPM, source: 'env' };
   }
 
-  const hit = DEFAULT_LIMITS.find(entry => entry.test.test(model || ''));
-  return { tpm: hit?.tpm ?? 150000, source: 'default' };
+  const modelConfig = runtimeContext.rateLimits?.models?.[model] || {};
+  const configuredTPM = Number(modelConfig.tpm);
+  if (Number.isFinite(configuredTPM) && configuredTPM > 0) {
+    return { tpm: configuredTPM, source: 'config' };
+  }
+  const fallbackTPM = Number(modelConfig.fallbackTPM);
+  return {
+    tpm: Number.isFinite(fallbackTPM) && fallbackTPM > 0
+      ? fallbackTPM
+      : getFallbackTPM(runtimeContext.provider, model),
+    source: Number.isFinite(fallbackTPM) && fallbackTPM > 0 ? 'config-fallback' : 'default',
+  };
+}
+
+function resolveExplicitLimitForModel(model) {
+  const envTPM = getEnvironmentTPMOverride();
+  if (envTPM) return { tpm: envTPM, source: 'env' };
+  const configuredTPM = Number(runtimeContext.rateLimits?.models?.[model]?.tpm);
+  if (Number.isFinite(configuredTPM) && configuredTPM > 0) {
+    return { tpm: configuredTPM, source: 'config' };
+  }
+  return null;
 }
 
 /**
@@ -205,6 +212,8 @@ function resolveDefaultLimitForModel(model) {
  */
 export function getCurrentTokenLimit(model) {
   const m = (model || '').trim() || 'unknown';
+  const explicit = resolveExplicitLimitForModel(m);
+  if (explicit) return explicit;
   const live = liveLimits.get(m);
   if (live?.tpm) return live;
 
@@ -225,7 +234,11 @@ export function updateLimitsFromHeaders(model, headers = {}) {
 
   // Normalize header keys to lowercase
   const lower = {};
-  for (const k of Object.keys(headers || {})) lower[k.toLowerCase()] = headers[k];
+  if (typeof headers?.forEach === 'function') {
+    headers.forEach((value, key) => { lower[String(key).toLowerCase()] = value; });
+  } else {
+    for (const k of Object.keys(headers || {})) lower[k.toLowerCase()] = headers[k];
+  }
 
   const tokenCap =
     Number(lower['x-ratelimit-limit-tokens']) ||
@@ -236,12 +249,38 @@ export function updateLimitsFromHeaders(model, headers = {}) {
   if (tokenCap && Number.isFinite(tokenCap) && tokenCap > 0) {
     const cur = getCurrentTokenLimit(model);
     if (cur.tpm !== tokenCap || cur.source !== 'header') {
-      liveLimits.set(model, { tpm: tokenCap, source: 'header' });
+      liveLimits.set(model, {
+        tpm: tokenCap,
+        remainingTokens: numericHeader(lower['x-ratelimit-remaining-tokens']),
+        resetTokens: lower['x-ratelimit-reset-tokens'],
+        rpm: numericHeader(lower['x-ratelimit-limit-requests']),
+        remainingRequests: numericHeader(lower['x-ratelimit-remaining-requests']),
+        resetRequests: lower['x-ratelimit-reset-requests'],
+        source: 'header',
+        observedAt: Date.now(),
+      });
       console.log(`📏 Updated TPM for ${model}: ${tokenCap} (from headers)`);
     }
   }
 }
 
+function numericHeader(value) {
+  const parsed = Number(value);
+  return Number.isFinite(parsed) && parsed >= 0 ? parsed : undefined;
+}
+
+export function updateLimitsFromError(model, error = {}) {
+  const headers = error.headers || error.response?.headers || {};
+  updateLimitsFromHeaders(model, headers);
+  const retryAfter = typeof headers?.get === 'function'
+    ? headers.get('retry-after')
+    : headers?.['retry-after'];
+  const message = String(error.message || error.error?.message || '');
+  const retryMatch = message.match(/retry(?: in| after)?\s+([0-9.]+)s/i);
+  const seconds = Number(retryAfter ?? retryMatch?.[1]);
+  return { retryAfterMs: Number.isFinite(seconds) && seconds >= 0 ? seconds * 1000 : undefined };
+}
+
 /* ---------------- Cooloff / backoff logic ---------------- */
 
 /**
@@ -252,14 +291,18 @@ export function updateLimitsFromHeaders(model, headers = {}) {
  * @param {string} model - model id for TPM lookup
  * @returns {Promise<{shouldBackoff:boolean,totalTokensUsed:number,turnTimestamps:Array}>}
  */
-export const tokenUseCoolOff = async (totalTokensUsed, turnTimestamps, model) => {
+export const tokenUseCoolOff = async (totalTokensUsed, turnTimestamps, model, projectedTokens = 0) => {
   const { tpm } = getCurrentTokenLimit(model);
-  if (totalTokensUsed > tpm) {
-    const msToWait = await getDynamicBackoffMs(turnTimestamps, tpm);
-    console.warn(`⚠️ Token throttle risk (${totalTokensUsed}/${tpm}) → Waiting ${Math.ceil((msToWait || 1000)/1000)}s...`);
+  const margin = Number(runtimeContext.rateLimits?.safetyMargin);
+  const effectiveLimit = tpm * (Number.isFinite(margin) && margin > 0 && margin <= 1 ? margin : 0.9);
+  const projectedTotal = totalTokensUsed + Math.max(0, Number(projectedTokens) || 0);
+  if (projectedTotal > effectiveLimit && turnTimestamps.length) {
+    const msToWait = await getDynamicBackoffMs(turnTimestamps, Math.max(0, effectiveLimit - projectedTokens));
+    console.warn(`⚠️ Token throttle risk (${Math.ceil(projectedTotal)}/${Math.floor(effectiveLimit)}) → Waiting ${Math.ceil((msToWait || 1000)/1000)}s...`);
     await wait(msToWait || 1000);
     console.log('✅ Backoff complete, resuming...');
-    return { shouldBackoff: true, totalTokensUsed: 0, turnTimestamps: [] };
+    const refreshed = pruneOldTokenUsage(turnTimestamps);
+    return { shouldBackoff: true, ...refreshed };
   }
   return { shouldBackoff: false, totalTokensUsed, turnTimestamps };
 };
@@ -298,14 +341,13 @@ export const pruneOldTokenUsage = (turnTimestamps, windowMs = 60000) => {
  */
 const getDynamicBackoffMs = async (turnTimestamps, tokenLimit) => {
   const now = Date.now();
-  let runningTotal = 0;
-
   const sorted = [...turnTimestamps].sort((a, b) => a[0] - b[0]);
+  let remainingTotal = sorted.reduce((sum, [, tokens]) => sum + tokens, 0);
   for (let i = 0; i < sorted.length; i++) {
-    runningTotal += sorted[i][1];
-    if (runningTotal > tokenLimit) {
-      const [timestampOfExcess] = sorted[i];
-      const msUntilSafe = 60000 - (now - timestampOfExcess);
+    remainingTotal -= sorted[i][1];
+    if (remainingTotal <= tokenLimit) {
+      const [timestampToExpire] = sorted[i];
+      const msUntilSafe = 60000 - (now - timestampToExpire);
       return Math.max(msUntilSafe, 1000); // at least 1s
     }
   }
@@ -320,4 +362,6 @@ const getDynamicBackoffMs = async (turnTimestamps, tokenLimit) => {
 export function __resetTokenControlForTests() {
   liveLimits.clear();
   warnedModels.clear();
+  runtimeContext = { provider: undefined, rateLimits: undefined };
+  overrideNoticeShown = false;
 }