environment_variables.py

from __future__ import annotations

import os
from enum import Enum
from pathlib import Path

from dotenv import load_dotenv

from src.utils.logging import make_logger
from src.utils.model_utils import BaseModel

PROJECT_ROOT = Path(__file__).resolve().parents[2]
logger = make_logger(__name__)


class EnvVarKeys(str, Enum):
    ENVIRONMENT = "ENVIRONMENT"
    OPENAI_API_KEY = "OPENAI_API_KEY"
    DATABASE_URL = "DATABASE_URL"
    READ_ONLY_DATABASE_URL = "READ_ONLY_DATABASE_URL"
    TEMPORAL_ADDRESS = "TEMPORAL_ADDRESS"
    TEMPORAL_NAMESPACE = "TEMPORAL_NAMESPACE"
    REDIS_URL = "REDIS_URL"
    AGENTEX_BASE_URL = "AGENTEX_BASE_URL"
    TEMPORAL_WORKER_ACTIVITY_THREAD_POOL_SIZE = (
        "TEMPORAL_WORKER_ACTIVITY_THREAD_POOL_SIZE"
    )
    TEMPORAL_WORKER_MAX_ACTIVITIES_PER_WORKER = (
        "TEMPORAL_WORKER_MAX_ACTIVITIES_PER_WORKER"
    )
    BUILD_REGISTRY_URL = "BUILD_REGISTRY_URL"
    BUILD_CONTEXTS_PATH = "BUILD_CONTEXTS_PATH"
    BUILD_CONTEXT_PVC_NAME = "BUILD_CONTEXT_PVC_NAME"
    BUILD_REGISTRY_SECRET_NAME = "BUILD_REGISTRY_SECRET_NAME"
    AGENTS_NAMESPACE = "AGENTS_NAMESPACE"
    MONGODB_URI = "MONGODB_URI"
    MONGODB_DATABASE_NAME = "MONGODB_DATABASE_NAME"
    MONGODB_MAX_POOL_SIZE = "MONGODB_MAX_POOL_SIZE"
    MONGODB_MIN_POOL_SIZE = "MONGODB_MIN_POOL_SIZE"
    MONGODB_OIDC_REFRESH_INTERVAL_SECONDS = "MONGODB_OIDC_REFRESH_INTERVAL_SECONDS"
    REDIS_MAX_CONNECTIONS = "REDIS_MAX_CONNECTIONS"
    REDIS_CONNECTION_TIMEOUT = "REDIS_CONNECTION_TIMEOUT"
    REDIS_SOCKET_TIMEOUT = "REDIS_SOCKET_TIMEOUT"
    REDIS_STREAM_MAXLEN = "REDIS_STREAM_MAXLEN"
    REDIS_STREAM_TTL_SECONDS = "REDIS_STREAM_TTL_SECONDS"
    IMAGE_PULL_SECRET_NAME = "IMAGE_PULL_SECRET_NAME"
    AGENTEX_AUTH_URL = "AGENTEX_AUTH_URL"
    ALLOWED_ORIGINS = "ALLOWED_ORIGINS"
    DD_AGENT_HOST = "DD_AGENT_HOST"
    DD_STATSD_PORT = "DD_STATSD_PORT"
    HTTPX_MAX_CONNECTIONS = "HTTPX_MAX_CONNECTIONS"
    HTTPX_MAX_KEEPALIVE_CONNECTIONS = "HTTPX_MAX_KEEPALIVE_CONNECTIONS"
    HTTPX_CONNECT_TIMEOUT = "HTTPX_CONNECT_TIMEOUT"
    HTTPX_READ_TIMEOUT = "HTTPX_READ_TIMEOUT"
    HTTPX_WRITE_TIMEOUT = "HTTPX_WRITE_TIMEOUT"
    HTTPX_POOL_TIMEOUT = "HTTPX_POOL_TIMEOUT"
    HTTPX_STREAMING_READ_TIMEOUT = "HTTPX_STREAMING_READ_TIMEOUT"
    SSE_KEEPALIVE_PING_INTERVAL = "SSE_KEEPALIVE_PING_INTERVAL"
    AGENTEX_SERVER_TASK_QUEUE = "AGENTEX_SERVER_TASK_QUEUE"
    ENABLE_HEALTH_CHECK_WORKFLOW = "ENABLE_HEALTH_CHECK_WORKFLOW"
    WEBHOOK_REQUEST_TIMEOUT = "WEBHOOK_REQUEST_TIMEOUT"
    RETENTION_CLEANUP_ENABLED = "RETENTION_CLEANUP_ENABLED"
    RETENTION_CLEANUP_AGENT_ALLOWLIST = "RETENTION_CLEANUP_AGENT_ALLOWLIST"
    RETENTION_CLEANUP_IDLE_DAYS = "RETENTION_CLEANUP_IDLE_DAYS"
    RETENTION_CLEANUP_CRON = "RETENTION_CLEANUP_CRON"
    RETENTION_CLEANUP_PAGE_SIZE = "RETENTION_CLEANUP_PAGE_SIZE"
    RETENTION_CLEANUP_MAX_IN_FLIGHT = "RETENTION_CLEANUP_MAX_IN_FLIGHT"
    RETENTION_CLEANUP_DRY_RUN = "RETENTION_CLEANUP_DRY_RUN"


class Environment(str, Enum):
    DEV = "development"
    STAGING = "staging"
    PROD = "production"


refreshed_environment_variables = None


class EnvironmentVariables(BaseModel):
    ENVIRONMENT: str | None = Environment.DEV
    OPENAI_API_KEY: str | None
    DATABASE_URL: str | None
    READ_ONLY_DATABASE_URL: str | None = None
    TEMPORAL_ADDRESS: str | None
    TEMPORAL_NAMESPACE: str | None
    REDIS_URL: str | None
    AGENTEX_BASE_URL: str | None
    TEMPORAL_WORKER_ACTIVITY_THREAD_POOL_SIZE: int = 4  # Default 4 for local dev
    TEMPORAL_WORKER_MAX_ACTIVITIES_PER_WORKER: int = 10  # Default 10 for local dev
    BUILD_REGISTRY_URL: str | None = None
    BUILD_CONTEXTS_PATH: str | None = None
    BUILD_CONTEXT_PVC_NAME: str | None = None
    BUILD_REGISTRY_SECRET_NAME: str | None = None
    AGENTS_NAMESPACE: str | None = None
    MONGODB_URI: str | None = None
    MONGODB_DATABASE_NAME: str | None = "agentex"
    MONGODB_MAX_POOL_SIZE: int = 50
    MONGODB_MIN_POOL_SIZE: int = 5
    # Rebuild the Mongo client on this interval to renew GCP OIDC credentials.
    # pymongo caches the OIDC token for the life of the client and never refreshes
    # it proactively, so a long-lived client fails auth once the ~1h GCP token
    # expires. Only applied to MONGODB-OIDC URIs; 0 disables. Default 45 min.
    MONGODB_OIDC_REFRESH_INTERVAL_SECONDS: int = 2700
    REDIS_MAX_CONNECTIONS: int = 50  # Increased for SSE streaming
    REDIS_CONNECTION_TIMEOUT: int = 60  # Connection timeout in seconds
    REDIS_SOCKET_TIMEOUT: int = 30  # Socket timeout in seconds
    REDIS_STREAM_MAXLEN: int = (
        10000  # Max entries per Redis stream to prevent unbounded growth
    )
    REDIS_STREAM_TTL_SECONDS: int = (
        3600  # Sliding TTL on stream keys (seconds). 0 disables.
    )
    IMAGE_PULL_SECRET_NAME: str | None = None
    AGENTEX_AUTH_URL: str | None = None
    ALLOWED_ORIGINS: str | None = None
    HTTPX_MAX_CONNECTIONS: int = 200  # Max total connections allowed
    HTTPX_MAX_KEEPALIVE_CONNECTIONS: int = 100  # Max connections to keep alive
    HTTPX_CONNECT_TIMEOUT: float = 10.0  # HTTPX connection timeout in seconds
    HTTPX_READ_TIMEOUT: float = 30.0  # HTTPX read timeout in seconds
    HTTPX_WRITE_TIMEOUT: float = 30.0  # HTTPX write timeout in seconds
    HTTPX_POOL_TIMEOUT: float = 10.0  # HTTPX pool timeout in seconds
    HTTPX_STREAMING_READ_TIMEOUT: float = (
        300.0  # HTTPX streaming read timeout in seconds (5 minutes)
    )
    SSE_KEEPALIVE_PING_INTERVAL: int = 15  # SSE keepalive ping interval in seconds
    AGENTEX_SERVER_TASK_QUEUE: str | None = None
    ENABLE_HEALTH_CHECK_WORKFLOW: bool = False
    WEBHOOK_REQUEST_TIMEOUT: float = 15.0  # Webhook request timeout in seconds
    RETENTION_CLEANUP_ENABLED: bool = False
    RETENTION_CLEANUP_AGENT_ALLOWLIST: list[str] = []
    RETENTION_CLEANUP_IDLE_DAYS: int = 7
    RETENTION_CLEANUP_CRON: str = "0 4 * * *"
    RETENTION_CLEANUP_PAGE_SIZE: int = 200
    RETENTION_CLEANUP_MAX_IN_FLIGHT: int = 20
    RETENTION_CLEANUP_DRY_RUN: bool = True

    @classmethod
    def refresh(cls, force_refresh: bool = False) -> EnvironmentVariables | None:
        global refreshed_environment_variables
        if refreshed_environment_variables is not None and not force_refresh:
            return refreshed_environment_variables

        if os.environ.get(EnvVarKeys.ENVIRONMENT) == Environment.DEV:
            load_dotenv(dotenv_path=Path(PROJECT_ROOT / ".env"), override=True)
        environment_variables = EnvironmentVariables(
            ENVIRONMENT=os.environ.get(EnvVarKeys.ENVIRONMENT),
            OPENAI_API_KEY=os.environ.get(EnvVarKeys.OPENAI_API_KEY),
            DATABASE_URL=os.environ.get(EnvVarKeys.DATABASE_URL),
            READ_ONLY_DATABASE_URL=os.environ.get(EnvVarKeys.READ_ONLY_DATABASE_URL),
            TEMPORAL_ADDRESS=os.environ.get(EnvVarKeys.TEMPORAL_ADDRESS),
            TEMPORAL_NAMESPACE=os.environ.get(EnvVarKeys.TEMPORAL_NAMESPACE),
            REDIS_URL=os.environ.get(EnvVarKeys.REDIS_URL),
            AGENTEX_BASE_URL=os.environ.get(EnvVarKeys.AGENTEX_BASE_URL),
            BUILD_REGISTRY_URL=os.environ.get(EnvVarKeys.BUILD_REGISTRY_URL),
            BUILD_CONTEXTS_PATH=os.environ.get(EnvVarKeys.BUILD_CONTEXTS_PATH),
            BUILD_CONTEXT_PVC_NAME=os.environ.get(EnvVarKeys.BUILD_CONTEXT_PVC_NAME),
            BUILD_REGISTRY_SECRET_NAME=os.environ.get(
                EnvVarKeys.BUILD_REGISTRY_SECRET_NAME
            ),
            AGENTS_NAMESPACE=os.environ.get(EnvVarKeys.AGENTS_NAMESPACE),
            MONGODB_URI=os.environ.get(EnvVarKeys.MONGODB_URI),
            MONGODB_DATABASE_NAME=os.environ.get(
                EnvVarKeys.MONGODB_DATABASE_NAME, "agentex"
            ),
            MONGODB_MAX_POOL_SIZE=int(
                os.environ.get(EnvVarKeys.MONGODB_MAX_POOL_SIZE, "50")
            ),
            MONGODB_MIN_POOL_SIZE=int(
                os.environ.get(EnvVarKeys.MONGODB_MIN_POOL_SIZE, "5")
            ),
            MONGODB_OIDC_REFRESH_INTERVAL_SECONDS=int(
                os.environ.get(EnvVarKeys.MONGODB_OIDC_REFRESH_INTERVAL_SECONDS, "2700")
            ),
            REDIS_MAX_CONNECTIONS=int(
                os.environ.get(EnvVarKeys.REDIS_MAX_CONNECTIONS, "100")
            ),
            REDIS_CONNECTION_TIMEOUT=int(
                os.environ.get(EnvVarKeys.REDIS_CONNECTION_TIMEOUT, "20")
            ),
            REDIS_SOCKET_TIMEOUT=int(
                os.environ.get(EnvVarKeys.REDIS_SOCKET_TIMEOUT, "30")
            ),
            REDIS_STREAM_MAXLEN=int(
                os.environ.get(EnvVarKeys.REDIS_STREAM_MAXLEN, "10000")
            ),
            REDIS_STREAM_TTL_SECONDS=int(
                os.environ.get(EnvVarKeys.REDIS_STREAM_TTL_SECONDS, "3600")
            ),
            IMAGE_PULL_SECRET_NAME=os.environ.get(EnvVarKeys.IMAGE_PULL_SECRET_NAME),
            AGENTEX_AUTH_URL=os.environ.get(EnvVarKeys.AGENTEX_AUTH_URL),
            ALLOWED_ORIGINS=os.environ.get(EnvVarKeys.ALLOWED_ORIGINS, "*"),
            DD_AGENT_HOST=os.environ.get(EnvVarKeys.DD_AGENT_HOST),
            DD_STATSD_PORT=os.environ.get(EnvVarKeys.DD_STATSD_PORT),
            HTTPX_MAX_CONNECTIONS=int(
                os.environ.get(EnvVarKeys.HTTPX_MAX_CONNECTIONS, "200")
            ),
            HTTPX_MAX_KEEPALIVE_CONNECTIONS=int(
                os.environ.get(EnvVarKeys.HTTPX_MAX_KEEPALIVE_CONNECTIONS, "100")
            ),
            HTTPX_CONNECT_TIMEOUT=float(
                os.environ.get(EnvVarKeys.HTTPX_CONNECT_TIMEOUT, "10.0")
            ),
            HTTPX_READ_TIMEOUT=float(
                os.environ.get(EnvVarKeys.HTTPX_READ_TIMEOUT, "30.0")
            ),
            HTTPX_WRITE_TIMEOUT=float(
                os.environ.get(EnvVarKeys.HTTPX_WRITE_TIMEOUT, "30.0")
            ),
            HTTPX_POOL_TIMEOUT=float(
                os.environ.get(EnvVarKeys.HTTPX_POOL_TIMEOUT, "10.0")
            ),
            HTTPX_STREAMING_READ_TIMEOUT=float(
                os.environ.get(EnvVarKeys.HTTPX_STREAMING_READ_TIMEOUT, "300.0")
            ),
            SSE_KEEPALIVE_PING_INTERVAL=int(
                os.environ.get(EnvVarKeys.SSE_KEEPALIVE_PING_INTERVAL, "15")
            ),
            AGENTEX_SERVER_TASK_QUEUE=os.environ.get(
                EnvVarKeys.AGENTEX_SERVER_TASK_QUEUE
            ),
            ENABLE_HEALTH_CHECK_WORKFLOW=(
                os.environ.get(EnvVarKeys.ENABLE_HEALTH_CHECK_WORKFLOW, "false")
                == "true"
            ),
            WEBHOOK_REQUEST_TIMEOUT=float(
                os.environ.get(EnvVarKeys.WEBHOOK_REQUEST_TIMEOUT, "15.0")
            ),
            RETENTION_CLEANUP_ENABLED=(
                os.environ.get(EnvVarKeys.RETENTION_CLEANUP_ENABLED, "false") == "true"
            ),
            RETENTION_CLEANUP_AGENT_ALLOWLIST=[
                name.strip()
                for name in os.environ.get(
                    EnvVarKeys.RETENTION_CLEANUP_AGENT_ALLOWLIST, ""
                ).split(",")
                if name.strip()
            ],
            RETENTION_CLEANUP_IDLE_DAYS=int(
                os.environ.get(EnvVarKeys.RETENTION_CLEANUP_IDLE_DAYS, "7")
            ),
            RETENTION_CLEANUP_CRON=os.environ.get(
                EnvVarKeys.RETENTION_CLEANUP_CRON, "0 4 * * *"
            ),
            RETENTION_CLEANUP_PAGE_SIZE=int(
                os.environ.get(EnvVarKeys.RETENTION_CLEANUP_PAGE_SIZE, "200")
            ),
            RETENTION_CLEANUP_MAX_IN_FLIGHT=int(
                os.environ.get(EnvVarKeys.RETENTION_CLEANUP_MAX_IN_FLIGHT, "20")
            ),
            RETENTION_CLEANUP_DRY_RUN=(
                os.environ.get(EnvVarKeys.RETENTION_CLEANUP_DRY_RUN, "true") == "true"
            ),
        )
        refreshed_environment_variables = environment_variables
        return refreshed_environment_variables

    @classmethod
    def clear_cache(cls):
        """Clear the cached environment variables to force refresh on next access"""
        global refreshed_environment_variables
        refreshed_environment_variables = None