ctxoptimizer 0.4.1__py3-none-any.whl

ctxoptimizer/ids.py

@@ -0,0 +1,73 @@
+"""Time-sortable ULID identifiers, using only the standard library.
+
+A ULID is 26 characters of Crockford base32: a 48-bit millisecond timestamp
+followed by 80 bits of randomness. ULIDs sort lexicographically by creation
+time, which keeps request ids cheap to order in SQLite without a separate
+counter.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+import time
+
+CROCKFORD = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
+ULID_LENGTH = 26
+_TIMESTAMP_CHARS = 10
+_RANDOM_CHARS = 16
+_RANDOM_BITS = 80
+
+_lock = threading.Lock()
+_last_ms = 0
+_last_random = 0
+
+
+def new_ulid(*, now_ms: int | None = None) -> str:
+    """Return a new ULID, monotonic within a single millisecond.
+
+    Passing ``now_ms`` is intended for tests; production callers omit it.
+    """
+    global _last_ms, _last_random
+    with _lock:
+        current = int(time.time() * 1000) if now_ms is None else now_ms
+        if current > _last_ms:
+            _last_ms = current
+            _last_random = int.from_bytes(os.urandom(10), "big")
+        else:
+            _last_random += 1
+            if _last_random >= 1 << _RANDOM_BITS:
+                _last_ms += 1
+                _last_random = 0
+        timestamp = _last_ms
+        randomness = _last_random
+    return _encode(timestamp, _TIMESTAMP_CHARS) + _encode(randomness, _RANDOM_CHARS)
+
+
+def ulid_timestamp_ms(value: str) -> int:
+    """Return the millisecond timestamp encoded in *value*."""
+    if len(value) != ULID_LENGTH:
+        raise ValueError(f"ULID must be {ULID_LENGTH} characters, got {len(value)}")
+    return _decode(value[:_TIMESTAMP_CHARS])
+
+
+def _encode(value: int, length: int) -> str:
+    digits: list[str] = []
+    for _ in range(length):
+        digits.append(CROCKFORD[value & 0x1F])
+        value >>= 5
+    return "".join(reversed(digits))
+
+
+def _decode(value: str) -> int:
+    result = 0
+    for char in value:
+        try:
+            digit = CROCKFORD.index(char)
+        except ValueError as exc:
+            raise ValueError(f"invalid Crockford base32 character: {char!r}") from exc
+        result = (result << 5) | digit
+    return result
+
+
+__all__ = ["CROCKFORD", "ULID_LENGTH", "new_ulid", "ulid_timestamp_ms"]

ctxoptimizer/log.py

@@ -0,0 +1,73 @@
+"""Logging setup and secret sanitization (SPEC section 17, SEC-2).
+
+Nothing in this proxy may write credentials to a log. Call
+:func:`sanitize_headers` and :func:`sanitize_url` at every log site that can
+see client traffic; the values stay untouched on the forwarding path.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+REDACTED = "REDACTED"
+
+SENSITIVE_HEADERS = frozenset(
+    {
+        "authorization",
+        "proxy-authorization",
+        "x-api-key",
+        "x-goog-api-key",
+        "api-key",
+        "cookie",
+        "set-cookie",
+        "x-ctxoptimizer-admin-token",
+    }
+)
+
+SENSITIVE_QUERY_KEYS = frozenset({"key", "api_key", "apikey", "access_token", "token", "auth"})
+
+
+def sanitize_headers(headers: Mapping[str, str]) -> dict[str, str]:
+    """Return a copy of *headers* with credential values redacted."""
+    return {
+        name: REDACTED if name.lower() in SENSITIVE_HEADERS else value
+        for name, value in headers.items()
+    }
+
+
+def sanitize_url(url: str) -> str:
+    """Return a log-safe form of *url* with credentials and keys redacted."""
+    parts = urlsplit(url)
+    netloc = parts.netloc
+    if "@" in netloc:
+        netloc = f"{REDACTED}@{netloc.rsplit('@', 1)[1]}"
+    query = [
+        (name, REDACTED if name.lower() in SENSITIVE_QUERY_KEYS else value)
+        for name, value in parse_qsl(parts.query, keep_blank_values=True)
+    ]
+    return urlunsplit((parts.scheme, netloc, parts.path, urlencode(query), parts.fragment))
+
+
+def configure_logging(level: str = "info") -> None:
+    """Configure the root logger once, at the requested level."""
+    numeric = logging.getLevelNamesMapping().get(level.upper())
+    if numeric is None:
+        raise ValueError(f"unknown log level: {level!r}")
+    root = logging.getLogger()
+    if not root.handlers:
+        handler = logging.StreamHandler()
+        handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s"))
+        root.addHandler(handler)
+    root.setLevel(numeric)
+
+
+__all__ = [
+    "REDACTED",
+    "SENSITIVE_HEADERS",
+    "SENSITIVE_QUERY_KEYS",
+    "configure_logging",
+    "sanitize_headers",
+    "sanitize_url",
+]

ctxoptimizer/mcp.py

@@ -0,0 +1,124 @@
+"""Minimal MCP stdio server exposing ``ctxoptimizer_expand`` (SPEC 12.13)."""
+
+from __future__ import annotations
+
+import json
+import sys
+from collections.abc import Callable
+from typing import Any
+
+from ctxoptimizer import __version__
+from ctxoptimizer.analysis.tokens import estimate_tokens
+from ctxoptimizer.config import data_dir
+from ctxoptimizer.telemetry.blobs import BlobStore
+from ctxoptimizer.telemetry.expand import record_expand
+
+_PROTOCOL_VERSION = "2024-11-05"
+
+
+def handle_message(
+    message: dict[str, Any],
+    blobs: BlobStore,
+    on_expand: Callable[[str, bool], None] | None = None,
+) -> dict[str, Any] | None:
+    """Handle one JSON-RPC message. Notifications return None."""
+    request_id = message.get("id")
+    method = message.get("method")
+    if request_id is None:
+        return None
+    if method == "initialize":
+        result: dict[str, Any] = {
+            "protocolVersion": _PROTOCOL_VERSION,
+            "capabilities": {"tools": {}},
+            "serverInfo": {"name": "ctxoptimizer", "version": __version__},
+        }
+        return _result(request_id, result)
+    if method == "tools/list":
+        return _result(
+            request_id,
+            {
+                "tools": [
+                    {
+                        "name": "ctxoptimizer_expand",
+                        "description": ("Retrieve the original bytes behind a ctxoptimizer ref."),
+                        "inputSchema": {
+                            "type": "object",
+                            "properties": {
+                                "ref": {"type": "string"},
+                                "max_tokens": {"type": "integer"},
+                            },
+                            "required": ["ref"],
+                        },
+                    }
+                ]
+            },
+        )
+    if method == "tools/call":
+        params = message.get("params") or {}
+        if params.get("name") != "ctxoptimizer_expand":
+            return _error(request_id, -32602, "unknown tool")
+        arguments = params.get("arguments") or {}
+        ref = str(arguments.get("ref") or "")
+        data = blobs.get(ref)
+        if data is None:
+            if on_expand is not None:
+                on_expand(ref, False)
+            return _result(
+                request_id,
+                {
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": (
+                                f"error: no stored blob for ref {ref!r}; "
+                                "re-run the tool that produced it"
+                            ),
+                        }
+                    ],
+                    "isError": True,
+                },
+            )
+        text = data.decode("utf-8", "replace")
+        max_tokens = arguments.get("max_tokens")
+        if isinstance(max_tokens, int) and max_tokens > 0 and estimate_tokens(text) > max_tokens:
+            head = text[: max_tokens * 2]
+            tail = text[-max_tokens * 2 :]
+            text = f"{head}\n[ctxoptimizer: middle omitted; ref={ref}]\n{tail}"
+        if on_expand is not None:
+            on_expand(ref, True)
+        return _result(request_id, {"content": [{"type": "text", "text": text}]})
+    return _error(request_id, -32601, f"unknown method: {method}")
+
+
+def run_stdio() -> None:
+    """Serve JSON-RPC over stdin/stdout until EOF."""
+    blobs = BlobStore(data_dir() / "blobs")
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            message = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if not isinstance(message, dict):
+            continue
+        response = handle_message(message, blobs, on_expand=record_expand)
+        if response is not None:
+            sys.stdout.write(json.dumps(response) + "\n")
+            sys.stdout.flush()
+
+
+def _result(request_id: Any, result: dict[str, Any]) -> dict[str, Any]:
+    return {"jsonrpc": "2.0", "id": request_id, "result": result}
+
+
+def _error(request_id: Any, code: int, message: str) -> dict[str, Any]:
+    return {
+        "jsonrpc": "2.0",
+        "id": request_id,
+        "error": {"code": code, "message": message},
+    }
+
+
+__all__ = ["handle_message", "run_stdio"]

ctxoptimizer/models.py

@@ -0,0 +1,145 @@
+"""Core dataclasses and enums (SPEC section 10).
+
+These are the provider-agnostic shapes every layer passes around. Payload
+mutation stays behind :class:`ToolResultRef`, so transforms never need to know
+a provider's JSON layout (INV-7).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Literal
+
+from ctxoptimizer.config import Config
+
+
+class Provider(StrEnum):
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    GEMINI = "gemini"
+    OLLAMA = "ollama"
+
+
+class Endpoint(StrEnum):
+    MESSAGES = "messages"
+    CHAT = "chat"
+    RESPONSES = "responses"
+    GENERATE = "generate"
+    PASSTHROUGH = "passthrough"
+
+
+class Mode(StrEnum):
+    PASSTHROUGH = "passthrough"
+    SHADOW = "shadow"
+    ACTIVE = "active"
+
+
+@dataclass(frozen=True, slots=True)
+class RequestContext:
+    """Everything the pipeline needs about one request, frozen at ingress."""
+
+    request_id: str
+    provider: Provider
+    endpoint: Endpoint
+    upstream_name: str
+    model: str | None
+    stream: bool
+    mode: Mode
+    session_id: str
+    project: str | None
+    client_name: str | None
+    config: Config
+    config_hash: str
+    started_monotonic: float
+
+
+@dataclass(slots=True)
+class ToolResultRef:
+    """Mutable handle to one tool-result text block (INV-7: text only).
+
+    Adapters implement :meth:`get_text` and :meth:`set_text` for their payload
+    shape. Everything downstream treats the block as an opaque string.
+    """
+
+    message_index: int
+    block_index: int | None
+    tool_use_id: str
+    tool_name: str | None
+    is_error: bool
+
+    def get_text(self) -> str | None:
+        """Return the text, or None when the block is not text."""
+        raise NotImplementedError
+
+    def set_text(self, new_text: str) -> None:
+        """Replace the block text in place."""
+        raise NotImplementedError
+
+
+@dataclass(slots=True)
+class TransformReport:
+    """Outcome of one transform run, persisted per request."""
+
+    transform_id: str
+    ran: bool = False
+    changed: bool = False
+    items_touched: int = 0
+    bytes_before: int = 0
+    bytes_after: int = 0
+    est_tokens_before: int = 0
+    est_tokens_after: int = 0
+    duration_ms: float = 0.0
+    discarded_reason: str | None = None
+    error: str | None = None
+    notes: dict[str, object] = field(default_factory=dict)
+
+
+@dataclass(slots=True)
+class Usage:
+    """Normalized provider usage (SPEC section 14.1)."""
+
+    fresh_input: int | None = None
+    cache_read: int | None = None
+    cache_write_5m: int | None = None
+    cache_write_1h: int | None = None
+    cache_write_30m: int | None = None
+    cache_write_12h: int | None = None
+    output: int | None = None
+    reasoning: int | None = None
+    provider_cost_usd: float | None = None
+    source: Literal["provider", "estimated", "none"] = "none"
+    raw: dict[str, object] | None = None
+    assistant_text_hashes: list[bytes] | None = None
+
+
+@dataclass(slots=True)
+class WasteFinding:
+    """An advisory, evidence-linked waste finding (INV-16, INV-18)."""
+
+    finding_id: str
+    ts: float
+    session_id: str
+    project: str | None
+    detector: str
+    key_hash: bytes
+    request_id: str | None
+    occurrences: int
+    first_seen_ts: float
+    last_seen_ts: float
+    evidence: dict[str, object]
+    est_tokens_wasted: int | None
+    est_cost_usd: float | None
+    status: Literal["new", "ack", "fixed", "skipped"] = "new"
+
+
+__all__ = [
+    "Endpoint",
+    "Mode",
+    "Provider",
+    "RequestContext",
+    "ToolResultRef",
+    "TransformReport",
+    "Usage",
+    "WasteFinding",
+]

ctxoptimizer/pipeline.py

@@ -0,0 +1,126 @@
+"""Transform pipeline: ordering, journal, budget and no-regress (SPEC 12.7)."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from ctxoptimizer.analysis.segments import Segment
+from ctxoptimizer.config import Config
+from ctxoptimizer.models import Mode, RequestContext, ToolResultRef, TransformReport
+from ctxoptimizer.providers.base import ProviderAdapter
+from ctxoptimizer.transforms.base import Transform
+
+_CACHE_MARKER_TRANSFORMS = frozenset({"anthropic_cache", "openai_hints"})
+
+
+@dataclass(frozen=True, slots=True)
+class PipelineResult:
+    reports: list[TransformReport]
+    original_segments: list[Segment]
+    final_segments: list[Segment]
+    orig_est_tokens: int
+    final_est_tokens: int
+    changed: bool
+    shadow: bool
+    changed_blocks: list[tuple[str, str]]
+    original_texts: list[tuple[str | None, str | None, str]]
+
+
+def run_pipeline(
+    ctx: RequestContext,
+    payload: dict[str, Any],
+    adapter: ProviderAdapter,
+    transforms: Sequence[Transform],
+    config: Config,
+) -> PipelineResult:
+    """Run transforms in order, reverting failures and non-saving changes."""
+    original_segments = adapter.segments(payload)
+    original_texts: list[tuple[str | None, str | None, str]] = [
+        (ref.tool_use_id, ref.tool_name, text)
+        for ref in adapter.iter_tool_results(payload)
+        if (text := ref.get_text()) is not None
+    ]
+    orig_tokens = _token_total(original_segments)
+    reports: list[TransformReport] = []
+    changed_blocks: list[tuple[str, str]] = []
+    started = time.perf_counter()
+    budget = config.pipeline.time_budget_ms
+
+    for transform in transforms:
+        elapsed_ms = (time.perf_counter() - started) * 1000
+        if elapsed_ms > budget:
+            reports.append(
+                TransformReport(transform_id=transform.id, ran=False, discarded_reason="budget")
+            )
+            continue
+        if not transform.applies(ctx, payload, adapter):
+            continue
+        before_tokens = _token_total(adapter.segments(payload))
+        snapshot = _snapshot(adapter, payload)
+        try:
+            report = transform.apply(ctx, payload, adapter, config)
+        except Exception as exc:
+            _restore(snapshot)
+            reports.append(
+                TransformReport(
+                    transform_id=transform.id,
+                    ran=True,
+                    discarded_reason="error",
+                    error=str(exc),
+                )
+            )
+            continue
+        after_tokens = _token_total(adapter.segments(payload))
+        if (
+            report.changed
+            and transform.id not in _CACHE_MARKER_TRANSFORMS
+            and before_tokens - after_tokens < config.pipeline.min_saving_tokens
+        ):
+            _restore(snapshot)
+            report.changed = False
+            report.discarded_reason = "no_regress"
+            report.est_tokens_after = report.est_tokens_before
+            report.bytes_after = report.bytes_before
+        if report.changed and report.discarded_reason is None:
+            for ref, previous in snapshot:
+                current = ref.get_text()
+                if current is not None and current != previous:
+                    changed_blocks.append((previous, current))
+        reports.append(report)
+
+    final_segments = adapter.segments(payload)
+    return PipelineResult(
+        reports=reports,
+        original_segments=original_segments,
+        final_segments=final_segments,
+        orig_est_tokens=orig_tokens,
+        final_est_tokens=_token_total(final_segments),
+        changed=any(report.changed and report.discarded_reason is None for report in reports),
+        shadow=ctx.mode is Mode.SHADOW,
+        changed_blocks=changed_blocks,
+        original_texts=original_texts,
+    )
+
+
+def _snapshot(adapter: ProviderAdapter, payload: dict[str, Any]) -> list[tuple[ToolResultRef, str]]:
+    snapshot: list[tuple[ToolResultRef, str]] = []
+    for ref in adapter.iter_tool_results(payload):
+        text = ref.get_text()
+        if text is not None:
+            snapshot.append((ref, text))
+    return snapshot
+
+
+def _restore(snapshot: list[tuple[ToolResultRef, str]]) -> None:
+    for ref, text in snapshot:
+        ref.set_text(text)
+
+
+def _token_total(segments: list[Segment]) -> int:
+    return sum(segment.est_tokens for segment in segments)
+
+
+__all__ = ["PipelineResult", "run_pipeline"]

ctxoptimizer/profiles.py

@@ -0,0 +1,152 @@
+"""Per-upstream profiles (ADR 0006).
+
+A profile answers the questions the proxy needs in order to treat a host
+correctly: which usage shape it speaks, how its cache behaves, whether cost comes
+from our pricing table or the provider, and whether cache hints may be injected.
+
+Protocol detection stays in :mod:`ctxoptimizer.routing`; this module only resolves
+names to URLs and profiles.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from urllib.parse import urlsplit
+
+from ctxoptimizer.config import (
+    CachePolicy,
+    Config,
+    CostSource,
+    UpstreamEntry,
+    UpstreamKind,
+    UsageStyle,
+)
+
+
+class UnknownUpstream(LookupError):
+    """Raised when a named upstream is not configured."""
+
+
+@dataclass(frozen=True, slots=True)
+class UpstreamProfile:
+    kind: UpstreamKind | None
+    usage: UsageStyle
+    cache_policy: CachePolicy
+    cost_source: CostSource
+    cache_hints: bool
+    inject_stream_usage: bool
+    pricing_provider: str
+
+
+@dataclass(frozen=True, slots=True)
+class ResolvedUpstream:
+    name: str
+    url: str
+    profile: UpstreamProfile
+
+
+BUILTIN_PROFILES: dict[str, UpstreamProfile] = {
+    "anthropic": UpstreamProfile(
+        "anthropic",
+        "anthropic",
+        "explicit-anthropic",
+        "pricing",
+        True,
+        False,
+        "anthropic",
+    ),
+    "openai": UpstreamProfile(
+        "openai", "openai", "implicit-openai", "pricing", True, False, "openai"
+    ),
+    "gemini": UpstreamProfile(
+        "gemini", "gemini", "gemini-implicit", "pricing", False, False, "gemini"
+    ),
+    "ollama": UpstreamProfile("ollama", "ollama", "none", "none", False, False, "ollama"),
+    "deepseek": UpstreamProfile(
+        "openai", "deepseek", "automatic-deepseek", "pricing", False, False, "deepseek"
+    ),
+    "deepseek-anthropic": UpstreamProfile(
+        "anthropic",
+        "anthropic",
+        "automatic-deepseek",
+        "pricing",
+        False,
+        False,
+        "deepseek",
+    ),
+    "openrouter": UpstreamProfile(
+        "openai", "openrouter", "implicit-openai", "provider", False, False, "openrouter"
+    ),
+    "gemini-openai": UpstreamProfile(
+        "openai", "gemini-openai", "gemini-implicit", "pricing", False, False, "gemini"
+    ),
+    "ollama-openai": UpstreamProfile(
+        "openai", "ollama-openai", "none", "none", False, False, "ollama"
+    ),
+}
+
+
+def resolve_upstream(name: str, config: Config) -> ResolvedUpstream:
+    """Resolve a configured upstream name to its URL and effective profile."""
+    entry = config.upstreams.get(name)
+    if entry is None:
+        raise UnknownUpstream(f"unknown upstream: {name}")
+    if isinstance(entry, str):
+        profile = BUILTIN_PROFILES.get(name) or _default_profile(None, entry, name)
+        return ResolvedUpstream(name=name, url=entry, profile=profile)
+    profile = _profile_for_entry(name, entry)
+    return ResolvedUpstream(name=name, url=entry.url, profile=profile)
+
+
+def _profile_for_entry(name: str, entry: UpstreamEntry) -> UpstreamProfile:
+    builtin = BUILTIN_PROFILES.get(name)
+    base = builtin or _default_profile(entry.kind, entry.url, name)
+    return UpstreamProfile(
+        kind=entry.kind,
+        usage=entry.usage or base.usage,
+        cache_policy=entry.cache_policy or base.cache_policy,
+        cost_source=entry.cost_source or base.cost_source,
+        cache_hints=base.cache_hints if entry.cache_hints is None else entry.cache_hints,
+        inject_stream_usage=(
+            base.inject_stream_usage
+            if entry.inject_stream_usage is None
+            else entry.inject_stream_usage
+        ),
+        pricing_provider=entry.pricing_provider or base.pricing_provider,
+    )
+
+
+def _default_profile(kind: UpstreamKind | None, url: str, name: str) -> UpstreamProfile:
+    host = (urlsplit(url).hostname or "").lower()
+    if kind == "anthropic":
+        return UpstreamProfile(
+            kind,
+            "anthropic",
+            "explicit-anthropic",
+            "pricing",
+            host == "api.anthropic.com",
+            False,
+            name,
+        )
+    if kind == "gemini":
+        return UpstreamProfile(kind, "gemini", "gemini-implicit", "pricing", False, False, name)
+    if kind == "ollama":
+        return UpstreamProfile(kind, "ollama", "none", "none", False, False, name)
+    return UpstreamProfile(
+        kind,
+        "openai",
+        "implicit-openai",
+        "pricing",
+        host == "api.openai.com",
+        False,
+        name,
+    )
+
+
+__all__ = [
+    "BUILTIN_PROFILES",
+    "ResolvedUpstream",
+    "UnknownUpstream",
+    "UpstreamProfile",
+    "resolve_upstream",
+]