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

toolplane/core/session_context.py

@@ -0,0 +1,514 @@
+"""Session context implementation."""
+
+import asyncio
+import logging
+import time
+import uuid
+from typing import Any, Callable, Dict, List, Optional
+
+from toolplane.utils.schema import generate_schema_from_function
+
+from ..common.constants import WAIT_TIMEOUT_MAX_SECONDS
+from .connection import ConnectionManager
+from .errors import (
+    ToolplaneAPIError,
+    ToolplaneCancelledError,
+    ToolplaneError,
+    ToolplaneInvalidArgumentError,
+    ToolplaneTimeoutError,
+    normalize_status_name,
+)
+from .machine import MachineManager
+from .request import RequestManager
+from .session import SessionManager
+from .tool import ToolManager
+
+logger = logging.getLogger(__name__)
+
+
+def derive_wait_for(timeout_seconds: int, wait_timeout: Optional[int]) -> int:
+    """Resolve the local wait budget that also drives the server long-poll.
+
+    Derived waits (``timeout_seconds + 15``, else 60) clamp to the server's
+    wait ceiling so a legal timeout at the ceiling is not rejected. An
+    explicit wait passes through untouched: the server rejects over-max
+    values loudly (OUT_OF_RANGE) instead of the SDK silently shrinking a
+    caller's chosen budget.
+    """
+    if wait_timeout is not None:
+        return wait_timeout
+    derived = timeout_seconds + 15 if timeout_seconds > 0 else 60
+    return min(derived, WAIT_TIMEOUT_MAX_SECONDS)
+
+
+class SessionContext:
+    """
+    Represents a session context with its own tools and machine registration.
+    Encapsulates all session-specific state and operations.
+    """
+
+    def __init__(
+        self,
+        session_id: str,
+        connection_manager: ConnectionManager,
+        machine_manager: MachineManager,
+        tool_manager: ToolManager,
+        request_manager: RequestManager,
+        session_manager: SessionManager,
+    ):
+        """Initialize session context."""
+        self.session_id = session_id
+        self.connection_manager = connection_manager
+        self.machine_manager = machine_manager
+        self.tool_manager = tool_manager
+        self.request_manager = request_manager
+        self.session_manager = session_manager
+
+        self.machine_id: Optional[str] = None
+
+    def register_machine(self) -> bool:
+        """Register a machine for this session."""
+        try:
+            self.machine_id = self.machine_manager.register_machine(self.session_id)
+            return True
+        except Exception as e:
+            logger.warning(
+                "Error registering machine for session %s: %s", self.session_id, e
+            )
+            return False
+
+    def register_tool(
+        self,
+        name: str,
+        func: Callable,
+        schema: Optional[Dict] = None,
+        description: Optional[str] = None,
+        stream: bool = False,
+        tags: Optional[List[str]] = None,
+    ):
+        """Register a tool for this session."""
+        if not self.machine_id:
+            raise ToolplaneError(
+                f"Session {self.session_id} has no machine registration. Use ProviderRuntime.attach_session(...) or register_machine() before registering tools."
+            )
+
+        try:
+            self.tool_manager.register_tool(
+                self.session_id,
+                self.machine_id,
+                name,
+                func,
+                schema,
+                description,
+                stream,
+                tags,
+            )
+            logger.debug("Registered tool '%s' for session %s", name, self.session_id)
+        except Exception as e:
+            raise ToolplaneError(
+                f"Failed to register tool {name} for session {self.session_id}: {e}"
+            )
+
+    def invoke(
+        self,
+        tool_name: str,
+        timeout_seconds: int = 0,
+        wait_timeout: Optional[int] = None,
+        **params,
+    ) -> Any:
+        """Invoke a tool in this session and return the tool's result value.
+
+        timeout_seconds sets the request's absolute per-attempt execution
+        timeout on the wire (0 keeps the server default). The local wait ends
+        after wait_timeout seconds (default: timeout_seconds + 15, else 60).
+        """
+        try:
+            # Wait budget: also drives the server-side long-poll, so the
+            # common case returns terminal inside the submit call itself.
+            wait_for = derive_wait_for(timeout_seconds, wait_timeout)
+
+            deadline = time.monotonic() + wait_for
+
+            request_id, terminal_status, result = self.tool_manager.execute_tool(
+                self.session_id,
+                tool_name,
+                params,
+                timeout_seconds=timeout_seconds,
+                wait_timeout_seconds=wait_for,
+            )
+
+            if terminal_status == "done":
+                return result
+            if terminal_status == "cancelled":
+                raise ToolplaneCancelledError(
+                    f"Request was cancelled (request_id={request_id})"
+                )
+            if terminal_status == "failed":
+                raise ToolplaneError(f"Tool execution failed (request_id={request_id})")
+
+            # Still in flight after the server-side wait: poll for the
+            # remaining budget only, so the documented cap is a total.
+            remaining = max(1, int(deadline - time.monotonic()))
+            status = self._wait_for_completion(request_id, timeout=remaining)
+
+            # Unwrap: callers want the tool's return value, not the envelope.
+            return status.get("result")
+
+        except ToolplaneTimeoutError:
+            raise
+        except ToolplaneAPIError:
+            # Typed server errors keep their identity for the caller.
+            raise
+
+    async def ainvoke(
+        self,
+        tool_name: str,
+        timeout_seconds: int = 0,
+        **params,
+    ) -> str:
+        """Submit a tool invocation without blocking the caller.
+
+        Returns the request ID once the server accepts the work; poll
+        get_request_status (or await astream) for the outcome. The blocking
+        submission runs in a worker thread, so this is safe to await from a
+        running event loop.
+        """
+        try:
+            request_id, _, _ = await asyncio.to_thread(
+                self.tool_manager.execute_tool,
+                self.session_id,
+                tool_name,
+                params,
+                timeout_seconds=timeout_seconds,
+            )
+            return request_id
+        except Exception as e:
+            raise ToolplaneError(f"Failed to async invoke tool {tool_name}: {e}")
+
+    async def astream(
+        self, tool_name: str, callback: Callable[[Any, bool], None], **params
+    ):
+        """Awaitable stream: runs the blocking stream loop in a worker
+        thread and resolves with the collected chunks."""
+        return await asyncio.to_thread(self.stream, tool_name, callback, **params)
+
+    def stream(
+        self,
+        tool_name: str,
+        callback: Callable[[Any, bool], None],
+        timeout_seconds: int = 0,
+        **params,
+    ):
+        """Stream tool execution.
+
+        The stream is resumable: if the direct stream fails mid-flight, the
+        fallback resumes from the last received sequence number (or, when it
+        died before the first chunk, re-invokes under the same idempotency
+        key so the server returns the original request). A tool that
+        completed with an error is never re-executed.
+        """
+        idempotency_key = uuid.uuid4().hex
+        all_chunks = []
+        request_id = None
+        last_seq = 0
+
+        try:
+            try:
+                for chunk in self.tool_manager.stream_tool(
+                    self.session_id,
+                    tool_name,
+                    params,
+                    idempotency_key,
+                    timeout_seconds=timeout_seconds,
+                ):
+                    if chunk.request_id:
+                        request_id = chunk.request_id
+                    if chunk.seq:
+                        last_seq = max(last_seq, chunk.seq)
+                    callback(chunk.chunk, chunk.is_final)
+                    all_chunks.append(chunk.chunk)
+
+                    if chunk.error:
+                        raise ToolplaneError(f"Streaming error: {chunk.error}")
+
+                    if chunk.is_final:
+                        break
+
+                return all_chunks
+
+            except ToolplaneError:
+                # The tool itself failed (or the stream completed with an
+                # error marker): re-executing it would be a side-effect
+                # duplication, so surface the failure.
+                raise
+            except Exception:
+                # Transport failure mid-stream: resume from the last sequence
+                # number the server acknowledges.
+                if request_id:
+                    return self._resume_stream(
+                        request_id,
+                        last_seq,
+                        callback,
+                        all_chunks,
+                        tool_name,
+                        params,
+                        idempotency_key,
+                    )
+                # The stream died before any chunk arrived. The server may
+                # already have created the request; re-invoking under the same
+                # idempotency key returns that request instead of executing
+                # the tool a second time.
+                return self._stream_via_polling(
+                    tool_name,
+                    callback,
+                    params,
+                    idempotency_key,
+                    timeout_seconds=timeout_seconds,
+                )
+
+        except ToolplaneError:
+            raise
+        except Exception as e:
+            raise ToolplaneError(f"Failed to stream tool {tool_name}: {e}")
+
+    def _resume_stream(
+        self,
+        request_id: str,
+        last_seq: int,
+        callback: Callable,
+        all_chunks: List,
+        tool_name: str,
+        params: Dict,
+        idempotency_key: str,
+        timeout_seconds: int = 0,
+    ):
+        """Continue a broken stream from the last acknowledged sequence."""
+        try:
+            for chunk in self.request_manager.resume_stream(
+                self.session_id, request_id, last_seq
+            ):
+                value = chunk["chunk"]
+                if value not in ("", None):
+                    callback(value, chunk["is_final"])
+                    all_chunks.append(value)
+                if chunk["error"]:
+                    raise ToolplaneError(f"Streaming error: {chunk['error']}")
+                if chunk["is_final"]:
+                    return all_chunks
+        except ToolplaneInvalidArgumentError:
+            # The retained window moved past our position; the full result is
+            # still fetchable by polling the original request.
+            self._stream_via_polling(
+                tool_name,
+                callback,
+                params,
+                idempotency_key,
+                request_id=request_id,
+                skip=len(all_chunks),
+                accumulate=all_chunks,
+                timeout_seconds=timeout_seconds,
+            )
+        return all_chunks
+
+    def _stream_via_polling(
+        self,
+        tool_name: str,
+        callback: Callable,
+        params: Dict,
+        idempotency_key: str = "",
+        request_id: Optional[str] = None,
+        skip: int = 0,
+        accumulate: Optional[List] = None,
+        timeout_seconds: int = 0,
+    ):
+        """Stream via polling fallback.
+
+        When request_id is given, polls that request; otherwise invokes the
+        tool (under idempotency_key when provided) and polls the result.
+        skip suppresses the first skip chunks, which the caller already
+        delivered. When accumulate is given, polled chunks append to it so
+        callers keep the chunks they already delivered.
+        """
+        if request_id is None:
+            request_id = self.tool_manager.execute_tool(
+                self.session_id,
+                tool_name,
+                params,
+                idempotency_key,
+                timeout_seconds=timeout_seconds,
+            )[0]
+
+        all_chunks = accumulate if accumulate is not None else []
+        last_chunk_count = skip
+
+        while True:
+            status = self.get_request_status(request_id)
+
+            if "streamResults" in status:
+                chunks = status["streamResults"]
+
+                # Process new chunks
+                for i in range(last_chunk_count, len(chunks)):
+                    callback(chunks[i], False)
+                    all_chunks.append(chunks[i])
+
+                last_chunk_count = len(chunks)
+
+            if normalize_status_name(status["status"]) == "done":
+                callback("", True)
+                break
+
+            if normalize_status_name(status["status"]) == "cancelled":
+                raise ToolplaneCancelledError(
+                    "Request was cancelled "
+                    f"(request_id={status.get('requestId', request_id)})"
+                )
+
+            if normalize_status_name(status["status"]) == "failure":
+                raise ToolplaneError(
+                    f"Streaming failed: {status.get('error', 'Unknown error')}"
+                )
+
+            import time
+
+            time.sleep(0.5)
+
+        return all_chunks
+
+    def get_request_status(self, request_id: str) -> Dict[str, Any]:
+        """Get request status."""
+        return self.request_manager.get_request_status(self.session_id, request_id)
+
+    def get_available_tools(self) -> Dict[str, Any]:
+        """Get available tools for this session."""
+        return self.tool_manager.get_available_tools(self.session_id)
+
+    def list_tools(self) -> List[Dict[str, Any]]:
+        """List tools for this session."""
+        return self.tool_manager.list_tools(self.session_id)
+
+    def get_tool_by_id(self, tool_id: str) -> Dict[str, Any]:
+        """Get a tool by ID for this session."""
+        return self.tool_manager.get_tool_by_id(self.session_id, tool_id)
+
+    def get_tool_by_name(self, tool_name: str) -> Dict[str, Any]:
+        """Get a tool by name for this session."""
+        return self.tool_manager.get_tool_by_name(self.session_id, tool_name)
+
+    def delete_tool(self, tool_id: str) -> bool:
+        """Delete a tool by ID for this session."""
+        return self.tool_manager.delete_tool(self.session_id, tool_id)
+
+    def tool(self, name=None, description=None, stream=False, tags=None):
+        """Decorator for registering tools."""
+        if tags is None:
+            tags = []
+
+        def decorator(func):
+            tool_name = name or func.__name__
+            tool_schema = generate_schema_from_function(func)
+
+            if description:
+                tool_schema["description"] = description
+
+            self.register_tool(tool_name, func, tool_schema, description, stream, tags)
+            return func
+
+        return decorator
+
+    def _wait_for_completion(
+        self, request_id: str, timeout: int = 60
+    ) -> Dict[str, Any]:
+        """Wait for request completion and return the full response/status dict."""
+        import time
+
+        start_time = time.time()
+        last_status: Dict[str, Any] = {}
+
+        while time.time() - start_time < timeout:
+            status = self.get_request_status(request_id)
+            last_status = status
+
+            if normalize_status_name(status["status"]) == "done":
+                # Ensure result is JSON-parsed if possible (RequestManager already attempts this)
+                return status
+
+            if normalize_status_name(status["status"]) == "cancelled":
+                raise ToolplaneCancelledError(
+                    f"Request was cancelled (request_id={request_id})"
+                )
+
+            if normalize_status_name(status["status"]) == "failure":
+                raise ToolplaneError(
+                    f"Tool execution failed: {status.get('error', 'Unknown error')}"
+                )
+
+            time.sleep(0.5)
+
+        raise ToolplaneTimeoutError(
+            f"Tool execution timed out after {timeout}s (request_id={request_id}, "
+            f"status={last_status.get('status', 'unknown')})"
+        )
+
+    def cleanup(self):
+        """Cleanup this session."""
+        try:
+            # Cleanup tools
+            self.tool_manager.cleanup_session_tools(self.session_id)
+
+            # Unregister machine
+            if self.machine_id:
+                self.machine_manager.unregister_machine(
+                    self.session_id, reason="session_context_cleanup"
+                )
+
+            # Remove from session manager
+            if hasattr(self.session_manager, "cleanup_session_context"):
+                self.session_manager.cleanup_session_context(self.session_id)
+
+        except Exception as e:
+            logger.warning("Error cleaning up session %s: %s", self.session_id, e)
+
+    def poll_requests(self):
+        """Poll for requests in this session."""
+        if not self.machine_id:
+            return
+
+        try:
+            tools = self.tool_manager.get_session_tools(self.session_id)
+            streaming_tools = self.tool_manager.streaming_tools.get(
+                self.session_id, set()
+            )
+
+            self.request_manager.poll_session_requests(
+                self.session_id, self.machine_id, tools, streaming_tools
+            )
+        except Exception as e:
+            logger.warning(
+                "Error polling requests for session %s: %s", self.session_id, e
+            )
+
+    # New methods for user session management
+    def list_user_sessions(
+        self, user_id: str, page_size: int = 10, page_token: int = 0, filter: str = ""
+    ) -> Dict[str, Any]:
+        """List user sessions with pagination and filtering."""
+        return self.session_manager.list_user_sessions(
+            user_id, page_size, page_token, filter
+        )
+
+    def bulk_delete_sessions(
+        self, user_id: str, session_ids: Optional[List[str]] = None, filter: str = ""
+    ) -> Dict[str, Any]:
+        """Bulk delete sessions for a user."""
+        return self.session_manager.bulk_delete_sessions(
+            user_id, session_ids or [], filter
+        )
+
+    def get_session_stats(self, user_id: str) -> Dict[str, int]:
+        """Get session statistics for a user."""
+        return self.session_manager.get_session_stats(user_id)
+
+    def invalidate_session(self, reason: str = "") -> bool:
+        """Invalidate a session."""
+        return self.session_manager.invalidate_session(self.session_id, reason)

toolplane/core/task.py

@@ -0,0 +1,130 @@
+"""Task management for Toolplane gRPC client."""
+
+from typing import Any, Dict, List
+
+import grpc
+
+from toolplane.proto.service_pb2 import (
+    CancelTaskRequest,
+    CreateTaskRequest,
+    GetTaskRequest,
+    ListTasksRequest,
+    TaskStatus,
+)
+
+from ..common.utils import proto_enum_name, timestamp_to_iso
+from .connection import ConnectionManager
+from .errors import TaskError, api_error_from_rpc_error, normalize_status_name
+
+
+class TaskManager:
+    """Manages task lifecycle for the gRPC client."""
+
+    def __init__(self, connection_manager: ConnectionManager):
+        """Initialize task manager."""
+        self.connection_manager = connection_manager
+
+    def _normalize_task(self, task: Any) -> Dict[str, Any]:
+        return {
+            "id": task.id,
+            "session_id": task.session_id,
+            "tool_name": task.tool_name,
+            "status": normalize_status_name(proto_enum_name(task.status, TaskStatus)),
+            "input": task.input,
+            "result": task.result,
+            "result_type": task.result_type,
+            "error": task.error,
+            "created_at": timestamp_to_iso(task.created_at),
+            "updated_at": timestamp_to_iso(task.updated_at),
+            "completed_at": timestamp_to_iso(getattr(task, "completed_at", None)),
+        }
+
+    def create_task(
+        self,
+        session_id: str,
+        tool_name: str,
+        input_data: str,
+        idempotency_key: str = "",
+    ) -> Dict[str, Any]:
+        """Create a task for a session.
+
+        idempotency_key, when set, dedups creates within the session:
+        retrying with the same key returns the original task without
+        re-executing it.
+        """
+        self.connection_manager.ensure_connected()
+
+        try:
+            request = CreateTaskRequest(
+                session_id=session_id,
+                tool_name=tool_name,
+                input=input_data,
+                idempotency_key=idempotency_key,
+            )
+            response = self.connection_manager.tasks_stub.CreateTask(
+                request, metadata=self.connection_manager.get_metadata()
+            )
+            return self._normalize_task(response)
+        except grpc.RpcError as rpc_error:
+            raise api_error_from_rpc_error(
+                rpc_error, context=f"Failed to create task for session {session_id}"
+            ) from rpc_error
+        except Exception as exc:
+            raise TaskError(f"Failed to create task for session {session_id}: {exc}")
+
+    def get_task(self, session_id: str, task_id: str) -> Dict[str, Any]:
+        """Get a task for a session."""
+        self.connection_manager.ensure_connected()
+
+        try:
+            request = GetTaskRequest(session_id=session_id, task_id=task_id)
+            response = self.connection_manager.tasks_stub.GetTask(
+                request, metadata=self.connection_manager.get_metadata()
+            )
+            return self._normalize_task(response)
+        except grpc.RpcError as rpc_error:
+            raise api_error_from_rpc_error(
+                rpc_error,
+                context=f"Failed to get task {task_id} for session {session_id}",
+            ) from rpc_error
+        except Exception as exc:
+            raise TaskError(
+                f"Failed to get task {task_id} for session {session_id}: {exc}"
+            )
+
+    def list_tasks(self, session_id: str) -> List[Dict[str, Any]]:
+        """List tasks for a session."""
+        self.connection_manager.ensure_connected()
+
+        try:
+            request = ListTasksRequest(session_id=session_id)
+            response = self.connection_manager.tasks_stub.ListTasks(
+                request, metadata=self.connection_manager.get_metadata()
+            )
+            return [self._normalize_task(task) for task in response.tasks]
+        except grpc.RpcError as rpc_error:
+            raise api_error_from_rpc_error(
+                rpc_error, context=f"Failed to list tasks for session {session_id}"
+            ) from rpc_error
+        except Exception as exc:
+            raise TaskError(f"Failed to list tasks for session {session_id}: {exc}")
+
+    def cancel_task(self, session_id: str, task_id: str) -> bool:
+        """Cancel a task for a session."""
+        self.connection_manager.ensure_connected()
+
+        try:
+            request = CancelTaskRequest(session_id=session_id, task_id=task_id)
+            response = self.connection_manager.tasks_stub.CancelTask(
+                request, metadata=self.connection_manager.get_metadata()
+            )
+            return response.success
+        except grpc.RpcError as rpc_error:
+            raise api_error_from_rpc_error(
+                rpc_error,
+                context=f"Failed to cancel task {task_id} for session {session_id}",
+            ) from rpc_error
+        except Exception as exc:
+            raise TaskError(
+                f"Failed to cancel task {task_id} for session {session_id}: {exc}"
+            )