connectonion 0.6.4__py3-none-any.whl → 0.6.5__py3-none-any.whl

connectonion/network/asgi/http.py

@@ -107,6 +107,7 @@ async def handle_http(
     route_handlers: dict,
     storage,
     trust: str,
+    trust_config: dict | None = None,
     start_time: float,
     blacklist: list | None = None,
     whitelist: list | None = None,
@@ -120,6 +121,7 @@ async def handle_http(
         route_handlers: Dict of route handler functions (input, session, sessions, health, info, auth)
         storage: SessionStorage instance
         trust: Trust level (open/careful/strict)
+        trust_config: Parsed YAML config from trust policy (for /info onboard)
         start_time: Server start time
         blacklist: Blocked identities
         whitelist: Allowed identities
@@ -133,26 +135,130 @@ async def handle_http(
         await send({"type": "http.response.body", "body": b""})
         return
 
-    # Admin endpoints require API key auth
+    # Admin endpoints
     if path.startswith("/admin"):
-        headers = dict(scope.get("headers", []))
-        auth = headers.get(b"authorization", b"").decode()
-        expected = os.environ.get("OPENONION_API_KEY", "")
-        if not expected or not auth.startswith("Bearer ") or not hmac.compare_digest(auth[7:], expected):
-            await send_json(send, {"error": "unauthorized"}, 401)
-            return
-
-        if method == "GET" and path == "/admin/logs":
-            result = route_handlers["admin_logs"]()
-            if "error" in result:
-                await send_json(send, result, 404)
+        # Legacy admin endpoints (Bearer token auth)
+        if path in ["/admin/logs", "/admin/sessions"]:
+            headers_dict = dict(scope.get("headers", []))
+            auth = headers_dict.get(b"authorization", b"").decode()
+            expected = os.environ.get("OPENONION_API_KEY", "")
+            if not expected or not auth.startswith("Bearer ") or not hmac.compare_digest(auth[7:], expected):
+                await send_json(send, {"error": "unauthorized"}, 401)
+                return
+
+            if method == "GET" and path == "/admin/logs":
+                result = route_handlers["admin_logs"]()
+                if "error" in result:
+                    await send_json(send, result, 404)
+                else:
+                    await send_text(send, result["content"])
+                return
+
+            if method == "GET" and path == "/admin/sessions":
+                await send_json(send, route_handlers["admin_sessions"]())
+                return
+
+        # Admin trust routes (signed request + admin check)
+        if path.startswith("/admin/trust/") or path.startswith("/superadmin/"):
+            trust_agent = route_handlers["trust_agent"]
+
+            # For GET requests, allow auth via headers (bodies often stripped by proxies)
+            headers_dict = dict(scope.get("headers", []))
+            header_from = headers_dict.get(b"x-from", b"").decode()
+            header_sig = headers_dict.get(b"x-signature", b"").decode()
+            header_ts = headers_dict.get(b"x-timestamp", b"").decode()
+
+            if method == "GET" and header_from and header_sig and header_ts:
+                # Build signed request from headers
+                try:
+                    data = {
+                        "payload": {"timestamp": float(header_ts)},
+                        "from": header_from,
+                        "signature": header_sig
+                    }
+                except ValueError:
+                    await send_json(send, {"error": "Invalid X-Timestamp header"}, 400)
+                    return
             else:
-                await send_text(send, result["content"])
-            return
-
-        if method == "GET" and path == "/admin/sessions":
-            await send_json(send, route_handlers["admin_sessions"]())
-            return
+                # Read from body (POST or GET without headers)
+                body = await read_body(receive)
+                try:
+                    data = json.loads(body) if body else {}
+                except json.JSONDecodeError:
+                    await send_json(send, {"error": "Invalid JSON"}, 400)
+                    return
+
+            # Authenticate signed request (reuse same auth as /input, but with "open" trust)
+            _, identity, sig_valid, err = route_handlers["auth"](data, "open")
+            if err or not sig_valid:
+                await send_json(send, {"error": err or "unauthorized: invalid signature"}, 401)
+                return
+
+            # Check admin permission
+            if path.startswith("/superadmin/"):
+                # Super admin only (self address)
+                if not trust_agent.is_super_admin(identity):
+                    await send_json(send, {"error": "forbidden: super admin only"}, 403)
+                    return
+            else:
+                # Any admin
+                if not trust_agent.is_admin(identity):
+                    await send_json(send, {"error": "forbidden: admin only"}, 403)
+                    return
+
+            payload = data.get("payload", {})
+            client_id = payload.get("client_id")
+
+            # Route to handlers
+            if method == "POST" and path == "/admin/trust/promote":
+                if not client_id:
+                    await send_json(send, {"error": "client_id required"}, 400)
+                    return
+                await send_json(send, route_handlers["admin_trust_promote"](client_id))
+                return
+
+            if method == "POST" and path == "/admin/trust/demote":
+                if not client_id:
+                    await send_json(send, {"error": "client_id required"}, 400)
+                    return
+                await send_json(send, route_handlers["admin_trust_demote"](client_id))
+                return
+
+            if method == "POST" and path == "/admin/trust/block":
+                if not client_id:
+                    await send_json(send, {"error": "client_id required"}, 400)
+                    return
+                reason = payload.get("reason", "")
+                await send_json(send, route_handlers["admin_trust_block"](client_id, reason))
+                return
+
+            if method == "POST" and path == "/admin/trust/unblock":
+                if not client_id:
+                    await send_json(send, {"error": "client_id required"}, 400)
+                    return
+                await send_json(send, route_handlers["admin_trust_unblock"](client_id))
+                return
+
+            if method == "GET" and path.startswith("/admin/trust/level/"):
+                client_id = path[len("/admin/trust/level/"):]
+                await send_json(send, route_handlers["admin_trust_level"](client_id))
+                return
+
+            if method == "POST" and path == "/superadmin/add":
+                admin_id = payload.get("admin_id")
+                if not admin_id:
+                    await send_json(send, {"error": "admin_id required"}, 400)
+                    return
+                await send_json(send, route_handlers["admin_admins_add"](admin_id))
+                return
+
+            if method == "POST" and path == "/superadmin/remove":
+                admin_id = payload.get("admin_id")
+                if not admin_id:
+                    await send_json(send, {"error": "admin_id required"}, 400)
+                    return
+                await send_json(send, route_handlers["admin_admins_remove"](admin_id))
+                return
 
         await send_json(send, {"error": "not found"}, 404)
         return
@@ -189,7 +295,7 @@ async def handle_http(
         await send_json(send, route_handlers["health"](start_time))
 
     elif method == "GET" and path == "/info":
-        await send_json(send, route_handlers["info"](trust))
+        await send_json(send, route_handlers["info"](trust, trust_config))
 
     elif method == "GET" and path == "/docs":
         # Serve static docs page

connectonion/network/asgi/websocket.py

@@ -1,13 +1,13 @@
 """
-Purpose: WebSocket bidirectional communication for ASGI server with real-time agent I/O
+Purpose: WebSocket bidirectional communication for ASGI server with real-time agent I/O and keep-alive
 LLM-Note:
   Dependencies: imports from [network/io/websocket.py, network/asgi/http.py pydantic_json_encoder, asyncio, json, queue, threading] | imported by [network/asgi/__init__.py] | tested by [tests/network/test_asgi_websocket.py]
-  Data flow: handle_websocket() accepts connection → receives INPUT message with prompt+session → authenticates via route_handlers["auth"] → starts agent in background thread with WebSocketIO → agent sends events via io.log()/send() → forwards to client via websocket.send → client sends ASK_USER_RESPONSE for approvals → io receives via queue → agent resumes → returns OUTPUT with result+session_id
-  State/Effects: maintains WebSocket connection during agent execution | runs agent in daemon thread (non-blocking) | uses queue.Queue for thread-safe I/O between agent and WebSocket | no persistent state (connection-scoped)
-  Integration: exposes handle_websocket(scope, receive, send, route_handlers, storage, trust, blacklist, whitelist) | uses WebSocketIO for bidirectional I/O | supports session continuation (same as HTTP) | message types: INPUT (client), OUTPUT/ERROR (server), ASK_USER_RESPONSE (client), trace events (server)
-  Performance: async WebSocket handling | agent runs in separate thread to avoid blocking | queue-based message passing | streams events in real-time (thinking, tool_result, approval_needed)
-  Errors: sends ERROR message for invalid JSON, auth failures, missing prompt | closes connection with code 4004 for wrong path | catches exceptions in agent thread and sends ERROR
-WebSocket handling for ASGI.
+  Data flow: handle_websocket() accepts connection → receives INPUT message with prompt+session → authenticates via route_handlers["auth"] → starts agent in background thread with WebSocketIO → sends PING every 30s for keep-alive → agent sends events via io.log()/send() → forwards to client via websocket.send → client responds with PONG to PING → client sends ASK_USER_RESPONSE for approvals → io receives via queue → agent resumes → returns OUTPUT with result+session_id → result saved to SessionStorage even if client disconnects
+  State/Effects: maintains WebSocket connection during agent execution | runs agent in daemon thread (non-blocking) | uses queue.Queue for thread-safe I/O between agent and WebSocket | no persistent state (connection-scoped) | results persisted to .co/session_results.jsonl for 24h TTL
+  Integration: exposes handle_websocket(scope, receive, send, route_handlers, storage, trust, blacklist, whitelist) | uses WebSocketIO for bidirectional I/O | supports session continuation (same as HTTP) | message types: INPUT (client), OUTPUT/ERROR (server), PING (server), PONG (client), ASK_USER_RESPONSE (client), trace events (server) | clients can poll GET /sessions/{id} to recover results after disconnect
+  Performance: async WebSocket handling | agent runs in separate thread to avoid blocking | queue-based message passing | streams events in real-time (thinking, tool_result, approval_needed) | PING task keeps connection alive through proxies/firewalls
+  Errors: sends ERROR message for invalid JSON, auth failures, missing prompt | closes connection with code 4004 for wrong path | catches exceptions in agent thread and sends ERROR | saves result to storage even on disconnect (enables polling recovery)
+WebSocket handling for ASGI with keep-alive and session recovery.
 
 ASGI Protocol Types (not our custom types - this is the ASGI spec):
 - websocket.connect    : ASGI sends when client wants to connect
@@ -18,10 +18,21 @@ ASGI Protocol Types (not our custom types - this is the ASGI spec):
 - websocket.close      : We send to close the connection
 
 Our Application Types (sent inside websocket.send text payload):
-- INPUT              : Client sends prompt
-- OUTPUT             : Server sends final result
+- INPUT              : Client sends prompt with session_id
+- OUTPUT             : Server sends final result with session_id
 - ERROR              : Server sends error message
+- PING               : Server keep-alive (every 30s)
+- PONG               : Client acknowledges keep-alive
 - ASK_USER_RESPONSE  : Client responds to approval request
+- ONBOARD_REQUIRED   : Server sends when stranger needs to onboard
+- ONBOARD_SUBMIT     : Client sends invite_code or payment
+- ONBOARD_SUCCESS    : Server confirms onboard complete
+- ADMIN_PROMOTE      : Client (admin) requests promote
+- ADMIN_DEMOTE       : Client (admin) requests demote
+- ADMIN_BLOCK        : Client (admin) requests block
+- ADMIN_UNBLOCK      : Client (admin) requests unblock
+- ADMIN_GET_LEVEL    : Client (admin) requests client level
+- ADMIN_RESULT       : Server sends admin action result
 - (trace events)     : thinking, tool_result, approval_needed, etc.
 """
 
@@ -30,7 +41,10 @@ import json
 import queue
 import threading
 
+from rich.console import Console
 from ..io import WebSocketIO
+
+console = Console()
 # Import pydantic_json_encoder for serializing Pydantic models (e.g., TokenUsage) in WebSocket responses
 from .http import pydantic_json_encoder
 
@@ -78,10 +92,37 @@ async def handle_websocket(
                            "text": json.dumps({"type": "ERROR", "message": "Invalid JSON"})})
                 continue
 
-            if data.get("type") == "INPUT":  # Our app type: client wants to run agent
+            msg_type = data.get("type")
+
+            # Handle ONBOARD_SUBMIT (stranger submitting invite code or payment)
+            if msg_type == "ONBOARD_SUBMIT":
+                await _handle_onboard_submit(data, send, route_handlers)
+                continue
+
+            # Handle ADMIN messages (promote, demote, block, unblock, get_level)
+            if msg_type and msg_type.startswith("ADMIN_"):
+                await _handle_admin_message(data, send, route_handlers)
+                continue
+
+            if msg_type == "INPUT":  # Our app type: client wants to run agent
                 prompt, identity, sig_valid, err = route_handlers["auth"](
                     data, trust, blacklist=blacklist, whitelist=whitelist
                 )
+
+                # Check if stranger needs onboarding
+                if err and "forbidden" in err.lower():
+                    # Get onboard requirements from trust config
+                    trust_agent = route_handlers["trust_agent"]
+                    onboard_info = _get_onboard_requirements(trust_agent)
+                    if onboard_info:
+                        await send({"type": "websocket.send",
+                                   "text": json.dumps({
+                                       "type": "ONBOARD_REQUIRED",
+                                       "identity": identity,
+                                       **onboard_info
+                                   })})
+                        continue
+
                 if err:
                     await send({"type": "websocket.send",
                                "text": json.dumps({"type": "ERROR", "message": err})})
@@ -139,16 +180,23 @@ async def handle_websocket(
 
 
 async def _pump_messages(ws_receive, ws_send, io: WebSocketIO, agent_done: threading.Event) -> bool:
-    """Pump messages between WebSocket and IO queues.
+    """Pump messages between WebSocket and IO queues with keep-alive.
 
     Runs until agent completes. Handles:
-    - Outgoing: io._outgoing queue -> WebSocket
-    - Incoming: WebSocket -> io._incoming queue (for approval responses)
+    - Outgoing: io._outgoing queue -> WebSocket (agent events)
+    - Incoming: WebSocket -> io._incoming queue (approval responses, PONG)
+    - Keep-alive: PING messages every 30s to detect dead connections
 
     Returns:
         True if client disconnected before agent completed, False otherwise.
         When True, caller should skip sending final OUTPUT (client is gone,
-        but result is already saved to SessionStorage).
+        but result is already saved to SessionStorage for polling recovery).
+
+    Keep-alive mechanism:
+        - send_ping() task sends PING every 30s
+        - Client responds with PONG (handled in receive_incoming)
+        - Keeps connection alive through proxies/firewalls
+        - Client can detect dead connection if no PING for 60s
 
     Implementation note:
         Uses asyncio.Event for signaling disconnect between nested async functions.
@@ -191,7 +239,11 @@ async def _pump_messages(ws_receive, ws_send, io: WebSocketIO, agent_done: threa
                 if msg["type"] == "websocket.receive":
                     try:
                         data = json.loads(msg.get("text", "{}"))
-                        io._incoming.put(data)
+                        # Handle PONG responses (client responding to PING)
+                        if data.get("type") == "PONG":
+                            pass  # Just acknowledge, no action needed
+                        else:
+                            io._incoming.put(data)
                     except json.JSONDecodeError:
                         pass
                 elif msg["type"] == "websocket.disconnect":
@@ -201,17 +253,226 @@ async def _pump_messages(ws_receive, ws_send, io: WebSocketIO, agent_done: threa
             except asyncio.TimeoutError:
                 continue
 
+    async def send_ping():
+        """Send PING messages every 30 seconds to keep connection alive."""
+        while not agent_done.is_set() and not disconnected.is_set():
+            await asyncio.sleep(30)  # Send ping every 30 seconds
+            if not agent_done.is_set() and not disconnected.is_set():
+                try:
+                    await ws_send({"type": "websocket.send", "text": json.dumps({"type": "PING"})})
+                except Exception:
+                    # Connection likely closed, stop pinging
+                    break
+
     send_task = asyncio.create_task(send_outgoing())
     recv_task = asyncio.create_task(receive_incoming())
+    ping_task = asyncio.create_task(send_ping())
 
     while not agent_done.is_set():
         await asyncio.sleep(0.05)
 
+    # Cancel all tasks
     recv_task.cancel()
+    ping_task.cancel()
     try:
         await recv_task
     except asyncio.CancelledError:
         pass
+    try:
+        await ping_task
+    except asyncio.CancelledError:
+        pass
     await send_task
 
     return disconnected.is_set()
+
+
+def _get_onboard_requirements(trust_agent) -> dict | None:
+    """Get onboard requirements from trust config."""
+    config = trust_agent.config
+    onboard = config.get("onboard", {})
+    if not onboard:
+        return None
+
+    result = {"methods": []}
+    if "invite_code" in onboard:
+        result["methods"].append("invite_code")
+    if "payment" in onboard:
+        result["methods"].append("payment")
+        result["payment_amount"] = onboard["payment"]
+        # Include agent's address for payment transfer
+        payment_address = trust_agent.get_self_address()
+        if payment_address:
+            result["payment_address"] = payment_address
+
+    return result if result["methods"] else None
+
+
+async def _handle_onboard_submit(data: dict, send, route_handlers: dict):
+    """Handle ONBOARD_SUBMIT message from client."""
+    trust_agent = route_handlers["trust_agent"]
+
+    # Authenticate the request (signature check only, open trust)
+    _, identity, sig_valid, err = route_handlers["auth"](data, "open")
+    if err or not sig_valid:
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ERROR", "message": err or "invalid signature"})})
+        return
+
+    # Check if blocked BEFORE onboarding (don't waste invite codes/payments)
+    if trust_agent.is_blocked(identity):
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ERROR", "message": "forbidden: blocked"})})
+        return
+
+    payload = data.get("payload", {})
+    invite_code = payload.get("invite_code")
+    payment = payload.get("payment", 0)
+
+    # Try invite code
+    if invite_code:
+        if trust_agent.verify_invite(identity, invite_code):
+            # Get actual level after promotion (may differ if already promoted)
+            actual_level = trust_agent.get_level(identity)
+            # Log to server console for auditing
+            console.print(f"[green]✓[/green] Verified [bold]{identity[:16]}...[/bold] with invite code [cyan]{invite_code}[/cyan] → {actual_level}")
+            await send({"type": "websocket.send",
+                       "text": json.dumps({
+                           "type": "ONBOARD_SUCCESS",
+                           "identity": identity,
+                           "level": actual_level,
+                           "message": f"Invite code verified. You are now a {actual_level}."
+                       })})
+            return
+        else:
+            console.print(f"[red]✗[/red] Invalid invite code [cyan]{invite_code}[/cyan] from {identity[:16]}...")
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "Invalid invite code"})})
+            return
+
+    # Try payment
+    if payment > 0:
+        if trust_agent.verify_payment(identity, payment):
+            # Get actual level after promotion
+            actual_level = trust_agent.get_level(identity)
+            # Log to server console for auditing
+            console.print(f"[green]✓[/green] Verified [bold]{identity[:16]}...[/bold] with payment [cyan]${payment}[/cyan] → {actual_level}")
+            await send({"type": "websocket.send",
+                       "text": json.dumps({
+                           "type": "ONBOARD_SUCCESS",
+                           "identity": identity,
+                           "level": actual_level,
+                           "message": f"Payment verified. You are now a {actual_level}."
+                       })})
+            return
+        else:
+            console.print(f"[red]✗[/red] Insufficient payment [cyan]${payment}[/cyan] from {identity[:16]}...")
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "Insufficient payment"})})
+            return
+
+    await send({"type": "websocket.send",
+               "text": json.dumps({"type": "ERROR", "message": "invite_code or payment required"})})
+
+
+async def _handle_admin_message(data: dict, send, route_handlers: dict):
+    """Handle ADMIN_* messages from client."""
+    trust_agent = route_handlers["trust_agent"]
+    msg_type = data.get("type")
+
+    # Authenticate the request
+    _, identity, sig_valid, err = route_handlers["auth"](data, "open")
+    if err or not sig_valid:
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ERROR", "message": err or "invalid signature"})})
+        return
+
+    # Check admin permission
+    if not trust_agent.is_admin(identity):
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ERROR", "message": "forbidden: admin only"})})
+        return
+
+    payload = data.get("payload", {})
+    client_id = payload.get("client_id")
+
+    # Handle each admin action
+    if msg_type == "ADMIN_PROMOTE":
+        if not client_id:
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "client_id required"})})
+            return
+        result = route_handlers["admin_trust_promote"](client_id)
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ADMIN_RESULT", "action": "promote", **result})})
+
+    elif msg_type == "ADMIN_DEMOTE":
+        if not client_id:
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "client_id required"})})
+            return
+        result = route_handlers["admin_trust_demote"](client_id)
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ADMIN_RESULT", "action": "demote", **result})})
+
+    elif msg_type == "ADMIN_BLOCK":
+        if not client_id:
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "client_id required"})})
+            return
+        reason = payload.get("reason", "")
+        result = route_handlers["admin_trust_block"](client_id, reason)
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ADMIN_RESULT", "action": "block", **result})})
+
+    elif msg_type == "ADMIN_UNBLOCK":
+        if not client_id:
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "client_id required"})})
+            return
+        result = route_handlers["admin_trust_unblock"](client_id)
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ADMIN_RESULT", "action": "unblock", **result})})
+
+    elif msg_type == "ADMIN_GET_LEVEL":
+        if not client_id:
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "client_id required"})})
+            return
+        result = route_handlers["admin_trust_level"](client_id)
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ADMIN_RESULT", "action": "get_level", **result})})
+
+    elif msg_type == "ADMIN_ADD":
+        # Super admin only
+        if not trust_agent.is_super_admin(identity):
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "forbidden: super admin only"})})
+            return
+        admin_id = payload.get("admin_id")
+        if not admin_id:
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "admin_id required"})})
+            return
+        result = route_handlers["admin_admins_add"](admin_id)
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ADMIN_RESULT", "action": "add_admin", **result})})
+
+    elif msg_type == "ADMIN_REMOVE":
+        # Super admin only
+        if not trust_agent.is_super_admin(identity):
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "forbidden: super admin only"})})
+            return
+        admin_id = payload.get("admin_id")
+        if not admin_id:
+            await send({"type": "websocket.send",
+                       "text": json.dumps({"type": "ERROR", "message": "admin_id required"})})
+            return
+        result = route_handlers["admin_admins_remove"](admin_id)
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ADMIN_RESULT", "action": "remove_admin", **result})})
+
+    else:
+        await send({"type": "websocket.send",
+                   "text": json.dumps({"type": "ERROR", "message": f"Unknown admin action: {msg_type}"})})