Skip to main content

agent_tui_infra/infra/ipc/
client.rs

1#![expect(
2    clippy::print_stderr,
3    reason = "CLI status messages during daemon autostart"
4)]
5
6//! IPC client implementation.
7
8use std::path::PathBuf;
9use std::sync::Arc;
10use std::sync::atomic::AtomicBool;
11use std::sync::atomic::AtomicU64;
12use std::sync::atomic::Ordering;
13use std::time::Duration;
14use std::time::Instant;
15
16use serde::Deserialize;
17use serde::Serialize;
18use serde_json::Value;
19use tracing::debug;
20use tracing::trace;
21
22use crate::common::Colors;
23use crate::common::RpcId;
24use crate::common::error_codes;
25use crate::infra::ipc::error::ClientError;
26use crate::infra::ipc::process::ProcessIdentity;
27use crate::infra::ipc::socket::socket_path;
28use crate::infra::ipc::transport::ClientConnection;
29use crate::infra::ipc::transport::IpcTransport;
30use crate::infra::ipc::transport::UnixSocketTransport;
31use crate::infra::ipc::transport::default_transport;
32
33static REQUEST_ID: AtomicU64 = AtomicU64::new(1);
34const STREAM_POLL_TIMEOUT: Duration = Duration::from_millis(250);
35
36#[derive(Debug, Clone)]
37pub struct DaemonClientConfig {
38    read_timeout: Duration,
39    write_timeout: Duration,
40    max_retries: u32,
41    initial_retry_delay: Duration,
42}
43
44impl Default for DaemonClientConfig {
45    fn default() -> Self {
46        Self {
47            read_timeout: Duration::from_secs(60),
48            write_timeout: Duration::from_secs(10),
49            max_retries: 3,
50            initial_retry_delay: Duration::from_millis(100),
51        }
52    }
53}
54
55impl DaemonClientConfig {
56    pub fn read_timeout(&self) -> Duration {
57        self.read_timeout
58    }
59
60    pub fn write_timeout(&self) -> Duration {
61        self.write_timeout
62    }
63
64    pub fn max_retries(&self) -> u32 {
65        self.max_retries
66    }
67
68    pub fn initial_retry_delay(&self) -> Duration {
69        self.initial_retry_delay
70    }
71
72    pub fn with_read_timeout(mut self, timeout: Duration) -> Self {
73        self.read_timeout = timeout;
74        self
75    }
76
77    pub fn with_write_timeout(mut self, timeout: Duration) -> Self {
78        self.write_timeout = timeout;
79        self
80    }
81
82    pub fn with_max_retries(mut self, retries: u32) -> Self {
83        self.max_retries = retries;
84        self
85    }
86
87    pub fn with_initial_retry_delay(mut self, delay: Duration) -> Self {
88        self.initial_retry_delay = delay;
89        self
90    }
91}
92
93#[derive(Debug, Serialize)]
94struct Request {
95    jsonrpc: String,
96    id: RpcId,
97    method: String,
98    #[serde(skip_serializing_if = "Option::is_none")]
99    params: Option<Value>,
100}
101
102#[derive(Debug, Deserialize)]
103struct Response {
104    #[serde(rename = "jsonrpc")]
105    _jsonrpc: String,
106    #[serde(rename = "id")]
107    id: Option<RpcId>,
108    result: Option<Value>,
109    error: Option<RpcError>,
110}
111
112#[derive(Debug, Deserialize)]
113struct RpcError {
114    code: i32,
115    message: String,
116    #[serde(default)]
117    data: Option<Value>,
118}
119
120pub trait DaemonClient: Send + Sync {
121    fn call(&mut self, method: &str, params: Option<Value>) -> Result<Value, ClientError>;
122
123    fn call_with_config(
124        &mut self,
125        method: &str,
126        params: Option<Value>,
127        config: &DaemonClientConfig,
128    ) -> Result<Value, ClientError>;
129
130    fn call_stream(
131        &mut self,
132        method: &str,
133        params: Option<Value>,
134    ) -> Result<StreamResponse, ClientError> {
135        self.call_stream_with_config(method, params, &DaemonClientConfig::default())
136    }
137
138    fn call_stream_with_config(
139        &mut self,
140        _method: &str,
141        _params: Option<Value>,
142        _config: &DaemonClientConfig,
143    ) -> Result<StreamResponse, ClientError> {
144        Err(ClientError::UnexpectedResponse {
145            message: "Streaming RPC not supported by this client".to_string(),
146        })
147    }
148}
149
150pub struct UnixSocketClient {
151    transport: std::sync::Arc<dyn IpcTransport>,
152}
153
154impl UnixSocketClient {
155    pub fn connect() -> Result<Self, ClientError> {
156        Self::connect_with_transport(default_transport())
157    }
158
159    pub fn connect_local() -> Result<Self, ClientError> {
160        Self::connect_with_transport(std::sync::Arc::new(UnixSocketTransport))
161    }
162
163    pub(crate) fn connect_with_transport(
164        transport: std::sync::Arc<dyn IpcTransport>,
165    ) -> Result<Self, ClientError> {
166        let connection = transport.connect_connection()?;
167        drop(connection);
168
169        Ok(Self { transport })
170    }
171
172    pub fn is_daemon_running() -> bool {
173        default_transport().is_daemon_running()
174    }
175}
176
177pub struct StreamAbortHandle {
178    aborted: Arc<AtomicBool>,
179}
180
181impl StreamAbortHandle {
182    pub fn abort(&self) {
183        self.aborted.store(true, Ordering::Relaxed);
184    }
185}
186
187pub struct StreamResponse {
188    connection: ClientConnection,
189    expected_id: RpcId,
190    aborted: Arc<AtomicBool>,
191}
192
193impl StreamResponse {
194    pub fn next_result(&mut self) -> Result<Option<Value>, ClientError> {
195        loop {
196            if self.aborted.load(Ordering::Relaxed) {
197                let _ = self.connection.shutdown();
198                return Ok(None);
199            }
200
201            let response_line = match self.connection.read_message() {
202                Ok(Some(line)) => line,
203                Ok(None) => return Ok(None),
204                Err(err) if is_timeout_error(&err) => continue,
205                Err(err) => return Err(err),
206            };
207
208            let response: Response = serde_json::from_str(&response_line)?;
209            return response_to_result(response, &self.expected_id).map(Some);
210        }
211    }
212
213    pub fn abort_handle(&self) -> Option<StreamAbortHandle> {
214        Some(StreamAbortHandle {
215            aborted: Arc::clone(&self.aborted),
216        })
217    }
218}
219
220impl Drop for StreamResponse {
221    fn drop(&mut self) {
222        let _ = self.connection.shutdown();
223    }
224}
225
226fn is_timeout_error(error: &ClientError) -> bool {
227    match error {
228        ClientError::ConnectionFailed(io_err) => matches!(
229            io_err.kind(),
230            std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock
231        ),
232        _ => false,
233    }
234}
235
236fn is_retryable_call_error(error: &ClientError) -> bool {
237    match error {
238        ClientError::ConnectionFailed(_) => true,
239        ClientError::RpcError { retryable, .. } => *retryable,
240        _ => false,
241    }
242}
243
244fn retry_delay_ms_from_data(data: Option<&Value>) -> Option<u64> {
245    let data = data?;
246    ["retry_delay_ms", "retry_after_ms"]
247        .into_iter()
248        .find_map(|key| data.get(key).and_then(Value::as_u64))
249}
250
251fn next_retry_delay(
252    error: &ClientError,
253    default_delay: Duration,
254    current_delay: Duration,
255) -> Duration {
256    error
257        .retry_delay_ms()
258        .map(Duration::from_millis)
259        .unwrap_or_else(|| {
260            if current_delay.is_zero() {
261                default_delay
262            } else {
263                current_delay
264            }
265        })
266}
267
268fn validate_response_id(response: &Response, expected_id: &RpcId) -> Result<(), ClientError> {
269    match &response.id {
270        Some(actual_id) if actual_id == expected_id => Ok(()),
271        Some(_) => Err(ClientError::InvalidResponse),
272        None if response.error.is_some() => Ok(()),
273        None => Err(ClientError::InvalidResponse),
274    }
275}
276
277fn response_to_result(response: Response, expected_id: &RpcId) -> Result<Value, ClientError> {
278    validate_response_id(&response, expected_id)?;
279
280    if let Some(rpc_error) = response.error {
281        let (category, retryable, context, suggestion) = if let Some(data) = rpc_error.data.as_ref()
282        {
283            let cat = data
284                .get("category")
285                .and_then(|v| v.as_str())
286                .and_then(|s| s.parse::<error_codes::ErrorCategory>().ok());
287            let retry = data
288                .get("retryable")
289                .and_then(serde_json::Value::as_bool)
290                .unwrap_or_else(|| error_codes::is_retryable(rpc_error.code));
291            let ctx = data.get("context").cloned();
292            let sug = data
293                .get("suggestion")
294                .and_then(|v| v.as_str())
295                .map(String::from);
296            (cat, retry, ctx, sug)
297        } else {
298            (
299                Some(error_codes::category_for_code(rpc_error.code)),
300                error_codes::is_retryable(rpc_error.code),
301                None,
302                None,
303            )
304        };
305        let retry_delay_ms = retry_delay_ms_from_data(rpc_error.data.as_ref());
306
307        return Err(ClientError::RpcError {
308            code: rpc_error.code,
309            message: rpc_error.message,
310            category,
311            retryable,
312            retry_delay_ms,
313            context,
314            suggestion,
315        });
316    }
317
318    response.result.ok_or(ClientError::InvalidResponse)
319}
320
321impl DaemonClient for UnixSocketClient {
322    fn call(&mut self, method: &str, params: Option<Value>) -> Result<Value, ClientError> {
323        self.call_with_config(method, params, &DaemonClientConfig::default())
324    }
325
326    fn call_with_config(
327        &mut self,
328        method: &str,
329        params: Option<Value>,
330        config: &DaemonClientConfig,
331    ) -> Result<Value, ClientError> {
332        let request_id = REQUEST_ID.fetch_add(1, Ordering::SeqCst);
333        let rpc_id = RpcId::from(request_id);
334        let request = Request {
335            jsonrpc: "2.0".to_string(),
336            id: rpc_id.clone(),
337            method: method.to_string(),
338            params,
339        };
340        let request_json = serde_json::to_string(&request)?;
341        let start = Instant::now();
342        let mut attempt = 0;
343        let mut retry_delay = config.initial_retry_delay();
344
345        loop {
346            debug!(
347                request_id,
348                method = %method,
349                attempt,
350                read_timeout_ms = config.read_timeout().as_millis(),
351                write_timeout_ms = config.write_timeout().as_millis(),
352                "RPC call started"
353            );
354
355            let result = (|| {
356                let mut connection = self.transport.connect_connection()?;
357                connection.set_read_timeout(Some(config.read_timeout()))?;
358                connection.set_write_timeout(Some(config.write_timeout()))?;
359
360                trace!(
361                    request_id,
362                    attempt,
363                    bytes = request_json.len(),
364                    "RPC request serialized"
365                );
366                connection.send_message(&request_json)?;
367
368                let response_line = connection
369                    .read_message()?
370                    .ok_or(ClientError::InvalidResponse)?;
371                trace!(
372                    request_id,
373                    attempt,
374                    bytes = response_line.len(),
375                    "RPC response received"
376                );
377
378                let response: Response = serde_json::from_str(&response_line)?;
379                response_to_result(response, &rpc_id)
380            })();
381
382            match result {
383                Ok(value) => {
384                    debug!(
385                        request_id,
386                        method = %method,
387                        attempt,
388                        elapsed_ms = start.elapsed().as_millis(),
389                        "RPC call finished"
390                    );
391                    return Ok(value);
392                }
393                Err(err) if attempt < config.max_retries() && is_retryable_call_error(&err) => {
394                    let delay = next_retry_delay(&err, config.initial_retry_delay(), retry_delay);
395                    debug!(
396                        request_id,
397                        method = %method,
398                        attempt,
399                        retry_in_ms = delay.as_millis(),
400                        error = %err,
401                        "RPC call failed; retrying"
402                    );
403                    std::thread::park_timeout(delay);
404                    retry_delay = delay.saturating_mul(2);
405                    attempt += 1;
406                }
407                Err(err) => {
408                    debug!(
409                        request_id,
410                        method = %method,
411                        attempt,
412                        elapsed_ms = start.elapsed().as_millis(),
413                        error = %err,
414                        "RPC call finished with error"
415                    );
416                    return Err(err);
417                }
418            }
419        }
420    }
421
422    fn call_stream(
423        &mut self,
424        method: &str,
425        params: Option<Value>,
426    ) -> Result<StreamResponse, ClientError> {
427        self.call_stream_with_config(method, params, &DaemonClientConfig::default())
428    }
429
430    fn call_stream_with_config(
431        &mut self,
432        method: &str,
433        params: Option<Value>,
434        config: &DaemonClientConfig,
435    ) -> Result<StreamResponse, ClientError> {
436        let request_id = REQUEST_ID.fetch_add(1, Ordering::SeqCst);
437        let rpc_id = RpcId::from(request_id);
438        let request = Request {
439            jsonrpc: "2.0".to_string(),
440            id: rpc_id.clone(),
441            method: method.to_string(),
442            params,
443        };
444        let request_json = serde_json::to_string(&request)?;
445        let start = Instant::now();
446        let mut attempt = 0;
447        let mut retry_delay = config.initial_retry_delay();
448
449        loop {
450            debug!(
451                request_id,
452                method = %method,
453                attempt,
454                read_timeout_ms = config.read_timeout().as_millis(),
455                write_timeout_ms = config.write_timeout().as_millis(),
456                "RPC stream call started"
457            );
458
459            let result = (|| {
460                let mut connection = self.transport.connect_connection()?;
461                connection.set_read_timeout(Some(config.read_timeout()))?;
462                connection.set_write_timeout(Some(config.write_timeout()))?;
463
464                trace!(
465                    request_id,
466                    attempt,
467                    bytes = request_json.len(),
468                    "RPC stream request serialized"
469                );
470                connection.send_message(&request_json)?;
471                let response_line = connection
472                    .read_message()?
473                    .ok_or(ClientError::InvalidResponse)?;
474                trace!(
475                    request_id,
476                    attempt,
477                    bytes = response_line.len(),
478                    "RPC stream handshake received"
479                );
480                let response: Response = serde_json::from_str(&response_line)?;
481                let _ = response_to_result(response, &rpc_id)?;
482
483                connection.set_read_timeout(Some(STREAM_POLL_TIMEOUT))?;
484
485                Ok(StreamResponse {
486                    connection,
487                    expected_id: rpc_id.clone(),
488                    aborted: Arc::new(AtomicBool::new(false)),
489                })
490            })();
491
492            match result {
493                Ok(stream) => {
494                    debug!(
495                        request_id,
496                        method = %method,
497                        attempt,
498                        elapsed_ms = start.elapsed().as_millis(),
499                        "RPC stream call finished"
500                    );
501                    return Ok(stream);
502                }
503                Err(err) if attempt < config.max_retries() && is_retryable_call_error(&err) => {
504                    let delay = next_retry_delay(&err, config.initial_retry_delay(), retry_delay);
505                    debug!(
506                        request_id,
507                        method = %method,
508                        attempt,
509                        retry_in_ms = delay.as_millis(),
510                        error = %err,
511                        "RPC stream call failed; retrying"
512                    );
513                    std::thread::park_timeout(delay);
514                    retry_delay = delay.saturating_mul(2);
515                    attempt += 1;
516                }
517                Err(err) => {
518                    debug!(
519                        request_id,
520                        method = %method,
521                        attempt,
522                        elapsed_ms = start.elapsed().as_millis(),
523                        error = %err,
524                        "RPC stream call finished with error"
525                    );
526                    return Err(err);
527                }
528            }
529        }
530    }
531}
532
533pub fn ensure_daemon() -> Result<UnixSocketClient, ClientError> {
534    ensure_daemon_with_transport(default_transport())
535}
536
537pub(crate) fn ensure_daemon_with_transport(
538    transport: std::sync::Arc<dyn IpcTransport>,
539) -> Result<UnixSocketClient, ClientError> {
540    debug!("Ensuring daemon is running");
541    if !transport.is_daemon_running() {
542        debug!("Daemon not running");
543        if transport.supports_autostart() {
544            debug!("Attempting daemon autostart");
545            eprintln!("{} Starting daemon in background...", Colors::dim("Note:"));
546            transport.start_daemon_background()?;
547        } else {
548            return Err(ClientError::DaemonNotRunning);
549        }
550    }
551
552    UnixSocketClient::connect_with_transport(transport)
553}
554
555#[derive(Debug, Clone, PartialEq, Eq)]
556pub enum PidLookupResult {
557    Found(u32),
558    NotRunning,
559    Error(String),
560}
561
562#[derive(Debug, Clone, PartialEq, Eq)]
563pub enum DaemonProcessLookupResult {
564    Found(ProcessIdentity),
565    NotRunning,
566    InvalidState { path: PathBuf, message: String },
567}
568
569#[derive(Debug, Deserialize)]
570struct DaemonLockFile {
571    pid: u32,
572    #[serde(default)]
573    process_started_at: Option<u64>,
574}
575
576pub fn get_daemon_process_identity() -> DaemonProcessLookupResult {
577    let lock_path = socket_path().with_extension("lock");
578    if !lock_path.exists() {
579        return DaemonProcessLookupResult::NotRunning;
580    }
581
582    let content = match std::fs::read_to_string(&lock_path) {
583        Ok(content) => content,
584        Err(err) => {
585            return DaemonProcessLookupResult::InvalidState {
586                path: lock_path,
587                message: format!("failed to read lock file: {err}"),
588            };
589        }
590    };
591
592    parse_daemon_lock_file(&content).map_or_else(
593        |message| DaemonProcessLookupResult::InvalidState {
594            path: lock_path,
595            message,
596        },
597        DaemonProcessLookupResult::Found,
598    )
599}
600
601pub fn get_daemon_pid() -> PidLookupResult {
602    match get_daemon_process_identity() {
603        DaemonProcessLookupResult::Found(identity) => PidLookupResult::Found(identity.pid),
604        DaemonProcessLookupResult::NotRunning => PidLookupResult::NotRunning,
605        DaemonProcessLookupResult::InvalidState { path, message } => PidLookupResult::Error(
606            format!("Invalid daemon lock state {}: {message}", path.display()),
607        ),
608    }
609}
610
611fn parse_daemon_lock_file(content: &str) -> Result<ProcessIdentity, String> {
612    let trimmed = content.trim();
613    if trimmed.is_empty() {
614        return Err("lock file is empty".to_string());
615    }
616
617    if let Ok(pid) = trimmed.parse::<u32>() {
618        return Ok(ProcessIdentity {
619            pid,
620            started_at: None,
621        });
622    }
623
624    let payload: DaemonLockFile = serde_json::from_str(trimmed)
625        .map_err(|err| format!("lock file is not a valid daemon identity payload: {err}"))?;
626    Ok(ProcessIdentity {
627        pid: payload.pid,
628        started_at: payload.process_started_at,
629    })
630}
631
632#[cfg(test)]
633#[path = "client_tests.rs"]
634mod tests;