Skip to content

feat: make the router horizontally scalable - #100

Merged
jmlago merged 2 commits into
mainfrom
feat/horizontal-scaling
Aug 12, 2026
Merged

feat: make the router horizontally scalable#100
jmlago merged 2 commits into
mainfrom
feat/horizontal-scaling

Conversation

@jmlago

@jmlago jmlago commented Aug 12, 2026

Copy link
Copy Markdown
Member

What changes

  • moves consumer-key lookup and token-bucket admission to PostgreSQL so quotas remain global across replicas
  • adds bounded per-pod admission with retryable 503 responses and Prometheus autoscaling metrics
  • enforces AntSeed seller concurrency globally with renewable PostgreSQL leases for both JSON and streaming calls
  • reuses upstream HTTP pools and closes disconnected streams safely
  • adds an authenticated Codex broker so stateless replicas never receive refresh credentials
  • makes wallet keeper, synthetic probes, and startup maintenance single-owner tasks
  • documents the staged, reversible migration

Independence

This PR is self-contained on main and has no prerequisite PRs. The peer-concurrency, streaming cleanup, and client-lifecycle protections are included and tested directly in this branch.

Verification

  • Python: 759 passed, 2 skipped
  • Lua core: 637 passed

The GitOps rollout is intentionally separate: first a shadow pool, then a service-selector cutover after live validation.

Summary by CodeRabbit

  • New Features

    • Added support for horizontally scaled, stateless API deployments.
    • Added configurable request queuing, overload protection, connection limits, and distributed concurrency controls.
    • Added a secure Codex broker for shared authentication and streaming requests.
    • Added Prometheus metrics at /metrics, including capacity, rejection, latency, and availability data.
    • Added configuration controls for background workers, synthetic probes, and cost backfills.
  • Bug Fixes

    • Improved streaming cleanup and resource handling when clients disconnect.
  • Documentation

    • Added deployment guidance for horizontal scaling, broker usage, monitoring, migration, and rollback.

@coderabbitai

coderabbitai Bot commented Aug 12, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

The 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.

Changes

Horizontal deployment

Layer / File(s) Summary
Shared authentication, rate limits, and leases
host_store.py, auth_proxy.py, tests/test_host_store.py
Shared PostgreSQL state now stores consumer-key digests, global token buckets, and renewable peer leases. Consumer-key writes and revocations use atomic persistence before local mirror updates.
Proxy admission and metrics
auth_proxy.py, .env.example, tests/test_auth_proxy_scaling.py, tests/test_auth_proxy_dashboard_full.py
The proxy adds bounded active and pending request capacity, fail-closed shared-state errors, retry headers, streaming cleanup, cost-backfill control, and a Prometheus /metrics endpoint.
Codex broker and shared clients
codex_broker.py, remote_codex.py, codex_backend.py, serve.py, streaming.py, auth_proxy.py, tests/test_codex_broker.py
A bearer-authenticated Codex broker supports unary and NDJSON streaming calls. Router and provider paths reuse shared HTTP clients and close only internally owned clients.
Distributed provider capacity
provider_adapters/openai_compatible.py, .env.example, tests/test_antseed_concurrency.py
Provider calls use local or PostgreSQL-backed renewable peer-capacity slots. Unary and streaming paths release capacity and close owned clients on completion or interruption.
Replica ownership and runtime packaging
serve.py, Dockerfile.antseed, .env.example, docs/DEPLOY.md, docs/HORIZONTAL-SCALING.md, tests/test_sources.py
Environment flags control source refreshers, wallet keeper, synthetic probes, and cost backfills. Deployment guidance defines stateful control replicas, private brokers, and Node 22 dependency packaging.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

Suggested reviewers: muncleuscles

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 20.95% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes the primary change: enabling horizontal scalability for the router.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches 💡 1
📝 Generate docstrings 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/horizontal-scaling

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@jmlago
jmlago force-pushed the feat/horizontal-scaling branch from caa10ed to 343338c Compare August 12, 2026 19:27
@jmlago
jmlago marked this pull request as ready for review August 12, 2026 19:36
@jmlago
jmlago merged commit f70672e into main Aug 12, 2026
3 of 4 checks passed

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 9

🧹 Nitpick comments (14)
remote_codex.py (1)

33-36: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Use an identity check when you select the client.

Line 33 computes ownership with client is None, but line 34 selects with client 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_client stays False, so aclose never 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 win

One failing aclose skips the remaining shutdown work.

The loop at Lines 303-304 awaits each aclose() directly. If provider_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 win

Deduplicate the _enabled flag parser.

provider_adapters/openai_compatible.py Lines 249-253 defines a byte-identical _enabled helper. 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 example provider_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 win

Add coverage for the broker error paths in RemoteCodexClient.

The tests cover only 200 responses. _http_error maps 401/403 to auth_error, 429 to rate_limit, and everything else to server_error, and stream returns stream_interrupted when deltas were already emitted before the wire ended. Those branches carry the router's fallback semantics, so a small MockTransport case 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 value

Consider failing fast when app.state.http is not yet initialized.

create_app assigns app.state.http = client at line 60, so the value is None until the lifespan runs. If any caller mounts this app without running the lifespan (for example a bare TestClient(app) call without the context manager), /v1/call and /v1/stream pass client=None to the backends, which then create a per-request httpx.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 value

Clear _PEER_GATE_TIMESTAMPS with the gate dict.

The fixture resets H._PEER_GATES only. 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 win

Add 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_lease to return (False, False) and assert the result is network_error with the "shared peer gate unavailable" detail. Add a second case returning (False, True) with a short PEER_GATE_WAIT_S and assert rate_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 value

Document or key the cached lock per event loop.

_PEER_GATE_LOCK and _PEER_GATES live for the process lifetime. If a second event loop uses them (common in test suites that create one loop per test), a contended acquire() 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 by id(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 win

Back off between lease attempts instead of polling every 100 ms.

Each waiter runs one try_acquire_peer_lease transaction per 100 ms, and that function takes pg_advisory_xact_lock on 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 random at 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 | 🔵 Trivial

Consider a periodic sweep for abandoned peer leases.

try_acquire_peer_lease deletes expired rows only for the peer that is being acquired. If a peer stops receiving traffic, its expired rows stay in peer_concurrency_leases indefinitely. The row count is bounded by distinct peers multiplied by their cap, so this is small, but a scheduled delete on expires_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 win

Use the database clock for bucket timestamps.

stamp comes from each replica's wall clock. With several replicas, clock skew changes the refill amount: a replica whose clock is behind stores an older updated_at, and the next caller's max(0.0, stamp - row[1]) clamps to 0.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 use EXTRACT(EPOCH FROM clock_timestamp()), so the two shared-state mechanisms disagree on the time authority.

Keep the now parameter 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 lift

Consider a short-TTL cache for the shared auth lookups.

Each authenticated request now makes two PostgreSQL round trips here (consumer_for_digest and get_consumer_key), and _rate_ok adds a third with a row lock. The host-store pool is bounded to HOST_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 win

Seed 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 win

Mock host_store.consumer_for_digest for 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

📥 Commits

Reviewing files that changed from the base of the PR and between a0a1503 and bcad36c.

📒 Files selected for processing (18)
  • .env.example
  • Dockerfile.antseed
  • auth_proxy.py
  • codex_backend.py
  • codex_broker.py
  • docs/DEPLOY.md
  • docs/HORIZONTAL-SCALING.md
  • host_store.py
  • provider_adapters/openai_compatible.py
  • remote_codex.py
  • serve.py
  • streaming.py
  • tests/test_antseed_concurrency.py
  • tests/test_auth_proxy_dashboard_full.py
  • tests/test_auth_proxy_scaling.py
  • tests/test_codex_broker.py
  • tests/test_host_store.py
  • tests/test_sources.py

Comment thread auth_proxy.py
Comment on lines +3940 to +3948
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"}})

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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.

Comment thread codex_backend.py
Comment on lines +347 to +369
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()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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 apply before_first_output to the first aiter_lines read; close resp explicitly if the bound expires.
  • streaming.py#L99-L105: apply the same restructuring to the client.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.

Comment thread Dockerfile.antseed
Comment on lines +6 to +16
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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:


🏁 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
done

Repository: 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:


🌐 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:


🏁 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
done

Repository: 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
done

Repository: 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' || true

Repository: 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' || true

Repository: 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.

Comment on lines +43 to +56
- 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.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 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.

Comment thread host_store.py
Comment on lines +158 to +165
"""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)",

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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 with FOREIGN 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 hold pg_advisory_xact_lock across 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.

Comment on lines +232 to +264
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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.

Comment on lines +304 to +329
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()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Make the heartbeat resilient and stop release() from re-raising heartbeat errors.

Two problems in this pair of methods:

  1. Line 312: one failed renewal ends the heartbeat permanently. host_store.renew_peer_lease returns False for 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.
  2. Lines 320-322: release() awaits the heartbeat task and suppresses only CancelledError. If the task failed (for example import host_store raised, or asyncio.to_thread raised), the stored exception is re-raised here. release() is called from finally blocks 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.

Suggested change
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.

Comment on lines +626 to +632
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()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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.

Suggested change
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.

Comment on lines +181 to +192
@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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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.

Suggested change
@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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant