diff --git a/packages/prime-tunnel/LICENSE b/packages/prime-tunnel/LICENSE new file mode 100644 index 00000000..44a4d7b1 --- /dev/null +++ b/packages/prime-tunnel/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Prime Intellect + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/prime-tunnel/README.md b/packages/prime-tunnel/README.md new file mode 100644 index 00000000..62e7c6fb --- /dev/null +++ b/packages/prime-tunnel/README.md @@ -0,0 +1,51 @@ +# Prime Tunnel SDK + +Expose local services via secure tunnels on Prime infrastructure. + +## Installation + +```bash +uv pip install prime-tunnel +``` + +Or with pip: + +```bash +pip install prime-tunnel +``` + +## Quick Start + +```python +from prime_tunnel import Tunnel + +# Create and start a tunnel +async with Tunnel(local_port=8765) as tunnel: + print(f"Tunnel URL: {tunnel.url}") + # Your local service on port 8765 is now accessible at tunnel.url +``` + +## CLI Usage + +```bash +# Start a tunnel +prime tunnel start --port 8765 + +# List active tunnels +prime tunnel list + +# Get tunnel status +prime tunnel status +``` + +## Documentation + +Full API reference: https://github.com/PrimeIntellect-ai/prime-cli/tree/main/packages/prime-tunnel + +## Related Packages + +- **`prime`** - Full CLI + SDK with pods, inference, and more (includes this package) + +## License + +MIT License - see LICENSE file for details diff --git a/packages/prime-tunnel/pyproject.toml b/packages/prime-tunnel/pyproject.toml new file mode 100644 index 00000000..62808361 --- /dev/null +++ b/packages/prime-tunnel/pyproject.toml @@ -0,0 +1,62 @@ +[project] +name = "prime-tunnel" +# Version is single-sourced from src/prime_tunnel/__init__.py via Hatch +dynamic = ["version"] +description = "Prime Intellect Tunnel SDK - Expose local services via secure tunnels" +readme = "README.md" +requires-python = ">=3.10" +license = {text = "MIT"} +authors = [ + { name = "Prime Intellect", email = "contact@primeintellect.ai" } +] +dependencies = [ + "httpx>=0.25.0", + "pydantic>=2.0.0", + "tenacity>=8.0.0", +] +keywords = ["tunnel", "reverse-proxy", "networking", "frp", "sdk"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Software Development :: Libraries :: Python Modules" +] + +[project.urls] +Homepage = "https://github.com/PrimeIntellect-ai/prime-cli" +Documentation = "https://github.com/PrimeIntellect-ai/prime-cli/tree/main/packages/prime-tunnel" +Repository = "https://github.com/PrimeIntellect-ai/prime-cli.git" + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "pytest-xdist>=3.0.0", + "pytest-asyncio>=0.23.0", + "ruff>=0.13.1", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.version] +path = "src/prime_tunnel/__init__.py" + +[tool.hatch.build.targets.wheel] +packages = ["src/prime_tunnel"] + +[tool.pytest.ini_options] +addopts = "-v" +testpaths = ["tests"] + +[tool.ruff] +line-length = 100 +target-version = "py310" + +[tool.ruff.lint] +extend-select = ["E", "F", "I"] diff --git a/packages/prime-tunnel/src/prime_tunnel/__init__.py b/packages/prime-tunnel/src/prime_tunnel/__init__.py new file mode 100644 index 00000000..5d045e66 --- /dev/null +++ b/packages/prime-tunnel/src/prime_tunnel/__init__.py @@ -0,0 +1,29 @@ +"""Prime Tunnel SDK - Expose local services via secure tunnels.""" + +__version__ = "0.1.0" + +from prime_tunnel.core import Config, TunnelClient +from prime_tunnel.exceptions import ( + TunnelAuthError, + TunnelConnectionError, + TunnelError, + TunnelTimeoutError, +) +from prime_tunnel.models import TunnelInfo +from prime_tunnel.tunnel import Tunnel + +__all__ = [ + "__version__", + # Core + "Config", + "TunnelClient", + # Main interface + "Tunnel", + # Models + "TunnelInfo", + # Exceptions + "TunnelError", + "TunnelAuthError", + "TunnelConnectionError", + "TunnelTimeoutError", +] diff --git a/packages/prime-tunnel/src/prime_tunnel/binary.py b/packages/prime-tunnel/src/prime_tunnel/binary.py new file mode 100644 index 00000000..6910734b --- /dev/null +++ b/packages/prime-tunnel/src/prime_tunnel/binary.py @@ -0,0 +1,155 @@ +import hashlib +import os +import platform +import shutil +import stat +import tarfile +import tempfile +from pathlib import Path + +import httpx + +from prime_tunnel.core.config import Config +from prime_tunnel.exceptions import BinaryDownloadError + +FRPC_VERSION = "0.66.0" + +FRPC_CHECKSUMS = { + ("Darwin", "arm64"): "eb24c3c172a20056d83379496500b92600a992f68e8ae2e27d128ce1f36d7a92", + ("Darwin", "x86_64"): "9558d55a9d8bc40e22018379ea645251f803f9e2d69e7a7a2fd1588f98f8ef43", + ("Linux", "x86_64"): "317a17a7adac2e6bed2d7a83dc077da91ced0d110e1636373ece8ae5ac8b578b", + ("Linux", "aarch64"): "196ddaa51b716c2e99aeb2916b0a2bf55bb317494c4acdcefab36c383de950ba", +} + +FRPC_URLS = { + ( + "Darwin", + "arm64", + ): f"https://github.com/fatedier/frp/releases/download/v{FRPC_VERSION}/frp_{FRPC_VERSION}_darwin_arm64.tar.gz", + ( + "Darwin", + "x86_64", + ): f"https://github.com/fatedier/frp/releases/download/v{FRPC_VERSION}/frp_{FRPC_VERSION}_darwin_amd64.tar.gz", + ( + "Linux", + "x86_64", + ): f"https://github.com/fatedier/frp/releases/download/v{FRPC_VERSION}/frp_{FRPC_VERSION}_linux_amd64.tar.gz", + ( + "Linux", + "aarch64", + ): f"https://github.com/fatedier/frp/releases/download/v{FRPC_VERSION}/frp_{FRPC_VERSION}_linux_arm64.tar.gz", +} + + +def _get_platform_key() -> tuple[str, str]: + system = platform.system() + machine = platform.machine() + + if machine in ("AMD64", "x86_64"): + machine = "x86_64" + elif machine in ("arm64", "aarch64"): + machine = "arm64" if system == "Darwin" else "aarch64" + + return (system, machine) + + +def _verify_checksum(file_path: Path, expected_checksum: str) -> None: + """Verify SHA256 checksum of downloaded file.""" + sha256 = hashlib.sha256() + with open(file_path, "rb") as f: + for chunk in iter(lambda: f.read(8192), b""): + sha256.update(chunk) + actual_checksum = sha256.hexdigest() + if actual_checksum != expected_checksum: + raise BinaryDownloadError( + f"Checksum verification failed: expected {expected_checksum}, got {actual_checksum}" + ) + + +def _download_frpc(dest: Path) -> None: + platform_key = _get_platform_key() + url = FRPC_URLS.get(platform_key) + expected_checksum = FRPC_CHECKSUMS.get(platform_key) + + if not url: + raise BinaryDownloadError(f"Unsupported platform: {platform_key[0]} {platform_key[1]}") + if not expected_checksum: + raise BinaryDownloadError( + f"No checksum available for platform: {platform_key[0]} {platform_key[1]}" + ) + + with tempfile.TemporaryDirectory() as tmpdir: + tmpdir_path = Path(tmpdir) + archive_path = tmpdir_path / "frp.tar.gz" + + try: + with httpx.stream("GET", url, follow_redirects=True, timeout=120.0) as response: + response.raise_for_status() + with open(archive_path, "wb") as f: + for chunk in response.iter_bytes(chunk_size=8192): + f.write(chunk) + + except httpx.HTTPError as e: + raise BinaryDownloadError(f"Failed to download frpc: {e}") from e + + _verify_checksum(archive_path, expected_checksum) + + try: + with tarfile.open(archive_path, "r:gz") as tar: + for member in tar.getmembers(): + if member.name.endswith("/frpc") or member.name == "frpc": + member.name = "frpc" + tar.extract(member, tmpdir_path) + break + else: + raise BinaryDownloadError("frpc binary not found in archive") + + except tarfile.TarError as e: + raise BinaryDownloadError(f"Failed to extract frpc: {e}") from e + + extracted_path = tmpdir_path / "frpc" + if not extracted_path.exists(): + raise BinaryDownloadError("frpc binary not found after extraction") + + dest.parent.mkdir(parents=True, exist_ok=True) + + # Set executable permissions on extracted file before moving + extracted_path.chmod( + extracted_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + ) + + # Copy to temp file in same directory, then rename + # This prevents corruption if multiple processes download simultaneously + temp_dest = dest.parent / f".frpc.{os.getpid()}.tmp" + try: + shutil.copy2(extracted_path, temp_dest) + os.replace(temp_dest, dest) # Atomic on POSIX + finally: + # Clean up temp file if rename failed + if temp_dest.exists(): + temp_dest.unlink() + + +def get_frpc_path() -> Path: + config = Config() + frpc_path = config.bin_dir / "frpc" + version_file = config.bin_dir / ".frpc_version" + + if frpc_path.exists(): + if version_file.exists(): + current_version = version_file.read_text().strip() + if current_version == FRPC_VERSION: + return frpc_path + + _download_frpc(frpc_path) + + # Write to temp file in same directory, then rename to prevent partial reads + temp_version = version_file.parent / f".frpc_version.{os.getpid()}.tmp" + try: + temp_version.write_text(FRPC_VERSION) + os.replace(temp_version, version_file) + finally: + if temp_version.exists(): + temp_version.unlink() + + return frpc_path diff --git a/packages/prime-tunnel/src/prime_tunnel/core/__init__.py b/packages/prime-tunnel/src/prime_tunnel/core/__init__.py new file mode 100644 index 00000000..9e728ea4 --- /dev/null +++ b/packages/prime-tunnel/src/prime_tunnel/core/__init__.py @@ -0,0 +1,6 @@ +"""Prime Tunnel Core - HTTP client and configuration.""" + +from prime_tunnel.core.client import TunnelClient +from prime_tunnel.core.config import Config + +__all__ = ["Config", "TunnelClient"] diff --git a/packages/prime-tunnel/src/prime_tunnel/core/client.py b/packages/prime-tunnel/src/prime_tunnel/core/client.py new file mode 100644 index 00000000..33093113 --- /dev/null +++ b/packages/prime-tunnel/src/prime_tunnel/core/client.py @@ -0,0 +1,258 @@ +import sys +from typing import Any, Dict, Optional + +import httpx +from tenacity import ( + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) + +from prime_tunnel.core.config import Config +from prime_tunnel.exceptions import TunnelAuthError, TunnelError, TunnelTimeoutError +from prime_tunnel.models import TunnelInfo + +# Retry configuration for transient connection errors +RETRYABLE_EXCEPTIONS = ( + httpx.RemoteProtocolError, + httpx.ConnectError, + httpx.PoolTimeout, +) + + +def _default_user_agent() -> str: + """Build default User-Agent string.""" + from prime_tunnel import __version__ + + python_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" + return f"prime-tunnel/{__version__} python/{python_version}" + + +class TunnelClient: + """Client for interacting with Prime Tunnel API.""" + + def __init__( + self, + api_key: Optional[str] = None, + user_agent: Optional[str] = None, + timeout: float = 30.0, + ): + """ + Initialize the tunnel client. + + Args: + api_key: Optional API key (defaults to config) + user_agent: Optional custom User-Agent string + timeout: Request timeout in seconds + """ + self.config = Config() + self.api_key = api_key or self.config.api_key + self.base_url = self.config.base_url + self._timeout = timeout + + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + headers["User-Agent"] = user_agent if user_agent else _default_user_agent() + + self._client: Optional[httpx.AsyncClient] = None + self._headers = headers + + def _check_auth_required(self) -> None: + """Check if API key is configured.""" + if not self.api_key: + raise TunnelError("No API key configured. Set PRIME_API_KEY environment variable.") + + async def _get_client(self) -> httpx.AsyncClient: + """Get or create async HTTP client.""" + if self._client is None or self._client.is_closed: + self._client = httpx.AsyncClient( + timeout=self._timeout, + headers=self._headers, + follow_redirects=True, + ) + return self._client + + async def close(self) -> None: + """Close the HTTP client.""" + if self._client and not self._client.is_closed: + await self._client.aclose() + self._client = None + + @retry( + retry=retry_if_exception_type(RETRYABLE_EXCEPTIONS), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=0.1, min=0.1, max=2), + reraise=True, + ) + async def _request_with_retry( + self, + method: str, + url: str, + json: Optional[Dict[str, Any]] = None, + ) -> httpx.Response: + """Make async HTTP request with retry on transient connection errors.""" + client = await self._get_client() + return await client.request(method, url, json=json) + + async def _handle_response(self, response: httpx.Response, operation: str) -> Dict[str, Any]: + """Handle response and raise appropriate errors.""" + if response.status_code == 401: + raise TunnelAuthError("API key unauthorized. Check PRIME_API_KEY.") + elif response.status_code == 402: + raise TunnelAuthError("Payment required. Check billing status.") + elif response.status_code == 404: + return {} # Handle 404 specially in callers + elif response.status_code >= 400: + try: + error_detail = response.json().get("detail", response.text) + except Exception: + error_detail = response.text + raise TunnelError(f"Failed to {operation}: {error_detail}") + + if response.status_code == 204: + return {} + + return response.json() + + async def create_tunnel( + self, + local_port: int, + name: Optional[str] = None, + ) -> TunnelInfo: + """ + Register a new tunnel with the backend. + + Args: + local_port: Local port the tunnel will forward to + name: Optional friendly name for the tunnel + + Returns: + TunnelInfo with connection details + + Raises: + TunnelAuthError: If authentication fails + TunnelError: If registration fails + """ + self._check_auth_required() + + url = f"{self.base_url}/api/v1/tunnel" + payload: Dict[str, Any] = {"local_port": local_port} + if name: + payload["name"] = name + + try: + response = await self._request_with_retry("POST", url, json=payload) + except httpx.TimeoutException as e: + raise TunnelTimeoutError(f"Request timed out: {e}") from e + except httpx.RequestError as e: + raise TunnelError(f"Failed to connect to API: {e}") from e + + data = await self._handle_response(response, "create tunnel") + if not data: + raise TunnelError("Failed to create tunnel: unexpected empty response") + + return TunnelInfo(**data) + + async def get_tunnel(self, tunnel_id: str) -> Optional[TunnelInfo]: + """ + Get tunnel status by ID. + + Args: + tunnel_id: The tunnel identifier + + Returns: + TunnelInfo if found, None otherwise + """ + self._check_auth_required() + + url = f"{self.base_url}/api/v1/tunnel/{tunnel_id}" + + try: + response = await self._request_with_retry("GET", url) + except httpx.TimeoutException as e: + raise TunnelTimeoutError(f"Request timed out: {e}") from e + except httpx.RequestError as e: + raise TunnelError(f"Failed to connect to API: {e}") from e + + if response.status_code == 404: + return None + + data = await self._handle_response(response, "get tunnel") + return TunnelInfo( + tunnel_id=data["tunnel_id"], + hostname=data["hostname"], + url=data["url"], + frp_token="", # Token not returned on status check + server_host="", + server_port=7000, + expires_at=data["expires_at"], + ) + + async def delete_tunnel(self, tunnel_id: str) -> bool: + """ + Delete a tunnel. + + Args: + tunnel_id: The tunnel identifier + + Returns: + True if deleted successfully + """ + self._check_auth_required() + + url = f"{self.base_url}/api/v1/tunnel/{tunnel_id}" + + try: + response = await self._request_with_retry("DELETE", url) + except httpx.TimeoutException as e: + raise TunnelTimeoutError(f"Request timed out: {e}") from e + except httpx.RequestError as e: + raise TunnelError(f"Failed to connect to API: {e}") from e + + if response.status_code == 404: + return False + + await self._handle_response(response, "delete tunnel") + return True + + async def list_tunnels(self) -> list[TunnelInfo]: + """ + List all tunnels for the current user. + + Returns: + List of TunnelInfo objects + """ + self._check_auth_required() + + url = f"{self.base_url}/api/v1/tunnel" + + try: + response = await self._request_with_retry("GET", url) + except httpx.TimeoutException as e: + raise TunnelTimeoutError(f"Request timed out: {e}") from e + except httpx.RequestError as e: + raise TunnelError(f"Failed to connect to API: {e}") from e + + data = await self._handle_response(response, "list tunnels") + tunnels = [] + for t in data.get("tunnels", []): + tunnels.append( + TunnelInfo( + tunnel_id=t["tunnel_id"], + hostname=t["hostname"], + url=t["url"], + frp_token="", + server_host="", + server_port=7000, + expires_at=t["expires_at"], + ) + ) + return tunnels + + async def __aenter__(self) -> "TunnelClient": + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + await self.close() diff --git a/packages/prime-tunnel/src/prime_tunnel/core/config.py b/packages/prime-tunnel/src/prime_tunnel/core/config.py new file mode 100644 index 00000000..b2982822 --- /dev/null +++ b/packages/prime-tunnel/src/prime_tunnel/core/config.py @@ -0,0 +1,70 @@ +import json +import os +from pathlib import Path +from typing import Optional + + +class Config: + """Minimal configuration class for Prime Tunnel SDK. + + Reads from ~/.prime/config.json and environment variables. + This is a simplified version that doesn't write configs. + """ + + DEFAULT_BASE_URL: str = "https://api.primeintellect.ai" + + def __init__(self) -> None: + self.config_dir = Path.home() / ".prime" + self.config_file = self.config_dir / "config.json" + self._load_config() + + def _load_config(self) -> None: + """Load configuration from file.""" + if self.config_file.exists(): + try: + config_data = json.loads(self.config_file.read_text()) + self.config = config_data if isinstance(config_data, dict) else {} + except (json.JSONDecodeError, IOError): + self.config = {} + else: + self.config = {} + + @staticmethod + def _strip_api_v1(url: str) -> str: + return url.rstrip("/").removesuffix("/api/v1") + + @property + def api_key(self) -> str: + """Get API key with precedence: env > file > empty.""" + return os.getenv("PRIME_API_KEY") or self.config.get("api_key", "") + + @property + def team_id(self) -> Optional[str]: + """Get team ID with precedence: env > file > None.""" + team_id = os.getenv("PRIME_TEAM_ID") + if team_id is not None: + return team_id + return self.config.get("team_id") or None + + @property + def user_id(self) -> Optional[str]: + """Get user ID with precedence: env > file > None.""" + user_id = os.getenv("PRIME_USER_ID") + if user_id is not None: + return user_id + return self.config.get("user_id") or None + + @property + def base_url(self) -> str: + """Get API base URL with precedence: env > file > default.""" + env_val = os.getenv("PRIME_API_BASE_URL") or os.getenv("PRIME_BASE_URL") + if env_val: + return self._strip_api_v1(env_val) + return self._strip_api_v1(self.config.get("base_url", self.DEFAULT_BASE_URL)) + + @property + def bin_dir(self) -> Path: + """Directory for binary files (frpc).""" + path = self.config_dir / "bin" + path.mkdir(parents=True, exist_ok=True) + return path diff --git a/packages/prime-tunnel/src/prime_tunnel/exceptions.py b/packages/prime-tunnel/src/prime_tunnel/exceptions.py new file mode 100644 index 00000000..e95d2992 --- /dev/null +++ b/packages/prime-tunnel/src/prime_tunnel/exceptions.py @@ -0,0 +1,28 @@ +class TunnelError(Exception): + """Base exception for tunnel errors.""" + + pass + + +class TunnelAuthError(TunnelError): + """Authentication failed when registering tunnel.""" + + pass + + +class TunnelConnectionError(TunnelError): + """Failed to establish tunnel connection.""" + + pass + + +class TunnelTimeoutError(TunnelError): + """Tunnel operation timed out.""" + + pass + + +class BinaryDownloadError(TunnelError): + """Failed to download frpc binary.""" + + pass diff --git a/packages/prime-tunnel/src/prime_tunnel/models.py b/packages/prime-tunnel/src/prime_tunnel/models.py new file mode 100644 index 00000000..7b1c8ee5 --- /dev/null +++ b/packages/prime-tunnel/src/prime_tunnel/models.py @@ -0,0 +1,19 @@ +from datetime import datetime + +from pydantic import BaseModel, Field + + +class TunnelInfo(BaseModel): + """Information about a registered tunnel.""" + + tunnel_id: str = Field(..., description="Unique tunnel identifier") + hostname: str = Field(..., description="Tunnel hostname") + url: str = Field(..., description="Full HTTPS URL") + frp_token: str = Field(..., description="Authentication token for frpc") + binding_secret: str = Field("", description="Per-tunnel secret for frpc metadata") + server_host: str = Field(..., description="frps server hostname") + server_port: int = Field(7000, description="frps server port") + expires_at: datetime = Field(..., description="Token expiration time") + + class Config: + from_attributes = True diff --git a/packages/prime-tunnel/src/prime_tunnel/py.typed b/packages/prime-tunnel/src/prime_tunnel/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/packages/prime-tunnel/src/prime_tunnel/tunnel.py b/packages/prime-tunnel/src/prime_tunnel/tunnel.py new file mode 100644 index 00000000..86994609 --- /dev/null +++ b/packages/prime-tunnel/src/prime_tunnel/tunnel.py @@ -0,0 +1,367 @@ +import asyncio +import fcntl +import os +import subprocess +import tempfile +import threading +import time +from pathlib import Path +from typing import Optional + +from prime_tunnel.binary import get_frpc_path +from prime_tunnel.core.client import TunnelClient +from prime_tunnel.exceptions import TunnelConnectionError, TunnelError, TunnelTimeoutError +from prime_tunnel.models import TunnelInfo + + +class Tunnel: + """Tunnel interface for exposing local services.""" + + def __init__( + self, + local_port: int, + local_addr: str = "127.0.0.1", + name: Optional[str] = None, + connection_timeout: float = 30.0, + log_level: str = "info", + ): + """ + Initialize a tunnel. + + Args: + local_port: Local port to tunnel + local_addr: Local address to tunnel (default: 127.0.0.1) + name: Optional friendly name for the tunnel + connection_timeout: Timeout for establishing connection (seconds) + log_level: frpc log level (trace, debug, info, warn, error) + """ + self.local_port = local_port + self.local_addr = local_addr + self.name = name + self.connection_timeout = connection_timeout + self.log_level = log_level + + self._client = TunnelClient() + self._process: Optional[subprocess.Popen] = None + self._tunnel_info: Optional[TunnelInfo] = None + self._config_file: Optional[Path] = None + self._started = False + self._output_lines: list[str] = [] + + @property + def tunnel_id(self) -> Optional[str]: + """Get the tunnel ID.""" + return self._tunnel_info.tunnel_id if self._tunnel_info else None + + @property + def url(self) -> Optional[str]: + """Get the tunnel URL.""" + return self._tunnel_info.url if self._tunnel_info else None + + @property + def hostname(self) -> Optional[str]: + """Get the tunnel hostname.""" + return self._tunnel_info.hostname if self._tunnel_info else None + + @property + def is_running(self) -> bool: + """Check if the tunnel is running.""" + if self._process is None: + return False + return self._process.poll() is None + + async def start(self) -> str: + """ + Start the tunnel. + + Returns: + The tunnel URL + + Raises: + TunnelError: If tunnel registration fails + TunnelConnectionError: If frpc fails to connect + TunnelTimeoutError: If connection times out + """ + if self._started: + raise TunnelError("Tunnel is already started") + + # 1. Get frpc binary + frpc_path = await asyncio.to_thread(get_frpc_path) + + # 2. Register tunnel with backend + try: + self._tunnel_info = await self._client.create_tunnel( + local_port=self.local_port, + name=self.name, + ) + except BaseException as e: + await self._cleanup() + if isinstance(e, asyncio.CancelledError): + raise + raise TunnelError(f"Failed to register tunnel: {e}") from e + + # 3. Generate frpc config + try: + self._config_file = self._write_frpc_config() + except BaseException as e: + await self._cleanup() + if isinstance(e, asyncio.CancelledError): + raise + raise TunnelError(f"Failed to write frpc config: {e}") from e + + # 4. Start frpc process + try: + self._process = subprocess.Popen( + [str(frpc_path), "-c", str(self._config_file)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + except BaseException as e: + await self._cleanup() + if isinstance(e, asyncio.CancelledError): + raise + raise TunnelConnectionError(f"Failed to start frpc: {e}") from e + + # 5. Wait for connection + try: + await self._wait_for_connection() + except BaseException: + await self._cleanup() + raise + + # 6. Start background thread to drain pipes (prevents buffer exhaustion) + try: + self._start_pipe_drain() + except BaseException as e: + await self._cleanup() + if isinstance(e, asyncio.CancelledError): + raise + raise TunnelConnectionError(f"Failed to start pipe drain: {e}") from e + + self._started = True + + return self.url + + async def stop(self) -> None: + """Stop the tunnel and cleanup resources.""" + if not self._started: + return + + await self._cleanup() + self._started = False + + async def _cleanup(self) -> None: + """Clean up tunnel resources.""" + # Stop frpc process (this will cause drain threads to exit via EOF) + if self._process is not None: + try: + self._process.terminate() + try: + self._process.wait(timeout=5) + except subprocess.TimeoutExpired: + self._process.kill() + self._process.wait(timeout=2) + except Exception: + pass + finally: + self._process = None + + # Delete tunnel registration + if self._tunnel_info is not None: + try: + await self._client.delete_tunnel(self._tunnel_info.tunnel_id) + except Exception: + pass + finally: + self._tunnel_info = None + + # Clean up config file + if self._config_file is not None: + try: + if self._config_file.exists(): + self._config_file.unlink() + except Exception: + pass + finally: + self._config_file = None + + # Close HTTP client + try: + await self._client.close() + except Exception: + pass + + def _start_pipe_drain(self) -> None: + """Start background threads to drain subprocess pipes. + + This prevents the pipe buffer from filling up and blocking frpc + when it produces output (logs, reconnection attempts, etc.). + """ + if self._process is None: + return + + def drain_pipe(pipe): + """Read and discard output from a pipe until EOF.""" + if pipe is None: + return + try: + for _ in pipe: + pass # Discard all output + except (OSError, ValueError): + pass # Pipe closed + + # Use separate threads for stdout/stderr to avoid blocking on one + stdout_thread = threading.Thread( + target=drain_pipe, args=(self._process.stdout,), daemon=True + ) + stderr_thread = threading.Thread( + target=drain_pipe, args=(self._process.stderr,), daemon=True + ) + stdout_thread.start() + stderr_thread.start() + + def _write_frpc_config(self) -> Path: + """Generate and write frpc configuration file.""" + if self._tunnel_info is None: + raise TunnelError("Tunnel not registered") + + server_host = self._tunnel_info.server_host + server_port = self._tunnel_info.server_port + + # Generate config content + config = f"""# Prime Tunnel frpc configuration +# Tunnel ID: {self._tunnel_info.tunnel_id} + +serverAddr = "{server_host}" +serverPort = {server_port} + +# Authentication +user = "{self._tunnel_info.tunnel_id}" +auth.method = "token" +auth.token = "{self._tunnel_info.frp_token}" + +# Per-tunnel binding secret +metadatas.binding_secret = "{self._tunnel_info.binding_secret}" + +# Transport settings +transport.tcpMux = true +transport.tcpMuxKeepaliveInterval = 30 +transport.poolCount = 5 + +# Logging - always use console so we can detect connection via stdout +log.to = "console" +log.level = "{self.log_level}" + +# HTTP proxy configuration +[[proxies]] +name = "http" +type = "http" +localIP = "{self.local_addr}" +localPort = {self.local_port} +subdomain = "{self._tunnel_info.tunnel_id}" +""" + + # Write to temp file + config_dir = Path(tempfile.gettempdir()) / "prime-tunnel" + config_dir.mkdir(parents=True, exist_ok=True) + config_file = config_dir / f"{self._tunnel_info.tunnel_id}.toml" + + # Create file with 0600 permissions + fd = os.open(str(config_file), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + try: + os.write(fd, config.encode()) + finally: + os.close(fd) + + return config_file + + async def _wait_for_connection(self) -> None: + """Wait for frpc to establish connection.""" + start_time = time.time() + self._output_lines = [] + + while time.time() - start_time < self.connection_timeout: + if self._process is None: + raise TunnelConnectionError("frpc process not running") + + return_code = self._process.poll() + if return_code is not None: + remaining_output = [] + if self._process.stdout: + remaining_output.extend(self._process.stdout.readlines()) + if self._process.stderr: + remaining_output.extend(self._process.stderr.readlines()) + self._output_lines.extend(line.strip() for line in remaining_output if line.strip()) + + # Build detailed error message + output_text = ( + "\n".join(self._output_lines) if self._output_lines else "(no output captured)" + ) + raise TunnelConnectionError( + f"frpc exited with code {return_code}\n" + f"--- frpc output ---\n{output_text}\n-------------------" + ) + + if os.name == "posix": + # Set both pipes to non-blocking mode to drain them without deadlock + pipes_to_drain = [] + original_flags = {} + + for pipe in (self._process.stdout, self._process.stderr): + if pipe: + fd = pipe.fileno() + fl = fcntl.fcntl(fd, fcntl.F_GETFL) + original_flags[fd] = fl + fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK) + pipes_to_drain.append(pipe) + + try: + # Drain both stdout and stderr to prevent buffer exhaustion + for pipe in pipes_to_drain: + try: + while True: + line = pipe.readline() + if not line: + break + line = line.strip() + if line: + self._output_lines.append(line) + # Check for success/failure indicators + if "start proxy success" in line.lower(): + return + if "login failed" in line.lower(): + raise TunnelConnectionError(f"frpc login failed: {line}") + if "authorization failed" in line.lower(): + raise TunnelConnectionError( + f"frpc authorization failed: {line}" + ) + except (BlockingIOError, IOError): + pass # No more data available on this pipe + finally: + # Restore original flags + for fd, fl in original_flags.items(): + try: + fcntl.fcntl(fd, fcntl.F_SETFL, fl) + except (OSError, ValueError): + pass # Pipe may have closed + + await asyncio.sleep(0.1) + + # Timeout - include any captured output + output_text = ( + "\n".join(self._output_lines) if self._output_lines else "(no output captured)" + ) + raise TunnelTimeoutError( + f"Tunnel connection timed out after {self.connection_timeout}s\n" + f"--- frpc output ---\n{output_text}\n-------------------" + ) + + async def __aenter__(self) -> "Tunnel": + """Async context manager entry.""" + await self.start() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + """Async context manager exit.""" + await self.stop() diff --git a/packages/prime-tunnel/tests/test_tunnel.py b/packages/prime-tunnel/tests/test_tunnel.py new file mode 100644 index 00000000..aff22966 --- /dev/null +++ b/packages/prime-tunnel/tests/test_tunnel.py @@ -0,0 +1,58 @@ +from prime_tunnel import Config, Tunnel, TunnelClient + + +def test_tunnel_init(): + """Test Tunnel initialization.""" + tunnel = Tunnel(local_port=8080) + assert tunnel.local_port == 8080 + assert tunnel.local_addr == "127.0.0.1" + assert tunnel.name is None + assert not tunnel.is_running + + +def test_tunnel_init_with_name(): + """Test Tunnel initialization with name.""" + tunnel = Tunnel(local_port=9000, name="my-tunnel") + assert tunnel.local_port == 9000 + assert tunnel.name == "my-tunnel" + + +def test_config_default_base_url(): + """Test Config default base URL.""" + config = Config() + assert config.DEFAULT_BASE_URL == "https://api.primeintellect.ai" + + +def test_config_base_url_from_env(monkeypatch): + """Test Config base_url from environment variable.""" + monkeypatch.setenv("PRIME_BASE_URL", "https://custom.example.com") + config = Config() + assert config.base_url == "https://custom.example.com" + + +def test_config_base_url_strips_api_v1(monkeypatch): + """Test Config strips /api/v1 from base URL.""" + monkeypatch.setenv("PRIME_BASE_URL", "https://example.com/api/v1") + config = Config() + assert config.base_url == "https://example.com" + + +def test_config_api_key_from_env(monkeypatch): + """Test Config api_key from environment variable.""" + monkeypatch.setenv("PRIME_API_KEY", "test-key-123") + config = Config() + assert config.api_key == "test-key-123" + + +def test_config_bin_dir(): + """Test Config bin_dir property.""" + config = Config() + assert config.bin_dir.name == "bin" + assert ".prime" in str(config.bin_dir) + + +def test_tunnel_client_init(): + """Test TunnelClient initialization.""" + client = TunnelClient(api_key="test-key") + assert client.api_key == "test-key" + assert client.base_url == Config.DEFAULT_BASE_URL diff --git a/packages/prime/pyproject.toml b/packages/prime/pyproject.toml index e1e4267f..2d78900c 100644 --- a/packages/prime/pyproject.toml +++ b/packages/prime/pyproject.toml @@ -12,6 +12,7 @@ authors = [ dependencies = [ "prime-sandboxes>=0.1.0", "prime-evals>=0.1.3", + "prime-tunnel>=0.1.0", "httpx>=0.25.0", "pydantic>=2.0.0", "typer>=0.9.0", @@ -55,6 +56,7 @@ dev = [ [tool.uv.sources] prime-sandboxes = { workspace = true } prime-evals = { workspace = true } +prime-tunnel = { workspace = true } [build-system] requires = ["hatchling"] diff --git a/packages/prime/src/prime_cli/commands/tunnel.py b/packages/prime/src/prime_cli/commands/tunnel.py new file mode 100644 index 00000000..d06280da --- /dev/null +++ b/packages/prime/src/prime_cli/commands/tunnel.py @@ -0,0 +1,151 @@ +import asyncio +import signal +from typing import Optional + +import typer +from prime_tunnel import Tunnel +from prime_tunnel.core.client import TunnelClient +from rich.console import Console +from rich.table import Table + +app = typer.Typer(help="Manage tunnels for exposing local services", no_args_is_help=True) +console = Console() + + +@app.command("start") +def start_tunnel( + port: int = typer.Option(8765, "--port", "-p", help="Local port to tunnel"), + name: Optional[str] = typer.Option(None, "--name", "-n", help="Friendly name for the tunnel"), +) -> None: + """Start a tunnel to expose a local port.""" + + async def run_tunnel(): + tunnel = Tunnel(local_port=port, name=name) + + shutdown_event = asyncio.Event() + + def signal_handler(): + console.print("\n[yellow]Shutting down tunnel...[/yellow]") + shutdown_event.set() + + loop = asyncio.get_event_loop() + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, signal_handler) + except NotImplementedError: + # Windows doesn't support add_signal_handler + pass + + try: + url = await tunnel.start() + console.print("\n[green]Tunnel started successfully![/green]") + console.print(f"[bold]URL:[/bold] {url}") + console.print(f"[bold]Tunnel ID:[/bold] {tunnel.tunnel_id}") + console.print(f"\n[dim]Forwarding to localhost:{port}[/dim]") + console.print("[dim]Press Ctrl+C to stop the tunnel[/dim]\n") + + await shutdown_event.wait() + + except Exception as e: + console.print(f"[red]Error:[/red] {e}", style="bold") + raise typer.Exit(1) + finally: + await tunnel.stop() + console.print("[green]Tunnel stopped[/green]") + + try: + asyncio.run(run_tunnel()) + except KeyboardInterrupt: + pass + + +@app.command("list") +def list_tunnels() -> None: + """List active tunnels.""" + + async def fetch_tunnels(): + client = TunnelClient() + try: + tunnels = await client.list_tunnels() + return tunnels + finally: + await client.close() + + try: + tunnels = asyncio.run(fetch_tunnels()) + except Exception as e: + console.print(f"[red]Error:[/red] {e}", style="bold") + raise typer.Exit(1) + + if not tunnels: + console.print("[dim]No active tunnels[/dim]") + return + + table = Table(title="Active Tunnels") + table.add_column("Tunnel ID", style="cyan") + table.add_column("URL", style="green") + table.add_column("Expires At") + + for tunnel in tunnels: + table.add_row( + tunnel.tunnel_id, + tunnel.url, + str(tunnel.expires_at), + ) + + console.print(table) + + +@app.command("status") +def tunnel_status( + tunnel_id: str = typer.Argument(..., help="Tunnel ID to check"), +) -> None: + """Get status of a specific tunnel.""" + + async def fetch_status(): + client = TunnelClient() + try: + return await client.get_tunnel(tunnel_id) + finally: + await client.close() + + try: + tunnel = asyncio.run(fetch_status()) + except Exception as e: + console.print(f"[red]Error:[/red] {e}", style="bold") + raise typer.Exit(1) + + if not tunnel: + console.print(f"[red]Tunnel not found:[/red] {tunnel_id}") + raise typer.Exit(1) + + console.print(f"[bold]Tunnel ID:[/bold] {tunnel.tunnel_id}") + console.print(f"[bold]URL:[/bold] {tunnel.url}") + console.print(f"[bold]Hostname:[/bold] {tunnel.hostname}") + console.print(f"[bold]Expires At:[/bold] {tunnel.expires_at}") + + +@app.command("stop") +def stop_tunnel( + tunnel_id: str = typer.Argument(..., help="Tunnel ID to stop"), +) -> None: + """Stop and delete a tunnel.""" + + async def delete_tunnel(): + client = TunnelClient() + try: + return await client.delete_tunnel(tunnel_id) + finally: + await client.close() + + try: + success = asyncio.run(delete_tunnel()) + except Exception as e: + console.print(f"[red]Error:[/red] {e}", style="bold") + raise typer.Exit(1) + + if success: + console.print(f"[green]Tunnel deleted:[/green] {tunnel_id}") + else: + console.print(f"[red]Tunnel not found:[/red] {tunnel_id}") + raise typer.Exit(1) diff --git a/packages/prime/src/prime_cli/main.py b/packages/prime/src/prime_cli/main.py index bbafcd10..20c92431 100644 --- a/packages/prime/src/prime_cli/main.py +++ b/packages/prime/src/prime_cli/main.py @@ -19,6 +19,7 @@ from .commands.rl import app as rl_app from .commands.sandbox import app as sandbox_app from .commands.teams import app as teams_app +from .commands.tunnel import app as tunnel_app from .commands.upgrade import app as upgrade_app from .commands.whoami import app as whoami_app from .core import Config @@ -51,6 +52,7 @@ app.add_typer(sandbox_app, name="sandbox", rich_help_panel="Compute") app.add_typer(images_app, name="images", rich_help_panel="Compute") app.add_typer(registry_app, name="registry", rich_help_panel="Compute") +app.add_typer(tunnel_app, name="tunnel", rich_help_panel="Compute") app.add_typer(inference_app, name="inference", rich_help_panel="Compute") diff --git a/uv.lock b/uv.lock index b2bdb53c..8f6773e6 100644 --- a/uv.lock +++ b/uv.lock @@ -13,6 +13,7 @@ members = [ "prime-evals", "prime-mcp-server", "prime-sandboxes", + "prime-tunnel", ] [manifest.dependency-groups] @@ -1661,6 +1662,7 @@ dependencies = [ { name = "httpx" }, { name = "prime-evals" }, { name = "prime-sandboxes" }, + { name = "prime-tunnel" }, { name = "pydantic" }, { name = "rich" }, { name = "toml" }, @@ -1683,6 +1685,7 @@ requires-dist = [ { name = "httpx", specifier = ">=0.25.0" }, { name = "prime-evals", editable = "packages/prime-evals" }, { name = "prime-sandboxes", editable = "packages/prime-sandboxes" }, + { name = "prime-tunnel", editable = "packages/prime-tunnel" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, { name = "rich", specifier = ">=13.3.1" }, @@ -1778,6 +1781,35 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "prime-tunnel" +source = { editable = "packages/prime-tunnel" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, + { name = "tenacity" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, + { name = "pytest-asyncio" }, + { name = "pytest-xdist" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.25.0" }, + { name = "pydantic", specifier = ">=2.0.0" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, + { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0" }, + { name = "pytest-xdist", marker = "extra == 'dev'", specifier = ">=3.0.0" }, + { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.13.1" }, + { name = "tenacity", specifier = ">=8.0.0" }, +] +provides-extras = ["dev"] + [[package]] name = "propcache" version = "0.4.1"