ctxoptimizer 0.4.1__py3-none-any.whl

ctxoptimizer/service.py

@@ -0,0 +1,150 @@
+"""Autostart service generation (SPEC section 21, P7-1).
+
+Generates a per-user launchd/systemd/scheduled-task registration. The pure
+``service_plan`` helper is testable without touching the host.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True, slots=True)
+class ServicePlan:
+    platform: str
+    path: Path
+    content: str
+    commands: tuple[tuple[str, ...], ...]
+    remove_commands: tuple[tuple[str, ...], ...]
+
+
+def service_plan(
+    *, platform: str | None = None, home: Path | None = None, config_path: str | None = None
+) -> ServicePlan:
+    name = platform or _platform()
+    home = home or Path.home()
+    args = [sys.executable, "-m", "ctxoptimizer", "run"]
+    if config_path:
+        args.extend(["--config", config_path])
+    command = " ".join(_quote(part) for part in args)
+
+    if name == "windows":
+        task = "ctxoptimizer"
+        return ServicePlan(
+            platform=name,
+            path=home / "AppData" / "Local" / "ctxoptimizer" / "ctxoptimizer-service.txt",
+            content=f"Scheduled task {task}: {command}\n",
+            commands=(
+                (
+                    "schtasks",
+                    "/Create",
+                    "/SC",
+                    "ONLOGON",
+                    "/TN",
+                    task,
+                    "/TR",
+                    command,
+                    "/F",
+                ),
+            ),
+            remove_commands=(("schtasks", "/Delete", "/TN", task, "/F"),),
+        )
+
+    if name == "darwin":
+        plist = home / "Library" / "LaunchAgents" / "com.ctxoptimizer.plist"
+        content = f"""<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0"><dict>
+  <key>Label</key><string>com.ctxoptimizer</string>
+  <key>ProgramArguments</key><array>
+    {"".join(f"<string>{_xml(part)}</string>" for part in args)}
+  </array>
+  <key>RunAtLoad</key><true/>
+  <key>KeepAlive</key><true/>
+</dict></plist>
+"""
+        return ServicePlan(
+            platform=name,
+            path=plist,
+            content=content,
+            commands=(("launchctl", "load", str(plist)),),
+            remove_commands=(("launchctl", "unload", str(plist)),),
+        )
+
+    unit = home / ".config" / "systemd" / "user" / "ctxoptimizer.service"
+    content = (
+        "[Unit]\n"
+        "Description=ctxoptimizer local LLM proxy\n"
+        "After=network-online.target\n\n"
+        "[Service]\n"
+        "Type=simple\n"
+        f"ExecStart={command}\n"
+        "Restart=on-failure\n\n"
+        "[Install]\n"
+        "WantedBy=default.target\n"
+    )
+    return ServicePlan(
+        platform="linux",
+        path=unit,
+        content=content,
+        commands=(
+            ("systemctl", "--user", "daemon-reload"),
+            ("systemctl", "--user", "enable", "--now", "ctxoptimizer.service"),
+        ),
+        remove_commands=(("systemctl", "--user", "disable", "--now", "ctxoptimizer.service"),),
+    )
+
+
+def install_service(
+    *,
+    platform: str | None = None,
+    home: Path | None = None,
+    config_path: str | None = None,
+    dry_run: bool = False,
+) -> ServicePlan:
+    plan = service_plan(platform=platform, home=home, config_path=config_path)
+    if dry_run:
+        return plan
+    plan.path.parent.mkdir(parents=True, exist_ok=True)
+    plan.path.write_text(plan.content, encoding="utf-8")
+    for command in plan.commands:
+        subprocess.run(command, check=False)  # noqa: S603
+    return plan
+
+
+def uninstall_service(
+    *,
+    platform: str | None = None,
+    home: Path | None = None,
+    dry_run: bool = False,
+) -> ServicePlan:
+    plan = service_plan(platform=platform, home=home)
+    if dry_run:
+        return plan
+    for command in plan.remove_commands:
+        subprocess.run(command, check=False)  # noqa: S603
+    plan.path.unlink(missing_ok=True)
+    return plan
+
+
+def _platform() -> str:
+    if os.name == "nt":
+        return "windows"
+    if sys.platform == "darwin":
+        return "darwin"
+    return "linux"
+
+
+def _quote(value: str) -> str:
+    return f'"{value}"' if " " in value else value
+
+
+def _xml(value: str) -> str:
+    return value.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
+
+
+__all__ = ["ServicePlan", "install_service", "service_plan", "uninstall_service"]

ctxoptimizer/telemetry/__init__.py

@@ -0,0 +1 @@
+"""Telemetry: usage parsing, pricing, storage and redaction (SPEC section 14)."""

ctxoptimizer/telemetry/blobs.py

@@ -0,0 +1,57 @@
+"""Content-addressed blob storage for omitted/duplicated originals (SPEC 12.13)."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import zlib
+from hashlib import sha256
+from pathlib import Path
+
+
+class BlobStore:
+    """zlib-compressed files named by sha256; supports short-ref lookup."""
+
+    def __init__(self, root: Path) -> None:
+        self.root = Path(root)
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def put(self, data: bytes) -> str:
+        digest = sha256(data).hexdigest()
+        path = self._path(digest)
+        if not path.exists():
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.write_bytes(zlib.compress(data, 6))
+            if os.name == "posix":
+                os.chmod(path, 0o600)
+        return digest
+
+    def get(self, ref: str) -> bytes | None:
+        if not ref:
+            return None
+        if len(ref) >= 64:
+            path = self._path(ref)
+        else:
+            matches = sorted(self.root.glob(f"*/{ref}*"))
+            if len(matches) != 1:
+                return None
+            path = matches[0]
+        if not path.is_file():
+            return None
+        try:
+            return zlib.decompress(path.read_bytes())
+        except zlib.error:
+            return None
+
+    def purge(self) -> None:
+        shutil.rmtree(self.root, ignore_errors=True)
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def count(self) -> int:
+        return sum(1 for _ in self.root.glob("*/*"))
+
+    def _path(self, digest: str) -> Path:
+        return self.root / digest[:2] / digest
+
+
+__all__ = ["BlobStore"]

ctxoptimizer/telemetry/expand.py

@@ -0,0 +1,62 @@
+"""Expand-usage log (SPEC 12.13).
+
+The MCP server runs in a separate process, so it writes to the same SQLite file
+with a tiny direct connection. Content is never logged, only the ref and result.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+import time
+from pathlib import Path
+from typing import Any
+
+from ctxoptimizer.telemetry.store import default_db_path
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS expand_events(
+  id INTEGER PRIMARY KEY AUTOINCREMENT,
+  ts REAL NOT NULL,
+  ref TEXT NOT NULL,
+  ok INTEGER NOT NULL
+);
+CREATE INDEX IF NOT EXISTS idx_expand_ts ON expand_events(ts);
+"""
+
+
+def record_expand(ref: str, ok: bool, *, db_path: Path | None = None) -> None:
+    path = db_path or default_db_path()
+    try:
+        with sqlite3.connect(path) as conn:
+            conn.executescript(_SCHEMA)
+            conn.execute(
+                "INSERT INTO expand_events(ts, ref, ok) VALUES(?, ?, ?)",
+                (time.time(), ref, 1 if ok else 0),
+            )
+            conn.commit()
+    except sqlite3.Error:
+        # Expand must still return its result if telemetry is unavailable.
+        return
+
+
+def expand_stats(*, db_path: Path | None = None) -> dict[str, Any]:
+    path = db_path or default_db_path()
+    try:
+        with sqlite3.connect(path) as conn:
+            conn.executescript(_SCHEMA)
+            row = conn.execute(
+                "SELECT COUNT(*), COALESCE(SUM(ok), 0), MAX(ts) FROM expand_events"
+            ).fetchone()
+    except sqlite3.Error:
+        return {"total": 0, "ok": 0, "errors": 0, "last_ts": None}
+    total = int(row[0])
+    ok = int(row[1])
+    return {
+        "total": total,
+        "ok": ok,
+        "errors": total - ok,
+        "last_ts": row[2],
+    }
+
+
+__all__ = ["expand_stats", "record_expand"]

ctxoptimizer/telemetry/pricing.py

@@ -0,0 +1,168 @@
+"""Model pricing lookup and cost math (SPEC section 14.2).
+
+Provider-reported cost (OpenRouter) wins when present. Otherwise the bundled
+table is used; an unknown model or a missing rate yields ``priced=False`` and a
+NULL cost rather than an invented number (INV-13).
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from fnmatch import fnmatchcase
+from importlib import resources
+from pathlib import Path
+from typing import Any, Literal
+
+from ctxoptimizer.models import Usage
+
+PriceSource = Literal["provider", "pricing", "none"]
+
+
+@dataclass(frozen=True, slots=True)
+class ModelPrice:
+    provider: str
+    match: str
+    input: float
+    output: float
+    cache_read: float = 0.0
+    cache_write_5m: float | None = None
+    cache_write_1h: float | None = None
+    cache_write_30m: float | None = None
+    cache_write_12h: float | None = None
+    verified: bool = False
+    source: str = ""
+    notes: str = ""
+
+
+@dataclass(frozen=True, slots=True)
+class PriceResult:
+    cost_usd: float | None
+    priced: bool
+    source: PriceSource
+    price: ModelPrice | None
+
+
+class PricingTable:
+    def __init__(
+        self,
+        models: list[ModelPrice],
+        *,
+        per: int = 1_000_000,
+        currency: str = "USD",
+        version: str = "",
+    ) -> None:
+        self.models = tuple(models)
+        self.per = per
+        self.currency = currency
+        self.version = version
+
+    def lookup(self, provider: str, model: str) -> ModelPrice | None:
+        candidates = [
+            price
+            for price in self.models
+            if price.provider == provider and fnmatchcase(model, price.match)
+        ]
+        if not candidates:
+            return None
+        return max(candidates, key=lambda price: (_literal_prefix(price.match), len(price.match)))
+
+    def cost(
+        self,
+        usage: Usage,
+        price: ModelPrice | None,
+        *,
+        prefer_provider_cost: bool = True,
+    ) -> PriceResult:
+        if prefer_provider_cost and usage.provider_cost_usd is not None:
+            return PriceResult(usage.provider_cost_usd, True, "provider", price)
+        if price is None:
+            return PriceResult(None, False, "none", None)
+        total = 0.0
+        buckets = (
+            (usage.fresh_input, price.input),
+            (usage.cache_read, price.cache_read),
+            (usage.cache_write_5m, price.cache_write_5m),
+            (usage.cache_write_1h, price.cache_write_1h),
+            (usage.cache_write_30m, price.cache_write_30m),
+            (usage.cache_write_12h, price.cache_write_12h),
+            (usage.output, price.output),
+        )
+        for tokens, rate in buckets:
+            if not tokens:
+                continue
+            if rate is None:
+                return PriceResult(None, False, "none", price)
+            total += tokens * rate
+        return PriceResult(total / self.per, True, "pricing", price)
+
+
+def load_pricing(user_file: Path | None = None) -> PricingTable:
+    """Load the bundled table and merge an optional user override."""
+    bundled = _read_json(
+        resources.files("ctxoptimizer").joinpath("data", "pricing.json").read_text("utf-8")
+    )
+    models = [_parse_price(raw) for raw in bundled.get("models", [])]
+    if user_file is not None:
+        user = _read_json(user_file.read_text(encoding="utf-8"))
+        overrides = {(entry.provider, entry.match): entry for entry in models}
+        for raw in user.get("models", []):
+            parsed = _parse_price(raw)
+            overrides[(parsed.provider, parsed.match)] = parsed
+        models = list(overrides.values())
+    return PricingTable(
+        models,
+        per=int(bundled.get("per", 1_000_000)),
+        currency=str(bundled.get("currency", "USD")),
+        version=str(bundled.get("version", "")),
+    )
+
+
+def price_usage(usage: Usage, table: PricingTable, provider: str, model: str | None) -> PriceResult:
+    """Price one request. Unknown models stay unpriced, never guessed."""
+    price = table.lookup(provider, model) if model else None
+    return table.cost(usage, price)
+
+
+def _parse_price(raw: dict[str, Any]) -> ModelPrice:
+    return ModelPrice(
+        provider=str(raw["provider"]),
+        match=str(raw["match"]),
+        input=float(raw["input"]),
+        output=float(raw["output"]),
+        cache_read=float(raw.get("cache_read", 0.0)),
+        cache_write_5m=_optional_float(raw.get("cache_write_5m")),
+        cache_write_1h=_optional_float(raw.get("cache_write_1h")),
+        cache_write_30m=_optional_float(raw.get("cache_write_30m")),
+        cache_write_12h=_optional_float(raw.get("cache_write_12h")),
+        verified=bool(raw.get("verified", False)),
+        source=str(raw.get("source", "")),
+        notes=str(raw.get("notes", "")),
+    )
+
+
+def _read_json(text: str) -> dict[str, Any]:
+    value = json.loads(text)
+    if not isinstance(value, dict):
+        raise ValueError("pricing file must contain a JSON object")
+    return value
+
+
+def _optional_float(value: Any) -> float | None:
+    return None if value is None else float(value)
+
+
+def _literal_prefix(pattern: str) -> int:
+    for index, char in enumerate(pattern):
+        if char in "*?[":
+            return index
+    return len(pattern)
+
+
+__all__ = [
+    "ModelPrice",
+    "PriceResult",
+    "PricingTable",
+    "load_pricing",
+    "price_usage",
+]

ctxoptimizer/telemetry/redact.py

@@ -0,0 +1,53 @@
+"""Secret redaction for stored copies (SPEC section 14.4).
+
+Redaction applies only to blobs/exports written to disk, never to forwarded
+traffic. Patterns are linear-time and conservative: when in doubt, redact.
+"""
+
+from __future__ import annotations
+
+import re
+
+_PATTERNS: tuple[re.Pattern[str], ...] = (
+    re.compile(r"sk-ant-[A-Za-z0-9_\-]{8,}"),
+    re.compile(r"sk-[A-Za-z0-9_\-]{16,}"),
+    re.compile(r"AIza[0-9A-Za-z_\-]{20,}"),
+    re.compile(r"gh[pousr]_[A-Za-z0-9]{20,}"),
+    re.compile(r"github_pat_[A-Za-z0-9_\-]{20,}"),
+    re.compile(r"AKIA[0-9A-Z]{16}"),
+    re.compile(r"xox[baprs]-[A-Za-z0-9\-]{10,}"),
+    re.compile(
+        r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----",
+        re.DOTALL,
+    ),
+    re.compile(r"(?i)\bBearer\s+[A-Za-z0-9._\-+/=]{8,}"),
+    re.compile(
+        r"(?im)^([ \t]*(?:export[ \t]+)?[A-Za-z0-9_.-]*"
+        r"(?:secret|token|password|passwd|api[_-]?key)[A-Za-z0-9_.-]*[ \t]*=[ \t]*)"
+        r"(\"[^\"]*\"|'[^']*'|[^\s#]+)"
+    ),
+)
+
+_REDACTED = "[REDACTED]"
+
+
+def redact_text(text: str) -> str:
+    """Return *text* with recognized credentials replaced."""
+    if not text:
+        return text
+    for pattern in _PATTERNS[:-1]:
+        text = pattern.sub(_REDACTED, text)
+    text = _PATTERNS[-1].sub(lambda match: match.group(1) + _REDACTED, text)
+    return text
+
+
+def redact_bytes(data: bytes) -> bytes:
+    """Best-effort UTF-8 decoding so binary blobs are never damaged."""
+    try:
+        text = data.decode("utf-8")
+    except UnicodeDecodeError:
+        return data
+    return redact_text(text).encode("utf-8")
+
+
+__all__ = ["redact_bytes", "redact_text"]

ctxoptimizer/telemetry/sse.py

@@ -0,0 +1,145 @@
+"""Incremental Server-Sent Events decoder (WHATWG event-stream grammar).
+
+Byte-buffered so a chunk can split a line or a multi-byte UTF-8 character, and
+bounded so a malicious stream cannot grow memory without limit. The decoder only
+frames events; provider parsers interpret the JSON.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+DEFAULT_MAX_LINE_BYTES = 4 * 1024 * 1024
+DEFAULT_MAX_EVENT_BYTES = 8 * 1024 * 1024
+
+
+class SSEFrameError(ValueError):
+    """Raised when a frame exceeds the configured memory bound."""
+
+
+@dataclass(frozen=True, slots=True)
+class SSEEvent:
+    event: str | None
+    data: str
+    id: str | None = None
+    retry: int | None = None
+
+
+class SSEDecoder:
+    """Feed it bytes, receive complete events. Tolerates CR, LF and CRLF."""
+
+    def __init__(
+        self,
+        *,
+        max_line_bytes: int = DEFAULT_MAX_LINE_BYTES,
+        max_event_bytes: int = DEFAULT_MAX_EVENT_BYTES,
+    ) -> None:
+        self._buffer = bytearray()
+        self._event: str | None = None
+        self._data: list[str] = []
+        self._data_bytes = 0
+        self._last_id: str | None = None
+        self._retry: int | None = None
+        self._first_line = True
+        self._max_line = max_line_bytes
+        self._max_event = max_event_bytes
+
+    def feed(self, chunk: bytes) -> list[SSEEvent]:
+        """Append bytes and return every event completed by this chunk."""
+        self._buffer.extend(chunk)
+        events: list[SSEEvent] = []
+        while True:
+            line = self._pop_line()
+            if line is None:
+                break
+            event = self._process_line(line)
+            if event is not None:
+                events.append(event)
+        if len(self._buffer) > self._max_line:
+            raise SSEFrameError("SSE line exceeds the configured limit")
+        return events
+
+    def finish(self) -> list[SSEEvent]:
+        """Flush a trailing partial line and any pending data at EOF."""
+        events: list[SSEEvent] = []
+        if self._buffer:
+            raw = bytes(self._buffer).rstrip(b"\r")
+            self._buffer.clear()
+            event = self._process_line(raw)
+            if event is not None:
+                events.append(event)
+        event = self._dispatch()
+        if event is not None:
+            events.append(event)
+        return events
+
+    def _pop_line(self) -> bytes | None:
+        cr = self._buffer.find(b"\r")
+        lf = self._buffer.find(b"\n")
+        candidates = [index for index in (cr, lf) if index >= 0]
+        if not candidates:
+            return None
+        index = min(candidates)
+        if self._buffer[index] == 13:
+            if index + 1 == len(self._buffer):
+                return None
+            end = index + 2 if self._buffer[index + 1] == 10 else index + 1
+        else:
+            end = index + 1
+        line = bytes(self._buffer[:index])
+        del self._buffer[:end]
+        return line
+
+    def _process_line(self, raw: bytes) -> SSEEvent | None:
+        line = raw.decode("utf-8", errors="replace")
+        if self._first_line:
+            line = line.removeprefix("\ufeff")
+            self._first_line = False
+        if line == "":
+            return self._dispatch()
+        if line.startswith(":"):
+            return None
+        field, separator, value = line.partition(":")
+        if separator and value.startswith(" "):
+            value = value[1:]
+        if field == "event":
+            self._event = value
+        elif field == "data":
+            self._data.append(value)
+            self._data_bytes += len(value.encode("utf-8"))
+            if self._data_bytes > self._max_event:
+                raise SSEFrameError("SSE event exceeds the configured limit")
+        elif field == "id":
+            if "\0" not in value:
+                self._last_id = value
+        elif field == "retry":
+            try:
+                self._retry = int(value)
+            except ValueError:
+                pass
+        return None
+
+    def _dispatch(self) -> SSEEvent | None:
+        if not self._data and self._event is None:
+            self._retry = None
+            return None
+        event = SSEEvent(
+            event=self._event,
+            data="\n".join(self._data),
+            id=self._last_id,
+            retry=self._retry,
+        )
+        self._event = None
+        self._data = []
+        self._data_bytes = 0
+        self._retry = None
+        return event
+
+
+__all__ = [
+    "DEFAULT_MAX_EVENT_BYTES",
+    "DEFAULT_MAX_LINE_BYTES",
+    "SSEDecoder",
+    "SSEEvent",
+    "SSEFrameError",
+]