synapse-sdk 1.0.0a11__py3-none-any.whl → 2026.1.1b2__py3-none-any.whl

synapse_sdk/clients/agent/__init__.py

@@ -1,24 +1,42 @@
-from synapse_sdk.clients.agent.core import CoreClientMixin
+from __future__ import annotations
+
+from synapse_sdk.clients.agent.async_ray import AsyncRayClientMixin
+from synapse_sdk.clients.agent.container import ContainerClientMixin
+from synapse_sdk.clients.agent.plugin import PluginClientMixin
 from synapse_sdk.clients.agent.ray import RayClientMixin
-from synapse_sdk.clients.agent.service import ServiceClientMixin
-from synapse_sdk.clients.exceptions import ClientError
+from synapse_sdk.clients.base import AsyncBaseClient, BaseClient
+
 
+class AgentClient(ContainerClientMixin, PluginClientMixin, RayClientMixin, BaseClient):
+    """Sync client for synapse-agent API."""
 
-class AgentClient(CoreClientMixin, RayClientMixin, ServiceClientMixin):
     name = 'Agent'
-    agent_token = None
-    user_token = None
-    tenant = None
-    long_poll_handler = None
 
-    def __init__(self, base_url, agent_token, user_token=None, tenant=None, long_poll_handler=None):
-        super().__init__(base_url)
+    def __init__(
+        self,
+        base_url: str,
+        agent_token: str,
+        *,
+        user_token: str | None = None,
+        tenant: str | None = None,
+        timeout: dict[str, int] | None = None,
+    ):
+        """Initialize the agent client.
+
+        Args:
+            base_url: Agent API base URL.
+            agent_token: Agent authentication token.
+            user_token: Optional user token for impersonation.
+            tenant: Optional tenant identifier.
+            timeout: Request timeout dict with 'connect' and 'read' keys.
+        """
+        super().__init__(base_url, timeout=timeout or {'connect': 3, 'read': 10})
         self.agent_token = agent_token
         self.user_token = user_token
         self.tenant = tenant
-        self.long_poll_handler = long_poll_handler
 
-    def _get_headers(self):
+    def _get_headers(self) -> dict[str, str]:
+        """Return authentication headers."""
         headers = {'Authorization': self.agent_token}
         if self.user_token:
             headers['SYNAPSE-User'] = f'Token {self.user_token}'
@@ -26,26 +44,63 @@ class AgentClient(CoreClientMixin, RayClientMixin, ServiceClientMixin):
             headers['SYNAPSE-Tenant'] = f'Token {self.tenant}'
         return headers
 
-    def _request(self, method, path, **kwargs):
-        if self.long_poll_handler:
-            return self._request_long_poll(method, path, **kwargs)
-        return super()._request(method, path, **kwargs)
+    def health_check(self) -> dict:
+        """Check agent health."""
+        path = 'health/'
+        return self._get(path)
+
+
+class AsyncAgentClient(AsyncRayClientMixin, AsyncBaseClient):
+    """Async client for synapse-agent API.
+
+    Provides async/await interface for all agent operations including
+    WebSocket and HTTP streaming for job log tailing.
+
+    Example:
+        >>> async with AsyncAgentClient(base_url, agent_token) as client:
+        ...     jobs = await client.list_jobs()
+        ...     async for line in client.tail_job_logs('job-123'):
+        ...         print(line)
+    """
+
+    name = 'Agent'
 
-    def _request_long_poll(self, method, path, **kwargs):
-        headers = self._get_headers()
+    def __init__(
+        self,
+        base_url: str,
+        agent_token: str,
+        *,
+        user_token: str | None = None,
+        tenant: str | None = None,
+        timeout: float | None = None,
+    ):
+        """Initialize the async agent client.
 
-        if kwargs.get('files'):
-            raise ClientError(400, 'file is not allowed when long polling')
+        Args:
+            base_url: Agent API base URL.
+            agent_token: Agent authentication token.
+            user_token: Optional user token for impersonation.
+            tenant: Optional tenant identifier.
+            timeout: Request timeout in seconds.
+        """
+        super().__init__(base_url, timeout=timeout)
+        self.agent_token = agent_token
+        self.user_token = user_token
+        self.tenant = tenant
 
-        headers['Content-Type'] = 'application/json'
+    def _get_headers(self) -> dict[str, str]:
+        """Return authentication headers."""
+        headers = {'Authorization': self.agent_token}
+        if self.user_token:
+            headers['SYNAPSE-User'] = f'Token {self.user_token}'
+        if self.tenant:
+            headers['SYNAPSE-Tenant'] = f'Token {self.tenant}'
+        return headers
 
-        request_id = self.long_poll_handler.set_request({'method': method, 'path': path, 'headers': headers, **kwargs})
-        try:
-            response = self.long_poll_handler.get_response(request_id)
-        except TimeoutError:
-            raise ClientError(408, f'{self.name} is not responding')
+    async def health_check(self) -> dict:
+        """Check agent health."""
+        path = 'health/'
+        return await self._get(path)
 
-        if 400 <= response['status'] < 600:
-            raise ClientError(response['status'], response.json() if response['status'] == 400 else response['reason'])
 
-        return response['data']
+__all__ = ['AgentClient', 'AsyncAgentClient']

synapse_sdk/clients/agent/async_ray.py

@@ -0,0 +1,289 @@
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, AsyncGenerator, Literal
+
+from synapse_sdk.exceptions import ClientError
+from synapse_sdk.utils.network import (
+    StreamLimits,
+    http_to_websocket_url,
+    sanitize_error_message,
+    validate_resource_id,
+    validate_timeout,
+)
+
+if TYPE_CHECKING:
+    from synapse_sdk.clients.protocols import AsyncClientProtocol
+
+StreamProtocol = Literal['websocket', 'http', 'auto']
+
+
+class AsyncRayClientMixin:
+    """Async mixin for Ray cluster management endpoints."""
+
+    _stream_limits: StreamLimits | None = None
+
+    @property
+    def stream_limits(self) -> StreamLimits:
+        """Get stream limits configuration."""
+        if self._stream_limits is None:
+            self._stream_limits = StreamLimits()
+        return self._stream_limits
+
+    @stream_limits.setter
+    def stream_limits(self, value: StreamLimits) -> None:
+        """Set stream limits configuration."""
+        self._stream_limits = value
+
+    # -------------------------------------------------------------------------
+    # Jobs
+    # -------------------------------------------------------------------------
+
+    async def list_jobs(self: AsyncClientProtocol) -> list[dict]:
+        """List all Ray jobs."""
+        return await self._get('ray/jobs/')
+
+    async def get_job(self: AsyncClientProtocol, job_id: str) -> dict:
+        """Get a Ray job by ID."""
+        return await self._get(f'ray/jobs/{job_id}')
+
+    async def get_job_logs(self: AsyncClientProtocol, job_id: str) -> str:
+        """Get all logs for a job (non-streaming)."""
+        return await self._get(f'ray/jobs/{job_id}/logs')
+
+    async def stop_job(self: AsyncClientProtocol, job_id: str) -> dict:
+        """Stop a running job."""
+        return await self._post(f'ray/jobs/{job_id}/stop/')
+
+    async def websocket_tail_job_logs(
+        self: AsyncClientProtocol,
+        job_id: str,
+        timeout: float = 30.0,
+    ) -> AsyncGenerator[str, None]:
+        """Stream job logs via WebSocket protocol (async).
+
+        Establishes an async WebSocket connection for real-time log streaming.
+
+        Args:
+            job_id: The Ray job ID to tail logs for.
+            timeout: Connection and read timeout in seconds.
+
+        Yields:
+            Log message strings.
+
+        Raises:
+            ClientError: On connection, protocol, or validation errors.
+
+        Example:
+            >>> async for line in client.websocket_tail_job_logs('raysubmit_abc123'):
+            ...     print(line)
+        """
+        validated_id = validate_resource_id(job_id, 'job')
+        validated_timeout = validate_timeout(timeout)
+
+        url = self._get_url(f'ray/jobs/{validated_id}/logs/')
+        ws_url = http_to_websocket_url(f'{self.base_url}/{url}')
+        headers = self._get_headers()
+
+        try:
+            import websockets
+        except ImportError:
+            raise ClientError(500, 'websockets package required for async WebSocket streaming')
+
+        limits = self.stream_limits
+        message_count = 0
+
+        try:
+            async with websockets.connect(
+                ws_url,
+                additional_headers=headers,
+                close_timeout=validated_timeout,
+                ping_timeout=validated_timeout,
+            ) as ws:
+                async for data in ws:
+                    if not data:
+                        break
+
+                    message_count += 1
+                    if message_count > limits.max_messages:
+                        raise ClientError(429, 'Stream message limit exceeded')
+
+                    if len(str(data)) > limits.max_message_size:
+                        continue
+
+                    try:
+                        event = json.loads(data)
+                    except json.JSONDecodeError:
+                        event = {'message': data}
+
+                    event_type = event.get('type')
+                    if event_type == 'error':
+                        raise ClientError(500, event.get('message', 'Unknown error'))
+                    elif event_type == 'complete':
+                        return
+
+                    if msg := event.get('message'):
+                        yield msg
+
+        except ClientError:
+            raise
+        except Exception as e:
+            if 'ConnectionClosed' in type(e).__name__:
+                return  # Normal close
+            raise ClientError(503, sanitize_error_message(str(e), 'WebSocket error'))
+
+    async def stream_tail_job_logs(
+        self: AsyncClientProtocol,
+        job_id: str,
+        timeout: float = 30.0,
+    ) -> AsyncGenerator[str, None]:
+        """Stream job logs via HTTP chunked transfer (async).
+
+        Uses HTTP streaming as an alternative when WebSocket is unavailable.
+
+        Args:
+            job_id: The Ray job ID to tail logs for.
+            timeout: Connection timeout in seconds.
+
+        Yields:
+            Log lines as strings.
+
+        Raises:
+            ClientError: On connection, protocol, or validation errors.
+
+        Example:
+            >>> async for line in client.stream_tail_job_logs('raysubmit_abc123'):
+            ...     print(line)
+        """
+        validated_id = validate_resource_id(job_id, 'job')
+        validated_timeout = validate_timeout(timeout)
+
+        url = self._get_url(f'ray/jobs/{validated_id}/logs/')
+        headers = self._get_headers()
+
+        client = await self._get_client()
+        limits = self.stream_limits
+        line_count = 0
+        total_bytes = 0
+
+        try:
+            async with client.stream(
+                'GET',
+                url,
+                headers=headers,
+                timeout=validated_timeout,
+            ) as response:
+                if not response.is_success:
+                    raise ClientError(response.status_code, 'HTTP streaming failed')
+
+                async for line in response.aiter_lines():
+                    if line:
+                        line_count += 1
+                        total_bytes += len(line.encode('utf-8'))
+
+                        if line_count > limits.max_lines:
+                            raise ClientError(429, 'Stream line limit exceeded')
+
+                        if total_bytes > limits.max_bytes:
+                            raise ClientError(429, 'Stream size limit exceeded')
+
+                        if len(line) > limits.max_message_size:
+                            continue
+
+                        yield line
+
+        except ClientError:
+            raise
+        except Exception as e:
+            raise ClientError(503, sanitize_error_message(str(e), 'HTTP streaming error'))
+
+    async def tail_job_logs(
+        self: AsyncClientProtocol,
+        job_id: str,
+        timeout: float = 30.0,
+        *,
+        protocol: StreamProtocol = 'auto',
+    ) -> AsyncGenerator[str, None]:
+        """Stream job logs with automatic protocol selection (async).
+
+        Unified method that supports WebSocket, HTTP, and auto-selection.
+
+        Args:
+            job_id: The Ray job ID to tail logs for.
+            timeout: Connection timeout in seconds.
+            protocol: Protocol to use:
+                - 'websocket': Use WebSocket only
+                - 'http': Use HTTP streaming only
+                - 'auto': Try WebSocket, fall back to HTTP on connection failure
+
+        Yields:
+            Log message strings.
+
+        Raises:
+            ClientError: On connection, protocol, or validation errors.
+
+        Example:
+            >>> async for line in client.tail_job_logs('raysubmit_abc123'):
+            ...     print(line)
+        """
+        validate_resource_id(job_id, 'job')
+        validate_timeout(timeout)
+
+        if protocol == 'websocket':
+            async for msg in self.websocket_tail_job_logs(job_id, timeout):
+                yield msg
+        elif protocol == 'http':
+            async for msg in self.stream_tail_job_logs(job_id, timeout):
+                yield msg
+        elif protocol == 'auto':
+            try:
+                async for msg in self.websocket_tail_job_logs(job_id, timeout):
+                    yield msg
+            except ClientError as e:
+                if e.status_code in (500, 503):
+                    async for msg in self.stream_tail_job_logs(job_id, timeout):
+                        yield msg
+                else:
+                    raise
+        else:
+            raise ClientError(400, f'Invalid protocol: {protocol}')
+
+    # -------------------------------------------------------------------------
+    # Nodes
+    # -------------------------------------------------------------------------
+
+    async def list_nodes(self: AsyncClientProtocol) -> list[dict]:
+        """List all Ray nodes."""
+        return await self._get('ray/nodes/')
+
+    async def get_node(self: AsyncClientProtocol, node_id: str) -> dict:
+        """Get a Ray node by ID."""
+        return await self._get(f'ray/nodes/{node_id}')
+
+    # -------------------------------------------------------------------------
+    # Tasks
+    # -------------------------------------------------------------------------
+
+    async def list_tasks(self: AsyncClientProtocol) -> list[dict]:
+        """List all Ray tasks."""
+        return await self._get('ray/tasks/')
+
+    async def get_task(self: AsyncClientProtocol, task_id: str) -> dict:
+        """Get a Ray task by ID."""
+        return await self._get(f'ray/tasks/{task_id}')
+
+    # -------------------------------------------------------------------------
+    # Serve Applications
+    # -------------------------------------------------------------------------
+
+    async def list_serve_applications(self: AsyncClientProtocol) -> list[dict]:
+        """List all Ray Serve applications."""
+        return await self._get('ray/serve_applications/')
+
+    async def get_serve_application(self: AsyncClientProtocol, name: str) -> dict:
+        """Get a Ray Serve application by name."""
+        return await self._get(f'ray/serve_applications/{name}')
+
+    async def delete_serve_application(self: AsyncClientProtocol, name: str) -> None:
+        """Delete a Ray Serve application."""
+        await self._delete(f'ray/serve_applications/{name}')

synapse_sdk/clients/agent/container.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from synapse_sdk.clients.protocols import ClientProtocol
+
+
+class ContainerClientMixin:
+    """Mixin for container management endpoints."""
+
+    # Docker containers
+    def list_docker_containers(self: ClientProtocol) -> list[dict]:
+        """List all Docker containers on the host."""
+        return self._get('containers/docker/')
+
+    def get_docker_container(self: ClientProtocol, container_id: str) -> dict:
+        """Get a specific Docker container by ID."""
+        return self._get(f'containers/docker/{container_id}')
+
+    def create_docker_container(
+        self: ClientProtocol,
+        plugin_release: str,
+        *,
+        params: dict[str, Any] | None = None,
+        envs: dict[str, str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        labels: list[str] | None = None,
+    ) -> dict:
+        """Build and run a Docker container for a plugin.
+
+        Args:
+            plugin_release: Plugin identifier (e.g., "plugin_code@version").
+            params: Parameters forwarded to the plugin.
+            envs: Environment variables injected into the container.
+            metadata: Additional metadata stored with the container record.
+            labels: Container labels for display or filtering.
+        """
+        data = {'plugin_release': plugin_release}
+        if params is not None:
+            data['params'] = params
+        if envs is not None:
+            data['envs'] = envs
+        if metadata is not None:
+            data['metadata'] = metadata
+        if labels is not None:
+            data['labels'] = labels
+
+        return self._post('containers/docker/', data=data)
+
+    def delete_docker_container(self: ClientProtocol, container_id: str) -> None:
+        """Stop and remove a Docker container."""
+        self._delete(f'containers/docker/{container_id}')
+
+    # Database container records
+    def list_containers(
+        self: ClientProtocol,
+        params: dict | None = None,
+        *,
+        list_all: bool = False,
+    ) -> dict | tuple[Any, int]:
+        """List tracked containers from database."""
+        return self._list('containers/', params=params, list_all=list_all)
+
+    def get_container(self: ClientProtocol, container_id: int) -> dict:
+        """Get a tracked container by database ID."""
+        return self._get(f'containers/{container_id}')
+
+    def update_container(
+        self: ClientProtocol,
+        container_id: int,
+        *,
+        status: str | None = None,
+    ) -> dict:
+        """Update a tracked container's status."""
+        data = {}
+        if status is not None:
+            data['status'] = status
+        return self._patch(f'containers/{container_id}', data=data)
+
+    def delete_container(self: ClientProtocol, container_id: int) -> None:
+        """Delete a tracked container (stops Docker container too)."""
+        self._delete(f'containers/{container_id}')

synapse_sdk/clients/agent/plugin.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from synapse_sdk.clients.protocols import ClientProtocol
+
+
+class PluginClientMixin:
+    """Mixin for plugin release endpoints."""
+
+    def list_plugin_releases(
+        self: ClientProtocol,
+        params: dict | None = None,
+        *,
+        list_all: bool = False,
+    ) -> dict | tuple[Any, int]:
+        """List all plugin releases."""
+        return self._list('plugin_releases/', params=params, list_all=list_all)
+
+    def get_plugin_release(self: ClientProtocol, lookup: str) -> dict:
+        """Get a plugin release by ID or code@version."""
+        return self._get(f'plugin_releases/{lookup}')
+
+    def create_plugin_release(
+        self: ClientProtocol,
+        plugin: str,
+        version: str,
+    ) -> dict:
+        """Fetch and cache a plugin release."""
+        return self._post('plugin_releases/', data={'plugin': plugin, 'version': version})
+
+    def delete_plugin_release(self: ClientProtocol, lookup: str) -> None:
+        """Delete a plugin release."""
+        self._delete(f'plugin_releases/{lookup}')
+
+    def run_plugin_release(
+        self: ClientProtocol,
+        lookup: str,
+        action: str,
+        params: dict[str, Any] | None = None,
+        *,
+        requirements: list[str] | None = None,
+        job_id: str | None = None,
+    ) -> Any:
+        """Run a plugin release action.
+
+        Args:
+            lookup: Plugin identifier (ID or "plugin@version").
+            action: Action name to execute.
+            params: Parameters to pass to the action.
+            requirements: Additional pip requirements.
+            job_id: Optional job ID for tracking.
+        """
+        data: dict[str, Any] = {'action': action}
+        if params is not None:
+            data['params'] = params
+        if requirements is not None:
+            data['requirements'] = requirements
+        if job_id is not None:
+            data['job_id'] = job_id
+
+        return self._post(f'plugin_releases/{lookup}/run/', data=data)
+
+    def run_debug_plugin_release(
+        self: ClientProtocol,
+        action: str,
+        params: dict[str, Any] | None = None,
+        *,
+        plugin_path: str | None = None,
+        config: dict[str, Any] | None = None,
+        modules: dict[str, str] | None = None,
+        requirements: list[str] | None = None,
+        job_id: str | None = None,
+    ) -> Any:
+        """Run a plugin in debug mode (from source path).
+
+        Args:
+            action: Action name to execute.
+            params: Parameters to pass to the action.
+            plugin_path: Path to the plugin source directory.
+            config: Plugin configuration override.
+            modules: Module source code mapping.
+            requirements: Additional pip requirements.
+            job_id: Optional job ID for tracking.
+        """
+        data: dict[str, Any] = {'action': action}
+        if params is not None:
+            data['params'] = params
+        if plugin_path is not None:
+            data['plugin_path'] = plugin_path
+        if config is not None:
+            data['config'] = config
+        if modules is not None:
+            data['modules'] = modules
+        if requirements is not None:
+            data['requirements'] = requirements
+        if job_id is not None:
+            data['job_id'] = job_id
+
+        return self._post('plugin_releases/run_debug/', data=data)