1use crate::bundle::types::{DefinitionSnapshot, Manifest, RunState};
6use crate::protocol::{apply_patch, ClientMessage, ServerMessage, PROTOCOL_ID};
7use anyhow::{Context, Result};
8use futures_util::{SinkExt, StreamExt};
9use serde_json::Value;
10use std::collections::{HashMap, HashSet};
11use std::sync::{Arc, Mutex};
12use std::thread::JoinHandle;
13use std::time::Duration;
14use tokio::sync::mpsc;
15use tokio_tungstenite::tungstenite::Message;
16
17pub struct RemoteView {
18 pub revision: u64,
19 generation: u64,
20 pub manifest: Manifest,
21 pub state: RunState,
22 pub snapshot: Option<DefinitionSnapshot>,
23 pub events: Vec<Value>,
24 pub session_binding: Option<Value>,
25 pub session_entries: Vec<Value>,
26 pub session_events: Vec<Value>,
27 pub session_events_malformed: bool,
28 pub session_events_torn_tail: bool,
29 pub session_capture: Option<Value>,
30 pub live: bool,
31 pub possibly_interrupted: bool,
32}
33
34fn decode_view(revision: u64, generation: u64, raw: &Value) -> Option<RemoteView> {
35 let manifest: Manifest = serde_json::from_value(raw.get("manifest")?.clone()).ok()?;
36 let state: RunState = serde_json::from_value(raw.get("state")?.clone()).ok()?;
37 let snapshot: Option<DefinitionSnapshot> = raw
38 .get("workflow")
39 .and_then(|value| serde_json::from_value(value.clone()).ok());
40 let events = raw
41 .get("events")
42 .and_then(Value::as_array)
43 .cloned()
44 .unwrap_or_default();
45 let session_binding = raw.pointer("/session/binding").cloned();
46 let session_entries = raw
47 .pointer("/session/entries")
48 .and_then(Value::as_array)
49 .cloned()
50 .unwrap_or_default();
51 let session_events = raw
52 .pointer("/session/events")
53 .and_then(Value::as_array)
54 .cloned()
55 .unwrap_or_default();
56 let session_events_malformed = raw
57 .pointer("/session/eventsMalformed")
58 .and_then(Value::as_bool)
59 .unwrap_or(false);
60 let session_events_torn_tail = raw
61 .pointer("/session/eventsTornTail")
62 .and_then(Value::as_bool)
63 .unwrap_or(false);
64 let session_capture = raw.pointer("/session/capture").cloned();
65 Some(RemoteView {
66 revision,
67 generation,
68 manifest,
69 state,
70 snapshot,
71 events,
72 session_binding,
73 session_entries,
74 session_events,
75 session_events_malformed,
76 session_events_torn_tail,
77 session_capture,
78 live: raw.get("live").and_then(Value::as_bool).unwrap_or(false),
79 possibly_interrupted: raw
80 .get("possiblyInterrupted")
81 .and_then(Value::as_bool)
82 .unwrap_or(false),
83 })
84}
85
86#[derive(Debug, Clone)]
87enum ArtifactEntry {
88 Loading,
89 Ready(String),
90 Error(String),
91}
92
93#[derive(Default)]
94struct Shared {
95 connected: bool,
96 connecting: bool,
97 reconnect_attempt: u32,
98 error: Option<String>,
99 summaries: Vec<Value>,
100 raw_views: HashMap<String, (u64, u64, Value)>,
101 next_view_generation: u64,
102 watched: HashSet<String>,
103 artifacts: HashMap<(String, String), ArtifactEntry>,
104}
105
106pub struct RemoteRuns {
107 shared: Arc<Mutex<Shared>>,
108 wake: Option<mpsc::UnboundedSender<()>>,
109 worker: Option<JoinHandle<()>>,
110 decoded: HashMap<String, RemoteView>,
111}
112
113impl RemoteRuns {
114 pub fn connect(url: &str) -> Result<Self> {
115 let shared = Arc::new(Mutex::new(Shared {
116 connecting: true,
117 ..Shared::default()
118 }));
119 let (wake_tx, wake_rx) = mpsc::unbounded_channel();
120 let task_shared = Arc::clone(&shared);
121 let url = url.to_string();
122 let worker = std::thread::spawn(move || {
123 let runtime = tokio::runtime::Builder::new_current_thread()
124 .enable_all()
125 .build()
126 .expect("tokio runtime");
127 runtime.block_on(run_reconnecting(&url, task_shared, wake_rx));
128 });
129 let _ = wake_tx.send(());
130 Ok(Self {
131 shared,
132 wake: Some(wake_tx),
133 worker: Some(worker),
134 decoded: HashMap::new(),
135 })
136 }
137
138 pub fn connected(&self) -> bool {
139 self.shared.lock().unwrap().connected
140 }
141
142 pub fn status_label(&self) -> &'static str {
143 let shared = self.shared.lock().unwrap();
144 if shared.connected {
145 "connected"
146 } else if shared.reconnect_attempt > 0 {
147 "reconnecting"
148 } else if shared.connecting {
149 "connecting"
150 } else {
151 "disconnected"
152 }
153 }
154
155 pub fn error(&self) -> Option<String> {
156 self.shared.lock().unwrap().error.clone()
157 }
158
159 pub fn summaries(&self) -> Vec<Value> {
160 self.shared.lock().unwrap().summaries.clone()
161 }
162
163 pub fn watch(&mut self, run_id: &str) {
164 let previous: Vec<String> = {
165 let mut shared = self.shared.lock().unwrap();
166 if shared.watched.len() == 1 && shared.watched.contains(run_id) {
167 return;
168 }
169 let previous: Vec<String> = shared.watched.drain().collect();
170 for old in &previous {
171 shared.raw_views.remove(old);
172 }
173 shared.watched.insert(run_id.to_string());
174 previous
175 };
176 for old in previous {
177 self.decoded.remove(&old);
178 }
179 self.wake();
180 }
181
182 pub fn request_artifact(&self, run_id: &str, path: &str) {
183 let inserted = {
184 let mut shared = self.shared.lock().unwrap();
185 let key = (run_id.to_string(), path.to_string());
186 if let std::collections::hash_map::Entry::Vacant(entry) = shared.artifacts.entry(key) {
187 entry.insert(ArtifactEntry::Loading);
188 true
189 } else {
190 false
191 }
192 };
193 if inserted {
194 self.wake();
195 }
196 }
197
198 pub fn artifact_content(&self, run_id: &str, path: &str) -> Option<Result<String, String>> {
199 match self
200 .shared
201 .lock()
202 .unwrap()
203 .artifacts
204 .get(&(run_id.to_string(), path.to_string()))
205 .cloned()?
206 {
207 ArtifactEntry::Loading => None,
208 ArtifactEntry::Ready(content) => Some(Ok(content)),
209 ArtifactEntry::Error(error) => Some(Err(error)),
210 }
211 }
212
213 pub fn artifact_snapshot(&self, run_id: &str) -> HashMap<String, Result<String, String>> {
214 self.shared
215 .lock()
216 .unwrap()
217 .artifacts
218 .iter()
219 .filter_map(|((candidate_run, path), entry)| {
220 if candidate_run != run_id {
221 return None;
222 }
223 match entry {
224 ArtifactEntry::Loading => None,
225 ArtifactEntry::Ready(content) => Some((path.clone(), Ok(content.clone()))),
226 ArtifactEntry::Error(error) => Some((path.clone(), Err(error.clone()))),
227 }
228 })
229 .collect()
230 }
231
232 pub fn view(&mut self, run_id: &str) -> Option<&RemoteView> {
233 let raw = {
234 let shared = self.shared.lock().unwrap();
235 let (revision, generation, raw) = shared.raw_views.get(run_id)?;
236 let cached = self.decoded.get(run_id);
237 if cached
238 .is_some_and(|view| view.revision == *revision && view.generation == *generation)
239 {
240 None
241 } else {
242 Some((*revision, *generation, raw.clone()))
243 }
244 };
245 if let Some((revision, generation, raw)) = raw {
246 if let Some(view) = decode_view(revision, generation, &raw) {
247 self.decoded.insert(run_id.to_string(), view);
248 }
249 }
250 self.decoded.get(run_id)
251 }
252
253 fn wake(&self) {
254 if let Some(wake) = &self.wake {
255 let _ = wake.send(());
256 }
257 }
258}
259
260impl Drop for RemoteRuns {
261 fn drop(&mut self) {
262 self.wake.take();
263 if let Some(worker) = self.worker.take() {
264 let _ = worker.join();
265 }
266 }
267}
268
269async fn run_reconnecting(
270 url: &str,
271 shared: Arc<Mutex<Shared>>,
272 mut wake: mpsc::UnboundedReceiver<()>,
273) {
274 let mut attempt = 0u32;
275 loop {
276 {
277 let mut shared = shared.lock().unwrap();
278 shared.connected = false;
279 shared.connecting = true;
280 shared.reconnect_attempt = attempt;
281 }
282 let connection = tokio::select! {
283 connection = tokio_tungstenite::connect_async(url) => connection,
284 message = wake.recv() => {
285 if message.is_none() {
286 return;
287 }
288 continue;
289 }
290 };
291 match connection {
292 Ok((socket, _)) => {
293 attempt = 0;
294 let result = run_socket(socket, Arc::clone(&shared), &mut wake).await;
295 let mut state = shared.lock().unwrap();
296 state.connected = false;
297 state.connecting = false;
298 if state.error.is_none() {
299 state.error = Some(match result {
300 Ok(()) => "connection closed".to_string(),
301 Err(error) => format!("{error:#}"),
302 });
303 }
304 }
305 Err(error) => {
306 let mut state = shared.lock().unwrap();
307 state.connected = false;
308 state.connecting = false;
309 state.error = Some(format!("connecting to {url}: {error}"));
310 }
311 }
312 attempt = attempt.saturating_add(1);
313 {
314 let mut state = shared.lock().unwrap();
315 state.reconnect_attempt = attempt;
316 }
317 let base_ms = (250u64.saturating_mul(1u64 << attempt.min(5))).min(10_000);
318 let jitter_ms = (u64::from(attempt).wrapping_mul(137)) % 251;
319 tokio::select! {
320 _ = tokio::time::sleep(Duration::from_millis(base_ms + jitter_ms)) => {}
321 message = wake.recv() => {
322 if message.is_none() {
323 return;
324 }
325 }
326 }
327 if wake.is_closed() {
328 return;
329 }
330 }
331}
332
333async fn run_socket(
334 socket: tokio_tungstenite::WebSocketStream<
335 tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
336 >,
337 shared: Arc<Mutex<Shared>>,
338 wake: &mut mpsc::UnboundedReceiver<()>,
339) -> Result<()> {
340 let (mut sink, mut reads) = socket.split();
341 let mut hello_received = false;
342 let mut subscribed = HashSet::new();
343 let mut submitted_artifacts = HashSet::new();
344
345 loop {
346 tokio::select! {
347 wake_message = wake.recv() => {
348 if wake_message.is_none() {
349 return Ok(());
350 }
351 if hello_received {
352 reconcile_desired(
353 &mut sink,
354 &shared,
355 &mut subscribed,
356 &mut submitted_artifacts,
357 ).await?;
358 }
359 }
360 incoming = reads.next() => {
361 let Some(incoming) = incoming else { return Ok(()) };
362 let text = match incoming? {
363 Message::Text(text) => text,
364 Message::Close(_) => return Ok(()),
365 _ => continue,
366 };
367 let Ok(message) = serde_json::from_str::<ServerMessage>(&text) else {
368 continue;
369 };
370 let mut resubscribe = None;
371 match message {
372 ServerMessage::Hello { protocol } => {
373 if protocol != PROTOCOL_ID {
374 anyhow::bail!("unsupported protocol {protocol}");
375 }
376 {
377 let mut state = shared.lock().unwrap();
378 state.connected = true;
379 state.connecting = false;
380 state.reconnect_attempt = 0;
381 state.error = None;
382 }
383 hello_received = true;
384 send_message(&mut sink, &ClientMessage::WatchRuns).await?;
385 reconcile_desired(
386 &mut sink,
387 &shared,
388 &mut subscribed,
389 &mut submitted_artifacts,
390 ).await?;
391 }
392 ServerMessage::Runs { runs } => {
393 shared.lock().unwrap().summaries = runs;
394 }
395 ServerMessage::RunSnapshot { run_id, revision, view } => {
396 let mut state = shared.lock().unwrap();
397 if state.watched.contains(&run_id) {
398 state.next_view_generation = state.next_view_generation.wrapping_add(1);
399 let generation = state.next_view_generation;
400 state.raw_views.insert(run_id, (revision, generation, view));
401 }
402 }
403 ServerMessage::RunPatch { run_id, revision, patch } => {
404 let mut state = shared.lock().unwrap();
405 match state.raw_views.get_mut(&run_id) {
406 Some((current, _, view)) if revision == *current + 1 => {
407 if apply_patch(view, &patch).is_ok() {
408 *current = revision;
409 } else {
410 resubscribe = Some(run_id);
411 }
412 }
413 Some(_) => resubscribe = Some(run_id),
414 None => {}
415 }
416 }
417 ServerMessage::Artifact { run_id, path, content } => {
418 let key = (run_id, path);
419 submitted_artifacts.remove(&key);
420 shared
421 .lock()
422 .unwrap()
423 .artifacts
424 .insert(key, ArtifactEntry::Ready(content));
425 }
426 ServerMessage::Error { message, run_id } => {
427 let mut state = shared.lock().unwrap();
428 if let Some(run_id) = run_id {
429 if let Some(key) = submitted_artifacts
430 .iter()
431 .find(|(candidate_run, _)| candidate_run == &run_id)
432 .cloned()
433 {
434 submitted_artifacts.remove(&key);
435 state.artifacts.insert(key, ArtifactEntry::Error(message));
436 } else {
437 state.error = Some(message);
438 }
439 } else {
440 state.error = Some(message);
441 }
442 }
443 }
444 if hello_received {
445 reconcile_desired(
446 &mut sink,
447 &shared,
448 &mut subscribed,
449 &mut submitted_artifacts,
450 ).await?;
451 }
452 if let Some(run_id) = resubscribe {
453 send_message(&mut sink, &ClientMessage::WatchRun { run_id }).await?;
454 }
455 }
456 }
457 }
458}
459
460async fn reconcile_desired(
461 sink: &mut (impl SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
462 shared: &Arc<Mutex<Shared>>,
463 subscribed: &mut HashSet<String>,
464 submitted_artifacts: &mut HashSet<(String, String)>,
465) -> Result<()> {
466 let (desired, artifacts) = {
467 let state = shared.lock().unwrap();
468 let desired = state.watched.clone();
469 let artifacts = state
470 .artifacts
471 .iter()
472 .filter_map(|(key, entry)| {
473 matches!(entry, ArtifactEntry::Loading).then_some(key.clone())
474 })
475 .collect::<Vec<_>>();
476 (desired, artifacts)
477 };
478 let removals: Vec<String> = subscribed.difference(&desired).cloned().collect();
479 let additions: Vec<String> = desired.difference(subscribed).cloned().collect();
480 for run_id in removals {
481 send_message(
482 sink,
483 &ClientMessage::UnwatchRun {
484 run_id: run_id.clone(),
485 },
486 )
487 .await?;
488 subscribed.remove(&run_id);
489 }
490 for run_id in additions {
491 send_message(
492 sink,
493 &ClientMessage::WatchRun {
494 run_id: run_id.clone(),
495 },
496 )
497 .await?;
498 subscribed.insert(run_id);
499 }
500 if submitted_artifacts.is_empty() {
501 if let Some((run_id, path)) = artifacts.into_iter().next() {
502 let key = (run_id.clone(), path.clone());
503 submitted_artifacts.insert(key);
504 send_message(sink, &ClientMessage::FetchArtifact { run_id, path }).await?;
505 }
506 }
507 Ok(())
508}
509
510async fn send_message(
511 sink: &mut (impl SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
512 message: &ClientMessage,
513) -> Result<()> {
514 let text = serde_json::to_string(message).context("encoding client message")?;
515 sink.send(Message::Text(text.into())).await?;
516 Ok(())
517}