Use db in token redis plugin

Author: ajhai

llmstack/common/runner/server.py

@@ -317,11 +317,18 @@ def main():
     redis_client.ping()
 
     display_pool = VirtualDisplayPool(
-        redis_client, hostname=args.hostname, max_displays=args.max_displays, start_display=args.start_display, display_res=args.display_res, rfb_start_port=args.rfb_start_port)
+        redis_client, hostname=args.hostname, max_displays=args.max_displays,
+        start_display=args.start_display, display_res=args.display_res,
+        rfb_start_port=args.rfb_start_port)
 
     # Start websockify server
-    websockify_process = subprocess.Popen(['websockify', f'{args.wss_port}', '--web', '/usr/share/www/html', '--token-plugin=TokenRedis', f'--token-source={args.redis_host}:{args.redis_port}',
-                                           '-v', '--auth-plugin=llmstack.common.runner.auth.BasicHTTPAuthWithRedis', f'--auth-source={args.redis_host}:{args.redis_port}{f":{args.redis_password}" if args.redis_password else ""}'], close_fds=True)
+    websockify_process = subprocess.Popen(
+        ['websockify', f'{args.wss_port}', '--web', '/usr/share/www/html',
+         '--token-plugin=llmstack.common.runner.token.TokenRedis',
+         f'--token-source={args.redis_host}:{args.redis_port}:{args.redis_db}{f":{args.redis_password}" if args.redis_password else ""}',
+         '-v', '--auth-plugin=llmstack.common.runner.auth.BasicHTTPAuthWithRedis',
+         f'--auth-source={args.redis_host}:{args.redis_port}:{args.redis_db}{f":{args.redis_password}" if args.redis_password else ""}'],
+        close_fds=True)
 
     server = grpc_server(futures.ThreadPoolExecutor(max_workers=10))
     runner = Runner(display_pool=display_pool)

llmstack/common/runner/token.py

@@ -0,0 +1,52 @@
+import json
+import logging
+import re
+import sys
+
+import redis
+
+logger = logging.getLogger(__name__)
+
+
+class TokenRedis():
+    def __init__(self, src):
+        try:
+            # Split parts and use defaults for missing fields
+            parts = src.split(":")
+            self._server = parts[0]
+            self._port = int(parts[1]) if len(parts) > 1 and parts[1] else 6379
+            self._db = int(parts[2]) if len(parts) > 2 and parts[2] else 0
+            self._password = parts[3] if len(parts) > 3 else None
+
+            logger.info(
+                f'TokenRedis initialized ({self._server}:{self._port})')
+        except (ValueError, IndexError):
+            logger.error(f'Invalid format: {src}')
+            sys.exit()
+
+    def lookup(self, token):
+        logger.info(f'resolving token "{token}"')
+        client = redis.Redis(host=self._server, port=self._port,
+                             db=self._db, password=self._password)
+
+        stuff = client.get(token)
+        if stuff is None:
+            return None
+
+        response_str = stuff.decode("utf-8").strip()
+        logger.debug(f'response from redis: {response_str}')
+
+        try:
+            # Attempt to load JSON
+            if response_str.startswith("{"):
+                data = json.loads(response_str)
+                host, port = data["host"].split(":")
+            # Or parse simple format
+            elif re.match(r'\S+:\S+', response_str):
+                host, port = response_str.split(":")
+            else:
+                raise ValueError('Unable to parse token')
+            return [host, port]
+        except (ValueError, json.JSONDecodeError) as e:
+            logger.error(f'Unable to process token: {response_str}')
+            return None

runner/Dockerfile

@@ -6,9 +6,7 @@ ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && apt-get install -y \
     python3 \
     python3-pip \
-    xdotool \
     x11vnc \
-    websockify \
     novnc \
     redis-server \
     && rm -rf /var/lib/apt/lists/*