feat: add ChromeDriver proxy for WebDriver and BiDi protocol support

images/chromium-headful/Dockerfile

@@ -261,6 +261,15 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=$CACHEIDPREFIX-ap
     apt -y install chromium && \
     apt --no-install-recommends -y install sqlite3;
 
+# Install ChromeDriver matching the installed Chromium version
+RUN set -eux; \
+    CHROMIUM_VERSION=$(chromium --version | awk '{print $2}'); \
+    curl -fsSL "https://storage.googleapis.com/chrome-for-testing-public/${CHROMIUM_VERSION}/linux64/chromedriver-linux64.zip" -o /tmp/cd.zip; \
+    unzip /tmp/cd.zip -d /tmp; \
+    mv /tmp/chromedriver-linux64/chromedriver /usr/local/bin/chromedriver; \
+    chmod +x /usr/local/bin/chromedriver; \
+    rm -rf /tmp/cd.zip /tmp/chromedriver-linux64
+
 # Copy Chromium policy configuration
 RUN mkdir -p /etc/chromium/policies/managed
 COPY shared/chromium-policies/managed/policy.json /etc/chromium/policies/managed/policy.json

images/chromium-headful/run-docker.sh

@@ -57,6 +57,7 @@ RUN_ARGS=(
   -v "$HOST_RECORDINGS_DIR:/recordings"
   --memory 8192m
   -p 9222:9222
+  -p 9224:9224
   -p 444:10001
   -e DISPLAY_NUM=1
   -e HEIGHT=1080

images/chromium-headful/run-unikernel.sh

@@ -69,6 +69,7 @@ deploy_args=(
   --vcpus ${VCPUS:-4}
   -M 4096
   -p 9222:9222/tls
+  -p 9224:9224/tls
   -p 444:10001/tls
   -e DISPLAY_NUM=1
   -e HEIGHT=1080

images/chromium-headful/supervisor/services/chromedriver.conf

@@ -0,0 +1,7 @@
+[program:chromedriver]
+command=/usr/local/bin/chromedriver --port=9225 --allowed-ips=127.0.0.1 --log-level=INFO
+autostart=false
+autorestart=true
+startsecs=2
+stdout_logfile=/var/log/supervisord/chromedriver
+redirect_stderr=true

images/chromium-headful/wrapper.sh

@@ -143,6 +143,7 @@ cleanup () {
   echo "[wrapper] Cleaning up..."
   # Re-enable scale-to-zero if the script terminates early
   enable_scale_to_zero
+  supervisorctl -c /etc/supervisor/supervisord.conf stop chromedriver || true
   supervisorctl -c /etc/supervisor/supervisord.conf stop chromium || true
   supervisorctl -c /etc/supervisor/supervisord.conf stop kernel-images-api || true
   supervisorctl -c /etc/supervisor/supervisord.conf stop dbus || true
@@ -242,6 +243,15 @@ API_OUTPUT_DIR="${KERNEL_IMAGES_API_OUTPUT_DIR:-/recordings}"
 # Start via supervisord (env overrides are read by the service's command)
 supervisorctl -c /etc/supervisor/supervisord.conf start kernel-images-api
 
+echo "[wrapper] Starting ChromeDriver via supervisord"
+supervisorctl -c /etc/supervisor/supervisord.conf start chromedriver
+for i in {1..50}; do
+  if nc -z 127.0.0.1 9225 2>/dev/null; then
+    break
+  fi
+  sleep 0.2
+done
+
 echo "[wrapper] Starting PulseAudio daemon via supervisord"
 supervisorctl -c /etc/supervisor/supervisord.conf start pulseaudio
 

images/chromium-headless/image/Dockerfile

@@ -153,7 +153,16 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=$CACHEIDPREFIX-ap
     --mount=type=cache,target=/var/lib/apt,sharing=locked,id=$CACHEIDPREFIX-apt-lib \
     apt-get update -y && \
     apt-get -y install chromium && \
-    apt-get --no-install-recommends -y install sqlite3;
+    apt-get --no-install-recommends -y install sqlite3 unzip;
+
+# Install ChromeDriver matching the installed Chromium version
+RUN set -eux; \
+    CHROMIUM_VERSION=$(chromium --version | awk '{print $2}'); \
+    curl -fsSL "https://storage.googleapis.com/chrome-for-testing-public/${CHROMIUM_VERSION}/linux64/chromedriver-linux64.zip" -o /tmp/cd.zip; \
+    unzip /tmp/cd.zip -d /tmp; \
+    mv /tmp/chromedriver-linux64/chromedriver /usr/local/bin/chromedriver; \
+    chmod +x /usr/local/bin/chromedriver; \
+    rm -rf /tmp/cd.zip /tmp/chromedriver-linux64
 
 # Copy Chromium policy configuration
 RUN mkdir -p /etc/chromium/policies/managed

images/chromium-headless/image/supervisor/services/chromedriver.conf

@@ -0,0 +1,7 @@
+[program:chromedriver]
+command=/usr/local/bin/chromedriver --port=9225 --allowed-ips=127.0.0.1 --log-level=INFO
+autostart=false
+autorestart=true
+startsecs=2
+stdout_logfile=/var/log/supervisord/chromedriver
+redirect_stderr=true

images/chromium-headless/image/wrapper.sh

@@ -184,6 +184,7 @@ cleanup () {
   echo "[wrapper] Cleaning up..."
   # Re-enable scale-to-zero if the script terminates early
   enable_scale_to_zero
+  supervisorctl -c /etc/supervisor/supervisord.conf stop chromedriver || true
   supervisorctl -c /etc/supervisor/supervisord.conf stop chromium || true
   supervisorctl -c /etc/supervisor/supervisord.conf stop xvfb || true
   supervisorctl -c /etc/supervisor/supervisord.conf stop dbus || true
@@ -245,6 +246,15 @@ while ! (echo >/dev/tcp/127.0.0.1/"${API_PORT}") >/dev/null 2>&1; do
   sleep 0.5
 done
 
+echo "[wrapper] Starting ChromeDriver via supervisord"
+supervisorctl -c /etc/supervisor/supervisord.conf start chromedriver
+for i in {1..50}; do
+  if (echo >/dev/tcp/127.0.0.1/9225) >/dev/null 2>&1; then
+    break
+  fi
+  sleep 0.2
+done
+
 echo "[wrapper] startup complete!"
 # Re-enable scale-to-zero once startup has completed (when not under Docker)
 if [[ -z "${WITHDOCKER:-}" ]]; then

images/chromium-headless/run-docker.sh

@@ -15,6 +15,7 @@ RUN_ARGS=(
   --privileged
   --tmpfs /dev/shm:size=2g
   -p 9222:9222
+  -p 9224:9224
   -p 444:10001
   -v "$HOST_RECORDINGS_DIR:/recordings"
 )

images/chromium-headless/run-unikernel.sh

@@ -20,6 +20,7 @@ deploy_args=(
   -e RUN_AS_ROOT="$RUN_AS_ROOT"
   -e LOG_CDP_MESSAGES=true
   -p 9222:9222/tls
+  -p 9224:9224/tls
   -p 444:10001/tls
   -n "$NAME"
 )

server/cmd/api/api/display.go

@@ -390,4 +390,3 @@ func (s *ApiService) setResolutionViaNeko(ctx context.Context, width, height, re
 	log.Info("successfully changed resolution via Neko API", "width", width, "height", height, "refresh_rate", refreshRate)
 	return nil
 }
-

server/cmd/api/main.go

@@ -22,6 +22,7 @@ import (
 	serverpkg "github.com/onkernel/kernel-images/server"
 	"github.com/onkernel/kernel-images/server/cmd/api/api"
 	"github.com/onkernel/kernel-images/server/cmd/config"
+	"github.com/onkernel/kernel-images/server/lib/chromedriverproxy"
 	"github.com/onkernel/kernel-images/server/lib/devtoolsproxy"
 	"github.com/onkernel/kernel-images/server/lib/logger"
 	"github.com/onkernel/kernel-images/server/lib/nekoclient"
@@ -158,98 +159,48 @@ func main() {
 		},
 		scaletozero.Middleware(stz),
 	)
-	// Expose /json/version endpoint so clients that attempt to resolve a browser
-	// websocket URL via HTTP can succeed. We map the upstream path onto this
-	// proxy's host:port so clients connect back to us.
-	// Note: Playwright's connectOverCDP requests /json/version/ with trailing slash,
-	// so we register both variants to avoid 426 errors from the WebSocket handler.
-	jsonVersionHandler := func(w http.ResponseWriter, r *http.Request) {
-		current := upstreamMgr.Current()
-		if current == "" {
-			http.Error(w, "upstream not ready", http.StatusServiceUnavailable)
-			return
-		}
-		proxyWSURL := (&url.URL{Scheme: "ws", Host: r.Host}).String()
-		w.Header().Set("Content-Type", "application/json")
-		_ = json.NewEncoder(w).Encode(map[string]string{
-			"webSocketDebuggerUrl": proxyWSURL,
-		})
-	}
+	// Proxy /json/version and /json/list to upstream Chrome with URL rewriting.
+	// Playwright's connectOverCDP requests these with trailing slashes,
+	// so we register both variants.
+	jsonVersionHandler := chromeJSONProxyHandler(upstreamMgr, slogger, "/json/version")
 	rDevtools.Get("/json/version", jsonVersionHandler)
 	rDevtools.Get("/json/version/", jsonVersionHandler)
 
-	// Handler for /json and /json/list - proxies to Chrome and rewrites URLs.
-	// This is needed for Playwright's connectOverCDP which fetches /json for target discovery.
-	jsonTargetHandler := func(w http.ResponseWriter, r *http.Request) {
-		current := upstreamMgr.Current()
-		if current == "" {
-			http.Error(w, "upstream not ready", http.StatusServiceUnavailable)
-			return
-		}
-
-		// Parse upstream URL to get Chrome's host (e.g., ws://127.0.0.1:9223/...)
-		parsed, err := url.Parse(current)
-		if err != nil {
-			http.Error(w, "invalid upstream URL", http.StatusInternalServerError)
-			return
-		}
-
-		// Fetch /json from Chrome
-		chromeJSONURL := fmt.Sprintf("http://%s/json", parsed.Host)
-		resp, err := http.Get(chromeJSONURL)
-		if err != nil {
-			slogger.Error("failed to fetch /json from Chrome", "err", err, "url", chromeJSONURL)
-			http.Error(w, "failed to fetch target list from browser", http.StatusBadGateway)
-			return
-		}
-		defer resp.Body.Close()
-
-		// Verify Chrome returned a successful response
-		if resp.StatusCode != http.StatusOK {
-			slogger.Error("Chrome /json returned non-200 status", "status", resp.StatusCode, "url", chromeJSONURL)
-			http.Error(w, fmt.Sprintf("browser returned status %d", resp.StatusCode), http.StatusBadGateway)
-			return
-		}
-
-		// Read and parse the JSON response
-		var targets []map[string]interface{}
-		if err := json.NewDecoder(resp.Body).Decode(&targets); err != nil {
-			slogger.Error("failed to decode /json response", "err", err)
-			http.Error(w, "failed to parse target list", http.StatusBadGateway)
-			return
-		}
-
-		// Rewrite URLs to use this proxy's host instead of Chrome's
-		proxyHost := r.Host
-		chromeHost := parsed.Host
-		for i := range targets {
-			// Rewrite webSocketDebuggerUrl
-			if wsURL, ok := targets[i]["webSocketDebuggerUrl"].(string); ok {
-				targets[i]["webSocketDebuggerUrl"] = rewriteWSURL(wsURL, chromeHost, proxyHost)
-			}
-			// Rewrite devtoolsFrontendUrl if present
-			if frontendURL, ok := targets[i]["devtoolsFrontendUrl"].(string); ok {
-				targets[i]["devtoolsFrontendUrl"] = rewriteWSURL(frontendURL, chromeHost, proxyHost)
-			}
-		}
-
-		w.Header().Set("Content-Type", "application/json")
-		_ = json.NewEncoder(w).Encode(targets)
-	}
+	jsonTargetHandler := chromeJSONProxyHandler(upstreamMgr, slogger, "/json")
 	rDevtools.Get("/json", jsonTargetHandler)
 	rDevtools.Get("/json/", jsonTargetHandler)
 	rDevtools.Get("/json/list", jsonTargetHandler)
 	rDevtools.Get("/json/list/", jsonTargetHandler)
-
 	rDevtools.Get("/*", func(w http.ResponseWriter, r *http.Request) {
 		devtoolsproxy.WebSocketProxyHandler(upstreamMgr, slogger, config.LogCDPMessages, stz).ServeHTTP(w, r)
 	})
 
 	srvDevtools := &http.Server{
-		Addr:    "0.0.0.0:9222",
+		Addr:    fmt.Sprintf("0.0.0.0:%d", config.DevToolsProxyPort),
 		Handler: rDevtools,
 	}
 
+	// ChromeDriver proxy: intercepts POST /session to inject debuggerAddress,
+	// proxies WebSocket (BiDi) and all other HTTP to the internal ChromeDriver.
+	rChromeDriver := chi.NewRouter()
+	rChromeDriver.Use(
+		chiMiddleware.Logger,
+		chiMiddleware.Recoverer,
+		func(next http.Handler) http.Handler {
+			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				ctxWithLogger := logger.AddToContext(r.Context(), slogger)
+				next.ServeHTTP(w, r.WithContext(ctxWithLogger))
+			})
+		},
+		scaletozero.Middleware(stz),
+	)
+	rChromeDriver.Handle("/*", chromedriverproxy.Handler(slogger))
+
+	srvChromeDriver := &http.Server{
+		Addr:    fmt.Sprintf("0.0.0.0:%d", config.ChromeDriverProxyPort),
+		Handler: rChromeDriver,
+	}
+
 	go func() {
 		slogger.Info("http server starting", "addr", srv.Addr)
 		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
@@ -266,6 +217,14 @@ func main() {
 		}
 	}()
 
+	go func() {
+		slogger.Info("chromedriver proxy starting", "addr", srvChromeDriver.Addr)
+		if err := srvChromeDriver.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			slogger.Error("chromedriver proxy failed", "err", err)
+			stop()
+		}
+	}()
+
 	// graceful shutdown
 	<-ctx.Done()
 	slogger.Info("shutdown signal received")
@@ -284,6 +243,9 @@ func main() {
 		upstreamMgr.Stop()
 		return srvDevtools.Shutdown(shutdownCtx)
 	})
+	g.Go(func() error {
+		return srvChromeDriver.Shutdown(shutdownCtx)
+	})
 
 	if err := g.Wait(); err != nil {
 		slogger.Error("server failed to shutdown", "err", err)
@@ -297,6 +259,68 @@ func mustFFmpeg() {
 	}
 }
 
+// chromeJSONProxyHandler returns a handler that proxies a JSON endpoint from
+// Chrome's DevTools API and rewrites WebSocket/DevTools URLs to point to this proxy.
+func chromeJSONProxyHandler(upstreamMgr *devtoolsproxy.UpstreamManager, slogger *slog.Logger, chromePath string) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		current := upstreamMgr.Current()
+		if current == "" {
+			http.Error(w, "upstream not ready", http.StatusServiceUnavailable)
+			return
+		}
+
+		parsed, err := url.Parse(current)
+		if err != nil {
+			http.Error(w, "invalid upstream URL", http.StatusInternalServerError)
+			return
+		}
+
+		chromeURL := fmt.Sprintf("http://%s%s", parsed.Host, chromePath)
+		resp, err := http.Get(chromeURL)
+		if err != nil {
+			slogger.Error("failed to fetch from Chrome", "err", err, "url", chromeURL)
+			http.Error(w, "failed to fetch from browser", http.StatusBadGateway)
+			return
+		}
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			slogger.Error("Chrome returned non-200 status", "status", resp.StatusCode, "url", chromeURL)
+			http.Error(w, fmt.Sprintf("browser returned status %d", resp.StatusCode), http.StatusBadGateway)
+			return
+		}
+
+		var raw interface{}
+		if err := json.NewDecoder(resp.Body).Decode(&raw); err != nil {
+			slogger.Error("failed to decode Chrome JSON response", "err", err, "path", chromePath)
+			http.Error(w, "failed to parse browser response", http.StatusBadGateway)
+			return
+		}
+
+		rewriteChromeURLs(raw, parsed.Host, r.Host)
+
+		w.Header().Set("Content-Type", "application/json")
+		_ = json.NewEncoder(w).Encode(raw)
+	}
+}
+
+var chromeURLFields = []string{"webSocketDebuggerUrl", "devtoolsFrontendUrl"}
+
+func rewriteChromeURLs(v interface{}, chromeHost, proxyHost string) {
+	switch val := v.(type) {
+	case map[string]interface{}:
+		for _, field := range chromeURLFields {
+			if s, ok := val[field].(string); ok {
+				val[field] = rewriteWSURL(s, chromeHost, proxyHost)
+			}
+		}
+	case []interface{}:
+		for _, item := range val {
+			rewriteChromeURLs(item, chromeHost, proxyHost)
+		}
+	}
+}
+
 // rewriteWSURL replaces the Chrome host with the proxy host in WebSocket URLs.
 // It handles two cases:
 // 1. Direct WebSocket URLs: "ws://127.0.0.1:9223/devtools/page/..." -> "ws://127.0.0.1:9222/devtools/page/..."

server/cmd/config/config.go

@@ -21,7 +21,11 @@ type Config struct {
 	PathToFFmpeg string `envconfig:"FFMPEG_PATH" default:"ffmpeg"`
 
 	// DevTools proxy configuration
-	LogCDPMessages bool `envconfig:"LOG_CDP_MESSAGES" default:"false"`
+	DevToolsProxyPort int  `envconfig:"DEVTOOLS_PROXY_PORT" default:"9222"`
+	LogCDPMessages    bool `envconfig:"LOG_CDP_MESSAGES" default:"false"`
+
+	// ChromeDriver proxy: external port where the proxy listens.
+	ChromeDriverProxyPort int `envconfig:"CHROMEDRIVER_PROXY_PORT" default:"9224"`
 }
 
 // Load loads configuration from environment variables