toolplane-python-client 0.1.0__py3-none-any.whl

toolplane/common/base_session_manager.py

@@ -0,0 +1,321 @@
+"""Base session manager class shared between gRPC and HTTP clients."""
+
+import threading
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Set
+
+try:
+    from ..core.errors import SessionError
+except ImportError:
+    # Fallback for standalone testing
+    class SessionError(Exception):
+        """Session error for standalone testing."""
+
+        pass
+
+
+from .utils import (
+    generate_session_id,
+    sanitize_input,
+    validate_description,
+    validate_session_name,
+)
+
+
+class BaseSessionManager(ABC):
+    """Base class for session management."""
+
+    def __init__(self, connection_manager):
+        """Initialize base session manager."""
+        self.connection_manager = connection_manager
+        self.sessions: Dict[str, Any] = {}
+        self.sessions_lock = threading.RLock()
+        self._owned_sessions: Set[str] = set()
+
+    def create_session(
+        self,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        namespace: Optional[str] = None,
+        api_key: Optional[str] = None,
+    ) -> str:
+        """Create a new session."""
+        # Validate inputs
+        if name and not validate_session_name(name):
+            raise SessionError(f"Invalid session name: {name}")
+
+        if description and not validate_description(description):
+            raise SessionError(
+                f"Session description too long: {len(description)} characters"
+            )
+
+        # Sanitize inputs
+        if name:
+            name = sanitize_input(name, max_length=100)
+        if description:
+            description = sanitize_input(description, max_length=500)
+        if namespace:
+            namespace = sanitize_input(namespace, max_length=100)
+
+        # Generate session ID if not provided
+        if not session_id:
+            session_id = generate_session_id()
+
+        try:
+            self.connection_manager.ensure_connected()
+            created_session_id = self._create_session_on_server(
+                session_id=session_id,
+                user_id=user_id,
+                name=name,
+                description=description,
+                namespace=namespace,
+                api_key=api_key,
+            )
+
+            # Store session info locally
+            with self.sessions_lock:
+                self.sessions[created_session_id] = {
+                    "id": created_session_id,
+                    "user_id": user_id,
+                    "name": name,
+                    "description": description,
+                    "namespace": namespace,
+                    "api_key": "",
+                    "created_at": threading.current_thread().ident,  # Basic tracking
+                }
+                self._owned_sessions.add(created_session_id)
+
+            return created_session_id
+
+        except Exception as e:
+            raise SessionError(f"Failed to create session: {e}")
+
+    @abstractmethod
+    def _create_session_on_server(
+        self,
+        session_id: str,
+        user_id: Optional[str],
+        name: Optional[str],
+        description: Optional[str],
+        namespace: Optional[str],
+        api_key: Optional[str],
+    ) -> str:
+        """Create session on server (protocol-specific)."""
+        pass
+
+    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """Get session information."""
+        try:
+            self.connection_manager.ensure_connected()
+
+            # Try to get from local cache first
+            with self.sessions_lock:
+                local_session = self.sessions.get(session_id)
+                if local_session:
+                    return local_session
+
+            # Get from server
+            server_session = self._get_session_from_server(session_id)
+
+            # Cache locally
+            if server_session:
+                with self.sessions_lock:
+                    self.sessions[session_id] = server_session
+
+            return server_session
+
+        except Exception as e:
+            raise SessionError(f"Failed to get session {session_id}: {e}")
+
+    @abstractmethod
+    def _get_session_from_server(self, session_id: str) -> Optional[Dict[str, Any]]:
+        """Get session from server (protocol-specific)."""
+        pass
+
+    def delete_session(self, session_id: str):
+        """Delete a session."""
+        try:
+            self.connection_manager.ensure_connected()
+
+            # Delete from server
+            self._delete_session_on_server(session_id)
+
+            # Remove from local cache
+            with self.sessions_lock:
+                self.sessions.pop(session_id, None)
+                self._owned_sessions.discard(session_id)
+
+        except Exception as e:
+            raise SessionError(f"Failed to delete session {session_id}: {e}")
+
+    @abstractmethod
+    def _delete_session_on_server(self, session_id: str):
+        """Delete session on server (protocol-specific)."""
+        pass
+
+    def list_sessions(self) -> List[Dict[str, Any]]:
+        """List all sessions."""
+        try:
+            self.connection_manager.ensure_connected()
+            sessions = self._list_sessions_on_server()
+
+            # Update local cache
+            with self.sessions_lock:
+                for session in sessions:
+                    session_id = session.get("id")
+                    if session_id:
+                        self.sessions[session_id] = session
+
+            return sessions
+
+        except Exception as e:
+            raise SessionError(f"Failed to list sessions: {e}")
+
+    @abstractmethod
+    def _list_sessions_on_server(self) -> List[Dict[str, Any]]:
+        """List sessions on server (protocol-specific)."""
+        pass
+
+    def session_exists(self, session_id: str) -> bool:
+        """Check if a session exists."""
+        try:
+            session = self.get_session(session_id)
+            return session is not None
+        except Exception:
+            return False
+
+    def get_session_info(self, session_id: str) -> Dict[str, Any]:
+        """Get detailed session information."""
+        session = self.get_session(session_id)
+        if not session:
+            raise SessionError(f"Session {session_id} not found")
+
+        return {
+            "id": session.get("id"),
+            "name": session.get("name"),
+            "description": session.get("description"),
+            "namespace": session.get("namespace"),
+            "user_id": session.get("user_id"),
+            "created_at": session.get("created_at"),
+            "status": session.get("status", "active"),
+        }
+
+    def update_session(
+        self,
+        session_id: str,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        namespace: Optional[str] = None,
+    ):
+        """Update session information."""
+        # Validate inputs
+        if name and not validate_session_name(name):
+            raise SessionError(f"Invalid session name: {name}")
+
+        if description and not validate_description(description):
+            raise SessionError(
+                f"Session description too long: {len(description)} characters"
+            )
+
+        # Sanitize inputs
+        if name:
+            name = sanitize_input(name, max_length=100)
+        if description:
+            description = sanitize_input(description, max_length=500)
+        if namespace:
+            namespace = sanitize_input(namespace, max_length=100)
+
+        try:
+            self.connection_manager.ensure_connected()
+
+            # Update on server
+            updated_session = self._update_session_on_server(
+                session_id, name, description, namespace
+            )
+
+            # Update local cache
+            with self.sessions_lock:
+                if updated_session:
+                    self.sessions[session_id] = updated_session
+                elif session_id in self.sessions:
+                    if name:
+                        self.sessions[session_id]["name"] = name
+                    if description:
+                        self.sessions[session_id]["description"] = description
+                    if namespace:
+                        self.sessions[session_id]["namespace"] = namespace
+
+            return updated_session
+
+        except Exception as e:
+            raise SessionError(f"Failed to update session {session_id}: {e}")
+
+    @abstractmethod
+    def _update_session_on_server(
+        self,
+        session_id: str,
+        name: Optional[str],
+        description: Optional[str],
+        namespace: Optional[str],
+    ) -> Dict[str, Any]:
+        """Update session on server (protocol-specific)."""
+        pass
+
+    def cleanup_session(self, session_id: str):
+        """Cleanup a session and its resources."""
+        try:
+            # Delete session only if this client created it
+            owned = False
+            with self.sessions_lock:
+                owned = session_id in self._owned_sessions
+
+            if owned:
+                self.delete_session(session_id)
+
+            # Remove from local cache
+            with self.sessions_lock:
+                self.sessions.pop(session_id, None)
+                self._owned_sessions.discard(session_id)
+
+        except Exception:
+            # Log error but don't raise for cleanup
+            pass
+
+    def cleanup_all(self):
+        """Cleanup all sessions."""
+        with self.sessions_lock:
+            session_ids = list(self.sessions.keys())
+            for session_id in session_ids:
+                self.cleanup_session(session_id)
+
+    def get_session_count(self) -> int:
+        """Get the number of active sessions."""
+        with self.sessions_lock:
+            return len(self.sessions)
+
+    def get_session_stats(self) -> Dict[str, Any]:
+        """Get session statistics."""
+        with self.sessions_lock:
+            sessions = list(self.sessions.values())
+
+            stats = {
+                "total_sessions": len(sessions),
+                "session_ids": list(self.sessions.keys()),
+                "sessions_by_namespace": {},
+                "sessions_by_user": {},
+            }
+
+            for session in sessions:
+                namespace = session.get("namespace", "default")
+                user_id = session.get("user_id", "unknown")
+
+                stats["sessions_by_namespace"][namespace] = (
+                    stats["sessions_by_namespace"].get(namespace, 0) + 1
+                )
+                stats["sessions_by_user"][user_id] = (
+                    stats["sessions_by_user"].get(user_id, 0) + 1
+                )
+
+            return stats

toolplane/common/base_tool_manager.py

@@ -0,0 +1,347 @@
+"""Base tool manager class shared between gRPC and HTTP clients."""
+
+import logging
+from abc import ABC, abstractmethod
+from threading import RLock
+from typing import Any, Callable, Dict, List, Optional, Set
+
+try:
+    from toolplane.utils.schema import generate_schema_from_function
+except ImportError:
+    # Fallback for standalone testing
+    def generate_schema_from_function(func):
+        """Fallback schema generator for standalone testing."""
+        return {"type": "object", "properties": {}}
+
+
+try:
+    from ..core.errors import ToolError, ToolplaneAPIError
+except ImportError:
+    # Fallback for standalone testing
+    class ToolError(Exception):
+        """Tool error for standalone testing."""
+
+        pass
+
+    class ToolplaneAPIError(Exception):
+        """API error for standalone testing."""
+
+        pass
+
+
+from .constants import CACHE_TTL_SECONDS
+from .utils import (
+    cache_with_ttl,
+    sanitize_input,
+    validate_tool_name,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class BaseToolManager(ABC):
+    """Base class for tool management."""
+
+    def __init__(self, connection_manager):
+        """Initialize base tool manager."""
+        self.connection_manager = connection_manager
+        self.tools: Dict[str, Dict[str, Callable]] = (
+            {}
+        )  # session_id -> {tool_name: func}
+        self.tool_schemas: Dict[str, Dict[str, Dict]] = (
+            {}
+        )  # session_id -> {tool_name: schema}
+        self.streaming_tools: Dict[str, Set[str]] = {}  # session_id -> {tool_names}
+        self._tool_cache: Dict[str, tuple] = {}  # session_id -> (timestamp, tools)
+        self._lock = RLock()
+
+    def register_tool(
+        self,
+        session_id: str,
+        machine_id: str,
+        name: str,
+        func: Callable,
+        schema: Optional[Dict] = None,
+        description: Optional[str] = None,
+        stream: bool = False,
+        tags: Optional[List[str]] = None,
+    ):
+        """Register a tool for a session."""
+        # Validate inputs
+        if not validate_tool_name(name):
+            raise ToolError(f"Invalid tool name: {name}")
+
+        if tags is None:
+            tags = []
+
+        # Sanitize description
+        if description:
+            description = sanitize_input(description, max_length=500)
+
+        # Generate schema if not provided
+        if schema is None:
+            try:
+                schema = generate_schema_from_function(func)
+            except Exception as e:
+                raise ToolError(f"Failed to generate schema for tool {name}: {e}")
+
+        # Add description and tags to schema
+        if description:
+            schema["description"] = description
+        schema["tags"] = tags
+
+        # Store tool locally
+        with self._lock:
+            if session_id not in self.tools:
+                self.tools[session_id] = {}
+                self.tool_schemas[session_id] = {}
+                self.streaming_tools[session_id] = set()
+
+            self.tools[session_id][name] = func
+            self.tool_schemas[session_id][name] = schema
+
+            if stream:
+                self.streaming_tools[session_id].add(name)
+
+        # Register with server
+        self._register_tool_with_server(session_id, machine_id, name, schema)
+
+    @abstractmethod
+    def _register_tool_with_server(
+        self, session_id: str, machine_id: str, name: str, schema: Dict
+    ):
+        """Register tool with server (protocol-specific)."""
+        pass
+
+    def unregister_tool(self, session_id: str, name: str):
+        """Unregister a tool from a session."""
+        try:
+            # Remove from local storage
+            with self._lock:
+                if session_id in self.tools:
+                    self.tools[session_id].pop(name, None)
+                    self.tool_schemas[session_id].pop(name, None)
+                    self.streaming_tools[session_id].discard(name)
+
+            # Remove from server
+            self._unregister_tool_from_server(session_id, name)
+
+        except Exception as e:
+            raise ToolError(f"Failed to unregister tool {name}: {e}")
+
+    @abstractmethod
+    def _unregister_tool_from_server(self, session_id: str, name: str):
+        """Unregister tool from server (protocol-specific)."""
+        pass
+
+    def get_tool(self, session_id: str, name: str) -> Optional[Callable]:
+        """Get a tool function by name."""
+        with self._lock:
+            return self.tools.get(session_id, {}).get(name)
+
+    def is_streaming_tool(self, session_id: str, name: str) -> bool:
+        """Check if a tool is a streaming tool."""
+        with self._lock:
+            return name in self.streaming_tools.get(session_id, set())
+
+    def get_session_tools(self, session_id: str) -> Dict[str, Callable]:
+        """Get all tools for a session."""
+        with self._lock:
+            return self.tools.get(session_id, {}).copy()
+
+    def get_tool_schema(self, session_id: str, name: str) -> Optional[Dict]:
+        """Get tool schema by name."""
+        with self._lock:
+            return self.tool_schemas.get(session_id, {}).get(name)
+
+    def list_session_tools(self, session_id: str) -> List[str]:
+        """List all tool names for a session."""
+        with self._lock:
+            return list(self.tools.get(session_id, {}).keys())
+
+    @cache_with_ttl(CACHE_TTL_SECONDS)
+    def get_available_tools(self, session_id: str) -> Dict[str, Any]:
+        """Get available tools from server (cached)."""
+        try:
+            self.connection_manager.ensure_connected()
+            return self._get_available_tools_from_server(session_id)
+
+        except Exception as e:
+            raise ToolError(f"Failed to get available tools: {e}")
+
+    @abstractmethod
+    def _get_available_tools_from_server(self, session_id: str) -> Dict[str, Any]:
+        """Get available tools from server (protocol-specific)."""
+        pass
+
+    def execute_tool(
+        self,
+        session_id: str,
+        tool_name: str,
+        params: Dict,
+        idempotency_key: str = "",
+        timeout_seconds: int = 0,
+        wait_timeout_seconds: int = 0,
+    ):
+        """Execute a tool and return ``(request_id, terminal_status, result)``.
+
+        timeout_seconds overrides the request's absolute per-attempt timeout
+        on the wire (0 keeps the server default). wait_timeout_seconds asks
+        the server to long-poll until the request is terminal; then
+        terminal_status names the final state and result carries the parsed
+        result value for a successful run, and both are ``(None, None)``
+        when the call returned while the request is still in flight.
+        """
+        try:
+            # Validate tool name
+            if not validate_tool_name(tool_name):
+                raise ToolError(f"Invalid tool name: {tool_name}")
+
+            self.connection_manager.ensure_connected()
+            return self._execute_tool_on_server(
+                session_id,
+                tool_name,
+                params,
+                idempotency_key,
+                timeout_seconds,
+                wait_timeout_seconds,
+            )
+
+        except ToolplaneAPIError:
+            # Typed server errors keep their identity for the caller.
+            raise
+        except Exception as e:
+            raise ToolError(f"Failed to execute tool {tool_name}: {e}")
+
+    @abstractmethod
+    def _execute_tool_on_server(
+        self,
+        session_id: str,
+        tool_name: str,
+        params: Dict,
+        idempotency_key: str = "",
+        timeout_seconds: int = 0,
+    ) -> str:
+        """Execute tool on server (protocol-specific)."""
+        pass
+
+    def stream_tool(
+        self,
+        session_id: str,
+        tool_name: str,
+        params: Dict,
+        idempotency_key: str = "",
+        timeout_seconds: int = 0,
+    ):
+        """Stream tool execution."""
+        try:
+            # Validate tool name
+            if not validate_tool_name(tool_name):
+                raise ToolError(f"Invalid tool name: {tool_name}")
+
+            self.connection_manager.ensure_connected()
+            yield from self._stream_tool_on_server(
+                session_id, tool_name, params, idempotency_key, timeout_seconds
+            )
+
+        except ToolplaneAPIError:
+            raise
+        except Exception as e:
+            raise ToolError(f"Failed to stream tool {tool_name}: {e}")
+
+    @abstractmethod
+    def _stream_tool_on_server(
+        self,
+        session_id: str,
+        tool_name: str,
+        params: Dict,
+        idempotency_key: str = "",
+        timeout_seconds: int = 0,
+    ):
+        """Stream tool execution on server (protocol-specific)."""
+        pass
+
+    def cleanup_session_tools(self, session_id: str):
+        """Cleanup all tools for a session."""
+        with self._lock:
+            if session_id in self.tools:
+                tool_names = list(self.tools[session_id].keys())
+                for tool_name in tool_names:
+                    try:
+                        self.unregister_tool(session_id, tool_name)
+                    except Exception:
+                        # Ignore cleanup errors
+                        pass
+
+                # Clear local storage
+                self.tools.pop(session_id, None)
+                self.tool_schemas.pop(session_id, None)
+                self.streaming_tools.pop(session_id, None)
+
+    def cleanup_all(self):
+        """Cleanup all tools."""
+        with self._lock:
+            session_ids = list(self.tools.keys())
+            for session_id in session_ids:
+                self.cleanup_session_tools(session_id)
+
+    def get_tool_stats(self, session_id: str) -> Dict[str, Any]:
+        """Get tool statistics for a session."""
+        with self._lock:
+            session_tools = self.tools.get(session_id, {})
+            streaming_tools = self.streaming_tools.get(session_id, set())
+
+            return {
+                "total_tools": len(session_tools),
+                "streaming_tools": len(streaming_tools),
+                "regular_tools": len(session_tools) - len(streaming_tools),
+                "tool_names": list(session_tools.keys()),
+                "streaming_tool_names": list(streaming_tools),
+            }
+
+    def reregister_session_tools(self, session_id: str, machine_id: Optional[str]):
+        """Re-register all tools for a session after reconnecting."""
+        if not machine_id:
+            return
+
+        with self._lock:
+            schemas = list(self.tool_schemas.get(session_id, {}).items())
+
+        for tool_name, schema in schemas:
+            try:
+                self._register_tool_with_server(
+                    session_id, machine_id, tool_name, schema
+                )
+            except Exception as exc:
+                # Registration is an upsert: same-machine re-register never
+                # conflicts, so every failure here is real (including a name
+                # owned by another live machine, which the server reports as
+                # FAILED_PRECONDITION) and surfaces as a warning.
+                logger.warning(
+                    "Failed to re-register tool %s for session %s: %s",
+                    tool_name,
+                    session_id,
+                    exc,
+                )
+
+    def validate_tool_params(
+        self, session_id: str, tool_name: str, params: Dict
+    ) -> bool:
+        """Validate tool parameters against schema."""
+        schema = self.get_tool_schema(session_id, tool_name)
+        if not schema:
+            return True  # No schema to validate against
+
+        try:
+            # Basic validation - can be extended with jsonschema
+            tool_schema = schema.get("schema", {})
+            if "properties" in tool_schema:
+                required_fields = tool_schema.get("required", [])
+                for field in required_fields:
+                    if field not in params:
+                        return False
+
+            return True
+
+        except Exception:
+            return False