feat(pool): 增加 OAuth 账号池管理功能

- 新增 pool manager / strategy / health_policy / cost_tracker / redis_ops 等核心模块
- 新增 pool admin API 路由与 schemas
- 新增 OAuth 账号类型解析 (oauth_plan)
- 前端增加 PoolManagement 页面、PoolConfigDialog、PoolImportDialog、PoolStatusCard 组件
- 补充 pool config / cost tracker / health policy / manager / strategy / trace 等测试

src/api/admin/pool/__init__.py

@@ -0,0 +1,5 @@
+"""Pool management admin API."""
+
+from .routes import router
+
+__all__ = ["router"]

src/api/admin/pool/routes.py

@@ -0,0 +1,450 @@
+"""Pool management admin API routes.
+
+Provides endpoints for managing account pools at scale:
+- Overview of all pool-enabled providers
+- Paginated key listing with search/filter
+- Batch import / batch actions
+"""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+from fastapi import APIRouter, Depends, Query, Request
+from sqlalchemy.orm import Session
+
+from src.api.base.admin_adapter import AdminApiAdapter
+from src.api.base.context import ApiRequestContext
+from src.api.base.pipeline import ApiRequestPipeline
+from src.core.crypto import crypto_service
+from src.core.exceptions import NotFoundException
+from src.core.logger import logger
+from src.database import get_db
+from src.models.database import Provider, ProviderAPIKey
+from src.services.provider.pool import redis_ops as pool_redis
+from src.services.provider.pool.config import parse_pool_config
+
+from .schemas import (
+    BatchActionRequest,
+    BatchActionResponse,
+    BatchImportError,
+    BatchImportRequest,
+    BatchImportResponse,
+    PoolKeyDetail,
+    PoolKeysPageResponse,
+    PoolOverviewItem,
+    PoolOverviewResponse,
+)
+
+router = APIRouter(prefix="/api/admin/pool", tags=["pool-management"])
+pipeline = ApiRequestPipeline()
+
+
+# ---------------------------------------------------------------------------
+# GET /api/admin/pool/overview
+# ---------------------------------------------------------------------------
+
+
+@router.get("/overview", response_model=PoolOverviewResponse)
+async def pool_overview(
+    request: Request,
+    db: Session = Depends(get_db),
+) -> PoolOverviewResponse:
+    """Return all pool-enabled providers with summary stats."""
+    adapter = AdminPoolOverviewAdapter()
+    return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
+
+
+# ---------------------------------------------------------------------------
+# GET /api/admin/pool/{provider_id}/keys
+# ---------------------------------------------------------------------------
+
+
+@router.get("/{provider_id}/keys", response_model=PoolKeysPageResponse)
+async def list_pool_keys(
+    provider_id: str,
+    request: Request,
+    page: int = Query(1, ge=1),
+    page_size: int = Query(50, ge=1, le=200),
+    search: str = Query("", description="Search by key name"),
+    status: str = Query("all", description="all/active/cooldown/inactive"),
+    db: Session = Depends(get_db),
+) -> PoolKeysPageResponse:
+    """Server-side paginated account list for a pool-enabled provider."""
+    adapter = AdminListPoolKeysAdapter(
+        provider_id=provider_id,
+        page=page,
+        page_size=page_size,
+        search=search,
+        status=status,
+    )
+    return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
+
+
+# ---------------------------------------------------------------------------
+# POST /api/admin/pool/{provider_id}/keys/batch-import
+# ---------------------------------------------------------------------------
+
+
+@router.post("/{provider_id}/keys/batch-import", response_model=BatchImportResponse)
+async def batch_import_keys(
+    provider_id: str,
+    body: BatchImportRequest,
+    request: Request,
+    db: Session = Depends(get_db),
+) -> BatchImportResponse:
+    """Batch import keys into a provider's pool."""
+    adapter = AdminBatchImportKeysAdapter(provider_id=provider_id, body=body)
+    return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
+
+
+# ---------------------------------------------------------------------------
+# POST /api/admin/pool/{provider_id}/keys/batch-action
+# ---------------------------------------------------------------------------
+
+ALLOWED_ACTIONS = {"enable", "disable", "delete", "clear_cooldown", "reset_cost"}
+
+
+@router.post("/{provider_id}/keys/batch-action", response_model=BatchActionResponse)
+async def batch_action_keys(
+    provider_id: str,
+    body: BatchActionRequest,
+    request: Request,
+    db: Session = Depends(get_db),
+) -> BatchActionResponse:
+    """Batch enable/disable/delete/clear_cooldown/reset_cost on pool keys."""
+    adapter = AdminBatchActionKeysAdapter(provider_id=provider_id, body=body)
+    return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
+
+
+# ---------------------------------------------------------------------------
+# Adapters
+# ---------------------------------------------------------------------------
+
+
+class AdminPoolOverviewAdapter(AdminApiAdapter):
+    async def handle(self, context: ApiRequestContext) -> Any:  # type: ignore[override]
+        db = context.db
+        providers = (
+            db.query(Provider)
+            .filter(Provider.is_active.is_(True))
+            .order_by(Provider.provider_priority.asc())
+            .all()
+        )
+
+        items: list[PoolOverviewItem] = []
+        for p in providers:
+            pid = str(p.id)
+            pcfg = parse_pool_config(getattr(p, "config", None))
+
+            # Non-pool providers: skip Redis + key queries entirely.
+            if pcfg is None:
+                items.append(
+                    PoolOverviewItem(
+                        provider_id=pid,
+                        provider_name=p.name,
+                        provider_type=str(getattr(p, "provider_type", "custom") or "custom"),
+                        pool_enabled=False,
+                    )
+                )
+                continue
+
+            keys = db.query(ProviderAPIKey).filter(ProviderAPIKey.provider_id == pid).all()
+            key_ids = [str(k.id) for k in keys]
+
+            cooldown_count = 0
+            if key_ids:
+                cooldowns = await pool_redis.batch_get_cooldowns(pid, key_ids)
+                cooldown_count = sum(1 for v in cooldowns.values() if v is not None)
+
+            items.append(
+                PoolOverviewItem(
+                    provider_id=pid,
+                    provider_name=p.name,
+                    provider_type=str(getattr(p, "provider_type", "custom") or "custom"),
+                    total_keys=len(keys),
+                    active_keys=sum(1 for k in keys if k.is_active),
+                    cooldown_count=cooldown_count,
+                    pool_enabled=True,
+                )
+            )
+
+        return PoolOverviewResponse(items=items)
+
+
+@dataclass
+class AdminListPoolKeysAdapter(AdminApiAdapter):
+    provider_id: str = ""
+    page: int = 1
+    page_size: int = 50
+    search: str = ""
+    status: str = "all"
+
+    async def handle(self, context: ApiRequestContext) -> Any:  # type: ignore[override]
+        db = context.db
+        provider = db.query(Provider).filter(Provider.id == self.provider_id).first()
+        if not provider:
+            raise NotFoundException("Provider not found", "provider")
+
+        pcfg = parse_pool_config(getattr(provider, "config", None))
+        pid = str(provider.id)
+
+        # Base query
+        q = db.query(ProviderAPIKey).filter(ProviderAPIKey.provider_id == pid)
+
+        if self.search:
+            escaped = self.search.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
+            q = q.filter(ProviderAPIKey.name.ilike(f"%{escaped}%"))
+
+        if self.status == "active":
+            q = q.filter(ProviderAPIKey.is_active.is_(True))
+        elif self.status == "inactive":
+            q = q.filter(ProviderAPIKey.is_active.is_(False))
+        # "cooldown" filtering is done post-query (Redis state)
+
+        total = q.count()
+
+        # For cooldown filtering we need to fetch all, then filter, then paginate.
+        # Limit scan range to avoid loading the entire table into memory.
+        if self.status == "cooldown":
+            _max_scan = 2000
+            all_keys = q.order_by(ProviderAPIKey.created_at.desc()).limit(_max_scan).all()
+            key_ids = [str(k.id) for k in all_keys]
+            cooldowns = await pool_redis.batch_get_cooldowns(pid, key_ids) if key_ids else {}
+            all_keys = [k for k in all_keys if cooldowns.get(str(k.id)) is not None]
+            total = len(all_keys)
+            offset = (self.page - 1) * self.page_size
+            keys = all_keys[offset : offset + self.page_size]
+        else:
+            offset = (self.page - 1) * self.page_size
+            keys = (
+                q.order_by(ProviderAPIKey.created_at.desc())
+                .offset(offset)
+                .limit(self.page_size)
+                .all()
+            )
+
+        # Batch fetch Redis state (parallel where possible)
+        key_ids = [str(k.id) for k in keys]
+        if key_ids:
+            _lru_coro = (
+                pool_redis.get_lru_scores(pid, key_ids)
+                if pcfg and pcfg.lru_enabled
+                else asyncio.sleep(0, result={})
+            )
+            _cost_coro = (
+                pool_redis.batch_get_cost_totals(pid, key_ids, pcfg.cost_window_seconds)
+                if pcfg
+                else asyncio.sleep(0, result={})
+            )
+            cooldowns, cooldown_ttls, lru_scores, cost_totals = await asyncio.gather(
+                pool_redis.batch_get_cooldowns(pid, key_ids),
+                pool_redis.batch_get_cooldown_ttls(pid, key_ids),
+                _lru_coro,
+                _cost_coro,
+            )
+        else:
+            cooldowns, cooldown_ttls, lru_scores, cost_totals = {}, {}, {}, {}
+
+        # Sticky session count per key is expensive (SCAN+MGET per key).
+        # Only compute when the page is small enough to avoid timeout.
+        sticky_counts: dict[str, int] = {}
+        if key_ids and len(key_ids) <= 30:
+            counts = await asyncio.gather(
+                *(pool_redis.get_key_sticky_count(pid, kid) for kid in key_ids)
+            )
+            sticky_counts = dict(zip(key_ids, counts))
+
+        key_details: list[PoolKeyDetail] = []
+        for k in keys:
+            kid = str(k.id)
+            cd_reason = cooldowns.get(kid)
+            cd_ttl = cooldown_ttls.get(kid) if cd_reason else None
+
+            key_details.append(
+                PoolKeyDetail(
+                    key_id=kid,
+                    key_name=k.name or "",
+                    is_active=bool(k.is_active),
+                    auth_type=str(getattr(k, "auth_type", "api_key") or "api_key"),
+                    cooldown_reason=cd_reason,
+                    cooldown_ttl_seconds=cd_ttl,
+                    cost_window_usage=cost_totals.get(kid, 0),
+                    cost_limit=pcfg.cost_limit_per_key_tokens if pcfg else None,
+                    sticky_sessions=sticky_counts.get(kid, 0),
+                    lru_score=lru_scores.get(kid),
+                    created_at=(
+                        k.created_at.isoformat() if getattr(k, "created_at", None) else None
+                    ),
+                    last_used_at=(
+                        k.last_used_at.isoformat() if getattr(k, "last_used_at", None) else None
+                    ),
+                )
+            )
+
+        return PoolKeysPageResponse(
+            total=total,
+            page=self.page,
+            page_size=self.page_size,
+            keys=key_details,
+        )
+
+
+@dataclass
+class AdminBatchImportKeysAdapter(AdminApiAdapter):
+    provider_id: str = ""
+    body: BatchImportRequest = field(default_factory=lambda: BatchImportRequest(keys=[]))
+
+    async def handle(self, context: ApiRequestContext) -> Any:  # type: ignore[override]
+        db = context.db
+        provider = db.query(Provider).filter(Provider.id == self.provider_id).first()
+        if not provider:
+            raise NotFoundException("Provider not found", "provider")
+
+        imported = 0
+        skipped = 0
+        errors: list[BatchImportError] = []
+        now = datetime.now(timezone.utc)
+
+        for idx, item in enumerate(self.body.keys):
+            if not item.api_key.strip():
+                errors.append(BatchImportError(index=idx, reason="api_key is empty"))
+                continue
+
+            try:
+                encrypted_key = crypto_service.encrypt(item.api_key)
+                new_key = ProviderAPIKey(
+                    id=str(uuid.uuid4()),
+                    provider_id=self.provider_id,
+                    name=item.name or f"imported-{idx}",
+                    api_key=encrypted_key,
+                    auth_type=item.auth_type or "api_key",
+                    is_active=True,
+                    created_at=now,
+                    updated_at=now,
+                )
+                db.add(new_key)
+                imported += 1
+            except Exception as exc:
+                logger.warning("batch import key #{} failed: {}", idx, exc)
+                errors.append(BatchImportError(index=idx, reason=str(exc)))
+
+        if imported > 0:
+            try:
+                db.commit()
+            except Exception as exc:
+                db.rollback()
+                logger.error("batch import commit failed: {}", exc)
+                return BatchImportResponse(
+                    imported=0,
+                    skipped=skipped,
+                    errors=[BatchImportError(index=-1, reason=f"commit failed: {exc}")],
+                )
+
+        admin_name = context.user.username if context.user else "admin"
+        logger.info(
+            "Pool batch import by {}: provider={}, imported={}, skipped={}, errors={}",
+            admin_name,
+            self.provider_id[:8],
+            imported,
+            skipped,
+            len(errors),
+        )
+
+        return BatchImportResponse(imported=imported, skipped=skipped, errors=errors)
+
+
+@dataclass
+class AdminBatchActionKeysAdapter(AdminApiAdapter):
+    provider_id: str = ""
+    body: BatchActionRequest = field(
+        default_factory=lambda: BatchActionRequest(key_ids=[], action="")
+    )
+
+    async def handle(self, context: ApiRequestContext) -> Any:  # type: ignore[override]
+        from fastapi import HTTPException
+
+        if self.body.action not in ALLOWED_ACTIONS:
+            raise HTTPException(
+                status_code=400,
+                detail=(
+                    f"Invalid action: {self.body.action}. "
+                    f"Allowed: {', '.join(sorted(ALLOWED_ACTIONS))}"
+                ),
+            )
+
+        db = context.db
+        provider = db.query(Provider).filter(Provider.id == self.provider_id).first()
+        if not provider:
+            raise NotFoundException("Provider not found", "provider")
+
+        pid = str(provider.id)
+        affected = 0
+
+        keys = (
+            db.query(ProviderAPIKey)
+            .filter(
+                ProviderAPIKey.provider_id == pid,
+                ProviderAPIKey.id.in_(self.body.key_ids),
+            )
+            .all()
+        )
+
+        for key in keys:
+            kid = str(key.id)
+
+            if self.body.action == "enable":
+                key.is_active = True
+                affected += 1
+
+            elif self.body.action == "disable":
+                key.is_active = False
+                affected += 1
+
+            elif self.body.action == "delete":
+                db.delete(key)
+                affected += 1
+
+            elif self.body.action == "clear_cooldown":
+                await pool_redis.clear_cooldown(pid, kid)
+                affected += 1
+
+            elif self.body.action == "reset_cost":
+                await pool_redis.clear_cost(pid, kid)
+                affected += 1
+
+        if self.body.action in {"enable", "disable", "delete"}:
+            try:
+                db.commit()
+            except Exception as exc:
+                db.rollback()
+                logger.error("batch action commit failed: {}", exc)
+                return BatchActionResponse(affected=0, message=f"commit failed: {exc}")
+
+        action_labels = {
+            "enable": "enabled",
+            "disable": "disabled",
+            "delete": "deleted",
+            "clear_cooldown": "cooldown cleared",
+            "reset_cost": "cost reset",
+        }
+
+        admin_name = context.user.username if context.user else "admin"
+        affected_ids = [str(k.id)[:8] for k in keys]
+        logger.info(
+            "Pool batch action by {}: provider={}, action={}, affected={}, key_ids={}",
+            admin_name,
+            self.provider_id[:8],
+            self.body.action,
+            affected,
+            affected_ids,
+        )
+
+        return BatchActionResponse(
+            affected=affected,
+            message=f"{affected} keys {action_labels.get(self.body.action, self.body.action)}",
+        )

src/api/admin/pool/schemas.py

@@ -0,0 +1,103 @@
+"""Pydantic schemas for Pool management API."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict, Field
+
+# ---------------------------------------------------------------------------
+# Overview
+# ---------------------------------------------------------------------------
+
+
+class PoolOverviewItem(BaseModel):
+    """One Provider in the overview list."""
+
+    provider_id: str
+    provider_name: str
+    provider_type: str = "custom"
+    total_keys: int = 0
+    active_keys: int = 0
+    cooldown_count: int = 0
+    pool_enabled: bool = False
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class PoolOverviewResponse(BaseModel):
+    items: list[PoolOverviewItem] = Field(default_factory=list)
+
+
+# ---------------------------------------------------------------------------
+# Paginated key list
+# ---------------------------------------------------------------------------
+
+
+class PoolKeyDetail(BaseModel):
+    """Detailed status of a single pool key."""
+
+    key_id: str
+    key_name: str
+    is_active: bool
+    auth_type: str = "api_key"
+    cooldown_reason: str | None = None
+    cooldown_ttl_seconds: int | None = None
+    cost_window_usage: int = 0
+    cost_limit: int | None = None
+    sticky_sessions: int = 0
+    lru_score: float | None = None
+    created_at: str | None = None
+    last_used_at: str | None = None
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class PoolKeysPageResponse(BaseModel):
+    """Server-side paginated key list."""
+
+    total: int
+    page: int
+    page_size: int
+    keys: list[PoolKeyDetail] = Field(default_factory=list)
+
+
+# ---------------------------------------------------------------------------
+# Batch import
+# ---------------------------------------------------------------------------
+
+
+class PoolKeyImportItem(BaseModel):
+    """Single key to import."""
+
+    name: str
+    api_key: str
+    auth_type: str = "api_key"
+
+
+class BatchImportRequest(BaseModel):
+    keys: list[PoolKeyImportItem] = Field(..., max_length=500)
+
+
+class BatchImportError(BaseModel):
+    index: int
+    reason: str
+
+
+class BatchImportResponse(BaseModel):
+    imported: int = 0
+    skipped: int = 0
+    errors: list[BatchImportError] = Field(default_factory=list)
+
+
+# ---------------------------------------------------------------------------
+# Batch action
+# ---------------------------------------------------------------------------
+
+
+class BatchActionRequest(BaseModel):
+    key_ids: list[str] = Field(..., max_length=500)
+    action: str  # enable / disable / delete / clear_cooldown / reset_cost
+
+
+class BatchActionResponse(BaseModel):
+    affected: int = 0
+    message: str = ""

src/core/oauth_plan.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from src.core.crypto import crypto_service
+
+
+def normalize_oauth_plan_type(plan_type: Any) -> str | None:
+    if not isinstance(plan_type, str):
+        return None
+    normalized = plan_type.strip().lower()
+    if not normalized:
+        return None
+    return normalized
+
+
+def extract_oauth_plan_type_from_auth_config_data(auth_config: Any) -> str | None:
+    if not isinstance(auth_config, dict):
+        return None
+
+    # Codex: plan_type (free/plus/team/enterprise)
+    plan_type = normalize_oauth_plan_type(auth_config.get("plan_type"))
+    if plan_type:
+        return plan_type
+
+    # Antigravity: tier (PAID/FREE/...)
+    tier = normalize_oauth_plan_type(auth_config.get("tier"))
+    if tier:
+        return tier
+
+    return None
+
+
+def decrypt_auth_config_to_dict(
+    encrypted_auth_config: str | None,
+    *,
+    silent: bool = True,
+) -> dict[str, Any] | None:
+    if not encrypted_auth_config:
+        return None
+    try:
+        decrypted = crypto_service.decrypt(encrypted_auth_config, silent=silent)
+        parsed = json.loads(decrypted)
+        if isinstance(parsed, dict):
+            return parsed
+    except Exception:
+        return None
+    return None
+
+
+def _strip_provider_prefix(value: str, provider_type: str | None = None) -> str:
+    normalized = value.strip()
+    if not normalized:
+        return ""
+
+    prefixes: list[str] = []
+    if isinstance(provider_type, str) and provider_type.strip():
+        prefixes.append(provider_type.strip())
+    # Kiro 目前将套餐信息记录为 "KIRO FREE" / "KIRO PRO+"。
+    if "kiro" not in {p.lower() for p in prefixes}:
+        prefixes.append("kiro")
+
+    upper = normalized.upper()
+    for prefix in prefixes:
+        prefix_upper = prefix.upper()
+        if upper == prefix_upper:
+            return ""
+        if upper.startswith(f"{prefix_upper} "):
+            return normalized[len(prefix) :].strip()
+
+    return normalized
+
+
+def extract_oauth_plan_type_from_upstream_metadata(
+    upstream_metadata: Any,
+    *,
+    provider_type: str | None = None,
+) -> str | None:
+    if not isinstance(upstream_metadata, dict):
+        return None
+
+    kiro_meta = upstream_metadata.get("kiro")
+    if isinstance(kiro_meta, dict):
+        subscription_title = kiro_meta.get("subscription_title")
+        if isinstance(subscription_title, str):
+            normalized = _strip_provider_prefix(subscription_title, provider_type=provider_type)
+            return normalize_oauth_plan_type(normalized)
+
+    return None
+
+
+def extract_oauth_plan_type(
+    encrypted_auth_config: str | None,
+    *,
+    upstream_metadata: Any = None,
+    provider_type: str | None = None,
+    silent: bool = True,
+) -> str | None:
+    auth_config = decrypt_auth_config_to_dict(encrypted_auth_config, silent=silent)
+    plan_type = extract_oauth_plan_type_from_auth_config_data(auth_config)
+    if plan_type:
+        return plan_type
+    return extract_oauth_plan_type_from_upstream_metadata(
+        upstream_metadata, provider_type=provider_type
+    )
+
+
+def normalize_antigravity_tier(raw_tier: Any) -> str | None:
+    normalized = normalize_oauth_plan_type(raw_tier)
+    if not normalized:
+        return None
+    if "ultra" in normalized:
+        return "ultra"
+    if "pro" in normalized or "paid" in normalized:
+        return "pro"
+    if "free" in normalized or "legacy" in normalized:
+        return "free"
+    return normalized
+
+
+def _extract_antigravity_tier_raw(tier_obj: Any) -> str | None:
+    if isinstance(tier_obj, str):
+        stripped = tier_obj.strip()
+        return stripped or None
+    if isinstance(tier_obj, dict):
+        for key in ("id", "tierType"):
+            value = tier_obj.get(key)
+            if isinstance(value, str):
+                stripped = value.strip()
+                if stripped:
+                    return stripped
+    return None
+
+
+def _format_antigravity_tier_label(
+    normalized_tier: str,
+    *,
+    fallback_raw: str | None = None,
+) -> str:
+    if normalized_tier == "ultra":
+        return "Ultra"
+    if normalized_tier == "pro":
+        return "Pro"
+    if normalized_tier == "free":
+        return "Free"
+    if fallback_raw:
+        return fallback_raw
+    return normalized_tier
+
+
+def extract_antigravity_tier_from_code_assist(code_assist: Any) -> str:
+    if not isinstance(code_assist, dict):
+        return "Free"
+
+    paid_tier_raw = _extract_antigravity_tier_raw(code_assist.get("paidTier"))
+    paid_tier = normalize_antigravity_tier(paid_tier_raw)
+    if paid_tier:
+        return _format_antigravity_tier_label(paid_tier, fallback_raw=paid_tier_raw)
+
+    current_tier_raw = _extract_antigravity_tier_raw(code_assist.get("currentTier"))
+    current_tier = normalize_antigravity_tier(current_tier_raw)
+    if current_tier:
+        return _format_antigravity_tier_label(current_tier, fallback_raw=current_tier_raw)
+
+    return "Free"

src/services/provider/adapters/claude_code/pool_config.py

@@ -0,0 +1,6 @@
+"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.config."""
+
+from src.services.provider.pool.config import *  # noqa: F401,F403
+from src.services.provider.pool.config import PoolConfig, UnschedulableRule, parse_pool_config
+
+__all__ = ["PoolConfig", "UnschedulableRule", "parse_pool_config"]

src/services/provider/adapters/claude_code/pool_cost_tracker.py

@@ -0,0 +1,10 @@
+"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.cost_tracker."""
+
+from src.services.provider.pool.cost_tracker import (  # noqa: F401
+    get_window_usage,
+    is_approaching_limit,
+    is_at_limit,
+    record_usage,
+)
+
+__all__ = ["record_usage", "get_window_usage", "is_at_limit", "is_approaching_limit"]

src/services/provider/adapters/claude_code/pool_health_policy.py

@@ -0,0 +1,6 @@
+"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.health_policy."""
+
+from src.services.provider.pool.health_policy import *  # noqa: F401,F403
+from src.services.provider.pool.health_policy import apply_health_policy  # noqa: F811
+
+__all__ = ["apply_health_policy"]

src/services/provider/adapters/claude_code/pool_hook.py

@@ -0,0 +1,27 @@
+"""Claude Code pool scheduling hook."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class ClaudeCodePoolHook:
+    """Pool scheduling hook for Claude Code providers.
+
+    Extracts the session UUID from ``metadata.user_id`` which follows the
+    pattern ``<user>_session_<uuid>``.
+    """
+
+    name = "claude_code"
+
+    def extract_session_uuid(self, request_body: dict[str, Any]) -> str | None:
+        metadata = request_body.get("metadata")
+        if isinstance(metadata, dict):
+            user_id = metadata.get("user_id")
+            if isinstance(user_id, str) and "_session_" in user_id:
+                idx = user_id.rfind("_session_")
+                return user_id[idx + len("_session_") :].strip() or None
+        return None
+
+
+claude_code_pool_hook = ClaudeCodePoolHook()

src/services/provider/adapters/claude_code/pool_manager.py

@@ -0,0 +1,8 @@
+"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.manager."""
+
+from src.services.provider.pool.manager import *  # noqa: F401,F403
+from src.services.provider.pool.manager import PoolManager
+
+ClaudeCodePoolManager = PoolManager  # noqa: F811
+
+__all__ = ["PoolManager", "ClaudeCodePoolManager"]

src/services/provider/adapters/claude_code/pool_oauth.py

@@ -0,0 +1,9 @@
+"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.oauth_cache."""
+
+from src.services.provider.pool.oauth_cache import (  # noqa: F401
+    cache_token,
+    get_cached_token,
+    invalidate_token,
+)
+
+__all__ = ["get_cached_token", "cache_token", "invalidate_token"]

src/services/provider/adapters/claude_code/pool_redis.py

@@ -0,0 +1,23 @@
+"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.redis_ops."""
+
+from src.services.provider.pool.redis_ops import *  # noqa: F401,F403
+from src.services.provider.pool.redis_ops import (
+    add_cost_entry,
+    batch_get_cooldowns,
+    cache_oauth_token,
+    clear_cooldown,
+    clear_cost,
+    delete_sticky_binding,
+    get_cached_oauth_token,
+    get_cooldown,
+    get_cooldown_ttl,
+    get_cost_window_total,
+    get_key_sticky_count,
+    get_lru_scores,
+    get_sticky_binding,
+    get_sticky_session_count,
+    invalidate_oauth_token_cache,
+    set_cooldown,
+    set_sticky_binding,
+    touch_lru,
+)

src/services/provider/pool/__init__.py

@@ -0,0 +1,14 @@
+"""Generic Account Pool management for any Provider type.
+
+Re-exports the main public API for convenience.
+"""
+
+from src.services.provider.pool.config import PoolConfig, UnschedulableRule, parse_pool_config
+from src.services.provider.pool.manager import PoolManager
+
+__all__ = [
+    "PoolConfig",
+    "PoolManager",
+    "UnschedulableRule",
+    "parse_pool_config",
+]

src/services/provider/pool/config.py

@@ -0,0 +1,138 @@
+"""Account Pool configuration (provider-agnostic)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from src.core.logger import logger
+
+
+@dataclass(frozen=True, slots=True)
+class UnschedulableRule:
+    """Keyword-based temporary unschedule rule."""
+
+    keyword: str
+    duration_minutes: int = 5
+
+
+@dataclass(frozen=True, slots=True)
+class PoolConfig:
+    """Parsed pool configuration for any Provider.
+
+    All transient state lives in Redis; this dataclass only holds
+    the *configuration* that controls pool behaviour.
+    """
+
+    # -- Sticky Session -------------------------------------------------------
+    sticky_session_ttl_seconds: int = 3600  # 1 hour
+
+    # -- Load-Aware Selection -------------------------------------------------
+    load_threshold_percent: int = 80
+
+    # -- LRU ------------------------------------------------------------------
+    lru_enabled: bool = True
+
+    # -- Rolling-Window Cost Tracking -----------------------------------------
+    cost_window_seconds: int = 18000  # 5 hours
+    cost_limit_per_key_tokens: int | None = None  # None = unlimited
+    cost_soft_threshold_percent: int = 80
+
+    # -- Cooldown Defaults ----------------------------------------------------
+    rate_limit_cooldown_seconds: int = 300  # 429
+    overload_cooldown_seconds: int = 30  # 529
+
+    # -- OAuth Proactive Refresh ----------------------------------------------
+    proactive_refresh_seconds: int = 180  # 3 minutes before expiry
+
+    # -- Health Policy --------------------------------------------------------
+    health_policy_enabled: bool = True
+
+    # -- Temporary Unschedulable Rules ----------------------------------------
+    unschedulable_rules: list[UnschedulableRule] = field(default_factory=list)
+
+    # -- Pluggable Strategies -------------------------------------------------
+    strategies: tuple[str, ...] = ()
+
+
+def parse_pool_config(provider_config: Any) -> PoolConfig | None:
+    """Parse PoolConfig from ``Provider.config``.
+
+    Only looks for the explicit ``pool_advanced`` key.  Returns ``None``
+    when the provider has no pool section configured, meaning the caller
+    should use the normal (non-pool) scheduling path.
+    """
+    config_dict = provider_config if isinstance(provider_config, dict) else {}
+
+    raw_advanced = config_dict.get("pool_advanced")
+    if raw_advanced is None:
+        return None
+
+    if not isinstance(raw_advanced, dict):
+        # Could be a pre-validated Pydantic model; grab its dict.
+        try:
+            raw_advanced = raw_advanced.model_dump()  # type: ignore[union-attr]
+        except Exception:
+            logger.warning(
+                "PoolConfig: advanced config type invalid ({}), falling back to defaults",
+                type(raw_advanced).__name__,
+            )
+            return PoolConfig()
+
+    rules: list[UnschedulableRule] = []
+    raw_rules = raw_advanced.get("unschedulable_rules")
+    if isinstance(raw_rules, list):
+        for r in raw_rules:
+            if isinstance(r, dict) and isinstance(r.get("keyword"), str):
+                rules.append(
+                    UnschedulableRule(
+                        keyword=r["keyword"],
+                        duration_minutes=int(r.get("duration_minutes", 5)),
+                    )
+                )
+
+    def _int_or(key: str, default: int) -> int:
+        v = raw_advanced.get(key)
+        if v is None:
+            return default
+        try:
+            return int(v)
+        except (TypeError, ValueError):
+            return default
+
+    def _bool_or(key: str, default: bool) -> bool:
+        v = raw_advanced.get(key)
+        if v is None:
+            return default
+        return bool(v)
+
+    def _opt_int(key: str) -> int | None:
+        v = raw_advanced.get(key)
+        if v is None:
+            return None
+        try:
+            return int(v)
+        except (TypeError, ValueError):
+            return None
+
+    return PoolConfig(
+        sticky_session_ttl_seconds=_int_or("sticky_session_ttl_seconds", 3600),
+        load_threshold_percent=_int_or("load_threshold_percent", 80),
+        lru_enabled=_bool_or("lru_enabled", True),
+        cost_window_seconds=_int_or("cost_window_seconds", 18000),
+        cost_limit_per_key_tokens=_opt_int("cost_limit_per_key_tokens"),
+        cost_soft_threshold_percent=_int_or("cost_soft_threshold_percent", 80),
+        rate_limit_cooldown_seconds=_int_or("rate_limit_cooldown_seconds", 300),
+        overload_cooldown_seconds=_int_or("overload_cooldown_seconds", 30),
+        proactive_refresh_seconds=_int_or("proactive_refresh_seconds", 180),
+        health_policy_enabled=_bool_or("health_policy_enabled", True),
+        unschedulable_rules=rules,
+        strategies=_parse_strategies(raw_advanced.get("strategies")),
+    )
+
+
+def _parse_strategies(raw: Any) -> tuple[str, ...]:
+    """Parse strategy names from config (list[str] -> tuple[str, ...])."""
+    if not isinstance(raw, list):
+        return ()
+    return tuple(str(s) for s in raw if isinstance(s, str) and s)

src/services/provider/pool/cost_tracker.py

@@ -0,0 +1,66 @@
+"""Rolling-window cost tracking for the Account Pool.
+
+Each key has a configurable token budget per rolling window (e.g. 5 hours).
+When the budget is exhausted the key is marked as unschedulable by the pool
+manager.  A "soft threshold" (default 80 %) causes the pool to *prefer*
+other keys but still allows traffic if no alternatives exist.
+
+All state is stored in Redis sorted sets via :mod:`redis_ops`.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from src.services.provider.pool import redis_ops
+
+if TYPE_CHECKING:
+    from src.services.provider.pool.config import PoolConfig
+
+
+async def record_usage(
+    provider_id: str,
+    key_id: str,
+    tokens: int,
+    config: PoolConfig,
+) -> None:
+    """Record *tokens* used by *key_id* in the rolling cost window."""
+    if tokens <= 0:
+        return
+    if config.cost_limit_per_key_tokens is None:
+        return  # cost tracking disabled
+    await redis_ops.add_cost_entry(provider_id, key_id, tokens, config.cost_window_seconds)
+
+
+async def get_window_usage(
+    provider_id: str,
+    key_id: str,
+    config: PoolConfig,
+) -> int:
+    """Return total tokens used by *key_id* within the current window."""
+    return await redis_ops.get_cost_window_total(provider_id, key_id, config.cost_window_seconds)
+
+
+async def is_at_limit(
+    provider_id: str,
+    key_id: str,
+    config: PoolConfig,
+) -> bool:
+    """Return ``True`` if the key has exhausted its budget."""
+    if config.cost_limit_per_key_tokens is None:
+        return False
+    total = await get_window_usage(provider_id, key_id, config)
+    return total >= config.cost_limit_per_key_tokens
+
+
+async def is_approaching_limit(
+    provider_id: str,
+    key_id: str,
+    config: PoolConfig,
+) -> bool:
+    """Return ``True`` if the key is above the soft threshold."""
+    if config.cost_limit_per_key_tokens is None:
+        return False
+    total = await get_window_usage(provider_id, key_id, config)
+    threshold = config.cost_limit_per_key_tokens * config.cost_soft_threshold_percent / 100
+    return total >= threshold

src/services/provider/pool/health_policy.py

@@ -0,0 +1,204 @@
+"""Account Pool health policy: error code classification and key state management.
+
+Maps upstream HTTP status codes to pool-level actions:
+
+| Code | Action                                                       |
+|------|--------------------------------------------------------------|
+| 401  | Invalidate OAuth token cache -> attempt refresh -> disable   |
+| 402  | Auto-disable key (payment issue)                             |
+| 403  | Auto-disable key (suspended/banned)                          |
+| 400  | Check body for "organization has been disabled" -> disable   |
+| 429  | Set cooldown (retry-after header or config default)          |
+| 529  | Set cooldown (config default)                                |
+| *    | Check unschedulable_rules keyword matching                   |
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING
+
+from src.core.logger import logger
+from src.services.provider.pool import redis_ops
+
+if TYPE_CHECKING:
+    from src.services.provider.pool.config import PoolConfig
+
+# Patterns in 400 error body that indicate account-level issues.
+_ACCOUNT_DISABLE_PATTERNS = (
+    "organization has been disabled",
+    "organization_disabled",
+    "account has been disabled",
+    "account_disabled",
+)
+
+
+def _parse_retry_after(headers: dict[str, str] | None) -> int | None:
+    """Extract retry-after seconds from response headers."""
+    if not headers:
+        return None
+    raw = headers.get("retry-after") or headers.get("Retry-After")
+    if not raw:
+        return None
+    try:
+        val = int(raw)
+        return max(1, min(val, 3600))
+    except (ValueError, TypeError):
+        return None
+
+
+def _extract_error_message(error_body: str | None) -> str:
+    """Best-effort extraction of error message from JSON body."""
+    if not error_body:
+        return ""
+    try:
+        data = json.loads(error_body)
+        if isinstance(data, dict):
+            error_obj = data.get("error")
+            if isinstance(error_obj, dict):
+                return str(error_obj.get("message", ""))
+            if isinstance(error_obj, str):
+                return error_obj
+            return str(data.get("message", ""))
+    except (json.JSONDecodeError, TypeError):
+        pass
+    return error_body[:500]
+
+
+async def apply_health_policy(
+    *,
+    provider_id: str,
+    key_id: str,
+    status_code: int,
+    error_body: str | None,
+    response_headers: dict[str, str] | None,
+    config: PoolConfig,
+) -> None:
+    """Apply health policy for an upstream error.
+
+    This is fire-and-forget; exceptions are caught and logged.
+    """
+    if not config.health_policy_enabled:
+        return
+
+    try:
+        await _apply(
+            provider_id=provider_id,
+            key_id=key_id,
+            status_code=status_code,
+            error_body=error_body,
+            response_headers=response_headers,
+            config=config,
+        )
+    except Exception as exc:
+        logger.warning(
+            "Pool health policy failed for key {}: {}",
+            key_id[:8],
+            str(exc),
+        )
+
+
+async def _apply(
+    *,
+    provider_id: str,
+    key_id: str,
+    status_code: int,
+    error_body: str | None,
+    response_headers: dict[str, str] | None,
+    config: PoolConfig,
+) -> None:
+    error_msg = _extract_error_message(error_body)
+
+    # --- 401 Unauthorized ---------------------------------------------------
+    if status_code == 401:
+        await redis_ops.invalidate_oauth_token_cache(key_id)
+        # Set a short cooldown to avoid hammering while refresh happens.
+        await redis_ops.set_cooldown(provider_id, key_id, "auth_failed_401", ttl=60)
+        logger.info(
+            "Pool[{}]: key {} got 401, token cache invalidated + 60s cooldown",
+            provider_id[:8],
+            key_id[:8],
+        )
+        return
+
+    # --- 402 Payment Required ------------------------------------------------
+    if status_code == 402:
+        await redis_ops.set_cooldown(provider_id, key_id, "payment_required_402", ttl=3600)
+        logger.warning(
+            "Pool[{}]: key {} got 402 (payment required), cooldown 1h",
+            provider_id[:8],
+            key_id[:8],
+        )
+        return
+
+    # --- 403 Forbidden -------------------------------------------------------
+    if status_code == 403:
+        await redis_ops.set_cooldown(provider_id, key_id, "forbidden_403", ttl=3600)
+        logger.warning(
+            "Pool[{}]: key {} got 403 (forbidden/suspended), cooldown 1h",
+            provider_id[:8],
+            key_id[:8],
+        )
+        return
+
+    # --- 400 with account-disable pattern ------------------------------------
+    if status_code == 400:
+        error_lower = error_msg.lower()
+        for pattern in _ACCOUNT_DISABLE_PATTERNS:
+            if pattern in error_lower:
+                await redis_ops.set_cooldown(
+                    provider_id, key_id, f"account_disabled_400:{pattern}", ttl=3600
+                )
+                logger.warning(
+                    "Pool[{}]: key {} got 400 with '{}', cooldown 1h",
+                    provider_id[:8],
+                    key_id[:8],
+                    pattern,
+                )
+                return
+
+    # --- 429 Rate Limited ----------------------------------------------------
+    if status_code == 429:
+        retry_after = _parse_retry_after(response_headers)
+        ttl = retry_after or config.rate_limit_cooldown_seconds
+        await redis_ops.set_cooldown(provider_id, key_id, "rate_limited_429", ttl=ttl)
+        logger.info(
+            "Pool[{}]: key {} got 429, cooldown {}s",
+            provider_id[:8],
+            key_id[:8],
+            ttl,
+        )
+        return
+
+    # --- 529 Overloaded ------------------------------------------------------
+    if status_code == 529:
+        ttl = config.overload_cooldown_seconds
+        await redis_ops.set_cooldown(provider_id, key_id, "overloaded_529", ttl=ttl)
+        logger.info(
+            "Pool[{}]: key {} got 529, cooldown {}s",
+            provider_id[:8],
+            key_id[:8],
+            ttl,
+        )
+        return
+
+    # --- Keyword-based unschedulable rules -----------------------------------
+    if config.unschedulable_rules and error_msg:
+        error_lower = error_msg.lower()
+        for rule in config.unschedulable_rules:
+            if rule.keyword.lower() in error_lower:
+                ttl = max(60, rule.duration_minutes * 60)
+                await redis_ops.set_cooldown(
+                    provider_id,
+                    key_id,
+                    f"rule:{rule.keyword}",
+                    ttl=ttl,
+                )
+                logger.info(
+                    "Pool[{}]: key {} matched rule '{}', cooldown {}m",
+                    provider_id[:8],
+                    key_id[:8],
+                    rule.keyword,
+                    rule.duration_minutes,
+                )
+                return

src/services/provider/pool/hooks.py

@@ -0,0 +1,91 @@
+"""Pool scheduling hooks -- provider-type-specific pool behaviour.
+
+Some provider types need custom logic during pool scheduling (e.g. extracting
+a session UUID for sticky binding).  This module provides a small Protocol +
+registry so the pool layer stays generic while provider-specific behaviour
+lives alongside each adapter.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Protocol, runtime_checkable
+
+# ---------------------------------------------------------------------------
+# Protocol
+# ---------------------------------------------------------------------------
+
+
+@runtime_checkable
+class PoolSchedulingHook(Protocol):
+    """Provider-type-specific pool scheduling behaviour.
+
+    Each provider type can optionally register a hook to customize:
+    - Session UUID extraction (for sticky sessions)
+    - Post-success / post-error callbacks
+
+    Optional methods (checked via ``hasattr`` by callers):
+    - ``on_pool_success``
+    - ``on_pool_error``
+    """
+
+    name: str
+
+    def extract_session_uuid(self, request_body: dict[str, Any]) -> str | None:
+        """Extract a session UUID for sticky binding from the request body."""
+        ...
+
+    # -- Optional lifecycle callbacks -----------------------------------------
+    # These are checked via ``hasattr`` so existing implementations that
+    # don't define them will continue to work.
+
+    def on_pool_success(
+        self,
+        *,
+        key_id: str,
+        session_uuid: str | None,
+        context: dict[str, Any],
+    ) -> None:
+        """Called after a successful pool request (provider-specific logic)."""
+        ...
+
+    def on_pool_error(
+        self,
+        *,
+        key_id: str,
+        status_code: int,
+        context: dict[str, Any],
+    ) -> None:
+        """Called after a failed pool request (provider-specific logic)."""
+        ...
+
+
+# ---------------------------------------------------------------------------
+# Registry
+# ---------------------------------------------------------------------------
+
+_hook_registry: dict[str, PoolSchedulingHook] = {}
+_registry_lock = threading.Lock()
+
+
+def register_pool_hook(provider_type: str, hook: PoolSchedulingHook) -> None:
+    """Register a pool scheduling hook for a provider type."""
+    from src.core.provider_types import normalize_provider_type
+
+    pt = normalize_provider_type(provider_type)
+    with _registry_lock:
+        _hook_registry[pt] = hook
+
+
+def get_pool_hook(provider_type: str | None) -> PoolSchedulingHook | None:
+    """Return the pool scheduling hook for a provider type, or ``None``."""
+    if not provider_type:
+        return None
+    from src.services.provider.envelope import ensure_providers_bootstrapped
+
+    ensure_providers_bootstrapped()
+
+    from src.core.provider_types import normalize_provider_type
+
+    pt = normalize_provider_type(provider_type)
+    return _hook_registry.get(pt)

src/services/provider/pool/manager.py

@@ -0,0 +1,545 @@
+"""Account Pool Manager (provider-agnostic).
+
+Stateless facade that coordinates pool operations for any Provider with
+pool configuration enabled.  All state lives in Redis via :mod:`redis_ops`.
+
+Usage::
+
+    mgr = PoolManager(provider_id, pool_config)
+    reordered = await mgr.reorder_candidates(session_uuid, candidates)
+    # ... execute request ...
+    await mgr.on_request_success(session_uuid=..., key_id=..., tokens_used=...)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from src.core.logger import logger
+from src.services.provider.pool import redis_ops
+from src.services.provider.pool.config import PoolConfig
+from src.services.provider.pool.trace import PoolCandidateTrace, PoolSchedulingTrace
+
+if TYPE_CHECKING:
+    from src.models.database import ProviderAPIKey
+    from src.services.scheduling.schemas import ProviderCandidate
+
+
+class PoolManager:
+    """Coordinate pool-level scheduling for a single Provider."""
+
+    __slots__ = ("provider_id", "config")
+
+    def __init__(self, provider_id: str, config: PoolConfig) -> None:
+        self.provider_id = provider_id
+        self.config = config
+
+    # ------------------------------------------------------------------
+    # Core scheduling: reorder candidate list for pool-aware selection
+    # ------------------------------------------------------------------
+
+    async def reorder_candidates(
+        self,
+        session_uuid: str | None,
+        candidates: list[ProviderCandidate],
+    ) -> list[ProviderCandidate]:
+        """Reorder *candidates* according to pool rules.
+
+        The returned list keeps the same elements but in a new order:
+
+        1. **Sticky session hit** -- if the session is already bound to a key
+           and that key appears in *candidates* and is not in cooldown, move it
+           to position 0.
+        2. **Filter** out keys in cooldown or cost-exhausted state (mark
+           ``is_skipped``).
+        3. **LRU sort** -- among remaining candidates at the same priority
+           level, sort by least-recently-used.
+        4. **Random tiebreak** -- among candidates with identical LRU score.
+
+        Also builds a :class:`PoolSchedulingTrace` and attaches per-candidate
+        trace data via ``_pool_extra_data`` / ``_pool_scheduling_trace``
+        attributes on candidate objects.
+        """
+        if not candidates:
+            return candidates
+
+        pid = self.provider_id
+
+        # Build trace
+        trace = PoolSchedulingTrace(
+            provider_id=pid,
+            total_keys=len(candidates),
+            session_uuid=session_uuid[:8] if session_uuid else None,
+        )
+
+        # --- Strategy: before_select ----------------------------------
+        strategies = _get_active_strategies(self.config)
+        key_ids = [str(c.key.id) for c in candidates]
+        strategy_context: dict[str, Any] = {"session_uuid": session_uuid}
+        for strategy in strategies:
+            if hasattr(strategy, "on_before_select"):
+                try:
+                    filtered = strategy.on_before_select(
+                        provider_id=pid,
+                        key_ids=key_ids,
+                        config=self.config,
+                        context=strategy_context,
+                    )
+                    if filtered is not None:
+                        key_ids = filtered
+                except Exception:
+                    logger.opt(exception=True).debug(
+                        "Pool[{}]: strategy before_select failed", pid[:8]
+                    )
+
+        # --- 1. Sticky session ----------------------------------------
+        sticky_key_id: str | None = None
+        if session_uuid and self.config.sticky_session_ttl_seconds > 0:
+            sticky_key_id = await redis_ops.get_sticky_binding(
+                pid, session_uuid, self.config.sticky_session_ttl_seconds
+            )
+
+        # --- 2. Batch fetch pool state (parallel) ---------------------
+        all_key_ids = [str(c.key.id) for c in candidates]
+
+        # Fire independent Redis queries concurrently.
+        _cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=True)
+        _cost_coro = (
+            redis_ops.batch_get_cost_totals(pid, all_key_ids, self.config.cost_window_seconds)
+            if self.config.cost_limit_per_key_tokens is not None
+            else None
+        )
+        _lru_coro = redis_ops.get_lru_scores(pid, all_key_ids) if self.config.lru_enabled else None
+
+        # Gather all non-None coroutines in parallel.
+        coros: list[Any] = [_cooldown_coro]
+        _cost_idx = -1
+        _lru_idx = -1
+        if _cost_coro is not None:
+            _cost_idx = len(coros)
+            coros.append(_cost_coro)
+        if _lru_coro is not None:
+            _lru_idx = len(coros)
+            coros.append(_lru_coro)
+
+        gathered = await asyncio.gather(*coros)
+
+        cooldowns_raw = gathered[0]
+        # cooldowns_raw: dict[str, tuple[str | None, int | None]]
+        cooldowns: dict[str, str | None] = {}
+        cooldown_ttls: dict[str, int | None] = {}
+        for kid, val in cooldowns_raw.items():
+            if isinstance(val, tuple):
+                cooldowns[kid] = val[0]
+                cooldown_ttls[kid] = val[1]
+            else:
+                cooldowns[kid] = val
+                cooldown_ttls[kid] = None
+
+        # Cost check
+        cost_exhausted: set[str] = set()
+        cost_soft: set[str] = set()
+        cost_totals: dict[str, int] = {}
+        if _cost_idx >= 0:
+            cost_totals = gathered[_cost_idx]
+            limit = self.config.cost_limit_per_key_tokens
+            assert limit is not None  # guarded by _cost_idx >= 0
+            for kid, total in cost_totals.items():
+                if total >= limit:
+                    cost_exhausted.add(kid)
+                elif total >= limit * self.config.cost_soft_threshold_percent / 100:
+                    cost_soft.add(kid)
+
+        # LRU scores
+        lru_scores: dict[str, float] = {}
+        if _lru_idx >= 0:
+            lru_scores = gathered[_lru_idx]
+
+        # --- Strategy: compute_score ----------------------------------
+        for strategy in strategies:
+            if hasattr(strategy, "compute_score"):
+                for kid in all_key_ids:
+                    try:
+                        custom = strategy.compute_score(
+                            key_id=kid,
+                            config=self.config,
+                            context=strategy_context,
+                        )
+                        if custom is not None:
+                            lru_scores[kid] = custom
+                    except Exception:
+                        pass
+
+        # --- 3. Classify candidates -----------------------------------
+        sticky_candidate: ProviderCandidate | None = None
+        available: list[ProviderCandidate] = []
+        skipped: list[ProviderCandidate] = []
+
+        for c in candidates:
+            kid = str(c.key.id)
+            ct = PoolCandidateTrace(key_id=kid)
+
+            # Already skipped upstream?
+            if c.is_skipped:
+                skipped.append(c)
+                ct.skipped = True
+                ct.skip_type = "upstream"
+                trace.candidate_traces[kid] = ct
+                continue
+
+            # Cooldown?
+            cd_reason = cooldowns.get(kid)
+            if cd_reason is not None:
+                c.is_skipped = True
+                c.skip_reason = f"pool cooldown: {cd_reason}"
+                skipped.append(c)
+                ct.skipped = True
+                ct.skip_type = "cooldown"
+                ct.cooldown_reason = cd_reason
+                ct.cooldown_ttl = cooldown_ttls.get(kid)
+                _attach_pool_extra(c, ct)
+                trace.candidate_traces[kid] = ct
+                continue
+
+            # Cost exhausted?
+            if kid in cost_exhausted:
+                c.is_skipped = True
+                c.skip_reason = "pool cost limit reached"
+                skipped.append(c)
+                ct.skipped = True
+                ct.skip_type = "cost_exhausted"
+                ct.cost_window_usage = cost_totals.get(kid, 0)
+                ct.cost_limit = self.config.cost_limit_per_key_tokens
+                _attach_pool_extra(c, ct)
+                trace.candidate_traces[kid] = ct
+                continue
+
+            # Sticky hit?
+            if sticky_key_id and kid == sticky_key_id:
+                sticky_candidate = c
+                ct.reason = "sticky"
+                ct.sticky_hit = True
+                trace.sticky_session_used = True
+            else:
+                available.append(c)
+                ct.reason = "lru" if lru_scores.get(kid, 0) > 0 else "random"
+
+            ct.lru_score = lru_scores.get(kid, 0.0)
+            ct.cost_window_usage = cost_totals.get(kid, 0)
+            ct.cost_limit = self.config.cost_limit_per_key_tokens
+            if kid in cost_soft:
+                ct.cost_soft_threshold = True
+            _attach_pool_extra(c, ct)
+            trace.candidate_traces[kid] = ct
+
+        # --- 4. Sort available by LRU ---------------------------------
+        if lru_scores and available:
+            available.sort(key=lambda c: lru_scores.get(str(c.key.id), 0.0))
+
+        # Random tiebreak among candidates with the same LRU score
+        if len(available) > 1 and lru_scores:
+            _shuffle_same_score_groups(available, lru_scores)
+
+        # --- 5. Assemble final order ----------------------------------
+        result: list[ProviderCandidate] = []
+        if sticky_candidate is not None:
+            result.append(sticky_candidate)
+        result.extend(available)
+        result.extend(skipped)
+
+        if sticky_candidate:
+            logger.debug(
+                "Pool[{}]: sticky hit key={}",
+                pid[:8],
+                sticky_key_id and sticky_key_id[:8],
+            )
+
+        # --- Strategy: after_select -----------------------------------
+        if result:
+            first_kid = str(result[0].key.id)
+            first_trace = trace.candidate_traces.get(first_kid)
+            for strategy in strategies:
+                if hasattr(strategy, "on_after_select") and first_trace:
+                    try:
+                        strategy.on_after_select(
+                            provider_id=pid,
+                            selected_key_id=first_kid,
+                            trace=first_trace,
+                            config=self.config,
+                            context=strategy_context,
+                        )
+                    except Exception:
+                        pass
+
+        # Attach the full trace to the first candidate for downstream use.
+        if result:
+            setattr(result[0], "_pool_scheduling_trace", trace)
+
+        return result
+
+    # ------------------------------------------------------------------
+    # Single-key selection (used by CandidateBuilder for pooled providers)
+    # ------------------------------------------------------------------
+
+    async def select_key(
+        self,
+        session_uuid: str | None,
+        keys: list[ProviderAPIKey],
+    ) -> ProviderAPIKey | None:
+        """Select the best key from *keys* according to pool rules.
+
+        Same logic as :meth:`reorder_candidates` but operates directly on
+        :class:`ProviderAPIKey` objects instead of candidates:
+
+        1. Sticky session hit (if bound and still healthy).
+        2. Filter out keys in cooldown or cost-exhausted.
+        3. LRU sort among remaining keys.
+        4. Random tiebreak for identical LRU scores.
+        5. Return the first available key, or ``None``.
+        """
+        if not keys:
+            return None
+
+        pid = self.provider_id
+
+        # --- 1. Sticky session ------------------------------------------------
+        sticky_key_id: str | None = None
+        if session_uuid and self.config.sticky_session_ttl_seconds > 0:
+            sticky_key_id = await redis_ops.get_sticky_binding(
+                pid, session_uuid, self.config.sticky_session_ttl_seconds
+            )
+
+        # --- 2. Batch fetch pool state (parallel) -----------------------------
+        key_ids = [str(k.id) for k in keys]
+
+        _cooldown_coro = redis_ops.batch_get_cooldowns(pid, key_ids)
+        _cost_coro = (
+            redis_ops.batch_get_cost_totals(pid, key_ids, self.config.cost_window_seconds)
+            if self.config.cost_limit_per_key_tokens is not None
+            else None
+        )
+        _lru_coro = redis_ops.get_lru_scores(pid, key_ids) if self.config.lru_enabled else None
+
+        coros_sk: list[Any] = [_cooldown_coro]
+        _cost_idx_sk = -1
+        _lru_idx_sk = -1
+        if _cost_coro is not None:
+            _cost_idx_sk = len(coros_sk)
+            coros_sk.append(_cost_coro)
+        if _lru_coro is not None:
+            _lru_idx_sk = len(coros_sk)
+            coros_sk.append(_lru_coro)
+
+        gathered_sk = await asyncio.gather(*coros_sk)
+
+        cooldowns = gathered_sk[0]
+
+        cost_exhausted: set[str] = set()
+        if _cost_idx_sk >= 0:
+            cost_totals = gathered_sk[_cost_idx_sk]
+            for kid, total in cost_totals.items():
+                if total >= self.config.cost_limit_per_key_tokens:  # type: ignore[operator]
+                    cost_exhausted.add(kid)
+
+        lru_scores: dict[str, float] = {}
+        if _lru_idx_sk >= 0:
+            lru_scores = gathered_sk[_lru_idx_sk]
+
+        # --- Strategy: compute_score ------------------------------------------
+        strategies = _get_active_strategies(self.config)
+        strategy_context: dict[str, Any] = {"session_uuid": session_uuid}
+        for strategy in strategies:
+            if hasattr(strategy, "compute_score"):
+                for kid in key_ids:
+                    try:
+                        custom = strategy.compute_score(
+                            key_id=kid,
+                            config=self.config,
+                            context=strategy_context,
+                        )
+                        if custom is not None:
+                            lru_scores[kid] = custom
+                    except Exception:
+                        pass
+
+        # --- 3. Classify keys -------------------------------------------------
+        sticky_key: ProviderAPIKey | None = None
+        available: list[ProviderAPIKey] = []
+
+        for k in keys:
+            kid = str(k.id)
+
+            if cooldowns.get(kid) is not None:
+                continue
+            if kid in cost_exhausted:
+                continue
+
+            if sticky_key_id and kid == sticky_key_id:
+                sticky_key = k
+                continue
+
+            available.append(k)
+
+        # --- 4. Sort by LRU ---------------------------------------------------
+        if lru_scores and available:
+            available.sort(key=lambda k: lru_scores.get(str(k.id), 0.0))
+
+        # Random tiebreak within same-score groups
+        if len(available) > 1 and lru_scores:
+            _shuffle_same_score_keys(available, lru_scores)
+
+        # --- 5. Pick the winner -----------------------------------------------
+        if sticky_key is not None:
+            logger.debug(
+                "Pool[{}]: sticky select key={}",
+                pid[:8],
+                sticky_key_id and sticky_key_id[:8],
+            )
+            return sticky_key
+
+        return available[0] if available else None
+
+    # ------------------------------------------------------------------
+    # Post-request hooks
+    # ------------------------------------------------------------------
+
+    async def on_request_success(
+        self,
+        *,
+        session_uuid: str | None,
+        key_id: str,
+        tokens_used: int = 0,
+    ) -> None:
+        """Called after a successful upstream request."""
+        pid = self.provider_id
+
+        # Bind sticky session
+        if session_uuid and self.config.sticky_session_ttl_seconds > 0:
+            await redis_ops.set_sticky_binding(
+                pid, session_uuid, key_id, self.config.sticky_session_ttl_seconds
+            )
+
+        # Touch LRU
+        if self.config.lru_enabled:
+            await redis_ops.touch_lru(pid, key_id)
+
+        # Record cost
+        if tokens_used > 0 and self.config.cost_limit_per_key_tokens is not None:
+            await redis_ops.add_cost_entry(
+                pid, key_id, tokens_used, self.config.cost_window_seconds
+            )
+
+    async def on_request_error(
+        self,
+        *,
+        key_id: str,
+        status_code: int,
+        error_body: str | None = None,
+        response_headers: dict[str, str] | None = None,
+    ) -> None:
+        """Called after an upstream error. Delegates to health policy."""
+        # Import lazily to avoid circular deps
+        from src.services.provider.pool.health_policy import apply_health_policy
+
+        await apply_health_policy(
+            provider_id=self.provider_id,
+            key_id=key_id,
+            status_code=status_code,
+            error_body=error_body,
+            response_headers=response_headers,
+            config=self.config,
+        )
+
+    # ------------------------------------------------------------------
+    # Key schedulability check (used by candidate_builder)
+    # ------------------------------------------------------------------
+
+    async def is_key_schedulable(self, key_id: str) -> tuple[bool, str | None]:
+        """Check if *key_id* is currently schedulable (not in cooldown, not
+        cost-exhausted).  Returns ``(True, None)`` or ``(False, reason)``.
+        """
+        pid = self.provider_id
+
+        # Cooldown check
+        cd = await redis_ops.get_cooldown(pid, key_id)
+        if cd is not None:
+            return False, f"pool cooldown: {cd}"
+
+        # Cost check
+        if self.config.cost_limit_per_key_tokens is not None:
+            total = await redis_ops.get_cost_window_total(
+                pid, key_id, self.config.cost_window_seconds
+            )
+            if total >= self.config.cost_limit_per_key_tokens:
+                return False, "pool cost limit reached"
+
+        return True, None
+
+
+# Backward-compatible alias
+ClaudeCodePoolManager = PoolManager
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+_T = TypeVar("_T")
+
+
+def _attach_pool_extra(candidate: Any, ct: PoolCandidateTrace) -> None:
+    """Attach pool trace extra_data onto a candidate object."""
+    existing = getattr(candidate, "_pool_extra_data", None) or {}
+    existing.update(ct.to_extra_data())
+    setattr(candidate, "_pool_extra_data", existing)
+
+
+def _get_active_strategies(config: PoolConfig) -> list[Any]:
+    """Get active strategies for the given config (lazy import)."""
+    if not config.strategies:
+        return []
+    try:
+        from src.services.provider.pool.strategy import get_active_strategies
+
+        return get_active_strategies(config.strategies)
+    except Exception:
+        return []
+
+
+def _shuffle_same_score(
+    items: list[_T],
+    lru_scores: dict[str, float],
+    key_fn: Callable[[_T], str],
+) -> None:
+    """In-place random shuffle within groups that share the same LRU score."""
+    if len(items) <= 1:
+        return
+
+    i = 0
+    while i < len(items):
+        score_i = lru_scores.get(key_fn(items[i]), 0.0)
+        j = i + 1
+        while j < len(items) and lru_scores.get(key_fn(items[j]), 0.0) == score_i:
+            j += 1
+        if j - i > 1:
+            group = items[i:j]
+            random.shuffle(group)
+            items[i:j] = group
+        i = j
+
+
+def _shuffle_same_score_groups(
+    candidates: list[ProviderCandidate],
+    lru_scores: dict[str, float],
+) -> None:
+    _shuffle_same_score(candidates, lru_scores, lambda c: str(c.key.id))
+
+
+def _shuffle_same_score_keys(
+    keys: list[ProviderAPIKey],
+    lru_scores: dict[str, float],
+) -> None:
+    _shuffle_same_score(keys, lru_scores, lambda k: str(k.id))

src/services/provider/pool/oauth_cache.py

@@ -0,0 +1,42 @@
+"""OAuth token Redis cache for the Account Pool.
+
+Additions over the base ``auth.py`` refresh flow:
+
+- **Redis token cache**: Avoids repeated DB decryption for hot keys.
+  Cache key: ``provider_oauth_token_cache:{key_id}``
+- **Configurable proactive refresh skew**: Default 180 s (3 min) instead
+  of the base 120 s, configurable via ``PoolConfig.proactive_refresh_seconds``.
+- **401 immediate invalidation**: Clears the Redis cache so the next request
+  triggers a fresh refresh.
+
+This module does NOT replace ``auth.py``; it adds a caching layer that
+``auth.py`` can consult before decrypting from DB.
+"""
+
+from __future__ import annotations
+
+from src.core.logger import logger
+from src.services.provider.pool import redis_ops
+
+
+async def get_cached_token(key_id: str) -> str | None:
+    """Return cached access token from Redis, or None."""
+    return await redis_ops.get_cached_oauth_token(key_id)
+
+
+async def cache_token(key_id: str, token: str, expires_in_seconds: int) -> None:
+    """Cache an access token in Redis.
+
+    *expires_in_seconds* is the remaining lifetime of the token.  We shave
+    off 60 s so the cache expires slightly before the token itself, giving
+    the refresh flow time to act.
+    """
+    ttl = max(1, expires_in_seconds - 60)
+    await redis_ops.cache_oauth_token(key_id, token, ttl)
+    logger.debug("Pool OAuth: cached token for key {} (TTL={}s)", key_id[:8], ttl)
+
+
+async def invalidate_token(key_id: str) -> None:
+    """Invalidate the cached token (e.g. after a 401)."""
+    await redis_ops.invalidate_oauth_token_cache(key_id)
+    logger.debug("Pool OAuth: invalidated token cache for key {}", key_id[:8])

src/services/provider/pool/redis_ops.py

@@ -0,0 +1,470 @@
+"""Redis operations for the Account Pool (provider-agnostic).
+
+All pool transient state is stored in Redis. This module centralises key
+naming, Lua scripts, and graceful fallbacks so that the rest of the pool
+layer is free of Redis specifics.
+
+Key schema
+----------
+ap:{pid}:sticky:{session_uuid}             STRING  -> key_id   (TTL: config)
+ap:{pid}:lru                               ZSET    member=key_id, score=unix_ts
+ap:{pid}:cooldown:{key_id}                 STRING  -> reason   (TTL: error-specific)
+ap:{pid}:cost:{key_id}                     ZSET    member=req_id, score=unix_ts
+provider_oauth_token_cache:{key_id}        STRING  -> access_token (TTL: expires - 60)
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+from typing import TYPE_CHECKING
+
+from src.clients.redis_client import get_redis_client
+from src.core.logger import logger
+
+if TYPE_CHECKING:
+    import redis.asyncio as aioredis
+
+PREFIX = "ap"
+
+
+def _sticky_key(provider_id: str, session_uuid: str) -> str:
+    return f"{PREFIX}:{provider_id}:sticky:{session_uuid}"
+
+
+def _lru_key(provider_id: str) -> str:
+    return f"{PREFIX}:{provider_id}:lru"
+
+
+def _cooldown_key(provider_id: str, key_id: str) -> str:
+    return f"{PREFIX}:{provider_id}:cooldown:{key_id}"
+
+
+def _cost_key(provider_id: str, key_id: str) -> str:
+    return f"{PREFIX}:{provider_id}:cost:{key_id}"
+
+
+def _oauth_cache_key(key_id: str) -> str:
+    return f"provider_oauth_token_cache:{key_id}"
+
+
+# ---------------------------------------------------------------------------
+# Lua scripts
+# ---------------------------------------------------------------------------
+
+# Sticky-select: GET binding, verify it's not in cooldown, refresh TTL.
+# KEYS[1] = sticky key, KEYS[2] = cooldown key prefix (ap:{pid}:cooldown:)
+# ARGV[1] = ttl
+# Returns: key_id or nil
+_STICKY_SELECT_LUA = """
+local binding = redis.call("GET", KEYS[1])
+if not binding then
+    return nil
+end
+-- Check cooldown for the bound key
+local cooldown_key = KEYS[2] .. binding
+local in_cooldown = redis.call("EXISTS", cooldown_key)
+if in_cooldown == 1 then
+    redis.call("DEL", KEYS[1])
+    return nil
+end
+redis.call("EXPIRE", KEYS[1], tonumber(ARGV[1]))
+return binding
+"""
+
+# Cost window cleanup + sum tokens in a single round-trip.
+# KEYS[1] = cost zset key, ARGV[1] = window_start timestamp
+# Returns total token count within the window.
+_COST_WINDOW_SUM_LUA = """
+local key = KEYS[1]
+local window_start = tonumber(ARGV[1])
+redis.call("ZREMRANGEBYSCORE", key, "-inf", window_start)
+local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
+local total = 0
+for _, m in ipairs(members) do
+    local colon = string.find(m, ":", 1, true)
+    if colon then
+        local n = tonumber(string.sub(m, colon + 1))
+        if n then total = total + n end
+    end
+end
+return total
+"""
+
+
+async def _get_redis() -> "aioredis.Redis | None":
+    return await get_redis_client(require_redis=False)
+
+
+# ---------------------------------------------------------------------------
+# Sticky session
+# ---------------------------------------------------------------------------
+
+
+async def get_sticky_binding(provider_id: str, session_uuid: str, ttl: int) -> str | None:
+    """Get and refresh sticky session binding. Returns key_id or None."""
+    redis = await _get_redis()
+    if redis is None:
+        return None
+    try:
+        result = await redis.eval(
+            _STICKY_SELECT_LUA,
+            2,
+            _sticky_key(provider_id, session_uuid),
+            f"{PREFIX}:{provider_id}:cooldown:",
+            str(ttl),
+        )
+        if result:
+            return result.decode() if isinstance(result, bytes) else str(result)
+        return None
+    except Exception:
+        logger.debug("Pool: sticky GET failed for session {}", session_uuid[:8])
+        return None
+
+
+async def set_sticky_binding(provider_id: str, session_uuid: str, key_id: str, ttl: int) -> None:
+    """Create or update sticky session binding."""
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        await redis.setex(_sticky_key(provider_id, session_uuid), ttl, key_id)
+    except Exception:
+        logger.debug("Pool: sticky SET failed for session {}", session_uuid[:8])
+
+
+async def delete_sticky_binding(provider_id: str, session_uuid: str) -> None:
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        await redis.delete(_sticky_key(provider_id, session_uuid))
+    except Exception:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# LRU
+# ---------------------------------------------------------------------------
+
+
+async def get_lru_scores(provider_id: str, key_ids: list[str]) -> dict[str, float]:
+    """Batch-fetch LRU timestamps. Missing keys get score 0 (highest priority)."""
+    redis = await _get_redis()
+    if redis is None:
+        return {}
+    try:
+        lru_k = _lru_key(provider_id)
+        scores = await redis.zmscore(lru_k, key_ids)
+        result: dict[str, float] = {}
+        for kid, score in zip(key_ids, scores):
+            result[kid] = float(score) if score is not None else 0.0
+        return result
+    except Exception:
+        logger.debug("Pool: LRU ZMSCORE failed for provider {}", provider_id[:8])
+        return {}
+
+
+async def touch_lru(provider_id: str, key_id: str) -> None:
+    """Update last-used timestamp."""
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        await redis.zadd(_lru_key(provider_id), {key_id: time.time()})
+    except Exception:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# Cooldown
+# ---------------------------------------------------------------------------
+
+
+async def set_cooldown(provider_id: str, key_id: str, reason: str, ttl: int) -> None:
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        await redis.setex(_cooldown_key(provider_id, key_id), ttl, reason)
+        logger.info(
+            "Pool: key {} cooldown set: {} ({}s)",
+            key_id[:8],
+            reason,
+            ttl,
+        )
+    except Exception:
+        logger.debug("Pool: cooldown SET failed for key {}", key_id[:8])
+
+
+async def get_cooldown(provider_id: str, key_id: str) -> str | None:
+    redis = await _get_redis()
+    if redis is None:
+        return None
+    try:
+        val = await redis.get(_cooldown_key(provider_id, key_id))
+        if val:
+            return val.decode() if isinstance(val, bytes) else str(val)
+        return None
+    except Exception:
+        return None
+
+
+async def clear_cooldown(provider_id: str, key_id: str) -> None:
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        await redis.delete(_cooldown_key(provider_id, key_id))
+    except Exception:
+        pass
+
+
+async def batch_get_cooldowns(
+    provider_id: str,
+    key_ids: list[str],
+    *,
+    include_ttl: bool = False,
+) -> dict[str, str | None] | dict[str, tuple[str | None, int | None]]:
+    """Batch check cooldown status for multiple keys.
+
+    When *include_ttl* is ``True``, each value is a ``(reason, ttl_seconds)``
+    tuple instead of a plain reason string.  The TTL commands are batched in
+    the same pipeline so there is no extra round-trip.
+    """
+    redis = await _get_redis()
+    if redis is None:
+        if include_ttl:
+            return {k: (None, None) for k in key_ids}
+        return {k: None for k in key_ids}
+    try:
+        pipe = redis.pipeline()
+        for kid in key_ids:
+            ck = _cooldown_key(provider_id, kid)
+            pipe.get(ck)
+            if include_ttl:
+                pipe.ttl(ck)
+        results = await pipe.execute()
+
+        if include_ttl:
+            out_ttl: dict[str, tuple[str | None, int | None]] = {}
+            # results interleave GET/TTL: [val0, ttl0, val1, ttl1, ...]
+            for i, kid in enumerate(key_ids):
+                val = results[i * 2]
+                ttl_val = results[i * 2 + 1]
+                reason: str | None = None
+                if val:
+                    reason = val.decode() if isinstance(val, bytes) else str(val)
+                ttl_sec: int | None = None
+                if isinstance(ttl_val, int) and ttl_val > 0:
+                    ttl_sec = ttl_val
+                out_ttl[kid] = (reason, ttl_sec)
+            return out_ttl
+
+        out: dict[str, str | None] = {}
+        for kid, val in zip(key_ids, results):
+            if val:
+                out[kid] = val.decode() if isinstance(val, bytes) else str(val)
+            else:
+                out[kid] = None
+        return out
+    except Exception:
+        if include_ttl:
+            return {k: (None, None) for k in key_ids}
+        return {k: None for k in key_ids}
+
+
+# ---------------------------------------------------------------------------
+# Cost tracking
+# ---------------------------------------------------------------------------
+
+
+async def add_cost_entry(provider_id: str, key_id: str, tokens: int, window_seconds: int) -> None:
+    """Record a cost entry (tokens used) with automatic window expiry."""
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        now = time.time()
+        cost_k = _cost_key(provider_id, key_id)
+        member = f"{uuid.uuid4().hex}:{tokens}"
+        pipe = redis.pipeline()
+        pipe.zadd(cost_k, {member: now})
+        # Set a TTL slightly larger than the window to auto-clean abandoned keys.
+        pipe.expire(cost_k, window_seconds + 600)
+        await pipe.execute()
+    except Exception:
+        logger.debug("Pool: cost ADD failed for key {}", key_id[:8])
+
+
+async def get_cost_window_total(provider_id: str, key_id: str, window_seconds: int) -> int:
+    """Sum tokens used within the rolling window (single key)."""
+    redis = await _get_redis()
+    if redis is None:
+        return 0
+    try:
+        now = time.time()
+        window_start = now - window_seconds
+        cost_k = _cost_key(provider_id, key_id)
+        result = await redis.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
+        return int(result) if result else 0
+    except Exception:
+        logger.debug("Pool: cost SUM failed for key {}", key_id[:8])
+        return 0
+
+
+async def batch_get_cost_totals(
+    provider_id: str, key_ids: list[str], window_seconds: int
+) -> dict[str, int]:
+    """Batch-fetch cost totals for multiple keys using pipeline + Lua."""
+    redis = await _get_redis()
+    if redis is None:
+        return {k: 0 for k in key_ids}
+    try:
+        now = time.time()
+        window_start = now - window_seconds
+        pipe = redis.pipeline()
+        for kid in key_ids:
+            cost_k = _cost_key(provider_id, kid)
+            pipe.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
+        results = await pipe.execute()
+        out: dict[str, int] = {}
+        for kid, val in zip(key_ids, results):
+            out[kid] = int(val) if val else 0
+        return out
+    except Exception:
+        logger.debug("Pool: batch cost SUM failed for provider {}", provider_id[:8])
+        return {k: 0 for k in key_ids}
+
+
+async def clear_cost(provider_id: str, key_id: str) -> None:
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        await redis.delete(_cost_key(provider_id, key_id))
+    except Exception:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# OAuth token cache
+# ---------------------------------------------------------------------------
+
+
+async def cache_oauth_token(key_id: str, token: str, ttl: int) -> None:
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        if ttl > 0:
+            await redis.setex(_oauth_cache_key(key_id), ttl, token)
+    except Exception:
+        pass
+
+
+async def get_cached_oauth_token(key_id: str) -> str | None:
+    redis = await _get_redis()
+    if redis is None:
+        return None
+    try:
+        val = await redis.get(_oauth_cache_key(key_id))
+        if val:
+            return val.decode() if isinstance(val, bytes) else str(val)
+        return None
+    except Exception:
+        return None
+
+
+async def invalidate_oauth_token_cache(key_id: str) -> None:
+    redis = await _get_redis()
+    if redis is None:
+        return
+    try:
+        await redis.delete(_oauth_cache_key(key_id))
+    except Exception:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# Pool status query (admin)
+# ---------------------------------------------------------------------------
+
+
+async def get_sticky_session_count(provider_id: str) -> int:
+    """Approximate count of active sticky sessions (via SCAN, for admin display only)."""
+    redis = await _get_redis()
+    if redis is None:
+        return 0
+    try:
+        pattern = f"{PREFIX}:{provider_id}:sticky:*"
+        count = 0
+        async for _ in redis.scan_iter(match=pattern, count=100):
+            count += 1
+        return count
+    except Exception:
+        return 0
+
+
+async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
+    """Count sticky sessions bound to a specific key (admin only).
+
+    Uses batched SCAN + pipeline MGET to reduce Redis round-trips.
+    """
+    redis = await _get_redis()
+    if redis is None:
+        return 0
+    try:
+        pattern = f"{PREFIX}:{provider_id}:sticky:*"
+        count = 0
+        batch: list[bytes | str] = []
+        async for k in redis.scan_iter(match=pattern, count=200):
+            batch.append(k)
+            if len(batch) >= 200:
+                vals = await redis.mget(batch)
+                for val in vals:
+                    if val:
+                        bound_id = val.decode() if isinstance(val, bytes) else str(val)
+                        if bound_id == key_id:
+                            count += 1
+                batch.clear()
+        if batch:
+            vals = await redis.mget(batch)
+            for val in vals:
+                if val:
+                    bound_id = val.decode() if isinstance(val, bytes) else str(val)
+                    if bound_id == key_id:
+                        count += 1
+        return count
+    except Exception:
+        return 0
+
+
+async def get_cooldown_ttl(provider_id: str, key_id: str) -> int | None:
+    """Get remaining cooldown TTL in seconds. None = no cooldown."""
+    redis = await _get_redis()
+    if redis is None:
+        return None
+    try:
+        ttl = await redis.ttl(_cooldown_key(provider_id, key_id))
+        return ttl if ttl > 0 else None
+    except Exception:
+        return None
+
+
+async def batch_get_cooldown_ttls(provider_id: str, key_ids: list[str]) -> dict[str, int | None]:
+    """Batch-fetch cooldown TTLs for multiple keys using pipeline."""
+    redis = await _get_redis()
+    if redis is None:
+        return {k: None for k in key_ids}
+    try:
+        pipe = redis.pipeline()
+        for kid in key_ids:
+            pipe.ttl(_cooldown_key(provider_id, kid))
+        results = await pipe.execute()
+        out: dict[str, int | None] = {}
+        for kid, ttl in zip(key_ids, results):
+            out[kid] = int(ttl) if isinstance(ttl, int) and ttl > 0 else None
+        return out
+    except Exception:
+        return {k: None for k in key_ids}

src/services/provider/pool/strategy.py

@@ -0,0 +1,107 @@
+"""Pluggable pool scheduling strategies.
+
+Strategies allow customising pool-level candidate selection without
+modifying the core :class:`PoolManager`.  Each strategy is an object
+that implements one or more optional methods defined by the
+:class:`PoolSchedulingStrategy` protocol.
+
+Registration uses a thread-safe global registry (same pattern as
+:mod:`~src.services.provider.pool.hooks`).
+
+Usage::
+
+    from src.services.provider.pool.strategy import register_pool_strategy
+
+    class MyStrategy:
+        name = "usage_weight"
+
+        def compute_score(self, *, key_id, config, context):
+            ...
+
+    register_pool_strategy("usage_weight", MyStrategy())
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from src.services.provider.pool.config import PoolConfig
+    from src.services.provider.pool.trace import PoolCandidateTrace
+
+
+@runtime_checkable
+class PoolSchedulingStrategy(Protocol):
+    """Pluggable pool scheduling strategy.
+
+    All methods are optional -- callers check via ``hasattr``.
+    Strategies are activated per-provider through ``PoolConfig.strategies``.
+    """
+
+    name: str
+
+    def on_before_select(
+        self,
+        *,
+        provider_id: str,
+        key_ids: list[str],
+        config: PoolConfig,
+        context: dict[str, Any],
+    ) -> list[str] | None:
+        """Filter / reorder *key_ids* before selection.
+
+        Return ``None`` to leave the list unchanged.
+        """
+        ...
+
+    def on_after_select(
+        self,
+        *,
+        provider_id: str,
+        selected_key_id: str,
+        trace: PoolCandidateTrace,
+        config: PoolConfig,
+        context: dict[str, Any],
+    ) -> None:
+        """Called after a key has been selected (for logging / metrics)."""
+        ...
+
+    def compute_score(
+        self,
+        *,
+        key_id: str,
+        config: PoolConfig,
+        context: dict[str, Any],
+    ) -> float | None:
+        """Return a custom sort score.  ``None`` means "do not override"."""
+        ...
+
+
+# ---------------------------------------------------------------------------
+# Registry
+# ---------------------------------------------------------------------------
+
+_strategy_registry: dict[str, PoolSchedulingStrategy] = {}
+_strategy_lock = threading.Lock()
+
+
+def register_pool_strategy(name: str, strategy: PoolSchedulingStrategy) -> None:
+    """Register a pool scheduling strategy globally."""
+    with _strategy_lock:
+        _strategy_registry[name] = strategy
+
+
+def get_pool_strategy(name: str) -> PoolSchedulingStrategy | None:
+    """Return a registered strategy by *name*, or ``None``."""
+    return _strategy_registry.get(name)
+
+
+def get_active_strategies(names: tuple[str, ...] | list[str]) -> list[PoolSchedulingStrategy]:
+    """Return registered strategies whose names appear in *names*."""
+    result: list[PoolSchedulingStrategy] = []
+    for n in names:
+        s = _strategy_registry.get(n)
+        if s is not None:
+            result.append(s)
+    return result

src/services/provider/pool/trace.py

@@ -0,0 +1,96 @@
+"""Pool scheduling trace -- per-candidate decision records.
+
+Collects scheduling decisions made during pool-level candidate selection
+without adding any extra Redis round-trips.  Trace data is later written
+to ``RequestCandidate.extra_data`` and ``Usage.request_metadata``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(slots=True)
+class PoolCandidateTrace:
+    """Single candidate scheduling decision in a pool context."""
+
+    key_id: str
+    reason: str = ""  # sticky / lru / random / tiebreak
+    sticky_hit: bool = False
+    lru_score: float = 0.0
+    cost_window_usage: int = 0
+    cost_limit: int | None = None
+    cost_soft_threshold: bool = False
+    skipped: bool = False
+    skip_type: str | None = None  # cooldown / cost_exhausted
+    cooldown_reason: str | None = None
+    cooldown_ttl: int | None = None
+
+    def to_extra_data(self) -> dict[str, Any]:
+        """Build dict to merge into ``RequestCandidate.extra_data``."""
+        if self.skipped:
+            skip_info: dict[str, Any] = {"type": self.skip_type}
+            if self.cooldown_reason is not None:
+                skip_info["cooldown_reason"] = self.cooldown_reason
+            if self.cooldown_ttl is not None:
+                skip_info["cooldown_ttl"] = self.cooldown_ttl
+            if self.cost_window_usage:
+                skip_info["cost_window_usage"] = self.cost_window_usage
+            return {"pool_skip": skip_info}
+
+        sel: dict[str, Any] = {"reason": self.reason}
+        if self.sticky_hit:
+            sel["sticky_hit"] = True
+        if self.lru_score:
+            sel["lru_score"] = self.lru_score
+        if self.cost_window_usage:
+            sel["cost_window_usage"] = self.cost_window_usage
+        if self.cost_limit is not None:
+            sel["cost_limit"] = self.cost_limit
+        if self.cost_soft_threshold:
+            sel["cost_soft_threshold"] = True
+        return {"pool_selection": sel}
+
+
+@dataclass(slots=True)
+class PoolSchedulingTrace:
+    """Aggregated scheduling trace for one pool-provider dispatch."""
+
+    provider_id: str
+    total_keys: int = 0
+    sticky_session_used: bool = False
+    session_uuid: str | None = None
+    candidate_traces: dict[str, PoolCandidateTrace] = field(default_factory=dict)
+
+    def build_summary(self, success_key_id: str | None = None) -> dict[str, Any]:
+        """Build compact dict for ``Usage.request_metadata["pool_summary"]``."""
+        skipped_cooldown = 0
+        skipped_cost = 0
+        attempted = 0
+        for t in self.candidate_traces.values():
+            if t.skipped:
+                if t.skip_type == "cooldown":
+                    skipped_cooldown += 1
+                elif t.skip_type == "cost_exhausted":
+                    skipped_cost += 1
+            else:
+                attempted += 1
+
+        success_reason: str | None = None
+        if success_key_id and success_key_id in self.candidate_traces:
+            success_reason = self.candidate_traces[success_key_id].reason
+
+        summary: dict[str, Any] = {
+            "enabled": True,
+            "total_keys": self.total_keys,
+            "attempted": attempted,
+            "skipped_cooldown": skipped_cooldown,
+            "skipped_cost": skipped_cost,
+            "sticky_session": self.sticky_session_used,
+        }
+        if success_key_id:
+            summary["success_key_id"] = success_key_id[:8]
+        if success_reason:
+            summary["success_reason"] = success_reason
+        return summary