generic-ml-wrapper 0.1.0__py3-none-any.whl

generic_ml_wrapper/adapter/outbound/gateway/relay.py

@@ -0,0 +1,402 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""A local metering relay: forward a client's API traffic and record per-turn usage.
+
+The relay listens on ``127.0.0.1``; the caller points the client at it via the
+client's official base-URL override (Claude: ``ANTHROPIC_BASE_URL``). Each request
+is forwarded verbatim to the upstream API and the response is streamed straight
+back to the client, so the interactive experience is unchanged. While streaming,
+the relay tees the response to read the turn's token usage and record it.
+
+The upstream forwarder is injectable so the relay is testable against a fake
+upstream; the default talks HTTPS to the real API. When a transcript port is
+injected, each metered turn's request, response, and usage are handed to it, so the
+run is recorded per call.
+"""
+
+from __future__ import annotations
+
+import http.client
+import itertools
+import secrets
+import threading
+import time
+from dataclasses import dataclass
+from http import HTTPStatus
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import TYPE_CHECKING, cast
+from urllib.parse import urlsplit
+
+from generic_ml_wrapper.adapter.outbound.gateway.anthropic_sse import read_usage as _anthropic_usage
+from generic_ml_wrapper.application.domain.model.turn_usage import TurnUsage
+from generic_ml_wrapper.application.domain.service.interceptor_chain import InterceptorChain
+from generic_ml_wrapper.application.port.outbound.transcript import TranscriptCall
+from generic_ml_wrapper.common.log import log
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterable, Mapping
+
+    from generic_ml_wrapper.adapter.outbound.gateway.anthropic_sse import StreamUsage
+    from generic_ml_wrapper.application.port.outbound.per_turn_metering import PerTurnMeteringPort
+    from generic_ml_wrapper.application.port.outbound.transcript import TranscriptPort
+
+    UsageReader = Callable[[str], "StreamUsage | None"]
+    MeteredPredicate = Callable[[str, str], bool]
+    PathMap = Callable[[str], str]
+
+_ANTHROPIC = "https://api.anthropic.com"
+
+
+def _claude_metered(method: str, path: str) -> bool:
+    """The default metered-turn predicate: a Claude ``POST /v1/messages``."""
+    return method == "POST" and path.split("?", 1)[0].endswith("/v1/messages")
+
+
+# Headers that describe one hop's connection, not the payload; never forwarded.
+_HOP_BY_HOP = frozenset(
+    {
+        "host",
+        "content-length",
+        "connection",
+        "keep-alive",
+        "transfer-encoding",
+        "te",
+        "trailers",
+        "upgrade",
+        "proxy-authorization",
+    }
+)
+
+
+@dataclass(frozen=True)
+class UpstreamResponse:
+    """An upstream response: a status, headers, and a streaming body."""
+
+    status: int
+    headers: list[tuple[str, str]]
+    body: Iterable[bytes]
+
+
+class MeteringRelay:
+    """Forward a client's traffic to an upstream API and record each turn's usage."""
+
+    def __init__(  # noqa: PLR0913  (per-client relay knobs, all keyword-only)
+        self,
+        *,
+        job: str,
+        session: str,
+        metering: PerTurnMeteringPort,
+        client: str = "claude",
+        transcript: TranscriptPort | None = None,
+        forwarder: Callable[[str, str, Mapping[str, str], bytes], UpstreamResponse] | None = None,
+        upstream_base: str = _ANTHROPIC,
+        path_map: PathMap | None = None,
+        usage_reader: UsageReader | None = None,
+        is_metered: MeteredPredicate | None = None,
+        interceptors: InterceptorChain | None = None,
+        clock: Callable[[], float] = time.time,
+    ) -> None:
+        """Bind the relay to a run and its metering store.
+
+        Args:
+            job: The job the turns belong to.
+            session: The session the turns belong to.
+            metering: Where per-turn usage is recorded.
+            client: The client name, placed in the capability URL for observability.
+            transcript: Where each call's request/response/usage is recorded, or ``None``.
+            forwarder: The upstream forwarder; defaults to real HTTPS to ``upstream_base``.
+            upstream_base: The upstream API base URL for the default forwarder.
+            path_map: Maps the incoming request path to the upstream path for the
+                default forwarder (e.g. Codex ``/v1/x`` -> ``/backend-api/codex/x``);
+                identity when ``None``.
+            usage_reader: Reads a turn's usage from a response body; defaults to the
+                Anthropic reader. A per-client reader plugs in a different wire shape.
+            is_metered: Predicate ``(method, path) -> bool`` selecting the turn
+                endpoint; defaults to Claude's ``POST /v1/messages``.
+            interceptors: The interceptor chain applied to the wire — ``request`` to
+                the outbound body, ``response`` to the captured reply; empty when ``None``.
+            clock: Returns the current time (epoch seconds); injectable for tests.
+                Used to stamp each turn's timestamp and duration.
+        """
+        self._job = job
+        self._session = session
+        self._metering = metering
+        self._transcript = transcript
+        self._client = client
+        # A per-run capability token: the client is pointed at
+        # http://127.0.0.1:<port>/<client>/<token>, so only a caller that was handed
+        # this base URL (our client) can drive the relay. This is the sole auth guard;
+        # the client segment is for observability + a routing sanity check.
+        self._token = secrets.token_urlsafe(16)
+        self._forward = forwarder or _https_forwarder(upstream_base, path_map)
+        self._read_usage = usage_reader or _anthropic_usage
+        self._metered = is_metered or _claude_metered
+        self._interceptors = interceptors or InterceptorChain(())
+        self._clock = clock
+        self._server: ThreadingHTTPServer | None = None
+        # A thread-safe capture counter: the server is threaded, so two concurrent turns
+        # must not read-modify-write a shared int (which could collide on a filename).
+        self._turns = itertools.count(1)
+
+    def start(self) -> None:
+        """Bind an ephemeral local port and serve in a background thread."""
+        self._server = _RelayServer(("127.0.0.1", 0), _Handler, self)
+        threading.Thread(target=self._server.serve_forever, daemon=True).start()
+
+    def stop(self) -> None:
+        """Stop serving and release the port."""
+        if self._server is not None:
+            self._server.shutdown()
+            self._server.server_close()
+            self._server = None
+
+    @property
+    def port(self) -> int:
+        """The bound port (only valid after ``start``)."""
+        if self._server is None:
+            message = "relay is not started"
+            raise RuntimeError(message)
+        return self._server.server_address[1]
+
+    @property
+    def base_url(self) -> str:
+        """The capability base URL to point the client at (only valid after ``start``).
+
+        Carries the run's client name (observability) and the per-run token (auth):
+        ``http://127.0.0.1:<port>/<client>/<token>``. The client appends its API path
+        (e.g. ``/v1/messages``) to this, so requests arrive already authenticated.
+        """
+        return f"http://127.0.0.1:{self.port}/{self._client}/{self._token}"
+
+    def authorize(self, path: str) -> str | None:
+        """Validate a request path's ``/<client>/<token>`` prefix and strip it.
+
+        Args:
+            path: The incoming request path.
+
+        Returns:
+            The upstream path (prefix removed), or ``None`` if the client segment or
+            token does not match this run -- in which case the request is refused.
+        """
+        prefix = f"/{self._client}/{self._token}"
+        if path == prefix:
+            return "/"
+        if path.startswith(prefix + "/"):
+            return path[len(prefix) :]
+        return None
+
+    def forward(
+        self, method: str, path: str, headers: Mapping[str, str], body: bytes
+    ) -> UpstreamResponse:
+        """Forward one request upstream."""
+        return self._forward(method, path, headers, body)
+
+    def now(self) -> float:
+        """The current time, from the injected clock."""
+        return self._clock()
+
+    def is_metered(self, method: str, path: str) -> bool:
+        """Whether a request to ``path`` is a metered turn for this client."""
+        return self._metered(method, path)
+
+    def wants_body(self, method: str, path: str) -> bool:
+        """Whether the full response body must be captured this request.
+
+        True only when something consumes it -- a metered turn, a ``response``
+        interceptor, or raw capture. When False the relay streams the response
+        straight through without buffering it in memory.
+        """
+        return self._metered(method, path) or self._interceptors.has("response")
+
+    def intercept_request(self, body: bytes) -> bytes:
+        """Apply ``request`` interceptors to the outbound body (identity if none).
+
+        The body is left byte-for-byte untouched when no ``request`` interceptor is
+        configured, so the common path never re-encodes.
+
+        Args:
+            body: The raw request body from the client.
+
+        Returns:
+            The (possibly transformed) body to forward upstream.
+        """
+        if not self._interceptors.has("request"):
+            return body
+        transformed = self._interceptors.apply("request", body.decode("utf-8", "replace"))
+        return transformed.encode("utf-8")
+
+    def intercept_response(self, captured: bytes) -> None:
+        """Run ``response`` interceptors over the captured body (observe only).
+
+        The response has already streamed to the client, so this is for logging or
+        redacting-for-logs; the returned text is intentionally discarded.
+
+        Args:
+            captured: The full captured response body.
+        """
+        if self._interceptors.has("response"):
+            self._interceptors.apply("response", captured.decode("utf-8", "replace"))
+
+    def record(self, request: bytes, captured: bytes, started_at: float) -> None:
+        """Record the turn's usage (if any) and its transcript (if configured).
+
+        Args:
+            request: The request body forwarded upstream.
+            captured: The full response body.
+            started_at: The turn's start time (from :meth:`now`), for timestamp/duration.
+        """
+        usage = self._read_usage(captured.decode("utf-8", "replace"))
+        turn: TurnUsage | None = None
+        if usage is not None:
+            turn = TurnUsage(
+                self._session,
+                usage.input_tokens,
+                usage.output_tokens,
+                None,
+                usage.model,
+                cache_creation_tokens=usage.cache_creation_tokens,
+                cache_read_tokens=usage.cache_read_tokens,
+                timestamp=started_at,
+                duration_s=round(self._clock() - started_at, 3),
+                turn_id=usage.turn_id,
+            )
+            self._metering.record(self._job, turn)
+        if self._transcript is not None:
+            self._transcript.record(
+                TranscriptCall(self._job, self._session, next(self._turns), request, captured, turn)
+            )
+
+
+class _RelayServer(ThreadingHTTPServer):
+    def __init__(
+        self,
+        address: tuple[str, int],
+        handler: type[BaseHTTPRequestHandler],
+        relay: MeteringRelay,
+    ) -> None:
+        super().__init__(address, handler)
+        self.relay = relay
+
+
+class _Handler(BaseHTTPRequestHandler):
+    def do_GET(self) -> None:
+        self._proxy()
+
+    def do_POST(self) -> None:
+        self._proxy()
+
+    def log_message(self, format: str, *args: object) -> None:  # noqa: A002  (stdlib signature)
+        log.debug("gateway " + (format % args if args else format))
+
+    def _proxy(self) -> None:
+        relay = cast("_RelayServer", self.server).relay
+        # A browser cross-origin POST carries Origin; a real client SDK does not. And
+        # any request whose /<client>/<token>/ prefix doesn't match this run is refused,
+        # so a stray local process can't drive the relay.
+        if self.headers.get("Origin") is not None:
+            self.send_error(HTTPStatus.FORBIDDEN)
+            return
+        upstream_path = relay.authorize(self.path)
+        if upstream_path is None:
+            self.send_error(HTTPStatus.NOT_FOUND)
+            return
+        length = int(self.headers.get("Content-Length") or 0)
+        body = relay.intercept_request(self.rfile.read(length) if length else b"")
+        # Force an uncompressed upstream response so usage can be read off it; the
+        # client accepts identity fine (only the localhost hop loses compression).
+        request_headers = {
+            name: value for name, value in self.headers.items() if name.lower() != "accept-encoding"
+        }
+        request_headers["Accept-Encoding"] = "identity"
+        started_at = relay.now()
+        response = relay.forward(self.command, upstream_path, request_headers, body)
+
+        self.close_connection = True  # signal end-of-stream by closing (no content-length)
+        try:
+            self.send_response(response.status)
+            for name, value in response.headers:
+                if name.lower() not in _HOP_BY_HOP:
+                    self.send_header(name, value)
+            self.send_header("Connection", "close")
+            self.end_headers()
+        except OSError:
+            pass  # client already gone; still drain upstream below to record usage
+
+        if relay.wants_body(self.command, upstream_path):
+            captured = _tee(response.body, self._write_chunk)
+            relay.intercept_response(captured)
+            if relay.is_metered(self.command, upstream_path):
+                relay.record(body, captured, started_at)
+        else:
+            _stream(response.body, self._write_chunk)
+
+    def _write_chunk(self, chunk: bytes) -> None:
+        self.wfile.write(chunk)
+        self.wfile.flush()
+
+
+def _stream(chunks: Iterable[bytes], sink: Callable[[bytes], None]) -> None:
+    """Stream ``chunks`` to the client without buffering (nothing consumes the body)."""
+    for chunk in chunks:
+        try:
+            sink(chunk)
+        except OSError:
+            return  # client hung up; no body to record or capture, so stop
+
+
+def _tee(chunks: Iterable[bytes], sink: Callable[[bytes], None]) -> bytes:
+    """Stream ``chunks`` to ``sink`` while capturing them all for usage extraction.
+
+    If the client hangs up mid-stream (e.g. it exits or cancels), writes raise and
+    are swallowed, but the upstream is still drained fully so the turn's usage —
+    which arrives at the end — is captured.
+
+    Args:
+        chunks: The upstream response body chunks.
+        sink: Writes a chunk to the client (may raise ``OSError`` once it hangs up).
+
+    Returns:
+        The full captured body.
+    """
+    captured = bytearray()
+    client_gone = False
+    for chunk in chunks:
+        captured.extend(chunk)
+        if not client_gone:
+            try:
+                sink(chunk)
+            except OSError:
+                client_gone = True
+    return bytes(captured)
+
+
+def _https_forwarder(
+    base: str, path_map: PathMap | None = None
+) -> Callable[[str, str, Mapping[str, str], bytes], UpstreamResponse]:
+    parsed = urlsplit(base)
+    host = parsed.hostname or ""
+    port = parsed.port or 443
+
+    def forward(
+        method: str, path: str, headers: Mapping[str, str], body: bytes
+    ) -> UpstreamResponse:  # pragma: no cover  (real network; verified by a live round-trip)
+        connection = http.client.HTTPSConnection(host, port, timeout=600)
+        sent = {name: value for name, value in headers.items() if name.lower() not in _HOP_BY_HOP}
+        sent["Host"] = host
+        up_path = path_map(path) if path_map is not None else path
+        connection.request(method, up_path, body=body or None, headers=sent)
+        upstream = connection.getresponse()
+
+        def stream() -> Iterable[bytes]:
+            try:
+                while True:
+                    chunk = upstream.read(8192)
+                    if not chunk:
+                        break
+                    yield chunk
+            finally:
+                connection.close()
+
+        return UpstreamResponse(upstream.status, list(upstream.getheaders()), stream())
+
+    return forward

generic_ml_wrapper/adapter/outbound/interceptor/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""Built-in interceptors that plug into the ``[[interceptors]]`` chain."""

generic_ml_wrapper/adapter/outbound/interceptor/compressor.py

@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""A context compressor interceptor that compresses through generic-ml-cache."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+from generic_ml_cache_bootstrap.application import build_application_api
+from generic_ml_cache_core.application.domain.model.execution.artifact import ArtifactType
+from generic_ml_cache_core.application.domain.model.execution.execution_kind import ExecutionKind
+from generic_ml_cache_core.application.domain.model.run.cache_mode import CacheMode
+from generic_ml_cache_core.application.port.inbound.run_ml_execution import (
+    run_ml_execution_command as _run_cmd,
+)
+from generic_ml_cache_core.application.usecase.select_adapter_for_execution_service import (
+    SelectAdapterForExecutionService,
+)
+
+from generic_ml_wrapper.application.port.outbound.interceptor import InterceptorPort
+from generic_ml_wrapper.common import config, paths
+from generic_ml_wrapper.common.log import log
+
+if TYPE_CHECKING:
+    from generic_ml_cache_core.application.domain.model.execution.ml_execution import MlExecution
+    from generic_ml_cache_core.application.port.outbound.adapter_catalog_port import (
+        AdapterCatalogPort,
+    )
+    from generic_ml_cache_core.application.port.outbound.adapter_resolver_port import (
+        AdapterResolverPort,
+    )
+    from generic_ml_cache_core.application.port.outbound.registered_adapter_port import (
+        RegisteredAdapterPort,
+    )
+
+
+class CompressorInterceptor(InterceptorPort):
+    """Compress a context section through the generic-ml-cache record/replay cache.
+
+    Runs the configured model (default gpt-5.4 at low effort via the cursor adapter)
+    with the user's compression prompt, keyed and cached by content — so compressing
+    the same context replays for free and returns an identical result. Off until
+    ``[compress] prompt`` names a prompt file; non-destructive — any failure leaves
+    the text uncompressed.
+    """
+
+    def intercept(self, text: str, target: str) -> str:  # noqa: ARG002  (target-agnostic)
+        """Return the compressed text, or the input unchanged when off or on failure.
+
+        Args:
+            text: The section text to compress.
+            target: The target it is running for (unused; a compressor is target-agnostic).
+
+        Returns:
+            The compressed text, or ``text`` unchanged.
+        """
+        settings = config.compress()
+        if settings.prompt is None:
+            return text  # compression off until [compress] prompt is set
+        try:
+            prompt = Path(settings.prompt).read_text(encoding="utf-8")
+        except OSError as error:
+            log.warning(f"cannot read compress prompt {settings.prompt!r} ({error}); skipping")
+            return text
+        try:
+            execution = self._compress(text, prompt, settings)
+        except Exception as error:  # noqa: BLE001  (a compile must never die on the cache/LLM)
+            log.warning(f"compression failed ({error}); leaving context uncompressed")
+            return text
+        return _stdout(execution) or text
+
+    def _compress(  # pragma: no cover  (real cache/LLM round-trip; verified by a live run)
+        self, context: str, prompt: str, settings: config.CompressSettings
+    ) -> MlExecution:
+        command = _run_cmd.RunMlExecutionCommand(
+            execution_kind=ExecutionKind.LOCAL_MANAGED,
+            client=settings.adapter,
+            model=settings.model,
+            effort=settings.effort,
+            context=context,
+            prompt=prompt,
+            cache_mode=CacheMode.CACHE,
+        )
+
+        def runners(
+            catalog: AdapterCatalogPort, resolver: AdapterResolverPort
+        ) -> dict[str, RegisteredAdapterPort]:
+            descriptor = SelectAdapterForExecutionService(catalog).select(
+                settings.adapter, ExecutionKind.LOCAL_MANAGED
+            )
+            client = resolver.resolve_local_client(descriptor.adapter_id)
+            return {settings.adapter: cast("RegisteredAdapterPort", client)}
+
+        paths.COMPRESS_CACHE.mkdir(parents=True, exist_ok=True)
+        api = build_application_api(paths.COMPRESS_CACHE, runners)
+        return api.run_ml.execute(command)
+
+
+def _stdout(execution: MlExecution) -> str | None:
+    """Return the STDOUT artifact's text, or ``None`` on failure or if absent."""
+    if execution.failure is not None:
+        return None
+    for artifact in execution.artifacts:
+        if artifact.artifact_type is ArtifactType.STDOUT and artifact.content is not None:
+            return artifact.content.decode(artifact.encoding or "utf-8")
+    return None

generic_ml_wrapper/adapter/outbound/interceptor/size_logger.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""A reference interceptor that logs each message's size and returns it unchanged."""
+
+from __future__ import annotations
+
+from generic_ml_wrapper.application.port.outbound.interceptor import InterceptorPort
+from generic_ml_wrapper.common.log import log
+
+
+class MessageSizeLogger(InterceptorPort):
+    """Log the size of every intercepted message; a non-transforming observer.
+
+    A reference interceptor and the simplest example of the plugin contract: it logs
+    how many characters flow through each target it is bound to — e.g. ``request`` for
+    outbound prompts, ``response`` for the replies — and returns the text untouched.
+    Configure it under ``[[interceptors]]`` on the targets you want traced; put it on
+    both ``request`` and ``response`` to log sizes in and out.
+    """
+
+    def intercept(self, text: str, target: str) -> str:
+        """Log the message size for ``target`` and return the text unchanged.
+
+        Args:
+            text: The intercepted message.
+            target: The target it is running for (e.g. ``request``, ``response``).
+
+        Returns:
+            The input text, unchanged.
+        """
+        log.info(f"[interceptor] {target}: {len(text)} chars")
+        return text

generic_ml_wrapper/adapter/outbound/status/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""Client status-payload parsing adapters."""

generic_ml_wrapper/adapter/outbound/status/claude_status_parser.py

@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: 2026 Daniel Slobozian
+# SPDX-License-Identifier: Apache-2.0
+"""Parse Claude Code's status-line payload into a ``ClientStatus``."""
+
+from __future__ import annotations
+
+from typing import cast
+
+from generic_ml_wrapper.application.domain.model.client_status import ClientStatus
+from generic_ml_wrapper.application.port.outbound.client_status import ClientStatusParserPort
+
+
+def _dig(payload: object, *keys: str) -> object:
+    current: object = payload
+    for key in keys:
+        if not isinstance(current, dict):
+            return None
+        current = cast("dict[str, object]", current).get(key)
+    return current
+
+
+class ClaudeStatusParser(ClientStatusParserPort):
+    """Read model, context fill, and session cost from Claude Code's payload."""
+
+    def parse(self, payload: dict[str, object]) -> ClientStatus:
+        """Parse Claude Code's status payload.
+
+        Args:
+            payload: The decoded JSON Claude Code pipes to the status-line command
+                (``model.display_name``, ``context_window.used_percentage``,
+                ``cost.total_cost_usd``, and ``rate_limits`` buckets).
+
+        Returns:
+            The parsed status.
+        """
+        return ClientStatus(
+            model=_as_str(_dig(payload, "model", "display_name")),
+            context_pct=_as_pct(_dig(payload, "context_window", "used_percentage")),
+            session_cost_usd=_as_float(_dig(payload, "cost", "total_cost_usd")),
+            extras=_extras(payload),
+        )
+
+
+# Claude reports two rolling allowance windows; each answers "am I burning it?".
+_QUOTA_WINDOWS = (("5h", "five_hour"), ("wk", "seven_day"))
+
+
+def _extras(payload: dict[str, object]) -> tuple[str, ...]:
+    """Claude's client-specific status blocks (currently its rate-limit quota)."""
+    quota = _quota(_dig(payload, "rate_limits"))
+    return (f"quota {quota}",) if quota else ()
+
+
+def _quota(rate_limits: object) -> str | None:
+    parts: list[str] = []
+    for label, key in _QUOTA_WINDOWS:
+        pct = _as_pct(_dig(rate_limits, key, "used_percentage"))
+        if pct is not None:
+            parts.append(f"{label} {pct}%")
+    return " · ".join(parts) if parts else None
+
+
+def _as_str(value: object) -> str | None:
+    return value if isinstance(value, str) and value else None
+
+
+def _as_pct(value: object) -> int | None:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        return None
+    return round(value)
+
+
+def _as_float(value: object) -> float | None:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        return None
+    return float(value)