WebSocket API Guide¶

import asyncio
import json
from arcp import ARCPClient

async def public_websocket_example():
    client = ARCPClient("http://localhost:8001")

    print("Connecting to public WebSocket...")

    message_count = 0
    async for message in client.websocket_public():
        message_count += 1
        print(f"Message {message_count}: {message.get('type', 'unknown')}")

        # Handle different message types
        if message.get('type') == 'welcome':
            print(f"  Welcome: {message.get('message', '')}")
            print(f"  Features: {message.get('features', [])}")

        elif message.get('type') == 'stats_update':
            data = message.get('data', {})
            print(f"  System stats: {data.get('total_agents', 0)} agents")
            print(f"  Alive agents: {data.get('alive_agents', 0)}")

        elif message.get('type') == 'discovery_data':
            data = message.get('data', {})
            pagination = data.get('pagination', {})
            print(f"  Discovery: {pagination.get('total_agents', 0)} total agents")

        elif message.get('type') == 'agents_update':
            data = message.get('data', {})
            print(f"  Agents update: {data.get('total_count', 0)} agents")

        elif message.get('type') == 'pong':
            print(f"  Pong received at: {message.get('timestamp', '')}")

        # Stop after 10 messages
        if message_count >= 10:
            break

    print(f"Received {message_count} messages")

asyncio.run(public_websocket_example())

import asyncio
from arcp import ARCPClient

async def agent_websocket_example():
    client = ARCPClient("http://localhost:8001")

    # Register the agent
    registration_response = await client.register_agent(
        agent_id="my-agent-001",
        name="My testing Agent",
        agent_type="testing",
        endpoint="http://localhost:8010",
        capabilities=["processing", "analysis"],
        context_brief="A testing agent for WebSocket demo",
        version="1.0.0",
        owner="Test User",
        public_key="test-public-key-that-is-at-least-32-characters-long-for-validation",
        communication_mode="remote",
        metadata={"environment": "test", "purpose": "websocket-demo"},
        agent_key="test-agent-001"
    )

    print(f"Agent registered successfully: {registration_response.agent_id}")

    print("Connecting to agent WebSocket...")

    message_count = 0
    async for message in client.websocket_agent():
        message_count += 1
        print(f"Agent message {message_count}: {message}")

        # Handle different message types - message might be a list or dict
        if isinstance(message, list):
            print(f"  Received agent list with {len(message)} agents")
            for agent in message:
                if isinstance(agent, dict) and 'agent_id' in agent:
                    print(f"    - Agent: {agent.get('agent_id')} ({agent.get('status', 'unknown')})")
        elif isinstance(message, dict):
            # Handle authentication and control messages
            if message.get('type') == 'auth_required':
                print(f"  Authentication required: {message.get('message')}")

            elif message.get('type') == 'auth_success':
                print(f"  Authentication successful: {message.get('message')}")

            elif message.get('type') == 'auth_error':
                print(f"  Authentication failed: {message.get('message')}")

            elif message.get('type') == 'pong':
                print(f"  Pong received at: {message.get('timestamp', '')}")

            else:
                print(f"  Other message type: {message.get('type', 'unknown')}")

        # Stop after 5 messages
        if message_count >= 5:
            break

    print(f"Received {message_count} agent messages")

asyncio.run(agent_websocket_example())

import asyncio
import json
import websockets
from typing import Dict, Any, Optional

class ARCPWebSocketClient:
    def __init__(self, base_url: str, token: Optional[str] = None):
        self.base_url = base_url.replace('http', 'ws')
        self.token = token
        self.websocket = None
        self.running = False

    async def connect_public(self):
        """Connect to public WebSocket"""
        url = f"{self.base_url}/public/ws"
        self.websocket = await websockets.connect(url)
        self.running = True

        # Send initial ping
        await self.websocket.send("ping")

        # Listen for messages
        async for message in self.websocket:
            if not self.running:
                break

            try:
                data = json.loads(message)
                await self.handle_message(data)
            except json.JSONDecodeError:
                # Handle non-JSON messages (like pong)
                if message == "pong":
                    print("Received pong")

    async def connect_agent(self):
        """Connect to agent WebSocket with authentication"""
        url = f"{self.base_url}/agents/ws"
        self.websocket = await websockets.connect(url)
        self.running = True

        # Wait for auth request
        auth_request = await self.websocket.recv()
        auth_data = json.loads(auth_request)

        if auth_data.get('type') == 'auth_required':
            # Send authentication token
            await self.websocket.send(json.dumps({
                'token': self.token
            }))

            # Wait for auth response
            auth_response = await self.websocket.recv()
            auth_result = json.loads(auth_response)

            if auth_result.get('type') == 'auth_success':
                print("Agent WebSocket authenticated")
            else:
                print(f"Authentication failed: {auth_result.get('message')}")
                return

        # Listen for messages
        async for message in self.websocket:
            if not self.running:
                break

            try:
                data = json.loads(message)
                await self.handle_message(data)
            except json.JSONDecodeError:
                if message == "ping":
                    await self.websocket.send("pong")

    async def handle_message(self, message: Dict[str, Any]):
        """Handle incoming WebSocket messages"""
        message_type = message.get('type', 'unknown')

        if message_type == 'stats_update':
            await self.handle_stats_update(message.get('data', {}))
        elif message_type == 'agents_update':
            await self.handle_agents_update(message.get('data', {}))
        elif message_type == 'auth_required':
            await self.handle_auth_required(message)
        elif message_type == 'auth_success':
            await self.handle_auth_success(message)
        elif message_type == 'auth_error':
            await self.handle_auth_error(message)
        else:
            print(f"Unknown message type: {message_type}")

    async def handle_stats_update(self, data: Dict[str, Any]):
        """Handle stats update messages"""
        print(f"Stats update: {data.get('total_agents', 0)} agents")
        print(f"   Alive: {data.get('alive_agents', 0)}")
        print(f"   Dead: {data.get('dead_agents', 0)}")

    async def handle_agents_update(self, data: Dict[str, Any]):
        """Handle agents update messages"""
        print(f"Agents update: {data.get('total_count', 0)} agents")

    async def handle_auth_required(self, message: Dict[str, Any]):
        """Handle authentication required messages"""
        print(f"Authentication required: {message.get('message')}")

    async def handle_auth_success(self, message: Dict[str, Any]):
        """Handle authentication success messages"""
        print(f"Authentication successful: {message.get('message')}")

    async def handle_auth_error(self, message: Dict[str, Any]):
        """Handle authentication error messages"""
        print(f"Authentication failed: {message.get('message')}")

    async def send_message(self, message: Dict[str, Any]):
        """Send a message through the WebSocket"""
        if self.websocket and self.running:
            await self.websocket.send(json.dumps(message))

    async def disconnect(self):
        """Disconnect from WebSocket"""
        self.running = False
        if self.websocket:
            await self.websocket.close()

# Usage
async def custom_websocket_example():
    client = ARCPWebSocketClient("http://localhost:8001")

    try:
        # Connect to public WebSocket
        await client.connect_public()

        # Keep running for 30 seconds
        await asyncio.sleep(30)

    finally:
        await client.disconnect()

asyncio.run(custom_websocket_example())

import asyncio
import json
import websockets
from typing import Optional, Dict, Any

class ReconnectingWebSocketClient:
    def __init__(self, base_url: str, token: Optional[str] = None):
        self.base_url = base_url.replace('http', 'ws')
        self.token = token
        self.websocket = None
        self.running = False
        self.reconnect_delay = 5
        self.max_reconnect_delay = 60

    async def connect_with_reconnect(self):
        """Connect with automatic reconnection"""
        while self.running:
            try:
                await self.connect()
                await self.listen()
            except websockets.exceptions.ConnectionClosed:
                print("WebSocket connection closed, reconnecting...")
                await asyncio.sleep(self.reconnect_delay)
                self.reconnect_delay = min(self.reconnect_delay * 2, self.max_reconnect_delay)
            except Exception as e:
                print(f"WebSocket error: {e}")
                await asyncio.sleep(self.reconnect_delay)

    async def connect(self):
        """Establish WebSocket connection"""
        self.websocket = await websockets.connect(self.base_url)
        print("WebSocket connected")

    async def listen(self):
        """Listen for messages"""
        async for message in self.websocket:
            if not self.running:
                break

            try:
                data = json.loads(message)
                await self.handle_message(data)
            except json.JSONDecodeError:
                if message == "ping":
                    await self.websocket.send("pong")

    async def handle_message(self, message: Dict[str, Any]):
        """Handle incoming messages"""
        print(f"Message: {message.get('type', 'unknown')}")

    def start(self):
        """Start the WebSocket client"""
        self.running = True
        return asyncio.create_task(self.connect_with_reconnect())

    def stop(self):
        """Stop the WebSocket client"""
        self.running = False
        if self.websocket:
            asyncio.create_task(self.websocket.close())

# Usage
async def reconnecting_example():
    client = ReconnectingWebSocketClient("ws://localhost:8001/public/ws")

    # Start client
    task = client.start()

    try:
        # Keep running for 2 minutes
        await asyncio.sleep(120)
    finally:
        client.stop()
        task.cancel()

asyncio.run(reconnecting_example())

import asyncio
import json
import websockets

class WebSocketTester:
    def __init__(self, base_url: str):
        self.base_url = base_url.replace('http', 'ws')
        self.results = []

    async def test_public_websocket(self):
        """Test public WebSocket connection"""
        print("Testing public WebSocket...")

        try:
            async with websockets.connect(f"{self.base_url}/public/ws") as ws:
                ping_pong_success = False

                # Wait for messages
                message_count = 0
                async for message in ws:
                    message_count += 1
                    try:
                        data = json.loads(message)
                        print(f"Received message: {data.get('type', 'unknown')}")
                        self.results.append(("public_websocket", "message_received", "success"))
                    except json.JSONDecodeError:
                        print(f"Received non-JSON: {message}")
                        if message == "ping":
                            # Respond to server ping with pong
                            await ws.send("pong")
                            print("Responded to ping with pong")
                            ping_pong_success = True

                    if message_count >= 3:
                        break

                if ping_pong_success:
                    print("Ping/pong successful")
                    self.results.append(("public_websocket", "ping_pong", "success"))
                else:
                    print("No ping received from server")
                    self.results.append(("public_websocket", "ping_pong", "failed"))

                print(f"Received {message_count} messages")

        except Exception as e:
            print(f"Public WebSocket test failed: {e}")
            self.results.append(("public_websocket", "connection", "failed"))

    async def test_agent_websocket(self, token: str):
        """Test agent WebSocket connection"""
        print("Testing agent WebSocket...")

        try:
            async with websockets.connect(f"{self.base_url}/agents/ws") as ws:
                # Wait for auth request
                auth_request = await ws.recv()
                auth_data = json.loads(auth_request)

                if auth_data.get('type') == 'auth_required':
                    # Send token
                    await ws.send(json.dumps({'token': token}))

                    # Wait for auth response
                    auth_response = await ws.recv()
                    auth_result = json.loads(auth_response)

                    if auth_result.get('type') == 'auth_success':
                        print("Agent WebSocket authentication successful")
                        self.results.append(("agent_websocket", "authentication", "success"))
                    else:
                        print("Agent WebSocket authentication failed")
                        self.results.append(("agent_websocket", "authentication", "failed"))
                        return

                # Wait for messages
                message_count = 0
                async for message in ws:
                    message_count += 1
                    try:
                        data = json.loads(message)
                        print(f"Received agent message: {data.get('type', 'unknown')}")
                        self.results.append(("agent_websocket", "message_received", "success"))
                    except json.JSONDecodeError:
                        if message == "ping":
                            await ws.send("pong")

                    if message_count >= 3:
                        break

                print(f"Received {message_count} agent messages")

        except Exception as e:
            print(f"Agent WebSocket test failed: {e}")
            self.results.append(("agent_websocket", "connection", "failed"))

    def print_results(self):
        """Print test results"""
        print("\nWebSocket Test Results:")
        print("=" * 50)

        for test_type, test_name, result in self.results:
            print(f"{test_type}: {test_name} - {result}")

# Usage
async def test_websockets():
    tester = WebSocketTester("http://localhost:8001")

    # Test public WebSocket
    await tester.test_public_websocket()

    # Test agent WebSocket (requires valid token)
    # await tester.test_agent_websocket("your-jwt-token")

    # Print results
    tester.print_results()

asyncio.run(test_websockets())

import asyncio
import websockets
from typing import List

class WebSocketPool:
    def __init__(self, base_url: str, pool_size: int = 5):
        self.base_url = base_url.replace('http', 'ws')
        self.pool_size = pool_size
        self.connections: List[websockets.WebSocketServerProtocol] = []
        self.available: asyncio.Queue = asyncio.Queue()

    async def initialize(self):
        """Initialize connection pool"""
        for _ in range(self.pool_size):
            ws = await websockets.connect(f"{self.base_url}/public/ws")
            self.connections.append(ws)
            await self.available.put(ws)

    async def get_connection(self):
        """Get an available connection from the pool"""
        return await self.available.get()

    async def return_connection(self, ws):
        """Return a connection to the pool"""
        await self.available.put(ws)

    async def close_all(self):
        """Close all connections"""
        for ws in self.connections:
            await ws.close()

# Usage
async def websocket_pool_example():
    pool = WebSocketPool("http://localhost:8001", pool_size=3)

    try:
        await pool.initialize()

        # Use connections from pool
        for i in range(5):
            ws = await pool.get_connection()
            try:
                await ws.send("ping")
                response = await ws.recv()
                print(f"Request {i+1}: {response}")
            finally:
                await pool.return_connection(ws)

    finally:
        await pool.close_all()

asyncio.run(websocket_pool_example())

# Check token validity
import asyncio
from arcp import ARCPClient

async def check_token():
    client = ARCPClient("http://localhost:8001")
    result = await client.validate_token("fewfwfewfewfew")
    print(f"Token valid: {result}")

asyncio.run(check_token())

# Handle non-JSON messages gracefully
async def handle_message_safely(message: str):
    try:
        data = json.loads(message)
        return data
    except json.JSONDecodeError:
        # Handle non-JSON messages (like pong)
        if message == "pong":
            return {"type": "pong"}
        elif message == "ping":
            return {"type": "ping"}
        else:
            print(f"Unknown non-JSON message: {message}")
            return None