Skip to main content

reovim_server/session/
capture.rs

1//! TUI capture request tracking for CLI → Server → TUI → Server → CLI relay.
2//!
3//! Manages pending capture requests and correlates responses from TUI clients.
4//!
5//! # Flow
6//!
7//! 1. CLI sends `GetScreenContent` RPC request to server
8//! 2. Server creates a pending request with a unique ID and `oneshot::Sender`
9//! 3. Server sends `capture_request` notification to TUI client
10//! 4. TUI client processes request and sends `capture_response` notification
11//! 5. Server receives notification, finds pending request, sends result via channel
12//! 6. Original handler receives result and responds to CLI
13
14use std::{
15    collections::HashMap,
16    sync::atomic::{AtomicU64, Ordering},
17    time::Duration,
18};
19
20use {parking_lot::RwLock, tokio::sync::oneshot};
21
22/// Default timeout for capture requests.
23pub const CAPTURE_TIMEOUT_SECS: u64 = 5;
24
25/// Error type for capture operations.
26#[derive(Debug)]
27pub enum CaptureError {
28    /// No TUI client connected to handle the capture.
29    NoTuiClient,
30    /// Capture request timed out.
31    Timeout,
32    /// TUI client disconnected during capture.
33    Disconnected,
34    /// Invalid response from TUI.
35    InvalidResponse(String),
36}
37
38impl std::fmt::Display for CaptureError {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        match self {
41            Self::NoTuiClient => {
42                write!(f, "No TUI client connected. Start TUI first with `reovim tui`")
43            }
44            Self::Timeout => {
45                write!(f, "TUI capture timeout ({CAPTURE_TIMEOUT_SECS}s). TUI may be frozen")
46            }
47            Self::Disconnected => write!(f, "TUI client disconnected during capture"),
48            Self::InvalidResponse(msg) => write!(f, "Invalid capture response: {msg}"),
49        }
50    }
51}
52
53impl std::error::Error for CaptureError {}
54
55/// Captured screen content result.
56#[derive(Debug, Clone)]
57pub struct CaptureResult {
58    /// Frame width.
59    pub width: u64,
60    /// Frame height.
61    pub height: u64,
62    /// Format used for capture.
63    pub format: String,
64    /// Captured frame content.
65    pub content: String,
66}
67
68/// Tracks pending capture requests and correlates responses.
69///
70/// Thread-safe structure that allows concurrent request creation and response delivery.
71/// Uses atomic counters for unique request IDs and `RwLock` for the pending request map.
72pub struct CaptureTracker {
73    /// Next request ID.
74    next_id: AtomicU64,
75    /// Pending requests waiting for responses.
76    pending: RwLock<HashMap<u64, PendingCapture>>,
77}
78
79/// A pending capture request waiting for a response.
80struct PendingCapture {
81    /// Channel to deliver the result.
82    sender: oneshot::Sender<CaptureResult>,
83}
84
85impl CaptureTracker {
86    /// Create a new capture tracker.
87    #[must_use]
88    pub fn new() -> Self {
89        Self {
90            next_id: AtomicU64::new(1),
91            pending: RwLock::new(HashMap::new()),
92        }
93    }
94
95    /// Create a pending capture request.
96    ///
97    /// Returns the request ID and a receiver for the result.
98    pub fn create_pending(&self) -> (u64, oneshot::Receiver<CaptureResult>) {
99        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
100        let (tx, rx) = oneshot::channel();
101        self.pending
102            .write()
103            .insert(id, PendingCapture { sender: tx });
104        (id, rx)
105    }
106
107    /// Deliver a capture response.
108    ///
109    /// Returns `true` if the response was delivered, `false` if no matching request.
110    pub fn deliver_response(&self, request_id: u64, result: CaptureResult) -> bool {
111        // Extract pending request before processing to avoid holding lock
112        let pending = self.pending.write().remove(&request_id);
113        if let Some(pending) = pending {
114            // Ignore send error (receiver dropped = timeout/cancelled)
115            let _ = pending.sender.send(result);
116            true
117        } else {
118            tracing::warn!("Received capture response for unknown request {request_id}");
119            false
120        }
121    }
122
123    /// Cancel a pending request (e.g., on timeout).
124    ///
125    /// Returns `true` if the request was found and cancelled.
126    pub fn cancel(&self, request_id: u64) -> bool {
127        self.pending.write().remove(&request_id).is_some()
128    }
129
130    /// Get the number of pending requests.
131    #[must_use]
132    pub fn pending_count(&self) -> usize {
133        self.pending.read().len()
134    }
135}
136
137impl Default for CaptureTracker {
138    fn default() -> Self {
139        Self::new()
140    }
141}
142
143impl std::fmt::Debug for CaptureTracker {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        f.debug_struct("CaptureTracker")
146            .field("pending_count", &self.pending_count())
147            .finish()
148    }
149}
150
151/// Wait for a capture response with timeout.
152///
153/// # Arguments
154///
155/// * `rx` - The receiver for the capture result
156///
157/// # Errors
158///
159/// Returns `CaptureError::Timeout` if the timeout expires.
160/// Returns `CaptureError::Disconnected` if the sender was dropped.
161pub async fn wait_for_capture(
162    rx: oneshot::Receiver<CaptureResult>,
163) -> Result<CaptureResult, CaptureError> {
164    match tokio::time::timeout(Duration::from_secs(CAPTURE_TIMEOUT_SECS), rx).await {
165        Ok(Ok(result)) => Ok(result),
166        Ok(Err(_)) => Err(CaptureError::Disconnected),
167        Err(_) => Err(CaptureError::Timeout),
168    }
169}
170
171#[cfg(test)]
172#[path = "capture_tests.rs"]
173mod tests;