diff --git a/custom_components/ha_text_ai/__init__.py b/custom_components/ha_text_ai/__init__.py index d5e6d3e..5782a57 100644 --- a/custom_components/ha_text_ai/__init__.py +++ b/custom_components/ha_text_ai/__init__.py @@ -9,7 +9,7 @@ The HA Text AI integration. from __future__ import annotations import logging -from typing import Any, Dict +from typing import Any import asyncio @@ -20,12 +20,11 @@ from homeassistant.const import CONF_API_KEY, CONF_NAME from homeassistant.core import HomeAssistant, ServiceCall, SupportsResponse from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError from homeassistant.helpers import config_validation as cv -from homeassistant.helpers import aiohttp_client from homeassistant.util import dt as dt_util from .coordinator import HATextAICoordinator from .api_client import APIClient -from .utils import normalize_name, safe_log_data, validate_endpoint +from .utils import create_pinned_session, normalize_name, safe_log_data, validate_endpoint from .providers import get_default_endpoint, get_default_model, build_auth_headers from .const import ( DOMAIN, @@ -110,7 +109,7 @@ def get_coordinator_by_instance(hass: HomeAssistant, instance: str) -> HATextAIC raise HomeAssistantError(f"Instance {instance} not found") -async def async_setup(hass: HomeAssistant, config: Dict[str, Any]) -> bool: +async def async_setup(hass: HomeAssistant, config: dict[str, Any]) -> bool: """Set up the Home Assistant Text AI component.""" # Initialize domain data storage hass.data.setdefault(DOMAIN, {}) @@ -274,22 +273,26 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: _LOGGER.error("API provider not specified") raise ConfigEntryNotReady("API provider is required") - session = aiohttp_client.async_get_clientsession(hass) - model = config.get(CONF_MODEL, get_default_model(api_provider)) raw_endpoint = config.get(CONF_API_ENDPOINT, get_default_endpoint(api_provider)) allow_local = config.get(CONF_ALLOW_LOCAL_NETWORK, DEFAULT_ALLOW_LOCAL_NETWORK) if allow_local: - _LOGGER.warning( + _LOGGER.info( "Local network mode enabled for endpoint %s — " - "SSRF protection disabled, API credentials may be sent without TLS", + "SSRF protection relaxed for self-hosted proxies", raw_endpoint, ) try: - endpoint = await validate_endpoint(hass, raw_endpoint, allow_local=allow_local) + endpoint, resolved_ips = await validate_endpoint( + hass, raw_endpoint, allow_local=allow_local + ) except ValueError as err: _LOGGER.error("Invalid API endpoint: %s", err) raise ConfigEntryNotReady(f"Invalid API endpoint: {err}") from err + + # Pinned session closes DNS-rebinding TOCTOU and isolates cookies + # from other integrations sharing the same endpoint hostname. + session = create_pinned_session(hass, endpoint, resolved_ips) # API key can now be updated via options api_key = config.get(CONF_API_KEY, entry.data.get(CONF_API_KEY)) instance_name = entry.data.get(CONF_NAME, entry.entry_id) @@ -363,7 +366,6 @@ async def async_update_options(hass: HomeAssistant, entry: ConfigEntry) -> None: _LOGGER.info("Options updated for %s, reloading integration", entry.title) await hass.config_entries.async_reload(entry.entry_id) - async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" try: diff --git a/custom_components/ha_text_ai/api_client.py b/custom_components/ha_text_ai/api_client.py index 9f266a5..dbc5737 100644 --- a/custom_components/ha_text_ai/api_client.py +++ b/custom_components/ha_text_ai/api_client.py @@ -8,9 +8,11 @@ API Client for HA Text AI. """ from __future__ import annotations -import logging import asyncio -from typing import Any, Dict, List, Optional +import json +import logging +import re +from typing import Any from aiohttp import ClientSession, ClientTimeout from homeassistant.exceptions import HomeAssistantError @@ -29,7 +31,6 @@ from .const import ( _LOGGER = logging.getLogger(__name__) - class APIClient: """API Client for OpenAI and Anthropic.""" @@ -37,11 +38,11 @@ class APIClient: self, session: ClientSession, endpoint: str, - headers: Dict[str, str], + headers: dict[str, str], api_provider: str, model: str, api_timeout: int = DEFAULT_API_TIMEOUT, - api_key: Optional[str] = None, + api_key: str | None = None, ) -> None: """Initialize API client.""" self.session = session @@ -89,8 +90,8 @@ class APIClient: async def _make_request( self, url: str, - payload: Dict[str, Any], - ) -> Dict[str, Any]: + payload: dict[str, Any], + ) -> dict[str, Any]: """Make API request with retry logic for transient errors only. Retries on: @@ -174,7 +175,7 @@ class APIClient: raise HomeAssistantError("API request failed after all retries") @staticmethod - def _parse_retry_after(value: Optional[str]) -> Optional[float]: + def _parse_retry_after(value: str | None) -> float | None: """Parse Retry-After header (seconds). Caps at 60s to avoid long stalls.""" if not value: return None @@ -189,13 +190,13 @@ class APIClient: async def create( self, model: str, - messages: List[Dict[str, str]], + messages: list[dict[str, str]], temperature: float, max_tokens: int, structured_output: bool = False, - json_schema: Optional[str] = None, + json_schema: str | None = None, disable_thinking: bool = False, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Create completion using appropriate API.""" try: self._validate_parameters(temperature, max_tokens) @@ -224,43 +225,60 @@ class APIClient: _LOGGER.error("API request failed: %s", str(e)) raise HomeAssistantError(f"API request failed: {str(e)}") from e - @staticmethod - def _is_openai_reasoning_model(model: str) -> bool: - """Detect OpenAI reasoning models (o-series and GPT-5 family). + # Non-reasoning variants whose names otherwise overlap with the + # reasoning prefix set (e.g. "gpt-5-chat-latest" is classic chat). + _OPENAI_NON_REASONING_PATTERNS: tuple[str, ...] = ("gpt-5-chat",) + _OPENAI_REASONING_REGEX = re.compile( + r"^(?:o\d+|gpt-[5-9](?:\.\d+)?)(?:[-_].*)?$" + ) + + @classmethod + def _is_openai_reasoning_model(cls, model: str) -> bool: + """Detect OpenAI reasoning models (o-series and GPT-5+ family). Reasoning models require max_completion_tokens (not max_tokens), do not accept custom temperature, and use "developer" role instead - of "system". Cutoff: models released 2025-09 and later are all - reasoning-by-default (o3, o4-mini, gpt-5, gpt-5-mini, gpt-5-nano). + of "system". Uses a regex so future o5/gpt-6 releases are caught + without code change. Explicitly excludes chat-variants + (e.g. gpt-5-chat-latest) which are classic chat models. """ if not model: return False - m = model.lower().lstrip() - # Match bare model names and dated variants (o3-2025-04-16 etc.) - return ( - m.startswith(("o1", "o3", "o4-mini", "o4")) - or m.startswith(("gpt-5", "gpt5")) - ) + m = model.strip().lower() + # OpenRouter-style "openai/o3" prefix — strip provider namespace. + if "/" in m: + m = m.rsplit("/", 1)[-1] + for non_reasoning in cls._OPENAI_NON_REASONING_PATTERNS: + if m.startswith(non_reasoning): + return False + return bool(cls._OPENAI_REASONING_REGEX.match(m)) @staticmethod def _convert_system_to_developer( - messages: List[Dict[str, str]], - ) -> List[Dict[str, str]]: + messages: list[dict[str, str]], + ) -> list[dict[str, str]]: """Rename role "system" to "developer" for OpenAI reasoning models.""" return [ {**m, "role": "developer"} if m.get("role") == "system" else m for m in messages ] - @staticmethod + # Matches /no_think only as a standalone soft-switch token (word-bounded), + # not when users discuss the concept ("discuss /no_think semantics"). + _NO_THINK_TOKEN_RE = re.compile(r"(?:^|\s)/no_think(?:\s|$)") + + @classmethod def _apply_no_think_tag( - messages: List[Dict[str, str]], - ) -> List[Dict[str, str]]: + cls, + messages: list[dict[str, str]], + ) -> list[dict[str, str]]: """Append Qwen-style /no_think soft switch to the last user message. Why: Qwen3 reasoning models treat "/no_think" in the last user turn as a request to skip thinking. Non-Qwen models ignore the trailing token harmlessly, so this is safe to apply to all OpenAI-compatible backends. + Uses word-boundary regex for dedup so that user content mentioning + "/no_think" mid-sentence isn't mistaken for an existing soft switch. """ if not messages: return messages @@ -268,8 +286,8 @@ class APIClient: for i in range(len(patched) - 1, -1, -1): if patched[i].get("role") == "user": content = patched[i].get("content", "") - if "/no_think" not in content: - patched[i]["content"] = f"{content} /no_think".strip() + if not cls._NO_THINK_TOKEN_RE.search(content): + patched[i]["content"] = f"{content.rstrip()} /no_think".lstrip() break return patched @@ -285,7 +303,6 @@ class APIClient: """ if not text or "" not in text: return text - import re pattern = re.compile(r".*?", flags=re.DOTALL) cleaned = text # Iterative pass: each iteration peels one layer of nested tags. @@ -303,14 +320,13 @@ class APIClient: @staticmethod def _apply_structured_output( - payload: Dict[str, Any], + payload: dict[str, Any], structured_output: bool, - json_schema: Optional[str], + json_schema: str | None, ) -> None: """Apply OpenAI-compatible structured output to payload in-place.""" if not (structured_output and json_schema): return - import json try: schema = json.loads(json_schema) payload["response_format"] = { @@ -328,16 +344,28 @@ class APIClient: async def _create_deepseek_completion( self, model: str, - messages: List[Dict[str, str]], + messages: list[dict[str, str]], temperature: float, max_tokens: int, structured_output: bool = False, - json_schema: Optional[str] = None, + json_schema: str | None = None, disable_thinking: bool = False, - ) -> Dict[str, Any]: - """Create completion using DeepSeek API.""" + ) -> dict[str, Any]: + """Create completion using DeepSeek API. + + DeepSeek-reasoner (R1) is a reasoning model: it ignores /no_think + (thinking is always on by design) and emits reasoning_content as a + separate field alongside content. We skip the no_think append for + this model and preserve reasoning_content in the response payload + so it's available for logging/debug. + """ url = f"{self.endpoint}/chat/completions" - final_messages = self._apply_no_think_tag(messages) if disable_thinking else messages + is_reasoner = "reasoner" in model.lower() + final_messages = ( + self._apply_no_think_tag(messages) + if (disable_thinking and not is_reasoner) + else messages + ) payload = { "model": model, "messages": final_messages, @@ -348,13 +376,18 @@ class APIClient: self._apply_structured_output(payload, structured_output, json_schema) data = await self._make_request(url, payload) - content = data["choices"][0]["message"]["content"] - if disable_thinking: + message = data["choices"][0]["message"] + content = message.get("content", "") + reasoning = message.get("reasoning_content") + if disable_thinking and not is_reasoner: content = self._strip_think_blocks(content) return { "choices": [ { - "message": {"content": content}, + "message": { + "content": content, + **({"reasoning_content": reasoning} if reasoning else {}), + }, } ], "usage": { @@ -367,13 +400,13 @@ class APIClient: async def _create_openai_completion( self, model: str, - messages: List[Dict[str, str]], + messages: list[dict[str, str]], temperature: float, max_tokens: int, structured_output: bool = False, - json_schema: Optional[str] = None, + json_schema: str | None = None, disable_thinking: bool = False, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Create completion using OpenAI API. Reasoning models (o-series, gpt-5 family) require a different payload @@ -388,13 +421,16 @@ class APIClient: if is_reasoning: prepared_messages = self._convert_system_to_developer(messages) - payload: Dict[str, Any] = { + payload: dict[str, Any] = { "model": model, "messages": prepared_messages, "max_completion_tokens": max_tokens, } if disable_thinking: - payload["reasoning_effort"] = "low" + # gpt-5+ supports "minimal" (cheapest, lowest-CoT). o-series + # rejects "minimal" and accepts low/medium/high — fall back to "low". + effort = "minimal" if model.lower().startswith(("gpt-5", "gpt5")) else "low" + payload["reasoning_effort"] = effort else: prepared_messages = ( self._apply_no_think_tag(messages) if disable_thinking else messages @@ -429,13 +465,13 @@ class APIClient: async def _create_anthropic_completion( self, model: str, - messages: List[Dict[str, str]], + messages: list[dict[str, str]], temperature: float, max_tokens: int, structured_output: bool = False, - json_schema: Optional[str] = None, + json_schema: str | None = None, disable_thinking: bool = False, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Create completion using Anthropic API.""" url = f"{self.endpoint}/v1/messages" @@ -455,10 +491,9 @@ class APIClient: # schema strings (built from templates/webhook data) could otherwise # break out of the JSON fence and rewrite the system instruction. if structured_output and json_schema: - import json as _json try: - _json.loads(json_schema) - except _json.JSONDecodeError as err: + json.loads(json_schema) + except json.JSONDecodeError as err: _LOGGER.warning( "Anthropic: invalid JSON schema, ignoring structured_output: %s", err ) @@ -489,10 +524,14 @@ class APIClient: payload["system"] = system_prompt data = await self._make_request(url, payload) - # Anthropic returns text in "content" array; extended thinking arrives - # as a separate thinking-type block (not tags), so no strip - # is required. Extended thinking is opt-in — we simply never request it. - content = data["content"][0]["text"] + # Anthropic returns an array of content blocks; if extended thinking + # is ever enabled the first block may be type="thinking". Find the + # first text-type block instead of hardcoding index [0]. + content = "" + for block in data.get("content", []): + if block.get("type") == "text": + content = block.get("text", "") + break return { "choices": [ { @@ -509,13 +548,13 @@ class APIClient: async def _create_gemini_completion( self, model: str, - messages: List[Dict[str, str]], + messages: list[dict[str, str]], temperature: float, max_tokens: int, structured_output: bool = False, - json_schema: Optional[str] = None, + json_schema: str | None = None, disable_thinking: bool = False, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Create completion using Gemini API with google-genai library. Args: @@ -566,7 +605,6 @@ class APIClient: parsed_schema = None if structured_output and json_schema: try: - import json parsed_schema = json.loads(json_schema) _LOGGER.debug("Gemini structured output enabled with schema") except json.JSONDecodeError as e: @@ -589,10 +627,14 @@ class APIClient: config.response_mime_type = "application/json" config.response_schema = parsed_schema - # Disable thinking for Gemini 2.5+ models (ignored by 2.0 and earlier) + # Disable thinking for Gemini 2.5+ models. Flash variants accept + # thinking_budget=0 (fully off); Pro variants reject 0 and require + # at least 128 tokens of thinking. 2.0 and earlier ignore the field. if disable_thinking: + m_lower = model.lower() + budget = 128 if "2.5-pro" in m_lower else 0 try: - config.thinking_config = types.ThinkingConfig(thinking_budget=0) + config.thinking_config = types.ThinkingConfig(thinking_budget=budget) except (AttributeError, TypeError) as err: _LOGGER.debug( "ThinkingConfig not supported by this google-genai version: %s", err @@ -685,10 +727,12 @@ class APIClient: except ImportError as e: _LOGGER.error("Google Gemini library not installed: %s", e) - raise HomeAssistantError("Missing dependency: google-genai. Please install it.") + raise HomeAssistantError( + "Missing dependency: google-genai. Please install it." + ) from e except Exception as e: _LOGGER.error("Gemini API error: %s", e) - raise HomeAssistantError("Gemini API request failed") + raise HomeAssistantError(f"Gemini API request failed: {e}") from e async def shutdown(self) -> None: """Shutdown API client.""" diff --git a/custom_components/ha_text_ai/config_flow.py b/custom_components/ha_text_ai/config_flow.py index c744c36..7eca6ca 100644 --- a/custom_components/ha_text_ai/config_flow.py +++ b/custom_components/ha_text_ai/config_flow.py @@ -9,14 +9,13 @@ Config flow for HA text AI integration. from __future__ import annotations import logging -from typing import Any, Dict, Optional +from typing import Any import voluptuous as vol from homeassistant import config_entries from homeassistant.const import CONF_API_KEY, CONF_NAME from homeassistant.core import callback from homeassistant.config_entries import ConfigFlowResult -from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers import selector from .const import ( @@ -61,13 +60,17 @@ from .const import ( ) from homeassistant.util import dt as dt_util -from .utils import normalize_name, safe_log_data, validate_endpoint +from .utils import ( + create_pinned_session, + normalize_name, + safe_log_data, + validate_endpoint, +) from .providers import get_default_endpoint, get_default_model, build_auth_headers _LOGGER = logging.getLogger(__name__) - -def _build_parameter_schema(data: Dict[str, Any]) -> dict: +def _build_parameter_schema(data: dict[str, Any]) -> dict: """Build shared parameter schema fields used by both ConfigFlow and OptionsFlow.""" return { vol.Optional( @@ -100,7 +103,6 @@ def _build_parameter_schema(data: Dict[str, Any]) -> dict: ): bool, } - class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): """Handle a config flow for HA text AI.""" @@ -112,7 +114,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self._data = {} self._provider = None - async def async_step_user(self, user_input: Optional[Dict[str, Any]] = None) -> ConfigFlowResult: + async def async_step_user(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult: """Handle the initial step.""" if user_input is None: return self.async_show_form( @@ -131,7 +133,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return await self.async_step_provider() def _build_provider_schema( - self, data: Optional[Dict[str, Any]] = None + self, data: dict[str, Any] | None = None ) -> vol.Schema: """Build provider configuration schema with optional defaults from data.""" defaults = data or {} @@ -150,7 +152,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): schema_dict.update(_build_parameter_schema(defaults)) return vol.Schema(schema_dict) - async def async_step_provider(self, user_input: Optional[Dict[str, Any]] = None) -> ConfigFlowResult: + async def async_step_provider(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult: """Handle provider configuration step.""" self._errors = {} @@ -229,7 +231,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return normalized - async def _async_validate_api(self, user_input: Dict[str, Any]) -> bool: + async def _async_validate_api(self, user_input: dict[str, Any]) -> bool: """Validate API connection using provider registry.""" try: if CONF_API_KEY not in user_input: @@ -239,7 +241,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): try: allow_local = user_input.get(CONF_ALLOW_LOCAL_NETWORK, DEFAULT_ALLOW_LOCAL_NETWORK) - endpoint = await validate_endpoint( + endpoint, resolved_ips = await validate_endpoint( self.hass, user_input[CONF_API_ENDPOINT], allow_local=allow_local ) except ValueError as err: @@ -253,7 +255,9 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return False return True - session = async_get_clientsession(self.hass) + # Pinned session ensures the reachability check goes to the same + # IP that will later be used by api_client (no DNS rebinding). + session = create_pinned_session(self.hass, endpoint, resolved_ips) headers = build_auth_headers(self._provider, user_input[CONF_API_KEY]) from .providers import get_provider_config @@ -274,7 +278,7 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self._errors["base"] = "cannot_connect" return False - async def _create_entry(self, user_input: Dict[str, Any]) -> ConfigFlowResult: + async def _create_entry(self, user_input: dict[str, Any]) -> ConfigFlowResult: """Create the config entry with unique_id deduplication.""" instance_name = user_input[CONF_NAME] normalized_name = normalize_name(instance_name) @@ -314,7 +318,6 @@ class HATextAIConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): """Get the options flow for this handler.""" return OptionsFlowHandler() - class OptionsFlowHandler(config_entries.OptionsFlow): """Handle options flow.""" @@ -326,7 +329,9 @@ class OptionsFlowHandler(config_entries.OptionsFlow): return False try: - endpoint = await validate_endpoint(self.hass, endpoint, allow_local=allow_local) + endpoint, resolved_ips = await validate_endpoint( + self.hass, endpoint, allow_local=allow_local + ) except ValueError as err: _LOGGER.error("Endpoint validation failed: %s", err) self._errors["base"] = "cannot_connect" @@ -335,7 +340,7 @@ class OptionsFlowHandler(config_entries.OptionsFlow): if provider == API_PROVIDER_GEMINI: return True - session = async_get_clientsession(self.hass) + session = create_pinned_session(self.hass, endpoint, resolved_ips) headers = build_auth_headers(provider, api_key) from .providers import get_provider_config @@ -356,11 +361,11 @@ class OptionsFlowHandler(config_entries.OptionsFlow): self._errors["base"] = "cannot_connect" return False - async def async_step_init(self, user_input: Optional[Dict[str, Any]] = None) -> ConfigFlowResult: + async def async_step_init(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult: """Handle provider selection step.""" if not hasattr(self, "_errors"): self._errors: dict[str, str] = {} - self._selected_provider: Optional[str] = None + self._selected_provider: str | None = None current_data = {**self.config_entry.data, **self.config_entry.options} current_provider = current_data.get(CONF_API_PROVIDER, API_PROVIDER_OPENAI) @@ -386,7 +391,7 @@ class OptionsFlowHandler(config_entries.OptionsFlow): } ) - async def async_step_settings(self, user_input: Optional[Dict[str, Any]] = None) -> ConfigFlowResult: + async def async_step_settings(self, user_input: dict[str, Any] | None = None) -> ConfigFlowResult: """Handle settings configuration step.""" self._errors = {} current_data = {**self.config_entry.data, **self.config_entry.options} @@ -477,8 +482,8 @@ class OptionsFlowHandler(config_entries.OptionsFlow): def _get_settings_schema( self, provider: str, - current_data: Dict[str, Any], - user_input: Optional[Dict[str, Any]], + current_data: dict[str, Any], + user_input: dict[str, Any] | None, default_endpoint: str, default_model: str, ) -> vol.Schema: diff --git a/custom_components/ha_text_ai/const.py b/custom_components/ha_text_ai/const.py index 97f08ba..f3ab014 100644 --- a/custom_components/ha_text_ai/const.py +++ b/custom_components/ha_text_ai/const.py @@ -29,7 +29,7 @@ API_PROVIDERS: Final = [ API_PROVIDER_GEMINI ] -VERSION: Final = "2.5.0" +VERSION: Final = "2.5.1" # Default endpoints DEFAULT_OPENAI_ENDPOINT: Final = "https://api.openai.com/v1" diff --git a/custom_components/ha_text_ai/coordinator.py b/custom_components/ha_text_ai/coordinator.py index 104d08a..a6b5132 100644 --- a/custom_components/ha_text_ai/coordinator.py +++ b/custom_components/ha_text_ai/coordinator.py @@ -12,7 +12,7 @@ import asyncio import logging import os from datetime import timedelta -from typing import Any, Dict, List, Optional +from typing import Any from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant @@ -40,7 +40,6 @@ from .utils import normalize_name _LOGGER = logging.getLogger(__name__) - class HATextAICoordinator(DataUpdateCoordinator): """Home Assistant Text AI Conversation Coordinator.""" @@ -101,9 +100,9 @@ class HATextAICoordinator(DataUpdateCoordinator): self._is_rate_limited = False self._is_maintenance = False self.endpoint_status = "ready" - self._system_prompt: Optional[str] = None + self._system_prompt: str | None = None - self._last_response: Dict[str, Any] = { + self._last_response: dict[str, Any] = { "timestamp": dt_util.utcnow().isoformat(), "question": "", "response": "", @@ -132,7 +131,7 @@ class HATextAICoordinator(DataUpdateCoordinator): # Convenience accessors for backward compatibility # ------------------------------------------------------------------ @property - def _conversation_history(self) -> List[Dict[str, Any]]: + def _conversation_history(self) -> list[dict[str, Any]]: return self._history.conversation_history @property @@ -155,12 +154,12 @@ class HATextAICoordinator(DataUpdateCoordinator): # Last response # ------------------------------------------------------------------ @property - def last_response(self) -> Dict[str, Any]: + def last_response(self) -> dict[str, Any]: """Get the last response.""" return self._last_response @last_response.setter - def last_response(self, value: Dict[str, Any]) -> None: + def last_response(self, value: dict[str, Any]) -> None: self._last_response = value # ------------------------------------------------------------------ @@ -173,7 +172,7 @@ class HATextAICoordinator(DataUpdateCoordinator): except Exception as err: _LOGGER.error("Error updating HA state for %s: %s", self.instance_name, err) - async def _async_update_data(self) -> Dict[str, Any]: + async def _async_update_data(self) -> dict[str, Any]: """Update coordinator data.""" try: current_state = self._get_current_state() @@ -209,14 +208,14 @@ class HATextAICoordinator(DataUpdateCoordinator): async def async_ask_question( self, question: str, - model: Optional[str] = None, - temperature: Optional[float] = None, - max_tokens: Optional[int] = None, - system_prompt: Optional[str] = None, - context_messages: Optional[int] = None, + model: str | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + system_prompt: str | None = None, + context_messages: int | None = None, structured_output: bool = False, - json_schema: Optional[str] = None, - disable_thinking: Optional[bool] = None, + json_schema: str | None = None, + disable_thinking: bool | None = None, ) -> dict: """Process question with context management.""" if self.client is None: @@ -279,11 +278,11 @@ class HATextAICoordinator(DataUpdateCoordinator): self, question: str, model: str, - messages: List[Dict[str, str]], + messages: list[dict[str, str]], temperature: float, max_tokens: int, structured_output: bool = False, - json_schema: Optional[str] = None, + json_schema: str | None = None, disable_thinking: bool = False, ) -> dict: """Send request to AI provider and return structured response. @@ -348,12 +347,12 @@ class HATextAICoordinator(DataUpdateCoordinator): async def async_get_history( self, - limit: Optional[int] = None, - filter_model: Optional[str] = None, - start_date: Optional[str] = None, + limit: int | None = None, + filter_model: str | None = None, + start_date: str | None = None, include_metadata: bool = False, sort_order: str = "newest", - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Get conversation history with optional filtering.""" return await self._history.async_get_history( limit=limit, @@ -383,7 +382,7 @@ class HATextAICoordinator(DataUpdateCoordinator): return STATE_ERROR return STATE_READY - def _get_safe_initial_state(self) -> Dict[str, Any]: + def _get_safe_initial_state(self) -> dict[str, Any]: return { "state": STATE_ERROR, "metrics": {}, @@ -403,7 +402,7 @@ class HATextAICoordinator(DataUpdateCoordinator): "normalized_name": self.normalized_name, } - def _get_sanitized_last_response(self) -> Dict[str, Any]: + def _get_sanitized_last_response(self) -> dict[str, Any]: """Get sanitized version of last response with truncation.""" response = self.last_response.copy() @@ -422,7 +421,7 @@ class HATextAICoordinator(DataUpdateCoordinator): def _calculate_uptime(self) -> float: return (dt_util.utcnow() - self._start_time).total_seconds() - def _get_truncated_system_prompt(self) -> Optional[str]: + def _get_truncated_system_prompt(self) -> str | None: if not self._system_prompt: return None if len(self._system_prompt) <= 4096: @@ -430,7 +429,7 @@ class HATextAICoordinator(DataUpdateCoordinator): return self._system_prompt[:4096] + TRUNCATION_INDICATOR @staticmethod - def _validate_update_data(data: Dict[str, Any]) -> None: + def _validate_update_data(data: dict[str, Any]) -> None: for key in ("state", "metrics", "last_response"): if key not in data: raise ValueError(f"Missing required key: {key}") diff --git a/custom_components/ha_text_ai/history.py b/custom_components/ha_text_ai/history.py index cd440d7..303f5a1 100644 --- a/custom_components/ha_text_ai/history.py +++ b/custom_components/ha_text_ai/history.py @@ -14,7 +14,7 @@ import os import shutil import traceback from datetime import datetime -from typing import Any, Dict, List, Optional +from typing import Any import aiofiles @@ -34,7 +34,6 @@ MAX_ARCHIVE_FILES = 3 _LOGGER = logging.getLogger(__name__) - class AsyncFileHandler: """Async context manager for file operations.""" @@ -49,7 +48,6 @@ class AsyncFileHandler: async def __aexit__(self, exc_type, exc_val, exc_tb): await self.file.close() - def _assert_not_symlink(path: str) -> None: """Refuse to operate on a path that resolves to a symlink. @@ -61,7 +59,6 @@ def _assert_not_symlink(path: str) -> None: if os.path.islink(path): raise OSError(f"Refusing to operate on symlink: {path}") - class HistoryManager: """Manages conversation history for an instance.""" @@ -84,10 +81,10 @@ class HistoryManager: history_dir, f"{normalized_name}_history.json" ) self._max_history_file_size = MAX_HISTORY_FILE_SIZE - self._conversation_history: List[Dict[str, Any]] = [] + self._conversation_history: list[dict[str, Any]] = [] @property - def conversation_history(self) -> List[Dict[str, Any]]: + def conversation_history(self) -> list[dict[str, Any]]: return self._conversation_history @property @@ -388,13 +385,13 @@ class HistoryManager: async def async_get_history( self, - limit: Optional[int] = None, - filter_model: Optional[str] = None, - start_date: Optional[str] = None, + limit: int | None = None, + filter_model: str | None = None, + start_date: str | None = None, include_metadata: bool = False, sort_order: str = "newest", default_model: str = "", - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Get conversation history with optional filtering and sorting.""" try: history = self._conversation_history.copy() @@ -425,8 +422,11 @@ class HistoryManager: else: history.sort(key=lambda x: x.get("timestamp", ""), reverse=True) + # Clamp limit to ABSOLUTE_MAX_HISTORY_SIZE to prevent pathological + # caller requests from producing multi-MB service payloads. if limit and limit > 0: - history = history[:limit] + effective_limit = min(int(limit), ABSOLUTE_MAX_HISTORY_SIZE) + history = history[:effective_limit] if include_metadata: enriched = [] @@ -447,7 +447,7 @@ class HistoryManager: _LOGGER.error("Error getting history: %s", e) return [] - def get_limited_history(self, max_display: int = 5) -> Dict[str, Any]: + def get_limited_history(self, max_display: int = 5) -> dict[str, Any]: """Get limited conversation history for sensor attributes. Returns last `max_display` entries with truncated text for HA state. diff --git a/custom_components/ha_text_ai/manifest.json b/custom_components/ha_text_ai/manifest.json index 92e6698..e6ec8bf 100644 --- a/custom_components/ha_text_ai/manifest.json +++ b/custom_components/ha_text_ai/manifest.json @@ -15,5 +15,5 @@ "google-genai>=1.16.0,<2.0.0" ], "single_config_entry": false, - "version": "2.5.0" + "version": "2.5.1" } diff --git a/custom_components/ha_text_ai/metrics.py b/custom_components/ha_text_ai/metrics.py index 6369ef8..bad17f6 100644 --- a/custom_components/ha_text_ai/metrics.py +++ b/custom_components/ha_text_ai/metrics.py @@ -13,7 +13,7 @@ import logging import os import re import traceback -from typing import Any, Dict +from typing import Any from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError @@ -21,7 +21,7 @@ from homeassistant.util import dt as dt_util _LOGGER = logging.getLogger(__name__) -DEFAULT_METRICS: Dict[str, Any] = { +DEFAULT_METRICS: dict[str, Any] = { "total_tokens": 0, "prompt_tokens": 0, "completion_tokens": 0, @@ -33,7 +33,6 @@ DEFAULT_METRICS: Dict[str, Any] = { "min_latency": 0, } - class MetricsManager: """Manages performance metrics for an instance.""" @@ -46,10 +45,10 @@ class MetricsManager: self.hass = hass self.instance_name = instance_name self._metrics_file = metrics_file - self._performance_metrics: Dict[str, Any] = DEFAULT_METRICS.copy() + self._performance_metrics: dict[str, Any] = DEFAULT_METRICS.copy() @property - def metrics(self) -> Dict[str, Any]: + def metrics(self) -> dict[str, Any]: return self._performance_metrics async def async_initialize(self) -> None: @@ -57,7 +56,7 @@ class MetricsManager: loaded = await self._load_metrics() self._performance_metrics = loaded or DEFAULT_METRICS.copy() - async def _load_metrics(self) -> Dict[str, Any] | None: + async def _load_metrics(self) -> dict[str, Any] | None: try: exists = await self.hass.async_add_executor_job( os.path.exists, self._metrics_file @@ -108,7 +107,7 @@ class MetricsManager: await self._save_metrics() - async def get_current_metrics(self) -> Dict[str, Any]: + async def get_current_metrics(self) -> dict[str, Any]: """Get current performance metrics.""" return self._performance_metrics.copy() @@ -116,7 +115,7 @@ class MetricsManager: self, error: Exception, model: str, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Record an error in metrics and return error details.""" self._performance_metrics["total_errors"] += 1 self._performance_metrics["failed_requests"] += 1 @@ -146,7 +145,7 @@ class MetricsManager: if len(error_msg) > 256: error_msg = error_msg[:256] + "..." - error_details: Dict[str, Any] = { + error_details: dict[str, Any] = { "timestamp": dt_util.utcnow().isoformat(), "model": model, "instance": self.instance_name, diff --git a/custom_components/ha_text_ai/providers.py b/custom_components/ha_text_ai/providers.py index 1e5f959..8ad6030 100644 --- a/custom_components/ha_text_ai/providers.py +++ b/custom_components/ha_text_ai/providers.py @@ -62,7 +62,6 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { }, } - def get_provider_config(provider: str) -> dict[str, Any]: """Get full provider configuration. @@ -73,17 +72,14 @@ def get_provider_config(provider: str) -> dict[str, Any]: raise ValueError(f"Unknown API provider: {provider}") return PROVIDER_REGISTRY[provider] - def get_default_endpoint(provider: str) -> str: """Get default API endpoint for a provider.""" return get_provider_config(provider)["default_endpoint"] - def get_default_model(provider: str) -> str: """Get default model for a provider.""" return get_provider_config(provider)["default_model"] - def build_auth_headers(provider: str, api_key: str) -> dict[str, str]: """Build authentication headers for a provider.""" config = get_provider_config(provider) diff --git a/custom_components/ha_text_ai/sensor.py b/custom_components/ha_text_ai/sensor.py index 20c50c6..45c357b 100644 --- a/custom_components/ha_text_ai/sensor.py +++ b/custom_components/ha_text_ai/sensor.py @@ -10,7 +10,7 @@ from __future__ import annotations import logging import math -from typing import Any, Dict +from typing import Any from homeassistant.components.sensor import ( SensorEntity, SensorEntityDescription, @@ -79,7 +79,6 @@ _LOGGER = logging.getLogger(__name__) _ATTR_TEXT_LIMIT = 2048 _ATTR_PROMPT_LIMIT = 512 - async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, @@ -185,7 +184,7 @@ class HATextAISensor(CoordinatorEntity, SensorEntity): return None return value - def _sanitize_attributes(self, attributes: Dict[str, Any]) -> Dict[str, Any]: + def _sanitize_attributes(self, attributes: dict[str, Any]) -> dict[str, Any]: """Sanitize all attributes for JSON serialization.""" sanitized = { key: self._sanitize_value(value) @@ -231,7 +230,7 @@ class HATextAISensor(CoordinatorEntity, SensorEntity): return ENTITY_ICON @property - def extra_state_attributes(self) -> Dict[str, Any]: + def extra_state_attributes(self) -> dict[str, Any]: """Return entity specific state attributes.""" if not self.coordinator.data: return {} diff --git a/custom_components/ha_text_ai/utils.py b/custom_components/ha_text_ai/utils.py index 1e6d106..1965aec 100644 --- a/custom_components/ha_text_ai/utils.py +++ b/custom_components/ha_text_ai/utils.py @@ -8,14 +8,22 @@ Utility functions for HA Text AI integration. """ from __future__ import annotations +import hashlib import ipaddress +import logging import socket from typing import Any from urllib.parse import urlparse +import aiohttp +from aiohttp.abc import AbstractResolver +from aiohttp.resolver import DefaultResolver + from homeassistant.const import CONF_API_KEY from homeassistant.core import HomeAssistant +from homeassistant.helpers.aiohttp_client import async_create_clientsession +_LOGGER = logging.getLogger(__name__) def normalize_name(name: str) -> str: """Normalize name to conform to HA naming convention using underscores. @@ -27,12 +35,10 @@ def normalize_name(name: str) -> str: normalized = ''.join(c if c.isalnum() or c == '_' else '_' for c in name) normalized = '_'.join(filter(None, normalized.split('_'))).lower() if not normalized: - import hashlib digest = hashlib.sha256(name.encode("utf-8", errors="replace")).hexdigest()[:8] normalized = f"instance_{digest}" return normalized - def safe_log_data( data: dict[str, Any], sensitive_keys: tuple[str, ...] = (CONF_API_KEY,), @@ -40,11 +46,9 @@ def safe_log_data( """Filter sensitive keys from data for safe logging.""" return {k: "***" if k in sensitive_keys else v for k, v in data.items()} - class _RestrictedIPError(ValueError): """Raised when an IP address is in a restricted range.""" - def _check_ip_restricted(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: """Check if an IP address is in a restricted range.""" return ( @@ -56,14 +60,136 @@ def _check_ip_restricted(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> or addr.is_unspecified ) +def _is_cloud_metadata_or_unsafe( + addr: ipaddress.IPv4Address | ipaddress.IPv6Address, +) -> bool: + """Block link-local and cloud instance-metadata addresses. -async def validate_endpoint(hass: HomeAssistant, endpoint: str, *, allow_local: bool = False) -> str: - """Validate API endpoint URL for security. + These must be blocked even in allow_local_network mode: + - IPv4 link-local (169.254.0.0/16) covers AWS/GCP/Azure IMDS 169.254.169.254. + - IPv6 link-local (fe80::/10). + - Multicast and unspecified addresses. + Without this check, a self-hosted HA running on a cloud VM could be + tricked into exfiltrating cloud credentials via IMDS. + """ + return addr.is_multicast or addr.is_unspecified or addr.is_link_local + +class _PinnedResolver(AbstractResolver): + """aiohttp resolver that returns pre-validated IPs for a single hostname. + + Why: prevents DNS-rebinding attacks. After validate_endpoint has + confirmed the hostname resolves to a safe IP, we pin that IP in the + resolver used by the aiohttp session. aiohttp then skips its own + DNS lookup on each request and uses the pinned IP, closing the + TOCTOU gap between validation and actual HTTP call. + """ + + def __init__( + self, + pinned: dict[str, list[tuple[str, int]]], + fallback: AbstractResolver | None = None, + ) -> None: + self._pinned = pinned + self._fallback = fallback or DefaultResolver() + + async def resolve( + self, + host: str, + port: int = 0, + family: int = socket.AF_INET, + ) -> list[dict[str, Any]]: + entries = self._pinned.get(host.lower()) + if entries is None: + return await self._fallback.resolve(host, port, family) + return [ + { + "hostname": host, + "host": ip, + "port": port or default_port, + "family": _family_for(ip), + "proto": 0, + "flags": 0, + } + for ip, default_port in entries + ] + + async def close(self) -> None: + await self._fallback.close() + +def _family_for(ip: str) -> int: + """Return AF_INET or AF_INET6 based on the IP literal.""" + try: + return socket.AF_INET6 if ":" in ip else socket.AF_INET + except Exception: + return socket.AF_INET + +async def resolve_hostname_ips( + hass: HomeAssistant, + hostname: str, +) -> list[str]: + """Resolve hostname to all its IPs via the event-loop-safe executor. + + Returns a list of IP strings (may contain both IPv4 and IPv6). + Raises ValueError on resolution failure. + """ + try: + addrinfos = await hass.async_add_executor_job( + socket.getaddrinfo, hostname, None + ) + except socket.gaierror as err: + raise ValueError(f"Cannot resolve hostname: {hostname}") from err + ips = [] + seen: set[str] = set() + for _family, _type, _proto, _canonname, sockaddr in addrinfos: + ip = sockaddr[0] + if ip not in seen: + seen.add(ip) + ips.append(ip) + if not ips: + raise ValueError(f"No IPs for hostname: {hostname}") + return ips + +def create_pinned_session( + hass: HomeAssistant, + endpoint: str, + resolved_ips: list[str], +) -> aiohttp.ClientSession: + """Create an isolated aiohttp session with pinned DNS and no cookie jar. + + Addresses two issues at once: + - DNS rebinding: aiohttp will reuse the pinned IPs from validate_endpoint + rather than re-resolving the hostname on each request. + - Cookie pollution: DummyCookieJar prevents cookies from leaking between + this integration and other HA components sharing the same domain. + """ + parsed = urlparse(endpoint) + hostname = (parsed.hostname or "").lower() + port = parsed.port or (443 if parsed.scheme == "https" else 80) + pinned: dict[str, list[tuple[str, int]]] = { + hostname: [(ip, port) for ip in resolved_ips] + } + connector = aiohttp.TCPConnector(resolver=_PinnedResolver(pinned)) + return async_create_clientsession( + hass, + cookie_jar=aiohttp.DummyCookieJar(), + connector=connector, + ) + +async def validate_endpoint( + hass: HomeAssistant, + endpoint: str, + *, + allow_local: bool = False, +) -> tuple[str, list[str]]: + """Validate API endpoint URL for security and pin resolved IPs. Ensures HTTPS-only and blocks private/reserved IP ranges (SSRF protection). When allow_local is True, permits private IPs and HTTP scheme for self-hosted proxies. Uses async DNS resolution to avoid blocking the event loop. - Returns the validated endpoint stripped of trailing slashes. + + Returns: (validated_endpoint_without_trailing_slash, resolved_ips). + The resolved IPs are intended for pinning in aiohttp resolver via + create_pinned_session(), closing the DNS-rebinding TOCTOU gap. Raises: ValueError: If the endpoint fails validation. @@ -81,54 +207,47 @@ async def validate_endpoint(hass: HomeAssistant, endpoint: str, *, allow_local: if not hostname: raise ValueError("Invalid endpoint URL: no hostname") - if allow_local: - # Even in local mode, block multicast and unspecified addresses + resolved_ips: list[str] = [] + + # Collect and check all resolved IPs (or IP literal directly). + def _collect(ips: list[str]) -> None: + for ip in ips: + if ip not in resolved_ips: + resolved_ips.append(ip) + + try: + addr = ipaddress.ip_address(hostname) _is_ip_literal = True - try: - addr = ipaddress.ip_address(hostname) - except ValueError: - _is_ip_literal = False + except ValueError: + addr = None + _is_ip_literal = False - if _is_ip_literal: - if addr.is_multicast or addr.is_unspecified: - raise ValueError("Multicast and unspecified addresses are not allowed") - else: - # Not an IP literal — resolve and check - try: - addrinfos = await hass.async_add_executor_job( - socket.getaddrinfo, hostname, None - ) - for family, _type, _proto, _canonname, sockaddr in addrinfos: - resolved = ipaddress.ip_address(sockaddr[0]) - if resolved.is_multicast or resolved.is_unspecified: - raise ValueError( - "Hostname resolves to multicast/unspecified address" - ) - except socket.gaierror as err: - raise ValueError(f"Cannot resolve hostname: {hostname}") from err + if _is_ip_literal: + _collect([hostname]) else: - # Full SSRF protection — block private/reserved IPs try: - addr = ipaddress.ip_address(hostname) - if _check_ip_restricted(addr): - raise _RestrictedIPError("Private/reserved IP addresses are not allowed") - except _RestrictedIPError: - raise - except ValueError: - # Not an IP literal — resolve hostname and check all resolved IPs - # to prevent DNS rebinding attacks - try: - addrinfos = await hass.async_add_executor_job( - socket.getaddrinfo, hostname, None - ) - for family, _type, _proto, _canonname, sockaddr in addrinfos: - ip_str = sockaddr[0] - resolved_addr = ipaddress.ip_address(ip_str) - if _check_ip_restricted(resolved_addr): - raise ValueError( - "Hostname resolves to a restricted IP range" - ) - except socket.gaierror as err: - raise ValueError(f"Cannot resolve hostname: {hostname}") from err + addrinfos = await hass.async_add_executor_job( + socket.getaddrinfo, hostname, None + ) + except socket.gaierror as err: + raise ValueError(f"Cannot resolve hostname: {hostname}") from err + _collect([sockaddr[0] for (*_, sockaddr) in addrinfos]) - return endpoint.rstrip("/") + if not resolved_ips: + raise ValueError(f"No IPs resolved for hostname: {hostname}") + + # Validate each resolved IP against the selected policy. + for ip_str in resolved_ips: + ip_obj = ipaddress.ip_address(ip_str) + if allow_local: + if _is_cloud_metadata_or_unsafe(ip_obj): + raise ValueError( + "Link-local/metadata/multicast addresses are not allowed" + ) + else: + if _check_ip_restricted(ip_obj): + raise _RestrictedIPError( + "Private/reserved IP addresses are not allowed" + ) + + return endpoint.rstrip("/"), resolved_ips