Merge pull request #9 from knoop7/main

Enhance streaming with first-chunk timeout and auto-retry logic

Author: Desmond-Dong

custom_components/ai_hub/__init__.py

@@ -109,59 +109,41 @@ async def async_setup_entry(hass: HomeAssistant, entry: AIHubConfigEntry) -> boo
     """Set up AI Hub from a config entry."""
 
     # Get API key (may be None if not provided)
+    # No startup validation - each entity validates on actual use
     api_key = get_configured_api_key(entry)
 
-    # Validate API key by testing API connection only if provided
-    if api_key and api_key.strip():
-        try:
-            from .config_flow_validation import validate_input
-
-            await validate_input(hass, {CONF_API_KEY: api_key})
-        except aiohttp.ClientError as err:
-            _LOGGER.error("Failed to connect to API: %s", err)
-            raise ConfigEntryNotReady(f"Failed to connect: {err}") from err
-        except ValueError as err:
-            reason = str(err)
-            if reason == "invalid_auth":
-                raise ConfigEntryAuthFailed("Invalid API key") from err
-            if reason == "cannot_connect":
-                raise ConfigEntryNotReady("API test failed") from err
-            if reason.startswith("cannot_connect:"):
-                detail = reason.split(":", 1)[1].strip()
-                raise ConfigEntryNotReady(f"API test failed: {detail}") from err
-            _LOGGER.error("API validation failed: %s", err)
-            raise ConfigEntryNotReady(f"API validation failed: {err}") from err
-        except ConfigEntryAuthFailed:
-            raise
-        except Exception as err:
-            _LOGGER.error("API validation failed: %s", err)
-            raise ConfigEntryNotReady(f"API validation failed: {err}") from err
-
     # Initialize runtime data in hass.data
     ai_hub_data = get_or_create_ai_hub_data(hass)
     ai_hub_data.api_key = api_key
 
     # Store in entry.runtime_data
     entry.runtime_data = api_key
 
-    # Sync auto-generated intent lists before loading local intent config
-    from .intents.loader import async_sync_intent_lists
-    await async_sync_intent_lists(hass)
-
-    # Set up intent handlers
-    from .intents import async_setup_intents
-    await async_setup_intents(hass)
-
-    # Set up services
-    from .services import async_setup_services
-    await async_setup_services(hass, entry)
-
-    # Forward setup to platforms last. If any initialization above fails and Home
-    # Assistant retries the config entry, delaying platform setup prevents the
-    # entity component from seeing the same entry as already configured.
-    _LOGGER.debug("Setting up AI Hub platforms: %s", PLATFORMS)
-    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
-    _LOGGER.debug("Platforms setup completed")
+    # Each step is independent - one failure does not block others
+    try:
+        from .intents.loader import async_sync_intent_lists
+        await async_sync_intent_lists(hass)
+    except Exception as err:
+        _LOGGER.warning("Intent list sync failed (non-fatal): %s", err)
+
+    try:
+        from .intents import async_setup_intents
+        await async_setup_intents(hass)
+    except Exception as err:
+        _LOGGER.warning("Intent handlers setup failed (non-fatal): %s", err)
+
+    try:
+        from .services import async_setup_services
+        await async_setup_services(hass, entry)
+    except Exception as err:
+        _LOGGER.warning("Services setup failed (non-fatal): %s", err)
+
+    # Forward setup to platforms individually - one failure should not block others
+    for platform in PLATFORMS:
+        try:
+            await hass.config_entries.async_forward_entry_setups(entry, [platform])
+        except Exception as err:
+            _LOGGER.warning("Platform %s setup failed (others continue): %s", platform, err)
 
     # Listen for options updates
     entry.async_on_unload(entry.add_update_listener(async_update_options))
@@ -182,9 +164,15 @@ async def async_unload_entry(hass: HomeAssistant, entry: AIHubConfigEntry) -> bo
         if config_entry.entry_id != entry.entry_id
     ]
 
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if not unload_ok:
-        return False
+    all_ok = True
+    for platform in PLATFORMS:
+        try:
+            if not await hass.config_entries.async_unload_platforms(entry, [platform]):
+                _LOGGER.warning("Failed to unload platform %s", platform)
+                all_ok = False
+        except Exception as err:
+            _LOGGER.warning("Error unloading platform %s: %s", platform, err)
+            all_ok = False
 
     from .services import async_unload_services
     await async_unload_services(hass, entry.entry_id)
@@ -195,7 +183,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: AIHubConfigEntry) -> bo
             ai_hub_data.cleanup()
             hass.data.pop(DOMAIN, None)
 
-    return True
+    return all_ok
 
 
 async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:

custom_components/ai_hub/entity.py

@@ -16,6 +16,7 @@
 from homeassistant.helpers.entity import Entity
 from homeassistant.util import ulid
 from .providers import LLMMessage, create_provider
+from .http import FirstChunkTimeoutError
 from .consts import (
     CONF_CHAT_MODEL,
     CONF_CHAT_URL,
@@ -397,20 +398,33 @@ async def _async_run_provider_stream(
         provider: Any,
         llm_messages: list[LLMMessage],
         tools: list[dict[str, Any]] | None = None,
+        _retry_count: int = 0,
     ) -> None:
         """Run a provider in streaming mode for text-only turns."""
         _LOGGER.debug(
-            "Invoking provider=%s stream=%s tools=%d",
+            "Invoking provider=%s stream=%s tools=%d retry=%d",
             provider_name,
             True,
             len(tools or []),
+            _retry_count,
         )
         has_output = False
-        async for _ in chat_log.async_add_delta_content_stream(
-            self.entity_id,
-            self._transform_provider_stream(provider.complete_stream(llm_messages, tools=tools)),
-        ):
-            has_output = True
+        try:
+            async for _ in chat_log.async_add_delta_content_stream(
+                self.entity_id,
+                self._transform_provider_stream(provider.complete_stream(llm_messages, tools=tools)),
+            ):
+                has_output = True
+        except FirstChunkTimeoutError:
+            if _retry_count < 1:
+                _LOGGER.warning(
+                    "Provider %s first chunk timeout, retrying...",
+                    provider_name,
+                )
+                return await self._async_run_provider_stream(
+                    chat_log, provider_name, provider, llm_messages, tools, _retry_count + 1
+                )
+            raise
 
         if not has_output:
             _LOGGER.warning(

custom_components/ai_hub/http.py

@@ -39,6 +39,10 @@ async def async_post_json(
             return await response.json()
 
 
+class FirstChunkTimeoutError(TimeoutError):
+    """Raised when the first chunk is not received within the timeout."""
+
+
 async def async_stream_response_text(
     url: str,
     *,
@@ -47,15 +51,33 @@ async def async_stream_response_text(
     ssl: bool | None,
     timeout: float,
     error_label: str,
+    first_chunk_timeout: float = 15.0,
 ) -> AsyncGenerator[str, None]:
     """POST JSON and yield decoded response chunks."""
+    import asyncio
+
     async with aiohttp.ClientSession(timeout=client_timeout(timeout)) as session:
         async with session.post(url, json=payload, headers=headers, ssl=ssl) as response:
             if response.status != 200:
                 error_text = await response.text()
                 raise RuntimeError(f"{error_label}: {error_text}")
 
-            async for chunk in response.content:
+            content_iter = response.content.__aiter__()
+            try:
+                first_chunk = await asyncio.wait_for(
+                    content_iter.__anext__(), timeout=first_chunk_timeout
+                )
+            except asyncio.TimeoutError as err:
+                raise FirstChunkTimeoutError(
+                    f"No response from server within {first_chunk_timeout}s"
+                ) from err
+            except StopAsyncIteration:
+                return
+
+            if first_chunk:
+                yield first_chunk.decode("utf-8", errors="ignore")
+
+            async for chunk in content_iter:
                 if chunk:
                     yield chunk.decode("utf-8", errors="ignore")
 

custom_components/ai_hub/manifest.json

@@ -9,5 +9,5 @@
   "iot_class": "cloud_polling",
   "issue_tracker": "https://github.com/ha-china/ai_hub/issues",
   "requirements": ["edge-tts==7.2.7", "aiofiles", "aiohttp"],
-  "version": "v2026.04.12"
+  "version": "v2026.04.13"
 }