feat: Enhance proxy clearance management with support for clearance origins and host handling

app/control/proxy/__init__.py

@@ -6,6 +6,7 @@
 """
 
 import asyncio
+from urllib.parse import urlparse
 
 from app.platform.logging.logger import logger
 from app.platform.config.snapshot import get_config
@@ -26,6 +27,14 @@
 from .providers.manual import ManualClearanceProvider
 from .providers.flaresolverr import FlareSolverrClearanceProvider
 
+_DEFAULT_CLEARANCE_ORIGIN = "https://grok.com"
+BundleKey = tuple[str, str]
+
+
+def _clearance_host(clearance_origin: str | None) -> str:
+    host = urlparse(clearance_origin or _DEFAULT_CLEARANCE_ORIGIN).hostname
+    return (host or "grok.com").lower()
+
 
 class ProxyDirectory:
     """Owns egress nodes and clearance bundles.
@@ -36,11 +45,11 @@ class ProxyDirectory:
     def __init__(self) -> None:
         self._nodes: list[EgressNode] = []
         self._resource_nodes: list[EgressNode] = []  # for media downloads
-        self._bundles: dict[str, ClearanceBundle] = {}
+        self._bundles: dict[BundleKey, ClearanceBundle] = {}
         self._lock = asyncio.Lock()
-        # Single-flight guard: at most one FlareSolverr call per affinity key.
+        # Single-flight guard: at most one FlareSolverr call per proxy+host key.
         # Other coroutines wait on the Event until the active refresh completes.
-        self._refresh_events: dict[str, asyncio.Event] = {}
+        self._refresh_events: dict[BundleKey, asyncio.Event] = {}
         self._manual = ManualClearanceProvider()
         self._flare = FlareSolverrClearanceProvider()
         self._egress_mode: EgressMode = EgressMode.DIRECT
@@ -117,12 +126,12 @@ async def load(self) -> None:
             self._bundles = {
                 key: bundle.model_copy(update={"state": ClearanceBundleState.INVALID})
                 for key, bundle in self._bundles.items()
-                if key in valid_affinities
+                if key[0] in valid_affinities
             }
             self._refresh_events = {
                 key: event
                 for key, event in self._refresh_events.items()
-                if key in valid_affinities
+                if key[0] in valid_affinities
             }
             self._config_sig = config_sig
 
@@ -144,23 +153,28 @@ async def acquire(
         scope: ProxyScope = ProxyScope.APP,
         kind: RequestKind = RequestKind.HTTP,
         resource: bool = False,
+        clearance_origin: str | None = None,
     ) -> ProxyLease:
         """Return a ProxyLease for the next request.
 
         For DIRECT mode, returns a lease with no proxy or clearance.
         """
         proxy_url = await self._pick_proxy_url(resource=resource)
         affinity = proxy_url or "direct"
+        clearance_host = _clearance_host(clearance_origin)
 
         bundle = await self._get_or_build_bundle(
-            affinity_key=affinity, proxy_url=proxy_url or ""
+            affinity_key=affinity,
+            proxy_url=proxy_url or "",
+            clearance_origin=clearance_origin or _DEFAULT_CLEARANCE_ORIGIN,
         )
 
         return ProxyLease(
             lease_id=next_hex(),
             proxy_url=proxy_url,
             cf_cookies=bundle.cf_cookies if bundle else "",
             user_agent=bundle.user_agent if bundle else "",
+            clearance_host=clearance_host,
             scope=scope,
             kind=kind,
             acquired_at=now_ms(),
@@ -173,13 +187,13 @@ async def feedback(self, lease: ProxyLease, result: ProxyFeedback) -> None:
             ProxyFeedbackKind.UNAUTHORIZED,
         ):
             # Invalidate associated clearance bundle.
-            affinity = lease.proxy_url or "direct"
+            key = (lease.proxy_url or "direct", lease.clearance_host)
             async with self._lock:
-                bundle = self._bundles.get(affinity)
-                if bundle:
-                    from .models import ClearanceBundleState
+                from .models import ClearanceBundleState
 
-                    self._bundles[affinity] = bundle.model_copy(
+                bundle = self._bundles.get(key)
+                if bundle:
+                    self._bundles[key] = bundle.model_copy(
                         update={"state": ClearanceBundleState.INVALID}
                     )
 
@@ -233,41 +247,48 @@ async def _get_or_build_bundle(
         *,
         affinity_key: str,
         proxy_url: str,
+        clearance_origin: str,
     ) -> ClearanceBundle | None:
         if self._clearance_mode == ClearanceMode.NONE:
             return None
+        clearance_host = _clearance_host(clearance_origin)
+        key: BundleKey = (affinity_key, clearance_host)
 
-        # Single-flight: only one coroutine fetches clearance per affinity key.
+        # Single-flight: only one coroutine fetches clearance per proxy+host key.
         # Concurrent callers wait on the Event and retry once it fires.
         while True:
             async with self._lock:
-                bundle = self._bundles.get(affinity_key)
+                bundle = self._bundles.get(key)
                 if bundle and bundle.state.value == 0:  # VALID
                     return bundle
-                event = self._refresh_events.get(affinity_key)
+                event = self._refresh_events.get(key)
                 if event is None:
                     # This coroutine wins the right to refresh.
                     event = asyncio.Event()
-                    self._refresh_events[affinity_key] = event
+                    self._refresh_events[key] = event
                     break
             # Another coroutine is already refreshing — wait for it, then retry.
             await event.wait()
 
         try:
             if self._clearance_mode == ClearanceMode.MANUAL:
-                bundle = self._manual.build_bundle(affinity_key=affinity_key)
+                bundle = self._manual.build_bundle(
+                    affinity_key=affinity_key,
+                    clearance_host=clearance_host,
+                )
             else:
                 bundle = await self._flare.refresh_bundle(
                     affinity_key=affinity_key,
                     proxy_url=proxy_url,
+                    target_url=clearance_origin,
                 )
             if bundle:
                 async with self._lock:
-                    self._bundles[affinity_key] = bundle
+                    self._bundles[key] = bundle
             return bundle
         finally:
             async with self._lock:
-                self._refresh_events.pop(affinity_key, None)
+                self._refresh_events.pop(key, None)
             event.set()  # Wake all waiters so they retry with the new bundle.
 
     # ------------------------------------------------------------------
@@ -308,6 +329,7 @@ async def warm_up(self) -> None:
             await self._get_or_build_bundle(
                 affinity_key=affinity,
                 proxy_url=proxy_url,
+                clearance_origin=_DEFAULT_CLEARANCE_ORIGIN,
             )
 
     async def refresh_clearance_safe(self) -> None:
@@ -322,34 +344,45 @@ async def refresh_clearance_safe(self) -> None:
             return
         async with self._lock:
             nodes = list(self._nodes)
-            existing = set(self._bundles.keys())
+            existing = list(self._bundles.keys())
 
-        affinity_items = (
+        refresh_targets: dict[BundleKey, tuple[str, str]] = {}
+        default_items = (
             [(n.proxy_url or "direct", n.proxy_url or "") for n in nodes]
             if nodes
             else [("direct", "")]
         )
-        # Also refresh bundles for keys that no longer have a matching node
-        # (e.g. pool was reconfigured) so stale entries get cleaned up.
-        all_keys = {a for a, _ in affinity_items} | existing
+        for affinity, proxy_url in default_items:
+            key: BundleKey = (affinity, _clearance_host(_DEFAULT_CLEARANCE_ORIGIN))
+            refresh_targets[key] = (proxy_url, _DEFAULT_CLEARANCE_ORIGIN)
+        for key in existing:
+            affinity, clearance_host = key
+            refresh_targets.setdefault(
+                key,
+                ("" if affinity == "direct" else affinity, f"https://{clearance_host}"),
+            )
 
-        for affinity in all_keys:
-            proxy_url = "" if affinity == "direct" else affinity
+        for key, (proxy_url, clearance_origin) in refresh_targets.items():
+            affinity, clearance_host = key
             if self._clearance_mode == ClearanceMode.MANUAL:
-                new_bundle = self._manual.build_bundle(affinity_key=affinity)
+                new_bundle = self._manual.build_bundle(
+                    affinity_key=affinity,
+                    clearance_host=clearance_host,
+                )
             else:
                 new_bundle = await self._flare.refresh_bundle(
                     affinity_key=affinity,
                     proxy_url=proxy_url,
+                    target_url=clearance_origin,
                 )
             if new_bundle:
                 async with self._lock:
-                    self._bundles[affinity] = new_bundle
-                logger.debug("clearance bundle refreshed: affinity={}", affinity)
+                    self._bundles[key] = new_bundle
+                logger.debug("clearance bundle refreshed: bundle={}", key)
             else:
                 logger.warning(
-                    "clearance refresh failed, keeping old bundle: affinity={}",
-                    affinity,
+                    "clearance refresh failed, keeping old bundle: bundle={}",
+                    key,
                 )
 
     # ------------------------------------------------------------------
@@ -374,7 +407,7 @@ def nodes(self) -> list[EgressNode]:
         return list(self._nodes)
 
     @property
-    def bundles(self) -> dict[str, ClearanceBundle]:
+    def bundles(self) -> dict[BundleKey, ClearanceBundle]:
         """Read-only snapshot of the current clearance bundles."""
         return dict(self._bundles)
 

app/control/proxy/models.py

@@ -1,9 +1,9 @@
 """Control-plane proxy domain models."""
 
 from enum import IntEnum, StrEnum
-from typing import Any, Self
+from typing import Self
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
 
 class ProxyScope(StrEnum):
@@ -76,6 +76,7 @@ class ClearanceBundle(BaseModel):
     user_agent:      str            = ""
     state:           ClearanceBundleState = ClearanceBundleState.VALID
     affinity_key:    str            = ""  # associates bundle with an egress node
+    clearance_host:  str            = "grok.com"
     last_refresh_at: int | None     = None  # ms
 
 
@@ -84,6 +85,7 @@ class ProxyLease(BaseModel):
     proxy_url:   str | None    = None
     cf_cookies:  str           = ""
     user_agent:  str           = ""
+    clearance_host: str        = "grok.com"
     scope:       ProxyScope    = ProxyScope.APP
     kind:        RequestKind   = RequestKind.HTTP
     acquired_at: int           = 0   # ms

app/control/proxy/providers/flaresolverr.py

@@ -2,9 +2,9 @@
 
 import asyncio
 import json
-import re
 from urllib import request as urllib_request
 from urllib.error import HTTPError, URLError
+from urllib.parse import urlparse
 
 from app.platform.logging.logger import logger
 from app.platform.config.snapshot import get_config
@@ -15,18 +15,6 @@ def _extract_all_cookies(cookies: list[dict]) -> str:
     return "; ".join(f"{c.get('name')}={c.get('value')}" for c in cookies)
 
 
-def _extract_cookie_value(cookies: list[dict], name: str) -> str:
-    for c in cookies:
-        if c.get("name") == name:
-            return c.get("value") or ""
-    return ""
-
-
-def _browser_profile(user_agent: str) -> str:
-    m = re.search(r"Chrome/(\d+)", user_agent)
-    return f"chrome{m.group(1)}" if m else "chrome120"
-
-
 class FlareSolverrClearanceProvider:
     """Refresh CF clearance bundles via a FlareSolverr instance."""
 
@@ -35,6 +23,7 @@ async def refresh_bundle(
         *,
         affinity_key: str,
         proxy_url:    str,
+        target_url:   str = "https://grok.com",
     ) -> ClearanceBundle | None:
         cfg = get_config()
         mode = ClearanceMode.parse(cfg.get_str("proxy.clearance.mode", "none"))
@@ -49,19 +38,22 @@ async def refresh_bundle(
             fs_url      = fs_url,
             proxy_url   = proxy_url,
             timeout_sec = timeout_sec,
+            target_url  = target_url,
         )
         if not result:
             logger.warning(
-                "flaresolverr clearance refresh failed: affinity={} proxy={}",
-                affinity_key, proxy_url or "<direct>",
+                "flaresolverr clearance refresh failed: affinity={} proxy={} target={}",
+                affinity_key, proxy_url or "<direct>", target_url,
             )
             return None
+        host = result.get("clearance_host", "grok.com")
 
         return ClearanceBundle(
-            bundle_id    = f"flaresolverr:{affinity_key}",
+            bundle_id    = f"flaresolverr:{affinity_key}@{host}",
             cf_cookies   = result.get("cookies", ""),
             user_agent   = result.get("user_agent", ""),
             affinity_key = affinity_key,
+            clearance_host = host,
         )
 
     async def _solve(
@@ -70,10 +62,12 @@ async def _solve(
         fs_url:      str,
         proxy_url:   str,
         timeout_sec: int,
+        target_url:  str,
     ) -> dict[str, str] | None:
+        target = target_url.strip() or "https://grok.com"
         payload: dict = {
             "cmd":        "request.get",
-            "url":        "https://grok.com",
+            "url":        target,
             "maxTimeout": timeout_sec * 1000,
         }
         if proxy_url:
@@ -107,10 +101,16 @@ def _post() -> dict:
                 return None
 
             ua = solution.get("userAgent", "") or ""
+            host = (urlparse(target).hostname or "").lower()
+            filtered = [
+                cookie for cookie in cookies
+                if not host or not cookie.get("domain") or host.endswith(str(cookie.get("domain", "")).lstrip(".").lower())
+            ]
+            chosen = filtered or cookies
             return {
-                "cookies":    _extract_all_cookies(cookies),
+                "cookies":    _extract_all_cookies(chosen),
                 "user_agent": ua,
-                "browser":    _browser_profile(ua),
+                "clearance_host": host or "grok.com",
             }
 
         except HTTPError as exc:

app/control/proxy/providers/manual.py

@@ -8,17 +8,23 @@
 class ManualClearanceProvider:
     """Build a ClearanceBundle from static config values."""
 
-    def build_bundle(self, *, affinity_key: str) -> ClearanceBundle | None:
+    def build_bundle(
+        self,
+        *,
+        affinity_key: str,
+        clearance_host: str = "grok.com",
+    ) -> ClearanceBundle | None:
         cfg = get_config()
         mode = ClearanceMode.parse(cfg.get_str("proxy.clearance.mode", "none"))
         if mode != ClearanceMode.MANUAL:
             return None
         clearance = resolve_clearance_config(cfg)
         return ClearanceBundle(
-            bundle_id=f"manual:{affinity_key}",
+            bundle_id=f"manual:{affinity_key}@{clearance_host}",
             cf_cookies=clearance.cf_cookies,
             user_agent=clearance.user_agent,
             affinity_key=affinity_key,
+            clearance_host=clearance_host,
         )
 
 

app/dataplane/proxy/__init__.py

@@ -23,8 +23,14 @@ async def acquire(
         scope:    ProxyScope  = ProxyScope.APP,
         kind:     RequestKind = RequestKind.HTTP,
         resource: bool        = False,
+        clearance_origin: str | None = None,
     ) -> ProxyLease:
-        return await self._dir.acquire(scope=scope, kind=kind, resource=resource)
+        return await self._dir.acquire(
+            scope=scope,
+            kind=kind,
+            resource=resource,
+            clearance_origin=clearance_origin,
+        )
 
     async def feedback(self, lease: ProxyLease, result: ProxyFeedback) -> None:
         await self._dir.feedback(lease, result)