refactor(infra): split provider usage
This commit is contained in:
Peter Steinberger
195
src/infra/provider-usage.auth.ts
Normal file
195
src/infra/provider-usage.auth.ts
Normal 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
196
src/infra/provider-usage.fetch.claude.ts
Normal 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
102
src/infra/provider-usage.fetch.codex.ts
Normal 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
70
src/infra/provider-usage.fetch.copilot.ts
Normal 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
81
src/infra/provider-usage.fetch.gemini.ts
Normal 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
14
src/infra/provider-usage.fetch.shared.ts
Normal 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
5
src/infra/provider-usage.fetch.ts
Normal 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
97
src/infra/provider-usage.fetch.zai.ts
Normal 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
91
src/infra/provider-usage.format.ts
Normal 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
102
src/infra/provider-usage.load.ts
Normal 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
61
src/infra/provider-usage.shared.ts
Normal 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);
|
||||||
|
}
|
||||||
|
};
|
||||||
@@ -1,964 +1,12 @@
|
|||||||
import fs from "node:fs";
|
export {
|
||||||
import os from "node:os";
|
formatUsageReportLines,
|
||||||
import path from "node:path";
|
formatUsageSummaryLine,
|
||||||
|
} from "./provider-usage.format.js";
|
||||||
import {
|
export { loadProviderUsageSummary } from "./provider-usage.load.js";
|
||||||
CLAUDE_CLI_PROFILE_ID,
|
export { resolveUsageProviderId } from "./provider-usage.shared.js";
|
||||||
ensureAuthProfileStore,
|
export type {
|
||||||
listProfilesForProvider,
|
ProviderUsageSnapshot,
|
||||||
resolveApiKeyForProfile,
|
UsageProviderId,
|
||||||
resolveAuthProfileOrder,
|
UsageSummary,
|
||||||
} from "../agents/auth-profiles.js";
|
UsageWindow,
|
||||||
import {
|
} from "./provider-usage.types.js";
|
||||||
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 };
|
|
||||||
}
|
|
||||||
|
|||||||
26
src/infra/provider-usage.types.ts
Normal file
26
src/infra/provider-usage.types.ts
Normal 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";
|
||||||
@@ -1,520 +1,8 @@
|
|||||||
import fs from "node:fs/promises";
|
export {
|
||||||
import { createRequire } from "node:module";
|
MemoryIndexManager,
|
||||||
import path from "node:path";
|
type MemorySearchResult,
|
||||||
|
} from "./manager.js";
|
||||||
import type { DatabaseSync } from "node:sqlite";
|
export {
|
||||||
import chokidar, { type FSWatcher } from "chokidar";
|
getMemorySearchManager,
|
||||||
|
type MemorySearchManagerResult,
|
||||||
import {
|
} from "./search-manager.js";
|
||||||
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 };
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|||||||
481
src/memory/manager.ts
Normal file
481
src/memory/manager.ts
Normal 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
20
src/memory/search-manager.ts
Normal 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
22
src/memory/sqlite.ts
Normal 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user