comfygit-deploy 0.3.20__tar.gz → 0.4.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/.gitignore +5 -2
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/PKG-INFO +3 -3
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/cli.py +2 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/commands/worker.py +61 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/config.py +34 -1
- comfygit_deploy-0.4.0/comfygit_deploy/tunnel/__init__.py +6 -0
- comfygit_deploy-0.4.0/comfygit_deploy/tunnel/client.py +191 -0
- comfygit_deploy-0.4.0/comfygit_deploy/tunnel/handler.py +472 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/worker/native_manager.py +72 -8
- comfygit_deploy-0.4.0/comfygit_deploy/worker/server.py +1514 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/pyproject.toml +4 -4
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_native_manager.py +13 -3
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_worker_server.py +112 -0
- comfygit_deploy-0.3.20/comfygit_deploy/worker/server.py +0 -511
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/README.md +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/__init__.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/commands/__init__.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/commands/custom.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/commands/dev.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/commands/instances.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/commands/runpod.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/providers/__init__.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/providers/custom.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/providers/runpod.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/startup/__init__.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/startup/scripts.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/worker/__init__.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/worker/mdns.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/comfygit_deploy/worker/state.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/docs/architecture.md +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/__init__.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/conftest.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_cli.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_config.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_custom_client.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_logs_streaming.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_mdns.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_mdns_scanner.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_runpod_client.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_startup_script.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_unified_instances.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_worker_commands.py +0 -0
- {comfygit_deploy-0.3.20 → comfygit_deploy-0.4.0}/tests/test_worker_state.py +0 -0
|
@@ -29,6 +29,7 @@ node_modules/
|
|
|
29
29
|
npm-debug.log*
|
|
30
30
|
yarn-debug.log*
|
|
31
31
|
yarn-error.log*
|
|
32
|
+
*.tsbuildinfo
|
|
32
33
|
|
|
33
34
|
# Environment variables
|
|
34
35
|
.env
|
|
@@ -94,5 +95,7 @@ docs/contexts
|
|
|
94
95
|
# MEOW agent orchestration runtime
|
|
95
96
|
.meow/
|
|
96
97
|
|
|
97
|
-
#
|
|
98
|
-
.beads/
|
|
98
|
+
# Local/retired agent task systems
|
|
99
|
+
.beads/
|
|
100
|
+
.threads/
|
|
101
|
+
.ctxr/
|
|
@@ -1,13 +1,13 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: comfygit-deploy
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.4.0
|
|
4
4
|
Summary: ComfyGit Deploy - Remote deployment and worker management CLI
|
|
5
5
|
Project-URL: Documentation, https://docs.comfygit.org/
|
|
6
6
|
Project-URL: Repository, https://github.com/comfygit-ai/comfygit
|
|
7
7
|
Project-URL: Issues, https://github.com/comfygit-ai/comfygit/issues
|
|
8
8
|
Requires-Python: >=3.10
|
|
9
|
-
Requires-Dist: aiohttp>=3.
|
|
10
|
-
Requires-Dist: comfygit==0.
|
|
9
|
+
Requires-Dist: aiohttp>=3.13.4
|
|
10
|
+
Requires-Dist: comfygit==0.4.0
|
|
11
11
|
Requires-Dist: zeroconf>=0.131.0
|
|
12
12
|
Description-Content-Type: text/markdown
|
|
13
13
|
|
|
@@ -187,6 +187,8 @@ def create_parser() -> argparse.ArgumentParser:
|
|
|
187
187
|
)
|
|
188
188
|
up_parser.add_argument("--broadcast", action="store_true", help="Enable mDNS broadcast")
|
|
189
189
|
up_parser.add_argument("--port-range", default="8200:8210", help="Instance port range")
|
|
190
|
+
up_parser.add_argument("--cloud", help="Remote coordination service base URL or websocket URL")
|
|
191
|
+
up_parser.add_argument("--token", help="Tunnel auth token for the remote coordination service")
|
|
190
192
|
up_parser.add_argument("--dev", action="store_true", help="Use saved dev config (from 'dev setup')")
|
|
191
193
|
up_parser.add_argument("--dev-core", metavar="PATH", help="Use local comfygit-core (editable)")
|
|
192
194
|
up_parser.add_argument("--dev-manager", metavar="PATH", help="Use local comfygit-manager")
|
|
@@ -4,10 +4,14 @@ Commands for setting up and managing the worker server on GPU machines.
|
|
|
4
4
|
"""
|
|
5
5
|
|
|
6
6
|
import argparse
|
|
7
|
+
import asyncio
|
|
8
|
+
import contextlib
|
|
7
9
|
import json
|
|
8
10
|
import secrets
|
|
9
11
|
from pathlib import Path
|
|
10
12
|
|
|
13
|
+
from ..config import DeployConfig
|
|
14
|
+
|
|
11
15
|
WORKER_CONFIG_PATH = Path.home() / ".config" / "comfygit" / "deploy" / "worker.json"
|
|
12
16
|
|
|
13
17
|
|
|
@@ -59,6 +63,28 @@ def save_worker_config(config: dict) -> None:
|
|
|
59
63
|
WORKER_CONFIG_PATH.write_text(json.dumps(config, indent=2))
|
|
60
64
|
|
|
61
65
|
|
|
66
|
+
def _resolve_cloud_settings(args: argparse.Namespace) -> tuple[str | None, str | None]:
|
|
67
|
+
config = DeployConfig()
|
|
68
|
+
|
|
69
|
+
cloud_arg = getattr(args, "cloud", None)
|
|
70
|
+
token_arg = getattr(args, "token", None)
|
|
71
|
+
|
|
72
|
+
cloud_url = cloud_arg or config.cloud_url
|
|
73
|
+
cloud_token = token_arg or config.cloud_token
|
|
74
|
+
|
|
75
|
+
updated = False
|
|
76
|
+
if cloud_arg:
|
|
77
|
+
config.cloud_url = cloud_arg
|
|
78
|
+
updated = True
|
|
79
|
+
if token_arg:
|
|
80
|
+
config.cloud_token = token_arg
|
|
81
|
+
updated = True
|
|
82
|
+
if updated:
|
|
83
|
+
config.save()
|
|
84
|
+
|
|
85
|
+
return cloud_url, cloud_token
|
|
86
|
+
|
|
87
|
+
|
|
62
88
|
def is_worker_running() -> bool:
|
|
63
89
|
"""Check if worker server is running."""
|
|
64
90
|
# Simple check - look for PID file or try to connect
|
|
@@ -156,15 +182,23 @@ def handle_up(args: argparse.Namespace) -> int:
|
|
|
156
182
|
manager_link.symlink_to(dev_manager)
|
|
157
183
|
print(f"Dev mode: manager -> {dev_manager}")
|
|
158
184
|
|
|
185
|
+
cloud_url, cloud_token = _resolve_cloud_settings(args)
|
|
186
|
+
if bool(cloud_url) != bool(cloud_token):
|
|
187
|
+
print("Cloud tunnel requires both --cloud and --token, or both values saved in config.")
|
|
188
|
+
return 1
|
|
189
|
+
|
|
159
190
|
print(f"Starting worker server on {args.host}:{args.port}...")
|
|
160
191
|
print(f" Mode: {args.mode}")
|
|
161
192
|
print(f" Instance ports: {port_start}-{port_end}")
|
|
162
193
|
print(f" Broadcast: {args.broadcast}")
|
|
194
|
+
if cloud_url:
|
|
195
|
+
print(f" Cloud tunnel: {cloud_url}")
|
|
163
196
|
print()
|
|
164
197
|
print("Press Ctrl+C to stop.")
|
|
165
198
|
|
|
166
199
|
from aiohttp import web
|
|
167
200
|
|
|
201
|
+
from ..tunnel.client import TunnelClient
|
|
168
202
|
from ..worker.server import create_worker_app
|
|
169
203
|
|
|
170
204
|
app = create_worker_app(
|
|
@@ -175,6 +209,33 @@ def handle_up(args: argparse.Namespace) -> int:
|
|
|
175
209
|
port_range_end=port_end,
|
|
176
210
|
)
|
|
177
211
|
|
|
212
|
+
if cloud_url and cloud_token:
|
|
213
|
+
try:
|
|
214
|
+
tunnel_client = TunnelClient(
|
|
215
|
+
cloud_url=cloud_url,
|
|
216
|
+
token=cloud_token,
|
|
217
|
+
worker_server=app["worker"],
|
|
218
|
+
)
|
|
219
|
+
except ValueError as exc:
|
|
220
|
+
print(f"Invalid cloud URL: {exc}")
|
|
221
|
+
return 1
|
|
222
|
+
|
|
223
|
+
app["tunnel_client"] = tunnel_client
|
|
224
|
+
|
|
225
|
+
async def _start_tunnel(app: web.Application) -> None:
|
|
226
|
+
app["tunnel_task"] = asyncio.create_task(tunnel_client.run())
|
|
227
|
+
|
|
228
|
+
async def _stop_tunnel(app: web.Application) -> None:
|
|
229
|
+
await tunnel_client.close()
|
|
230
|
+
tunnel_task = app.get("tunnel_task")
|
|
231
|
+
if tunnel_task:
|
|
232
|
+
tunnel_task.cancel()
|
|
233
|
+
with contextlib.suppress(asyncio.CancelledError):
|
|
234
|
+
await tunnel_task
|
|
235
|
+
|
|
236
|
+
app.on_startup.append(_start_tunnel)
|
|
237
|
+
app.on_cleanup.append(_stop_tunnel)
|
|
238
|
+
|
|
178
239
|
# Save PID file
|
|
179
240
|
pid_file = WORKER_CONFIG_PATH.parent / "worker.pid"
|
|
180
241
|
import os
|
|
@@ -30,7 +30,12 @@ class DeployConfig:
|
|
|
30
30
|
path: Config file path. Defaults to ~/.config/comfygit/deploy/config.json
|
|
31
31
|
"""
|
|
32
32
|
self.path = path or _get_default_config_path()
|
|
33
|
-
self._data: dict[str, Any] = {
|
|
33
|
+
self._data: dict[str, Any] = {
|
|
34
|
+
"version": "1",
|
|
35
|
+
"providers": {},
|
|
36
|
+
"workers": {},
|
|
37
|
+
"cloud": {},
|
|
38
|
+
}
|
|
34
39
|
self._load()
|
|
35
40
|
|
|
36
41
|
def _load(self) -> None:
|
|
@@ -69,6 +74,34 @@ class DeployConfig:
|
|
|
69
74
|
"""Get custom workers registry."""
|
|
70
75
|
return self._data.get("workers", {})
|
|
71
76
|
|
|
77
|
+
@property
|
|
78
|
+
def cloud_url(self) -> str | None:
|
|
79
|
+
"""Get the saved remote service URL for tunnel mode."""
|
|
80
|
+
return self._data.get("cloud", {}).get("url")
|
|
81
|
+
|
|
82
|
+
@cloud_url.setter
|
|
83
|
+
def cloud_url(self, value: str | None) -> None:
|
|
84
|
+
if "cloud" not in self._data:
|
|
85
|
+
self._data["cloud"] = {}
|
|
86
|
+
if value is None:
|
|
87
|
+
self._data["cloud"].pop("url", None)
|
|
88
|
+
else:
|
|
89
|
+
self._data["cloud"]["url"] = value
|
|
90
|
+
|
|
91
|
+
@property
|
|
92
|
+
def cloud_token(self) -> str | None:
|
|
93
|
+
"""Get the saved remote service auth token for tunnel mode."""
|
|
94
|
+
return self._data.get("cloud", {}).get("token")
|
|
95
|
+
|
|
96
|
+
@cloud_token.setter
|
|
97
|
+
def cloud_token(self, value: str | None) -> None:
|
|
98
|
+
if "cloud" not in self._data:
|
|
99
|
+
self._data["cloud"] = {}
|
|
100
|
+
if value is None:
|
|
101
|
+
self._data["cloud"].pop("token", None)
|
|
102
|
+
else:
|
|
103
|
+
self._data["cloud"]["token"] = value
|
|
104
|
+
|
|
72
105
|
def add_worker(
|
|
73
106
|
self,
|
|
74
107
|
name: str,
|
|
@@ -0,0 +1,191 @@
|
|
|
1
|
+
"""Worker-side WebSocket tunnel client."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
import contextlib
|
|
7
|
+
import json
|
|
8
|
+
from typing import Any
|
|
9
|
+
from urllib.parse import urlsplit, urlunsplit
|
|
10
|
+
|
|
11
|
+
import aiohttp
|
|
12
|
+
|
|
13
|
+
from ..worker.server import WorkerServer
|
|
14
|
+
from .handler import TunnelHandler
|
|
15
|
+
|
|
16
|
+
MAX_MESSAGE_BYTES = 1_000_000
|
|
17
|
+
PING_INTERVAL_SECONDS = 30.0
|
|
18
|
+
AUTH_TIMEOUT_SECONDS = 10.0
|
|
19
|
+
MAX_RECONNECT_SECONDS = 30.0
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def normalize_cloud_ws_url(cloud_url: str) -> str:
|
|
23
|
+
parsed = urlsplit(cloud_url)
|
|
24
|
+
if parsed.scheme not in {"http", "https", "ws", "wss"}:
|
|
25
|
+
raise ValueError("Cloud URL must start with http://, https://, ws://, or wss://")
|
|
26
|
+
if not parsed.netloc:
|
|
27
|
+
raise ValueError("Cloud URL must include a hostname")
|
|
28
|
+
|
|
29
|
+
if parsed.scheme == "http":
|
|
30
|
+
scheme = "ws"
|
|
31
|
+
elif parsed.scheme == "https":
|
|
32
|
+
scheme = "wss"
|
|
33
|
+
else:
|
|
34
|
+
scheme = parsed.scheme
|
|
35
|
+
|
|
36
|
+
path = parsed.path or ""
|
|
37
|
+
if path in {"", "/"}:
|
|
38
|
+
path = "/api/workers/ws"
|
|
39
|
+
elif not path.endswith("/api/workers/ws"):
|
|
40
|
+
path = path.rstrip("/") + "/api/workers/ws"
|
|
41
|
+
|
|
42
|
+
return urlunsplit((scheme, parsed.netloc, path, "", ""))
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class TunnelClient:
|
|
46
|
+
"""Maintains an outbound worker tunnel connection to a remote coordination service."""
|
|
47
|
+
|
|
48
|
+
def __init__(self, cloud_url: str, token: str, worker_server: WorkerServer):
|
|
49
|
+
self.websocket_url = normalize_cloud_ws_url(cloud_url)
|
|
50
|
+
self.token = token
|
|
51
|
+
self.worker_server = worker_server
|
|
52
|
+
self.handler = TunnelHandler(worker_server)
|
|
53
|
+
self._stop_event = asyncio.Event()
|
|
54
|
+
self._session: aiohttp.ClientSession | None = None
|
|
55
|
+
self._ws: aiohttp.ClientWebSocketResponse | None = None
|
|
56
|
+
self.worker_id: str | None = None
|
|
57
|
+
|
|
58
|
+
async def run(self) -> None:
|
|
59
|
+
backoff = 1.0
|
|
60
|
+
|
|
61
|
+
while not self._stop_event.is_set():
|
|
62
|
+
try:
|
|
63
|
+
await self._run_once()
|
|
64
|
+
backoff = 1.0
|
|
65
|
+
except asyncio.CancelledError:
|
|
66
|
+
raise
|
|
67
|
+
except Exception as exc:
|
|
68
|
+
if self._stop_event.is_set():
|
|
69
|
+
break
|
|
70
|
+
|
|
71
|
+
print(f"Cloud tunnel disconnected: {exc}")
|
|
72
|
+
print(f"Reconnecting in {int(backoff)}s...")
|
|
73
|
+
try:
|
|
74
|
+
await asyncio.wait_for(self._stop_event.wait(), timeout=backoff)
|
|
75
|
+
except asyncio.TimeoutError:
|
|
76
|
+
pass
|
|
77
|
+
backoff = min(backoff * 2, MAX_RECONNECT_SECONDS)
|
|
78
|
+
|
|
79
|
+
await self.close()
|
|
80
|
+
|
|
81
|
+
async def close(self) -> None:
|
|
82
|
+
self._stop_event.set()
|
|
83
|
+
if self._ws is not None and not self._ws.closed:
|
|
84
|
+
await self._ws.close()
|
|
85
|
+
if self._session is not None and not self._session.closed:
|
|
86
|
+
await self._session.close()
|
|
87
|
+
|
|
88
|
+
async def _run_once(self) -> None:
|
|
89
|
+
self._session = aiohttp.ClientSession()
|
|
90
|
+
try:
|
|
91
|
+
self._ws = await self._session.ws_connect(
|
|
92
|
+
self.websocket_url,
|
|
93
|
+
autoping=False,
|
|
94
|
+
heartbeat=None,
|
|
95
|
+
max_msg_size=MAX_MESSAGE_BYTES,
|
|
96
|
+
)
|
|
97
|
+
await self._send_json({"type": "auth", "token": self.token})
|
|
98
|
+
|
|
99
|
+
auth_message = await self._receive_json(timeout=AUTH_TIMEOUT_SECONDS)
|
|
100
|
+
if auth_message.get("type") == "auth_error":
|
|
101
|
+
raise RuntimeError(str(auth_message.get("message") or "Tunnel auth failed"))
|
|
102
|
+
if auth_message.get("type") != "auth_ok":
|
|
103
|
+
raise RuntimeError("Tunnel auth handshake returned an unexpected response")
|
|
104
|
+
|
|
105
|
+
self.worker_id = str(auth_message.get("worker_id") or "")
|
|
106
|
+
print(f"Connected to {self.websocket_url}")
|
|
107
|
+
|
|
108
|
+
ping_task = asyncio.create_task(self._ping_loop())
|
|
109
|
+
try:
|
|
110
|
+
while not self._stop_event.is_set():
|
|
111
|
+
message = await self._receive_json()
|
|
112
|
+
await self.handle_message(message)
|
|
113
|
+
finally:
|
|
114
|
+
ping_task.cancel()
|
|
115
|
+
with contextlib.suppress(asyncio.CancelledError):
|
|
116
|
+
await ping_task
|
|
117
|
+
finally:
|
|
118
|
+
self.worker_id = None
|
|
119
|
+
if self._ws is not None and not self._ws.closed:
|
|
120
|
+
await self._ws.close()
|
|
121
|
+
if self._session is not None and not self._session.closed:
|
|
122
|
+
await self._session.close()
|
|
123
|
+
self._ws = None
|
|
124
|
+
self._session = None
|
|
125
|
+
|
|
126
|
+
async def _ping_loop(self) -> None:
|
|
127
|
+
while not self._stop_event.is_set():
|
|
128
|
+
await asyncio.sleep(PING_INTERVAL_SECONDS)
|
|
129
|
+
await self._send_json({"type": "ping"})
|
|
130
|
+
|
|
131
|
+
async def _send_json(self, payload: dict[str, Any]) -> None:
|
|
132
|
+
if self._ws is None:
|
|
133
|
+
raise RuntimeError("Tunnel websocket is not connected")
|
|
134
|
+
|
|
135
|
+
encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=True)
|
|
136
|
+
if len(encoded.encode("utf-8")) > MAX_MESSAGE_BYTES:
|
|
137
|
+
raise RuntimeError("Tunnel message exceeds the 1MB size limit")
|
|
138
|
+
await self._ws.send_str(encoded)
|
|
139
|
+
|
|
140
|
+
async def _receive_json(self, *, timeout: float | None = None) -> dict[str, Any]:
|
|
141
|
+
if self._ws is None:
|
|
142
|
+
raise RuntimeError("Tunnel websocket is not connected")
|
|
143
|
+
|
|
144
|
+
message = await self._ws.receive(timeout=timeout)
|
|
145
|
+
if message.type != aiohttp.WSMsgType.TEXT:
|
|
146
|
+
if message.type in {aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED}:
|
|
147
|
+
raise RuntimeError("Cloud tunnel closed the connection")
|
|
148
|
+
if message.type == aiohttp.WSMsgType.ERROR:
|
|
149
|
+
raise RuntimeError("Cloud tunnel websocket reported an error")
|
|
150
|
+
raise RuntimeError(f"Unexpected websocket message type: {message.type}")
|
|
151
|
+
|
|
152
|
+
raw = message.data
|
|
153
|
+
if len(raw.encode("utf-8")) > MAX_MESSAGE_BYTES:
|
|
154
|
+
raise RuntimeError("Tunnel message exceeds the 1MB size limit")
|
|
155
|
+
|
|
156
|
+
try:
|
|
157
|
+
payload = json.loads(raw)
|
|
158
|
+
except json.JSONDecodeError as exc:
|
|
159
|
+
raise RuntimeError("Cloud tunnel sent invalid JSON") from exc
|
|
160
|
+
|
|
161
|
+
if not isinstance(payload, dict):
|
|
162
|
+
raise RuntimeError("Cloud tunnel sent a non-object message")
|
|
163
|
+
return payload
|
|
164
|
+
|
|
165
|
+
async def handle_message(self, message: dict[str, Any]) -> None:
|
|
166
|
+
message_type = str(message.get("type") or "")
|
|
167
|
+
|
|
168
|
+
if message_type == "pong":
|
|
169
|
+
return
|
|
170
|
+
|
|
171
|
+
if message_type == "ping":
|
|
172
|
+
await self._send_json({"type": "pong"})
|
|
173
|
+
return
|
|
174
|
+
|
|
175
|
+
try:
|
|
176
|
+
response = await self.handler.handle_message(message)
|
|
177
|
+
except Exception as exc:
|
|
178
|
+
request_id = message.get("request_id")
|
|
179
|
+
if isinstance(request_id, str) and request_id:
|
|
180
|
+
await self._send_json(
|
|
181
|
+
{
|
|
182
|
+
"type": "error",
|
|
183
|
+
"request_id": request_id,
|
|
184
|
+
"message": str(exc),
|
|
185
|
+
}
|
|
186
|
+
)
|
|
187
|
+
else:
|
|
188
|
+
print(f"Tunnel handler error: {exc}")
|
|
189
|
+
return
|
|
190
|
+
|
|
191
|
+
await self._send_json(response)
|