@jxgame2020/dsh-token-quota 0.1.34 → 0.1.36

package/lib/index.js

@@ -23,6 +23,10 @@ import { LlmError } from "@deepseek-ai/dsh-llm";
 const TOKEN_QUOTA_NAMESPACE = "token-quota";
 /** Stable machine code thrown from `agent/request` when a model is over its daily cap. */
 const TOKEN_QUOTA_EXCEEDED_CODE = "TOKEN_QUOTA_EXCEEDED";
+const TOKEN_QUOTA_DEFAULT_BALANCE = {
+	enabled: true,
+	pollMinutes: 5
+};
 /** Default music settings: silent until the user flips the panel switch. */
 const TOKEN_QUOTA_DEFAULT_MUSIC = {
 	enabled: false,
@@ -236,7 +240,11 @@ const TOKEN_QUOTA_SETTINGS_SCHEMA = z.object({
 			"pentatonic"
 		]).default(TOKEN_QUOTA_DEFAULT_MUSIC.style),
 		onlyCurrentSession: z.boolean().default(TOKEN_QUOTA_DEFAULT_MUSIC.onlyCurrentSession)
-	}).default({ ...TOKEN_QUOTA_DEFAULT_MUSIC })
+	}).default({ ...TOKEN_QUOTA_DEFAULT_MUSIC }),
+	balance: z.object({
+		enabled: z.boolean().default(TOKEN_QUOTA_DEFAULT_BALANCE.enabled),
+		pollMinutes: z.number().min(1).max(120).default(TOKEN_QUOTA_DEFAULT_BALANCE.pollMinutes)
+	}).default({ ...TOKEN_QUOTA_DEFAULT_BALANCE })
 });
 /** Two-digit zero-pad helper. */
 function pad2(value) {
@@ -274,6 +282,17 @@ function usageTokens(usage) {
 * Mount it beside the other rows (`@jxgame2020/dsh-token-quota`) and write
 * per-model limits through the `token-quota` settings namespace.
 */
+/**
+* Providers the host can query for account balance, keyed by route key. Each
+* entry names the credential reference (environment-variable name), the
+* balance endpoint base, and the optional `*_BASE_URL` override variable.
+* Directory providers outside this map render as "not integrated for query".
+*/
+const BALANCE_SOURCES = { deepseek: {
+	keyEnv: "DEEPSEEK_API_KEY",
+	baseURL: "https://api.deepseek.com",
+	baseEnv: "DEEPSEEK_BASE_URL"
+} };
 var TokenQuotaService = class TokenQuotaService extends Service {
 	static Config = z.object({ storagePath: z.string().default("") });
 	storagePath;
@@ -295,7 +314,8 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 		dimWhenIdle: false,
 		music: { ...TOKEN_QUOTA_DEFAULT_MUSIC },
 		order: [],
-		notes: []
+		notes: [],
+		balance: { ...TOKEN_QUOTA_DEFAULT_BALANCE }
 	});
 	/** Whether update checks are enabled (mirrors the settings document). */
 	checkUpdates = true;
@@ -305,6 +325,14 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 	upgradeError = null;
 	/** Timer for the periodic update check. */
 	updateTimer;
+	/** Whether the account balance is enabled (mirrors the settings document). */
+	balanceEnabled = TOKEN_QUOTA_DEFAULT_BALANCE.enabled;
+	/** Cached per-provider account balances, served in the snapshot. */
+	balances = /* @__PURE__ */ new Map();
+	/** Timer for the periodic balance poll. */
+	balanceTimer;
+	/** In-flight balance refreshes by provider, to coalesce overlapping triggers. */
+	balanceFetches = /* @__PURE__ */ new Map();
 	/** In-flight update check, to coalesce the periodic and manual triggers. */
 	upgradeCheck;
 	/** Per-session folded model key from the latest `request/header`. */
@@ -323,6 +351,8 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 	disposeRouteCheck;
 	/** Disposer for the optional clear-log route (`POST /token-quota/clear-log`). */
 	disposeRouteClear;
+	/** Disposer for the optional balance-refresh route (`POST /token-quota/refresh-balance`). */
+	disposeRouteBalance;
 	/** Disposer for the optional music-event route (`GET /token-quota/events`, SSE). */
 	disposeRouteEvents;
 	/** Live SSE subscribers receiving streamed music actions. */
@@ -349,6 +379,10 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 							checkUpdates: doc.checkUpdates ?? true,
 							dimWhenIdle: doc.dimWhenIdle ?? false,
 							order: Array.isArray(doc.order) ? doc.order : [],
+							balance: {
+								enabled: doc.balance?.enabled ?? TOKEN_QUOTA_DEFAULT_BALANCE.enabled,
+								pollMinutes: typeof doc.balance?.pollMinutes === "number" ? Math.max(1, Math.min(120, Math.round(doc.balance.pollMinutes))) : TOKEN_QUOTA_DEFAULT_BALANCE.pollMinutes
+							},
 							notes: Array.isArray(doc.notes) ? doc.notes.slice(0, 5).flatMap((note) => {
 								if (typeof note.id !== "string") return [];
 								return [{
@@ -385,6 +419,12 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 					this.monitored = doc.monitored !== void 0 && doc.monitored.length > 0 ? new Set(doc.monitored) : void 0;
 					this.checkUpdates = doc.checkUpdates;
 					this.syncUpdateChecking();
+					const balanceEnabled = doc.balance?.enabled ?? TOKEN_QUOTA_DEFAULT_BALANCE.enabled;
+					const turnedOn = balanceEnabled && !this.balanceEnabled;
+					this.balanceEnabled = balanceEnabled;
+					if (!balanceEnabled) this.balances.clear();
+					this.syncBalancePolling();
+					if (turnedOn) this.refreshAllBalances();
 					if (doc.reset !== void 0 && typeof doc.reset === "object" && doc.reset !== null && typeof doc.reset.offsetHours === "number" && typeof doc.reset.hour === "number" && typeof doc.reset.minute === "number") {
 						const configured = doc.reset;
 						this.reset = {
@@ -415,6 +455,8 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 			agent.ctx.on("agent/request", async (payload, next) => this.onRequest(payload, next), { prepend: true });
 		});
 		this.syncUpdateChecking();
+		this.syncBalancePolling();
+		if (this.balanceEnabled) this.refreshAllBalances();
 		ctx.effect(() => () => {
 			this.disposeLocal();
 		}, "token-quota: flush on unload");
@@ -461,6 +503,39 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 				});
 			}
 		});
+		this.disposeRouteBalance = server.register({
+			kind: "exact",
+			path: "/token-quota/refresh-balance",
+			handler: (req, res) => {
+				if (req.method !== "POST") {
+					res.writeHead(405);
+					res.end();
+					return;
+				}
+				const chunks = [];
+				req.on("data", (chunk) => {
+					chunks.push(Buffer.from(chunk));
+				});
+				req.on("end", () => {
+					let provider;
+					try {
+						const raw = Buffer.concat(chunks).toString("utf8");
+						if (raw !== "") {
+							const parsed = JSON.parse(raw);
+							if (typeof parsed === "object" && parsed !== null && typeof parsed.provider === "string") provider = parsed.provider;
+						}
+					} catch {}
+					(provider !== void 0 && provider in BALANCE_SOURCES ? this.refreshProviderBalance(provider) : this.refreshAllBalances()).finally(() => {
+						const body = JSON.stringify(this.readSnapshot());
+						res.writeHead(200, {
+							"content-type": "application/json",
+							"cache-control": "no-store"
+						});
+						res.end(body);
+					});
+				});
+			}
+		});
 		this.disposeRouteClear = server.register({
 			kind: "exact",
 			path: "/token-quota/clear-log",
@@ -544,7 +619,7 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 		return this.upgradeCheck;
 	}
 	async doRefreshUpgrade() {
-		const current = "0.1.34";
+		const current = "0.1.36";
 		let latest;
 		try {
 			const response = await fetch(REGISTRY_LATEST_URL, {
@@ -595,6 +670,119 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 			this.syncUpdateChecking();
 		}, UPDATE_CHECK_INTERVAL_MS);
 	}
+	/**
+	* Start or stop the periodic balance poll to match the current setting. The
+	* interval comes from the settings document; a poll chain re-arms itself
+	* after each refresh. No fetch happens here — callers decide (startup and
+	* the enabled-transition in `onChange` fetch immediately).
+	*/
+	syncBalancePolling() {
+		if (this.balanceTimer !== void 0) {
+			clearTimeout(this.balanceTimer);
+			this.balanceTimer = void 0;
+		}
+		if (!this.balanceEnabled) return;
+		const minutes = this.settingsSource().balance?.pollMinutes ?? TOKEN_QUOTA_DEFAULT_BALANCE.pollMinutes;
+		const schedule = () => {
+			if (!this.balanceEnabled) return;
+			this.balanceTimer = setTimeout(() => {
+				this.balanceTimer = void 0;
+				this.refreshAllBalances().finally(schedule);
+			}, Math.max(6e4, minutes * 6e4));
+		};
+		schedule();
+	}
+	/** Refresh every provider the host knows how to query. */
+	refreshAllBalances() {
+		return Promise.all(Object.keys(BALANCE_SOURCES).map((provider) => this.refreshProviderBalance(provider))).then(() => void 0);
+	}
+	/**
+	* Refresh one provider's cached account balance. The API key is resolved
+	* through the harness credential seam (`credentials.resolve`) with the
+	* environment as a fallback, so the Models page key is used with nothing to
+	* type. A failed fetch keeps the previous value and marks it as an error.
+	*/
+	async refreshProviderBalance(provider) {
+		const source = BALANCE_SOURCES[provider];
+		if (source === void 0) return;
+		const existing = this.balanceFetches.get(provider);
+		if (existing !== void 0) return existing;
+		const run = async () => {
+			if (!this.balanceEnabled) return;
+			const credentials = this.ctx.get("credentials");
+			let key;
+			if (credentials !== void 0) try {
+				key = (await credentials.resolve(source.keyEnv))?.value;
+			} catch {}
+			if (key === void 0 || key === "") key = process.env[source.keyEnv];
+			if (key === void 0 || key === "") {
+				this.balances.set(provider, {
+					provider,
+					currency: "CNY",
+					total: 0,
+					isAvailable: false,
+					fetchedAt: Date.now(),
+					status: "unconfigured"
+				});
+				return;
+			}
+			const base = source.baseEnv !== void 0 ? process.env[source.baseEnv] ?? source.baseURL : source.baseURL;
+			try {
+				const controller = new AbortController();
+				const timer = setTimeout(() => {
+					controller.abort();
+				}, 1e4);
+				let res;
+				try {
+					res = await fetch(`${base}/user/balance`, {
+						headers: {
+							authorization: `Bearer ${key}`,
+							accept: "application/json"
+						},
+						signal: controller.signal
+					});
+				} finally {
+					clearTimeout(timer);
+				}
+				if (!res.ok) throw new Error(`HTTP ${res.status}`);
+				const data = await res.json();
+				const first = Array.isArray(data?.balance_infos) ? data.balance_infos[0] : void 0;
+				const num = (value) => {
+					const n = typeof value === "number" ? value : typeof value === "string" && value !== "" ? Number(value) : NaN;
+					return Number.isFinite(n) ? n : 0;
+				};
+				this.balances.set(provider, {
+					provider,
+					currency: typeof first?.currency === "string" && first.currency !== "" ? first.currency : "CNY",
+					total: num(first?.total_balance),
+					granted: num(first?.granted_balance),
+					toppedUp: num(first?.topped_up_balance),
+					isAvailable: data?.is_available !== false,
+					fetchedAt: Date.now(),
+					status: "ok"
+				});
+			} catch (error) {
+				const previous = this.balances.get(provider);
+				const next = {
+					provider,
+					currency: previous?.currency ?? "CNY",
+					total: previous?.total ?? 0,
+					isAvailable: false,
+					fetchedAt: previous?.fetchedAt ?? Date.now(),
+					status: "error",
+					error: error instanceof Error ? error.message : String(error)
+				};
+				if (previous?.granted !== void 0) next.granted = previous.granted;
+				if (previous?.toppedUp !== void 0) next.toppedUp = previous.toppedUp;
+				this.balances.set(provider, next);
+			}
+		};
+		const promise = run().finally(() => {
+			this.balanceFetches.delete(provider);
+		});
+		this.balanceFetches.set(provider, promise);
+		return promise;
+	}
 	/** Read the full per-cycle usage history (log dialog data). */
 	readLog() {
 		this.rollCycleIfNeeded();
@@ -654,12 +842,14 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 			});
 		}
 		entries.sort((left, right) => left.key.localeCompare(right.key));
-		return {
+		const snapshot = {
 			day: this.cycle,
 			entries,
 			upgrade: this.upgrade,
 			upgradeError: this.upgradeError
 		};
+		if (this.balanceEnabled && this.balances.size > 0) snapshot.balances = [...this.balances.values()];
+		return snapshot;
 	}
 	/** Today's used tokens for one model key, or `0`. */
 	usedToday(key) {
@@ -896,6 +1086,11 @@ var TokenQuotaService = class TokenQuotaService extends Service {
 		if (this.disposeRouteLog !== void 0) this.disposeRouteLog();
 		if (this.disposeRouteCheck !== void 0) this.disposeRouteCheck();
 		if (this.disposeRouteClear !== void 0) this.disposeRouteClear();
+		if (this.disposeRouteBalance !== void 0) this.disposeRouteBalance();
+		if (this.balanceTimer !== void 0) {
+			clearTimeout(this.balanceTimer);
+			this.balanceTimer = void 0;
+		}
 		if (this.disposeRouteEvents !== void 0) this.disposeRouteEvents();
 		for (const res of this.musicClients) try {
 			res.end();

package/lib/types/client/TokenQuotaPanel.d.ts

@@ -24,6 +24,10 @@ export interface TokenQuotaPanelInjected {
     setModelOrder: (order: string[]) => void;
     /** Persist the floating scratchpad notes (content, title, placement). */
     saveNotes: (notes: TokenQuotaNote[]) => void;
+    /** Persist the account-balance display settings. */
+    setBalanceSettings: (enabled: boolean, pollMinutes: number) => void;
+    /** Ask the Host to refresh account balances (one provider, or all). */
+    refreshBalance: (provider?: string) => void;
     /** Toggle the live music output (requires a user gesture to start audio). */
     setMusicEnabled: (enabled: boolean) => void;
     /** Change the live-music master volume 0..1. */
@@ -44,5 +48,5 @@ export type TokenQuotaPanelComponentProps = PropsRuntime<'shell.overlay'> & Prop
  * @param props - composed slot props.
  * @returns the panel element tree.
  */
-export declare function TokenQuotaPanel({ t, load, setLimit, selectModel, setMonitored, setOnFull, setReset, setCheckUpdates, checkUpdatesNow, clearLogBefore, setDimWhenIdle, setModelOrder, saveNotes, setMusicEnabled, setMusicVolume, setMusicStyle, setMusicOnlyCurrentSession, useStore, actions, useSessions, }: TokenQuotaPanelComponentProps): import("react").JSX.Element;
+export declare function TokenQuotaPanel({ t, load, setLimit, selectModel, setMonitored, setOnFull, setReset, setCheckUpdates, checkUpdatesNow, clearLogBefore, setDimWhenIdle, setModelOrder, saveNotes, setBalanceSettings, refreshBalance, setMusicEnabled, setMusicVolume, setMusicStyle, setMusicOnlyCurrentSession, useStore, actions, useSessions, }: TokenQuotaPanelComponentProps): import("react").JSX.Element;
 //# sourceMappingURL=TokenQuotaPanel.d.ts.map