From a821f4316a7e5983b230f315327535b38109b74f Mon Sep 17 00:00:00 2001 From: can1357 Date: Fri, 14 Aug 2026 07:23:02 +0200 Subject: [PATCH] fix(ai): sized usage requests by unfiltered account batch and updated cache methods - Updated `RemoteAuthCredentialStore` to track broker usage accounts across snapshots and streams before account-pool filtering. - Replaced `#countUsageAccounts` with dynamic tracking methods `#replaceBrokerUsageAccounts`, `#upsertBrokerUsageAccount`, and `#removeBrokerUsageAccount`. - Refactored `AuthStorage.#fetchUsageCached` to accept an options object for `timeoutMs` and `forceRefresh`. - Updated dockerignore patterns to exclude `**/.venv/` and ensure depth-agnostic `.env` secret exclusion. --- .dockerignore | 8 ++- Dockerfile.dockerignore | 8 ++- Dockerfile.robomp.dockerignore | 8 ++- packages/ai/CHANGELOG.md | 2 +- packages/ai/src/auth-broker/remote-store.ts | 56 ++++++++++++++++----- packages/ai/src/auth-storage.ts | 18 ++++--- 6 files changed, 73 insertions(+), 27 deletions(-) diff --git a/.dockerignore b/.dockerignore index e271c3e56..a22a783a0 100644 --- a/.dockerignore +++ b/.dockerignore @@ -74,5 +74,9 @@ out.jsonl out.html pi-*.html -# Secrets. Should never be in an image regardless. -.env +# Host virtualenvs — the image installs its own interpreter deps. +**/.venv/ + +# Secrets. Should never be in an image regardless. Depth-agnostic: a bare +# `.env` misses nested ones such as `python/robomp/.env`. +**/.env diff --git a/Dockerfile.dockerignore b/Dockerfile.dockerignore index 05fa0925b..d8175d469 100644 --- a/Dockerfile.dockerignore +++ b/Dockerfile.dockerignore @@ -75,5 +75,9 @@ out.jsonl out.html pi-*.html -# Secrets. Should never be in the image regardless. -.env +# Host virtualenvs — the image installs its own interpreter deps. +**/.venv/ + +# Secrets. Should never be in the image regardless. Depth-agnostic: a bare +# `.env` misses `python/robomp/.env`, which `COPY . /pi/` would bake in. +**/.env diff --git a/Dockerfile.robomp.dockerignore b/Dockerfile.robomp.dockerignore index 21ac7290d..0596a88fe 100644 --- a/Dockerfile.robomp.dockerignore +++ b/Dockerfile.robomp.dockerignore @@ -78,8 +78,12 @@ out.jsonl out.html pi-*.html -# Secrets. Should never be in the image regardless. -.env +# Host virtualenvs — the image installs its own interpreter deps. +**/.venv/ + +# Secrets. Should never be in the image regardless. Depth-agnostic: a bare +# `.env` misses `python/robomp/.env`, which sits next to the copied src tree. +**/.env # Robomp-only excludes. Natives + wheel + python + bun + rustup all come # from PI_BASE; the web-builder stage only needs root manifests + the diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 88bf6b7f5..1b4932ed0 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -4,7 +4,7 @@ ### Fixed -- Fixed `omp usage invalidate` to discard stale OAuth and API-key usage snapshots, then force a cache-bypassing, per-provider serialized refresh so upgraded subscriptions do not silently retain pre-change quota data. +- Fixed `omp usage invalidate` to discard stale OAuth and API-key usage snapshots, then force a cache-bypassing, per-provider serialized refresh with a broker request budget sized for the full unfiltered account batch, so upgraded subscriptions do not silently retain pre-change quota data. - Fixed quota reporting and Cookie capture guidance for China (Beijing) Alibaba Token Plan credentials ([#8509](https://github.com/can1357/oh-my-pi/issues/8509)). ## [17.3.3] - 2026-08-14 diff --git a/packages/ai/src/auth-broker/remote-store.ts b/packages/ai/src/auth-broker/remote-store.ts index f02994481..229cbc380 100644 --- a/packages/ai/src/auth-broker/remote-store.ts +++ b/packages/ai/src/auth-broker/remote-store.ts @@ -253,7 +253,9 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { #usageInflight?: Promise; #credentialBlockReconcileAfter: Map = new Map(); #usageCacheEpoch = 0; - #maxUsageAccountsPerProvider = 1; + /** Raw broker credentials retained to size aggregate usage requests before account-pool filtering. */ + #brokerUsageProviderByCredentialId = new Map(); + #brokerUsageAccountCounts = new Map(); /** Per-snapshot lookup of oauth credentials by provider; rebuilt when `#snapshot` is replaced. */ #usageFilterLookup?: { snapshot: SnapshotResponse; byProvider: Map }; /** Memoized `#filterUsageReports` output, keyed on (input identity, lookup identity). */ @@ -297,7 +299,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { #applySnapshot(snapshot: SnapshotResponse, generation: number, protectNewBlocks = true): void { const nowMs = Date.now(); - this.#maxUsageAccountsPerProvider = this.#countUsageAccounts(snapshot.credentials); + this.#replaceBrokerUsageAccounts(snapshot.credentials); const previousCredentials = this.#snapshot.credentials; const credentials = snapshot.credentials .filter(entry => isCredentialInAccountPool(entry, this.#accountPool)) @@ -429,8 +431,9 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { generation: number, serverNowMs: number, ): void { + this.#upsertBrokerUsageAccount(entry); if (!isCredentialInAccountPool(entry, this.#accountPool)) { - this.#removeStreamCredential(entry.id, refresher, generation, serverNowMs); + this.#removeStreamCredential(entry.id, refresher, generation, serverNowMs, { retainBrokerUsageAccount: true }); return; } const incoming = this.#normalizeSnapshotEntryBlocks(entry, Date.now()); @@ -448,7 +451,14 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { this.#snapshotReceivedAt = Date.now(); } - #removeStreamCredential(id: number, refresher: RefresherSchedule, generation: number, serverNowMs: number): void { + #removeStreamCredential( + id: number, + refresher: RefresherSchedule, + generation: number, + serverNowMs: number, + options?: { retainBrokerUsageAccount?: boolean }, + ): void { + if (!options?.retainBrokerUsageAccount) this.#removeBrokerUsageAccount(id); const removed = this.#snapshot.credentials.find(entry => entry.id === id); if (removed?.blocks && removed.blocks.length > 0) this.#invalidateUsageCache(); const credentials = this.#snapshot.credentials.filter(entry => entry.id !== id); @@ -1068,14 +1078,36 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { }); } - #countUsageAccounts(entries: readonly SnapshotEntry[]): number { - const counts = new Map(); - let maximum = 1; - for (const entry of entries) { - const count = (counts.get(entry.provider) ?? 0) + 1; - counts.set(entry.provider, count); - maximum = Math.max(maximum, count); + #replaceBrokerUsageAccounts(entries: readonly SnapshotEntry[]): void { + this.#brokerUsageProviderByCredentialId.clear(); + this.#brokerUsageAccountCounts.clear(); + for (const entry of entries) this.#upsertBrokerUsageAccount(entry); + } + + #upsertBrokerUsageAccount(entry: Pick): void { + const previous = this.#brokerUsageProviderByCredentialId.get(entry.id); + if (previous === entry.provider) return; + if (previous !== undefined) { + const count = this.#brokerUsageAccountCounts.get(previous) ?? 0; + if (count <= 1) this.#brokerUsageAccountCounts.delete(previous); + else this.#brokerUsageAccountCounts.set(previous, count - 1); } + this.#brokerUsageProviderByCredentialId.set(entry.id, entry.provider); + this.#brokerUsageAccountCounts.set(entry.provider, (this.#brokerUsageAccountCounts.get(entry.provider) ?? 0) + 1); + } + + #removeBrokerUsageAccount(id: number): void { + const provider = this.#brokerUsageProviderByCredentialId.get(id); + if (provider === undefined) return; + this.#brokerUsageProviderByCredentialId.delete(id); + const count = this.#brokerUsageAccountCounts.get(provider) ?? 0; + if (count <= 1) this.#brokerUsageAccountCounts.delete(provider); + else this.#brokerUsageAccountCounts.set(provider, count - 1); + } + + #maxBrokerUsageAccounts(): number { + let maximum = 1; + for (const count of this.#brokerUsageAccountCounts.values()) maximum = Math.max(maximum, count); return maximum; } @@ -1087,7 +1119,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { if (this.#usageInflight) return this.#usageInflight; const epoch = this.#usageCacheEpoch; const inflight = this.#client - .fetchUsage({ maxAccountsPerProvider: this.#maxUsageAccountsPerProvider }) + .fetchUsage({ maxAccountsPerProvider: this.#maxBrokerUsageAccounts() }) .then(body => { if (epoch !== this.#usageCacheEpoch) return this.#loadUsageReports(); this.#usageCache = { reports: body.reports, fetchedAt: Date.now() }; diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 43a0f9211..625f6b6d4 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -3274,9 +3274,10 @@ export class AuthStorage { async #fetchUsageCached( request: UsageRequestDescriptor, - timeoutMs?: number, - forceRefresh = false, + options: { timeoutMs?: number; forceRefresh?: boolean } = {}, ): Promise { + const timeoutMs = options.timeoutMs; + const forceRefresh = options.forceRefresh ?? false; const cacheKey = this.#buildUsageReportCacheKey(request); const now = Date.now(); const cached = forceRefresh ? undefined : this.#usageCache.get(cacheKey); @@ -3797,10 +3798,9 @@ export class AuthStorage { if (!resolvedApiKey) return null; usageCredential.apiKey = resolvedApiKey; } - return this.#fetchUsageCached( - this.#buildUsageRequest(provider, usageCredential, options?.baseUrl), - options?.timeoutMs ?? this.#usageRequestTimeoutMs, - ); + return this.#fetchUsageCached(this.#buildUsageRequest(provider, usageCredential, options?.baseUrl), { + timeoutMs: options?.timeoutMs ?? this.#usageRequestTimeoutMs, + }); } /** @@ -4007,10 +4007,12 @@ export class AuthStorage { requests.map(request => { const forceRefresh = serializedProviders.has(request.provider); if (!forceRefresh) { - return this.#fetchUsageCached(request, this.#usageRequestTimeoutMs); + return this.#fetchUsageCached(request, { timeoutMs: this.#usageRequestTimeoutMs }); } const tail = tails.get(request.provider) ?? Promise.resolve(); - const current = tail.then(() => this.#fetchUsageCached(request, this.#usageRequestTimeoutMs, true)); + const current = tail.then(() => + this.#fetchUsageCached(request, { timeoutMs: this.#usageRequestTimeoutMs, forceRefresh: true }), + ); tails.set( request.provider, current.then(