let5see/clawdbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(infra): split provider usage

Browse Source
This commit is contained in:
Peter Steinberger
2026-01-14 05:40:03 +00:00
parent e2f8909982
commit 3e0e608110
21 changed files with 1583 additions and 1484 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
src/infra/bridge/server.part-1.test.ts → src/infra/bridge/server.enables-keepalive-sockets.test.ts
View File

0
src/infra/bridge/server.part-2.test.ts → src/infra/bridge/server.supports-invoke-roundtrip-connected-node.test.ts
View File

0
src/infra/heartbeat-runner.part-2.test.ts → src/infra/heartbeat-runner.respects-ackmaxchars-heartbeat-acks.test.ts
View File

0
src/infra/heartbeat-runner.part-1.test.ts → src/infra/heartbeat-runner.returns-default-unset.test.ts
View File

195
src/infra/provider-usage.auth.ts Normal file
View File

@@ -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<ProviderAuth | null> {
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<ProviderAuth[]> {
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;
}

196
src/infra/provider-usage.fetch.claude.ts Normal file
View File

@@ -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<ProviderUsageSnapshot | null> {
const headers: Record<string, string> = {
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<ProviderUsageSnapshot> {
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,
};
}

102
src/infra/provider-usage.fetch.codex.ts Normal file
View File

@@ -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<ProviderUsageSnapshot> {
const headers: Record<string, string> = {
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,
};
}

70
src/infra/provider-usage.fetch.copilot.ts Normal file
View File

@@ -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<ProviderUsageSnapshot> {
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,
};
}

81
src/infra/provider-usage.fetch.gemini.ts Normal file
View File

@@ -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<ProviderUsageSnapshot> {
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<string, number> = {};
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 };
}

14
src/infra/provider-usage.fetch.shared.ts Normal file
View File

@@ -0,0 +1,14 @@
export async function fetchJson(
url: string,
init: RequestInit,
timeoutMs: number,
fetchFn: typeof fetch,
): Promise<Response> {
const controller = new AbortController();
const timer = setTimeout(() => controller.abort(), timeoutMs);
try {
return await fetchFn(url, { ...init, signal: controller.signal });
} finally {
clearTimeout(timer);
}
}

5
src/infra/provider-usage.fetch.ts Normal file
View File

@@ -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";

97
src/infra/provider-usage.fetch.zai.ts Normal file
View File

@@ -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<ProviderUsageSnapshot> {
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,
};
}

91
src/infra/provider-usage.format.ts Normal file
View File

@@ -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;
}

102
src/infra/provider-usage.load.ts Normal file
View File

@@ -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<UsageSummary> {
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<ProviderUsageSnapshot> => {
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 };
}

61
src/infra/provider-usage.shared.ts Normal file
View File

@@ -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<UsageProviderId, string> = {
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 <T>(
work: Promise<T>,
ms: number,
fallback: T,
): Promise<T> => {
let timeout: NodeJS.Timeout | undefined;
try {
return await Promise.race([
work,
new Promise<T>((resolve) => {
timeout = setTimeout(() => resolve(fallback), ms);
}),
]);
} finally {
if (timeout) clearTimeout(timeout);
}
};

976
src/infra/provider-usage.ts
View File

@@ -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<UsageProviderId, string> = {
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 <T>(
work: Promise<T>,
ms: number,
fallback: T,
): Promise<T> => {
let timeout: NodeJS.Timeout | undefined;
try {
return await Promise.race([
work,
new Promise<T>((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<Response> {
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<ProviderUsageSnapshot> {
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<ProviderUsageSnapshot | null> {
const headers: Record<string, string> = {
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<ProviderUsageSnapshot> {
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<ProviderUsageSnapshot> {
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<string, number> = {};
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<ProviderUsageSnapshot> {
const headers: Record<string, string> = {
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<ProviderUsageSnapshot> {
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<ProviderAuth | null> {
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<ProviderAuth[]> {
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<UsageSummary> {
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<ProviderUsageSnapshot> => {
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";

26
src/infra/provider-usage.types.ts Normal file
View File

@@ -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";

528
src/memory/index.ts
View File

@@ -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<string, MemoryIndexManager>();
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<string>();
private syncing: Promise<void> | null = null;
static async get(params: {
cfg: ClawdbotConfig;
agentId: string;
}): Promise<MemoryIndexManager | null> {
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<void> {
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<MemorySearchResult[]> {
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<void> {
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<void> {
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<MemorySearchManagerResult> {
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";

481
src/memory/manager.ts Normal file
View File

@@ -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<string, MemoryIndexManager>();
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<string>();
private syncing: Promise<void> | null = null;
static async get(params: {
cfg: ClawdbotConfig;
agentId: string;
}): Promise<MemoryIndexManager | null> {
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<void> {
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<MemorySearchResult[]> {
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<void> {
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<void> {
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);
}
}

20
src/memory/search-manager.ts Normal file
View File

@@ -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<MemorySearchManagerResult> {
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 };
}
}

22
src/memory/sqlite.ts Normal file
View File

@@ -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);
}
}