add changes to go & py clients

seanmcguire12 · seanmcguire12 · commit d5d52769e7aa · 2026-05-20T11:31:57.000-07:00
diff --git a/go/modcdp/client/ModCDPClient.go b/go/modcdp/client/ModCDPClient.go
@@ -54,6 +54,7 @@ const DefaultServiceWorkerReadyTimeoutMS = 60_000
 const DefaultServiceWorkerPollIntervalMS = 100
 const DefaultTargetSessionPollIntervalMS = 20
 const DefaultWSConnectErrorSettleTimeoutMS = 250
+const DefaultClientHeartbeatIntervalMS = 250
 
 func boolPointer(value bool) *bool {
 	return &value
@@ -180,14 +181,16 @@ func freePort() (int, error) {
 // --- public types --------------------------------------------------------
 
 type ServerConfig struct {
-	ServerLoopbackCDPURL                    string            `json:"server_loopback_cdp_url,omitempty"`
-	ServerRoutes                            map[string]string `json:"server_routes,omitempty"`
-	ServerBrowserToken                      string            `json:"server_browser_token,omitempty"`
-	ServerCDPSendTimeoutMS                  int               `json:"server_cdp_send_timeout_ms,omitempty"`
-	ServerLoopbackExecutionContextTimeoutMS int               `json:"server_loopback_execution_context_timeout_ms,omitempty"`
-	ServerWSConnectErrorSettleTimeoutMS     int               `json:"server_ws_connect_error_settle_timeout_ms,omitempty"`
-	Options                                 map[string]any    `json:"-"`
-	disabled                                bool
+	ServerLoopbackCDPURL                     string            `json:"server_loopback_cdp_url,omitempty"`
+	ServerRoutes                             map[string]string `json:"server_routes,omitempty"`
+	ServerBrowserToken                       string            `json:"server_browser_token,omitempty"`
+	ServerCDPSendTimeoutMS                   int               `json:"server_cdp_send_timeout_ms,omitempty"`
+	ServerLoopbackExecutionContextTimeoutMS  int               `json:"server_loopback_execution_context_timeout_ms,omitempty"`
+	ServerWSConnectErrorSettleTimeoutMS      int               `json:"server_ws_connect_error_settle_timeout_ms,omitempty"`
+	ServerDownstreamClientTimeoutMS          int               `json:"server_downstream_client_timeout_ms,omitempty"`
+	ServerCloseBrowserOnDownstreamDisconnect *bool             `json:"server_close_browser_on_downstream_disconnect,omitempty"`
+	Options                                  map[string]any    `json:"-"`
+	disabled                                 bool
 }
 
 var ServerNone = &ServerConfig{disabled: true}
@@ -254,6 +257,7 @@ type ClientConfig struct {
 	ClientMirrorUpstreamEvents *bool             `json:"client_mirror_upstream_events,omitempty"`
 	ClientCDPSendTimeoutMS     int               `json:"client_cdp_send_timeout_ms,omitempty"`
 	ClientEventWaitTimeoutMS   int               `json:"client_event_wait_timeout_ms,omitempty"`
+	ClientHeartbeatIntervalMS  int               `json:"client_heartbeat_interval_ms,omitempty"`
 }
 
 type Options struct {
@@ -406,6 +410,7 @@ type ModCDPClient struct {
 	launchedBrowser          *LaunchedBrowser
 	extensionInjectors       []extensionInjector
 	configuredPeerGeneration int64
+	heartbeatStop            chan struct{}
 }
 
 type extensionInjector interface {
@@ -501,6 +506,9 @@ func New(opts Options) *ModCDPClient {
 	if opts.Client.ClientEventWaitTimeoutMS == 0 {
 		opts.Client.ClientEventWaitTimeoutMS = DefaultEventWaitTimeoutMS
 	}
+	if opts.Client.ClientHeartbeatIntervalMS == 0 {
+		opts.Client.ClientHeartbeatIntervalMS = DefaultClientHeartbeatIntervalMS
+	}
 	if opts.Injector.InjectorExecutionContextTimeoutMS == 0 {
 		opts.Injector.InjectorExecutionContextTimeoutMS = DefaultExecutionContextTimeoutMS
 	}
@@ -575,7 +583,10 @@ func (c *ModCDPClient) Connect() error {
 		return fmt.Errorf("upstream transport did not connect")
 	}
 	c.transport.OnRecv(func(message map[string]any) { c.handleMessage(message) })
-	c.transport.OnClose(func(err error) { c.rejectAll(err) })
+	c.transport.OnClose(func(err error) {
+		c.stopHeartbeat()
+		c.rejectAll(err)
+	})
 	if transportpkg.EndpointKindForUpstream(c.Upstream.UpstreamMode) == UpstreamEndpointKindModCDPServer {
 		if err := c.transport.WaitForPeer(); err != nil {
 			c.Close()
@@ -588,6 +599,7 @@ func (c *ModCDPClient) Connect() error {
 			}
 			c.configuredPeerGeneration = c.transport.PeerGeneration()
 		}
+		c.startHeartbeat()
 		c.startPingLatencyMeasurement()
 		connectedAt := time.Now().UnixMilli()
 		c.ConnectTiming = map[string]any{
@@ -683,6 +695,7 @@ func (c *ModCDPClient) Connect() error {
 			return fmt.Errorf("Mod.configure: %w", err)
 		}
 	}
+	c.startHeartbeat()
 	c.startPingLatencyMeasurement()
 	connectedAt := time.Now().UnixMilli()
 	c.ConnectTiming = map[string]any{
@@ -899,6 +912,7 @@ func (c *ModCDPClient) serverConfigureParams(customCommands []map[string]any, cu
 		"server_cdp_send_timeout_ms":                   c.Client.ClientCDPSendTimeoutMS,
 		"server_loopback_execution_context_timeout_ms": c.Injector.InjectorExecutionContextTimeoutMS,
 		"server_ws_connect_error_settle_timeout_ms":    c.Upstream.UpstreamWSConnectErrorSettleTimeoutMS,
+		"server_downstream_client_timeout_ms":          maxInt(c.Client.ClientHeartbeatIntervalMS*4, 1_000),
 	}
 	if c.Server != nil {
 		server["server_loopback_cdp_url"] = c.Server.ServerLoopbackCDPURL
@@ -915,6 +929,12 @@ func (c *ModCDPClient) serverConfigureParams(customCommands []map[string]any, cu
 		if c.Server.ServerWSConnectErrorSettleTimeoutMS != 0 {
 			server["server_ws_connect_error_settle_timeout_ms"] = c.Server.ServerWSConnectErrorSettleTimeoutMS
 		}
+		if c.Server.ServerDownstreamClientTimeoutMS != 0 {
+			server["server_downstream_client_timeout_ms"] = c.Server.ServerDownstreamClientTimeoutMS
+		}
+		if c.Server.ServerCloseBrowserOnDownstreamDisconnect != nil {
+			server["server_close_browser_on_downstream_disconnect"] = *c.Server.ServerCloseBrowserOnDownstreamDisconnect
+		}
 		for key, value := range c.Server.Options {
 			server[key] = value
 		}
@@ -1410,6 +1430,7 @@ func handlerPointer(handler Handler) uintptr {
 }
 
 func (c *ModCDPClient) Close() {
+	c.stopHeartbeat()
 	if c.launchedBrowser != nil {
 		c.launchedBrowser.Close()
 		c.launchedBrowser = nil
@@ -1664,6 +1685,48 @@ func (c *ModCDPClient) startPingLatencyMeasurement() {
 	}()
 }
 
+func (c *ModCDPClient) startHeartbeat() {
+	c.stopHeartbeat()
+	if c.Server == nil || c.Server.ServerCloseBrowserOnDownstreamDisconnect == nil || !*c.Server.ServerCloseBrowserOnDownstreamDisconnect {
+		return
+	}
+	interval := c.Client.ClientHeartbeatIntervalMS
+	if interval <= 0 {
+		return
+	}
+	stop := make(chan struct{})
+	c.heartbeatStop = stop
+	go func() {
+		ticker := time.NewTicker(time.Duration(interval) * time.Millisecond)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ticker.C:
+				if _, err := c.Send("Mod.ping", map[string]any{"sent_at": time.Now().UnixMilli()}); err != nil {
+					return
+				}
+			case <-stop:
+				return
+			}
+		}
+	}()
+}
+
+func (c *ModCDPClient) stopHeartbeat() {
+	if c.heartbeatStop == nil {
+		return
+	}
+	close(c.heartbeatStop)
+	c.heartbeatStop = nil
+}
+
+func maxInt(left int, right int) int {
+	if left > right {
+		return left
+	}
+	return right
+}
+
 func numberAsInt64(value any) (int64, bool) {
 	switch v := value.(type) {
 	case int64:
diff --git a/python/modcdp/client/ModCDPClient.py b/python/modcdp/client/ModCDPClient.py
@@ -187,6 +187,7 @@ def ping(self, **params: Any) -> AwaitableDict | AwaitableValue:
 DEFAULT_SERVICE_WORKER_POLL_INTERVAL_MS = 100
 DEFAULT_TARGET_SESSION_POLL_INTERVAL_MS = 20
 DEFAULT_WS_CONNECT_ERROR_SETTLE_TIMEOUT_MS = 250
+DEFAULT_CLIENT_HEARTBEAT_INTERVAL_MS = 250
 
 
 class _RawCDP:
@@ -313,6 +314,9 @@ def __init__(
             "client_mirror_upstream_events": bool(client_input.get("client_mirror_upstream_events", True)),
             "client_cdp_send_timeout_ms": int(_defaulted(client_input.get("client_cdp_send_timeout_ms"), DEFAULT_CDP_SEND_TIMEOUT_MS)),
             "client_event_wait_timeout_ms": int(_defaulted(client_input.get("client_event_wait_timeout_ms"), DEFAULT_EVENT_WAIT_TIMEOUT_MS)),
+            "client_heartbeat_interval_ms": int(
+                _defaulted(client_input.get("client_heartbeat_interval_ms"), DEFAULT_CLIENT_HEARTBEAT_INTERVAL_MS)
+            ),
         }
         self.cdp_url: str | None = cast(str | None, self.upstream.get("upstream_cdp_url"))
         if server is DEFAULT_SERVER:
@@ -362,6 +366,8 @@ def __init__(
         self._closed = False
         self._launched_browser: Any | None = None
         self._extension_injectors: list[ExtensionInjector] = []
+        self._heartbeat_stop: threading.Event | None = None
+        self._heartbeat_thread: threading.Thread | None = None
         self._cdp = _RawCDP(self)
         self._hydrate_native_protocol_schemas()
         self._hydrate_custom_surface()
@@ -374,13 +380,14 @@ def connect(self) -> "ModCDPClient":
         if self.transport is None:
             raise RuntimeError("upstream transport did not connect.")
         self.transport.onRecv(lambda message: self._on_recv(cast(CdpMessage, message)))
-        self.transport.onClose(lambda error: self._reject_all(error))
+        self.transport.onClose(lambda error: self._handle_transport_close(error))
 
         if self.upstream_endpoint_kind == "modcdp_server":
             self.transport.waitForPeer()
             if self.server is not None:
                 self._send_message("Mod.configure", cast(ProtocolParams, self._server_configure_params()))
             threading.Thread(target=self._measure_ping_latency, daemon=True).start()
+            self._start_heartbeat()
             connected_at = int(time.time() * 1000)
             self.connect_timing = cast(ModCDPConnectTiming, {
                 "started_at": connect_started_at,
@@ -420,6 +427,7 @@ def connect(self) -> "ModCDPClient":
                 routes=cast(ModCDPRoutes, self.client["client_routes"]),
                 cdp_session_id=self.ext_session_id,
             ))
+        self._start_heartbeat()
         threading.Thread(target=self._measure_ping_latency, daemon=True).start()
         connected_at = int(time.time() * 1000)
         self.connect_timing = cast(ModCDPConnectTiming, {
@@ -617,6 +625,10 @@ def _server_configure_params(self) -> ModCDPServerConfig:
             "server_ws_connect_error_settle_timeout_ms",
             self.upstream["upstream_ws_connect_error_settle_timeout_ms"],
         )
+        server_downstream_client_timeout_ms = server.pop(
+            "server_downstream_client_timeout_ms",
+            max(int(self.client["client_heartbeat_interval_ms"]) * 4, 1_000),
+        )
         custom_events: list[ModCDPAddCustomEventObjectParams] = []
         for event in self.custom_events:
             custom_events.append(
@@ -645,6 +657,7 @@ def _server_configure_params(self) -> ModCDPServerConfig:
                 "server_cdp_send_timeout_ms": server_cdp_send_timeout_ms,
                 "server_loopback_execution_context_timeout_ms": server_loopback_execution_context_timeout_ms,
                 "server_ws_connect_error_settle_timeout_ms": server_ws_connect_error_settle_timeout_ms,
+                "server_downstream_client_timeout_ms": server_downstream_client_timeout_ms,
                 **({"server_routes": server_routes} if server_routes is not None else {}),
                 **({"server_loopback_cdp_url": server_loopback_cdp_url} if server_loopback_cdp_url is not None else {}),
                 **({"server_browser_token": server_browser_token} if server_browser_token is not None else {}),
@@ -659,6 +672,7 @@ def close(self) -> None:
         if self._closed:
             return
         self._closed = True
+        self._stop_heartbeat()
         if self._launched_browser is not None:
             self._launched_browser["close"]()
             self._launched_browser = None
@@ -675,6 +689,35 @@ def close(self) -> None:
                 pass
         self._extension_injectors = []
 
+    def _handle_transport_close(self, error: Exception) -> None:
+        self._stop_heartbeat()
+        self._reject_all(error)
+
+    def _start_heartbeat(self) -> None:
+        self._stop_heartbeat()
+        if not self.server or self.server.get("server_close_browser_on_downstream_disconnect") is not True:
+            return
+        interval_ms = int(self.client["client_heartbeat_interval_ms"])
+        stop = threading.Event()
+        self._heartbeat_stop = stop
+
+        def run() -> None:
+            while not stop.wait(interval_ms / 1000):
+                try:
+                    self.send("Mod.ping", {"sent_at": int(time.time() * 1000)})
+                except Exception:
+                    return
+
+        self._heartbeat_thread = threading.Thread(target=run, daemon=True)
+        self._heartbeat_thread.start()
+
+    def _stop_heartbeat(self) -> None:
+        stop = self._heartbeat_stop
+        self._heartbeat_stop = None
+        if stop is not None:
+            stop.set()
+        self._heartbeat_thread = None
+
     def _session_id_for_target(self, target_id: str, timeout: float = 0) -> str | None:
         if timeout <= 0:
             return self.auto_sessions.sessionIdForTarget(target_id)
diff --git a/python/modcdp/types/modcdp.py b/python/modcdp/types/modcdp.py
@@ -92,6 +92,8 @@ class ModCDPServerConfig(TypedDict, total=False):
     server_cdp_send_timeout_ms: int
     server_loopback_execution_context_timeout_ms: int
     server_ws_connect_error_settle_timeout_ms: int
+    server_downstream_client_timeout_ms: int
+    server_close_browser_on_downstream_disconnect: bool
     server_browser_token: str | None
     custom_commands: list[ModCDPAddCustomCommandParams]
     custom_events: list[ModCDPAddCustomEventObjectParams]
