testronaut 1.0.26 → 1.1.0

package/bin/cli.js

@@ -26,6 +26,27 @@ const apiBase = 'http://api.testronaut.app'; // Replace with your actual API bas
 
 
 const args = process.argv.slice(2);
+
+// Look for --model=<id> or --model <id>
+let modelOverride;
+const modelFlagIndex = args.findIndex(a => a === '--model' || a.startsWith('--model='));
+if (modelFlagIndex >= 0) {
+  if (args[modelFlagIndex].includes('=')) {
+    modelOverride = args[modelFlagIndex].split('=')[1];
+  } else if (args[modelFlagIndex + 1]) {
+    modelOverride = args[modelFlagIndex + 1];
+  }
+
+  if (modelOverride) {
+    process.env.TESTRONAUT_MODEL = modelOverride.trim();
+    console.log(`🧠 Model override: ${process.env.TESTRONAUT_MODEL}`);
+  }
+
+  // Remove flag & value from args so they don't look like mission filenames
+  args.splice(modelFlagIndex, modelOverride ? 2 : 1);
+}
+
+
 const missionsDir = path.resolve(process.cwd(), 'missions');
 
 const allResults = [];
@@ -165,10 +186,35 @@ const flatMissions = allResults.flatMap(entry => {
   }));
 });
 
+// Read provider/model from config (allow env override for model)
+const cfgPath = path.resolve(process.cwd(), 'testronaut-config.json');
+let llmProvider = 'openai';
+let llmModel = 'gpt-4o';
+
+if (fs.existsSync(cfgPath)) {
+  try {
+    const cfg = JSON.parse(fs.readFileSync(cfgPath, 'utf8'));
+    if (cfg?.provider) llmProvider = String(cfg.provider);
+    if (cfg?.model) llmModel = String(cfg.model);
+  } catch (e) {
+    console.warn('⚠️ Could not read testronaut-config.json:', e.message);
+  }
+}
+
+// Env override takes precedence for the model (matches runtime behavior)
+if (process.env.TESTRONAUT_MODEL?.trim()) {
+  llmModel = process.env.TESTRONAUT_MODEL.trim();
+}
+
+
 const report = {
   runId,
   startTime: startTime.toISOString(),
   endTime: endTime.toISOString(),
+  llm: {
+    provider: llmProvider,
+    model: llmModel,
+  },
   summary: {
     totalMissions: flatMissions.length,
     passed: flatMissions.filter(m => m.status === 'passed').length,

package/bin/init.js

@@ -13,33 +13,76 @@ export async function initializeTestronautProject() {
   let config = {};
   if (fs.existsSync(configPath)) {
     config = JSON.parse(fs.readFileSync(configPath, 'utf-8'));
+
+    // --- Back-compat migration: old configs stored provider in `model`
+    if (!config.provider && config.model === 'openai') {
+      config.provider = 'openai';
+      // don't overwrite model yet; we'll re-ask below only if not initialized
+    }
   }
 
   // If not initialized before, walk through setup
   if (!config.initialized) {
-    // Prompt for LLM
-    const response = await prompts({
+    // 1) Provider
+    const { llmProvider } = await prompts({
       type: 'select',
       name: 'llmProvider',
       message: 'Which LLM provider would you like to use?',
       choices: [
-        { title: 'OpenAI (GPT-4o)', value: 'openai' },
+        { title: 'OpenAI (GPT family)', value: 'openai' },
         // Add others like Anthropic, Mistral, etc. later
       ]
     });
 
-    config.model = response.llmProvider;
+    config.provider = llmProvider; // store provider explicitly
+
+    // 2) If OpenAI, ask for model
+    const { openaiModel } = await prompts({
+      type: 'select',
+      name: 'openaiModel',
+      message: 'Select an OpenAI model for agentic workflows (function/tool calling):',
+      choices: [
+        { title: 'GPT-4.1 (general purpose)', value: 'gpt-4.1' },
+        { title: 'GPT-4.1 mini (faster, cheaper)', value: 'gpt-4.1-mini' },
+        { title: 'GPT-4o (multimodal, tool use)', value: 'gpt-4o' },
+        { title: 'GPT-4o mini (speed/cost optimized)', value: 'gpt-4o-mini' },
+        { title: 'o3 (reasoning w/ native tool use)', value: 'o3' },
+        { title: 'o4-mini (reasoning, cost-effective)', value: 'o4-mini' },
+
+        // 👇 GPT-5 family
+        { title: 'GPT-5 (latest reasoning model)', value: 'gpt-5' },
+        { title: 'GPT-5 mini (faster)', value: 'gpt-5-mini' },
+        { title: 'GPT-5 nano (lightweight)', value: 'gpt-5-nano' },
+      ],
+      initial: 1
+    });
+
+    config.model = openaiModel;
+
+    // Warn about GPT-5 availability
+    if (openaiModel.startsWith('gpt-5')) {
+      console.log(`
+    ⚠️  Warning: GPT-5 support may not be available for all users.
+      - Access depends on your OpenAI account/region
+      - Some capabilities may be gated or rate-limited
+      `);
+    }
+
+    // Default config fields
     config.initialized = true;
-    config.outputDir = 'missions/mission_reports';
-    config.projectName = path.basename(root);
-    config.maxTurns = 20;
+    config.outputDir = config.outputDir || 'missions/mission_reports';
+    config.projectName = config.projectName || path.basename(root);
+    config.maxTurns = config.maxTurns ?? 20;
+
     fs.writeFileSync(configPath, JSON.stringify(config, null, 2));
     console.log('📝 Created testronaut-config.json');
 
-    if (response.llmProvider === 'openai') {
-      // Display .env instructions
+    if (llmProvider === 'openai') {
+      // Display .env instructions for OpenAI
       if (!fs.existsSync(envPath)) {
-        const exampleEnv = `# Add your OpenAI API key below\nOPENAI_API_KEY=sk-...\n`;
+        const exampleEnv = `# Add your OpenAI API key below
+OPENAI_API_KEY=sk-...
+`;
         fs.writeFileSync(envPath, exampleEnv);
         console.log('🔐 Created .env file');
       }
@@ -51,7 +94,7 @@ export async function initializeTestronautProject() {
 
     OPENAI_API_KEY=sk-...
 
-📌 The key must have permissions to call the OpenAI Chat Completions API.
+📌 The key must have permissions to call the OpenAI API.
 
 ✅ You’re now ready to write missions in \`missions/\` and run your first mission using:
 

package/core/turnLoop.js

@@ -11,8 +11,15 @@ import {
   tokenEstimate, 
   tokenUseCoolOff, 
   recordTokenUsage, 
-  pruneOldTokenUsage 
+  pruneOldTokenUsage,
+  getCurrentTokenLimit,
+  updateLimitsFromHeaders
 } from '../tools/tokenControl.js';
+import { resolveModel } from '../openAI/modelResolver.js';
+
+
+const MODEL_ID = resolveModel();
+console.log(`🧠 Using OpenAI model: ${MODEL_ID}`);
 
 const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
 
@@ -21,6 +28,14 @@ let turnTimestamps = [];
 let shouldBackoff;
 let steps = [];
 
+function formatBytes(n) {
+  if (!Number.isFinite(n)) return `${n}`;
+  const u = ['B','KB','MB','GB','TB'];
+  let i = 0, v = n;
+  while (v >= 1024 && i < u.length - 1) { v /= 1024; i++; }
+  return `${v.toFixed(v >= 10 || i === 0 ? 0 : 1)} ${u[i]}`;
+}
+
 const pushDOMAssistant = async (browser, messages, agentMemory, { skipIfLastTool } = {}) => {
   if (skipIfLastTool && skipIfLastTool.includes(messages.at(-1)?.name)) {
     console.log(`[skip] Skipping DOM push after redundant tool: ${messages.at(-1)?.name}`);
@@ -44,7 +59,7 @@ const pushDOMAssistant = async (browser, messages, agentMemory, { skipIfLastTool
   });
 
   const domHtml = await CHROME_TOOL_MAP.get_dom(browser, { limit: 100000, exclude: true }, agentMemory);
-  await tokenEstimate('gpt-4o', domHtml);
+  await tokenEstimate( MODEL_ID, domHtml);
   // fs.writeFileSync(`missions/mission_reports/debug-expanded-${Date.now()}.html`, domHtml);
 
   messages.push({
@@ -81,7 +96,7 @@ export const turnLoop = async (
       ({ turnTimestamps, totalTokensUsed } = pruneOldTokenUsage(turnTimestamps));
 
       // Cool off check
-      ({ totalTokensUsed, turnTimestamps, shouldBackoff } = await tokenUseCoolOff(totalTokensUsed, turnTimestamps));
+      ({ totalTokensUsed, turnTimestamps, shouldBackoff } = await tokenUseCoolOff(totalTokensUsed, turnTimestamps, MODEL_ID));
       if(shouldBackoff) return await turnLoop(browser, messages, maxTurns, turn, currentStep);
 
       //Tool Validator
@@ -115,12 +130,14 @@ export const turnLoop = async (
       
 
       response = await openai.chat.completions.create({
-        model: 'gpt-4o',
+        model: MODEL_ID,
         messages,
         tools: toolsSchema,
       });
     } catch (err) {
       if (err.status === 429) {
+        // If headers include limits, learn them for this model
+        try { updateLimitsFromHeaders(MODEL_ID, err.headers || err.response?.headers || {}); } catch {}
         const delay = Math.min(60000, 2 ** retryCount * 2000);
         console.warn(`⚠️ Rate limited. Retrying in ${delay / 1000}s... (retry ${retryCount + 1})`);
         await wait(delay);
@@ -187,6 +204,25 @@ export const turnLoop = async (
           result = errorMessage;
         }
 
+        // ------ Capture file upload/download events for mission report ------
+        try {
+          const maybeJson = JSON.parse(result);
+          if (maybeJson && maybeJson._testronaut_file_event) {
+            currentStep.files = currentStep.files || [];
+            currentStep.files.push(maybeJson);
+
+            if (maybeJson._testronaut_file_event === 'upload') {
+              const msgLine = `📤 Uploaded "${maybeJson.fileName}" (${formatBytes(maybeJson.bytes)}) via ${maybeJson.method}`;
+              currentStep.events.push(msgLine);
+              console.log(msgLine);
+            } else if (maybeJson._testronaut_file_event === 'download') {
+              const msgLine = `📥 Downloaded "${maybeJson.fileName}" (${formatBytes(maybeJson.bytes)}) via ${maybeJson.mode}`;
+              currentStep.events.push(msgLine);
+              console.log(msgLine);
+            }
+          }
+        } catch { /* not JSON / ignore */ }
+
         console.log(`[tool ] ← ${fnName} result:`, errorMessage ? '❌ Failed' : '✅ Success');
         currentStep.events.push(`[tool ] ← ${fnName} result: ${errorMessage ? '❌ Failed' : '✅ Success'}`)
         const truncated = result.length > 100 ? result.slice(0, 100) + '…' : result;
@@ -221,7 +257,7 @@ export const turnLoop = async (
             exclude: true,
             focus: [],
           }, agentMemory);
-          await tokenEstimate('gpt-4o', domHtml);
+          await tokenEstimate(MODEL_ID, domHtml);
           await pushDOMAssistant(browser, messages, agentMemory, {
             skipIfLastTool: ['get_dom', 'check_text']
           });

package/openAI/modelResolver.js

@@ -0,0 +1,33 @@
+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';
+}
+
+

package/package.json

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

package/tools/chromeBrowser.js

@@ -19,8 +19,26 @@ import {
   removeObfuscatedClassNames, 
 } from './domControl.js';
 import fs from 'fs';
+import path from 'path';
 import { ensureBrowsers } from '../tools/playwrightSetup.js';
 
+const FILES_DIR = path.join('missions', 'files');
+const REPORTS_DIR = path.join('missions', 'mission_reports');
+const FILES_LOG = path.join(REPORTS_DIR, 'files.jsonl');
+
+function ensureDir(p) {
+  try { fs.mkdirSync(p, { recursive: true }); } catch {}
+}
+
+function appendFileEvent(eventObj = {}) {
+  try {
+    ensureDir(REPORTS_DIR);
+    fs.appendFileSync(FILES_LOG, JSON.stringify(eventObj) + '\n', 'utf8');
+  } catch (e) {
+    console.warn('⚠️ Could not append file event:', e?.message);
+  }
+}
+
 export class ChromeBrowser {
   constructor() {
     this.browser = null;
@@ -549,7 +567,13 @@ export class ChromeBrowser {
       });
       // return await pw.chromium.launch({ headless: true });
     }
-    this.context = await this.browser.newContext();
+
+    ensureDir(FILES_DIR); // ⬅️ make sure missions/files exists
+
+    this.context = await this.browser.newContext({
+      acceptDownloads: true // ⬅️ important for Playwright download handling
+    });
+    // this.context = await this.browser.newContext();
     const first = await this.context.newPage();
     const id = this._addAndFocus(first);
 
@@ -991,7 +1015,127 @@ export class ChromeBrowser {
     return `Screenshot saved at: ./screenshots/${name}_${timestamp}.png`;
   }
   
-  
+  async upload_file({ selector, fileName, useChooser = false, timeoutMs = 15000 }) {
+    if (!selector) throw new Error('upload_file: selector is required');
+    if (!fileName) throw new Error('upload_file: fileName is required');
+
+    ensureDir(FILES_DIR);
+    const filePath = path.resolve(FILES_DIR, fileName);
+    if (!fs.existsSync(filePath)) {
+      throw new Error(`upload_file: file not found at ${filePath}`);
+    }
+
+    await this.page.waitForSelector(selector, { timeout: timeoutMs });
+
+    let method = 'input';
+    if (useChooser) {
+      method = 'chooser';
+      const [fileChooser] = await Promise.all([
+        this.page.waitForEvent('filechooser', { timeout: timeoutMs }),
+        this.page.click(selector, { timeout: timeoutMs })
+      ]);
+      await fileChooser.setFiles(filePath);
+    } else {
+      const input = this.page.locator(selector).first();
+      await input.waitFor({ state: 'visible', timeout: timeoutMs }).catch(() => {});
+      await input.setInputFiles(filePath, { timeout: timeoutMs });
+    }
+
+    const stats = fs.statSync(filePath);
+    const payload = {
+      _testronaut_file_event: 'upload',
+      method,                       // 'input' | 'chooser'
+      fileName,
+      sourcePath: filePath,
+      bytes: stats.size,
+      selector,
+      ts: new Date().toISOString()
+    };
+
+    appendFileEvent(payload);
+    return JSON.stringify(payload);
+  }
+
+
+  async download_file({ selector, url, fileName, timeoutMs = 30000 }) {
+    ensureDir(FILES_DIR);
+
+    if (!selector && !url) {
+      throw new Error('download_file: provide either selector or url');
+    }
+
+    const sanitize = (name) => name.replace(/[\\/:*?"<>|]+/g, '_').trim();
+
+    // Mode 1: page click triggers a browser download
+    if (selector) {
+      await this.page.waitForSelector(selector, { timeout: timeoutMs });
+
+      const [download] = await Promise.all([
+        this.page.waitForEvent('download', { timeout: timeoutMs }),
+        this.page.click(selector)
+      ]);
+
+      const suggested = sanitize(download.suggestedFilename?.() || 'download.bin');
+      const targetPath = path.join(FILES_DIR, sanitize(fileName || suggested));
+      await download.saveAs(targetPath);
+
+      const stats = fs.statSync(targetPath);
+      const payload = {
+        _testronaut_file_event: 'download',
+        mode: 'click',
+        fileName: path.basename(targetPath),
+        destPath: targetPath,
+        bytes: stats.size,
+        selector,
+        url: null,
+        ts: new Date().toISOString()
+      };
+
+      appendFileEvent(payload);
+      return JSON.stringify(payload);
+    }
+
+    // Mode 2: direct URL fetch (no page interaction)
+    const res = await fetch(url);
+    if (!res.ok) {
+      throw new Error(`download_file: GET ${url} failed with ${res.status}`);
+    }
+
+    // Try to infer a filename if not provided
+    const cd = res.headers.get('content-disposition') || '';
+    let inferred = 'download.bin';
+    const match = /filename\*=UTF-8''([^;]+)|filename="([^"]+)"/i.exec(cd);
+    if (match) {
+      inferred = decodeURIComponent(match[1] || match[2]);
+    } else {
+      try {
+        const u = new URL(url);
+        const tail = u.pathname.split('/').filter(Boolean).pop();
+        if (tail) inferred = tail;
+      } catch {}
+    }
+
+    const targetPath = path.join(FILES_DIR, sanitize(fileName || inferred));
+    const buf = Buffer.from(await res.arrayBuffer());
+    fs.writeFileSync(targetPath, buf);
+
+    const stats = fs.statSync(targetPath);
+    const payload = {
+      _testronaut_file_event: 'download',
+      mode: 'url',
+      fileName: path.basename(targetPath),
+      destPath: targetPath,
+      bytes: stats.size,
+      selector: null,
+      url,
+      ts: new Date().toISOString()
+    };
+
+    appendFileEvent(payload);
+    return JSON.stringify(payload);
+  }
+
+
   async close() {
     await this.browser?.close();
   }
@@ -1016,4 +1160,6 @@ export const CHROME_TOOL_MAP = {
   click_and_follow_popup: (b, args) => b.click_and_follow_popup(args),
   switch_to_page: (b, args) => b.switch_to_page(args),
   close_current_page: (b, args) => b.close_current_page(args),
+  upload_file: (b, args) => b.upload_file(args),
+  download_file: (b, args) => b.download_file(args),
 };

package/tools/generateHtmlReport.js

@@ -3,7 +3,7 @@ import fs from 'fs';
 import path from 'path';
 
 export function generateHtmlReport(report, outputPath) {
-  const { runId, startTime, endTime, missions = [], summary = {} } = report;
+  const { runId, startTime, endTime, missions = [], summary = {}, llm = {} } = report;
   const durationSec =
     (startTime && endTime)
       ? ((new Date(endTime) - new Date(startTime)) / 1000).toFixed(2)
@@ -241,6 +241,7 @@ export function generateHtmlReport(report, outputPath) {
       <div><strong>Run ID:</strong> ${esc(runId ?? '—')}</div>
       <div><strong>Start:</strong> ${esc(startTime ?? '—')}</div>
       <div><strong>End:</strong> ${esc(endTime ?? '—')} • <strong>Duration:</strong> ${durationSec}s</div>
+      <div><strong>LLM:</strong> ${esc(llm.provider ?? '—')} • <strong>Model:</strong> ${esc(llm.model ?? '—')}</div>
     </div>
   </div>
 
@@ -248,6 +249,7 @@ export function generateHtmlReport(report, outputPath) {
     <div class="pill">Missions: ${esc(summary.totalMissions ?? totals.total)}</div>
     <div class="pill ok">Passed: ${esc(summary.passed ?? totals.passed)}</div>
     <div class="pill bad">Failed: ${esc(summary.failed ?? totals.failed)}</div>
+    <div class="pill">LLM: ${esc(llm.provider ?? '—')} • ${esc(llm.model ?? '')}</div>
   </div>
 
   <div class="container">

package/tools/tokenControl.js

@@ -1,33 +1,164 @@
-import { encoding_for_model } from '@dqbd/tiktoken';
+import { encoding_for_model, get_encoding } from '@dqbd/tiktoken';
 import { wait } from './turnLoopUtils.js';
 
-const TOKEN_BACKOFF_MS = 60000;
-const TOKEN_LIMIT_PER_MIN = 90000;
+/**
+ * Dynamic token-per-minute limits by model family.
+ * - You can tweak these defaults anytime.
+ * - 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).
+ */
+const DEFAULT_LIMITS = [
+  // Newer families first (checked in order)
+  { 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: /^o3(-|$)/i,                  tpm: 300000 },
+  { test: /^o4-mini(-|$)/i,             tpm: 600000 },
+
+  // Older / fallback
+  { test: /^gpt-4(-|$)/i,               tpm: 150000 },
+  { test: /^gpt-3\.5(-|$)/i,            tpm: 600000 },
+
+  // Ultimate fallback for anything else
+  { test: /.*/,                         tpm: 150000 },
+];
+
+// 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();
+
+/* ---------------- Tokenizer helpers (same idea as before) ---------------- */
+
+function getTokenizer(model) {
+  try {
+    return encoding_for_model(model);
+  } catch (_) {
+    const m = (model || '').toLowerCase();
+    const useO200k =
+      m.startsWith('gpt-5') ||
+      m.startsWith('gpt-4o') ||
+      m.startsWith('gpt-4.1') ||
+      m.startsWith('o') ||
+      m.includes('omni');
+
+    try {
+      return get_encoding(useO200k ? 'o200k_base' : 'cl100k_base');
+    } catch (e2) {
+      if (!warnedModels.has(model)) {
+        console.warn(`⚠️ Could not load tokenizer for "${model}". Will approximate by bytes/4.`, e2?.message || e2);
+        warnedModels.add(model);
+      }
+      return null;
+    }
+  }
+}
 
 export const tokenEstimate = async (model, text) => {
-  const encoding = encoding_for_model(model);
-  const tokenCount = encoding.encode(text).length;
-  console.log(`🧠 Estimated token count: ${tokenCount}`);
-  return tokenCount;
+  const str = typeof text === 'string' ? text : JSON.stringify(text ?? '');
+  const enc = getTokenizer(model);
+
+  if (enc) {
+    try {
+      const tokenCount = enc.encode(str).length;
+      enc.free?.();
+      console.log(`🧠 Estimated token count (${model}): ${tokenCount}`);
+      return tokenCount;
+    } catch (e) {
+      enc.free?.();
+      if (!warnedModels.has(model)) {
+        console.warn(`⚠️ Tokenizer encode failed for "${model}". Falling back to bytes/4.`, e?.message || e);
+        warnedModels.add(model);
+      }
+    }
+  }
+
+  // Final fallback: rough heuristic (UTF-8 bytes / 4)
+  const bytes = Buffer.from(str, 'utf8').length;
+  const approx = Math.ceil(bytes / 4);
+  console.log(`🧠 Estimated token count (approx, ${model}): ${approx}`);
+  return approx;
+};
+
+/* ---------------- Dynamic limit resolution ---------------- */
+
+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;
+  if (envTPM && Number.isFinite(envTPM) && envTPM > 0) {
+    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.
+ */
+export function getCurrentTokenLimit(model) {
+  const m = (model || '').trim() || 'unknown';
+  const live = liveLimits.get(m);
+  if (live?.tpm) return live;
+
+  const resolved = resolveDefaultLimitForModel(m);
+  liveLimits.set(m, resolved);
+  return resolved;
 }
 
-export const tokenUseCoolOff = async (totalTokensUsed, turnTimestamps) => {
-  // console.log("total tokens used: ", totalTokensUsed);
-  // console.log("turnTimestamps: ", turnTimestamps);
-  
-  if (totalTokensUsed > TOKEN_LIMIT_PER_MIN) {
-    const msToWait = await getDynamicBackoffMs(turnTimestamps);
-    console.warn(`⚠️ Token throttle risk → Waiting ${msToWait / 1000}s to cool off...`);
-    await wait(msToWait);
-    // totalTokensUsed = 0; // Reset token count after backoff
-    // turnTimestamps = []; // Reset timestamps after backoff
+/**
+ * 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').
+ */
+export function updateLimitsFromHeaders(model, headers = {}) {
+  if (!model) return;
+
+  // Normalize header keys to lowercase
+  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']) ||
+    Number(lower['x-ratelimit-limit-token']) ||
+    undefined;
+
+  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' });
+      console.log(`📏 Updated TPM for ${model}: ${tokenCap} (from headers)`);
+    }
+  }
+}
+
+/* ---------------- Cooloff / backoff logic ---------------- */
+
+export const tokenUseCoolOff = async (totalTokensUsed, turnTimestamps, model) => {
+  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...`);
+    await wait(msToWait || 1000);
     console.log('✅ Backoff complete, resuming...');
-    // return true;
     return { shouldBackoff: true, totalTokensUsed: 0, turnTimestamps: [] };
   }
-  // return false;
   return { shouldBackoff: false, totalTokensUsed, turnTimestamps };
-}
+};
 
 export const recordTokenUsage = (turnTimestamps, tokensUsed) => {
   const now = Date.now();
@@ -41,22 +172,18 @@ export const pruneOldTokenUsage = (turnTimestamps, windowMs = 60000) => {
   return { turnTimestamps: recentEntries, totalTokensUsed };
 };
 
-const getDynamicBackoffMs = async (turnTimestamps, tokenLimit = TOKEN_LIMIT_PER_MIN) => {
+const getDynamicBackoffMs = async (turnTimestamps, tokenLimit) => {
   const now = Date.now();
   let runningTotal = 0;
 
-  // Sort oldest → newest to find the earliest point to fall below the limit
   const sorted = [...turnTimestamps].sort((a, b) => a[0] - b[0]);
-
   for (let i = 0; i < sorted.length; i++) {
-    runningTotal += sorted[i][1]; // add tokensUsed
-
+    runningTotal += sorted[i][1];
     if (runningTotal > tokenLimit) {
       const [timestampOfExcess] = sorted[i];
       const msUntilSafe = 60000 - (now - timestampOfExcess);
-      return Math.max(msUntilSafe, 1000); // wait at least 1s
+      return Math.max(msUntilSafe, 1000); // at least 1s
     }
   }
-
-  return 0; // Safe to proceed
-}
+  return 0;
+};

package/tools/toolSchema.js

@@ -174,59 +174,118 @@ const toolsSchema = [
     },
   },
   {
-  type: 'function',
-  function: {
-    name: 'click_and_follow_popup',
-    description:
-      'Click a selector that opens a new tab/window (e.g., OAuth) and switch control to it. Falls back to same-tab navigation if no popup appears.',
-    parameters: {
-      type: 'object',
-      properties: {
-        selector: { type: 'string', description: 'CSS or text selector to click' },
-        expectedUrlHost: {
-          type: 'string',
-          description: "Optional hostname hint like 'accounts.google.com'"
+    type: 'function',
+    function: {
+      name: 'click_and_follow_popup',
+      description:
+        'Click a selector that opens a new tab/window (e.g., OAuth) and switch control to it. Falls back to same-tab navigation if no popup appears.',
+      parameters: {
+        type: 'object',
+        properties: {
+          selector: { type: 'string', description: 'CSS or text selector to click' },
+          expectedUrlHost: {
+            type: 'string',
+            description: "Optional hostname hint like 'accounts.google.com'"
+          },
+          timeoutMs: {
+            type: 'number',
+            default: 20000,
+            description: 'How long to wait for the popup or navigation'
+          }
         },
-        timeoutMs: {
-          type: 'number',
-          default: 20000,
-          description: 'How long to wait for the popup or navigation'
-        }
-      },
-      required: ['selector']
+        required: ['selector']
+      }
     }
-  }
-},
-{
-  type: 'function',
-  function: {
-    name: 'switch_to_page',
-    description:
-      "Switch agent focus between known tabs/pages (e.g., back to 'main' or to the latest popup). Use a pageId previously returned by click_and_follow_popup to target a specific tab.",
-    parameters: {
-      type: 'object',
-      properties: {
-        target: {
-          type: 'string',
-          description: "One of: 'main', 'latest', or a pageId returned by click_and_follow_popup"
-        }
-      },
-      required: ['target']
+  },
+  {
+    type: 'function',
+    function: {
+      name: 'switch_to_page',
+      description:
+        "Switch agent focus between known tabs/pages (e.g., back to 'main' or to the latest popup). Use a pageId previously returned by click_and_follow_popup to target a specific tab.",
+      parameters: {
+        type: 'object',
+        properties: {
+          target: {
+            type: 'string',
+            description: "One of: 'main', 'latest', or a pageId returned by click_and_follow_popup"
+          }
+        },
+        required: ['target']
+      }
+    }
+  },
+  {
+    type: 'function',
+    function: {
+      name: 'close_current_page',
+      description:
+        'Close the current tab/page and automatically switch back to another open page if available.',
+      parameters: { type: 'object', properties: {} }
+    }
+  },
+  {
+    type: 'function',
+    function: {
+      name: 'upload_file',
+      description:
+        'Upload a file into an <input type="file"> or via a file chooser. Defaults to missions/files as the file source.',
+      parameters: {
+        type: 'object',
+        properties: {
+          selector: {
+            type: 'string',
+            description: 'CSS selector for the file input or a button that opens a file chooser'
+          },
+          fileName: {
+            type: 'string',
+            description: 'File name located in missions/files (e.g., "sample.pdf")'
+          },
+          useChooser: {
+            type: 'boolean',
+            default: false,
+            description: 'If true, click the selector and fulfill the file chooser instead of setInputFiles'
+          },
+          timeoutMs: {
+            type: 'number',
+            default: 15000
+          }
+        },
+        required: ['selector', 'fileName']
+      }
+    }
+  },
+  {
+    type: 'function',
+    function: {
+      name: 'download_file',
+      description:
+        'Download a file to missions/files. Either click a selector that triggers a download or fetch a direct URL.',
+      parameters: {
+        type: 'object',
+        properties: {
+          selector: {
+            type: 'string',
+            description: 'An element to click that triggers a download'
+          },
+          url: {
+            type: 'string',
+            description: 'Direct URL to download without clicking the page'
+          },
+          fileName: {
+            type: 'string',
+            description: 'Optional override for the saved filename'
+          },
+          timeoutMs: {
+            type: 'number',
+            default: 30000
+          }
+        },
+        // At least one of selector or url should be provided (enforced in code)
+        required: []
+      }
     }
   }
-},
-{
-  type: 'function',
-  function: {
-    name: 'close_current_page',
-    description:
-      'Close the current tab/page and automatically switch back to another open page if available.',
-    parameters: { type: 'object', properties: {} }
-  }
-}
-
-  
-    
 ];
 
 export default toolsSchema;