HScarb · bigxin0124 · Aug 8, 2025 · coderabbitai · Aug 8, 2025 · coderabbitai
-    async def on_message_send_streaming(
-        self, 
-        params: MessageSendParams, 
-        context=None
-    ) -> AsyncGenerator[Message | Task, None]:
-        """处理流式消息发送请求。"""
+    async def on_message_send_stream(
+        self, 
+        params: MessageSendParams, 
+        context=None
+    ) -> AsyncGenerator[Message | Task, None]:
+        """处理流式消息发送请求。"""
-    async def get(self) -> Any:
-        """Get the next item from the stream."""
-        if self._exception:
-            raise self._exception
-        return await self.queue.get()
-    
-    def set_exception(self, exception: Exception) -> None:
-        """Set an exception for the stream."""
-        self._exception = exception
-        self._done = True
-    
-    def set_done(self) -> None:
-        """Mark the stream as complete."""
-        self._done = True
-    
+    async def get(self) -> Any:
+        """Get the next item from the stream."""
+        if self._exception:
+            raise self._exception
+
+        if self._done and self.queue.empty():
+            # Signal end-of-stream
+            raise StopAsyncIteration
+
+        return await self.queue.get()
+
+    def set_exception(self, exception: Exception) -> None:
+        """Set an exception for the stream."""
+        self._exception = exception
+        self._done = True
+
+    def set_done(self) -> None:
+        """Mark the stream as complete."""
+        self._done = True
+        # Unblock any waiters when no more data will arrive.
+        if self.queue.empty():
+            self.queue.put_nowait(None)  # sentinel
-        params = MessageSendParams(
-            content="测试流式消息",
-            role="user"
-        )
-        events = []
-        async for event in handler.on_message_send_stream(params):
-            events.append(event)
-            print(f"✓ 收到流式事件: {event.content}")
-        print(f"✓ on_message_send_stream 正常工作，收到 {len(events)} 个事件")
+        params = MessageSendParams(
+            content="测试流式消息",
+            role="user"
+        )
+        events = []
+        async for event in handler.on_message_send_stream(params):
+            events.append(event)
+            print(f"✓ 收到流式事件: {type(event).__name__}")
+        print(f"✓ on_message_send_stream 正常工作，收到 {len(events)} 个事件")
-sys.path.append('src')
-
-from a2a.client.transports.kafka import KafkaClientTransport
+sys.path.append('src')
+
+-from a2a.client.transports.kafka import KafkaClientTransport
+from a2a.client.transports import KafkaClientTransport
+
+if KafkaClientTransport is None:
+    print("Kafka dependencies not installed; skipping Kafka client instantiation test.")
+    sys.exit(0)
-    # 创建 Kafka 客户端传输
-    transport = KafkaClientTransport(
-        agent_card=agent_card,
-        bootstrap_servers="localhost:9092",
-        request_topic="a2a-requests"
-    )
-    
+    # 创建 Kafka 客户端传输
+    transport = KafkaClientTransport(
+        agent_card=agent_card,
+        bootstrap_servers="localhost:9092",
+        request_topic="a2a-requests"
+    )
+
+    if os.getenv("TEST_KAFKA"):
+        try:
+            import socket
+            conn = socket.create_connection(("localhost", 9092), timeout=2)
+            conn.close()
+        except Exception:
+            pytest.skip("Kafka broker not reachable")
+
+        try:
+            await transport.start()
+        finally:
+            await transport.stop()
-    message_params = MessageSendParams(
-        content="测试消息",
-        role="user"
-    )
-    print(f"消息参数创建成功: {message_params.content}")
+    import uuid
+    from a2a.types import Message, TextPart, Part, Role
+
+    message_params = MessageSendParams(
+        message=Message(
+            message_id=str(uuid.uuid4()),
+            role=Role.user,
+            parts=[Part(TextPart(text="测试消息"))],
+        )
+    )
+    print(f"消息参数创建成功: {message_params.message.parts[0].root.text}")
-        with patch.object(kafka_transport.correlation_manager, 'generate_correlation_id') as mock_gen_id, \
-             patch.object(kafka_transport.correlation_manager, 'register') as mock_register:
-            
-            mock_gen_id.return_value = "test-correlation-id"
-            
-            # Create a future that resolves to a response
-            response = Message(content="test response", role="assistant")
-            future = asyncio.Future()
-            future.set_result(response)
-            mock_register.return_value = future
-            
+        with patch.object(kafka_transport.correlation_manager, 'generate_correlation_id') as mock_gen_id, \
+             patch.object(kafka_transport.correlation_manager, 'register', new_callable=AsyncMock) as mock_register:
+            
+            mock_gen_id.return_value = "test-correlation-id"
+            
+            # Create a future that resolves to a response
+            response = Message(content="test response", role="assistant")
+            future = asyncio.Future()
+            future.set_result(response)
+            mock_register.return_value = future
-        if not self._running:
-            return
++ b/tests/client/test_kafka_client.py
+@@ -156,7 +156,7 @@ async def test_kafka_transport_start_and_publish(...):
+    await kafka_transport.start()
+     # ... assertions that transport is running ...
+
+@@ -169,7 +169,7 @@ async def test_kafka_transport_stop_and_cleanup(...):
+     # ... prepare for shutdown ...
+    await kafka_transport.stop()
+     # ... assertions that transport is no longer running ...
-            if context.trace_id:
-                headers.append(('trace_id', context.trace_id.encode('utf-8')))
+        if context:
+            # Add context headers if needed
+            if context.state.get('trace_id'):
+                headers.append(('trace_id', context.state['trace_id'].encode('utf-8')))
-            if context and context.timeout:
-                timeout = context.timeout
+            if context and context.state.get('timeout'):
+                timeout = context.state['timeout']
               __all__ = [
                   'ClientTransport',
                   'GrpcTransport',
                   'JsonRpcTransport',
+                  'KafkaClientTransport',
                   'RestTransport',
               ]
+                      correlation_id = self.correlation_manager.generate_correlation_id()
+                      # Prepare request message
+                      request_data = {
+                          'method': method,
+                          'params': params.model_dump() if hasattr(params, 'model_dump') else params,
+                          'streaming': streaming,
+                          'agent_card': self.agent_card.model_dump(),
+                      }
+                      # Prepare headers
+                      headers = [
+                          ('correlation_id', correlation_id.encode('utf-8')),
+                          ('reply_topic', (self.reply_topic or '').encode('utf-8')),
+                          ('agent_id', self.agent_card.name.encode('utf-8')),
+                      ]
+                      if context:
+                          # Add context headers if needed
+                          if context.trace_id:
+                              headers.append(('trace_id', context.trace_id.encode('utf-8')))
+                      try:
+                          await self.producer.send_and_wait(
+                              self.request_topic,
+                              value=request_data,
+                              headers=headers
+                          )
+                          return correlation_id
+                      except KafkaError as e:
+                          raise A2AClientError(f"Failed to send Kafka message: {e}") from e
+                  async def send_message(
+                      self,
+                      request: MessageSendParams,
+                      *,
+                      context: ClientCallContext | None = None,
+                  ) -> Task | Message:
+                      """Send a non-streaming message request to the agent."""
+                      correlation_id = await self._send_request('message_send', request, context, streaming=False)
+                      # Register and wait for response
+                      future = await self.correlation_manager.register(correlation_id)
+                      try:
+                          # Wait for response with timeout
+                          timeout = 30.0  # Default timeout
+                          if context and context.timeout:
+                              timeout = context.timeout
+                          result = await asyncio.wait_for(future, timeout=timeout)
+                          return result
+                      except asyncio.TimeoutError:
+                          await self.correlation_manager.complete_with_exception(
+                              correlation_id,
+                              A2AClientError(f"Request timed out after {timeout} seconds")
+                          )
+                          raise A2AClientError(f"Request timed out after {timeout} seconds")
		sys.path.append('src')

		from a2a.client.transports.kafka import KafkaClientTransport
		if context.trace_id:
		headers.append(('trace_id', context.trace_id.encode('utf-8')))
		task_request = TaskQueryParams(task_id=request.task_id)
		task = await self.get_task(task_request, context=context)
+                      mock_consumer_class.return_value = mock_consumer
+                      # Start transport using internal method
+                      await kafka_transport._start()
		@@ -0,0 +1,526 @@
		# A2A on Kafka

		## 1. 概要设计 (High-Level Design)
		@@ -0,0 +1,142 @@
		"""示例演示 A2A Kafka 传输使用方法。"""
+                      context: ClientCallContext | None = None,
+                  ) -> TaskPushNotificationConfig:
+                      """Set task push notification configuration."""
+                      # For Kafka, we can store the callback configuration locally
+                      context: ClientCallContext | None = None,
+                  ) -> AgentCard:
+                      """Retrieve the agent card."""
+                      # For Kafka transport, we return the local agent card