blaxel 0.2.36__py3-none-any.whl → 0.2.38__py3-none-any.whl

blaxel/core/sandbox/types.py

@@ -4,7 +4,14 @@ from typing import Any, Callable, Dict, List, TypeVar, Union
 import httpx
 from attrs import define as _attrs_define
 
-from ..client.models import Port, Sandbox, SandboxLifecycle, VolumeAttachment
+from ..client.models import (
+    Env,
+    Port,
+    PortProtocol,
+    Sandbox,
+    SandboxLifecycle,
+    VolumeAttachment,
+)
 from ..client.types import UNSET
 from .client.models.process_request import ProcessRequest
 from .client.models.process_response import ProcessResponse
@@ -205,7 +212,7 @@ class SandboxCreateConfiguration:
             if isinstance(port, Port):
                 # If it's already a Port object, ensure protocol defaults to HTTP
                 if port.protocol is UNSET or not port.protocol:
-                    port.protocol = "HTTP"
+                    port.protocol = PortProtocol.HTTP
                 port_objects.append(port)
             elif isinstance(port, dict):
                 # Convert dict to Port object with HTTP as default protocol
@@ -218,20 +225,22 @@ class SandboxCreateConfiguration:
 
         return port_objects
 
-    def _normalize_envs(self) -> List[Dict[str, str]] | None:
+    def _normalize_envs(self) -> List[Env] | None:
         """Convert envs to list of dicts with name and value keys."""
         if not self.envs:
             return None
 
         env_objects = []
         for env in self.envs:
-            if isinstance(env, dict):
+            if isinstance(env, Env):
+                env_objects.append(env)
+            elif isinstance(env, dict):
                 # Validate that the dict has the required keys
                 if "name" not in env or "value" not in env:
                     raise ValueError(
                         f"Environment variable dict must have 'name' and 'value' keys: {env}"
                     )
-                env_objects.append({"name": env["name"], "value": env["value"]})
+                env_objects.append(Env(name=env["name"], value=env["value"]))
             else:
                 raise ValueError(
                     f"Invalid env type: {type(env)}. Expected dict with 'name' and 'value' keys."
@@ -385,3 +394,201 @@ class Context:
     @classmethod
     def from_json(cls, data: Dict[str, Any]) -> "Context":
         return cls(id=str(data.get("id") or data.get("context_id") or ""))
+
+
+class StreamHandle:
+    """Handle for managing a streaming operation (sync version).
+
+    Can be used as a context manager for automatic cleanup:
+
+        with sandbox.process.stream_logs(name, options) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.process.stream_logs(name, options)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the stream and stop receiving data."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the stream handle has been closed."""
+        return self._closed
+
+    def __enter__(self) -> "StreamHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class AsyncStreamHandle:
+    """Handle for managing a streaming operation (async version).
+
+    Can be used as an async context manager for automatic cleanup:
+
+        async with sandbox.process.stream_logs(name, options) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.process.stream_logs(name, options)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the stream and stop receiving data."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the stream handle has been closed."""
+        return self._closed
+
+    async def __aenter__(self) -> "AsyncStreamHandle":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Also support sync context manager for convenience
+    def __enter__(self) -> "AsyncStreamHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class WatchHandle:
+    """Handle for managing a file system watch operation (sync version).
+
+    Can be used as a context manager for automatic cleanup:
+
+        with sandbox.fs.watch(path, callback) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.fs.watch(path, callback)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the watch and stop receiving events."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the watch handle has been closed."""
+        return self._closed
+
+    def __enter__(self) -> "WatchHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)
+
+
+class AsyncWatchHandle:
+    """Handle for managing a file system watch operation (async version).
+
+    Can be used as an async context manager for automatic cleanup:
+
+        async with sandbox.fs.watch(path, callback) as handle:
+            # do something
+        # handle is automatically closed
+
+    Or used manually:
+
+        handle = sandbox.fs.watch(path, callback)
+        try:
+            # do something
+        finally:
+            handle.close()
+    """
+
+    def __init__(self, close_func: Callable[[], None]):
+        self._close_func = close_func
+        self._closed = False
+
+    def close(self) -> None:
+        """Close the watch and stop receiving events."""
+        if not self._closed:
+            self._close_func()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        """Returns True if the watch handle has been closed."""
+        return self._closed
+
+    async def __aenter__(self) -> "AsyncWatchHandle":
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Also support sync context manager for convenience
+    def __enter__(self) -> "AsyncWatchHandle":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+    # Backward compatibility: support dict-like access
+    def __getitem__(self, key: str) -> Callable[[], None]:
+        if key == "close":
+            return self.close
+        raise KeyError(key)

blaxel/core/tools/__init__.py

@@ -254,6 +254,10 @@ class PersistentMcpClient:
                 await self.client_exit_stack.aclose()
             except Exception as e:
                 logger.debug(f"Error closing client exit stack: {e}")
+            # Create fresh exit stacks so that future initialize() calls
+            # don't reuse stacks tainted by old cancel scopes
+            self.session_exit_stack = AsyncExitStack()
+            self.client_exit_stack = AsyncExitStack()
             logger.debug("WebSocket connection closed due to inactivity.")
 
 

blaxel/core/volume/volume.py

@@ -1,3 +1,5 @@
+import asyncio
+import time
 import uuid
 from typing import Callable, Dict, List, Union
 
@@ -135,6 +137,8 @@ class VolumeCreateConfiguration:
 
 
 class VolumeInstance:
+    delete: "_AsyncDeleteDescriptor"
+
     def __init__(self, volume: Volume):
         self.volume = volume
 
@@ -273,6 +277,8 @@ class VolumeInstance:
 
 
 class SyncVolumeInstance:
+    delete: "_SyncDeleteDescriptor"
+
     """Synchronous volume instance for managing persistent storage."""
 
     def __init__(self, volume: Volume):
@@ -500,6 +506,8 @@ async def _update_volume_by_name(
         status_code = int(response.code) if response.code is not UNSET else None
         message = response.message if response.message is not UNSET else response.error
         raise VolumeAPIError(message, status_code=status_code, code=response.error)
+    # This is for safe update
+    await asyncio.sleep(0.5)
     return VolumeInstance(response)
 
 
@@ -576,6 +584,8 @@ def _update_volume_by_name_sync(
         status_code = int(response.code) if response.code is not UNSET else None
         message = response.message if response.message is not UNSET else response.error
         raise VolumeAPIError(message, status_code=status_code, code=response.error)
+    # This is for safe update
+    time.sleep(0.5)
     return SyncVolumeInstance(response)
 
 

blaxel/crewai/model.py

@@ -1,6 +1,8 @@
 from logging import getLogger
 
-from crewai import LLM
+import httpx
+from crewai import LLM  # type: ignore[import-not-found]
+from crewai.llms.hooks.base import BaseInterceptor  # type: ignore[import-not-found]
 
 from blaxel.core import bl_model as bl_model_core
 from blaxel.core import settings
@@ -8,62 +10,97 @@ from blaxel.core import settings
 logger = getLogger(__name__)
 
 
-class AuthenticatedLLM(LLM):
-    def call(self, *args, **kwargs):
-        self.additional_params["extra_headers"] = settings.auth.get_headers()
-        return super().call(*args, **kwargs)
+class AuthInterceptor(BaseInterceptor[httpx.Request, httpx.Response]):
+    """Interceptor that injects dynamic auth headers into every HTTP request.
+
+    Used for crewai native providers (OpenAI, Anthropic, Gemini, etc.)
+    where the LLM.__new__ factory returns a provider-specific instance
+    and subclass overrides are not possible.
+    """
+
+    def on_outbound(self, message: httpx.Request) -> httpx.Request:
+        auth_headers = settings.auth.get_headers()
+        # Remove the SDK's default "Authorization: Bearer replaced" header
+        # when our auth uses a different header (e.g. X-Blaxel-Authorization with API keys)
+        if "Authorization" not in auth_headers:
+            message.headers.pop("Authorization", None)
+            message.headers.pop("authorization", None)
+        for key, value in auth_headers.items():
+            message.headers[key] = value
+        return message
+
+    def on_inbound(self, message: httpx.Response) -> httpx.Response:
+        return message
+
+    async def aon_outbound(self, message: httpx.Request) -> httpx.Request:
+        return self.on_outbound(message)
+
+    async def aon_inbound(self, message: httpx.Response) -> httpx.Response:
+        return message
+
+
+# Provider types that crewai routes to native SDK implementations.
+# These support the interceptor mechanism for auth.
+_NATIVE_PROVIDER_PREFIXES = {"openai", "anthropic", "gemini", "azure", "bedrock"}
+
+
+def _is_native_route(provider_prefix: str) -> bool:
+    """Check if a provider prefix will be routed to a native SDK by crewai."""
+    return provider_prefix.lower() in _NATIVE_PROVIDER_PREFIXES
 
 
 async def bl_model(name: str, **kwargs):
     url, type, model = await bl_model_core(name).get_parameters()
+
+    # Map blaxel model types to crewai provider prefixes and base URLs
     if type == "mistral":
-        return AuthenticatedLLM(
-            model=f"mistral/{model}",
-            api_key="replaced",
-            base_url=f"{url}/v1",
-            **kwargs,
-        )
+        provider_prefix = "mistral"
+        base_url = f"{url}/v1"
     elif type == "xai":
-        return AuthenticatedLLM(
-            model=f"groq/{model}",
-            api_key="replaced",
-            base_url=f"{url}/v1",
-            **kwargs,
-        )
+        provider_prefix = "groq"
+        base_url = f"{url}/v1"
     elif type == "deepseek":
-        return AuthenticatedLLM(
-            model=f"openai/{model}",
-            api_key="replaced",
-            base_url=f"{url}/v1",
-            **kwargs,
-        )
+        provider_prefix = "openai"
+        base_url = f"{url}/v1"
     elif type == "anthropic":
-        return AuthenticatedLLM(
-            model=f"anthropic/{model}",
-            api_key="replaced",
-            base_url=url,
-            **kwargs,
-        )
+        provider_prefix = "anthropic"
+        base_url = url
     elif type == "gemini":
-        return AuthenticatedLLM(
-            model=f"gemini/{model}",
-            api_key="replaced",
-            base_url=f"{url}/v1beta/models/{model}",
-            **kwargs,
-        )
+        provider_prefix = "gemini"
+        base_url = f"{url}/v1beta/models/{model}"
     elif type == "cerebras":
-        return AuthenticatedLLM(
-            model=f"cerebras/{model}",
-            api_key="replaced",
-            base_url=f"{url}/v1",
-            **kwargs,
-        )
+        provider_prefix = "cerebras"
+        base_url = f"{url}/v1"
     else:
         if type != "openai":
             logger.warning(f"Model {model} is not supported by CrewAI, defaulting to OpenAI")
-        return AuthenticatedLLM(
-            model=f"openai/{model}",
+        provider_prefix = "openai"
+        base_url = f"{url}/v1"
+
+    model_string = f"{provider_prefix}/{model}"
+    auth_headers = settings.auth.get_headers()
+
+    if _is_native_route(provider_prefix):
+        # Native providers: use interceptor for dynamic auth headers.
+        # Always pass api_key="replaced" because crewai's native providers
+        # require a non-None api_key. The AuthInterceptor handles stripping
+        # the dummy Authorization header and injecting the real auth.
+        return LLM(
+            model=model_string,
             api_key="replaced",
-            base_url=f"{url}/v1",
+            base_url=base_url,
+            interceptor=AuthInterceptor(),
+            **kwargs,
+        )
+    else:
+        # LiteLLM fallback: pass auth headers via extra_headers param.
+        # Omit api_key when auth uses X-Blaxel-Authorization to prevent
+        # litellm from adding "Authorization: Bearer replaced".
+        llm_api_key = "replaced" if "Authorization" in auth_headers else None
+        return LLM(
+            model=model_string,
+            api_key=llm_api_key,
+            base_url=base_url,
+            extra_headers=auth_headers,
             **kwargs,
         )

blaxel/crewai/tools.py

@@ -1,10 +1,86 @@
+from typing import Any
+
 from crewai.tools import BaseTool
+from pydantic import BaseModel
 
 from blaxel.core.tools import bl_tools as bl_tools_core
-from blaxel.core.tools.common import create_model_from_json_schema
 from blaxel.core.tools.types import Tool
 
 
+def _clean_schema_for_openai(schema: dict) -> dict:
+    """Clean JSON schema to be compatible with OpenAI strict mode.
+
+    Recursively resolves anyOf patterns, ensures all schemas have type keys,
+    removes additionalProperties and $schema, and ensures object types have
+    properties and required fields.
+    """
+    if not isinstance(schema, dict):
+        return schema
+
+    cleaned = schema.copy()
+
+    # Remove unsupported keys
+    cleaned.pop("$schema", None)
+    cleaned.pop("additionalProperties", None)
+
+    # Resolve anyOf: pick the non-null type
+    if "anyOf" in cleaned:
+        any_of = cleaned.pop("anyOf")
+        non_null = [s for s in any_of if s.get("type") != "null"]
+        if non_null:
+            # Merge the first non-null variant into current schema
+            resolved = _clean_schema_for_openai(non_null[0])
+            cleaned.update(resolved)
+        else:
+            cleaned["type"] = "string"
+
+    # Ensure type exists
+    if "type" not in cleaned and "properties" in cleaned:
+        cleaned["type"] = "object"
+
+    # Handle object types
+    if cleaned.get("type") == "object":
+        if "properties" not in cleaned:
+            cleaned["properties"] = {}
+        if "required" not in cleaned:
+            cleaned["required"] = list(cleaned["properties"].keys())
+
+    # Recursively clean properties
+    if "properties" in cleaned:
+        cleaned["properties"] = {
+            k: _clean_schema_for_openai(v) for k, v in cleaned["properties"].items()
+        }
+
+    # Recursively clean array items
+    if "items" in cleaned:
+        cleaned["items"] = _clean_schema_for_openai(cleaned["items"])
+        # Ensure items has a type
+        if "type" not in cleaned["items"]:
+            cleaned["items"]["type"] = "string"
+
+    return cleaned
+
+
+def _make_clean_args_schema(tool: Tool) -> type[BaseModel]:
+    """Create a Pydantic model whose JSON schema returns the pre-cleaned schema.
+
+    CrewAI calls model_json_schema() on args_schema to build the OpenAI tool
+    parameters. By overriding model_json_schema we ensure the cleaned schema
+    is used directly, avoiding issues with Pydantic re-introducing anyOf or
+    dropping type keys on array items.
+    """
+    clean = _clean_schema_for_openai(tool.input_schema)
+
+    class CleanArgsSchema(BaseModel):
+        @classmethod
+        def model_json_schema(cls, *args: Any, **kwargs: Any) -> dict[str, Any]:
+            return clean
+
+    CleanArgsSchema.__name__ = f"{tool.name}Schema"
+    CleanArgsSchema.__qualname__ = f"{tool.name}Schema"
+    return CleanArgsSchema
+
+
 class CrewAITool(BaseTool):
     _tool: Tool
 
@@ -12,13 +88,20 @@ class CrewAITool(BaseTool):
         super().__init__(
             name=tool.name,
             description=tool.description,
-            args_schema=create_model_from_json_schema(tool.input_schema),
+            args_schema=_make_clean_args_schema(tool),
         )
         self._tool = tool
 
     def _run(self, *args, **kwargs):
+        if not self._tool.sync_coroutine:
+            raise ValueError(f"Tool {self._tool.name} does not have a sync_coroutine defined")
         return self._tool.sync_coroutine(**kwargs)
 
+    async def _arun(self, *args, **kwargs):
+        if not self._tool.coroutine:
+            raise ValueError(f"Tool {self._tool.name} does not have a coroutine defined")
+        return await self._tool.coroutine(**kwargs)
+
 
 async def bl_tools(tools_names: list[str], **kwargs) -> list[BaseTool]:
     tools = bl_tools_core(tools_names, **kwargs)

blaxel/googleadk/model.py

@@ -1,6 +1,9 @@
 from logging import getLogger
 
-from google.adk.models.lite_llm import LiteLlm, LiteLLMClient
+from google.adk.models.lite_llm import (  # type: ignore[import-not-found]
+    LiteLlm,
+    LiteLLMClient,
+)
 
 from blaxel.core import bl_model as bl_model_core
 from blaxel.core import settings
@@ -23,7 +26,15 @@ class AuthenticatedLiteLLMClient(LiteLLMClient):
         Returns:
           The model response as a message.
         """
-        kwargs["extra_headers"] = settings.auth.get_headers()
+        auth_headers = settings.auth.get_headers()
+        extra = dict(auth_headers)
+        # When auth uses X-Blaxel-Authorization (API keys), override the
+        # Authorization header that litellm sets from api_key or OPENAI_API_KEY
+        # env var. Without this, the server sees an invalid Authorization header
+        # and rejects the request.
+        if "Authorization" not in auth_headers:
+            extra["Authorization"] = ""
+        kwargs["extra_headers"] = extra
         return await super().acompletion(
             model=model,
             messages=messages,
@@ -44,7 +55,15 @@ class AuthenticatedLiteLLMClient(LiteLLMClient):
         Returns:
           The response from the model.
         """
-        kwargs["extra_headers"] = settings.auth.get_headers()
+        auth_headers = settings.auth.get_headers()
+        extra = dict(auth_headers)
+        # When auth uses X-Blaxel-Authorization (API keys), override the
+        # Authorization header that litellm sets from api_key or OPENAI_API_KEY
+        # env var. Without this, the server sees an invalid Authorization header
+        # and rejects the request.
+        if "Authorization" not in auth_headers:
+            extra["Authorization"] = ""
+        kwargs["extra_headers"] = extra
         return super().completion(
             model=model,
             messages=messages,

blaxel/googleadk/tools.py

@@ -1,8 +1,8 @@
 import inspect
-from typing import Any, override
+from typing import Any
 
-from google.adk.tools import BaseTool, ToolContext
-from google.genai import types
+from google.adk.tools import BaseTool, ToolContext  # type: ignore[import-not-found]
+from google.genai import types  # type: ignore[import-not-found]
 
 from blaxel.core.tools import bl_tools as bl_tools_core
 from blaxel.core.tools.types import Tool
@@ -31,15 +31,33 @@ class GoogleADKTool(BaseTool):
         if "additionalProperties" in cleaned_schema:
             del cleaned_schema["additionalProperties"]
 
+        # Google genai Schema expects type as a single enum string (e.g. "STRING"),
+        # not a JSON Schema union list like ["null", "string"].
+        if "type" in cleaned_schema and isinstance(cleaned_schema["type"], list):
+            type_list = [t for t in cleaned_schema["type"] if t != "null"]
+            cleaned_schema["type"] = type_list[0].upper() if type_list else "STRING"
+            # Mark as nullable if "null" was in the original list
+            if "null" in schema["type"]:
+                cleaned_schema["nullable"] = True
+        elif "type" in cleaned_schema and isinstance(cleaned_schema["type"], str):
+            cleaned_schema["type"] = cleaned_schema["type"].upper()
+
+        # Ensure object types have properties
+        if cleaned_schema.get("type") == "OBJECT" and "properties" not in cleaned_schema:
+            cleaned_schema["properties"] = {}
+
         # Recursively clean properties if they exist
         if "properties" in cleaned_schema:
             cleaned_schema["properties"] = {
                 k: self._clean_schema(v) for k, v in cleaned_schema["properties"].items()
             }
 
+        # Recursively clean items for array types
+        if "items" in cleaned_schema and isinstance(cleaned_schema["items"], dict):
+            cleaned_schema["items"] = self._clean_schema(cleaned_schema["items"])
+
         return cleaned_schema
 
-    @override
     def _get_declaration(self) -> types.FunctionDeclaration | None:
         # Clean the schema recursively
         schema = self._clean_schema(self._tool.input_schema)
@@ -48,14 +66,15 @@ class GoogleADKTool(BaseTool):
             types.FunctionDeclaration(
                 name=self._tool.name,
                 description=self._tool.description,
-                parameters=schema,
+                parameters=types.Schema(**schema),
             )
         )
 
         return function_decl
 
-    @override
     async def run_async(self, *, args: dict[str, Any], tool_context: ToolContext) -> Any:
+        if not self._tool.coroutine:
+            raise ValueError(f"Tool {self._tool.name} does not have a coroutine defined")
         args_to_call = args.copy()
         signature = inspect.signature(self._tool.coroutine)
         if "tool_context" in signature.parameters: