diff --git a/src/worker_runner.rs b/src/worker_runner.rs index 2e7d112..abed451 100644 --- a/src/worker_runner.rs +++ b/src/worker_runner.rs @@ -102,6 +102,19 @@ async fn connect_and_run(server_url: &str, worker_name: &str, llm_config: &crate let svc_mgr = ServiceManager::new(9100); let ws_tx = Arc::new(tokio::sync::Mutex::new(ws_tx)); + // Ping task to keep connection alive + let ping_tx = ws_tx.clone(); + tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(30)); + loop { + interval.tick().await; + let mut tx = ping_tx.lock().await; + if tx.send(Message::Ping(vec![].into())).await.is_err() { + break; + } + } + }); + // Channel for forwarding comments to the running workflow let comment_tx: Arc>>> = Arc::new(tokio::sync::Mutex::new(None)); @@ -111,6 +124,7 @@ async fn connect_and_run(server_url: &str, worker_name: &str, llm_config: &crate let text = match msg? { Message::Text(t) => t, Message::Close(_) => break, + Message::Pong(_) => continue, _ => continue, };