From 04c0bb20d6b4cb88df4c3c33d990147f09627338 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Wed, 17 Sep 2025 04:30:15 -0400
Subject: [PATCH] AI Task to store generated images in media dir (#152463)

---
 .../components/ai_task/media_source.py        | 14 ++++++++--
 homeassistant/components/backup/const.py      |  1 -
 tests/components/ai_task/test_media_source.py | 27 +++++++++++++++++++
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/ai_task/media_source.py b/homeassistant/components/ai_task/media_source.py
index 2906acf7a2d..61a212be5b0 100644
--- a/homeassistant/components/ai_task/media_source.py
+++ b/homeassistant/components/ai_task/media_source.py
@@ -2,21 +2,31 @@
 
 from __future__ import annotations
 
+from pathlib import Path
+
 from homeassistant.components.media_source import MediaSource, local_source
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 
 from .const import DATA_MEDIA_SOURCE, DOMAIN, IMAGE_DIR
 
 
 async def async_get_media_source(hass: HomeAssistant) -> MediaSource:
     """Set up local media source."""
-    media_dir = hass.config.path(f"{DOMAIN}/{IMAGE_DIR}")
+    media_dirs = list(hass.config.media_dirs.values())
+
+    if not media_dirs:
+        raise HomeAssistantError(
+            "AI Task media source requires at least one media directory configured"
+        )
+
+    media_dir = Path(media_dirs[0]) / DOMAIN / IMAGE_DIR
 
     hass.data[DATA_MEDIA_SOURCE] = source = local_source.LocalSource(
         hass,
         DOMAIN,
         "AI Generated Images",
-        {IMAGE_DIR: media_dir},
+        {IMAGE_DIR: str(media_dir)},
         f"/{DOMAIN}",
     )
     return source
diff --git a/homeassistant/components/backup/const.py b/homeassistant/components/backup/const.py
index 1cfb796bd2e..773deaef174 100644
--- a/homeassistant/components/backup/const.py
+++ b/homeassistant/components/backup/const.py
@@ -26,7 +26,6 @@ EXCLUDE_FROM_BACKUP = [
     "tmp_backups/*.tar",
     "OZW_Log.txt",
     "tts/*",
-    "ai_task/*",
 ]
 
 EXCLUDE_DATABASE_FROM_BACKUP = [
diff --git a/tests/components/ai_task/test_media_source.py b/tests/components/ai_task/test_media_source.py
index 18f1834e082..fd3aa0bdaae 100644
--- a/tests/components/ai_task/test_media_source.py
+++ b/tests/components/ai_task/test_media_source.py
@@ -1,7 +1,11 @@
 """Test ai_task media source."""
 
+import pytest
+
 from homeassistant.components import media_source
+from homeassistant.components.ai_task.media_source import async_get_media_source
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 
 
 async def test_local_media_source(hass: HomeAssistant, init_components: None) -> None:
@@ -9,3 +13,26 @@ async def test_local_media_source(hass: HomeAssistant, init_components: None) ->
     item = await media_source.async_browse_media(hass, "media-source://")
 
     assert any(c.title == "AI Generated Images" for c in item.children)
+
+    source = await async_get_media_source(hass)
+    assert isinstance(source, media_source.local_source.LocalSource)
+    assert source.name == "AI Generated Images"
+    assert source.domain == "ai_task"
+    assert list(source.media_dirs) == ["image"]
+    # Depending on Docker, the default is one of the two paths
+    assert source.media_dirs["image"] in (
+        "/media/ai_task/image",
+        hass.config.path("media/ai_task/image"),
+    )
+    assert source.url_prefix == "/ai_task"
+
+    hass.config.media_dirs = {}
+
+    with pytest.raises(
+        HomeAssistantError,
+        match="AI Task media source requires at least one media directory configured",
+    ):
+        await async_get_media_source(hass)
+
+
+# The following is from media_source/__init__.py for reference