1use crate::channels::tmux;
13use crate::channels::transport::{session_key, Transport};
14use crate::channels::OutputChunk;
15use chrono::{DateTime, Utc};
16use sha2::{Digest, Sha256};
17use std::collections::HashMap;
18use std::sync::Arc;
19use tokio::sync::RwLock;
20
21pub(crate) const MAX_OUTPUT_BUFFER_BYTES: usize = 1024 * 1024;
22
23#[derive(Debug, Clone)]
25pub struct OutputBuffer {
26 pub repo: String,
28 pub session: String,
30 pub task_id: String,
32 pub last_content: String,
34 pub last_len: usize,
36 pub last_hash: Option<[u8; 32]>,
38 pub last_capture: DateTime<Utc>,
40 pub seen_alive: bool,
46 pub registered_at: DateTime<Utc>,
48 pub has_output: bool,
53 pub generation: u64,
58}
59
60pub struct CaptureService {
62 buffers: Arc<RwLock<HashMap<String, OutputBuffer>>>,
64 transport: Arc<Transport>,
66 interval: std::time::Duration,
68}
69
70impl CaptureService {
71 pub fn new(transport: Arc<Transport>) -> Self {
73 Self {
74 buffers: Arc::new(RwLock::new(HashMap::new())),
75 transport,
76 interval: std::time::Duration::from_secs(2),
77 }
78 }
79
80 pub async fn register_session(&self, repo: &str, task_id: &str, session: &str) {
82 let now = Utc::now();
83 let skey = session_key(repo, task_id);
84
85 let generation = {
89 let mut buffers = self.buffers.write().await;
90 let generation = buffers.get(&skey).map(|b| b.generation + 1).unwrap_or(0);
91 let buffer = OutputBuffer {
92 repo: repo.to_string(),
93 session: session.to_string(),
94 task_id: task_id.to_string(),
95 last_content: String::new(),
96 last_len: 0,
97 last_hash: None,
98 last_capture: now,
99 seen_alive: false,
100 registered_at: now,
101 has_output: false,
102 generation,
103 };
104 buffers.insert(skey, buffer);
105 generation
106 };
107 tracing::debug!(
108 repo,
109 task_id,
110 session,
111 generation,
112 "session registered for capture"
113 );
114 }
115
116 pub async fn unregister_session(&self, repo: &str, task_id: &str) {
118 let skey = session_key(repo, task_id);
119 if let Some(buffer) = self.buffers.write().await.remove(&skey) {
120 tracing::debug!(
121 repo = buffer.repo,
122 task_id = buffer.task_id,
123 session = buffer.session,
124 "session unregistered"
125 );
126 }
127 self.transport.unbind(repo, task_id).await;
128 }
129
130 pub async fn start(&self) {
134 tracing::info!(
135 interval_secs = self.interval.as_secs(),
136 "capture service started"
137 );
138 let mut interval = tokio::time::interval(self.interval);
139
140 loop {
141 interval.tick().await;
142 self.tick().await;
143 }
144 }
145
146 pub async fn run(&self) {
151 let mut interval = tokio::time::interval(self.interval);
152
153 loop {
154 interval.tick().await;
155
156 let has_sessions = !self.buffers.read().await.is_empty();
158 if !has_sessions {
159 tracing::debug!("no sessions registered, capture loop exiting");
160 break;
161 }
162
163 self.tick().await;
164 }
165 }
166
167 pub async fn get_silent_sessions_for_repo(
183 &self,
184 repo: &str,
185 grace_period: std::time::Duration,
186 ) -> Vec<(String, String, i64)> {
187 let now = Utc::now();
188 let buffers = self.buffers.read().await;
189 let mut candidates: Vec<OutputBuffer> = buffers
190 .values()
191 .filter(|buf| buf.repo == repo && !buf.has_output)
192 .cloned()
193 .collect();
194 drop(buffers);
195
196 let mut silent = Vec::new();
197 for buf in candidates.drain(..) {
198 if buf.seen_alive {
199 if !tmux::is_session_dead(&buf.session).await {
206 continue;
207 }
208 tracing::debug!(
209 task_id = buf.task_id,
210 session = buf.session,
211 "session seen-alive but now dead with no output, treating as silent exit"
212 );
213 }
214
215 let skey = session_key(&buf.repo, &buf.task_id);
219 let current_generation = {
220 let buffers = self.buffers.read().await;
221 buffers.get(&skey).map(|b| b.generation)
222 };
223 if current_generation != Some(buf.generation) {
224 tracing::debug!(
225 task_id = buf.task_id,
226 session = buf.session,
227 old_gen = buf.generation,
228 new_gen = current_generation,
229 "session generation changed, skipping silence detection (was re-registered)"
230 );
231 continue;
232 }
233
234 let raw_age = now.signed_duration_since(buf.registered_at);
235 let suspended = crate::engine::suspend::suspended_duration_since(buf.registered_at);
241 let age = (raw_age - suspended).max(chrono::Duration::zero());
242 if age.num_seconds() > grace_period.as_secs() as i64 {
243 silent.push((buf.task_id.clone(), buf.session.clone(), age.num_seconds()));
244 }
245 }
246 silent
247 }
248
249 async fn tick(&self) {
256 let buffers = self.buffers.read().await;
257 let sessions: Vec<(String, OutputBuffer)> = buffers
258 .iter()
259 .map(|(k, v)| (k.clone(), v.clone()))
260 .collect();
261 drop(buffers);
262
263 for (skey, buffer) in sessions {
264 let current_content = match tmux::capture_pane(&buffer.session).await {
266 Ok(s) => s,
267 Err(e) => {
268 let current = {
274 let buffers = self.buffers.read().await;
275 buffers.get(&skey).cloned()
276 };
277
278 let should_finalize = current
284 .as_ref()
285 .map(|c| {
286 c.generation == buffer.generation
287 && c.seen_alive
288 && c.session == buffer.session
289 })
290 .unwrap_or(false);
291
292 if should_finalize && tmux::is_session_dead(&buffer.session).await {
293 tracing::info!(
294 task_id = buffer.task_id,
295 session = buffer.session,
296 generation = buffer.generation,
297 "session ended, sending final chunk"
298 );
299 let chunk = OutputChunk {
300 content: String::new(),
301 is_final: true,
302 };
303 self.transport
304 .push_output(&buffer.repo, &buffer.task_id, chunk)
305 .await;
306 self.unregister_session(&buffer.repo, &buffer.task_id).await;
307 } else if current.is_none() {
308 tracing::trace!(
309 task_id = buffer.task_id,
310 session = buffer.session,
311 "session already unregistered"
312 );
313 } else {
314 tracing::trace!(
315 task_id = buffer.task_id,
316 session = buffer.session,
317 ?e,
318 "capture failed (transient)"
319 );
320 }
321 continue;
322 }
323 };
324
325 let new_content = {
326 let mut buffers = self.buffers.write().await;
327 if let Some(buf) = buffers.get_mut(&skey) {
328 buf.seen_alive = true;
329 buf.diff_and_update(¤t_content)
330 } else {
331 None
332 }
333 };
334
335 if let Some(new_content) = new_content {
336 let chunk = OutputChunk {
337 content: new_content,
338 is_final: false,
339 };
340 self.transport
341 .push_output(&buffer.repo, &buffer.task_id, chunk)
342 .await;
343 }
344 }
345 }
346}
347
348impl OutputBuffer {
349 pub(crate) fn diff_and_update(&mut self, current_content: &str) -> Option<String> {
350 let current_bytes = current_content.as_bytes();
351 let current_len = current_bytes.len();
352 let current_hash = hash_bytes(current_bytes);
353
354 if self.last_len == current_len && self.last_hash == Some(current_hash) {
355 self.last_capture = Utc::now();
356 return None;
357 }
358
359 let new_content = if self.last_len > current_len {
360 String::new()
363 } else if self.last_len == current_len {
364 let common_prefix = self
367 .last_content
368 .as_bytes()
369 .iter()
370 .zip(current_content.as_bytes())
371 .take_while(|(a, b)| a == b)
372 .count();
373 let mut start = common_prefix;
374 while start < current_len && !current_content.is_char_boundary(start) {
375 start += 1;
376 }
377 current_content[start..].to_string()
378 } else {
379 let mut offset = self.last_len;
380 while offset < current_len && !current_content.is_char_boundary(offset) {
381 offset += 1;
382 }
383 String::from_utf8_lossy(¤t_bytes[offset..]).to_string()
384 };
385
386 self.last_len = current_len;
387 self.last_hash = Some(current_hash);
388 self.last_content = cap_content(current_content);
389 self.last_capture = Utc::now();
390
391 if new_content.trim().is_empty() {
392 None
393 } else {
394 self.has_output = true;
395 Some(new_content)
396 }
397 }
398}
399
400fn hash_bytes(bytes: &[u8]) -> [u8; 32] {
401 let mut hasher = Sha256::new();
402 hasher.update(bytes);
403 hasher.finalize().into()
404}
405
406fn cap_content(content: &str) -> String {
407 if content.len() <= MAX_OUTPUT_BUFFER_BYTES {
408 return content.to_string();
409 }
410
411 let mut start = content.len() - MAX_OUTPUT_BUFFER_BYTES;
412 while start < content.len() && !content.is_char_boundary(start) {
413 start += 1;
414 }
415 content[start..].to_string()
416}
417
418impl CaptureService {
419 #[cfg(test)]
422 pub(crate) async fn set_buffer_state_for_test(
423 &self,
424 repo: &str,
425 task_id: &str,
426 seen_alive: bool,
427 has_output: bool,
428 registered_at: chrono::DateTime<Utc>,
429 ) {
430 let skey = crate::channels::transport::session_key(repo, task_id);
431 let mut buffers = self.buffers.write().await;
432 let buf = buffers
433 .get_mut(&skey)
434 .unwrap_or_else(|| panic!("no buffer for {repo}/{task_id}"));
435 buf.seen_alive = seen_alive;
436 buf.has_output = has_output;
437 buf.registered_at = registered_at;
438 }
439}
440
441#[cfg(test)]
442mod tests {
443 use super::*;
444 use crate::channels::transport::Transport;
445
446 fn make_buffer() -> OutputBuffer {
447 OutputBuffer {
448 repo: "owner/repo".to_string(),
449 session: "test-session".to_string(),
450 task_id: "task-1".to_string(),
451 last_content: String::new(),
452 last_len: 0,
453 last_hash: None,
454 last_capture: Utc::now(),
455 seen_alive: false,
456 registered_at: Utc::now(),
457 has_output: false,
458 generation: 0,
459 }
460 }
461
462 #[test]
463 fn normal_append_returns_new_suffix() {
464 let mut buf = make_buffer();
465 let result = buf.diff_and_update("hello");
467 assert_eq!(result.as_deref(), Some("hello"));
468 let result = buf.diff_and_update("hello world");
470 assert_eq!(result.as_deref(), Some(" world"));
471 }
472
473 #[test]
474 fn no_change_returns_none() {
475 let mut buf = make_buffer();
476 buf.diff_and_update("hello");
477 let result = buf.diff_and_update("hello");
479 assert_eq!(result, None);
480 }
481
482 #[test]
483 fn terminal_clear_shrink_returns_none_not_full_content() {
484 let mut buf = make_buffer();
485 buf.diff_and_update("line1\nline2\nline3\n");
487 let result = buf.diff_and_update("prompt$ ");
489 assert_eq!(
491 result, None,
492 "shrinking pane content should produce no output, not a duplicate of visible screen"
493 );
494 }
495
496 #[test]
497 fn terminal_clear_to_empty_returns_none() {
498 let mut buf = make_buffer();
499 buf.diff_and_update("lots of output here\nmore output\n");
500 let result = buf.diff_and_update("");
502 assert_eq!(result, None);
503 }
504
505 #[test]
506 fn same_length_content_change_is_broadcast() {
507 let mut buf = make_buffer();
508 buf.diff_and_update("aaaa");
509
510 let result = buf.diff_and_update("bbbb");
511
512 assert_eq!(result.as_deref(), Some("bbbb"));
513 }
514
515 #[test]
516 fn same_length_partial_change_broadcasts_only_changed_suffix() {
517 let mut buf = make_buffer();
518 let prior = format!("{}/", "line\n".repeat(50));
520 buf.diff_and_update(&prior);
521
522 let updated = format!("{}-", "line\n".repeat(50));
524 assert_eq!(
525 prior.len(),
526 updated.len(),
527 "lengths must match for this test"
528 );
529
530 let result = buf.diff_and_update(&updated);
531
532 assert_eq!(
534 result.as_deref(),
535 Some("-"),
536 "should return only the changed suffix, not the full content"
537 );
538 }
539
540 #[test]
541 fn empty_initial_capture_returns_none() {
542 let mut buf = make_buffer();
543 let result = buf.diff_and_update("");
544 assert_eq!(result, None);
545 }
546
547 #[test]
548 fn has_output_set_on_meaningful_content() {
549 let mut buf = make_buffer();
550 assert!(!buf.has_output);
551 buf.diff_and_update("hello");
552 assert!(buf.has_output);
553 }
554
555 #[test]
556 fn has_output_not_set_on_empty_content() {
557 let mut buf = make_buffer();
558 buf.diff_and_update("");
559 assert!(!buf.has_output, "empty content should not set has_output");
560 }
561
562 #[test]
563 fn has_output_stays_true_after_silence() {
564 let mut buf = make_buffer();
565 buf.diff_and_update("hello");
566 assert!(buf.has_output);
567 buf.diff_and_update("hello");
569 assert!(buf.has_output, "has_output should not revert once set");
570 }
571
572 #[tokio::test]
573 async fn get_silent_sessions_returns_only_no_output_past_grace() {
574 let transport = Arc::new(Transport::new());
575 let svc = CaptureService::new(transport);
576 let repo = "owner/repo";
577
578 svc.register_session(repo, "silent-task", "orch-test-silent")
581 .await;
582 svc.set_buffer_state_for_test(
583 repo,
584 "silent-task",
585 false,
586 false,
587 Utc::now() - chrono::Duration::seconds(200),
588 )
589 .await;
590
591 svc.register_session(repo, "active-task", "orch-test-active")
593 .await;
594 svc.set_buffer_state_for_test(
595 repo,
596 "active-task",
597 false,
598 true, Utc::now() - chrono::Duration::seconds(200),
600 )
601 .await;
602
603 svc.register_session(repo, "new-task", "orch-test-new")
605 .await;
606
607 let grace = std::time::Duration::from_secs(120);
608 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
609 assert_eq!(silent.len(), 1);
610 assert_eq!(silent[0].0, "silent-task");
611 }
612
613 #[tokio::test]
630 async fn get_silent_sessions_seen_alive_dead_session_is_silent() {
631 let transport = Arc::new(Transport::new());
632 let svc = CaptureService::new(transport);
633 let repo = "owner/repo";
634
635 svc.register_session(repo, "never-seen-task", "orch-never-seen")
637 .await;
638 svc.set_buffer_state_for_test(
639 repo,
640 "never-seen-task",
641 false, false, Utc::now() - chrono::Duration::seconds(500),
644 )
645 .await;
646
647 svc.register_session(repo, "seen-but-dead-task", "orch-seen-dead")
653 .await;
654 svc.set_buffer_state_for_test(
655 repo,
656 "seen-but-dead-task",
657 true, false, Utc::now() - chrono::Duration::seconds(500),
660 )
661 .await;
662
663 let grace = std::time::Duration::from_secs(120);
664 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
665 assert_eq!(
668 silent.len(),
669 2,
670 "both never-alive and dead-with-no-output sessions must be silenced"
671 );
672 let ids: std::collections::HashSet<&str> =
673 silent.iter().map(|(id, _, _)| id.as_str()).collect();
674 assert!(ids.contains("never-seen-task"));
675 assert!(ids.contains("seen-but-dead-task"));
676 }
677
678 #[tokio::test]
679 async fn get_silent_sessions_filters_by_repo() {
680 let transport = Arc::new(Transport::new());
681 let svc = CaptureService::new(transport);
682
683 svc.register_session("owner/repo-a", "task-a", "orch-a")
684 .await;
685 svc.register_session("owner/repo-b", "task-b", "orch-b")
686 .await;
687 svc.set_buffer_state_for_test(
688 "owner/repo-a",
689 "task-a",
690 false,
691 false,
692 Utc::now() - chrono::Duration::seconds(200),
693 )
694 .await;
695 svc.set_buffer_state_for_test(
696 "owner/repo-b",
697 "task-b",
698 false,
699 false,
700 Utc::now() - chrono::Duration::seconds(200),
701 )
702 .await;
703
704 let grace = std::time::Duration::from_secs(120);
705 let silent_a = svc
706 .get_silent_sessions_for_repo("owner/repo-a", grace)
707 .await;
708 assert_eq!(silent_a.len(), 1);
709 assert_eq!(silent_a[0].0, "task-a");
710
711 let silent_b = svc
712 .get_silent_sessions_for_repo("owner/repo-b", grace)
713 .await;
714 assert_eq!(silent_b.len(), 1);
715 assert_eq!(silent_b[0].0, "task-b");
716 }
717
718 #[tokio::test]
720 async fn same_task_id_different_repos_capture_buffers_isolated() {
721 let transport = Arc::new(Transport::new());
722 let svc = CaptureService::new(transport);
723
724 svc.register_session("owner/repo-a", "42", "orch-a-42")
725 .await;
726 svc.register_session("owner/repo-b", "42", "orch-b-42")
727 .await;
728
729 let skey_a = session_key("owner/repo-a", "42");
730 let skey_b = session_key("owner/repo-b", "42");
731
732 let buffers = svc.buffers.read().await;
733 let buf_a = buffers
734 .get(&skey_a)
735 .expect("buffer should exist for repo-a session");
736 let buf_b = buffers
737 .get(&skey_b)
738 .expect("buffer should exist for repo-b session");
739
740 assert_eq!(buf_a.session, "orch-a-42");
741 assert_eq!(buf_b.session, "orch-b-42");
742 assert_ne!(buf_a.session, buf_b.session);
743 assert_eq!(buf_a.repo, "owner/repo-a");
744 assert_eq!(buf_b.repo, "owner/repo-b");
745 }
746
747 #[test]
748 fn multibyte_boundary_no_bytes_dropped() {
749 let mut buf = make_buffer();
752 buf.diff_and_update("Hello ");
753 let full = "Hello 世界";
755 let result = buf.diff_and_update(full);
756 let s = result.expect("should return new content");
758 assert!(s.contains('世'), "first CJK char must not be dropped");
759 assert!(s.contains('界'), "second CJK char must not be dropped");
760 }
761
762 #[test]
763 fn whitespace_only_new_content_returns_none() {
764 let mut buf = make_buffer();
765 buf.diff_and_update("hello");
766 let result = buf.diff_and_update("hello \n");
768 assert_eq!(result, None);
769 }
770
771 #[tokio::test]
779 async fn get_silent_sessions_skips_stale_generation() {
780 let transport = Arc::new(Transport::new());
781 let svc = CaptureService::new(transport);
782 let repo = "owner/repo";
783
784 svc.register_session(repo, "task-42", "orch-old-session")
786 .await;
787 svc.set_buffer_state_for_test(
788 repo,
789 "task-42",
790 false,
791 false,
792 Utc::now() - chrono::Duration::seconds(500),
793 )
794 .await;
795
796 svc.register_session(repo, "task-42", "orch-new-session")
798 .await;
799
800 let grace = std::time::Duration::from_secs(120);
801 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
802
803 assert!(
807 silent.is_empty(),
808 "re-registered session should not be silenced; stale entry should be skipped"
809 );
810 }
811}