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

toolplane/core/connection.py

@@ -0,0 +1,285 @@
+"""Connection management for Toolplane client."""
+
+import logging
+import random
+import threading
+import time
+from pathlib import Path
+from typing import Optional
+
+import grpc
+
+from toolplane.proto.service_pb2_grpc import (
+    MachinesServiceStub,
+    RequestsServiceStub,
+    SessionsServiceStub,
+    TasksServiceStub,
+    ToolServiceStub,
+)
+
+from ..common.constants import (
+    GRPC_MAX_RECEIVE_MESSAGE_LENGTH,
+    GRPC_MAX_SEND_MESSAGE_LENGTH,
+)
+from .config import ClientConfig
+from .errors import ConnectionError, ToolplaneConnectionError
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectionManager:
+    """Manages gRPC connections and stubs."""
+
+    def __init__(self, config: ClientConfig):
+        """Initialize connection manager with configuration."""
+        self.config = config
+        self.channel: Optional[grpc.Channel] = None
+        self.tool_stub: Optional[ToolServiceStub] = None
+        self.session_stub: Optional[SessionsServiceStub] = None
+        self.machine_stub: Optional[MachinesServiceStub] = None
+        self.requests_stub: Optional[RequestsServiceStub] = None
+        self.tasks_stub: Optional[TasksServiceStub] = None
+        self.connected = False
+        self._connect_lock = threading.RLock()
+
+        # Retry configuration
+        self.max_retries = getattr(config, "max_retries", 3)
+        self.retry_base_delay = getattr(config, "retry_base_delay", 1.0)
+        self.retry_max_delay = getattr(config, "retry_max_delay", 60.0)
+        self.retry_backoff_factor = getattr(config, "retry_backoff_factor", 2.0)
+        self.connection_state = (
+            "disconnected"  # disconnected, connecting, connected, failed
+        )
+        self.consecutive_failures = 0
+
+    def _cleanup_failed_connection(self):
+        """Reset stubs and close any existing channel after a failed attempt."""
+        if self.channel is not None:
+            try:
+                self.channel.close()
+            except Exception as e:
+                logger.debug("Error closing channel: %s", e)
+        self._reset_stubs()
+        self.connected = False
+        self.connection_state = "failed"
+
+    def _calculate_delay(self, attempt: int) -> float:
+        """Calculate delay with exponential backoff and jitter."""
+        delay = min(
+            self.retry_base_delay * (self.retry_backoff_factor**attempt),
+            self.retry_max_delay,
+        )
+        # Add jitter to prevent thundering herd
+        jitter = random.uniform(0, 0.1 * delay)
+        return delay + jitter
+
+    def _should_retry(self, attempt: int, exception: Exception) -> bool:
+        """Determine if we should retry based on exception type and attempt count."""
+        if attempt >= self.max_retries:
+            return False
+
+        # Retry on specific gRPC status codes
+        if isinstance(exception, grpc.RpcError):
+            retryable_codes = [
+                grpc.StatusCode.UNAVAILABLE,
+                grpc.StatusCode.DEADLINE_EXCEEDED,
+                grpc.StatusCode.RESOURCE_EXHAUSTED,
+            ]
+            return exception.code() in retryable_codes
+
+        # Retry on network-related exceptions
+        if isinstance(
+            exception, (ToolplaneConnectionError, ConnectionError, TimeoutError)
+        ):
+            return True
+
+        return False
+
+    def _load_tls_bytes(self, path_value: Optional[str], label: str) -> Optional[bytes]:
+        """Load TLS material from disk when configured."""
+        if not path_value:
+            return None
+
+        try:
+            return Path(path_value).expanduser().read_bytes()
+        except OSError as exc:
+            raise ToolplaneConnectionError(f"Failed to read {label}: {exc}") from exc
+
+    def _channel_options(self) -> list[tuple[str, object]]:
+        """Build gRPC channel options: the message-size ladder plus TLS
+        overrides.
+
+        The size options lift the client above the gRPC 4 MiB receive
+        default, which rejects a full 8 MiB replay-window read at the
+        client before it reaches application code.
+        """
+        options: list[tuple[str, object]] = [
+            (
+                "grpc.max_receive_message_length",
+                GRPC_MAX_RECEIVE_MESSAGE_LENGTH,
+            ),
+            ("grpc.max_send_message_length", GRPC_MAX_SEND_MESSAGE_LENGTH),
+        ]
+        server_name = getattr(self.config, "tls_server_name", None)
+        if server_name:
+            options.extend(
+                [
+                    ("grpc.ssl_target_name_override", server_name),
+                    ("grpc.default_authority", server_name),
+                ]
+            )
+        return options
+
+    def _create_channel(self, target: str) -> grpc.Channel:
+        """Create a secure or insecure gRPC channel from client configuration."""
+        use_tls = bool(
+            getattr(self.config, "use_tls", False)
+            or getattr(self.config, "tls_ca_cert_path", None)
+            or getattr(self.config, "tls_server_name", None)
+            or getattr(self.config, "tls_cert_path", None)
+            or getattr(self.config, "tls_key_path", None)
+        )
+        options = self._channel_options()
+        if not use_tls:
+            return grpc.insecure_channel(target, options=options)
+
+        certificate_chain = self._load_tls_bytes(
+            getattr(self.config, "tls_cert_path", None),
+            "TLS client certificate",
+        )
+        private_key = self._load_tls_bytes(
+            getattr(self.config, "tls_key_path", None),
+            "TLS client key",
+        )
+        if bool(certificate_chain) != bool(private_key):
+            raise ToolplaneConnectionError(
+                "TLS client authentication requires both certificate and key files"
+            )
+
+        root_certificates = self._load_tls_bytes(
+            getattr(self.config, "tls_ca_cert_path", None),
+            "TLS CA certificate",
+        )
+        credentials = grpc.ssl_channel_credentials(
+            root_certificates=root_certificates,
+            private_key=private_key,
+            certificate_chain=certificate_chain,
+        )
+        return grpc.secure_channel(target, credentials, options=options)
+
+    def connect(self) -> bool:
+        """Establish connection to gRPC server with retry logic."""
+        with self._connect_lock:
+            if self.connected and self.channel is not None:
+                return True
+
+            last_exception = None
+            ready_timeout = max(1, min(self.config.request_timeout, 10))
+
+            for attempt in range(self.max_retries + 1):
+                try:
+                    if attempt > 0:
+                        delay = self._calculate_delay(attempt - 1)
+                        time.sleep(delay)
+                        logger.debug(
+                            "Retry attempt %d/%d after %.2fs delay",
+                            attempt,
+                            self.max_retries,
+                            delay,
+                        )
+
+                    self.connection_state = "connecting"
+                    target = self.config.get_server_address()
+                    self.channel = self._create_channel(target)
+
+                    # Wait until the channel reports ready or timeout occurs
+                    grpc.channel_ready_future(self.channel).result(
+                        timeout=ready_timeout
+                    )
+
+                    self._initialize_stubs()
+                    self.connected = True
+                    self.connection_state = "connected"
+                    self.consecutive_failures = 0
+                    logger.info("Successfully connected to %s", target)
+                    return True
+
+                except Exception as e:
+                    last_exception = e
+                    self.consecutive_failures += 1
+                    logger.warning("Connection attempt %d failed: %s", attempt + 1, e)
+                    self._cleanup_failed_connection()
+
+                    if attempt < self.max_retries and self._should_retry(attempt, e):
+                        continue
+                    break
+
+            raise ToolplaneConnectionError(
+                f"Failed to connect after {self.max_retries + 1} attempts. Last error: {last_exception}"
+            )
+
+    def disconnect(self):
+        """Close connection to gRPC server."""
+        with self._connect_lock:
+            if self.channel is not None:
+                try:
+                    self.channel.close()
+                except Exception as e:
+                    logger.debug("Error closing channel: %s", e)
+            self._reset_stubs()
+            self.connected = False
+            self.connection_state = "disconnected"
+
+    def _initialize_stubs(self):
+        """Initialize all gRPC service stubs."""
+        if self.channel is None:
+            raise ToolplaneConnectionError("Channel is not initialized")
+
+        self.tool_stub = ToolServiceStub(self.channel)
+        self.session_stub = SessionsServiceStub(self.channel)
+        self.machine_stub = MachinesServiceStub(self.channel)
+        self.requests_stub = RequestsServiceStub(self.channel)
+        self.tasks_stub = TasksServiceStub(self.channel)
+
+    def _reset_stubs(self):
+        """Reset all service stubs."""
+        self.tool_stub = None
+        self.session_stub = None
+        self.machine_stub = None
+        self.requests_stub = None
+        self.tasks_stub = None
+        self.channel = None
+
+    def ensure_connected(self):
+        """Ensure connection is established."""
+        if not self.connected or self.channel is None:
+            if not self.connect():
+                raise ToolplaneConnectionError("Failed to establish connection")
+
+    def get_metadata(self):
+        """Get metadata for gRPC calls."""
+        metadata = self.config.get_metadata()
+        # Per-machine credential for provide-scoped RPCs (set by the machine
+        # registration that minted it). Sent on every call; the server only
+        # consults it on provider operations.
+        if getattr(self, "machine_token", ""):
+            metadata = list(metadata) + [
+                ("x-toolplane-machine-token", self.machine_token)
+            ]
+        return metadata
+
+    def set_machine_token(self, token: str):
+        """Store the per-machine credential minted at registration."""
+        self.machine_token = token
+
+    def mark_unhealthy(self):
+        """Mark the current channel as unhealthy so the next call reconnects."""
+        with self._connect_lock:
+            if self.channel is not None:
+                try:
+                    self.channel.close()
+                except Exception as e:
+                    logger.debug("Error closing unhealthy channel: %s", e)
+            self._reset_stubs()
+            self.connected = False
+            self.connection_state = "disconnected"

toolplane/core/errors.py

@@ -0,0 +1,298 @@
+"""Error handling for Toolplane client."""
+
+import json
+from typing import Any, Optional
+
+from grpc import StatusCode
+
+
+class ToolplaneError(Exception):
+    """Base exception for Toolplane client errors."""
+
+    pass
+
+
+class ToolplaneConnectionError(ToolplaneError):
+    """Error related to gRPC connection.
+
+    Named to avoid shadowing the builtin ConnectionError: modules that
+    imported the old name silently lost the ability to catch (or retry on)
+    real socket connection errors.
+    """
+
+    pass
+
+
+# Backwards-compatible alias for the pre-1.0 name. New code should prefer
+# ToolplaneConnectionError.
+ConnectionError = ToolplaneConnectionError
+
+
+class ToolError(ToolplaneError):
+    """Error related to tool operations."""
+
+    pass
+
+
+class SessionError(ToolplaneError):
+    """Error related to session operations."""
+
+    pass
+
+
+class MachineError(ToolplaneError):
+    """Error related to machine operations."""
+
+    pass
+
+
+class RequestError(ToolplaneError):
+    """Error related to request operations."""
+
+    pass
+
+
+class TaskError(ToolplaneError):
+    """Error related to task operations."""
+
+    pass
+
+
+# gRPC codes where re-issuing the same call can plausibly succeed (the server
+# or a proxy was momentarily unavailable, or a capacity limit will clear).
+# Everything else — bad credentials, bad arguments, state conflicts,
+# missing entities — is deterministic: retrying re-runs the same failure.
+RETRYABLE_GRPC_CODES = frozenset(
+    {StatusCode.UNAVAILABLE, StatusCode.RESOURCE_EXHAUSTED}
+)
+
+# The JSON gateway translates gRPC codes to HTTP statuses and carries the
+# numeric gRPC code in the error body; this table covers responses whose body
+# is not that envelope (bare proxies, load balancers).
+_HTTP_STATUS_FALLBACK = {
+    400: StatusCode.INVALID_ARGUMENT,
+    401: StatusCode.UNAUTHENTICATED,
+    403: StatusCode.PERMISSION_DENIED,
+    404: StatusCode.NOT_FOUND,
+    409: StatusCode.ALREADY_EXISTS,
+    422: StatusCode.INVALID_ARGUMENT,
+    429: StatusCode.RESOURCE_EXHAUSTED,
+    500: StatusCode.INTERNAL,
+    502: StatusCode.UNAVAILABLE,
+    503: StatusCode.UNAVAILABLE,
+    504: StatusCode.DEADLINE_EXCEEDED,
+}
+
+# grpc.StatusCode members are keyed by their wire strings ("failed
+# precondition"), not by the protocol's numeric codes, so the numbers from
+# the gateway error envelope are mapped explicitly.
+_GRPC_NUMERIC_CODE = {
+    0: StatusCode.OK,
+    1: StatusCode.CANCELLED,
+    2: StatusCode.UNKNOWN,
+    3: StatusCode.INVALID_ARGUMENT,
+    4: StatusCode.DEADLINE_EXCEEDED,
+    5: StatusCode.NOT_FOUND,
+    6: StatusCode.ALREADY_EXISTS,
+    7: StatusCode.PERMISSION_DENIED,
+    8: StatusCode.RESOURCE_EXHAUSTED,
+    9: StatusCode.FAILED_PRECONDITION,
+    10: StatusCode.ABORTED,
+    11: StatusCode.OUT_OF_RANGE,
+    12: StatusCode.UNIMPLEMENTED,
+    13: StatusCode.INTERNAL,
+    14: StatusCode.UNAVAILABLE,
+    15: StatusCode.DATA_LOSS,
+    16: StatusCode.UNAUTHENTICATED,
+}
+
+
+class ToolplaneAPIError(ToolplaneError):
+    """A server-reported failure carrying a machine-readable status code.
+
+    Attributes:
+        code: gRPC status name, e.g. "NOT_FOUND" or "FAILED_PRECONDITION".
+        retryable: True when re-issuing the call can plausibly succeed
+            (UNAVAILABLE or RESOURCE_EXHAUSTED).
+        request_id: The request the failing call operated on, when known.
+        status: Last-known request status for poll-style calls, when known.
+        details: Raw transport error (gRPC trailing status or HTTP body).
+    """
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        code: str,
+        retryable: bool,
+        request_id: Optional[str] = None,
+        status: Optional[str] = None,
+        details: Any = None,
+    ):
+        super().__init__(message)
+        self.code = code
+        self.retryable = retryable
+        self.request_id = request_id
+        self.status = status
+        self.details = details
+
+
+class ToolplaneNotFoundError(ToolplaneAPIError):
+    """The targeted session, request, machine, tool, key, or task is missing."""
+
+
+class ToolplaneInvalidArgumentError(ToolplaneAPIError):
+    """The server rejected the request payload (INVALID_ARGUMENT)."""
+
+
+class ToolplaneFailedPreconditionError(ToolplaneAPIError):
+    """The operation conflicts with server state: a lost claim race, a
+    stale lease, a draining machine, a terminal request, or no provider
+    registered for the tool (FAILED_PRECONDITION)."""
+
+
+class ToolplaneResourceExhaustedError(ToolplaneAPIError):
+    """A capacity limit was hit (RESOURCE_EXHAUSTED); retryable with backoff."""
+
+
+class ToolplaneUnauthenticatedError(ToolplaneAPIError):
+    """The presented API key is missing, unknown, or revoked
+    (UNAUTHENTICATED). Retrying cannot succeed."""
+
+
+class ToolplanePermissionDeniedError(ToolplaneAPIError):
+    """The caller lacks the required capability, or the per-machine
+    credential check failed (PERMISSION_DENIED)."""
+
+
+class ToolplaneAlreadyExistsError(ToolplaneAPIError):
+    """The created entity already exists (ALREADY_EXISTS)."""
+
+
+class ToolplaneTimeoutError(ToolplaneAPIError):
+    """The call exceeded its deadline (DEADLINE_EXCEEDED)."""
+
+
+class ToolplaneUnavailableError(ToolplaneAPIError):
+    """The server or a proxy was momentarily unreachable (UNAVAILABLE);
+    retryable with backoff."""
+
+
+class ToolplaneCancelledError(ToolplaneAPIError):
+    """The call was cancelled before completing (CANCELLED)."""
+
+
+class ToolplaneInternalError(ToolplaneAPIError):
+    """An unexpected server-side failure, or an unrecognized status code."""
+
+
+_GRPC_CODE_TO_ERROR = {
+    StatusCode.NOT_FOUND: ToolplaneNotFoundError,
+    StatusCode.INVALID_ARGUMENT: ToolplaneInvalidArgumentError,
+    StatusCode.FAILED_PRECONDITION: ToolplaneFailedPreconditionError,
+    StatusCode.OUT_OF_RANGE: ToolplaneInvalidArgumentError,
+    StatusCode.RESOURCE_EXHAUSTED: ToolplaneResourceExhaustedError,
+    StatusCode.UNAUTHENTICATED: ToolplaneUnauthenticatedError,
+    StatusCode.PERMISSION_DENIED: ToolplanePermissionDeniedError,
+    StatusCode.ALREADY_EXISTS: ToolplaneAlreadyExistsError,
+    StatusCode.DEADLINE_EXCEEDED: ToolplaneTimeoutError,
+    StatusCode.UNAVAILABLE: ToolplaneUnavailableError,
+    StatusCode.CANCELLED: ToolplaneCancelledError,
+    StatusCode.INTERNAL: ToolplaneInternalError,
+    StatusCode.UNKNOWN: ToolplaneInternalError,
+    StatusCode.UNIMPLEMENTED: ToolplaneInternalError,
+    StatusCode.ABORTED: ToolplaneFailedPreconditionError,
+}
+
+
+def api_error_from_rpc_error(
+    rpc_error: Any,
+    context: str = "",
+    request_id: Optional[str] = None,
+    status: Optional[str] = None,
+) -> ToolplaneAPIError:
+    """Translate a grpc.RpcError into a typed ToolplaneAPIError subclass."""
+    code = rpc_error.code()
+    message = rpc_error.details() or str(rpc_error)
+    if context:
+        message = f"{context}: {message}"
+    error_class = _GRPC_CODE_TO_ERROR.get(code, ToolplaneInternalError)
+    return error_class(
+        message,
+        code=code.name,
+        retryable=code in RETRYABLE_GRPC_CODES,
+        request_id=request_id,
+        status=status,
+        details=rpc_error,
+    )
+
+
+def api_error_from_http_response(
+    http_status: int,
+    body: str,
+    context: str = "",
+    request_id: Optional[str] = None,
+) -> ToolplaneAPIError:
+    """Translate an HTTP error response into a typed ToolplaneAPIError.
+
+    The JSON gateway emits ``{"code": <grpc code number>, "message": ...}``
+    for unary endpoints and wraps the status in an ``{"error": {...}}``
+    frame for streaming endpoints; the numeric code is authoritative when
+    present in either shape. Otherwise the HTTP status is mapped through
+    _HTTP_STATUS_FALLBACK.
+    """
+    message = ""
+    code: Optional[StatusCode] = None
+    try:
+        parsed = json.loads(body) if body else None
+        if isinstance(parsed, dict):
+            envelope = parsed
+            nested = parsed.get("error")
+            if isinstance(nested, dict):
+                envelope = nested
+            message = str(envelope.get("message") or "")
+            raw_code = envelope.get("code")
+            if raw_code is None:
+                raw_code = parsed.get("code")
+            if isinstance(raw_code, int):
+                code = _GRPC_NUMERIC_CODE.get(raw_code)
+    except (ValueError, TypeError):
+        pass
+
+    if code is None:
+        code = _HTTP_STATUS_FALLBACK.get(http_status, StatusCode.UNKNOWN)
+    if not message:
+        message = f"HTTP {http_status} {body or ''}".strip()
+    if context:
+        message = f"{context}: {message}"
+
+    error_class = _GRPC_CODE_TO_ERROR.get(code, ToolplaneInternalError)
+    return error_class(
+        message,
+        code=code.name,
+        retryable=code in RETRYABLE_GRPC_CODES,
+        request_id=request_id,
+        details=body,
+    )
+
+
+def status_for_wire(name):
+    """Map a friendly status ("done", "pending") onto the v1 enum name
+    ("REQUEST_STATUS_DONE") used on requests with status filters. An empty
+    value maps to 0 (UNSPECIFIED), which gRPC accepts and means "no filter".
+    """
+    if not isinstance(name, str) or not name:
+        return 0
+    return "REQUEST_STATUS_" + name.upper()
+
+
+def normalize_status_name(status):
+    """Map a v1 enum status name (e.g. "REQUEST_STATUS_DONE") onto the
+    friendly lowercase form ("done") clients have always seen. Passes
+    through anything else unchanged."""
+    if not isinstance(status, str):
+        return status
+    for prefix in ("REQUEST_STATUS_", "TASK_STATUS_"):
+        if status.startswith(prefix):
+            return status[len(prefix) :].lower()
+    return status