fix: increment request_count in least_used pool strategy

The least_used strategy selected entries via min(request_count) but
never incremented the counter. All entries stayed at count=0, so the
strategy degenerated to fill_first behavior with no actual load balancing.

Now increments request_count after each selection and persists the update.
This commit is contained in:
vominh1919 2026-04-23 23:09:59 +07:00 committed by Teknium
parent b3aed6cfd8
commit 461899894e
2 changed files with 34 additions and 1 deletions

View file

@ -739,8 +739,11 @@ class CredentialPool:
if self._strategy == STRATEGY_LEAST_USED and len(available) > 1:
entry = min(available, key=lambda e: e.request_count)
# Increment usage counter so subsequent selections distribute load
updated = replace(entry, request_count=entry.request_count + 1)
self._replace_entry(entry, updated)
self._current_id = entry.id
return entry
return updated
if self._strategy == STRATEGY_ROUND_ROBIN and len(available) > 1:
entry = available[0]

View file

@ -1171,3 +1171,33 @@ def test_nous_seed_from_singletons_preserves_obtained_at_timestamps(tmp_path, mo
assert e.agent_key_expires_in == 86400
assert e.agent_key_reused is False
class TestLeastUsedStrategy:
"""Regression: least_used strategy must increment request_count on select."""
def test_request_count_increments(self):
"""Each select() call should increment the chosen entry's request_count."""
from unittest.mock import patch as _patch
from agent.credential_pool import CredentialPool, PooledCredential, STRATEGY_LEAST_USED
entries = [
PooledCredential(provider="test", id="a", label="a", auth_type="api_key",
source="a", access_token="tok-a", priority=0, request_count=0),
PooledCredential(provider="test", id="b", label="b", auth_type="api_key",
source="b", access_token="tok-b", priority=1, request_count=0),
]
with _patch("agent.credential_pool.get_pool_strategy", return_value=STRATEGY_LEAST_USED):
pool = CredentialPool("test", entries)
# First select should pick entry with lowest count (both 0 → first)
e1 = pool.select()
assert e1 is not None
count_after_first = e1.request_count
assert count_after_first == 1, f"Expected 1 after first select, got {count_after_first}"
# Second select should pick the OTHER entry (now has lower count)
e2 = pool.select()
assert e2 is not None
assert e2.id != e1.id or e2.request_count == 2, (
"least_used should alternate or increment"
)