From bb5ce098a99e4ea8f36c6a725290d5858c36460f Mon Sep 17 00:00:00 2001 From: Adam Malczewski 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/api/src/agent-manager.ts | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) (limited to 'packages/api/src/agent-manager.ts') diff --git a/packages/api/src/agent-manager.ts b/packages/api/src/agent-manager.ts index 2532efa..d1f81ed 100644 --- a/packages/api/src/agent-manager.ts +++ b/packages/api/src/agent-manager.ts @@ -13,6 +13,7 @@ import { clearSpillForTab, configToRuleset, createConfigWatcher, + createKeyUsageTool, createListFilesTool, createLspTool, createReadFileSliceTool, @@ -84,6 +85,8 @@ const TOOL_DESCRIPTIONS: Record = { search_code: "Search the codebase by query using the 'cs' code search engine (relevance-ranked, structure-aware). Returns the most relevant files first with matching snippets and line numbers. Better than grep/find for exploratory 'where is X / how does Y work' searches; use run_shell with rg for exhaustive exact-match lists.", todo: "Create/maintain a todo list to plan and track work. Declarative whole-list write: send the entire list in `todos` each call (it replaces the previous list). Statuses: pending, in_progress, completed, cancelled.", + key_usage: + "Report current usage levels for configured API keys: provider, active/exhausted status, remaining rate-limit headroom and reset times per window (5-hour, weekly, monthly where available), and whether the figures are live or cached. Pass key_id for one key; omit to report all. Supported for anthropic and opencode-go keys.", summon: "Spawn a child agent to work on a task independently. By default blocks until the child finishes. Set background=true to return immediately with an agent_id for later retrieval.", retrieve: @@ -515,10 +518,11 @@ export class AgentManager { const permReadTab = getSetting("perm_read_tab") === "allow"; const permWebSearch = getSetting("perm_web_search") === "allow"; const permSearchCode = getSetting("perm_search_code") === "allow"; + const permKeyUsage = getSetting("perm_key_usage") === "allow"; const permYoutubeTranscribe = getSetting("perm_youtube_transcribe") === "allow"; const permLsp = getSetting("perm_lsp") === "allow"; const sysPrompt = getSetting("system_prompt") ?? ""; - const permKey = `${permRead}:${permEdit}:${permBash}:${permSummon}:${permUserAgent}:${permSendToTab}:${permReadTab}:${permWebSearch}:${permYoutubeTranscribe}:${permSearchCode}:${permLsp}:${sysPrompt}`; + const permKey = `${permRead}:${permEdit}:${permBash}:${permSummon}:${permUserAgent}:${permSendToTab}:${permReadTab}:${permWebSearch}:${permYoutubeTranscribe}:${permSearchCode}:${permKeyUsage}:${permLsp}:${sysPrompt}`; // If the override differs or permissions changed, invalidate the cached agent if ( @@ -610,6 +614,9 @@ export class AgentManager { if (allowed.has("web_search")) { toolEntries.push({ name: "web_search", tool: createWebSearchTool() }); } + if (allowed.has("key_usage")) { + toolEntries.push({ name: "key_usage", tool: this.buildKeyUsageTool() }); + } if (allowed.has("lsp") && lspServers.length > 0) { toolEntries.push({ name: "lsp", @@ -715,6 +722,9 @@ export class AgentManager { if (permWebSearch) { toolEntries.push({ name: "web_search", tool: createWebSearchTool() }); } + if (permKeyUsage) { + toolEntries.push({ name: "key_usage", tool: this.buildKeyUsageTool() }); + } // The `lsp` tool exposes diagnostics + navigation on demand. It is // gated by `perm_lsp` AND requires at least one server configured // in the working directory's `dispatch.toml`. @@ -1404,6 +1414,19 @@ export class AgentManager { // running→queue / idle→new-turn routing that `POST /chat` uses (see // `deliverMessage`), so an agent message behaves identically to a user one. + /** + * Build the `key_usage` tool, wired to the live model registry (key states) + * and the discovered Claude accounts. The tool fetches usage live with a + * cache fallback (anthropic) or a live scrape (opencode-go), reporting + * remaining headroom, reset times, and data freshness per key. + */ + private buildKeyUsageTool(): ReturnType { + return createKeyUsageTool({ + listKeys: () => this.modelRegistry?.getKeys() ?? [], + listClaudeAccounts: () => this.claudeAccounts, + }); + } + /** * Build the `send_to_tab` + `read_tab` tool entries for `tabId`. Shared by * both tool-construction paths (child whitelist + permission-gated parent). -- cgit v1.2.3