feat: make the router horizontally scalable - #100
Conversation
📝 WalkthroughWalkthroughThe change adds horizontal deployment controls, shared PostgreSQL authentication and admission state, bounded router capacity, distributed peer leases, a private Codex broker, shared HTTP clients, conditional replica workers, and related tests and deployment documentation. ChangesHorizontal deployment
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
caa10ed to
343338c
Compare
There was a problem hiding this comment.
Actionable comments posted: 9
🧹 Nitpick comments (14)
remote_codex.py (1)
33-36: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winUse an identity check when you select the client.
Line 33 computes ownership with
client is None, but line 34 selects withclient or .... If a caller passes a client object that evaluates as falsy (a test double that defines__bool__or__len__), the constructor creates an internal client while_owns_clientstaysFalse, soaclosenever closes it. Align both lines on the same identity test.♻️ Proposed fix
self._owns_client = client is None - self._client = client or httpx.AsyncClient( - limits=httpx.Limits(max_connections=100, - max_keepalive_connections=40)) + self._client = client if client is not None else httpx.AsyncClient( + limits=httpx.Limits(max_connections=100, + max_keepalive_connections=40))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@remote_codex.py` around lines 33 - 36, Update the client selection in the constructor to use the same identity check as `_owns_client`: reuse the caller-provided client whenever `client is not None`, and only create an `httpx.AsyncClient` when it is actually `None`, preserving consistent ownership and cleanup behavior.serve.py (2)
296-304: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winOne failing
acloseskips the remaining shutdown work.The loop at Lines 303-304 awaits each
aclose()directly. Ifprovider_http.aclose()raises,remote_codex.aclose()never runs and the exception escapes the lifespan during shutdown. Close every closeable independently.♻️ Proposed fix
for closeable in closeables: - await closeable.aclose() + with contextlib.suppress(Exception): + await closeable.aclose()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@serve.py` around lines 296 - 304, Update the lifespan cleanup block to close every item in closeables independently, ensuring one closeable’s aclose failure does not prevent later closeables from being closed or escape shutdown cleanup. Preserve the existing task cancellation and gathering behavior, and use the existing closeables collection.
249-253: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeduplicate the
_enabledflag parser.
provider_adapters/openai_compatible.pyLines 249-253 defines a byte-identical_enabledhelper. Two copies of the same environment-flag semantics drift easily; a change to the accepted false values in one file would silently not apply to the other. Move it to a shared module (for exampleprovider_adapters/common.py) and import it in both places.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@serve.py` around lines 249 - 253, Move the duplicate _enabled environment-flag parser from serve.py into a shared provider_adapters.common module, then import and reuse that helper in both serve.py and provider_adapters.openai_compatible.py. Remove the local duplicate definitions while preserving the existing default handling and false-value semantics.tests/test_codex_broker.py (1)
100-142: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the broker error paths in
RemoteCodexClient.The tests cover only 200 responses.
_http_errormaps 401/403 toauth_error, 429 torate_limit, and everything else toserver_error, andstreamreturnsstream_interruptedwhen deltas were already emitted before the wire ended. Those branches carry the router's fallback semantics, so a smallMockTransportcase per branch is valuable.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_codex_broker.py` around lines 100 - 142, Add MockTransport-based async tests for RemoteCodexClient covering _http_error mappings for 401/403 to auth_error, 429 to rate_limit, and other non-success responses to server_error, plus stream returning stream_interrupted when the wire ends after at least one emitted delta. Assert each returned error preserves the expected fallback semantics and close the AsyncClient in every test.codex_broker.py (1)
43-60: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider failing fast when
app.state.httpis not yet initialized.
create_appassignsapp.state.http = clientat line 60, so the value isNoneuntil the lifespan runs. If any caller mounts this app without running the lifespan (for example a bareTestClient(app)call without the context manager),/v1/calland/v1/streampassclient=Noneto the backends, which then create a per-requesthttpx.AsyncClient. That silently loses pooling instead of erroring. A short guard in the endpoints, or a distinct sentinel, makes the misuse visible.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@codex_broker.py` around lines 43 - 60, Ensure the HTTP-dependent endpoints fail fast when app.state.http has not been initialized by the lifespan, instead of passing None to backend calls that create per-request clients. Add the guard in the /v1/call and /v1/stream endpoint handlers, using the existing application state and returning the established internal-error response or exception path.tests/test_antseed_concurrency.py (2)
67-73: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueClear
_PEER_GATE_TIMESTAMPSwith the gate dict.The fixture resets
H._PEER_GATESonly. Stale timestamps for removed peers survive between tests, so the two module dicts drift apart and the eviction sweep sees entries with no gate.♻️ Proposed refactor
H._PEER_GATES.clear() + H._PEER_GATE_TIMESTAMPS.clear() yield H._PEER_GATES.clear() + H._PEER_GATE_TIMESTAMPS.clear()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_antseed_concurrency.py` around lines 67 - 73, Update the _clear_gates fixture to clear H._PEER_GATE_TIMESTAMPS alongside H._PEER_GATES both before yielding and during teardown, keeping the two peer-state dictionaries synchronized between tests.
193-195: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd cases for the store-unavailable and saturated classifications.
The new tests cover the happy path and lease cleanup. The error contract is untested. Patch
host_store.try_acquire_peer_leaseto return(False, False)and assert the result isnetwork_errorwith the "shared peer gate unavailable" detail. Add a second case returning(False, True)with a shortPEER_GATE_WAIT_Sand assertrate_limit. Both cases need no database. Do you want me to draft them?🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_antseed_concurrency.py` around lines 193 - 195, Extend the tests around the peer concurrency classification flow to cover unavailable and saturated outcomes. Mock host_store.try_acquire_peer_lease to return (False, False) and assert network_error with “shared peer gate unavailable”; add a separate case returning (False, True), configure a short PEER_GATE_WAIT_S, and assert rate_limit. Keep both cases database-free.provider_adapters/openai_compatible.py (2)
215-229: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueDocument or key the cached lock per event loop.
_PEER_GATE_LOCKand_PEER_GATESlive for the process lifetime. If a second event loop uses them (common in test suites that create one loop per test), a contendedacquire()on the cached lock can fail with a cross-loop future error, and semaphores created in a closed loop are reused. Either key both caches byid(asyncio.get_running_loop()), or state the single-loop assumption in the docstring so the constraint is explicit.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@provider_adapters/openai_compatible.py` around lines 215 - 229, Update _peer_gate_lock and the related _PEER_GATES cache to be keyed per running event loop, using the loop identity so locks and semaphores are never reused across loops; preserve lazy initialization and ensure stale per-loop entries can still be evicted as intended.
359-372: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBack off between lease attempts instead of polling every 100 ms.
Each waiter runs one
try_acquire_peer_leasetransaction per 100 ms, and that function takespg_advisory_xact_lockon the same peer key plus a delete and a count. When a popular peer is saturated, all waiters serialize on that key and multiply database load precisely when the system is under pressure. Use jittered backoff bounded by the remaining budget.♻️ Proposed refactor
deadline = time.monotonic() + wait_s + delay = 0.02 while True: @@ remaining = deadline - time.monotonic() if remaining <= 0: return None, "saturated" - await asyncio.sleep(min(0.1, remaining)) + await asyncio.sleep(min(delay * (0.5 + random.random()), remaining)) + delay = min(0.5, delay * 2)Add
import randomat the top of the module.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@provider_adapters/openai_compatible.py` around lines 359 - 372, Update the lease-wait loop containing try_acquire_peer_lease to use jittered exponential backoff instead of polling every 100 ms. Add the random module import, increase the delay after each failed acquisition up to a reasonable cap, and bound every sleep by the remaining deadline so saturated waiters reduce database contention without exceeding wait_s.host_store.py (2)
843-857: 📐 Maintainability & Code Quality | 🔵 TrivialConsider a periodic sweep for abandoned peer leases.
try_acquire_peer_leasedeletes expired rows only for the peer that is being acquired. If a peer stops receiving traffic, its expired rows stay inpeer_concurrency_leasesindefinitely. The row count is bounded by distinct peers multiplied by their cap, so this is small, but a scheduled delete onexpires_at <= now()keeps the table and its index clean. The existing retention sweep in this module is a natural owner.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@host_store.py` around lines 843 - 857, Extend the existing retention sweep in this module to periodically delete all rows from peer_concurrency_leases whose expires_at is at or before the current database time, rather than relying only on try_acquire_peer_lease. Keep the per-peer cleanup in try_acquire_peer_lease unchanged and place the global cleanup in the established scheduled sweep.
789-808: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winUse the database clock for bucket timestamps.
stampcomes from each replica's wall clock. With several replicas, clock skew changes the refill amount: a replica whose clock is behind stores an olderupdated_at, and the next caller'smax(0.0, stamp - row[1])clamps to0.0, so the elapsed time is lost. A replica whose clock is ahead grants extra tokens. The peer-lease functions in the same file already useEXTRACT(EPOCH FROM clock_timestamp()), so the two shared-state mechanisms disagree on the time authority.Keep the
nowparameter as a test override, and use the server clock when it is not supplied.♻️ Suggested change to use the server clock
- stamp = float(time.time() if now is None else now) refill_per_s = rate / 60.0 with _get_pool().connection() as conn: + stamp = (float(now) if now is not None else float(conn.execute( + "SELECT EXTRACT(EPOCH FROM clock_timestamp())").fetchone()[0])) # Establish the row before locking it. Concurrent first requests # serialize on the unique key via ON CONFLICT, then FOR UPDATE.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@host_store.py` around lines 789 - 808, Update the rate-bucket timestamp logic around the refill operation to use the database server clock when now is not provided, matching the peer-lease functions’ EXTRACT(EPOCH FROM clock_timestamp()) approach. Preserve now as a test override, and ensure the same server-derived timestamp is used for refill calculations, inserts, and updates.auth_proxy.py (1)
714-728: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftConsider a short-TTL cache for the shared auth lookups.
Each authenticated request now makes two PostgreSQL round trips here (
consumer_for_digestandget_consumer_key), and_rate_okadds a third with a row lock. The host-store pool is bounded toHOST_STORE_POOL_MAX_SIZE(default 8), so a burst of requests queues on the pool and can turn into pool timeouts under load. Unauthenticated traffic with random bearer tokens also reaches the digest query.A small TTL cache (one to two seconds) for the digest-to-consumer and consumer-to-metadata results keeps revocation latency acceptable and removes most of the per-request round trips. Keep the fail-closed behavior on a store error, and do not cache negative store results.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@auth_proxy.py` around lines 714 - 728, Add a one-to-two-second TTL cache around the shared lookups in the caller authentication flow, covering digest-to-consumer resolution and consumer-to-metadata retrieval before _rate_ok runs. Cache only successful store results, never negative or error responses, and preserve fail-closed behavior by returning the existing *_auth_unavailable responses when lookups fail; use the existing host_store and _normalize_consumer_record flow without changing revocation handling.tests/test_host_store.py (1)
238-242: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSeed a record so the atomicity claim is tested.
The test asserts an empty table after the rejected write. The fixture already starts empty, so the assertion passes even if no write were attempted. Write a record first. Then the empty-versus-preserved distinction proves that the rejected batch did not execute the
DELETE FROM consumer_keys.♻️ Suggested test change
def test_consumer_digest_validation_fails_as_one_write(store): + assert store.set_consumer_keys({"kept": {"status": "active"}}) is True assert store.set_consumer_keys( {"crm": {"status": "active"}}, key_digests={"not-a-sha256": "crm"}) is False - assert store.get_consumer_keys() == ({}, True) + assert store.get_consumer_keys() == ({"kept": {"status": "active"}}, True)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_host_store.py` around lines 238 - 242, Update test_consumer_digest_validation_fails_as_one_write to seed an existing consumer key record before calling set_consumer_keys with the invalid digest, then assert that the rejected operation preserves the seeded record while retaining the failure result.tests/test_auth_proxy_dashboard_full.py (1)
1843-1844: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMock
host_store.consumer_for_digestfor unknown-token requests.This test has no database fixture. An unavailable store makes the unknown-token request return 503 instead of the asserted 401.
Suggested test change
monkeypatch.setattr(host_store, "get_consumer_key", lambda _consumer: ({"status": "active"}, True)) + monkeypatch.setattr(host_store, "consumer_for_digest", + lambda _digest: (None, True))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_auth_proxy_dashboard_full.py` around lines 1843 - 1844, Update the unknown-token test around the existing host_store.get_consumer_key mock to also mock host_store.consumer_for_digest, returning the unavailable/absent consumer result expected for an unknown token so the request reaches the 401 assertion instead of failing with 503.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@auth_proxy.py`:
- Around line 3940-3948: The request flow currently consumes a rate token in
_rate_ok before _capacity_acquire can reject overloaded requests. Reorder the
checks to acquire capacity before performing _rate_ok, or otherwise refund the
token when the router_overloaded response is returned, ensuring shed 503
requests do not consume caller quota while preserving normal rate-limit
behavior.
In `@codex_backend.py`:
- Around line 347-369: Restructure both streaming call sites so
`client.stream("POST", url, ...)` is entered through `AsyncExitStack` before
applying the first-output timeout, then apply `before_first_output` to the first
`resp.aiter_lines()` read and explicitly close `resp` when that timeout expires.
Update `codex_backend.py` lines 347-369 around the existing `owns_client`,
`AsyncExitStack`, and `before_first_output` flow, and apply the same change in
`streaming.py` lines 99-105; preserve shared-client reuse and existing
timeout/error notification behavior.
In `@Dockerfile.antseed`:
- Around line 6-16: Update the dependency installation in the
antseed-dependencies build stage to include the cmake package alongside python3,
make, and g++. Keep the runtime stage unchanged and preserve the existing
package cleanup.
In `@docs/HORIZONTAL-SCALING.md`:
- Around line 43-56: Update the horizontal-scaling invariants documentation to
include per-role values for RUN_SOURCE_REFRESHERS, CODEX_BROKER_URL, and
DISTRIBUTED_PEER_GATES. Explicitly state whether API replicas enable source
refreshers and how their paid or rate-limited polling is controlled, along with
the required broker URL and peer-gate settings for control and API replicas.
In `@host_store.py`:
- Around line 158-165: The consumer digest and record tables can become
inconsistent, and concurrent full-table replacement can overwrite another
replica’s updates. In host_store.py lines 158-165, enforce
consumer_key_digests.consumer integrity with a cascading foreign key or delete
digests for consumers absent from the replacement set within the same
transaction; in auth_proxy.py lines 3741-3751, replace the full-table
read-modify-write with per-consumer upserts/deletes or hold
pg_advisory_xact_lock across the read and write.
In `@provider_adapters/openai_compatible.py`:
- Around line 626-632: Update the cleanup logic in the surrounding provider
method’s finally block so slot.release() executes even when the owned client’s
aclose() raises. Isolate the client close and capacity-slot release into
independent cleanup paths, while preserving ownership checks and the existing
release behavior for non-null slots.
- Around line 232-264: Update the peer-capacity lifecycle in
_acquire_peer_capacity and _PeerCapacitySlot.release to maintain a per-peer
_PEER_GATE_HOLDERS count, incrementing only after gate.acquire succeeds and
decrementing on release. During _peer_gate eviction, skip entries with active
holders so an in-use semaphore is never removed; preserve the existing
cached-by-peer behavior for cap changes.
- Around line 304-329: Update _renew_loop to retry failed renew_peer_lease calls
a bounded number of times before returning, while preserving the existing
renewal interval and stopping once renewal succeeds or retries are exhausted. In
release, await the cancelled heartbeat while suppressing both CancelledError and
any other heartbeat exception, ensuring lease/semaphore cleanup continues and
release never re-raises heartbeat task failures.
In `@tests/test_antseed_concurrency.py`:
- Around line 181-192: Increase the PEER_GATE_WAIT_S value in
test_distributed_peer_gate_uses_shared_store to provide enough time for all four
request waves to complete under loaded CI conditions, while keeping the
shared-store concurrency assertions unchanged.
---
Nitpick comments:
In `@auth_proxy.py`:
- Around line 714-728: Add a one-to-two-second TTL cache around the shared
lookups in the caller authentication flow, covering digest-to-consumer
resolution and consumer-to-metadata retrieval before _rate_ok runs. Cache only
successful store results, never negative or error responses, and preserve
fail-closed behavior by returning the existing *_auth_unavailable responses when
lookups fail; use the existing host_store and _normalize_consumer_record flow
without changing revocation handling.
In `@codex_broker.py`:
- Around line 43-60: Ensure the HTTP-dependent endpoints fail fast when
app.state.http has not been initialized by the lifespan, instead of passing None
to backend calls that create per-request clients. Add the guard in the /v1/call
and /v1/stream endpoint handlers, using the existing application state and
returning the established internal-error response or exception path.
In `@host_store.py`:
- Around line 843-857: Extend the existing retention sweep in this module to
periodically delete all rows from peer_concurrency_leases whose expires_at is at
or before the current database time, rather than relying only on
try_acquire_peer_lease. Keep the per-peer cleanup in try_acquire_peer_lease
unchanged and place the global cleanup in the established scheduled sweep.
- Around line 789-808: Update the rate-bucket timestamp logic around the refill
operation to use the database server clock when now is not provided, matching
the peer-lease functions’ EXTRACT(EPOCH FROM clock_timestamp()) approach.
Preserve now as a test override, and ensure the same server-derived timestamp is
used for refill calculations, inserts, and updates.
In `@provider_adapters/openai_compatible.py`:
- Around line 215-229: Update _peer_gate_lock and the related _PEER_GATES cache
to be keyed per running event loop, using the loop identity so locks and
semaphores are never reused across loops; preserve lazy initialization and
ensure stale per-loop entries can still be evicted as intended.
- Around line 359-372: Update the lease-wait loop containing
try_acquire_peer_lease to use jittered exponential backoff instead of polling
every 100 ms. Add the random module import, increase the delay after each failed
acquisition up to a reasonable cap, and bound every sleep by the remaining
deadline so saturated waiters reduce database contention without exceeding
wait_s.
In `@remote_codex.py`:
- Around line 33-36: Update the client selection in the constructor to use the
same identity check as `_owns_client`: reuse the caller-provided client whenever
`client is not None`, and only create an `httpx.AsyncClient` when it is actually
`None`, preserving consistent ownership and cleanup behavior.
In `@serve.py`:
- Around line 296-304: Update the lifespan cleanup block to close every item in
closeables independently, ensuring one closeable’s aclose failure does not
prevent later closeables from being closed or escape shutdown cleanup. Preserve
the existing task cancellation and gathering behavior, and use the existing
closeables collection.
- Around line 249-253: Move the duplicate _enabled environment-flag parser from
serve.py into a shared provider_adapters.common module, then import and reuse
that helper in both serve.py and provider_adapters.openai_compatible.py. Remove
the local duplicate definitions while preserving the existing default handling
and false-value semantics.
In `@tests/test_antseed_concurrency.py`:
- Around line 67-73: Update the _clear_gates fixture to clear
H._PEER_GATE_TIMESTAMPS alongside H._PEER_GATES both before yielding and during
teardown, keeping the two peer-state dictionaries synchronized between tests.
- Around line 193-195: Extend the tests around the peer concurrency
classification flow to cover unavailable and saturated outcomes. Mock
host_store.try_acquire_peer_lease to return (False, False) and assert
network_error with “shared peer gate unavailable”; add a separate case returning
(False, True), configure a short PEER_GATE_WAIT_S, and assert rate_limit. Keep
both cases database-free.
In `@tests/test_auth_proxy_dashboard_full.py`:
- Around line 1843-1844: Update the unknown-token test around the existing
host_store.get_consumer_key mock to also mock host_store.consumer_for_digest,
returning the unavailable/absent consumer result expected for an unknown token
so the request reaches the 401 assertion instead of failing with 503.
In `@tests/test_codex_broker.py`:
- Around line 100-142: Add MockTransport-based async tests for RemoteCodexClient
covering _http_error mappings for 401/403 to auth_error, 429 to rate_limit, and
other non-success responses to server_error, plus stream returning
stream_interrupted when the wire ends after at least one emitted delta. Assert
each returned error preserves the expected fallback semantics and close the
AsyncClient in every test.
In `@tests/test_host_store.py`:
- Around line 238-242: Update test_consumer_digest_validation_fails_as_one_write
to seed an existing consumer key record before calling set_consumer_keys with
the invalid digest, then assert that the rejected operation preserves the seeded
record while retaining the failure result.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 65360688-2c9f-4581-956f-93fbe0460587
📒 Files selected for processing (18)
.env.exampleDockerfile.antseedauth_proxy.pycodex_backend.pycodex_broker.pydocs/DEPLOY.mddocs/HORIZONTAL-SCALING.mdhost_store.pyprovider_adapters/openai_compatible.pyremote_codex.pyserve.pystreaming.pytests/test_antseed_concurrency.pytests/test_auth_proxy_dashboard_full.pytests/test_auth_proxy_scaling.pytests/test_codex_broker.pytests/test_host_store.pytests/test_sources.py
| if not await _capacity_acquire(): | ||
| _record_reject(reason="router_overloaded", path="/" + path, | ||
| caller=caller, status=503, route=requested_route) | ||
| return JSONResponse( | ||
| status_code=503, | ||
| headers={"Retry-After": str(max(1, math.ceil(CAPACITY_QUEUE_TIMEOUT_S)))}, | ||
| content={"error": { | ||
| "message": "router capacity is temporarily exhausted; retry shortly", | ||
| "type": "server_error", "code": "router_overloaded"}}) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
The shed path consumes the caller's rate token.
_rate_ok consumes one global token at Line 3889, and _capacity_acquire runs later at Line 3940. When the pod sheds load, the caller receives 503 and the token is already gone. Under sustained overload a consumer loses quota for requests that were never served, and the retry advised by Retry-After spends another token.
Consider acquiring capacity before the rate check, or refunding the token on the shed path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@auth_proxy.py` around lines 3940 - 3948, The request flow currently consumes
a rate token in _rate_ok before _capacity_acquire can reject overloaded
requests. Reorder the checks to acquire capacity before performing _rate_ok, or
otherwise refund the token when the router_overloaded response is returned,
ensuring shed 503 requests do not consume caller quota while preserving normal
rate-limit behavior.
| owns_client = client is None | ||
| try: | ||
| async with httpx.AsyncClient(timeout=timeout) as c: | ||
| async with AsyncExitStack() as stack: | ||
| async with AsyncExitStack() as stack: | ||
| c = (await stack.enter_async_context( | ||
| httpx.AsyncClient(timeout=timeout)) | ||
| if owns_client else client) | ||
| stream_kwargs = { | ||
| "json": body, | ||
| "headers": headers, | ||
| } | ||
| # A shared client has no request-specific default; an owned | ||
| # client was constructed with this timeout above. Keeping the | ||
| # latter call shape also supports lightweight test transports. | ||
| if not owns_client: | ||
| stream_kwargs["timeout"] = timeout | ||
| try: | ||
| resp = await before_first_output(stack.enter_async_context( | ||
| c.stream("POST", url, **stream_kwargs)), | ||
| first_timeout_s, t0, _saw_output) | ||
| except (asyncio.TimeoutError, TimeoutError): | ||
| if not status_seen: | ||
| _notify(0) | ||
| return _timeout_err() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Cancelling stack.enter_async_context(...stream(...)) loses the response handle. Both call sites wrap the stream-context entry in before_first_output, which uses asyncio.wait_for. If the first-token bound expires after the request is sent but before the entry returns, the response is never registered on the AsyncExitStack, so it is never closed. Both sites now receive the shared provider_http client from serve.py, so each occurrence checks out a pooled connection permanently instead of dying with a per-call client.
codex_backend.py#L347-L369: enter the stream context outside the timeout bound, then applybefore_first_outputto the firstaiter_linesread; closerespexplicitly if the bound expires.streaming.py#L99-L105: apply the same restructuring to theclient.stream("POST", url, ...)entry so the response is closed on the timeout path.
📍 Affects 2 files
codex_backend.py#L347-L369(this comment)streaming.py#L99-L105
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@codex_backend.py` around lines 347 - 369, Restructure both streaming call
sites so `client.stream("POST", url, ...)` is entered through `AsyncExitStack`
before applying the first-output timeout, then apply `before_first_output` to
the first `resp.aiter_lines()` read and explicitly close `resp` when that
timeout expires. Update `codex_backend.py` lines 347-369 around the existing
`owns_client`, `AsyncExitStack`, and `before_first_output` flow, and apply the
same change in `streaming.py` lines 99-105; preserve shared-client reuse and
existing timeout/error notification behavior.
| FROM node:22-slim AS antseed-dependencies | ||
|
|
||
| # Native dependencies normally download a prebuilt binary. Keep a compiler | ||
| # fallback in the build stage so a transient prebuild-host failure cannot make | ||
| # releases undeployable; none of this toolchain is copied into the runtime. | ||
| RUN apt-get update \ | ||
| && apt-get install -y --no-install-recommends python3 make g++ \ | ||
| && npm install -g @antseed/cli@0.1.128 pg@8.16.3 \ | ||
| && rm -rf /var/lib/apt/lists/* | ||
|
|
||
| FROM node:22-slim |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
@antseed/cli supported Node.js versions engines
💡 Result:
The @antseed/cli package requires Node.js version 20 or higher [1][2]. The project documentation explicitly lists the runtime requirement as "Node.js >= 20" [1][2]. Additionally, deployment scripts for the provider infrastructure default to using Node.js version 22 [3].
Citations:
- 1: https://github.com/AntSeed/antseed/tree/b12b80fecbdd482eee313ac548494ac1dfbbccfc
- 2: https://github.com/AntSeed/antseed/tree/96548e525ec4884c83f1791958c34fe10c76107b
- 3: https://github.com/AntSeed/antseed/blob/main/scripts/deploy-provider.sh
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- Dockerfile and package references ---'
rg -n -C 3 '`@antseed/`(cli|node)|pg@|node:2[25]|antseed' Dockerfile.antseed package.json package-lock.json npm-shrinkwrap.json pnpm-lock.yaml yarn.lock 2>/dev/null || true
printf '%s\n' '--- npm metadata for the pinned packages ---'
for spec in '`@antseed/cli`@0.1.128' 'pg@8.16.3' '`@antseed/node`'; do
echo "### $spec"
npm view "$spec" version engines dependencies optionalDependencies peerDependencies dist-tags --json 2>/dev/null || true
doneRepository: genlayerlabs/unhardcoded
Length of output: 4534
🌐 Web query:
"@antseed/node" 0.1.128 Node 22 ABI prebuild
💡 Result:
The package @antseed/node is a core SDK for the AntSeed peer-to-peer AI services network, designed for TypeScript and Node.js environments (with a minimum requirement of Node.js 20.0.0) [1][2]. Regarding your query about version 0.1.128 and Node 22 ABI prebuilds: 1. Package Version: The version 0.1.128 is an older release of @antseed/node [3]. Development of the package has continued significantly beyond this version, with current versions reaching into the 0.2.x series [3][4]. 2. Node 22 ABI Compatibility: Node.js 22 uses ABI (Application Binary Interface) version 127 [5]. If @antseed/node or any of its dependencies include native addons, they must be compiled against ABI 127 to be compatible with Node 22 [6]. 3. Prebuilds: Whether @antseed/node version 0.1.128 specifically contains prebuilt binaries for ABI 127 depends on whether those binaries were bundled at the time of that release [6]. Given that Node 22 was released after 0.1.128, it is highly likely that this specific older version does not contain prebuilt binaries for ABI 127 [3][5]. If you are experiencing compatibility issues when running @antseed/node 0.1.128 on Node 22, it is recommended that you update to the latest available version of the package to benefit from updated dependencies and broader platform support [3][4]. If you must use 0.1.128 and encounter native module errors, you may need to rebuild the package from source using node-gyp or similar tools to generate the correct ABI 127 binaries for your environment [6].
Citations:
- 1: https://github.com/AntSeed/antseed/blob/fb5f3ea0b3b8d7aeb8598fb32d24d7ee2bea6342/packages/node/README.md
- 2: https://github.com/antseed/antseed
- 3: https://github.com/AntSeed/antseed/blob/main/CHANGELOG.md
- 4: https://socket.dev/npm/package/@antseed/node/license/0.2.96
- 5: Support for Node v22 / ABI version 127 homebridge/node-pty-prebuilt-multiarch#42
- 6: https://nodejs.org/api/n-api.html
🌐 Web query:
site:npmjs.com/package/@antseed/node @antseed/node Node.js 22
💡 Result:
@antseed/node is the core protocol SDK for the AntSeed network, a decentralized, peer-to-peer (P2P) marketplace for AI inference [1][2][3]. Regarding its compatibility with Node.js 22: - The project's technical documentation specifies a runtime requirement of Node.js >= 20 [1][2][3]. - As Node.js 22 is a later version that maintains backward compatibility with Node.js 20, the package is compatible with Node.js 22, provided your environment meets the minimum Node.js 20 requirement [1][2][3]. Key details about the package: - Purpose: It handles P2P discovery, protocol implementation, metering, and payment processing for the AntSeed network [1][2][3]. - Tech Stack: It is built with TypeScript 5.x and utilizes ES modules [1][2][3]. - Usage: It serves as a core library within the AntSeed monorepo, which also includes CLI tools, provider plugins (e.g., for Anthropic, OpenAI, and local LLMs), and router infrastructure [1][2][3]. You can install it via npm using: npm install @antseed/node [4]
Citations:
- 1: https://github.com/AntSeed/antseed/tree/b12b80fecbdd482eee313ac548494ac1dfbbccfc
- 2: https://github.com/antseed/antseed
- 3: https://github.com/AntSeed/antseed/tree/ff879acb75b318f46b407d124b148a438f6ac08f
- 4: https://github.com/AntSeed/antseed/blob/fb5f3ea0b3b8d7aeb8598fb32d24d7ee2bea6342/packages/node/README.md
🏁 Script executed:
#!/bin/bash
set -eu
tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT
get_json() {
curl -fsSL "https://registry.npmjs.org/$1/$2"
}
for spec in \
'`@antseed/node/0.2.92`' \
'better-sqlite3/12.6.2' \
'node-datachannel/0.7.0' \
'keytar/7.9.0'
do
echo "### $spec"
get_json "$spec" | jq '{
name, version, engines, os, cpu, scripts, gypfile,
dependencies, optionalDependencies, peerDependencies,
binary, dist: {tarball, integrity}
}'
done
echo '### `@antseed/node` exact dependency metadata'
get_json '`@antseed/node/0.2.92`' | jq '{name, version, engines, dependencies, optionalDependencies}'
echo '### native package release metadata'
for spec in \
'better-sqlite3/12.6.2' \
'node-datachannel/0.7.0' \
'keytar/7.9.0'
do
tarball="$(get_json "$spec" | jq -r '.dist.tarball')"
name="${spec%%/*}"
curl -fsSL "$tarball" -o "$tmp/$name.tgz"
echo "## $spec package contents and native markers"
tar -tzf "$tmp/$name.tgz" | grep -E '(^|/)(prebuilds|build|src|binding\.gyp|prebuild-install|node-gyp|\.node)' | head -80 || true
doneRepository: genlayerlabs/unhardcoded
Length of output: 648
🏁 Script executed:
#!/bin/bash
set -eu
tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT
get_json() {
curl -fsSL "https://registry.npmjs.org/$1"
}
for spec in \
'`@antseed/node/0.2.92`' \
'better-sqlite3/12.6.2' \
'node-datachannel/0.7.0' \
'keytar/7.9.0'
do
echo "### $spec"
get_json "$spec" | jq '{
name, version, engines, os, cpu, scripts, gypfile,
dependencies, optionalDependencies, peerDependencies,
binary, dist: {tarball, integrity}
}'
done
echo '### `@antseed/node` exact dependency metadata'
get_json '`@antseed/node/0.2.92`' | jq '{name, version, engines, dependencies, optionalDependencies}'
echo '### native package release metadata'
for spec in \
'better-sqlite3/12.6.2' \
'node-datachannel/0.7.0' \
'keytar/7.9.0'
do
tarball="$(get_json "$spec" | jq -r '.dist.tarball')"
name="${spec%%/*}"
curl -fsSL "$tarball" -o "$tmp/$name.tgz"
echo "## $spec package contents and native markers"
tar -tzf "$tmp/$name.tgz" | grep -E '(^|/)(prebuilds|build|src|binding\.gyp|prebuild-install|node-gyp|\.node)' | head -80 || true
doneRepository: genlayerlabs/unhardcoded
Length of output: 6983
🏁 Script executed:
#!/bin/bash
set -eu
tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT
for spec in \
'`@antseed/node`@0.2.92' \
'better-sqlite3@12.6.2' \
'node-datachannel@0.7.0' \
'keytar@7.9.0'
do
name="${spec%@*}"
version="${spec##*@}"
echo "### $spec"
tarball="$(npm view "$spec" dist.tarball --json)"
echo "tarball=$tarball"
curl -fsSL "$tarball" -o "$tmp/pkg.tgz"
tar -xOf "$tmp/pkg.tgz" package/package.json | jq '{
name, version, engines, scripts, binary, dependencies, optionalDependencies
}'
done
echo '### better-sqlite3 v12.6.2 release assets'
curl -fsSL \
'https://api.github.com/repos/WiseLibs/better-sqlite3/releases/tags/v12.6.2' |
jq -r '.assets[].name' | grep -E 'node-v127|napi|linux.*x64|glibc' || true
echo '### node-datachannel v0.7.0 release assets'
curl -fsSL \
'https://api.github.com/repos/node-webrtc/node-datachannel/releases/tags/v0.7.0' |
jq -r '.assets[].name' | grep -E 'linux|napi|x64|node' || true
echo '### keytar v7.9.0 release assets'
curl -fsSL \
'https://api.github.com/repos/atom/node-keytar/releases/tags/v7.9.0' |
jq -r '.assets[].name' | grep -E 'linux|napi|x64|node' || trueRepository: genlayerlabs/unhardcoded
Length of output: 644
🏁 Script executed:
#!/bin/bash
set -u
tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT
fetch() {
curl --noproxy '*' -fsSL "$1"
}
for url in \
'https://registry.npmjs.org/@antseed%2fnode/0.2.92' \
'https://registry.npmjs.org/better-sqlite3/12.6.2' \
'https://registry.npmjs.org/node-datachannel/0.7.0' \
'https://registry.npmjs.org/keytar/7.9.0'
do
echo "### $url"
if json="$(fetch "$url")"; then
printf '%s' "$json" | jq '{
name, version, engines, scripts, binary, dependencies, optionalDependencies,
dist: {tarball, integrity}
}'
else
echo 'request failed'
fi
done
echo '### better-sqlite3 v12.6.2 release assets'
fetch 'https://api.github.com/repos/WiseLibs/better-sqlite3/releases/tags/v12.6.2' |
jq -r '.assets[].name' | grep -E 'node-v127|napi|linux.*x64|glibc' || true
echo '### node-datachannel v0.7.0 release assets'
fetch 'https://api.github.com/repos/node-webrtc/node-datachannel/releases/tags/v0.7.0' |
jq -r '.assets[].name' | grep -E 'linux|napi|x64|node' || true
echo '### keytar v7.9.0 release assets'
fetch 'https://api.github.com/repos/atom/node-keytar/releases/tags/v7.9.0' |
jq -r '.assets[].name' | grep -E 'linux|napi|x64|node' || trueRepository: genlayerlabs/unhardcoded
Length of output: 4181
🏁 Script executed:
#!/bin/bash
set -eu
tmp="$(mktemp -d)"
trap 'rm -rf "$tmp"' EXIT
for spec in \
'`@antseed/node/-/node-0.2.92.tgz`' \
'better-sqlite3/-/better-sqlite3-12.6.2.tgz' \
'node-datachannel/-/node-datachannel-0.7.0.tgz' \
'keytar/-/keytar-7.9.0.tgz'
do
url="https://registry.npmjs.org/$spec"
file="$tmp/$(basename "$spec")"
echo "### $url"
curl --noproxy '*' -fsSL "$url" -o "$file"
tar -xOf "$file" package/package.json | jq '{
name, version, repository, homepage, engines, scripts, binary,
dependencies, optionalDependencies
}'
done
echo '### available CMake-related executables'
command -v cmake || true
command -v cmake-js || true
command -v ninja || true
echo '### node-datachannel source-build references'
curl --noproxy '*' -fsSL \
'https://registry.npmjs.org/node-datachannel/-/node-datachannel-0.7.0.tgz' |
tar -xzOf - package/package.json |
jq -r '.scripts.install, .scripts._prebuild, .devDependencies'Repository: genlayerlabs/unhardcoded
Length of output: 4864
Add CMake to the build stage.
Node 22 satisfies the declared engines. better-sqlite3@12.6.2 explicitly supports Node 22, and node-datachannel@0.7.0 uses N-API 8. If its prebuild is unavailable, its fallback invokes cmake-js, but this image installs no CMake executable. Add cmake to the build-stage packages.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@Dockerfile.antseed` around lines 6 - 16, Update the dependency installation
in the antseed-dependencies build stage to include the cmake package alongside
python3, make, and g++. Keep the runtime stage unchanged and preserve the
existing package cleanup.
| - Codex credentials stay on the control PVC. `codex_broker.py` selects an | ||
| account and performs calls behind `CODEX_BROKER_TOKEN`; responses contain only | ||
| canonical call results and quota headers, never tokens. | ||
| - AntSeed API replicas use the private AntSeed Service URL. Only the control | ||
| replica runs the wallet keeper (`RUN_WALLET_KEEPER=1`); every API replica sets | ||
| it to `0`. | ||
| - AntSeed seller concurrency is enforced with renewable PostgreSQL leases on | ||
| both JSON and streaming calls. A peer's advertised cap therefore stays | ||
| global as replicas are added, and a crashed pod's lease expires automatically. | ||
| - Synthetic probes have one owner. API replicas set | ||
| `DASHBOARD_SYNTHETIC_PROBES_ENABLED=0`, so scaling does not multiply paid | ||
| requests. | ||
| - Startup cost backfill has one owner. API replicas set `RUN_COST_BACKFILL=0` | ||
| so a scale-out event cannot fan out the same maintenance query. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Document the remaining replica flags.
The invariants list names RUN_WALLET_KEEPER, DASHBOARD_SYNTHETIC_PROBES_ENABLED, and RUN_COST_BACKFILL, but this PR also adds RUN_SOURCE_REFRESHERS (serve.py Line 287), CODEX_BROKER_URL (serve.py Line 133), and DISTRIBUTED_PEER_GATES (provider_adapters/openai_compatible.py Line 347). An operator following only this page cannot tell whether API replicas must run the source refreshers, and the source-refresh loops also poll paid or rate-limited endpoints. Add the per-role value for each flag.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/HORIZONTAL-SCALING.md` around lines 43 - 56, Update the
horizontal-scaling invariants documentation to include per-role values for
RUN_SOURCE_REFRESHERS, CODEX_BROKER_URL, and DISTRIBUTED_PEER_GATES. Explicitly
state whether API replicas enable source refreshers and how their paid or
rate-limited polling is controlled, along with the required broker URL and
peer-gate settings for control and API replicas.
| """CREATE TABLE IF NOT EXISTS consumer_key_digests ( | ||
| digest TEXT PRIMARY KEY, | ||
| consumer TEXT NOT NULL, | ||
| created_at BIGINT NOT NULL, | ||
| updated_at BIGINT NOT NULL | ||
| )""", | ||
| "CREATE INDEX IF NOT EXISTS idx_consumer_key_digests_consumer" | ||
| " ON consumer_key_digests(consumer)", |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
consumer_keys and consumer_key_digests can disagree. set_consumer_keys replaces the whole consumer_keys table from a caller-supplied snapshot, while consumer_key_digests is maintained incrementally with upserts and prefix deletes. Nothing keeps the two tables consistent, so a digest can outlive its consumer record. consumer_for_digest then resolves the digest, get_consumer_key returns (None, True), and _normalize_consumer_record turns the missing record into an active consumer with default limits.
host_store.py#L158-L165: relate the digest table to the consumer record, for example withFOREIGN KEY (consumer) REFERENCES consumer_keys(consumer) ON DELETE CASCADE, or delete digests for consumers absent from the replacement set inside the same transaction.auth_proxy.py#L3741-L3751: stop the full-table read-modify-write that drops records written by another replica. Use per-consumer upsert and delete operations, or holdpg_advisory_xact_lockacross the read and the write.
📍 Affects 2 files
host_store.py#L158-L165(this comment)auth_proxy.py#L3741-L3751
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@host_store.py` around lines 158 - 165, The consumer digest and record tables
can become inconsistent, and concurrent full-table replacement can overwrite
another replica’s updates. In host_store.py lines 158-165, enforce
consumer_key_digests.consumer integrity with a cascading foreign key or delete
digests for consumers absent from the replacement set within the same
transaction; in auth_proxy.py lines 3741-3751, replace the full-table
read-modify-write with per-consumer upserts/deletes or hold
pg_advisory_xact_lock across the read and write.
| async def _peer_gate(peer_id: str, cap: int) -> asyncio.Semaphore: | ||
| """Return a per-peer concurrency semaphore, creating one if needed. | ||
|
|
||
| Uses a lock to prevent a race where two concurrent coroutines both see | ||
| ``None`` and create separate semaphores for the same peer_id — the | ||
| second would overwrite the first, silently dropping any coroutines | ||
| already waiting on it. | ||
|
|
||
| Old entries are evicted after ``_PEER_GATE_TTL`` seconds of inactivity | ||
| so that the dict does not grow without bound when AntSeed peers are | ||
| dynamically discovered and later retired. | ||
| """ | ||
| import time as _time | ||
|
|
||
| lock = _peer_gate_lock() | ||
| async with lock: | ||
| # Evict stale entries first (best-effort, not on every call to | ||
| # avoid quadratic sweep cost — only when the lock is already held | ||
| # and the dict is getting large). | ||
| if len(_PEER_GATES) > 256: | ||
| now = _time.monotonic() | ||
| stale = [k for k, ts in _PEER_GATE_TIMESTAMPS.items() | ||
| if now - ts > _PEER_GATE_TTL] | ||
| for k in stale: | ||
| _PEER_GATES.pop(k, None) | ||
| _PEER_GATE_TIMESTAMPS.pop(k, None) | ||
|
|
||
| sem = _PEER_GATES.get(peer_id) | ||
| if sem is None: | ||
| sem = asyncio.Semaphore(cap) | ||
| _PEER_GATES[peer_id] = sem | ||
| _PEER_GATE_TIMESTAMPS[peer_id] = _time.monotonic() | ||
| return sem |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Eviction can drop a gate that still has held permits and breach the seller cap.
_PEER_GATE_TIMESTAMPS[peer_id] is refreshed only at acquisition (Line 263). A long-running call (a stream can exceed _PEER_GATE_TTL = 600 s) leaves its peer entry idle by this measure. If _PEER_GATES then exceeds 256 entries, the sweep at Lines 251-257 removes that gate while permits are still held. The next request for the same peer creates a new asyncio.Semaphore(cap), so in-flight plus new calls exceed max_concurrency. The earlier holder later releases an orphaned semaphore.
Track outstanding holders and skip eviction while a gate is in use. Note that a changed cap for an existing peer is also ignored, which is a related consequence of caching by peer_id alone.
🐛 Proposed fix: refuse eviction for in-use gates
_PEER_GATE_TTL: float = 600.0
_PEER_GATE_TIMESTAMPS: dict[str, float] = {}
+_PEER_GATE_HOLDERS: dict[str, int] = {} if len(_PEER_GATES) > 256:
now = _time.monotonic()
stale = [k for k, ts in _PEER_GATE_TIMESTAMPS.items()
- if now - ts > _PEER_GATE_TTL]
+ if now - ts > _PEER_GATE_TTL
+ and _PEER_GATE_HOLDERS.get(k, 0) <= 0]
for k in stale:
_PEER_GATES.pop(k, None)
_PEER_GATE_TIMESTAMPS.pop(k, None)
+ _PEER_GATE_HOLDERS.pop(k, None)Increment _PEER_GATE_HOLDERS[peer_id] after a successful gate.acquire() in _acquire_peer_capacity, and decrement it in _PeerCapacitySlot.release().
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@provider_adapters/openai_compatible.py` around lines 232 - 264, Update the
peer-capacity lifecycle in _acquire_peer_capacity and _PeerCapacitySlot.release
to maintain a per-peer _PEER_GATE_HOLDERS count, incrementing only after
gate.acquire succeeds and decrementing on release. During _peer_gate eviction,
skip entries with active holders so an in-use semaphore is never removed;
preserve the existing cached-by-peer behavior for cap changes.
| async def _renew_loop(self) -> None: | ||
| import host_store | ||
|
|
||
| while True: | ||
| await asyncio.sleep(max(5.0, self._ttl_s / 3.0)) | ||
| ok = await asyncio.to_thread( | ||
| host_store.renew_peer_lease, | ||
| self._peer_id, self._lease_id, ttl_s=self._ttl_s) | ||
| if not ok: | ||
| return | ||
|
|
||
| async def release(self) -> None: | ||
| if self._released: | ||
| return | ||
| self._released = True | ||
| if self._heartbeat is not None: | ||
| self._heartbeat.cancel() | ||
| with suppress(asyncio.CancelledError): | ||
| await self._heartbeat | ||
| if self._peer_id and self._lease_id: | ||
| import host_store | ||
| await asyncio.to_thread( | ||
| host_store.release_peer_lease, | ||
| self._peer_id, self._lease_id) | ||
| elif self._semaphore is not None: | ||
| self._semaphore.release() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Make the heartbeat resilient and stop release() from re-raising heartbeat errors.
Two problems in this pair of methods:
- Line 312: one failed renewal ends the heartbeat permanently.
host_store.renew_peer_leasereturnsFalsefor a transient database failure as well as for a lost lease. The lease then expires while the call is still in flight, and another replica admits work over the seller cap. Retry a bounded number of times before giving up. - Lines 320-322:
release()awaits the heartbeat task and suppresses onlyCancelledError. If the task failed (for exampleimport host_storeraised, orasyncio.to_threadraised), the stored exception is re-raised here.release()is called fromfinallyblocks at Lines 461-462 and 631-632, so that exception replaces the provider result and skips the lease cleanup below it.
🐛 Proposed fix
async def _renew_loop(self) -> None:
import host_store
+ failures = 0
while True:
await asyncio.sleep(max(5.0, self._ttl_s / 3.0))
ok = await asyncio.to_thread(
host_store.renew_peer_lease,
self._peer_id, self._lease_id, ttl_s=self._ttl_s)
- if not ok:
- return
+ if ok:
+ failures = 0
+ continue
+ failures += 1
+ if failures >= 3:
+ return
async def release(self) -> None:
if self._released:
return
self._released = True
if self._heartbeat is not None:
self._heartbeat.cancel()
- with suppress(asyncio.CancelledError):
+ with suppress(BaseException):
await self._heartbeat📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async def _renew_loop(self) -> None: | |
| import host_store | |
| while True: | |
| await asyncio.sleep(max(5.0, self._ttl_s / 3.0)) | |
| ok = await asyncio.to_thread( | |
| host_store.renew_peer_lease, | |
| self._peer_id, self._lease_id, ttl_s=self._ttl_s) | |
| if not ok: | |
| return | |
| async def release(self) -> None: | |
| if self._released: | |
| return | |
| self._released = True | |
| if self._heartbeat is not None: | |
| self._heartbeat.cancel() | |
| with suppress(asyncio.CancelledError): | |
| await self._heartbeat | |
| if self._peer_id and self._lease_id: | |
| import host_store | |
| await asyncio.to_thread( | |
| host_store.release_peer_lease, | |
| self._peer_id, self._lease_id) | |
| elif self._semaphore is not None: | |
| self._semaphore.release() | |
| async def _renew_loop(self) -> None: | |
| import host_store | |
| failures = 0 | |
| while True: | |
| await asyncio.sleep(max(5.0, self._ttl_s / 3.0)) | |
| ok = await asyncio.to_thread( | |
| host_store.renew_peer_lease, | |
| self._peer_id, self._lease_id, ttl_s=self._ttl_s) | |
| if ok: | |
| failures = 0 | |
| continue | |
| failures += 1 | |
| if failures >= 3: | |
| return | |
| async def release(self) -> None: | |
| if self._released: | |
| return | |
| self._released = True | |
| if self._heartbeat is not None: | |
| self._heartbeat.cancel() | |
| with suppress(BaseException): | |
| await self._heartbeat | |
| if self._peer_id and self._lease_id: | |
| import host_store | |
| await asyncio.to_thread( | |
| host_store.release_peer_lease, | |
| self._peer_id, self._lease_id) | |
| elif self._semaphore is not None: | |
| self._semaphore.release() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@provider_adapters/openai_compatible.py` around lines 304 - 329, Update
_renew_loop to retry failed renew_peer_lease calls a bounded number of times
before returning, while preserving the existing renewal interval and stopping
once renewal succeeds or retries are exhausted. In release, await the cancelled
heartbeat while suppressing both CancelledError and any other heartbeat
exception, ensuring lease/semaphore cleanup continues and release never
re-raises heartbeat task failures.
| finally: | ||
| # Close the client if we created it, to prevent connection leaks. | ||
| # Caller-owned clients are their responsibility. | ||
| if _owns_client: | ||
| await client.aclose() | ||
| if slot is not None: | ||
| await slot.release() |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Release the capacity slot even if client.aclose() raises.
aclose() runs first in this finally. If it raises, slot.release() is skipped. In distributed mode the PostgreSQL lease then survives until PEER_LEASE_TTL_S expires, so the seller cap stays reduced for up to two minutes. Isolate the two cleanup actions.
🐛 Proposed fix
finally:
- # Close the client if we created it, to prevent connection leaks.
- # Caller-owned clients are their responsibility.
- if _owns_client:
- await client.aclose()
- if slot is not None:
- await slot.release()
+ try:
+ # Close the client if we created it, to prevent connection leaks.
+ # Caller-owned clients are their responsibility.
+ if _owns_client:
+ await client.aclose()
+ finally:
+ if slot is not None:
+ await slot.release()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| finally: | |
| # Close the client if we created it, to prevent connection leaks. | |
| # Caller-owned clients are their responsibility. | |
| if _owns_client: | |
| await client.aclose() | |
| if slot is not None: | |
| await slot.release() | |
| finally: | |
| try: | |
| # Close the client if we created it, to prevent connection leaks. | |
| # Caller-owned clients are their responsibility. | |
| if _owns_client: | |
| await client.aclose() | |
| finally: | |
| if slot is not None: | |
| await slot.release() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@provider_adapters/openai_compatible.py` around lines 626 - 632, Update the
cleanup logic in the surrounding provider method’s finally block so
slot.release() executes even when the owned client’s aclose() raises. Isolate
the client close and capacity-slot release into independent cleanup paths, while
preserving ownership checks and the existing release behavior for non-null
slots.
| @pytest.mark.asyncio | ||
| async def test_distributed_peer_gate_uses_shared_store( | ||
| monkeypatch, host_store_clean): | ||
| monkeypatch.setenv("DISTRIBUTED_PEER_GATES", "1") | ||
| monkeypatch.setenv("PEER_GATE_WAIT_S", "1") | ||
| client = _FakeClient(delay=0.04) | ||
| call = make_async_call_provider(client=client) | ||
|
|
||
| results = await asyncio.gather(*( | ||
| call(_req("peer-global", 2)) for _ in range(8))) | ||
| assert all(result["ok"] for result in results) | ||
| assert client.max_in_flight == 2 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Raise the gate wait budget to avoid a flaky assertion.
PEER_GATE_WAIT_S is 1. With cap 2 and 8 calls, the requests serialize into four waves, and each losing waiter re-polls PostgreSQL every 100 ms while holding an advisory lock per attempt. On a loaded CI database the last request can exhaust the 1 s budget, return the rate_limit peer-capacity error, and fail assert all(result["ok"] ...). A larger budget still exercises the shared-store path.
💚 Proposed fix
- monkeypatch.setenv("PEER_GATE_WAIT_S", "1")
+ monkeypatch.setenv("PEER_GATE_WAIT_S", "5")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @pytest.mark.asyncio | |
| async def test_distributed_peer_gate_uses_shared_store( | |
| monkeypatch, host_store_clean): | |
| monkeypatch.setenv("DISTRIBUTED_PEER_GATES", "1") | |
| monkeypatch.setenv("PEER_GATE_WAIT_S", "1") | |
| client = _FakeClient(delay=0.04) | |
| call = make_async_call_provider(client=client) | |
| results = await asyncio.gather(*( | |
| call(_req("peer-global", 2)) for _ in range(8))) | |
| assert all(result["ok"] for result in results) | |
| assert client.max_in_flight == 2 | |
| @pytest.mark.asyncio | |
| async def test_distributed_peer_gate_uses_shared_store( | |
| monkeypatch, host_store_clean): | |
| monkeypatch.setenv("DISTRIBUTED_PEER_GATES", "1") | |
| monkeypatch.setenv("PEER_GATE_WAIT_S", "5") | |
| client = _FakeClient(delay=0.04) | |
| call = make_async_call_provider(client=client) | |
| results = await asyncio.gather(*( | |
| call(_req("peer-global", 2)) for _ in range(8))) | |
| assert all(result["ok"] for result in results) | |
| assert client.max_in_flight == 2 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_antseed_concurrency.py` around lines 181 - 192, Increase the
PEER_GATE_WAIT_S value in test_distributed_peer_gate_uses_shared_store to
provide enough time for all four request waves to complete under loaded CI
conditions, while keeping the shared-store concurrency assertions unchanged.
What changes
Independence
This PR is self-contained on
mainand has no prerequisite PRs. The peer-concurrency, streaming cleanup, and client-lifecycle protections are included and tested directly in this branch.Verification
The GitOps rollout is intentionally separate: first a shadow pool, then a service-selector cutover after live validation.
Summary by CodeRabbit
New Features
/metrics, including capacity, rejection, latency, and availability data.Bug Fixes
Documentation