swarmkit 0.1.34__py3-none-any.whl

swarmkit/bridge.py

@@ -0,0 +1,475 @@
+"""Node.js bridge subprocess manager for JSON-RPC communication."""
+
+import asyncio
+import atexit
+import json
+import logging
+import os
+import signal
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+
+logger = logging.getLogger(__name__)
+
+# Global registry of bridge process PIDs for atexit cleanup
+# We store PIDs instead of process objects because asyncio.subprocess.Process
+# doesn't have sync wait() - we use os.kill/os.waitpid for cleanup
+_bridge_pids: List[int] = []
+
+
+def _atexit_cleanup():
+    """Kill all bridge processes on Python exit."""
+    for pid in _bridge_pids:
+        try:
+            os.kill(pid, signal.SIGTERM)
+        except (OSError, ProcessLookupError):
+            pass  # Process already dead
+        try:
+            # Non-blocking wait to reap zombie
+            os.waitpid(pid, os.WNOHANG)
+        except (OSError, ChildProcessError):
+            pass
+    _bridge_pids.clear()
+
+
+atexit.register(_atexit_cleanup)
+
+
+class SandboxNotFoundError(Exception):
+    """Raised when sandbox is not found (expired or killed)."""
+    pass
+
+
+class BridgeConnectionError(Exception):
+    """Raised when bridge process fails to start or dies unexpectedly."""
+    pass
+
+
+class BridgeBuildError(Exception):
+    """Raised when bridge build (npm install/build) fails."""
+    pass
+
+
+class BridgeManager:
+    """Manages Node.js subprocess running the JSON-RPC bridge.
+
+    Uses asyncio.create_subprocess_exec for native async I/O.
+    This allows proper cancellation and clean shutdown without blocking threads.
+    """
+
+    def __init__(self):
+        self.process: Optional[asyncio.subprocess.Process] = None
+        self.request_id = 0
+        self.pending_requests: Dict[int, asyncio.Future] = {}
+        # Serialize writes to stdin to avoid interleaved JSON-RPC requests
+        self._write_lock = asyncio.Lock()
+        # Default RPC timeout to prevent hangs (overridden per-call for long runs)
+        self.default_call_timeout_s: float = 120.0
+        self.stderr_task: Optional[asyncio.Task] = None
+        # Buffers for chunked stdout/stderr events from the bridge.
+        # The Node bridge may emit {seq, done} chunks for oversized streams.
+        self._stream_buffers: Dict[str, List[str]] = {"stdout": [], "stderr": []}
+        self.event_callbacks: Dict[str, List[Callable]] = {
+            'stdout': [],     # str data
+            'stderr': [],     # str data
+            'content': [],    # dict params
+        }
+        self.reader_task: Optional[asyncio.Task] = None
+        self._pid: Optional[int] = None
+
+    async def start(self):
+        """Start the Node.js bridge process."""
+        if self.process is not None:
+            return
+
+        # Find bridge script (bundled version for distribution)
+        bridge_dir = Path(__file__).parent.parent / 'bridge'
+        bridge_script = bridge_dir / 'dist' / 'bridge.bundle.cjs'
+
+        # Fallback to unbundled version for development
+        if not bridge_script.exists():
+            bridge_script = bridge_dir / 'dist' / 'bridge.js'
+
+        # Auto-build bridge if missing (turnkey experience)
+        if not bridge_script.exists():
+            await self._build_bridge(bridge_dir)
+
+        # Start Node.js process with native asyncio subprocess
+        self.process = await asyncio.create_subprocess_exec(
+            'node', str(bridge_script),
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        self._pid = self.process.pid
+
+        # Register PID for atexit cleanup
+        _bridge_pids.append(self._pid)
+
+        # Start reading responses (native async - no blocking threads)
+        self.reader_task = asyncio.create_task(self._read_responses())
+        # Drain stderr to avoid pipe backpressure
+        if self.process.stderr is not None:
+            self.stderr_task = asyncio.create_task(self._drain_stderr())
+
+    async def _build_bridge(self, bridge_dir: Path):
+        """Build the bridge if missing (first run experience)."""
+        import shutil
+        import subprocess
+
+        # Check if Node.js and npm are installed
+        if not shutil.which('node'):
+            raise BridgeBuildError(
+                "Bridge build failed: Node.js not found in PATH.\n"
+                "SwarmKit requires Node.js 18+ to run the TypeScript bridge.\n"
+                "Install from https://nodejs.org/ or run 'make build-dev' manually from packages/sdk-py/."
+            )
+
+        if not shutil.which('npm'):
+            raise BridgeBuildError(
+                "Bridge build failed: npm not found in PATH.\n"
+                "npm is usually installed with Node.js - check your Node.js installation.\n"
+                "Alternatively, run 'make build-dev' manually from packages/sdk-py/."
+            )
+
+        logger.info("First run: building Node.js bridge...")
+        try:
+            # Run npm install/build in executor to avoid blocking event loop
+            loop = asyncio.get_running_loop()
+
+            # Install dependencies
+            await loop.run_in_executor(
+                None,
+                lambda: subprocess.run(
+                    ['npm', 'install'],
+                    cwd=bridge_dir,
+                    check=True,
+                    capture_output=True,
+                    text=True
+                )
+            )
+
+            # Build bridge (dev mode for readable debugging)
+            await loop.run_in_executor(
+                None,
+                lambda: subprocess.run(
+                    ['npm', 'run', 'build:dev'],
+                    cwd=bridge_dir,
+                    check=True,
+                    capture_output=True,
+                    text=True
+                )
+            )
+            logger.info("Bridge built successfully")
+        except subprocess.CalledProcessError as e:
+            raise BridgeBuildError(
+                f"Bridge build failed during npm execution.\n"
+                f"Error: {e.stderr}\n"
+                f"Try running 'make build-dev' or 'make build-prod' manually from packages/sdk-py/ to see the full error."
+            ) from e
+
+    async def stop(self):
+        """Stop the Node.js bridge process."""
+        if self.process is None:
+            return
+
+        # Note: We intentionally do NOT clear event_callbacks here
+        # User-registered callbacks should persist across bridge restarts
+
+        # Remove from atexit registry
+        if self._pid and self._pid in _bridge_pids:
+            _bridge_pids.remove(self._pid)
+
+        # Cancel reader tasks first (they will exit cleanly now that process is terminating)
+        if self.reader_task:
+            self.reader_task.cancel()
+            try:
+                await self.reader_task
+            except asyncio.CancelledError:
+                pass
+
+        if self.stderr_task:
+            self.stderr_task.cancel()
+            try:
+                await self.stderr_task
+            except asyncio.CancelledError:
+                pass
+            self.stderr_task = None
+
+        # Terminate process
+        try:
+            self.process.terminate()
+            await asyncio.wait_for(self.process.wait(), timeout=5)
+        except asyncio.TimeoutError:
+            self.process.kill()
+            await self.process.wait()
+
+        self.process = None
+        self.reader_task = None
+        self._pid = None
+
+    def on(self, event_type: str, callback: Callable):
+        """Register event callback."""
+        if event_type in self.event_callbacks:
+            self.event_callbacks[event_type].append(callback)
+
+    async def call(
+        self,
+        method: str,
+        params: Optional[Dict[str, Any]] = None,
+        timeout_s: Optional[float] = None,
+    ) -> Any:
+        """Call a JSON-RPC method and wait for response.
+
+        Args:
+            method: JSON-RPC method name
+            params: JSON-RPC params dict
+            timeout_s: Optional timeout in seconds. If None, uses default_call_timeout_s.
+        """
+        if self.process is None or self.process.stdin is None:
+            raise BridgeConnectionError("Bridge not started. Call start() first.")
+
+        async with self._write_lock:
+            self.request_id += 1
+            request_id = self.request_id
+
+            request = {
+                'jsonrpc': '2.0',
+                'method': method,
+                'params': params or {},
+                'id': request_id,
+            }
+
+            # Create future for response
+            future: asyncio.Future = asyncio.Future()
+            self.pending_requests[request_id] = future
+
+            # Send request (native async write)
+            payload = json.dumps(request).encode('utf-8')
+            frame = len(payload).to_bytes(4, byteorder='big') + payload
+            self.process.stdin.write(frame)
+            await self.process.stdin.drain()
+
+        # Wait for response (error handling done in _handle_response)
+        timeout = timeout_s if timeout_s is not None else self.default_call_timeout_s
+        try:
+            return await asyncio.wait_for(future, timeout=timeout)
+        except asyncio.TimeoutError as e:
+            # Drop pending request to avoid leaks; late response will be ignored.
+            self.pending_requests.pop(request_id, None)
+            raise BridgeConnectionError(
+                f"Bridge call timed out after {timeout:.1f}s: {method}"
+            ) from e
+
+    async def _read_responses(self):
+        """Read framed JSON-RPC messages from bridge stdout.
+
+        Uses native asyncio reads - fully cancellable, no blocking threads.
+        """
+        if self.process is None or self.process.stdout is None:
+            return
+
+        # 50MB cap on incoming frames. This applies to ALL bridge responses including
+        # RPC results (run().stdout, get_output_files(), etc.). If a response exceeds
+        # 50MB, the bridge connection fails. This is stricter than the TS SDK which
+        # has no response size limit. For very large outputs, consider streaming via
+        # stdout/stderr events or fetching files individually.
+        max_frame_bytes = 50 * 1024 * 1024
+
+        try:
+            async def read_exact(n: int) -> Optional[bytes]:
+                """Read exactly n bytes from stdout."""
+                chunks: List[bytes] = []
+                remaining = n
+                while remaining > 0:
+                    chunk = await self.process.stdout.read(remaining)
+                    if not chunk:
+                        return None
+                    chunks.append(chunk)
+                    remaining -= len(chunk)
+                return b"".join(chunks)
+
+            while True:
+                header = await read_exact(4)
+                if not header:
+                    break
+                length = int.from_bytes(header, byteorder='big')
+                if length <= 0 or length > max_frame_bytes:
+                    logger.error(f"Invalid frame length from bridge: {length}")
+                    break
+
+                payload = await read_exact(length)
+                if payload is None:
+                    break
+
+                try:
+                    text = payload.decode('utf-8')
+                    message = json.loads(text)
+                except Exception:
+                    logger.exception("Failed to parse bridge frame")
+                    continue
+
+                if isinstance(message, dict) and message.get('method') == 'event':
+                    self._handle_event(message.get('params') or {})
+                elif isinstance(message, dict) and 'id' in message:
+                    self._handle_response(message)
+
+        except asyncio.CancelledError:
+            # Clean cancellation - expected during stop()
+            raise
+        except Exception as e:
+            logger.error(f"Bridge reader died: {e}")
+        finally:
+            # Fail all pending requests so callers don't hang
+            error = BridgeConnectionError("Bridge process terminated unexpectedly")
+            for request_id, future in list(self.pending_requests.items()):
+                if not future.done():
+                    future.set_exception(error)
+            self.pending_requests.clear()
+
+    async def _drain_stderr(self):
+        """Drain bridge stderr to prevent blocking.
+
+        Uses native asyncio reads - fully cancellable.
+        """
+        if self.process is None or self.process.stderr is None:
+            return
+        try:
+            while True:
+                line = await self.process.stderr.readline()
+                if not line:
+                    break
+                try:
+                    text = line.decode("utf-8", errors="ignore").rstrip()
+                except Exception:
+                    text = str(line).rstrip()
+                if text:
+                    logger.debug(f"[bridge stderr] {text}")
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            logger.debug(f"Bridge stderr drain died: {e}")
+
+    def _handle_event(self, params: Dict[str, Any]):
+        """Handle event notification from bridge."""
+        event_type = params.get('type')
+        callbacks = self.event_callbacks.get(event_type, [])
+
+        if event_type in ('stdout', 'stderr'):
+            data = params.get('data', '')
+            seq = params.get('seq')
+            done = params.get('done')
+
+            # If chunk metadata is present, reassemble to preserve "NDJSON line" semantics.
+            if seq is not None or done is not None:
+                buf = self._stream_buffers.setdefault(event_type, [])
+                if seq == 0 and buf:
+                    # Best-effort flush of previous incomplete sequence.
+                    prev = "".join(buf)
+                    buf.clear()
+                    for callback in callbacks:
+                        try:
+                            callback(prev)
+                        except Exception:
+                            logger.exception("Error in %s callback", event_type)
+
+                buf.append(data)
+                if done:
+                    full = "".join(buf)
+                    buf.clear()
+                    for callback in callbacks:
+                        try:
+                            callback(full)
+                        except Exception:
+                            logger.exception("Error in %s callback", event_type)
+                return
+
+            # No chunk metadata → emit directly.
+            for callback in callbacks:
+                try:
+                    callback(data)
+                except Exception:
+                    logger.exception("Error in %s callback", event_type)
+        elif event_type == 'content':
+            for callback in callbacks:
+                try:
+                    callback(params)
+                except Exception:
+                    logger.exception("Error in content callback")
+
+    def _handle_response(self, message: Dict[str, Any]):
+        """Handle JSON-RPC response."""
+        request_id = message.get('id')
+        if request_id is None or request_id not in self.pending_requests:
+            return
+
+        future = self.pending_requests.pop(request_id)
+
+        if 'error' in message:
+            error = message['error']
+            error_code = error.get('code', -32603)
+            error_message = error.get('message', 'Unknown error')
+
+            # Check for NotFoundError (code -32001 or message pattern)
+            if error_code == -32001 or 'not found' in error_message.lower():
+                future.set_exception(SandboxNotFoundError(error_message))
+            else:
+                future.set_exception(Exception(error_message))
+        else:
+            future.set_result(message.get('result'))
+
+    # =========================================================================
+    # MULTI-INSTANCE METHODS (for Swarm)
+    # =========================================================================
+
+    async def create_instance(
+        self,
+        instance_id: str,
+        params: Dict[str, Any],
+        timeout_s: Optional[float] = None,
+    ) -> Any:
+        """Create a new SwarmKit instance in the bridge."""
+        return await self.call(
+            'create_instance',
+            {'instance_id': instance_id, **params},
+            timeout_s=timeout_s,
+        )
+
+    async def run_on_instance(
+        self,
+        instance_id: str,
+        prompt: str,
+        timeout_ms: Optional[int] = None,
+        call_timeout_s: Optional[float] = None,
+    ) -> Any:
+        """Run prompt on a specific SwarmKit instance."""
+        params = {'instance_id': instance_id, 'prompt': prompt}
+        if timeout_ms is not None:
+            params['timeout_ms'] = timeout_ms
+        return await self.call('run_on_instance', params, timeout_s=call_timeout_s)
+
+    async def get_output_on_instance(
+        self,
+        instance_id: str,
+        recursive: bool = False,
+        timeout_s: Optional[float] = None,
+    ) -> Any:
+        """Get output files from a specific SwarmKit instance."""
+        return await self.call(
+            'get_output_on_instance',
+            {'instance_id': instance_id, 'recursive': recursive},
+            timeout_s=timeout_s,
+        )
+
+    async def kill_instance(
+        self,
+        instance_id: str,
+        timeout_s: Optional[float] = None,
+    ) -> Any:
+        """Kill and remove a specific SwarmKit instance."""
+        return await self.call(
+            'kill_instance',
+            {'instance_id': instance_id},
+            timeout_s=timeout_s,
+        )

swarmkit/config.py

@@ -0,0 +1,92 @@
+"""Configuration types for SwarmKit SDK."""
+
+from dataclasses import dataclass
+from typing import List, Literal, Optional, Protocol, runtime_checkable
+
+
+AgentType = Literal['codex', 'claude', 'gemini', 'qwen']
+WorkspaceMode = Literal['knowledge', 'swe']
+ReasoningEffort = Literal['low', 'medium', 'high', 'xhigh']
+ValidationMode = Literal['strict', 'loose']
+
+
+@dataclass
+class SchemaOptions:
+    """Validation options for schema validation.
+
+    Args:
+        mode: Validation mode - 'strict' (exact types) or 'loose' (coerce types, default)
+    """
+    mode: ValidationMode = 'loose'
+
+
+@dataclass
+class AgentConfig:
+    """Agent configuration.
+
+    All fields are optional - TS SDK resolves defaults from environment:
+    - type defaults to 'claude'
+    - api_key defaults to SWARMKIT_API_KEY env var
+
+    Args:
+        type: Agent type (codex, claude, gemini, qwen) - defaults to 'claude'
+        api_key: SwarmKit API key from https://dashboard.swarmlink.ai (defaults to SWARMKIT_API_KEY env var)
+        model: Model name (optional - uses agent's default if not specified)
+        reasoning_effort: Reasoning effort for Codex models (optional)
+        betas: Beta headers for Claude (Sonnet 4.5 only; e.g. ["context-1m-2025-08-07"] for 1M context)
+    """
+    type: Optional[AgentType] = None
+    api_key: Optional[str] = None
+    model: Optional[str] = None
+    reasoning_effort: Optional[ReasoningEffort] = None
+    betas: Optional[List[str]] = None
+
+
+@runtime_checkable
+class SandboxProvider(Protocol):
+    """Sandbox provider protocol.
+
+    Any sandbox provider must implement this protocol.
+    Currently supported: E2BProvider
+
+    To add a new provider:
+    1. Create a class with `type` and `config` properties
+    2. Add handling in bridge/src/adapter.ts
+    """
+
+    @property
+    def type(self) -> str:
+        """Provider type identifier (e.g., 'e2b')."""
+        ...
+
+    @property
+    def config(self) -> dict:
+        """Provider configuration dict for the bridge."""
+        ...
+
+
+@dataclass
+class E2BProvider:
+    """E2B sandbox provider configuration.
+
+    Args:
+        api_key: E2B API key (defaults to E2B_API_KEY env var)
+        timeout_ms: Sandbox timeout in milliseconds (default: 3600000 = 1 hour)
+    """
+    api_key: Optional[str] = None
+    timeout_ms: int = 3600000
+
+    @property
+    def type(self) -> Literal['e2b']:
+        """Provider type."""
+        return 'e2b'
+
+    @property
+    def config(self) -> dict:
+        """Provider configuration dict."""
+        result = {}
+        if self.api_key:
+            result['apiKey'] = self.api_key
+        if self.timeout_ms:
+            result['defaultTimeoutMs'] = self.timeout_ms
+        return result

swarmkit/pipeline/__init__.py

@@ -0,0 +1,59 @@
+"""Pipeline - Fluent API for chaining Swarm operations."""
+
+from .pipeline import Pipeline, TerminalPipeline
+from .types import (
+    # Step configs
+    MapConfig,
+    FilterConfig,
+    ReduceConfig,
+    # Results
+    StepResult,
+    PipelineResult,
+    # Events
+    PipelineEvents,
+    PipelineEventMap,
+    StepEvent,
+    StepStartEvent,
+    StepCompleteEvent,
+    StepErrorEvent,
+    ItemRetryEvent,
+    WorkerCompleteEvent,
+    VerifierCompleteEvent,
+    CandidateCompleteEvent,
+    JudgeCompleteEvent,
+    # Types
+    EmitOption,
+    EventName,
+    Step,
+    StepType,
+)
+
+__all__ = [
+    # Main classes
+    'Pipeline',
+    'TerminalPipeline',
+    # Step configs
+    'MapConfig',
+    'FilterConfig',
+    'ReduceConfig',
+    # Results
+    'StepResult',
+    'PipelineResult',
+    # Events
+    'PipelineEvents',
+    'PipelineEventMap',
+    'StepEvent',
+    'StepStartEvent',
+    'StepCompleteEvent',
+    'StepErrorEvent',
+    'ItemRetryEvent',
+    'WorkerCompleteEvent',
+    'VerifierCompleteEvent',
+    'CandidateCompleteEvent',
+    'JudgeCompleteEvent',
+    # Types
+    'EmitOption',
+    'EventName',
+    'Step',
+    'StepType',
+]