fix(terminal-paths): refund observability, chunk-key lifecycle, 499 classification, writer deadline races#494
fix(terminal-paths): refund observability, chunk-key lifecycle, 499 classification, writer deadline races#494Gajesh2007 wants to merge 5 commits into
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: COMMENT
Security — ⚠️ SKIPPED
Error: OpenRouter HTTP 402: {"error":{"message":"Insufficient credits. Add more using https://openrouter.ai/settings/credits","code":402}}
Performance — ⚠️ SKIPPED
Error: OpenRouter HTTP 402: {"error":{"message":"Insufficient credits. Add more using https://openrouter.ai/settings/credits","code":402}}
Type_diligence — ⚠️ SKIPPED
Error: OpenRouter HTTP 402: {"error":{"message":"Insufficient credits. Add more using https://openrouter.ai/settings/credits","code":402}}
Additive_complexity — ⚠️ SKIPPED
Error: OpenRouter HTTP 402: {"error":{"message":"Insufficient credits. Add more using https://openrouter.ai/settings/credits","code":402}}
✅ All four passes clean. No issues found.
🤖 Automated review by Centaur · DAR-186
|
This PR hardens key-material lifecycle (chunk-decryption key zeroing and cleanup) and fixes silently-swallowed refund errors; it is a net security improvement with one nuance worth verifying. Trust boundaries touched
Threat-by-threat assessmentT-005 (Repudiation — consumer billing) ✅ Strengthens T-008 (Information Disclosure — plaintext SSE chunk fallback) ✅ Strengthens T-014 (Information Disclosure — operator inspects process memory) ✅ Strengthens (incidental) T-006 (Spoofing — WebSocket registry flood) ℹ️ Neutral T-030 (Tampering — duplicate Stripe webhook) ℹ️ Neutral T-040 (Information Disclosure — Host header injection) ℹ️ Neutral T-031 (Information Disclosure — unauthenticated earnings endpoint) ℹ️ Neutral New attack surface / concerns not covered by an existing threat1. 2. 3. SEC-* findings resolved by this PRNone of the open SEC-* findings (SEC-007, SEC-012, SEC-034, etc.) are closed by this diff. The changes are hygiene and key-material hardening on top of existing infrastructure. Summary for reviewers: The zeroing upgrade ( 🔐 Threat model: |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d00fcbfcc5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| func (s *Server) cancelDispatchAndForget(provider *registry.Provider, pr *registry.PendingRequest) { | ||
| s.cancelDispatch(provider, pr) | ||
| if pr != nil { | ||
| s.chunkKeys.forget(pr.SessionPrivKey) |
There was a problem hiding this comment.
In the cancellation/retry paths where the provider can still be processing a late chunk, this single delete can race with chunkKeyCache.sharedKey: that function drops c.mu while computing the shared key, then stores the result afterward. If forget runs during that compute window while the entry is still absent, it is a no-op and the read loop then inserts a cache entry for a request that has already been removed from pendingReqs; no later terminal handler can find it, so the session key remains pinned until the wholesale cap reset. These cross-goroutine terminal paths need a tombstone/generation check or a second removal after the compute/store path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c94ed04. Went with the tombstone approach: forget/forgetAndZero now mark the priv pointer in a bounded tombstone set (chunkKeyCache.dead, capped at chunkKeyDeadMax with wholesale drop), and sharedKey re-checks it under the re-acquired lock before storing. A forget racing the unlocked X25519 compute window — or a late straggler chunk decrypted entirely after a cross-goroutine forget — now gets a working key back for that one decrypt without re-inserting an entry that no terminal path remains to clean up.
This is safe because a forgotten priv is never legitimately re-cached: every dispatch attempt generates fresh session keys (dispatchPrimary forgets the old pointer before overwriting it at dispatch.go:911-912), so post-forget decrypts are only stragglers of abandoned attempts — they recompute per chunk instead of caching. Tombstones hold no secret material and only need to outlive the ms-scale race window, so the wholesale cap drop is harmless.
Tests: TestChunkKeyCacheForgetPreventsRecache, TestChunkKeyCacheForgetBeforeFirstUsePreventsCache (cancel-before-first-chunk case), TestChunkKeyCacheDeadTombstoneCapResets; existing concurrent test passes under -race.
| // entries until the cap-reset). Zeroing is safe here: this defer runs | ||
| // on the provider's read-loop goroutine — the only goroutine that ever | ||
| // decrypts with these keys — after conn.Read has returned for good. | ||
| s.forgetProviderPendingKeys(provider) |
There was a problem hiding this comment.
This cleanup only snapshots pending keys when the read loop exits before calling registry.Disconnect, but several production disconnect paths call Registry.Disconnect directly first (stale eviction, duplicate-serial eviction, and forced provider removal in registry.go). Those calls clear p.pendingReqs and close the socket; when this defer later runs after CloseNow unblocks the read loop, PendingSessionKeys() is empty, so any cached chunk keys for in-flight requests on that provider remain orphaned until the cap reset. The key snapshot/forget needs to happen before every registry disconnect that wipes pending requests, not only this read-loop-owned path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c94ed04, via the cache rather than the registry (registry.go is frozen on this branch for the in-flight v0.6.31 work, so hooking every direct Registry.Disconnect caller — evictStale, DisconnectDuplicatesBySerial, RemoveProviderBySerial(force) — wasn't an option here, and would also leave future callers unprotected).
Instead, forgetProviderPendingKeys now ends with a new chunkKeys.forgetPeer(provider.PublicKey) sweep that drops every cache entry stored under this provider's public key (the cache already stores peerPub per entry, so it needs no pending-map state at all). Since Disconnect's closeWriterNow always unblocks conn.Read, the read-loop defer always runs — including for all three registry-initiated paths — and the sweep catches exactly the entries whose requests were wiped from pendingReqs before the per-request snapshot could see them.
The sweep is plain-delete only: no zeroing and no tombstoning, because the same peerPub can serve a newer same-keypair session of the same machine (duplicate-serial eviction) whose read loop may be decrypting with a matching entry right now — a swept live entry just recomputes on its next chunk. The per-request forgetAndZero snapshot stays first so the common read-loop-owned path still scrubs the shared secrets.
Test: TestRegistryInitiatedDisconnectSweepsChunkKeys (registers over WS with a public key, calls reg.Disconnect directly to wipe the pending map first, asserts the orphan is swept and NOT zeroed), plus TestChunkKeyCacheForgetPeerSweepsWithoutZeroingOrTombstoning.
| // chunkKeyCache's zeroing policy). Unconditional: even when the | ||
| // refund below no-ops (already-finalized dup park), no terminal | ||
| // handler ever claimed this record to forget the key. | ||
| s.chunkKeys.forget(expired.SessionPrivKey) |
There was a problem hiding this comment.
The new expiry forget is below the IsReservationFinalized() early return, so post-commit paths that already finalized the reservation before their cleanup defer runs (for example stream timeout/provider-incomplete branches that call refundReservedBalance, then holdForSettlement and RemovePending) skip both holding and forgetting. Those requests have already decrypted chunks, but after RemovePending no provider terminal can find the PendingRequest, leaving the cached shared key orphaned until the wholesale cache reset; drop the chunk key before returning for finalized records as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c94ed04. holdForSettlement's IsReservationFinalized() early return now does s.chunkKeys.forget(pr.SessionPrivKey) before returning: as you noted, the relay's stream-timeout/provider-incomplete branches finalize via refundReservedBalance without a read-loop terminal ever running, and since the record is never parked, no late handleComplete/handleInferenceError can claim it to drop the key. For records that WERE settled by a read-loop terminal first, the terminal's forgetAndZero already ran and this forget is a no-op. Plain forget (no zeroing) since this runs on the consumer goroutine while a late chunk decrypt may be in flight on the provider read loop.
Combined with the tombstone change in the same commit (see the dispatch.go thread), the forget also prevents a late chunk decrypted between this cleanup and RemovePending from re-caching the key.
Test: TestHoldForSettlementFinalizedForgetsChunkKey — finalized record's key is forgotten immediately, not zeroed, and the record is still not parked/refunded.
#494 1. Re-cache race (dispatch.go cancel/retry forgets): forget/forgetAndZero now TOMBSTONE the priv pointer; sharedKey re-checks the tombstone set under the re-acquired lock before storing, so a forget racing the unlocked X25519 compute window — or a late straggler chunk decrypted after a cross-goroutine forget — hands back a working key without re-inserting an entry no terminal path remains to clean up. Safe because every dispatch attempt generates fresh session keys, so a forgotten priv is never legitimately re-cached. Tombstones are bounded (chunkKeyDeadMax, wholesale drop) and hold no secrets. 2. Registry-initiated disconnects (stale eviction, duplicate-serial eviction, forced removal call Registry.Disconnect directly, wiping pendingReqs BEFORE the read-loop defer can snapshot session keys): forgetProviderPendingKeys now also sweeps the cache by the provider's public key (new chunkKeyCache.forgetPeer), which needs no pending-map state — the defer always runs once CloseNow unblocks conn.Read. Plain delete, no zeroing, no tombstoning: a same-keypair replacement session (dup-serial) may be decrypting with a matching entry on its own read loop; it just recomputes on its next chunk. 3. holdForSettlement's IsReservationFinalized early return now forgets the chunk key first: relay stream-timeout/provider-incomplete refunds finalize without a read-loop terminal ever running, and the record is never parked, so nothing else would ever drop the entry. Addresses all three P2 comments on PR #494 (chatgpt-codex-connector).
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — 1 finding(s)
- 🔵 [INFO]
coordinator/api/reservations.go:100-111— Refund failure logging exposes account_id in error logs- Suggestion: Consider redacting or hashing account_id in error logs to prevent potential PII exposure
Performance — 1 finding(s) (1 blocking)
- 🟡 [MEDIUM]
coordinator/api/chunk_key_cache.go:96-117— X25519 computation performed while holding mutex lock- Suggestion: Move the e2e.PrecomputeSharedKey call outside the mutex-protected section to avoid blocking other cache operations during the ~40-60µs scalar multiplication
Type_diligence — ✅ No issues found
Additive_complexity — 3 finding(s) (1 blocking)
- 🟡 [MEDIUM]
coordinator/api/chunk_key_cache.go:30-52— Complex lifecycle policy with multiple forget variants creates cognitive overhead- Suggestion: Consider consolidating the forget/forgetAndZero distinction into a single method with a boolean parameter, or use a clearer naming convention that makes the safety constraints obvious at call sites
- 🔵 [INFO]
coordinator/api/chunk_key_cache.go:96-117— Tombstone mechanism adds state tracking complexity for race protection- Suggestion: Document whether the tombstone approach is simpler than alternative solutions like reference counting or explicit lifecycle management
- 🔵 [INFO]
coordinator/api/dispatch.go:928-945— New cancelDispatchAndForget wrapper adds indirection layer- Suggestion: Consider whether the chunk key cleanup could be integrated directly into cancelDispatch rather than requiring a separate wrapper function
5 finding(s) total, 2 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
| // Financial: a failed refund over-charges the consumer. Never swallow it. | ||
| // No inline retry — this path can run on latency-sensitive request paths; | ||
| // the ERROR log + billing.refund_failures metric is the alerting hook. | ||
| if err := s.store.Credit(accountID, amount, store.LedgerRefund, "reservation_refund"); err != nil { | ||
| s.logger.Error("failed to credit reservation refund to consumer", | ||
| "account_id", accountID, | ||
| "model", model, | ||
| "refund_micro_usd", amount, | ||
| "error", err, | ||
| ) | ||
| s.ddIncr("billing.refund_failures", tags) | ||
| } |
There was a problem hiding this comment.
🔵 [INFO] 🔒 Refund failure logging exposes account_id in error logs
💡 Suggestion: Consider redacting or hashing account_id in error logs to prevent potential PII exposure
📊 Score: 2×3 = 6 · Category: missing-input-validation
There was a problem hiding this comment.
Account IDs here are opaque internal identifiers (Privy DIDs / key-hash labels), not direct PII, and logging them on billing paths matches established practice elsewhere in the coordinator (e.g. provider-account linking logs). Operationally the raw ID is required: a refund-failure alert is only actionable if on-call can locate the account to credit it manually. Happy to revisit if we adopt a repo-wide ID-redaction policy.
| shared := e2e.PrecomputeSharedKey(&pub, priv) | ||
|
|
||
| c.mu.Lock() | ||
| if _, gone := c.dead[priv]; gone { | ||
| // The request hit a terminal path while the lock was dropped for the | ||
| // compute (or this is a late chunk of an already-forgotten request): | ||
| // hand the key back for this one decrypt but do NOT cache it — no | ||
| // terminal path remains to forget the entry, so it would pin the | ||
| // session key until the cap reset. | ||
| c.mu.Unlock() | ||
| return shared, nil | ||
| } | ||
| if c.m == nil || len(c.m) >= chunkKeyCacheMax { | ||
| // Wholesale reset (abandoned-entry safety net). Deliberately does NOT | ||
| // zero the dropped keys: entries may belong to live streams whose read | ||
| // loops are decrypting with them right now (see the zeroing policy on | ||
| // the type comment). | ||
| c.m = make(map[*[32]byte]chunkKeyEntry) | ||
| } | ||
| c.m[priv] = chunkKeyEntry{peerPub: peerPub, shared: shared} | ||
| c.mu.Unlock() | ||
| return shared, nil |
There was a problem hiding this comment.
🟡 [MEDIUM] ⚡ X25519 computation performed while holding mutex lock
💡 Suggestion: Move the e2e.PrecomputeSharedKey call outside the mutex-protected section to avoid blocking other cache operations during the ~40-60µs scalar multiplication
📊 Score: 3×4 = 12 · Category: blocking_io
There was a problem hiding this comment.
The X25519 compute already runs with the mutex dropped — sharedKey unlocks before e2e.PrecomputeSharedKey and re-locks only to insert (chunk_key_cache.go: unlock at the miss, compute, re-lock). That lock-drop window is precisely why the tombstone re-check under the re-acquired lock exists. The section under the lock is map ops only.
| // key change (paranoia: reconnect races) invalidates the entry instead of | ||
| // decrypting with a stale shared key. | ||
| // | ||
| // Lifecycle / zeroing policy. Every request-terminal path must drop its entry | ||
| // (orphans are otherwise pinned until the chunkKeyCacheMax wholesale reset), | ||
| // but only SOME of those paths may zero the key material: | ||
| // | ||
| // - forgetAndZero (delete + zero): ONLY from the provider read-loop | ||
| // goroutine that performs this request's chunk decryption — | ||
| // handleComplete, handleInferenceError, and providerReadLoop's disconnect | ||
| // cleanup. On that goroutine no decrypt with the key can be concurrently | ||
| // in flight, so zeroing is safe and scrubs the shared secret from the | ||
| // heap promptly. | ||
| // - forget (delete only, NO zeroing): every CROSS-GOROUTINE terminal — | ||
| // dispatch retry key reassignment, the settlement-grace expiry timer, | ||
| // dispatch-loop cancel/abandon sites. Zeroing there would race a | ||
| // possibly-in-flight box.OpenAfterPrecomputation on the read loop; a | ||
| // corrupted decrypt is reported as an invalid encrypted chunk and | ||
| // triggers MarkUntrusted against an innocent provider. The deleted | ||
| // array is left intact for the GC. | ||
| // - The wholesale cap-reset in sharedKey must NOT zero either: the dropped | ||
| // entries may belong to live streams that will recompute and keep | ||
| // decrypting with the same private key. |
There was a problem hiding this comment.
🟡 [MEDIUM] 🧩 Complex lifecycle policy with multiple forget variants creates cognitive overhead
💡 Suggestion: Consider consolidating the forget/forgetAndZero distinction into a single method with a boolean parameter, or use a clearer naming convention that makes the safety constraints obvious at call sites
📊 Score: 3×4 = 12 · Category: over-abstraction
There was a problem hiding this comment.
Keeping the two named variants deliberately: the split IS the safety contract — forgetAndZero may only run on the owning provider read-loop goroutine (concurrent zeroing races an in-flight box.OpenAfterPrecomputation and would MarkUntrusted an innocent provider), while forget is safe anywhere. A forget(priv, zero bool) hides that constraint at call sites; the method name makes every call site auditable at a glance. The full policy is documented on the type.
| } | ||
|
|
||
| // cancelDispatchAndForget abandons a dispatched attempt (cancelDispatch: | ||
| // remove pending, idle the provider, WS-cancel, refund the attempt's top-up) | ||
| // and drops the attempt's memoized chunk-decryption key. Every dispatch-loop | ||
| // abandon site (speculative losers, failover retries, deadline/client-gone | ||
| // exits) funnels through here so an abandoned attempt cannot pin its session | ||
| // key in the cache until the cap-reset. Plain forget, NO zeroing: these sites | ||
| // run on the consumer/dispatch goroutine, not the provider read loop — the | ||
| // read loop may still be decrypting a late in-flight chunk with this key, and | ||
| // zeroing under it would corrupt that decrypt (and MarkUntrusted an innocent | ||
| // provider). See chunkKeyCache's zeroing policy. | ||
| func (s *Server) cancelDispatchAndForget(provider *registry.Provider, pr *registry.PendingRequest) { | ||
| s.cancelDispatch(provider, pr) | ||
| if pr != nil { | ||
| s.chunkKeys.forget(pr.SessionPrivKey) | ||
| } | ||
| } |
There was a problem hiding this comment.
🔵 [INFO] 🧩 New cancelDispatchAndForget wrapper adds indirection layer
💡 Suggestion: Consider whether the chunk key cleanup could be integrated directly into cancelDispatch rather than requiring a separate wrapper function
📊 Score: 2×2 = 4 · Category: additive complexity
There was a problem hiding this comment.
Agreed — folding the forget into cancelDispatch itself is the end state, but cancelDispatch lives in consumer.go, which carries in-flight v0.6.31 work on another branch (frozen for this PR series). The wrapper keeps all 38 dispatch.go sites covered now; absorbing it into cancelDispatch is on the documented post-v0.6.31 follow-up list (PR description, Deferred section).
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c94ed04536
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // entries pin their 32-byte priv arrays from the GC until that drop — | ||
| // bounded and secret-free (the SHARED key is what the cache zeroes; the | ||
| // priv array lives on in its PendingRequest anyway). | ||
| dead map[*[32]byte]struct{} |
There was a problem hiding this comment.
When any request reaches a terminal path, forget/forgetAndZero add its SessionPrivKey pointer to this map, and the map key keeps the 32-byte X25519 private-key object live until chunkKeyDeadMax later causes a wholesale reset. In normal successful traffic this means the new lifecycle cleanup still retains up to 8192 recently used response-decryption private keys in heap memory after their requests are gone, which undermines the crypto-hygiene goal of dropping per-request key material immediately; use a non-secret generation/tombstone key instead of retaining the private-key object itself.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in f008eee — tombstones are now keyed by the priv's address (uintptr), so they carry identity without pinning the 32-byte array; dead requests' key material is GC-collectable immediately. The documented trade is a benign ABA (a reused address suppresses caching for the new stream until the wholesale drop — correctness unaffected since sharedKey always returns a freshly computed key). The cap-reset test itself demonstrated the effect: its fillers now need runtime.KeepAlive to reach the cap.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f008eee3ac
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // and drops the attempt's memoized chunk-decryption key. Every dispatch-loop | ||
| // abandon site (speculative losers, failover retries, deadline/client-gone | ||
| // exits) funnels through here so an abandoned attempt cannot pin its session | ||
| // key in the cache until the cap-reset. Plain forget, NO zeroing: these sites |
There was a problem hiding this comment.
When /v1/completions or /v1/messages is in handleGenericInference, its pre-commit timeout/client-gone branches still remove pending directly (for example consumer.go:3710-3728 and the accepted wait at consumer.go:3781-3805) instead of using this new forget path. If a first chunk was already decrypted and buffered but the select takes the timeout/cancel case, decryptTextResponseChunk has inserted a cache entry; RemovePending makes the later provider terminal unable to find the request, so the session key remains cached until the wholesale reset. Add the same s.chunkKeys.forget(pr.SessionPrivKey)/wrapper to those generic abort paths.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6779cb5 — consumer.go is no longer frozen, so this closed properly: the forget moved INTO cancelDispatch itself (wrapper deleted), and every direct RemovePending abort site — including the generic /v1/completions and /v1/messages pre-commit timeout/client-gone/ErrorCh exits you flagged — now forgets via the RemovePending return value (13 sites converted). No abandoned attempt can pin its session key until the cap reset anymore.
| "refund_micro_usd", amount, | ||
| "error", err, | ||
| ) | ||
| s.ddIncr("billing.refund_failures", tags) |
There was a problem hiding this comment.
When store.Credit fails here, the function emits billing.refund_failures but then falls through to billing.reservation_refunds and billing.reservation_releases below. In the DB-failure case this change is meant to surface, those counters now report a successful ledger refund/release even though the consumer remains debited, making refund dashboards and release/refund reconciliation undercount the failure; return after the failure metric or tag the later counters as attempts instead of successes.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 6779cb5 — the failure branch now returns before billing.reservation_refunds/reservation_releases, so refund dashboards only count actual ledger credits. Same treatment applied to all three reservation_extra_refund sites in consumer.go (also newly unfrozen), which additionally skip the ReservedMicroUSD reset on failure so a later settlement path can retry the refund.
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — 1 finding(s)
- 🔵 [INFO]
coordinator/api/chunk_key_cache.go:4— Import of unsafe package for address-based tombstoning- Suggestion: Consider if the ABA issue described in comments could be avoided with a different design that doesn't require unsafe pointer arithmetic
Performance — 2 finding(s) (1 blocking)
- 🟡 [MEDIUM]
coordinator/api/chunk_key_cache.go:110-131— X25519 computation performed under mutex lock- Suggestion: Move the e2e.PrecomputeSharedKey call outside the lock, then reacquire the lock to check tombstones and update the cache. This prevents blocking all cache operations during the ~40-60µs X25519 computation.
- 🔵 [INFO]
coordinator/api/chunk_key_cache.go:172-182— forgetPeer iterates over all cache entries under lock- Suggestion: Consider adding early termination if the number of entries to check exceeds a threshold, or document that this is acceptable given the chunkKeyCacheMax bound of 8192 entries.
Type_diligence — 1 finding(s)
- 🔵 [INFO]
coordinator/api/chunk_key_cache.go:4— unsafe package imported but only used for pointer-to-uintptr conversion- Suggestion: Consider using reflect.ValueOf(priv).Pointer() instead of unsafe.Pointer for the tombstone key to avoid unsafe package dependency
Additive_complexity — ✅ No issues found
4 finding(s) total, 1 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
| @@ -2,6 +2,7 @@ package api | |||
|
|
|||
| import ( | |||
| "sync" | |||
There was a problem hiding this comment.
🔵 [INFO] 🔒 Import of unsafe package for address-based tombstoning
💡 Suggestion: Consider if the ABA issue described in comments could be avoided with a different design that doesn't require unsafe pointer arithmetic
📊 Score: 2×2 = 4 · Category: unsafe-usage
There was a problem hiding this comment.
Addressed in 6779cb5 — the unsafe import is gone: the tombstone key is now derived via reflect.ValueOf(priv).Pointer() with identical non-pinning semantics. On the broader design question: a pointer-keyed (pinning) map was the alternative without ABA, and it's exactly what the previous round flagged as retaining dead key material — the benign, documented ABA is the deliberate trade for immediate GC of key material.
| func (c *chunkKeyCache) forgetAndZero(priv *[32]byte) { | ||
| if priv == nil { | ||
| return | ||
| } | ||
| c.mu.Lock() | ||
| if e, ok := c.m[priv]; ok { | ||
| if e.shared != nil { | ||
| *e.shared = [32]byte{} | ||
| } | ||
| delete(c.m, priv) | ||
| } |
There was a problem hiding this comment.
🔵 [INFO] ⚡ forgetPeer iterates over all cache entries under lock
💡 Suggestion: Consider adding early termination if the number of entries to check exceeds a threshold, or document that this is acceptable given the chunkKeyCacheMax bound of 8192 entries.
📊 Score: 2×3 = 6 · Category: unbounded_loop
There was a problem hiding this comment.
This is already documented on forgetPeer: 'O(len(m)) ≤ chunkKeyCacheMax under the lock, on rare disconnects' — the map is hard-capped at 8192 entries and the sweep runs only on provider disconnect, so the worst case is a bounded, rare, in-memory scan. Early termination isn't possible: multiple requests can share the peer key, so the sweep must visit every entry.
…lassification, writer deadline races 1. releaseInitialReservation no longer swallows failed refund credits: ERROR log + billing.refund_failures metric (platform-fee-path policy). 2. chunkKeys session-key cache is now forgotten on EVERY terminal path: retry reassignment, settlement grace expiry, provider disconnect (new Provider.PendingSessionKeys snapshot taken before Disconnect wipes the pending map), and all 38 dispatch.go cancelDispatch sites via a cancelDispatchAndForget wrapper. Keys are zeroed only on read-loop- goroutine forgets; cross-goroutine forgets delete without zeroing — zeroing there races an in-flight decrypt and a corrupted open would MarkUntrusted an innocent provider (policy documented on the type). 3. Consumer-cancel terminals (overflow-499, client cancel) no longer book as provider_error_after_commit in route outcomes: shared terminal_class.go constants/helper replace the magic 499/'request cancelled' literals; postCommitProviderErrorOutcome short-circuits to partial_success/consumer_cancel_after_commit, AdmittedButFailed=false. Real 502s still book as provider faults (pinned). 4. provider_writer deadline/shutdown races: writeDeadline's value is now a CAS token — writeFrame and the watchdog race to claim each expired frame and exactly one wins; a write that lost the claim returns errProviderWriteTimeout even on nil err so dispatch fails over immediately instead of burning the TTFT budget on a dead socket. The sticky writeTimedOut latch is deleted (per-frame attribution). writeLane drains req.done before returning errProviderWriterStopped, fixing the ~50% select race that double-dispatched successfully-written frames. Deferred (in-flight v0.6.31 files): consumer.go refund swallows (:244, :717, :3305) and cancel/cleanup forget sites; registry.go control-lane rerouting for load/prefetch/desired-models.
#494 1. Re-cache race (dispatch.go cancel/retry forgets): forget/forgetAndZero now TOMBSTONE the priv pointer; sharedKey re-checks the tombstone set under the re-acquired lock before storing, so a forget racing the unlocked X25519 compute window — or a late straggler chunk decrypted after a cross-goroutine forget — hands back a working key without re-inserting an entry no terminal path remains to clean up. Safe because every dispatch attempt generates fresh session keys, so a forgotten priv is never legitimately re-cached. Tombstones are bounded (chunkKeyDeadMax, wholesale drop) and hold no secrets. 2. Registry-initiated disconnects (stale eviction, duplicate-serial eviction, forced removal call Registry.Disconnect directly, wiping pendingReqs BEFORE the read-loop defer can snapshot session keys): forgetProviderPendingKeys now also sweeps the cache by the provider's public key (new chunkKeyCache.forgetPeer), which needs no pending-map state — the defer always runs once CloseNow unblocks conn.Read. Plain delete, no zeroing, no tombstoning: a same-keypair replacement session (dup-serial) may be decrypting with a matching entry on its own read loop; it just recomputes on its next chunk. 3. holdForSettlement's IsReservationFinalized early return now forgets the chunk key first: relay stream-timeout/provider-incomplete refunds finalize without a read-loop terminal ever running, and the record is never parked, so nothing else would ever drop the entry. Addresses all three P2 comments on PR #494 (chatgpt-codex-connector).
…d session keys Codex P2 on c94ed04: the tombstone set was keyed by the *[32]byte pointer, so every terminated request's 32-byte X25519 private-key array stayed GC-reachable until the chunkKeyDeadMax wholesale reset — up to 8192 dead requests' key material retained, undermining the forget-on-terminal hygiene the cache exists for. Tombstones are now keyed by the priv's ADDRESS (uintptr): identity without pinning. Documented trade: after the GC reclaims a tombstoned array, a future SessionPrivKey allocated at the same address is falsely treated as dead — its stream still decrypts correctly (sharedKey always returns a freshly computed key), it just recomputes per chunk until the wholesale drop clears the stale tombstone. Perf-only, rare, self-healing. Also documents why a tombstone beats refcounting here (review ask): the racing readers are fire-and-forget late chunks with no owner to release against; a lock-only O(1) self-expiring set is the smallest mechanism. Test fix: the cap-reset test's fillers must be kept alive during the fill loop — with non-pinning tombstones the GC could otherwise reuse filler addresses and the set would never reach the cap (the ABA effect, demonstrated by the test itself).
…nd 3 consumer.go is no longer frozen (v0.6.31-era WIP has long since landed), so the deferred items from this PR's original scope now close: - cancelDispatch itself now drops the attempt's chunk-decryption key (plain forget, consumer-goroutine); the cancelDispatchAndForget wrapper is deleted and all 37 dispatch.go sites call cancelDispatch directly — consumer.go's own cancel paths get the forget for free (Codex P2 + earlier Centaur suggestion, both now unblocked). - Every direct RemovePending abort site (13: pre-commit timeout, client-gone, ErrorCh exits, reservation-failure paths in both the chat and generic pipelines) forgets via the RemovePending return value — covers the generic /v1/completions //v1/messages aborts Codex flagged. - All three reservation_extra_refund swallows now log ERROR + emit billing.refund_failures, and skip the success counters + the ReservedMicroUSD reset on failure so a later settlement can retry. - releaseInitialReservation returns after a failed refund instead of falling through to billing.reservation_refunds/releases — a failure counted as a success would undercount exactly what the alert surfaces (Codex P2). - Tombstone key now derived via reflect.ValueOf(priv).Pointer() — drops the unsafe import (review nit); identical non-pinning semantics. Rebased onto master b78c97f (post #512/#510/#502/#471); full api suite + registry/protocol green, race-clean on the touched suites.
Gajesh2007
force-pushed
the
fix/terminal-path-correctness
branch
from
f008eee to
6779cb5
Compare
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — ✅ No issues found
Performance — 1 finding(s) (1 blocking)
- 🟡 [MEDIUM]
coordinator/api/chunk_key_cache.go:110-131— X25519 computation (e2e.PrecomputeSharedKey) performed while holding mutex lock- Suggestion: Move the X25519 computation outside the critical section. Acquire lock, check cache/tombstone, release lock, compute key, re-acquire lock to store result.
Type_diligence — ✅ No issues found
Additive_complexity — ✅ No issues found
1 finding(s) total, 1 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
| shared := e2e.PrecomputeSharedKey(&pub, priv) | ||
|
|
||
| c.mu.Lock() | ||
| if _, gone := c.dead[tombstoneKey(priv)]; gone { | ||
| // The request hit a terminal path while the lock was dropped for the | ||
| // compute (or this is a late chunk of an already-forgotten request): | ||
| // hand the key back for this one decrypt but do NOT cache it — no | ||
| // terminal path remains to forget the entry, so it would pin the | ||
| // session key until the cap reset. | ||
| c.mu.Unlock() | ||
| return shared, nil | ||
| } | ||
| if c.m == nil || len(c.m) >= chunkKeyCacheMax { | ||
| // Wholesale reset (abandoned-entry safety net). Deliberately does NOT | ||
| // zero the dropped keys: entries may belong to live streams whose read | ||
| // loops are decrypting with them right now (see the zeroing policy on | ||
| // the type comment). | ||
| c.m = make(map[*[32]byte]chunkKeyEntry) | ||
| } | ||
| c.m[priv] = chunkKeyEntry{peerPub: peerPub, shared: shared} | ||
| c.mu.Unlock() | ||
| return shared, nil |
There was a problem hiding this comment.
🟡 [MEDIUM] ⚡ X25519 computation (e2e.PrecomputeSharedKey) performed while holding mutex lock
💡 Suggestion: Move the X25519 computation outside the critical section. Acquire lock, check cache/tombstone, release lock, compute key, re-acquire lock to store result.
📊 Score: 3×4 = 12 · Category: blocking_io
There was a problem hiding this comment.
This finding contradicts its own resolution: the same tool resolved the identical claim 21 seconds earlier on this same line ('Looks addressed in 6779cb5 — X25519 computation is now performed outside the mutex lock with proper tombstone checking after reacquiring the lock. Resolving.') and then re-filed it. The code does exactly what the suggestion asks — lock, check cache, unlock, compute, re-lock, tombstone-check, store — and has since the original design. No change needed; third time this has been filed.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6779cb558d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @@ -1 +1 @@ | |||
| Subproject commit a71c8018c02c5da6097bfd28c35c2b162c6fe5c6 | |||
| Subproject commit e5d0a94460d3ee10316a93690a87ea1a802c6b89 | |||
There was a problem hiding this comment.
This change advances the libs/mlx-swift-lm submodule even though the rest of the patch is coordinator-only terminal-path work. Anyone building the provider from this commit will pull a different MLX Swift LM dependency without any corresponding provider code or release-flow changes in this PR, which can unexpectedly change inference behavior or break reproducible provider builds; please revert this gitlink unless that dependency update is intentional and tested here.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — fixed in c5effc4. The rebase worktree's submodule was still checked out at the old master pin and a broad git add staged the stale gitlink, silently downgrading the engine from the v0.7.x fused-cache pin (a71c8018) to v0.6.30-era e5d0a944. Restored to master's exact pin; this PR is coordinator-only again (verified: git diff origin/master...HEAD -- libs/ is now empty).
6779cb5 unintentionally staged the worktree's stale submodule checkout, downgrading the engine pin from the v0.7.x fused-cache engine back to the v0.6.30-era e5d0a944 — this PR is coordinator-only and must not touch the provider engine (Codex catch).
ethenotethan
left a comment
ethenotethan
left a comment
There was a problem hiding this comment.
Automated Code Review — Layr-Labs/d-inference#
Verdict: REQUEST_CHANGES
Security — 1 finding(s)
- 🔵 [INFO]
coordinator/api/chunk_key_cache.go:89-90— reflect.ValueOf().Pointer() used for memory address extraction- Suggestion: Consider using unsafe.Pointer() with proper documentation of the ABA safety guarantees, or add explicit comments about the security implications of using reflection for pointer arithmetic
Performance — 2 finding(s) (1 blocking)
- 🟡 [MEDIUM]
coordinator/api/chunk_key_cache.go:110-121— X25519 computation performed while holding mutex lock- Suggestion: Move the e2e.PrecomputeSharedKey call outside the critical section to avoid blocking other cache operations during the ~40-60µs computation
- 🔵 [INFO]
coordinator/api/chunk_key_cache.go:172-180— O(n) iteration over cache entries under mutex lock- Suggestion: Consider batching the deletion or using a more efficient data structure if this becomes a bottleneck with large cache sizes
Type_diligence — ✅ No issues found
Additive_complexity — ✅ No issues found
3 finding(s) total, 1 blocking. Verdict: REQUEST_CHANGES.
🤖 Automated review by Centaur · DAR-186
|
|
||
| type chunkKeyEntry struct { |
There was a problem hiding this comment.
🔵 [INFO] 🔒 reflect.ValueOf().Pointer() used for memory address extraction
💡 Suggestion: Consider using unsafe.Pointer() with proper documentation of the ABA safety guarantees, or add explicit comments about the security implications of using reflection for pointer arithmetic
📊 Score: 2×3 = 6 · Category: unsafe-deserialization
| shared := e2e.PrecomputeSharedKey(&pub, priv) | ||
|
|
||
| c.mu.Lock() | ||
| if _, gone := c.dead[tombstoneKey(priv)]; gone { | ||
| // The request hit a terminal path while the lock was dropped for the | ||
| // compute (or this is a late chunk of an already-forgotten request): | ||
| // hand the key back for this one decrypt but do NOT cache it — no | ||
| // terminal path remains to forget the entry, so it would pin the | ||
| // session key until the cap reset. | ||
| c.mu.Unlock() | ||
| return shared, nil | ||
| } |
There was a problem hiding this comment.
🟡 [MEDIUM] ⚡ X25519 computation performed while holding mutex lock
💡 Suggestion: Move the e2e.PrecomputeSharedKey call outside the critical section to avoid blocking other cache operations during the ~40-60µs computation
📊 Score: 3×4 = 12 · Category: blocking_io
| func (c *chunkKeyCache) forgetAndZero(priv *[32]byte) { | ||
| if priv == nil { | ||
| return | ||
| } | ||
| c.mu.Lock() | ||
| if e, ok := c.m[priv]; ok { | ||
| if e.shared != nil { | ||
| *e.shared = [32]byte{} | ||
| } |
There was a problem hiding this comment.
🔵 [INFO] ⚡ O(n) iteration over cache entries under mutex lock
💡 Suggestion: Consider batching the deletion or using a more efficient data structure if this becomes a bottleneck with large cache sizes
📊 Score: 2×3 = 6 · Category: unbounded_iteration
|
Found 1 test failure on Blacksmith runners: Failure
|
2 participants
Summary
Terminal-path correctness fixes — the top follow-ups from the #491 review plus the P0 swallowed-refund bug. Four fixes, all about what happens when a request or connection ends: money that silently vanished, session keys that outlived their requests, backpressure aborts booked as provider faults, and two writer races at the deadline/shutdown boundary.
Coordinator-only; no wire-format changes;
consumer.go/registry.go/scheduler.gountouched (in-flight v0.6.31 work) — deferred sites listed below.Before / After
Behavior — request/connection terminals
flowchart LR subgraph Before A1[refund credit fails] --> B1["silently dropped<br/>(_ = store.Credit)"] C1[request abandoned:<br/>retry / cancel / disconnect / grace-expiry] --> D1["session key stays cached<br/>until 8192 cap-reset (maybe never)"] E1[overflow-499 / client cancel<br/>after commit] --> F1["route outcome: provider_error_after_commit<br/>AdmittedButFailed=true — pollutes calibration"] G1["write completes just as<br/>deadline expires"] --> H1["reported SUCCESS on a socket the<br/>watchdog killed -> TTFT burned + fake fault;<br/>writeTimedOut latch misattributes next frame;<br/>~50% 'stopped' for successfully-written frames"] end subgraph After A2[refund credit fails] --> B2["ERROR log + billing.refund_failures metric<br/>(alerting hook)"] C2[ALL terminal paths] --> D2["key forgotten immediately;<br/>zeroed when on the read-loop goroutine"] E2[consumer-cancel terminal] --> F2["partial_success / consumer_cancel_after_commit<br/>AdmittedButFailed=false"] G2[deadline boundary] --> H2["CAS handoff: exactly one winner;<br/>lost write returns errProviderWriteTimeout even on nil err;<br/>per-frame attribution; final done-drain before 'stopped'"] endCode — control flow
flowchart LR subgraph Before a1["cancelDispatch x38 (dispatch.go)<br/>grace expiry, disconnect, retry-reassign"] --> b1["no chunkKeys cleanup"] c1["magic literals: 499,<br/>'request cancelled' (2 sites)"] --> d1["postCommitProviderErrorOutcome:<br/>always provider fault"] e1["writeFrame: bare Load + sticky<br/>writeTimedOut bool"] --> f1["writeLane: random select between<br/>req.done and w.done"] end subgraph After a2["cancelDispatchAndForget wrapper;<br/>settlement onExpiry forget;<br/>forgetProviderPendingKeys (new<br/>registry/pending_iter.go accessor);<br/>forget-before-reassign"] --> b2["chunkKeyCache: forget (cross-goroutine)<br/>vs forgetAndZero (read-loop only,<br/>zeroing policy documented)"] c2["terminal_class.go:<br/>statusClientClosedRequest,<br/>isConsumerCancelTerminal()"] --> d2["postCommitProviderErrorOutcome<br/>short-circuits consumer-cancels"] e2["writeDeadline value IS the CAS token<br/>(nonzero->0); watchdog and writeFrame<br/>race to claim; writeTimedOut deleted"] --> f2["awaitWriteResult: non-blocking<br/>req.done drain before stoppedResult"] endThe fixes
1. Swallowed refund credit (money)
releaseInitialReservationdropped consumer refunds on store failure with_ =. Now: ERROR log (account/model/amount) +billing.refund_failuresmetric, matching the platform-fee path's documented "never swallow it" policy. No inline retry — this runs on latency-sensitive paths; the metric is the alerting hook.2.
chunkKeyslifecycle (crypto hygiene — highest-agreement #491 review finding)Only
handleComplete/handleInferenceErrorforgot the memoized X25519 shared key; every consumer-side terminal orphaned entries, pinning session keys until a cap-reset that may never fire. Now covered: retry reassignment (forget before theSessionPrivKeypointer overwrite), settlement grace expiry, provider disconnect (newProvider.PendingSessionKeys()snapshot taken beforeDisconnectwipes the map), and all 38cancelDispatchsites in dispatch.go via acancelDispatchAndForgetwrapper.Zeroing policy (deliberate, documented on the type): keys are zeroed only on read-loop-goroutine forgets (
handleComplete,handleInferenceError, disconnect cleanup — same goroutine as chunk decryption, no concurrent use possible). Cross-goroutine forgets (retry, grace timer) delete without zeroing: zeroing there races an in-flightbox.OpenAfterPrecomputation, and a corrupted decrypt wouldMarkUntrustedan innocent provider.3. Overflow-499 telemetry classification
The #491 overflow abort was correctly classified for reputation but
postCommitProviderErrorOutcomestill booked it asprovider_error_after_commit+AdmittedButFailed=true— polluting exactly the calibration signals the 499 was meant to protect. Now: sharedterminal_class.goconstants/helper (statusClientClosedRequest,isConsumerCancelTerminal) replace the magic literals in bothprovider.goandroute_outcome.go, and consumer-cancel terminals short-circuit topartial_success/consumer_cancel_after_commit/AdmittedButFailed=false. Real 502s still book as provider faults (pinned by test).4. Writer deadline/shutdown races
writeDeadline's value is now the CAS token (nonzero UnixNano → 0).writeFrameand the watchdog race to claim it; exactly one wins. A write that lost the claim returnserrProviderWriteTimeouteven on nil err (the socket is dead — fail over now, not after the TTFT budget). The stickywriteTimedOutlatch is deleted; attribution is per-frame.req.doneandw.donewere both ready, Go's select returnederrProviderWriterStopped~50% of the time for a frame that was on the wire — duplicate dispatch to a second provider. Now a final non-blockingreq.donedrain runs before the stopped fallback. Regression test proven to have teeth: reverting the fix fails 6/10 runs.Deferred (blocked on in-flight v0.6.31 files)
consumer.go:244/:717/:3305— same refund-swallow fixconsumer.gocancel/cleanup sites (cancelDispatchitself could absorb the forget once unfrozen;cleanupPendingclosures; generic-handler exits — full list in review notes)registry.go— routingSendLoadModel/SendPrefetchModel/SendDesiredModelsthrough the control laneVerification
gofmtclean,go vet ./...clean, fullgo test ./...green (api 104s)go test ./registry/ -race -count=3green; race-sensitive writer tests at-race -count=50Need help on this PR? Tag
/codesmithwith what you need. Autofix is disabled.