fix: make Agent run_async work with async streaming_callback (#9824)

* fix Agent streaming_callback requires_async

* add tests

* fix

* relnote

Author: anakin87

haystack/components/agents/agent.py

@@ -663,13 +663,13 @@ async def run_async(
 
         if snapshot:
             exe_context = self._initialize_from_snapshot(
-                snapshot=snapshot, streaming_callback=streaming_callback, requires_async=False, tools=tools
+                snapshot=snapshot, streaming_callback=streaming_callback, requires_async=True, tools=tools
             )
         else:
             exe_context = self._initialize_fresh_execution(
                 messages=messages,
                 streaming_callback=streaming_callback,
-                requires_async=False,
+                requires_async=True,
                 system_prompt=system_prompt,
                 tools=tools,
                 **kwargs,

releasenotes/notes/fix-agent-streaming-cb-605a497fb9fd2b24.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    Fix Agent `run_async` method to correctly handle async streaming callbacks.
+    This previously triggered errors due to a bug.

test/components/agents/test_agent.py

@@ -29,7 +29,13 @@
 from haystack.utils import Secret, serialize_callable
 
 
-def streaming_callback_for_serde(chunk: StreamingChunk):
+def sync_streaming_callback(chunk: StreamingChunk) -> None:
+    """A synchronous streaming callback."""
+    pass
+
+
+async def async_streaming_callback(chunk: StreamingChunk) -> None:
+    """An asynchronous streaming callback."""
     pass
 
 
@@ -501,18 +507,16 @@ def test_serde_with_streaming_callback(self, weather_tool, component_tool, monke
         monkeypatch.setenv("FAKE_OPENAI_KEY", "fake-key")
         generator = OpenAIChatGenerator(api_key=Secret.from_env_var("FAKE_OPENAI_KEY"))
         agent = Agent(
-            chat_generator=generator,
-            tools=[weather_tool, component_tool],
-            streaming_callback=streaming_callback_for_serde,
+            chat_generator=generator, tools=[weather_tool, component_tool], streaming_callback=sync_streaming_callback
         )
 
         serialized_agent = agent.to_dict()
 
         init_parameters = serialized_agent["init_parameters"]
-        assert init_parameters["streaming_callback"] == "test_agent.streaming_callback_for_serde"
+        assert init_parameters["streaming_callback"] == "test_agent.sync_streaming_callback"
 
         deserialized_agent = Agent.from_dict(serialized_agent)
-        assert deserialized_agent.streaming_callback is streaming_callback_for_serde
+        assert deserialized_agent.streaming_callback is sync_streaming_callback
 
     def test_exit_conditions_validation(self, weather_tool, component_tool, monkeypatch):
         monkeypatch.setenv("FAKE_OPENAI_KEY", "fake-key")
@@ -928,6 +932,36 @@ def streaming_callback(chunk: StreamingChunk) -> None:
         assert result["last_message"] is not None
         assert streaming_callback_called
 
+    @pytest.mark.asyncio
+    async def test_run_async_with_async_streaming_callback(self, weather_tool):
+        chat_generator = MockChatGenerator()
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=async_streaming_callback)
+        agent.warm_up()
+
+        # This should not raise any exception
+        result = await agent.run_async([ChatMessage.from_user("Hello")])
+
+        assert "messages" in result
+        assert len(result["messages"]) == 2
+        assert result["messages"][1].text == "Hello from run_async"
+
+    def test_run_with_async_streaming_callback_fails(self, weather_tool):
+        chat_generator = MockChatGenerator()
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=async_streaming_callback)
+        agent.warm_up()
+
+        with pytest.raises(ValueError, match="The init callback cannot be a coroutine"):
+            agent.run([ChatMessage.from_user("Hello")])
+
+    @pytest.mark.asyncio
+    async def test_run_async_with_sync_streaming_callback_fails(self, weather_tool):
+        chat_generator = MockChatGenerator()
+        agent = Agent(chat_generator=chat_generator, tools=[weather_tool], streaming_callback=sync_streaming_callback)
+        agent.warm_up()
+
+        with pytest.raises(ValueError, match="The init callback must be async compatible"):
+            await agent.run_async([ChatMessage.from_user("Hello")])
+
 
 class TestAgentTracing:
     def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool):