fix: Coordinator split, Gemini chat context, review findings fixes

- Extract HistoryManager (history.py) and MetricsManager (metrics.py) from coordinator
- Fix Gemini chat: use client.chats.create(history=...) instead of sequential send_message
- Fix float("inf") in metrics breaking json.dump persistence
- Fix _get_current_state checking wrong error key ("error" vs "error_message")
- Fix API key re-entry requirement when endpoint/provider changes in options flow
- Make CONF_API_KEY optional in options schema (stored key used as fallback)
- Add DNS rebinding protection via socket.getaddrinfo in validate_endpoint
- Remove mass assignment vulnerability in config_flow._create_entry
- Add description_placeholders to all error re-show paths
- Fix history migration overwriting existing JSON data
- Return list copy from get_limited_history to prevent reference leaks
- Add history_info to _get_safe_initial_state for data shape consistency
- Add defensive None check in _get_sanitized_last_response
- Move dt_util import to module level in config_flow
- Remove dead fallback branches in coordinator.last_response property
- Fix sensor min_latency display after float("inf") removal
- Fix history directory permissions from 0o777 to 0o755
- Deduplicate schema definitions via _build_provider_schema()
- Use centralized build_auth_headers from providers.py

custom_components/ha_text_ai/api_client.py

@@ -429,7 +429,6 @@ class APIClient:
             def generate_content():
                 # For single message without history, use generate_content
                 if len(contents) <= 1:
-                    # If we have no content yet, create a simple prompt
                     if not contents:
                         prompt = "I need your assistance."
                     else:
@@ -441,16 +440,30 @@ class APIClient:
                         config=config
                     )
                 else:
-                    # For multi-turn conversations, use chat
+                    # For multi-turn conversations, pass history to chat
-                    chat = client.chats.create(model=model, config=config)
+                    # and only send the last user message
+                    last_user_msg = None
+                    history = []
 
-                    # Send all messages in sequence
+                    # Find the last user message — that's the new query
-                    for content in contents:
+                    for i in range(len(contents) - 1, -1, -1):
-                        if content["role"] == "user":
+                        if contents[i]["role"] == "user":
-                            response = chat.send_message(content["parts"][0]["text"])
+                            last_user_msg = contents[i]["parts"][0]["text"]
-                            # We don't send assistant messages as they're already part of the history
+                            history = contents[:i]
+                            break
 
-                    return response
+                    if last_user_msg is None:
+                        # No user messages at all — shouldn't happen, but handle gracefully
+                        return client.models.generate_content(
+                            model=model,
+                            contents="I need your assistance.",
+                            config=config
+                        )
+
+                    chat = client.chats.create(
+                        model=model, config=config, history=history
+                    )
+                    return chat.send_message(last_user_msg)
 
             response = await asyncio.to_thread(generate_content)
 

custom_components/ha_text_ai/config_flow.py

@@ -8,7 +8,6 @@ Config flow for HA text AI integration.
 """
 import logging
 from typing import Any, Dict, Optional
-from datetime import datetime, timedelta
 
 import voluptuous as vol
 from homeassistant import config_entries
@@ -49,8 +48,10 @@ from .const import (
     DEFAULT_MAX_HISTORY,
     CONF_MAX_HISTORY_SIZE,
 )
-from .utils import normalize_name, safe_log_data, validate_endpoint  # noqa: F401 — re-exported for backward compat
+from homeassistant.util import dt as dt_util
-from .providers import get_default_endpoint, get_default_model
+
+from .utils import normalize_name, safe_log_data, validate_endpoint
+from .providers import get_default_endpoint, get_default_model, build_auth_headers
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -84,52 +85,56 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         self._provider = user_input[CONF_API_PROVIDER]
         return await self.async_step_provider()
 
+    def _build_provider_schema(
+        self, data: Optional[Dict[str, Any]] = None
+    ) -> vol.Schema:
+        """Build provider configuration schema with optional defaults from data."""
+        defaults = data or {}
+        return vol.Schema({
+            vol.Required(CONF_NAME, default=defaults.get(CONF_NAME, "my_assistant")): str,
+            vol.Required(CONF_API_KEY): str,
+            vol.Required(CONF_MODEL, default=defaults.get(CONF_MODEL, get_default_model(self._provider))): str,
+            vol.Required(CONF_API_ENDPOINT, default=defaults.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str,
+            vol.Optional(CONF_TEMPERATURE, default=defaults.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All(
+                vol.Coerce(float),
+                vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
+            ),
+            vol.Optional(CONF_MAX_TOKENS, default=defaults.get(CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS)): vol.All(
+                vol.Coerce(int),
+                vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
+            ),
+            vol.Optional(CONF_REQUEST_INTERVAL, default=defaults.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All(
+                vol.Coerce(float),
+                vol.Range(min=MIN_REQUEST_INTERVAL)
+            ),
+            vol.Optional(CONF_API_TIMEOUT, default=defaults.get(CONF_API_TIMEOUT, DEFAULT_API_TIMEOUT)): vol.All(
+                vol.Coerce(int),
+                vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT)
+            ),
+            vol.Optional(
+                CONF_CONTEXT_MESSAGES,
+                default=defaults.get(CONF_CONTEXT_MESSAGES, DEFAULT_CONTEXT_MESSAGES)
+            ): vol.All(
+                vol.Coerce(int),
+                vol.Range(min=1, max=20)
+            ),
+            vol.Optional(
+                CONF_MAX_HISTORY_SIZE,
+                default=defaults.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY)
+            ): vol.All(
+                vol.Coerce(int),
+                vol.Range(min=1, max=100)
+            ),
+        })
+
     async def async_step_provider(self, user_input: Optional[Dict[str, Any]] = None) -> FlowResult:
         """Handle provider configuration step."""
         self._errors = {}
 
         if user_input is None:
-            default_endpoint = get_default_endpoint(self._provider)
-            default_model = get_default_model(self._provider)
-
             return self.async_show_form(
                 step_id="provider",
-                data_schema=vol.Schema({
+                data_schema=self._build_provider_schema(),
-                    vol.Required(CONF_NAME, default="my_assistant"): str,
-                    vol.Required(CONF_API_KEY): str,
-                    vol.Required(CONF_MODEL, default=default_model): str,
-                    vol.Required(CONF_API_ENDPOINT, default=default_endpoint): str,
-                    vol.Optional(CONF_TEMPERATURE, default=DEFAULT_TEMPERATURE): vol.All(
-                        vol.Coerce(float),
-                        vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
-                    ),
-                    vol.Optional(CONF_MAX_TOKENS, default=DEFAULT_MAX_TOKENS): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
-                    ),
-                    vol.Optional(CONF_REQUEST_INTERVAL, default=DEFAULT_REQUEST_INTERVAL): vol.All(
-                        vol.Coerce(float),
-                        vol.Range(min=MIN_REQUEST_INTERVAL)
-                    ),
-                    vol.Optional(CONF_API_TIMEOUT, default=DEFAULT_API_TIMEOUT): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT)
-                    ),
-                    vol.Optional(
-                        CONF_CONTEXT_MESSAGES,
-                        default=DEFAULT_CONTEXT_MESSAGES
-                    ): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=1, max=20)
-                    ),
-                    vol.Optional(
-                        CONF_MAX_HISTORY_SIZE,
-                        default=DEFAULT_MAX_HISTORY
-                    ): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=1, max=100)
-                    ),
-                })
             )
 
         _LOGGER.debug("Provider step input data: %s", safe_log_data(user_input))
@@ -139,7 +144,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         # Check if CONF_NAME exists in input_copy and ensure it's not empty
         if CONF_NAME not in input_copy or not input_copy[CONF_NAME]:
             _LOGGER.warning("Missing name in configuration input: %s", safe_log_data(input_copy))
-            input_copy[CONF_NAME] = f"assistant_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+            input_copy[CONF_NAME] = f"assistant_{dt_util.utcnow().strftime('%Y%m%d_%H%M%S')}"
             _LOGGER.info("Auto-generated name: %s", input_copy[CONF_NAME])
 
         # Ensure API key is present
@@ -148,168 +153,35 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             _LOGGER.error("API validation error: 'api_key'")
             return self.async_show_form(
                 step_id="provider",
-                data_schema=vol.Schema({
+                data_schema=self._build_provider_schema(input_copy),
-                    vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str,
-                    vol.Required(CONF_API_KEY): str,
-                    vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str,
-                    vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str,
-                    vol.Optional(CONF_TEMPERATURE, default=input_copy.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All(
-                        vol.Coerce(float),
-                        vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
-                    ),
-                    vol.Optional(CONF_MAX_TOKENS, default=input_copy.get(CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS)): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
-                    ),
-                    vol.Optional(CONF_REQUEST_INTERVAL, default=input_copy.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All(
-                        vol.Coerce(float),
-                        vol.Range(min=MIN_REQUEST_INTERVAL)
-                    ),
-                    vol.Optional(CONF_API_TIMEOUT, default=input_copy.get(CONF_API_TIMEOUT, DEFAULT_API_TIMEOUT)): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT)
-                    ),
-                    vol.Optional(
-                        CONF_CONTEXT_MESSAGES,
-                        default=input_copy.get(CONF_CONTEXT_MESSAGES, DEFAULT_CONTEXT_MESSAGES)
-                    ): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=1, max=20)
-                    ),
-                    vol.Optional(
-                        CONF_MAX_HISTORY_SIZE,
-                        default=input_copy.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY)
-                    ): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=1, max=100)
-                    ),
-                }),
                 errors=self._errors
             )
 
         try:
-            # Validate and normalize the name
             normalized_name = self._validate_and_normalize_name(input_copy[CONF_NAME])
             input_copy[CONF_NAME] = normalized_name
         except ValueError as e:
             return self.async_show_form(
                 step_id="provider",
-                data_schema=vol.Schema({
+                data_schema=self._build_provider_schema(input_copy),
-                    vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str,
-                    vol.Required(CONF_API_KEY): str,
-                    vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str,
-                    vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str,
-                    vol.Optional(CONF_TEMPERATURE, default=input_copy.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All(
-                        vol.Coerce(float),
-                        vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
-                    ),
-                    vol.Optional(CONF_MAX_TOKENS, default=input_copy.get(CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS)): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
-                    ),
-                    vol.Optional(CONF_REQUEST_INTERVAL, default=input_copy.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All(
-                        vol.Coerce(float),
-                        vol.Range(min=MIN_REQUEST_INTERVAL)
-                    ),
-                    vol.Optional(CONF_API_TIMEOUT, default=input_copy.get(CONF_API_TIMEOUT, DEFAULT_API_TIMEOUT)): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT)
-                    ),
-                    vol.Optional(
-                        CONF_CONTEXT_MESSAGES,
-                        default=input_copy.get(CONF_CONTEXT_MESSAGES, DEFAULT_CONTEXT_MESSAGES)
-                    ): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=1, max=20)
-                    ),
-                    vol.Optional(
-                        CONF_MAX_HISTORY_SIZE,
-                        default=input_copy.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY)
-                    ): vol.All(
-                        vol.Coerce(int),
-                        vol.Range(min=1, max=100)
-                    ),
-                }),
                 errors={"name": str(e)}
             )
 
         try:
-            # Special handling for Gemini API validation
+            if not await self._async_validate_api(input_copy):
-            if self._provider == API_PROVIDER_GEMINI:
+                return self.async_show_form(
-                # For Gemini, we just check if API key is present as there's no simple endpoint to validate
+                    step_id="provider",
-                if not input_copy.get(CONF_API_KEY):
+                    data_schema=self._build_provider_schema(input_copy),
-                    self._errors["base"] = "invalid_auth"
+                    errors=self._errors
-                    _LOGGER.error("API validation error: 'api_key'")
+                )
-                    return self.async_show_form(
+        except Exception:
-                        step_id="provider",
-                        data_schema=vol.Schema({
-                            vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str,
-                            vol.Required(CONF_API_KEY): str,
-                            vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str,
-                            vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str,
-                            # Other fields remain the same
-                        }),
-                        errors=self._errors
-                    )
-            else:
-                # For other providers, validate API connection
-                if not await self._async_validate_api(input_copy):
-                    return self.async_show_form(
-                        step_id="provider",
-                        data_schema=vol.Schema({
-                            vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str,
-                            vol.Required(CONF_API_KEY): str,
-                            vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str,
-                            vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str,
-                            vol.Optional(CONF_TEMPERATURE, default=input_copy.get(CONF_TEMPERATURE, DEFAULT_TEMPERATURE)): vol.All(
-                                vol.Coerce(float),
-                                vol.Range(min=MIN_TEMPERATURE, max=MAX_TEMPERATURE)
-                            ),
-                            vol.Optional(CONF_MAX_TOKENS, default=input_copy.get(CONF_MAX_TOKENS, DEFAULT_MAX_TOKENS)): vol.All(
-                                vol.Coerce(int),
-                                vol.Range(min=MIN_MAX_TOKENS, max=MAX_MAX_TOKENS)
-                            ),
-                            vol.Optional(CONF_REQUEST_INTERVAL, default=input_copy.get(CONF_REQUEST_INTERVAL, DEFAULT_REQUEST_INTERVAL)): vol.All(
-                                vol.Coerce(float),
-                                vol.Range(min=MIN_REQUEST_INTERVAL)
-                            ),
-                            vol.Optional(CONF_API_TIMEOUT, default=input_copy.get(CONF_API_TIMEOUT, DEFAULT_API_TIMEOUT)): vol.All(
-                                vol.Coerce(int),
-                                vol.Range(min=MIN_API_TIMEOUT, max=MAX_API_TIMEOUT)
-                            ),
-                            vol.Optional(
-                                CONF_CONTEXT_MESSAGES,
-                                default=input_copy.get(CONF_CONTEXT_MESSAGES, DEFAULT_CONTEXT_MESSAGES)
-                            ): vol.All(
-                                vol.Coerce(int),
-                                vol.Range(min=1, max=20)
-                            ),
-                            vol.Optional(
-                                CONF_MAX_HISTORY_SIZE,
-                                default=input_copy.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY)
-                            ): vol.All(
-                                vol.Coerce(int),
-                                vol.Range(min=1, max=100)
-                            ),
-                        }),
-                        errors=self._errors
-                    )
-        except Exception as e:
-            # Handle any unexpected exceptions during validation
             _LOGGER.exception("Unexpected error during API validation")
             return self.async_show_form(
                 step_id="provider",
-                data_schema=vol.Schema({
+                data_schema=self._build_provider_schema(input_copy),
-                    vol.Required(CONF_NAME, default=input_copy.get(CONF_NAME, "my_assistant")): str,
-                    vol.Required(CONF_API_KEY): str,
-                    vol.Required(CONF_MODEL, default=input_copy.get(CONF_MODEL, get_default_model(self._provider))): str,
-                    vol.Required(CONF_API_ENDPOINT, default=input_copy.get(CONF_API_ENDPOINT, get_default_endpoint(self._provider))): str,
-                    # Other fields remain the same
-                }),
                 errors={"base": "unknown"}
             )
 
-        # All validation passed, create the entry
         return await self._create_entry(input_copy)
 
     def _validate_and_normalize_name(self, name: str) -> str:
@@ -345,14 +217,13 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         return normalized
 
     async def _async_validate_api(self, user_input: Dict[str, Any]) -> bool:
-        """Validate API connection."""
+        """Validate API connection using provider registry."""
         try:
             if CONF_API_KEY not in user_input:
                 _LOGGER.error("API validation error: 'api_key'")
                 self._errors["base"] = "invalid_auth"
                 return False
 
-            # Validate endpoint URL (HTTPS-only, SSRF protection)
             try:
                 endpoint = validate_endpoint(user_input[CONF_API_ENDPOINT])
             except ValueError as err:
@@ -360,57 +231,33 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 self._errors["base"] = "cannot_connect"
                 return False
 
-            session = async_get_clientsession(self.hass)
-            headers = self._get_api_headers(user_input)
-
             if self._provider == API_PROVIDER_GEMINI:
                 if not user_input[CONF_API_KEY]:
                     self._errors["base"] = "invalid_auth"
                     return False
                 return True
-            else:
-                check_url = (
-                    f"{endpoint}/v1/models" if self._provider == API_PROVIDER_ANTHROPIC
-                    else f"{endpoint}/models"
-                )
 
-                async with session.get(check_url, headers=headers) as response:
+            session = async_get_clientsession(self.hass)
-                    if response.status == 401:
+            headers = build_auth_headers(self._provider, user_input[CONF_API_KEY])
-                        self._errors["base"] = "invalid_auth"
+
-                        return False
+            from .providers import get_provider_config
-                    elif response.status != 200:
+            check_path = get_provider_config(self._provider).get("check_path", "/models")
-                        self._errors["base"] = "cannot_connect"
+            check_url = f"{endpoint}{check_path}"
-                        return False
+
-                    return True
+            async with session.get(check_url, headers=headers) as response:
+                if response.status == 401:
+                    self._errors["base"] = "invalid_auth"
+                    return False
+                elif response.status != 200:
+                    self._errors["base"] = "cannot_connect"
+                    return False
+                return True
 
         except Exception as err:
             _LOGGER.error("API validation error: %s", str(err))
             self._errors["base"] = "cannot_connect"
             return False
 
-    def _get_api_headers(self, user_input: Dict[str, Any]) -> Dict[str, str]:
-        """Get API headers based on provider."""
-        if CONF_API_KEY not in user_input:
-            return {"Content-Type": "application/json"}
-
-        api_key = user_input[CONF_API_KEY]
-
-        if self._provider == API_PROVIDER_ANTHROPIC:
-            return {
-                "x-api-key": api_key,
-                "anthropic-version": "2023-06-01",
-                "Content-Type": "application/json"
-            }
-        elif self._provider == API_PROVIDER_GEMINI:
-            return {
-                "Authorization": f"Bearer {api_key}",
-                "Content-Type": "application/json"
-            }
-        return {
-            "Authorization": f"Bearer {api_key}",
-            "Content-Type": "application/json"
-        }
-
     async def _create_entry(self, user_input: Dict[str, Any]) -> FlowResult:
         """Create the config entry with comprehensive data preservation."""
         instance_name = user_input[CONF_NAME]
@@ -436,10 +283,6 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             CONF_MAX_HISTORY_SIZE: user_input.get(CONF_MAX_HISTORY_SIZE, DEFAULT_MAX_HISTORY),
         }
 
-        for key, value in user_input.items():
-            if key not in entry_data:
-                entry_data[key] = value
-
         _LOGGER.debug("Creating config entry with data: %s", safe_log_data(entry_data))
 
         return self.async_create_entry(
@@ -462,35 +305,13 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
         self._errors = {}
         self._selected_provider = None
 
-    def _get_default_endpoint(self, provider: str) -> str:
-        """Get default endpoint for provider."""
-        return get_default_endpoint(provider)
-
-    def _get_default_model(self, provider: str) -> str:
-        """Get default model for provider."""
-        return get_default_model(provider)
-
-    def _get_api_headers(self, api_key: str, provider: str) -> Dict[str, str]:
-        """Get API headers based on provider."""
-        if provider == API_PROVIDER_ANTHROPIC:
-            return {
-                "x-api-key": api_key,
-                "anthropic-version": "2023-06-01",
-                "Content-Type": "application/json"
-            }
-        return {
-            "Authorization": f"Bearer {api_key}",
-            "Content-Type": "application/json"
-        }
-
     async def _async_validate_api(self, provider: str, api_key: str, endpoint: str) -> bool:
-        """Validate API connection."""
+        """Validate API connection using provider registry."""
         try:
             if not api_key:
                 self._errors["base"] = "invalid_auth"
                 return False
 
-            # Validate endpoint URL (HTTPS-only, SSRF protection)
             try:
                 endpoint = validate_endpoint(endpoint)
             except ValueError as err:
@@ -498,17 +319,15 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
                 self._errors["base"] = "cannot_connect"
                 return False
 
-            # For Gemini, just check if API key is present
             if provider == API_PROVIDER_GEMINI:
                 return True
 
             session = async_get_clientsession(self.hass)
-            headers = self._get_api_headers(api_key, provider)
+            headers = build_auth_headers(provider, api_key)
 
-            check_url = (
+            from .providers import get_provider_config
-                f"{endpoint}/v1/models" if provider == API_PROVIDER_ANTHROPIC
+            check_path = get_provider_config(provider).get("check_path", "/models")
-                else f"{endpoint}/models"
+            check_url = f"{endpoint}{check_path}"
-            )
 
             async with session.get(check_url, headers=headers) as response:
                 if response.status == 401:
@@ -562,22 +381,45 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
 
         # Use new defaults if provider changed, otherwise use current values
         if provider_changed:
-            default_endpoint = self._get_default_endpoint(provider)
+            default_endpoint = get_default_endpoint(provider)
-            default_model = self._get_default_model(provider)
+            default_model = get_default_model(provider)
         else:
-            default_endpoint = current_data.get(CONF_API_ENDPOINT, self._get_default_endpoint(provider))
+            default_endpoint = current_data.get(CONF_API_ENDPOINT, get_default_endpoint(provider))
-            default_model = current_data.get(CONF_MODEL, self._get_default_model(provider))
+            default_model = current_data.get(CONF_MODEL, get_default_model(provider))
 
         if user_input is not None:
-            # Validate API connection
+            api_key = user_input.get(CONF_API_KEY, "").strip()
-            api_key = user_input.get(CONF_API_KEY, current_data.get(CONF_API_KEY, ""))
             endpoint = user_input.get(CONF_API_ENDPOINT, default_endpoint)
 
+            # Require API key re-entry when endpoint or provider changed
+            stored_endpoint = current_data.get(CONF_API_ENDPOINT, "")
+            endpoint_changed = endpoint != stored_endpoint
+            if not api_key and (provider_changed or endpoint_changed):
+                self._errors["base"] = "invalid_auth"
+                return self.async_show_form(
+                    step_id="settings",
+                    data_schema=self._get_settings_schema(
+                        provider=provider,
+                        current_data=current_data,
+                        user_input=user_input,
+                        default_endpoint=default_endpoint,
+                        default_model=default_model,
+                    ),
+                    errors=self._errors,
+                    description_placeholders={
+                        "provider": provider
+                    }
+                )
+
+            # Fall back to stored key if not re-entered and endpoint unchanged
+            if not api_key:
+                api_key = current_data.get(CONF_API_KEY, "")
+
             if await self._async_validate_api(provider, api_key, endpoint):
-                # Merge with provider selection
                 final_data = {
                     CONF_API_PROVIDER: provider,
-                    **user_input
+                    **user_input,
+                    CONF_API_KEY: api_key,
                 }
                 return self.async_create_entry(title="", data=final_data)
 
@@ -591,7 +433,10 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
                     default_endpoint=default_endpoint,
                     default_model=default_model,
                 ),
-                errors=self._errors
+                errors=self._errors,
+                description_placeholders={
+                    "provider": provider
+                }
             )
 
         return self.async_show_form(
@@ -620,7 +465,7 @@ class OptionsFlowHandler(config_entries.OptionsFlow):
         data = user_input or current_data
 
         return vol.Schema({
-            vol.Required(CONF_API_KEY): str,
+            vol.Optional(CONF_API_KEY, default=""): str,
             vol.Required(
                 CONF_API_ENDPOINT,
                 default=data.get(CONF_API_ENDPOINT, default_endpoint)

custom_components/ha_text_ai/history.py

@@ -0,0 +1,460 @@
+"""
+History management for HA Text AI integration.
+
+@license: CC BY-NC-SA 4.0 International
+@author: SMKRV
+@github: https://github.com/smkrv/ha-text-ai
+@source: https://github.com/smkrv/ha-text-ai
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+import shutil
+import traceback
+from typing import Any, Dict, List, Optional
+
+import aiofiles
+
+from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
+
+from .const import (
+    ABSOLUTE_MAX_HISTORY_SIZE,
+    MAX_ATTRIBUTE_SIZE,
+    MAX_HISTORY_FILE_SIZE,
+    TRUNCATION_INDICATOR,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class AsyncFileHandler:
+    """Async context manager for file operations."""
+
+    def __init__(self, file_path: str, mode: str = "a"):
+        self.file_path = file_path
+        self.mode = mode
+
+    async def __aenter__(self):
+        self.file = await aiofiles.open(self.file_path, self.mode)
+        return self.file
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.file.close()
+
+
+class HistoryManager:
+    """Manages conversation history for an instance."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        instance_name: str,
+        normalized_name: str,
+        history_dir: str,
+        max_history_size: int,
+    ) -> None:
+        self.hass = hass
+        self.instance_name = instance_name
+        self.normalized_name = normalized_name
+        self._history_dir = history_dir
+        self.max_history_size = min(
+            max(1, max_history_size), ABSOLUTE_MAX_HISTORY_SIZE
+        )
+        self._history_file = os.path.join(
+            history_dir, f"{normalized_name}_history.json"
+        )
+        self._max_history_file_size = MAX_HISTORY_FILE_SIZE
+        self._conversation_history: List[Dict[str, Any]] = []
+
+    @property
+    def conversation_history(self) -> List[Dict[str, Any]]:
+        return self._conversation_history
+
+    @property
+    def history_size(self) -> int:
+        return len(self._conversation_history)
+
+    async def async_initialize(self) -> None:
+        """Initialize history: directories, file, migration."""
+        await self._create_history_dir()
+        await self._check_history_directory()
+        await self._initialize_history_file()
+        await self._migrate_history_from_txt_to_json()
+
+    async def _file_exists(self, path: str) -> bool:
+        try:
+            return await self.hass.async_add_executor_job(os.path.exists, path)
+        except Exception as e:
+            _LOGGER.error("Error checking file existence for %s: %s", path, e)
+            return False
+
+    async def _create_history_dir(self) -> None:
+        try:
+            await self.hass.async_add_executor_job(
+                os.makedirs, self._history_dir, 0o755, True
+            )
+        except PermissionError:
+            _LOGGER.error("Permission denied creating history directory: %s", self._history_dir)
+            raise
+        except OSError as e:
+            _LOGGER.error("Error creating history directory %s: %s", self._history_dir, e)
+            raise
+
+    async def _check_history_directory(self) -> None:
+        """Check history directory permissions and writability."""
+        try:
+            await self._create_history_dir()
+            test_file_path = os.path.join(self._history_dir, ".write_test")
+            await self.hass.async_add_executor_job(
+                self._sync_test_directory_write, test_file_path
+            )
+        except PermissionError:
+            _LOGGER.error("No write permissions for history directory: %s", self._history_dir)
+        except Exception as e:
+            _LOGGER.error("Error checking history directory: %s", e)
+
+    @staticmethod
+    def _sync_test_directory_write(test_file_path: str) -> None:
+        try:
+            with open(test_file_path, "w") as f:
+                f.write("Permission test")
+            os.remove(test_file_path)
+        except Exception as e:
+            _LOGGER.error("Directory write test failed: %s", e)
+
+    async def _initialize_history_file(self) -> None:
+        """Initialize history file and load existing history."""
+        try:
+            await self._create_history_dir()
+
+            if await self._file_exists(self._history_file):
+                async with AsyncFileHandler(self._history_file, "r") as f:
+                    content = await f.read()
+                    if content:
+                        history = json.loads(content)
+                        if isinstance(history, list):
+                            self._conversation_history = history[
+                                -self.max_history_size :
+                            ]
+                            _LOGGER.debug(
+                                "Loaded %d history entries for %s",
+                                len(self._conversation_history),
+                                self.instance_name,
+                            )
+            else:
+                async with AsyncFileHandler(self._history_file, "w") as f:
+                    await f.write(json.dumps([]))
+
+            await self._check_history_size()
+        except Exception as e:
+            _LOGGER.error("Could not initialize history file: %s", e)
+            _LOGGER.debug(traceback.format_exc())
+
+    async def update_history(self, question: str, response: dict) -> None:
+        """Update conversation history with size validation."""
+        try:
+            history_entry = {
+                "timestamp": dt_util.utcnow().isoformat(),
+                "question": self._truncate_text(question, MAX_ATTRIBUTE_SIZE),
+                "response": self._truncate_text(
+                    response.get("content", ""), MAX_ATTRIBUTE_SIZE
+                ),
+            }
+
+            entry_size = len(json.dumps(history_entry).encode("utf-8"))
+            current_size = await self._check_file_size(self._history_file)
+
+            if current_size + entry_size > MAX_HISTORY_FILE_SIZE:
+                _LOGGER.warning(
+                    "History size limit approaching. Current: %d, Entry: %d, Max: %d",
+                    current_size, entry_size, MAX_HISTORY_FILE_SIZE,
+                )
+                await self._rotate_history()
+
+            self._conversation_history.append(history_entry)
+
+            while len(self._conversation_history) > self.max_history_size:
+                self._conversation_history.pop(0)
+
+            await self._write_history_entry(history_entry)
+        except Exception as e:
+            _LOGGER.error("Error updating history: %s", e)
+            _LOGGER.debug(traceback.format_exc())
+
+    async def _write_history_entry(self, entry: dict) -> None:
+        """Write history entry with file size checks."""
+        try:
+            if not await self._file_exists(self._history_dir):
+                await self._create_history_dir()
+
+            current_size = 0
+            if await self._file_exists(self._history_file):
+                current_size = await self.hass.async_add_executor_job(
+                    os.path.getsize, self._history_file
+                )
+
+            entry_size = len(json.dumps(entry).encode("utf-8"))
+            if current_size + entry_size > MAX_HISTORY_FILE_SIZE:
+                _LOGGER.warning(
+                    "History file size limit reached. Current: %d, Entry: %d, Max: %d",
+                    current_size, entry_size, MAX_HISTORY_FILE_SIZE,
+                )
+                await self._rotate_history()
+
+            history = []
+            if await self._file_exists(self._history_file):
+                async with AsyncFileHandler(self._history_file, "r") as f:
+                    content = await f.read()
+                    if content:
+                        history = json.loads(content)
+
+            history.append(entry)
+            if len(history) > self.max_history_size:
+                history = history[-self.max_history_size :]
+
+            async with AsyncFileHandler(self._history_file, "w") as f:
+                await f.write(json.dumps(history, indent=2))
+        except Exception as e:
+            _LOGGER.error("Error writing history entry: %s", e)
+            _LOGGER.debug(traceback.format_exc())
+
+    async def _check_history_size(self) -> None:
+        if len(self._conversation_history) > self.max_history_size:
+            _LOGGER.warning(
+                "History size (%d) exceeds maximum (%d). Trimming...",
+                len(self._conversation_history), self.max_history_size,
+            )
+            self._conversation_history = self._conversation_history[
+                -self.max_history_size :
+            ]
+
+    async def _check_file_size(self, file_path: str) -> int:
+        try:
+            if await self._file_exists(file_path):
+                return await self.hass.async_add_executor_job(
+                    os.path.getsize, file_path
+                )
+            return 0
+        except Exception as e:
+            _LOGGER.error("Error checking file size for %s: %s", file_path, e)
+            return 0
+
+    async def _rotate_history(self) -> None:
+        try:
+            _LOGGER.debug("Starting history rotation for %s", self._history_file)
+            await self._rotate_history_files()
+        except Exception as e:
+            _LOGGER.error("Error rotating history: %s", e)
+            _LOGGER.debug(traceback.format_exc())
+
+    async def _rotate_history_files(self) -> None:
+        """Rotate history files with size validation."""
+        try:
+            if await self._file_exists(self._history_file):
+                current_size = await self._check_file_size(self._history_file)
+
+                if current_size > MAX_HISTORY_FILE_SIZE:
+                    _LOGGER.info(
+                        "Rotating history file. Current size: %d, Max: %d",
+                        current_size, MAX_HISTORY_FILE_SIZE,
+                    )
+
+                    archive_file = os.path.join(
+                        self._history_dir,
+                        f"{self.normalized_name}_history_{dt_util.utcnow().strftime('%Y%m%d_%H%M%S')}.json",
+                    )
+
+                    await self.hass.async_add_executor_job(
+                        shutil.move, self._history_file, archive_file
+                    )
+
+                    async with AsyncFileHandler(self._history_file, "w") as f:
+                        await f.write(
+                            json.dumps(
+                                self._conversation_history[
+                                    -self.max_history_size :
+                                ],
+                                indent=2,
+                            )
+                        )
+
+                    _LOGGER.info("History file rotated to: %s", archive_file)
+        except Exception as e:
+            _LOGGER.error("History rotation failed: %s", e)
+            _LOGGER.debug(traceback.format_exc())
+
+    async def _migrate_history_from_txt_to_json(self) -> None:
+        """Migrate old .txt history to .json format."""
+        try:
+            old_history_file = os.path.join(
+                self._history_dir, f"{self.normalized_name}_history.txt"
+            )
+
+            if not await self._file_exists(old_history_file):
+                return
+
+            # Skip migration if JSON history already has entries
+            if self._conversation_history:
+                _LOGGER.debug(
+                    "JSON history already has %d entries for %s, skipping txt migration",
+                    len(self._conversation_history), self.instance_name,
+                )
+                return
+
+            _LOGGER.info(
+                "Found old history file for %s, migrating to JSON", self.instance_name
+            )
+
+            history_entries = []
+            async with AsyncFileHandler(old_history_file, "r") as f:
+                content = await f.read()
+
+            for line in content.split("\n"):
+                if not line or line.startswith("History initialized at:"):
+                    continue
+                try:
+                    parts = line.split(": ", 1)
+                    if len(parts) != 2:
+                        continue
+                    timestamp = parts[0]
+                    content_parts = parts[1].split(" - ")
+                    if len(content_parts) != 2:
+                        continue
+                    question = content_parts[0].replace("Question: ", "")
+                    response = content_parts[1].replace("Response: ", "")
+                    history_entries.append(
+                        {
+                            "timestamp": timestamp,
+                            "question": question,
+                            "response": response,
+                        }
+                    )
+                except Exception as e:
+                    _LOGGER.warning("Error parsing history line: %s. Error: %s", line, e)
+                    continue
+
+            if history_entries:
+                async with AsyncFileHandler(self._history_file, "w") as f:
+                    await f.write(json.dumps(history_entries, indent=2))
+
+                backup_file = old_history_file + ".backup"
+                await self.hass.async_add_executor_job(
+                    shutil.move, old_history_file, backup_file
+                )
+
+                _LOGGER.info(
+                    "Migrated %d entries from txt to JSON for %s. Old file: %s",
+                    len(history_entries), self.instance_name, backup_file,
+                )
+
+                self._conversation_history = history_entries
+        except Exception as e:
+            _LOGGER.error("Error during history migration for %s: %s", self.instance_name, e)
+            _LOGGER.debug(traceback.format_exc())
+
+    async def async_clear_history(self) -> None:
+        """Clear conversation history."""
+        try:
+            self._conversation_history = []
+            if await self._file_exists(self._history_file):
+                await self.hass.async_add_executor_job(os.remove, self._history_file)
+            _LOGGER.info("History for %s cleared", self.instance_name)
+        except Exception as e:
+            _LOGGER.error("Error clearing history: %s", e)
+            _LOGGER.debug(traceback.format_exc())
+
+    async def async_get_history(
+        self,
+        limit: Optional[int] = None,
+        filter_model: Optional[str] = None,
+        start_date: Optional[str] = None,
+        include_metadata: bool = False,
+        sort_order: str = "newest",
+        default_model: str = "",
+    ) -> List[Dict[str, Any]]:
+        """Get conversation history with optional filtering and sorting."""
+        try:
+            history = self._conversation_history.copy()
+
+            if filter_model:
+                history = [
+                    entry for entry in history if entry.get("model") == filter_model
+                ]
+
+            if start_date:
+                try:
+                    from datetime import datetime
+
+                    start_dt = datetime.fromisoformat(
+                        start_date.replace("Z", "+00:00")
+                    )
+                    history = [
+                        entry
+                        for entry in history
+                        if datetime.fromisoformat(
+                            entry["timestamp"].replace("Z", "+00:00")
+                        )
+                        >= start_dt
+                    ]
+                except (ValueError, KeyError) as e:
+                    _LOGGER.warning("Invalid start_date format: %s. Error: %s", start_date, e)
+
+            if sort_order == "oldest":
+                history.sort(key=lambda x: x.get("timestamp", ""))
+            else:
+                history.sort(key=lambda x: x.get("timestamp", ""), reverse=True)
+
+            if limit and limit > 0:
+                history = history[:limit]
+
+            if include_metadata:
+                for entry in history:
+                    entry["metadata"] = {
+                        "entry_size": len(str(entry)),
+                        "question_length": len(entry.get("question", "")),
+                        "response_length": len(entry.get("response", "")),
+                        "model_used": entry.get("model", default_model),
+                        "instance": self.instance_name,
+                    }
+
+            return history
+        except Exception as e:
+            _LOGGER.error("Error getting history: %s", e)
+            return []
+
+    def get_limited_history(self) -> Dict[str, Any]:
+        """Get limited conversation history showing only last Q&A."""
+        limited_history = []
+
+        if self._conversation_history:
+            last_entry = self._conversation_history[-1]
+            limited_entry = {
+                "timestamp": last_entry["timestamp"],
+                "question": self._truncate_text(last_entry["question"], 4096),
+                "response": self._truncate_text(last_entry["response"], 4096),
+            }
+            limited_history.append(limited_entry)
+
+        history_info = {
+            "total_entries": len(self._conversation_history),
+            "displayed_entries": len(limited_history),
+            "full_history_available": True,
+        }
+
+        return {
+            "entries": limited_history,
+            "full_history": list(self._conversation_history),
+            "info": history_info,
+        }
+
+    @staticmethod
+    def _truncate_text(text: str, max_length: int = MAX_ATTRIBUTE_SIZE) -> str:
+        """Safely truncate text to maximum length with indicator."""
+        if not text:
+            return ""
+        if len(text) <= max_length:
+            return text
+        return text[:max_length] + TRUNCATION_INDICATOR

custom_components/ha_text_ai/metrics.py

@@ -0,0 +1,152 @@
+"""
+Metrics management for HA Text AI integration.
+
+@license: CC BY-NC-SA 4.0 International
+@author: SMKRV
+@github: https://github.com/smkrv/ha-text-ai
+@source: https://github.com/smkrv/ha-text-ai
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+import traceback
+from typing import Any, Dict
+
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.util import dt as dt_util
+
+_LOGGER = logging.getLogger(__name__)
+
+DEFAULT_METRICS: Dict[str, Any] = {
+    "total_tokens": 0,
+    "prompt_tokens": 0,
+    "completion_tokens": 0,
+    "successful_requests": 0,
+    "failed_requests": 0,
+    "total_errors": 0,
+    "average_latency": 0,
+    "max_latency": 0,
+    "min_latency": 0,
+}
+
+
+class MetricsManager:
+    """Manages performance metrics for an instance."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        instance_name: str,
+        metrics_file: str,
+    ) -> None:
+        self.hass = hass
+        self.instance_name = instance_name
+        self._metrics_file = metrics_file
+        self._performance_metrics: Dict[str, Any] = DEFAULT_METRICS.copy()
+
+    @property
+    def metrics(self) -> Dict[str, Any]:
+        return self._performance_metrics
+
+    async def async_initialize(self) -> None:
+        """Load metrics from storage or create defaults."""
+        loaded = await self._load_metrics()
+        self._performance_metrics = loaded or DEFAULT_METRICS.copy()
+
+    async def _load_metrics(self) -> Dict[str, Any] | None:
+        try:
+            exists = await self.hass.async_add_executor_job(
+                os.path.exists, self._metrics_file
+            )
+            if exists:
+                def read_metrics():
+                    with open(self._metrics_file, "r") as f:
+                        try:
+                            return json.load(f)
+                        except json.JSONDecodeError:
+                            _LOGGER.warning("Metrics file corrupted, creating new")
+                            return None
+
+                return await self.hass.async_add_executor_job(read_metrics)
+        except Exception as e:
+            _LOGGER.warning("Failed to load metrics: %s", e)
+        return None
+
+    async def _save_metrics(self) -> None:
+        try:
+            def write_metrics():
+                with open(self._metrics_file, "w") as f:
+                    json.dump(self._performance_metrics, f)
+
+            await self.hass.async_add_executor_job(write_metrics)
+        except Exception as e:
+            _LOGGER.warning("Failed to save metrics: %s", e)
+
+    async def update_metrics(self, latency: float, response: dict) -> None:
+        """Update performance metrics after a successful request."""
+        metrics = self._performance_metrics
+        tokens = response.get("tokens", {})
+
+        metrics["total_tokens"] += tokens.get("total", 0)
+        metrics["prompt_tokens"] += tokens.get("prompt", 0)
+        metrics["completion_tokens"] += tokens.get("completion", 0)
+        metrics["successful_requests"] += 1
+
+        metrics["average_latency"] = (
+            (metrics["average_latency"] * (metrics["successful_requests"] - 1) + latency)
+            / metrics["successful_requests"]
+        )
+        metrics["max_latency"] = max(metrics["max_latency"], latency)
+        if metrics["min_latency"] == 0:
+            metrics["min_latency"] = latency
+        else:
+            metrics["min_latency"] = min(metrics["min_latency"], latency)
+
+        await self._save_metrics()
+
+    async def get_current_metrics(self) -> Dict[str, Any]:
+        """Get current performance metrics."""
+        return self._performance_metrics.copy()
+
+    async def handle_error(
+        self,
+        error: Exception,
+        model: str,
+    ) -> Dict[str, Any]:
+        """Record an error in metrics and return error details."""
+        self._performance_metrics["total_errors"] += 1
+        self._performance_metrics["failed_requests"] += 1
+        await self._save_metrics()
+
+        error_details: Dict[str, Any] = {
+            "timestamp": dt_util.utcnow().isoformat(),
+            "model": model,
+            "instance": self.instance_name,
+            "error_message": str(error),
+            "error_type": type(error).__name__,
+            "traceback": traceback.format_exc()
+            if _LOGGER.isEnabledFor(logging.DEBUG)
+            else None,
+        }
+
+        error_mapping = {
+            HomeAssistantError: {"is_ha_error": True},
+            ConnectionError: {"is_connection_error": True},
+            TimeoutError: {"is_timeout": True},
+            PermissionError: {"is_permission_denied": True},
+            ValueError: {"is_validation_error": True},
+        }
+
+        for error_type, error_flags in error_mapping.items():
+            if isinstance(error, error_type):
+                error_details.update(error_flags)
+                break
+
+        _LOGGER.error("AI Processing Error: %s", error_details)
+        if _LOGGER.isEnabledFor(logging.DEBUG):
+            _LOGGER.debug("Full Error Traceback: %s", error_details.get("traceback"))
+
+        return error_details

custom_components/ha_text_ai/sensor.py

@@ -9,8 +9,6 @@ Sensor platform for HA Text AI.
 import logging
 import math
 from typing import Any, Dict
-from datetime import datetime, timedelta
-
 from homeassistant.components.sensor import (
     SensorEntity,
     SensorEntityDescription,
@@ -277,9 +275,7 @@ class HATextAISensor(CoordinatorEntity, SensorEntity):
                     METRIC_FAILED_REQUESTS: metrics.get("failed_requests", 0),
                     METRIC_AVERAGE_LATENCY: round(metrics.get("average_latency", 0), 2),
                     METRIC_MAX_LATENCY: round(metrics.get("max_latency", 0), 2),
-                    METRIC_MIN_LATENCY: (metrics.get("min_latency")
+                    METRIC_MIN_LATENCY: metrics.get("min_latency", 0) or None,
-                                       if metrics.get("min_latency") != float("inf")
-                                       else None),
                 })
 
             # Last response handling

custom_components/ha_text_ai/utils.py

@@ -7,6 +7,7 @@ Utility functions for HA Text AI integration.
 @source: https://github.com/smkrv/ha-text-ai
 """
 import ipaddress
+import socket
 from typing import Any
 from urllib.parse import urlparse
 
@@ -47,7 +48,7 @@ def validate_endpoint(endpoint: str) -> str:
     if not hostname:
         raise ValueError("Invalid endpoint URL: no hostname")
 
-    # Block private/reserved IPs
+    # Block private/reserved IPs (direct IP or resolved hostname)
     try:
         addr = ipaddress.ip_address(hostname)
         if addr.is_private or addr.is_reserved or addr.is_loopback or addr.is_link_local:
@@ -55,6 +56,23 @@ def validate_endpoint(endpoint: str) -> str:
     except ValueError as e:
         if "not allowed" in str(e):
             raise
-        # Not an IP address — it's a hostname, which is fine
+        # Not an IP literal — resolve hostname and check all resolved IPs
+        # to prevent DNS rebinding attacks
+        try:
+            addrinfos = socket.getaddrinfo(hostname, None)
+            for family, _type, _proto, _canonname, sockaddr in addrinfos:
+                ip_str = sockaddr[0]
+                resolved_addr = ipaddress.ip_address(ip_str)
+                if (
+                    resolved_addr.is_private
+                    or resolved_addr.is_reserved
+                    or resolved_addr.is_loopback
+                    or resolved_addr.is_link_local
+                ):
+                    raise ValueError(
+                        f"Hostname {hostname} resolves to private/reserved IP {ip_str}"
+                    )
+        except socket.gaierror:
+            raise ValueError(f"Cannot resolve hostname: {hostname}")
 
     return endpoint.rstrip("/")