endpoint_registry.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# oc_chat/core/endpoint_registry.py
 
"""
Endpoint registry & LAN policy (Phase O2).
 
Purpose
-------
A provider may serve multiple named base-URL targets (e.g. an Ollama
install on ``localhost`` plus a LAN GPU box). :class:`EndpointRegistry`
collects them from three sources and exposes a runtime-switchable active
selection consumed by :class:`oc_chat.core.client_interface.ClientInterface`.
 
Responsibilities
----------------
- :func:`classify_lan_policy` — pure helper mapping a base URL to one of
  ``"local"`` / ``"lan"`` / ``"remote"`` / ``"blocked"``.
- :class:`EndpointRegistry` — provider-scoped collection of
  :class:`oc_chat.core.provider_plugin.ProviderEndpoint` entries with
  three loaders (constructor seed from ``provider.endpoints()``,
  :meth:`load_env` for ``OLLAMA_BASE_URLS``, :meth:`load_toml` for
  ``~/.oc_chat/endpoints.toml``) and a runtime :meth:`switch` that
  enforces LAN policy.
 
Sources merge in this order: provider defaults → env var → TOML. Later
definitions override earlier ones for the same name.
 
Diagnostics
-----------
Domain: OC_CHAT.MAIN
Levels:
    L1 — malformed env / TOML entries skipped (override=True)
    L2 — registry load lifecycle, endpoint switches
    L3 — per-entry LAN classification
    L4 — module load
 
Contracts
---------
- Constructor seeds ``entries`` from ``provider.endpoints()``. After
  construction, callers may call :meth:`load_env` and/or :meth:`load_toml`
  to add more.
- :meth:`switch` raises :class:`KeyError` on unknown name and
  :class:`PermissionError` on a ``"blocked"`` entry.
- :meth:`active` returns the currently selected entry; defaults to the
  one named ``"default"`` if it exists, else the first registered.
- ``"blocked"`` classification is consulted via
  ``OC_CHAT_ALLOW_REMOTE_ENDPOINTS`` at load time; toggling the env var
  after load does not re-classify existing entries until a loader is
  re-run.
 
Dependencies
------------
- oc_chat.core.provider_plugin (ProviderEndpoint, ProviderPlugin and the provider default endpoints)
- oc_diagnostics (structured debug logging via _dbg)
- tomllib (stdlib TOML parser for the endpoints file on Python 3.11 and later)
- oc_chat.core.egress_guard (forbidden provider host check for LAN policy and endpoint selection)
- tomli (TOML parser fallback for the endpoints file on older Pythons)
"""
from __future__ import annotations
 
import ipaddress
import os
import re
from collections.abc import Mapping
from pathlib import Path
from typing import Literal, TYPE_CHECKING
from urllib.parse import urlparse
 
try:
    import tomllib  # type: ignore[import-not-found]
except ModuleNotFoundError:  # pragma: no cover - Python < 3.11 fallback
    import tomli as tomllib  # type: ignore[import-not-found,no-redef]
 
from oc_diagnostics import _dbg
 
from oc_chat.core.egress_guard import is_forbidden_provider_host
from oc_chat.core.provider_plugin import ProviderEndpoint, endpoints as _provider_endpoints
 
if TYPE_CHECKING:
    from oc_chat.core.provider_plugin import ProviderPlugin
 
_dbg("OC_CHAT.MAIN", "module", "loaded", 4)
 
 
LANPolicy = Literal["local", "lan", "remote", "blocked"]
 
 
_LOCAL_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "0.0.0.0"})
 
_ENV_VAR_BASE_URLS = "OLLAMA_BASE_URLS"
_ENV_VAR_ALLOW_REMOTE = "OC_CHAT_ALLOW_REMOTE_ENDPOINTS"
 
_ENV_ENTRY_RE = re.compile(r"^\s*([A-Za-z0-9_.\-]+)\s*=\s*(\S.*?)\s*$")
 
 
def _remote_allowed() -> bool:
    """Return ``True`` when ``OC_CHAT_ALLOW_REMOTE_ENDPOINTS=1`` is set."""
    return os.environ.get(_ENV_VAR_ALLOW_REMOTE, "") == "1"
 
 
def classify_lan_policy(base_url: str) -> LANPolicy:
    """Classify *base_url* as ``local`` / ``lan`` / ``remote`` / ``blocked``.
 
    Loopback and ``localhost`` are ``"local"``. RFC1918 private space is
    ``"lan"``. Any other host is ``"remote"`` if
    ``OC_CHAT_ALLOW_REMOTE_ENDPOINTS=1``, otherwise ``"blocked"``.
    """
    func = "classify_lan_policy"
    parsed = urlparse(base_url)
    host = (parsed.hostname or "").strip().lower()
    if not host:
        return "blocked"
    if host in _LOCAL_HOSTS:
        _dbg("OC_CHAT.MAIN", func, f"{host} -> local", 3)
        return "local"
    # M8.0 WU-4c: link-local (incl. IMDS 169.254.169.254) / multicast / reserved
    # are never a legitimate endpoint -> blocked (was wrongly "lan": link-local is
    # in ipaddress.is_private). Checked before the is_private->lan branch.
    if is_forbidden_provider_host(host):
        _dbg("OC_CHAT.MAIN", func, f"{host} -> blocked (forbidden class)", 2)
        return "blocked"
    try:
        ip = ipaddress.ip_address(host)
        if ip.is_loopback:
            return "local"
        if ip.is_private:
            _dbg("OC_CHAT.MAIN", func, f"{host} -> lan", 3)
            return "lan"
    except ValueError:
        pass
    policy: LANPolicy = "remote" if _remote_allowed() else "blocked"
    _dbg("OC_CHAT.MAIN", func, f"{host} -> {policy}", 3)
    return policy
 
 
class EndpointRegistry:
    """Collects named :class:`ProviderEndpoint` entries for one provider.
 
    The registry is provider-scoped: each :class:`ProviderPlugin`
    instance gets its own registry built at startup from its own
    ``endpoints()`` plus the env-var + TOML overlays.
    """
 
    def __init__(self, provider: "ProviderPlugin") -> None:
        func = "EndpointRegistry.__init__"
        self.provider = provider
        self.entries: dict[str, ProviderEndpoint] = dict(_provider_endpoints(provider))
        if "default" in self.entries:
            self.active_name = "default"
        elif self.entries:
            self.active_name = next(iter(self.entries))
        else:
            self.active_name = ""
        # Plan 152 R11-03: the ACTIVE endpoint is the committed OBJECT, not just a
        # name. Resolving `entries[active_name]` afresh on every call let an
        # allowed post-freeze overlay swap the object under a committed
        # activation. `None` here means "not yet explicitly activated" — `active()`
        # falls back to name resolution until the first activation commits.
        self._active_endpoint: ProviderEndpoint | None = None
        self._default_endpoint: ProviderEndpoint | None = None
        self._default_frozen = False
        _dbg(
            "OC_CHAT.MAIN",
            func,
            f"provider={provider.name} entries={list(self.entries)} active={self.active_name}",
            2,
        )
 
    # ── Committed-object activation + frozen default (plan 152 ADR-2) ───────
 
    def _check_policy(self, endpoint: ProviderEndpoint, label: str) -> ProviderEndpoint:
        """Re-run both M8.0 WU-4c checks on *endpoint*. Fails closed.
 
        Factored out so every access path — :meth:`active`, :meth:`switch`,
        :meth:`resolve` — applies the SAME policy, and so the committed-object
        model preserves the "re-evaluated on every call (no allow caching)"
        property rather than caching an allow decision alongside the object.
        """
        if endpoint.lan_policy_status == "blocked":
            raise PermissionError(
                f"{label} endpoint {endpoint.name!r} ({endpoint.base_url}) is blocked "
                f"by LAN policy. Set {_ENV_VAR_ALLOW_REMOTE}=1 to opt in."
            )
        if is_forbidden_provider_host(endpoint.base_url):
            raise PermissionError(
                f"{label} endpoint {endpoint.name!r} ({endpoint.base_url}) is a "
                f"forbidden internal address; refused."
            )
        return endpoint
 
    def resolve(self, name: str) -> ProviderEndpoint:
        """Resolve + policy-check *name* WITHOUT activating it (PREPARE).
 
        The non-mutating half of activation: unknown / missing / blocked /
        forbidden-host are all rejected HERE, in the owner, so a transition can
        validate a candidate endpoint during PREPARE without touching
        ``active_name`` or the live client. An empty *name* resolves to the
        frozen default (:meth:`default_name`) rather than to whatever the
        previous session's :meth:`switch` happened to leave active.
 
        This exists so callers never hand-roll policy or reach for
        ``entries[name]`` + ``rebind_to`` — that path silently bypasses the
        ``blocked`` LAN check, which ``rebind_to`` does not enforce.
        """
        func = "EndpointRegistry.resolve"
        if not self.entries:
            raise KeyError(f"No endpoints registered for provider {self.provider.name!r}")
        target_name = name or self.default_name
        if target_name not in self.entries:
            available = ", ".join(sorted(self.entries)) or "(none)"
            raise KeyError(
                f"Unknown endpoint {target_name!r} for provider {self.provider.name!r}. "
                f"Available: {available}"
            )
        endpoint = self.entries[target_name]
        _dbg("OC_CHAT.MAIN", func, f"resolved {target_name!r} (non-mutating)", 3)
        return self._check_policy(endpoint, "Resolved")
 
    def freeze_default(self) -> None:
        """Capture the default endpoint OBJECT once, after the overlays.
 
        Idempotent and first-call-wins: repeated calls are a no-op, so the frozen
        default can never change after startup. Called from :meth:`load_toml`'s
        ``finally`` — the LAST overlay — so it runs on every exit path including
        that method's three early returns (missing file, malformed TOML,
        non-table ``[endpoints]``), each of which is a normal production case.
 
        The OBJECT is captured, not merely the name, because :meth:`add`,
        :meth:`load_env` and :meth:`load_toml` all REPLACE the entry stored under
        a name — a frozen name alone would let a later overlay silently swap the
        default out from under a session that selected it.
 
        An empty registry freezes to ``None``; :meth:`default_name` and
        :meth:`resolve` then fail closed with ``KeyError``, matching
        :meth:`active`.
        """
        func = "EndpointRegistry.freeze_default"
        if self._default_frozen:
            return                                  # first-call-wins
        self._default_frozen = True
        if "default" in self.entries:
            self._default_endpoint = self.entries["default"]
        elif self.entries:
            self._default_endpoint = next(iter(self.entries.values()))
        else:
            self._default_endpoint = None
        name = self._default_endpoint.name if self._default_endpoint is not None else ""
        _dbg("OC_CHAT.MAIN", func, f"froze default={name!r}", 2)
 
    @property
    def default_name(self) -> str:
        """The frozen default endpoint name.
 
        Before :meth:`freeze_default` runs this returns a PROVISIONAL value (the
        construction-time ``active_name``); afterwards it is stable for the
        process lifetime. Raises ``KeyError`` when the registry has no entries.
        """
        if self._default_frozen:
            if self._default_endpoint is None:
                raise KeyError(
                    f"No endpoints registered for provider {self.provider.name!r}"
                )
            return self._default_endpoint.name
        if not self.entries:
            raise KeyError(f"No endpoints registered for provider {self.provider.name!r}")
        return self.active_name or next(iter(self.entries))
 
    def commit_endpoint_activation(self, endpoint: ProviderEndpoint) -> None:
        """Publish *endpoint* as the committed active object — no-fail.
 
        The COMMIT half of activation. Assigns the OBJECT handle and its name
        together, so ``active()`` returns exactly this object until the next
        explicit activation. Pure assignment: policy already ran in
        :meth:`resolve`, and ``active()`` re-checks on every read anyway.
        """
        self._active_endpoint = endpoint
        self.active_name = endpoint.name
 
    def active(self) -> ProviderEndpoint:
        """Return the currently active endpoint.
 
        Raises :class:`KeyError` when the registry is empty, and
        :class:`PermissionError` when the active endpoint resolves to a forbidden
        internal address (IMDS / link-local / multicast / reserved) **or** carries
        a ``lan_policy_status == "blocked"`` LAN policy. M8.0 WU-4c fails closed
        even on the default / ``active_name`` path (closing the
        ``active_name="default"``-by-direct-assignment bypass), not only on
        :meth:`switch`. The two checks are complementary: ``is_forbidden_provider_host``
        catches the IP-literal internal class (incl. remote IP literals only when
        forbidden), while the ``blocked`` guard catches a remote HOSTNAME that
        ``is_forbidden_provider_host`` cannot parse as an IP — closing the
        AC-V12-3 direct-``active_name``-assignment bypass of the
        ``OC_CHAT_ALLOW_REMOTE_ENDPOINTS`` opt-in that :meth:`switch` already
        enforces. Re-evaluated on every call (no allow caching).
 
        **Committed-object model (plan 152 R11-03).** Once an explicit activation
        has committed — :meth:`switch`, the session-transition seam, or
        :meth:`commit_endpoint_activation` — this returns THAT object until the
        next activation commits another. It no longer re-resolves
        ``entries[active_name]`` on every call, because a post-freeze overlay
        replacing ``entries["default"]`` would otherwise hand the next turn's
        rebind an endpoint nobody activated (and a blocked replacement would
        raise ``PermissionError`` where the caller only catches ``KeyError``,
        aborting the turn). The NAME selects at activation time; the committed
        OBJECT is authoritative between activations.
 
        Policy is still re-run on the committed object on EVERY call, so the
        WU-4c "no allow caching" property is preserved.
        """
        if not self.entries:
            raise KeyError(f"No endpoints registered for provider {self.provider.name!r}")
        if self._active_endpoint is not None:
            return self._check_policy(self._active_endpoint, "Active")
        # No explicit activation yet — fall back to name resolution (startup).
        if self.active_name and self.active_name in self.entries:
            endpoint = self.entries[self.active_name]
        else:
            self.active_name = next(iter(self.entries))
            endpoint = self.entries[self.active_name]
        return self._check_policy(endpoint, "Active")
 
    def switch(self, name: str) -> ProviderEndpoint:
        """Activate the endpoint named *name*.
 
        Raises :class:`KeyError` when no such endpoint exists, and
        :class:`PermissionError` when the endpoint's LAN policy is
        ``"blocked"``.
        """
        func = "EndpointRegistry.switch"
        if name not in self.entries:
            available = ", ".join(sorted(self.entries)) or "(none)"
            raise KeyError(
                f"Unknown endpoint {name!r} for provider {self.provider.name!r}. "
                f"Available: {available}"
            )
        target = self.entries[name]
        if target.lan_policy_status == "blocked":
            raise PermissionError(
                f"Endpoint {name!r} ({target.base_url}) is blocked by LAN policy. "
                f"Set {_ENV_VAR_ALLOW_REMOTE}=1 to opt in."
            )
        # An explicit activation: commit the OBJECT together with its name so
        # `active()` cannot be re-pointed by a later overlay (plan 152 R11-03).
        self.commit_endpoint_activation(target)
        _dbg(
            "OC_CHAT.MAIN",
            func,
            f"provider={self.provider.name} -> {name} ({target.base_url})",
            2,
        )
        return target
 
    def add(self, endpoint: ProviderEndpoint) -> None:
        """Insert or replace *endpoint* by name.
 
        Classification is the caller's responsibility — pass a
        :class:`ProviderEndpoint` with ``lan_policy_status`` already
        set.
        """
        self.entries[endpoint.name] = endpoint
 
    def load_env(self, env: Mapping[str, str]) -> int:
        """Merge entries parsed from ``env['OLLAMA_BASE_URLS']``.
 
        Format is ``name=url[,name=url...]``. Whitespace tolerant;
        malformed entries are logged at L1 and skipped. Returns the
        number of successfully parsed entries.
        """
        func = "EndpointRegistry.load_env"
        raw = env.get(_ENV_VAR_BASE_URLS, "").strip()
        if not raw:
            return 0
        count = 0
        for part in raw.split(","):
            part = part.strip()
            if not part:
                continue
            match = _ENV_ENTRY_RE.match(part)
            if not match:
                _dbg(
                    "OC_CHAT.MAIN",
                    func,
                    f"malformed entry {part!r} skipped",
                    1,
                    override=True,
                )
                continue
            name, url = match.group(1), match.group(2)
            policy = classify_lan_policy(url)
            self.entries[name] = ProviderEndpoint(name=name, base_url=url, lan_policy_status=policy)
            count += 1
        _dbg(
            "OC_CHAT.MAIN",
            func,
            f"loaded {count} entries from {_ENV_VAR_BASE_URLS}",
            2,
        )
        return count
 
    def load_toml(self, path: Path) -> int:
        """Merge entries from a ``[endpoints]`` TOML table at *path*.
 
        Missing file is a silent no-op (returns 0). Malformed TOML logs
        at L1 and returns 0. Each ``name = "url"`` becomes one entry;
        non-string values are skipped with an L1 log.
 
        **This is the LAST overlay**, so it is also where the default endpoint is
        frozen (plan 152 ADR-2 / R10-11). The freeze lives in a ``finally`` around
        the private parse/apply helper so it runs on EVERY exit path — including
        the three early returns (missing file, malformed TOML / I-O failure,
        non-table ``[endpoints]``), the first of which is the common production
        case because ``~/.oc_chat/endpoints.toml`` usually does not exist. A
        literal "call it at the end" would skip exactly those paths and leave
        ``default_name`` provisional forever.
 
        Malformed-file policy: freeze the VALID constructor+env state rather than
        leaving the default provisional. Post-freeze calls may still REPLACE
        entries — runtime endpoints stay mutable — but can never change the
        frozen default OBJECT; the freeze itself is first-call-wins.
        """
        try:
            return self._load_toml_entries(path)
        finally:
            self.freeze_default()
 
    def _load_toml_entries(self, path: Path) -> int:
        """Parse + apply the ``[endpoints]`` table. See :meth:`load_toml`.
 
        Split out so the freeze can wrap every exit in a ``finally``; the three
        early returns below are exactly why.
        """
        func = "EndpointRegistry.load_toml"
        if not path.exists():
            return 0
        try:
            with open(path, "rb") as fh:
                data = tomllib.load(fh)
        except (tomllib.TOMLDecodeError, OSError) as exc:
            _dbg(
                "OC_CHAT.MAIN",
                func,
                f"malformed TOML at {path}: {exc!r}",
                1,
                override=True,
            )
            return 0
        section = data.get("endpoints", {})
        if not isinstance(section, dict):
            _dbg(
                "OC_CHAT.MAIN",
                func,
                f"[endpoints] is not a table at {path}",
                1,
                override=True,
            )
            return 0
        count = 0
        for name, url in section.items():
            if not isinstance(url, str) or not url:
                _dbg(
                    "OC_CHAT.MAIN",
                    func,
                    f"non-string url for entry {name!r} skipped",
                    1,
                    override=True,
                )
                continue
            policy = classify_lan_policy(url)
            self.entries[name] = ProviderEndpoint(name=name, base_url=url, lan_policy_status=policy)
            count += 1
        _dbg(
            "OC_CHAT.MAIN",
            func,
            f"loaded {count} entries from {path}",
            2,
        )
        return count
 
 
__all__ = [
    "EndpointRegistry",
    "ProviderEndpoint",
    "classify_lan_policy",
    "LANPolicy",
]