kailash 0.8.5__py3-none-any.whl → 0.8.6__py3-none-any.whl

kailash/mcp_server/transports.py

@@ -1084,6 +1084,310 @@ class WebSocketTransport(BaseTransport):
             self.websocket = None
 
 
+class WebSocketServerTransport(BaseTransport):
+    """WebSocket server transport for accepting MCP connections."""
+
+    def __init__(
+        self,
+        host: str = "0.0.0.0",
+        port: int = 3001,
+        message_handler: Optional[
+            Callable[[Dict[str, Any], str], Dict[str, Any]]
+        ] = None,
+        ping_interval: float = 20.0,
+        ping_timeout: float = 20.0,
+        max_message_size: int = 10 * 1024 * 1024,  # 10MB
+        **kwargs,
+    ):
+        """Initialize WebSocket server transport.
+
+        Args:
+            host: Host to bind to
+            port: Port to listen on
+            message_handler: Handler for incoming messages
+            ping_interval: Ping interval in seconds
+            ping_timeout: Ping timeout in seconds
+            max_message_size: Maximum message size in bytes
+            **kwargs: Base transport arguments
+        """
+        super().__init__("websocket_server", **kwargs)
+
+        self.host = host
+        self.port = port
+        self.message_handler = message_handler
+        self.ping_interval = ping_interval
+        self.ping_timeout = ping_timeout
+        self.max_message_size = max_message_size
+
+        # Server state
+        self.server: Optional[websockets.WebSocketServer] = None
+        self._clients: Dict[str, Any] = {}  # websockets.WebSocketServerProtocol
+        self._client_sessions: Dict[str, Dict[str, Any]] = {}
+        self._server_task: Optional[asyncio.Task] = None
+
+    async def connect(self) -> None:
+        """Start the WebSocket server."""
+        if self._connected:
+            return
+
+        try:
+            # Create handler that works with new websockets API
+            async def connection_handler(websocket):
+                # Get path from the websocket's request path
+                path = websocket.path if hasattr(websocket, "path") else "/"
+                await self.handle_client(websocket, path)
+
+            # Start WebSocket server
+            self.server = await websockets.serve(
+                connection_handler,
+                self.host,
+                self.port,
+                ping_interval=self.ping_interval,
+                ping_timeout=self.ping_timeout,
+                max_size=self.max_message_size,
+            )
+
+            self._connected = True
+            self._update_metrics("connections_total")
+
+            logger.info(f"WebSocket server listening on {self.host}:{self.port}")
+
+        except Exception as e:
+            self._update_metrics("connections_failed")
+            raise TransportError(
+                f"Failed to start WebSocket server: {e}",
+                transport_type="websocket_server",
+            )
+
+    async def disconnect(self) -> None:
+        """Stop the WebSocket server."""
+        if not self._connected:
+            return
+
+        self._connected = False
+
+        # Close all client connections
+        clients = list(self._clients.values())
+        for client in clients:
+            await client.close()
+
+        # Stop server
+        if self.server:
+            self.server.close()
+            await self.server.wait_closed()
+            self.server = None
+
+        # Clear client tracking
+        self._clients.clear()
+        self._client_sessions.clear()
+
+        logger.info("WebSocket server stopped")
+
+    async def send_message(
+        self, message: Dict[str, Any], client_id: Optional[str] = None
+    ) -> None:
+        """Send message to specific client or broadcast to all.
+
+        Args:
+            message: Message to send
+            client_id: Target client ID (None for broadcast)
+        """
+        if not self._connected:
+            raise TransportError(
+                "Transport not connected", transport_type="websocket_server"
+            )
+
+        message_data = json.dumps(message)
+
+        try:
+            if client_id:
+                # Send to specific client
+                if client_id in self._clients:
+                    await self._clients[client_id].send(message_data)
+                    self._update_metrics("messages_sent")
+                    self._update_metrics("bytes_sent", len(message_data))
+                else:
+                    raise TransportError(
+                        f"Client {client_id} not found",
+                        transport_type="websocket_server",
+                    )
+            else:
+                # Broadcast to all clients
+                if self._clients:
+                    await asyncio.gather(
+                        *[
+                            client.send(message_data)
+                            for client in self._clients.values()
+                        ],
+                        return_exceptions=True,
+                    )
+                    self._update_metrics("messages_sent", len(self._clients))
+                    self._update_metrics(
+                        "bytes_sent", len(message_data) * len(self._clients)
+                    )
+
+        except Exception as e:
+            self._update_metrics("errors_total")
+            raise TransportError(
+                f"Failed to send message: {e}", transport_type="websocket_server"
+            )
+
+    async def receive_message(self) -> Dict[str, Any]:
+        """Not implemented for server transport."""
+        raise NotImplementedError(
+            "Server transport doesn't support receive_message. "
+            "Messages are handled via handle_client callback."
+        )
+
+    async def handle_client(self, websocket, path: str):
+        """Handle a client connection.
+
+        Args:
+            websocket: WebSocket connection
+            path: Request path
+        """
+        client_id = str(uuid.uuid4())
+        self._clients[client_id] = websocket
+        self._client_sessions[client_id] = {
+            "connected_at": time.time(),
+            "path": path,
+            "remote_address": websocket.remote_address,
+        }
+
+        logger.info(f"Client {client_id} connected from {websocket.remote_address}")
+        self._update_metrics("connections_total")
+
+        try:
+            async for message in websocket:
+                try:
+                    # Parse message
+                    request = json.loads(message)
+
+                    # Update metrics
+                    self._update_metrics("messages_received")
+                    self._update_metrics("bytes_received", len(message))
+
+                    # Handle message
+                    if self.message_handler:
+                        response = await self._handle_message_safely(request, client_id)
+                    else:
+                        response = {
+                            "jsonrpc": "2.0",
+                            "error": {
+                                "code": -32601,
+                                "message": "No message handler configured",
+                            },
+                            "id": request.get("id"),
+                        }
+
+                    # Send response
+                    await websocket.send(json.dumps(response))
+                    self._update_metrics("messages_sent")
+                    self._update_metrics("bytes_sent", len(json.dumps(response)))
+
+                except json.JSONDecodeError as e:
+                    logger.error(f"Invalid JSON from client {client_id}: {e}")
+                    self._update_metrics("errors_total")
+
+                    error_response = {
+                        "jsonrpc": "2.0",
+                        "error": {
+                            "code": -32700,
+                            "message": "Parse error: Invalid JSON",
+                        },
+                        "id": None,
+                    }
+                    await websocket.send(json.dumps(error_response))
+
+                except Exception as e:
+                    logger.error(f"Error handling message from client {client_id}: {e}")
+                    self._update_metrics("errors_total")
+
+        except websockets.exceptions.ConnectionClosed:
+            logger.info(f"Client {client_id} disconnected")
+        except Exception as e:
+            logger.error(f"Error in client handler for {client_id}: {e}")
+        finally:
+            # Clean up client
+            del self._clients[client_id]
+            del self._client_sessions[client_id]
+
+    async def _handle_message_safely(
+        self, request: Dict[str, Any], client_id: str
+    ) -> Dict[str, Any]:
+        """Handle message with error handling.
+
+        Args:
+            request: JSON-RPC request
+            client_id: Client identifier
+
+        Returns:
+            JSON-RPC response
+        """
+        try:
+            if asyncio.iscoroutinefunction(self.message_handler):
+                return await self.message_handler(request, client_id)
+            else:
+                return self.message_handler(request, client_id)
+        except Exception as e:
+            logger.error(f"Message handler error: {e}")
+            return {
+                "jsonrpc": "2.0",
+                "error": {
+                    "code": -32603,
+                    "message": f"Internal error: {str(e)}",
+                },
+                "id": request.get("id"),
+            }
+
+    def get_client_info(self, client_id: str) -> Optional[Dict[str, Any]]:
+        """Get information about a connected client.
+
+        Args:
+            client_id: Client identifier
+
+        Returns:
+            Client information or None
+        """
+        if client_id not in self._client_sessions:
+            return None
+
+        session = self._client_sessions[client_id]
+        return {
+            "client_id": client_id,
+            "connected_at": session["connected_at"],
+            "connection_duration": time.time() - session["connected_at"],
+            "path": session["path"],
+            "remote_address": session["remote_address"],
+        }
+
+    def list_clients(self) -> List[Dict[str, Any]]:
+        """List all connected clients.
+
+        Returns:
+            List of client information
+        """
+        return [self.get_client_info(client_id) for client_id in self._client_sessions]
+
+    async def close_client(
+        self, client_id: str, code: int = 1000, reason: str = ""
+    ) -> bool:
+        """Close a specific client connection.
+
+        Args:
+            client_id: Client to disconnect
+            code: WebSocket close code
+            reason: Close reason
+
+        Returns:
+            True if client was closed
+        """
+        if client_id in self._clients:
+            await self._clients[client_id].close(code, reason)
+            return True
+        return False
+
+
 class TransportManager:
     """Manager for MCP transport instances."""
 
@@ -1095,6 +1399,7 @@ class TransportManager:
             "sse": SSETransport,
             "streamable_http": StreamableHTTPTransport,
             "websocket": WebSocketTransport,
+            "websocket_server": WebSocketServerTransport,
         }
 
     def register_transport_factory(self, transport_type: str, factory: Callable):

kailash/middleware/gateway/event_store.py

@@ -125,6 +125,7 @@ class EventStore:
             self._flush_task = asyncio.create_task(self._flush_loop())
         except RuntimeError:
             # If no event loop is running, defer task creation
+            # Don't create the coroutine here as it will never be awaited
             self._flush_task = None
 
     async def _ensure_flush_task(self):

kailash/nodes/base.py

@@ -255,6 +255,54 @@ class Node(ABC):
                 f"Failed to initialize node '{self.id}': {e}"
             ) from e
 
+    def get_workflow_context(self, key: str, default: Any = None) -> Any:
+        """Get a value from the workflow context.
+
+        This method allows nodes to retrieve shared state from the workflow
+        execution context. The workflow context is managed by the runtime
+        and provides a way for nodes to share data within a single workflow
+        execution.
+
+        Args:
+            key: The key to retrieve from the workflow context
+            default: Default value to return if key is not found
+
+        Returns:
+            The value from the workflow context, or default if not found
+
+        Example:
+            >>> # In a transaction node
+            >>> connection = self.get_workflow_context('transaction_connection')
+            >>> if connection:
+            >>>     # Use the shared connection
+            >>>     result = await connection.execute(query)
+        """
+        if not hasattr(self, "_workflow_context"):
+            self._workflow_context = {}
+        return self._workflow_context.get(key, default)
+
+    def set_workflow_context(self, key: str, value: Any) -> None:
+        """Set a value in the workflow context.
+
+        This method allows nodes to store shared state in the workflow
+        execution context. Other nodes in the same workflow execution
+        can retrieve this data using get_workflow_context().
+
+        Args:
+            key: The key to store the value under
+            value: The value to store in the workflow context
+
+        Example:
+            >>> # In a transaction scope node
+            >>> connection = await self.get_connection()
+            >>> transaction = await connection.begin()
+            >>> self.set_workflow_context('transaction_connection', connection)
+            >>> self.set_workflow_context('active_transaction', transaction)
+        """
+        if not hasattr(self, "_workflow_context"):
+            self._workflow_context = {}
+        self._workflow_context[key] = value
+
     @abstractmethod
     def get_parameters(self) -> dict[str, NodeParameter]:
         """Define the parameters this node accepts.
@@ -467,9 +515,23 @@ class Node(ABC):
                 # Skip type checking for Any type
                 if param_def.type is Any:
                     continue
+                # Skip validation for template expressions like ${variable_name}
+                if isinstance(value, str) and self._is_template_expression(value):
+                    continue
                 if not isinstance(value, param_def.type):
                     try:
-                        self.config[param_name] = param_def.type(value)
+                        # Special handling for datetime conversion from ISO strings
+                        if param_def.type.__name__ == "datetime" and isinstance(
+                            value, str
+                        ):
+                            from datetime import datetime
+
+                            # Try to parse ISO format string
+                            self.config[param_name] = datetime.fromisoformat(
+                                value.replace("Z", "+00:00")
+                            )
+                        else:
+                            self.config[param_name] = param_def.type(value)
                     except (ValueError, TypeError) as e:
                         raise NodeConfigurationError(
                             f"Configuration parameter '{param_name}' must be of type "
@@ -477,6 +539,20 @@ class Node(ABC):
                             f"Conversion failed: {e}"
                         ) from e
 
+    def _is_template_expression(self, value: str) -> bool:
+        """Check if a string value is a template expression like ${variable_name}.
+
+        Args:
+            value: String value to check
+
+        Returns:
+            True if the value is a template expression, False otherwise
+        """
+        import re
+
+        # Match template expressions like ${variable_name} or ${node.output}
+        return bool(re.match(r"^\$\{[^}]+\}$", value))
+
     def _get_cached_parameters(self) -> dict[str, NodeParameter]:
         """Get cached parameter definitions.
 

kailash/nodes/code/python.py

@@ -389,7 +389,9 @@ class CodeExecutor:
                 f"Error position: {' ' * (e.offset - 1) if e.offset else ''}^"
             )
 
-    def execute_code(self, code: str, inputs: dict[str, Any]) -> dict[str, Any]:
+    def execute_code(
+        self, code: str, inputs: dict[str, Any], node_instance=None
+    ) -> dict[str, Any]:
         """Execute Python code with given inputs.
 
         Args:
@@ -476,6 +478,43 @@ class CodeExecutor:
                 except ImportError:
                     logger.warning(f"Module {module_name} not available")
 
+        # Add global utility functions to namespace
+        try:
+            from kailash.utils.data_paths import (
+                get_data_path,
+                get_input_data_path,
+                get_output_data_path,
+            )
+
+            namespace["get_input_data_path"] = get_input_data_path
+            namespace["get_output_data_path"] = get_output_data_path
+            namespace["get_data_path"] = get_data_path
+        except ImportError:
+            logger.warning(
+                "Could not import data path utilities - functions will not be available in PythonCodeNode execution"
+            )
+
+        # Add workflow context functions if node instance is available
+        if node_instance and hasattr(node_instance, "get_workflow_context"):
+            # Bind the actual node methods
+            namespace["get_workflow_context"] = node_instance.get_workflow_context
+            namespace["set_workflow_context"] = node_instance.set_workflow_context
+        else:
+            # Add placeholder functions that warn about unavailability
+            def _get_workflow_context(key: str, default=None):
+                logger.warning(
+                    "get_workflow_context() is not available in PythonCodeNode execution context. Node instance not provided."
+                )
+                return default
+
+            def _set_workflow_context(key: str, value):
+                logger.warning(
+                    "set_workflow_context() is not available in PythonCodeNode execution context. Node instance not provided."
+                )
+
+            namespace["get_workflow_context"] = _get_workflow_context
+            namespace["set_workflow_context"] = _set_workflow_context
+
         # Add sanitized inputs
         namespace.update(sanitized_inputs)
 
@@ -1222,7 +1261,9 @@ class PythonCodeNode(Node):
         try:
             if self.code:
                 # Execute code string
-                outputs = self.executor.execute_code(self.code, kwargs)
+                outputs = self.executor.execute_code(
+                    self.code, kwargs, node_instance=self
+                )
                 # Return 'result' variable if it exists, otherwise all outputs
                 if "result" in outputs:
                     return {"result": outputs["result"]}
@@ -1454,7 +1495,7 @@ class PythonCodeNode(Node):
         """
         # Execute directly based on execution type
         if self.code:
-            outputs = self.executor.execute_code(self.code, inputs)
+            outputs = self.executor.execute_code(self.code, inputs, node_instance=self)
             return outputs.get("result", outputs)
         elif self.function:
             wrapper = FunctionWrapper(self.function, self.executor)

kailash/nodes/data/async_sql.py

@@ -489,6 +489,7 @@ class PostgreSQLAdapter(DatabaseAdapter):
                     or "DELETE" in query_upper
                     or "INSERT" in query_upper
                 )
+                and "SELECT" not in query_upper
                 and "RETURNING" not in query_upper
                 and fetch_mode == FetchMode.ALL
             ):
@@ -527,6 +528,7 @@ class PostgreSQLAdapter(DatabaseAdapter):
                         or "DELETE" in query_upper
                         or "INSERT" in query_upper
                     )
+                    and "SELECT" not in query_upper
                     and "RETURNING" not in query_upper
                     and fetch_mode == FetchMode.ALL
                 ):
@@ -2617,33 +2619,53 @@ class AsyncSQLDatabaseNode(AsyncNode):
 
     async def cleanup(self):
         """Clean up database connections."""
+        try:
+            # Check if we have a running event loop
+            loop = asyncio.get_running_loop()
+            if loop.is_closed():
+                # Event loop is closing, skip cleanup
+                return
+        except RuntimeError:
+            # No event loop, skip cleanup
+            return
+
         # Rollback any active transaction
         if self._active_transaction and self._adapter:
             try:
-                await self._adapter.rollback_transaction(self._active_transaction)
-            except Exception:
+                await asyncio.wait_for(
+                    self._adapter.rollback_transaction(self._active_transaction),
+                    timeout=1.0,
+                )
+            except (Exception, asyncio.TimeoutError):
                 pass  # Best effort cleanup
             self._active_transaction = None
 
         if self._adapter and self._connected:
-            if self._share_pool and self._pool_key:
-                # Decrement reference count for shared pool
-                async with self._get_pool_lock():
-                    if self._pool_key in self._shared_pools:
-                        adapter, ref_count = self._shared_pools[self._pool_key]
-                        if ref_count > 1:
-                            # Others still using the pool
-                            self._shared_pools[self._pool_key] = (
-                                adapter,
-                                ref_count - 1,
-                            )
-                        else:
-                            # Last reference, close the pool
-                            del self._shared_pools[self._pool_key]
-                            await adapter.disconnect()
-            else:
-                # Dedicated pool, close directly
-                await self._adapter.disconnect()
+            try:
+                if self._share_pool and self._pool_key:
+                    # Decrement reference count for shared pool with timeout
+                    async with await asyncio.wait_for(
+                        self._get_pool_lock(), timeout=1.0
+                    ):
+                        if self._pool_key in self._shared_pools:
+                            adapter, ref_count = self._shared_pools[self._pool_key]
+                            if ref_count > 1:
+                                # Others still using the pool
+                                self._shared_pools[self._pool_key] = (
+                                    adapter,
+                                    ref_count - 1,
+                                )
+                            else:
+                                # Last reference, close the pool
+                                del self._shared_pools[self._pool_key]
+                                await asyncio.wait_for(
+                                    adapter.disconnect(), timeout=1.0
+                                )
+                else:
+                    # Dedicated pool, close directly
+                    await asyncio.wait_for(self._adapter.disconnect(), timeout=1.0)
+            except (Exception, asyncio.TimeoutError):
+                pass  # Best effort cleanup
 
             self._connected = False
             self._adapter = None

kailash/nodes/edge/edge_migration_node.py

@@ -8,6 +8,7 @@ import asyncio
 from datetime import datetime
 from typing import Any, Dict, List, Optional
 
+from kailash.edge.migration.edge_migration_service import EdgeMigrationService
 from kailash.edge.migration.edge_migrator import (
     EdgeMigrator,
     MigrationPhase,
@@ -59,18 +60,21 @@ class EdgeMigrationNode(AsyncNode):
         """Initialize edge migration node."""
         super().__init__(**kwargs)
 
-        # Extract configuration
-        checkpoint_interval = kwargs.get("checkpoint_interval", 60)
-        sync_batch_size = kwargs.get("sync_batch_size", 1000)
-        bandwidth_limit_mbps = kwargs.get("bandwidth_limit_mbps")
-        enable_compression = kwargs.get("enable_compression", True)
-
-        # Initialize migrator
-        self.migrator = EdgeMigrator(
-            checkpoint_interval=checkpoint_interval,
-            sync_batch_size=sync_batch_size,
-            bandwidth_limit_mbps=bandwidth_limit_mbps,
-            enable_compression=enable_compression,
+        # Extract node-specific configuration
+        self.node_config = {
+            "checkpoint_interval": kwargs.get("checkpoint_interval", 60),
+            "sync_batch_size": kwargs.get("sync_batch_size", 1000),
+            "bandwidth_limit_mbps": kwargs.get("bandwidth_limit_mbps"),
+            "enable_compression": kwargs.get("enable_compression", True),
+        }
+
+        # Get reference to shared migration service
+        self.migration_service = EdgeMigrationService(self.node_config)
+
+        # Get migrator instance from shared service with node-specific config
+        self.node_id = f"edge_migration_node_{id(self)}"
+        self.migrator = self.migration_service.get_migrator_for_node(
+            self.node_id, self.node_config
         )
 
         self._migrator_started = False