""" SnakeBridge Snakepit Adapter Provides the Snakepit adapter interface for SnakeBridge. This module is called by the Snakepit runtime to execute Python functions and return results to Elixir. Main function: snakebridge_call(module: str, function: str, args: dict) -> dict The adapter: 1. Imports the specified Python module 2. Gets the specified function from the module 3. Decodes the arguments from SnakeBridge format 4. Calls the function with the decoded arguments 5. Encodes the result back to SnakeBridge format 6. Returns a success/error response Graceful Serialization: When encoding results, this adapter preserves container structure (lists, dicts, tuples, sets) even when they contain non-serializable objects. Non-serializable leaf values are ref-wrapped in place. No lossy conversion is performed (e.g., no model_dump/to_dict calls); non-serializable leaves become refs that can be used for further Python operations. This enables users to access all serializable fields while still having refs for non-serializable objects. """ import sys import importlib import importlib.util import inspect import traceback import uuid import os import glob import hashlib import time import threading import types from contextlib import nullcontext from typing import Any, Dict, List, Tuple, Optional # Import the SnakeBridge type encoding system try: from snakebridge_types import ( decode, encode, encode_error, _is_generator_or_iterator, _get_stream_type, Atom, SCHEMA_VERSION, ) except ImportError: # If running as a script, try relative import import os sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from snakebridge_types import ( decode, encode, encode_error, _is_generator_or_iterator, _get_stream_type, Atom, SCHEMA_VERSION, ) try: from snakepit_bridge import telemetry as snakepit_telemetry except Exception: snakepit_telemetry = None # Module cache to avoid repeated imports _module_cache: Dict[str, Any] = {} _instance_registry: Dict[str, Any] = {} _helper_registry: Dict[str, Any] = {} _helper_registry_key: Optional[Tuple[Any, ...]] = None _helper_registry_index: List[Dict[str, Any]] = [] # Thread locks for global state _module_cache_lock = threading.RLock() _registry_lock = threading.RLock() _helper_lock = threading.RLock() PROTOCOL_VERSION = 1 MIN_SUPPORTED_VERSION = 1 REF_SCHEMA_VERSION = 1 DEFAULT_REF_TTL_SECONDS = 0.0 DEFAULT_REF_MAX_SIZE = 10000 ALLOW_LEGACY_PROTOCOL = os.getenv("SNAKEBRIDGE_ALLOW_LEGACY_PROTOCOL", "false").lower() in ( "1", "true", "yes", ) class SnakeBridgeHelperNotFoundError(Exception): pass class SnakeBridgeHelperLoadError(Exception): pass class SnakeBridgeSerializationError(Exception): pass class SnakeBridgeProtocolError(Exception): def __init__(self, message: str, details: Optional[Dict[str, Any]] = None): super().__init__(message) self.details = details or {} def _protocol_compatibility(arguments: Dict[str, Any]) -> None: protocol_version = arguments.get("protocol_version") min_supported_version = arguments.get("min_supported_version") if protocol_version is None and min_supported_version is None and ALLOW_LEGACY_PROTOCOL: protocol_version = MIN_SUPPORTED_VERSION min_supported_version = 0 else: if protocol_version is None: protocol_version = 0 if min_supported_version is None: min_supported_version = 0 if protocol_version < MIN_SUPPORTED_VERSION or min_supported_version > PROTOCOL_VERSION: details = { "caller_protocol_version": protocol_version, "caller_min_supported_version": min_supported_version, "adapter_protocol_version": PROTOCOL_VERSION, "adapter_min_supported_version": MIN_SUPPORTED_VERSION, } message = ( "SnakeBridge protocol version mismatch " f"(caller_protocol_version={protocol_version}, " f"caller_min_supported_version={min_supported_version}, " f"adapter_protocol_version={PROTOCOL_VERSION}, " f"adapter_min_supported_version={MIN_SUPPORTED_VERSION}). " "Ensure Elixir and Python SnakeBridge versions are compatible." ) raise SnakeBridgeProtocolError( message, details=details, ) def _registry_limits() -> Tuple[float, int]: ttl_env = os.getenv("SNAKEBRIDGE_REF_TTL_SECONDS") max_env = os.getenv("SNAKEBRIDGE_REF_MAX") try: ttl = float(ttl_env) if ttl_env is not None else DEFAULT_REF_TTL_SECONDS except ValueError: ttl = DEFAULT_REF_TTL_SECONDS try: max_size = int(max_env) if max_env is not None else DEFAULT_REF_MAX_SIZE except ValueError: max_size = DEFAULT_REF_MAX_SIZE return ttl, max_size def _entry_last_access(entry: Any) -> float: if isinstance(entry, dict): return float(entry.get("last_access") or entry.get("created_at") or 0.0) return 0.0 def _touch_entry(entry: Any) -> None: if isinstance(entry, dict): entry["last_access"] = time.time() def _prune_registry() -> None: with _registry_lock: ttl_seconds, max_size = _registry_limits() now = time.time() if ttl_seconds and ttl_seconds > 0: for key, entry in list(_instance_registry.items()): if now - _entry_last_access(entry) > ttl_seconds: del _instance_registry[key] if max_size and max_size > 0 and len(_instance_registry) > max_size: overflow = len(_instance_registry) - max_size oldest = sorted(_instance_registry.items(), key=lambda item: _entry_last_access(item[1])) for key, _entry in oldest[:overflow]: del _instance_registry[key] def _store_ref(key: str, obj: Any) -> None: now = time.time() with _registry_lock: _instance_registry[key] = {"obj": obj, "created_at": now, "last_access": now} def _registry_key(session_id: str, ref_id: str) -> str: """ Construct a registry key from session_id and ref_id. This is the single source of truth for registry key format. All code that constructs or parses registry keys should use this function. """ return f"{session_id}:{ref_id}" def _registry_key_prefix(session_id: str) -> str: """ Get the registry key prefix for a session. Used for session-scoped operations like release_session. Delegates to _registry_key to maintain single source of truth for key format. """ return _registry_key(session_id, "") def _extract_ref_identity(ref: dict, session_id: str) -> Tuple[str, str]: if ref.get("__type__") == "ref": ref_id = ref.get("id") or ref.get("ref_id") ref_session = ref.get("session_id") or session_id elif ref.get("__snakebridge_ref__"): ref_id = ref.get("ref_id") ref_session = ref.get("session_id") or session_id else: raise ValueError("Invalid SnakeBridge reference payload") if not ref_id: raise ValueError("SnakeBridge reference missing id") if ref_session and session_id and ref_session != session_id: raise ValueError("SnakeBridge reference session mismatch") return ref_id, ref_session def _call_telemetry_span(metadata: Dict[str, Any]): if snakepit_telemetry is None: return nullcontext() try: return snakepit_telemetry.span("python.call", metadata) except Exception: return nullcontext() def _call_function_name(call_type: str, function: Optional[str], arguments: Dict[str, Any]) -> str: if function: return str(function) if call_type == "class": return str(arguments.get("class") or arguments.get("class_name") or "unknown") if call_type in ("get_attr", "set_attr", "module_attr"): return str(arguments.get("attr") or "unknown") if call_type == "helper": return str(arguments.get("helper") or "unknown") return "unknown" def _call_metadata( call_type: str, library: Optional[str], python_module: Optional[str], function: Optional[str], arguments: Dict[str, Any], ) -> Dict[str, Any]: metadata = { "library": str(library or "unknown"), "function": _call_function_name(call_type, function, arguments), "call_type": str(call_type), } if python_module: metadata["python_module"] = str(python_module) return metadata def snakebridge_call(module: str, function: str, args: dict) -> dict: """ Call a Python function from SnakeBridge. This is the main entry point called by the Snakepit adapter. Args: module: The Python module name (e.g., 'math', 'numpy') function: The function name to call (e.g., 'sqrt', 'array') args: Dictionary of argument names to values (in SnakeBridge encoded format) Returns: Dictionary with either: - {"success": True, "result": } - {"success": False, "error": , "error_type": } Examples: >>> snakebridge_call('math', 'sqrt', {'x': 16}) {'success': True, 'result': 4.0} >>> snakebridge_call('math', 'gcd', {'a': 48, 'b': 18}) {'success': True, 'result': 6} >>> snakebridge_call('statistics', 'mean', {'data': [1, 2, 3, 4, 5]}) {'success': True, 'result': 3.0} """ try: # Import the module (use cache if available) try: mod = _import_module(module) except ImportError as e: return encode_error(ImportError(f"Failed to import module '{module}': {str(e)}")) # Get the function from the module if not hasattr(mod, function): return encode_error(AttributeError(f"Module '{module}' has no function '{function}'")) func = getattr(mod, function) # Check if it's callable if not callable(func): return encode_error(TypeError(f"'{module}.{function}' is not callable")) # Decode arguments from SnakeBridge format try: decoded_args = {name: decode(value) for name, value in args.items()} except Exception as e: return encode_error(ValueError(f"Failed to decode arguments: {str(e)}")) # Call the function # Try to determine if we should use positional or keyword arguments try: # First, try with keyword arguments (most flexible) try: result = func(**decoded_args) except TypeError as e: error_msg = str(e) # If it fails because it doesn't accept keyword arguments, # try with positional arguments instead if "keyword argument" in error_msg.lower(): # Try to get the signature to determine argument order try: sig = inspect.signature(func) # Create positional args in parameter order positional_args = [] for param_name in sig.parameters.keys(): if param_name in decoded_args: positional_args.append(decoded_args[param_name]) # If we didn't find any matching parameters, it might be a *args function # Fall back to using values in insertion order if not positional_args: positional_args = list(decoded_args.values()) result = func(*positional_args) except (ValueError, TypeError): # Can't get signature, use values in insertion order (Python 3.7+ dicts) positional_args = list(decoded_args.values()) result = func(*positional_args) else: # Re-raise if it's a different kind of TypeError raise except TypeError as e: # Provide helpful error message for argument mismatches error_msg = str(e) return encode_error(TypeError(f"Argument error calling {module}.{function}: {error_msg}")) except Exception as e: # Return any exception from the function call return encode_error(e) # Encode and return the result library = module.split(".")[0] if module else "unknown" return { "success": True, "result": encode_result(result, "default", module, library), } except Exception as e: # Catch any unexpected errors error_info = { "success": False, "error": str(e), "error_type": type(e).__name__, "traceback": traceback.format_exc() } return error_info def snakebridge_batch_call(calls: list) -> list: """ Execute multiple function calls in a batch. Args: calls: List of call specifications, each with 'module', 'function', and 'args' Returns: List of results corresponding to each call Example: >>> snakebridge_batch_call([ ... {'module': 'math', 'function': 'sqrt', 'args': {'x': 16}}, ... {'module': 'math', 'function': 'gcd', 'args': {'a': 48, 'b': 18}} ... ]) [{'success': True, 'result': 4.0}, {'success': True, 'result': 6}] """ results = [] for call in calls: try: module = call['module'] function = call['function'] args = call.get('args', {}) result = snakebridge_call(module, function, args) results.append(result) except Exception as e: results.append(encode_error(e)) return results def snakebridge_get_attribute(module: str, attribute: str) -> dict: """ Get an attribute or constant from a module. Args: module: The Python module name attribute: The attribute name to get Returns: Dictionary with either: - {"success": True, "result": } - {"success": False, "error": } Example: >>> snakebridge_get_attribute('math', 'pi') {'success': True, 'result': 3.141592653589793} """ try: # Import the module try: mod = _import_module(module) except ImportError as e: return encode_error(ImportError(f"Failed to import module '{module}': {str(e)}")) # Get the attribute if not hasattr(mod, attribute): return encode_error(AttributeError(f"Module '{module}' has no attribute '{attribute}'")) value = getattr(mod, attribute) # Encode and return the value library = module.split(".")[0] if module else "unknown" return { "success": True, "result": encode_result(value, "default", module, library), } except Exception as e: return encode_error(e) def snakebridge_create_instance(module: str, class_name: str, args: dict) -> dict: """ Create an instance of a class. Args: module: The Python module name class_name: The class name to instantiate args: Dictionary of constructor arguments Returns: Dictionary with either success or error Note: Instance objects cannot be serialized, so this is mainly useful for testing or when combined with a session/state system. """ try: # Import the module try: mod = _import_module(module) except ImportError as e: return encode_error(ImportError(f"Failed to import module '{module}': {str(e)}")) # Get the class if not hasattr(mod, class_name): return encode_error(AttributeError(f"Module '{module}' has no class '{class_name}'")) cls = getattr(mod, class_name) # Check if it's a class if not isinstance(cls, type): return encode_error(TypeError(f"'{module}.{class_name}' is not a class")) # Decode arguments decoded_args = {name: decode(value) for name, value in args.items()} # Create instance instance = cls(**decoded_args) # Encode and return (note: complex objects may not serialize well) library = module.split(".")[0] if module else "unknown" return { "success": True, "result": encode_result(instance, "default", module, library), } except Exception as e: return encode_error(e) def _import_module(module_name: str) -> Any: with _module_cache_lock: if module_name in _module_cache: return _module_cache[module_name] mod = importlib.import_module(module_name) _module_cache[module_name] = mod return mod def _make_ref(session_id: str, obj: Any, python_module: str, library: str) -> dict: ref_id = uuid.uuid4().hex key = _registry_key(session_id, ref_id) _prune_registry() _store_ref(key, obj) type_name = type(obj).__name__ return { "__type__": "ref", "__schema__": REF_SCHEMA_VERSION, "id": ref_id, "session_id": session_id, "python_module": python_module, "library": library, "type_name": type_name, "__type_name__": type_name, } def _make_stream_ref( session_id: str, obj: Any, python_module: str, library: str, stream_type: str, ) -> dict: ref_id = uuid.uuid4().hex key = _registry_key(session_id, ref_id) _prune_registry() _store_ref(key, obj) type_name = type(obj).__name__ return { "__type__": "stream_ref", "__schema__": REF_SCHEMA_VERSION, "id": ref_id, "session_id": session_id, "python_module": python_module, "library": library, "stream_type": stream_type, "type_name": type_name, "__type_name__": type_name, } def _truncate_message(msg: str, max_len: int = 200) -> str: """Truncate a message to max_len, stripping newlines for log safety.""" msg = msg.replace('\n', ' ').replace('\r', '') if len(msg) > max_len: return msg[:max_len] + "..." return msg def _cleanup_created_refs(created_keys: list) -> None: """Remove refs from registry that were created during a failed encode.""" with _registry_lock: for key in created_keys: _instance_registry.pop(key, None) def encode_result(result: Any, session_id: str, python_module: str, library: str) -> Any: """ Encode a Python result for transmission to Elixir. INVARIANT: Always returns JSON-safe values. Any non-serializable object is represented by a ref or stream_ref payload (lossless - the object remains accessible via Python operations). The returned value is ALWAYS one of: 1. A JSON-safe primitive (None, bool, int, float, str) 2. A JSON-safe list/dict (recursively containing the above types) 3. A tagged value with __type__ that Elixir can decode 4. A ref ({"__type__": "ref", ...}) for non-serializable objects 5. A stream_ref ({"__type__": "stream_ref", ...}) for sync iterators This function preserves container structure: if a list/dict contains non-serializable items, only those items become refs, not the entire container. This enables "graceful serialization" where users can access all serializable fields even when some fields are non-serializable. No lossy conversion is performed (e.g., no model_dump/to_dict calls). Repeated references to the same non-serializable object yield the same ref payload; repeated references to the same iterator yield the same stream_ref payload. Async generators become refs (not stream_refs) since they cannot be consumed via sync iteration. On encoding failure, any refs created during partial encoding are cleaned up to prevent unreachable ref leakage in the registry. """ # Use set of object ids to detect cycles in_progress = set() # Memo table: id(obj) -> ref/stream_ref payload for deduplication ref_memo = {} # Track registry keys created during this encode for cleanup on failure created_keys = [] try: encoded = _encode_result_recursive( result, session_id, python_module, library, in_progress, ref_memo, created_keys ) except Exception as e: # Any exception during encoding (RecursionError, RuntimeError from # concurrent mutation, weird __iter__ implementations, etc.) # Clean up any refs created during partial encoding _cleanup_created_refs(created_keys) # Fall back to ref-wrapping the entire result (lossless) _log_warning( f"{type(e).__name__} encoding type {type(result).__name__}: " f"{_truncate_message(str(e))}, wrapping in ref as fallback" ) return _make_ref(session_id, result, python_module, library) # Safety net: validate result is actually JSON-safe # This catches any edge cases missed by the recursive encoder try: is_safe = _is_json_safe(encoded) except RecursionError: # _is_json_safe itself can recurse deeply on near-limit structures # Clean up any refs created during encoding _cleanup_created_refs(created_keys) _log_warning( f"RecursionError in _is_json_safe for type {type(result).__name__}, " "wrapping in ref as safety fallback" ) return _make_ref(session_id, result, python_module, library) if not is_safe: # Clean up any refs created during encoding _cleanup_created_refs(created_keys) _log_warning( f"encode_result produced non-JSON-safe result for type {type(result).__name__}, " "wrapping in ref as safety fallback" ) return _make_ref(session_id, result, python_module, library) return encoded def _encode_result_recursive( value: Any, session_id: str, python_module: str, library: str, in_progress: set, ref_memo: dict, created_keys: list, ) -> Any: """ Recursively encode a Python value, creating refs for non-serializable leaves. Args: value: The Python value to encode session_id: Session ID for ref tracking python_module: Module name for ref metadata library: Library name for ref metadata in_progress: Set of object ids currently being encoded (for cycle detection) ref_memo: Dict mapping id(obj) -> ref payload for deduplication created_keys: List to append registry keys of created refs (for cleanup on failure) Returns: A JSON-safe value with refs for non-serializable objects """ import base64 import math from datetime import datetime, date, time # Handle None if value is None: return None # Handle booleans (must check before int - bool is subclass of int) if isinstance(value, bool): return value # Handle integers if isinstance(value, int): return value # Handle float (check for special values) if isinstance(value, float): if math.isinf(value): return { "__type__": "special_float", "__schema__": SCHEMA_VERSION, "value": "infinity" if value > 0 else "neg_infinity", } if math.isnan(value): return { "__type__": "special_float", "__schema__": SCHEMA_VERSION, "value": "nan", } return value # Handle strings if isinstance(value, str): return value # Handle bytes/bytearray - always tagged if isinstance(value, (bytes, bytearray)): return { "__type__": "bytes", "__schema__": SCHEMA_VERSION, "data": base64.b64encode(bytes(value)).decode("ascii"), } # Handle complex if isinstance(value, complex): return { "__type__": "complex", "__schema__": SCHEMA_VERSION, "real": value.real, "imag": value.imag, } # Handle datetime types if isinstance(value, datetime): return { "__type__": "datetime", "__schema__": SCHEMA_VERSION, "value": value.isoformat(), } if isinstance(value, date): return { "__type__": "date", "__schema__": SCHEMA_VERSION, "value": value.isoformat(), } if isinstance(value, time): return { "__type__": "time", "__schema__": SCHEMA_VERSION, "value": value.isoformat(), } # Handle Atom - tagged atom value for Elixir interop if isinstance(value, Atom): return { "__type__": "atom", "__schema__": SCHEMA_VERSION, "value": value.value, } # Check for async generators FIRST - these cannot be consumed via next() # so we treat them as regular refs, not stream_refs if hasattr(types, 'AsyncGeneratorType') and isinstance(value, types.AsyncGeneratorType): obj_id = id(value) if obj_id in ref_memo: return ref_memo[obj_id] ref_payload = _make_ref(session_id, value, python_module, library) created_keys.append(_registry_key(session_id, ref_payload['id'])) ref_memo[obj_id] = ref_payload return ref_payload # Check for sync generators/iterators - create stream_ref in-place if _is_generator_or_iterator(value): obj_id = id(value) if obj_id in ref_memo: return ref_memo[obj_id] stream_type = _get_stream_type(value) stream_ref_payload = _make_stream_ref(session_id, value, python_module, library, stream_type) created_keys.append(_registry_key(session_id, stream_ref_payload['id'])) ref_memo[obj_id] = stream_ref_payload return stream_ref_payload # For container types, check for cycles obj_id = id(value) # Cycle detection: if we're already encoding this object, create a ref # Use memo if available to ensure same ref for same object if obj_id in in_progress: if obj_id in ref_memo: return ref_memo[obj_id] ref_payload = _make_ref(session_id, value, python_module, library) created_keys.append(_registry_key(session_id, ref_payload['id'])) ref_memo[obj_id] = ref_payload return ref_payload # Handle list - recursively encode elements, refs for non-encodable # Snapshot to avoid "list changed size during iteration" errors if isinstance(value, list): in_progress.add(obj_id) try: items = list(value) # Snapshot return [ _encode_result_recursive(item, session_id, python_module, library, in_progress, ref_memo, created_keys) for item in items ] finally: in_progress.discard(obj_id) # Handle tuple - recursively encode elements # Tuples are immutable, but snapshot for consistency if isinstance(value, tuple): in_progress.add(obj_id) try: items = tuple(value) # Snapshot (tuples are immutable but be consistent) elements = [ _encode_result_recursive(item, session_id, python_module, library, in_progress, ref_memo, created_keys) for item in items ] return { "__type__": "tuple", "__schema__": SCHEMA_VERSION, "elements": elements, } finally: in_progress.discard(obj_id) # Handle set - recursively encode elements (no sorting; sets have no order semantics) # Snapshot to avoid "set changed size during iteration" errors if isinstance(value, set): in_progress.add(obj_id) try: items = list(value) # Snapshot elements = [ _encode_result_recursive(item, session_id, python_module, library, in_progress, ref_memo, created_keys) for item in items ] return { "__type__": "set", "__schema__": SCHEMA_VERSION, "elements": elements, } finally: in_progress.discard(obj_id) # Handle frozenset - recursively encode elements (no sorting; sets have no order semantics) # Frozensets are immutable, but snapshot for consistency if isinstance(value, frozenset): in_progress.add(obj_id) try: items = list(value) # Snapshot elements = [ _encode_result_recursive(item, session_id, python_module, library, in_progress, ref_memo, created_keys) for item in items ] return { "__type__": "frozenset", "__schema__": SCHEMA_VERSION, "elements": elements, } finally: in_progress.discard(obj_id) # Handle dict - recursively encode values (and keys if non-string) # Snapshot to avoid "dictionary changed size during iteration" errors if isinstance(value, dict): in_progress.add(obj_id) try: items = list(value.items()) # Snapshot all_string_keys = all(isinstance(k, str) for k, _ in items) if all_string_keys: # Plain dict with string keys return { k: _encode_result_recursive(v, session_id, python_module, library, in_progress, ref_memo, created_keys) for k, v in items } else: # Tagged dict with pairs for non-string keys pairs = [ [ _encode_result_recursive(k, session_id, python_module, library, in_progress, ref_memo, created_keys), _encode_result_recursive(v, session_id, python_module, library, in_progress, ref_memo, created_keys), ] for k, v in items ] return { "__type__": "dict", "__schema__": SCHEMA_VERSION, "pairs": pairs, } finally: in_progress.discard(obj_id) # EVERYTHING ELSE - create a ref (with memoization) # This includes custom objects, functions, modules, etc. if obj_id in ref_memo: return ref_memo[obj_id] ref_payload = _make_ref(session_id, value, python_module, library) created_keys.append(_registry_key(session_id, ref_payload['id'])) ref_memo[obj_id] = ref_payload return ref_payload def _is_streamable(value: Any) -> bool: """ Check if value can be streamed (consumed via sync iteration). NOTE: Async generators are explicitly EXCLUDED - they cannot be consumed via next() and must be treated as regular refs. """ # Explicitly exclude async generators - they cannot be consumed via next() if hasattr(types, 'AsyncGeneratorType') and isinstance(value, types.AsyncGeneratorType): return False if _is_generator_or_iterator(value): return True if hasattr(value, "__iter__"): if isinstance(value, (str, bytes, bytearray, list, tuple, dict, set, frozenset)): return False # Prefer regular refs for context managers (e.g. file objects) if hasattr(value, "__enter__") and hasattr(value, "__exit__"): return False return True return False def _dynamic_stream_iterator( result: Any, session_id: str, python_module: str, library: str, ): if _is_streamable(result): iterator = result if not hasattr(iterator, "__next__") and hasattr(iterator, "__iter__"): iterator = iter(iterator) for item in iterator: yield encode_result(item, session_id, python_module, library) else: yield encode_result(result, session_id, python_module, library) def _is_json_safe(value: Any) -> bool: """ Verify a value is safe to serialize as JSON. This is a safety check after encoding - if encode() is correct, this should always return True. """ import math if value is None: return True if isinstance(value, bool): return True if isinstance(value, (int, float)): # Check for non-JSON floats (should be tagged already) if isinstance(value, float) and (math.isinf(value) or math.isnan(value)): return False return True if isinstance(value, str): return True if isinstance(value, list): return all(_is_json_safe(item) for item in value) if isinstance(value, dict): # Check for valid tagged types type_tag = value.get("__type__") if type_tag in ( "bytes", "tuple", "set", "frozenset", "complex", "datetime", "date", "time", "special_float", "atom", "dict", "ref", "stream_ref", "callback", "stop_iteration" ): # Tagged values are safe - validate values recursively return all(_is_json_safe(v) for v in value.values()) # Regular dict - keys must be strings if not all(isinstance(k, str) for k in value.keys()): return False return all(_is_json_safe(v) for v in value.values()) return False def _log_warning(message: str) -> None: """Log a warning message.""" print(f"[SnakeBridge WARNING] {message}", file=sys.stderr) def _is_ref_payload(value: Any) -> bool: if not isinstance(value, dict): return False type_tag = value.get("__type__") if type_tag is not None: return type_tag == "ref" return "id" in value and ("session_id" in value or "ref_id" in value) def _resolve_refs(value: Any, session_id: str) -> Any: if isinstance(value, dict): if _is_ref_payload(value): return _resolve_ref(value, session_id) return {k: _resolve_refs(v, session_id) for k, v in value.items()} if isinstance(value, list): return [_resolve_refs(item, session_id) for item in value] if isinstance(value, tuple): return tuple(_resolve_refs(item, session_id) for item in value) return value def _resolve_ref(ref: dict, session_id: str) -> Any: if not isinstance(ref, dict): raise ValueError("Invalid SnakeBridge reference payload") with _registry_lock: _prune_registry() ref_id, ref_session = _extract_ref_identity(ref, session_id) key = _registry_key(ref_session, ref_id) if key not in _instance_registry: raise KeyError(f"Unknown SnakeBridge reference: {ref_id}") entry = _instance_registry[key] if isinstance(entry, dict): _touch_entry(entry) return entry.get("obj") return entry def _release_ref(ref: dict, session_id: str) -> bool: if not isinstance(ref, dict): raise ValueError("Invalid SnakeBridge reference payload") with _registry_lock: _prune_registry() ref_id, ref_session = _extract_ref_identity(ref, session_id) key = _registry_key(ref_session, ref_id) if key in _instance_registry: del _instance_registry[key] return True return False def _release_session(session_id: str) -> int: if not session_id: return 0 with _registry_lock: removed = 0 prefix = _registry_key_prefix(session_id) for key in list(_instance_registry.keys()): if key.startswith(prefix): del _instance_registry[key] removed += 1 return removed def _default_helper_config() -> Dict[str, Any]: return { "helper_paths": ["priv/python/helpers"], "helper_pack_enabled": True, "helper_allowlist": "all" } def _normalize_helper_config(config: Optional[Dict[str, Any]]) -> Dict[str, Any]: normalized = _default_helper_config() if isinstance(config, dict): normalized.update(config) helper_paths = normalized.get("helper_paths") or [] if isinstance(helper_paths, str): helper_paths = [helper_paths] helper_paths = [os.path.abspath(path) for path in helper_paths if path] normalized["helper_paths"] = helper_paths normalized["helper_pack_enabled"] = bool(normalized.get("helper_pack_enabled", True)) allowlist = normalized.get("helper_allowlist", "all") if allowlist in [None, "all", ":all"]: normalized["helper_allowlist"] = "all" elif isinstance(allowlist, (list, tuple, set)): normalized["helper_allowlist"] = [str(item) for item in allowlist] else: normalized["helper_allowlist"] = [str(allowlist)] return normalized def _helper_config_key(config: Dict[str, Any]) -> Tuple[Any, ...]: helper_paths = tuple(config.get("helper_paths", [])) allowlist = config.get("helper_allowlist", "all") allowlist_key = "all" if allowlist == "all" else tuple(allowlist) return (helper_paths, allowlist_key, bool(config.get("helper_pack_enabled", True))) def _resolve_helper_paths(config: Dict[str, Any]) -> List[str]: paths: List[str] = [] if config.get("helper_pack_enabled", True): pack_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "helpers") paths.append(pack_path) paths.extend(config.get("helper_paths", [])) return paths def _list_helper_files(paths: List[str]) -> List[str]: files: List[str] = [] for path in paths: if not path: continue if os.path.isdir(path): for file_path in sorted(glob.glob(os.path.join(path, "*.py"))): base = os.path.basename(file_path) if base == "__init__.py" or base.startswith("_"): continue files.append(file_path) elif os.path.isfile(path): base = os.path.basename(path) if base != "__init__.py" and not base.startswith("_"): files.append(path) return files def _import_helper_module(path: str) -> Any: module_name = f"snakebridge_helper_{hashlib.md5(path.encode('utf-8')).hexdigest()}" if module_name in sys.modules: return sys.modules[module_name] spec = importlib.util.spec_from_file_location(module_name, path) if spec is None or spec.loader is None: raise SnakeBridgeHelperLoadError(f"Unable to load helper module: {path}") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) sys.modules[module_name] = module return module def _extract_helpers_from_module(module: Any) -> Dict[str, Any]: if hasattr(module, "__snakebridge_helpers__"): helpers = getattr(module, "__snakebridge_helpers__") elif hasattr(module, "snakebridge_helpers"): helpers = module.snakebridge_helpers() elif hasattr(module, "HELPERS"): helpers = getattr(module, "HELPERS") else: return {} if not isinstance(helpers, dict): raise SnakeBridgeHelperLoadError("Helper registry must be a dict of name => callable") for name, func in helpers.items(): if not isinstance(name, str): raise SnakeBridgeHelperLoadError("Helper names must be strings") if not callable(func): raise SnakeBridgeHelperLoadError(f"Helper '{name}' is not callable") return helpers def _apply_allowlist(helpers: Dict[str, Any], allowlist: Any) -> Dict[str, Any]: if allowlist == "all": return helpers if not allowlist: return {} return {name: func for name, func in helpers.items() if name in allowlist} def _load_helper_registry(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: global _helper_registry, _helper_registry_key, _helper_registry_index with _helper_lock: normalized = _normalize_helper_config(config) key = _helper_config_key(normalized) if key == _helper_registry_key: return _helper_registry registry: Dict[str, Any] = {} for path in _resolve_helper_paths(normalized): if not path or not os.path.exists(path): continue for file_path in _list_helper_files([path]): module = _import_helper_module(file_path) helpers = _extract_helpers_from_module(module) registry.update(helpers) registry = _apply_allowlist(registry, normalized.get("helper_allowlist", "all")) _helper_registry = registry _helper_registry_key = key _helper_registry_index = _build_helper_index(registry) return registry def _format_annotation(annotation: Any) -> Optional[str]: if annotation is inspect.Signature.empty: return None if hasattr(annotation, "__name__"): return annotation.__name__ return str(annotation) def _param_info(param: inspect.Parameter) -> Dict[str, Any]: info: Dict[str, Any] = {"name": param.name, "kind": param.kind.name} if param.default is not inspect.Parameter.empty: info["default"] = repr(param.default) if param.annotation is not inspect.Parameter.empty: info["annotation"] = _format_annotation(param.annotation) return info def _build_helper_index(helpers: Dict[str, Any]) -> List[Dict[str, Any]]: index: List[Dict[str, Any]] = [] for name, func in helpers.items(): entry = {"name": name} try: sig = inspect.signature(func) entry["parameters"] = [_param_info(p) for p in sig.parameters.values()] except (ValueError, TypeError): entry["parameters"] = [] doc = inspect.getdoc(func) or "" if doc: entry["docstring"] = doc[:8000] index.append(entry) index.sort(key=lambda item: item.get("name", "")) return index def helper_registry_index(config: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: _load_helper_registry(config or {}) return _helper_registry_index class SnakeBridgeAdapter: def __init__(self): self.session_context = None def set_session_context(self, session_context): self.session_context = session_context def execute_tool(self, tool_name: str, arguments: dict, context): if tool_name == "snakebridge.helpers": helper_config = {} if isinstance(arguments, dict): helper_config = arguments.get("helper_config") or arguments return helper_registry_index(helper_config) if tool_name not in [ "snakebridge.call", "snakebridge.stream", "snakebridge.release_ref", "snakebridge.release_session", ]: raise AttributeError(f"Tool '{tool_name}' not supported by SnakeBridgeAdapter") if isinstance(arguments, dict): _protocol_compatibility(arguments) session_id = None if isinstance(arguments, dict) and arguments.get("session_id"): session_id = arguments.get("session_id") elif context is not None and hasattr(context, "session_id"): session_id = context.session_id elif self.session_context is not None: session_id = self.session_context.session_id else: session_id = "default" if tool_name == "snakebridge.release_ref": ref = arguments.get("ref") if isinstance(arguments, dict) else None if ref is None: raise ValueError("snakebridge.release_ref requires ref") return _release_ref(ref, session_id) if tool_name == "snakebridge.release_session": return _release_session(session_id) call_type = arguments.get("call_type") or "function" module_path = arguments.get("module_path") python_module = arguments.get("python_module") or arguments.get("module") function = arguments.get("function") args = arguments.get("args") or [] kwargs = arguments.get("kwargs") or {} if call_type in ("dynamic", "dynamic_stream") and not python_module: python_module = module_path library = arguments.get("library") or (python_module.split(".")[0] if python_module else None) if call_type not in ("helper", "stream_next") and not python_module: raise ValueError("snakebridge.call requires python_module") if not python_module: python_module = library or "unknown" if not library: library = python_module.split(".")[0] if python_module else "unknown" metadata = _call_metadata(call_type, library, python_module, function, arguments) decoded_args = [decode(item, session_id=session_id, context=context) for item in args] decoded_kwargs = { key: decode(value, session_id=session_id, context=context) for key, value in kwargs.items() } decoded_args = [_resolve_refs(item, session_id) for item in decoded_args] decoded_kwargs = {key: _resolve_refs(value, session_id) for key, value in decoded_kwargs.items()} with _call_telemetry_span(metadata): if call_type == "stream_next": stream_ref_payload = arguments.get("stream_ref") if stream_ref_payload is None: raise ValueError("snakebridge.call requires stream_ref for stream_next") stream_ref = decode(stream_ref_payload, session_id=session_id, context=context) iterator = _resolve_ref(stream_ref, session_id) python_module = "" library = "" if isinstance(stream_ref_payload, dict): python_module = stream_ref_payload.get("python_module", "") or "" library = stream_ref_payload.get("library", "") or "" try: item = next(iterator) return encode_result(item, session_id, python_module, library) except StopIteration: return {"__type__": "stop_iteration"} if call_type == "dynamic_stream": module_path = arguments.get("module_path") or python_module if not module_path: raise ValueError("snakebridge.call requires module_path for dynamic calls") mod = _import_module(module_path) func = getattr(mod, function) result = func(*decoded_args, **decoded_kwargs) return _dynamic_stream_iterator(result, session_id, module_path, library) if call_type == "dynamic": module_path = arguments.get("module_path") or python_module if not module_path: raise ValueError("snakebridge.call requires module_path for dynamic calls") mod = _import_module(module_path) func = getattr(mod, function) result = func(*decoded_args, **decoded_kwargs) return encode_result(result, session_id, module_path, library) if call_type == "class": class_name = arguments.get("class") or arguments.get("class_name") mod = _import_module(python_module) cls = getattr(mod, class_name) instance = cls(*decoded_args, **decoded_kwargs) return encode_result(instance, session_id, python_module, library) if call_type == "method": instance_payload = arguments.get("instance") instance = _resolve_ref(decode(instance_payload), session_id) method = getattr(instance, function) result = method(*decoded_args, **decoded_kwargs) return encode_result(result, session_id, python_module, library) if call_type == "get_attr": instance_payload = arguments.get("instance") instance = _resolve_ref(decode(instance_payload), session_id) attr = arguments.get("attr") or function result = getattr(instance, attr) return encode_result(result, session_id, python_module, library) if call_type == "module_attr": attr = arguments.get("attr") or function mod = _import_module(python_module) result = getattr(mod, attr) return encode_result(result, session_id, python_module, library) if call_type == "set_attr": instance_payload = arguments.get("instance") instance = _resolve_ref(decode(instance_payload), session_id) attr = arguments.get("attr") or function value = decoded_args[0] if decoded_args else None setattr(instance, attr, value) return encode_result(True, session_id, python_module, library) if call_type == "helper": helper_name = arguments.get("helper") or function helper_config = arguments.get("helper_config") or {} if not helper_name: raise SnakeBridgeHelperNotFoundError("Helper name is required") registry = _load_helper_registry(helper_config) if helper_name not in registry: raise SnakeBridgeHelperNotFoundError(f"Helper '{helper_name}' not found") result = registry[helper_name](*decoded_args, **decoded_kwargs) return encode_result(result, session_id, python_module, library) mod = _import_module(python_module) func = getattr(mod, function) result = func(*decoded_args, **decoded_kwargs) return encode_result(result, session_id, python_module, library) # Make the module callable for testing if __name__ == "__main__": import json # Simple test runner if len(sys.argv) > 1: # Test with command-line arguments # Usage: python snakebridge_adapter.py if len(sys.argv) >= 4: module = sys.argv[1] function = sys.argv[2] args_json = sys.argv[3] args = json.loads(args_json) result = snakebridge_call(module, function, args) print(json.dumps(result, indent=2)) else: print("Usage: python snakebridge_adapter.py ") else: # Run built-in tests print("Running SnakeBridge adapter tests...\n") # Test 1: math.sqrt print("Test 1: math.sqrt(16)") result = snakebridge_call('math', 'sqrt', {'x': 16}) print(json.dumps(result, indent=2)) assert result['success'] == True assert result['result'] == 4.0 print("PASS\n") # Test 2: math.gcd print("Test 2: math.gcd(48, 18)") result = snakebridge_call('math', 'gcd', {'a': 48, 'b': 18}) print(json.dumps(result, indent=2)) assert result['success'] == True assert result['result'] == 6 print("PASS\n") # Test 3: Error handling - module not found print("Test 3: Error handling - nonexistent module") result = snakebridge_call('nonexistent_module', 'func', {}) print(json.dumps(result, indent=2)) assert result['success'] == False print("PASS\n") # Test 4: Error handling - function not found print("Test 4: Error handling - nonexistent function") result = snakebridge_call('math', 'nonexistent_function', {}) print(json.dumps(result, indent=2)) assert result['success'] == False print("PASS\n") # Test 5: Complex types - tuple encoding print("Test 5: Complex types - math.gcd with large numbers") result = snakebridge_call('math', 'gcd', {'a': 1071, 'b': 462}) print(json.dumps(result, indent=2)) assert result['success'] == True assert result['result'] == 21 print("PASS\n") # Test 6: Get attribute print("Test 6: Get math.pi") result = snakebridge_get_attribute('math', 'pi') print(json.dumps(result, indent=2)) assert result['success'] == True assert abs(result['result'] - 3.141592653589793) < 0.0001 print("PASS\n") print("All tests passed!")