@agentguard-run/spend 0.2.2 → 0.3.0

package/src/bindings/bedrock.ts

@@ -0,0 +1,200 @@
+/** Native AWS Bedrock binding for AgentGuard Spend. */
+
+import type { CapabilityTier, CallContext, Provider, SpendPolicy, SpendScope } from '../types';
+import {
+  AgentGuardBlockedError,
+  SpendGuard,
+  type SpendGuardConfig,
+  extractText,
+  usageCountsFromObject,
+  wrapUsageStream,
+  type UsageCounts,
+} from '../spend-guard';
+
+export interface BedrockBindingOptions {
+  policy: SpendPolicy;
+  scope: SpendScope;
+  capabilityClaim?: CapabilityTier;
+  config?: Omit<SpendGuardConfig, 'policy'>;
+}
+
+type BedrockCommand = {
+  input?: Record<string, unknown>;
+  constructor?: { name?: string };
+};
+
+export function withSpendGuardBedrock(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  client: any,
+  opts: BedrockBindingOptions,
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+): any {
+  const guard = new SpendGuard({ policy: opts.policy, ...(opts.config ?? {}) });
+  const originalSend = client?.send?.bind(client);
+  if (!originalSend) throw new Error('withSpendGuardBedrock: expected client.send');
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  client.send = async (command: BedrockCommand, ...args: any[]) => {
+    if (!isInvokeModelCommand(command)) return originalSend(command, ...args);
+
+    const call = buildBedrockCall(command, guard, opts);
+    const { decision } = await guard.decide(call);
+    if (decision.action === 'block') throw new AgentGuardBlockedError(decision, opts.scope, opts.config?.locale);
+
+    if (decision.action === 'downgrade' && decision.modelResolved !== decision.modelRequested && command.input) {
+      command.input = { ...command.input, modelId: decision.modelResolved };
+    }
+
+    const response = await originalSend(command, ...args);
+    if (isStreamingCommand(command) && response?.body) {
+      return {
+        ...response,
+        body: wrapBedrockStream(response.body, guard, decision.decisionId, call),
+      };
+    }
+
+    const usage = extractBedrockResponseUsage(response, call.model);
+    if (usage) await guard.settleStreamUsage(decision.decisionId, usage.inputTokens, usage.outputTokens);
+    return response;
+  };
+
+  client.__agentguard = guard;
+  return client;
+}
+
+function isInvokeModelCommand(command: BedrockCommand): boolean {
+  const name = command?.constructor?.name;
+  return name === 'InvokeModelCommand' || name === 'InvokeModelWithResponseStreamCommand';
+}
+
+function isStreamingCommand(command: BedrockCommand): boolean {
+  return command?.constructor?.name === 'InvokeModelWithResponseStreamCommand';
+}
+
+function buildBedrockCall(command: BedrockCommand, guard: SpendGuard, opts: BedrockBindingOptions): CallContext {
+  const input = command.input ?? {};
+  const modelId = String(input.modelId ?? 'unknown');
+  const body = parseJsonBody(input.body);
+  const inputText = extractText(body);
+  return {
+    provider: 'bedrock' as Provider,
+    model: modelId,
+    inputTokens: guard.estimateTokens(inputText),
+    outputTokens: outputTokensFromBody(body),
+    scope: opts.scope,
+    capabilityClaim: opts.capabilityClaim,
+  };
+}
+
+function wrapBedrockStream(
+  stream: AsyncIterable<unknown>,
+  guard: SpendGuard,
+  decisionId: string,
+  call: CallContext,
+): AsyncIterable<unknown> {
+  const tracker = new BedrockStreamUsageTracker(guard, call);
+  return wrapUsageStream(stream, {
+    onChunk: (event) => tracker.observe(event),
+    settle: (partial, reason) => {
+      const usage = tracker.usage();
+      return guard.settleStreamUsage(decisionId, usage.inputTokens, usage.outputTokens, { partial, reason });
+    },
+  });
+}
+
+class BedrockStreamUsageTracker {
+  private input: number | null = null;
+  private output: number | null = null;
+  private content = '';
+
+  constructor(private guard: SpendGuard, private call: CallContext) {}
+
+  observe(event: unknown): void {
+    const payload = bedrockEventPayload(event);
+    if (!payload) return;
+    const usage = usageFromBedrockPayload(payload);
+    if (usage) {
+      this.input = usage.inputTokens;
+      this.output = usage.outputTokens;
+    }
+    this.content += textFromBedrockPayload(payload);
+  }
+
+  usage(): UsageCounts {
+    return {
+      inputTokens: this.input ?? this.call.inputTokens,
+      outputTokens: this.output ?? this.guard.estimateTokens(this.content),
+    };
+  }
+}
+
+function extractBedrockResponseUsage(response: unknown, modelId: string): UsageCounts | null {
+  const body = (response as { body?: unknown } | null)?.body;
+  const parsed = parseJsonBody(body);
+  const usage = usageFromBedrockPayload(parsed);
+  if (usage) return usage;
+  if (modelId.startsWith('amazon.nova.')) return usageCountsFromObject((parsed as Record<string, unknown> | null)?.usage);
+  return null;
+}
+
+function usageFromBedrockPayload(payload: unknown): UsageCounts | null {
+  if (!payload || typeof payload !== 'object') return null;
+  const obj = payload as Record<string, unknown>;
+  const direct = usageCountsFromObject(obj.usage);
+  if (direct) return direct;
+
+  const bedrockMetrics = obj['amazon-bedrock-invocationMetrics'];
+  if (bedrockMetrics && typeof bedrockMetrics === 'object') {
+    const metrics = bedrockMetrics as Record<string, unknown>;
+    const input = metrics.inputTokenCount;
+    const output = metrics.outputTokenCount;
+    if (Number.isSafeInteger(input) && Number.isSafeInteger(output)) {
+      return { inputTokens: input as number, outputTokens: output as number };
+    }
+  }
+
+  const nested = usageCountsFromObject((obj.message as Record<string, unknown> | undefined)?.usage);
+  if (nested) return nested;
+  return null;
+}
+
+function textFromBedrockPayload(payload: unknown): string {
+  if (!payload || typeof payload !== 'object') return '';
+  const obj = payload as Record<string, unknown>;
+  const delta = obj.delta as Record<string, unknown> | undefined;
+  if (typeof delta?.text === 'string') return delta.text;
+  const contentBlockDelta = obj.contentBlockDelta as Record<string, unknown> | undefined;
+  const nestedDelta = contentBlockDelta?.delta as Record<string, unknown> | undefined;
+  if (typeof nestedDelta?.text === 'string') return nestedDelta.text;
+  const output = obj.output as Record<string, unknown> | undefined;
+  if (output) return extractText(output);
+  return '';
+}
+
+function outputTokensFromBody(body: unknown): number {
+  if (!body || typeof body !== 'object') return 512;
+  const obj = body as Record<string, unknown>;
+  const value = obj.max_tokens ?? obj.maxTokens ?? obj.max_new_tokens ??
+    (obj.inferenceConfig as Record<string, unknown> | undefined)?.maxTokens;
+  return Number.isSafeInteger(value) && (value as number) >= 0 ? value as number : 512;
+}
+
+function parseJsonBody(body: unknown): unknown {
+  if (!body) return {};
+  try {
+    if (typeof body === 'string') return JSON.parse(body);
+    if (body instanceof Uint8Array) return JSON.parse(new TextDecoder().decode(body));
+    if (typeof Buffer !== 'undefined' && Buffer.isBuffer(body)) return JSON.parse(body.toString('utf8'));
+    if (typeof body === 'object') return body;
+  } catch {
+    return {};
+  }
+  return {};
+}
+
+function bedrockEventPayload(event: unknown): unknown {
+  const obj = event as Record<string, unknown> | null;
+  const chunk = obj?.chunk as Record<string, unknown> | undefined;
+  const bytes = chunk?.bytes ?? chunk?.payload ?? obj?.payload;
+  return parseJsonBody(bytes);
+}

package/src/cli/auth.ts

@@ -0,0 +1,145 @@
+import * as crypto from 'crypto';
+import * as fs from 'fs';
+import * as http from 'http';
+import * as https from 'https';
+import * as path from 'path';
+import * as readline from 'readline';
+import { AGENTGUARD_SPEND_VERSION } from '../index';
+import { agentguardHome, banner, dim, green, statusBar, yellow } from './colors';
+
+const OPENROUTER_AUTH_URL = 'https://openrouter.ai/api/v1/auth/key';
+
+export function openRouterKeyPath(): string {
+  return path.join(agentguardHome(), 'openrouter-key');
+}
+
+export function keyFingerprint(key: string): string {
+  return crypto.createHash('sha256').update(key).digest('hex').slice(0, 16);
+}
+
+export async function validateOpenRouterKey(key: string, endpoint = process.env.AGENTGUARD_OPENROUTER_AUTH_URL || OPENROUTER_AUTH_URL): Promise<boolean> {
+  if (!key.trim()) return false;
+  if (endpoint.startsWith('mock://')) {
+    const expected = new URL(endpoint).searchParams.get('key');
+    return expected === key.trim();
+  }
+  try {
+    const status = await requestStatus(endpoint, key.trim());
+    return status >= 200 && status < 300;
+  } catch {
+    return false;
+  }
+}
+
+export function readOpenRouterKey(): string | null {
+  if (process.env.OPENROUTER_API_KEY) return process.env.OPENROUTER_API_KEY;
+  try {
+    const value = fs.readFileSync(openRouterKeyPath(), 'utf8').trim();
+    return value.length > 0 ? value : null;
+  } catch {
+    return null;
+  }
+}
+
+export function writeOpenRouterKey(key: string): void {
+  const filePath = openRouterKeyPath();
+  fs.mkdirSync(path.dirname(filePath), { recursive: true });
+  fs.writeFileSync(filePath, key.trim() + '\n', { mode: 0o600 });
+  try { fs.chmodSync(filePath, 0o600); } catch { return; }
+}
+
+export async function runAuth(argv: string[]): Promise<number> {
+  if (argv.includes('--help') || argv.includes('-h')) {
+    console.log('agentguard auth openrouter | status | clear openrouter');
+    return 0;
+  }
+  const sub = argv[0] ?? 'status';
+  if (sub === 'status') return authStatus();
+  if (sub === 'clear' && argv[1] === 'openrouter') return clearOpenRouter();
+  if (sub === 'openrouter') return configureOpenRouter(argv);
+  console.error(`agentguard auth: unknown command '${argv.join(' ')}'`);
+  return 2;
+}
+
+async function configureOpenRouter(argv: string[]): Promise<number> {
+  const flagKey = valueAfter(argv, '--key');
+  const key = flagKey ?? await promptForKey();
+  if (!key) {
+    console.error('No key provided.');
+    return 2;
+  }
+  const ok = await validateOpenRouterKey(key);
+  if (!ok) {
+    console.error('OpenRouter key validation failed.');
+    return 1;
+  }
+  writeOpenRouterKey(key);
+  console.log('');
+  console.log('  ' + banner(AGENTGUARD_SPEND_VERSION));
+  console.log('');
+  console.log(`  ${green('configured')} openrouter  fingerprint ${keyFingerprint(key)}`);
+  console.log(`  ${dim('saved')} ${openRouterKeyPath()}`);
+  console.log('');
+  return 0;
+}
+
+function authStatus(): number {
+  const key = readOpenRouterKey();
+  console.log('');
+  console.log('  ' + banner(AGENTGUARD_SPEND_VERSION));
+  console.log('');
+  if (key) {
+    console.log(`  ${green('openrouter')} configured  fingerprint ${keyFingerprint(key)}`);
+  } else {
+    console.log(`  ${yellow('openrouter')} not configured`);
+  }
+  console.log('');
+  console.log('  ' + statusBar());
+  console.log('');
+  return 0;
+}
+
+function clearOpenRouter(): number {
+  try { fs.unlinkSync(openRouterKeyPath()); } catch { return 0; }
+  console.log('OpenRouter key removed.');
+  return 0;
+}
+
+async function promptForKey(): Promise<string> {
+  const rl = readline.createInterface({ input: process.stdin, output: process.stdout });
+  return new Promise((resolve) => {
+    rl.question('OpenRouter API key: ', (answer) => {
+      rl.close();
+      resolve(answer.trim());
+    });
+  });
+}
+
+function requestStatus(endpoint: string, key: string): Promise<number> {
+  return new Promise((resolve, reject) => {
+    const url = new URL(endpoint);
+    const client = url.protocol === 'http:' ? http : https;
+    const req = client.request(
+      {
+        method: 'GET',
+        hostname: url.hostname,
+        port: url.port,
+        path: url.pathname + url.search,
+        headers: { authorization: `Bearer ${key}`, accept: 'application/json' },
+        timeout: 5000,
+      },
+      (res) => {
+        res.resume();
+        res.on('end', () => resolve(res.statusCode ?? 0));
+      },
+    );
+    req.on('error', reject);
+    req.on('timeout', () => req.destroy(new Error('OpenRouter auth validation timed out')));
+    req.end();
+  });
+}
+
+function valueAfter(argv: string[], flag: string): string | undefined {
+  const index = argv.indexOf(flag);
+  return index >= 0 ? argv[index + 1] : undefined;
+}

package/src/cli/models.ts

@@ -0,0 +1,236 @@
+import * as fs from 'fs';
+import * as path from 'path';
+import * as readline from 'readline';
+import { AGENTGUARD_SPEND_VERSION } from '../index';
+import { fetchCatalog, getCachedCatalog, syncPricingIntoCostTable, type OpenRouterModel } from '../openrouter-catalog';
+import { getModelCost, knownModelCosts } from '../cost-table';
+import { TASK_TEMPLATES, getTaskTemplate } from '../templates';
+import { agentguardHome, banner, cyanBold, dim, green, statusBar, yellow } from './colors';
+
+export interface ModelRow {
+  id: string;
+  name: string;
+  inputCentsPerMtok: number | null;
+  outputCentsPerMtok: number | null;
+  capability: string;
+}
+
+export function rowsFromCatalog(models: OpenRouterModel[]): ModelRow[] {
+  return models.map((model) => {
+    const input = Number(model.pricing?.prompt);
+    const output = Number(model.pricing?.completion);
+    return {
+      id: model.id,
+      name: model.name ?? model.id,
+      inputCentsPerMtok: Number.isFinite(input) ? input * 100 * 1_000_000 : null,
+      outputCentsPerMtok: Number.isFinite(output) ? output * 100 * 1_000_000 : null,
+      capability: capabilityForModel(model.id),
+    };
+  });
+}
+
+export function rowsFromCostTable(): ModelRow[] {
+  return Object.entries(knownModelCosts()).map(([id, cost]) => ({
+    id,
+    name: id,
+    inputCentsPerMtok: cost.inputCentsPerKtok * 1000,
+    outputCentsPerMtok: cost.outputCentsPerKtok * 1000,
+    capability: capabilityForModel(id),
+  }));
+}
+
+export function filterRows(rows: ModelRow[], opts: { search?: string; capability?: string; task?: string }): ModelRow[] {
+  let out = rows;
+  if (opts.task) {
+    const template = getTaskTemplate(opts.task);
+    if (template) {
+      const allowed = new Set(template.allowedModels);
+      out = out.filter((row) => allowed.has(row.id) || row.capability === template.requiredCapability);
+    }
+  }
+  if (opts.capability) out = out.filter((row) => row.capability === opts.capability);
+  if (opts.search) {
+    const needle = opts.search.toLowerCase();
+    out = out.filter((row) => row.id.toLowerCase().includes(needle) || row.name.toLowerCase().includes(needle));
+  }
+  return out.sort((a, b) => displayCost(a) - displayCost(b) || a.id.localeCompare(b.id));
+}
+
+export function suggestedAssignments(task: string): Record<string, string[]> {
+  const template = getTaskTemplate(task);
+  if (!template) return {};
+  return {
+    primary: [template.primaryModel],
+    fallback: [template.fallbackModel],
+    allowed: template.allowedModels,
+  };
+}
+
+export async function runModels(argv: string[]): Promise<number> {
+  if (argv.includes('--help') || argv.includes('-h')) {
+    console.log('agentguard models [--sync-pricing] [--search <term>] [--capability <tier>] [--task <template>] [--json]');
+    return 0;
+  }
+
+  const sync = argv.includes('--sync-pricing');
+  const json = argv.includes('--json');
+  const search = argValue(argv, '--search');
+  const capability = argValue(argv, '--capability');
+  const task = argValue(argv, '--task');
+
+  if (sync) {
+    const result = await syncPricingIntoCostTable({ force: true });
+    if (json) {
+      console.log(JSON.stringify(result, null, 2));
+    } else {
+      console.log('');
+      console.log('  ' + banner(AGENTGUARD_SPEND_VERSION));
+      console.log('');
+      console.log(`  ${green('applied')} ${result.count} OpenRouter pricing overrides`);
+      console.log(`  ${dim('saved')} ${agentguardHome()}/cost-overrides.json`);
+      console.log('');
+    }
+    return 0;
+  }
+
+  const rows = await loadRows();
+  const filtered = filterRows(rows, { search, capability, task });
+  if (json) {
+    console.log(JSON.stringify({ models: filtered, assignments: task ? suggestedAssignments(task) : undefined }, null, 2));
+    return 0;
+  }
+
+  if (process.stdin.isTTY && !search && !capability && !task) {
+    const selected = await interactivePick(filtered);
+    if (selected.length > 0) writeSelectedModels(selected);
+    return 0;
+  }
+
+  printRows(filtered, task);
+  return 0;
+}
+
+async function loadRows(): Promise<ModelRow[]> {
+  const cached = getCachedCatalog();
+  if (cached) return rowsFromCatalog(cached.data);
+  try {
+    const catalog = await fetchCatalog();
+    return rowsFromCatalog(catalog.data);
+  } catch {
+    return rowsFromCostTable();
+  }
+}
+
+function printRows(rows: ModelRow[], task?: string): void {
+  console.log('');
+  console.log('  ' + banner(AGENTGUARD_SPEND_VERSION));
+  console.log('');
+  if (task && TASK_TEMPLATES[task]) {
+    const assignment = suggestedAssignments(task);
+    console.log('  ' + cyanBold(`task: ${task}`));
+    console.log(`  primary: ${assignment.primary?.join(', ') ?? 'n/a'}`);
+    console.log(`  fallback: ${assignment.fallback?.join(', ') ?? 'n/a'}`);
+    console.log('');
+  }
+  for (const row of rows.slice(0, 80)) {
+    console.log(`  ${green(row.id.padEnd(38))} ${formatCost(row).padEnd(22)} ${dim(row.capability)}`);
+  }
+  if (rows.length > 80) console.log(`  ${yellow(String(rows.length - 80) + ' more')}  refine with agentguard models --search <term>`);
+  console.log('');
+  console.log('  ' + statusBar());
+  console.log('');
+}
+
+async function interactivePick(rows: ModelRow[]): Promise<string[]> {
+  const selected = new Set<string>();
+  let cursor = 0;
+  let filtered = rows.slice(0, 80);
+  let query = '';
+  readline.emitKeypressEvents(process.stdin);
+  if (process.stdin.isTTY) process.stdin.setRawMode(true);
+
+  const render = () => {
+    process.stdout.write('\x1b[2J\x1b[H');
+    console.log(banner(AGENTGUARD_SPEND_VERSION));
+    console.log('');
+    console.log(`Search: ${query || dim('press / to filter')}   Selected: ${selected.size}`);
+    console.log('');
+    filtered.slice(0, 18).forEach((row, index) => {
+      const marker = selected.has(row.id) ? '[x]' : '[ ]';
+      const pointer = index === cursor ? '>' : ' ';
+      console.log(`${pointer} ${marker} ${row.id.padEnd(40)} ${formatCost(row).padEnd(20)} ${row.capability}`);
+    });
+    console.log('');
+    console.log(dim('[/] search  [space] select  [enter] confirm  [q] quit'));
+  };
+
+  render();
+  return new Promise((resolve) => {
+    const onKey = (_chunk: string, key: readline.Key) => {
+      if (key.name === 'q' || (key.ctrl && key.name === 'c')) finish([]);
+      else if (key.name === 'return') finish([...selected]);
+      else if (key.name === 'down') { cursor = Math.min(cursor + 1, Math.max(0, filtered.length - 1)); render(); }
+      else if (key.name === 'up') { cursor = Math.max(cursor - 1, 0); render(); }
+      else if (key.name === 'space') {
+        const row = filtered[cursor];
+        if (row) selected.has(row.id) ? selected.delete(row.id) : selected.add(row.id);
+        render();
+      } else if (key.name === 'backspace') {
+        query = query.slice(0, -1);
+        filtered = filterRows(rows, { search: query });
+        cursor = 0;
+        render();
+      } else if (key.sequence && key.sequence.length === 1 && key.sequence !== '/') {
+        query += key.sequence;
+        filtered = filterRows(rows, { search: query });
+        cursor = 0;
+        render();
+      }
+    };
+    const finish = (value: string[]) => {
+      process.stdin.off('keypress', onKey);
+      if (process.stdin.isTTY) process.stdin.setRawMode(false);
+      process.stdout.write('\n');
+      resolve(value);
+    };
+    process.stdin.on('keypress', onKey);
+  });
+}
+
+function writeSelectedModels(models: string[]): void {
+  const filePath = path.join(agentguardHome(), 'policy.yaml');
+  fs.mkdirSync(path.dirname(filePath), { recursive: true });
+  let existing = '';
+  try { existing = fs.readFileSync(filePath, 'utf8'); } catch { existing = defaultPolicy(); }
+  const block = ['allowedModels:', ...models.map((model) => `  - ${model}`)].join('\n') + '\n';
+  const next = existing.match(/^allowedModels:/m)
+    ? existing.replace(/^allowedModels:\n(?:  - .*\n)*/m, block)
+    : existing.trimEnd() + '\n' + block;
+  fs.writeFileSync(filePath, next);
+  console.log(`${green('saved')} ${filePath}`);
+}
+
+function defaultPolicy(): string {
+  return 'id: local-policy-v1\nname: Local AgentGuard policy\nversion: 1\neffectiveFrom: "' + new Date().toISOString() + '"\nmode: enforce\nscope:\n  tenantId: local\ncaps:\n  - amountCents: 100\n    window: per_call\n    action: allow\n';
+}
+
+function formatCost(row: ModelRow): string {
+  if (row.inputCentsPerMtok === null || row.outputCentsPerMtok === null) return '$n/a per 1M';
+  return `$${(row.inputCentsPerMtok / 100).toFixed(2)}/$${(row.outputCentsPerMtok / 100).toFixed(2)} per 1M`;
+}
+
+function displayCost(row: ModelRow): number {
+  return row.inputCentsPerMtok ?? Number.MAX_SAFE_INTEGER;
+}
+
+function capabilityForModel(model: string): string {
+  if (model.includes('gpt-5') || model.includes('opus') || model.includes('sonnet')) return 'payment_initiate';
+  if (model.includes('mini') || model.includes('haiku') || model.includes('flash')) return 'read_only';
+  if (getModelCost(model)) return 'data_write';
+  return 'read_only';
+}
+
+function argValue(argv: string[], flag: string): string | undefined {
+  const index = argv.indexOf(flag);
+  return index >= 0 ? argv[index + 1] : undefined;
+}