Forward service call context to entity in ai_task generate services (#179537)

Co-authored-by: Claude <noreply@anthropic.com>

homeassistant/components/ai_task/__init__.py

@@ -146,13 +146,15 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_service_generate_data(call: ServiceCall) -> ServiceResponse:
     """Run the data task service."""
-    result = await async_generate_data(hass=call.hass, **call.data)
+    result = await async_generate_data(
+        hass=call.hass, context=call.context, **call.data
+    )
     return result.as_dict()
 
 
 async def async_service_generate_image(call: ServiceCall) -> ServiceResponse:
     """Run the image task service."""
-    return await async_generate_image(hass=call.hass, **call.data)
+    return await async_generate_image(hass=call.hass, context=call.context, **call.data)
 
 
 class AITaskPreferences:

homeassistant/components/ai_task/entity.py

@@ -12,6 +12,7 @@ from homeassistant.components.conversation import (
     async_get_chat_log,
 )
 from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN
+from homeassistant.core import Context
 from homeassistant.helpers import llm
 from homeassistant.helpers.chat_session import ChatSession
 from homeassistant.helpers.restore_state import RestoreEntity
@@ -61,6 +62,7 @@ class AITaskEntity(RestoreEntity):
         self,
         session: ChatSession,
         task: GenDataTask | GenImageTask,
+        context: Context | None,
     ) -> AsyncGenerator[ChatLog]:
         """Context manager used to manage the ChatLog used during an AI Task."""
         user_llm_hass_api: llm.API | None = None
@@ -78,7 +80,7 @@ class AITaskEntity(RestoreEntity):
             await chat_log.async_provide_llm_data(
                 llm.LLMContext(
                     platform=self.platform.domain,
-                    context=None,
+                    context=context,
                     language=None,
                     assistant=DOMAIN,
                     device_id=None,
@@ -98,11 +100,14 @@ class AITaskEntity(RestoreEntity):
         self,
         session: ChatSession,
         task: GenDataTask,
+        context: Context | None = None,
     ) -> GenDataTaskResult:
         """Run a gen data task."""
+        if context is not None:
+            self.async_set_context(context)
         self.__last_activity = dt_util.utcnow().isoformat()
         self.async_write_ha_state()
-        async with self._async_get_ai_task_chat_log(session, task) as chat_log:
+        async with self._async_get_ai_task_chat_log(session, task, context) as chat_log:
             return await self._async_generate_data(task, chat_log)
 
     async def _async_generate_data(
@@ -118,11 +123,14 @@ class AITaskEntity(RestoreEntity):
         self,
         session: ChatSession,
         task: GenImageTask,
+        context: Context | None = None,
     ) -> GenImageTaskResult:
         """Run a gen image task."""
+        if context is not None:
+            self.async_set_context(context)
         self.__last_activity = dt_util.utcnow().isoformat()
         self.async_write_ha_state()
-        async with self._async_get_ai_task_chat_log(session, task) as chat_log:
+        async with self._async_get_ai_task_chat_log(session, task, context) as chat_log:
             return await self._async_generate_image(task, chat_log)
 
     async def _async_generate_image(

homeassistant/components/ai_task/task.py

@@ -12,7 +12,7 @@ import voluptuous as vol
 
 from homeassistant.components import camera, conversation, image, media_source
 from homeassistant.components.http.auth import async_sign_path
-from homeassistant.core import HomeAssistant, ServiceResponse, callback
+from homeassistant.core import Context, HomeAssistant, ServiceResponse, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import llm
 from homeassistant.helpers.chat_session import ChatSession, async_get_chat_session
@@ -119,6 +119,7 @@ async def async_generate_data(
     structure: vol.Schema | None = None,
     attachments: list[dict] | None = None,
     llm_api: llm.API | None = None,
+    context: Context | None = None,
 ) -> GenDataTaskResult:
     """Run a data generation task in the AI Task integration."""
     if entity_id is None:
@@ -156,6 +157,7 @@ async def async_generate_data(
                 attachments=resolved_attachments or None,
                 llm_api=llm_api,
             ),
+            context,
         )
 
 
@@ -166,6 +168,7 @@ async def async_generate_image(
     entity_id: str | None = None,
     instructions: str,
     attachments: list[dict] | None = None,
+    context: Context | None = None,
 ) -> ServiceResponse:
     """Run an image generation task in the AI Task integration."""
     if entity_id is None:
@@ -201,6 +204,7 @@ async def async_generate_image(
                 instructions=instructions,
                 attachments=resolved_attachments or None,
             ),
+            context,
         )
 
     service_result = task_result.as_dict()

tests/components/ai_task/conftest.py

@@ -48,12 +48,14 @@ class MockAITaskEntity(AITaskEntity):
         super().__init__()
         self.mock_generate_data_tasks = []
         self.mock_generate_image_tasks = []
+        self.mock_chat_logs = []
 
     async def _async_generate_data(
         self, task: GenDataTask, chat_log: ChatLog
     ) -> GenDataTaskResult:
         """Mock handling of generate data task."""
         self.mock_generate_data_tasks.append(task)
+        self.mock_chat_logs.append(chat_log)
         if task.structure is not None:
             data = {"name": "Tracy Chen", "age": 30}
             data_chat_log = json.dumps(data)

tests/components/ai_task/test_init.py

@@ -15,7 +15,7 @@ from homeassistant.components.ai_task.const import (
     DATA_PREFERENCES,
     DOMAIN,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Context, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import selector
 
@@ -87,6 +87,7 @@ async def test_generate_data_service(
     mock_ai_task_entity: MockAITaskEntity,
 ) -> None:
     """Test the generate data service."""
+    context = Context()
     preferences = hass.data[DATA_PREFERENCES]
     preferences.async_set_preferences(**set_preferences)
 
@@ -108,9 +109,11 @@ async def test_generate_data_service(
             | msg_extra,
             blocking=True,
             return_response=True,
+            context=context,
         )
 
     assert result["data"] == "Mock result"
+    assert hass.states.get(TEST_ENTITY_ID).context is context
 
     assert len(mock_ai_task_entity.mock_generate_data_tasks) == 1
     task = mock_ai_task_entity.mock_generate_data_tasks[0]
@@ -317,6 +320,7 @@ async def test_generate_image_service(
     mock_ai_task_entity: MockAITaskEntity,
 ) -> None:
     """Test the generate image service."""
+    context = Context()
     preferences = hass.data[DATA_PREFERENCES]
     preferences.async_set_preferences(**set_preferences)
 
@@ -335,9 +339,11 @@ async def test_generate_image_service(
             | msg_extra,
             blocking=True,
             return_response=True,
+            context=context,
         )
 
     mock_upload_media.assert_called_once()
+    assert hass.states.get(TEST_ENTITY_ID).context is context
     assert "image_data" not in result
     assert (
         result["media_source_id"]

tests/components/ai_task/test_task.py

@@ -18,7 +18,7 @@ from homeassistant.components.camera import Image
 from homeassistant.components.conversation import async_get_chat_log
 from homeassistant.components.llm import AssistAPI
 from homeassistant.const import STATE_UNKNOWN
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Context, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import chat_session
 from homeassistant.util import dt as dt_util
@@ -78,14 +78,21 @@ async def test_generate_data_preferred_entity(
     assert state is not None
     assert state.state == STATE_UNKNOWN
 
+    context = Context()
     llm_api = AssistAPI(hass)
     result = await async_generate_data(
         hass,
         task_name="Test Task",
         instructions="Test prompt",
         llm_api=llm_api,
+        context=context,
     )
     assert result.data == "Mock result"
+
+    # The LLM API uses the context to check permissions when calling tools
+    chat_log = mock_ai_task_entity.mock_chat_logs[0]
+    assert chat_log.llm_api is not None
+    assert chat_log.llm_api.llm_context.context is context
     as_dict = result.as_dict()
     assert as_dict["conversation_id"] == result.conversation_id
     assert as_dict["data"] == "Mock result"