use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::sync::mpsc;
use crate::adapter::Emitter;
use crate::agent::SessionOptions;
use crate::event::DiagnosticLevel;
use crate::process::{Child, SharedStdin};
pub(crate) const FRAME_BUFFER: usize = 64;
pub(crate) struct LineWire {
stdin: SharedStdin,
pub frames: mpsc::Receiver<Value>,
recorder: Option<WireRecorder>,
}
impl LineWire {
pub(crate) fn over(child: &mut Child, recorder: Option<WireRecorder>) -> Self {
let stdin = child.stdin.clone();
let stdout = child.stdout.take().expect("piped stdout");
let (tx, frames) = mpsc::channel(FRAME_BUFFER);
let reader_recorder = recorder.clone();
tokio::spawn(async move {
let mut lines = BufReader::new(stdout).lines();
while let Ok(Some(line)) = lines.next_line().await {
let Ok(frame) = serde_json::from_str::<Value>(&line) else {
continue;
};
if let Some(recorder) = &reader_recorder {
recorder.record("in", &frame);
}
if tx.send(frame).await.is_err() {
break;
}
}
});
Self {
stdin,
frames,
recorder,
}
}
pub(crate) async fn write(&mut self, frame: Value) -> std::io::Result<()> {
if let Some(recorder) = &self.recorder {
recorder.record("out", &frame);
}
let mut line = frame.to_string();
line.push('\n');
match self.stdin.lock().await.as_mut() {
Some(stdin) => stdin.write_all(line.as_bytes()).await,
None => Err(std::io::ErrorKind::BrokenPipe.into()),
}
}
}
#[derive(Clone)]
pub(crate) struct WireRecorder {
lines: mpsc::UnboundedSender<Vec<u8>>,
}
impl WireRecorder {
pub(crate) async fn for_session(options: &SessionOptions, events: &Emitter) -> Option<Self> {
let path = options.record_wire.as_deref()?;
let events = events.clone();
let file = match tokio::fs::OpenOptions::new()
.create(true)
.append(true)
.open(path)
.await
{
Ok(file) => file,
Err(e) => {
let _ = events
.diagnostic(
DiagnosticLevel::Warning,
format!("wire recording is off: {e}"),
)
.await;
return None;
}
};
let (lines, mut rx) = mpsc::unbounded_channel::<Vec<u8>>();
tokio::spawn(async move {
let mut file = file;
while let Some(bytes) = rx.recv().await {
if let Err(e) = append(&mut file, &bytes).await {
let _ = events
.diagnostic(
DiagnosticLevel::Warning,
format!("wire recording stopped: {e}"),
)
.await;
break;
}
}
});
Some(Self { lines })
}
pub(crate) fn record(&self, dir: &'static str, frame: &Value) {
let mut line = json!({ "dir": dir, "frame": frame }).to_string();
line.push('\n');
let _ = self.lines.send(line.into_bytes());
}
}
async fn append(file: &mut tokio::fs::File, bytes: &[u8]) -> std::io::Result<()> {
file.write_all(bytes).await?;
file.flush().await
}