1use crate::channels::OutputChunk;
13use crate::channels::tmux;
14use crate::channels::transport::{Transport, session_key};
15use crate::engine::suspend;
16use chrono::{DateTime, Utc};
17use sha2::{Digest, Sha256};
18use std::cmp::Ordering;
19use std::collections::HashMap;
20use std::fmt;
21use std::fmt::Debug;
22use std::fmt::Formatter;
23use std::sync::Arc;
24use std::time::Duration;
25use tokio::sync::RwLock;
26use tokio::time;
27
28pub(crate) const MAX_OUTPUT_BUFFER_BYTES: usize = 1024 * 1024;
29
30#[derive(Debug, Clone)]
32pub struct OutputBuffer {
33 pub repo: String,
35 pub session: String,
37 pub task_id: String,
39 pub last_content: String,
41 pub last_len: usize,
43 pub last_hash: Option<[u8; 32]>,
45 pub last_capture: DateTime<Utc>,
47 pub seen_alive: bool,
53 pub registered_at: DateTime<Utc>,
55 pub has_output: bool,
60 pub generation: u64,
65}
66
67pub struct CaptureService {
69 buffers: Arc<RwLock<HashMap<String, OutputBuffer>>>,
71 transport: Arc<Transport>,
73 interval: Duration,
75}
76
77impl Debug for CaptureService {
78 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
79 f.debug_struct("CaptureService")
80 .field("interval", &self.interval)
81 .finish_non_exhaustive()
82 }
83}
84
85impl CaptureService {
86 pub fn new(transport: Arc<Transport>) -> Self {
88 Self {
89 buffers: Arc::new(RwLock::new(HashMap::new())),
90 transport,
91 interval: Duration::from_secs(2),
92 }
93 }
94
95 pub async fn register_session(&self, repo: &str, task_id: &str, session: &str) {
97 let now = Utc::now();
98 let skey = session_key(repo, task_id);
99
100 let generation = {
104 let mut buffers = self.buffers.write().await;
105 let generation = buffers.get(&skey).map_or(0, |b| b.generation + 1);
106 let buffer = OutputBuffer {
107 repo: repo.to_string(),
108 session: session.to_string(),
109 task_id: task_id.to_string(),
110 last_content: String::new(),
111 last_len: 0,
112 last_hash: None,
113 last_capture: now,
114 seen_alive: false,
115 registered_at: now,
116 has_output: false,
117 generation,
118 };
119 buffers.insert(skey, buffer);
120 generation
121 };
122 tracing::debug!(
123 repo,
124 task_id,
125 session,
126 generation,
127 "session registered for capture"
128 );
129 }
130
131 pub async fn unregister_session(&self, repo: &str, task_id: &str) {
133 let skey = session_key(repo, task_id);
134 if let Some(buffer) = self.buffers.write().await.remove(&skey) {
135 tracing::debug!(
136 repo = buffer.repo,
137 task_id = buffer.task_id,
138 session = buffer.session,
139 "session unregistered"
140 );
141 }
142 self.transport.unbind(repo, task_id).await;
143 }
144
145 pub async fn start(&self) {
149 tracing::info!(
150 interval_secs = self.interval.as_secs(),
151 "capture service started"
152 );
153 let mut interval = time::interval(self.interval);
154
155 loop {
156 interval.tick().await;
157 self.tick().await;
158 }
159 }
160
161 pub async fn run(&self) {
166 let mut interval = time::interval(self.interval);
167
168 loop {
169 interval.tick().await;
170
171 let has_sessions = !self.buffers.read().await.is_empty();
173 if !has_sessions {
174 tracing::debug!("no sessions registered, capture loop exiting");
175 break;
176 }
177
178 self.tick().await;
179 }
180 }
181
182 pub async fn get_silent_sessions_for_repo(
198 &self,
199 repo: &str,
200 grace_period: Duration,
201 ) -> Vec<(String, String, i64)> {
202 let now = Utc::now();
203 let buffers = self.buffers.read().await;
204 let mut candidates: Vec<OutputBuffer> = buffers
205 .values()
206 .filter(|buf| buf.repo == repo && !buf.has_output)
207 .cloned()
208 .collect();
209 drop(buffers);
210
211 let mut silent = Vec::new();
212 for buf in candidates.drain(..) {
213 if buf.seen_alive {
214 if !tmux::is_session_dead(&buf.session).await {
221 continue;
222 }
223 tracing::debug!(
224 task_id = buf.task_id,
225 session = buf.session,
226 "session seen-alive but now dead with no output, treating as silent exit"
227 );
228 }
229
230 let skey = session_key(&buf.repo, &buf.task_id);
234 let current_generation = {
235 let buffers = self.buffers.read().await;
236 buffers.get(&skey).map(|b| b.generation)
237 };
238 if current_generation != Some(buf.generation) {
239 tracing::debug!(
240 task_id = buf.task_id,
241 session = buf.session,
242 old_gen = buf.generation,
243 new_gen = current_generation,
244 "session generation changed, skipping silence detection (was re-registered)"
245 );
246 continue;
247 }
248
249 let raw_age = now.signed_duration_since(buf.registered_at);
250 let suspended = suspend::suspended_duration_since(buf.registered_at);
256 let age = (raw_age - suspended).max(chrono::Duration::zero());
257 if age.num_seconds() > i64::try_from(grace_period.as_secs()).unwrap_or(i64::MAX) {
258 silent.push((buf.task_id.clone(), buf.session.clone(), age.num_seconds()));
259 }
260 }
261 silent
262 }
263
264 async fn tick(&self) {
271 let buffers = self.buffers.read().await;
272 let sessions: Vec<(String, OutputBuffer)> = buffers
273 .iter()
274 .map(|(k, v)| (k.clone(), v.clone()))
275 .collect();
276 drop(buffers);
277
278 for (skey, buffer) in sessions {
279 let current_content = match tmux::capture_pane(&buffer.session).await {
281 Ok(s) => s,
282 Err(e) => {
283 let current = {
289 let buffers = self.buffers.read().await;
290 buffers.get(&skey).cloned()
291 };
292
293 let should_finalize = current.as_ref().is_some_and(|c| {
299 c.generation == buffer.generation
300 && c.seen_alive
301 && c.session == buffer.session
302 });
303
304 if should_finalize && tmux::is_session_dead(&buffer.session).await {
305 tracing::info!(
306 task_id = buffer.task_id,
307 session = buffer.session,
308 generation = buffer.generation,
309 "session ended, sending final chunk"
310 );
311 let chunk = OutputChunk {
312 content: String::new(),
313 is_final: true,
314 };
315 self.transport
316 .push_output(&buffer.repo, &buffer.task_id, chunk)
317 .await;
318 self.unregister_session(&buffer.repo, &buffer.task_id).await;
319 } else if current.is_none() {
320 tracing::trace!(
321 task_id = buffer.task_id,
322 session = buffer.session,
323 "session already unregistered"
324 );
325 } else {
326 tracing::trace!(
327 task_id = buffer.task_id,
328 session = buffer.session,
329 ?e,
330 "capture failed (transient)"
331 );
332 }
333 continue;
334 }
335 };
336
337 let new_content = {
338 let mut buffers = self.buffers.write().await;
339 if let Some(buf) = buffers.get_mut(&skey) {
340 buf.seen_alive = true;
341 buf.diff_and_update(¤t_content)
342 } else {
343 None
344 }
345 };
346
347 if let Some(new_content) = new_content {
348 let chunk = OutputChunk {
349 content: new_content,
350 is_final: false,
351 };
352 self.transport
353 .push_output(&buffer.repo, &buffer.task_id, chunk)
354 .await;
355 }
356 }
357 }
358}
359
360impl OutputBuffer {
361 pub(crate) fn diff_and_update(&mut self, current_content: &str) -> Option<String> {
362 let current_bytes = current_content.as_bytes();
363 let current_len = current_bytes.len();
364 let current_hash = hash_bytes(current_bytes);
365
366 if self.last_len == current_len && self.last_hash == Some(current_hash) {
367 self.last_capture = Utc::now();
368 return None;
369 }
370
371 let new_content = match self.last_len.cmp(¤t_len) {
372 Ordering::Greater => String::new(),
375 Ordering::Equal => {
376 let common_prefix = self
379 .last_content
380 .as_bytes()
381 .iter()
382 .zip(current_content.as_bytes())
383 .take_while(|(a, b)| a == b)
384 .count();
385 let mut start = common_prefix;
386 while start < current_len && !current_content.is_char_boundary(start) {
387 start += 1;
388 }
389 current_content[start..].to_string()
390 }
391 Ordering::Less => {
392 let mut offset = self.last_len;
393 while offset < current_len && !current_content.is_char_boundary(offset) {
394 offset += 1;
395 }
396 String::from_utf8_lossy(¤t_bytes[offset..]).to_string()
397 }
398 };
399
400 self.last_len = current_len;
401 self.last_hash = Some(current_hash);
402 self.last_content = cap_content(current_content);
403 self.last_capture = Utc::now();
404
405 if new_content.trim().is_empty() {
406 None
407 } else {
408 self.has_output = true;
409 Some(new_content)
410 }
411 }
412}
413
414fn hash_bytes(bytes: &[u8]) -> [u8; 32] {
415 let mut hasher = Sha256::new();
416 hasher.update(bytes);
417 hasher.finalize().into()
418}
419
420fn cap_content(content: &str) -> String {
421 if content.len() <= MAX_OUTPUT_BUFFER_BYTES {
422 return content.to_string();
423 }
424
425 let mut start = content.len() - MAX_OUTPUT_BUFFER_BYTES;
426 while start < content.len() && !content.is_char_boundary(start) {
427 start += 1;
428 }
429 content[start..].to_string()
430}
431
432impl CaptureService {
433 #[cfg(test)]
436 pub(crate) async fn set_buffer_state_for_test(
437 &self,
438 repo: &str,
439 task_id: &str,
440 seen_alive: bool,
441 has_output: bool,
442 registered_at: chrono::DateTime<Utc>,
443 ) {
444 let skey = session_key(repo, task_id);
445 let mut buffers = self.buffers.write().await;
446 let buf = buffers
447 .get_mut(&skey)
448 .unwrap_or_else(|| panic!("no buffer for {repo}/{task_id}"));
449 buf.seen_alive = seen_alive;
450 buf.has_output = has_output;
451 buf.registered_at = registered_at;
452 }
453}
454
455#[cfg(test)]
456mod tests {
457 use std::collections::HashSet;
458
459 use super::*;
460 use crate::channels::transport::Transport;
461
462 fn make_buffer() -> OutputBuffer {
463 OutputBuffer {
464 repo: "owner/repo".to_string(),
465 session: "test-session".to_string(),
466 task_id: "task-1".to_string(),
467 last_content: String::new(),
468 last_len: 0,
469 last_hash: None,
470 last_capture: Utc::now(),
471 seen_alive: false,
472 registered_at: Utc::now(),
473 has_output: false,
474 generation: 0,
475 }
476 }
477
478 #[test]
479 fn normal_append_returns_new_suffix() {
480 let mut buf = make_buffer();
481 let result = buf.diff_and_update("hello");
483 assert_eq!(result.as_deref(), Some("hello"));
484 let result = buf.diff_and_update("hello world");
486 assert_eq!(result.as_deref(), Some(" world"));
487 }
488
489 #[test]
490 fn no_change_returns_none() {
491 let mut buf = make_buffer();
492 buf.diff_and_update("hello");
493 let result = buf.diff_and_update("hello");
495 assert_eq!(result, None);
496 }
497
498 #[test]
499 fn terminal_clear_shrink_returns_none_not_full_content() {
500 let mut buf = make_buffer();
501 buf.diff_and_update("line1\nline2\nline3\n");
503 let result = buf.diff_and_update("prompt$ ");
505 assert_eq!(
507 result, None,
508 "shrinking pane content should produce no output, not a duplicate of visible screen"
509 );
510 }
511
512 #[test]
513 fn terminal_clear_to_empty_returns_none() {
514 let mut buf = make_buffer();
515 buf.diff_and_update("lots of output here\nmore output\n");
516 let result = buf.diff_and_update("");
518 assert_eq!(result, None);
519 }
520
521 #[test]
522 fn same_length_content_change_is_broadcast() {
523 let mut buf = make_buffer();
524 buf.diff_and_update("aaaa");
525
526 let result = buf.diff_and_update("bbbb");
527
528 assert_eq!(result.as_deref(), Some("bbbb"));
529 }
530
531 #[test]
532 fn same_length_partial_change_broadcasts_only_changed_suffix() {
533 let mut buf = make_buffer();
534 let prior = format!("{}/", "line\n".repeat(50));
536 buf.diff_and_update(&prior);
537
538 let updated = format!("{}-", "line\n".repeat(50));
540 assert_eq!(
541 prior.len(),
542 updated.len(),
543 "lengths must match for this test"
544 );
545
546 let result = buf.diff_and_update(&updated);
547
548 assert_eq!(
550 result.as_deref(),
551 Some("-"),
552 "should return only the changed suffix, not the full content"
553 );
554 }
555
556 #[test]
557 fn empty_initial_capture_returns_none() {
558 let mut buf = make_buffer();
559 let result = buf.diff_and_update("");
560 assert_eq!(result, None);
561 }
562
563 #[test]
564 fn has_output_set_on_meaningful_content() {
565 let mut buf = make_buffer();
566 assert!(!buf.has_output);
567 buf.diff_and_update("hello");
568 assert!(buf.has_output);
569 }
570
571 #[test]
572 fn has_output_not_set_on_empty_content() {
573 let mut buf = make_buffer();
574 buf.diff_and_update("");
575 assert!(!buf.has_output, "empty content should not set has_output");
576 }
577
578 #[test]
579 fn has_output_stays_true_after_silence() {
580 let mut buf = make_buffer();
581 buf.diff_and_update("hello");
582 assert!(buf.has_output);
583 buf.diff_and_update("hello");
585 assert!(buf.has_output, "has_output should not revert once set");
586 }
587
588 #[tokio::test]
589 async fn get_silent_sessions_returns_only_no_output_past_grace() {
590 let transport = Arc::new(Transport::new());
591 let svc = CaptureService::new(transport);
592 let repo = "owner/repo";
593
594 svc.register_session(repo, "silent-task", "orch-test-silent")
597 .await;
598 svc.set_buffer_state_for_test(
599 repo,
600 "silent-task",
601 false,
602 false,
603 Utc::now() - chrono::Duration::seconds(200),
604 )
605 .await;
606
607 svc.register_session(repo, "active-task", "orch-test-active")
609 .await;
610 svc.set_buffer_state_for_test(
611 repo,
612 "active-task",
613 false,
614 true, Utc::now() - chrono::Duration::seconds(200),
616 )
617 .await;
618
619 svc.register_session(repo, "new-task", "orch-test-new")
621 .await;
622
623 let grace = Duration::from_secs(120);
624 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
625 assert_eq!(silent.len(), 1);
626 assert_eq!(silent[0].0, "silent-task");
627 }
628
629 #[tokio::test]
646 async fn get_silent_sessions_seen_alive_dead_session_is_silent() {
647 let transport = Arc::new(Transport::new());
648 let svc = CaptureService::new(transport);
649 let repo = "owner/repo";
650
651 svc.register_session(repo, "never-seen-task", "orch-never-seen")
653 .await;
654 svc.set_buffer_state_for_test(
655 repo,
656 "never-seen-task",
657 false, false, Utc::now() - chrono::Duration::seconds(500),
660 )
661 .await;
662
663 svc.register_session(repo, "seen-but-dead-task", "orch-seen-dead")
669 .await;
670 svc.set_buffer_state_for_test(
671 repo,
672 "seen-but-dead-task",
673 true, false, Utc::now() - chrono::Duration::seconds(500),
676 )
677 .await;
678
679 let grace = Duration::from_secs(120);
680 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
681 assert_eq!(
684 silent.len(),
685 2,
686 "both never-alive and dead-with-no-output sessions must be silenced"
687 );
688 let ids: HashSet<&str> = silent.iter().map(|(id, _, _)| id.as_str()).collect();
689 assert!(ids.contains("never-seen-task"));
690 assert!(ids.contains("seen-but-dead-task"));
691 }
692
693 #[tokio::test]
694 async fn get_silent_sessions_filters_by_repo() {
695 let transport = Arc::new(Transport::new());
696 let svc = CaptureService::new(transport);
697
698 svc.register_session("owner/repo-a", "task-a", "orch-a")
699 .await;
700 svc.register_session("owner/repo-b", "task-b", "orch-b")
701 .await;
702 svc.set_buffer_state_for_test(
703 "owner/repo-a",
704 "task-a",
705 false,
706 false,
707 Utc::now() - chrono::Duration::seconds(200),
708 )
709 .await;
710 svc.set_buffer_state_for_test(
711 "owner/repo-b",
712 "task-b",
713 false,
714 false,
715 Utc::now() - chrono::Duration::seconds(200),
716 )
717 .await;
718
719 let grace = Duration::from_secs(120);
720 let silent_a = svc
721 .get_silent_sessions_for_repo("owner/repo-a", grace)
722 .await;
723 assert_eq!(silent_a.len(), 1);
724 assert_eq!(silent_a[0].0, "task-a");
725
726 let silent_b = svc
727 .get_silent_sessions_for_repo("owner/repo-b", grace)
728 .await;
729 assert_eq!(silent_b.len(), 1);
730 assert_eq!(silent_b[0].0, "task-b");
731 }
732
733 #[tokio::test]
735 async fn same_task_id_different_repos_capture_buffers_isolated() {
736 let transport = Arc::new(Transport::new());
737 let svc = CaptureService::new(transport);
738
739 svc.register_session("owner/repo-a", "42", "orch-a-42")
740 .await;
741 svc.register_session("owner/repo-b", "42", "orch-b-42")
742 .await;
743
744 let skey_a = session_key("owner/repo-a", "42");
745 let skey_b = session_key("owner/repo-b", "42");
746
747 let buffers = svc.buffers.read().await;
748 let buf_a = buffers
749 .get(&skey_a)
750 .expect("buffer should exist for repo-a session");
751 let buf_b = buffers
752 .get(&skey_b)
753 .expect("buffer should exist for repo-b session");
754
755 assert_eq!(buf_a.session, "orch-a-42");
756 assert_eq!(buf_b.session, "orch-b-42");
757 assert_ne!(buf_a.session, buf_b.session);
758 assert_eq!(buf_a.repo, "owner/repo-a");
759 assert_eq!(buf_b.repo, "owner/repo-b");
760 }
761
762 #[test]
763 fn multibyte_boundary_no_bytes_dropped() {
764 let mut buf = make_buffer();
767 buf.diff_and_update("Hello ");
768 let full = "Hello 世界";
770 let result = buf.diff_and_update(full);
771 let s = result.expect("should return new content");
773 assert!(s.contains('世'), "first CJK char must not be dropped");
774 assert!(s.contains('界'), "second CJK char must not be dropped");
775 }
776
777 #[test]
778 fn whitespace_only_new_content_returns_none() {
779 let mut buf = make_buffer();
780 buf.diff_and_update("hello");
781 let result = buf.diff_and_update("hello \n");
783 assert_eq!(result, None);
784 }
785
786 #[tokio::test]
794 async fn get_silent_sessions_skips_stale_generation() {
795 let transport = Arc::new(Transport::new());
796 let svc = CaptureService::new(transport);
797 let repo = "owner/repo";
798
799 svc.register_session(repo, "task-42", "orch-old-session")
801 .await;
802 svc.set_buffer_state_for_test(
803 repo,
804 "task-42",
805 false,
806 false,
807 Utc::now() - chrono::Duration::seconds(500),
808 )
809 .await;
810
811 svc.register_session(repo, "task-42", "orch-new-session")
813 .await;
814
815 let grace = Duration::from_secs(120);
816 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
817
818 assert!(
822 silent.is_empty(),
823 "re-registered session should not be silenced; stale entry should be skipped"
824 );
825 }
826}