Skip to content

A
aisbf
Commit 674031f5 authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

Add "shodan" last-resort discovery/failover provider (private) 

A new special provider, type "shodan", routes OpenAI-compatible requests to a
pool of open model endpoints (Ollama/vLLM/llama.cpp/...) that are auto-discovered
and validated in the background, failing over across them transparently. It is
credential-free by design: client credentials are never forwarded to the
(untrusted) discovered upstreams.

Discovery/validation/live-registry/failover engine is vendored, essentially
verbatim, from Pasquale Minervini's open-router project (@pminervini,
https://github.com/pminervini/open-router, upstream commit 0378a08); full credit
for that engine is his. See aisbf/providers/shodan/open_router_engine.py for
provenance/attribution and docs/shodan-provider.md. AISBF uses only the engine
(not its aiohttp server/CLI layer) via a small adapter (runtime.py, handler.py).

Rotation integration: a new general per-entry flag "last_resort: true" holds an
entry aside and engages it ONLY when every normal provider in the rotation is
unavailable (disabled/rate-limited — the case that otherwise returns HTTP 429)
or has failed during the request's retry loop. The flag works for any provider
type, not just shodan. The special model name "auto" routes to whatever the
discovery pool currently advertises.

- aisbf/providers/shodan/: vendored engine + runtime adapter + handler
- aisbf/providers/__init__.py: register 'shodan'
- aisbf/config.py: ProviderConfig.shodan_config
- aisbf/handlers.py: last_resort partition/engagement + shodan SSE streaming
- requirements.txt: aiohttp>=3.14
- docs/shodan-provider.md: usage, config, security, credit
- tests/providers/test_shodan_provider.py: 21 tests (routing failover,
  model resolution, config build, last_resort partition, singleton)

Full suite: 141 passed (providers + routes + license).
Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
parent b9fe99b1
Show whitespace changes
Inline Side-by-side
Showing with 3484 additions and 2 deletions
aisbf/config.py
View file @ 674031f5
......@@ -125,6 +125,7 @@ class ProviderConfig(BaseModel):
qwen_config: Optional[Dict] = None # Optional Qwen-specific configuration - DEPRECATED
coderai_config: Optional[Dict] = None # Optional CoderAI-specific configuration
runpod_config: Optional[Dict] = None # Optional RunPod-specific configuration
shodan_config: Optional[Dict] = None # Optional "shodan" last-resort discovery configuration
# Default settings for models in this provider
default_rate_limit: Optional[float] = None
default_max_request_tokens: Optional[int] = None
......
aisbf/handlers.py
View file @ 674031f5
......@@ -168,6 +168,20 @@ def _provider_in_availability_cooldown(handler) -> bool:
return bool((failure_until and failure_until > now) or (usage_until and usage_until > now))
def _partition_last_resort(models):
"""Split rotation model entries into (normal, last_resort) by the
``_last_resort`` tag set during the provider scan.
Last-resort entries (any provider marked ``last_resort: true`` in the
rotation, e.g. the discovery-backed 'shodan' provider) are held aside and
engaged only when every normal provider is unavailable or has failed —
preserving each list's original order.
"""
normal = [m for m in models if not m.get('_last_resort')]
last_resort = [m for m in models if m.get('_last_resort')]
return normal, last_resort
# Per-provider client rate limiting — a global sliding window (all clients
# combined) that bounds how many requests are forwarded to a given provider,
# protecting the upstream from hammering (e.g. a client exhausting a per-IP
......@@ -3374,6 +3388,14 @@ class RotationHandler:
logger.info(f"")
logger.info(f"--- Processing provider: {provider_id} ---")
# A "last_resort" entry (any provider type, e.g. the discovery-backed
# 'shodan' provider) is held aside and engaged only when every normal
# provider in this rotation is unavailable or has failed. Its models
# are tagged here and partitioned out of the normal pool after the scan.
is_last_resort = bool(provider.get('last_resort'))
if is_last_resort:
logger.info(f" [LAST RESORT] {provider_id} will be used only if all normal providers are unavailable")
# Meta-provider: this entry forwards to another rotation/autoselect.
# Each model 'name' is the id of the target rotation/autoselect.
delegate_kind = self._delegation_kind(provider_id)
......@@ -3405,6 +3427,7 @@ class RotationHandler:
'api_key': None,
'_delegate_kind': delegate_kind,
'_delegate_target': tname,
'_last_resort': is_last_resort,
})
total_models_considered += 1
logger.info(f" [DELEGATION] -> {delegate_kind} '{tname}' (weight {target_weight})")
......@@ -3521,6 +3544,7 @@ class RotationHandler:
model_with_provider = model.copy()
model_with_provider['provider_id'] = provider_id
model_with_provider['api_key'] = api_key
model_with_provider['_last_resort'] = is_last_resort
available_models.append(model_with_provider)
logger.info(f"")
......@@ -3530,6 +3554,14 @@ class RotationHandler:
if skipped_providers:
logger.info(f"Skipped providers: {', '.join(skipped_providers)}")
logger.info(f"Total models considered: {total_models_considered}")
# Partition last-resort models out of the normal pool. They are engaged
# only when the normal pool is empty (all providers unavailable) or has
# been fully exhausted during the retry loop below.
available_models, last_resort_models = _partition_last_resort(available_models)
last_resort_engaged = False
if last_resort_models:
logger.info(f"Last-resort models held aside: {len(last_resort_models)} (normal pool: {len(available_models)})")
logger.info(f"Total models available: {len(available_models)}")
# Apply NSFW/Privacy content classification filtering
......@@ -3616,6 +3648,15 @@ class RotationHandler:
if rotation_privacy:
logger.info(f"Rotation allows Privacy content - keeping models that support it")
if not available_models and last_resort_models:
logger.warning(
f"All normal providers in rotation '{rotation_id}' are unavailable; "
f"engaging {len(last_resort_models)} last-resort model(s)"
)
available_models = last_resort_models
last_resort_models = []
last_resort_engaged = True
if not available_models:
logger.error("No models available in rotation (all providers may be rate limited)")
logger.error("All providers in this rotation are currently deactivated")
......@@ -3779,6 +3820,20 @@ class RotationHandler:
# Select a model that hasn't been tried yet, or retry a failed model with rate limiting
remaining_models = [m for m in available_models if m not in tried_models]
if not remaining_models:
# Normal pool exhausted — bring in the last-resort models now (once)
# so failover reaches them even when normal providers were initially
# available but all failed at request time.
if last_resort_models and not last_resort_engaged:
logger.warning(
f"All normal providers in rotation '{rotation_id}' have failed; "
f"engaging {len(last_resort_models)} last-resort model(s)"
)
available_models = available_models + last_resort_models
last_resort_models = []
last_resort_engaged = True
remaining_models = [m for m in available_models if m not in tried_models]
if not remaining_models:
logger.error(f"No more models available to try")
logger.error(f"All {len(available_models)} models have been attempted")
......@@ -4746,7 +4801,9 @@ class RotationHandler:
# Check if this is a Google or Kilo provider based on configuration
is_google_provider = provider_type == 'google'
is_kilo_provider = provider_type in ('kilo', 'kilocode')
is_coderai_provider = provider_type == 'coderai'
# shodan relays raw OpenAI-compatible SSE bytes from a discovered upstream,
# exactly like the coderai/kilo pass-through path, so it shares that branch.
is_coderai_provider = provider_type in ('coderai', 'shodan')
logger.info(f"Creating streaming response for provider type: {provider_type}, is_google: {is_google_provider}, is_kilo: {is_kilo_provider}, is_coderai: {is_coderai_provider}")
# Generate system_fingerprint for this request
......
aisbf/providers/__init__.py
View file @ 674031f5
......@@ -44,6 +44,7 @@ from .codex import CodexProviderHandler
from .coderai import CoderAIProviderHandler
from .qwen import QwenProviderHandler
from .runpod import RunpodProviderHandler
from .shodan import ShodanProviderHandler
from .preconfigured_openai import (
# Major inference API providers
GroqProviderHandler,
......@@ -143,6 +144,7 @@ PROVIDER_HANDLERS = {
'coderai': CoderAIProviderHandler,
'qwen': QwenProviderHandler,
'runpod': RunpodProviderHandler,
'shodan': ShodanProviderHandler,
# --- pre-configured OpenAI-compatible providers ---
'groq': GroqProviderHandler,
......
aisbf/providers/shodan/__init__.py 0 → 100644
View file @ 674031f5
"""
Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
AISBF - AI Service Broker Framework || AI Should Be Free
"shodan" last-resort provider package.
The discovery/validation/failover engine in ``open_router_engine.py`` is
vendored from Pasquale Minervini's open-router project
(https://github.com/pminervini/open-router); full credit for that engine is his.
The AISBF-specific adapter (runtime.py, handler.py) wraps it as a provider.
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
from .handler import ShodanProviderHandler
__all__ = ["ShodanProviderHandler"]
aisbf/providers/shodan/handler.py 0 → 100644
View file @ 674031f5
"""
Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
AISBF - AI Service Broker Framework || AI Should Be Free
"shodan" provider handler — a credential-free, last-resort failover provider.
It routes OpenAI-compatible requests to a pool of open model endpoints that are
discovered and validated by the vendored open-router engine (see
open_router_engine.py, by Pasquale Minervini / @pminervini). It is intended to
be added to a rotation as a ``last_resort`` entry: the rotation engages it only
when every normal provider is disabled or rate limited.
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
import logging
from typing import Dict, List, Optional, Union
from ..base import BaseProviderHandler
from ...config import config
from ...models import Model
from .runtime import get_runtime
logger = logging.getLogger(__name__)
class ShodanProviderHandler(BaseProviderHandler):
"""Last-resort provider backed by auto-discovered open model endpoints.
Credential-free by design: the underlying engine deliberately never forwards
client credentials to the (untrusted) discovered upstreams, and the provider
itself needs no API key.
"""
def __init__(
self,
provider_id: str,
api_key: Optional[str] = None,
user_id: Optional[int] = None,
provider_config=None,
):
self.provider_config = (
provider_config if provider_config is not None else config.providers.get(provider_id)
)
super().__init__(provider_id, api_key, user_id=user_id)
def _shodan_config(self) -> Dict:
pc = self.provider_config
if pc is None:
return {}
if isinstance(pc, dict):
return pc.get("shodan_config") or {}
return getattr(pc, "shodan_config", None) or {}
def validate_credentials(self) -> bool:
# No credentials are used or required — the provider is always "valid".
return True
async def handle_request(
self,
model: str,
messages: List[Dict],
max_tokens: Optional[int] = None,
temperature: Optional[float] = 1.0,
stream: Optional[bool] = False,
tools: Optional[List[Dict]] = None,
tool_choice: Optional[Union[str, Dict]] = None,
) -> Union[Dict, object]:
if self.is_rate_limited():
raise Exception("Provider rate limited")
await self.apply_rate_limit()
body: Dict = {
"model": model,
"messages": messages,
"stream": bool(stream),
}
if max_tokens is not None:
body["max_tokens"] = max_tokens
if temperature is not None:
body["temperature"] = temperature
if tools is not None:
body["tools"] = tools
if tool_choice is not None:
body["tool_choice"] = tool_choice
try:
runtime = await get_runtime(self._shodan_config())
result = await runtime.route("chat/completions", body, bool(stream))
except Exception as e:
logger.error("ShodanProviderHandler: request failed: %s", e)
self.record_failure()
raise
# Streaming returns a lazy async generator; the rotation primes and
# consumes it and records success itself (mirrors the OpenAI handler).
if not stream:
self.record_success()
return result
async def get_models(self) -> List[Model]:
try:
runtime = await get_runtime(self._shodan_config())
model_ids = await runtime.available_models()
except Exception as e:
logger.error("ShodanProviderHandler: get_models failed: %s", e)
raise
return [
Model(id=mid, name=mid, provider_id=self.provider_id)
for mid in model_ids
]
aisbf/providers/shodan/open_router_engine.py 0 → 100755
View file @ 674031f5
#!/usr/bin/env python3
# ---------------------------------------------------------------------------
# VENDORED THIRD-PARTY CODE — DO NOT REWRITE IN AISBF STYLE.
#
# This module is vendored, essentially verbatim, from the "open-router" project
# by Pasquale Minervini (@pminervini):
#
# https://github.com/pminervini/open-router
# upstream commit 0378a0811050640910fd771897d010e17abdccc4
#
# Full credit for the discovery, validation, live-registry, and failover-routing
# engine below belongs to Pasquale Minervini. AISBF uses it as the engine behind
# its "shodan" last-resort provider (see aisbf/providers/shodan/handler.py). The
# aiohttp *server* layer (RouterHttpApp, the CLI, JSONL request logging) is
# retained but intentionally never instantiated by AISBF; only the engine pieces
# — DiscoveryConfig, ProbeConfig, endpoint discovery/probing, EndpointRegistry,
# and RouterRuntime — are imported by the AISBF adapter.
#
# Keep this file in sync with upstream rather than editing it locally, so credit
# and provenance stay clear. The upstream repository carries no LICENSE file; it
# is included here with attribution and is not relicensed under AISBF's GPL.
# ---------------------------------------------------------------------------
"""Small OpenAI-compatible failover router.
The service discovers and validates credential-free model endpoints, maintains
an in-memory live registry, and retries generation requests across compatible
upstreams. It exposes model and endpoint inventories, Chat Completions,
Responses, and unauthenticated health checks. Optional client bearer tokens are
never forwarded upstream.
See README.md for deployment, discovery, routing, security, and API details.
"""
import argparse
import asyncio
import gzip
import hmac
import ipaddress
import json
import logging
import math
import os
import re
import sys
import time
import uuid
from dataclasses import dataclass, replace as dataclass_replace
from datetime import UTC, datetime
from email.utils import parsedate_to_datetime
from html.parser import HTMLParser
from http import HTTPStatus
from pathlib import Path
from typing import Any, Awaitable, Callable, Iterable, Mapping, Sequence
from urllib.parse import urlencode
import aiohttp
from aiohttp import web
from aiohttp.http_exceptions import HttpProcessingError
from multidict import CIMultiDict
VERSION = "2.33.0"
DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 8000
DEFAULT_HEALTH_INTERVAL = 86_400.0
DEFAULT_DISCOVERY_INTERVAL = 86_400.0
DEFAULT_PROBE_TIMEOUT = 5.0
DEFAULT_READINESS_TIMEOUT = 30.0
DEFAULT_READINESS_MODEL_ATTEMPTS = 3
DEFAULT_UPSTREAM_TIMEOUT = 5 * 60.0
DEFAULT_REQUEST_HEADER_TIMEOUT = 10.0
DEFAULT_REQUEST_BODY_TIMEOUT = 30.0
DEFAULT_KEEPALIVE_TIMEOUT = 30.0
GENERATION_LATENCY_EMA_ALPHA = 0.25
DEFAULT_RETRIEVAL_TIMEOUT = 30.0
DEFAULT_RETRIEVAL_RETRIES = 3
DEFAULT_WORKERS = 8
DEFAULT_PROGRESS_EVERY = 1024
DEFAULT_BLACKLIST_DURATION = 24 * 60 * 60.0
DEFAULT_LOG_LEVEL = "INFO"
LOG_LEVEL_NAMES = ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL")
MAX_MODELS_RESPONSE_BYTES = 8 * 1024 * 1024
COPY_BUFFER_BYTES = 64 * 1024
SHODAN_PORT = 11434
SHODAN_QUERY = f"port:{SHODAN_PORT}"
# Public IP facets do not identify which matching port produced an address, so
# keep each port/query pair separate and attach its port after parsing.
SHODAN_PUBLIC_TARGETS: tuple[tuple[int, str], ...] = (
(11434, 'port:11434 product:"Ollama"'),
(11434, 'port:11434 "Ollama"'),
(8000, 'port:8000 "vLLM"'),
(8000, 'port:8000 "llama.cpp"'),
(8080, 'port:8080 "llama.cpp"'),
(1234, "port:1234 lmstudio"),
(4891, "port:4891"),
(8000, 'port:8000 "LangChain"'),
(8000, 'port:8000 "langserve"'),
)
SHODAN_API_URL = "https://api.shodan.io/shodan/host/search"
SHODAN_PUBLIC_FACET_URL = "https://www.shodan.io/search/facet"
SHODAN_RETRIABLE_HTTP_CODES = {429, 500, 502, 503, 504}
SHODAN_USER_AGENT = (
"Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
"(KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
)
NON_GENERATION_MODEL_MARKERS = (
"embed",
"embedding",
"bge-",
"nomic-",
"snowflake-arctic",
"flux",
"rce_test",
"leaktest",
"pwn",
)
ROUTES = {
"/responses": "responses",
"/v1/responses": "responses",
"/chat/completions": "chat/completions",
"/v1/chat/completions": "chat/completions",
}
MODEL_ROUTES = {"/models", "/v1/models"}
ENDPOINT_ROUTES = {"/endpoints"}
REQUEST_ID_KEY = web.RequestKey("request_id", str)
HOP_BY_HOP_HEADERS = {
"connection",
"keep-alive",
"proxy-authenticate",
"proxy-authorization",
"te",
"trailer",
"transfer-encoding",
"upgrade",
}
SENSITIVE_HEADERS = {
"api-key",
"authorization",
"cookie",
"cookie2",
"openai-organization",
"openai-project",
"proxy-authorization",
"set-cookie",
"x-api-key",
"x-access-token",
"x-auth-key",
"x-auth-token",
"x-goog-api-key",
}
UPSTREAM_FAILURE_STATUSES = {
HTTPStatus.NOT_FOUND,
HTTPStatus.METHOD_NOT_ALLOWED,
HTTPStatus.REQUEST_TIMEOUT,
HTTPStatus.TOO_EARLY,
HTTPStatus.TOO_MANY_REQUESTS,
}
ENDPOINT_FAILURE_STATUSES = {
HTTPStatus.UNAUTHORIZED,
}
MODEL_FAILURE_STATUSES = {
HTTPStatus.FORBIDDEN,
HTTPStatus.GONE,
}
MODEL_LOADING_ERROR_MARKERS = (
"llm server loading model",
"model is loading",
)
MODEL_UNAVAILABLE_ERROR_MARKERS = (
"requires more system memory",
"unable to load model",
)
@dataclass(frozen=True)
class DiscoveryConfig:
shodan_api_key: str | None = None
shodan_mode: str = "auto"
retrieval_timeout: float = DEFAULT_RETRIEVAL_TIMEOUT
retrieval_retries: int = DEFAULT_RETRIEVAL_RETRIES
shodan_max_results: int | None = None
masscan_list_paths: tuple[Path, ...] = ()
interval: float = DEFAULT_DISCOVERY_INTERVAL
blacklist_duration: float = DEFAULT_BLACKLIST_DURATION
@dataclass(frozen=True)
class ProbeConfig:
workers: int = DEFAULT_WORKERS
timeout: float = DEFAULT_PROBE_TIMEOUT
readiness_timeout: float = DEFAULT_READINESS_TIMEOUT
readiness_model_attempts: int = DEFAULT_READINESS_MODEL_ATTEMPTS
health_interval: float = DEFAULT_HEALTH_INTERVAL
progress_every: int = DEFAULT_PROGRESS_EVERY
@dataclass(frozen=True)
class ServerOptions:
host: str = DEFAULT_HOST
port: int = DEFAULT_PORT
api_keys: tuple[str, ...] = ()
log_file: Path | None = None
log_requests: Path | None = None
upstream_timeout: float = DEFAULT_UPSTREAM_TIMEOUT
log_level: str = DEFAULT_LOG_LEVEL
quiet: bool = False
@dataclass(frozen=True)
class RouterConfig:
discovery: DiscoveryConfig
probe: ProbeConfig
server: ServerOptions
def _is_sensitive_header(name: str) -> bool:
lower = name.lower()
return (
lower in SENSITIVE_HEADERS
or lower.endswith("-api-key")
or lower.endswith("-auth-token")
or lower.endswith("-authorization")
)
def _is_upstream_failure_status(status: int) -> bool:
return status in UPSTREAM_FAILURE_STATUSES or 500 <= status <= 599
def _upstream_error_message(raw: bytes) -> str | None:
try:
payload = json.loads(raw)
except (UnicodeDecodeError, json.JSONDecodeError):
return None
if not isinstance(payload, Mapping):
return None
error = payload.get("error", payload.get("message"))
if isinstance(error, Mapping):
error = error.get("message")
return error.lower() if isinstance(error, str) else None
def _is_model_loading_error(status: int, raw: bytes) -> bool:
if status not in {HTTPStatus.INTERNAL_SERVER_ERROR, HTTPStatus.SERVICE_UNAVAILABLE}:
return False
message = _upstream_error_message(raw)
return message is not None and any(
marker in message for marker in MODEL_LOADING_ERROR_MARKERS
)
def _is_model_unavailable_error(status: int, raw: bytes) -> bool:
if status != HTTPStatus.INTERNAL_SERVER_ERROR:
return False
message = _upstream_error_message(raw)
return message is not None and any(
marker in message for marker in MODEL_UNAVAILABLE_ERROR_MARKERS
)
class ConfigurationError(ValueError):
"""Endpoint discovery or CLI configuration is invalid."""
class RequestError(Exception):
"""An incoming request cannot be routed."""
def __init__(
self,
status: HTTPStatus,
message: str,
*,
code: str,
param: str | None = None,
) -> None:
super().__init__(message)
self.status = status
self.message = message
self.code = code
self.param = param
class RetrievalError(RuntimeError):
"""A Shodan request or response could not be handled."""
def _concise_http_parser_error(record: logging.LogRecord) -> bool:
"""Replace expected malformed-request tracebacks with one warning."""
if record.exc_info is None or not isinstance(
record.exc_info[1],
HttpProcessingError,
):
return True
error = record.exc_info[1]
prefix = "Error handling request from "
original_message = record.getMessage()
remote = (
original_message.removeprefix(prefix)
if original_message.startswith(prefix)
else "unknown client"
)
detail = " ".join(error.message.split()).rstrip(".")
record.msg = (
f"Rejected malformed HTTP request from {remote}: "
f"HTTP {error.code}, {detail}"
)
record.args = ()
record.levelno = logging.WARNING
record.levelname = logging.getLevelName(logging.WARNING)
record.exc_info = None
record.exc_text = None
record.stack_info = None
return True
def configure_logger(
level_name: str,
*,
quiet: bool,
log_file: Path | None = None,
) -> logging.Logger:
"""Configure operational stderr logging separately from request auditing."""
level = logging.getLevelNamesMapping().get(level_name.upper())
if not isinstance(level, int):
raise ConfigurationError(f"unsupported log level: {level_name!r}")
if quiet:
level = max(level, logging.WARNING)
logger = logging.getLogger("open_router")
logger.setLevel(level)
logger.propagate = False
logger.filters.clear()
logger.addFilter(_concise_http_parser_error)
for existing_handler in list(logger.handlers):
logger.removeHandler(existing_handler)
existing_handler.close()
handler = logging.StreamHandler(sys.stderr)
handler.setLevel(level)
formatter = logging.Formatter(
"{asctime}Z {levelname} {name}: {message}",
datefmt="%Y-%m-%dT%H:%M:%S",
style="{",
)
formatter.converter = time.gmtime
handler.setFormatter(formatter)
logger.addHandler(handler)
if log_file is not None:
path = log_file.expanduser().resolve()
path.parent.mkdir(parents=True, exist_ok=True)
file_handler = logging.FileHandler(path, encoding="utf-8")
file_handler.setLevel(level)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
return logger
class _HeaderTimeoutAppRunner(web.AppRunner):
"""Close connections that never complete their first HTTP headers."""
def __init__(
self,
app: web.Application,
*,
header_timeout: float = DEFAULT_REQUEST_HEADER_TIMEOUT,
**kwargs: Any,
) -> None:
self.header_timeout = header_timeout
kwargs.setdefault("handler_cancellation", True)
kwargs.setdefault("keepalive_timeout", DEFAULT_KEEPALIVE_TIMEOUT)
super().__init__(app, **kwargs)
async def _make_server(self) -> Any:
server = await super()._make_server()
connection_made = server.connection_made
loop = asyncio.get_running_loop()
def connection_made_with_timeout(
handler: Any,
transport: asyncio.Transport,
) -> None:
connection_made(handler, transport)
loop.call_later(
self.header_timeout,
self._close_incomplete_headers,
handler,
transport,
)
setattr(server, "connection_made", connection_made_with_timeout)
return server
@staticmethod
def _close_incomplete_headers(
handler: Any,
transport: asyncio.Transport,
) -> None:
if not transport.is_closing() and handler._request_count == 0:
transport.close()
@dataclass(frozen=True)
class ShodanService:
address: ipaddress.IPv4Address | ipaddress.IPv6Address
port: int
transport: str = ""
class ShodanFacetParser(HTMLParser):
"""Extract facet IPs and the advertised total from Shodan's HTML."""
def __init__(self) -> None:
super().__init__(convert_charrefs=True)
self._inside_strong = False
self._text: list[str] = []
self.addresses: set[ipaddress.IPv4Address | ipaddress.IPv6Address] = set()
self.no_information_available: bool = False
def handle_starttag(self, tag: str, _attrs: list[tuple[str, str | None]]) -> None:
if tag.lower() == "strong":
self._inside_strong = True
self._text = []
def handle_data(self, data: str) -> None:
if "No information available" in data:
self.no_information_available = True
if self._inside_strong:
self._text.append(data)
def handle_endtag(self, tag: str) -> None:
if tag.lower() != "strong" or not self._inside_strong:
return
candidate = "".join(self._text).strip()
try:
self.addresses.add(ipaddress.ip_address(candidate))
except ValueError:
pass
self._inside_strong = False
self._text = []
def _shodan_error_message(status: int, reason: str, raw: bytes) -> str:
try:
body = raw.decode("utf-8", errors="replace")
payload = json.loads(body)
if isinstance(payload, dict) and payload.get("error"):
return str(payload["error"])
if body.strip():
return body.strip()[:500]
except (UnicodeError, json.JSONDecodeError):
pass
return reason or f"HTTP {status}"
def _retry_delay(attempt: int, retry_after: str | None) -> float:
delay = float(2**attempt)
if not retry_after:
return delay
try:
requested = float(retry_after)
except ValueError:
try:
requested = (
parsedate_to_datetime(retry_after) - datetime.now(UTC)
).total_seconds()
except (TypeError, ValueError):
return delay
return max(delay, requested, 0.0)
async def _fetch_shodan(
session: aiohttp.ClientSession,
url: str,
*,
timeout: float,
retries: int,
) -> bytes:
headers = {
"Accept": "application/json,text/html;q=0.9,*/*;q=0.8",
"Accept-Encoding": "identity",
"Accept-Language": "en-US,en;q=0.9",
"User-Agent": SHODAN_USER_AGENT,
}
last_error = "Unable to reach Shodan"
attempt = 0
while True:
retriable = True
retry_after: str | None = None
try:
async with session.get(
url,
headers=headers,
timeout=aiohttp.ClientTimeout(total=timeout),
allow_redirects=False,
) as response:
raw = await response.read()
if response.status < 400:
return raw
message = _shodan_error_message(
response.status,
response.reason or "",
raw,
)
last_error = f"Shodan returned HTTP {response.status}: {message}"
retriable = response.status in SHODAN_RETRIABLE_HTTP_CODES
retry_after = response.headers.get("Retry-After")
except asyncio.TimeoutError:
last_error = "The Shodan request timed out"
except aiohttp.ClientError as error:
last_error = f"Unable to reach Shodan: {error}"
if not retriable or attempt >= retries:
raise RetrievalError(last_error)
await asyncio.sleep(_retry_delay(attempt, retry_after))
attempt += 1
def _parse_shodan_api_response(
raw: bytes, page: int
) -> tuple[list[dict[str, object]], int]:
try:
payload = json.loads(raw)
except (UnicodeDecodeError, json.JSONDecodeError) as error:
raise RetrievalError(f"Shodan API page {page} returned invalid JSON") from error
if not isinstance(payload, dict):
raise RetrievalError(f"Shodan API page {page} returned an unexpected response")
if payload.get("error"):
raise RetrievalError(f"Shodan API error: {payload['error']}")
matches = payload.get("matches")
total = payload.get("total")
if not isinstance(matches, list) or not isinstance(total, int) or total < 0:
raise RetrievalError(f"Shodan API page {page} omitted matches or total")
if not all(isinstance(match, dict) for match in matches):
raise RetrievalError(f"Shodan API page {page} contains a malformed match")
return matches, total
def _shodan_service_from_match(
match: dict[str, object],
) -> ShodanService | None:
raw_address = match.get("ip_str", match.get("ip"))
raw_port = match.get("port")
if not isinstance(raw_address, (str, int)) or not isinstance(raw_port, (str, int)):
return None
try:
address = ipaddress.ip_address(raw_address)
port = int(raw_port)
except (TypeError, ValueError):
return None
if not 1 <= port <= 65_535:
return None
raw_transport = match.get("transport", "")
transport = raw_transport if isinstance(raw_transport, str) else ""
return ShodanService(address, port, transport)
async def _retrieve_shodan_api(
session: aiohttp.ClientSession,
api_key: str,
*,
timeout: float,
retries: int,
max_results: int | None,
logger: logging.Logger,
) -> set[ShodanService]:
services: set[ShodanService] = set()
page = 1
expected_matches: int | None = None
pages = 1
matches_seen = 0
while page <= pages:
parameters = urlencode(
{
"key": api_key,
"query": SHODAN_QUERY,
"page": page,
"minify": "true",
"fields": "ip_str,port,transport",
}
)
matches, total = _parse_shodan_api_response(
await _fetch_shodan(
session,
f"{SHODAN_API_URL}?{parameters}",
timeout=timeout,
retries=retries,
),
page,
)
if expected_matches is None:
expected_matches = total
if max_results is not None:
expected_matches = min(expected_matches, max_results)
pages = max(1, math.ceil(expected_matches / 100))
remaining = expected_matches - matches_seen
selected_matches = matches[:remaining]
matches_seen += len(selected_matches)
for match in selected_matches:
service = _shodan_service_from_match(match)
if service is not None:
services.add(service)
if page == 1 or page == pages or page % 10 == 0:
logger.info(
f"Retrieved API page {page}/{pages} "
f"({matches_seen}/{expected_matches} matches)"
)
if matches_seen >= expected_matches:
break
if not matches:
raise RetrievalError(
f"Shodan returned no matches on page {page} before all results "
"were retrieved"
)
page += 1
if expected_matches is not None and matches_seen < expected_matches:
raise RetrievalError(
f"Shodan returned only {matches_seen} of "
f"{expected_matches} expected matches"
)
return services
async def _retrieve_shodan_public_facet(
session: aiohttp.ClientSession,
*,
timeout: float,
retries: int,
logger: logging.Logger,
) -> set[ShodanService]:
services: set[ShodanService] = set()
for port, query in SHODAN_PUBLIC_TARGETS:
parameters = urlencode({"query": query, "facet": "ip"})
try:
raw = await _fetch_shodan(
session,
f"{SHODAN_PUBLIC_FACET_URL}?{parameters}",
timeout=timeout,
retries=retries,
)
document = raw.decode("utf-8")
except RetrievalError as error:
logger.warning(f"Public Shodan query {query!r} failed: {error}")
continue
except UnicodeDecodeError:
logger.warning(f"Public Shodan query {query!r} returned invalid UTF-8")
continue
parser = ShodanFacetParser()
parser.feed(document)
if not parser.addresses:
if parser.no_information_available:
logger.debug(f"Public Shodan query {query!r} returned no matches")
continue
logger.warning(
f"Public Shodan query {query!r} returned no IP facets; the page "
"may have changed or been blocked"
)
continue
services.update(ShodanService(address, port) for address in parser.addresses)
logger.debug(
f"Public Shodan query {query!r} returned "
f"{len(parser.addresses)} IP facets"
)
if not services:
raise RetrievalError("No public Shodan query returned usable IP facets")
logger.info(
f"Retrieved {len(services)} unique public IP/port facets from "
f"{len(SHODAN_PUBLIC_TARGETS)} targeted queries."
)
return services
@dataclass(frozen=True)
class EndpointConfig:
"""Configuration for one credential-free OpenAI-compatible upstream."""
base_url: str
def api_url(self, resource: str) -> str:
return f"{self.base_url}/v1/{resource}"
@dataclass(frozen=True)
class EndpointState:
"""Last successful probe result for an upstream."""
config: EndpointConfig
models: frozenset[str]
latency_seconds: float
checked_at: float
@dataclass(frozen=True)
class BufferedUpstreamResponse:
status: int
reason: str
headers: tuple[tuple[str, str], ...]
body: bytes
endpoint: EndpointConfig
class JsonlLogger:
"""Event-loop-owned, append-only request logger."""
def __init__(self, path: Path) -> None:
self.path = path.expanduser().resolve()
self.path.parent.mkdir(parents=True, exist_ok=True)
self._handle = self.path.open("a", encoding="utf-8", buffering=1)
def append(self, record: Mapping[str, Any]) -> None:
self._handle.write(
json.dumps(
record,
separators=(",", ":"),
allow_nan=False,
)
+ "\n"
)
self._handle.flush()
def close(self) -> None:
self._handle.close()
def _ordered_masscan_list_paths(
list_paths: Sequence[Path] | None,
) -> list[Path]:
candidates = list(list_paths or ())
paths_with_mtime: dict[Path, int] = {}
for candidate in candidates:
path = candidate.expanduser().resolve()
if path.suffix.lower() != ".list":
raise ConfigurationError(f"masscan input must be a .list file: {path}")
try:
stat = path.stat()
except OSError as error:
raise ConfigurationError(
f"unable to inspect masscan list {path}: {error}"
) from error
if not path.is_file():
raise ConfigurationError(f"masscan input is not a regular file: {path}")
paths_with_mtime[path] = stat.st_mtime_ns
return [
path
for path, _ in sorted(
paths_with_mtime.items(),
key=lambda item: (-item[1], str(item[0])),
)
]
def _iter_file_lines_reversed(path: Path) -> Iterable[str]:
"""Yield UTF-8 lines from a file's end without loading it all into memory."""
with path.open("rb") as handle:
handle.seek(0, os.SEEK_END)
position = handle.tell()
file_size = position
remainder = b""
at_end = True
while position > 0:
chunk_size = min(COPY_BUFFER_BYTES, position)
position -= chunk_size
handle.seek(position)
block = handle.read(chunk_size) + remainder
lines = block.split(b"\n")
remainder = lines[0]
complete_lines = lines[1:]
if at_end and block.endswith(b"\n"):
complete_lines = complete_lines[:-1]
at_end = False
for raw_line in reversed(complete_lines):
yield raw_line.rstrip(b"\r").decode("utf-8")
if file_size:
yield remainder.rstrip(b"\r").decode("utf-8")
def _parse_masscan_line(line: str) -> EndpointConfig | None:
stripped = line.strip()
if not stripped or stripped.startswith("#"):
return None
fields = stripped.split()
if len(fields) < 4:
return None
state, transport, raw_port, raw_address = fields[:4]
try:
port = int(raw_port)
address = ipaddress.ip_address(raw_address)
except (TypeError, ValueError):
return None
if state.lower() != "open" or transport.lower() != "tcp" or port != SHODAN_PORT:
return None
formatted_address = f"[{address}]" if address.version == 6 else str(address)
return EndpointConfig(base_url=f"http://{formatted_address}:{port}")
def load_masscan_endpoint_configs(
list_paths: Sequence[Path] | None = None,
*,
logger: logging.Logger,
) -> list[EndpointConfig]:
"""Load TCP/11434 results newest-file-first and newest-entry-first."""
paths = _ordered_masscan_list_paths(list_paths)
if not paths:
return []
configs: dict[str, EndpointConfig] = {}
for path in paths:
try:
for line in _iter_file_lines_reversed(path):
config = _parse_masscan_line(line)
if config is not None:
configs.setdefault(config.base_url, config)
except (OSError, UnicodeError) as error:
raise ConfigurationError(
f"unable to read masscan list {path}: {error}"
) from error
logger.info(
f"Loaded {len(configs)} unique endpoints from {len(paths)} masscan list "
"file(s), newest files and entries first."
)
return list(configs.values())
def _resolve_shodan_mode(settings: DiscoveryConfig) -> str:
if settings.shodan_mode not in {"auto", "api", "public"}:
raise ConfigurationError(
f"unsupported Shodan retrieval mode: {settings.shodan_mode!r}"
)
selected_mode = (
"api"
if settings.shodan_mode == "auto" and settings.shodan_api_key
else settings.shodan_mode
)
if selected_mode == "auto":
selected_mode = "public"
if selected_mode == "api" and not settings.shodan_api_key:
raise ConfigurationError(
"Shodan API mode requires --shodan-api-key or SHODAN_API_KEY"
)
if selected_mode == "public" and settings.shodan_max_results is not None:
raise ConfigurationError(
"--shodan-max-results is only supported in Shodan API mode"
)
return selected_mode
async def retrieve_shodan_configs(
session: aiohttp.ClientSession,
settings: DiscoveryConfig,
*,
logger: logging.Logger,
) -> list[EndpointConfig]:
"""Retrieve, normalize, filter, and order Shodan endpoint candidates."""
if _resolve_shodan_mode(settings) == "api":
api_key = settings.shodan_api_key
if api_key is None: # Guarded by _resolve_shodan_mode.
raise AssertionError("Shodan API mode was selected without an API key")
services = await _retrieve_shodan_api(
session,
api_key,
timeout=settings.retrieval_timeout,
retries=settings.retrieval_retries,
max_results=settings.shodan_max_results,
logger=logger,
)
else:
services = await _retrieve_shodan_public_facet(
session,
timeout=settings.retrieval_timeout,
retries=settings.retrieval_retries,
logger=logger,
)
configs: dict[str, EndpointConfig] = {}
for service in sorted(
services,
key=lambda item: (
item.address.version,
int(item.address),
item.port,
item.transport,
),
):
if service.transport not in {"", "tcp"}:
continue
address = str(service.address)
if service.address.version == 6:
address = f"[{address}]"
base_url = f"http://{address}:{service.port}"
configs.setdefault(base_url, EndpointConfig(base_url=base_url))
return list(configs.values())
def merge_prioritized_configs(
shodan_configs: Sequence[EndpointConfig],
masscan_configs: Sequence[EndpointConfig],
*,
logger: logging.Logger,
) -> list[EndpointConfig]:
"""Deduplicate candidates while retaining Shodan and freshness priority."""
configs: dict[str, EndpointConfig] = {}
for config in shodan_configs:
configs.setdefault(config.base_url, config)
shodan_count = len(configs)
for config in masscan_configs:
configs.setdefault(config.base_url, config)
if not configs:
raise ConfigurationError("endpoint retrieval returned no TCP endpoints")
if masscan_configs:
logger.info(
f"Validation priority: {shodan_count} Shodan endpoint(s), then "
f"{len(configs) - shodan_count} masscan-only endpoint(s)."
)
return list(configs.values())
async def retrieve_endpoint_configs(
session: aiohttp.ClientSession,
settings: DiscoveryConfig,
*,
logger: logging.Logger,
) -> list[EndpointConfig]:
"""Orchestrate Shodan retrieval, masscan loading, and fallback behavior."""
shodan_error: RetrievalError | None
try:
shodan_configs = await retrieve_shodan_configs(
session,
settings,
logger=logger,
)
except RetrievalError as error:
shodan_configs = []
shodan_error = error
else:
shodan_error = None
masscan_configs = await asyncio.to_thread(
load_masscan_endpoint_configs,
settings.masscan_list_paths,
logger=logger,
)
if shodan_error is not None:
if not masscan_configs:
raise ConfigurationError(
f"unable to retrieve Shodan endpoints: {shodan_error}"
) from shodan_error
logger.warning(
f"Shodan retrieval failed; continuing with masscan lists: "
f"{shodan_error}"
)
return merge_prioritized_configs(
shodan_configs,
masscan_configs,
logger=logger,
)
def _upstream_headers() -> dict[str, str]:
return {
"Accept": "application/json",
"Accept-Encoding": "identity",
"User-Agent": f"open-router-cli/{VERSION}",
}
def _readiness_models(models: Iterable[str], limit: int) -> list[str]:
candidates = [
model
for model in models
if not any(marker in model.lower() for marker in NON_GENERATION_MODEL_MARKERS)
]
def score(model: str) -> tuple[float, int, str]:
lower = model.lower()
sizes = [float(value) for value in re.findall(r"(\d+(?:\.\d+)?)b\b", lower)]
estimated_size = min(sizes) if sizes else 50.0
remote_penalty = 1 if "cloud" in lower else 0
return estimated_size, remote_penalty, lower
return sorted(candidates, key=score)[:limit]
async def _probe_generation_request(
session: aiohttp.ClientSession,
config: EndpointConfig,
model: str,
resource: str,
*,
timeout: float,
) -> bool:
if resource == "chat/completions":
payload = {
"model": model,
"messages": [{"role": "user", "content": "Reply with OK."}],
"temperature": 0,
"max_tokens": 1,
"stream": False,
}
expected_object = "chat.completion"
expected_array = "choices"
else:
payload = {
"model": model,
"input": "Reply with OK.",
"temperature": 0,
"max_output_tokens": 1,
"stream": False,
}
expected_object = "response"
expected_array = "output"
body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
headers = _upstream_headers()
headers["Content-Type"] = "application/json"
deadline = time.monotonic() + timeout
loading_attempt = 0
while True:
remaining = deadline - time.monotonic()
if remaining <= 0:
return False
try:
async with session.post(
config.api_url(resource),
data=body,
headers=headers,
timeout=aiohttp.ClientTimeout(total=remaining),
allow_redirects=False,
) as response:
raw = await response.content.read(MAX_MODELS_RESPONSE_BYTES + 1)
retry_after = response.headers.get("Retry-After")
status = response.status
except (aiohttp.ClientError, asyncio.TimeoutError, ValueError):
return False
if _is_model_loading_error(status, raw):
delay = min(
_retry_delay(loading_attempt, retry_after),
max(deadline - time.monotonic(), 0),
)
loading_attempt += 1
await asyncio.sleep(delay)
continue
if status != HTTPStatus.OK or len(raw) > MAX_MODELS_RESPONSE_BYTES:
return False
try:
response_payload = json.loads(raw)
except (UnicodeDecodeError, json.JSONDecodeError):
return False
return (
isinstance(response_payload, Mapping)
and response_payload.get("object") == expected_object
and isinstance(response_payload.get(expected_array), list)
)
async def _probe_generation_readiness(
session: aiohttp.ClientSession,
config: EndpointConfig,
models: Iterable[str],
*,
timeout: float,
model_attempts: int,
) -> bool:
for model in _readiness_models(models, model_attempts):
chat_ready = await _probe_generation_request(
session,
config,
model,
"chat/completions",
timeout=timeout,
)
if not chat_ready:
continue
if await _probe_generation_request(
session,
config,
model,
"responses",
timeout=timeout,
):
return True
return False
async def probe_endpoint(
session: aiohttp.ClientSession,
config: EndpointConfig,
timeout: float,
*,
readiness_timeout: float | None = None,
readiness_model_attempts: int = DEFAULT_READINESS_MODEL_ATTEMPTS,
) -> EndpointState | None:
"""Validate model discovery and optionally a minimal generation request."""
started = time.monotonic()
try:
async with session.get(
config.api_url("models"),
headers=_upstream_headers(),
timeout=aiohttp.ClientTimeout(total=timeout),
allow_redirects=False,
) as response:
raw = await response.content.read(MAX_MODELS_RESPONSE_BYTES + 1)
if response.status != HTTPStatus.OK:
return None
if len(raw) > MAX_MODELS_RESPONSE_BYTES:
return None
try:
payload = json.loads(raw)
except (UnicodeDecodeError, json.JSONDecodeError):
return None
if not isinstance(payload, Mapping) or payload.get("object") != "list":
return None
data = payload.get("data")
if not isinstance(data, list):
return None
models = frozenset(
item["id"]
for item in data
if isinstance(item, Mapping)
and isinstance(item.get("id"), str)
and item["id"]
and not any(
marker in item["id"].lower()
for marker in NON_GENERATION_MODEL_MARKERS
)
)
if not models:
return None
if readiness_timeout is not None and not await _probe_generation_readiness(
session,
config,
models,
timeout=readiness_timeout,
model_attempts=readiness_model_attempts,
):
return None
return EndpointState(
config=config,
models=models,
latency_seconds=time.monotonic() - started,
checked_at=time.time(),
)
except (aiohttp.ClientError, asyncio.TimeoutError, ValueError):
return None
def _model_stats(states: Iterable[EndpointState]) -> tuple[int, int]:
"""Return unique-model and endpoint/model-assignment counts."""
snapshot = tuple(states)
models = {model for state in snapshot for model in state.models}
return len(models), sum(len(state.models) for state in snapshot)
async def probe_endpoints(
session: aiohttp.ClientSession,
configs: Sequence[EndpointConfig],
settings: ProbeConfig,
logger: logging.Logger,
*,
check_readiness: bool = False,
progress_label: str | None = None,
state_callback: Callable[[EndpointState], object] | None = None,
progress_model_stats: Callable[[], tuple[int, int]] | None = None,
) -> tuple[list[EndpointState], int]:
if not configs:
return [], 0
states: list[EndpointState] = []
failures = 0
total = len(configs)
if progress_label is not None and settings.progress_every > 0:
logger.info(
f"{progress_label}: probing {total} endpoints with "
f"{min(settings.workers, total)} workers; progress every "
f"{settings.progress_every} completions"
)
queue: asyncio.Queue[EndpointConfig] = asyncio.Queue()
for config in configs:
queue.put_nowait(config)
completed = 0
async def worker() -> None:
nonlocal completed, failures
while True:
try:
config = queue.get_nowait()
except asyncio.QueueEmpty:
return
try:
result = await probe_endpoint(
session,
config,
settings.timeout,
readiness_timeout=(
settings.readiness_timeout if check_readiness else None
),
readiness_model_attempts=settings.readiness_model_attempts,
)
except Exception: # pragma: no cover - defensive isolation
result = None
completed += 1
if result is not None:
states.append(result)
if state_callback is not None:
state_callback(result)
else:
failures += 1
if (
progress_label is not None
and settings.progress_every > 0
and completed % settings.progress_every == 0
and completed < total
):
model_count, model_assignments = (
progress_model_stats()
if progress_model_stats is not None
else _model_stats(states)
)
logger.info(
f"{progress_label} progress: {completed}/{total} checked "
f"({len(states)} live, {failures} failed); "
f"models={model_count} assignments={model_assignments}"
)
workers = [
asyncio.create_task(worker(), name=f"endpoint-probe-{index}")
for index in range(min(settings.workers, total))
]
try:
await asyncio.gather(*workers)
finally:
for task in workers:
if not task.done():
task.cancel()
if workers:
await asyncio.gather(*workers, return_exceptions=True)
return states, failures
class EndpointRegistry:
"""Event-loop-owned snapshots of currently usable endpoints."""
def __init__(self, blacklist_duration: float = DEFAULT_BLACKLIST_DURATION) -> None:
self.blacklist_duration = blacklist_duration
self._states: dict[str, EndpointState] = {}
self._endpoint_quarantine: dict[str, float] = {}
self._model_quarantine: dict[tuple[str, str], float] = {}
self._model_latencies: dict[tuple[str, str], float] = {}
def _clear_model_latencies(
self,
base_url: str,
model: str | None = None,
) -> None:
if model is not None:
self._model_latencies.pop((base_url, model), None)
return
self._model_latencies = {
pair: latency
for pair, latency in self._model_latencies.items()
if pair[0] != base_url
}
def _prune_expired_quarantines(self, now: float) -> None:
self._endpoint_quarantine = {
url: expires_at
for url, expires_at in self._endpoint_quarantine.items()
if expires_at > now
}
self._model_quarantine = {
pair: expires_at
for pair, expires_at in self._model_quarantine.items()
if expires_at > now
}
def _filtered_state(
self,
state: EndpointState,
models: Iterable[str] | None = None,
) -> EndpointState | None:
"""Apply active endpoint/model quarantines to an immutable state."""
url = state.config.base_url
if url in self._endpoint_quarantine:
return None
candidate_models = state.models if models is None else frozenset(models)
eligible_models = frozenset(
model
for model in candidate_models
if (url, model) not in self._model_quarantine
)
if not eligible_models:
return None
return dataclass_replace(state, models=eligible_models)
def upsert(self, state: EndpointState) -> bool:
"""Publish one successful probe without waiting for its probe batch."""
self._prune_expired_quarantines(time.monotonic())
self._clear_model_latencies(state.config.base_url)
filtered = self._filtered_state(state)
if filtered is None:
return False
self._states[state.config.base_url] = filtered
return True
def clear_endpoint_blacklist(self) -> int:
"""Make failed endpoints eligible for the next source-based refresh."""
self._prune_expired_quarantines(time.monotonic())
cleared = len(self._endpoint_quarantine)
self._endpoint_quarantine.clear()
return cleared
def apply_health_results(
self,
states: Iterable[EndpointState],
probed_configs: Iterable[EndpointConfig],
) -> None:
successful = {state.config.base_url: state for state in states}
probed_urls = {config.base_url for config in probed_configs}
self._prune_expired_quarantines(time.monotonic())
for url in probed_urls:
if url in successful:
previous = self._states.get(url)
refreshed = successful[url]
models = (
refreshed.models & previous.models
if previous is not None
else refreshed.models
)
filtered = self._filtered_state(refreshed, models)
if filtered is not None:
self._states[url] = filtered
for model in set(self.model_latencies(url)) - filtered.models:
self._clear_model_latencies(url, model)
else:
self._states.pop(url, None)
self._clear_model_latencies(url)
else:
self._states.pop(url, None)
self._clear_model_latencies(url)
def discard(self, base_url: str) -> bool:
now = time.monotonic()
self._prune_expired_quarantines(now)
newly_quarantined = False
if self.blacklist_duration > 0:
newly_quarantined = base_url not in self._endpoint_quarantine
self._endpoint_quarantine[base_url] = now + self.blacklist_duration
removed = self._states.pop(base_url, None) is not None
self._clear_model_latencies(base_url)
return newly_quarantined or removed
def discard_model(self, base_url: str, model: str) -> bool:
now = time.monotonic()
self._prune_expired_quarantines(now)
pair = (base_url, model)
newly_quarantined = False
if self.blacklist_duration > 0:
newly_quarantined = pair not in self._model_quarantine
self._model_quarantine[pair] = now + self.blacklist_duration
self._clear_model_latencies(base_url, model)
state = self._states.get(base_url)
if state is None or model not in state.models:
return newly_quarantined
remaining_models = state.models - {model}
if not remaining_models:
self._states.pop(base_url, None)
else:
self._states[base_url] = dataclass_replace(
state,
models=frozenset(remaining_models),
)
return True
def candidates(self, model: str) -> list[EndpointState]:
candidates = [state for state in self._states.values() if model in state.models]
if "cloud" in model.lower():
return sorted(
candidates,
key=lambda state: (state.latency_seconds, state.config.base_url),
)
def priority(state: EndpointState) -> tuple[bool, float, str]:
latency = self._model_latencies.get((state.config.base_url, model))
return (
latency is None,
state.latency_seconds if latency is None else latency,
state.config.base_url,
)
return sorted(candidates, key=priority)
def record_generation_latency(
self,
base_url: str,
model: str,
latency_seconds: float,
) -> None:
if "cloud" in model.lower():
return
state = self._states.get(base_url)
if state is None or model not in state.models:
return
pair = (base_url, model)
previous = self._model_latencies.get(pair)
self._model_latencies[pair] = (
latency_seconds
if previous is None
else previous + GENERATION_LATENCY_EMA_ALPHA * (latency_seconds - previous)
)
def model_latencies(self, base_url: str) -> dict[str, float]:
return dict(
sorted(
(
(model, latency)
for (url, model), latency in self._model_latencies.items()
if url == base_url
)
)
)
def configs(self) -> list[EndpointConfig]:
return [state.config for state in self._states.values()]
def models(self) -> list[str]:
models = {model for state in self._states.values() for model in state.models}
return sorted(models)
def states(self) -> list[EndpointState]:
states = list(self._states.values())
return sorted(
states,
key=lambda state: (state.latency_seconds, state.config.base_url),
)
def model_stats(self) -> tuple[int, int]:
return _model_stats(self._states.values())
def endpoint_count(self) -> int:
return len(self._states)
class RouterRuntime:
"""Own startup discovery and the health/discovery background workers."""
def __init__(
self,
config: RouterConfig,
logger: logging.Logger,
session: aiohttp.ClientSession,
) -> None:
self.config = config
self.logger = logger
self.session = session
self.registry = EndpointRegistry(config.discovery.blacklist_duration)
self._refresh_lock = asyncio.Lock()
self._stop_event = asyncio.Event()
self._initial_discovery_done = asyncio.Event()
self._tasks: list[asyncio.Task[None]] = []
def _report(self, label: str, total: int, live: int, failed: int) -> None:
model_count, model_assignments = self.registry.model_stats()
self.logger.info(
f"{label}: {live}/{total} endpoints live ({failed} removed); "
f"models={model_count} assignments={model_assignments}"
)
async def discover(self) -> tuple[int, int]:
async with self._refresh_lock:
discovery = self.config.discovery
cleared = self.registry.clear_endpoint_blacklist()
if cleared:
self.logger.info(
f"Refresh made {cleared} blacklisted endpoints eligible for "
"Shodan/masscan import."
)
configs = await retrieve_endpoint_configs(
self.session,
discovery,
logger=self.logger,
)
label = (
"Shodan + masscan discovery"
if discovery.masscan_list_paths
else "Shodan discovery"
)
states, failed = await probe_endpoints(
self.session,
configs,
self.config.probe,
self.logger,
check_readiness=True,
progress_label=label,
state_callback=self.registry.upsert,
progress_model_stats=self.registry.model_stats,
)
active = self.registry.endpoint_count()
model_count, model_assignments = self.registry.model_stats()
self.logger.info(
f"{label}: {len(states)}/{len(configs)} candidates passed "
f"({failed} failed); registry={active} endpoints; "
f"models={model_count} assignments={model_assignments}"
)
return active, failed
async def health_check(self) -> tuple[int, int]:
async with self._refresh_lock:
configs = self.registry.configs()
states, failed = await probe_endpoints(
self.session,
configs,
self.config.probe,
self.logger,
progress_label="Health check",
progress_model_stats=self.registry.model_stats,
)
self.registry.apply_health_results(states, configs)
self._report("Health check", len(configs), len(states), failed)
return len(states), failed
async def _run_background_action(
self,
name: str,
action: Callable[[], Awaitable[object]],
) -> None:
try:
await action()
except asyncio.CancelledError:
raise
except ConfigurationError as error:
self.logger.warning(f"{name} skipped: {error}")
except Exception: # pragma: no cover - defensive isolation
self.logger.exception(f"{name} failed")
async def _initial_discovery(self) -> None:
try:
await self._run_background_action(
"Initial endpoint discovery",
self.discover,
)
finally:
self._initial_discovery_done.set()
async def _background_loop(
self,
name: str,
interval: float,
action: Callable[[], Awaitable[object]],
) -> None:
await self._initial_discovery_done.wait()
while not self._stop_event.is_set():
try:
await asyncio.wait_for(self._stop_event.wait(), timeout=interval)
return
except TimeoutError:
await self._run_background_action(name, action)
def start_background(self) -> None:
initial_task = asyncio.create_task(
self._initial_discovery(),
name="initial-endpoint-discovery",
)
self._tasks.append(initial_task)
loops = (
(
"endpoint health check",
self.config.probe.health_interval,
self.health_check,
),
(
"endpoint discovery",
self.config.discovery.interval,
self.discover,
),
)
for name, interval, action in loops:
if interval <= 0:
continue
task = asyncio.create_task(
self._background_loop(name, interval, action),
name=name.replace(" ", "-"),
)
self._tasks.append(task)
async def stop(self) -> None:
self._stop_event.set()
self._initial_discovery_done.set()
for task in self._tasks:
task.cancel()
if self._tasks:
await asyncio.gather(*self._tasks, return_exceptions=True)
self._tasks.clear()
def _strict_json_loads(raw: bytes) -> dict[str, Any]:
def reject_constant(value: str) -> None:
raise ValueError(f"invalid JSON constant {value}")
try:
payload = json.loads(raw, parse_constant=reject_constant)
except (
UnicodeDecodeError,
json.JSONDecodeError,
RecursionError,
ValueError,
) as error:
raise RequestError(
HTTPStatus.BAD_REQUEST,
f"Invalid JSON body: {error}",
code="invalid_json",
) from error
if not isinstance(payload, dict):
raise RequestError(
HTTPStatus.BAD_REQUEST,
"The request body must be a JSON object.",
code="invalid_request_error",
)
return payload
def _router_headers(
request_id: str,
upstream: str | None = None,
) -> CIMultiDict[str]:
headers = CIMultiDict(
{
"Access-Control-Allow-Origin": "*",
"Cache-Control": "no-store",
"Server": f"OpenAI-Compatible-Router/{VERSION}",
"Via": f"1.1 open-router-cli/{VERSION}",
"X-Open-Router-Request-ID": request_id,
}
)
if upstream is not None:
headers["X-Open-Router-Upstream"] = upstream
return headers
def _json_response(
status: HTTPStatus,
body: Mapping[str, Any],
request_id: str,
*,
headers: Mapping[str, str] | None = None,
) -> web.Response:
response_headers = _router_headers(request_id)
if headers is not None:
response_headers.update(headers)
return web.json_response(
body,
status=int(status),
headers=response_headers,
dumps=lambda value: json.dumps(
value,
separators=(",", ":"),
allow_nan=False,
),
)
def _error_response(error: RequestError, request_id: str) -> web.Response:
error_type = (
"invalid_request_error"
if error.status < HTTPStatus.INTERNAL_SERVER_ERROR
else "server_error"
)
headers = (
{"WWW-Authenticate": 'Bearer realm="OpenAI-compatible API"'}
if error.status == HTTPStatus.UNAUTHORIZED
else None
)
return _json_response(
error.status,
{
"error": {
"message": error.message,
"type": error_type,
"param": error.param,
"code": error.code,
}
},
request_id,
headers=headers,
)
class RouterHttpApp:
"""aiohttp routes and credential-isolating upstream proxy."""
def __init__(
self,
runtime: RouterRuntime,
options: ServerOptions,
logger: logging.Logger,
session: aiohttp.ClientSession,
request_logger: JsonlLogger | None = None,
) -> None:
self.runtime = runtime
self.options = options
self.logger = logger
self.session = session
self.request_logger = request_logger
def create(self) -> web.Application:
app = web.Application(
client_max_size=0,
handler_args={"auto_decompress": False, "lingering_time": 0},
middlewares=[self.errors_and_auth],
)
self._add_routes(app, "GET", {"/health", "/healthz"}, self.health)
self._add_routes(app, "HEAD", {"/health", "/healthz"}, self.health)
self._add_routes(app, "GET", MODEL_ROUTES, self.models)
self._add_routes(app, "GET", ENDPOINT_ROUTES, self.endpoints)
self._add_routes(app, "POST", set(ROUTES), self.generate)
app.router.add_route("OPTIONS", "/{path:.*}", self.options_request)
app.router.add_route("*", "/{path:.*}", self.unknown)
return app
@staticmethod
def _add_routes(
app: web.Application,
method: str,
paths: Iterable[str],
handler: Callable[[web.Request], Awaitable[web.StreamResponse]],
) -> None:
for path in paths:
app.router.add_route(method, path, handler)
app.router.add_route(method, f"{path}/", handler)
@staticmethod
def _path(request: web.Request) -> str:
return request.path.rstrip("/") if request.path != "/" else "/"
def _check_authentication(self, request: web.Request) -> None:
expected = self.options.api_keys
if not expected:
return
authorization = request.headers.get("Authorization", "")
scheme, separator, token = authorization.partition(" ")
valid = bool(separator and scheme.lower() == "bearer")
matches = [hmac.compare_digest(token, api_key) for api_key in expected]
if not valid or not any(matches):
raise RequestError(
HTTPStatus.UNAUTHORIZED,
"Incorrect API key provided.",
code="invalid_api_key",
)
@web.middleware
async def errors_and_auth(
self,
request: web.Request,
handler: Callable[[web.Request], Awaitable[web.StreamResponse]],
) -> web.StreamResponse:
request[REQUEST_ID_KEY] = f"req_{uuid.uuid4().hex}"
try:
if request.method == "POST" or self._path(request) in (
MODEL_ROUTES | ENDPOINT_ROUTES
):
self._check_authentication(request)
return await handler(request)
except RequestError as error:
response = _error_response(error, request[REQUEST_ID_KEY])
if error.code == "request_timeout":
response.force_close()
return response
except OSError as error:
return _error_response(
RequestError(
HTTPStatus.INTERNAL_SERVER_ERROR,
f"Unable to log the request: {error}",
code="logging_error",
),
request[REQUEST_ID_KEY],
)
async def options_request(self, request: web.Request) -> web.Response:
headers = _router_headers(request[REQUEST_ID_KEY])
headers.update(
{
"Access-Control-Allow-Methods": "GET, POST, OPTIONS",
"Access-Control-Allow-Headers": (
"Authorization, Content-Encoding, Content-Type, OpenAI-Beta, "
"OpenAI-Organization, OpenAI-Project, X-Request-ID"
),
"Access-Control-Max-Age": "86400",
}
)
return web.Response(status=HTTPStatus.NO_CONTENT, headers=headers)
async def health(self, request: web.Request) -> web.Response:
return _json_response(
HTTPStatus.OK,
{
"status": "ok",
"endpoints": self.runtime.registry.endpoint_count(),
"models": len(self.runtime.registry.models()),
},
request[REQUEST_ID_KEY],
)
async def models(self, request: web.Request) -> web.Response:
models = self.runtime.registry.models()
return _json_response(
HTTPStatus.OK,
{
"object": "list",
"data": [
{
"id": model,
"object": "model",
"created": 0,
"owned_by": "open-router",
}
for model in models
],
},
request[REQUEST_ID_KEY],
)
async def endpoints(self, request: web.Request) -> web.Response:
states = self.runtime.registry.states()
model_count, model_assignments = _model_stats(states)
return _json_response(
HTTPStatus.OK,
{
"object": "list",
"data": [
{
"object": "endpoint",
"base_url": state.config.base_url,
"probe_latency_seconds": state.latency_seconds,
"model_generation_latencies_seconds": (
self.runtime.registry.model_latencies(state.config.base_url)
),
"checked_at": datetime.fromtimestamp(
state.checked_at,
UTC,
)
.isoformat()
.replace("+00:00", "Z"),
"model_count": len(state.models),
"models": sorted(state.models),
}
for state in states
],
"stats": {
"endpoint_count": len(states),
"model_count": model_count,
"model_assignments": model_assignments,
},
},
request[REQUEST_ID_KEY],
)
async def unknown(self, request: web.Request) -> web.Response:
return _error_response(
RequestError(
HTTPStatus.NOT_FOUND,
f"Unknown endpoint: {self._path(request)}",
code="not_found",
),
request[REQUEST_ID_KEY],
)
async def _read_request_body(self, request: web.Request) -> bytes:
content_encoding = request.headers.get("Content-Encoding", "identity").lower()
if content_encoding not in {"", "identity", "gzip"}:
raise RequestError(
HTTPStatus.UNSUPPORTED_MEDIA_TYPE,
f"Unsupported Content-Encoding: {content_encoding}",
code="unsupported_content_encoding",
)
if (
not request.headers.get("Transfer-Encoding")
and request.content_length is None
):
raise RequestError(
HTTPStatus.LENGTH_REQUIRED,
"Content-Length or chunked Transfer-Encoding is required.",
code="length_required",
)
try:
body = await asyncio.wait_for(
request.read(),
timeout=DEFAULT_REQUEST_BODY_TIMEOUT,
)
except TimeoutError as error:
raise RequestError(
HTTPStatus.REQUEST_TIMEOUT,
f"The request body was not received within "
f"{DEFAULT_REQUEST_BODY_TIMEOUT:g} seconds.",
code="request_timeout",
) from error
except web.RequestPayloadError as error:
message = (
"The gzip request body is invalid."
if content_encoding == "gzip"
else "The request body ended before it was complete."
)
code = (
"invalid_content_encoding"
if content_encoding == "gzip"
else "incomplete_body"
)
raise RequestError(HTTPStatus.BAD_REQUEST, message, code=code) from error
if content_encoding == "gzip":
try:
body = gzip.decompress(body)
except (EOFError, OSError) as error:
raise RequestError(
HTTPStatus.BAD_REQUEST,
"The gzip request body is invalid.",
code="invalid_content_encoding",
) from error
return body
def _log_request(
self,
request: web.Request,
request_id: str,
resource: str,
payload: Mapping[str, Any],
candidates: Sequence[EndpointState],
) -> None:
if self.request_logger is None:
return
peer = (
request.transport.get_extra_info("peername") if request.transport else None
)
client_host = peer[0] if isinstance(peer, tuple) and peer else request.remote
client_port = peer[1] if isinstance(peer, tuple) and len(peer) > 1 else None
headers = {
name: "[REDACTED]" if _is_sensitive_header(name) else value
for name, value in request.headers.items()
}
self.request_logger.append(
{
"timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"),
"request_id": request_id,
"endpoint": resource.replace("/", "."),
"method": "POST",
"path": request.path_qs,
"model": payload["model"],
"candidate_count": len(candidates),
"client": {"host": client_host, "port": client_port},
"headers": headers,
"body": payload,
}
)
@staticmethod
def _forward_headers(request: web.Request, request_id: str) -> dict[str, str]:
forwarded: dict[str, tuple[str, str]] = {}
for name, value in request.headers.items():
lower = name.lower()
if lower in HOP_BY_HOP_HEADERS | {
"host",
"content-encoding",
"content-length",
} or _is_sensitive_header(lower):
continue
forwarded[lower] = (name, value)
forwarded["x-open-router-request-id"] = (
"X-Open-Router-Request-ID",
request_id,
)
return {name: value for name, value in forwarded.values()}
def _discard_upstream(self, endpoint: EndpointConfig, reason: str) -> None:
if self.runtime.registry.discard(endpoint.base_url):
duration = self.runtime.config.discovery.blacklist_duration
action = (
f"Blacklisted upstream {endpoint.base_url} for {duration:g} seconds"
if duration > 0
else f"Removed upstream {endpoint.base_url} until the next discovery"
)
self.logger.warning(f"{action}: {reason}")
def _discard_upstream_model(
self,
endpoint: EndpointConfig,
model: str,
reason: str,
) -> None:
if self.runtime.registry.discard_model(endpoint.base_url, model):
duration = self.runtime.config.discovery.blacklist_duration
action = (
f"Blacklisted model {model!r} on upstream {endpoint.base_url} for "
f"{duration:g} seconds"
if duration > 0
else f"Removed model {model!r} from upstream {endpoint.base_url} "
"until the next discovery"
)
self.logger.warning(f"{action}: {reason}")
async def generate(self, request: web.Request) -> web.StreamResponse:
request_id = request[REQUEST_ID_KEY]
resource = ROUTES[self._path(request)]
body = await self._read_request_body(request)
payload = _strict_json_loads(body)
model = payload.get("model")
if not isinstance(model, str) or not model:
raise RequestError(
HTTPStatus.BAD_REQUEST,
"The model field is required and must be a non-empty string.",
code="missing_required_parameter",
param="model",
)
candidates = self.runtime.registry.candidates(model)
if not candidates:
raise RequestError(
HTTPStatus.SERVICE_UNAVAILABLE,
f"No live endpoint currently advertises model {model!r}.",
code="model_not_available",
param="model",
)
self._log_request(request, request_id, resource, payload, candidates)
return await self._proxy_request(
request,
request_id=request_id,
resource=resource,
model=model,
body=body,
stream=payload.get("stream") is True,
candidates=candidates,
)
@staticmethod
def _relay_headers(
source: Iterable[tuple[str, str]],
request_id: str,
upstream: str,
*,
buffered: bool,
) -> CIMultiDict[str]:
headers: CIMultiDict[str] = CIMultiDict()
for name, value in source:
lower = name.lower()
if lower in HOP_BY_HOP_HEADERS or (buffered and lower == "content-length"):
continue
if lower in {
"access-control-allow-origin",
"cache-control",
"server",
"via",
"x-open-router-request-id",
"x-open-router-upstream",
}:
continue
headers.add(name, value)
headers.update(_router_headers(request_id, upstream))
return headers
async def _proxy_request(
self,
request: web.Request,
*,
request_id: str,
resource: str,
model: str,
body: bytes,
stream: bool,
candidates: Sequence[EndpointState],
) -> web.StreamResponse:
buffered_error: BufferedUpstreamResponse | None = None
timeout_seconds = self.options.upstream_timeout
timeout = aiohttp.ClientTimeout(
total=None,
sock_connect=timeout_seconds,
sock_read=timeout_seconds,
)
for state in candidates:
endpoint = state.config
generation_started = time.monotonic()
upstream_url = endpoint.api_url(resource)
if request.query_string:
upstream_url = f"{upstream_url}?{request.query_string}"
loading_deadline: float | None = None
loading_attempt = 0
while True:
if (
loading_deadline is not None
and time.monotonic() >= loading_deadline
):
self.logger.warning(
f"Upstream {endpoint.base_url} was still loading model "
f"{model!r} after {timeout_seconds:g} seconds; trying the "
"next candidate"
)
break
try:
response = await self.session.post(
upstream_url,
data=body,
headers=self._forward_headers(request, request_id),
timeout=timeout,
allow_redirects=False,
)
except (aiohttp.ClientError, asyncio.TimeoutError, ValueError) as error:
self._discard_upstream(endpoint, f"request failed: {error}")
break
if (
_is_upstream_failure_status(response.status)
or response.status in ENDPOINT_FAILURE_STATUSES
or response.status in MODEL_FAILURE_STATUSES
):
try:
error_body = await response.content.read(
MAX_MODELS_RESPONSE_BYTES + 1
)
except (aiohttp.ClientError, asyncio.TimeoutError) as error:
self._discard_upstream(
endpoint,
f"failed while reading HTTP {response.status}: {error}",
)
response.close()
break
buffered_error = BufferedUpstreamResponse(
status=response.status,
reason=response.reason or "",
headers=tuple(response.headers.items()),
body=error_body,
endpoint=endpoint,
)
if _is_model_loading_error(response.status, error_body):
response.close()
now = time.monotonic()
if loading_deadline is None:
loading_deadline = now + timeout_seconds
remaining = loading_deadline - now
if remaining <= 0:
continue
delay = min(
_retry_delay(
loading_attempt,
response.headers.get("Retry-After"),
),
remaining,
)
loading_attempt += 1
self.logger.info(
f"Upstream {endpoint.base_url} is loading model "
f"{model!r}; retrying the same endpoint in {delay:g} "
"seconds"
)
await asyncio.sleep(delay)
continue
if (
response.status in MODEL_FAILURE_STATUSES
or _is_model_unavailable_error(
response.status,
error_body,
)
):
self._discard_upstream_model(
endpoint,
model,
f"returned HTTP {response.status}",
)
else:
self._discard_upstream(
endpoint,
f"returned HTTP {response.status}",
)
response.close()
break
if not stream:
try:
response_body = await response.read()
except (aiohttp.ClientError, asyncio.TimeoutError) as error:
self._discard_upstream(
endpoint,
f"stream ended early: {error}",
)
response.close()
break
buffered = BufferedUpstreamResponse(
status=response.status,
reason=response.reason or "",
headers=tuple(response.headers.items()),
body=response_body,
endpoint=endpoint,
)
if response.status == HTTPStatus.OK:
self.runtime.registry.record_generation_latency(
endpoint.base_url,
model,
time.monotonic() - generation_started,
)
response.close()
return self._relay_buffered(buffered, request_id)
return await self._relay_upstream(
request,
response,
endpoint,
request_id,
model,
generation_started,
)
if buffered_error is not None:
return self._relay_buffered(buffered_error, request_id)
detail = (
f"{len(candidates)} upstream connection attempts failed."
if candidates
else "All upstream candidates failed."
)
return _error_response(
RequestError(
HTTPStatus.BAD_GATEWAY,
f"Unable to connect to a live upstream: {detail}",
code="upstream_unavailable",
),
request_id,
)
async def _relay_upstream(
self,
request: web.Request,
response: aiohttp.ClientResponse,
endpoint: EndpointConfig,
request_id: str,
model: str,
generation_started: float,
) -> web.StreamResponse:
headers = self._relay_headers(
response.headers.items(),
request_id,
endpoint.base_url,
buffered=False,
)
downstream = web.StreamResponse(
status=response.status,
reason=response.reason,
headers=headers,
)
await downstream.prepare(request)
try:
async for chunk in response.content.iter_any():
await downstream.write(chunk)
await downstream.write_eof()
if response.status == HTTPStatus.OK:
self.runtime.registry.record_generation_latency(
endpoint.base_url,
model,
time.monotonic() - generation_started,
)
except ConnectionResetError:
self.logger.debug(
f"Client disconnected while streaming {endpoint.base_url}"
)
except (aiohttp.ClientError, asyncio.TimeoutError) as error:
self._discard_upstream(
endpoint,
f"stream ended early: {error}",
)
if request.transport is not None:
request.transport.abort()
finally:
response.close()
return downstream
def _relay_buffered(
self,
response: BufferedUpstreamResponse,
request_id: str,
) -> web.Response:
headers = self._relay_headers(
response.headers,
request_id,
response.endpoint.base_url,
buffered=True,
)
return web.Response(
status=response.status,
reason=response.reason,
body=response.body,
headers=headers,
)
def non_negative_int(value: str) -> int:
try:
parsed = int(value)
except ValueError as error:
raise argparse.ArgumentTypeError("must be an integer") from error
if parsed < 0:
raise argparse.ArgumentTypeError("must be zero or greater")
return parsed
def positive_int(value: str) -> int:
parsed = non_negative_int(value)
if parsed < 1:
raise argparse.ArgumentTypeError("must be at least 1")
return parsed
def non_negative_float(value: str) -> float:
try:
parsed = float(value)
except ValueError as error:
raise argparse.ArgumentTypeError("must be a number") from error
if parsed < 0:
raise argparse.ArgumentTypeError("must be zero or greater")
return parsed
def positive_float(value: str) -> float:
parsed = non_negative_float(value)
if parsed <= 0:
raise argparse.ArgumentTypeError("must be greater than zero")
return parsed
def http_safe_api_key(value: str) -> str:
if not value:
raise argparse.ArgumentTypeError("must not be empty")
if "\r" in value or "\n" in value:
raise argparse.ArgumentTypeError("must not contain CR or LF characters")
return value
def parse_args(argv: list[str] | None = None) -> RouterConfig:
parser = argparse.ArgumentParser(
description="Route OpenAI-compatible requests across validated endpoints.",
allow_abbrev=False,
)
# fmt: off
argument_specs: tuple[tuple[str, dict[str, Any]], ...] = (
("--shodan-api-key", {"help": "Shodan API key (defaults to SHODAN_API_KEY)"}),
("--shodan-mode", {"choices": ("auto", "api", "public"), "default": "auto", "help": "endpoint retrieval mode; auto uses the API when a key is available (default: auto)"}),
("--shodan-max-results", {"type": positive_int, "help": "stop Shodan API retrieval after this many raw matches (default: all)"}),
("--masscan-lists", {"dest": "masscan_list_paths", "nargs": "+", "type": Path, "metavar": "FILE", "help": "masscan .list files to import; shell globs are supported (default: none)"}),
("--retrieval-retries", {"type": int, "choices": range(0, 11), "default": DEFAULT_RETRIEVAL_RETRIES, "metavar": "0..10", "help": "Shodan transient-error retries (default: 3)"}),
("--host", {"default": DEFAULT_HOST, "help": f"listen host (default: {DEFAULT_HOST})"}),
("--retrieval-timeout", {"type": positive_float, "default": DEFAULT_RETRIEVAL_TIMEOUT, "help": "Shodan request timeout in seconds (default: 30)"}),
("--port", {"type": non_negative_int, "default": DEFAULT_PORT, "help": "listen port; use 0 for an OS-selected port (default: 8000)"}),
("--workers", {"type": positive_int, "default": DEFAULT_WORKERS, "help": f"parallel endpoint probe workers (default: {DEFAULT_WORKERS})"}),
("--probe-timeout", {"type": positive_float, "default": DEFAULT_PROBE_TIMEOUT, "help": "model discovery timeout in seconds (default: 5)"}),
("--readiness-timeout", {"type": positive_float, "default": DEFAULT_READINESS_TIMEOUT, "help": "generation readiness timeout per model in seconds (default: 30)"}),
("--readiness-model-attempts", {"type": positive_int, "default": DEFAULT_READINESS_MODEL_ATTEMPTS, "help": "models tried while validating an endpoint (default: 3)"}),
("--upstream-timeout", {"type": positive_float, "default": DEFAULT_UPSTREAM_TIMEOUT, "help": f"per-upstream attempt timeout in seconds (default: {DEFAULT_UPSTREAM_TIMEOUT:g})"}),
("--health-interval", {"type": non_negative_float, "default": DEFAULT_HEALTH_INTERVAL, "help": f"seconds between live endpoint checks; 0 disables (default: {DEFAULT_HEALTH_INTERVAL:g})"}),
("--discovery-interval", {"type": non_negative_float, "default": DEFAULT_DISCOVERY_INTERVAL, "help": f"seconds between endpoint retrievals; 0 disables (default: {DEFAULT_DISCOVERY_INTERVAL:g})"}),
("--blacklist-duration", {"type": non_negative_float, "default": DEFAULT_BLACKLIST_DURATION, "help": f"maximum blacklist seconds; endpoint entries reset during source refresh, model pairings retain this TTL; 0 disables (default: {DEFAULT_BLACKLIST_DURATION:g})"}),
("--progress-every", {"type": non_negative_int, "default": DEFAULT_PROGRESS_EVERY, "metavar": "N", "help": f"print progress/model totals every N completed endpoints; 0 disables (default: {DEFAULT_PROGRESS_EVERY})"}),
("--log-file", {"type": Path, "help": "append operational and access logs"}),
("--log-requests", {"type": Path, "help": "append routed request audits as JSONL"}),
("--api-keys", {"nargs": "+", "type": http_safe_api_key, "default": (), "metavar": "KEY", "help": "authorized inbound bearer tokens; supplying this option enables authentication"}),
("--log-level", {"choices": LOG_LEVEL_NAMES, "default": DEFAULT_LOG_LEVEL, "help": "operational log level (default: INFO)"}),
("--quiet", {"action": "store_true", "help": "suppress DEBUG/INFO operational logs; warnings and errors remain"}),
("--version", {"action": "version", "version": f"{parser.prog} {VERSION}"}),
)
# fmt: on
for flag, kwargs in argument_specs:
parser.add_argument(flag, **kwargs)
args = parser.parse_args(argv)
if args.port > 65_535:
parser.error("--port must be between 0 and 65535")
args.shodan_api_key = args.shodan_api_key or os.environ.get("SHODAN_API_KEY")
if args.shodan_mode == "api" and not args.shodan_api_key:
parser.error("--shodan-mode api needs --shodan-api-key or SHODAN_API_KEY")
uses_public_retrieval = args.shodan_mode == "public" or (
args.shodan_mode == "auto" and not args.shodan_api_key
)
if uses_public_retrieval and args.shodan_max_results is not None:
parser.error("--shodan-max-results requires Shodan API mode")
if (
args.log_file is not None
and args.log_requests is not None
and args.log_file.expanduser().resolve()
== args.log_requests.expanduser().resolve()
):
parser.error("--log-file and --log-requests must use different files")
return RouterConfig(
discovery=DiscoveryConfig(
shodan_api_key=args.shodan_api_key,
shodan_mode=args.shodan_mode,
retrieval_timeout=args.retrieval_timeout,
retrieval_retries=args.retrieval_retries,
shodan_max_results=args.shodan_max_results,
masscan_list_paths=tuple(args.masscan_list_paths or ()),
interval=args.discovery_interval,
blacklist_duration=args.blacklist_duration,
),
probe=ProbeConfig(
workers=args.workers,
timeout=args.probe_timeout,
readiness_timeout=args.readiness_timeout,
readiness_model_attempts=args.readiness_model_attempts,
health_interval=args.health_interval,
progress_every=args.progress_every,
),
server=ServerOptions(
host=args.host,
port=args.port,
api_keys=tuple(dict.fromkeys(args.api_keys)),
log_file=args.log_file,
log_requests=args.log_requests,
upstream_timeout=args.upstream_timeout,
log_level=args.log_level,
quiet=args.quiet,
),
)
async def _run_server_async(
config: RouterConfig,
logger: logging.Logger,
request_logger: JsonlLogger | None,
) -> None:
options = config.server
connector = aiohttp.TCPConnector(limit=0, force_close=True)
async with aiohttp.ClientSession(
connector=connector,
cookie_jar=aiohttp.DummyCookieJar(),
auto_decompress=False,
trust_env=False,
) as session:
runtime = RouterRuntime(config, logger, session)
application = RouterHttpApp(
runtime,
options,
logger,
session,
request_logger,
).create()
runner = _HeaderTimeoutAppRunner(
application,
access_log=logger,
logger=logger,
shutdown_timeout=max(options.upstream_timeout, 1.0),
)
await runner.setup()
site = web.TCPSite(runner, options.host, options.port)
try:
await site.start()
address, port = runner.addresses[0][:2]
logger.info(
f"OpenAI router listening on http://{address}:{port}/v1 with "
f"{runtime.registry.endpoint_count()} live endpoints; initial "
"discovery is running in the background"
)
if request_logger is not None:
logger.info(f"Appending requests to {request_logger.path}")
runtime.start_background()
await asyncio.Event().wait()
finally:
await runtime.stop()
await runner.cleanup()
def run_server(config: RouterConfig) -> int:
options = config.server
try:
logger = configure_logger(
options.log_level,
quiet=options.quiet,
log_file=options.log_file,
)
except (ConfigurationError, OSError) as error:
print(f"Unable to configure logging: {error}", file=sys.stderr)
return 1
request_logger: JsonlLogger | None = None
try:
if options.log_file is not None:
logger.info(
f"Appending operational logs to "
f"{options.log_file.expanduser().resolve()}"
)
if options.log_requests is not None:
request_logger = JsonlLogger(options.log_requests)
asyncio.run(_run_server_async(config, logger, request_logger))
except KeyboardInterrupt:
logger.info("Shutting down.")
except (ConfigurationError, OSError) as error:
logger.error(f"{error}")
return 1
finally:
if request_logger is not None:
request_logger.close()
return 0
def main(argv: list[str] | None = None) -> int:
return run_server(parse_args(argv))
if __name__ == "__main__":
raise SystemExit(main())
aisbf/providers/shodan/runtime.py 0 → 100644
View file @ 674031f5
"""
Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
AISBF - AI Service Broker Framework || AI Should Be Free
Process-wide runtime for the "shodan" last-resort provider.
This module owns a single shared discovery/failover engine (vendored from
Pasquale Minervini's open-router project, see open_router_engine.py) and adapts
it to AISBF's provider contract. It builds the engine's DiscoveryConfig /
ProbeConfig from an AISBF provider's ``shodan_config`` dict, runs the background
discovery + health loops once per process, and exposes a small OpenAI-compatible
failover routing helper that returns AISBF-friendly types (a plain dict for
non-streaming, an async generator of OpenAI SSE bytes for streaming).
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
import asyncio
import json
import logging
import time
from http import HTTPStatus
from pathlib import Path
from typing import Any, AsyncIterator, Dict, List, Optional, Sequence
import aiohttp
from . import open_router_engine as engine
logger = logging.getLogger(__name__)
# Model names that mean "route to whatever the discovery pool currently offers"
# rather than a specific upstream model id (which for discovered Ollama/vLLM
# hosts is unpredictable and changes over time).
AUTO_MODEL_NAMES = frozenset({"", "auto", "any", "*", "shodan/auto", "shodan"})
class ShodanRuntime:
"""Owns one engine.RouterRuntime + shared upstream session for the process."""
def __init__(self, router_config: "engine.RouterConfig", upstream_timeout: float) -> None:
self._router_config = router_config
self._upstream_timeout = upstream_timeout
self._session: Optional[aiohttp.ClientSession] = None
self._runtime: Optional["engine.RouterRuntime"] = None
self._started = False
async def _ensure_started(self) -> None:
if self._started:
return
# The credential-free upstream client. Discovered hosts are untrusted, so
# we never attach auth headers (engine._upstream_headers already omits
# them); this session only carries content negotiation headers per call.
self._session = aiohttp.ClientSession()
self._runtime = engine.RouterRuntime(self._router_config, logger, self._session)
self._runtime.start_background()
self._started = True
logger.info(
"ShodanRuntime started: discovery mode=%s, masscan lists=%d, "
"discovery interval=%.0fs, health interval=%.0fs",
self._router_config.discovery.shodan_mode,
len(self._router_config.discovery.masscan_list_paths),
self._router_config.discovery.interval,
self._router_config.probe.health_interval,
)
@property
def registry(self) -> "engine.EndpointRegistry":
assert self._runtime is not None
return self._runtime.registry
async def available_models(self) -> List[str]:
await self._ensure_started()
return self.registry.models()
async def endpoint_count(self) -> int:
await self._ensure_started()
return self.registry.endpoint_count()
def _resolve_model(self, requested: str) -> Optional[str]:
"""Map a requested model to a concrete model id the pool can serve.
For an explicit model name we honour it as-is (candidates() will match).
For an ``auto`` alias we pick the fastest-servable model currently in the
registry: prefer a non-"cloud" model with at least one live candidate.
Returns None when the pool is empty.
"""
requested_norm = (requested or "").strip()
if requested_norm.lower() not in AUTO_MODEL_NAMES:
return requested_norm
models = self.registry.models()
if not models:
return None
# Prefer a model that actually has a live candidate right now, and prefer
# local (non-"cloud") models so we don't lean on remote paid services.
best: Optional[str] = None
for model in models:
if not self.registry.candidates(model):
continue
if "cloud" not in model.lower():
return model
best = best or model
return best or models[0]
async def route(
self,
resource: str,
body: Dict[str, Any],
stream: bool,
) -> Any:
"""Try the discovered upstreams in priority order until one answers.
Returns a plain OpenAI-compatible dict for non-streaming requests, or an
async generator yielding OpenAI-compatible SSE byte chunks for streaming.
Raises on total failure (no endpoint served the request) so the caller
can record a provider failure and, in a rotation, fail over.
"""
await self._ensure_started()
requested_model = body.get("model", "")
model = self._resolve_model(requested_model)
if model is None:
raise RuntimeError(
"shodan: no discovered endpoints are available yet "
"(discovery may still be warming up)"
)
# Send the resolved concrete model id upstream.
body = {**body, "model": model}
payload = json.dumps(body).encode("utf-8")
candidates = self.registry.candidates(model)
if not candidates:
raise RuntimeError(
f"shodan: no live upstream advertises model {model!r}"
)
t = self._upstream_timeout
timeout = aiohttp.ClientTimeout(total=None, sock_connect=t, sock_read=t)
headers = {
"content-type": "application/json",
"accept": "text/event-stream" if stream else "application/json",
}
last_error: Optional[str] = None
for state in candidates:
endpoint = state.config
url = endpoint.api_url(resource)
started = time.monotonic()
try:
response = await self._session.post(
url,
data=payload,
headers=headers,
timeout=timeout,
allow_redirects=False,
)
except (aiohttp.ClientError, asyncio.TimeoutError, ValueError) as error:
last_error = f"{endpoint.base_url}: request failed: {error}"
self.registry.discard(endpoint.base_url)
logger.warning("shodan upstream %s discarded: %s", endpoint.base_url, error)
continue
if self._is_failure_status(response.status):
error_body = await self._safe_read(response)
response.close()
last_error = f"{endpoint.base_url}: HTTP {response.status}"
self._quarantine_on_status(endpoint, model, response.status)
logger.warning(
"shodan upstream %s returned HTTP %s for model %s; trying next",
endpoint.base_url, response.status, model,
)
continue
# Success — this endpoint served the request.
if stream:
logger.info("shodan routing stream to %s (model %s)", endpoint.base_url, model)
return self._relay_stream(response, endpoint.base_url, model, started)
data = await self._safe_read(response)
response.close()
self.registry.record_generation_latency(
endpoint.base_url, model, time.monotonic() - started
)
logger.info("shodan served non-stream from %s (model %s)", endpoint.base_url, model)
try:
return json.loads(data)
except json.JSONDecodeError as error:
last_error = f"{endpoint.base_url}: invalid JSON ({error})"
self.registry.discard(endpoint.base_url)
continue
raise RuntimeError(
f"shodan: all {len(candidates)} candidate upstream(s) failed for "
f"model {model!r} (last error: {last_error})"
)
async def _relay_stream(
self,
response: aiohttp.ClientResponse,
base_url: str,
model: str,
started: float,
) -> AsyncIterator[bytes]:
"""Relay an upstream SSE body chunk-by-chunk as OpenAI-compatible bytes.
A failure after the first bytes cannot be retried on another endpoint
(the client stream has already begun), matching the upstream engine's
behaviour: we quarantine the endpoint and let the error surface.
"""
try:
async for chunk in response.content.iter_any():
if chunk:
yield chunk
self.registry.record_generation_latency(
base_url, model, time.monotonic() - started
)
except (aiohttp.ClientError, asyncio.TimeoutError) as error:
self.registry.discard(base_url)
logger.warning("shodan stream from %s failed mid-flight: %s", base_url, error)
raise
finally:
response.close()
@staticmethod
def _is_failure_status(status: int) -> bool:
return (
engine._is_upstream_failure_status(status)
or status in engine.ENDPOINT_FAILURE_STATUSES
or status in engine.MODEL_FAILURE_STATUSES
or status >= 500
)
def _quarantine_on_status(self, endpoint: "engine.EndpointConfig", model: str, status: int) -> None:
# 401 means the endpoint demands credentials we deliberately never send:
# drop the whole endpoint. 403/410 blacklist just this endpoint/model
# pairing. Everything else drops the endpoint for this pass.
if status in engine.MODEL_FAILURE_STATUSES:
self.registry.discard_model(endpoint.base_url, model)
else:
self.registry.discard(endpoint.base_url)
@staticmethod
async def _safe_read(response: aiohttp.ClientResponse) -> bytes:
try:
return await response.content.read(engine.MAX_MODELS_RESPONSE_BYTES + 1)
except (aiohttp.ClientError, asyncio.TimeoutError):
return b""
async def stop(self) -> None:
if self._runtime is not None:
await self._runtime.stop()
if self._session is not None:
await self._session.close()
self._started = False
# --- Process-wide singleton -------------------------------------------------
_runtime_singleton: Optional[ShodanRuntime] = None
_runtime_lock = asyncio.Lock()
def _build_router_config(shodan_config: Optional[Dict[str, Any]]) -> "engine.RouterConfig":
cfg = dict(shodan_config or {})
masscan_paths = tuple(
Path(p).expanduser() for p in (cfg.get("masscan_lists") or []) if str(p).strip()
)
discovery = engine.DiscoveryConfig(
shodan_api_key=cfg.get("shodan_api_key") or None,
shodan_mode=str(cfg.get("shodan_mode", "auto")),
retrieval_timeout=float(cfg.get("retrieval_timeout", engine.DEFAULT_RETRIEVAL_TIMEOUT)),
retrieval_retries=int(cfg.get("retrieval_retries", engine.DEFAULT_RETRIEVAL_RETRIES)),
shodan_max_results=cfg.get("shodan_max_results"),
masscan_list_paths=masscan_paths,
interval=float(cfg.get("discovery_interval", engine.DEFAULT_DISCOVERY_INTERVAL)),
blacklist_duration=float(cfg.get("blacklist_duration", engine.DEFAULT_BLACKLIST_DURATION)),
)
probe = engine.ProbeConfig(
workers=int(cfg.get("workers", engine.DEFAULT_WORKERS)),
timeout=float(cfg.get("probe_timeout", engine.DEFAULT_PROBE_TIMEOUT)),
readiness_timeout=float(cfg.get("readiness_timeout", engine.DEFAULT_READINESS_TIMEOUT)),
readiness_model_attempts=int(
cfg.get("readiness_model_attempts", engine.DEFAULT_READINESS_MODEL_ATTEMPTS)
),
health_interval=float(cfg.get("health_interval", engine.DEFAULT_HEALTH_INTERVAL)),
progress_every=int(cfg.get("progress_every", engine.DEFAULT_PROGRESS_EVERY)),
)
server = engine.ServerOptions(
upstream_timeout=float(cfg.get("upstream_timeout", engine.DEFAULT_UPSTREAM_TIMEOUT)),
)
return engine.RouterConfig(discovery=discovery, probe=probe, server=server)
async def get_runtime(shodan_config: Optional[Dict[str, Any]]) -> ShodanRuntime:
"""Return the process-wide ShodanRuntime, creating it on first use.
The first caller's ``shodan_config`` configures discovery for the whole
process; later callers reuse the same runtime (a single discovery pool is
the intent — one broker, one shared set of discovered endpoints).
"""
global _runtime_singleton
if _runtime_singleton is not None:
return _runtime_singleton
async with _runtime_lock:
if _runtime_singleton is None:
router_config = _build_router_config(shodan_config)
_runtime_singleton = ShodanRuntime(
router_config, router_config.server.upstream_timeout
)
return _runtime_singleton
docs/shodan-provider.md 0 → 100644
View file @ 674031f5
# Shodan last-resort provider
The `shodan` provider is a **credential-free, last-resort failover provider**. It
routes OpenAI-compatible requests to a pool of open model endpoints (Ollama,
vLLM, llama.cpp, LM Studio, …) that are automatically discovered and validated in
the background. It is meant to be added to a rotation as a `last_resort` entry so
that the rotation can still answer when every normal provider is disabled or
rate limited.
## Credit
The discovery, validation, live-registry, and failover-routing **engine** is
vendored, essentially verbatim, from the **open-router** project by
**Pasquale Minervini (@pminervini)**:
> https://github.com/pminervini/open-router
Full credit for that engine belongs to Pasquale Minervini. It lives at
`aisbf/providers/shodan/open_router_engine.py` with provenance and attribution
in its header, and is kept in sync with upstream rather than edited locally. The
upstream repository carries no LICENSE file; it is included here with
attribution and is **not** relicensed under AISBF's GPL. Only the engine (not its
aiohttp server / CLI layer) is used; AISBF wraps it with a small adapter
(`runtime.py`, `handler.py`).
## How it works
1. On first use, a single process-wide runtime starts the engine's background
discovery + health loops (Shodan facet queries and/or masscan list imports).
2. Each candidate endpoint is probed and only published after it answers
`/v1/models`, advertises a text-generation model, and completes a minimal
Chat Completions and Responses request.
3. A request routed to `shodan` is tried against the validated endpoints in
latency order, failing over transparently to the next endpoint on error.
4. Discovered upstreams are always queried **without credentials** — client
`Authorization`/API-key/cookie headers are never forwarded.
## Security note
The embedded engine issues Shodan queries and (optionally) imports masscan scan
results **from the host running AISBF**. Shodan/masscan results are untrusted
Internet hosts. Only scan networks you own or are explicitly authorized to test.
Prefer running discovery on a host and network where this is acceptable.
## Configuring the provider
Add a provider of `type: "shodan"`. It needs no API key. All discovery settings
live under `shodan_config` and are optional (sensible defaults apply):
```json
{
"id": "shodan",
"name": "Shodan last-resort",
"type": "shodan",
"endpoint": "",
"api_key_required": false,
"rate_limit": 0,
"shodan_config": {
"shodan_mode": "public",
"shodan_api_key": null,
"shodan_max_results": null,
"masscan_lists": [],
"discovery_interval": 3600,
"health_interval": 300,
"blacklist_duration": 86400,
"workers": 8,
"probe_timeout": 5.0,
"readiness_timeout": 30.0,
"readiness_model_attempts": 3,
"upstream_timeout": 300.0
}
}
```
`shodan_config` keys:
| Key | Default | Meaning |
| --- | --- | --- |
| `shodan_mode` | `auto` | `public` (facet scrape, no key), `api` (needs `shodan_api_key`), or `auto` |
| `shodan_api_key` | `null` | Shodan API key; only used in `api`/`auto` mode; sent only to Shodan |
| `shodan_max_results` | `null` | Cap on Shodan API results |
| `masscan_lists` | `[]` | Paths to masscan `list`-format outputs to import (port 11434) |
| `discovery_interval` | `86400` | Seconds between discovery refreshes (`0` disables the loop) |
| `health_interval` | `86400` | Seconds between liveness passes (`0` disables the loop) |
| `blacklist_duration` | `86400` | Seconds a failing endpoint/model stays quarantined |
| `workers` | `8` | Bounded concurrent probe queue size |
| `upstream_timeout` | `300` | Per-request upstream socket timeout, in seconds |
## Using it as a rotation last-resort
Mark the rotation entry with `"last_resort": true`. The special `auto` model name
means "use whatever the discovery pool currently advertises" (recommended, since
discovered model names change over time):
```json
{
"model_name": "lisa",
"providers": [
{ "provider_id": "codex", "weight": 10, "enabled": true,
"models": [{ "name": "gpt-5.6-luna" }] },
{ "provider_id": "kilo-stefy", "weight": 3, "enabled": true,
"models": [{ "name": "kilo-auto/free" }] },
{ "provider_id": "shodan", "last_resort": true,
"models": [{ "name": "auto" }] }
]
}
```
Behavior:
- While **any** normal provider (here `codex`, `kilo-stefy`) is available, the
`shodan` entry is held aside and never selected.
- When **all** normal providers are unavailable (disabled in the rotation via
`enabled: false`, or in a failure/usage cooldown) — the exact case that
otherwise returns HTTP 429 — the rotation engages the `last_resort` entries
instead.
- The last-resort pool is also engaged if all normal providers were initially
available but every one failed during the request's retry loop.
`last_resort` is a general rotation-entry flag: any provider type can carry it,
not just `shodan`.
## Notes / limitations
- Discovery warms up in the background; immediately after start the pool may be
empty and `shodan` will report no available upstream (the rotation then behaves
as if it had no last resort). Give discovery a little time, or seed it with
`masscan_lists`.
- Because discovered upstreams are heterogeneous open servers, capabilities
(tools, long context, streaming quirks) vary. It is a best-effort safety net,
not a quality-guaranteed provider.
requirements.txt
View file @ 674031f5
......@@ -43,3 +43,6 @@ mnemonic>=0.20
bitcoinlib>=0.6.14
web3>=6.0.0
eth-account>=0.9.0
# Required by the "shodan" last-resort provider discovery/failover engine
aiohttp>=3.14
tests/providers/test_shodan_provider.py 0 → 100644
View file @ 674031f5
"""Tests for the "shodan" last-resort provider adapter.
These cover the AISBF-specific adapter (model resolution, failover routing,
handler return contract, config building) around the vendored open-router
engine. The vendored engine itself is exercised by its upstream project.
"""
import json
import pytest
from aisbf.providers.shodan.runtime import (
ShodanRuntime,
_build_router_config,
get_runtime,
)
from aisbf.providers.shodan import open_router_engine as engine
# --- Fakes ------------------------------------------------------------------
class FakeConfig:
def __init__(self, base_url):
self.base_url = base_url
def api_url(self, resource):
return f"{self.base_url}/v1/{resource}"
class FakeState:
def __init__(self, base_url, models, latency=0.1):
self.config = FakeConfig(base_url)
self.models = frozenset(models)
self.latency_seconds = latency
class FakeRegistry:
"""Minimal stand-in for engine.EndpointRegistry."""
def __init__(self, states):
self._states = {s.config.base_url: s for s in states}
self.discarded = []
self.discarded_models = []
self.recorded_latencies = []
def models(self):
return sorted({m for s in self._states.values() for m in s.models})
def candidates(self, model):
return [s for s in self._states.values() if model in s.models]
def discard(self, base_url):
self.discarded.append(base_url)
self._states.pop(base_url, None)
return True
def discard_model(self, base_url, model):
self.discarded_models.append((base_url, model))
return True
def record_generation_latency(self, base_url, model, latency):
self.recorded_latencies.append((base_url, model, latency))
def endpoint_count(self):
return len(self._states)
class FakeContent:
def __init__(self, body: bytes, chunks=None, raise_mid=False):
self._body = body
self._chunks = chunks or []
self._raise_mid = raise_mid
async def read(self, _limit=None):
return self._body
async def iter_any(self):
import aiohttp
for i, ch in enumerate(self._chunks):
if self._raise_mid and i == 1:
raise aiohttp.ClientError("mid-stream failure")
yield ch
class FakeResponse:
def __init__(self, status=200, body=b"", chunks=None, raise_mid=False):
self.status = status
self.reason = "OK"
self.content = FakeContent(body, chunks, raise_mid)
self.closed = False
def close(self):
self.closed = True
class FakeSession:
"""Records POSTs and returns queued responses per URL host."""
def __init__(self, responses_by_host):
# responses_by_host: dict base_host -> list[FakeResponse | Exception]
self._by_host = {k: list(v) for k, v in responses_by_host.items()}
self.posts = []
async def post(self, url, data=None, headers=None, timeout=None, allow_redirects=False):
self.posts.append(url)
host = url.split("/v1/")[0]
queue = self._by_host.get(host)
if not queue:
raise AssertionError(f"no queued response for {host}")
item = queue.pop(0)
if isinstance(item, Exception):
raise item
return item
def _runtime_with(registry, session, upstream_timeout=5.0):
rt = ShodanRuntime.__new__(ShodanRuntime)
rt._router_config = None
rt._upstream_timeout = upstream_timeout
rt._session = session
rt._started = True
class _R:
def __init__(self, reg):
self.registry = reg
rt._runtime = _R(registry)
return rt
# --- Config builder ---------------------------------------------------------
class TestBuildRouterConfig:
def test_defaults_when_empty(self):
rc = _build_router_config(None)
assert rc.discovery.shodan_mode == "auto"
assert rc.discovery.masscan_list_paths == ()
assert rc.probe.workers == engine.DEFAULT_WORKERS
def test_overrides_applied(self):
rc = _build_router_config({
"shodan_mode": "public",
"discovery_interval": 1800,
"health_interval": 300,
"workers": 32,
"masscan_lists": ["/tmp/a.list", " ", "/tmp/b.list"],
"upstream_timeout": 120,
})
assert rc.discovery.shodan_mode == "public"
assert rc.discovery.interval == 1800
assert rc.probe.health_interval == 300
assert rc.probe.workers == 32
# blank entries are dropped
assert len(rc.discovery.masscan_list_paths) == 2
assert rc.server.upstream_timeout == 120
# --- Model resolution -------------------------------------------------------
class TestResolveModel:
def test_explicit_model_passthrough(self):
rt = _runtime_with(FakeRegistry([FakeState("http://h:1", ["llama3"])]), None)
assert rt._resolve_model("llama3") == "llama3"
def test_auto_prefers_non_cloud_with_candidate(self):
reg = FakeRegistry([
FakeState("http://h:1", ["some-cloud-model"]),
FakeState("http://h:2", ["qwen2.5:7b"]),
])
rt = _runtime_with(reg, None)
assert rt._resolve_model("auto") == "qwen2.5:7b"
def test_auto_empty_pool_returns_none(self):
rt = _runtime_with(FakeRegistry([]), None)
assert rt._resolve_model("auto") is None
@pytest.mark.parametrize("alias", ["", "auto", "any", "*", "shodan", "shodan/auto"])
def test_alias_names(self, alias):
reg = FakeRegistry([FakeState("http://h:1", ["m1"])])
rt = _runtime_with(reg, None)
assert rt._resolve_model(alias) == "m1"
# --- Failover routing -------------------------------------------------------
class TestRouteFailover:
@pytest.mark.asyncio
async def test_non_stream_success(self):
reg = FakeRegistry([FakeState("http://good:1", ["m1"])])
payload = {"id": "x", "choices": [{"message": {"content": "hi"}}]}
session = FakeSession({"http://good:1": [FakeResponse(200, json.dumps(payload).encode())]})
rt = _runtime_with(reg, session)
result = await rt.route("chat/completions", {"model": "m1", "messages": []}, stream=False)
assert result == payload
assert reg.recorded_latencies # latency recorded on success
@pytest.mark.asyncio
async def test_failover_to_second_endpoint(self):
reg = FakeRegistry([
FakeState("http://bad:1", ["m1"], latency=0.01), # tried first (faster)
FakeState("http://good:2", ["m1"], latency=0.02),
])
payload = {"ok": True}
session = FakeSession({
"http://bad:1": [FakeResponse(500, b"boom")],
"http://good:2": [FakeResponse(200, json.dumps(payload).encode())],
})
rt = _runtime_with(reg, session)
result = await rt.route("chat/completions", {"model": "m1", "messages": []}, stream=False)
assert result == payload
assert "http://bad:1" in reg.discarded # 5xx dropped the bad endpoint
@pytest.mark.asyncio
async def test_401_discards_endpoint_403_discards_model(self):
reg = FakeRegistry([FakeState("http://auth:1", ["m1"])])
session = FakeSession({"http://auth:1": [FakeResponse(401, b"nope")]})
rt = _runtime_with(reg, session)
with pytest.raises(RuntimeError):
await rt.route("chat/completions", {"model": "m1", "messages": []}, stream=False)
assert "http://auth:1" in reg.discarded
reg2 = FakeRegistry([FakeState("http://f:1", ["m1"])])
session2 = FakeSession({"http://f:1": [FakeResponse(403, b"forbidden")]})
rt2 = _runtime_with(reg2, session2)
with pytest.raises(RuntimeError):
await rt2.route("chat/completions", {"model": "m1", "messages": []}, stream=False)
assert ("http://f:1", "m1") in reg2.discarded_models
@pytest.mark.asyncio
async def test_all_fail_raises(self):
reg = FakeRegistry([FakeState("http://x:1", ["m1"])])
session = FakeSession({"http://x:1": [FakeResponse(503, b"down")]})
rt = _runtime_with(reg, session)
with pytest.raises(RuntimeError, match="candidate upstream"):
await rt.route("chat/completions", {"model": "m1", "messages": []}, stream=False)
@pytest.mark.asyncio
async def test_no_candidates_for_model_raises(self):
reg = FakeRegistry([FakeState("http://x:1", ["other"])])
session = FakeSession({})
rt = _runtime_with(reg, session)
with pytest.raises(RuntimeError, match="no live upstream"):
await rt.route("chat/completions", {"model": "m1", "messages": []}, stream=False)
@pytest.mark.asyncio
async def test_stream_relays_chunks(self):
reg = FakeRegistry([FakeState("http://s:1", ["m1"])])
chunks = [b"data: {\"a\":1}\n\n", b"data: [DONE]\n\n"]
session = FakeSession({"http://s:1": [FakeResponse(200, chunks=chunks)]})
rt = _runtime_with(reg, session)
gen = await rt.route("chat/completions", {"model": "m1", "messages": []}, stream=True)
received = [c async for c in gen]
assert received == chunks
# --- Rotation last-resort partitioning --------------------------------------
class TestPartitionLastResort:
def test_partition_splits_and_preserves_order(self):
from aisbf.handlers import _partition_last_resort
models = [
{"name": "a", "_last_resort": False},
{"name": "lr1", "_last_resort": True},
{"name": "b"}, # untagged == normal
{"name": "lr2", "_last_resort": True},
]
normal, last_resort = _partition_last_resort(models)
assert [m["name"] for m in normal] == ["a", "b"]
assert [m["name"] for m in last_resort] == ["lr1", "lr2"]
def test_no_last_resort(self):
from aisbf.handlers import _partition_last_resort
models = [{"name": "a"}, {"name": "b"}]
normal, last_resort = _partition_last_resort(models)
assert [m["name"] for m in normal] == ["a", "b"]
assert last_resort == []
def test_all_last_resort(self):
from aisbf.handlers import _partition_last_resort
models = [{"name": "a", "_last_resort": True}]
normal, last_resort = _partition_last_resort(models)
assert normal == []
assert [m["name"] for m in last_resort] == ["a"]
# --- Singleton --------------------------------------------------------------
class TestSingleton:
@pytest.mark.asyncio
async def test_get_runtime_is_singleton(self):
import aisbf.providers.shodan.runtime as rt_mod
rt_mod._runtime_singleton = None
a = await get_runtime({"shodan_mode": "public"})
b = await get_runtime({"shodan_mode": "api"}) # later config ignored
assert a is b
rt_mod._runtime_singleton = None # cleanup for other tests
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment