From 49d3359532ff1fe951c521bd173c95b9a8b47b22 Mon Sep 17 00:00:00 2001 From: Georgi Gaydarov Date: Sun, 31 May 2026 16:58:14 -0400 Subject: [PATCH 1/2] feat(llm): shared LLM egress (CROSS-1) + model-agnostic adapter (#4) [PR-1] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One egress point for every LLM call, with a 4-event hook bus so #2 (caching), #4 (model-identity), and #1 (Reflexion) attach without forking the egress — the dependency spine the whole upgrade set rides on. - core/llm/egress.py: LLMEgress (ordered local-first fallback chain, retry/backoff/cooldown = LiteLLM logic, no dep) + HookBus (isolated dispatch: a broken hook never breaks a call) + frozen LLMRequest - core/llm/backends.py: OpenAICompatLocalBackend + AnthropicCloudBackend via RAW httpx (NO `import anthropic` — keeps the air-gap invariant provable) + frozen LLMResult / ModelIdentity / LLMAttempt - core/llm/{errors,airgap,config,identity}.py: typed taxonomy + is_retryable; air-gap fail-closed (is_loopback; refuse cloud/non-loopback before a socket); env-driven immutable config (local-first; cloud dropped in air-gap); local weights-sha256 vs honest cloud identity-claim (never a faked weight hash) - core/llm/adapter.py: LLMAdapter facade delegating to the ONE egress (no 2nd chain) - tests/llm_egress_test.py: 28 tests (taxonomy, is_loopback, air-gap, immutability, wired fallback/retry/cooldown/fail-closed, hook isolation, httpx MockTransport backend over no sockets, single-egress anthropic scan, no-litellm). 28/28 green. Deferred to the cutover PR: the "no /v1/chat/completions outside core/llm" scan (HttpBackend.generate_config still owns its POST). All 6 existing suites stay green. --- core/llm/__init__.py | 73 +++++++ core/llm/adapter.py | 64 ++++++ core/llm/airgap.py | 69 ++++++ core/llm/backends.py | 291 +++++++++++++++++++++++++ core/llm/config.py | 114 ++++++++++ core/llm/egress.py | 225 +++++++++++++++++++ core/llm/errors.py | 128 +++++++++++ core/llm/identity.py | 86 ++++++++ tests/llm_egress_test.py | 454 +++++++++++++++++++++++++++++++++++++++ 9 files changed, 1504 insertions(+) create mode 100644 core/llm/__init__.py create mode 100644 core/llm/adapter.py create mode 100644 core/llm/airgap.py create mode 100644 core/llm/backends.py create mode 100644 core/llm/config.py create mode 100644 core/llm/egress.py create mode 100644 core/llm/errors.py create mode 100644 core/llm/identity.py create mode 100644 tests/llm_egress_test.py diff --git a/core/llm/__init__.py b/core/llm/__init__.py new file mode 100644 index 0000000..3f3a265 --- /dev/null +++ b/core/llm/__init__.py @@ -0,0 +1,73 @@ +"""aegis/core/llm/ — the single shared LLM egress (CROSS-1) + model-agnostic adapter (#4). + +The ONLY import surface other modules should use. The single-egress test asserts that no +module outside aegis/core/llm/ imports `anthropic` or POSTs to an LLM completion URL — so +everyone consumes the egress through these exports. +""" +from __future__ import annotations + +from .adapter import LLMAdapter +from .backends import ( + AnthropicCloudBackend, + Capability, + HashKind, + LLMAttempt, + LLMBackend, + LLMResult, + ModelIdentity, + OpenAICompatLocalBackend, + Provider, + Tier, + build_backend, +) +from .config import AdapterConfig, BackendSpec, capabilities_for +from .egress import HookBus, LLMEgress, LLMRequest +from .errors import ( + AdapterError, + AuthError, + ContextExceeded, + Permanent, + RateLimited, + Transient, + classify, + is_retryable, +) +from .identity import resolve_model_identity, unknown_identity + +__all__ = [ + # egress (CROSS-1) + "LLMEgress", + "LLMRequest", + "HookBus", + # adapter façade + "LLMAdapter", + # backends + "LLMBackend", + "OpenAICompatLocalBackend", + "AnthropicCloudBackend", + "build_backend", + # data structures + "LLMResult", + "LLMAttempt", + "ModelIdentity", + "Provider", + "Capability", + "HashKind", + "Tier", + # config + "AdapterConfig", + "BackendSpec", + "capabilities_for", + # errors + "AdapterError", + "AuthError", + "RateLimited", + "ContextExceeded", + "Transient", + "Permanent", + "classify", + "is_retryable", + # identity + "resolve_model_identity", + "unknown_identity", +] diff --git a/core/llm/adapter.py b/core/llm/adapter.py new file mode 100644 index 0000000..8633c75 --- /dev/null +++ b/core/llm/adapter.py @@ -0,0 +1,64 @@ +"""aegis/core/llm/adapter.py — back-compat façade over the single egress. + +Part of PR-1 (CROSS-1 shared LLM egress + #4 model-agnostic adapter). + +The fallback-chain logic lives in ONE place — egress.py (LLMEgress). LLMAdapter is a thin +keyword-args façade preserving the signature the legacy sync callers expect +(HttpBackend.generate_config, multivendor _llm_query), so there is exactly one chain +implementation to audit — not two. This is the anti-fragmentation invariant the whole +proposal set exists to enforce. +""" +from __future__ import annotations + +from .backends import LLMResult, Tier +from .egress import HookBus, LLMEgress, LLMRequest + + +class LLMAdapter: + """Keyword-args façade. Delegates the chain + retry/cooldown to LLMEgress.""" + + def __init__(self, egress: LLMEgress) -> None: + self._egress = egress + + @classmethod + def from_env(cls, hooks: HookBus | None = None) -> "LLMAdapter": + """Build config + backends from env (air-gap checks fire at backend construction).""" + return cls(LLMEgress.from_env(hooks=hooks)) + + async def complete( + self, + *, + system: str, + user: str, + max_tokens: int = 500, + temperature: float = 0.1, + tier: Tier = "cheap", + ) -> LLMResult: + return await self._egress.complete( + LLMRequest( + system=system, + user=user, + max_tokens=max_tokens, + temperature=temperature, + tier=tier, + ) + ) + + def complete_sync( + self, + *, + system: str, + user: str, + max_tokens: int = 500, + temperature: float = 0.1, + tier: Tier = "cheap", + ) -> LLMResult: + return self._egress.complete_sync( + LLMRequest( + system=system, + user=user, + max_tokens=max_tokens, + temperature=temperature, + tier=tier, + ) + ) diff --git a/core/llm/airgap.py b/core/llm/airgap.py new file mode 100644 index 0000000..c8cb818 --- /dev/null +++ b/core/llm/airgap.py @@ -0,0 +1,69 @@ +"""aegis/core/llm/airgap.py — air-gap fail-closed invariant (THE WEDGE, part 1). + +Part of PR-1 (CROSS-1 shared LLM egress + #4 model-agnostic adapter). + +A HARD code invariant that fails BEFORE a request is built — strictly stronger than +``network_mode: none`` at the container layer: + 1. refuses to construct the anthropic-cloud backend at all, + 2. refuses any non-loopback base URL, + 3. asserts the `anthropic` SDK was never imported into the process. + +Import-safe, fully type-hinted, no network — unit-testable on day one. No secrets touched. + +NOTE: AEGIS's cloud backend uses raw httpx (no `import anthropic`) precisely so that +invariant #3 holds even when the cloud path IS exercised in a non-air-gapped build. +""" +from __future__ import annotations + +import sys +from urllib.parse import urlparse + +_CLOUD_PROVIDER = "anthropic-cloud" + +# Loopback hosts allowed even in air-gap mode (in-perimeter only). +_LOOPBACK_HOSTS: frozenset[str] = frozenset({"localhost", "127.0.0.1", "::1", ""}) + + +def is_loopback(url: str) -> bool: + """True only for loopback / .local hosts. An empty host (relative URL) counts as + loopback because the existing http_backend posts to absolute in-perimeter URLs only; + a bare host means same-host.""" + host = (urlparse(url).hostname or "").lower() + if host in _LOOPBACK_HOSTS: + return True + if host.endswith(".local"): + return True + if host.startswith("127."): # 127.0.0.0/8 is all loopback + return True + return False + + +def assert_airgap_ok(provider: str, base_url: str) -> None: + """HARD invariant — raises RuntimeError BEFORE any socket opens. + + Call at backend *construction* time so an air-gap misconfiguration fails closed at + startup, not mid-pipeline. A failure here means the pipeline produces NO bundle — the + correct fail-closed behavior (AEGIS never seals evidence for a change a model could + not have produced). + """ + if provider == _CLOUD_PROVIDER: + raise RuntimeError( + "AEGIS_AIRGAP=1: cloud backend (anthropic-cloud) is forbidden in air-gap mode" + ) + if not is_loopback(base_url): + raise RuntimeError(f"AEGIS_AIRGAP=1: non-loopback egress forbidden: {base_url!r}") + if "anthropic" in sys.modules: + raise RuntimeError( + "AEGIS_AIRGAP=1: the anthropic SDK is loaded in-process — aborting " + "(in-process exposure violates the air-gap trust surface)" + ) + + +def assert_no_cloud_sdk_imported() -> None: + """Standalone check for the single-egress test and a startup guard: in air-gap mode + the cloud SDK must never be importable into the live process.""" + if "anthropic" in sys.modules: + raise RuntimeError( + "air-gap invariant: `anthropic` is present in sys.modules — the cloud SDK " + "must stay unimported (AEGIS uses raw httpx for the cloud path)" + ) diff --git a/core/llm/backends.py b/core/llm/backends.py new file mode 100644 index 0000000..12578c9 --- /dev/null +++ b/core/llm/backends.py @@ -0,0 +1,291 @@ +"""aegis/core/llm/backends.py — backend Protocol + the two concrete backends. + +Part of PR-1 (CROSS-1 shared LLM egress + #4 model-agnostic adapter). + +Two backends, no more (NOT a 100+-provider router): + - OpenAICompatLocalBackend : POST {base}/v1/chat/completions (vLLM/Ollama/llama.cpp/LM Studio) + - AnthropicCloudBackend : POST {base}/v1/messages (RAW httpx — NO SDK import) + +Why raw httpx for the cloud path: importing the `anthropic` SDK would put it in +sys.modules and break the air-gap invariant (airgap.assert_no_cloud_sdk_imported). A raw +POST keeps the trust surface minimal and the air-gap claim provable. + +The local backend's request body is BYTE-COMPATIBLE with what HttpBackend.generate_config() +posts today (http_backend.py:166-177): {"model","messages":[system,user],"temperature", +"stream":false} (+ max_tokens, which OpenAI-compat runtimes accept). + +All data structures (LLMResult / ModelIdentity / LLMAttempt) are frozen/immutable. The +backends accept an optional httpx transport so tests drive them with httpx.MockTransport +(no sockets). +""" +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable + +import httpx + +from .config import BackendSpec +from .errors import AdapterError, classify + +if TYPE_CHECKING: # avoid importing for runtime; type hints only + pass + +Provider = Literal["openai-compatible-local", "anthropic-cloud"] +Capability = Literal["chat", "json", "tools"] +HashKind = Literal["weights-sha256", "identity-claim"] +Tier = Literal["cheap", "deep"] + + +@dataclass(frozen=True) +class ModelIdentity: + """Goes verbatim into change.model_identity in the sealed bundle (THE WEDGE). + + LOCAL backend: model_hash is a real SHA-256 of the served weights file + (model_hash_kind="weights-sha256"). CLOUD backend: model_hash is None with + kind="identity-claim" — NEVER a faked weight hash for a hosted API. + """ + + provider: Provider + model: str + model_hash: str | None + model_hash_kind: HashKind + api_version: str | None + capabilities: tuple[Capability, ...] # tuple => immutable + hashable + resolved_at_utc: str # ISO-8601 + + def to_bundle_dict(self) -> dict[str, object]: + """Exact shape sealed into bundle['change']['model_identity'] (schema v1.1).""" + return { + "provider": self.provider, + "model": self.model, + "model_hash": self.model_hash, + "model_hash_kind": self.model_hash_kind, + "api_version": self.api_version, + "capabilities": list(self.capabilities), + "resolved_at_utc": self.resolved_at_utc, + } + + +@dataclass(frozen=True) +class LLMAttempt: + """One try in the fallback chain — recorded for the optional audit trail.""" + + provider: Provider + model: str + ok: bool + status: int | None + error_kind: str | None # name of the typed error, if any + latency_ms: int + + +@dataclass(frozen=True) +class LLMResult: + text: str + identity: ModelIdentity + attempts: tuple[LLMAttempt, ...] + total_latency_ms: int + + +@runtime_checkable +class LLMBackend(Protocol): + """The minimal surface the egress needs from any backend.""" + + provider: Provider + + async def complete( + self, *, system: str, user: str, max_tokens: int, temperature: float = 0.1 + ) -> tuple[str, ModelIdentity]: + """Single non-streamed completion. Raises a typed AdapterError on failure.""" + ... + + def identity(self) -> ModelIdentity: + """Resolve identity WITHOUT making a request (used for air-gap pre-checks).""" + ... + + +def _retry_after_seconds(resp: httpx.Response) -> float | None: + """Prefer the HTTP Retry-After header (seconds form).""" + raw = resp.headers.get("retry-after") + if raw is None: + return None + try: + return float(raw) + except ValueError: + return None + + +class OpenAICompatLocalBackend: + """POST {base}/v1/chat/completions — the universal local contract. + + Construction runs the air-gap pre-check (loopback enforced when AEGIS_AIRGAP=1). + """ + + provider: Provider = "openai-compatible-local" + + def __init__( + self, + spec: BackendSpec, + timeout_s: float = 60.0, + transport: httpx.BaseTransport | None = None, + ) -> None: + from .airgap import assert_airgap_ok + + if os.environ.get("AEGIS_AIRGAP", "") == "1": + assert_airgap_ok(self.provider, spec.base_url) + self._spec = spec + self._timeout_s = timeout_s + self._transport = transport # injected only in tests (httpx.MockTransport) + + def identity(self) -> ModelIdentity: + from .identity import resolve_model_identity + + return resolve_model_identity(self._spec) + + async def complete( + self, *, system: str, user: str, max_tokens: int, temperature: float = 0.1 + ) -> tuple[str, ModelIdentity]: + body = { + "model": self._spec.model, + "messages": [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ], + "temperature": temperature, + "stream": False, + "max_tokens": max_tokens, + } + url = f"{self._spec.base_url}/v1/chat/completions" + try: + async with httpx.AsyncClient( + transport=self._transport, timeout=self._timeout_s + ) as client: + resp = await client.post(url, json=body) + except httpx.RequestError as exc: # conn reset / DNS / timeout -> Transient + raise classify(None, str(exc), self.provider) from exc + + if resp.status_code >= 400: + err = classify(resp.status_code, resp.text, self.provider) + from .errors import RateLimited + + if isinstance(err, RateLimited) and err.retry_after_s is None: + hdr = _retry_after_seconds(resp) + if hdr is not None: + err = RateLimited( + message=err.message, + provider=err.provider, + status=err.status, + retry_after_s=hdr, + ) + raise err + + try: + text = resp.json()["choices"][0]["message"]["content"] + except (KeyError, IndexError, ValueError) as exc: + from .errors import Permanent + + raise Permanent( + message="malformed OpenAI-compat completion response", + provider=self.provider, + status=resp.status_code, + ) from exc + return text, self.identity() + + +class AnthropicCloudBackend: + """POST {base}/v1/messages via RAW httpx — air-gap forbidden, no SDK import. + + Reads ANTHROPIC_API_KEY at call time; the key is NEVER stored or logged. + """ + + provider: Provider = "anthropic-cloud" + + def __init__( + self, + spec: BackendSpec, + timeout_s: float = 60.0, + transport: httpx.BaseTransport | None = None, + ) -> None: + from .airgap import assert_airgap_ok + + # In air-gap mode this RAISES — the cloud backend must never construct. + if os.environ.get("AEGIS_AIRGAP", "") == "1": + assert_airgap_ok(self.provider, spec.base_url) + self._spec = spec + self._timeout_s = timeout_s + self._transport = transport + + def identity(self) -> ModelIdentity: + from .identity import resolve_model_identity + + return resolve_model_identity(self._spec) + + async def complete( + self, *, system: str, user: str, max_tokens: int, temperature: float = 0.1 + ) -> tuple[str, ModelIdentity]: + from .errors import AuthError, Permanent + + api_key = os.environ.get("ANTHROPIC_API_KEY") + if not api_key: + raise AuthError( + message="ANTHROPIC_API_KEY is not set", provider=self.provider, status=401 + ) + headers = { + "x-api-key": api_key, + "anthropic-version": self._spec.api_version or "2023-06-01", + "content-type": "application/json", + } + body = { + "model": self._spec.model, + "max_tokens": max_tokens, + "temperature": temperature, + "system": system, + "messages": [{"role": "user", "content": user}], + } + url = f"{self._spec.base_url}/v1/messages" + try: + async with httpx.AsyncClient( + transport=self._transport, timeout=self._timeout_s + ) as client: + resp = await client.post(url, json=body, headers=headers) + except httpx.RequestError as exc: + raise classify(None, str(exc), self.provider) from exc + + if resp.status_code >= 400: + err = classify(resp.status_code, resp.text, self.provider) + from .errors import RateLimited + + if isinstance(err, RateLimited) and err.retry_after_s is None: + hdr = _retry_after_seconds(resp) + if hdr is not None: + err = RateLimited( + message=err.message, + provider=err.provider, + status=err.status, + retry_after_s=hdr, + ) + raise err + + try: + text = resp.json()["content"][0]["text"] + except (KeyError, IndexError, ValueError) as exc: + raise Permanent( + message="malformed Anthropic messages response", + provider=self.provider, + status=resp.status_code, + ) from exc + return text, self.identity() + + +def build_backend( + spec: BackendSpec, + timeout_s: float = 60.0, + transport: httpx.BaseTransport | None = None, +) -> LLMBackend: + """Factory: construct the right backend for a spec (air-gap checks fire here).""" + if spec.provider == "openai-compatible-local": + return OpenAICompatLocalBackend(spec, timeout_s=timeout_s, transport=transport) + if spec.provider == "anthropic-cloud": + return AnthropicCloudBackend(spec, timeout_s=timeout_s, transport=transport) + raise AdapterError(message=f"unknown provider {spec.provider!r}", provider=spec.provider) diff --git a/core/llm/config.py b/core/llm/config.py new file mode 100644 index 0000000..fc4a989 --- /dev/null +++ b/core/llm/config.py @@ -0,0 +1,114 @@ +"""aegis/core/llm/config.py — immutable, env-driven adapter configuration. + +Part of PR-1 (CROSS-1 shared LLM egress + #4 model-agnostic adapter). + +Cost-aware defaults: the fallback chain is LOCAL-FIRST. Cloud is fallback only, and is +DROPPED from the chain entirely in air-gap mode (defense in depth — also refused at +construction by airgap.py). + +Env vars read: + AEGIS_AIRGAP "1" => air-gap mode (drop cloud, refuse non-loopback) + AEGIS_LLM_URL local OpenAI-compat base (default http://localhost:11434) + AEGIS_LLM_MODEL local model id (default ai/qwen3:latest) + AEGIS_LLM_WEIGHTS path to local weights file for sha256 attestation (optional) + MODEL_RUNNER_URL alt local base (Docker Model Runner :12434), optional + ANTHROPIC_API_KEY presence enables the cloud fallback (NEVER logged/echoed) + ANTHROPIC_MODEL cloud model id (default claude-haiku-4-5-20251001) + ANTHROPIC_VERSION cloud api version (default 2023-06-01) + +Import-safe, fully type-hinted. `from_env()` is side-effect-free (reads env, builds frozen +specs). NO secret material is stored on the config — only a derived "cloud available" bool. +""" +from __future__ import annotations + +import os +from dataclasses import dataclass + +# Static capability map keyed by model string. Manual refresh only — auto-probing a model +# would risk egress in air-gap mode. Conservative default ('chat',) for unknown models. +_CAPABILITIES: dict[str, tuple[str, ...]] = { + "ai/qwen3:latest": ("chat", "json"), + "gemma4:latest": ("chat",), + "claude-haiku-4-5-20251001": ("chat", "json", "tools"), +} + +_DEFAULT_LOCAL_URL = "http://localhost:11434" +_DEFAULT_LOCAL_MODEL = "ai/qwen3:latest" +_DEFAULT_CLOUD_MODEL = "claude-haiku-4-5-20251001" +_DEFAULT_CLOUD_VERSION = "2023-06-01" +_CLOUD_BASE_URL = "https://api.anthropic.com" + + +@dataclass(frozen=True) +class BackendSpec: + """One ordered link in the fallback chain. Immutable.""" + + provider: str # "openai-compatible-local" | "anthropic-cloud" + base_url: str + model: str + api_version: str | None = None + weights_path: str | None = None # local only; enables weights-sha256 attestation + + +@dataclass(frozen=True) +class AdapterConfig: + """Immutable adapter configuration. Carries NO secret material.""" + + chain: tuple[BackendSpec, ...] # ordered fallback: local first (cost-aware) + airgap: bool + cloud_available: bool # derived from ANTHROPIC_API_KEY presence; never the key itself + max_retries: int = 2 + base_backoff_s: float = 0.5 + cooldown_s: float = 60.0 # LiteLLM's failing-provider cooldown idea + timeout_s: float = 60.0 + + @staticmethod + def from_env() -> "AdapterConfig": + """Build an immutable config from the environment. + + In air-gap mode the anthropic-cloud spec is DROPPED from the chain here (defense + in depth) and ALSO refused at construction by airgap.assert_airgap_ok(). + """ + airgap = os.environ.get("AEGIS_AIRGAP", "") == "1" + + local_url = ( + os.environ.get("AEGIS_LLM_URL") + or os.environ.get("MODEL_RUNNER_URL") + or _DEFAULT_LOCAL_URL + ) + local_model = os.environ.get("AEGIS_LLM_MODEL", _DEFAULT_LOCAL_MODEL) + weights_path = os.environ.get("AEGIS_LLM_WEIGHTS") or None + + local_spec = BackendSpec( + provider="openai-compatible-local", + base_url=local_url.rstrip("/"), + model=local_model, + api_version=None, + weights_path=weights_path, + ) + + chain: list[BackendSpec] = [local_spec] + + # presence-only check — the key is NEVER stored on the config object + cloud_available = bool(os.environ.get("ANTHROPIC_API_KEY")) + if cloud_available and not airgap: + chain.append( + BackendSpec( + provider="anthropic-cloud", + base_url=_CLOUD_BASE_URL, + model=os.environ.get("ANTHROPIC_MODEL", _DEFAULT_CLOUD_MODEL), + api_version=os.environ.get("ANTHROPIC_VERSION", _DEFAULT_CLOUD_VERSION), + weights_path=None, + ) + ) + + return AdapterConfig( + chain=tuple(chain), + airgap=airgap, + cloud_available=cloud_available, + ) + + +def capabilities_for(model: str) -> tuple[str, ...]: + """Conservative default ('chat',) for any model not in the static map.""" + return _CAPABILITIES.get(model, ("chat",)) diff --git a/core/llm/egress.py b/core/llm/egress.py new file mode 100644 index 0000000..9de05e0 --- /dev/null +++ b/core/llm/egress.py @@ -0,0 +1,225 @@ +"""aegis/core/llm/egress.py — the single shared LLM egress + hook bus (CROSS-1). + +ONE place every LLM call leaves the process. An ordered fallback chain (local-first, +cost-aware) with retry / exponential backoff / per-provider cooldown — LiteLLM's decision +*logic* with none of its dependency tree — wrapped in a 4-event hook bus so #2 +(cache_control), #4 (model-identity), and #1 (Reflexion) attach WITHOUT forking the egress. + +The core knows nothing about caching, identity schemas, or reflexion — it only emits events. +This is the load-bearing decoupling: the chain stays small and auditable; features are tenants. + +Immutability: LLMRequest / LLMResult / ModelIdentity / LLMAttempt are all frozen. The ONE +mutable bit is the per-provider cooldown map (process-local, never part of a returned result). +""" +from __future__ import annotations + +import asyncio +import logging +import time +from collections import defaultdict +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from typing import Literal + +from .backends import LLMAttempt, LLMBackend, LLMResult, ModelIdentity, build_backend +from .config import AdapterConfig +from .errors import AdapterError, RateLimited, is_retryable + +Tier = Literal["cheap", "deep"] +HookEvent = Literal["before_request", "build_payload", "after_response", "on_result"] + +_log = logging.getLogger("aegis.llm.egress") + + +@dataclass(frozen=True) +class LLMRequest: + """Immutable call envelope. trace_id lets #1 correlate generate→verify→repair.""" + + system: str + user: str + max_tokens: int = 500 + temperature: float = 0.1 + tier: Tier = "cheap" + trace_id: str | None = None + # opaque per-call hints (e.g. #2's prefix_kind); the core never interprets these. + hints: Mapping[str, object] = field(default_factory=dict) + + +class HookBus: + """Deterministic, in-process, synchronous hook dispatch — the ENTIRE extension surface. + + Transform hooks (before_request / build_payload) may return a replacement value; + observe hooks (after_response / on_result) return None. A hook that raises is ISOLATED: + the egress logs and continues. A broken cache/identity hook must NEVER break an LLM call. + """ + + def __init__(self) -> None: + self._hooks: dict[HookEvent, list[Callable]] = defaultdict(list) + + def register(self, event: HookEvent, fn: Callable) -> None: + self._hooks[event].append(fn) + + def transform(self, event: HookEvent, value, *args): + for fn in self._hooks[event]: + try: + out = fn(value, *args) + if out is not None: + value = out + except Exception as exc: # noqa: BLE001 — isolation is the point + _log_hook_error(event, fn, exc) + return value + + def observe(self, event: HookEvent, *args) -> None: + for fn in self._hooks[event]: + try: + fn(*args) + except Exception as exc: # noqa: BLE001 + _log_hook_error(event, fn, exc) + + +class LLMEgress: + """The single egress. Ordered fallback chain + retry/cooldown + hook bus.""" + + def __init__( + self, + config: AdapterConfig, + backends: tuple[LLMBackend, ...], + hooks: HookBus | None = None, + ) -> None: + self._cfg = config + self._backends = backends # already air-gap-checked at construction + self._hooks = hooks or HookBus() + self._cooldown_until: dict[str, float] = {} # the ONE mutable bit (process-local) + + @classmethod + def from_env(cls, hooks: HookBus | None = None) -> "LLMEgress": + cfg = AdapterConfig.from_env() # drops the cloud backend when AEGIS_AIRGAP=1 + backends = tuple(build_backend(s, timeout_s=cfg.timeout_s) for s in cfg.chain) + return cls(cfg, backends, hooks) + + def on(self, event: HookEvent, fn: Callable) -> "LLMEgress": + """Register a hook. #2/#4/#1 call this once at wire-up. Fluent for chaining. + + NOTE: ``build_payload`` is registerable today and is fired by #2 once the backends + accept a payload hook; PR-1 fires before_request / after_response / on_result. + """ + self._hooks.register(event, fn) + return self + + async def complete(self, req: LLMRequest) -> LLMResult: + req = self._hooks.transform("before_request", req) # ← #2 / #1 hook point + ordered = self._ordered_backends(req.tier) + attempts: list[LLMAttempt] = [] + run_t0 = time.perf_counter() + + for backend in ordered: + if self._is_cooled_down(backend.provider): + continue + outcome = await self._try_backend(backend, req, attempts) + if outcome is not None: + text, identity = outcome + self._hooks.observe("after_response", backend.provider, identity) # ← #4 + total_ms = int((time.perf_counter() - run_t0) * 1000) + result = LLMResult( + text=text, + identity=identity, + attempts=tuple(attempts), + total_latency_ms=total_ms, + ) + self._hooks.observe("on_result", result) # ← #1 / #4 hook point + return result + + total_ms = int((time.perf_counter() - run_t0) * 1000) + raise AdapterError( + message=f"all backends failed ({len(attempts)} attempts, {total_ms}ms)", + provider="egress", + status=None, + ) + + def complete_sync(self, req: LLMRequest) -> LLMResult: + """Sync bridge for legacy callers. Loop-safe: if already inside an event loop + (Flask under an async server) offload to a worker thread instead of asyncio.run + (which would raise 'event loop already running').""" + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(self.complete(req)) + import concurrent.futures + + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(lambda: asyncio.run(self.complete(req))).result() + + # ── internals ──────────────────────────────────────────────────────── + def _ordered_backends(self, tier: Tier) -> tuple[LLMBackend, ...]: + """Cost-aware ordering. 'cheap' => chain order (local first). 'deep' => prefer a + cloud backend first when present + not air-gapped (in air-gap 'deep' stays local).""" + if tier == "deep" and not self._cfg.airgap: + cloud = tuple(b for b in self._backends if b.provider == "anthropic-cloud") + rest = tuple(b for b in self._backends if b.provider != "anthropic-cloud") + return cloud + rest + return self._backends + + def _is_cooled_down(self, provider: str) -> bool: + until = self._cooldown_until.get(provider) + return until is not None and time.monotonic() < until + + def _cool_down(self, provider: str) -> None: + self._cooldown_until[provider] = time.monotonic() + self._cfg.cooldown_s + + async def _try_backend( + self, backend: LLMBackend, req: LLMRequest, attempts: list[LLMAttempt] + ) -> tuple[str, ModelIdentity] | None: + """Retry one backend ≤ max_retries. Append an LLMAttempt per try. Returns + (text, identity) on success, else None (skip to the next backend in the chain).""" + for attempt in range(self._cfg.max_retries + 1): + t0 = time.perf_counter() + try: + text, identity = await backend.complete( + system=req.system, + user=req.user, + max_tokens=req.max_tokens, + temperature=req.temperature, + ) + attempts.append( + LLMAttempt( + provider=backend.provider, + model=identity.model, + ok=True, + status=200, + error_kind=None, + latency_ms=int((time.perf_counter() - t0) * 1000), + ) + ) + return text, identity + except AdapterError as err: + attempts.append( + LLMAttempt( + provider=backend.provider, + model="unknown", + ok=False, + status=err.status, + error_kind=type(err).__name__, + latency_ms=int((time.perf_counter() - t0) * 1000), + ) + ) + if not is_retryable(err) or attempt >= self._cfg.max_retries: + if isinstance(err, RateLimited): + self._cool_down(backend.provider) + return None + await asyncio.sleep(self._backoff_s(attempt, err)) + return None + + def _backoff_s(self, attempt: int, err: AdapterError) -> float: + """Exponential backoff; honor Retry-After on 429 when present.""" + if isinstance(err, RateLimited) and err.retry_after_s is not None: + return err.retry_after_s + return self._cfg.base_backoff_s * (2**attempt) + + +def _log_hook_error(event: HookEvent, fn: Callable, exc: Exception) -> None: + _log.warning( + "hook %s on %s failed (isolated, call continues): %s", + getattr(fn, "__name__", fn), + event, + exc, + ) diff --git a/core/llm/errors.py b/core/llm/errors.py new file mode 100644 index 0000000..03956c6 --- /dev/null +++ b/core/llm/errors.py @@ -0,0 +1,128 @@ +"""aegis/core/llm/errors.py — typed error taxonomy + retry decision-logic. + +Part of PR-1 (CROSS-1 shared LLM egress + #4 model-agnostic adapter). + +We port LiteLLM's *ideas* — a small typed exception hierarchy, an is-retryable +predicate, and Retry-After honoring — WITHOUT its 100+-provider dependency tree. +Scoped to exactly the two backends AEGIS ships: anthropic-cloud + openai-compat-local. + +Import-safe, fully type-hinted, no network. `classify()` is real so the taxonomy is +unit-testable on day one. +""" +from __future__ import annotations + +import re +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AdapterError(Exception): + """Base typed error. Shape mirrors LiteLLM's exception (message/provider/status) + minus the dependency. Frozen => immutable per coding-style rule.""" + + message: str + provider: str + status: int | None = None + + def __str__(self) -> str: # keep secrets out of logs — only structural fields + return f"{type(self).__name__}[{self.provider} status={self.status}]: {self.message}" + + +@dataclass(frozen=True) +class AuthError(AdapterError): + """401 / 403 — bad or missing key. NEVER retried. Skip to next backend.""" + + +@dataclass(frozen=True) +class RateLimited(AdapterError): + """429 — back off, honor Retry-After, then cool the provider.""" + + retry_after_s: float | None = None + + +@dataclass(frozen=True) +class ContextExceeded(AdapterError): + """400 / 413 with context-window phrasing. NEVER retried (same prompt = same fail).""" + + +@dataclass(frozen=True) +class Transient(AdapterError): + """408 / 5xx / connection reset / timeout. RETRYABLE with backoff.""" + + +@dataclass(frozen=True) +class Permanent(AdapterError): + """400 / 404 / 422 (non-context). NEVER retried.""" + + +# Phrases that mark a 4xx body as a context-window failure. Deliberately small — +# scoped to the two runtimes we ship (Anthropic + OpenAI-compat), NOT a cross-provider +# string zoo. +_CONTEXT_PHRASES: tuple[str, ...] = ( + "context window", + "maximum context length", + "too many tokens", + "context_length_exceeded", + "prompt is too long", +) + + +def classify(status: int | None, body: str, provider: str) -> AdapterError: + """Map a raw HTTP status + response body into the small internal taxonomy. + + Fail-safe: an unknown 5xx / connection error is treated as retryable Transient. + 401/403 -> AuthError (never retry) + 429 -> RateLimited (retry, honor Retry-After) + 400/413 + context phrasing -> ContextExceeded (never retry) + other 4xx -> Permanent (never retry) + 408/5xx/None(conn error) -> Transient (retry) + """ + low = (body or "").lower() + + if status in (401, 403): + return AuthError(message="authentication failed", provider=provider, status=status) + + if status == 429: + return RateLimited( + message="rate limited", + provider=provider, + status=status, + retry_after_s=_parse_retry_after(body), + ) + + if status in (400, 413) and any(p in low for p in _CONTEXT_PHRASES): + return ContextExceeded( + message="context window exceeded", provider=provider, status=status + ) + + if status is not None and 400 <= status < 500: + return Permanent( + message=f"permanent client error {status}", provider=provider, status=status + ) + + # status is None (connection reset / DNS / timeout) OR a 5xx OR 408 -> retryable. + return Transient( + message=f"transient error {status if status is not None else 'conn'}", + provider=provider, + status=status, + ) + + +def is_retryable(err: AdapterError) -> bool: + """5-line decision: retry Transient + RateLimited; never the rest.""" + return isinstance(err, (Transient, RateLimited)) + + +def _parse_retry_after(body: str) -> float | None: + """Best-effort Retry-After extraction from a JSON-ish error body. + + NOTE: prefer the HTTP ``Retry-After`` *header* when the backend surfaces it; the + backends pass it in via the body fallback only when the header is absent. + """ + m = re.search(r'retry[_-]?after["\s:=]+(\d+(?:\.\d+)?)', (body or "").lower()) + if m: + try: + return float(m.group(1)) + except ValueError: + return None + return None diff --git a/core/llm/identity.py b/core/llm/identity.py new file mode 100644 index 0000000..700c6d7 --- /dev/null +++ b/core/llm/identity.py @@ -0,0 +1,86 @@ +"""aegis/core/llm/identity.py — model identity resolution (THE WEDGE, part 2). + +Part of PR-1 (CROSS-1 shared LLM egress + #4 model-agnostic adapter). + +LOCAL backend => a REAL sha256 of the served weights file ("weights-sha256"). +CLOUD backend => an honest "identity-claim" with model_hash=None (a weight hash for a + hosted API is physically impossible; we NEVER fake one). + +Multi-GB GGUF/safetensors files are hashed ONCE and cached by (path, mtime, size) so +generate_config() does not re-hash on every inference. + +Import-safe, fully type-hinted, no network. +""" +from __future__ import annotations + +import hashlib +import os +from datetime import datetime, timezone + +from .backends import ModelIdentity +from .config import BackendSpec, capabilities_for + +# (path, mtime, size) -> sha256 hex. Auto-invalidates on file change. +_HASH_CACHE: dict[tuple[str, float, int], str] = {} + +_LOCAL_PROVIDER = "openai-compatible-local" +_CHUNK = 1 << 20 # 1 MiB + + +def _hash_local_weights(path: str) -> str | None: + """SHA-256 of the served weights file, cached by (path, mtime, size). Returns None + if the path is empty/missing/unreadable — identity then downgrades to an honest + 'identity-claim' rather than failing the whole run.""" + if not path or not os.path.exists(path): + return None + try: + st = os.stat(path) + except OSError: + return None + key = (path, st.st_mtime, st.st_size) + cached = _HASH_CACHE.get(key) + if cached is not None: + return cached + h = hashlib.sha256() + try: + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(_CHUNK), b""): + h.update(chunk) + except OSError: + return None + digest = h.hexdigest() + _HASH_CACHE[key] = digest + return digest + + +def resolve_model_identity(spec: BackendSpec) -> ModelIdentity: + """Local => real weights-sha256 attestation when the weights file is present. + Cloud => identity-claim with model_hash=None (never faked).""" + is_local = spec.provider == _LOCAL_PROVIDER + weight_hash = _hash_local_weights(spec.weights_path or "") if is_local else None + return ModelIdentity( + provider=spec.provider, # type: ignore[arg-type] + model=spec.model, + model_hash=weight_hash, + model_hash_kind="weights-sha256" if weight_hash else "identity-claim", + api_version=spec.api_version, + capabilities=capabilities_for(spec.model), # type: ignore[arg-type] + resolved_at_utc=datetime.now(timezone.utc).isoformat(), + ) + + +def unknown_identity(created_utc: str) -> dict[str, object]: + """_UNKNOWN_IDENTITY for config_import runs (no LLM touched the change). + + Keeps every bundle schema-valid (v1.1 requires change.model_identity) and honest — + an operator-supplied config never claims a model produced it. + """ + return { + "provider": "none", + "model": "operator-supplied", + "model_hash": None, + "model_hash_kind": "identity-claim", + "api_version": None, + "capabilities": [], + "resolved_at_utc": created_utc, + } diff --git a/tests/llm_egress_test.py b/tests/llm_egress_test.py new file mode 100644 index 0000000..7594754 --- /dev/null +++ b/tests/llm_egress_test.py @@ -0,0 +1,454 @@ +"""aegis/tests/llm_egress_test.py — PR-1 single-egress (CROSS-1) + adapter (#4) tests. + +Mirrors the repo's 6-suite house style: pytest-discoverable `test_*` functions PLUS a +dependency-light `__main__` runner (CI runs `python3 -m aegis.tests.llm_egress_test`). +No network: the wired-path tests drive fake in-process backends and httpx.MockTransport. + +PR-1 gate (per the roadmap): unit taxonomy · is_loopback · immutability · the wired +fallback/retry/cooldown chain · hook-bus isolation · the air-gap fail-closed invariant · +no-litellm · the no-foreign-`anthropic`-import single-egress scan. + +DEFERRED to a later PR (documented, see test_single_egress_completion_post_DEFERRED): +the "no /v1/chat/completions outside core/llm" scan — HttpBackend.generate_config still +owns its POST until the cutover PR; enabling it now would be a false RED. + +Run: python3 -m aegis.tests.llm_egress_test +""" +from __future__ import annotations + +import ast +import asyncio +import dataclasses +import hashlib +import os +import sys +import tempfile +from pathlib import Path + +import httpx + +from aegis.core.llm import ( + AdapterConfig, + AdapterError, + AuthError, + BackendSpec, + ContextExceeded, + HookBus, + LLMEgress, + LLMRequest, + ModelIdentity, + OpenAICompatLocalBackend, + Permanent, + RateLimited, + Transient, + classify, + is_retryable, + resolve_model_identity, + unknown_identity, +) +from aegis.core.llm.airgap import assert_airgap_ok, is_loopback + + +# ── test doubles ────────────────────────────────────────────────────────────── + +def _ident(provider: str = "openai-compatible-local", model: str = "m") -> ModelIdentity: + return ModelIdentity( + provider=provider, # type: ignore[arg-type] + model=model, + model_hash=None, + model_hash_kind="identity-claim", + api_version=None, + capabilities=("chat",), + resolved_at_utc="2026-05-31T00:00:00+00:00", + ) + + +class _FakeBackend: + """In-process backend driven by a script of ('ok', text) | ('err', AdapterError).""" + + def __init__(self, provider: str, model: str = "m", script=None) -> None: + self.provider = provider + self._model = model + self._script = list(script or []) + self.calls = 0 + + def identity(self) -> ModelIdentity: + return _ident(self.provider, self._model) + + async def complete(self, *, system, user, max_tokens, temperature=0.1): + self.calls += 1 + kind, payload = self._script.pop(0) if self._script else ("ok", user) + if kind == "ok": + return payload, self.identity() + raise payload + + +def _cfg(max_retries: int = 2, base_backoff_s: float = 0.0, cooldown_s: float = 60.0) -> AdapterConfig: + return AdapterConfig( + chain=(), + airgap=False, + cloud_available=False, + max_retries=max_retries, + base_backoff_s=base_backoff_s, + cooldown_s=cooldown_s, + ) + + +def _egress(backends, hooks=None, **cfg_kw) -> LLMEgress: + return LLMEgress(_cfg(**cfg_kw), tuple(backends), hooks) + + +def _req(user: str = "hello") -> LLMRequest: + return LLMRequest(system="sys", user=user, max_tokens=64) + + +# ── (gate) error taxonomy + is_retryable ─────────────────────────────────────── + +def test_is_retryable_truth_table() -> None: + assert is_retryable(Transient("x", "p")) + assert is_retryable(RateLimited("x", "p")) + assert not is_retryable(AuthError("x", "p")) + assert not is_retryable(ContextExceeded("x", "p")) + assert not is_retryable(Permanent("x", "p")) + + +def test_classify_routing() -> None: + assert isinstance(classify(401, "", "cloud"), AuthError) + assert isinstance(classify(403, "", "cloud"), AuthError) + rl = classify(429, "retry_after: 5", "cloud") + assert isinstance(rl, RateLimited) and rl.retry_after_s == 5.0 + assert isinstance(classify(400, "maximum context length", "l"), ContextExceeded) + assert isinstance(classify(413, "prompt is too long", "l"), ContextExceeded) + assert isinstance(classify(422, "bad", "l"), Permanent) + assert isinstance(classify(404, "nope", "l"), Permanent) + assert isinstance(classify(503, "", "l"), Transient) + assert isinstance(classify(500, "", "l"), Transient) + assert isinstance(classify(None, "conn reset", "l"), Transient) # fail-safe + + +def test_adapter_error_str_hides_body() -> None: + e = AuthError(message="authentication failed", provider="cloud", status=401) + assert "cloud" in str(e) and "401" in str(e) + + +# ── (gate) is_loopback / air-gap fail-closed ─────────────────────────────────── + +def test_is_loopback() -> None: + assert is_loopback("http://127.0.0.1:11434") + assert is_loopback("http://localhost:11434") + assert is_loopback("http://127.5.5.5:1") + assert is_loopback("http://box.local") + assert not is_loopback("https://api.anthropic.com") + assert not is_loopback("http://10.0.0.5:11434") + + +def test_airgap_refuses_cloud_construction() -> None: + raised = False + try: + assert_airgap_ok("anthropic-cloud", "https://api.anthropic.com") + except RuntimeError: + raised = True + assert raised, "cloud backend must be forbidden in air-gap mode" + + +def test_airgap_refuses_non_loopback() -> None: + raised = False + try: + assert_airgap_ok("openai-compatible-local", "https://evil.example.com") + except RuntimeError: + raised = True + assert raised, "non-loopback egress must be forbidden in air-gap mode" + + +def test_airgap_allows_loopback_local() -> None: + assert_airgap_ok("openai-compatible-local", "http://127.0.0.1:11434") # no raise + + +# ── (gate) immutability ──────────────────────────────────────────────────────── + +def test_frozen_dataclasses_reject_mutation() -> None: + req = _req() + ident = _ident() + for obj, attr, val in ((req, "user", "x"), (ident, "model", "y")): + try: + setattr(obj, attr, val) + raise AssertionError(f"{type(obj).__name__}.{attr} mutated — must be frozen") + except dataclasses.FrozenInstanceError: + pass + + +# ── model identity: local weights-sha256 vs cloud identity-claim ─────────────── + +def test_local_identity_carries_real_weights_sha256() -> None: + with tempfile.NamedTemporaryFile("wb", suffix=".gguf", delete=False) as f: + f.write(b"fake-weights-payload-for-test") + wpath = f.name + try: + expected = hashlib.sha256(b"fake-weights-payload-for-test").hexdigest() + spec = BackendSpec( + provider="openai-compatible-local", + base_url="http://127.0.0.1:11434", + model="ai/qwen3:latest", + weights_path=wpath, + ) + ident = resolve_model_identity(spec) + assert ident.model_hash_kind == "weights-sha256" + assert ident.model_hash == expected + assert resolve_model_identity(spec).model_hash == expected # cache hit + finally: + os.unlink(wpath) + + +def test_cloud_identity_is_claim_with_null_hash() -> None: + spec = BackendSpec( + provider="anthropic-cloud", + base_url="https://api.anthropic.com", + model="claude-haiku-4-5-20251001", + api_version="2023-06-01", + ) + ident = resolve_model_identity(spec) + assert ident.model_hash is None + assert ident.model_hash_kind == "identity-claim" + assert ident.api_version == "2023-06-01" + + +def test_missing_weights_downgrades_to_identity_claim() -> None: + spec = BackendSpec( + provider="openai-compatible-local", + base_url="http://127.0.0.1:11434", + model="ai/qwen3:latest", + weights_path="/nonexistent/path/model.gguf", + ) + ident = resolve_model_identity(spec) + assert ident.model_hash is None + assert ident.model_hash_kind == "identity-claim" + + +def test_unknown_identity_shape() -> None: + ui = unknown_identity("2026-05-31T10:30:00+00:00") + assert set(ui) == { + "provider", "model", "model_hash", "model_hash_kind", + "api_version", "capabilities", "resolved_at_utc", + } + assert ui["provider"] == "none" and ui["model_hash"] is None + + +# ── (gate) the WIRED chain: success / fallback / retry / cooldown / fail-closed ─ + +def test_chain_returns_on_first_success() -> None: + b = _FakeBackend("openai-compatible-local", script=[("ok", "ANSWER")]) + res = asyncio.run(_egress([b]).complete(_req())) + assert res.text == "ANSWER" + assert len(res.attempts) == 1 and res.attempts[0].ok + assert res.identity.provider == "openai-compatible-local" + assert b.calls == 1 + + +def test_chain_falls_back_on_non_retryable() -> None: + b1 = _FakeBackend("openai-compatible-local", script=[("err", AuthError("no key", "local", 401))]) + b2 = _FakeBackend("anthropic-cloud", script=[("ok", "FROM_CLOUD")]) + res = asyncio.run(_egress([b1, b2]).complete(_req())) + assert res.text == "FROM_CLOUD" + assert b1.calls == 1, "non-retryable error must NOT retry the same backend" + assert [a.ok for a in res.attempts] == [False, True] + + +def test_chain_retries_transient_then_succeeds() -> None: + b = _FakeBackend( + "openai-compatible-local", + script=[("err", Transient("flaky", "local", 503)), ("ok", "RECOVERED")], + ) + res = asyncio.run(_egress([b], max_retries=2).complete(_req())) + assert res.text == "RECOVERED" + assert b.calls == 2, "transient error must retry the same backend" + + +def test_ratelimited_cools_provider_and_is_skipped_next_call() -> None: + b1 = _FakeBackend( + "openai-compatible-local", + script=[("err", RateLimited("429", "local", 429)), ("err", RateLimited("429", "local", 429))], + ) + b2 = _FakeBackend("anthropic-cloud", script=[("ok", "A"), ("ok", "B")]) + eg = _egress([b1, b2], max_retries=1, cooldown_s=60.0) + r1 = asyncio.run(eg.complete(_req())) + assert r1.text == "A" and b1.calls == 2 # tried twice (attempt 0 + retry), then cooled + r2 = asyncio.run(eg.complete(_req())) + assert r2.text == "B" and b1.calls == 2, "cooled-down provider must be skipped" + + +def test_chain_exhausted_raises_fail_closed() -> None: + b = _FakeBackend("openai-compatible-local", script=[("err", Permanent("nope", "local", 422))]) + raised = False + try: + asyncio.run(_egress([b]).complete(_req())) + except AdapterError: + raised = True + assert raised, "no fake result on full failure — must fail closed" + + +def test_deep_tier_prefers_cloud_first() -> None: + local = _FakeBackend("openai-compatible-local", script=[("ok", "LOCAL")]) + cloud = _FakeBackend("anthropic-cloud", script=[("ok", "CLOUD")]) + res = asyncio.run(_egress([local, cloud]).complete( + LLMRequest(system="s", user="u", tier="deep"))) + assert res.text == "CLOUD" and cloud.calls == 1 and local.calls == 0 + + +# ── (gate) hook bus: transform, observe, isolation ───────────────────────────── + +def test_before_request_hook_transforms() -> None: + echo = _FakeBackend("openai-compatible-local") # echoes req.user + hooks = HookBus() + hooks.register("before_request", lambda req: dataclasses.replace(req, user="REWRITTEN")) + res = asyncio.run(_egress([echo], hooks=hooks).complete(_req("original"))) + assert res.text == "REWRITTEN" + + +def test_after_response_and_on_result_observed() -> None: + b = _FakeBackend("openai-compatible-local", script=[("ok", "X")]) + seen: dict[str, object] = {} + hooks = HookBus() + hooks.register("after_response", lambda provider, identity: seen.update(provider=provider, ident=identity)) + hooks.register("on_result", lambda result: seen.update(result_text=result.text)) + asyncio.run(_egress([b], hooks=hooks).complete(_req())) + assert seen["provider"] == "openai-compatible-local" + assert isinstance(seen["ident"], ModelIdentity) + assert seen["result_text"] == "X" + + +def test_broken_hook_never_breaks_the_call() -> None: + b = _FakeBackend("openai-compatible-local", script=[("ok", "SURVIVED")]) + hooks = HookBus() + + def _boom(*_a): + raise RuntimeError("hook exploded") + + hooks.register("before_request", _boom) + hooks.register("after_response", _boom) + hooks.register("on_result", _boom) + res = asyncio.run(_egress([b], hooks=hooks).complete(_req())) + assert res.text == "SURVIVED", "a throwing hook must be isolated" + + +# ── (gate) wired backend over httpx.MockTransport (no sockets) ────────────────── + +def _local_backend(handler) -> OpenAICompatLocalBackend: + spec = BackendSpec( + provider="openai-compatible-local", + base_url="http://127.0.0.1:11434", + model="ai/qwen3:latest", + ) + return OpenAICompatLocalBackend(spec, timeout_s=5.0, transport=httpx.MockTransport(handler)) + + +def test_local_backend_http_happy_path() -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/v1/chat/completions" + return httpx.Response(200, json={"choices": [{"message": {"content": "hi-there"}}]}) + + text, ident = asyncio.run(_local_backend(handler).complete(system="s", user="u", max_tokens=16)) + assert text == "hi-there" + assert ident.provider == "openai-compatible-local" + + +def test_local_backend_429_classifies_ratelimited_with_header() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(429, headers={"Retry-After": "7"}, json={"error": "slow down"}) + + try: + asyncio.run(_local_backend(handler).complete(system="s", user="u", max_tokens=16)) + raise AssertionError("expected RateLimited") + except RateLimited as err: + assert err.retry_after_s == 7.0 + + +def test_local_backend_500_is_transient() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, text="boom") + + try: + asyncio.run(_local_backend(handler).complete(system="s", user="u", max_tokens=16)) + raise AssertionError("expected Transient") + except Transient: + pass + + +def test_local_backend_connection_error_is_transient() -> None: + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("refused") + + try: + asyncio.run(_local_backend(handler).complete(system="s", user="u", max_tokens=16)) + raise AssertionError("expected Transient") + except Transient: + pass + + +# ── single-egress invariant (anthropic import scan) + dep budget ─────────────── + +def test_single_egress_no_foreign_anthropic_import() -> None: + repo = _aegis_root() + offenders: list[str] = [] + for py in repo.rglob("*.py"): + rel = py.relative_to(repo).as_posix() + if rel.startswith("core/llm/") or "/tests/" in f"/{rel}": + continue + try: + tree = ast.parse(py.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError): + continue + for node in ast.walk(tree): + if isinstance(node, ast.Import) and any( + a.name.split(".")[0] == "anthropic" for a in node.names + ): + offenders.append(rel) + elif isinstance(node, ast.ImportFrom) and (node.module or "").split(".")[0] == "anthropic": + offenders.append(rel) + assert not offenders, f"non-llm modules import anthropic: {sorted(set(offenders))}" + + +def test_no_litellm_dependency() -> None: + req = _aegis_root() / "requirements.txt" + text = req.read_text(encoding="utf-8").lower() if req.exists() else "" + assert "litellm" not in text, "litellm must never appear in requirements" + + +def test_single_egress_completion_post_DEFERRED() -> None: + """DEFERRED to the cutover PR. HttpBackend.generate_config still owns its own POST to + /v1/chat/completions; enabling the 'no completion URL outside core/llm' scan now would + be a false RED. This placeholder documents the deferral and stays green.""" + repo = _aegis_root() + hb = repo / "core" / "backends" / "http_backend.py" + assert hb.exists(), "expected the pre-cutover HttpBackend to still exist in PR-1" + print(" NOTE: completion-POST single-egress scan deferred to the http_backend cutover PR") + + +# ── helpers + runner ─────────────────────────────────────────────────────────── + +def _aegis_root() -> Path: + here = Path(__file__).resolve() + for parent in here.parents: + if (parent / "core" / "llm").exists(): + return parent + return here.parent.parent + + +def _run_all() -> int: + funcs = [v for k, v in sorted(globals().items()) if k.startswith("test_") and callable(v)] + failures = 0 + for fn in funcs: + try: + fn() + print(f"PASS {fn.__name__}") + except AssertionError as e: + failures += 1 + print(f"FAIL {fn.__name__}: {e}") + except Exception as e: # noqa: BLE001 + failures += 1 + print(f"ERROR {fn.__name__}: {type(e).__name__}: {e}") + print(f"\n{len(funcs) - failures}/{len(funcs)} passed") + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(_run_all()) From 6e71d03f6501a4d1977dc87137d2b8c6f90d343d Mon Sep 17 00:00:00 2001 From: Georgi Gaydarov Date: Sun, 31 May 2026 17:05:11 -0400 Subject: [PATCH 2/2] fix(llm): harden is_loopback against parser-differential air-gap bypass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Automated security review (HIGH) flagged the string-prefix host check: a host that only STARTS WITH "127." (e.g. 127.evil.com) or contains "localhost" (localhost.evil.com) passed startswith("127.") yet resolves OFF-perimeter — a real air-gap bypass, since is_loopback gates the entire wedge. Fix: parse the host with ipaddress.ip_address(...).is_loopback (strict, literal-IP only) instead of a string prefix. Allow only the literal name localhost/ip6-localhost, genuine 127.0.0.0/8 + ::1, and the reserved non-routable .local mDNS suffix; reject empty/relative hosts. NO DNS resolution — resolving an attacker-supplied host would itself be egress, exactly what air-gap forbids. Adds test_is_loopback_rejects_spoofed_prefix (127.evil.com / 127.0.0.1.evil.com / localhost.evil.com / empty all rejected; ::1 + 127.5.5.5 allowed). Suite 29/29 green. --- core/llm/airgap.py | 37 ++++++++++++++++++++++++++----------- tests/llm_egress_test.py | 12 ++++++++++++ 2 files changed, 38 insertions(+), 11 deletions(-) diff --git a/core/llm/airgap.py b/core/llm/airgap.py index c8cb818..4533fde 100644 --- a/core/llm/airgap.py +++ b/core/llm/airgap.py @@ -12,30 +12,45 @@ NOTE: AEGIS's cloud backend uses raw httpx (no `import anthropic`) precisely so that invariant #3 holds even when the cloud path IS exercised in a non-air-gapped build. + +Security: `is_loopback` parses the host with `ipaddress` (NOT a string prefix), so a +spoofed hostname like ``127.evil.com`` or ``127.0.0.1.attacker.com`` — which would pass a +naive ``startswith("127.")`` yet resolve OFF-perimeter — is correctly rejected. No DNS +resolution is performed: resolving an attacker-supplied host would itself be egress, which +is exactly what air-gap mode forbids. """ from __future__ import annotations import sys +from ipaddress import ip_address from urllib.parse import urlparse _CLOUD_PROVIDER = "anthropic-cloud" -# Loopback hosts allowed even in air-gap mode (in-perimeter only). -_LOOPBACK_HOSTS: frozenset[str] = frozenset({"localhost", "127.0.0.1", "::1", ""}) +# Literal non-IP hostnames allowed in air-gap mode (loopback aliases only). +_LOOPBACK_NAMES: frozenset[str] = frozenset({"localhost", "ip6-localhost"}) def is_loopback(url: str) -> bool: - """True only for loopback / .local hosts. An empty host (relative URL) counts as - loopback because the existing http_backend posts to absolute in-perimeter URLs only; - a bare host means same-host.""" + """True only for a genuine loopback address (127.0.0.0/8 or ::1), the literal name + 'localhost', or a reserved, non-routable '.local' mDNS host (in-perimeter only). + + Uses strict `ipaddress` parsing rather than a string prefix, so a host that merely + *starts with* "127." but is actually a routable name (e.g. ``127.evil.com``) is + rejected. An empty/relative host is rejected (never auto-allowed). No DNS lookup is + performed — resolving an external host would itself violate the air-gap. + """ host = (urlparse(url).hostname or "").lower() - if host in _LOOPBACK_HOSTS: - return True - if host.endswith(".local"): - return True - if host.startswith("127."): # 127.0.0.0/8 is all loopback + if not host: + return False # reject empty / relative — never auto-allow + if host in _LOOPBACK_NAMES: return True - return False + try: + return ip_address(host).is_loopback # 127.0.0.0/8 and ::1, strictly + except ValueError: + # Not a literal IP. Only the reserved, non-internet-routable mDNS suffix is + # treated as in-perimeter (RFC 6762); everything else is non-loopback. + return host.endswith(".local") def assert_airgap_ok(provider: str, base_url: str) -> None: diff --git a/tests/llm_egress_test.py b/tests/llm_egress_test.py index 7594754..fda773c 100644 --- a/tests/llm_egress_test.py +++ b/tests/llm_egress_test.py @@ -425,6 +425,18 @@ def test_single_egress_completion_post_DEFERRED() -> None: # ── helpers + runner ─────────────────────────────────────────────────────────── +def test_is_loopback_rejects_spoofed_prefix() -> None: + """Regression for the parser-differential the security review caught: a host that only + *starts with* '127.' (or contains 'localhost') but resolves OFF-perimeter must NOT be + treated as loopback. Strict ipaddress parsing (no DNS) closes the bypass.""" + assert not is_loopback("http://127.evil.com") + assert not is_loopback("http://127.0.0.1.evil.com") + assert not is_loopback("http://localhost.evil.com") + assert not is_loopback("http://") # empty / relative host rejected + assert is_loopback("http://[::1]:5757") # genuine ipv6 loopback + assert is_loopback("http://127.5.5.5") # genuine 127.0.0.0/8 + + def _aegis_root() -> Path: here = Path(__file__).resolve() for parent in here.parents: