Persist credential cooldowns and usage snapshots across restarts - #35
maiphucgiang merged 8 commits into
Conversation
A credential that just failed authentication, or a model that just answered 429, was remembered only in memory. Restarting the gateway forgot both, so the very next request could be dispatched to an account the backend had just refused and waste a round trip on a predictable failure. Cooldowns are now written to credential-cooldowns.json next to the other managed state and restored on startup. Two properties matter more than the storage itself: * A restored deadline can never sit further out than its ceiling. A credential circuit breaker is bounded by an hour and a model quota cooldown by a day, and a deadline that is expired, unbounded, non-finite or in the past is discarded. Reopening the store therefore never extends a cooldown, so restarting repeatedly cannot park an account permanently. * What the writer produces is always something the reader accepts. A file the reader rejects discards every cooldown it held, not just the offending row, so the byte bound is enforced on write as well as on read: at capacity the rows nearest expiry are shed first until the payload fits. Records are keyed by the validated account identity rather than the credential path, so a path reused by a different account or product does not inherit the previous account's cooldowns. The in-memory table stays authoritative; the file is a mirror that is adopted once per identity. A same-account token refresh keeps its cooldowns, an explicit reload lifts the auth breaker while retaining model cooldowns, and deleting a credential forgets its rows. Storage is optional and advisory: when no path is configured the previous in-memory behaviour is unchanged, and a write failure is reported through cooldown_storage() while the running process keeps enforcing the cooldown it just recorded. Because a restart can no longer be used to clear a cooldown, clear_cooldowns() provides an explicit in-process reset.
Usage is collected by an hourly maintenance pass and kept only in memory, so right after a restart the dashboard and the billing totals were blank until the first pass completed. The snapshots are now cached in usage-snapshots.json next to the other managed state and reused as the starting point for the aggregate. Usage is presentation data. It is never consulted when choosing a credential, so the cache is deliberately not authoritative: a cached snapshot is only ever used to fill a gap the live aggregate does not already cover, and the next successful refresh replaces it. A missing, stale or unreadable file therefore costs nothing but a temporarily empty panel. Rows record the account identity they were fetched for, not just the path they were keyed by. A path reused by a different account has its cached usage dropped rather than displayed as the new account's, and deleting a credential forgets its snapshot. Snapshots older than a week are pruned, and the byte bound is enforced on write as well as on read: a file the reader rejects would discard every cached snapshot, so at capacity the oldest rows are shed first until the payload fits. A write failure is recorded in last_error and never interrupts aggregation.
Follow-up to the cooldown and usage-cache work, fixing defects found while reviewing the lifecycle paths rather than the happy path. Cooldowns: * An explicit reload no longer lifts the *incoming* account's breaker when a credential path is reused for a different account. Hydration now happens after the old account's state is cleared, so a replacement keeps its own cooldowns and does not inherit the previous account's. * Durable state is keyed only when the account's identifying components are actually present. account_key() hashes the profile and UID, so an account with no UID still produced a stable-looking hash shared by every other account in the same state; such an account now keeps its cooldown in memory only. * note_credential()/note_model() validate the deadline and timestamp before mutating anything. A huge integer raised OverflowError and None or a string raised TypeError, so a malformed upstream reset value could break the caller. * Expired rows are pruned before capacity is evaluated, and a write that had to shed rows to satisfy the byte bound now reports that instead of claiming the update was durable. * clear_cooldowns() reports changed_in_memory and durable separately, and a persistence failure is logged at a rate limit rather than only being readable from a diagnostics method. Usage snapshots: * Startup publishes the cache before serving, so the dashboard is populated from the first request instead of waiting for the first maintenance pass. * Rows keep the identity they were fetched for, and ownership is re-checked at publication time against the live account. Rows hydrated on an earlier pass are dropped too, so a reused path cannot keep showing the previous account's usage. * A restored snapshot is marked stale until a refresh confirms it, and rows older than the retention window are dropped rather than shown indefinitely. * The writer reports capacity shedding, stores the validated copy rather than the caller's mutable dictionaries, and rejects fractional request counts, impossible calendar dates and future timestamps instead of coercing them to zero.
Second review pass over the cooldown and usage-cache lifecycle. The theme is that a value must not claim more than it verified. Durability reporting: * clear_cooldowns() no longer infers durability from "nothing changed in memory". A model cooldown already absent from memory still had a row on disk, and a failed clear reported durable=True while the row survived; the two outcomes are now taken from the store's own result. _forget_credential_cooldown, deletion and pruning route their failures through the same bounded warning. * A write that had to shed rows to satisfy the byte bound reports that instead of claiming the update was durable. Separation of concerns: * Usage cleanup is its own method, forget_usage(), invoked alongside cooldown cleanup at the replacement and deletion points rather than being hidden inside forget_cooldowns(). It also removes the deleted path's live aggregate row, so deleting and re-adding a credential cannot resurrect the previous usage. Usage validation: * store() validates types before normalizing. `usage.get(...) or 0` and bool(partial) turned missing, empty-string, False and non-mapping values into plausible data; malformed input is now rejected, including fractional request counts and impossible calendar dates. Path length is bounded on load as well as write, and the validated copy is stored so later caller mutation cannot corrupt the cache. Freshness: * Expired rows are removed from the live aggregate, not merely skipped while summing. Only accounts with a usable UID own durable usage. `partial` remains the aggregate "this view is incomplete" flag while stale_accounts names which accounts, and a successful refresh clears only that account's staleness. Verification: * The startup tests now drive the real main() with synthetic credentials and a preseeded cache, asserting the aggregate is populated and every maintenance thread starts after hydration, instead of calling the helper directly. Hydration moved before the maintenance threads are started. * The token-refresh test now advances the manager generation, and the write-failure test exercises a real write.
Review pass over the failure paths. Clearing a cooldown and failing to write the
result were both reported as False, so a caller could not tell "there was nothing
to clear" from "the disk still holds a row that a restart will restore".
The cooldown and usage stores now return {"changed", "durable"} from clear_model,
clear_credential, forget and prune. A store also remembers that its last write
failed, so a later clear retries the write instead of reporting nothing to do
while stale rows sit on disk. Persistence failures from reload, deletion and
pruning are routed through the same bounded warning rather than being discarded.
An account whose identity components are incomplete still has its in-memory
cooldown cleared; only the durable half is skipped, because an account_key is a
hash of the profile and UID and an account with no UID would share one identity
with every other such account. Usage adoption applies the same rule, so a cached
snapshot is only owned by an account that can prove it.
Two tests were strengthened rather than added to: the token-refresh case now
rewrites the credential file and invalidates the manager so the generation
advances through the real path, and the usage write-failure case drives
_sync_usage() with a mocked upstream response and asserts the cache writer was
actually reached. A real-sync round trip asserts the cache is restorable.
Persisting cooldowns removed the old "restart the gateway to clear it"
workaround, so an operator had no supported way back when a circuit breaker or
a 429 cooldown had been recorded in error or when upstream had demonstrably
recovered.
POST /admin/credentials/{id}/reset-cooldown lifts every cooldown held by one
account: both the 401/403 circuit breaker and the per-model 429 cooldowns. It
resets the whole account rather than just the breaker because a credential that
is mid-429 is exactly the case an operator is trying to unstick, and lifting only
the breaker would leave it unusable.
The action is a local state change. It takes no request body, never refreshes a
token, never contacts upstream and never queues synchronization, so a subsequent
genuine failure is free to arm the cooldown again. It deliberately skips the
ledger requirement and the maintenance lock, and it works for a manually disabled
account, since a stale cooldown is worth clearing either way. The one durable
write is a forget() of the account's row, which also performs the retry when a
previous write failed.
The result reports changed_in_memory and durable separately, following the same
contract as the existing pool-level reset. ok is false when the write failed: the
in-memory reset is already effective, but a restart would restore the stored row,
so the request can simply be repeated. The failure is also visible in the audit
trail.
The credential-row button is intentionally not gated on the displayed cooldowns
or on enablement, because a failed durable clear has to stay retryable.
Reviewer's GuideThe PR adds defensive, identity-aware persistence for credential cooldowns and display-only usage snapshots, hydrates usage before serving after restart, and provides a fully audited local admin reset action with durable-failure reporting, UI access, and documentation. Sequence diagram for persisted credential cooldown lifecyclesequenceDiagram
participant Upstream
participant CredentialPool
participant CredentialCooldowns
participant Disk
Upstream-->>CredentialPool: 401/403 or 429 response
CredentialPool->>CredentialCooldowns: note_credential(identity, profile, until, reason)
CredentialCooldowns->>Disk: _save_locked()
Note over CredentialCooldowns,Disk: Absolute deadlines are validated and atomically persisted
CredentialPool->>CredentialCooldowns: restore(identity, profile)
CredentialCooldowns-->>CredentialPool: fail_until and model cooldowns
CredentialPool->>CredentialPool: _adopt_cooldowns(entry)
Sequence diagram for local cooldown resetsequenceDiagram
actor Admin
participant AdminAPI
participant CredentialActions
participant CredentialPool
participant CredentialCooldowns
participant AuditStore
Admin->>AdminAPI: POST /admin/credentials/{id}/reset-cooldown
AdminAPI->>CredentialActions: run(gateway, reset-cooldown, identity)
CredentialActions->>CredentialPool: reset_cooldowns_for(identity)
CredentialPool->>CredentialCooldowns: forget(identity)
CredentialCooldowns-->>CredentialPool: durable status
CredentialPool-->>CredentialActions: changed_in_memory and durable
CredentialActions->>AuditStore: event(admin, credential.reset-cooldown)
CredentialActions-->>AdminAPI: ok, results
AdminAPI-->>Admin: reset result
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 1 issue
Prompt for AI Agents
Please address the comments from this code review:
## Individual Comments
### Comment 1
<location path="converter.py" line_range="1522-1528" />
<code_context>
+ "by_day": usage["by_day"],
"total_credits": round(usage["total_credits"], 2),
"requests": usage["requests"],
"partial": bool(usage.get("partial")),
"fetched_at": time.time()}
+ snapshots = CONFIG.get("usage_snapshots")
+ if snapshots is not None and entry.get("account_key") and entry.get("uid"):
+ # Only a validated identity may key durable state.
+ snapshots.store(entry["id"], entry["account_key"], site, usage,
+ partial=bool(usage.get("partial")))
if not pool.apply_if_current(cm, generation, store):
stale.add(entry["id"])
</code_context>
<issue_to_address>
**issue (bug_risk):** The usage-sync path coerces any upstream `partial` value with `bool(...)` before passing it to the validated snapshot store, so malformed values such as `"false"`, `1`, or arbitrary non-empty objects are accepted and displayed as a legitimate partial-status flag instead of being rejected.
**Triggers:** When the upstream usage response contains a malformed non-boolean `partial` field.
**Suggested fix:** Pass the raw `partial` value to the store and only publish the live row after strict validation, or explicitly require `type(usage.get("partial")) is bool` before updating the aggregate.
</issue_to_address>Sourcery assessment
Needs a human reviewer. 1 finding to address first, and a wrong cooldown persistence or reset can cause extra credential/model requests after restart, while a wrong usage snapshot can display stale totals; both are bounded and can be repaired by clearing the local state or reverting the change. Reverting does not remove files already written, but those files are derived, non-authoritative state rather than irreversible account or billing data.
Blocking findings: converter.py:1528
The usage cache stored its completeness flag with bool(usage.get("partial")),
which masks a malformed value instead of rejecting it: bool("false") is True, so
a non-boolean flag would have been persisted as a legitimate one. The store
already requires exactly a bool, but the coercion ran first, so its validator
never saw the raw value.
Pass the value through instead, defaulting to False only when the key is absent.
This is a hardening of the durability boundary rather than a fix for a reachable
defect: fetch_request_usage computes partial locally and never reads it from the
upstream payload, so no malformed value is currently produced there. A test
asserts that each malformed value is rejected and leaves nothing durable behind,
while the live dashboard row still shows the freshly fetched figures.
|
Addressed in 82cdb79 — thanks, the underlying concern was worth fixing. One correction to the premise, though: That said, the coercion itself was still wrong at that boundary, so I changed it: snapshots.store(entry["id"], entry["account_key"], site, usage,
partial=usage.get("partial", False))The raw value now reaches Added Full backend suite is unchanged against the pre-existing baseline (0 new failures). |
|
@codex review |
maiphucgiang
left a comment
maiphucgiang
left a comment
There was a problem hiding this comment.
Please fix the concurrent usage-publication regression below before merging. I reviewed 82cdb79 and ran the focused cooldown, usage, credential-action, runtime, credential-lifecycle, and region-routing tests (209 passed). A deterministic two-thread reproduction, including with UsageSnapshots enabled, raises KeyError on this head while the equivalent main@426e8a5 sequence succeeds. The reproduction used synthetic credentials in an isolated sandbox with no upstream access.
| for cred_id in list(accounts): | ||
| snap = accounts[cred_id] |
There was a problem hiding this comment.
[P2] Snapshot usage rows under the same lock used by credential cleanup
list(accounts) only copies the keys; accounts[cred_id] still reads the shared mutable dictionary. This PR also adds forget_usage(), which removes entries from that dictionary during credential deletion or replacement under the pool lock, but this reader does not hold that lock. I reproduced the race by pausing publication just before this lookup, deleting the credential through pool.remove_file() on another thread, then resuming: publication raises KeyError. This also reproduces with the real usage-snapshot store enabled, whereas the equivalent sequence on main succeeds. It aborts the current usage-maintenance pass and can make a manual synchronization fail.
Please coordinate usage cleanup and snapshot acquisition with a common lock, then aggregate a consistent local snapshot instead of looking up copied keys in the live dictionary. _adopt_cached_usage() has the same keys-then-lookup pattern at lines 1564–1565 and needs the same protection. Add a deterministic concurrent deletion/replacement regression test.
There was a problem hiding this comment.
Fixed in 7eb3a8c. Confirmed as a regression from this branch: main@426e8a5 used a single accounts.items() pass and had no forget_usage(), so neither the prune-during-read nor the lookup-after-copy pattern existed there. _publish_usage_daily now holds pool._lock across the read and aggregates from a local copy; _adopt_cached_usage iterates list(accounts.items()) instead of indexing by copied key. Added test_a_concurrent_deletion_during_publication_does_not_raise, which reproduces your KeyError on 82cdb79 (converter.py:1600) and passes on the fix, plus a test asserting the lock is genuinely held. Lock ordering stays pool._lock -> snapshots._lock with no inverse.
_publish_usage_daily read the live usage map by copying its keys and then
indexing the dictionary per row:
for cred_id in list(accounts):
snap = accounts[cred_id]
forget_usage() pops from that same dictionary under the pool lock, from
remove_file() -> prune() -> forget_credential_state(). The reader held no lock,
so a credential deleted or replaced while a usage pass was running raised
KeyError, which aborted the pass and could fail a manual synchronization.
Publication now takes pool._lock for the whole read and aggregates from a local
copy of the map rather than indexing the live one, so a concurrent deletion can
no longer pull a row out from under it. _adopt_cached_usage() had the same
keys-then-lookup pattern and is fixed the same way, iterating over a copy of the
items instead of indexing by copied key.
This is a regression introduced by the caching in this branch: main used
accounts.items() over a single pass and had no forget_usage(), so neither the
prune-during-read nor the lookup-after-copy pattern existed there.
Lock ordering is unchanged and consistent: pool._lock -> snapshots._lock, with
UsageSnapshots referencing neither the pool nor CONFIG, so no inverse order
exists. _publish_usage_daily is called with no other lock held, and the log line
stays outside the lock.
Adds a deterministic regression test that reproduces the original KeyError by
deleting a row from inside the publication window, plus a test asserting the
pool lock is actually held while the map is read. Both fail if the lock is
removed; the KeyError reproduction is confirmed against the previous revision.
|
You're right, and thank you for the precise diagnosis — fixed in 7eb3a8c. Your reading is exactly correct. Fix. Confirmed as a regression from this branch. I verified on Regression test. Lock ordering is unchanged and consistent: Focused set you listed now reports 210 passed, 1 skipped (your 209 plus the new tests), and the full backend suite shows no new failures against the pre-existing baseline. |
Upstream added a usage-snapshot publish before `runtime_management.install` (PR maiphucgiang#35), which conflicts with the `SessionStoreError` guard added here. Keep both: publish cached usage first, then install under the guard.
2 participants
What this does
A backend restart currently discards two kinds of cooldown, so a credential that upstream just refused is retried immediately after a restart, and an operator has no supported way to clear a cooldown that was recorded in error.
This persists both, and adds an admin action to clear them.
1. Credential cooldowns survive a restart (
auth/credential-cooldowns.json)fail_until) and the per-model 429 cooldowns.2. Usage snapshots are cached across restarts (
auth/usage-snapshots.json)stale, revalidated against current ownership, and aged out.3.
POST /admin/credentials/{id}/reset-cooldownchanged_in_memoryanddurableseparately;okis false when the write failed, so a failed reset is never presented as done and the request stays retryable. Failures are also visible in the audit trail.What is intentionally not persisted
_stickysession bindings, login throttling, in-flight capacity leases, file locks, OAuth handshakes and poll cursors. Balances already persist through the existing ledger.Verification
tests/test_credential_cooldowns.py,tests/test_usage_snapshots.py,tests/test_credential_actions.py,tests/test_runtime_endpoints.pycover restart survival, key rotation, disabled keys, hostile snapshots (duplicate keys,Infinity/NaN, oversized fields, wrong types), symlinked files, write failure followed by a successful retry, ownership revalidation, CSRF/Origin/authentication rejection without mutation, and account isolation.main()and asserting the aggregate is populated before serving; the hydration step was removed to confirm the test fails without it.vp buildsucceeds.Only synthetic credentials were used. No upstream requests were made.
Scope
Backend plus a one-line credential-row button and the endpoint tables in
docs/advanced.md/docs/advanced.zh-CN.md. No new dependencies.Summary by Sourcery
Persist credential cooldowns and usage snapshots safely across restarts, and provide an administrative way to clear an account’s complete cooldown state.
New Features:
Bug Fixes:
Enhancements:
Documentation:
Tests: