Skip to main content

agent_tui_infra/infra/daemon/
repository.rs

1//! Session repository implementation.
2
3use crossbeam_channel as channel;
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::sync::Mutex;
7use std::time::Duration;
8
9use crate::common::mutex_lock_or_recover;
10use crate::domain::RestartOutput;
11use crate::domain::core::CursorPosition;
12use crate::domain::session_types::TerminalSize;
13use crate::usecases::ports::LivePreviewSnapshot;
14use crate::usecases::ports::SessionError;
15use crate::usecases::ports::SessionHandle;
16use crate::usecases::ports::SessionOps;
17use crate::usecases::ports::SessionRepository;
18use crate::usecases::ports::StreamCursor;
19use crate::usecases::ports::StreamRead;
20use crate::usecases::ports::StreamWaiterHandle;
21use crate::usecases::ports::TerminalError;
22
23use crate::infra::daemon::session::PUMP_FLUSH_TIMEOUT;
24use crate::infra::daemon::session::Session;
25use crate::infra::daemon::session::SessionId;
26use crate::infra::daemon::session::SessionInfo;
27use crate::infra::daemon::session::SessionManager;
28use crate::infra::daemon::session::StreamReader;
29
30struct SessionHandleImpl {
31    inner: Arc<Mutex<Session>>,
32    stream: StreamReader,
33    pty_cursor: Arc<Mutex<StreamCursor>>,
34}
35
36impl SessionHandleImpl {
37    fn new_handle(inner: Arc<Mutex<Session>>) -> SessionHandle {
38        let (stream, pty_cursor) = {
39            let session_guard = mutex_lock_or_recover(&inner);
40            (
41                session_guard.stream_reader(),
42                session_guard.pty_cursor_handle(),
43            )
44        };
45        Arc::new(Self {
46            inner,
47            stream,
48            pty_cursor,
49        })
50    }
51}
52
53fn wait_for_flush_ack(
54    ack: Option<channel::Receiver<()>>,
55    timeout: Duration,
56) -> Result<(), SessionError> {
57    let Some(ack) = ack else {
58        return Ok(());
59    };
60
61    ack.recv_timeout(timeout).map_err(|err| match err {
62        channel::RecvTimeoutError::Timeout => SessionError::Terminal(TerminalError::Read {
63            reason: "Timed out waiting for session state to flush".to_string(),
64            source: None,
65        }),
66        channel::RecvTimeoutError::Disconnected => SessionError::Terminal(TerminalError::Read {
67            reason: "Session state flush closed before it acknowledged".to_string(),
68            source: None,
69        }),
70    })
71}
72
73impl SessionOps for SessionHandleImpl {
74    fn update(&self) -> Result<(), SessionError> {
75        let ack = {
76            let session_guard = mutex_lock_or_recover(&self.inner);
77            session_guard.request_flush()
78        };
79        wait_for_flush_ack(ack, PUMP_FLUSH_TIMEOUT)
80    }
81
82    fn screen_text(&self) -> String {
83        let session_guard = mutex_lock_or_recover(&self.inner);
84        session_guard.screen_text()
85    }
86
87    fn screen_render(&self) -> String {
88        let session_guard = mutex_lock_or_recover(&self.inner);
89        session_guard.screen_render()
90    }
91
92    fn screen_render_compact(&self) -> String {
93        let session_guard = mutex_lock_or_recover(&self.inner);
94        session_guard.screen_render_compact()
95    }
96
97    fn terminal_write(&self, data: &[u8]) -> Result<(), SessionError> {
98        let mut session_guard = mutex_lock_or_recover(&self.inner);
99        session_guard.pty_write(data)
100    }
101
102    fn terminal_try_read(&self, buf: &mut [u8], timeout_ms: i32) -> Result<usize, SessionError> {
103        let mut cursor = self
104            .pty_cursor
105            .lock()
106            .unwrap_or_else(std::sync::PoisonError::into_inner);
107        let read = self.stream.read(&mut cursor, buf.len(), timeout_ms)?;
108        let bytes_read = read.data.len().min(buf.len());
109        buf[..bytes_read].copy_from_slice(&read.data[..bytes_read]);
110        Ok(bytes_read)
111    }
112
113    fn stream_read(
114        &self,
115        cursor: &mut StreamCursor,
116        max_bytes: usize,
117        timeout_ms: i32,
118    ) -> Result<StreamRead, SessionError> {
119        self.stream.read(cursor, max_bytes, timeout_ms)
120    }
121
122    fn stream_subscribe(&self) -> StreamWaiterHandle {
123        self.stream.subscribe()
124    }
125
126    fn keystroke(&self, key: &str) -> Result<(), SessionError> {
127        let mut session_guard = mutex_lock_or_recover(&self.inner);
128        session_guard.keystroke(key)
129    }
130
131    fn type_text(&self, text: &str) -> Result<(), SessionError> {
132        let mut session_guard = mutex_lock_or_recover(&self.inner);
133        session_guard.type_text(text)
134    }
135
136    fn keydown(&self, key: &str) -> Result<(), SessionError> {
137        let mut session_guard = mutex_lock_or_recover(&self.inner);
138        session_guard.keydown(key)
139    }
140
141    fn keyup(&self, key: &str) -> Result<(), SessionError> {
142        let mut session_guard = mutex_lock_or_recover(&self.inner);
143        session_guard.keyup(key)
144    }
145
146    fn is_running(&self) -> bool {
147        let mut session_guard = mutex_lock_or_recover(&self.inner);
148        session_guard.is_running()
149    }
150
151    fn resize(&self, size: TerminalSize) -> Result<(), SessionError> {
152        let mut session_guard = mutex_lock_or_recover(&self.inner);
153        session_guard.resize(size)
154    }
155
156    fn cursor(&self) -> CursorPosition {
157        let session_guard = mutex_lock_or_recover(&self.inner);
158        session_guard.cursor()
159    }
160
161    fn session_id(&self) -> SessionId {
162        let session_guard = mutex_lock_or_recover(&self.inner);
163        session_guard.id.clone()
164    }
165
166    fn command(&self) -> String {
167        let session_guard = mutex_lock_or_recover(&self.inner);
168        session_guard.command.clone()
169    }
170
171    fn size(&self) -> TerminalSize {
172        let session_guard = mutex_lock_or_recover(&self.inner);
173        session_guard.size()
174    }
175
176    fn live_preview_snapshot(&self) -> LivePreviewSnapshot {
177        let session_guard = mutex_lock_or_recover(&self.inner);
178        session_guard.live_preview_snapshot()
179    }
180}
181
182#[allow(clippy::too_many_arguments)]
183impl SessionRepository for SessionManager {
184    fn spawn(
185        &self,
186        command: &str,
187        args: &[String],
188        cwd: Option<&str>,
189        env: Option<&HashMap<String, String>>,
190        session_id: Option<SessionId>,
191        size: TerminalSize,
192    ) -> Result<(SessionId, u32), SessionError> {
193        SessionManager::spawn(self, command, args, cwd, env, session_id, size)
194    }
195
196    fn get(&self, session_id: &SessionId) -> Result<SessionHandle, SessionError> {
197        let session = SessionManager::get(self, session_id)?;
198        Ok(SessionHandleImpl::new_handle(session))
199    }
200
201    fn active(&self) -> Result<SessionHandle, SessionError> {
202        let session = SessionManager::active(self)?;
203        Ok(SessionHandleImpl::new_handle(session))
204    }
205
206    fn resolve(&self, session_id: Option<&SessionId>) -> Result<SessionHandle, SessionError> {
207        let session = SessionManager::resolve(self, session_id)?;
208        Ok(SessionHandleImpl::new_handle(session))
209    }
210
211    fn set_active(&self, session_id: &SessionId) -> Result<(), SessionError> {
212        SessionManager::set_active(self, session_id)
213    }
214
215    fn list(&self) -> Vec<SessionInfo> {
216        SessionManager::list(self)
217    }
218
219    fn kill(&self, session_id: &SessionId) -> Result<(), SessionError> {
220        SessionManager::kill(self, session_id)
221    }
222
223    fn restart(&self, session_id: Option<&SessionId>) -> Result<RestartOutput, SessionError> {
224        SessionManager::restart(self, session_id)
225    }
226
227    fn session_count(&self) -> usize {
228        SessionManager::session_count(self)
229    }
230
231    fn active_session_id(&self) -> Option<SessionId> {
232        SessionManager::active_session_id(self)
233    }
234}
235
236#[cfg(test)]
237#[path = "repository_tests.rs"]
238mod tests;