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