From 3e0e608110fc5ab91fbc901b4b77e7c597487a06 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Wed, 14 Jan 2026 05:40:03 +0000 Subject: [PATCH] refactor(infra): split provider usage --- ... server.enables-keepalive-sockets.test.ts} | 0 ...s-invoke-roundtrip-connected-node.test.ts} | 0 ...spects-ackmaxchars-heartbeat-acks.test.ts} | 0 ...beat-runner.returns-default-unset.test.ts} | 0 src/infra/provider-usage.auth.ts | 195 ++++ src/infra/provider-usage.fetch.claude.ts | 196 ++++ src/infra/provider-usage.fetch.codex.ts | 102 ++ src/infra/provider-usage.fetch.copilot.ts | 70 ++ src/infra/provider-usage.fetch.gemini.ts | 81 ++ src/infra/provider-usage.fetch.shared.ts | 14 + src/infra/provider-usage.fetch.ts | 5 + src/infra/provider-usage.fetch.zai.ts | 97 ++ src/infra/provider-usage.format.ts | 91 ++ src/infra/provider-usage.load.ts | 102 ++ src/infra/provider-usage.shared.ts | 61 ++ src/infra/provider-usage.ts | 976 +----------------- src/infra/provider-usage.types.ts | 26 + src/memory/index.ts | 528 +--------- src/memory/manager.ts | 481 +++++++++ src/memory/search-manager.ts | 20 + src/memory/sqlite.ts | 22 + 21 files changed, 1583 insertions(+), 1484 deletions(-) rename src/infra/bridge/{server.part-1.test.ts => server.enables-keepalive-sockets.test.ts} (100%) rename src/infra/bridge/{server.part-2.test.ts => server.supports-invoke-roundtrip-connected-node.test.ts} (100%) rename src/infra/{heartbeat-runner.part-2.test.ts => heartbeat-runner.respects-ackmaxchars-heartbeat-acks.test.ts} (100%) rename src/infra/{heartbeat-runner.part-1.test.ts => heartbeat-runner.returns-default-unset.test.ts} (100%) create mode 100644 src/infra/provider-usage.auth.ts create mode 100644 src/infra/provider-usage.fetch.claude.ts create mode 100644 src/infra/provider-usage.fetch.codex.ts create mode 100644 src/infra/provider-usage.fetch.copilot.ts create mode 100644 src/infra/provider-usage.fetch.gemini.ts create mode 100644 src/infra/provider-usage.fetch.shared.ts create mode 100644 src/infra/provider-usage.fetch.ts create mode 100644 src/infra/provider-usage.fetch.zai.ts create mode 100644 src/infra/provider-usage.format.ts create mode 100644 src/infra/provider-usage.load.ts create mode 100644 src/infra/provider-usage.shared.ts create mode 100644 src/infra/provider-usage.types.ts create mode 100644 src/memory/manager.ts create mode 100644 src/memory/search-manager.ts create mode 100644 src/memory/sqlite.ts diff --git a/src/infra/bridge/server.part-1.test.ts b/src/infra/bridge/server.enables-keepalive-sockets.test.ts similarity index 100% rename from src/infra/bridge/server.part-1.test.ts rename to src/infra/bridge/server.enables-keepalive-sockets.test.ts diff --git a/src/infra/bridge/server.part-2.test.ts b/src/infra/bridge/server.supports-invoke-roundtrip-connected-node.test.ts similarity index 100% rename from src/infra/bridge/server.part-2.test.ts rename to src/infra/bridge/server.supports-invoke-roundtrip-connected-node.test.ts diff --git a/src/infra/heartbeat-runner.part-2.test.ts b/src/infra/heartbeat-runner.respects-ackmaxchars-heartbeat-acks.test.ts similarity index 100% rename from src/infra/heartbeat-runner.part-2.test.ts rename to src/infra/heartbeat-runner.respects-ackmaxchars-heartbeat-acks.test.ts diff --git a/src/infra/heartbeat-runner.part-1.test.ts b/src/infra/heartbeat-runner.returns-default-unset.test.ts similarity index 100% rename from src/infra/heartbeat-runner.part-1.test.ts rename to src/infra/heartbeat-runner.returns-default-unset.test.ts diff --git a/src/infra/provider-usage.auth.ts b/src/infra/provider-usage.auth.ts new file mode 100644 index 000000000..b2dd5afb1 --- /dev/null +++ b/src/infra/provider-usage.auth.ts @@ -0,0 +1,195 @@ +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +import { + CLAUDE_CLI_PROFILE_ID, + ensureAuthProfileStore, + listProfilesForProvider, + resolveApiKeyForProfile, + resolveAuthProfileOrder, +} from "../agents/auth-profiles.js"; +import { + getCustomProviderApiKey, + resolveEnvApiKey, +} from "../agents/model-auth.js"; +import { normalizeProviderId } from "../agents/model-selection.js"; +import { loadConfig } from "../config/config.js"; +import type { UsageProviderId } from "./provider-usage.types.js"; + +export type ProviderAuth = { + provider: UsageProviderId; + token: string; + accountId?: string; +}; + +function parseGoogleToken(apiKey: string): { token: string } | null { + if (!apiKey) return null; + try { + const parsed = JSON.parse(apiKey) as { token?: unknown }; + if (parsed && typeof parsed.token === "string") { + return { token: parsed.token }; + } + } catch { + // ignore + } + return null; +} + +function resolveZaiApiKey(): string | undefined { + const envDirect = + process.env.ZAI_API_KEY?.trim() || process.env.Z_AI_API_KEY?.trim(); + if (envDirect) return envDirect; + + const envResolved = resolveEnvApiKey("zai"); + if (envResolved?.apiKey) return envResolved.apiKey; + + const cfg = loadConfig(); + const key = + getCustomProviderApiKey(cfg, "zai") || getCustomProviderApiKey(cfg, "z-ai"); + if (key) return key; + + const store = ensureAuthProfileStore(); + const apiProfile = [ + ...listProfilesForProvider(store, "zai"), + ...listProfilesForProvider(store, "z-ai"), + ].find((id) => store.profiles[id]?.type === "api_key"); + if (apiProfile) { + const cred = store.profiles[apiProfile]; + if (cred?.type === "api_key" && cred.key?.trim()) { + return cred.key.trim(); + } + } + + try { + const authPath = path.join(os.homedir(), ".pi", "agent", "auth.json"); + if (!fs.existsSync(authPath)) return undefined; + const data = JSON.parse(fs.readFileSync(authPath, "utf-8")) as Record< + string, + { access?: string } + >; + return data["z-ai"]?.access || data.zai?.access; + } catch { + return undefined; + } +} + +async function resolveOAuthToken(params: { + provider: UsageProviderId; + agentDir?: string; +}): Promise { + const cfg = loadConfig(); + const store = ensureAuthProfileStore(params.agentDir, { + allowKeychainPrompt: false, + }); + const order = resolveAuthProfileOrder({ + cfg, + store, + provider: params.provider, + }); + + // Claude CLI creds are the only Anthropic tokens that reliably include the + // `user:profile` scope required for the OAuth usage endpoint. + const candidates = + params.provider === "anthropic" ? [CLAUDE_CLI_PROFILE_ID, ...order] : order; + const deduped: string[] = []; + for (const entry of candidates) { + if (!deduped.includes(entry)) deduped.push(entry); + } + + for (const profileId of deduped) { + const cred = store.profiles[profileId]; + if (!cred || (cred.type !== "oauth" && cred.type !== "token")) continue; + try { + const resolved = await resolveApiKeyForProfile({ + // Usage snapshots should work even if config profile metadata is stale. + // (e.g. config says api_key but the store has a token profile.) + cfg: undefined, + store, + profileId, + agentDir: params.agentDir, + }); + if (!resolved?.apiKey) continue; + let token = resolved.apiKey; + if ( + params.provider === "google-gemini-cli" || + params.provider === "google-antigravity" + ) { + const parsed = parseGoogleToken(resolved.apiKey); + token = parsed?.token ?? resolved.apiKey; + } + return { + provider: params.provider, + token, + accountId: + cred.type === "oauth" && "accountId" in cred + ? (cred as { accountId?: string }).accountId + : undefined, + }; + } catch { + // ignore + } + } + + return null; +} + +function resolveOAuthProviders(agentDir?: string): UsageProviderId[] { + const store = ensureAuthProfileStore(agentDir, { + allowKeychainPrompt: false, + }); + const cfg = loadConfig(); + const providers = [ + "anthropic", + "github-copilot", + "google-gemini-cli", + "google-antigravity", + "openai-codex", + ] satisfies UsageProviderId[]; + const isOAuthLikeCredential = (id: string) => { + const cred = store.profiles[id]; + return cred?.type === "oauth" || cred?.type === "token"; + }; + return providers.filter((provider) => { + const profiles = listProfilesForProvider(store, provider).filter( + isOAuthLikeCredential, + ); + if (profiles.length > 0) return true; + const normalized = normalizeProviderId(provider); + const configuredProfiles = Object.entries(cfg.auth?.profiles ?? {}) + .filter( + ([, profile]) => normalizeProviderId(profile.provider) === normalized, + ) + .map(([id]) => id) + .filter(isOAuthLikeCredential); + return configuredProfiles.length > 0; + }); +} + +export async function resolveProviderAuths(params: { + providers: UsageProviderId[]; + auth?: ProviderAuth[]; + agentDir?: string; +}): Promise { + if (params.auth) return params.auth; + + const oauthProviders = resolveOAuthProviders(params.agentDir); + const auths: ProviderAuth[] = []; + + for (const provider of params.providers) { + if (provider === "zai") { + const apiKey = resolveZaiApiKey(); + if (apiKey) auths.push({ provider, token: apiKey }); + continue; + } + + if (!oauthProviders.includes(provider)) continue; + const auth = await resolveOAuthToken({ + provider, + agentDir: params.agentDir, + }); + if (auth) auths.push(auth); + } + + return auths; +} diff --git a/src/infra/provider-usage.fetch.claude.ts b/src/infra/provider-usage.fetch.claude.ts new file mode 100644 index 000000000..56ca7bcd5 --- /dev/null +++ b/src/infra/provider-usage.fetch.claude.ts @@ -0,0 +1,196 @@ +import { fetchJson } from "./provider-usage.fetch.shared.js"; +import { clampPercent, PROVIDER_LABELS } from "./provider-usage.shared.js"; +import type { + ProviderUsageSnapshot, + UsageWindow, +} from "./provider-usage.types.js"; + +type ClaudeUsageResponse = { + five_hour?: { utilization?: number; resets_at?: string }; + seven_day?: { utilization?: number; resets_at?: string }; + seven_day_sonnet?: { utilization?: number }; + seven_day_opus?: { utilization?: number }; +}; + +type ClaudeWebOrganizationsResponse = Array<{ + uuid?: string; + name?: string; +}>; + +type ClaudeWebUsageResponse = ClaudeUsageResponse; + +function resolveClaudeWebSessionKey(): string | undefined { + const direct = + process.env.CLAUDE_AI_SESSION_KEY?.trim() ?? + process.env.CLAUDE_WEB_SESSION_KEY?.trim(); + if (direct?.startsWith("sk-ant-")) return direct; + + const cookieHeader = process.env.CLAUDE_WEB_COOKIE?.trim(); + if (!cookieHeader) return undefined; + const stripped = cookieHeader.replace(/^cookie:\\s*/i, ""); + const match = stripped.match(/(?:^|;\\s*)sessionKey=([^;\\s]+)/i); + const value = match?.[1]?.trim(); + return value?.startsWith("sk-ant-") ? value : undefined; +} + +async function fetchClaudeWebUsage( + sessionKey: string, + timeoutMs: number, + fetchFn: typeof fetch, +): Promise { + const headers: Record = { + Cookie: `sessionKey=${sessionKey}`, + Accept: "application/json", + }; + + const orgRes = await fetchJson( + "https://claude.ai/api/organizations", + { headers }, + timeoutMs, + fetchFn, + ); + if (!orgRes.ok) return null; + + const orgs = (await orgRes.json()) as ClaudeWebOrganizationsResponse; + const orgId = orgs?.[0]?.uuid?.trim(); + if (!orgId) return null; + + const usageRes = await fetchJson( + `https://claude.ai/api/organizations/${orgId}/usage`, + { headers }, + timeoutMs, + fetchFn, + ); + if (!usageRes.ok) return null; + + const data = (await usageRes.json()) as ClaudeWebUsageResponse; + const windows: UsageWindow[] = []; + + if (data.five_hour?.utilization !== undefined) { + windows.push({ + label: "5h", + usedPercent: clampPercent(data.five_hour.utilization), + resetAt: data.five_hour.resets_at + ? new Date(data.five_hour.resets_at).getTime() + : undefined, + }); + } + + if (data.seven_day?.utilization !== undefined) { + windows.push({ + label: "Week", + usedPercent: clampPercent(data.seven_day.utilization), + resetAt: data.seven_day.resets_at + ? new Date(data.seven_day.resets_at).getTime() + : undefined, + }); + } + + const modelWindow = data.seven_day_sonnet || data.seven_day_opus; + if (modelWindow?.utilization !== undefined) { + windows.push({ + label: data.seven_day_sonnet ? "Sonnet" : "Opus", + usedPercent: clampPercent(modelWindow.utilization), + }); + } + + if (windows.length === 0) return null; + return { + provider: "anthropic", + displayName: PROVIDER_LABELS.anthropic, + windows, + }; +} + +export async function fetchClaudeUsage( + token: string, + timeoutMs: number, + fetchFn: typeof fetch, +): Promise { + const res = await fetchJson( + "https://api.anthropic.com/api/oauth/usage", + { + headers: { + Authorization: `Bearer ${token}`, + "User-Agent": "clawdbot", + Accept: "application/json", + "anthropic-version": "2023-06-01", + "anthropic-beta": "oauth-2025-04-20", + }, + }, + timeoutMs, + fetchFn, + ); + + if (!res.ok) { + let message: string | undefined; + try { + const data = (await res.json()) as { + error?: { message?: unknown } | null; + }; + const raw = data?.error?.message; + if (typeof raw === "string" && raw.trim()) message = raw.trim(); + } catch { + // ignore parse errors + } + + // Claude CLI setup-token yields tokens that can be used for inference, but may not + // include user:profile scope required by the OAuth usage endpoint. When a claude.ai + // browser sessionKey is available, fall back to the web API. + if ( + res.status === 403 && + message?.includes("scope requirement user:profile") + ) { + const sessionKey = resolveClaudeWebSessionKey(); + if (sessionKey) { + const web = await fetchClaudeWebUsage(sessionKey, timeoutMs, fetchFn); + if (web) return web; + } + } + + const suffix = message ? `: ${message}` : ""; + return { + provider: "anthropic", + displayName: PROVIDER_LABELS.anthropic, + windows: [], + error: `HTTP ${res.status}${suffix}`, + }; + } + + const data = (await res.json()) as ClaudeUsageResponse; + const windows: UsageWindow[] = []; + + if (data.five_hour?.utilization !== undefined) { + windows.push({ + label: "5h", + usedPercent: clampPercent(data.five_hour.utilization), + resetAt: data.five_hour.resets_at + ? new Date(data.five_hour.resets_at).getTime() + : undefined, + }); + } + + if (data.seven_day?.utilization !== undefined) { + windows.push({ + label: "Week", + usedPercent: clampPercent(data.seven_day.utilization), + resetAt: data.seven_day.resets_at + ? new Date(data.seven_day.resets_at).getTime() + : undefined, + }); + } + + const modelWindow = data.seven_day_sonnet || data.seven_day_opus; + if (modelWindow?.utilization !== undefined) { + windows.push({ + label: data.seven_day_sonnet ? "Sonnet" : "Opus", + usedPercent: clampPercent(modelWindow.utilization), + }); + } + + return { + provider: "anthropic", + displayName: PROVIDER_LABELS.anthropic, + windows, + }; +} diff --git a/src/infra/provider-usage.fetch.codex.ts b/src/infra/provider-usage.fetch.codex.ts new file mode 100644 index 000000000..d387013ad --- /dev/null +++ b/src/infra/provider-usage.fetch.codex.ts @@ -0,0 +1,102 @@ +import { fetchJson } from "./provider-usage.fetch.shared.js"; +import { clampPercent, PROVIDER_LABELS } from "./provider-usage.shared.js"; +import type { + ProviderUsageSnapshot, + UsageWindow, +} from "./provider-usage.types.js"; + +type CodexUsageResponse = { + rate_limit?: { + primary_window?: { + limit_window_seconds?: number; + used_percent?: number; + reset_at?: number; + }; + secondary_window?: { + limit_window_seconds?: number; + used_percent?: number; + reset_at?: number; + }; + }; + plan_type?: string; + credits?: { balance?: number | string | null }; +}; + +export async function fetchCodexUsage( + token: string, + accountId: string | undefined, + timeoutMs: number, + fetchFn: typeof fetch, +): Promise { + const headers: Record = { + Authorization: `Bearer ${token}`, + "User-Agent": "CodexBar", + Accept: "application/json", + }; + if (accountId) headers["ChatGPT-Account-Id"] = accountId; + + const res = await fetchJson( + "https://chatgpt.com/backend-api/wham/usage", + { method: "GET", headers }, + timeoutMs, + fetchFn, + ); + + if (res.status === 401 || res.status === 403) { + return { + provider: "openai-codex", + displayName: PROVIDER_LABELS["openai-codex"], + windows: [], + error: "Token expired", + }; + } + + if (!res.ok) { + return { + provider: "openai-codex", + displayName: PROVIDER_LABELS["openai-codex"], + windows: [], + error: `HTTP ${res.status}`, + }; + } + + const data = (await res.json()) as CodexUsageResponse; + const windows: UsageWindow[] = []; + + if (data.rate_limit?.primary_window) { + const pw = data.rate_limit.primary_window; + const windowHours = Math.round((pw.limit_window_seconds || 10800) / 3600); + windows.push({ + label: `${windowHours}h`, + usedPercent: clampPercent(pw.used_percent || 0), + resetAt: pw.reset_at ? pw.reset_at * 1000 : undefined, + }); + } + + if (data.rate_limit?.secondary_window) { + const sw = data.rate_limit.secondary_window; + const windowHours = Math.round((sw.limit_window_seconds || 86400) / 3600); + const label = windowHours >= 24 ? "Day" : `${windowHours}h`; + windows.push({ + label, + usedPercent: clampPercent(sw.used_percent || 0), + resetAt: sw.reset_at ? sw.reset_at * 1000 : undefined, + }); + } + + let plan = data.plan_type; + if (data.credits?.balance !== undefined && data.credits.balance !== null) { + const balance = + typeof data.credits.balance === "number" + ? data.credits.balance + : parseFloat(data.credits.balance) || 0; + plan = plan ? `${plan} ($${balance.toFixed(2)})` : `$${balance.toFixed(2)}`; + } + + return { + provider: "openai-codex", + displayName: PROVIDER_LABELS["openai-codex"], + windows, + plan, + }; +} diff --git a/src/infra/provider-usage.fetch.copilot.ts b/src/infra/provider-usage.fetch.copilot.ts new file mode 100644 index 000000000..fa4f243d3 --- /dev/null +++ b/src/infra/provider-usage.fetch.copilot.ts @@ -0,0 +1,70 @@ +import { fetchJson } from "./provider-usage.fetch.shared.js"; +import { clampPercent, PROVIDER_LABELS } from "./provider-usage.shared.js"; +import type { + ProviderUsageSnapshot, + UsageWindow, +} from "./provider-usage.types.js"; + +type CopilotUsageResponse = { + quota_snapshots?: { + premium_interactions?: { percent_remaining?: number | null }; + chat?: { percent_remaining?: number | null }; + }; + copilot_plan?: string; +}; + +export async function fetchCopilotUsage( + token: string, + timeoutMs: number, + fetchFn: typeof fetch, +): Promise { + const res = await fetchJson( + "https://api.github.com/copilot_internal/user", + { + headers: { + Authorization: `token ${token}`, + "Editor-Version": "vscode/1.96.2", + "User-Agent": "GitHubCopilotChat/0.26.7", + "X-Github-Api-Version": "2025-04-01", + }, + }, + timeoutMs, + fetchFn, + ); + + if (!res.ok) { + return { + provider: "github-copilot", + displayName: PROVIDER_LABELS["github-copilot"], + windows: [], + error: `HTTP ${res.status}`, + }; + } + + const data = (await res.json()) as CopilotUsageResponse; + const windows: UsageWindow[] = []; + + if (data.quota_snapshots?.premium_interactions) { + const remaining = + data.quota_snapshots.premium_interactions.percent_remaining; + windows.push({ + label: "Premium", + usedPercent: clampPercent(100 - (remaining ?? 0)), + }); + } + + if (data.quota_snapshots?.chat) { + const remaining = data.quota_snapshots.chat.percent_remaining; + windows.push({ + label: "Chat", + usedPercent: clampPercent(100 - (remaining ?? 0)), + }); + } + + return { + provider: "github-copilot", + displayName: PROVIDER_LABELS["github-copilot"], + windows, + plan: data.copilot_plan, + }; +} diff --git a/src/infra/provider-usage.fetch.gemini.ts b/src/infra/provider-usage.fetch.gemini.ts new file mode 100644 index 000000000..5c6fe244d --- /dev/null +++ b/src/infra/provider-usage.fetch.gemini.ts @@ -0,0 +1,81 @@ +import { fetchJson } from "./provider-usage.fetch.shared.js"; +import { clampPercent, PROVIDER_LABELS } from "./provider-usage.shared.js"; +import type { + ProviderUsageSnapshot, + UsageProviderId, + UsageWindow, +} from "./provider-usage.types.js"; + +type GeminiUsageResponse = { + buckets?: Array<{ modelId?: string; remainingFraction?: number }>; +}; + +export async function fetchGeminiUsage( + token: string, + timeoutMs: number, + fetchFn: typeof fetch, + provider: UsageProviderId, +): Promise { + const res = await fetchJson( + "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota", + { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + }, + body: "{}", + }, + timeoutMs, + fetchFn, + ); + + if (!res.ok) { + return { + provider, + displayName: PROVIDER_LABELS[provider], + windows: [], + error: `HTTP ${res.status}`, + }; + } + + const data = (await res.json()) as GeminiUsageResponse; + const quotas: Record = {}; + + for (const bucket of data.buckets || []) { + const model = bucket.modelId || "unknown"; + const frac = bucket.remainingFraction ?? 1; + if (!quotas[model] || frac < quotas[model]) quotas[model] = frac; + } + + const windows: UsageWindow[] = []; + let proMin = 1; + let flashMin = 1; + let hasPro = false; + let hasFlash = false; + + for (const [model, frac] of Object.entries(quotas)) { + const lower = model.toLowerCase(); + if (lower.includes("pro")) { + hasPro = true; + if (frac < proMin) proMin = frac; + } + if (lower.includes("flash")) { + hasFlash = true; + if (frac < flashMin) flashMin = frac; + } + } + + if (hasPro) + windows.push({ + label: "Pro", + usedPercent: clampPercent((1 - proMin) * 100), + }); + if (hasFlash) + windows.push({ + label: "Flash", + usedPercent: clampPercent((1 - flashMin) * 100), + }); + + return { provider, displayName: PROVIDER_LABELS[provider], windows }; +} diff --git a/src/infra/provider-usage.fetch.shared.ts b/src/infra/provider-usage.fetch.shared.ts new file mode 100644 index 000000000..3e8062277 --- /dev/null +++ b/src/infra/provider-usage.fetch.shared.ts @@ -0,0 +1,14 @@ +export async function fetchJson( + url: string, + init: RequestInit, + timeoutMs: number, + fetchFn: typeof fetch, +): Promise { + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), timeoutMs); + try { + return await fetchFn(url, { ...init, signal: controller.signal }); + } finally { + clearTimeout(timer); + } +} diff --git a/src/infra/provider-usage.fetch.ts b/src/infra/provider-usage.fetch.ts new file mode 100644 index 000000000..76d41dfb7 --- /dev/null +++ b/src/infra/provider-usage.fetch.ts @@ -0,0 +1,5 @@ +export { fetchClaudeUsage } from "./provider-usage.fetch.claude.js"; +export { fetchCodexUsage } from "./provider-usage.fetch.codex.js"; +export { fetchCopilotUsage } from "./provider-usage.fetch.copilot.js"; +export { fetchGeminiUsage } from "./provider-usage.fetch.gemini.js"; +export { fetchZaiUsage } from "./provider-usage.fetch.zai.js"; diff --git a/src/infra/provider-usage.fetch.zai.ts b/src/infra/provider-usage.fetch.zai.ts new file mode 100644 index 000000000..9b52561ea --- /dev/null +++ b/src/infra/provider-usage.fetch.zai.ts @@ -0,0 +1,97 @@ +import { fetchJson } from "./provider-usage.fetch.shared.js"; +import { clampPercent, PROVIDER_LABELS } from "./provider-usage.shared.js"; +import type { + ProviderUsageSnapshot, + UsageWindow, +} from "./provider-usage.types.js"; + +type ZaiUsageResponse = { + success?: boolean; + code?: number; + msg?: string; + data?: { + planName?: string; + plan?: string; + limits?: Array<{ + type?: string; + percentage?: number; + unit?: number; + number?: number; + nextResetTime?: string; + }>; + }; +}; + +export async function fetchZaiUsage( + apiKey: string, + timeoutMs: number, + fetchFn: typeof fetch, +): Promise { + const res = await fetchJson( + "https://api.z.ai/api/monitor/usage/quota/limit", + { + method: "GET", + headers: { + Authorization: `Bearer ${apiKey}`, + Accept: "application/json", + }, + }, + timeoutMs, + fetchFn, + ); + + if (!res.ok) { + return { + provider: "zai", + displayName: PROVIDER_LABELS.zai, + windows: [], + error: `HTTP ${res.status}`, + }; + } + + const data = (await res.json()) as ZaiUsageResponse; + if (!data.success || data.code !== 200) { + return { + provider: "zai", + displayName: PROVIDER_LABELS.zai, + windows: [], + error: data.msg || "API error", + }; + } + + const windows: UsageWindow[] = []; + const limits = data.data?.limits || []; + + for (const limit of limits) { + const percent = clampPercent(limit.percentage || 0); + const nextReset = limit.nextResetTime + ? new Date(limit.nextResetTime).getTime() + : undefined; + let windowLabel = "Limit"; + if (limit.unit === 1) windowLabel = `${limit.number}d`; + else if (limit.unit === 3) windowLabel = `${limit.number}h`; + else if (limit.unit === 5) windowLabel = `${limit.number}m`; + + if (limit.type === "TOKENS_LIMIT") { + windows.push({ + label: `Tokens (${windowLabel})`, + usedPercent: percent, + resetAt: nextReset, + }); + } else if (limit.type === "TIME_LIMIT") { + windows.push({ + label: "Monthly", + usedPercent: percent, + resetAt: nextReset, + }); + } + } + + const planName = data.data?.planName || data.data?.plan || undefined; + return { + provider: "zai", + displayName: PROVIDER_LABELS.zai, + windows, + plan: planName, + }; +} diff --git a/src/infra/provider-usage.format.ts b/src/infra/provider-usage.format.ts new file mode 100644 index 000000000..b65ef5af6 --- /dev/null +++ b/src/infra/provider-usage.format.ts @@ -0,0 +1,91 @@ +import { clampPercent } from "./provider-usage.shared.js"; +import type { UsageSummary, UsageWindow } from "./provider-usage.types.js"; + +function formatResetRemaining(targetMs?: number, now?: number): string | null { + if (!targetMs) return null; + const base = now ?? Date.now(); + const diffMs = targetMs - base; + if (diffMs <= 0) return "now"; + + const diffMins = Math.floor(diffMs / 60000); + if (diffMins < 60) return `${diffMins}m`; + + const hours = Math.floor(diffMins / 60); + const mins = diffMins % 60; + if (hours < 24) return mins > 0 ? `${hours}h ${mins}m` : `${hours}h`; + + const days = Math.floor(hours / 24); + if (days < 7) return `${days}d ${hours % 24}h`; + + return new Intl.DateTimeFormat("en-US", { + month: "short", + day: "numeric", + }).format(new Date(targetMs)); +} + +function pickPrimaryWindow(windows: UsageWindow[]): UsageWindow | undefined { + if (windows.length === 0) return undefined; + return windows.reduce((best, next) => + next.usedPercent > best.usedPercent ? next : best, + ); +} + +function formatWindowShort(window: UsageWindow, now?: number): string { + const remaining = clampPercent(100 - window.usedPercent); + const reset = formatResetRemaining(window.resetAt, now); + const resetSuffix = reset ? ` ⏱${reset}` : ""; + return `${remaining.toFixed(0)}% left (${window.label}${resetSuffix})`; +} + +export function formatUsageSummaryLine( + summary: UsageSummary, + opts?: { now?: number; maxProviders?: number }, +): string | null { + const providers = summary.providers + .filter((entry) => entry.windows.length > 0 && !entry.error) + .slice(0, opts?.maxProviders ?? summary.providers.length); + if (providers.length === 0) return null; + + const parts = providers + .map((entry) => { + const window = pickPrimaryWindow(entry.windows); + if (!window) return null; + return `${entry.displayName} ${formatWindowShort(window, opts?.now)}`; + }) + .filter(Boolean) as string[]; + + if (parts.length === 0) return null; + return `📊 Usage: ${parts.join(" · ")}`; +} + +export function formatUsageReportLines( + summary: UsageSummary, + opts?: { now?: number }, +): string[] { + if (summary.providers.length === 0) { + return ["Usage: no provider usage available."]; + } + + const lines: string[] = ["Usage:"]; + for (const entry of summary.providers) { + const planSuffix = entry.plan ? ` (${entry.plan})` : ""; + if (entry.error) { + lines.push(` ${entry.displayName}${planSuffix}: ${entry.error}`); + continue; + } + if (entry.windows.length === 0) { + lines.push(` ${entry.displayName}${planSuffix}: no data`); + continue; + } + lines.push(` ${entry.displayName}${planSuffix}`); + for (const window of entry.windows) { + const remaining = clampPercent(100 - window.usedPercent); + const reset = formatResetRemaining(window.resetAt, opts?.now); + const resetSuffix = reset ? ` · resets ${reset}` : ""; + lines.push( + ` ${window.label}: ${remaining.toFixed(0)}% left${resetSuffix}`, + ); + } + } + return lines; +} diff --git a/src/infra/provider-usage.load.ts b/src/infra/provider-usage.load.ts new file mode 100644 index 000000000..ce1f9c8f9 --- /dev/null +++ b/src/infra/provider-usage.load.ts @@ -0,0 +1,102 @@ +import { + type ProviderAuth, + resolveProviderAuths, +} from "./provider-usage.auth.js"; +import { + fetchClaudeUsage, + fetchCodexUsage, + fetchCopilotUsage, + fetchGeminiUsage, + fetchZaiUsage, +} from "./provider-usage.fetch.js"; +import { + DEFAULT_TIMEOUT_MS, + ignoredErrors, + PROVIDER_LABELS, + usageProviders, + withTimeout, +} from "./provider-usage.shared.js"; +import type { + ProviderUsageSnapshot, + UsageProviderId, + UsageSummary, +} from "./provider-usage.types.js"; + +type UsageSummaryOptions = { + now?: number; + timeoutMs?: number; + providers?: UsageProviderId[]; + auth?: ProviderAuth[]; + agentDir?: string; + fetch?: typeof fetch; +}; + +export async function loadProviderUsageSummary( + opts: UsageSummaryOptions = {}, +): Promise { + const now = opts.now ?? Date.now(); + const timeoutMs = opts.timeoutMs ?? DEFAULT_TIMEOUT_MS; + const fetchFn = opts.fetch ?? fetch; + + const auths = await resolveProviderAuths({ + providers: opts.providers ?? usageProviders, + auth: opts.auth, + agentDir: opts.agentDir, + }); + if (auths.length === 0) { + return { updatedAt: now, providers: [] }; + } + + const tasks = auths.map((auth) => + withTimeout( + (async (): Promise => { + switch (auth.provider) { + case "anthropic": + return await fetchClaudeUsage(auth.token, timeoutMs, fetchFn); + case "github-copilot": + return await fetchCopilotUsage(auth.token, timeoutMs, fetchFn); + case "google-gemini-cli": + case "google-antigravity": + return await fetchGeminiUsage( + auth.token, + timeoutMs, + fetchFn, + auth.provider, + ); + case "openai-codex": + return await fetchCodexUsage( + auth.token, + auth.accountId, + timeoutMs, + fetchFn, + ); + case "zai": + return await fetchZaiUsage(auth.token, timeoutMs, fetchFn); + default: + return { + provider: auth.provider, + displayName: PROVIDER_LABELS[auth.provider], + windows: [], + error: "Unsupported provider", + }; + } + })(), + timeoutMs + 1000, + { + provider: auth.provider, + displayName: PROVIDER_LABELS[auth.provider], + windows: [], + error: "Timeout", + }, + ), + ); + + const snapshots = await Promise.all(tasks); + const providers = snapshots.filter((entry) => { + if (entry.windows.length > 0) return true; + if (!entry.error) return true; + return !ignoredErrors.has(entry.error); + }); + + return { updatedAt: now, providers }; +} diff --git a/src/infra/provider-usage.shared.ts b/src/infra/provider-usage.shared.ts new file mode 100644 index 000000000..c0872dd43 --- /dev/null +++ b/src/infra/provider-usage.shared.ts @@ -0,0 +1,61 @@ +import { normalizeProviderId } from "../agents/model-selection.js"; +import type { UsageProviderId } from "./provider-usage.types.js"; + +export const DEFAULT_TIMEOUT_MS = 5000; + +export const PROVIDER_LABELS: Record = { + anthropic: "Claude", + "github-copilot": "Copilot", + "google-gemini-cli": "Gemini", + "google-antigravity": "Antigravity", + "openai-codex": "Codex", + zai: "z.ai", +}; + +export const usageProviders: UsageProviderId[] = [ + "anthropic", + "github-copilot", + "google-gemini-cli", + "google-antigravity", + "openai-codex", + "zai", +]; + +export function resolveUsageProviderId( + provider?: string | null, +): UsageProviderId | undefined { + if (!provider) return undefined; + const normalized = normalizeProviderId(provider); + return usageProviders.includes(normalized as UsageProviderId) + ? (normalized as UsageProviderId) + : undefined; +} + +export const ignoredErrors = new Set([ + "No credentials", + "No token", + "No API key", + "Not logged in", + "No auth", +]); + +export const clampPercent = (value: number) => + Math.max(0, Math.min(100, Number.isFinite(value) ? value : 0)); + +export const withTimeout = async ( + work: Promise, + ms: number, + fallback: T, +): Promise => { + let timeout: NodeJS.Timeout | undefined; + try { + return await Promise.race([ + work, + new Promise((resolve) => { + timeout = setTimeout(() => resolve(fallback), ms); + }), + ]); + } finally { + if (timeout) clearTimeout(timeout); + } +}; diff --git a/src/infra/provider-usage.ts b/src/infra/provider-usage.ts index 9e081dbd9..f88b89fbc 100644 --- a/src/infra/provider-usage.ts +++ b/src/infra/provider-usage.ts @@ -1,964 +1,12 @@ -import fs from "node:fs"; -import os from "node:os"; -import path from "node:path"; - -import { - CLAUDE_CLI_PROFILE_ID, - ensureAuthProfileStore, - listProfilesForProvider, - resolveApiKeyForProfile, - resolveAuthProfileOrder, -} from "../agents/auth-profiles.js"; -import { - getCustomProviderApiKey, - resolveEnvApiKey, -} from "../agents/model-auth.js"; -import { normalizeProviderId } from "../agents/model-selection.js"; -import { loadConfig } from "../config/config.js"; - -export type UsageWindow = { - label: string; - usedPercent: number; - resetAt?: number; -}; - -export type ProviderUsageSnapshot = { - provider: UsageProviderId; - displayName: string; - windows: UsageWindow[]; - plan?: string; - error?: string; -}; - -export type UsageSummary = { - updatedAt: number; - providers: ProviderUsageSnapshot[]; -}; - -export type UsageProviderId = - | "anthropic" - | "github-copilot" - | "google-gemini-cli" - | "google-antigravity" - | "openai-codex" - | "zai"; - -type ClaudeUsageResponse = { - five_hour?: { utilization?: number; resets_at?: string }; - seven_day?: { utilization?: number; resets_at?: string }; - seven_day_sonnet?: { utilization?: number }; - seven_day_opus?: { utilization?: number }; -}; - -type ClaudeWebOrganizationsResponse = Array<{ - uuid?: string; - name?: string; -}>; - -type ClaudeWebUsageResponse = ClaudeUsageResponse; - -type CopilotUsageResponse = { - quota_snapshots?: { - premium_interactions?: { percent_remaining?: number | null }; - chat?: { percent_remaining?: number | null }; - }; - copilot_plan?: string; -}; - -type GeminiUsageResponse = { - buckets?: Array<{ modelId?: string; remainingFraction?: number }>; -}; - -type CodexUsageResponse = { - rate_limit?: { - primary_window?: { - limit_window_seconds?: number; - used_percent?: number; - reset_at?: number; - }; - secondary_window?: { - limit_window_seconds?: number; - used_percent?: number; - reset_at?: number; - }; - }; - plan_type?: string; - credits?: { balance?: number | string | null }; -}; - -type ZaiUsageResponse = { - success?: boolean; - code?: number; - msg?: string; - data?: { - planName?: string; - plan?: string; - limits?: Array<{ - type?: string; - percentage?: number; - unit?: number; - number?: number; - nextResetTime?: string; - }>; - }; -}; - -type ProviderAuth = { - provider: UsageProviderId; - token: string; - accountId?: string; -}; - -type UsageSummaryOptions = { - now?: number; - timeoutMs?: number; - providers?: UsageProviderId[]; - auth?: ProviderAuth[]; - agentDir?: string; - fetch?: typeof fetch; -}; - -const DEFAULT_TIMEOUT_MS = 5000; - -const PROVIDER_LABELS: Record = { - anthropic: "Claude", - "github-copilot": "Copilot", - "google-gemini-cli": "Gemini", - "google-antigravity": "Antigravity", - "openai-codex": "Codex", - zai: "z.ai", -}; - -const usageProviders: UsageProviderId[] = [ - "anthropic", - "github-copilot", - "google-gemini-cli", - "google-antigravity", - "openai-codex", - "zai", -]; - -export function resolveUsageProviderId( - provider?: string | null, -): UsageProviderId | undefined { - if (!provider) return undefined; - const normalized = normalizeProviderId(provider); - return usageProviders.includes(normalized as UsageProviderId) - ? (normalized as UsageProviderId) - : undefined; -} - -const ignoredErrors = new Set([ - "No credentials", - "No token", - "No API key", - "Not logged in", - "No auth", -]); - -const clampPercent = (value: number) => - Math.max(0, Math.min(100, Number.isFinite(value) ? value : 0)); - -const withTimeout = async ( - work: Promise, - ms: number, - fallback: T, -): Promise => { - let timeout: NodeJS.Timeout | undefined; - try { - return await Promise.race([ - work, - new Promise((resolve) => { - timeout = setTimeout(() => resolve(fallback), ms); - }), - ]); - } finally { - if (timeout) clearTimeout(timeout); - } -}; - -function formatResetRemaining(targetMs?: number, now?: number): string | null { - if (!targetMs) return null; - const base = now ?? Date.now(); - const diffMs = targetMs - base; - if (diffMs <= 0) return "now"; - - const diffMins = Math.floor(diffMs / 60000); - if (diffMins < 60) return `${diffMins}m`; - - const hours = Math.floor(diffMins / 60); - const mins = diffMins % 60; - if (hours < 24) return mins > 0 ? `${hours}h ${mins}m` : `${hours}h`; - - const days = Math.floor(hours / 24); - if (days < 7) return `${days}d ${hours % 24}h`; - - return new Intl.DateTimeFormat("en-US", { - month: "short", - day: "numeric", - }).format(new Date(targetMs)); -} - -function resolveClaudeWebSessionKey(): string | undefined { - const direct = - process.env.CLAUDE_AI_SESSION_KEY?.trim() ?? - process.env.CLAUDE_WEB_SESSION_KEY?.trim(); - if (direct?.startsWith("sk-ant-")) return direct; - - const cookieHeader = process.env.CLAUDE_WEB_COOKIE?.trim(); - if (!cookieHeader) return undefined; - const stripped = cookieHeader.replace(/^cookie:\\s*/i, ""); - const match = stripped.match(/(?:^|;\\s*)sessionKey=([^;\\s]+)/i); - const value = match?.[1]?.trim(); - return value?.startsWith("sk-ant-") ? value : undefined; -} - -function pickPrimaryWindow(windows: UsageWindow[]): UsageWindow | undefined { - if (windows.length === 0) return undefined; - return windows.reduce((best, next) => - next.usedPercent > best.usedPercent ? next : best, - ); -} - -function formatWindowShort(window: UsageWindow, now?: number): string { - const remaining = clampPercent(100 - window.usedPercent); - const reset = formatResetRemaining(window.resetAt, now); - const resetSuffix = reset ? ` ⏱${reset}` : ""; - return `${remaining.toFixed(0)}% left (${window.label}${resetSuffix})`; -} - -export function formatUsageSummaryLine( - summary: UsageSummary, - opts?: { now?: number; maxProviders?: number }, -): string | null { - const providers = summary.providers - .filter((entry) => entry.windows.length > 0 && !entry.error) - .slice(0, opts?.maxProviders ?? summary.providers.length); - if (providers.length === 0) return null; - - const parts = providers - .map((entry) => { - const window = pickPrimaryWindow(entry.windows); - if (!window) return null; - return `${entry.displayName} ${formatWindowShort(window, opts?.now)}`; - }) - .filter(Boolean) as string[]; - - if (parts.length === 0) return null; - return `📊 Usage: ${parts.join(" · ")}`; -} - -export function formatUsageReportLines( - summary: UsageSummary, - opts?: { now?: number }, -): string[] { - if (summary.providers.length === 0) { - return ["Usage: no provider usage available."]; - } - - const lines: string[] = ["Usage:"]; - for (const entry of summary.providers) { - const planSuffix = entry.plan ? ` (${entry.plan})` : ""; - if (entry.error) { - lines.push(` ${entry.displayName}${planSuffix}: ${entry.error}`); - continue; - } - if (entry.windows.length === 0) { - lines.push(` ${entry.displayName}${planSuffix}: no data`); - continue; - } - lines.push(` ${entry.displayName}${planSuffix}`); - for (const window of entry.windows) { - const remaining = clampPercent(100 - window.usedPercent); - const reset = formatResetRemaining(window.resetAt, opts?.now); - const resetSuffix = reset ? ` · resets ${reset}` : ""; - lines.push( - ` ${window.label}: ${remaining.toFixed(0)}% left${resetSuffix}`, - ); - } - } - return lines; -} - -function parseGoogleToken(apiKey: string): { token: string } | null { - if (!apiKey) return null; - try { - const parsed = JSON.parse(apiKey) as { token?: unknown }; - if (parsed && typeof parsed.token === "string") { - return { token: parsed.token }; - } - } catch { - // ignore - } - return null; -} - -async function fetchJson( - url: string, - init: RequestInit, - timeoutMs: number, - fetchFn: typeof fetch, -): Promise { - const controller = new AbortController(); - const timer = setTimeout(() => controller.abort(), timeoutMs); - try { - return await fetchFn(url, { ...init, signal: controller.signal }); - } finally { - clearTimeout(timer); - } -} - -async function fetchClaudeUsage( - token: string, - timeoutMs: number, - fetchFn: typeof fetch, -): Promise { - const res = await fetchJson( - "https://api.anthropic.com/api/oauth/usage", - { - headers: { - Authorization: `Bearer ${token}`, - "User-Agent": "clawdbot", - Accept: "application/json", - "anthropic-version": "2023-06-01", - "anthropic-beta": "oauth-2025-04-20", - }, - }, - timeoutMs, - fetchFn, - ); - - if (!res.ok) { - let message: string | undefined; - try { - const data = (await res.json()) as { - error?: { message?: unknown } | null; - }; - const raw = data?.error?.message; - if (typeof raw === "string" && raw.trim()) message = raw.trim(); - } catch { - // ignore parse errors - } - - // Claude CLI setup-token yields tokens that can be used for inference - // but may not include user:profile scope required by the OAuth usage endpoint. - // When a claude.ai browser sessionKey is available, fall back to the web API. - if ( - res.status === 403 && - message?.includes("scope requirement user:profile") - ) { - const sessionKey = resolveClaudeWebSessionKey(); - if (sessionKey) { - const web = await fetchClaudeWebUsage(sessionKey, timeoutMs, fetchFn); - if (web) return web; - } - } - - const suffix = message ? `: ${message}` : ""; - return { - provider: "anthropic", - displayName: PROVIDER_LABELS.anthropic, - windows: [], - error: `HTTP ${res.status}${suffix}`, - }; - } - - const data = (await res.json()) as ClaudeUsageResponse; - const windows: UsageWindow[] = []; - - if (data.five_hour?.utilization !== undefined) { - windows.push({ - label: "5h", - usedPercent: clampPercent(data.five_hour.utilization), - resetAt: data.five_hour.resets_at - ? new Date(data.five_hour.resets_at).getTime() - : undefined, - }); - } - - if (data.seven_day?.utilization !== undefined) { - windows.push({ - label: "Week", - usedPercent: clampPercent(data.seven_day.utilization), - resetAt: data.seven_day.resets_at - ? new Date(data.seven_day.resets_at).getTime() - : undefined, - }); - } - - const modelWindow = data.seven_day_sonnet || data.seven_day_opus; - if (modelWindow?.utilization !== undefined) { - windows.push({ - label: data.seven_day_sonnet ? "Sonnet" : "Opus", - usedPercent: clampPercent(modelWindow.utilization), - }); - } - - return { - provider: "anthropic", - displayName: PROVIDER_LABELS.anthropic, - windows, - }; -} - -async function fetchClaudeWebUsage( - sessionKey: string, - timeoutMs: number, - fetchFn: typeof fetch, -): Promise { - const headers: Record = { - Cookie: `sessionKey=${sessionKey}`, - Accept: "application/json", - }; - - const orgRes = await fetchJson( - "https://claude.ai/api/organizations", - { headers }, - timeoutMs, - fetchFn, - ); - if (!orgRes.ok) return null; - - const orgs = (await orgRes.json()) as ClaudeWebOrganizationsResponse; - const orgId = orgs?.[0]?.uuid?.trim(); - if (!orgId) return null; - - const usageRes = await fetchJson( - `https://claude.ai/api/organizations/${orgId}/usage`, - { headers }, - timeoutMs, - fetchFn, - ); - if (!usageRes.ok) return null; - - const data = (await usageRes.json()) as ClaudeWebUsageResponse; - const windows: UsageWindow[] = []; - - if (data.five_hour?.utilization !== undefined) { - windows.push({ - label: "5h", - usedPercent: clampPercent(data.five_hour.utilization), - resetAt: data.five_hour.resets_at - ? new Date(data.five_hour.resets_at).getTime() - : undefined, - }); - } - - if (data.seven_day?.utilization !== undefined) { - windows.push({ - label: "Week", - usedPercent: clampPercent(data.seven_day.utilization), - resetAt: data.seven_day.resets_at - ? new Date(data.seven_day.resets_at).getTime() - : undefined, - }); - } - - const modelWindow = data.seven_day_sonnet || data.seven_day_opus; - if (modelWindow?.utilization !== undefined) { - windows.push({ - label: data.seven_day_sonnet ? "Sonnet" : "Opus", - usedPercent: clampPercent(modelWindow.utilization), - }); - } - - if (windows.length === 0) return null; - return { - provider: "anthropic", - displayName: PROVIDER_LABELS.anthropic, - windows, - }; -} - -async function fetchCopilotUsage( - token: string, - timeoutMs: number, - fetchFn: typeof fetch, -): Promise { - const res = await fetchJson( - "https://api.github.com/copilot_internal/user", - { - headers: { - Authorization: `token ${token}`, - "Editor-Version": "vscode/1.96.2", - "User-Agent": "GitHubCopilotChat/0.26.7", - "X-Github-Api-Version": "2025-04-01", - }, - }, - timeoutMs, - fetchFn, - ); - - if (!res.ok) { - return { - provider: "github-copilot", - displayName: PROVIDER_LABELS["github-copilot"], - windows: [], - error: `HTTP ${res.status}`, - }; - } - - const data = (await res.json()) as CopilotUsageResponse; - const windows: UsageWindow[] = []; - - if (data.quota_snapshots?.premium_interactions) { - const remaining = - data.quota_snapshots.premium_interactions.percent_remaining; - windows.push({ - label: "Premium", - usedPercent: clampPercent(100 - (remaining ?? 0)), - }); - } - - if (data.quota_snapshots?.chat) { - const remaining = data.quota_snapshots.chat.percent_remaining; - windows.push({ - label: "Chat", - usedPercent: clampPercent(100 - (remaining ?? 0)), - }); - } - - return { - provider: "github-copilot", - displayName: PROVIDER_LABELS["github-copilot"], - windows, - plan: data.copilot_plan, - }; -} - -async function fetchGeminiUsage( - token: string, - timeoutMs: number, - fetchFn: typeof fetch, - provider: UsageProviderId, -): Promise { - const res = await fetchJson( - "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota", - { - method: "POST", - headers: { - Authorization: `Bearer ${token}`, - "Content-Type": "application/json", - }, - body: "{}", - }, - timeoutMs, - fetchFn, - ); - - if (!res.ok) { - return { - provider, - displayName: PROVIDER_LABELS[provider], - windows: [], - error: `HTTP ${res.status}`, - }; - } - - const data = (await res.json()) as GeminiUsageResponse; - const quotas: Record = {}; - - for (const bucket of data.buckets || []) { - const model = bucket.modelId || "unknown"; - const frac = bucket.remainingFraction ?? 1; - if (!quotas[model] || frac < quotas[model]) quotas[model] = frac; - } - - const windows: UsageWindow[] = []; - let proMin = 1; - let flashMin = 1; - let hasPro = false; - let hasFlash = false; - - for (const [model, frac] of Object.entries(quotas)) { - const lower = model.toLowerCase(); - if (lower.includes("pro")) { - hasPro = true; - if (frac < proMin) proMin = frac; - } - if (lower.includes("flash")) { - hasFlash = true; - if (frac < flashMin) flashMin = frac; - } - } - - if (hasPro) { - windows.push({ - label: "Pro", - usedPercent: clampPercent((1 - proMin) * 100), - }); - } - if (hasFlash) { - windows.push({ - label: "Flash", - usedPercent: clampPercent((1 - flashMin) * 100), - }); - } - - return { provider, displayName: PROVIDER_LABELS[provider], windows }; -} - -async function fetchCodexUsage( - token: string, - accountId: string | undefined, - timeoutMs: number, - fetchFn: typeof fetch, -): Promise { - const headers: Record = { - Authorization: `Bearer ${token}`, - "User-Agent": "CodexBar", - Accept: "application/json", - }; - if (accountId) headers["ChatGPT-Account-Id"] = accountId; - - const res = await fetchJson( - "https://chatgpt.com/backend-api/wham/usage", - { method: "GET", headers }, - timeoutMs, - fetchFn, - ); - - if (res.status === 401 || res.status === 403) { - return { - provider: "openai-codex", - displayName: PROVIDER_LABELS["openai-codex"], - windows: [], - error: "Token expired", - }; - } - - if (!res.ok) { - return { - provider: "openai-codex", - displayName: PROVIDER_LABELS["openai-codex"], - windows: [], - error: `HTTP ${res.status}`, - }; - } - - const data = (await res.json()) as CodexUsageResponse; - const windows: UsageWindow[] = []; - - if (data.rate_limit?.primary_window) { - const pw = data.rate_limit.primary_window; - const windowHours = Math.round((pw.limit_window_seconds || 10800) / 3600); - windows.push({ - label: `${windowHours}h`, - usedPercent: clampPercent(pw.used_percent || 0), - resetAt: pw.reset_at ? pw.reset_at * 1000 : undefined, - }); - } - - if (data.rate_limit?.secondary_window) { - const sw = data.rate_limit.secondary_window; - const windowHours = Math.round((sw.limit_window_seconds || 86400) / 3600); - const label = windowHours >= 24 ? "Day" : `${windowHours}h`; - windows.push({ - label, - usedPercent: clampPercent(sw.used_percent || 0), - resetAt: sw.reset_at ? sw.reset_at * 1000 : undefined, - }); - } - - let plan = data.plan_type; - if (data.credits?.balance !== undefined && data.credits.balance !== null) { - const balance = - typeof data.credits.balance === "number" - ? data.credits.balance - : parseFloat(data.credits.balance) || 0; - plan = plan ? `${plan} ($${balance.toFixed(2)})` : `$${balance.toFixed(2)}`; - } - - return { - provider: "openai-codex", - displayName: PROVIDER_LABELS["openai-codex"], - windows, - plan, - }; -} - -async function fetchZaiUsage( - apiKey: string, - timeoutMs: number, - fetchFn: typeof fetch, -): Promise { - const res = await fetchJson( - "https://api.z.ai/api/monitor/usage/quota/limit", - { - method: "GET", - headers: { - Authorization: `Bearer ${apiKey}`, - Accept: "application/json", - }, - }, - timeoutMs, - fetchFn, - ); - - if (!res.ok) { - return { - provider: "zai", - displayName: PROVIDER_LABELS.zai, - windows: [], - error: `HTTP ${res.status}`, - }; - } - - const data = (await res.json()) as ZaiUsageResponse; - if (!data.success || data.code !== 200) { - return { - provider: "zai", - displayName: PROVIDER_LABELS.zai, - windows: [], - error: data.msg || "API error", - }; - } - - const windows: UsageWindow[] = []; - const limits = data.data?.limits || []; - - for (const limit of limits) { - const percent = clampPercent(limit.percentage || 0); - const nextReset = limit.nextResetTime - ? new Date(limit.nextResetTime).getTime() - : undefined; - let windowLabel = "Limit"; - if (limit.unit === 1) windowLabel = `${limit.number}d`; - else if (limit.unit === 3) windowLabel = `${limit.number}h`; - else if (limit.unit === 5) windowLabel = `${limit.number}m`; - - if (limit.type === "TOKENS_LIMIT") { - windows.push({ - label: `Tokens (${windowLabel})`, - usedPercent: percent, - resetAt: nextReset, - }); - } else if (limit.type === "TIME_LIMIT") { - windows.push({ - label: "Monthly", - usedPercent: percent, - resetAt: nextReset, - }); - } - } - - const planName = data.data?.planName || data.data?.plan || undefined; - return { - provider: "zai", - displayName: PROVIDER_LABELS.zai, - windows, - plan: planName, - }; -} - -function resolveZaiApiKey(): string | undefined { - const envDirect = - process.env.ZAI_API_KEY?.trim() || process.env.Z_AI_API_KEY?.trim(); - if (envDirect) return envDirect; - - const envResolved = resolveEnvApiKey("zai"); - if (envResolved?.apiKey) return envResolved.apiKey; - - const cfg = loadConfig(); - const key = - getCustomProviderApiKey(cfg, "zai") || getCustomProviderApiKey(cfg, "z-ai"); - if (key) return key; - - const store = ensureAuthProfileStore(); - const apiProfile = [ - ...listProfilesForProvider(store, "zai"), - ...listProfilesForProvider(store, "z-ai"), - ].find((id) => store.profiles[id]?.type === "api_key"); - if (apiProfile) { - const cred = store.profiles[apiProfile]; - if (cred?.type === "api_key" && cred.key?.trim()) { - return cred.key.trim(); - } - } - - try { - const authPath = path.join(os.homedir(), ".pi", "agent", "auth.json"); - if (!fs.existsSync(authPath)) return undefined; - const data = JSON.parse(fs.readFileSync(authPath, "utf-8")) as Record< - string, - { access?: string } - >; - return data["z-ai"]?.access || data.zai?.access; - } catch { - return undefined; - } -} - -async function resolveOAuthToken(params: { - provider: UsageProviderId; - agentDir?: string; -}): Promise { - const cfg = loadConfig(); - const store = ensureAuthProfileStore(params.agentDir, { - allowKeychainPrompt: false, - }); - const order = resolveAuthProfileOrder({ - cfg, - store, - provider: params.provider, - }); - - // Claude CLI creds are the only Anthropic tokens that reliably include the - // `user:profile` scope required for the OAuth usage endpoint. - const candidates = - params.provider === "anthropic" ? [CLAUDE_CLI_PROFILE_ID, ...order] : order; - const deduped: string[] = []; - for (const entry of candidates) { - if (!deduped.includes(entry)) deduped.push(entry); - } - - for (const profileId of deduped) { - const cred = store.profiles[profileId]; - if (!cred || (cred.type !== "oauth" && cred.type !== "token")) continue; - try { - const resolved = await resolveApiKeyForProfile({ - // Usage snapshots should work even if config profile metadata is stale. - // (e.g. config says api_key but the store has a token profile.) - cfg: undefined, - store, - profileId, - agentDir: params.agentDir, - }); - if (!resolved?.apiKey) continue; - let token = resolved.apiKey; - if ( - params.provider === "google-gemini-cli" || - params.provider === "google-antigravity" - ) { - const parsed = parseGoogleToken(resolved.apiKey); - token = parsed?.token ?? resolved.apiKey; - } - return { - provider: params.provider, - token, - accountId: - cred.type === "oauth" && "accountId" in cred - ? (cred as { accountId?: string }).accountId - : undefined, - }; - } catch {} - } - - return null; -} - -function resolveOAuthProviders(agentDir?: string): UsageProviderId[] { - const store = ensureAuthProfileStore(agentDir, { - allowKeychainPrompt: false, - }); - const cfg = loadConfig(); - const providers = usageProviders.filter((provider) => provider !== "zai"); - const isOAuthLikeCredential = (id: string) => { - const cred = store.profiles[id]; - return cred?.type === "oauth" || cred?.type === "token"; - }; - return providers.filter((provider) => { - const profiles = listProfilesForProvider(store, provider).filter( - isOAuthLikeCredential, - ); - if (profiles.length > 0) return true; - const normalized = normalizeProviderId(provider); - const configuredProfiles = Object.entries(cfg.auth?.profiles ?? {}) - .filter( - ([, profile]) => normalizeProviderId(profile.provider) === normalized, - ) - .map(([id]) => id) - .filter(isOAuthLikeCredential); - return configuredProfiles.length > 0; - }); -} - -async function resolveProviderAuths( - opts: UsageSummaryOptions, -): Promise { - if (opts.auth) return opts.auth; - - const targetProviders = opts.providers ?? usageProviders; - const oauthProviders = resolveOAuthProviders(opts.agentDir); - const auths: ProviderAuth[] = []; - - for (const provider of targetProviders) { - if (provider === "zai") { - const apiKey = resolveZaiApiKey(); - if (apiKey) auths.push({ provider, token: apiKey }); - continue; - } - - if (!oauthProviders.includes(provider)) continue; - const auth = await resolveOAuthToken({ provider, agentDir: opts.agentDir }); - if (auth) auths.push(auth); - } - - return auths; -} - -export async function loadProviderUsageSummary( - opts: UsageSummaryOptions = {}, -): Promise { - const now = opts.now ?? Date.now(); - const timeoutMs = opts.timeoutMs ?? DEFAULT_TIMEOUT_MS; - const fetchFn = opts.fetch ?? fetch; - - const auths = await resolveProviderAuths(opts); - if (auths.length === 0) { - return { updatedAt: now, providers: [] }; - } - - const tasks = auths.map((auth) => - withTimeout( - (async (): Promise => { - switch (auth.provider) { - case "anthropic": - return await fetchClaudeUsage(auth.token, timeoutMs, fetchFn); - case "github-copilot": - return await fetchCopilotUsage(auth.token, timeoutMs, fetchFn); - case "google-gemini-cli": - case "google-antigravity": - return await fetchGeminiUsage( - auth.token, - timeoutMs, - fetchFn, - auth.provider, - ); - case "openai-codex": - return await fetchCodexUsage( - auth.token, - auth.accountId, - timeoutMs, - fetchFn, - ); - case "zai": - return await fetchZaiUsage(auth.token, timeoutMs, fetchFn); - default: - return { - provider: auth.provider, - displayName: PROVIDER_LABELS[auth.provider], - windows: [], - error: "Unsupported provider", - }; - } - })(), - timeoutMs + 1000, - { - provider: auth.provider, - displayName: PROVIDER_LABELS[auth.provider], - windows: [], - error: "Timeout", - }, - ), - ); - - const snapshots = await Promise.all(tasks); - const providers = snapshots.filter((entry) => { - if (entry.windows.length > 0) return true; - if (!entry.error) return true; - return !ignoredErrors.has(entry.error); - }); - - return { updatedAt: now, providers }; -} +export { + formatUsageReportLines, + formatUsageSummaryLine, +} from "./provider-usage.format.js"; +export { loadProviderUsageSummary } from "./provider-usage.load.js"; +export { resolveUsageProviderId } from "./provider-usage.shared.js"; +export type { + ProviderUsageSnapshot, + UsageProviderId, + UsageSummary, + UsageWindow, +} from "./provider-usage.types.js"; diff --git a/src/infra/provider-usage.types.ts b/src/infra/provider-usage.types.ts new file mode 100644 index 000000000..0fb18a794 --- /dev/null +++ b/src/infra/provider-usage.types.ts @@ -0,0 +1,26 @@ +export type UsageWindow = { + label: string; + usedPercent: number; + resetAt?: number; +}; + +export type ProviderUsageSnapshot = { + provider: UsageProviderId; + displayName: string; + windows: UsageWindow[]; + plan?: string; + error?: string; +}; + +export type UsageSummary = { + updatedAt: number; + providers: ProviderUsageSnapshot[]; +}; + +export type UsageProviderId = + | "anthropic" + | "github-copilot" + | "google-gemini-cli" + | "google-antigravity" + | "openai-codex" + | "zai"; diff --git a/src/memory/index.ts b/src/memory/index.ts index 042801940..66bbadbe5 100644 --- a/src/memory/index.ts +++ b/src/memory/index.ts @@ -1,520 +1,8 @@ -import fs from "node:fs/promises"; -import { createRequire } from "node:module"; -import path from "node:path"; - -import type { DatabaseSync } from "node:sqlite"; -import chokidar, { type FSWatcher } from "chokidar"; - -import { - resolveAgentDir, - resolveAgentWorkspaceDir, -} from "../agents/agent-scope.js"; -import type { ResolvedMemorySearchConfig } from "../agents/memory-search.js"; -import { resolveMemorySearchConfig } from "../agents/memory-search.js"; -import type { ClawdbotConfig } from "../config/config.js"; -import { resolveUserPath, truncateUtf16Safe } from "../utils.js"; -import { - createEmbeddingProvider, - type EmbeddingProvider, - type EmbeddingProviderResult, -} from "./embeddings.js"; -import { - buildFileEntry, - chunkMarkdown, - cosineSimilarity, - ensureDir, - hashText, - isMemoryPath, - listMemoryFiles, - type MemoryFileEntry, - normalizeRelPath, - parseEmbedding, -} from "./internal.js"; - -const require = createRequire(import.meta.url); - -function requireNodeSqlite(): typeof import("node:sqlite") { - const onWarning = (warning: Error & { name?: string; message?: string }) => { - if ( - warning.name === "ExperimentalWarning" && - warning.message?.includes("SQLite is an experimental feature") - ) { - return; - } - process.stderr.write(`${warning.stack ?? warning.toString()}\n`); - }; - - process.on("warning", onWarning); - try { - return require("node:sqlite") as typeof import("node:sqlite"); - } finally { - process.off("warning", onWarning); - } -} - -export type MemorySearchResult = { - path: string; - startLine: number; - endLine: number; - score: number; - snippet: string; -}; - -type MemoryIndexMeta = { - model: string; - provider: string; - chunkTokens: number; - chunkOverlap: number; -}; - -const META_KEY = "memory_index_meta_v1"; -const SNIPPET_MAX_CHARS = 700; - -const INDEX_CACHE = new Map(); - -export class MemoryIndexManager { - private readonly cacheKey: string; - private readonly cfg: ClawdbotConfig; - private readonly agentId: string; - private readonly workspaceDir: string; - private readonly settings: ResolvedMemorySearchConfig; - private readonly provider: EmbeddingProvider; - private readonly requestedProvider: "openai" | "local"; - private readonly fallbackReason?: string; - private readonly db: DatabaseSync; - private watcher: FSWatcher | null = null; - private watchTimer: NodeJS.Timeout | null = null; - private intervalTimer: NodeJS.Timeout | null = null; - private closed = false; - private dirty = false; - private sessionWarm = new Set(); - private syncing: Promise | null = null; - - static async get(params: { - cfg: ClawdbotConfig; - agentId: string; - }): Promise { - const { cfg, agentId } = params; - const settings = resolveMemorySearchConfig(cfg, agentId); - if (!settings) return null; - const workspaceDir = resolveAgentWorkspaceDir(cfg, agentId); - const key = `${agentId}:${workspaceDir}:${JSON.stringify(settings)}`; - const existing = INDEX_CACHE.get(key); - if (existing) return existing; - const providerResult = await createEmbeddingProvider({ - config: cfg, - agentDir: resolveAgentDir(cfg, agentId), - provider: settings.provider, - remote: settings.remote, - model: settings.model, - fallback: settings.fallback, - local: settings.local, - }); - const manager = new MemoryIndexManager({ - cacheKey: key, - cfg, - agentId, - workspaceDir, - settings, - providerResult, - }); - INDEX_CACHE.set(key, manager); - return manager; - } - - private constructor(params: { - cacheKey: string; - cfg: ClawdbotConfig; - agentId: string; - workspaceDir: string; - settings: ResolvedMemorySearchConfig; - providerResult: EmbeddingProviderResult; - }) { - this.cacheKey = params.cacheKey; - this.cfg = params.cfg; - this.agentId = params.agentId; - this.workspaceDir = params.workspaceDir; - this.settings = params.settings; - this.provider = params.providerResult.provider; - this.requestedProvider = params.providerResult.requestedProvider; - this.fallbackReason = params.providerResult.fallbackReason; - this.db = this.openDatabase(); - this.ensureSchema(); - this.ensureWatcher(); - this.ensureIntervalSync(); - this.dirty = true; - } - - async warmSession(sessionKey?: string): Promise { - if (!this.settings.sync.onSessionStart) return; - const key = sessionKey?.trim() || ""; - if (key && this.sessionWarm.has(key)) return; - await this.sync({ reason: "session-start" }); - if (key) this.sessionWarm.add(key); - } - - async search( - query: string, - opts?: { - maxResults?: number; - minScore?: number; - sessionKey?: string; - }, - ): Promise { - await this.warmSession(opts?.sessionKey); - if (this.settings.sync.onSearch && this.dirty) { - await this.sync({ reason: "search" }); - } - const cleaned = query.trim(); - if (!cleaned) return []; - const queryVec = await this.provider.embedQuery(cleaned); - if (queryVec.length === 0) return []; - const candidates = this.listChunks(); - const scored = candidates - .map((chunk) => ({ - chunk, - score: cosineSimilarity(queryVec, chunk.embedding), - })) - .filter((entry) => Number.isFinite(entry.score)); - const minScore = opts?.minScore ?? this.settings.query.minScore; - const maxResults = opts?.maxResults ?? this.settings.query.maxResults; - return scored - .filter((entry) => entry.score >= minScore) - .sort((a, b) => b.score - a.score) - .slice(0, maxResults) - .map((entry) => ({ - path: entry.chunk.path, - startLine: entry.chunk.startLine, - endLine: entry.chunk.endLine, - score: entry.score, - snippet: truncateUtf16Safe(entry.chunk.text, SNIPPET_MAX_CHARS), - })); - } - - async sync(params?: { reason?: string; force?: boolean }): Promise { - if (this.syncing) return this.syncing; - this.syncing = this.runSync(params).finally(() => { - this.syncing = null; - }); - return this.syncing; - } - - async readFile(params: { - relPath: string; - from?: number; - lines?: number; - }): Promise<{ text: string; path: string }> { - const relPath = normalizeRelPath(params.relPath); - if (!relPath || !isMemoryPath(relPath)) { - throw new Error("path required"); - } - const absPath = path.resolve(this.workspaceDir, relPath); - if (!absPath.startsWith(this.workspaceDir)) { - throw new Error("path escapes workspace"); - } - const content = await fs.readFile(absPath, "utf-8"); - if (!params.from && !params.lines) { - return { text: content, path: relPath }; - } - const lines = content.split("\n"); - const start = Math.max(1, params.from ?? 1); - const count = Math.max(1, params.lines ?? lines.length); - const slice = lines.slice(start - 1, start - 1 + count); - return { text: slice.join("\n"), path: relPath }; - } - - status(): { - files: number; - chunks: number; - dirty: boolean; - workspaceDir: string; - dbPath: string; - provider: string; - model: string; - requestedProvider: string; - fallback?: { from: string; reason?: string }; - } { - const files = this.db.prepare(`SELECT COUNT(*) as c FROM files`).get() as { - c: number; - }; - const chunks = this.db - .prepare(`SELECT COUNT(*) as c FROM chunks`) - .get() as { - c: number; - }; - return { - files: files?.c ?? 0, - chunks: chunks?.c ?? 0, - dirty: this.dirty, - workspaceDir: this.workspaceDir, - dbPath: this.settings.store.path, - provider: this.provider.id, - model: this.provider.model, - requestedProvider: this.requestedProvider, - fallback: this.fallbackReason - ? { from: "local", reason: this.fallbackReason } - : undefined, - }; - } - - async close(): Promise { - if (this.closed) return; - this.closed = true; - if (this.watchTimer) { - clearTimeout(this.watchTimer); - this.watchTimer = null; - } - if (this.intervalTimer) { - clearInterval(this.intervalTimer); - this.intervalTimer = null; - } - if (this.watcher) { - await this.watcher.close(); - this.watcher = null; - } - this.db.close(); - INDEX_CACHE.delete(this.cacheKey); - } - - private openDatabase(): DatabaseSync { - const dbPath = resolveUserPath(this.settings.store.path); - const dir = path.dirname(dbPath); - ensureDir(dir); - const { DatabaseSync } = requireNodeSqlite(); - return new DatabaseSync(dbPath); - } - - private ensureSchema() { - this.db.exec(` - CREATE TABLE IF NOT EXISTS meta ( - key TEXT PRIMARY KEY, - value TEXT NOT NULL - ); - `); - this.db.exec(` - CREATE TABLE IF NOT EXISTS files ( - path TEXT PRIMARY KEY, - hash TEXT NOT NULL, - mtime INTEGER NOT NULL, - size INTEGER NOT NULL - ); - `); - this.db.exec(` - CREATE TABLE IF NOT EXISTS chunks ( - id TEXT PRIMARY KEY, - path TEXT NOT NULL, - start_line INTEGER NOT NULL, - end_line INTEGER NOT NULL, - hash TEXT NOT NULL, - model TEXT NOT NULL, - text TEXT NOT NULL, - embedding TEXT NOT NULL, - updated_at INTEGER NOT NULL - ); - `); - this.db.exec(`CREATE INDEX IF NOT EXISTS idx_chunks_path ON chunks(path);`); - } - - private ensureWatcher() { - if (!this.settings.sync.watch || this.watcher) return; - const watchPaths = [ - path.join(this.workspaceDir, "MEMORY.md"), - path.join(this.workspaceDir, "memory"), - ]; - this.watcher = chokidar.watch(watchPaths, { - ignoreInitial: true, - awaitWriteFinish: { - stabilityThreshold: this.settings.sync.watchDebounceMs, - pollInterval: 100, - }, - }); - const markDirty = () => { - this.dirty = true; - this.scheduleWatchSync(); - }; - this.watcher.on("add", markDirty); - this.watcher.on("change", markDirty); - this.watcher.on("unlink", markDirty); - } - - private ensureIntervalSync() { - const minutes = this.settings.sync.intervalMinutes; - if (!minutes || minutes <= 0 || this.intervalTimer) return; - const ms = minutes * 60 * 1000; - this.intervalTimer = setInterval(() => { - void this.sync({ reason: "interval" }); - }, ms); - } - - private scheduleWatchSync() { - if (!this.settings.sync.watch) return; - if (this.watchTimer) clearTimeout(this.watchTimer); - this.watchTimer = setTimeout(() => { - this.watchTimer = null; - void this.sync({ reason: "watch" }); - }, this.settings.sync.watchDebounceMs); - } - - private listChunks(): Array<{ - path: string; - startLine: number; - endLine: number; - text: string; - embedding: number[]; - }> { - const rows = this.db - .prepare( - `SELECT path, start_line, end_line, text, embedding FROM chunks WHERE model = ?`, - ) - .all(this.provider.model) as Array<{ - path: string; - start_line: number; - end_line: number; - text: string; - embedding: string; - }>; - return rows.map((row) => ({ - path: row.path, - startLine: row.start_line, - endLine: row.end_line, - text: row.text, - embedding: parseEmbedding(row.embedding), - })); - } - - private async runSync(params?: { reason?: string; force?: boolean }) { - const meta = this.readMeta(); - const needsFullReindex = - params?.force || - !meta || - meta.model !== this.provider.model || - meta.provider !== this.provider.id || - meta.chunkTokens !== this.settings.chunking.tokens || - meta.chunkOverlap !== this.settings.chunking.overlap; - if (needsFullReindex) { - this.resetIndex(); - } - - const files = await listMemoryFiles(this.workspaceDir); - const fileEntries = await Promise.all( - files.map(async (file) => buildFileEntry(file, this.workspaceDir)), - ); - const activePaths = new Set(fileEntries.map((entry) => entry.path)); - - for (const entry of fileEntries) { - const record = this.db - .prepare(`SELECT hash FROM files WHERE path = ?`) - .get(entry.path) as { hash: string } | undefined; - if (!needsFullReindex && record?.hash === entry.hash) { - continue; - } - await this.indexFile(entry); - } - - const staleRows = this.db.prepare(`SELECT path FROM files`).all() as Array<{ - path: string; - }>; - for (const stale of staleRows) { - if (activePaths.has(stale.path)) continue; - this.db.prepare(`DELETE FROM files WHERE path = ?`).run(stale.path); - this.db.prepare(`DELETE FROM chunks WHERE path = ?`).run(stale.path); - } - - this.writeMeta({ - model: this.provider.model, - provider: this.provider.id, - chunkTokens: this.settings.chunking.tokens, - chunkOverlap: this.settings.chunking.overlap, - }); - this.dirty = false; - } - - private resetIndex() { - this.db.exec(`DELETE FROM files`); - this.db.exec(`DELETE FROM chunks`); - } - - private readMeta(): MemoryIndexMeta | null { - const row = this.db - .prepare(`SELECT value FROM meta WHERE key = ?`) - .get(META_KEY) as { value: string } | undefined; - if (!row?.value) return null; - try { - return JSON.parse(row.value) as MemoryIndexMeta; - } catch { - return null; - } - } - - private writeMeta(meta: MemoryIndexMeta) { - const value = JSON.stringify(meta); - this.db - .prepare( - `INSERT INTO meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value=excluded.value`, - ) - .run(META_KEY, value); - } - - private async indexFile(entry: MemoryFileEntry) { - const content = await fs.readFile(entry.absPath, "utf-8"); - const chunks = chunkMarkdown(content, this.settings.chunking); - const embeddings = await this.provider.embedBatch( - chunks.map((chunk) => chunk.text), - ); - const now = Date.now(); - this.db.prepare(`DELETE FROM chunks WHERE path = ?`).run(entry.path); - for (let i = 0; i < chunks.length; i++) { - const chunk = chunks[i]; - const embedding = embeddings[i] ?? []; - const id = hashText( - `${entry.path}:${chunk.startLine}:${chunk.endLine}:${chunk.hash}:${this.provider.model}`, - ); - this.db - .prepare( - `INSERT INTO chunks (id, path, start_line, end_line, hash, model, text, embedding, updated_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) - ON CONFLICT(id) DO UPDATE SET - hash=excluded.hash, - model=excluded.model, - text=excluded.text, - embedding=excluded.embedding, - updated_at=excluded.updated_at`, - ) - .run( - id, - entry.path, - chunk.startLine, - chunk.endLine, - chunk.hash, - this.provider.model, - chunk.text, - JSON.stringify(embedding), - now, - ); - } - this.db - .prepare( - `INSERT INTO files (path, hash, mtime, size) VALUES (?, ?, ?, ?) - ON CONFLICT(path) DO UPDATE SET hash=excluded.hash, mtime=excluded.mtime, size=excluded.size`, - ) - .run(entry.path, entry.hash, entry.mtimeMs, entry.size); - } -} - -export type MemorySearchManagerResult = { - manager: MemoryIndexManager | null; - error?: string; -}; - -export async function getMemorySearchManager(params: { - cfg: ClawdbotConfig; - agentId: string; -}): Promise { - try { - const manager = await MemoryIndexManager.get(params); - return { manager }; - } catch (err) { - const message = err instanceof Error ? err.message : String(err); - return { manager: null, error: message }; - } -} +export { + MemoryIndexManager, + type MemorySearchResult, +} from "./manager.js"; +export { + getMemorySearchManager, + type MemorySearchManagerResult, +} from "./search-manager.js"; diff --git a/src/memory/manager.ts b/src/memory/manager.ts new file mode 100644 index 000000000..0cce53e76 --- /dev/null +++ b/src/memory/manager.ts @@ -0,0 +1,481 @@ +import fs from "node:fs/promises"; +import path from "node:path"; + +import type { DatabaseSync } from "node:sqlite"; +import chokidar, { type FSWatcher } from "chokidar"; + +import { + resolveAgentDir, + resolveAgentWorkspaceDir, +} from "../agents/agent-scope.js"; +import type { ResolvedMemorySearchConfig } from "../agents/memory-search.js"; +import { resolveMemorySearchConfig } from "../agents/memory-search.js"; +import type { ClawdbotConfig } from "../config/config.js"; +import { resolveUserPath, truncateUtf16Safe } from "../utils.js"; +import { + createEmbeddingProvider, + type EmbeddingProvider, + type EmbeddingProviderResult, +} from "./embeddings.js"; +import { + buildFileEntry, + chunkMarkdown, + cosineSimilarity, + ensureDir, + hashText, + isMemoryPath, + listMemoryFiles, + type MemoryFileEntry, + normalizeRelPath, + parseEmbedding, +} from "./internal.js"; +import { requireNodeSqlite } from "./sqlite.js"; + +export type MemorySearchResult = { + path: string; + startLine: number; + endLine: number; + score: number; + snippet: string; +}; + +type MemoryIndexMeta = { + model: string; + provider: string; + chunkTokens: number; + chunkOverlap: number; +}; + +const META_KEY = "memory_index_meta_v1"; +const SNIPPET_MAX_CHARS = 700; + +const INDEX_CACHE = new Map(); + +export class MemoryIndexManager { + private readonly cacheKey: string; + private readonly cfg: ClawdbotConfig; + private readonly agentId: string; + private readonly workspaceDir: string; + private readonly settings: ResolvedMemorySearchConfig; + private readonly provider: EmbeddingProvider; + private readonly requestedProvider: "openai" | "local"; + private readonly fallbackReason?: string; + private readonly db: DatabaseSync; + private watcher: FSWatcher | null = null; + private watchTimer: NodeJS.Timeout | null = null; + private intervalTimer: NodeJS.Timeout | null = null; + private closed = false; + private dirty = false; + private sessionWarm = new Set(); + private syncing: Promise | null = null; + + static async get(params: { + cfg: ClawdbotConfig; + agentId: string; + }): Promise { + const { cfg, agentId } = params; + const settings = resolveMemorySearchConfig(cfg, agentId); + if (!settings) return null; + const workspaceDir = resolveAgentWorkspaceDir(cfg, agentId); + const key = `${agentId}:${workspaceDir}:${JSON.stringify(settings)}`; + const existing = INDEX_CACHE.get(key); + if (existing) return existing; + const providerResult = await createEmbeddingProvider({ + config: cfg, + agentDir: resolveAgentDir(cfg, agentId), + provider: settings.provider, + remote: settings.remote, + model: settings.model, + fallback: settings.fallback, + local: settings.local, + }); + const manager = new MemoryIndexManager({ + cacheKey: key, + cfg, + agentId, + workspaceDir, + settings, + providerResult, + }); + INDEX_CACHE.set(key, manager); + return manager; + } + + private constructor(params: { + cacheKey: string; + cfg: ClawdbotConfig; + agentId: string; + workspaceDir: string; + settings: ResolvedMemorySearchConfig; + providerResult: EmbeddingProviderResult; + }) { + this.cacheKey = params.cacheKey; + this.cfg = params.cfg; + this.agentId = params.agentId; + this.workspaceDir = params.workspaceDir; + this.settings = params.settings; + this.provider = params.providerResult.provider; + this.requestedProvider = params.providerResult.requestedProvider; + this.fallbackReason = params.providerResult.fallbackReason; + this.db = this.openDatabase(); + this.ensureSchema(); + this.ensureWatcher(); + this.ensureIntervalSync(); + this.dirty = true; + } + + async warmSession(sessionKey?: string): Promise { + if (!this.settings.sync.onSessionStart) return; + const key = sessionKey?.trim() || ""; + if (key && this.sessionWarm.has(key)) return; + await this.sync({ reason: "session-start" }); + if (key) this.sessionWarm.add(key); + } + + async search( + query: string, + opts?: { + maxResults?: number; + minScore?: number; + sessionKey?: string; + }, + ): Promise { + await this.warmSession(opts?.sessionKey); + if (this.settings.sync.onSearch && this.dirty) { + await this.sync({ reason: "search" }); + } + const cleaned = query.trim(); + if (!cleaned) return []; + const queryVec = await this.provider.embedQuery(cleaned); + if (queryVec.length === 0) return []; + const candidates = this.listChunks(); + const scored = candidates + .map((chunk) => ({ + chunk, + score: cosineSimilarity(queryVec, chunk.embedding), + })) + .filter((entry) => Number.isFinite(entry.score)); + const minScore = opts?.minScore ?? this.settings.query.minScore; + const maxResults = opts?.maxResults ?? this.settings.query.maxResults; + return scored + .filter((entry) => entry.score >= minScore) + .sort((a, b) => b.score - a.score) + .slice(0, maxResults) + .map((entry) => ({ + path: entry.chunk.path, + startLine: entry.chunk.startLine, + endLine: entry.chunk.endLine, + score: entry.score, + snippet: truncateUtf16Safe(entry.chunk.text, SNIPPET_MAX_CHARS), + })); + } + + async sync(params?: { reason?: string; force?: boolean }): Promise { + if (this.syncing) return this.syncing; + this.syncing = this.runSync(params).finally(() => { + this.syncing = null; + }); + return this.syncing; + } + + async readFile(params: { + relPath: string; + from?: number; + lines?: number; + }): Promise<{ text: string; path: string }> { + const relPath = normalizeRelPath(params.relPath); + if (!relPath || !isMemoryPath(relPath)) { + throw new Error("path required"); + } + const absPath = path.resolve(this.workspaceDir, relPath); + if (!absPath.startsWith(this.workspaceDir)) { + throw new Error("path escapes workspace"); + } + const content = await fs.readFile(absPath, "utf-8"); + if (!params.from && !params.lines) { + return { text: content, path: relPath }; + } + const lines = content.split("\n"); + const start = Math.max(1, params.from ?? 1); + const count = Math.max(1, params.lines ?? lines.length); + const slice = lines.slice(start - 1, start - 1 + count); + return { text: slice.join("\n"), path: relPath }; + } + + status(): { + files: number; + chunks: number; + dirty: boolean; + workspaceDir: string; + dbPath: string; + provider: string; + model: string; + requestedProvider: string; + fallback?: { from: string; reason?: string }; + } { + const files = this.db.prepare(`SELECT COUNT(*) as c FROM files`).get() as { + c: number; + }; + const chunks = this.db + .prepare(`SELECT COUNT(*) as c FROM chunks`) + .get() as { + c: number; + }; + return { + files: files?.c ?? 0, + chunks: chunks?.c ?? 0, + dirty: this.dirty, + workspaceDir: this.workspaceDir, + dbPath: this.settings.store.path, + provider: this.provider.id, + model: this.provider.model, + requestedProvider: this.requestedProvider, + fallback: this.fallbackReason + ? { from: "local", reason: this.fallbackReason } + : undefined, + }; + } + + async close(): Promise { + if (this.closed) return; + this.closed = true; + if (this.watchTimer) { + clearTimeout(this.watchTimer); + this.watchTimer = null; + } + if (this.intervalTimer) { + clearInterval(this.intervalTimer); + this.intervalTimer = null; + } + if (this.watcher) { + await this.watcher.close(); + this.watcher = null; + } + this.db.close(); + INDEX_CACHE.delete(this.cacheKey); + } + + private openDatabase(): DatabaseSync { + const dbPath = resolveUserPath(this.settings.store.path); + const dir = path.dirname(dbPath); + ensureDir(dir); + const { DatabaseSync } = requireNodeSqlite(); + return new DatabaseSync(dbPath); + } + + private ensureSchema() { + this.db.exec(` + CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + `); + this.db.exec(` + CREATE TABLE IF NOT EXISTS files ( + path TEXT PRIMARY KEY, + hash TEXT NOT NULL, + mtime INTEGER NOT NULL, + size INTEGER NOT NULL + ); + `); + this.db.exec(` + CREATE TABLE IF NOT EXISTS chunks ( + id TEXT PRIMARY KEY, + path TEXT NOT NULL, + start_line INTEGER NOT NULL, + end_line INTEGER NOT NULL, + hash TEXT NOT NULL, + model TEXT NOT NULL, + text TEXT NOT NULL, + embedding TEXT NOT NULL, + updated_at INTEGER NOT NULL + ); + `); + this.db.exec(`CREATE INDEX IF NOT EXISTS idx_chunks_path ON chunks(path);`); + } + + private ensureWatcher() { + if (!this.settings.sync.watch || this.watcher) return; + const watchPaths = [ + path.join(this.workspaceDir, "MEMORY.md"), + path.join(this.workspaceDir, "memory"), + ]; + this.watcher = chokidar.watch(watchPaths, { + ignoreInitial: true, + awaitWriteFinish: { + stabilityThreshold: this.settings.sync.watchDebounceMs, + pollInterval: 100, + }, + }); + const markDirty = () => { + this.dirty = true; + this.scheduleWatchSync(); + }; + this.watcher.on("add", markDirty); + this.watcher.on("change", markDirty); + this.watcher.on("unlink", markDirty); + } + + private ensureIntervalSync() { + const minutes = this.settings.sync.intervalMinutes; + if (!minutes || minutes <= 0 || this.intervalTimer) return; + const ms = minutes * 60 * 1000; + this.intervalTimer = setInterval(() => { + void this.sync({ reason: "interval" }); + }, ms); + } + + private scheduleWatchSync() { + if (!this.settings.sync.watch) return; + if (this.watchTimer) clearTimeout(this.watchTimer); + this.watchTimer = setTimeout(() => { + this.watchTimer = null; + void this.sync({ reason: "watch" }); + }, this.settings.sync.watchDebounceMs); + } + + private listChunks(): Array<{ + path: string; + startLine: number; + endLine: number; + text: string; + embedding: number[]; + }> { + const rows = this.db + .prepare( + `SELECT path, start_line, end_line, text, embedding FROM chunks WHERE model = ?`, + ) + .all(this.provider.model) as Array<{ + path: string; + start_line: number; + end_line: number; + text: string; + embedding: string; + }>; + return rows.map((row) => ({ + path: row.path, + startLine: row.start_line, + endLine: row.end_line, + text: row.text, + embedding: parseEmbedding(row.embedding), + })); + } + + private async runSync(params?: { reason?: string; force?: boolean }) { + const meta = this.readMeta(); + const needsFullReindex = + params?.force || + !meta || + meta.model !== this.provider.model || + meta.provider !== this.provider.id || + meta.chunkTokens !== this.settings.chunking.tokens || + meta.chunkOverlap !== this.settings.chunking.overlap; + if (needsFullReindex) { + this.resetIndex(); + } + + const files = await listMemoryFiles(this.workspaceDir); + const fileEntries = await Promise.all( + files.map(async (file) => buildFileEntry(file, this.workspaceDir)), + ); + const activePaths = new Set(fileEntries.map((entry) => entry.path)); + + for (const entry of fileEntries) { + const record = this.db + .prepare(`SELECT hash FROM files WHERE path = ?`) + .get(entry.path) as { hash: string } | undefined; + if (!needsFullReindex && record?.hash === entry.hash) { + continue; + } + await this.indexFile(entry); + } + + const staleRows = this.db.prepare(`SELECT path FROM files`).all() as Array<{ + path: string; + }>; + for (const stale of staleRows) { + if (activePaths.has(stale.path)) continue; + this.db.prepare(`DELETE FROM files WHERE path = ?`).run(stale.path); + this.db.prepare(`DELETE FROM chunks WHERE path = ?`).run(stale.path); + } + + this.writeMeta({ + model: this.provider.model, + provider: this.provider.id, + chunkTokens: this.settings.chunking.tokens, + chunkOverlap: this.settings.chunking.overlap, + }); + this.dirty = false; + } + + private resetIndex() { + this.db.exec(`DELETE FROM files`); + this.db.exec(`DELETE FROM chunks`); + } + + private readMeta(): MemoryIndexMeta | null { + const row = this.db + .prepare(`SELECT value FROM meta WHERE key = ?`) + .get(META_KEY) as { value: string } | undefined; + if (!row?.value) return null; + try { + return JSON.parse(row.value) as MemoryIndexMeta; + } catch { + return null; + } + } + + private writeMeta(meta: MemoryIndexMeta) { + const value = JSON.stringify(meta); + this.db + .prepare( + `INSERT INTO meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value=excluded.value`, + ) + .run(META_KEY, value); + } + + private async indexFile(entry: MemoryFileEntry) { + const content = await fs.readFile(entry.absPath, "utf-8"); + const chunks = chunkMarkdown(content, this.settings.chunking); + const embeddings = await this.provider.embedBatch( + chunks.map((chunk) => chunk.text), + ); + const now = Date.now(); + this.db.prepare(`DELETE FROM chunks WHERE path = ?`).run(entry.path); + for (let i = 0; i < chunks.length; i++) { + const chunk = chunks[i]; + const embedding = embeddings[i] ?? []; + const id = hashText( + `${entry.path}:${chunk.startLine}:${chunk.endLine}:${chunk.hash}:${this.provider.model}`, + ); + this.db + .prepare( + `INSERT INTO chunks (id, path, start_line, end_line, hash, model, text, embedding, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + hash=excluded.hash, + model=excluded.model, + text=excluded.text, + embedding=excluded.embedding, + updated_at=excluded.updated_at`, + ) + .run( + id, + entry.path, + chunk.startLine, + chunk.endLine, + chunk.hash, + this.provider.model, + chunk.text, + JSON.stringify(embedding), + now, + ); + } + this.db + .prepare( + `INSERT INTO files (path, hash, mtime, size) VALUES (?, ?, ?, ?) + ON CONFLICT(path) DO UPDATE SET hash=excluded.hash, mtime=excluded.mtime, size=excluded.size`, + ) + .run(entry.path, entry.hash, entry.mtimeMs, entry.size); + } +} diff --git a/src/memory/search-manager.ts b/src/memory/search-manager.ts new file mode 100644 index 000000000..5ec0d11bc --- /dev/null +++ b/src/memory/search-manager.ts @@ -0,0 +1,20 @@ +import type { ClawdbotConfig } from "../config/config.js"; +import { MemoryIndexManager } from "./manager.js"; + +export type MemorySearchManagerResult = { + manager: MemoryIndexManager | null; + error?: string; +}; + +export async function getMemorySearchManager(params: { + cfg: ClawdbotConfig; + agentId: string; +}): Promise { + try { + const manager = await MemoryIndexManager.get(params); + return { manager }; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + return { manager: null, error: message }; + } +} diff --git a/src/memory/sqlite.ts b/src/memory/sqlite.ts new file mode 100644 index 000000000..0680259fe --- /dev/null +++ b/src/memory/sqlite.ts @@ -0,0 +1,22 @@ +import { createRequire } from "node:module"; + +const require = createRequire(import.meta.url); + +export function requireNodeSqlite(): typeof import("node:sqlite") { + const onWarning = (warning: Error & { name?: string; message?: string }) => { + if ( + warning.name === "ExperimentalWarning" && + warning.message?.includes("SQLite is an experimental feature") + ) { + return; + } + process.stderr.write(`${warning.stack ?? warning.toString()}\n`); + }; + + process.on("warning", onWarning); + try { + return require("node:sqlite") as typeof import("node:sqlite"); + } finally { + process.off("warning", onWarning); + } +}