From bb5ce098a99e4ea8f36c6a725290d5858c36460f Mon Sep 17 00:00:00 2001
From: Adam Malczewski <github@tradam.dev>
Date: Tue, 2 Jun 2026 22:50:00 +0900
Subject: feat(tools): add key_usage tool reporting API-key usage levels

Adds an agent-callable `key_usage` tool that reports current usage for
configured API keys so the agent can pick a key with headroom, warn before
hitting a rate limit, and diagnose exhausted-key failures.

Per key it reports: provider, active/exhausted status (with last error +
when it was exhausted), remaining rate-limit headroom and reset timestamp per
window (5-hour, weekly, and monthly where the provider exposes it), and
whether the figures are live or served from cache (with the cache's
last-fetched-from-source time). Supports anthropic and opencode-go keys
(live with cache fallback for anthropic; live scrape for opencode-go).
Optional `key_id` reports one key; omitted reports all.

Hard permission gate `perm_key_usage` (default off): when disabled the tool
is completely removed from the toolset/context. Registered in both the
parent permission-gated path and the child whitelist path, advertised in the
system prompt (TOOL_DESCRIPTIONS), grantable to subagents via the summon
enum, and exposed as a frontend tool-permission checkbox.

To report data freshness, claude.ts gains `getAccountUsageWithSource` +
`ClaudeUsageResult` (live vs cache + cachedAt from usage_cache.cached_at);
the existing `getAccountUsage` now delegates to it, preserving behavior.

Tests: core key-usage tool suite (windows, %-conversion, freshness, exhausted
status, unsupported/unavailable, filtering) + agent-manager perm-gate test.
---
 packages/core/src/credentials/claude.ts | 69 +++++++++++++++++++++++++++------
 packages/core/src/credentials/index.ts  |  2 +
 2 files changed, 60 insertions(+), 11 deletions(-)

(limited to 'packages/core/src/credentials')

diff --git a/packages/core/src/credentials/claude.ts b/packages/core/src/credentials/claude.ts
index 7818222..050a0fc 100644
--- a/packages/core/src/credentials/claude.ts
+++ b/packages/core/src/credentials/claude.ts
@@ -441,6 +441,22 @@ export interface ClaudeUsageReport {
 	orgId?: string;
 }
 
+/**
+ * A usage report paired with provenance: whether it came back from a fresh
+ * live fetch against Anthropic's `/api/oauth/usage` endpoint or was served
+ * from the local `usage_cache` table after a failed/skipped live fetch.
+ *
+ * `source: "cache"` carries `cachedAt` — the epoch-ms timestamp recording when
+ * that cached payload was last fetched FROM the source (the `usage_cache.cached_at`
+ * column). `source: "live"` omits `cachedAt` (the data is current as of now).
+ */
+export interface ClaudeUsageResult {
+	report: ClaudeUsageReport;
+	source: "live" | "cache";
+	/** Epoch-ms the cached report was last fetched from source. Only on `source: "cache"`. */
+	cachedAt?: number;
+}
+
 // ─── Well-known Anthropic models ──────────────────────────────
 
 /**
@@ -602,14 +618,23 @@ async function fetchClaudeUsage(accessToken: string): Promise<ClaudeUsageReport
 	}
 }
 
-function getCachedUsage(keyId: string): ClaudeUsageReport | null {
+/**
+ * Read a cached usage report plus the epoch-ms it was last fetched from source.
+ * Returns `null` when there is no cached row (or on any DB/parse error).
+ */
+function getCachedUsageWithMeta(
+	keyId: string,
+): { report: ClaudeUsageReport; cachedAt: number } | null {
 	try {
 		const db = getDatabase();
 		const row = db
-			.query("SELECT report_json FROM usage_cache WHERE key_id = $keyId")
-			.get({ $keyId: keyId }) as { report_json: string } | null;
+			.query("SELECT report_json, cached_at FROM usage_cache WHERE key_id = $keyId")
+			.get({ $keyId: keyId }) as { report_json: string; cached_at: number } | null;
 		if (!row) return null;
-		return JSON.parse(row.report_json) as ClaudeUsageReport;
+		return {
+			report: JSON.parse(row.report_json) as ClaudeUsageReport,
+			cachedAt: row.cached_at,
+		};
 	} catch {
 		return null;
 	}
@@ -635,13 +660,35 @@ function setCachedUsage(keyId: string, provider: string, report: ClaudeUsageRepo
 	}
 }
 
-export async function getAccountUsage(account: ClaudeAccount): Promise<ClaudeUsageReport | null> {
+/**
+ * Fetch an account's usage report along with its provenance (live vs cache).
+ *
+ * Resolution: refresh credentials and hit the live `/api/oauth/usage` endpoint;
+ * on success the fresh report is cached and returned as `source: "live"`. If
+ * credentials cannot be refreshed OR the live fetch returns nothing, fall back
+ * to the local `usage_cache` row and return it as `source: "cache"` with the
+ * `cachedAt` timestamp recording when that payload was last fetched from source.
+ * Returns `null` only when neither a live report nor a cached row is available.
+ */
+export async function getAccountUsageWithSource(
+	account: ClaudeAccount,
+): Promise<ClaudeUsageResult | null> {
 	const creds = await refreshAccountCredentialsAsync(account);
-	if (!creds) return getCachedUsage(account.id);
-	const report = await fetchClaudeUsage(creds.accessToken);
-	if (report) {
-		setCachedUsage(account.id, "anthropic", report);
-		return report;
+	if (creds) {
+		const report = await fetchClaudeUsage(creds.accessToken);
+		if (report) {
+			setCachedUsage(account.id, "anthropic", report);
+			return { report, source: "live" };
+		}
 	}
-	return getCachedUsage(account.id);
+	const cached = getCachedUsageWithMeta(account.id);
+	if (cached) {
+		return { report: cached.report, source: "cache", cachedAt: cached.cachedAt };
+	}
+	return null;
+}
+
+export async function getAccountUsage(account: ClaudeAccount): Promise<ClaudeUsageReport | null> {
+	const result = await getAccountUsageWithSource(account);
+	return result?.report ?? null;
 }
diff --git a/packages/core/src/credentials/index.ts b/packages/core/src/credentials/index.ts
index 5221dc6..131f035 100644
--- a/packages/core/src/credentials/index.ts
+++ b/packages/core/src/credentials/index.ts
@@ -15,9 +15,11 @@ export {
 	type ClaudeProfile,
 	type ClaudeUsageBucket,
 	type ClaudeUsageReport,
+	type ClaudeUsageResult,
 	discoverClaudeAccounts,
 	fetchAnthropicModels,
 	getAccountUsage,
+	getAccountUsageWithSource,
 	getAnthropicBetas,
 	getAnthropicHeaders,
 	getClaudeAccountsFromDB,
-- 
cgit v1.2.3