@oneciel-ai/ciel-runtime 0.1.1 → 0.2.0

package/ciel_runtime_support/router_process_lifecycle.py

@@ -0,0 +1,401 @@
+"""Router process shutdown and port-replacement application policy."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+import os
+from pathlib import Path
+import subprocess
+import sys
+import threading
+from typing import Any
+
+
+def schedule_router_restart(
+    delay: float,
+    script_path: Path,
+    log: Callable[[str, str], None],
+    *,
+    timer_factory: Callable[..., Any] = threading.Timer,
+    exec_process: Callable[[str, list[str]], Any] = os.execv,
+    executable: str = sys.executable,
+    stderr: Any = sys.stderr,
+) -> None:
+    def restart() -> None:
+        try:
+            log("INFO", "router_debug_restart execing router process")
+            exec_process(
+                executable, [executable, str(script_path), "serve"]
+            )
+        except Exception as exc:
+            try:
+                log(
+                    "ERROR",
+                    f"router_debug_restart_failed {type(exc).__name__}: {exc}",
+                )
+            except Exception as log_exc:
+                try:
+                    stderr.write(
+                        "router_debug_restart_diagnostic_failed "
+                        f"restart_error={type(exc).__name__}: {exc} "
+                        f"log_error={type(log_exc).__name__}: {log_exc}\n"
+                    )
+                    stderr.flush()
+                except (OSError, ValueError):
+                    return
+
+    timer = timer_factory(delay, restart)
+    timer.daemon = True
+    timer.start()
+
+
+@dataclass(frozen=True, slots=True)
+class RouterProcessConfig:
+    pid_path: Path
+    router_port: int
+    router_base: str
+    config_dir: Path
+
+
+@dataclass(frozen=True, slots=True)
+class RouterStatePorts:
+    health: Callable[[], dict[str, Any] | None]
+    foreign_config: Callable[[dict[str, Any] | None], bool]
+    current_config: Callable[[dict[str, Any] | None], bool]
+    log: Callable[[str, str], None]
+
+
+@dataclass(frozen=True, slots=True)
+class RouterTerminationPorts:
+    terminate_pid: Callable[[int, str, bool], bool]
+    terminate_pid_file: Callable[[Path, str, bool], bool]
+    terminate_health: Callable[[dict[str, Any] | None, bool], bool]
+    stop_processes: Callable[[bool], bool]
+    listener_pids: Callable[[], list[int]]
+
+
+@dataclass(frozen=True, slots=True)
+class ClockPorts:
+    now: Callable[[], float]
+    sleep: Callable[[float], None]
+
+
+@dataclass(frozen=True, slots=True)
+class RouterStartupIdentity:
+    version: str
+    source_fingerprint: str
+
+
+@dataclass(frozen=True, slots=True)
+class RouterStartupStatePorts:
+    health: Callable[[], dict[str, Any] | None]
+    active_client_pids: Callable[[], list[int]]
+    health_matches_current: Callable[[dict[str, Any] | None], bool]
+    health_config_matches_current: Callable[[dict[str, Any] | None], bool]
+    terminate_active_clients: Callable[..., bool]
+    ensure_port_available: Callable[..., None]
+    reuse_enabled: Callable[[], bool]
+    log: Callable[[str, str], None]
+
+
+@dataclass(frozen=True, slots=True)
+class RouterSpawnPorts:
+    popen: Callable[..., Any]
+    router_up: Callable[[], bool]
+    now: Callable[[], float]
+    sleep: Callable[[float], None]
+    process_id: Callable[[], int]
+    environment: Callable[[], dict[str, str]]
+
+
+def start_router_if_needed(
+    *,
+    replace_active_clients: bool,
+    config: RouterProcessConfig,
+    identity: RouterStartupIdentity,
+    state: RouterStartupStatePorts,
+    spawn: RouterSpawnPorts,
+    executable: str,
+    entrypoint: Path,
+    log_path: Path,
+    platform_name: str,
+) -> bool:
+    health = state.health()
+    if health is not None:
+        active_clients = state.active_client_pids()
+        if state.health_matches_current(health):
+            if active_clients:
+                if replace_active_clients:
+                    state.log(
+                        "WARN",
+                        "router_prelaunch_replace_active_clients "
+                        f"base={config.router_base} active_clients={','.join(map(str, active_clients))}",
+                    )
+                    state.terminate_active_clients("prelaunch_active_clients", active_clients, quiet=True)
+                    state.ensure_port_available("prelaunch_active_clients", health)
+                else:
+                    state.log(
+                        "INFO",
+                        "router_check_state running=True spawn=False "
+                        f"base={config.router_base} active_clients={','.join(map(str, active_clients))}",
+                    )
+                    return True
+            elif state.reuse_enabled():
+                state.log("INFO", f"router_check_state running=True spawn=False base={config.router_base} reuse=env")
+                return True
+            else:
+                state.log(
+                    "INFO",
+                    "router_prelaunch_replace "
+                    f"running_version={health.get('version') or '-'} current_version={identity.version} "
+                    f"running_source={health.get('source_fingerprint') or '-'} current_source={identity.source_fingerprint} "
+                    f"pid={health.get('pid') or '-'}",
+                )
+                state.ensure_port_available("prelaunch_replace", health)
+        elif state.health_config_matches_current(health) and active_clients:
+            if replace_active_clients:
+                state.log(
+                    "WARN",
+                    "router_version_mismatch_replace_active_clients "
+                    f"running_version={health.get('version') or '-'} current_version={identity.version} "
+                    f"active_clients={','.join(map(str, active_clients))}",
+                )
+                state.terminate_active_clients("version_mismatch_active_clients", active_clients, quiet=True)
+                state.ensure_port_available("version_mismatch_active_clients", health)
+            else:
+                raise RuntimeError(
+                    f"ciel-runtime router on {config.router_base} belongs to this config but has active clients "
+                    f"({','.join(map(str, active_clients))}) and differs from this launch "
+                    f"(running_version={health.get('version') or '-'}, current_version={identity.version}). "
+                    "Stop the other Claude Code session or launch this instance with a different "
+                    "CIEL_RUNTIME_ROUTER_PORT."
+                )
+        else:
+            state.log(
+                "WARN",
+                "router_version_mismatch_restart "
+                f"running_version={health.get('version') or '-'} current_version={identity.version} "
+                f"running_source={health.get('source_fingerprint') or '-'} current_source={identity.source_fingerprint}",
+            )
+            state.ensure_port_available("version_mismatch", health)
+    else:
+        state.ensure_port_available("pre_spawn", None)
+    config.config_dir.mkdir(parents=True, exist_ok=True)
+    command = [executable, str(entrypoint), "serve"]
+    kwargs: dict[str, Any] = {}
+    if platform_name == "nt":
+        flags = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
+        if flags:
+            kwargs["creationflags"] = flags
+    else:
+        kwargs["start_new_session"] = True
+    state.log("INFO", f"router_check_state running=False spawn=True base={config.router_base}")
+    environment = spawn.environment()
+    environment["CIEL_RUNTIME_MANAGED_ROUTER"] = "1"
+    environment["CIEL_RUNTIME_ROUTER_OWNER_PID"] = str(spawn.process_id())
+    with log_path.open("ab", buffering=0) as log:
+        spawn.popen(
+            command,
+            stdin=subprocess.DEVNULL,
+            stdout=log,
+            stderr=log,
+            env=environment,
+            **kwargs,
+        )
+    started_at = spawn.now()
+    deadline = started_at + 30
+    while spawn.now() < deadline:
+        if spawn.router_up():
+            state.log("INFO", f"router_spawned running=True base={config.router_base} elapsed={spawn.now()-started_at:.1f}s")
+            return True
+        spawn.sleep(0.5)
+    raise RuntimeError(f"ciel-runtime router did not start. See {log_path}")
+
+
+def terminate_pid_file(
+    path: Path,
+    label: str,
+    quiet: bool,
+    *,
+    terminate_pid: Callable[[int, str, bool], bool],
+    pid_is_running: Callable[[int], bool],
+) -> bool:
+    if not path.exists():
+        return False
+    try:
+        pid = int(path.read_text().strip())
+    except Exception:
+        try:
+            path.unlink()
+        except FileNotFoundError:
+            pass
+        return False
+    stopped = terminate_pid(pid, label, quiet)
+    if stopped or not pid_is_running(pid):
+        try:
+            path.unlink()
+        except FileNotFoundError:
+            pass
+    return stopped
+
+
+def terminate_health_pid(
+    health: dict[str, Any] | None,
+    quiet: bool,
+    *,
+    config: RouterProcessConfig,
+    state: RouterStatePorts,
+    terminate_pid: Callable[[int, str, bool], bool],
+    protected_pids: tuple[int, int],
+) -> bool:
+    if not isinstance(health, dict):
+        return False
+    if not state.current_config(health):
+        state.log(
+            "INFO",
+            "router_kill_skipped_foreign_config "
+            f"running_config={health.get('config_dir') or '-'} "
+            f"current_config={config.config_dir}",
+        )
+        return False
+    try:
+        pid = int(health.get("pid") or 0)
+    except Exception:
+        return False
+    if pid in protected_pids:
+        return False
+    return terminate_pid(pid, "ciel-runtime router", quiet)
+
+
+def stop_router_processes(
+    quiet: bool,
+    *,
+    config: RouterProcessConfig,
+    state: RouterStatePorts,
+    termination: RouterTerminationPorts,
+) -> bool:
+    stopped = termination.terminate_pid_file(config.pid_path, "ciel-runtime router", quiet)
+    health = state.health()
+    if state.foreign_config(health):
+        state.log(
+            "INFO",
+            "router_stop_skipped_foreign_config "
+            f"running_config={health.get('config_dir') or '-'} "
+            f"current_config={config.config_dir}",
+        )
+        return stopped
+    if state.current_config(health):
+        stopped = termination.terminate_health(health, True) or stopped
+    return stopped
+
+
+def ensure_port_available(
+    reason: str,
+    health: dict[str, Any] | None,
+    max_wait_seconds: float,
+    *,
+    config: RouterProcessConfig,
+    state: RouterStatePorts,
+    termination: RouterTerminationPorts,
+    clock: ClockPorts,
+) -> None:
+    if state.foreign_config(health):
+        raise RuntimeError(_foreign_port_message(config, health))
+    stopped = termination.terminate_health(health, True)
+    stopped = termination.stop_processes(True) or stopped
+    deadline = clock.now() + max(0.1, max_wait_seconds)
+    while clock.now() < deadline:
+        current_health = state.health()
+        if state.foreign_config(current_health):
+            raise RuntimeError(_foreign_port_message(config, current_health))
+        if current_health is None and not termination.listener_pids():
+            state.log(
+                "INFO",
+                f"router_port_clear reason={reason} port={config.router_port} "
+                f"stopped={str(stopped).lower()}",
+            )
+            return
+        if current_health is not None:
+            termination.terminate_health(current_health, True)
+        termination.stop_processes(True)
+        clock.sleep(0.1)
+    pids = termination.listener_pids()
+    current_health = state.health()
+    description = ""
+    if isinstance(current_health, dict):
+        description = (
+            f" version={current_health.get('version') or '-'}"
+            f" source={current_health.get('source_fingerprint') or '-'}"
+            f" pid={current_health.get('pid') or '-'}"
+        )
+    raise RuntimeError(
+        f"stale ciel-runtime router is still serving on {config.router_base}; "
+        f"port {config.router_port} listener_pids={pids or '-'}{description}; "
+        "run `ciel-runtime stop` and launch again."
+    )
+
+
+def stop_with_guarantee(
+    reason: str,
+    max_wait_seconds: float,
+    quiet: bool,
+    *,
+    config: RouterProcessConfig,
+    state: RouterStatePorts,
+    termination: RouterTerminationPorts,
+    clock: ClockPorts,
+) -> bool:
+    initial_health = state.health()
+    if initial_health is None:
+        state.log("INFO", f"router_kill_guarantee reason={reason} state=already_down")
+        return False
+    if state.foreign_config(initial_health):
+        state.log(
+            "INFO",
+            "router_kill_guarantee_skipped_foreign_config "
+            f"reason={reason} running_config={initial_health.get('config_dir') or '-'} "
+            f"current_config={config.config_dir}",
+        )
+        return False
+    termination.stop_processes(quiet)
+    deadline = clock.now() + max(0.1, max_wait_seconds)
+    while clock.now() < deadline:
+        health = state.health()
+        if state.foreign_config(health):
+            state.log(
+                "INFO",
+                "router_kill_guarantee_skipped_foreign_config "
+                f"reason={reason} running_config={health.get('config_dir') or '-'} "
+                f"current_config={config.config_dir}",
+            )
+            return False
+        if health is None:
+            elapsed_ms = int((max_wait_seconds - (deadline - clock.now())) * 1000)
+            state.log(
+                "INFO",
+                f"router_kill_guarantee reason={reason} state=killed elapsed_ms={elapsed_ms}",
+            )
+            return True
+        clock.sleep(0.1)
+    state.log(
+        "ERROR",
+        f"router_kill_guarantee reason={reason} state=still_up_after_{max_wait_seconds:.1f}s",
+    )
+    raise RuntimeError(
+        f"ciel-runtime router is still serving on {config.router_base} after a "
+        f"{max_wait_seconds:.1f}s shutdown attempt for '{reason}'. Aborting to prevent "
+        "the subsequent claude launch from accidentally routing through it. "
+        f"Investigate the PID at {config.pid_path} or use `ciel-runtime stop` manually."
+    )
+
+
+def _foreign_port_message(
+    config: RouterProcessConfig, health: dict[str, Any] | None
+) -> str:
+    foreign_config = health.get("config_dir") if isinstance(health, dict) else "-"
+    return (
+        f"ciel-runtime router port {config.router_port} is already used by another "
+        f"ciel-runtime config ({foreign_config}). Set CIEL_RUNTIME_ROUTER_PORT or "
+        "CIEL_RUNTIME_CONFIG_DIR for this instance instead of killing the other router."
+    )

package/ciel_runtime_support/router_rate_limit_service.py

@@ -0,0 +1,285 @@
+"""Router-wide provider rate-limit application service."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from threading import Lock, RLock
+from typing import Any
+
+from . import rate_limit_policy
+from .provider_limits import (
+    RateLimitApplyPolicy,
+    RateLimitApplyServices,
+    RateLimitBackoffPolicy,
+    RateLimitBackoffServices,
+    RateLimitLearningPolicy,
+    RateLimitLearningServices,
+    RateLimitStateStore,
+    apply_rate_limit,
+    learn_rate_limit_headers,
+    register_rate_limit_backoff,
+)
+from .rate_limit_repository import RateLimitRepository
+
+
+@dataclass(frozen=True, slots=True)
+class RouterRateLimitPaths:
+    config_dir: Path
+    state_path: Path
+    lock: Lock | RLock
+
+
+@dataclass(frozen=True, slots=True)
+class RouterRateLimitPorts:
+    current_model_id: Callable[[str, dict[str, Any]], str]
+    api_key_count: Callable[[str, dict[str, Any]], int]
+    positive_int: Callable[[Any], int | None]
+    log: Callable[[str, str], None]
+    now: Callable[[], float] = time.time
+    sleep: Callable[[float], None] = time.sleep
+
+
+@dataclass(frozen=True, slots=True)
+class RouterRateLimitService:
+    paths: RouterRateLimitPaths
+    repository: RateLimitRepository
+    ports: RouterRateLimitPorts
+
+    def legacy_key(self, provider: str, config: dict[str, Any], model: str | None) -> str:
+        return f"{provider}:{model or self.ports.current_model_id(provider, config)}"
+
+    def configured_rpm(self, provider: str, config: dict[str, Any]) -> int | None:
+        del provider
+        return rate_limit_policy.configured_rpm(config, self.ports.positive_int)
+
+    def rpm(self, provider: str, config: dict[str, Any]) -> int | None:
+        rpm = self.configured_rpm(provider, config)
+        return rpm if rpm and rpm > 0 else None
+
+    @staticmethod
+    def key(provider: str, config: dict[str, Any], model: str | None = None) -> str:
+        del config, model
+        return f"{provider}:__global__"
+
+    def state_entry(
+        self, provider: str, config: dict[str, Any], model: str | None = None
+    ) -> dict[str, Any]:
+        return self.repository.entry(
+            self.key(provider, config, model), self.legacy_key(provider, config, model)
+        )
+
+    def effective_rpm(
+        self, provider: str, config: dict[str, Any], model: str | None = None
+    ) -> int | None:
+        return self.repository.effective_rpm(
+            self.key(provider, config, model),
+            self.legacy_key(provider, config, model),
+            self.configured_rpm(provider, config),
+        )
+
+    @staticmethod
+    def capacity(rpm: int) -> int:
+        return rate_limit_policy.capacity(rpm)
+
+    @staticmethod
+    def recent(
+        timestamps: Any, now: float, window: float, *, include_future: bool
+    ) -> list[float]:
+        return rate_limit_policy.recent_timestamps(
+            timestamps, now, window, include_future=include_future
+        )
+
+    def usage(
+        self, provider: str, config: dict[str, Any], model: str | None = None
+    ) -> tuple[int, int | None]:
+        return self.repository.usage(
+            self.key(provider, config, model),
+            self.legacy_key(provider, config, model),
+            self.effective_rpm(provider, config, model),
+            self.recent,
+        )
+
+    def record_usage(
+        self,
+        provider: str,
+        config: dict[str, Any],
+        model: str | None,
+        rpm: int | None,
+    ) -> tuple[int, int | None]:
+        return self.repository.record_usage(
+            self.key(provider, config, model),
+            self.legacy_key(provider, config, model),
+            rpm,
+            self.recent,
+        )
+
+    def _state_store(self) -> RateLimitStateStore:
+        return RateLimitStateStore(
+            CONFIG_DIR=self.paths.config_dir,
+            RATE_LIMIT_STATE_PATH=self.paths.state_path,
+            _RATE_LIMIT_LOCK=self.paths.lock,
+            router_log=self.ports.log,
+        )
+
+    def learn_headers(
+        self, provider: str, config: dict[str, Any], model: str | None, headers: Any
+    ) -> None:
+        learn_rate_limit_headers(
+            provider,
+            config,
+            model,
+            headers,
+            services=RateLimitLearningServices(
+                state_store=self._state_store(),
+                policy=RateLimitLearningPolicy(
+                    current_upstream_model_id=self.ports.current_model_id,
+                    first_header=rate_limit_policy.first_header,
+                    first_int_in_header=rate_limit_policy.first_integer,
+                    provider_api_key_count=self.ports.api_key_count,
+                    rate_limit_reset_seconds=rate_limit_policy.reset_seconds,
+                    router_rate_limit_configured_rpm=self.configured_rpm,
+                    router_rate_limit_key=self.key,
+                    router_rate_limit_recent=self.recent,
+                ),
+            ),
+        )
+
+    def register_backoff(
+        self,
+        provider: str,
+        config: dict[str, Any],
+        model: str | None,
+        retry_after: str | None = None,
+    ) -> float:
+        return register_rate_limit_backoff(
+            provider,
+            config,
+            model,
+            retry_after,
+            services=RateLimitBackoffServices(
+                state_store=self._state_store(),
+                policy=RateLimitBackoffPolicy(
+                    current_upstream_model_id=self.ports.current_model_id,
+                    parse_retry_after_seconds=rate_limit_policy.retry_after_seconds,
+                    provider_api_key_count=self.ports.api_key_count,
+                    router_rate_limit_capacity=self.capacity,
+                    router_rate_limit_configured_rpm=self.configured_rpm,
+                    router_rate_limit_effective_rpm=self.effective_rpm,
+                    router_rate_limit_key=self.key,
+                    router_rate_limit_recent=self.recent,
+                ),
+            ),
+        )
+
+    def apply(
+        self, provider: str, config: dict[str, Any], model: str | None = None
+    ) -> tuple[float, int, int | None]:
+        return apply_rate_limit(
+            provider,
+            config,
+            model,
+            services=RateLimitApplyServices(
+                state_store=self._state_store(),
+                policy=RateLimitApplyPolicy(
+                    current_upstream_model_id=self.ports.current_model_id,
+                    provider_api_key_count=self.ports.api_key_count,
+                    record_router_rate_usage=self.record_usage,
+                    router_rate_limit_capacity=self.capacity,
+                    router_rate_limit_effective_rpm=self.effective_rpm,
+                    router_rate_limit_key=self.key,
+                    router_rate_limit_recent=self.recent,
+                    wait_for_router_rate_limit_penalty=self.wait_for_penalty,
+                ),
+            ),
+        )
+
+    def wait_for_penalty(
+        self,
+        provider: str,
+        config: dict[str, Any],
+        model: str | None,
+        rpm: int | None,
+    ) -> float:
+        multi_key = self.ports.api_key_count(provider, config) > 1
+        waited = 0.0
+        while True:
+            entry = self.state_entry(provider, config, model)
+            now = self.ports.now()
+            try:
+                penalty_until = (
+                    0.0 if multi_key else float(entry.get("penalty_until") or 0.0)
+                )
+            except Exception:
+                penalty_until = 0.0
+            wait = max(0.0, penalty_until - now)
+            if wait <= 0.001:
+                return waited
+            sleep_for = min(wait, 10.0)
+            self.ports.log(
+                "INFO",
+                f"rate_limit_penalty_wait provider={provider} model={model or ''} "
+                f"rpm={rpm if rpm is not None else 'auto'} wait={wait:.2f}s "
+                f"waited={waited:.2f}s",
+            )
+            self.ports.sleep(sleep_for)
+            waited += sleep_for
+
+
+@dataclass(frozen=True, slots=True)
+class RouterRateLimitApi:
+    """Explicit compatibility API for late-bound router rate-limit services."""
+
+    service_factory: Callable[[], RouterRateLimitService]
+
+    def legacy_key(self, provider: str, pcfg: dict[str, Any], model: str | None) -> str:
+        return self.service_factory().legacy_key(provider, pcfg, model)
+
+    def configured_rpm(self, provider: str, pcfg: dict[str, Any]) -> int | None:
+        return self.service_factory().configured_rpm(provider, pcfg)
+
+    def rpm(self, provider: str, pcfg: dict[str, Any]) -> int | None:
+        return self.service_factory().rpm(provider, pcfg)
+
+    def key(self, provider: str, pcfg: dict[str, Any], model: str | None = None) -> str:
+        return self.service_factory().key(provider, pcfg, model)
+
+    def state_entry(self, provider: str, pcfg: dict[str, Any], model: str | None = None) -> dict[str, Any]:
+        return self.service_factory().state_entry(provider, pcfg, model)
+
+    def effective_rpm(self, provider: str, pcfg: dict[str, Any], model: str | None = None) -> int | None:
+        return self.service_factory().effective_rpm(provider, pcfg, model)
+
+    def capacity(self, rpm: int) -> int:
+        return self.service_factory().capacity(rpm)
+
+    def recent(self, timestamps: Any, now: float, window: float, *, include_future: bool) -> list[float]:
+        return self.service_factory().recent(timestamps, now, window, include_future=include_future)
+
+    def usage(self, provider: str, pcfg: dict[str, Any], model: str | None = None) -> tuple[int, int | None]:
+        return self.service_factory().usage(provider, pcfg, model)
+
+    def record_usage(self, provider: str, pcfg: dict[str, Any], model: str | None, rpm: int | None) -> tuple[int, int | None]:
+        return self.service_factory().record_usage(provider, pcfg, model, rpm)
+
+    def learn_headers(self, provider: str, pcfg: dict[str, Any], model: str | None, headers: Any) -> None:
+        self.service_factory().learn_headers(provider, pcfg, model, headers)
+
+    def register_backoff(self, provider: str, pcfg: dict[str, Any], model: str | None, retry_after: str | None = None) -> float:
+        return self.service_factory().register_backoff(provider, pcfg, model, retry_after)
+
+    def apply(self, provider: str, pcfg: dict[str, Any], model: str | None = None) -> tuple[float, int, int | None]:
+        return self.service_factory().apply(provider, pcfg, model)
+
+    def wait_for_penalty(self, provider: str, pcfg: dict[str, Any], model: str | None, rpm: int | None) -> float:
+        return self.service_factory().wait_for_penalty(provider, pcfg, model, rpm)
+
+
+__all__ = [
+    "RouterRateLimitApi",
+    "RouterRateLimitPaths",
+    "RouterRateLimitPorts",
+    "RouterRateLimitService",
+]