1use crate::channels::OutputChunk;
13use crate::channels::tmux;
14use crate::channels::transport::{Transport, session_key};
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 std::fmt::Debug for CaptureService {
71 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72 f.debug_struct("CaptureService")
73 .field("interval", &self.interval)
74 .finish_non_exhaustive()
75 }
76}
77
78impl CaptureService {
79 pub fn new(transport: Arc<Transport>) -> Self {
81 Self {
82 buffers: Arc::new(RwLock::new(HashMap::new())),
83 transport,
84 interval: std::time::Duration::from_secs(2),
85 }
86 }
87
88 pub async fn register_session(&self, repo: &str, task_id: &str, session: &str) {
90 let now = Utc::now();
91 let skey = session_key(repo, task_id);
92
93 let generation = {
97 let mut buffers = self.buffers.write().await;
98 let generation = buffers.get(&skey).map_or(0, |b| b.generation + 1);
99 let buffer = OutputBuffer {
100 repo: repo.to_string(),
101 session: session.to_string(),
102 task_id: task_id.to_string(),
103 last_content: String::new(),
104 last_len: 0,
105 last_hash: None,
106 last_capture: now,
107 seen_alive: false,
108 registered_at: now,
109 has_output: false,
110 generation,
111 };
112 buffers.insert(skey, buffer);
113 generation
114 };
115 tracing::debug!(
116 repo,
117 task_id,
118 session,
119 generation,
120 "session registered for capture"
121 );
122 }
123
124 pub async fn unregister_session(&self, repo: &str, task_id: &str) {
126 let skey = session_key(repo, task_id);
127 if let Some(buffer) = self.buffers.write().await.remove(&skey) {
128 tracing::debug!(
129 repo = buffer.repo,
130 task_id = buffer.task_id,
131 session = buffer.session,
132 "session unregistered"
133 );
134 }
135 self.transport.unbind(repo, task_id).await;
136 }
137
138 pub async fn start(&self) {
142 tracing::info!(
143 interval_secs = self.interval.as_secs(),
144 "capture service started"
145 );
146 let mut interval = tokio::time::interval(self.interval);
147
148 loop {
149 interval.tick().await;
150 self.tick().await;
151 }
152 }
153
154 pub async fn run(&self) {
159 let mut interval = tokio::time::interval(self.interval);
160
161 loop {
162 interval.tick().await;
163
164 let has_sessions = !self.buffers.read().await.is_empty();
166 if !has_sessions {
167 tracing::debug!("no sessions registered, capture loop exiting");
168 break;
169 }
170
171 self.tick().await;
172 }
173 }
174
175 pub async fn get_silent_sessions_for_repo(
191 &self,
192 repo: &str,
193 grace_period: std::time::Duration,
194 ) -> Vec<(String, String, i64)> {
195 let now = Utc::now();
196 let buffers = self.buffers.read().await;
197 let mut candidates: Vec<OutputBuffer> = buffers
198 .values()
199 .filter(|buf| buf.repo == repo && !buf.has_output)
200 .cloned()
201 .collect();
202 drop(buffers);
203
204 let mut silent = Vec::new();
205 for buf in candidates.drain(..) {
206 if buf.seen_alive {
207 if !tmux::is_session_dead(&buf.session).await {
214 continue;
215 }
216 tracing::debug!(
217 task_id = buf.task_id,
218 session = buf.session,
219 "session seen-alive but now dead with no output, treating as silent exit"
220 );
221 }
222
223 let skey = session_key(&buf.repo, &buf.task_id);
227 let current_generation = {
228 let buffers = self.buffers.read().await;
229 buffers.get(&skey).map(|b| b.generation)
230 };
231 if current_generation != Some(buf.generation) {
232 tracing::debug!(
233 task_id = buf.task_id,
234 session = buf.session,
235 old_gen = buf.generation,
236 new_gen = current_generation,
237 "session generation changed, skipping silence detection (was re-registered)"
238 );
239 continue;
240 }
241
242 let raw_age = now.signed_duration_since(buf.registered_at);
243 let suspended = crate::engine::suspend::suspended_duration_since(buf.registered_at);
249 let age = (raw_age - suspended).max(chrono::Duration::zero());
250 if age.num_seconds() > i64::try_from(grace_period.as_secs()).unwrap_or(i64::MAX) {
251 silent.push((buf.task_id.clone(), buf.session.clone(), age.num_seconds()));
252 }
253 }
254 silent
255 }
256
257 async fn tick(&self) {
264 let buffers = self.buffers.read().await;
265 let sessions: Vec<(String, OutputBuffer)> = buffers
266 .iter()
267 .map(|(k, v)| (k.clone(), v.clone()))
268 .collect();
269 drop(buffers);
270
271 for (skey, buffer) in sessions {
272 let current_content = match tmux::capture_pane(&buffer.session).await {
274 Ok(s) => s,
275 Err(e) => {
276 let current = {
282 let buffers = self.buffers.read().await;
283 buffers.get(&skey).cloned()
284 };
285
286 let should_finalize = current.as_ref().is_some_and(|c| {
292 c.generation == buffer.generation
293 && c.seen_alive
294 && c.session == buffer.session
295 });
296
297 if should_finalize && tmux::is_session_dead(&buffer.session).await {
298 tracing::info!(
299 task_id = buffer.task_id,
300 session = buffer.session,
301 generation = buffer.generation,
302 "session ended, sending final chunk"
303 );
304 let chunk = OutputChunk {
305 content: String::new(),
306 is_final: true,
307 };
308 self.transport
309 .push_output(&buffer.repo, &buffer.task_id, chunk)
310 .await;
311 self.unregister_session(&buffer.repo, &buffer.task_id).await;
312 } else if current.is_none() {
313 tracing::trace!(
314 task_id = buffer.task_id,
315 session = buffer.session,
316 "session already unregistered"
317 );
318 } else {
319 tracing::trace!(
320 task_id = buffer.task_id,
321 session = buffer.session,
322 ?e,
323 "capture failed (transient)"
324 );
325 }
326 continue;
327 }
328 };
329
330 let new_content = {
331 let mut buffers = self.buffers.write().await;
332 if let Some(buf) = buffers.get_mut(&skey) {
333 buf.seen_alive = true;
334 buf.diff_and_update(¤t_content)
335 } else {
336 None
337 }
338 };
339
340 if let Some(new_content) = new_content {
341 let chunk = OutputChunk {
342 content: new_content,
343 is_final: false,
344 };
345 self.transport
346 .push_output(&buffer.repo, &buffer.task_id, chunk)
347 .await;
348 }
349 }
350 }
351}
352
353impl OutputBuffer {
354 pub(crate) fn diff_and_update(&mut self, current_content: &str) -> Option<String> {
355 let current_bytes = current_content.as_bytes();
356 let current_len = current_bytes.len();
357 let current_hash = hash_bytes(current_bytes);
358
359 if self.last_len == current_len && self.last_hash == Some(current_hash) {
360 self.last_capture = Utc::now();
361 return None;
362 }
363
364 let new_content = match self.last_len.cmp(¤t_len) {
365 std::cmp::Ordering::Greater => String::new(),
368 std::cmp::Ordering::Equal => {
369 let common_prefix = self
372 .last_content
373 .as_bytes()
374 .iter()
375 .zip(current_content.as_bytes())
376 .take_while(|(a, b)| a == b)
377 .count();
378 let mut start = common_prefix;
379 while start < current_len && !current_content.is_char_boundary(start) {
380 start += 1;
381 }
382 current_content[start..].to_string()
383 }
384 std::cmp::Ordering::Less => {
385 let mut offset = self.last_len;
386 while offset < current_len && !current_content.is_char_boundary(offset) {
387 offset += 1;
388 }
389 String::from_utf8_lossy(¤t_bytes[offset..]).to_string()
390 }
391 };
392
393 self.last_len = current_len;
394 self.last_hash = Some(current_hash);
395 self.last_content = cap_content(current_content);
396 self.last_capture = Utc::now();
397
398 if new_content.trim().is_empty() {
399 None
400 } else {
401 self.has_output = true;
402 Some(new_content)
403 }
404 }
405}
406
407fn hash_bytes(bytes: &[u8]) -> [u8; 32] {
408 let mut hasher = Sha256::new();
409 hasher.update(bytes);
410 hasher.finalize().into()
411}
412
413fn cap_content(content: &str) -> String {
414 if content.len() <= MAX_OUTPUT_BUFFER_BYTES {
415 return content.to_string();
416 }
417
418 let mut start = content.len() - MAX_OUTPUT_BUFFER_BYTES;
419 while start < content.len() && !content.is_char_boundary(start) {
420 start += 1;
421 }
422 content[start..].to_string()
423}
424
425impl CaptureService {
426 #[cfg(test)]
429 pub(crate) async fn set_buffer_state_for_test(
430 &self,
431 repo: &str,
432 task_id: &str,
433 seen_alive: bool,
434 has_output: bool,
435 registered_at: chrono::DateTime<Utc>,
436 ) {
437 let skey = crate::channels::transport::session_key(repo, task_id);
438 let mut buffers = self.buffers.write().await;
439 let buf = buffers
440 .get_mut(&skey)
441 .unwrap_or_else(|| panic!("no buffer for {repo}/{task_id}"));
442 buf.seen_alive = seen_alive;
443 buf.has_output = has_output;
444 buf.registered_at = registered_at;
445 }
446}
447
448#[cfg(test)]
449mod tests {
450 use super::*;
451 use crate::channels::transport::Transport;
452
453 fn make_buffer() -> OutputBuffer {
454 OutputBuffer {
455 repo: "owner/repo".to_string(),
456 session: "test-session".to_string(),
457 task_id: "task-1".to_string(),
458 last_content: String::new(),
459 last_len: 0,
460 last_hash: None,
461 last_capture: Utc::now(),
462 seen_alive: false,
463 registered_at: Utc::now(),
464 has_output: false,
465 generation: 0,
466 }
467 }
468
469 #[test]
470 fn normal_append_returns_new_suffix() {
471 let mut buf = make_buffer();
472 let result = buf.diff_and_update("hello");
474 assert_eq!(result.as_deref(), Some("hello"));
475 let result = buf.diff_and_update("hello world");
477 assert_eq!(result.as_deref(), Some(" world"));
478 }
479
480 #[test]
481 fn no_change_returns_none() {
482 let mut buf = make_buffer();
483 buf.diff_and_update("hello");
484 let result = buf.diff_and_update("hello");
486 assert_eq!(result, None);
487 }
488
489 #[test]
490 fn terminal_clear_shrink_returns_none_not_full_content() {
491 let mut buf = make_buffer();
492 buf.diff_and_update("line1\nline2\nline3\n");
494 let result = buf.diff_and_update("prompt$ ");
496 assert_eq!(
498 result, None,
499 "shrinking pane content should produce no output, not a duplicate of visible screen"
500 );
501 }
502
503 #[test]
504 fn terminal_clear_to_empty_returns_none() {
505 let mut buf = make_buffer();
506 buf.diff_and_update("lots of output here\nmore output\n");
507 let result = buf.diff_and_update("");
509 assert_eq!(result, None);
510 }
511
512 #[test]
513 fn same_length_content_change_is_broadcast() {
514 let mut buf = make_buffer();
515 buf.diff_and_update("aaaa");
516
517 let result = buf.diff_and_update("bbbb");
518
519 assert_eq!(result.as_deref(), Some("bbbb"));
520 }
521
522 #[test]
523 fn same_length_partial_change_broadcasts_only_changed_suffix() {
524 let mut buf = make_buffer();
525 let prior = format!("{}/", "line\n".repeat(50));
527 buf.diff_and_update(&prior);
528
529 let updated = format!("{}-", "line\n".repeat(50));
531 assert_eq!(
532 prior.len(),
533 updated.len(),
534 "lengths must match for this test"
535 );
536
537 let result = buf.diff_and_update(&updated);
538
539 assert_eq!(
541 result.as_deref(),
542 Some("-"),
543 "should return only the changed suffix, not the full content"
544 );
545 }
546
547 #[test]
548 fn empty_initial_capture_returns_none() {
549 let mut buf = make_buffer();
550 let result = buf.diff_and_update("");
551 assert_eq!(result, None);
552 }
553
554 #[test]
555 fn has_output_set_on_meaningful_content() {
556 let mut buf = make_buffer();
557 assert!(!buf.has_output);
558 buf.diff_and_update("hello");
559 assert!(buf.has_output);
560 }
561
562 #[test]
563 fn has_output_not_set_on_empty_content() {
564 let mut buf = make_buffer();
565 buf.diff_and_update("");
566 assert!(!buf.has_output, "empty content should not set has_output");
567 }
568
569 #[test]
570 fn has_output_stays_true_after_silence() {
571 let mut buf = make_buffer();
572 buf.diff_and_update("hello");
573 assert!(buf.has_output);
574 buf.diff_and_update("hello");
576 assert!(buf.has_output, "has_output should not revert once set");
577 }
578
579 #[tokio::test]
580 async fn get_silent_sessions_returns_only_no_output_past_grace() {
581 let transport = Arc::new(Transport::new());
582 let svc = CaptureService::new(transport);
583 let repo = "owner/repo";
584
585 svc.register_session(repo, "silent-task", "orch-test-silent")
588 .await;
589 svc.set_buffer_state_for_test(
590 repo,
591 "silent-task",
592 false,
593 false,
594 Utc::now() - chrono::Duration::seconds(200),
595 )
596 .await;
597
598 svc.register_session(repo, "active-task", "orch-test-active")
600 .await;
601 svc.set_buffer_state_for_test(
602 repo,
603 "active-task",
604 false,
605 true, Utc::now() - chrono::Duration::seconds(200),
607 )
608 .await;
609
610 svc.register_session(repo, "new-task", "orch-test-new")
612 .await;
613
614 let grace = std::time::Duration::from_secs(120);
615 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
616 assert_eq!(silent.len(), 1);
617 assert_eq!(silent[0].0, "silent-task");
618 }
619
620 #[tokio::test]
637 async fn get_silent_sessions_seen_alive_dead_session_is_silent() {
638 let transport = Arc::new(Transport::new());
639 let svc = CaptureService::new(transport);
640 let repo = "owner/repo";
641
642 svc.register_session(repo, "never-seen-task", "orch-never-seen")
644 .await;
645 svc.set_buffer_state_for_test(
646 repo,
647 "never-seen-task",
648 false, false, Utc::now() - chrono::Duration::seconds(500),
651 )
652 .await;
653
654 svc.register_session(repo, "seen-but-dead-task", "orch-seen-dead")
660 .await;
661 svc.set_buffer_state_for_test(
662 repo,
663 "seen-but-dead-task",
664 true, false, Utc::now() - chrono::Duration::seconds(500),
667 )
668 .await;
669
670 let grace = std::time::Duration::from_secs(120);
671 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
672 assert_eq!(
675 silent.len(),
676 2,
677 "both never-alive and dead-with-no-output sessions must be silenced"
678 );
679 let ids: std::collections::HashSet<&str> =
680 silent.iter().map(|(id, _, _)| id.as_str()).collect();
681 assert!(ids.contains("never-seen-task"));
682 assert!(ids.contains("seen-but-dead-task"));
683 }
684
685 #[tokio::test]
686 async fn get_silent_sessions_filters_by_repo() {
687 let transport = Arc::new(Transport::new());
688 let svc = CaptureService::new(transport);
689
690 svc.register_session("owner/repo-a", "task-a", "orch-a")
691 .await;
692 svc.register_session("owner/repo-b", "task-b", "orch-b")
693 .await;
694 svc.set_buffer_state_for_test(
695 "owner/repo-a",
696 "task-a",
697 false,
698 false,
699 Utc::now() - chrono::Duration::seconds(200),
700 )
701 .await;
702 svc.set_buffer_state_for_test(
703 "owner/repo-b",
704 "task-b",
705 false,
706 false,
707 Utc::now() - chrono::Duration::seconds(200),
708 )
709 .await;
710
711 let grace = std::time::Duration::from_secs(120);
712 let silent_a = svc
713 .get_silent_sessions_for_repo("owner/repo-a", grace)
714 .await;
715 assert_eq!(silent_a.len(), 1);
716 assert_eq!(silent_a[0].0, "task-a");
717
718 let silent_b = svc
719 .get_silent_sessions_for_repo("owner/repo-b", grace)
720 .await;
721 assert_eq!(silent_b.len(), 1);
722 assert_eq!(silent_b[0].0, "task-b");
723 }
724
725 #[tokio::test]
727 async fn same_task_id_different_repos_capture_buffers_isolated() {
728 let transport = Arc::new(Transport::new());
729 let svc = CaptureService::new(transport);
730
731 svc.register_session("owner/repo-a", "42", "orch-a-42")
732 .await;
733 svc.register_session("owner/repo-b", "42", "orch-b-42")
734 .await;
735
736 let skey_a = session_key("owner/repo-a", "42");
737 let skey_b = session_key("owner/repo-b", "42");
738
739 let buffers = svc.buffers.read().await;
740 let buf_a = buffers
741 .get(&skey_a)
742 .expect("buffer should exist for repo-a session");
743 let buf_b = buffers
744 .get(&skey_b)
745 .expect("buffer should exist for repo-b session");
746
747 assert_eq!(buf_a.session, "orch-a-42");
748 assert_eq!(buf_b.session, "orch-b-42");
749 assert_ne!(buf_a.session, buf_b.session);
750 assert_eq!(buf_a.repo, "owner/repo-a");
751 assert_eq!(buf_b.repo, "owner/repo-b");
752 }
753
754 #[test]
755 fn multibyte_boundary_no_bytes_dropped() {
756 let mut buf = make_buffer();
759 buf.diff_and_update("Hello ");
760 let full = "Hello 世界";
762 let result = buf.diff_and_update(full);
763 let s = result.expect("should return new content");
765 assert!(s.contains('世'), "first CJK char must not be dropped");
766 assert!(s.contains('界'), "second CJK char must not be dropped");
767 }
768
769 #[test]
770 fn whitespace_only_new_content_returns_none() {
771 let mut buf = make_buffer();
772 buf.diff_and_update("hello");
773 let result = buf.diff_and_update("hello \n");
775 assert_eq!(result, None);
776 }
777
778 #[tokio::test]
786 async fn get_silent_sessions_skips_stale_generation() {
787 let transport = Arc::new(Transport::new());
788 let svc = CaptureService::new(transport);
789 let repo = "owner/repo";
790
791 svc.register_session(repo, "task-42", "orch-old-session")
793 .await;
794 svc.set_buffer_state_for_test(
795 repo,
796 "task-42",
797 false,
798 false,
799 Utc::now() - chrono::Duration::seconds(500),
800 )
801 .await;
802
803 svc.register_session(repo, "task-42", "orch-new-session")
805 .await;
806
807 let grace = std::time::Duration::from_secs(120);
808 let silent = svc.get_silent_sessions_for_repo(repo, grace).await;
809
810 assert!(
814 silent.is_empty(),
815 "re-registered session should not be silenced; stale entry should be skipped"
816 );
817 }
818}