@ngocsangairvds/vsaf 3.2.14 → 3.2.15

package/tools/vds-scripts/audit_orchestrator/src/vds_audit_orchestrator/llm/codex_oauth.py

@@ -1,340 +0,0 @@
-"""Codex OAuth token lifecycle for the openai-codex protocol."""
-
-from __future__ import annotations
-
-import base64
-import contextlib
-import json
-import os
-import tempfile
-import time
-from dataclasses import dataclass
-from datetime import UTC, datetime
-from pathlib import Path
-from typing import Any
-
-import httpx
-from filelock import FileLock
-
-CODEX_AUTH_FILE = Path.home() / ".codex" / "auth.json"
-CODEX_TOKEN_ENDPOINT = "https://auth.openai.com/oauth/token"
-CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
-CODEX_SCOPE = "openid profile email"
-CODEX_REFRESH_STALE_SECONDS = 8 * 24 * 60 * 60
-CODEX_EXPIRY_SAFETY_SECONDS = 60
-CODEX_LOCK_TIMEOUT_SECONDS = 30.0
-
-_ALLOWED_CODEX_FIELDS: frozenset[str] = frozenset({"access_token", "account_id", "refresh_token"})
-
-
-class CodexOAuthError(ValueError):
-    """Raised when Codex OAuth credentials are missing or invalid."""
-
-
-@dataclass(slots=True, frozen=True)
-class CodexCredentials:
-    """Resolved Codex credentials for subscription-provider requests."""
-
-    access_token: str
-    account_id: str
-    is_refreshed: bool
-
-
-def load_codex_auth_json(auth_file: Path | None = None) -> dict[str, Any]:
-    """Load and minimally validate ~/.codex/auth.json."""
-    path = _resolve_auth_file(auth_file)
-    if not path.exists():
-        raise CodexOAuthError(_missing_auth_file_message(path))
-    try:
-        payload = json.loads(path.read_text(encoding="utf-8"))
-    except json.JSONDecodeError as exc:
-        raise CodexOAuthError(f"Codex auth file is invalid JSON: {path}") from exc
-    if not isinstance(payload, dict):
-        raise CodexOAuthError(f"Codex auth file must be a JSON object: {path}")
-    _validate_auth_dot_json_shape(payload=payload, auth_file=path)
-    return payload
-
-
-def get_codex_credentials(auth_file: Path | None = None) -> CodexCredentials:
-    """Return an access token + account id, auto-refreshing stale tokens."""
-    tokens, is_refreshed = _load_tokens_with_refresh(auth_file=auth_file)
-    access_token = tokens.get("access_token", "")
-    account_id = tokens.get("account_id", "")
-    if not access_token or not account_id:
-        path = _resolve_auth_file(auth_file)
-        raise CodexOAuthError(
-            f"Codex OAuth tokens are incomplete in {path}. Expected tokens.access_token and tokens.account_id. "
-            "Run 'codex --login' (or run 'codex' and choose Sign in with ChatGPT), or switch to "
-            "model.llm.type: openai with OPENAI_API_KEY."
-        )
-    return CodexCredentials(access_token=access_token, account_id=account_id, is_refreshed=is_refreshed)
-
-
-def force_refresh_codex_credentials(auth_file: Path | None = None) -> CodexCredentials:
-    """Return credentials after forcing a refresh, even when token is not stale."""
-    tokens, _ = _load_tokens_with_refresh(auth_file=auth_file, force_refresh=True)
-    access_token = tokens.get("access_token", "")
-    account_id = tokens.get("account_id", "")
-    if not access_token or not account_id:
-        path = _resolve_auth_file(auth_file)
-        raise CodexOAuthError(
-            f"Codex OAuth tokens are incomplete in {path}. Expected tokens.access_token and tokens.account_id. "
-            "Run 'codex --login' (or run 'codex' and choose Sign in with ChatGPT), or switch to "
-            "model.llm.type: openai with OPENAI_API_KEY."
-        )
-    return CodexCredentials(access_token=access_token, account_id=account_id, is_refreshed=True)
-
-
-class CodexCredentialResolverImpl:
-    """Audit-side implementation of vds_agent_core Codex credential resolution."""
-
-    def resolve(self, field: str) -> str:
-        return resolve_codex_reference(field)
-
-
-def resolve_codex_reference(field: str, *, auth_file: Path | None = None) -> str:
-    """Resolve `${codex:<field>}` placeholders for runtime profiles."""
-    normalized_field = str(field or "").strip().lower()
-    if normalized_field not in _ALLOWED_CODEX_FIELDS:
-        allowed = ", ".join(sorted(_ALLOWED_CODEX_FIELDS))
-        raise CodexOAuthError(f"Unsupported codex reference field '{field}'. Allowed fields: {allowed}.")
-
-    tokens, _ = _load_tokens_with_refresh(auth_file=auth_file)
-    value = tokens.get(normalized_field, "").strip()
-    if not value:
-        path = _resolve_auth_file(auth_file)
-        raise CodexOAuthError(
-            f"Codex auth file {path} is missing tokens.{normalized_field}. "
-            "Run 'codex --login' (or run 'codex' and choose Sign in with ChatGPT), or switch to "
-            "model.llm.type: openai with OPENAI_API_KEY."
-        )
-    return value
-
-
-def _load_tokens_with_refresh(
-    *, auth_file: Path | None = None, force_refresh: bool = False
-) -> tuple[dict[str, str], bool]:
-    path = _resolve_auth_file(auth_file)
-    lock_path = path.with_name(f"{path.name}.lock")
-    with FileLock(str(lock_path), timeout=CODEX_LOCK_TIMEOUT_SECONDS):
-        payload = load_codex_auth_json(path)
-        _validate_auth_mode(payload=payload, auth_file=path)
-
-        tokens = _normalize_tokens(payload.get("tokens"))
-        access_token = tokens.get("access_token", "")
-        account_id = tokens.get("account_id", "")
-        if not access_token:
-            raise CodexOAuthError(_missing_token_bundle_message(path))
-        if not account_id:
-            raise CodexOAuthError(
-                f"Codex auth file {path} is missing tokens.account_id. Run 'codex --login' to refresh auth state."
-            )
-
-        if not force_refresh and not _is_token_stale(
-            access_token=access_token, last_refresh=payload.get("last_refresh")
-        ):
-            return tokens, False
-
-        refresh_token = tokens.get("refresh_token", "")
-        if not refresh_token:
-            raise CodexOAuthError(_missing_token_bundle_message(path))
-
-        refreshed = _refresh_tokens(refresh_token=refresh_token)
-        tokens.update({key: value for key, value in refreshed.items() if value})
-        if not tokens.get("account_id"):
-            tokens["account_id"] = account_id
-
-        persisted = dict(payload)
-        existing_tokens = payload.get("tokens") if isinstance(payload.get("tokens"), dict) else {}
-        persisted["tokens"] = dict(existing_tokens) | tokens
-        persisted["last_refresh"] = datetime.now(UTC).isoformat()
-        _persist_auth_json_atomic(auth_file=path, payload=persisted)
-        return tokens, True
-
-
-def _resolve_auth_file(auth_file: Path | None) -> Path:
-    return (auth_file or CODEX_AUTH_FILE).expanduser()
-
-
-def _validate_auth_dot_json_shape(*, payload: dict[str, Any], auth_file: Path) -> None:
-    for key in ("auth_mode", "openai_api_key", "last_refresh"):
-        value = payload.get(key)
-        if value is None:
-            continue
-        if not isinstance(value, str):
-            raise CodexOAuthError(f"Codex auth file field '{key}' must be a string in {auth_file}.")
-
-    raw_tokens = payload.get("tokens")
-    if raw_tokens is None:
-        return
-    if not isinstance(raw_tokens, dict):
-        raise CodexOAuthError(f"Codex auth file field 'tokens' must be an object in {auth_file}.")
-
-    for token_key in ("access_token", "refresh_token", "account_id", "id_token"):
-        token_value = raw_tokens.get(token_key)
-        if token_value is None:
-            continue
-        if not isinstance(token_value, str):
-            raise CodexOAuthError(f"Codex auth file field 'tokens.{token_key}' must be a string in {auth_file}.")
-
-
-def _validate_auth_mode(*, payload: dict[str, Any], auth_file: Path) -> None:
-    auth_mode = str(payload.get("auth_mode") or "").strip().lower()
-    has_openai_api_key = bool(str(payload.get("openai_api_key") or "").strip())
-    if auth_mode == "api_key" or (has_openai_api_key and not isinstance(payload.get("tokens"), dict)):
-        raise CodexOAuthError(
-            f"{auth_file} uses API key mode, not ChatGPT OAuth. Either: "
-            "(1) run 'codex --login' (or run 'codex' and choose Sign in with ChatGPT) to authenticate with your "
-            "ChatGPT plan and use the 'openai-codex' protocol, or "
-            "(2) switch to the 'openai' protocol with model.llm.type: openai and set OPENAI_API_KEY."
-        )
-
-
-def _normalize_tokens(raw_tokens: Any) -> dict[str, str]:
-    if not isinstance(raw_tokens, dict):
-        return {}
-    normalized: dict[str, str] = {}
-    for key in ("access_token", "refresh_token", "account_id", "id_token"):
-        value = raw_tokens.get(key)
-        if isinstance(value, str) and value.strip():
-            normalized[key] = value.strip()
-    return normalized
-
-
-def _is_token_stale(*, access_token: str, last_refresh: Any) -> bool:
-    now = time.time()
-    exp = _extract_jwt_exp(access_token)
-    if exp is not None:
-        return exp <= (now + CODEX_EXPIRY_SAFETY_SECONDS)
-
-    parsed_last_refresh = _parse_last_refresh(last_refresh)
-    if parsed_last_refresh is None:
-        return True
-    return (now - parsed_last_refresh.timestamp()) >= CODEX_REFRESH_STALE_SECONDS
-
-
-def _extract_jwt_exp(access_token: str) -> int | None:
-    parts = str(access_token).split(".")
-    if len(parts) < 2:
-        return None
-    payload_segment = parts[1]
-    padding = "=" * ((4 - len(payload_segment) % 4) % 4)
-    try:
-        decoded = base64.urlsafe_b64decode(f"{payload_segment}{padding}".encode())
-        payload = json.loads(decoded.decode("utf-8"))
-    except (ValueError, json.JSONDecodeError, UnicodeDecodeError):
-        return None
-    if not isinstance(payload, dict):
-        return None
-    raw_exp = payload.get("exp")
-    try:
-        return int(raw_exp) if raw_exp is not None else None
-    except (TypeError, ValueError):
-        return None
-
-
-def _parse_last_refresh(value: Any) -> datetime | None:
-    if not isinstance(value, str) or not value.strip():
-        return None
-    normalized = value.strip()
-    if normalized.endswith("Z"):
-        normalized = f"{normalized[:-1]}+00:00"
-    try:
-        parsed = datetime.fromisoformat(normalized)
-    except ValueError:
-        return None
-    if parsed.tzinfo is None:
-        parsed = parsed.replace(tzinfo=UTC)
-    return parsed.astimezone(UTC)
-
-
-def _refresh_tokens(*, refresh_token: str) -> dict[str, str]:
-    payload = {
-        "grant_type": "refresh_token",
-        "client_id": CODEX_CLIENT_ID,
-        "refresh_token": refresh_token,
-        "scope": CODEX_SCOPE,
-    }
-    try:
-        response = httpx.post(CODEX_TOKEN_ENDPOINT, data=payload, timeout=30.0)
-    except httpx.HTTPError as exc:
-        raise CodexOAuthError(
-            "Failed to refresh Codex OAuth token due to network/transport error. "
-            "Retry later, or run 'codex --login' to re-authenticate."
-        ) from exc
-
-    if response.status_code == 401:
-        error_code = _read_refresh_error_code(response)
-        if error_code:
-            raise CodexOAuthError(
-                f"Codex OAuth refresh token was rejected ({error_code}). "
-                "Refresh tokens are single-use and may be expired/revoked/reused. Run 'codex --login'."
-            )
-        raise CodexOAuthError("Codex OAuth refresh token was rejected (401 Unauthorized). Run 'codex --login'.")
-
-    try:
-        response.raise_for_status()
-    except httpx.HTTPStatusError as exc:
-        raise CodexOAuthError(
-            f"Codex OAuth refresh failed with HTTP {response.status_code}. Run 'codex --login' if this persists."
-        ) from exc
-
-    try:
-        data = response.json()
-    except ValueError as exc:
-        raise CodexOAuthError("Codex OAuth refresh response was not valid JSON.") from exc
-    if not isinstance(data, dict):
-        raise CodexOAuthError("Codex OAuth refresh response must be a JSON object.")
-
-    access_token = str(data.get("access_token") or "").strip()
-    rotated_refresh_token = str(data.get("refresh_token") or "").strip()
-    id_token = str(data.get("id_token") or "").strip()
-    if not access_token or not rotated_refresh_token:
-        raise CodexOAuthError("Codex OAuth refresh response missing access_token/refresh_token. Run 'codex --login'.")
-
-    refreshed: dict[str, str] = {"access_token": access_token, "refresh_token": rotated_refresh_token}
-    if id_token:
-        refreshed["id_token"] = id_token
-    return refreshed
-
-
-def _read_refresh_error_code(response: httpx.Response) -> str | None:
-    with contextlib.suppress(ValueError):
-        payload = response.json()
-        if isinstance(payload, dict):
-            code = payload.get("error")
-            if isinstance(code, str) and code.strip():
-                return code.strip()
-    return None
-
-
-def _persist_auth_json_atomic(*, auth_file: Path, payload: dict[str, Any]) -> None:
-    auth_file.parent.mkdir(parents=True, exist_ok=True)
-    fd, temp_path = tempfile.mkstemp(prefix=f".{auth_file.name}.", suffix=".tmp", dir=str(auth_file.parent))
-    try:
-        with os.fdopen(fd, "w", encoding="utf-8") as handle:
-            json.dump(payload, handle, ensure_ascii=True, indent=2, sort_keys=True)
-            handle.write("\n")
-            handle.flush()
-            os.fsync(handle.fileno())
-        os.replace(temp_path, auth_file)
-    except Exception:
-        with contextlib.suppress(FileNotFoundError):
-            os.unlink(temp_path)
-        raise
-
-
-def _missing_auth_file_message(auth_file: Path) -> str:
-    return (
-        f"Codex OAuth file not found: {auth_file}. Run 'codex --login' (or run 'codex' and choose "
-        "Sign in with ChatGPT) to authenticate with your ChatGPT subscription, or switch to "
-        "model.llm.type: openai with OPENAI_API_KEY for API-key billing."
-    )
-
-
-def _missing_token_bundle_message(auth_file: Path) -> str:
-    return (
-        f"Codex OAuth tokens not found in {auth_file}. Run 'codex --login' (or run 'codex' and choose "
-        "Sign in with ChatGPT) to bootstrap refreshable tokens, or switch to model.llm.type: openai "
-        "with OPENAI_API_KEY."
-    )

package/tools/vds-scripts/audit_orchestrator/src/vds_audit_orchestrator/llm/cost_tracker.py

@@ -1,288 +0,0 @@
-"""LLM usage and cost tracking (TSK-043)."""
-
-from __future__ import annotations
-
-import os
-import time
-from datetime import datetime
-from typing import Any
-
-import structlog
-from pydantic import BaseModel, Field
-
-logger = structlog.get_logger(__name__)
-
-
-class ModelCost(BaseModel):
-    """Cost configuration for a model."""
-
-    model_name: str
-    input_cost_per_1m: float  # Cost per 1M input tokens
-    output_cost_per_1m: float  # Cost per 1M output tokens
-
-
-class UsageRecord(BaseModel):
-    """Single usage record."""
-
-    timestamp: datetime = Field(default_factory=datetime.now)
-    model: str
-    prompt_type: str
-    input_tokens: int
-    output_tokens: int
-    cost: float
-    cached: bool = False
-    duration_ms: int = 0
-
-
-class UsageSummary(BaseModel):
-    """Summary of usage and costs."""
-
-    total_requests: int = 0
-    cached_requests: int = 0
-    total_input_tokens: int = 0
-    total_output_tokens: int = 0
-    total_cost: float = 0.0
-    cost_by_model: dict[str, float] = Field(default_factory=dict)
-    cost_by_prompt_type: dict[str, float] = Field(default_factory=dict)
-    avg_duration_ms: float = 0.0
-
-
-class CostTracker:
-    """Track LLM usage and costs."""
-
-    @staticmethod
-    def _resolve_default_model() -> str:
-        """Resolve active model from environment, preferring runtime tiered model config."""
-        for env_key in (
-            "VDS_AUDIT_LLM__MODEL_STANDARD",
-            "VDS_AUDIT_LLM__MODEL",
-            "VDS_AUDIT_LLM__MODEL_COMPLEX",
-            "VDS_AUDIT_LLM__MODEL_SIMPLE",
-        ):
-            value = (os.getenv(env_key) or "").strip()
-            if value:
-                return value
-        return "unknown"
-
-    # Default model costs (per 1M tokens)
-    MODEL_COSTS: dict[str, ModelCost] = {
-        "claude-sonnet-4-20250514": ModelCost(
-            model_name="claude-sonnet-4-20250514",
-            input_cost_per_1m=3.00,
-            output_cost_per_1m=15.00,
-        ),
-        "claude-opus-4-20250514": ModelCost(
-            model_name="claude-opus-4-20250514",
-            input_cost_per_1m=15.00,
-            output_cost_per_1m=75.00,
-        ),
-        "claude-haiku-3-20240307": ModelCost(
-            model_name="claude-haiku-3-20240307",
-            input_cost_per_1m=0.25,
-            output_cost_per_1m=1.25,
-        ),
-        # Claude 3.x model names (still used in default config)
-        "claude-3-haiku-20240307": ModelCost(
-            model_name="claude-3-haiku-20240307",
-            input_cost_per_1m=0.25,
-            output_cost_per_1m=1.25,
-        ),
-        "claude-3-sonnet-20240229": ModelCost(
-            model_name="claude-3-sonnet-20240229",
-            input_cost_per_1m=3.00,
-            output_cost_per_1m=15.00,
-        ),
-        "claude-3-opus-20240229": ModelCost(
-            model_name="claude-3-opus-20240229",
-            input_cost_per_1m=15.00,
-            output_cost_per_1m=75.00,
-        ),
-        "claude-3-5-sonnet-20240620": ModelCost(
-            model_name="claude-3-5-sonnet-20240620",
-            input_cost_per_1m=3.00,
-            output_cost_per_1m=15.00,
-        ),
-        # Local models are free
-        "qwen3-coder:30b": ModelCost(
-            model_name="qwen3-coder:30b",
-            input_cost_per_1m=0.0,
-            output_cost_per_1m=0.0,
-        ),
-        "glm-4.7-flash": ModelCost(
-            model_name="glm-4.7-flash",
-            input_cost_per_1m=0.0,
-            output_cost_per_1m=0.0,
-        ),
-    }
-
-    def __init__(self, budget_limit: float | None = None):
-        self.budget_limit = budget_limit
-        self.records: list[UsageRecord] = []
-        self._start_time = time.time()
-        # Aggregate tracking attributes (used by get_summary and callers)
-        self.input_tokens: int = 0
-        self.output_tokens: int = 0
-        self.total_cost: float = 0.0
-        self.model_usage: dict[str, dict[str, int]] = {}
-
-    def get_model_cost(self, model: str) -> ModelCost:
-        """Get cost configuration for model."""
-        if model in self.MODEL_COSTS:
-            return self.MODEL_COSTS[model]
-        # Default to free for unknown models (likely local)
-        return ModelCost(
-            model_name=model,
-            input_cost_per_1m=0.0,
-            output_cost_per_1m=0.0,
-        )
-
-    def calculate_cost(
-        self,
-        model: str,
-        input_tokens: int,
-        output_tokens: int,
-    ) -> float:
-        """Calculate cost for a request."""
-        costs = self.get_model_cost(model)
-        input_cost = (input_tokens / 1_000_000) * costs.input_cost_per_1m
-        output_cost = (output_tokens / 1_000_000) * costs.output_cost_per_1m
-        return input_cost + output_cost
-
-    def record_usage(
-        self,
-        model: str,
-        prompt_type: str,
-        input_tokens: int,
-        output_tokens: int,
-        cached: bool = False,
-        duration_ms: int = 0,
-    ) -> UsageRecord:
-        """Record a usage event."""
-        cost = 0.0 if cached else self.calculate_cost(model, input_tokens, output_tokens)
-
-        record = UsageRecord(
-            model=model,
-            prompt_type=prompt_type,
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
-            cost=cost,
-            cached=cached,
-            duration_ms=duration_ms,
-        )
-
-        self.records.append(record)
-
-        # Update compatibility tracking
-        self.input_tokens += input_tokens
-        self.output_tokens += output_tokens
-        self.total_cost += cost
-
-        if model not in self.model_usage:
-            self.model_usage[model] = {"input": 0, "output": 0, "count": 0}
-        self.model_usage[model]["input"] += input_tokens
-        self.model_usage[model]["output"] += output_tokens
-        self.model_usage[model]["count"] += 1
-
-        logger.info(
-            "LLM usage recorded",
-            model=model,
-            tokens=input_tokens + output_tokens,
-            cost=f"${cost:.4f}",
-            cached=cached,
-        )
-
-        # Check budget
-        if self.budget_limit:
-            total = self.get_summary().total_cost
-            if total >= self.budget_limit:
-                logger.warning(
-                    "Budget limit reached",
-                    limit=self.budget_limit,
-                    current=total,
-                )
-
-        return record
-
-    def track_usage(self, model: str, input_tokens: int, output_tokens: int) -> float:
-        """Record usage and return cost for this call (compatibility wrapper)."""
-        record = self.record_usage(
-            model=model,
-            prompt_type="compat",
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
-        )
-        return record.cost
-
-    def get_summary(self) -> UsageSummary:
-        """Get usage summary."""
-        if not self.records:
-            return UsageSummary()
-
-        total_duration = sum(r.duration_ms for r in self.records if not r.cached)
-        non_cached = len([r for r in self.records if not r.cached])
-
-        cost_by_model: dict[str, float] = {}
-        cost_by_prompt: dict[str, float] = {}
-
-        for r in self.records:
-            cost_by_model[r.model] = cost_by_model.get(r.model, 0) + r.cost
-            cost_by_prompt[r.prompt_type] = cost_by_prompt.get(r.prompt_type, 0) + r.cost
-
-        return UsageSummary(
-            total_requests=len(self.records),
-            cached_requests=len([r for r in self.records if r.cached]),
-            total_input_tokens=sum(r.input_tokens for r in self.records),
-            total_output_tokens=sum(r.output_tokens for r in self.records),
-            total_cost=sum(r.cost for r in self.records),
-            cost_by_model=cost_by_model,
-            cost_by_prompt_type=cost_by_prompt,
-            avg_duration_ms=total_duration / non_cached if non_cached > 0 else 0,
-        )
-
-    def summary(self) -> dict[str, Any]:
-        """Get usage summary (compatibility wrapper)."""
-        return {
-            "total_input_tokens": self.input_tokens,
-            "total_output_tokens": self.output_tokens,
-            "total_cost_usd": round(self.total_cost, 4),
-            "model_breakdown": self.model_usage,
-        }
-
-    def estimate_audit_cost(
-        self,
-        num_repos: int,
-        checks_per_repo: int,
-        llm_check_ratio: float = 0.4,
-        avg_tokens_per_check: int = 2000,
-        model: str | None = None,
-    ) -> dict[str, Any]:
-        """Estimate cost for an audit run."""
-        llm_checks = int(num_repos * checks_per_repo * llm_check_ratio)
-        total_tokens = llm_checks * avg_tokens_per_check
-
-        selected_model = (model or "").strip() or self._resolve_default_model()
-        costs = self.get_model_cost(selected_model)
-        estimated_cost = (total_tokens / 1_000_000) * (costs.input_cost_per_1m + costs.output_cost_per_1m) / 2
-
-        return {
-            "num_repos": num_repos,
-            "checks_per_repo": checks_per_repo,
-            "llm_checks": llm_checks,
-            "estimated_tokens": total_tokens,
-            "estimated_cost": round(estimated_cost, 2),
-            "model": selected_model,
-        }
-
-    def estimate_remaining_budget(self, budget: float) -> float:
-        """Calculate remaining budget (compatibility wrapper)."""
-        return max(0.0, budget - self.total_cost)
-
-    def is_within_budget(self) -> bool:
-        """Check if current usage is within budget."""
-        if self.budget_limit is None:
-            return True
-        return self.get_summary().total_cost < self.budget_limit
-
-
-# Global tracker instance (compatibility wrapper)
-global_tracker = CostTracker()