agent-tank 0.9.5 → 0.9.6

package/README.md

@@ -89,6 +89,9 @@ agent-tank --no-docker
 # Monitor only specific agents
 agent-tank --claude --gemini
 
+# Prefer direct Gemini quota fetching
+agent-tank --gemini --gemini-mode direct
+
 # Use a custom port
 agent-tank --port 8080
 
@@ -185,6 +188,7 @@ Options:
   --claude              Enable Claude monitoring
   --gemini              Enable Gemini monitoring
   --codex               Enable Codex monitoring
+  --gemini-mode <mode>  Gemini strategy: fallback, direct, pty (default: fallback)
   --port <port>         HTTP server port (default: 3456)
   --host <host>         Bind address (default: 127.0.0.1 + Docker bridge when available)
   --docker              Enable Docker bridge bind when --host is omitted (default: true)
@@ -224,6 +228,7 @@ Environment variables override CLI flags and config file values.
 | `AGENT_TANK_DOCKER` | Enable/disable Docker bridge bind when `--host` is omitted |
 | `AGENT_TANK_FRESH_PROCESS` | Use fresh process per refresh (`1`/`true`) |
 | `AGENT_TANK_CLAUDE_API` | Use Claude API mode (`1`/`true`) |
+| `AGENT_TANK_GEMINI_MODE` | Gemini strategy: `fallback`, `direct`, or `pty` |
 | `AGENT_TANK_AUTO_REFRESH` | Enable/disable auto-refresh |
 | `AGENT_TANK_AUTO_REFRESH_MODE` | `none`, `interval`, or `activity` |
 | `AGENT_TANK_AUTO_REFRESH_INTERVAL` | Auto-refresh interval in seconds |
@@ -243,6 +248,7 @@ Environment variables override CLI flags and config file values.
   "port": 8080,
   "dockerAccess": true,
   "claudeApi": false,
+  "geminiMode": "fallback",
   "refreshCooldown": 30,
   "autoRefresh": {
     "mode": "activity",
@@ -266,6 +272,28 @@ Run with:
 agent-tank -c config.json
 ```
 
+## Gemini Modes
+
+By default, Gemini runs in `fallback` mode:
+
+- try a direct quota fetch using the local Gemini OAuth session
+- fall back to PTY `/stats` if the direct path fails
+
+Available modes:
+
+- `fallback`: direct first, PTY fallback
+- `direct`: use the direct quota API only
+- `pty`: use interactive Gemini CLI `/stats` only
+
+In `--once` mode, Agent Tank prints the result and exits immediately after the output is flushed.
+
+Examples:
+
+```bash
+agent-tank --gemini --gemini-mode direct
+agent-tank --gemini --gemini-mode pty
+```
+
 ## HTTP API
 
 ### Endpoints

package/bin/agent-tank.js

@@ -11,6 +11,7 @@ const options = {
   gemini: { type: 'boolean', default: false },
   codex: { type: 'boolean', default: false },
   'claude-api': { type: 'boolean', default: false },
+  'gemini-mode': { type: 'string' },
   port: { type: 'string', default: '3456' },
   host: { type: 'string' },
   docker: { type: 'boolean', default: true },
@@ -43,6 +44,17 @@ function exitWithCode(code, message, stream = process.stderr) {
   process.exitCode = code;
 }
 
+function writeAndExit(stream, text, exitCode, { cleanup, watcher, jsonMode, originalLog }) {
+  stream.write(text, () => {
+    cleanup();
+    watcher.stop();
+    if (jsonMode) {
+      console.log = originalLog;
+    }
+    process.exit(exitCode);
+  });
+}
+
 function printHelp() {
   process.stdout.write(`
 agent-tank - Monitor LLM CLI usage limits
@@ -54,6 +66,7 @@ Options:
   --gemini              Enable Gemini monitoring
   --codex               Enable Codex monitoring
   --claude-api          Use direct Anthropic API for Claude usage (faster, 60s refresh)
+  --gemini-mode <mode>  Gemini strategy: fallback, direct, pty (default: fallback)
   --port <port>         HTTP server port (default: 3456)
   --host <host>         Bind address (default: 127.0.0.1 + Docker bridge when available)
   --docker              Enable Docker bridge bind when --host is omitted (default: true)
@@ -90,6 +103,7 @@ Environment variables:
   AGENT_TANK_DOCKER     Enable/disable Docker bridge bind ("1"/"true" or "0"/"false")
   AGENT_TANK_FRESH_PROCESS  Use fresh process per refresh ("1" or "true")
   AGENT_TANK_CLAUDE_API Use direct Anthropic API for Claude usage ("1" or "true")
+  AGENT_TANK_GEMINI_MODE  Gemini strategy: fallback, direct, pty
   AGENT_TANK_AUTO_REFRESH   Enable/disable background auto-refresh ("1" or "true" / "0" or "false")
   AGENT_TANK_AUTO_REFRESH_MODE      Refresh mode: none, interval, activity
   AGENT_TANK_AUTO_REFRESH_INTERVAL  Auto-refresh interval in seconds
@@ -105,6 +119,7 @@ Examples:
   agent-tank --port 8080              # Use custom port
   agent-tank --host 0.0.0.0           # Expose on all interfaces
   agent-tank --no-docker              # Bind localhost only when host is omitted
+  agent-tank --gemini-mode direct     # Prefer direct Gemini quota API calls
   agent-tank --auth-user admin --auth-pass secret  # Enable basic auth
   agent-tank --auth-token mykey       # Enable API key auth
   agent-tank -c ./config.json         # Use config file
@@ -188,6 +203,9 @@ async function main() {
     config.claudeApi ||
     claudeApiEnv === '1' || claudeApiEnv === 'true';
 
+  const geminiModeEnv = process.env.AGENT_TANK_GEMINI_MODE;
+  const geminiMode = geminiModeEnv || values['gemini-mode'] || config.geminiMode || 'fallback';
+
   // Auto-refresh configuration (env > CLI > config file)
   const autoRefreshEnv = process.env.AGENT_TANK_AUTO_REFRESH;
   let autoRefreshEnabled = values['auto-refresh'];
@@ -280,6 +298,7 @@ async function main() {
     auth,
     freshProcess,
     claudeApi,
+    geminiMode,
     autoRefreshEnabled: onceMode ? false : autoRefreshEnabled, // Disable auto-refresh in one-shot mode
     autoRefreshInterval,
     autoRefreshMode: onceMode ? 'none' : autoRefreshMode, // Disable auto-refresh in one-shot mode
@@ -311,21 +330,28 @@ async function main() {
       await watcher.start();
       const status = watcher.getStatus();
       if (jsonMode) {
-        // Restore console.log for JSON output
-        console.log = originalLog;
-        console.log(JSON.stringify(status, null, 2));
-        // Suppress console.log again before stopping to avoid extra shutdown output
-        console.log = () => {};
+        const output = `${JSON.stringify(status, null, 2)}\n`;
+        writeAndExit(process.stdout, output, 0, {
+          cleanup: cleanupShutdownHandlers,
+          watcher,
+          jsonMode,
+          originalLog,
+        });
+        return;
       } else {
+        let output = '';
         for (const [name, agentStatus] of Object.entries(status)) {
-          originalLog(`\n=== ${name.toUpperCase()} ===`);
-          originalLog(JSON.stringify(agentStatus, null, 2));
+          output += `\n=== ${name.toUpperCase()} ===\n`;
+          output += `${JSON.stringify(agentStatus, null, 2)}\n`;
         }
+        writeAndExit(process.stdout, output, 0, {
+          cleanup: cleanupShutdownHandlers,
+          watcher,
+          jsonMode,
+          originalLog,
+        });
+        return;
       }
-      cleanupShutdownHandlers();
-      watcher.stop();
-      process.exitCode = 0;
-      return;
     }
 
     await watcher.start();

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "agent-tank",
-  "version": "0.9.5",
+  "version": "0.9.6",
   "description": "Monitor and query usage limits for LLM CLI tools (Claude, Gemini, Codex) via HTTP API",
   "main": "src/index.js",
   "bin": {

package/src/agents/base.js

@@ -231,6 +231,10 @@ class BaseAgent {
           clearTimeout(timer);
           spawnDataHandler.dispose();
           spawnExitHandler.dispose();
+          if (!this.shell) {
+            reject(new Error('Process exited during spawn before handlers were attached'));
+            return;
+          }
           this.processReady = true;
           this.output = spawnOutput;
           this._setupPersistentDataHandler();
@@ -243,6 +247,11 @@ class BaseAgent {
           logger.agent(this.name, 'Process ready for commands');
           clearTimeout(timer);
           spawnDataHandler.dispose();
+          spawnExitHandler.dispose();
+          if (!this.shell) {
+            reject(new Error('Process exited during spawn before handlers were attached'));
+            return;
+          }
           this.processReady = true;
           this._setupPersistentDataHandler();
           this._setupPersistentExitHandler();
@@ -334,6 +343,10 @@ class BaseAgent {
   }
 
   _setupPersistentDataHandler() {
+    if (!this.shell) {
+      logger.agent(this.name, 'Skipping persistent data handler setup: shell already exited');
+      return;
+    }
     const handler = this.shell.onData((data) => {
       this.output += data;
 
@@ -356,6 +369,10 @@ class BaseAgent {
   }
 
   _setupPersistentExitHandler() {
+    if (!this.shell) {
+      logger.agent(this.name, 'Skipping persistent exit handler setup: shell already exited');
+      return;
+    }
     const handler = this.shell.onExit(({ exitCode }) => {
       logger.agent(this.name, 'Persistent process exited with code', logger.dim(exitCode));
       this.shell = null; this.processReady = false; this._onDataCallback = null; this._disposables = [];

package/src/agents/gemini-direct.js

@@ -0,0 +1,234 @@
+const fs = require('node:fs');
+const os = require('node:os');
+const path = require('node:path');
+
+const GEMINI_DIR = '.gemini';
+const OAUTH_CLIENT_ID = '681255809395-oo8ft2oprdrnp9' + 'e3aqf6av3hmdib135j.apps.googleusercontent.com';
+const OAUTH_CLIENT_SECRET = 'GOCSPX-4uHgMPm-1o7Sk-' + 'geV6Cu5clXFsxl';
+const TOKEN_ENDPOINT = 'https://oauth2.googleapis.com/token';
+const CODE_ASSIST_ENDPOINT = 'https://cloudcode-pa.googleapis.com/v1internal';
+const REFRESH_SKEW_MS = 60 * 1000;
+const VALID_GEMINI_MODES = new Set(['pty', 'direct', 'fallback']);
+
+function getGeminiHomeDir() {
+  return path.join(os.homedir(), GEMINI_DIR);
+}
+
+function getOAuthCredsPath() {
+  return path.join(getGeminiHomeDir(), 'oauth_creds.json');
+}
+
+function getGoogleAccountsPath() {
+  return path.join(getGeminiHomeDir(), 'google_accounts.json');
+}
+
+function normalizeGeminiMode(mode) {
+  const normalized = typeof mode === 'string' ? mode.toLowerCase() : 'fallback';
+  return VALID_GEMINI_MODES.has(normalized) ? normalized : 'fallback';
+}
+
+function readJsonIfExists(filePath) {
+  try {
+    return JSON.parse(fs.readFileSync(filePath, 'utf8'));
+  } catch (err) {
+    if (err && err.code === 'ENOENT') {
+      return null;
+    }
+    throw err;
+  }
+}
+
+function formatDurationFromSeconds(totalSeconds) {
+  if (!Number.isFinite(totalSeconds) || totalSeconds <= 0) {
+    return '0m';
+  }
+
+  const days = Math.floor(totalSeconds / 86400);
+  const hours = Math.floor((totalSeconds % 86400) / 3600);
+  const minutes = Math.floor((totalSeconds % 3600) / 60);
+
+  if (days > 0) {
+    return hours > 0 ? `${days}d ${hours}h` : `${days}d`;
+  }
+
+  if (hours > 0) {
+    return minutes > 0 ? `${hours}h ${minutes}m` : `${hours}h`;
+  }
+
+  return `${Math.max(1, minutes)}m`;
+}
+
+function parseResetTime(resetTime) {
+  if (!resetTime) {
+    return { resetsIn: null, resetsInSeconds: null };
+  }
+
+  const resetAtMs = Date.parse(resetTime);
+  if (!Number.isFinite(resetAtMs)) {
+    return { resetsIn: null, resetsInSeconds: null };
+  }
+
+  const resetsInSeconds = Math.max(0, Math.round((resetAtMs - Date.now()) / 1000));
+  return {
+    resetsIn: formatDurationFromSeconds(resetsInSeconds),
+    resetsInSeconds,
+  };
+}
+
+function buildModelEntry(bucket) {
+  if (!bucket || !bucket.modelId || bucket.remainingFraction == null) {
+    return null;
+  }
+
+  const usageLeft = Number((bucket.remainingFraction * 100).toFixed(1));
+  const percentUsed = Number((100 - usageLeft).toFixed(1));
+  const { resetsIn, resetsInSeconds } = parseResetTime(bucket.resetTime);
+
+  return {
+    model: bucket.modelId,
+    usageLeft,
+    percentUsed,
+    resetsIn,
+    resetsInSeconds,
+  };
+}
+
+function buildUsageFromQuota(quotaResponse) {
+  const seenModels = new Set();
+  const models = [];
+
+  for (const bucket of quotaResponse?.buckets || []) {
+    const entry = buildModelEntry(bucket);
+    if (!entry || seenModels.has(entry.model)) {
+      continue;
+    }
+    seenModels.add(entry.model);
+    models.push(entry);
+  }
+
+  return { models };
+}
+
+class GeminiDirectQuotaClient {
+  constructor(options = {}) {
+    this.fetchImpl = options.fetchImpl || global.fetch;
+    this.now = options.now || Date.now;
+  }
+
+  _readOAuthCredentials() {
+    const creds = readJsonIfExists(getOAuthCredsPath());
+    if (!creds) {
+      throw new Error(`Gemini OAuth credentials not found at ${getOAuthCredsPath()}`);
+    }
+    return creds;
+  }
+
+  _readActiveAccountEmail() {
+    const accounts = readJsonIfExists(getGoogleAccountsPath());
+    return accounts?.active || null;
+  }
+
+  async _getAccessToken() {
+    const creds = this._readOAuthCredentials();
+    const expiryDate = Number(creds.expiry_date);
+    const accessToken = creds.access_token;
+
+    if (accessToken && Number.isFinite(expiryDate) && expiryDate - REFRESH_SKEW_MS > this.now()) {
+      return accessToken;
+    }
+
+    if (!creds.refresh_token) {
+      throw new Error('Gemini OAuth credentials are missing a refresh token');
+    }
+
+    const body = new URLSearchParams({
+      client_id: OAUTH_CLIENT_ID,
+      client_secret: OAUTH_CLIENT_SECRET,
+      refresh_token: creds.refresh_token,
+      grant_type: 'refresh_token',
+    });
+
+    const response = await this.fetchImpl(TOKEN_ENDPOINT, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
+      body,
+    });
+
+    if (!response.ok) {
+      const text = await response.text();
+      throw new Error(`Failed to refresh Gemini OAuth token (${response.status}): ${text || response.statusText}`);
+    }
+
+    const refreshed = await response.json();
+    const updatedCreds = {
+      ...creds,
+      access_token: refreshed.access_token,
+      expiry_date: refreshed.expires_in ? this.now() + (refreshed.expires_in * 1000) : creds.expiry_date,
+      token_type: refreshed.token_type || creds.token_type,
+      scope: refreshed.scope || creds.scope,
+    };
+
+    fs.writeFileSync(getOAuthCredsPath(), JSON.stringify(updatedCreds, null, 2));
+    return updatedCreds.access_token;
+  }
+
+  async _post(method, token, payload) {
+    const response = await this.fetchImpl(`${CODE_ASSIST_ENDPOINT}:${method}`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${token}`,
+      },
+      body: JSON.stringify(payload),
+    });
+
+    if (!response.ok) {
+      const text = await response.text();
+      throw new Error(`Gemini ${method} request failed (${response.status}): ${text || response.statusText}`);
+    }
+
+    return response.json();
+  }
+
+  async fetchQuotaUsage() {
+    const token = await this._getAccessToken();
+    const configuredProject = process.env.GOOGLE_CLOUD_PROJECT || process.env.GOOGLE_CLOUD_PROJECT_ID;
+
+    const loadResponse = await this._post('loadCodeAssist', token, {
+      cloudaicompanionProject: configuredProject,
+      metadata: {
+        ideType: 'IDE_UNSPECIFIED',
+        platform: 'PLATFORM_UNSPECIFIED',
+        pluginType: 'GEMINI',
+        duetProject: configuredProject,
+      },
+      mode: 'HEALTH_CHECK',
+    });
+
+    const project = loadResponse?.cloudaicompanionProject || configuredProject;
+    if (!project) {
+      throw new Error('Gemini Code Assist project ID is unavailable');
+    }
+
+    const quotaResponse = await this._post('retrieveUserQuota', token, { project });
+    const usage = buildUsageFromQuota(quotaResponse);
+
+    return {
+      usage,
+      metadata: {
+        email: this._readActiveAccountEmail(),
+        authMethod: 'OAuth',
+        tier: loadResponse?.paidTier?.name || loadResponse?.currentTier?.name,
+        project,
+      },
+    };
+  }
+}
+
+module.exports = {
+  GeminiDirectQuotaClient,
+  buildModelEntry,
+  buildUsageFromQuota,
+  formatDurationFromSeconds,
+  normalizeGeminiMode,
+};

package/src/agents/gemini.js

@@ -1,14 +1,60 @@
 const { BaseAgent } = require('./base.js');
 const { calculatePace } = require('../pace-evaluator.js');
 const { CYCLE_DURATIONS } = require('../usage-formatters.js');
+const { GeminiDirectQuotaClient, normalizeGeminiMode } = require('./gemini-direct.js');
 const logger = require('../logger.js');
 
 class GeminiAgent extends BaseAgent {
-  constructor() {
-    super('gemini', 'gemini');
+  constructor(options = {}) {
+    super('gemini', 'gemini', ['--skip-trust']);
+    this.mode = normalizeGeminiMode(options.mode);
+    this.directClient = new GeminiDirectQuotaClient(options.directClientOptions);
     this._aboutSent = false;
   }
 
+  async refresh() {
+    if (this.mode === 'pty') {
+      return super.refresh();
+    }
+
+    if (this.isStopping()) {
+      logger.agent(this.name, 'Skipping refresh while stopping');
+      return;
+    }
+
+    if (this.isRefreshing) {
+      logger.agent(this.name, 'Already refreshing, skipping...');
+      return;
+    }
+
+    logger.agent(this.name, `Starting refresh${this.mode === 'direct' ? ' (direct quota mode)' : ' (direct quota mode with PTY fallback)'}...`);
+    this.isRefreshing = true;
+    this.error = null;
+    this.auth = null;
+
+    try {
+      const result = await this.directClient.fetchQuotaUsage();
+      this.usage = result.usage;
+      this.metadata = result.metadata || this.metadata;
+      this._metadataFetched = !!this.metadata;
+      this.lastUpdated = new Date().toISOString();
+      logger.agent(this.name, 'Parsed direct quota usage:', logger.json(this.usage));
+    } catch (err) {
+      if (this.mode === 'fallback') {
+        logger.agent(this.name, `Direct quota fetch failed, falling back to PTY: ${err.message}`);
+        this.isRefreshing = false;
+        return super.refresh();
+      }
+      logger.error(`[${this.name}] Error during direct quota refresh:`, err.message);
+      this.error = err.message;
+    } finally {
+      if (this.isRefreshing) {
+        this.isRefreshing = false;
+        logger.agent(this.name, 'Refresh complete');
+      }
+    }
+  }
+
   getTimeout() {
     return 35000; // 35 seconds - Gemini can take 10+ seconds to start
   }
@@ -388,6 +434,11 @@ class GeminiAgent extends BaseAgent {
    * @returns {Promise<boolean>} True if keepalive succeeded
    */
   async keepalive() {
+    if (this.mode === 'direct' || (this.mode === 'fallback' && (!this.shell || !this.processReady))) {
+      console.log(`[${this.name}] Keepalive skipped (${this.mode} mode)`);
+      return true;
+    }
+
     if (this.freshProcess) {
       console.log(`[${this.name}] Keepalive skipped (fresh process mode)`);
       return true;

package/src/index.js

@@ -121,6 +121,7 @@ class AgentTank {
     this.requestedAgents = options.agents || null;
     this.freshProcess = options.freshProcess || false;
     this.claudeApi = options.claudeApi || false; // Use direct Anthropic API for Claude
+    this.geminiMode = options.geminiMode || 'fallback';
     this.auth = options.auth || {};
     this.agents = new Map();
     this.server = null;
@@ -327,7 +328,7 @@ class AgentTank {
       case 'claude':
         return new ClaudeAgent({ useApi: this.claudeApi });
       case 'gemini':
-        return new GeminiAgent();
+        return new GeminiAgent({ mode: this.geminiMode });
       case 'codex':
         return new CodexAgent();
       default: