feat: centralized environment handling, ensures ZEN_MCP_FORCE_ENV_OVERRIDE is honored correctly

fix: updated tests to override env variables they need instead of relying on the current values from .env
This commit is contained in:
Fahad
2025-10-04 14:28:56 +04:00
parent 4015e917ed
commit 2c534ac06e
24 changed files with 300 additions and 179 deletions

View File

@@ -28,35 +28,6 @@ from logging.handlers import RotatingFileHandler
from pathlib import Path
from typing import Any, Optional
# Try to load environment variables from .env file if dotenv is available
# This is optional - environment variables can still be passed directly
try:
from dotenv import dotenv_values, load_dotenv
# Load environment variables from .env file in the script's directory
# This ensures .env is loaded regardless of the current working directory
script_dir = Path(__file__).parent
env_file = script_dir / ".env"
# First load only to read ZEN_MCP_FORCE_ENV_OVERRIDE, then reload with proper override setting
# Use a temporary environment to read just this configuration variable
temp_env = {}
if env_file.exists():
temp_env = dotenv_values(env_file)
# Check if we should force override based on .env file content (not system env)
force_override = temp_env.get("ZEN_MCP_FORCE_ENV_OVERRIDE", "false").lower() == "true"
# Load .env file with appropriate override setting
load_dotenv(dotenv_path=env_file, override=force_override)
# Store override setting for logging after logger is configured
_zen_mcp_force_override = force_override
except ImportError:
# dotenv not available - this is fine, environment variables can still be passed directly
# This commonly happens when running via uvx or in minimal environments
pass
from mcp.server import Server # noqa: E402
from mcp.server.models import InitializationOptions # noqa: E402
from mcp.server.stdio import stdio_server # noqa: E402
@@ -95,10 +66,11 @@ from tools import ( # noqa: E402
VersionTool,
)
from tools.models import ToolOutput # noqa: E402
from utils.env import env_override_enabled, get_env # noqa: E402
# Configure logging for server operations
# Can be controlled via LOG_LEVEL environment variable (DEBUG, INFO, WARNING, ERROR)
log_level = os.getenv("LOG_LEVEL", "DEBUG").upper()
log_level = (get_env("LOG_LEVEL", "DEBUG") or "DEBUG").upper()
# Create timezone-aware formatter
@@ -177,19 +149,12 @@ except Exception as e:
logger = logging.getLogger(__name__)
# Log ZEN_MCP_FORCE_ENV_OVERRIDE configuration if it was set during dotenv loading
try:
if "_zen_mcp_force_override" in globals():
if _zen_mcp_force_override:
logger.info(
"ZEN_MCP_FORCE_ENV_OVERRIDE enabled - .env file values will override system environment variables"
)
logger.debug("Environment override prevents conflicts between different AI tools passing cached API keys")
else:
logger.debug("ZEN_MCP_FORCE_ENV_OVERRIDE disabled - system environment variables take precedence")
except NameError:
# _zen_mcp_force_override not defined, which means dotenv wasn't available or no .env file
pass
# Log ZEN_MCP_FORCE_ENV_OVERRIDE configuration for transparency
if env_override_enabled():
logger.info("ZEN_MCP_FORCE_ENV_OVERRIDE enabled - .env file values will override system environment variables")
logger.debug("Environment override prevents conflicts between different AI tools passing cached API keys")
else:
logger.debug("ZEN_MCP_FORCE_ENV_OVERRIDE disabled - system environment variables take precedence")
# Create the MCP server instance with a unique name identifier
@@ -208,7 +173,7 @@ def parse_disabled_tools_env() -> set[str]:
Returns:
Set of lowercase tool names to disable, empty set if none specified
"""
disabled_tools_env = os.getenv("DISABLED_TOOLS", "").strip()
disabled_tools_env = (get_env("DISABLED_TOOLS", "") or "").strip()
if not disabled_tools_env:
return set()
return {t.strip().lower() for t in disabled_tools_env.split(",") if t.strip()}
@@ -409,7 +374,7 @@ def configure_providers():
logger.debug("Checking environment variables for API keys...")
api_keys_to_check = ["OPENAI_API_KEY", "OPENROUTER_API_KEY", "GEMINI_API_KEY", "XAI_API_KEY", "CUSTOM_API_URL"]
for key in api_keys_to_check:
value = os.getenv(key)
value = get_env(key)
logger.debug(f" {key}: {'[PRESENT]' if value else '[MISSING]'}")
from providers import ModelProviderRegistry
from providers.custom import CustomProvider
@@ -427,14 +392,14 @@ def configure_providers():
has_custom = False
# Check for Gemini API key
gemini_key = os.getenv("GEMINI_API_KEY")
gemini_key = get_env("GEMINI_API_KEY")
if gemini_key and gemini_key != "your_gemini_api_key_here":
valid_providers.append("Gemini")
has_native_apis = True
logger.info("Gemini API key found - Gemini models available")
# Check for OpenAI API key
openai_key = os.getenv("OPENAI_API_KEY")
openai_key = get_env("OPENAI_API_KEY")
logger.debug(f"OpenAI key check: key={'[PRESENT]' if openai_key else '[MISSING]'}")
if openai_key and openai_key != "your_openai_api_key_here":
valid_providers.append("OpenAI")
@@ -447,21 +412,21 @@ def configure_providers():
logger.debug("OpenAI API key is placeholder value")
# Check for X.AI API key
xai_key = os.getenv("XAI_API_KEY")
xai_key = get_env("XAI_API_KEY")
if xai_key and xai_key != "your_xai_api_key_here":
valid_providers.append("X.AI (GROK)")
has_native_apis = True
logger.info("X.AI API key found - GROK models available")
# Check for DIAL API key
dial_key = os.getenv("DIAL_API_KEY")
dial_key = get_env("DIAL_API_KEY")
if dial_key and dial_key != "your_dial_api_key_here":
valid_providers.append("DIAL")
has_native_apis = True
logger.info("DIAL API key found - DIAL models available")
# Check for OpenRouter API key
openrouter_key = os.getenv("OPENROUTER_API_KEY")
openrouter_key = get_env("OPENROUTER_API_KEY")
logger.debug(f"OpenRouter key check: key={'[PRESENT]' if openrouter_key else '[MISSING]'}")
if openrouter_key and openrouter_key != "your_openrouter_api_key_here":
valid_providers.append("OpenRouter")
@@ -474,14 +439,14 @@ def configure_providers():
logger.debug("OpenRouter API key is placeholder value")
# Check for custom API endpoint (Ollama, vLLM, etc.)
custom_url = os.getenv("CUSTOM_API_URL")
custom_url = get_env("CUSTOM_API_URL")
if custom_url:
# IMPORTANT: Always read CUSTOM_API_KEY even if empty
# - Some providers (vLLM, LM Studio, enterprise APIs) require authentication
# - Others (Ollama) work without authentication (empty key)
# - DO NOT remove this variable - it's needed for provider factory function
custom_key = os.getenv("CUSTOM_API_KEY", "") # Default to empty (Ollama doesn't need auth)
custom_model = os.getenv("CUSTOM_MODEL_NAME", "llama3.2")
custom_key = get_env("CUSTOM_API_KEY", "") or "" # Default to empty (Ollama doesn't need auth)
custom_model = get_env("CUSTOM_MODEL_NAME", "llama3.2") or "llama3.2"
valid_providers.append(f"Custom API ({custom_url})")
has_custom = True
logger.info(f"Custom API endpoint found: {custom_url} with model {custom_model}")
@@ -517,7 +482,7 @@ def configure_providers():
# Factory function that creates CustomProvider with proper parameters
def custom_provider_factory(api_key=None):
# api_key is CUSTOM_API_KEY (can be empty for Ollama), base_url from CUSTOM_API_URL
base_url = os.getenv("CUSTOM_API_URL", "")
base_url = get_env("CUSTOM_API_URL", "") or ""
return CustomProvider(api_key=api_key or "", base_url=base_url) # Use provided API key or empty string
ModelProviderRegistry.register_provider(ProviderType.CUSTOM, custom_provider_factory)
@@ -674,7 +639,8 @@ async def handle_list_tools() -> list[Tool]:
)
# Log cache efficiency info
if os.getenv("OPENROUTER_API_KEY") and os.getenv("OPENROUTER_API_KEY") != "your_openrouter_api_key_here":
openrouter_key_for_cache = get_env("OPENROUTER_API_KEY")
if openrouter_key_for_cache and openrouter_key_for_cache != "your_openrouter_api_key_here":
logger.debug("OpenRouter registry cache used efficiently across all tool schemas")
logger.debug(f"Returning {len(tools)} tools to MCP client")