kailash 0.6.6__py3-none-any.whl → 0.8.0__py3-none-any.whl

kailash/channels/session.py

@@ -0,0 +1,423 @@
+"""Session management for cross-channel communication."""
+
+import asyncio
+import logging
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional, Set
+
+logger = logging.getLogger(__name__)
+
+
+class SessionStatus(Enum):
+    """Session status states."""
+
+    ACTIVE = "active"
+    IDLE = "idle"
+    EXPIRED = "expired"
+    TERMINATED = "terminated"
+
+
+@dataclass
+class CrossChannelSession:
+    """Represents a session that can span multiple channels."""
+
+    session_id: str
+    user_id: Optional[str] = None
+    created_at: float = field(default_factory=time.time)
+    last_activity: float = field(default_factory=time.time)
+    expires_at: Optional[float] = None
+    status: SessionStatus = SessionStatus.ACTIVE
+
+    # Channel tracking
+    active_channels: Set[str] = field(default_factory=set)
+    channel_contexts: Dict[str, Dict[str, Any]] = field(default_factory=dict)
+
+    # Session data
+    shared_data: Dict[str, Any] = field(default_factory=dict)
+    workflow_states: Dict[str, Any] = field(default_factory=dict)
+
+    # Event tracking
+    event_history: List[Dict[str, Any]] = field(default_factory=list)
+    max_history_size: int = 1000
+
+    def touch(self) -> None:
+        """Update last activity timestamp."""
+        self.last_activity = time.time()
+        if self.status == SessionStatus.IDLE:
+            self.status = SessionStatus.ACTIVE
+
+    def add_channel(
+        self, channel_name: str, initial_context: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Add a channel to this session.
+
+        Args:
+            channel_name: Name of the channel to add
+            initial_context: Initial context data for the channel
+        """
+        self.active_channels.add(channel_name)
+        if initial_context:
+            self.channel_contexts[channel_name] = initial_context.copy()
+        else:
+            self.channel_contexts[channel_name] = {}
+        self.touch()
+        logger.debug(f"Added channel {channel_name} to session {self.session_id}")
+
+    def remove_channel(self, channel_name: str) -> None:
+        """Remove a channel from this session.
+
+        Args:
+            channel_name: Name of the channel to remove
+        """
+        self.active_channels.discard(channel_name)
+        self.channel_contexts.pop(channel_name, None)
+        logger.debug(f"Removed channel {channel_name} from session {self.session_id}")
+
+    def update_channel_context(
+        self, channel_name: str, context_updates: Dict[str, Any]
+    ) -> None:
+        """Update context data for a specific channel.
+
+        Args:
+            channel_name: Name of the channel
+            context_updates: Context updates to apply
+        """
+        if channel_name not in self.channel_contexts:
+            self.channel_contexts[channel_name] = {}
+
+        self.channel_contexts[channel_name].update(context_updates)
+        self.touch()
+
+    def get_channel_context(self, channel_name: str) -> Dict[str, Any]:
+        """Get context data for a specific channel.
+
+        Args:
+            channel_name: Name of the channel
+
+        Returns:
+            Channel context data
+        """
+        return self.channel_contexts.get(channel_name, {}).copy()
+
+    def set_shared_data(self, key: str, value: Any) -> None:
+        """Set shared data accessible across all channels.
+
+        Args:
+            key: Data key
+            value: Data value
+        """
+        self.shared_data[key] = value
+        self.touch()
+
+    def get_shared_data(self, key: str, default: Any = None) -> Any:
+        """Get shared data.
+
+        Args:
+            key: Data key
+            default: Default value if key not found
+
+        Returns:
+            Shared data value
+        """
+        return self.shared_data.get(key, default)
+
+    def add_event(self, event: Dict[str, Any]) -> None:
+        """Add an event to the session history.
+
+        Args:
+            event: Event data to add
+        """
+        event_record = {
+            "timestamp": time.time(),
+            "session_id": self.session_id,
+            **event,
+        }
+
+        self.event_history.append(event_record)
+
+        # Maintain max history size
+        if len(self.event_history) > self.max_history_size:
+            self.event_history = self.event_history[-self.max_history_size :]
+
+        self.touch()
+
+    def is_expired(self, timeout: int = 3600) -> bool:
+        """Check if the session has expired.
+
+        Args:
+            timeout: Session timeout in seconds
+
+        Returns:
+            True if session has expired
+        """
+        if self.expires_at:
+            return time.time() > self.expires_at
+
+        return (time.time() - self.last_activity) > timeout
+
+    def extend_expiry(self, additional_seconds: int = 3600) -> None:
+        """Extend session expiry time.
+
+        Args:
+            additional_seconds: Additional seconds to extend
+        """
+        if self.expires_at:
+            self.expires_at += additional_seconds
+        else:
+            self.expires_at = time.time() + additional_seconds
+        self.touch()
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert session to dictionary for serialization."""
+        return {
+            "session_id": self.session_id,
+            "user_id": self.user_id,
+            "created_at": self.created_at,
+            "last_activity": self.last_activity,
+            "expires_at": self.expires_at,
+            "status": self.status.value,
+            "active_channels": list(self.active_channels),
+            "channel_contexts": self.channel_contexts,
+            "shared_data": self.shared_data,
+            "workflow_states": self.workflow_states,
+            "event_count": len(self.event_history),
+        }
+
+
+class SessionManager:
+    """Manages cross-channel sessions for the Nexus framework."""
+
+    def __init__(self, default_timeout: int = 3600, cleanup_interval: int = 300):
+        """Initialize session manager.
+
+        Args:
+            default_timeout: Default session timeout in seconds
+            cleanup_interval: Interval for cleanup task in seconds
+        """
+        self.default_timeout = default_timeout
+        self.cleanup_interval = cleanup_interval
+        self._sessions: Dict[str, CrossChannelSession] = {}
+        self._cleanup_task: Optional[asyncio.Task] = None
+        self._running = False
+
+    async def start(self) -> None:
+        """Start the session manager."""
+        if self._running:
+            return
+
+        self._running = True
+        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
+        logger.info("Session manager started")
+
+    async def stop(self) -> None:
+        """Stop the session manager."""
+        self._running = False
+
+        if self._cleanup_task:
+            self._cleanup_task.cancel()
+            try:
+                await self._cleanup_task
+            except asyncio.CancelledError:
+                pass
+
+        logger.info("Session manager stopped")
+
+    def create_session(
+        self,
+        user_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        timeout: Optional[int] = None,
+    ) -> CrossChannelSession:
+        """Create a new session.
+
+        Args:
+            user_id: Optional user ID for the session
+            session_id: Optional custom session ID
+            timeout: Optional custom timeout
+
+        Returns:
+            New CrossChannelSession instance
+        """
+        if session_id is None:
+            session_id = str(uuid.uuid4())
+
+        if session_id in self._sessions:
+            raise ValueError(f"Session {session_id} already exists")
+
+        session = CrossChannelSession(session_id=session_id, user_id=user_id)
+
+        if timeout:
+            session.extends_at = time.time() + timeout
+
+        self._sessions[session_id] = session
+        logger.info(f"Created session {session_id} for user {user_id}")
+
+        return session
+
+    def get_session(self, session_id: str) -> Optional[CrossChannelSession]:
+        """Get an existing session.
+
+        Args:
+            session_id: Session ID to retrieve
+
+        Returns:
+            CrossChannelSession if found, None otherwise
+        """
+        session = self._sessions.get(session_id)
+
+        if session and session.is_expired(self.default_timeout):
+            self.terminate_session(session_id)
+            return None
+
+        return session
+
+    def get_or_create_session(
+        self, session_id: str, user_id: Optional[str] = None
+    ) -> CrossChannelSession:
+        """Get existing session or create new one.
+
+        Args:
+            session_id: Session ID
+            user_id: Optional user ID for new sessions
+
+        Returns:
+            CrossChannelSession instance
+        """
+        session = self.get_session(session_id)
+        if session:
+            return session
+
+        return self.create_session(user_id=user_id, session_id=session_id)
+
+    def terminate_session(self, session_id: str) -> bool:
+        """Terminate a session.
+
+        Args:
+            session_id: Session ID to terminate
+
+        Returns:
+            True if session was terminated, False if not found
+        """
+        session = self._sessions.pop(session_id, None)
+        if session:
+            session.status = SessionStatus.TERMINATED
+            logger.info(f"Terminated session {session_id}")
+            return True
+        return False
+
+    def list_sessions(
+        self, user_id: Optional[str] = None, status: Optional[SessionStatus] = None
+    ) -> List[CrossChannelSession]:
+        """List sessions with optional filtering.
+
+        Args:
+            user_id: Filter by user ID
+            status: Filter by status
+
+        Returns:
+            List of matching sessions
+        """
+        sessions = []
+
+        for session in self._sessions.values():
+            if user_id and session.user_id != user_id:
+                continue
+            if status and session.status != status:
+                continue
+            sessions.append(session)
+
+        return sessions
+
+    def get_channel_sessions(self, channel_name: str) -> List[CrossChannelSession]:
+        """Get all sessions active on a specific channel.
+
+        Args:
+            channel_name: Name of the channel
+
+        Returns:
+            List of sessions active on the channel
+        """
+        return [
+            session
+            for session in self._sessions.values()
+            if channel_name in session.active_channels
+        ]
+
+    async def broadcast_to_channel(
+        self, channel_name: str, event: Dict[str, Any]
+    ) -> int:
+        """Broadcast an event to all sessions on a channel.
+
+        Args:
+            channel_name: Target channel name
+            event: Event data to broadcast
+
+        Returns:
+            Number of sessions that received the event
+        """
+        sessions = self.get_channel_sessions(channel_name)
+
+        for session in sessions:
+            session.add_event(
+                {"type": "broadcast", "channel": channel_name, "data": event}
+            )
+
+        logger.debug(
+            f"Broadcasted event to {len(sessions)} sessions on channel {channel_name}"
+        )
+        return len(sessions)
+
+    async def _cleanup_loop(self) -> None:
+        """Background task to clean up expired sessions."""
+        while self._running:
+            try:
+                await asyncio.sleep(self.cleanup_interval)
+                await self._cleanup_expired_sessions()
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Error in session cleanup: {e}")
+
+    async def _cleanup_expired_sessions(self) -> None:
+        """Clean up expired sessions."""
+        expired_sessions = []
+
+        for session_id, session in self._sessions.items():
+            if session.is_expired(self.default_timeout):
+                expired_sessions.append(session_id)
+
+        for session_id in expired_sessions:
+            self.terminate_session(session_id)
+
+        if expired_sessions:
+            logger.info(f"Cleaned up {len(expired_sessions)} expired sessions")
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get session manager statistics.
+
+        Returns:
+            Dictionary with session statistics
+        """
+        active_sessions = len(
+            [s for s in self._sessions.values() if s.status == SessionStatus.ACTIVE]
+        )
+        idle_sessions = len(
+            [s for s in self._sessions.values() if s.status == SessionStatus.IDLE]
+        )
+
+        channel_usage = {}
+        for session in self._sessions.values():
+            for channel in session.active_channels:
+                channel_usage[channel] = channel_usage.get(channel, 0) + 1
+
+        return {
+            "total_sessions": len(self._sessions),
+            "active_sessions": active_sessions,
+            "idle_sessions": idle_sessions,
+            "channel_usage": channel_usage,
+            "default_timeout": self.default_timeout,
+            "cleanup_interval": self.cleanup_interval,
+        }

kailash/mcp_server/discovery.py

@@ -486,7 +486,7 @@ class FileBasedDiscovery(DiscoveryBackend):
             raise
 
 
-class NetworkDiscovery:
+class NetworkDiscovery(asyncio.DatagramProtocol):
     """Network-based discovery using UDP broadcast/multicast."""
 
     DISCOVERY_PORT = 8765
@@ -535,6 +535,40 @@ class NetworkDiscovery:
             self._transport = None
             self._protocol = None
 
+    # AsyncIO DatagramProtocol methods
+    def connection_made(self, transport):
+        """Called when a connection is made."""
+        self._transport = transport
+        logger.info(f"Network discovery protocol connected on port {self.port}")
+
+    def datagram_received(self, data, addr):
+        """Called when a datagram is received."""
+        try:
+            message = json.loads(data.decode())
+            # Try to get current event loop
+            try:
+                asyncio.get_running_loop()
+                asyncio.create_task(self._handle_discovery_message(message, addr))
+            except RuntimeError:
+                # No event loop, run synchronously
+                asyncio.run(self._handle_discovery_message(message, addr))
+        except json.JSONDecodeError:
+            logger.warning(f"Invalid JSON received from {addr}")
+        except Exception as e:
+            logger.error(f"Error handling datagram from {addr}: {e}")
+
+    def error_received(self, exc):
+        """Called when an error is received."""
+        logger.error(f"Network discovery protocol error: {exc}")
+
+    def connection_lost(self, exc):
+        """Called when the connection is lost."""
+        if exc:
+            logger.error(f"Network discovery connection lost: {exc}")
+        else:
+            logger.info("Network discovery connection closed")
+        self.running = False
+
     async def start_discovery_listener(self):
         """Start listening for server announcements."""
         await self.start()
@@ -590,6 +624,27 @@ class NetworkDiscovery:
         except (json.JSONDecodeError, KeyError) as e:
             logger.debug(f"Invalid announcement from {addr[0]}: {e}")
 
+    async def _is_port_open(self, host: str, port: int, timeout: float = 1.0) -> bool:
+        """Check if a port is open on a host.
+
+        Args:
+            host: Host to check
+            port: Port to check
+            timeout: Connection timeout
+
+        Returns:
+            True if port is open, False otherwise
+        """
+        try:
+            # Create socket connection with timeout
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.settimeout(timeout)
+            result = sock.connect_ex((host, port))
+            sock.close()
+            return result == 0
+        except (OSError, socket.error, socket.timeout):
+            return False
+
     async def scan_network(
         self, network: str = "192.168.1.0/24", timeout: float = 5.0
     ) -> List[ServerInfo]:
@@ -756,22 +811,6 @@ class NetworkDiscovery:
         else:
             logger.debug(f"Unknown message type: {msg_type}")
 
-    def datagram_received(self, data: bytes, addr: tuple):
-        """Handle received datagram (part of asyncio protocol)."""
-        try:
-            message = json.loads(data.decode())
-            # Try to get current event loop
-            try:
-                loop = asyncio.get_running_loop()
-                asyncio.create_task(self._handle_discovery_message(message, addr))
-            except RuntimeError:
-                # No event loop, run synchronously
-                asyncio.run(self._handle_discovery_message(message, addr))
-        except json.JSONDecodeError:
-            logger.warning(f"Invalid JSON received from {addr}")
-        except Exception as e:
-            logger.error(f"Error handling datagram from {addr}: {e}")
-
 
 class ServiceRegistry:
     """Main service registry coordinating multiple discovery backends."""
@@ -967,7 +1006,7 @@ class ServiceRegistry:
     async def stop_health_monitoring(self):
         """Stop health monitoring (async version)."""
         if self.health_checker:
-            self.health_checker.stop()
+            await self.health_checker.stop()
 
     async def get_best_server_for_capability(
         self, capability: str

kailash/middleware/communication/api_gateway.py

@@ -331,9 +331,23 @@ class APIGateway:
     def _setup_session_routes(self):
         """Setup session management routes."""
 
+        # Create auth dependency
+        async def get_optional_current_user():
+            """Optional auth dependency - returns None if auth is disabled."""
+            if self.enable_auth and self.auth_manager:
+                # Use auth manager's dependency if available
+                try:
+                    # This would normally use the auth manager's get_current_user_dependency
+                    # For now, return None to avoid complex auth setup
+                    return None
+                except:
+                    return None
+            return None
+
         @self.app.post("/api/sessions", response_model=SessionResponse)
         async def create_session(
-            request: SessionCreateRequest, current_user: Dict[str, Any] = None
+            request: SessionCreateRequest,
+            current_user: Dict[str, Any] = Depends(get_optional_current_user),
         ):
             """Create a new session for a frontend client."""
             try:
@@ -583,7 +597,13 @@ class APIGateway:
             """Get schemas for available node types."""
             try:
                 # Get all registered nodes
-                available_nodes = self.node_registry.get_all_nodes()
+                # NodeRegistry doesn't have get_all_nodes, need to use _nodes directly
+                available_nodes = {}
+                if hasattr(self.node_registry, "_nodes"):
+                    available_nodes = self.node_registry._nodes.copy()
+                else:
+                    # Fallback - return empty dict
+                    available_nodes = {}
 
                 # Filter by requested types if specified
                 if request.node_types:
@@ -611,7 +631,7 @@ class APIGateway:
         @self.app.get("/api/schemas/nodes/{node_type}")
         async def get_node_schema(node_type: str):
             """Get schema for a specific node type."""
-            node_class = self.node_registry.get_node(node_type)
+            node_class = self.node_registry.get(node_type)
             if not node_class:
                 raise HTTPException(status_code=404, detail="Node type not found")
 

kailash/middleware/communication/realtime.py

@@ -177,6 +177,89 @@ class ConnectionManager:
             "active_users": len(self.user_connections),
             "total_messages_sent": total_messages,
         }
+    
+    def filter_events(self, events: List[BaseEvent], event_filter: EventFilter = None) -> List[BaseEvent]:
+        """Filter events based on event filter criteria."""
+        if not event_filter:
+            return events
+        
+        filtered = []
+        for event in events:
+            # Apply session filter
+            if event_filter.session_id and hasattr(event, 'session_id'):
+                if event.session_id != event_filter.session_id:
+                    continue
+            
+            # Apply user filter
+            if event_filter.user_id and hasattr(event, 'user_id'):
+                if event.user_id != event_filter.user_id:
+                    continue
+            
+            # Apply event type filter
+            if event_filter.event_types and event.event_type not in event_filter.event_types:
+                continue
+            
+            filtered.append(event)
+        
+        return filtered
+    
+    def set_event_filter(self, connection_id: str, event_filter: EventFilter):
+        """Set event filter for a specific connection."""
+        if connection_id in self.connections:
+            self.connections[connection_id]["event_filter"] = event_filter
+    
+    def get_event_filter(self, connection_id: str) -> Optional[EventFilter]:
+        """Get event filter for a specific connection."""
+        if connection_id in self.connections:
+            return self.connections[connection_id].get("event_filter")
+        return None
+    
+    # Alias methods for compatibility
+    def event_filter(self, events: List[BaseEvent], filter_criteria: EventFilter = None) -> List[BaseEvent]:
+        """Alias for filter_events method."""
+        return self.filter_events(events, filter_criteria)
+    
+    async def on_event(self, event: BaseEvent):
+        """Handle incoming event - route to appropriate connections."""
+        await self.handle_event(event)
+    
+    async def handle_event(self, event: BaseEvent):
+        """Handle and route event to matching connections."""
+        await self.process_event(event)
+    
+    async def process_event(self, event: BaseEvent):
+        """Process event and broadcast to matching connections."""
+        message = {
+            "type": "event",
+            "event_type": event.event_type.value if hasattr(event.event_type, 'value') else str(event.event_type),
+            "data": event.data,
+            "timestamp": event.timestamp.isoformat() if hasattr(event, 'timestamp') else datetime.now(timezone.utc).isoformat(),
+            "session_id": getattr(event, 'session_id', None),
+            "user_id": getattr(event, 'user_id', None),
+        }
+        
+        # Broadcast to all matching connections
+        for connection_id, connection in self.connections.items():
+            event_filter = connection.get("event_filter")
+            
+            # Check if this connection should receive this event
+            should_send = True
+            if event_filter:
+                # Apply session filter
+                if event_filter.session_id and connection["session_id"] != event_filter.session_id:
+                    should_send = False
+                
+                # Apply user filter
+                if event_filter.user_id and connection["user_id"] != event_filter.user_id:
+                    should_send = False
+                
+                # Apply event type filter
+                if hasattr(event_filter, 'event_types') and event_filter.event_types:
+                    if event.event_type not in event_filter.event_types:
+                        should_send = False
+            
+            if should_send:
+                await self.send_to_connection(connection_id, message)
 
 
 class SSEManager:

kailash/middleware/core/agent_ui.py

@@ -944,7 +944,7 @@ class AgentUIMiddleware:
     async def get_available_nodes(self) -> List[Dict[str, Any]]:
         """Get all available node types with their schemas."""
         nodes = []
-        for node_name, node_class in self.node_registry.nodes.items():
+        for node_name, node_class in self.node_registry._nodes.items():
             # Get node schema (would be implemented in schema.py)
             schema = await self._get_node_schema(node_class)
             nodes.append(