fix: address code review issues in HTTP Stream implementation

- Fix error propagation in write_stream method
- Add proper shutdown handling in SSE listener to prevent resource leaks
- Implement proper type annotations based on schema types
- Add _closing flag to gracefully stop SSE listener
- Create basic test to verify transport selection logic

Co-authored-by: Mervin Praison <[email protected]>

Author: claude[bot]

src/praisonai-agents/praisonaiagents/mcp/mcp_http_stream.py

@@ -50,13 +50,30 @@ def __init__(self, name: str, description: str, session: ClientSession, input_sc
         # Create a signature based on input schema
         params = []
         if input_schema and 'properties' in input_schema:
-            for param_name in input_schema['properties']:
+            for param_name, prop_schema in input_schema['properties'].items():
+                # Determine type annotation based on schema
+                prop_type = prop_schema.get('type', 'string') if isinstance(prop_schema, dict) else 'string'
+                if prop_type == 'string':
+                    annotation = str
+                elif prop_type == 'integer':
+                    annotation = int
+                elif prop_type == 'number':
+                    annotation = float
+                elif prop_type == 'boolean':
+                    annotation = bool
+                elif prop_type == 'array':
+                    annotation = list
+                elif prop_type == 'object':
+                    annotation = dict
+                else:
+                    annotation = Any
+                
                 params.append(
                     inspect.Parameter(
                         name=param_name,
                         kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
                         default=inspect.Parameter.empty if param_name in input_schema.get('required', []) else None,
-                        annotation=str  # Default to string
+                        annotation=annotation
                     )
                 )
 
@@ -174,12 +191,16 @@ def __init__(self, base_url: str, session_id: Optional[str] = None, options: Opt
         self._sse_task = None
         self._message_queue = asyncio.Queue()
         self._pending_requests = {}
+        self._closing = False
 
     async def __aenter__(self):
         self._session = aiohttp.ClientSession()
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
+        # Set closing flag to stop listener gracefully
+        self._closing = True
+        
         if self._sse_task:
             self._sse_task.cancel()
             try:
@@ -254,6 +275,10 @@ async def _sse_listener(self):
         """Background task to listen for SSE events."""
         while True:
             try:
+                # Check if we should stop
+                if hasattr(self, '_closing') and self._closing:
+                    break
+                    
                 url = self.base_url
                 if self.session_id:
                     # Add session as query parameter for SSE connection
@@ -269,6 +294,10 @@ async def _sse_listener(self):
                 async with self._session.get(url, headers=headers) as response:
                     buffer = ""
                     async for chunk in response.content:
+                        # Check if we should stop
+                        if hasattr(self, '_closing') and self._closing:
+                            break
+                            
                         buffer += chunk.decode('utf-8')
 
                         # Process complete SSE events
@@ -289,9 +318,15 @@ async def _sse_listener(self):
                                 except json.JSONDecodeError:
                                     logger.error(f"Failed to parse SSE event: {data}")
 
+            except asyncio.CancelledError:
+                # Proper shutdown
+                break
             except Exception as e:
-                logger.error(f"SSE listener error: {e}")
-                await asyncio.sleep(1)  # Reconnect after 1 second
+                if not (hasattr(self, '_closing') and self._closing):
+                    logger.error(f"SSE listener error: {e}")
+                    await asyncio.sleep(1)  # Reconnect after 1 second
+                else:
+                    break
 
     def read_stream(self):
         """Create a read stream for the ClientSession."""
@@ -306,7 +341,8 @@ def write_stream(self):
         async def _write(message):
             if hasattr(message, 'to_dict'):
                 message = message.to_dict()
-            await self.send_request(message)
+            response = await self.send_request(message)
+            return response
         return _write
 
 

src/praisonai-agents/praisonaiagents/mcp/mcp_sse.py

@@ -44,13 +44,30 @@ def __init__(self, name: str, description: str, session: ClientSession, input_sc
         # Create a signature based on input schema
         params = []
         if input_schema and 'properties' in input_schema:
-            for param_name in input_schema['properties']:
+            for param_name, prop_schema in input_schema['properties'].items():
+                # Determine type annotation based on schema
+                prop_type = prop_schema.get('type', 'string') if isinstance(prop_schema, dict) else 'string'
+                if prop_type == 'string':
+                    annotation = str
+                elif prop_type == 'integer':
+                    annotation = int
+                elif prop_type == 'number':
+                    annotation = float
+                elif prop_type == 'boolean':
+                    annotation = bool
+                elif prop_type == 'array':
+                    annotation = list
+                elif prop_type == 'object':
+                    annotation = dict
+                else:
+                    annotation = str  # Default to string for SSE
+                
                 params.append(
                     inspect.Parameter(
                         name=param_name,
                         kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
                         default=inspect.Parameter.empty if param_name in input_schema.get('required', []) else None,
-                        annotation=str  # Default to string
+                        annotation=annotation
                     )
                 )
 

src/praisonai-agents/tests/test_http_stream_basic.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+"""
+Basic test to verify HTTP Stream implementation is working correctly.
+"""
+
+import sys
+import os
+
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from praisonaiagents.mcp import MCP
+
+def test_transport_selection():
+    """Test that URLs are correctly routed to appropriate transports."""
+    
+    print("Testing transport selection logic...")
+    
+    # Test 1: SSE URL should use SSE transport
+    try:
+        mcp_sse = MCP("http://localhost:8080/sse")
+        assert mcp_sse.is_sse == True
+        assert mcp_sse.is_http_stream == False
+        print("✓ SSE URL correctly uses SSE transport")
+    except Exception as e:
+        print(f"✗ SSE URL test failed: {e}")
+    
+    # Test 2: Regular HTTP URL should use HTTP Stream transport
+    try:
+        mcp_http = MCP("http://localhost:8080")
+        assert mcp_http.is_sse == False
+        assert mcp_http.is_http_stream == True
+        print("✓ HTTP URL correctly uses HTTP Stream transport")
+    except Exception as e:
+        print(f"✗ HTTP URL test failed: {e}")
+    
+    # Test 3: Custom endpoint should use HTTP Stream transport
+    try:
+        mcp_custom = MCP("http://localhost:8080/custom")
+        assert mcp_custom.is_sse == False
+        assert mcp_custom.is_http_stream == True
+        print("✓ Custom endpoint correctly uses HTTP Stream transport")
+    except Exception as e:
+        print(f"✗ Custom endpoint test failed: {e}")
+    
+    # Test 4: Stdio transport should still work
+    try:
+        mcp_stdio = MCP("python /path/to/server.py")
+        assert mcp_stdio.is_sse == False
+        assert mcp_stdio.is_http_stream == False
+        print("✓ Stdio transport still works")
+    except Exception as e:
+        print(f"✗ Stdio transport test failed: {e}")
+    
+    print("\nAll transport selection tests completed!")
+
+if __name__ == "__main__":
+    test_transport_selection()