Add opt-in upstream pooling and account capacity limits

.env.example

@@ -35,6 +35,10 @@ CODEBUDDY2API_MAX_INBOUND_BYTES=67108864
 # Aggregated output bytes and concurrent inference requests; zero disables each limit.
 CODEBUDDY2API_MAX_COLLECT_BYTES=8388608
 CODEBUDDY2API_MAX_CONCURRENT=64
+# Optional bounded upstream connection reuse; restart to apply, disabled by default.
+# CODEBUDDY2API_UPSTREAM_KEEPALIVE=false
+# Per-account in-flight limit; zero preserves unlimited account capacity.
+# CODEBUDDY2API_MAX_INFLIGHT_PER_ACCOUNT=0
 
 # SQLite audit defaults to the data directory; this optional path enables a separate text log.
 CODEBUDDY2API_LOG=

app/inference_resources.py

@@ -0,0 +1,156 @@
+"""Own bounded upstream clients and account capacity for each inference request."""
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+from http.cookiejar import CookieJar, DefaultCookiePolicy
+import threading
+
+import httpx
+
+from app.site_routing import PROFILE_ENDPOINTS
+
+
+class _RejectCookies(DefaultCookiePolicy):
+    def set_ok(self, cookie, request):
+        return False
+
+    def return_ok(self, cookie, request):
+        return False
+
+
+class UpstreamClients:
+    """Keep at most one cookie-free HTTP/1.1 client per trusted origin and event loop."""
+
+    def __init__(self):
+        self._clients = {}
+        self._loop = asyncio.get_running_loop()
+        self._closed = False
+        self._origins = {self._origin(url) for url in PROFILE_ENDPOINTS.values()}
+
+    @staticmethod
+    def _origin(url):
+        value = httpx.URL(url)
+        return value.scheme, value.host, value.port
+
+    def get(self, url):
+        origin = self._origin(url)
+        if self._closed or asyncio.get_running_loop() is not self._loop or origin not in self._origins:
+            return None
+        if origin not in self._clients:
+            self._clients[origin] = httpx.AsyncClient(
+                cookies=CookieJar(policy=_RejectCookies()), http2=False,
+                limits=httpx.Limits(max_connections=64, max_keepalive_connections=16, keepalive_expiry=30))
+        return self._clients[origin]
+
+    async def aclose(self):
+        self._closed = True
+        clients, self._clients = list(self._clients.values()), {}
+        results = await asyncio.gather(*(client.aclose() for client in clients), return_exceptions=True)
+        errors = [result for result in results if isinstance(result, BaseException)]
+        if errors:
+            raise BaseExceptionGroup("Upstream client shutdown failed", errors)
+
+
+@asynccontextmanager
+async def inference_lifespan(app):
+    clients = UpstreamClients()
+    try:
+        yield {"upstream_clients": clients}
+    finally:
+        await clients.aclose()
+
+
+class CredentialLease(tuple):
+    """Retain the existing (manager, generation) lease shape with idempotent capacity release."""
+
+    def __new__(cls, manager, generation, release):
+        lease = super().__new__(cls, (manager, generation))
+        lease._release = release
+        lease._lock = threading.Lock()
+        return lease
+
+    def release(self):
+        with self._lock:
+            release, self._release = self._release, None
+        if release is not None:
+            release()
+
+
+def release_credential(credential):
+    if isinstance(credential, CredentialLease):
+        credential.release()
+
+
+class AccountCapacity:
+    """Count account leases independently of credential I/O locks."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._counts = {}
+
+    def count(self, identity):
+        with self._lock:
+            return self._counts.get(identity, 0)
+
+    def acquire(self, identity, limit, manager, generation):
+        with self._lock:
+            count = self._counts.get(identity, 0)
+            if limit and count >= limit:
+                return None
+            self._counts[identity] = count + 1
+        return CredentialLease(manager, generation, lambda: self._release(identity))
+
+    def _release(self, identity):
+        with self._lock:
+            count = self._counts[identity] - 1
+            if count:
+                self._counts[identity] = count
+            else:
+                del self._counts[identity]
+
+
+class RequestResources:
+    """Release even leases acquired by a worker after its request has already closed."""
+
+    def __init__(self, clients=None):
+        self.clients = clients
+        self._lock = threading.Lock()
+        self._leases = []
+        self._closed = False
+
+    def add(self, lease):
+        with self._lock:
+            if not self._closed:
+                self._leases.append(lease)
+                return
+        release_credential(lease)
+        raise asyncio.CancelledError()
+
+    def close(self):
+        with self._lock:
+            self._closed = True
+            leases, self._leases = self._leases, []
+        for lease in leases:
+            release_credential(lease)
+
+
+request_resources = ContextVar("inference_resources", default=None)
+
+
+class InferenceResourcesMiddleware:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if (scope["type"] != "http" or scope.get("method") != "POST" or
+                scope.get("path") not in ("/v1/chat/completions", "/v1/responses", "/v1/messages")):
+            return await self.app(scope, receive, send)
+        resources = RequestResources(scope.get("state", {}).get("upstream_clients"))
+        token = request_resources.set(resources)
+        try:
+            await self.app(scope, receive, send)
+        finally:
+            resources.close()
+            request_resources.reset(token)

app/settings.py

@@ -43,6 +43,10 @@ def _item(default, type_, label, *, mode="hot", env=None, minimum=None, maximum=
                           minimum=0, maximum=10),
     "retry_write_timeout": _item(False, "boolean", "写超时参与重放",
                                  env="CODEBUDDY2API_RETRY_WRITE_TIMEOUT"),
+    "upstream_keepalive": _item(False, "boolean", "上游连接复用", mode="restart",
+                                env="CODEBUDDY2API_UPSTREAM_KEEPALIVE"),
+    "max_inflight_per_account": _item(0, "integer", "单账号在途上限（0 不限制）",
+                                      env="CODEBUDDY2API_MAX_INFLIGHT_PER_ACCOUNT", minimum=0, maximum=10000),
     "audit_max_bytes": _item(256 * 1024 * 1024, "integer", "审计明细预算", minimum=1024**2, maximum=1024**4),
     "audit_retention_days": _item(30, "integer", "审计明细保留天数", minimum=1, maximum=36500),
     "audit_diagnostic_bytes": _item(8192, "integer", "失败诊断最大字节", minimum=0, maximum=8192),

app/upstream_io.py

@@ -211,9 +211,19 @@ async def read_bounded_error(response, limit: int = ERROR_BODY_LIMIT) -> bytes:
 WRITE_TIMEOUT = (httpx.WriteTimeout,)
 
 
+@asynccontextmanager
+async def _attempt_client(url, timeout, clients):
+    client = clients.get(url) if clients is not None else None
+    if client is not None:
+        yield client
+    else:
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            yield client
+
+
 @asynccontextmanager
 async def open_backend_stream(url, headers, body, *, read_timeout=300, on_retry=None,
-                              retry_write_timeout=False):
+                              retry_write_timeout=False, clients=None):
     """Retry connection failures once on a fresh client; write timeouts require explicit opt-in.
     Never replay after the upstream response opens.
     """
@@ -222,8 +232,8 @@ async def open_backend_stream(url, headers, body, *, read_timeout=300, on_retry=
     for attempt in range(2):
         opened = False
         try:
-            async with httpx.AsyncClient(timeout=timeout) as client:
-                async with client.stream("POST", url, headers=headers, json=body) as response:
+            async with _attempt_client(url, timeout, clients if attempt == 0 else None) as client:
+                async with client.stream("POST", url, headers=headers, json=body, timeout=timeout) as response:
                     opened = True
                     yield response
                     return

converter.py

@@ -59,6 +59,8 @@ def desensitize_body(body, roles=("system",), desensitize_harness_user=False,
                                credential_file_lock)
 from app.upstream_io import (ChatSSEAccumulator, UpstreamHTTPError, UpstreamResponseError,
                              open_backend_stream, parse_retry_after, read_bounded_error)
+from app.inference_resources import (AccountCapacity, InferenceResourcesMiddleware, inference_lifespan,
+                                     request_resources, release_credential)
 from app.inference_auth import require_api_key
 from app.content_filter import ContentFilterDetector, is_filter_error
 from app.request_limits import ImageLimitError, apply_image_policy
@@ -464,6 +466,7 @@ def __init__(self, paths: list[Path] | None = None, scan: bool = False,
         self._blocks = ModelBlocks(blocks_path, ttl_s=MODEL_SITE_BLOCK_S, max_ttl_s=MODEL_SITE_BLOCK_MAX_S)
         self._rr = {None: 0, "cn": 0, "intl": 0}
         self._ledger = None              # Prefer credits expiring sooner.
+        self._capacity = AccountCapacity()
         self._scan = scan                # Rescan credentials before selection.
         self._ignored_duplicates: set[str] = set()
         self._sync_pending: set[str] = set()
@@ -793,8 +796,19 @@ def _candidates(self, model: str | None, *, region=None, tried=()) -> list[dict]
         healthy.sort(key=lambda entry: (not self._model_free(entry, model), *self._expiry_rank(entry)))
         return healthy
 
+    @staticmethod
+    def _capacity_error():
+        return HTTPException(status_code=503, headers={"Retry-After": "3"}, detail={"error": {
+            "message": "符合当前路由和免费优先策略的账号在途名额已满，请稍后重试",
+            "type": "service_unavailable", "code": "credential_concurrency_limit"}})
+
+    @staticmethod
+    def _capacity_key(entry):
+        return entry.get("account_key") or entry["id"]
+
+
     def pick(self, skey: str | None, model: str | None = None, *, region=None,
-             tried=()) -> CredentialManager | None:
+             tried=(), with_capacity=False) -> CredentialManager | None:
         """Select a healthy sticky or round-robin credential, preferring eligible zero-rate accounts."""
         self._rescan()  # Reload and prune acquire their own locks.
         with self._lock:
@@ -804,6 +818,13 @@ def pick(self, skey: str | None, model: str | None = None, *, region=None,
                 if skey:
                     self._sticky.pop(skey, None)
                 return None
+            limit = CONFIG.get("max_inflight_per_account", 0)
+            if with_capacity and limit:
+                free = self._model_free(candidates[0], model)
+                candidates = [entry for entry in candidates if self._model_free(entry, model) == free
+                              and self._capacity.count(self._capacity_key(entry)) < limit]
+                if not candidates:
+                    raise self._capacity_error()
             best = candidates[0]
             free = self._model_free(best, model)
             top = [e for e in candidates if self._model_free(e, model) == free
@@ -822,10 +843,11 @@ def pick(self, skey: str | None, model: str | None = None, *, region=None,
             return e["cm"]
 
     def headers_for(self, skey: str | None, model: str | None = None, *, region=None,
-                    with_generation=False, tried=()):
-        """Recheck the credential generation and site before sending."""
+                    with_generation=False, tried=(), with_capacity=False):
+        """Recheck identity and atomically reserve account capacity before sending."""
+        capacity_race = False
         for _ in range(max(1, len(self._entries))):
-            cm = self.pick(skey, model, region=region, tried=tried)
+            cm = self.pick(skey, model, region=region, tried=tried, with_capacity=with_capacity)
             if cm is None:
                 return None
             reason = None
@@ -844,7 +866,16 @@ def headers_for(self, skey: str | None, model: str | None = None, *, region=None
                 entry = next((entry for entry in self._entries if entry["cm"] is cm), None)
                 if (entry is not None and cm._generation == generation and self._healthy(entry)
                         and self._eligible(entry, model, region=region, profile=profile) and self._model_healthy(entry, model)):
+                    if with_capacity:
+                        lease = self._capacity.acquire(self._capacity_key(entry),
+                            CONFIG.get("max_inflight_per_account", 0), cm, generation)
+                        if lease is None:
+                            capacity_race = True
+                            continue
+                        return lease, headers
                     return ((cm, generation) if with_generation else cm), headers
+        if capacity_race:
+            raise self._capacity_error()
         return None
 
     @staticmethod
@@ -1041,6 +1072,8 @@ def snapshot(self) -> list[dict]:
             out = []
             for e in self._entries:
                 s: dict = {"auth_file": e["id"], "healthy": self._healthy(e),
+                           "in_flight": self._capacity.count(self._capacity_key(e)),
+                           "max_in_flight": CONFIG.get("max_inflight_per_account", 0),
                            "model_cooldowns": {m: time.strftime("%m-%d %H:%M:%S", time.localtime(u))
                                                for (cid, m), u in self._model_fail.items()
                                                if cid == e["id"] and u > now},
@@ -1401,7 +1434,8 @@ def _housekeeper_loop(pool: CredentialPool, ledger) -> None:
 # FastAPI application
 # ---------------------------------------------------------------------------
 
-app = FastAPI(title="codebuddy2api", version=APP_VERSION)
+app = FastAPI(title="codebuddy2api", version=APP_VERSION, lifespan=inference_lifespan)
+app.add_middleware(InferenceResourcesMiddleware)
 
 # Anthropic error types: https://platform.claude.com/docs/en/api/errors
 _ANTHROPIC_ERROR_TYPES = {
@@ -1450,6 +1484,7 @@ async def _protocol_http_exception(request: Request, exc: HTTPException):
                 "max_request_bytes": 32 * 1024 * 1024, "log_body_limit": 65536,
                 "max_inbound_bytes": 64 * 1024 * 1024,
                 "max_collect_bytes": 8 * 1024 * 1024, "max_concurrent": 64,
+                "upstream_keepalive": False, "max_inflight_per_account": 0,
                 "failover_max": 0,     # Credential failovers allowed before the first response byte
                 "retry_write_timeout": False,  # Opt-in replay after incomplete writes
                 "usage_daily": None,     # Usage aggregated by date and model
@@ -1541,7 +1576,9 @@ def _cred_for(payload: dict, model: str | None = None, *, region=None, tried=())
     skey = model_policy.sticky_scope(CONFIG, skey, model)
     pool = CONFIG.get("cred_pool")
     if pool is not None:
-        picked = pool.headers_for(skey, model, region=region, with_generation=True, tried=tried)
+        resources = request_resources.get()
+        picked = pool.headers_for(skey, model, region=region, with_generation=True, tried=tried,
+                                  with_capacity=resources is not None)
         if picked is None:
             until = pool.model_cooldown_until(model, region=region)
             if until:
@@ -1564,6 +1601,8 @@ def _cred_for(payload: dict, model: str | None = None, *, region=None, tried=())
                                 detail={"error": {"message": "无可用凭证（未登录、目录/额度未就绪或全部熔断）",
                                                   "type": "auth_error"}})
         cm, headers = picked
+        if resources is not None:
+            resources.add(cm)
     else:
         cm = CONFIG["cred"]
         if cm is None or cm in {_cred_manager(item) for item in tried}:
@@ -2589,8 +2628,11 @@ def retry(error):
         _log(f"[{rid}] {'写超时重放' if timeout_on_write else '建连失败'}，重试 1/1 | {model_name}"
              f" | {_network_error_text(error)}{_replay_cost_note(error)}")
     try:
+        resources = request_resources.get()
+        clients = resources.clients if resources is not None and CONFIG.get("upstream_keepalive") else None
         async with open_backend_stream(url, headers, body, read_timeout=timeout, on_retry=retry,
-                                       retry_write_timeout=bool(CONFIG.get("retry_write_timeout"))) as response:
+                                       retry_write_timeout=bool(CONFIG.get("retry_write_timeout")),
+                                       clients=clients) as response:
             opened = True
             observe_attempt("upstream_http", status_code=response.status_code,
                             duration_ms=(time.monotonic() - started) * 1000)
@@ -2928,6 +2970,7 @@ async def _stream_plan(payload, canonical, model_name, rid, t0, make, routed, cr
             first = await _preflight_stream(stream, model_name, t0, rid)
         except _StreamFailure as failure:
             await _close_stream(stream)   # Release the failed upstream connection.
+            release_credential(cred)
             recovered = observe_failure_seq()   # Recover only this failure sequence.
             tried.append(cred)
             limit = _failover_limit()
@@ -2976,6 +3019,7 @@ async def _routed_fetch(payload, canonical, model_name, rid, t0, fetch, routed,
                 observe_recovery(recovered)
             return collected
         except (httpx.HTTPError, UpstreamResponseError) as error:
+            release_credential(cred)
             status, raw = _upstream_failure(error, model_name, t0, rid)
             recovered = observe_failure_seq()
             tried.append(cred)
@@ -3390,6 +3434,12 @@ def main():
     ap.add_argument("--max-concurrent", type=_nonnegative_int, metavar="N",
                     default=os.environ.get("CODEBUDDY2API_MAX_CONCURRENT", "64"),
                     help="推理端点并发上限（超出立即 503），默认 64；0 不限制")
+    ap.add_argument("--max-inflight-per-account", type=_nonnegative_int, metavar="N",
+                    default=os.environ.get("CODEBUDDY2API_MAX_INFLIGHT_PER_ACCOUNT", "0"),
+                    help="单账号在途上限，默认 0（不限制）；满载返回 503，不借容量切换到收费账号")
+    ap.add_argument("--upstream-keepalive", type=_boolean_arg, nargs="?", const=True,
+                    default=os.environ.get("CODEBUDDY2API_UPSTREAM_KEEPALIVE", "false"),
+                    help="按上游入口复用有界连接池，默认 false；重启生效，不改变超时或重放规则")
     ap.add_argument("--log-body-limit", type=_nonnegative_int, metavar="BYTES",
                     default=os.environ.get("CODEBUDDY2API_LOG_BODY_LIMIT", "65536"),
                     help="每条正文日志的预览字节上限，默认 64 KiB；0 只记录摘要")
@@ -3419,7 +3469,7 @@ def main():
 
     for key in ("max_images", "image_policy", "max_request_bytes", "log_body_limit",
                 "tool_call_max_retry", "max_inbound_bytes", "max_collect_bytes", "max_concurrent",
-                "failover_max", "retry_write_timeout"):
+                "failover_max", "retry_write_timeout", "upstream_keepalive", "max_inflight_per_account"):
         CONFIG[key] = getattr(args, key)
     CONFIG["api_key"] = args.api_key
     CONFIG["desensitize"] = args.desensitize