blaxel 0.2.26rc119__py3-none-any.whl → 0.2.28__py3-none-any.whl

blaxel/core/common/webhook.py

@@ -0,0 +1,187 @@
+"""Webhook signature verification for async-sidecar callbacks."""
+
+import hashlib
+import hmac
+import time
+from typing import Protocol
+
+
+class RequestLike(Protocol):
+    """Protocol for request-like objects with body and headers."""
+
+    @property
+    def body(self) -> bytes:
+        """Raw request body as bytes."""
+        ...
+
+    @property
+    def headers(self) -> dict[str, str]:
+        """Request headers as dictionary."""
+        ...
+
+
+class AsyncSidecarCallback:
+    """Callback payload from async-sidecar."""
+
+    def __init__(
+        self,
+        status_code: int,
+        response_body: str,
+        response_length: int,
+        timestamp: int,
+    ):
+        self.status_code = status_code
+        self.response_body = response_body
+        self.response_length = response_length
+        self.timestamp = timestamp
+
+
+def verify_webhook_signature(
+    body: bytes | str,
+    signature: str,
+    secret: str,
+    timestamp: str | None = None,
+    max_age: int = 300,
+) -> bool:
+    """
+    Verify the HMAC-SHA256 signature of a webhook callback from async-sidecar.
+
+    Args:
+        body: The raw request body (bytes or string)
+        signature: The X-Blaxel-Signature header value (format: "sha256=<hex_digest>")
+        secret: The secret key used to sign the webhook (same as CALLBACK_SECRET in async-sidecar)
+        timestamp: Optional X-Blaxel-Timestamp header value for replay attack prevention
+        max_age: Maximum age of the webhook in seconds (default: 300 = 5 minutes)
+
+    Returns:
+        True if the signature is valid, False otherwise
+
+    Example:
+        ```python
+        from blaxel.core import verify_webhook_signature
+        from flask import Flask, request
+
+        app = Flask(__name__)
+
+        @app.route('/webhook', methods=['POST'])
+        def webhook():
+            is_valid = verify_webhook_signature(
+                body=request.get_data(),
+                signature=request.headers.get('X-Blaxel-Signature', ''),
+                secret='your-callback-secret'
+            )
+
+            if not is_valid:
+                return {'error': 'Invalid signature'}, 401
+
+            data = request.json
+            # Process callback...
+            return {'received': True}
+        ```
+    """
+    if not body or not signature or not secret:
+        return False
+
+    try:
+        # Verify timestamp if provided (prevents replay attacks)
+        if timestamp:
+            request_time = int(timestamp)
+            current_time = int(time.time())
+            age = abs(current_time - request_time)
+
+            if age > max_age:
+                return False
+
+        # Convert body to bytes if string
+        body_bytes = body.encode("utf-8") if isinstance(body, str) else body
+
+        # Extract hex signature from "sha256=<hex>" format
+        expected_signature = signature.replace("sha256=", "")
+
+        # Compute HMAC-SHA256 signature
+        computed_signature = hmac.new(
+            secret.encode("utf-8"), body_bytes, hashlib.sha256
+        ).hexdigest()
+
+        # Timing-safe comparison to prevent timing attacks
+        return hmac.compare_digest(expected_signature, computed_signature)
+
+    except (ValueError, TypeError):
+        # Invalid signature format or other error
+        return False
+
+
+def verify_webhook_from_request(
+    request: RequestLike,
+    secret: str,
+    max_age: int = 300,
+) -> bool:
+    """
+    Helper to verify webhook from a request object (Flask, FastAPI, etc.).
+
+    Args:
+        request: Request object with `body` and `headers` attributes
+        secret: The callback secret
+        max_age: Optional maximum age in seconds (default: 300)
+
+    Returns:
+        True if the signature is valid, False otherwise
+
+    Example with Flask:
+        ```python
+        from blaxel.core import verify_webhook_from_request
+        from flask import Flask, request
+
+        app = Flask(__name__)
+
+        @app.route('/webhook', methods=['POST'])
+        def webhook():
+            if not verify_webhook_from_request(request, 'your-callback-secret'):
+                return {'error': 'Invalid signature'}, 401
+
+            data = request.json
+            print(f"Received callback: {data}")
+            return {'received': True}
+        ```
+
+    Example with FastAPI:
+        ```python
+        from blaxel.core import verify_webhook_signature
+        from fastapi import FastAPI, Request, HTTPException
+
+        app = FastAPI()
+
+        @app.post('/webhook')
+        async def webhook(request: Request):
+            body = await request.body()
+            signature = request.headers.get('x-blaxel-signature', '')
+
+            if not verify_webhook_signature(body, signature, 'your-callback-secret'):
+                raise HTTPException(status_code=401, detail='Invalid signature')
+
+            data = await request.json()
+            return {'received': True}
+        ```
+    """
+    signature = request.headers.get("x-blaxel-signature", "")
+    timestamp = request.headers.get("x-blaxel-timestamp")
+
+    if not signature:
+        return False
+
+    return verify_webhook_signature(
+        body=request.body,
+        signature=signature,
+        secret=secret,
+        timestamp=timestamp,
+        max_age=max_age,
+    )
+
+
+__all__ = [
+    "verify_webhook_signature",
+    "verify_webhook_from_request",
+    "AsyncSidecarCallback",
+    "RequestLike",
+]
+

blaxel/core/jobs/__init__.py

@@ -1,13 +1,21 @@
 import argparse
 import asyncio
 import os
-import sys
+import time
 from logging import getLogger
-from typing import Any, Awaitable, Callable, Dict
+from typing import Any, Callable, Dict, List
 
 import requests
 
 from ..client import client
+from ..client.api.jobs import (
+    create_job_execution,
+    delete_job_execution,
+    get_job_execution,
+    list_job_executions,
+)
+from ..client.models.create_job_execution_request import CreateJobExecutionRequest
+from ..client.models.job_execution import JobExecution
 from ..common.internal import get_forced_url, get_global_unique_hash
 from ..common.settings import settings
 
@@ -27,7 +35,7 @@ class BlJobWrapper:
                     args_dict[key] = unknown[i + 1]
             return args_dict
 
-        response = requests.get(os.getenv("BL_EXECUTION_DATA_URL"))
+        response = requests.get(os.getenv("BL_EXECUTION_DATA_URL") or "")
         data = response.json()
         tasks = data.get("tasks", [])
         return tasks[self.index] if self.index < len(tasks) else {}
@@ -54,8 +62,7 @@ class BlJobWrapper:
             else:
                 func(**parsed_args)
         except Exception as error:
-            print("Job execution failed:", error, file=sys.stderr)
-            sys.exit(1)
+            logger.error(f"Job execution failed: {error}")
 
 
 logger = getLogger(__name__)
@@ -96,7 +103,7 @@ class BlJob:
 
     def call(self, url, input_data, headers: dict = {}, params: dict = {}):
         body = {"tasks": input_data}
-        
+
         # Merge settings headers with provided headers
         merged_headers = {**settings.headers, "Content-Type": "application/json", **headers}
 
@@ -110,7 +117,7 @@ class BlJob:
     async def acall(self, url, input_data, headers: dict = {}, params: dict = {}):
         logger.debug(f"Job Calling: {self.name}")
         body = {"tasks": input_data}
-        
+
         # Merge settings headers with provided headers
         merged_headers = {**settings.headers, "Content-Type": "application/json", **headers}
 
@@ -136,7 +143,7 @@ class BlJob:
                 )
         return response.text
 
-    async def arun(self, input: Any, headers: dict = {}, params: dict = {}) -> Awaitable[str]:
+    async def arun(self, input: Any, headers: dict = {}, params: dict = {}) -> str:
         logger.debug(f"Job Calling: {self.name}")
         response = await self.acall(self.url, input, headers, params)
         if response.status_code >= 400:
@@ -151,6 +158,346 @@ class BlJob:
                 )
         return response.text
 
+    def create_execution(self, request: CreateJobExecutionRequest) -> str:
+        """
+        Create a new execution for this job and return the execution ID.
+
+        Args:
+            request: The job execution request containing tasks and optional execution ID
+
+        Returns:
+            str: The execution ID
+
+        Raises:
+            Exception: If no execution ID is returned or the request fails
+        """
+        logger.debug(f"Creating execution for job: {self.name}")
+
+        response = create_job_execution.sync_detailed(
+            job_id=self.name,
+            client=client,
+            body=request,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to create job execution: {response.status_code}")
+
+        # The API returns executionId at the root level
+        if response.parsed and hasattr(response.parsed, "to_dict"):
+            response_dict = response.parsed.to_dict()
+        else:
+            # Parse the raw response to get executionId
+            import json
+
+            response_dict = json.loads(response.content)
+
+        # Check for both camelCase (API) and snake_case (parsed) formats
+        execution_id = response_dict.get("execution_id") or response_dict.get("executionId")
+        if not execution_id:
+            raise Exception("No execution ID returned from create job execution")
+
+        logger.debug(f"Created execution: {execution_id}")
+        return execution_id
+
+    async def acreate_execution(self, request: CreateJobExecutionRequest) -> str:
+        """
+        Create a new execution for this job and return the execution ID (async).
+
+        Args:
+            request: The job execution request containing tasks and optional execution ID
+
+        Returns:
+            str: The execution ID
+
+        Raises:
+            Exception: If no execution ID is returned or the request fails
+        """
+        logger.debug(f"Creating execution for job: {self.name}")
+
+        response = await create_job_execution.asyncio_detailed(
+            job_id=self.name,
+            client=client,
+            body=request,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to create job execution: {response.status_code}")
+
+        # The API returns executionId at the root level
+        if response.parsed and hasattr(response.parsed, "to_dict"):
+            response_dict = response.parsed.to_dict()
+        else:
+            # Parse the raw response to get executionId
+            import json
+
+            response_dict = json.loads(response.content)
+
+        # Check for both camelCase (API) and snake_case (parsed) formats
+        execution_id = response_dict.get("execution_id") or response_dict.get("executionId")
+        if not execution_id:
+            raise Exception("No execution ID returned from create job execution")
+
+        logger.debug(f"Created execution: {execution_id}")
+        return execution_id
+
+    def get_execution(self, execution_id: str) -> JobExecution:
+        """
+        Get a specific execution by ID.
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            JobExecution: The execution object
+
+        Raises:
+            Exception: If the execution is not found or the request fails
+        """
+        logger.debug(f"Getting execution {execution_id} for job: {self.name}")
+
+        response = get_job_execution.sync_detailed(
+            job_id=self.name,
+            execution_id=execution_id,
+            client=client,
+        )
+
+        if response.status_code == 404:
+            raise Exception(f"Execution '{execution_id}' not found for job '{self.name}'")
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to get job execution: {response.status_code}")
+
+        if not response.parsed:
+            raise Exception("No execution data returned")
+
+        return response.parsed
+
+    async def aget_execution(self, execution_id: str) -> JobExecution:
+        """
+        Get a specific execution by ID (async).
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            JobExecution: The execution object
+
+        Raises:
+            Exception: If the execution is not found or the request fails
+        """
+        logger.debug(f"Getting execution {execution_id} for job: {self.name}")
+
+        response = await get_job_execution.asyncio_detailed(
+            job_id=self.name,
+            execution_id=execution_id,
+            client=client,
+        )
+
+        if response.status_code == 404:
+            raise Exception(f"Execution '{execution_id}' not found for job '{self.name}'")
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to get job execution: {response.status_code}")
+
+        if not response.parsed:
+            raise Exception("No execution data returned")
+
+        return response.parsed
+
+    def list_executions(self, limit: int = 20, offset: int = 0) -> List[JobExecution]:
+        """
+        List all executions for this job.
+
+        Args:
+            limit: Maximum number of executions to return
+            offset: Offset for pagination
+
+        Returns:
+            List[JobExecution]: List of execution objects
+        """
+        logger.debug(f"Listing executions for job: {self.name}")
+
+        response = list_job_executions.sync_detailed(
+            job_id=self.name,
+            client=client,
+            limit=limit,
+            offset=offset,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to list job executions: {response.status_code}")
+
+        return response.parsed or []
+
+    async def alist_executions(self, limit: int = 20, offset: int = 0) -> List[JobExecution]:
+        """
+        List all executions for this job (async).
+
+        Args:
+            limit: Maximum number of executions to return
+            offset: Offset for pagination
+
+        Returns:
+            List[JobExecution]: List of execution objects
+        """
+        logger.debug(f"Listing executions for job: {self.name}")
+
+        response = await list_job_executions.asyncio_detailed(
+            job_id=self.name,
+            client=client,
+            limit=limit,
+            offset=offset,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to list job executions: {response.status_code}")
+
+        return response.parsed or []
+
+    def get_execution_status(self, execution_id: str) -> str:
+        """
+        Get the status of a specific execution.
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            str: The execution status
+        """
+        execution = self.get_execution(execution_id)
+        return execution.status if execution.status else "UNKNOWN"
+
+    async def aget_execution_status(self, execution_id: str) -> str:
+        """
+        Get the status of a specific execution (async).
+
+        Args:
+            execution_id: The execution ID
+
+        Returns:
+            str: The execution status
+        """
+        execution = await self.aget_execution(execution_id)
+        return execution.status if execution.status else "UNKNOWN"
+
+    def cancel_execution(self, execution_id: str) -> None:
+        """
+        Cancel a specific execution.
+
+        Args:
+            execution_id: The execution ID
+
+        Raises:
+            Exception: If the cancellation fails
+        """
+        logger.debug(f"Cancelling execution {execution_id} for job: {self.name}")
+
+        response = delete_job_execution.sync_detailed(
+            job_id=self.name,
+            execution_id=execution_id,
+            client=client,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to cancel job execution: {response.status_code}")
+
+    async def acancel_execution(self, execution_id: str) -> None:
+        """
+        Cancel a specific execution (async).
+
+        Args:
+            execution_id: The execution ID
+
+        Raises:
+            Exception: If the cancellation fails
+        """
+        logger.debug(f"Cancelling execution {execution_id} for job: {self.name}")
+
+        response = await delete_job_execution.asyncio_detailed(
+            job_id=self.name,
+            execution_id=execution_id,
+            client=client,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to cancel job execution: {response.status_code}")
+
+    def wait_for_execution(
+        self,
+        execution_id: str,
+        max_wait: int = 360,
+        interval: int = 3,
+    ) -> JobExecution:
+        """
+        Wait for an execution to complete.
+
+        Args:
+            execution_id: The execution ID to wait for
+            max_wait: Maximum time to wait in seconds (default: 360 = 6 minutes)
+            interval: Polling interval in seconds (default: 3 seconds)
+
+        Returns:
+            JobExecution: The completed execution
+
+        Raises:
+            Exception: If the execution doesn't complete within max_wait seconds
+        """
+        logger.debug(f"Waiting for execution {execution_id} to complete (max {max_wait}s)")
+
+        start_time = time.time()
+
+        while time.time() - start_time < max_wait:
+            execution = self.get_execution(execution_id)
+            status = execution.status
+
+            # Terminal states (Kubernetes-style: succeeded, failed, cancelled)
+            if status in ["succeeded", "failed", "cancelled"]:
+                logger.debug(f"Execution {execution_id} finished with status: {status}")
+                return execution
+
+            # Wait before polling again
+            time.sleep(interval)
+
+        raise Exception(f"Execution {execution_id} did not complete within {max_wait}s")
+
+    async def await_for_execution(
+        self,
+        execution_id: str,
+        max_wait: int = 360,
+        interval: int = 3,
+    ) -> JobExecution:
+        """
+        Wait for an execution to complete (async).
+
+        Args:
+            execution_id: The execution ID to wait for
+            max_wait: Maximum time to wait in seconds (default: 360 = 6 minutes)
+            interval: Polling interval in seconds (default: 3 seconds)
+
+        Returns:
+            JobExecution: The completed execution
+
+        Raises:
+            Exception: If the execution doesn't complete within max_wait seconds
+        """
+        logger.debug(f"Waiting for execution {execution_id} to complete (max {max_wait}s)")
+
+        start_time = time.time()
+
+        while time.time() - start_time < max_wait:
+            execution = await self.aget_execution(execution_id)
+            status = execution.status
+
+            # Terminal states (Kubernetes-style: succeeded, failed, cancelled)
+            if status in ["succeeded", "failed", "cancelled"]:
+                logger.debug(f"Execution {execution_id} finished with status: {status}")
+                return execution
+
+            # Wait before polling again
+            await asyncio.sleep(interval)
+
+        raise Exception(f"Execution {execution_id} did not complete within {max_wait}s")
+
     def __str__(self):
         return f"Job {self.name}"
 

blaxel/core/sandbox/client/models/__init__.py

@@ -2,8 +2,6 @@
 
 from .apply_edit_request import ApplyEditRequest
 from .apply_edit_response import ApplyEditResponse
-from .content_search_match import ContentSearchMatch
-from .content_search_response import ContentSearchResponse
 from .delete_network_process_pid_monitor_response_200 import (
     DeleteNetworkProcessPidMonitorResponse200,
 )
@@ -15,10 +13,6 @@ from .file_with_content import FileWithContent
 from .filesystem_multipart_upload import FilesystemMultipartUpload
 from .filesystem_multipart_upload_parts import FilesystemMultipartUploadParts
 from .filesystem_uploaded_part import FilesystemUploadedPart
-from .find_match import FindMatch
-from .find_response import FindResponse
-from .fuzzy_search_match import FuzzySearchMatch
-from .fuzzy_search_response import FuzzySearchResponse
 from .get_network_process_pid_ports_response_200 import GetNetworkProcessPidPortsResponse200
 from .multipart_complete_request import MultipartCompleteRequest
 from .multipart_initiate_request import MultipartInitiateRequest
@@ -39,15 +33,11 @@ from .ranked_file import RankedFile
 from .reranking_response import RerankingResponse
 from .subdirectory import Subdirectory
 from .success_response import SuccessResponse
-from .tree_request import TreeRequest
-from .tree_request_files import TreeRequestFiles
 from .welcome_response import WelcomeResponse
 
 __all__ = (
     "ApplyEditRequest",
     "ApplyEditResponse",
-    "ContentSearchMatch",
-    "ContentSearchResponse",
     "DeleteNetworkProcessPidMonitorResponse200",
     "Directory",
     "ErrorResponse",
@@ -57,10 +47,6 @@ __all__ = (
     "FilesystemMultipartUploadParts",
     "FilesystemUploadedPart",
     "FileWithContent",
-    "FindMatch",
-    "FindResponse",
-    "FuzzySearchMatch",
-    "FuzzySearchResponse",
     "GetNetworkProcessPidPortsResponse200",
     "MultipartCompleteRequest",
     "MultipartInitiateRequest",
@@ -81,7 +67,5 @@ __all__ = (
     "RerankingResponse",
     "Subdirectory",
     "SuccessResponse",
-    "TreeRequest",
-    "TreeRequestFiles",
     "WelcomeResponse",
 )