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

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

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"