token-runtime 0.1.0a2__py3-none-any.whl

token_runtime/gemini_conformance.py

@@ -0,0 +1,351 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from copy import deepcopy
+from dataclasses import dataclass, replace
+
+from .capabilities import (
+    CapabilityDetector,
+    CapabilityKey,
+    CapabilityProfile,
+    CapabilityRegistry,
+)
+from .compatibility import (
+    CompatibilityRecord,
+    CompatibilityRegistry,
+    CompatibilityState,
+)
+from .conformance import ConformanceResult, require_conformance, run_conformance
+from .gemini_adapter import GeminiGenerateContentAdapter
+
+
+GEMINI_GENERATE_CONTENT_KEY = CapabilityKey(
+    "generic-google",
+    "gemini_generate_content",
+    "google",
+)
+EVIDENCE_ID = "token-gemini-adapter-1:offline-request-boundary-v1"
+
+
+@dataclass(frozen=True, slots=True)
+class GeminiConformanceBundle:
+    profile: CapabilityProfile
+    record: CompatibilityRecord
+    detector: CapabilityDetector
+
+
+def build_gemini_conformance() -> GeminiConformanceBundle:
+    profile = CapabilityProfile(
+        key=GEMINI_GENERATE_CONTENT_KEY,
+        server_cache="cached-content-reference-preserved",
+        tool_calls=True,
+        reasoning_state="exact-replay-required",
+        exact_byte_preservation=True,
+        evidence_id=EVIDENCE_ID,
+    )
+    record = CompatibilityRecord(
+        key=GEMINI_GENERATE_CONTENT_KEY,
+        state=CompatibilityState.PASSTHROUGH_ONLY,
+        evidence_id=EVIDENCE_ID,
+        reason="offline_conformance_only",
+    )
+    capabilities = CapabilityRegistry()
+    capabilities.register(profile)
+    compatibility = CompatibilityRegistry()
+    compatibility.register(record)
+    return GeminiConformanceBundle(
+        profile=profile,
+        record=record,
+        detector=CapabilityDetector(capabilities, compatibility),
+    )
+
+
+def _supported_fixture() -> dict[str, object]:
+    return {
+        "contents": [
+            {"role": "user", "parts": [{"text": "current task"}]},
+            {
+                "role": "model",
+                "parts": [
+                    {
+                        "functionCall": {
+                            "id": "call_1",
+                            "name": "lookup",
+                            "args": {"q": "status"},
+                        }
+                    }
+                ],
+            },
+        ],
+        "systemInstruction": {"parts": [{"text": "stable policy"}]},
+        "tools": [
+            {
+                "functionDeclarations": [
+                    {
+                        "name": "lookup",
+                        "description": "fixture",
+                        "parameters": {"type": "OBJECT"},
+                    }
+                ]
+            }
+        ],
+        "toolConfig": {"functionCallingConfig": {"mode": "AUTO"}},
+        "cachedContent": "cachedContents/cache-1",
+        "generationConfig": {"temperature": 0.2, "thinkingLevel": "LOW"},
+        "futureField": {"preserve": True},
+    }
+
+
+def _supported_roundtrip() -> bool:
+    payload = _supported_fixture()
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    return envelope.wire_safe and adapter.serialize(envelope) == payload
+
+
+def _roleless_text_roundtrip() -> bool:
+    payload = {"contents": [{"parts": [{"text": "hello"}]}]}
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    user = next((block for block in envelope.blocks if block.text == "hello"), None)
+    return (
+        envelope.wire_safe
+        and user is not None
+        and user.kind == "user"
+        and user.role == "user"
+        and user.metadata.get("path") == ("contents", 0, "parts", 0, "text")
+        and adapter.serialize(envelope) == payload
+    )
+
+
+def _cached_content_reference_preserved() -> bool:
+    payload = _supported_fixture()
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    serialized = adapter.serialize(envelope)
+    return (
+        envelope.wire_safe
+        and serialized == payload
+        and serialized["cachedContent"] == "cachedContents/cache-1"
+    )
+
+
+def _text_path_mutation_preserves_protocol() -> bool:
+    payload = _supported_fixture()
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    target = next(block for block in envelope.blocks if block.text == "current task")
+    changed = replace(
+        envelope,
+        blocks=tuple(
+            replace(block, text="updated task") if block is target else block
+            for block in envelope.blocks
+        ),
+    )
+    actual = adapter.serialize(changed)
+    expected = deepcopy(payload)
+    expected["contents"][0]["parts"][0]["text"] = "updated task"
+    return actual == expected
+
+
+def _system_instruction_text_mutation() -> bool:
+    payload = _supported_fixture()
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    target = next(block for block in envelope.blocks if block.kind == "system")
+    changed = replace(
+        envelope,
+        blocks=tuple(
+            replace(block, text="updated policy") if block is target else block
+            for block in envelope.blocks
+        ),
+    )
+    actual = adapter.serialize(changed)
+    expected = deepcopy(payload)
+    expected["systemInstruction"]["parts"][0]["text"] = "updated policy"
+    return actual == expected
+
+
+def _tool_classification() -> bool:
+    call_payload = _supported_fixture()
+    call_envelope = GeminiGenerateContentAdapter().parse(call_payload)
+    call_by_kind = {block.kind: block for block in call_envelope.blocks}
+    call_required = {"tool_schema", "tool_call"}
+
+    response_payload = {
+        "contents": [
+            {
+                "role": "user",
+                "parts": [
+                    {
+                        "functionResponse": {
+                            "id": "call_1",
+                            "name": "lookup",
+                            "response": {"status": "green"},
+                        }
+                    }
+                ],
+            }
+        ]
+    }
+    response_envelope = GeminiGenerateContentAdapter().parse(response_payload)
+    output = next(
+        block for block in response_envelope.blocks if block.kind == "tool_output"
+    )
+    return (
+        call_envelope.wire_safe
+        and call_required <= set(call_by_kind)
+        and all("path" not in call_by_kind[kind].metadata for kind in call_required)
+        and not response_envelope.wire_safe
+        and "path" not in output.metadata
+        and GeminiGenerateContentAdapter().serialize(response_envelope) == response_payload
+    )
+
+
+def _function_response_exact_passthrough() -> bool:
+    payload = {
+        "contents": [
+            {
+                "role": "user",
+                "parts": [
+                    {
+                        "functionResponse": {
+                            "name": "lookup",
+                            "response": {"status": "green"},
+                        }
+                    }
+                ],
+            }
+        ]
+    }
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    return (
+        not envelope.wire_safe
+        and any(block.kind == "tool_output" for block in envelope.blocks)
+        and adapter.serialize(envelope) == payload
+    )
+
+
+def _thought_signature_exact_passthrough() -> bool:
+    payload = {
+        "contents": [
+            {
+                "role": "model",
+                "parts": [
+                    {
+                        "functionCall": {"name": "lookup", "args": {"q": "x"}},
+                        "thoughtSignature": "opaque-signature",
+                    }
+                ],
+            }
+        ]
+    }
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    return (
+        not envelope.wire_safe
+        and adapter.serialize(envelope) == payload
+        and any(block.kind == "protocol_state" for block in envelope.blocks)
+    )
+
+
+def _unknown_multimodal_passthrough() -> bool:
+    payload = {
+        "contents": [
+            {
+                "role": "user",
+                "parts": [
+                    {"text": "describe"},
+                    {"inlineData": {"mimeType": "image/png", "data": "opaque"}},
+                ],
+            }
+        ]
+    }
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    return not envelope.wire_safe and adapter.serialize(envelope) == payload
+
+
+def _unsupported_tool_passthrough() -> bool:
+    payload = {
+        "contents": [{"role": "user", "parts": [{"text": "search"}]}],
+        "tools": [{"googleSearch": {}}],
+    }
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    return not envelope.wire_safe and adapter.serialize(envelope) == payload
+
+
+def _malformed_function_passthrough() -> bool:
+    payload = {
+        "contents": [
+            {"role": "model", "parts": [{"functionCall": {"args": {"q": "x"}}}]}
+        ]
+    }
+    adapter = GeminiGenerateContentAdapter()
+    envelope = adapter.parse(payload)
+    return not envelope.wire_safe and adapter.serialize(envelope) == payload
+
+
+def _streaming_unclaimed() -> bool:
+    bundle = build_gemini_conformance()
+    return (
+        bundle.profile.streaming == "unknown"
+        and GeminiGenerateContentAdapter.protocol_id == "gemini_generate_content"
+    )
+
+
+def _certification_state_passthrough_only() -> bool:
+    bundle = build_gemini_conformance()
+    return (
+        bundle.record.state is CompatibilityState.PASSTHROUGH_ONLY
+        and bundle.record.reason == "offline_conformance_only"
+        and bundle.profile.evidence_id == bundle.record.evidence_id
+    )
+
+
+def _near_match_non_inheritance() -> bool:
+    bundle = build_gemini_conformance()
+    keys = (
+        CapabilityKey(
+            "generic-google",
+            "gemini_generate_content",
+            "google",
+            model_family="gemini-future",
+        ),
+        CapabilityKey("generic-google", "gemini_generate_content", "google-v2"),
+        CapabilityKey("gemini-cli", "gemini_generate_content", "google"),
+    )
+    return all(
+        bundle.detector.detect(key).profile is None
+        and bundle.detector.detect(key).compatibility.state
+        is CompatibilityState.PASSTHROUGH_ONLY
+        and bundle.detector.detect(key).compatibility.reason == "unknown_capability"
+        for key in keys
+    )
+
+
+def gemini_conformance_cases() -> Mapping[str, Callable[[], bool]]:
+    return {
+        "cached_content_reference_preserved": _cached_content_reference_preserved,
+        "certification_state_passthrough_only": _certification_state_passthrough_only,
+        "function_response_exact_passthrough": _function_response_exact_passthrough,
+        "malformed_function_passthrough": _malformed_function_passthrough,
+        "near_match_non_inheritance": _near_match_non_inheritance,
+        "roleless_text_roundtrip": _roleless_text_roundtrip,
+        "streaming_unclaimed": _streaming_unclaimed,
+        "supported_roundtrip": _supported_roundtrip,
+        "system_instruction_text_mutation": _system_instruction_text_mutation,
+        "text_path_mutation_preserves_protocol": _text_path_mutation_preserves_protocol,
+        "thought_signature_exact_passthrough": _thought_signature_exact_passthrough,
+        "tool_classification": _tool_classification,
+        "unknown_multimodal_passthrough": _unknown_multimodal_passthrough,
+        "unsupported_tool_passthrough": _unsupported_tool_passthrough,
+    }
+
+
+def run_gemini_conformance() -> tuple[ConformanceResult, ...]:
+    results = run_conformance(gemini_conformance_cases())
+    require_conformance(results)
+    return results

token_runtime/integrations.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import hashlib
+import json
+import os
+from pathlib import Path
+import re
+from typing import Literal
+
+
+_BEGIN = b"# >>> TOKEN MANAGED >>>"
+_END = b"# <<< TOKEN MANAGED <<<"
+_BASE_RE = re.compile(br"(?m)^OPENAI_BASE_URL\s*=.*$")
+_MANAGED_RE = re.compile(
+    re.escape(_BEGIN) + br"\r?\n.*?\r?\n" + re.escape(_END) + br"(?:\r?\n)?",
+    re.DOTALL,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class Integration:
+    kind: Literal["codex", "openai_env"]
+    path: Path
+    root: Path
+
+
+@dataclass(frozen=True, slots=True)
+class InstallPlan:
+    integration: Integration
+    mode: Literal["sidecar", "managed_block", "managed_toml"]
+    target_path: Path
+    gateway_url: str
+    original_bytes: bytes | None
+    desired_bytes: bytes
+    safe_to_apply: bool = True
+    reason: str = "ready"
+    metadata_path: Path | None = None
+    metadata_bytes: bytes | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ApplyResult:
+    changed: bool
+    reason: str
+    target_path: Path
+
+
+def detect_integrations(root: str | Path) -> list[Integration]:
+    base = Path(root)
+    found: list[Integration] = []
+    codex = base / ".codex" / "config.toml"
+    if codex.is_file():
+        found.append(Integration("codex", codex, base))
+    env = base / ".env"
+    if env.is_file():
+        content = env.read_bytes()
+        if b"OPENAI_" in content or _BEGIN in content:
+            found.append(Integration("openai_env", env, base))
+    return found
+
+
+def _managed_block(gateway_url: str) -> bytes:
+    return (
+        _BEGIN + b"\n" +
+        f"OPENAI_BASE_URL={gateway_url}\n".encode("utf-8") +
+        _END + b"\n"
+    )
+
+
+def _replace_or_append_block(original: bytes, gateway_url: str) -> bytes:
+    block = _managed_block(gateway_url)
+    if _MANAGED_RE.search(original):
+        return _MANAGED_RE.sub(block, original, count=1)
+    separator = b"" if not original or original.endswith((b"\n", b"\r")) else b"\n"
+    return original + separator + block
+
+
+def _codex_provider_sections(content: bytes):
+    return list(re.finditer(br"(?m)^\[model_providers\.([A-Za-z0-9_-]+)\][ \t]*(?:\r?\n|$)", content))
+
+
+def _replace_codex_base_url(content: bytes, provider_id: str, new_url: str):
+    sections = _codex_provider_sections(content)
+    match = next((m for m in sections if m.group(1).decode() == provider_id), None)
+    if match is None:
+        return None
+    next_headers = [m.start() for m in re.finditer(br"(?m)^\[[^\r\n]+\][ \t]*(?:\r?\n|$)", content) if m.start() > match.start()]
+    end = min(next_headers) if next_headers else len(content)
+    section = content[match.end():end]
+    base = re.search(br'(?m)^[ \t]*base_url[ \t]*=[ \t]*"([^"\r\n]+)"', section)
+    if base is None:
+        return None
+    start = match.end() + base.start(1)
+    stop = match.end() + base.end(1)
+    old = content[start:stop].decode("utf-8")
+    desired = content[:start] + new_url.encode("utf-8") + content[stop:]
+    return old, desired
+
+
+def _plan_codex(integration: Integration, gateway_url: str) -> InstallPlan:
+    original = integration.path.read_bytes()
+    target = integration.path
+    metadata_path = integration.root / ".token" / "integrations" / "codex.json"
+    if metadata_path.exists():
+        try:
+            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+        except (OSError, ValueError, TypeError):
+            return InstallPlan(integration, "managed_toml", target, gateway_url, original, original, False, "invalid_codex_metadata", metadata_path)
+        current_hash = hashlib.sha256(original).hexdigest()
+        if metadata.get("strategy") == "managed_provider_base_url" and current_hash == metadata.get("desired_sha256"):
+            return InstallPlan(integration, "managed_toml", target, gateway_url, original, original, True, "already_managed", metadata_path, metadata_path.read_bytes())
+        return InstallPlan(integration, "managed_toml", target, gateway_url, original, original, False, "target_drift", metadata_path)
+    sections = _codex_provider_sections(original)
+    if len(sections) != 1:
+        reason = "ambiguous_codex_provider" if len(sections) > 1 else "missing_codex_provider"
+        return InstallPlan(integration, "managed_toml", target, gateway_url, original, original, False, reason, metadata_path)
+    provider_id = sections[0].group(1).decode("utf-8")
+    replaced = _replace_codex_base_url(original, provider_id, gateway_url)
+    if replaced is None:
+        return InstallPlan(integration, "managed_toml", target, gateway_url, original, original, False, "missing_codex_base_url", metadata_path)
+    old_url, desired = replaced
+    metadata = {
+        "kind": "codex", "strategy": "managed_provider_base_url",
+        "provider_id": provider_id, "gateway_url": gateway_url,
+        "original_base_url": old_url,
+        "original_sha256": hashlib.sha256(original).hexdigest(),
+        "desired_sha256": hashlib.sha256(desired).hexdigest(),
+    }
+    metadata_bytes = (json.dumps(metadata, indent=2, sort_keys=True) + "\n").encode("utf-8")
+    return InstallPlan(integration, "managed_toml", target, gateway_url, original, desired, True, "ready", metadata_path, metadata_bytes)
+
+
+def plan_install(integration: Integration, *, gateway_url: str) -> InstallPlan:
+    if integration.kind == "codex":
+        return _plan_codex(integration, gateway_url)
+
+    original = integration.path.read_bytes()
+    unmanaged = _BASE_RE.search(_MANAGED_RE.sub(b"", original))
+    safe = unmanaged is None
+    desired = _replace_or_append_block(original, gateway_url) if safe else original
+    metadata_path = integration.root / ".token" / "integrations" / "openai_env.json"
+    if metadata_path.exists():
+        metadata_bytes = metadata_path.read_bytes()
+    else:
+        separator_added = bool(original and not original.endswith((b"\n", b"\r")))
+        metadata_bytes = (json.dumps({
+            "kind": "openai_env",
+            "original_sha256": hashlib.sha256(original).hexdigest(),
+            "separator_added": separator_added,
+        }, indent=2, sort_keys=True) + "\n").encode("utf-8")
+    return InstallPlan(
+        integration, "managed_block", integration.path, gateway_url,
+        original, desired, safe_to_apply=safe,
+        reason="ready" if safe else "existing_openai_base_url_unmanaged",
+        metadata_path=metadata_path, metadata_bytes=metadata_bytes,
+    )
+
+
+def _atomic_write(path: Path, content: bytes) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_name(path.name + ".token-tmp")
+    tmp.write_bytes(content)
+    os.chmod(tmp, 0o600)
+    os.replace(tmp, path)
+
+
+def apply_plan(plan: InstallPlan, *, apply: bool = False) -> ApplyResult:
+    current = plan.target_path.read_bytes() if plan.target_path.exists() else None
+    changed = current != plan.desired_bytes
+    if not plan.safe_to_apply:
+        return ApplyResult(False, plan.reason, plan.target_path)
+    if not apply or not changed:
+        return ApplyResult(changed, "dry_run" if not apply else "unchanged", plan.target_path)
+    if current not in (plan.original_bytes, plan.desired_bytes):
+        return ApplyResult(False, "target_drift", plan.target_path)
+    try:
+        _atomic_write(plan.target_path, plan.desired_bytes)
+        if plan.metadata_path is not None and plan.metadata_bytes is not None:
+            _atomic_write(plan.metadata_path, plan.metadata_bytes)
+    except Exception:
+        if plan.original_bytes is None:
+            plan.target_path.unlink(missing_ok=True)
+        else:
+            _atomic_write(plan.target_path, plan.original_bytes)
+        if plan.metadata_path is not None:
+            plan.metadata_path.unlink(missing_ok=True)
+        raise
+    return ApplyResult(True, "applied", plan.target_path)
+
+
+def restore_plan(plan: InstallPlan, *, apply: bool = False) -> ApplyResult:
+    current = plan.target_path.read_bytes() if plan.target_path.exists() else None
+    desired = plan.original_bytes
+    metadata_exists = bool(plan.metadata_path and plan.metadata_path.exists())
+    changed = current != desired or metadata_exists
+    if not apply or not changed:
+        return ApplyResult(changed, "dry_run" if not apply else "unchanged", plan.target_path)
+    if current not in (plan.original_bytes, plan.desired_bytes):
+        return ApplyResult(False, "target_drift", plan.target_path)
+    if current != desired:
+        if desired is None:
+            plan.target_path.unlink(missing_ok=True)
+        else:
+            _atomic_write(plan.target_path, desired)
+    if plan.metadata_path is not None:
+        plan.metadata_path.unlink(missing_ok=True)
+    return ApplyResult(True, "restored", plan.target_path)
+
+
+def uninstall_integration(integration: Integration, *, apply: bool = False) -> ApplyResult:
+    if integration.kind == "codex":
+        metadata_path = integration.root / ".token" / "integrations" / "codex.json"
+        if not metadata_path.exists():
+            return ApplyResult(False, "unchanged" if apply else "dry_run", integration.path)
+        if not apply:
+            return ApplyResult(True, "dry_run", integration.path)
+        try:
+            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+        except (OSError, ValueError, TypeError):
+            return ApplyResult(False, "invalid_codex_metadata", integration.path)
+        if metadata.get("strategy") == "sidecar":
+            metadata_path.unlink()
+            return ApplyResult(True, "uninstalled", integration.path)
+        if metadata.get("strategy") != "managed_provider_base_url":
+            return ApplyResult(False, "invalid_codex_metadata", integration.path)
+        current = integration.path.read_bytes()
+        current_hash = hashlib.sha256(current).hexdigest()
+        if current_hash == metadata.get("original_sha256"):
+            metadata_path.unlink()
+            return ApplyResult(True, "metadata_removed", integration.path)
+        if current_hash != metadata.get("desired_sha256"):
+            return ApplyResult(False, "target_drift", integration.path)
+        restored = _replace_codex_base_url(current, str(metadata.get("provider_id")), str(metadata.get("original_base_url")))
+        if restored is None:
+            return ApplyResult(False, "restore_failed", integration.path)
+        _, desired = restored
+        if hashlib.sha256(desired).hexdigest() != metadata.get("original_sha256"):
+            return ApplyResult(False, "restore_hash_mismatch", integration.path)
+        _atomic_write(integration.path, desired)
+        metadata_path.unlink()
+        return ApplyResult(True, "uninstalled", integration.path)
+
+    target = integration.path
+    current = target.read_bytes()
+    match = _MANAGED_RE.search(current)
+    metadata_path = integration.root / ".token" / "integrations" / "openai_env.json"
+    metadata_exists = metadata_path.exists()
+    if match is None:
+        if apply and metadata_exists:
+            metadata_path.unlink()
+        changed = metadata_exists
+        return ApplyResult(changed, "metadata_removed" if apply and changed else ("dry_run" if not apply else "unchanged"), target)
+
+    separator_added = False
+    if metadata_exists:
+        try:
+            metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
+            separator_added = bool(metadata.get("separator_added", False))
+        except (OSError, ValueError, TypeError):
+            separator_added = False
+    start = match.start()
+    end = match.end()
+    desired = current[:start] + current[end:]
+    if separator_added and start > 0 and current[start - 1:start] == b"\n":
+        desired = current[:start - 1] + current[end:]
+
+    if not apply:
+        return ApplyResult(True, "dry_run", target)
+    _atomic_write(target, desired)
+    metadata_path.unlink(missing_ok=True)
+    return ApplyResult(True, "uninstalled", target)