Skip to main content

rmux_sdk/
trace.rs

1//! Minimal JSONL tracing for terminal automation workflows.
2
3use std::collections::VecDeque;
4use std::path::{Path, PathBuf};
5use std::sync::{Arc, Mutex};
6use std::time::{SystemTime, UNIX_EPOCH};
7
8use serde::Serialize;
9
10use crate::{Pane, Result, Rmux, RmuxError};
11
12const DEFAULT_MAX_TRACE_EVENTS: usize = 100_000;
13
14/// Builder for a minimal trace session.
15#[derive(Debug)]
16pub struct RmuxTraceBuilder<'a> {
17    rmux: &'a Rmux,
18    max_events: usize,
19}
20
21impl<'a> RmuxTraceBuilder<'a> {
22    pub(crate) const fn new(rmux: &'a Rmux) -> Self {
23        Self {
24            rmux,
25            max_events: DEFAULT_MAX_TRACE_EVENTS,
26        }
27    }
28
29    /// Caps the number of events retained in memory before [`Self::start`].
30    ///
31    /// When the cap is reached, the oldest event is dropped. The trace API is
32    /// intentionally minimal and in-memory; use a small cap for long-running
33    /// processes that only need recent automation context.
34    pub const fn max_events(mut self, max_events: usize) -> Self {
35        self.max_events = max_events;
36        self
37    }
38
39    /// Starts an in-memory trace buffer.
40    ///
41    /// The trace is written when [`TraceSession::stop`] is called. Events are
42    /// buffered in memory up to the configured cap. The SDK does not install
43    /// global hooks; callers record events explicitly.
44    pub async fn start(self) -> Result<TraceSession> {
45        let session = TraceSession {
46            endpoint: format!("{:?}", self.rmux.resolved_endpoint()?),
47            max_events: self.max_events,
48            events: Arc::new(Mutex::new(VecDeque::new())),
49        };
50        session.record("trace.start", None::<TracePayload>)?;
51        Ok(session)
52    }
53}
54
55/// Active minimal trace session.
56#[derive(Debug, Clone)]
57pub struct TraceSession {
58    endpoint: String,
59    max_events: usize,
60    events: Arc<Mutex<VecDeque<TraceEvent>>>,
61}
62
63impl TraceSession {
64    /// Records a free-form action event.
65    pub fn record_action(&self, action: impl Into<String>) -> Result<()> {
66        self.record(
67            "action",
68            Some(TracePayload {
69                action: Some(action.into()),
70                ..TracePayload::default()
71            }),
72        )
73    }
74
75    /// Records input sent to a pane.
76    pub fn record_input(&self, pane: &Pane, input: impl Into<String>) -> Result<()> {
77        self.record(
78            "input",
79            Some(TracePayload {
80                pane: Some(format!("{}", pane.target().to_proto())),
81                input: Some(input.into()),
82                ..TracePayload::default()
83            }),
84        )
85    }
86
87    /// Captures and records the pane's current visible snapshot text.
88    pub async fn record_snapshot(&self, pane: &Pane) -> Result<()> {
89        let snapshot = pane.snapshot().await?;
90        self.record(
91            "snapshot",
92            Some(TracePayload {
93                pane: Some(format!("{}", pane.target().to_proto())),
94                revision: Some(snapshot.revision),
95                snapshot: Some(snapshot.visible_text()),
96                ..TracePayload::default()
97            }),
98        )
99    }
100
101    /// Stops tracing and writes `trace.jsonl` into `directory`.
102    pub async fn stop(self, directory: impl AsRef<Path>) -> Result<PathBuf> {
103        self.record("trace.stop", None::<TracePayload>)?;
104        let directory = directory.as_ref();
105        tokio::fs::create_dir_all(directory)
106            .await
107            .map_err(trace_io_error)?;
108        let path = directory.join("trace.jsonl");
109        let events = {
110            let events = self.events.lock().map_err(lock_error)?;
111            events.iter().cloned().collect::<Vec<_>>()
112        };
113        let lines = events
114            .iter()
115            .map(serde_json::to_string)
116            .collect::<core::result::Result<Vec<_>, _>>()
117            .map_err(|error| {
118                RmuxError::protocol(rmux_proto::RmuxError::Server(format!(
119                    "failed to encode rmux trace event: {error}"
120                )))
121            })?
122            .join("\n");
123        tokio::fs::write(&path, format!("{lines}\n"))
124            .await
125            .map_err(trace_io_error)?;
126        Ok(path)
127    }
128
129    fn record(&self, kind: &'static str, payload: Option<TracePayload>) -> Result<()> {
130        let mut events = self.events.lock().map_err(lock_error)?;
131        if self.max_events == 0 {
132            return Ok(());
133        }
134        if events.len() == self.max_events {
135            events.pop_front();
136        }
137        events.push_back(TraceEvent {
138            timestamp_ms: timestamp_ms(),
139            endpoint: self.endpoint.clone(),
140            kind,
141            payload,
142        });
143        Ok(())
144    }
145}
146
147#[derive(Debug, Clone, Serialize)]
148struct TraceEvent {
149    timestamp_ms: u128,
150    endpoint: String,
151    kind: &'static str,
152    #[serde(skip_serializing_if = "Option::is_none")]
153    payload: Option<TracePayload>,
154}
155
156#[derive(Debug, Clone, Default, Serialize)]
157struct TracePayload {
158    #[serde(skip_serializing_if = "Option::is_none")]
159    action: Option<String>,
160    #[serde(skip_serializing_if = "Option::is_none")]
161    pane: Option<String>,
162    #[serde(skip_serializing_if = "Option::is_none")]
163    input: Option<String>,
164    #[serde(skip_serializing_if = "Option::is_none")]
165    revision: Option<u64>,
166    #[serde(skip_serializing_if = "Option::is_none")]
167    snapshot: Option<String>,
168}
169
170fn timestamp_ms() -> u128 {
171    SystemTime::now()
172        .duration_since(UNIX_EPOCH)
173        .map_or(0, |duration| duration.as_millis())
174}
175
176fn trace_io_error(error: std::io::Error) -> RmuxError {
177    RmuxError::transport("write rmux trace", error)
178}
179
180fn lock_error<T>(error: std::sync::PoisonError<T>) -> RmuxError {
181    RmuxError::protocol(rmux_proto::RmuxError::Server(format!(
182        "rmux trace lock poisoned: {error}"
183    )))
184}