1  /**
  2   * Agent-Claude API Tests
  3   *
  4   * Tests for centralized Claude API interface with budget tracking.
  5   * Uses pg-mock (in-memory SQLite via db.js mock) for PG compatibility.
  6   */
  7  
  8  import { describe, it, before, beforeEach, mock } from 'node:test';
  9  import assert from 'node:assert';
 10  import Database from 'better-sqlite3';
 11  import { createPgMock } from '../helpers/pg-mock.js';
 12  
 13  // ─── Create in-memory test DB ─────────────────────────────────────────────────
 14  
 15  const db = new Database(':memory:');
 16  
 17  db.exec(`
 18    CREATE TABLE agent_tasks (
 19      id INTEGER PRIMARY KEY AUTOINCREMENT,
 20      task_type TEXT NOT NULL,
 21      assigned_to TEXT NOT NULL,
 22      status TEXT DEFAULT 'pending',
 23      priority INTEGER DEFAULT 5,
 24      context_json TEXT,
 25      result_json TEXT,
 26      parent_task_id INTEGER,
 27      created_at DATETIME DEFAULT CURRENT_TIMESTAMP
 28    );
 29    CREATE TABLE agent_llm_usage (
 30      id INTEGER PRIMARY KEY AUTOINCREMENT,
 31      agent_name TEXT NOT NULL CHECK(agent_name IN (
 32        'developer', 'qa', 'architect', 'monitor', 'security', 'runner', 'triage'
 33      )),
 34      task_id INTEGER REFERENCES agent_tasks(id),
 35      model TEXT NOT NULL,
 36      prompt_tokens INTEGER NOT NULL,
 37      completion_tokens INTEGER NOT NULL,
 38      cost_usd REAL NOT NULL,
 39      created_at DATETIME DEFAULT CURRENT_TIMESTAMP
 40    );
 41    CREATE INDEX IF NOT EXISTS idx_agent_llm_usage_agent ON agent_llm_usage(agent_name);
 42    CREATE INDEX IF NOT EXISTS idx_agent_llm_usage_created ON agent_llm_usage(created_at);
 43    CREATE INDEX IF NOT EXISTS idx_agent_llm_usage_agent_date ON agent_llm_usage(agent_name, DATE(created_at));
 44  `);
 45  
 46  // ─── Mock db.js BEFORE importing agent-claude-api.js ─────────────────────────
 47  
 48  mock.module('../../src/utils/db.js', {
 49    namedExports: createPgMock(db),
 50  });
 51  
 52  mock.module('../../src/utils/logger.js', {
 53    defaultExport: class {
 54      info() {}
 55      warn() {}
 56      error() {}
 57      success() {}
 58      debug() {}
 59    },
 60  });
 61  
 62  mock.module('../../src/utils/llm-provider.js', {
 63    namedExports: {
 64      callLLM: async () => ({ content: '' }),
 65      getProvider: () => 'openrouter',
 66      getProviderDisplayName: () => 'OpenRouter',
 67    },
 68  });
 69  
 70  // Import AFTER mock.module
 71  const {
 72    selectModel,
 73    getTodaySpending,
 74    getHourlySpending,
 75    getUsageStats,
 76    analyzeCode,
 77  } = await import('../../src/agents/utils/agent-claude-api.js');
 78  
 79  // ─── Helpers ─────────────────────────────────────────────────────────────────
 80  
 81  function clearUsage() {
 82    db.prepare('DELETE FROM agent_llm_usage').run();
 83    db.prepare('DELETE FROM agent_tasks').run();
 84  }
 85  
 86  function insertUsage(agentName, costUsd, createdAt = null) {
 87    if (createdAt) {
 88      db.prepare(
 89        `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd, created_at)
 90         VALUES (?, NULL, 'claude-3-5-sonnet-20241022', 100, 50, ?, ?)`
 91      ).run(agentName, costUsd, createdAt);
 92    } else {
 93      db.prepare(
 94        `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd)
 95         VALUES (?, NULL, 'claude-3-5-sonnet-20241022', 100, 50, ?)`
 96      ).run(agentName, costUsd);
 97    }
 98  }
 99  
100  // ─── Tests ───────────────────────────────────────────────────────────────────
101  
102  describe('Agent-Claude API Module', () => {
103    before(() => {
104      process.env.ANTHROPIC_API_KEY = 'test-api-key-12345';
105      process.env.AGENT_DAILY_BUDGET = '10';
106    });
107  
108    describe('Database integration', () => {
109      beforeEach(() => clearUsage());
110  
111      it('should track daily spending', async () => {
112        // Initially zero
113        const initialSpending = await getTodaySpending();
114        assert.strictEqual(initialSpending, 0);
115  
116        // Add some usage
117        insertUsage('developer', 0.001);
118  
119        // Should now show spending
120        const currentSpending = await getTodaySpending();
121        assert.strictEqual(currentSpending, 0.001);
122      });
123  
124      it('should track hourly spending', async () => {
125        // Use SQLite's built-in datetime functions to insert at correct times
126        // so the comparison with datetime('now', '-1 hour') works correctly
127        db.prepare(
128          `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd, created_at)
129           VALUES ('developer', NULL, 'claude-3-5-sonnet-20241022', 100, 50, 0.002, datetime('now', '-30 minutes'))`
130        ).run();
131  
132        db.prepare(
133          `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd, created_at)
134           VALUES ('developer', NULL, 'claude-3-5-sonnet-20241022', 100, 50, 0.005, datetime('now', '-2 hours'))`
135        ).run();
136  
137        const hourlySpending = await getHourlySpending();
138        // 30-min-ago row should be counted (0.002), 2-hour-ago should not
139        assert.ok(hourlySpending >= 0.002, `Expected >= 0.002, got ${hourlySpending}`);
140        assert.ok(hourlySpending < 0.005, `Expected < 0.005, got ${hourlySpending}`);
141      });
142  
143      it('should return usage stats for all agents', async () => {
144        // Add usage for multiple agents
145        insertUsage('developer', 0.001);
146        insertUsage('developer', 0.002);
147        insertUsage('qa', 0.0015);
148  
149        const stats = await getUsageStats();
150  
151        assert.ok(Array.isArray(stats.agents));
152        assert.strictEqual(stats.agents.length, 2); // developer and qa
153        assert.ok(stats.totalCost > 0);
154        assert.strictEqual(stats.days, 7);
155  
156        // Check developer stats
157        const devStats = stats.agents.find(a => a.agent_name === 'developer');
158        assert.ok(devStats);
159        assert.strictEqual(Number(devStats.call_count), 2);
160        assert.ok(Math.abs(Number(devStats.total_cost_usd) - 0.003) < 0.000001);
161      });
162  
163      it('should filter stats by agent name', async () => {
164        insertUsage('developer', 0.001);
165        insertUsage('qa', 0.0015);
166  
167        const stats = await getUsageStats('developer');
168  
169        assert.ok(Array.isArray(stats.agents));
170        assert.strictEqual(stats.agents.length, 1);
171        assert.strictEqual(stats.agents[0].agent_name, 'developer');
172      });
173  
174      it('should support custom time range', async () => {
175        // Add recent usage (today)
176        insertUsage('developer', 0.001);
177  
178        // Add old usage outside range
179        const tenDaysAgo = new Date(Date.now() - 10 * 24 * 60 * 60 * 1000).toISOString();
180        insertUsage('developer', 0.001, tenDaysAgo);
181  
182        const stats = await getUsageStats(null, 1); // Last 1 day
183  
184        assert.strictEqual(stats.days, 1);
185        // Only today's record should be counted
186        if (stats.agents.length > 0) {
187          assert.strictEqual(Number(stats.agents[0].call_count), 1);
188        }
189      });
190  
191      it('should maintain referential integrity with agent_tasks', async () => {
192        // Create task
193        const taskResult = db.prepare(
194          `INSERT INTO agent_tasks (task_type, assigned_to) VALUES ('test_task', 'developer')`
195        ).run();
196        const taskId = taskResult.lastInsertRowid;
197  
198        // Add usage for that task
199        db.prepare(
200          `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd)
201           VALUES ('developer', ?, 'claude-3-5-sonnet-20241022', 100, 50, 0.001)`
202        ).run(taskId);
203  
204        const usage = db.prepare('SELECT * FROM agent_llm_usage WHERE task_id = ?').get(taskId);
205  
206        assert.strictEqual(usage.task_id, taskId);
207        assert.strictEqual(usage.agent_name, 'developer');
208      });
209  
210      it('should support null task_id', async () => {
211        db.prepare(
212          `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd)
213           VALUES ('developer', NULL, 'claude-3-5-sonnet-20241022', 100, 50, 0.001)`
214        ).run();
215  
216        const usage = db.prepare('SELECT * FROM agent_llm_usage WHERE task_id IS NULL').get();
217  
218        assert.strictEqual(usage.task_id, null);
219        assert.strictEqual(usage.agent_name, 'developer');
220      });
221    });
222  
223    describe('Budget enforcement', () => {
224      beforeEach(() => clearUsage());
225  
226      it('should calculate cost correctly for Sonnet', () => {
227        // Cost for 1M input + 1M output tokens
228        // (1M/1M * $3) + (1M/1M * $15) = $18
229        db.prepare(
230          `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd)
231           VALUES ('developer', NULL, 'claude-3-5-sonnet-20241022', 1000000, 1000000, 18.0)`
232        ).run();
233  
234        const usage = db.prepare('SELECT cost_usd FROM agent_llm_usage').get();
235  
236        assert.strictEqual(usage.cost_usd, 18.0);
237      });
238  
239      it('should track fractional costs accurately', () => {
240        // Cost = (1234/1M * $3) + (567/1M * $15)
241        const expectedCost = (1234 / 1_000_000) * 3 + (567 / 1_000_000) * 15;
242  
243        db.prepare(
244          `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd)
245           VALUES ('developer', NULL, 'claude-3-5-sonnet-20241022', 1234, 567, ?)`
246        ).run(expectedCost);
247  
248        const usage = db.prepare('SELECT cost_usd FROM agent_llm_usage').get();
249  
250        assert.ok(Math.abs(usage.cost_usd - expectedCost) < 0.000001);
251      });
252    });
253  
254    describe('Schema validation', () => {
255      beforeEach(() => clearUsage());
256  
257      it('should enforce agent_name CHECK constraint', () => {
258        // Valid agent name
259        db.prepare(
260          `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd)
261           VALUES ('developer', NULL, 'claude-3-5-sonnet-20241022', 100, 50, 0.001)`
262        ).run();
263  
264        // Invalid agent name should fail
265        assert.throws(
266          () => {
267            db.prepare(
268              `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd)
269               VALUES ('invalid_agent', NULL, 'claude-3-5-sonnet-20241022', 100, 50, 0.001)`
270            ).run();
271          },
272          {
273            message: /CHECK constraint failed/,
274          }
275        );
276      });
277  
278      it('should require all non-null fields', () => {
279        // Missing required fields should fail
280        assert.throws(
281          () => {
282            db.prepare(
283              `INSERT INTO agent_llm_usage (agent_name) VALUES ('developer')`
284            ).run();
285          },
286          /NOT NULL constraint failed/
287        );
288      });
289  
290      it('should have proper indexes', () => {
291        const indexes = db.prepare("SELECT name FROM sqlite_master WHERE type='index'").all();
292        const indexNames = indexes.map(idx => idx.name);
293        assert.ok(indexNames.includes('idx_agent_llm_usage_agent'));
294        assert.ok(indexNames.includes('idx_agent_llm_usage_created'));
295      });
296    });
297  
298    describe('API integration (mocked)', () => {
299      it('should validate OPENROUTER_API_KEY is required', async () => {
300        const savedKey = process.env.OPENROUTER_API_KEY;
301        delete process.env.OPENROUTER_API_KEY;
302        delete process.env.ANTHROPIC_API_KEY;
303  
304        try {
305          await analyzeCode('developer', 1, '/src/test.js', 'Test');
306          assert.fail('Should have thrown error for missing API key');
307        } catch (error) {
308          assert.match(error.message, /No LLM API key configured|OPENROUTER_API_KEY/);
309        } finally {
310          if (savedKey) process.env.OPENROUTER_API_KEY = savedKey;
311          process.env.ANTHROPIC_API_KEY = 'test-api-key-12345';
312        }
313      });
314    });
315  
316    describe('Performance', () => {
317      beforeEach(() => clearUsage());
318  
319      it('should handle large batch inserts efficiently', () => {
320        const insertStmt = db.prepare(
321          `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd)
322           VALUES (?, ?, 'claude-3-5-sonnet-20241022', 100, 50, 0.001)`
323        );
324  
325        const startTime = Date.now();
326  
327        const insertMany = db.transaction(records => {
328          for (const record of records) {
329            insertStmt.run(record.agent, record.taskId);
330          }
331        });
332  
333        const records = Array.from({ length: 1000 }, () => ({
334          agent: 'developer',
335          taskId: null,
336        }));
337  
338        insertMany(records);
339  
340        const duration = Date.now() - startTime;
341        const count = db.prepare('SELECT COUNT(*) as count FROM agent_llm_usage').get();
342  
343        assert.strictEqual(count.count, 1000);
344        assert.ok(duration < 1000, `Batch insert took ${duration}ms (should be <1000ms)`);
345      });
346  
347      it('should query by date efficiently with index', () => {
348        // Insert records across multiple days
349        for (let i = 0; i < 100; i++) {
350          const date = new Date(Date.now() - i * 24 * 60 * 60 * 1000).toISOString();
351          db.prepare(
352            `INSERT INTO agent_llm_usage (agent_name, task_id, model, prompt_tokens, completion_tokens, cost_usd, created_at)
353             VALUES ('developer', NULL, 'claude-3-5-sonnet-20241022', 100, 50, 0.001, ?)`
354          ).run(date);
355        }
356  
357        const startTime = Date.now();
358  
359        // Query recent records (should use index)
360        const recent = db.prepare(
361          `SELECT * FROM agent_llm_usage WHERE created_at >= datetime('now', '-7 days')`
362        ).all();
363  
364        const duration = Date.now() - startTime;
365  
366        assert.ok(recent.length <= 8); // 7 days + today
367        assert.ok(duration < 50, `Query took ${duration}ms (should be <50ms with index)`);
368      });
369    });
370  
371    describe('selectModel', () => {
372      it('returns haiku for simple task types', () => {
373        assert.match(selectModel('classify_issue'), /haiku/);
374        assert.match(selectModel('scan_logs'), /haiku/);
375        assert.match(selectModel('scan_secrets'), /haiku/);
376      });
377  
378      it('returns sonnet for complex task types', () => {
379        assert.match(selectModel('generateCode'), /sonnet/);
380        assert.match(selectModel('reviewArchitecture'), /sonnet/);
381        assert.match(selectModel('threat_model'), /sonnet/);
382      });
383  
384      it('returns sonnet for unknown task types', () => {
385        assert.match(selectModel('unknown_task_xyz'), /sonnet/);
386      });
387  
388      it('explicit simple complexity override returns haiku', () => {
389        assert.match(selectModel('generateCode', 'simple'), /haiku/);
390      });
391  
392      it('explicit complex complexity override returns sonnet', () => {
393        assert.match(selectModel('classify_issue', 'complex'), /sonnet/);
394      });
395  
396      it('AGENT_USE_HAIKU_FOR_SIMPLE_TASKS=false always returns sonnet', () => {
397        process.env.AGENT_USE_HAIKU_FOR_SIMPLE_TASKS = 'false';
398        assert.match(selectModel('classify_issue'), /sonnet/);
399        delete process.env.AGENT_USE_HAIKU_FOR_SIMPLE_TASKS;
400      });
401    });
402  });