From cfa7bbcc718b191724cf3daaa9a7caa75898bf50 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Sat, 26 Sep 2026 18:48:07 -0400 Subject: [PATCH 01/18] feat(terminal-streamer): decode a complete TRP recording A finished recording can now be read as asciicast lines without tailing it. The lines come out one packet at a time, so memory use does not grow with the recording. The packet decoding is shared with the live stream decoder. Events written before the terminal setup are still held until the setup arrives, but no more than 1 MiB of them: past that, the header goes out with the terminal size known so far, so a recording that never sends its setup is not held whole in memory. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/terminal-streamer/src/trp_decoder.rs | 409 ++++++++++++++++---- 1 file changed, 339 insertions(+), 70 deletions(-) diff --git a/crates/terminal-streamer/src/trp_decoder.rs b/crates/terminal-streamer/src/trp_decoder.rs index 751ec3d8c..77c9d09ed 100644 --- a/crates/terminal-streamer/src/trp_decoder.rs +++ b/crates/terminal-streamer/src/trp_decoder.rs @@ -1,3 +1,5 @@ +use std::collections::VecDeque; +use std::io::Read; use std::pin::Pin; use std::task::{Context, Poll}; @@ -80,107 +82,226 @@ impl AsyncRead for AsyncReadChannel { } } -async fn parse_trp_stream( - mut input_stream: impl AsyncRead + Unpin + Send + 'static, - mut tx: tokio::sync::mpsc::Sender>, -) -> anyhow::Result<()> { - let mut time = 0.0; - let mut before_setup_cache = Some(Vec::new()); - let mut header = AsciinemaHeader::default(); +/// Decodes a complete TRP recording into asciicast v2 lines, without their line ending. +/// +/// The recording is read one packet at a time, so memory use does not grow with its size. +/// A truncated last packet, as left by an interrupted recording, is ignored. +/// The lines end after the first error. +pub struct AsciicastLines { + input: R, + decoder: Option, + pending: VecDeque, + payload: Vec, +} - loop { - let mut packet_head_buffer = [0u8; 8]; - if let Err(e) = input_stream.read_exact(&mut packet_head_buffer).await { - if e.kind() == std::io::ErrorKind::UnexpectedEof { - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - continue; - } - anyhow::bail!(e); +impl AsciicastLines { + pub fn new(input: R) -> Self { + Self { + input, + decoder: Some(TrpDecoder::default()), + pending: VecDeque::new(), + payload: Vec::new(), } + } - let time_delta = u32::from_le_bytes(packet_head_buffer[0..4].try_into()?); - let event_type = u16::from_le_bytes(packet_head_buffer[4..6].try_into()?); - let size = u16::from_le_bytes(packet_head_buffer[6..8].try_into()?); + /// Decodes the next packet into `pending`; at the end of the input, flushes the decoder. + fn decode_next_packet(&mut self) -> anyhow::Result<()> { + let Some(decoder) = self.decoder.as_mut() else { + return Ok(()); + }; - time += f64::from(time_delta) / 1000.0; + let mut header = [0u8; PACKET_HEADER_SIZE]; + if read_complete(&mut self.input, &mut header)? { + let header = PacketHeader::parse(&header); + self.payload.resize(usize::from(header.size), 0); - let mut event_payload = vec![0u8; size as usize]; - if let Err(e) = input_stream.read_exact(&mut event_payload).await { - if e.kind() == std::io::ErrorKind::UnexpectedEof { - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - continue; + if read_complete(&mut self.input, &mut self.payload)? { + self.pending.extend(decoder.push(&header, &self.payload)?); + return Ok(()); } - anyhow::bail!(e); } - match event_type { - 0 => { - // Terminal output - let event_payload = String::from_utf8_lossy(&event_payload).into_owned(); - let event = AsciinemaEvent::TerminalOutput { - payload: event_payload, - time, - }; - match before_setup_cache { - Some(ref mut cache) => { - cache.push(event); - } - None => { - send(&mut tx, event.to_json()).await?; - } - } + if let Some(decoder) = self.decoder.take() { + self.pending.extend(decoder.finish()); + } + + Ok(()) + } +} + +impl Iterator for AsciicastLines { + type Item = anyhow::Result; + + fn next(&mut self) -> Option { + loop { + if let Some(line) = self.pending.pop_front() { + return Some(Ok(line)); + } + + self.decoder.as_ref()?; + + if let Err(error) = self.decode_next_packet() { + self.decoder = None; + return Some(Err(error)); } - 1 => { - let event_payload = String::from_utf8_lossy(&event_payload).into_owned(); - let event = AsciinemaEvent::UserInput { - payload: event_payload, - time, + } + } +} + +/// Fills `buffer`, or returns `false` if the input ends first. +fn read_complete(input: &mut impl Read, buffer: &mut [u8]) -> std::io::Result { + match input.read_exact(buffer) { + Ok(()) => Ok(true), + Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => Ok(false), + Err(error) => Err(error), + } +} + +const PACKET_HEADER_SIZE: usize = 8; + +struct PacketHeader { + time_delta: u32, + event_type: u16, + size: u16, +} + +impl PacketHeader { + fn parse(buffer: &[u8; PACKET_HEADER_SIZE]) -> Self { + Self { + time_delta: u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]), + event_type: u16::from_le_bytes([buffer[4], buffer[5]]), + size: u16::from_le_bytes([buffer[6], buffer[7]]), + } + } +} + +/// Most payload bytes held while waiting for the terminal setup; past it, the header is sent with the size known so far. +const MAX_BEFORE_SETUP_LEN: usize = 1024 * 1024; + +/// Turns TRP packets into asciicast lines, holding the events that come before the terminal setup. +struct TrpDecoder { + time: f64, + before_setup_cache: Option>, + before_setup_len: usize, + header: AsciinemaHeader, +} + +impl Default for TrpDecoder { + fn default() -> Self { + Self { + time: 0.0, + before_setup_cache: Some(Vec::new()), + before_setup_len: 0, + header: AsciinemaHeader::default(), + } + } +} + +impl TrpDecoder { + fn push(&mut self, packet: &PacketHeader, payload: &[u8]) -> anyhow::Result> { + self.time += f64::from(packet.time_delta) / 1000.0; + let time = self.time; + + let mut lines = Vec::new(); + + match packet.event_type { + 0 | 1 => { + let packet_len = payload.len(); + let payload = String::from_utf8_lossy(payload).into_owned(); + let event = if packet.event_type == 0 { + AsciinemaEvent::TerminalOutput { payload, time } + } else { + AsciinemaEvent::UserInput { payload, time } }; - match before_setup_cache { + match self.before_setup_cache { Some(ref mut cache) => { cache.push(event); + self.before_setup_len += packet_len; + + // A recording that never sends its terminal setup must not be held whole in memory. + if self.before_setup_len > MAX_BEFORE_SETUP_LEN { + lines.extend(self.flush_before_setup()); + } } - None => { - send(&mut tx, event.to_json()).await?; - } + None => lines.push(event.to_json()), } } 2 => { // Terminal size change. Payload is little-endian [columns, rows]. - if event_payload.len() < 4 { - anyhow::bail!( - "invalid terminal size change payload length (len={})", - event_payload.len() - ); + if payload.len() < 4 { + anyhow::bail!("invalid terminal size change payload length (len={})", payload.len()); } - header.col = u16::from_le_bytes(event_payload[0..2].try_into()?); - header.row = u16::from_le_bytes(event_payload[2..4].try_into()?); - if before_setup_cache.is_none() { + self.header.col = u16::from_le_bytes([payload[0], payload[1]]); + self.header.row = u16::from_le_bytes([payload[2], payload[3]]); + if self.before_setup_cache.is_none() { let event = AsciinemaEvent::Resize { - width: header.col, - height: header.row, + width: self.header.col, + height: self.header.row, time, }; - send(&mut tx, event.to_json()).await?; + lines.push(event.to_json()); } } 4 => { // Terminal setup - if before_setup_cache.is_some() { - let header_json = header.to_json(); - send(&mut tx, header_json).await?; - if let Some(ref mut cache) = before_setup_cache { - for event in cache.drain(..) { - send(&mut tx, event.to_json()).await?; - } - } - before_setup_cache = None; + if self.before_setup_cache.is_some() { + lines.extend(self.flush_before_setup()); } else { warn!("Received terminal setup event but cache is empty"); } } _ => {} } + + Ok(lines) + } + + /// Returns the header and the cached events, once; the events that follow are not cached anymore. + fn flush_before_setup(&mut self) -> Vec { + match self.before_setup_cache.take() { + Some(cache) => core::iter::once(self.header.to_json()) + .chain(cache.iter().map(AsciinemaEvent::to_json)) + .collect(), + None => Vec::new(), + } + } + + /// Flushes the cached events of a recording that never sent its terminal setup. + fn finish(mut self) -> Vec { + self.flush_before_setup() + } +} + +async fn parse_trp_stream( + mut input_stream: impl AsyncRead + Unpin + Send + 'static, + mut tx: tokio::sync::mpsc::Sender>, +) -> anyhow::Result<()> { + let mut decoder = TrpDecoder::default(); + + loop { + let mut packet_head_buffer = [0u8; PACKET_HEADER_SIZE]; + if let Err(e) = input_stream.read_exact(&mut packet_head_buffer).await { + if e.kind() == std::io::ErrorKind::UnexpectedEof { + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + continue; + } + anyhow::bail!(e); + } + + let packet = PacketHeader::parse(&packet_head_buffer); + + let mut event_payload = vec![0u8; usize::from(packet.size)]; + if let Err(e) = input_stream.read_exact(&mut event_payload).await { + if e.kind() == std::io::ErrorKind::UnexpectedEof { + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + continue; + } + anyhow::bail!(e); + } + + for line in decoder.push(&packet, &event_payload)? { + send(&mut tx, line).await?; + } } } @@ -189,3 +310,151 @@ async fn send(sender: &mut tokio::sync::mpsc::Sender>, mu sender.send(Ok(json)).await?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn packet(time_delta: u32, event_type: u16, payload: &[u8]) -> Vec { + let mut packet = Vec::new(); + packet.extend_from_slice(&time_delta.to_le_bytes()); + packet.extend_from_slice(&event_type.to_le_bytes()); + packet.extend_from_slice(&u16::try_from(payload.len()).expect("small payload").to_le_bytes()); + packet.extend_from_slice(payload); + packet + } + + fn decode(trp: &[u8]) -> anyhow::Result { + AsciicastLines::new(trp) + .map(|line| Ok(format!("{}\n", line?))) + .collect() + } + + #[test] + fn decodes_a_complete_recording_and_ignores_a_truncated_tail() { + let mut trp = Vec::new(); + trp.extend(packet(0, 2, &[100, 0, 30, 0])); + trp.extend(packet(500, 0, b"$ ")); + trp.extend(packet(0, 4, &[])); + trp.extend(packet(1000, 1, b"l")); + trp.extend(packet(250, 0, b"ls\r\n")); + trp.extend(&packet(10, 0, b"lost")[..9]); + + let cast = decode(&trp).expect("valid recording"); + + assert_eq!( + cast, + concat!( + "{\"version\": 2, \"width\": 100, \"height\": 30}\n", + "[0.5,\"o\",\"$ \"]\n", + "[1.5,\"i\",\"l\"]\n", + r#"[1.75,"o","ls\u000d\u000a"]"#, + "\n", + ) + ); + } + + #[tokio::test] + async fn complete_recording_matches_the_live_stream_decoder() { + use tokio::io::AsyncBufReadExt as _; + + let mut trp = Vec::new(); + trp.extend(packet(0, 2, &[120, 0, 40, 0])); + trp.extend(packet(100, 0, b"before setup")); + trp.extend(packet(0, 4, &[])); + trp.extend(packet(300, 1, b"x")); + trp.extend(packet(0, 2, &[80, 0, 24, 0])); + trp.extend(packet(42, 0, "é\u{1b}[0m".as_bytes())); + + let complete = decode(&trp).expect("valid recording"); + + let (task, live_reader) = decode_stream(std::io::Cursor::new(trp)).expect("start the live decoder"); + let mut live_lines = tokio::io::BufReader::new(live_reader).lines(); + let mut streamed = String::new(); + for _ in complete.lines() { + let line = live_lines + .next_line() + .await + .expect("read a live line") + .expect("a live line"); + streamed.push_str(&line); + streamed.push('\n'); + } + task.abort(); + + assert_eq!(complete, streamed); + } + + struct FailingReader; + + impl Read for FailingReader { + fn read(&mut self, _: &mut [u8]) -> std::io::Result { + Err(std::io::Error::other("disk failure")) + } + } + + #[test] + fn lines_come_out_before_the_rest_of_the_recording_is_read() { + let mut trp = Vec::new(); + trp.extend(packet(0, 4, &[])); + trp.extend(packet(1000, 0, b"first")); + + let mut lines = AsciicastLines::new(Read::chain(trp.as_slice(), FailingReader)); + + assert_eq!( + lines.next().expect("header").expect("valid header"), + "{\"version\": 2, \"width\": 80, \"height\": 24}" + ); + assert_eq!( + lines.next().expect("event").expect("valid event"), + "[1,\"o\",\"first\"]" + ); + let error = lines.next().expect("read error").expect_err("the rest fails"); + assert!(error.to_string().contains("disk failure"), "{error:#}"); + assert!(lines.next().is_none(), "nothing comes after an error"); + } + + #[test] + fn events_waiting_for_the_terminal_setup_are_bounded() { + let payload = vec![b'x'; 60_000]; + let packets = MAX_BEFORE_SETUP_LEN / payload.len() + 1; + let mut trp = Vec::new(); + for _ in 0..packets { + trp.extend(packet(0, 0, &payload)); + } + + // The setup never comes: the rest of the recording fails to read. + let lines = AsciicastLines::new(Read::chain(trp.as_slice(), FailingReader)).collect::>(); + + assert_eq!(lines.len(), 1 + packets + 1, "header, every event, then the read error"); + assert_eq!( + lines[0].as_ref().expect("valid header"), + "{\"version\": 2, \"width\": 80, \"height\": 24}" + ); + assert!(lines[1..=packets].iter().all(Result::is_ok)); + assert!(lines[packets + 1].is_err()); + } + + #[test] + fn malformed_input_fails_cleanly() { + let error = decode(&packet(0, 2, &[1, 2])).expect_err("short resize payload"); + assert!( + error + .to_string() + .contains("invalid terminal size change payload length") + ); + + let cast = decode(&[0xFF; 5]).expect("input shorter than a packet header"); + assert_eq!(cast, "{\"version\": 2, \"width\": 80, \"height\": 24}\n"); + } + + #[test] + fn events_of_a_recording_without_setup_are_kept() { + let cast = decode(&packet(2000, 0, b"x")).expect("valid recording"); + + assert_eq!( + cast, + "{\"version\": 2, \"width\": 80, \"height\": 24}\n[2,\"o\",\"x\"]\n" + ); + } +} From 2c333625383ba2d994b4e517eb9711f693d641f1 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 7 Oct 2026 18:00:08 -0400 Subject: [PATCH 02/18] feat(dgw): let the AI cite the input instead of writing times Action times now come from the input, never from the model. Each transcript line sent to `describe_session_actions` gets an id (`L12`), the model cites the line where an action starts, and the action takes the time of that line. Actions citing a line that is not in the transcript are skipped. The prompt version is now `session-actions-3`. A new purpose, `describe_screen_actions`, lists what the user did in screenshots of a session: each PNG is sent after a label with its id and time, the model cites the screenshot that shows each action, and the action takes its time. Screenshots marked as context show where the session was and are never cited. Requests can now carry images, for OpenAI chat completions and Anthropic Messages. Co-Authored-By: Claude Opus 5.5 (1M context) --- Cargo.lock | 1 + crates/devolutions-gateway-ai/Cargo.toml | 1 + crates/devolutions-gateway-ai/src/client.rs | 22 ++- crates/devolutions-gateway-ai/src/lib.rs | 3 +- .../src/screen_actions.rs | 172 ++++++++++++++++++ .../src/session_actions.rs | 142 ++++++++++----- .../src/wire/anthropic.rs | 12 +- crates/devolutions-gateway-ai/src/wire/mod.rs | 27 ++- .../devolutions-gateway-ai/src/wire/openai.rs | 14 +- .../devolutions-gateway-ai/tests/providers.rs | 120 ++++++++++-- 10 files changed, 440 insertions(+), 74 deletions(-) create mode 100644 crates/devolutions-gateway-ai/src/screen_actions.rs diff --git a/Cargo.lock b/Cargo.lock index b29538e97..da4f7d7c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1937,6 +1937,7 @@ dependencies = [ name = "devolutions-gateway-ai" version = "0.0.0" dependencies = [ + "base64 0.23.1", "reqwest", "secrecy", "serde", diff --git a/crates/devolutions-gateway-ai/Cargo.toml b/crates/devolutions-gateway-ai/Cargo.toml index b6c6f2178..91a268069 100644 --- a/crates/devolutions-gateway-ai/Cargo.toml +++ b/crates/devolutions-gateway-ai/Cargo.toml @@ -10,6 +10,7 @@ publish = false workspace = true [dependencies] +base64 = "0.23" reqwest = { version = "0.12", default-features = false, features = ["json"] } secrecy = "0.10" serde = { version = "1", features = ["derive"] } diff --git a/crates/devolutions-gateway-ai/src/client.rs b/crates/devolutions-gateway-ai/src/client.rs index ca15a1717..c58c7998c 100644 --- a/crates/devolutions-gateway-ai/src/client.rs +++ b/crates/devolutions-gateway-ai/src/client.rs @@ -208,11 +208,19 @@ impl fmt::Debug for AiClient { pub(crate) struct Prompt<'a> { /// Instructions of the purpose. pub(crate) system: &'a str, - /// Data the purpose works on, such as a session transcript. - pub(crate) input: &'a str, + /// Data the purpose works on, such as a session transcript or screenshots, in order. + pub(crate) input: &'a [Input<'a>], pub(crate) max_output_tokens: u32, } +/// Part of the data a purpose works on. +#[derive(Clone, Copy)] +pub(crate) enum Input<'a> { + Text(&'a str), + /// A PNG image. + Png(&'a [u8]), +} + impl AiClient { pub fn builder() -> AiClientBuilder { AiClientBuilder::default() @@ -231,7 +239,15 @@ impl AiClient { provider = ?self.provider, model = %self.model, base_url = %self.base_url, - input_len = prompt.input.len(), + input_len = prompt + .input + .iter() + .map(|input| match input { + Input::Text(text) => text.len(), + Input::Png(_) => 0, + }) + .sum::(), + images = prompt.input.iter().filter(|input| matches!(input, Input::Png(_))).count(), max_output_tokens = prompt.max_output_tokens, "Send AI completion request" ); diff --git a/crates/devolutions-gateway-ai/src/lib.rs b/crates/devolutions-gateway-ai/src/lib.rs index 99bf42f89..f92c3283e 100644 --- a/crates/devolutions-gateway-ai/src/lib.rs +++ b/crates/devolutions-gateway-ai/src/lib.rs @@ -1,7 +1,7 @@ //! High-level AI for Devolutions Gateway: one method per purpose, never a raw request to a model. //! //! A purpose is one AI job Gateway needs done, such as listing what a user did in a session transcript with -//! [`AiClient::describe_session_actions`]. +//! [`AiClient::describe_session_actions`], or in screenshots of a session with [`AiClient::describe_screen_actions`]. //! The crate does not expose low-level AI requests: there is no chat, no completion call, and no way to pass a prompt. //! A consumer states what it needs and gets a typed result; it is not expected to know, or care, how the result is //! obtained. @@ -31,6 +31,7 @@ mod client; mod error; mod response; +pub mod screen_actions; pub mod session_actions; mod wire; diff --git a/crates/devolutions-gateway-ai/src/screen_actions.rs b/crates/devolutions-gateway-ai/src/screen_actions.rs new file mode 100644 index 000000000..67b1b4282 --- /dev/null +++ b/crates/devolutions-gateway-ai/src/screen_actions.rs @@ -0,0 +1,172 @@ +//! Purpose: list what the user did in screenshots of a session, with [`AiClient::describe_screen_actions`]. + +use std::fmt; +use std::time::Duration; + +use crate::client::{AiClient, Input, Prompt}; +use crate::session_actions::{Action, DEFAULT_MAX_OUTPUT_TOKENS, parse_actions}; +use crate::{Error, Response}; + +/// Version of the prompt of this purpose. +/// +/// Results do not carry it: a caller that stores results should store it with them, so readers know which prompt +/// produced them. +/// Bump it whenever the prompt changes. +pub const PROMPT_VERSION: &str = "screen-actions-1"; + +const PROMPT: &str = r#"You see screenshots of a remote desktop session, in time order, and list what the user did. + +Input: each screenshot comes after a label with its id and the elapsed time since the session started, in seconds. Example: `[F00012 t=34.5s]`. A label that ends with `context` marks a screenshot from before this part of the session: it only shows where the session was. + +Output: JSON Lines only. Write one JSON object per line and nothing else: no prose, no Markdown, no code fences. +Each object has these fields: +- "frame": string. The id of the earliest screenshot that shows the action or its result, like "F00012". Never a context screenshot. +- "description": string. A short past-tense sentence naming the action, like "Reset the password of a user account". Say so when the screen shows that it failed. +- "object": string, optional. The main thing acted on, like a file path, host, service, user, or setting. +- "parameters": object, optional. Every value is a string. Important details, like the window, the exact command, or the new value. + +Example output line: +{"frame":"F00012","description":"Reset the password of a user account","object":"David Moreau","parameters":{"Window":"Active Directory Users and Computers"}} + +Rules: +- Write one line per meaningful user action, in time order, such as opening a tool, changing a setting, running a command, or confirming a dialog. +- Ignore noise, such as mouse movement, hovering, and screen redraws, and actions only visible in context screenshots. +- Report only what the screenshots show; never guess actions between them. +- Never copy passwords, secrets, or tokens. Write "[redacted]" instead. +- If the user did nothing, write only this line: {"noActions":true}"#; + +/// One screenshot of a session. +#[derive(Clone, Copy)] +pub struct Screenshot<'a> { + /// Short id the model cites, unique in the request, such as `F00012`. + pub id: &'a str, + /// Elapsed time since the start of the session. + pub offset: Duration, + /// The image, as PNG. + pub png: &'a [u8], + /// Whether the screenshot comes from before the part of the session being described, to show where it was; no + /// action is ever found in it. + pub context: bool, +} + +// Screenshots hold session data, so only their id and size are printed. +impl fmt::Debug for Screenshot<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Screenshot") + .field("id", &self.id) + .field("offset", &self.offset) + .field("png_len", &self.png.len()) + .field("context", &self.context) + .finish() + } +} + +impl AiClient { + /// Asks the model which actions the user performed in screenshots of a session, in time order. + /// + /// The model cites the screenshot that shows each action, and the action gets the time of that screenshot. + pub fn describe_screen_actions<'a>(&'a self, screenshots: &'a [Screenshot<'a>]) -> DescribeScreenActions<'a> { + DescribeScreenActions { + client: self, + screenshots, + max_output_tokens: DEFAULT_MAX_OUTPUT_TOKENS, + } + } +} + +/// Request built by [`AiClient::describe_screen_actions`]. +#[must_use = "the request is sent only by `send`"] +pub struct DescribeScreenActions<'a> { + client: &'a AiClient, + screenshots: &'a [Screenshot<'a>], + max_output_tokens: u32, +} + +impl fmt::Debug for DescribeScreenActions<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("DescribeScreenActions") + .field("client", self.client) + .field("screenshots", &self.screenshots.len()) + .field("max_output_tokens", &self.max_output_tokens) + .finish() + } +} + +impl DescribeScreenActions<'_> { + /// Upper bound of tokens in the answer; the default is 16000. + pub fn max_output_tokens(mut self, max_output_tokens: u32) -> Self { + self.max_output_tokens = max_output_tokens; + self + } + + /// Returns the actions in the order of the answer; the list is empty only when the model answered that the user did + /// nothing. + /// + /// Invalid lines in the answer are skipped with a warning, and so are actions citing an unknown or context + /// screenshot. + /// The errors are the same as for + /// [`DescribeSessionActions::send`](crate::session_actions::DescribeSessionActions::send): send fewer screenshots + /// when the answer is [`Error::Truncated`]. + pub async fn send(self) -> Result>, Error> { + let labels = self + .screenshots + .iter() + .map(|screenshot| { + let context = if screenshot.context { " context" } else { "" }; + format!("[{} t={:.1}s{context}]", screenshot.id, screenshot.offset.as_secs_f64()) + }) + .collect::>(); + + let input = self + .screenshots + .iter() + .zip(&labels) + .flat_map(|(screenshot, label)| [Input::Text(label), Input::Png(screenshot.png)]) + .collect::>(); + + let prompt = Prompt { + system: PROMPT, + input: &input, + max_output_tokens: self.max_output_tokens, + }; + + self.client.complete(&prompt).await?.try_map(|answer| { + parse_actions(&answer, "frame", |frame| { + self.screenshots + .iter() + .find(|screenshot| screenshot.id == frame && !screenshot.context) + .map(|screenshot| screenshot.offset) + }) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prompt_and_parser_agree_on_the_no_actions_line() { + let line = PROMPT + .lines() + .last() + .and_then(|rule| rule.split_once(": ")) + .map(|(_, line)| line); + + assert_eq!(line, Some(r#"{"noActions":true}"#)); + } + + #[test] + fn prompt_asks_for_the_parsed_fields() { + for field in [ + "\"frame\"", + "description", + "object", + "parameters", + "context", + "noActions", + ] { + assert!(PROMPT.contains(field), "prompt is missing {field}"); + } + } +} diff --git a/crates/devolutions-gateway-ai/src/session_actions.rs b/crates/devolutions-gateway-ai/src/session_actions.rs index 75ef471b3..eda10ef19 100644 --- a/crates/devolutions-gateway-ai/src/session_actions.rs +++ b/crates/devolutions-gateway-ai/src/session_actions.rs @@ -7,7 +7,7 @@ use std::time::Duration; use serde::Deserialize; use tracing::warn; -use crate::client::{AiClient, Prompt}; +use crate::client::{AiClient, Input, Prompt}; use crate::{Error, Response}; /// Version of the prompt of this purpose. @@ -15,21 +15,21 @@ use crate::{Error, Response}; /// Results do not carry it: a caller that stores results should store it with them, so readers know which prompt /// produced them. /// Bump it whenever the prompt changes. -pub const PROMPT_VERSION: &str = "session-actions-2"; +pub const PROMPT_VERSION: &str = "session-actions-3"; const PROMPT: &str = r#"You read the transcript of a remote session and list what the user did. -Input: each transcript line starts with the elapsed time since the session started, in seconds, between square brackets. Example: `[12.5] ls -la`. +Input: each transcript line starts with its id, then the elapsed time since the session started, in seconds, between square brackets. Example: `L12 [34.5] ls -la`. Output: JSON Lines only. Write one JSON object per line and nothing else: no prose, no Markdown, no code fences. Each object has these fields: -- "offsetSeconds": number. Elapsed seconds when the action started, taken from the transcript. +- "line": string. The id of the transcript line where the action starts, like "L12". - "description": string. A short past-tense sentence naming the action, like "Listed directory contents". - "object": string, optional. The main thing acted on, like a file path, host, service, or account. - "parameters": object, optional. Every value is a string. Important details, like the exact command. Example output line: -{"offsetSeconds":12.5,"description":"Listed directory contents","object":"/var/log","parameters":{"Command":"ls -la /var/log"}} +{"line":"L12","description":"Listed directory contents","object":"/var/log","parameters":{"Command":"ls -la /var/log"}} Rules: - Write one line per meaningful user action, in time order. Merge the keystrokes of one command into one action. @@ -38,12 +38,12 @@ Rules: - If the user did nothing, write only this line: {"noActions":true}"#; /// The same limit DVLS and RDM use for Claude. Reasoning models count their reasoning in it, so it cannot be small. -const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 16_000; +pub(crate) const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 16_000; -/// One user action found in a session transcript. +/// One user action found in a session. #[derive(Debug, Clone, PartialEq)] pub struct Action { - /// Elapsed time since the start of the session, as reported by the model. + /// Elapsed time since the start of the session, taken from the input the model cited, never written by the model. pub offset: Duration, /// Short past-tense sentence naming the action, never empty. pub description: String, @@ -58,6 +58,7 @@ impl AiClient { /// /// Each line of `transcript` must start with the elapsed time in seconds between square brackets, such as /// `[12.5] ls`. + /// The model cites the line where each action starts, and the action gets the time of that line. pub fn describe_session_actions<'a>(&'a self, transcript: &'a str) -> DescribeSessionActions<'a> { DescribeSessionActions { client: self, @@ -99,30 +100,53 @@ impl DescribeSessionActions<'_> { /// Returns the actions in the order of the answer; the list is empty only when the model answered that the user did /// nothing. /// - /// Invalid lines in the answer are skipped with a warning. + /// Invalid lines in the answer are skipped with a warning, and so are actions citing a line that is not in the + /// transcript or has no time. /// The answer is [`Error::InvalidOutput`] when it has no valid action and does not say that the user did nothing, /// such as an empty answer. /// An answer cut short by the output token limit or the context window is [`Error::Truncated`]: send a shorter /// transcript instead. /// A refusal of the provider is [`Error::Refused`], never an empty list. pub async fn send(self) -> Result>, Error> { + let (labeled, times) = label_lines(self.transcript); + let prompt = Prompt { system: PROMPT, - input: self.transcript, + input: &[Input::Text(&labeled)], max_output_tokens: self.max_output_tokens, }; - self.client - .complete(&prompt) - .await? - .try_map(|answer| parse_actions(&answer)) + self.client.complete(&prompt).await?.try_map(|answer| { + parse_actions(&answer, "line", |line| { + let index = line.strip_prefix('L')?.parse::().ok()?; + times.get(index.checked_sub(1)?).copied().flatten() + }) + }) } } +/// Prefixes every transcript line with its id, `L1` for the first, and returns the times of the lines. +fn label_lines(transcript: &str) -> (String, Vec>) { + let mut labeled = String::with_capacity(transcript.len() + transcript.len() / 8); + let mut times = Vec::new(); + + for line in transcript.lines() { + times.push(line_time(line)); + labeled.push_str(&format!("L{} {line}\n", times.len())); + } + + (labeled, times) +} + +/// Reads the elapsed time at the start of a transcript line, such as `[12.5] ls`. +fn line_time(line: &str) -> Option { + let (time, _) = line.strip_prefix('[')?.split_once(']')?; + Duration::try_from_secs_f64(time.trim().parse().ok()?).ok() +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] struct ActionLine { - offset_seconds: f64, description: String, #[serde(default)] object: Option, @@ -130,7 +154,13 @@ struct ActionLine { parameters: BTreeMap, } -fn parse_actions(answer: &str) -> Result, Error> { +/// Reads JSON Lines actions that each cite a part of the input in `citation`, such as a transcript line; `time` gives +/// the time of a cited part, or `None` when the input has no such part. +pub(crate) fn parse_actions( + answer: &str, + citation: &str, + time: impl Fn(&str) -> Option, +) -> Result, Error> { let mut actions = Vec::new(); let mut invalid_lines = 0usize; let mut no_actions = false; @@ -147,7 +177,7 @@ fn parse_actions(answer: &str) -> Result, Error> { continue; } - match parse_action_line(line) { + match parse_action_line(line, citation, &time) { Ok(action) => actions.push(action), Err(reason) => { invalid_lines += 1; @@ -193,11 +223,17 @@ fn is_no_actions_line(line: &str) -> bool { } // The reason never quotes the line, because the line may contain session data. -fn parse_action_line(line: &str) -> Result { - let parsed: ActionLine = serde_json::from_str(line) +fn parse_action_line(line: &str, citation: &str, time: &impl Fn(&str) -> Option) -> Result { + let value: serde_json::Value = serde_json::from_str(line) .map_err(|error| format!("{:?} error at column {}", error.classify(), error.column()))?; - let offset = Duration::try_from_secs_f64(parsed.offset_seconds).map_err(|_| "invalid offsetSeconds".to_owned())?; + let cited = value + .get(citation) + .and_then(serde_json::Value::as_str) + .ok_or_else(|| format!("missing {citation}"))?; + let offset = time(cited.trim()).ok_or_else(|| format!("{citation} is not in the input"))?; + + let parsed: ActionLine = serde_json::from_value(value).map_err(|_| "invalid action fields".to_owned())?; let description = parsed.description.trim(); if description.is_empty() { @@ -214,26 +250,33 @@ fn parse_action_line(line: &str) -> Result { parameters: parsed.parameters, }) } - #[cfg(test)] mod tests { use super::*; + /// Parses an answer about a transcript whose line `Ln` is at `n` seconds, for lines 1 to 9. + fn parse(answer: &str) -> Result, Error> { + parse_actions(answer, "line", |line| { + let index = line.strip_prefix('L')?.parse::().ok()?; + (1..10).contains(&index).then(|| Duration::from_secs(index)) + }) + } + #[test] - fn parses_valid_lines() { + fn parses_valid_lines_with_the_time_of_the_cited_line() { let answer = concat!( - "{\"offsetSeconds\":1.5,\"description\":\"Listed files\",\"object\":\"/var/log\",\"parameters\":{\"Command\":\"ls\"}}\n", + "{\"line\":\"L1\",\"description\":\"Listed files\",\"object\":\"/var/log\",\"parameters\":{\"Command\":\"ls\"}}\n", "\n", - "{\"offsetSeconds\":3,\"description\":\"Opened a shell\"}\n", + "{\"line\":\"L3\",\"description\":\"Opened a shell\"}\n", ); - let actions = parse_actions(answer).expect("valid answer"); + let actions = parse(answer).expect("valid answer"); assert_eq!( actions, vec![ Action { - offset: Duration::from_millis(1500), + offset: Duration::from_secs(1), description: "Listed files".to_owned(), object: Some("/var/log".to_owned()), parameters: BTreeMap::from([("Command".to_owned(), "ls".to_owned())]), @@ -249,23 +292,35 @@ mod tests { } #[test] - fn skips_code_fences_and_invalid_lines() { + fn skips_code_fences_invalid_lines_and_unknown_citations() { let answer = concat!( "```jsonl\n", - "{\"offsetSeconds\":1,\"description\":\"Listed files\"}\n", + "{\"line\":\"L1\",\"description\":\"Listed files\"}\n", "Here are the actions:\n", - "{\"offsetSeconds\":-1,\"description\":\"Negative offset\"}\n", - "{\"offsetSeconds\":2,\"description\":\" \"}\n", - "{\"offsetSeconds\":2,\"description\":\"Numeric parameter\",\"parameters\":{\"Count\":5}}\n", + "{\"line\":\"L42\",\"description\":\"Line not in the transcript\"}\n", + "{\"offsetSeconds\":2,\"description\":\"No citation\"}\n", + "{\"line\":\"L2\",\"description\":\" \"}\n", + "{\"line\":\"L2\",\"description\":\"Numeric parameter\",\"parameters\":{\"Count\":5}}\n", "```\n", ); - let actions = parse_actions(answer).expect("one valid line"); + let actions = parse(answer).expect("one valid line"); assert_eq!(actions.len(), 1); assert_eq!(actions[0].description, "Listed files"); } + #[test] + fn lines_are_labeled_with_their_time() { + let (labeled, times) = label_lines("[1.5] ls\n[2] cd /tmp\nno time\n"); + + assert_eq!(labeled, "L1 [1.5] ls\nL2 [2] cd /tmp\nL3 no time\n"); + assert_eq!( + times, + [Some(Duration::from_millis(1500)), Some(Duration::from_secs(2)), None] + ); + } + #[test] fn no_actions_line_means_no_action() { for answer in [ @@ -273,14 +328,14 @@ mod tests { "\n```jsonl\n{\"noActions\":true}\n```\n", " { \"noActions\" : true } ", ] { - assert_eq!(parse_actions(answer).expect("no actions"), Vec::new(), "{answer:?}"); + assert_eq!(parse(answer).expect("no actions"), Vec::new(), "{answer:?}"); } } #[test] fn empty_answer_is_invalid_output() { for answer in ["", "\n```\n```\n"] { - let error = parse_actions(answer).expect_err("empty answer"); + let error = parse(answer).expect_err("empty answer"); assert!(matches!(error, Error::InvalidOutput { .. }), "{answer:?}: {error:?}"); } @@ -293,7 +348,7 @@ mod tests { "{\"noActions\":true,\"extra\":1}", "{\"noActions\":\"yes\"}", ] { - let error = parse_actions(line).expect_err("not the no-actions line"); + let error = parse(line).expect_err("not the no-actions line"); assert!(error.to_string().contains("1 invalid lines"), "{line}: {error}"); } @@ -301,9 +356,9 @@ mod tests { #[test] fn actions_win_over_the_no_actions_line() { - let answer = "{\"noActions\":true}\n{\"offsetSeconds\":1,\"description\":\"Listed files\"}\n"; + let answer = "{\"noActions\":true}\n{\"line\":\"L1\",\"description\":\"Listed files\"}\n"; - let actions = parse_actions(answer).expect("one action"); + let actions = parse(answer).expect("one action"); assert_eq!(actions.len(), 1); } @@ -322,7 +377,7 @@ mod tests { #[test] fn answer_without_valid_line_is_invalid_output() { - let error = parse_actions("not json\n{\"description\":\"no offset\"}\n").expect_err("no valid line"); + let error = parse("not json\n{\"description\":\"no line\"}\n").expect_err("no valid line"); assert!(matches!(error, Error::InvalidOutput { .. }), "{error:?}"); assert!(error.to_string().contains("2 invalid lines"), "{error}"); @@ -330,9 +385,12 @@ mod tests { #[test] fn invalid_line_reason_does_not_quote_the_line() { - let reason = - parse_action_line("{\"offsetSeconds\":1,\"description\":\"secret-value\",\"parameters\":{\"a\":1}}") - .expect_err("numeric parameter"); + let reason = parse_action_line( + "{\"line\":\"L1\",\"description\":\"secret-value\",\"parameters\":{\"a\":1}}", + "line", + &|_: &str| Some(Duration::ZERO), + ) + .expect_err("numeric parameter"); assert!(!reason.contains("secret-value")); } @@ -340,7 +398,7 @@ mod tests { #[test] fn prompt_asks_for_the_parsed_fields() { for field in [ - "offsetSeconds", + "\"line\"", "description", "object", "parameters", diff --git a/crates/devolutions-gateway-ai/src/wire/anthropic.rs b/crates/devolutions-gateway-ai/src/wire/anthropic.rs index ac3302197..41be5241e 100644 --- a/crates/devolutions-gateway-ai/src/wire/anthropic.rs +++ b/crates/devolutions-gateway-ai/src/wire/anthropic.rs @@ -4,9 +4,9 @@ use serde::{Deserialize, Serialize}; use tracing::warn; use url::Url; -use super::{Completion, Message, Stop, endpoint, parse_body, usage}; +use super::{Completion, Content, Message, Stop, base64, endpoint, parse_body, usage}; use crate::Error; -use crate::client::Prompt; +use crate::client::{Input, Prompt}; const VERSION: &str = "2023-06-01"; @@ -26,7 +26,13 @@ pub(crate) fn request( system: prompt.system, messages: [Message { role: "user", - content: prompt.input, + content: Content::of(prompt.input, |input| match input { + Input::Text(text) => serde_json::json!({ "type": "text", "text": text }), + Input::Png(png) => serde_json::json!({ + "type": "image", + "source": { "type": "base64", "media_type": "image/png", "data": base64(png) }, + }), + }), }], max_tokens: prompt.max_output_tokens, }) diff --git a/crates/devolutions-gateway-ai/src/wire/mod.rs b/crates/devolutions-gateway-ai/src/wire/mod.rs index b35c7b95f..7887e176e 100644 --- a/crates/devolutions-gateway-ai/src/wire/mod.rs +++ b/crates/devolutions-gateway-ai/src/wire/mod.rs @@ -8,6 +8,7 @@ use serde::de::DeserializeOwned; use serde_json::Value; use url::Url; +use crate::client::Input; use crate::{Error, Usage}; /// HTTP API spoken by a provider. @@ -43,7 +44,31 @@ pub(crate) enum Stop { #[derive(Serialize)] struct Message<'a> { role: &'static str, - content: &'a str, + content: Content<'a>, +} + +/// Content of a message: plain text, or blocks in the format of the API when the input has images. +#[derive(Serialize)] +#[serde(untagged)] +enum Content<'a> { + Text(&'a str), + Blocks(Vec), +} + +impl<'a> Content<'a> { + /// Builds the content of input, writing each part with lock. + fn of(input: &[Input<'a>], block: impl Fn(&Input<'a>) -> Value) -> Self { + match input { + [Input::Text(text)] => Self::Text(text), + input => Self::Blocks(input.iter().map(block).collect()), + } + } +} + +fn base64(data: &[u8]) -> String { + use base64::Engine as _; + + base64::engine::general_purpose::STANDARD.encode(data) } /// Message and code of an error answer, when its body has them. diff --git a/crates/devolutions-gateway-ai/src/wire/openai.rs b/crates/devolutions-gateway-ai/src/wire/openai.rs index d066bbf6d..4588f5786 100644 --- a/crates/devolutions-gateway-ai/src/wire/openai.rs +++ b/crates/devolutions-gateway-ai/src/wire/openai.rs @@ -5,9 +5,9 @@ use serde_json::Value; use tracing::warn; use url::Url; -use super::{Completion, Message, Stop, endpoint, parse_body, usage}; +use super::{Completion, Content, Message, Stop, base64, endpoint, parse_body, usage}; use crate::Error; -use crate::client::Prompt; +use crate::client::{Input, Prompt}; /// Role of the instructions for every server, OpenAI included. /// @@ -48,11 +48,17 @@ pub(crate) fn request( messages: [ Message { role: INSTRUCTIONS_ROLE, - content: prompt.system, + content: Content::Text(prompt.system), }, Message { role: "user", - content: prompt.input, + content: Content::of(prompt.input, |input| match input { + Input::Text(text) => serde_json::json!({ "type": "text", "text": text }), + Input::Png(png) => serde_json::json!({ + "type": "image_url", + "image_url": { "url": format!("data:image/png;base64,{}", base64(png)) }, + }), + }), }, ], max_completion_tokens, diff --git a/crates/devolutions-gateway-ai/tests/providers.rs b/crates/devolutions-gateway-ai/tests/providers.rs index e76eda94e..1c0926111 100644 --- a/crates/devolutions-gateway-ai/tests/providers.rs +++ b/crates/devolutions-gateway-ai/tests/providers.rs @@ -7,6 +7,7 @@ use std::net::Ipv4Addr; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; +use devolutions_gateway_ai::screen_actions::Screenshot; use devolutions_gateway_ai::session_actions::Action; use devolutions_gateway_ai::{AiClient, Error, Provider, Usage}; use reqwest::StatusCode; @@ -21,7 +22,9 @@ const USAGE: Usage = Usage { input_tokens: 10, output_tokens: 20, }; -const ANSWER: &str = "{\"offsetSeconds\":1.5,\"description\":\"Listed files\",\"object\":\"/var/log\",\"parameters\":{\"Command\":\"ls\"}}\nnot an action\n{\"offsetSeconds\":4,\"description\":\"Opened a shell\"}"; +const ANSWER: &str = "{\"line\":\"L1\",\"description\":\"Listed files\",\"object\":\"/var/log\",\"parameters\":{\"Command\":\"ls\"}}\nnot an action\n{\"line\":\"L2\",\"description\":\"Opened a shell\"}"; +/// Transcript the requests send; its lines are at 1.5 s and 4 s. +const TRANSCRIPT: &str = "[1.5] ls\n[4] bash\n"; const NO_ACTIONS: &str = "{\"noActions\":true}"; /// Largest answer body the client reads, the same as in the crate. const MAX_RESPONSE_BYTES: usize = 8 * 1024 * 1024; @@ -207,7 +210,7 @@ async fn openai_chat_request_and_response() { let (base_url, captured) = spawn_provider(StatusCode::OK, openai_response(ANSWER)).await; let response = client(Provider::OpenAi, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .max_output_tokens(1234) .send() .await @@ -225,7 +228,10 @@ async fn openai_chat_request_and_response() { ); assert_eq!(header(&request.headers, "x-api-key"), None); assert_eq!(request.body["model"], MODEL); - assert!(request.body.to_string().contains("[1.5] ls")); + assert!( + request.body.to_string().contains("L1 [1.5] ls"), + "the transcript lines are labeled" + ); assert_eq!(request.body["max_completion_tokens"], 1234); assert_eq!(request.body["store"], false); } @@ -235,7 +241,7 @@ async fn anthropic_messages_request_and_response() { let (base_url, captured) = spawn_provider(StatusCode::OK, anthropic_response(ANSWER)).await; let response = client(Provider::Anthropic, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .max_output_tokens(1234) .send() .await @@ -252,7 +258,10 @@ async fn anthropic_messages_request_and_response() { assert_eq!(header(&request.headers, "authorization"), None); assert_eq!(request.body["model"], MODEL); assert_eq!(request.body["max_tokens"], 1234); - assert!(request.body.to_string().contains("[1.5] ls")); + assert!( + request.body.to_string().contains("L1 [1.5] ls"), + "the transcript lines are labeled" + ); } #[tokio::test] @@ -260,7 +269,7 @@ async fn openai_compatible_uses_the_given_base_url_and_max_tokens() { let (base_url, captured) = spawn_provider(StatusCode::OK, openai_response(ANSWER)).await; let response = client(Provider::OpenAiCompatible, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap(); @@ -283,7 +292,7 @@ async fn gemini_speaks_openai_chat_with_max_tokens() { let (base_url, captured) = spawn_provider(StatusCode::OK, openai_response(ANSWER)).await; let response = client(Provider::Gemini, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .max_output_tokens(1234) .send() .await @@ -318,7 +327,7 @@ async fn mistral_chunked_content_is_read_as_text() { let (base_url, captured) = spawn_provider(StatusCode::OK, answer).await; let response = client(Provider::Mistral, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .max_output_tokens(1234) .send() .await @@ -346,7 +355,7 @@ async fn answer_without_model_or_usage_is_accepted() { let (base_url, _captured) = spawn_provider(StatusCode::OK, answer).await; let response = client(Provider::OpenAiCompatible, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap(); @@ -362,7 +371,7 @@ async fn think_blocks_are_not_read_as_actions() { let (base_url, _captured) = spawn_provider(StatusCode::OK, openai_response(&answer)).await; let response = client(Provider::OpenAiCompatible, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap(); @@ -375,7 +384,7 @@ async fn think_tags_inside_an_action_are_kept() { const COMMAND: &str = "grep 'x' notes.txt"; let action = serde_json::json!({ - "offsetSeconds": 2, + "line": "L1", "description": "Searched notes", "parameters": { "Command": COMMAND } }); @@ -417,7 +426,7 @@ async fn answers_cut_at_the_token_limit_are_truncated() { let (base_url, _captured) = spawn_provider(StatusCode::OK, response).await; let error = client(provider, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap_err(); @@ -450,7 +459,7 @@ async fn refusals_are_permanent_errors() { let (base_url, _captured) = spawn_provider(StatusCode::OK, response).await; let error = client(provider, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap_err(); @@ -468,7 +477,7 @@ async fn provider_failure_while_answering_is_invalid_response() { let (base_url, _captured) = spawn_provider(StatusCode::OK, answer).await; let error = client(Provider::Mistral, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap_err(); @@ -506,7 +515,7 @@ async fn unexpected_answers_are_invalid_responses() { let (base_url, _captured) = spawn_provider(StatusCode::OK, response).await; let error = client(provider, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap_err(); @@ -565,7 +574,7 @@ async fn answer_without_any_action_line_is_invalid_output() { let (base_url, _captured) = spawn_provider(StatusCode::OK, openai_response("Sorry, I cannot help.")).await; let error = client(Provider::OpenAi, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap_err(); @@ -726,7 +735,7 @@ async fn instructions_are_a_system_message() { let (base_url, captured) = spawn_provider(StatusCode::OK, openai_response(ANSWER)).await; client(provider, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap(); @@ -742,7 +751,7 @@ async fn instructions_are_a_system_message() { "{provider:?}" ); assert_eq!(messages[1]["role"], "user", "{provider:?}"); - assert_eq!(messages[1]["content"], "[1.5] ls", "{provider:?}"); + assert_eq!(messages[1]["content"], "L1 [1.5] ls\nL2 [4] bash\n", "{provider:?}"); } } @@ -778,7 +787,7 @@ async fn answer_up_to_the_size_limit_is_read() { let (base_url, _captured) = spawn_mock(answer).await; let response = client(Provider::OpenAi, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap(); @@ -800,7 +809,7 @@ async fn answer_over_the_size_limit_is_invalid_response() { let (base_url, _captured) = spawn_mock(answer).await; let error = client(Provider::OpenAi, base_url) - .describe_session_actions("[1.5] ls") + .describe_session_actions(TRANSCRIPT) .send() .await .unwrap_err(); @@ -920,3 +929,74 @@ async fn silent_provider_times_out_as_transient() { assert!(error.is_transient()); assert!(started.elapsed() < Duration::from_secs(5), "{:?}", started.elapsed()); } + +#[tokio::test] +async fn screenshots_are_sent_as_images_and_actions_take_their_time() { + const PNG: &[u8] = b"\x89PNG not really"; + + let screenshots = [ + Screenshot { + id: "F00001", + offset: Duration::from_secs(10), + png: PNG, + context: true, + }, + Screenshot { + id: "F00002", + offset: Duration::from_millis(12_500), + png: PNG, + context: false, + }, + ]; + let answer = "{\"frame\":\"F00002\",\"description\":\"Opened a dialog\"}\n{\"frame\":\"F00001\",\"description\":\"Only in context\"}"; + + for (provider, response) in [ + (Provider::OpenAi, openai_response(answer)), + (Provider::Anthropic, anthropic_response(answer)), + ] { + let (base_url, captured) = spawn_provider(StatusCode::OK, response).await; + + let response = client(provider, base_url) + .describe_screen_actions(&screenshots) + .send() + .await + .unwrap(); + + assert_eq!( + response.output.len(), + 1, + "a context screenshot is never cited: {provider:?}" + ); + assert_eq!(response.output[0].description, "Opened a dialog"); + assert_eq!(response.output[0].offset, Duration::from_millis(12_500)); + + let request = take(&captured); + let content = match provider { + Provider::Anthropic => &request.body["messages"][0]["content"], + _ => &request.body["messages"][1]["content"], + }; + let blocks = content.as_array().unwrap(); + assert_eq!(blocks.len(), 4, "{provider:?}"); + assert_eq!(blocks[0]["text"], "[F00001 t=10.0s context]"); + assert_eq!(blocks[2]["text"], "[F00002 t=12.5s]"); + + match provider { + Provider::Anthropic => { + assert_eq!(blocks[1]["type"], "image"); + assert_eq!(blocks[1]["source"]["media_type"], "image/png"); + assert!(blocks[1]["source"]["data"].as_str().unwrap().starts_with("iVBOR")); + } + _ => { + assert_eq!(blocks[1]["type"], "image_url"); + assert!( + blocks[1]["image_url"]["url"] + .as_str() + .unwrap() + .starts_with("data:image/png;base64,iVBOR"), + "{}", + blocks[1] + ); + } + } + } +} From 5e51eee36d28c4f27018351b97bd8f418bea0cfb Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 7 Oct 2026 18:01:08 -0400 Subject: [PATCH 03/18] feat(dgw): analyze any recording in the AI crate `AiClient::analyze_recording` takes the parsed session manifest and its path and does the rest. A session is recorded either as a terminal or as a video, and the types hold it: `Recordings` is all terminal or all video, and each recording is accepted only with the extension of its kind (`Recordings::from_files` refuses a manifest that mixes them or has none). - terminal recordings (asciicast and TRP) are read as a stream into timed transcript chunks, described with `describe_session_actions`; - video recordings (WebM, VP8 or VP9) are decoded with XMF and sampled twice a second; a screenshot is kept when the screen changed, once it settles, and every 3 seconds during a long change (16x16 tiles of a half-resolution gray image; tiles that keep changing, such as a clock, are ignored). The screenshots are sent at most 1280x720, 40 per request, with the previous one as context, and described with `describe_screen_actions`. A truncated answer is asked again in two halves. Each chunk's answer is saved as a checkpoint, so analysing again after a failure only asks about the chunks left. The actions of all chunks are merged in time order into a Session Recording Log, left next to the manifest. The working files live in `.ai-analysis/` next to the manifest; files made with another model or prompt are not reused. A successful analysis removes them, and `recording_analysis::discard` removes them otherwise. The crate never writes `recording.json`. `start` runs the analysis on the Tokio runtime and returns a channel of events: `Progress` (reading the recordings, with the bytes read; asking the AI, chunk by chunk; writing the log), then `Complete` or `Failure`, last. Each `Progress` has a stable stage name, a percent and a short sentence, to be stored with a task or shown to whoever asks about it. Dropping the receiver cancels the analysis. The `analyze_recording` example runs it against a real provider and shows the progress: `cargo run -p devolutions-gateway-ai --example analyze_recording -- --recording-path `. Co-Authored-By: Claude Opus 5.5 (1M context) --- Cargo.lock | 31 + crates/devolutions-gateway-ai/Cargo.toml | 12 + .../examples/analyze_recording.rs | 396 +++++++ crates/devolutions-gateway-ai/src/client.rs | 5 + crates/devolutions-gateway-ai/src/lib.rs | 5 + .../src/recording_analysis/checkpoint.rs | 200 ++++ .../src/recording_analysis/mod.rs | 1020 +++++++++++++++++ .../src/recording_analysis/recordings.rs | 250 ++++ .../src/recording_analysis/slog.rs | 174 +++ .../src/recording_analysis/tests.rs | 501 ++++++++ .../src/recording_analysis/transcript.rs | 665 +++++++++++ .../src/recording_analysis/video.rs | 697 +++++++++++ 12 files changed, 3956 insertions(+) create mode 100644 crates/devolutions-gateway-ai/examples/analyze_recording.rs create mode 100644 crates/devolutions-gateway-ai/src/recording_analysis/checkpoint.rs create mode 100644 crates/devolutions-gateway-ai/src/recording_analysis/mod.rs create mode 100644 crates/devolutions-gateway-ai/src/recording_analysis/recordings.rs create mode 100644 crates/devolutions-gateway-ai/src/recording_analysis/slog.rs create mode 100644 crates/devolutions-gateway-ai/src/recording_analysis/tests.rs create mode 100644 crates/devolutions-gateway-ai/src/recording_analysis/transcript.rs create mode 100644 crates/devolutions-gateway-ai/src/recording_analysis/video.rs diff --git a/Cargo.lock b/Cargo.lock index da4f7d7c0..73f9702c8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1937,15 +1937,24 @@ dependencies = [ name = "devolutions-gateway-ai" version = "0.0.0" dependencies = [ + "anyhow", "base64 0.23.1", + "cadeau", + "camino", + "png", "reqwest", "secrecy", "serde", "serde_json", + "tempfile", + "terminal-streamer", "thiserror 2.0.20", + "time", "tokio 1.52.3", + "tokio-util", "tracing", "url", + "webm-iterable", ] [[package]] @@ -2454,6 +2463,15 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +[[package]] +name = "fdeflate" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c" +dependencies = [ + "simd-adler32", +] + [[package]] name = "ff" version = "0.14.0" @@ -5617,6 +5635,19 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "png" +version = "0.17.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82151a2fc869e011c153adc57cf2789ccb8d9906ce52c0b39a6b5697749d7526" +dependencies = [ + "bitflags 1.3.2", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide", +] + [[package]] name = "polling" version = "2.8.0" diff --git a/crates/devolutions-gateway-ai/Cargo.toml b/crates/devolutions-gateway-ai/Cargo.toml index 91a268069..77f7fb1da 100644 --- a/crates/devolutions-gateway-ai/Cargo.toml +++ b/crates/devolutions-gateway-ai/Cargo.toml @@ -10,14 +10,26 @@ publish = false workspace = true [dependencies] +anyhow = "1" base64 = "0.23" +cadeau = { version = "0.5.3", features = ["dlopen"] } +camino = "1.1" +png = "0.17" reqwest = { version = "0.12", default-features = false, features = ["json"] } secrecy = "0.10" serde = { version = "1", features = ["derive"] } serde_json = "1" +terminal-streamer.path = "../terminal-streamer" thiserror = "2" +time = "0.3" +tokio = { version = "1.52", features = ["fs", "macros", "rt", "time"] } +tokio-util = "0.7" tracing = "0.1" url = "2.5" +webm-iterable = "0.7" [dev-dependencies] +# The example builds its own HTTP client, so it needs TLS; Gateway passes one with its own TLS settings. +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots"] } +tempfile = "3.24" tokio = { version = "1.52", features = ["io-util", "macros", "net", "rt-multi-thread"] } diff --git a/crates/devolutions-gateway-ai/examples/analyze_recording.rs b/crates/devolutions-gateway-ai/examples/analyze_recording.rs new file mode 100644 index 000000000..94bf8b1d1 --- /dev/null +++ b/crates/devolutions-gateway-ai/examples/analyze_recording.rs @@ -0,0 +1,396 @@ +//! Analyzes a recorded session with a real AI provider and prints the log. +//! +//! ```text +//! cargo run -p devolutions-gateway-ai --example analyze_recording -- --recording-path [options] +//! ``` +//! +//! `--recording-path` is a session folder, its `recording.json`, or one recording file (`.cast`, `.trp`, `.webm`). +//! For a single file next to a `recording.json`, the manifest gives its times; without one, the session starts with it. +//! +//! Options: +//! - `--provider `: `openai` by default. +//! - `--model `: `gpt-6-luna` for OpenAI and `claude-sonnet-5-5` for Anthropic by default; required otherwise. +//! - `--base-url `: required for `openai-compatible`. +//! - `--api-key-env `: environment variable holding the API key; `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`, +//! `MISTRAL_API_KEY`, `GEMINI_API_KEY` or `AI_API_KEY` by default. +//! - `--xmf `: XMF library, needed for video; `DGATEWAY_LIB_XMF_PATH` by default. +//! - `--max-output-tokens `. +//! +//! The working files and the log are left next to the manifest, in `.ai-analysis/` and `.ai-analysis.slog`. + +#![expect( + clippy::print_stdout, + reason = "the example reports its progress and the log on the console" +)] + +use std::io::Write as _; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use anyhow::{Context as _, bail}; +use camino::{Utf8Path, Utf8PathBuf}; +use devolutions_gateway_ai::recording_analysis::{ + Progress, RecordingAnalysisEvent, RecordingFile, RecordingManifest, Recordings, +}; +use devolutions_gateway_ai::{AiClient, Provider}; + +const USAGE: &str = "usage: analyze_recording --recording-path \ +[--provider openai|anthropic|mistral|gemini|openai-compatible] [--model ] [--base-url ] \ +[--api-key-env ] [--xmf ] [--max-output-tokens ]"; + +#[derive(Default)] +struct Args { + recording_path: Option, + provider: Option, + model: Option, + base_url: Option, + api_key_env: Option, + xmf: Option, + max_output_tokens: Option, +} + +fn parse_args() -> anyhow::Result { + let mut args = Args::default(); + let mut raw = std::env::args().skip(1); + + while let Some(flag) = raw.next() { + let mut value = || raw.next().with_context(|| format!("{flag} needs a value\n{USAGE}")); + + match flag.as_str() { + "--recording-path" => args.recording_path = Some(value()?.into()), + "--provider" => args.provider = Some(value()?), + "--model" => args.model = Some(value()?), + "--base-url" => args.base_url = Some(value()?), + "--api-key-env" => args.api_key_env = Some(value()?), + "--xmf" => args.xmf = Some(value()?), + "--max-output-tokens" => args.max_output_tokens = Some(value()?.parse()?), + "--help" | "-h" => { + println!("{USAGE}"); + std::process::exit(0); + } + other => bail!("unknown argument {other}\n{USAGE}"), + } + } + + Ok(args) +} + +/// Finds the manifest and the recordings to analyze from a session folder, a `recording.json`, or one recording file. +fn manifest_for(path: &Utf8Path) -> anyhow::Result<(Utf8PathBuf, RecordingManifest)> { + let (manifest_path, only) = if path.is_dir() { + (path.join("recording.json"), None) + } else if path.extension() == Some("json") { + (path.to_owned(), None) + } else { + let dir = match path.parent() { + Some(dir) if !dir.as_str().is_empty() => dir, + _ => Utf8Path::new("."), + }; + let file_name = path.file_name().context("recording path has no file name")?; + (dir.join("recording.json"), Some(file_name.to_owned())) + }; + + if !manifest_path.exists() { + let file_name = only.with_context(|| format!("{manifest_path} not found"))?; + let recordings = Recordings::from_files([RecordingFile { + file_name, + start_time: 0, + }])?; + + // Without a manifest, the session starts with the recording; its length is unknown. + return Ok(( + manifest_path, + RecordingManifest { + start_time: 0, + duration: 0, + recordings, + }, + )); + } + + let json: serde_json::Value = + serde_json::from_slice(&std::fs::read(&manifest_path).with_context(|| format!("read {manifest_path}"))?) + .with_context(|| format!("parse {manifest_path}"))?; + let number = |value: &serde_json::Value, name: &str| value[name].as_i64().with_context(|| format!("no {name}")); + + let files = json["files"] + .as_array() + .context("manifest has no files")? + .iter() + .map(|file| { + Ok(RecordingFile { + file_name: file["fileName"].as_str().context("no fileName")?.to_owned(), + start_time: number(file, "startTime")?, + }) + }) + .collect::>>()?; + + let files = match &only { + Some(name) => files.into_iter().filter(|file| &file.file_name == name).collect(), + None => files, + }; + + Ok(( + manifest_path, + RecordingManifest { + start_time: number(&json, "startTime")?, + duration: number(&json, "duration")?, + recordings: Recordings::from_files(files)?, + }, + )) +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let args = parse_args()?; + let recording_path = args + .recording_path + .with_context(|| format!("--recording-path is required\n{USAGE}"))?; + + let (provider, default_model, default_key_env) = match args.provider.as_deref().unwrap_or("openai") { + "openai" => (Provider::OpenAi, Some("gpt-6-luna"), "OPENAI_API_KEY"), + "anthropic" => (Provider::Anthropic, Some("claude-sonnet-5-5"), "ANTHROPIC_API_KEY"), + "mistral" => (Provider::Mistral, None, "MISTRAL_API_KEY"), + "gemini" => (Provider::Gemini, None, "GEMINI_API_KEY"), + "openai-compatible" => (Provider::OpenAiCompatible, None, "AI_API_KEY"), + other => bail!("unknown provider {other}\n{USAGE}"), + }; + + let model = args + .model + .or_else(|| default_model.map(str::to_owned)) + .with_context(|| format!("--model is required for this provider\n{USAGE}"))?; + let model_label = model.clone(); + let key_env = args.api_key_env.unwrap_or_else(|| default_key_env.to_owned()); + let api_key = std::env::var(&key_env).with_context(|| format!("{key_env} is not set"))?; + + if let Some(xmf) = args.xmf.or_else(|| std::env::var("DGATEWAY_LIB_XMF_PATH").ok()) { + // SAFETY: The library is XMF, loaded once before any other XMF call. + unsafe { cadeau::xmf::init(&xmf) }.with_context(|| format!("load XMF from {xmf}"))?; + } + + let mut builder = AiClient::builder() + .provider(provider) + .model(model) + .api_key(api_key) + .http_client(reqwest::Client::new()); + if let Some(base_url) = args.base_url { + builder = builder.base_url(base_url.parse().context("invalid --base-url")?); + } + let client = builder.build()?; + + let (manifest_path, manifest) = manifest_for(&recording_path)?; + let kind = match &manifest.recordings { + Recordings::Terminal(recordings) => format!("{} terminal recording(s)", recordings.len()), + Recordings::Video(recordings) => format!("{} video recording(s)", recordings.len()), + }; + + println!(); + println!(" {}", style(BOLD, "AI analysis")); + println!(" {} {manifest_path}", style(DIM, "session ")); + println!(" {} {kind}", style(DIM, "input ")); + println!( + " {} {} / {}", + style(DIM, "model "), + provider_name(provider), + model_label + ); + println!(); + + let started = Instant::now(); + let mut request = client.analyze_recording(&manifest, &manifest_path); + if let Some(max_output_tokens) = args.max_output_tokens { + request = request.max_output_tokens(max_output_tokens); + } + let mut events = request.start(); + + let current = Arc::new(Mutex::new(None)); + let ticker = tokio::spawn(show_progress(Arc::clone(¤t), started)); + + let result = loop { + match events.recv().await { + Some(RecordingAnalysisEvent::Progress(progress)) => { + *current.lock().expect("not poisoned") = Some(progress); + } + Some(RecordingAnalysisEvent::Complete(analysis)) => break Ok(analysis), + Some(RecordingAnalysisEvent::Failure(error)) => break Err(error.to_string()), + Some(_) => {} + None => break Err("the analysis stopped unexpectedly".to_owned()), + } + }; + ticker.abort(); + clear_line(); + + let analysis = match result { + Ok(analysis) => analysis, + Err(error) => { + println!( + " {} after {:.1}s: {error}", + style(RED, "✗ failed"), + started.elapsed().as_secs_f64() + ); + println!( + " {}", + style(DIM, "the working files are kept: run again to resume where it stopped") + ); + std::process::exit(1); + } + }; + println!( + " {} {} actions in {:.1}s", + style(GREEN, "✓"), + analysis.actions, + started.elapsed().as_secs_f64() + ); + println!(); + print_actions(&std::fs::read_to_string(&analysis.log_path)?)?; + println!(); + + let usage = analysis.usage.map_or_else( + || "not reported".to_owned(), + |usage| format!("{} in, {} out", usage.input_tokens, usage.output_tokens), + ); + println!(" {} {}", style(DIM, "model "), analysis.model); + println!(" {} {}", style(DIM, "prompt "), analysis.prompt_version); + println!(" {} {usage}", style(DIM, "tokens ")); + println!(" {} {}", style(DIM, "log "), analysis.log_path); + println!(); + + Ok(()) +} + +/// Redraws the progress line ten times a second, from the last progress event, until aborted. +async fn show_progress(current: Arc>>, started: Instant) { + const SPINNER: [char; 10] = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']; + const BAR: usize = 24; + + for tick in 0usize.. { + let text = match *current.lock().expect("not poisoned") { + None => "starting".to_owned(), + Some(progress @ Progress::Writing) => progress.to_string(), + Some(progress) => { + let filled = usize::from(progress.percent()) * BAR / 100; + format!("[{}{}] {progress}", "█".repeat(filled), "░".repeat(BAR - filled)) + } + }; + + clear_line(); + print!( + " {} {} {text}", + style(CYAN, &SPINNER[tick % SPINNER.len()].to_string()), + style(DIM, &format!("{:>6.1}s", started.elapsed().as_secs_f64())) + ); + let _ = std::io::stdout().flush(); + + tokio::time::sleep(Duration::from_millis(100)).await; + } +} + +fn clear_line() { + print!("\r\x1b[2K"); + let _ = std::io::stdout().flush(); +} + +/// Prints the actions of a `.slog`: their time since the session started, then their object and details. +fn print_actions(slog: &str) -> anyhow::Result<()> { + let entries = slog + .lines() + .map(serde_json::from_str::) + .collect::, _>>()?; + + let session_start = entries + .first() + .and_then(|entry| entry["timestamp"].as_str()) + .and_then(unix_millis) + .context("log has no start")?; + + let actions = entries + .iter() + .filter(|entry| entry["event"] == "session.action") + .collect::>(); + + if actions.is_empty() { + println!(" {}", style(DIM, "the AI found no user action")); + return Ok(()); + } + + for action in actions { + let at = action["timestamp"] + .as_str() + .and_then(unix_millis) + .unwrap_or(session_start); + let seconds = (at - session_start).max(0) / 1000; + let time = format!("{:02}:{:02}", seconds / 60, seconds % 60); + let description = action["description"].as_str().unwrap_or_default(); + + println!(" {} {}", style(CYAN, &time), style(BOLD, description)); + + if let Some(object) = action["object"].as_str() { + println!(" {} {object}", style(DIM, "→")); + } + + if let Some(parameters) = action["parameters"].as_object() { + for (name, value) in parameters { + println!( + " {} {}", + style(DIM, &format!("{name}:")), + value.as_str().unwrap_or_default() + ); + } + } + } + + Ok(()) +} + +/// Reads a `.slog` timestamp, such as `2026-10-02T16:41:48.001Z`, as milliseconds since the Unix epoch. +fn unix_millis(timestamp: &str) -> Option { + let (date, time) = timestamp.strip_suffix('Z')?.split_once('T')?; + let mut date = date.splitn(3, '-').map(str::parse::); + let (year, month, day) = (date.next()?.ok()?, date.next()?.ok()?, date.next()?.ok()?); + let mut time = time.splitn(3, ':'); + let (hours, minutes) = (time.next()?.parse::().ok()?, time.next()?.parse::().ok()?); + let seconds = time.next()?.parse::().ok()?; + + // Days since 1970-01-01 for a proleptic Gregorian date. + let (year, month) = if month <= 2 { + (year - 1, month + 9) + } else { + (year, month - 3) + }; + let era = year.div_euclid(400); + let year_of_era = year - era * 400; + let day_of_year = (153 * month + 2) / 5 + day - 1; + let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + let days = era * 146_097 + day_of_era - 719_468; + + #[expect(clippy::cast_possible_truncation, reason = "milliseconds of one minute fit in i64")] + let millis = (seconds * 1000.0).round() as i64; + + Some(((days * 24 + hours) * 60 + minutes) * 60_000 + millis) +} + +const BOLD: &str = "\x1b[1m"; +const DIM: &str = "\x1b[2m"; +const GREEN: &str = "\x1b[32m"; +const RED: &str = "\x1b[31m"; +const CYAN: &str = "\x1b[36m"; + +/// Wraps `text` in an ANSI style, unless `NO_COLOR` is set. +fn style(code: &str, text: &str) -> String { + if std::env::var_os("NO_COLOR").is_some() { + text.to_owned() + } else { + format!("{code}{text}\x1b[0m") + } +} + +fn provider_name(provider: Provider) -> &'static str { + match provider { + Provider::OpenAi => "OpenAI", + Provider::Anthropic => "Anthropic", + Provider::Mistral => "Mistral", + Provider::Gemini => "Gemini", + Provider::OpenAiCompatible => "OpenAI-compatible", + } +} diff --git a/crates/devolutions-gateway-ai/src/client.rs b/crates/devolutions-gateway-ai/src/client.rs index c58c7998c..7316d2fe0 100644 --- a/crates/devolutions-gateway-ai/src/client.rs +++ b/crates/devolutions-gateway-ai/src/client.rs @@ -226,6 +226,11 @@ impl AiClient { AiClientBuilder::default() } + /// Model requested from the provider. + pub(crate) fn model(&self) -> &str { + &self.model + } + /// Sends one completion request and returns the text of the answer. /// /// An answer cut short by the output token limit or the context window is [`Error::Truncated`], because no purpose diff --git a/crates/devolutions-gateway-ai/src/lib.rs b/crates/devolutions-gateway-ai/src/lib.rs index f92c3283e..7ccbf9a13 100644 --- a/crates/devolutions-gateway-ai/src/lib.rs +++ b/crates/devolutions-gateway-ai/src/lib.rs @@ -27,9 +27,14 @@ //! //! A new purpose is a module like [`session_actions`]: its prompt and `PROMPT_VERSION`, a request builder returned by //! a new [`AiClient`] method, and the parser of the answer. +//! +//! A purpose may also build on others and on files: [`recording_analysis`] reads the recordings of a session, terminal +//! or video, describes them chunk by chunk with [`session_actions`] and [`screen_actions`], and writes the result as a +//! session log, keeping its working files next to the recordings. mod client; mod error; +pub mod recording_analysis; mod response; pub mod screen_actions; pub mod session_actions; diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/checkpoint.rs b/crates/devolutions-gateway-ai/src/recording_analysis/checkpoint.rs new file mode 100644 index 000000000..e5a3a4c4e --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/checkpoint.rs @@ -0,0 +1,200 @@ +//! What the AI answered for one transcript chunk, saved in the working files so analysing again does not ask the AI again. + +use std::collections::BTreeMap; +use std::io::{BufReader, BufWriter}; +use std::time::Duration; + +use anyhow::Context as _; +use camino::{Utf8Path, Utf8PathBuf}; +use serde::{Deserialize, Serialize}; + +use crate::session_actions::Action; +use crate::{Response, Usage}; + +/// Tokens counted by the AI provider, as kept in checkpoints. +/// +/// It mirrors [`Usage`], so a change of [`Usage`] never changes the checkpoint format. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct TokenUsage { + input_tokens: u64, + output_tokens: u64, +} + +impl From for TokenUsage { + fn from(usage: Usage) -> Self { + Self { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + } + } +} + +impl From for Usage { + fn from(usage: TokenUsage) -> Self { + Self { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + } + } +} + +pub(super) fn path(workspace: &Utf8Path, index: usize) -> Utf8PathBuf { + workspace.join(format!("chunk-{index:04}.json")) +} + +/// Actions found in one chunk, with what the AI provider reported for the requests about it. +#[derive(Debug, Clone, PartialEq)] +pub(super) struct DescribedChunk { + pub(super) actions: Vec, + /// Model that answered, as reported by the provider. + pub(super) model: Option, + /// Tokens of every request about the chunk, answers cut and asked again included; `None` when one was not reported. + pub(super) usage: Option, +} + +impl Default for DescribedChunk { + /// No action yet, and no token counted yet. + fn default() -> Self { + Self { + actions: Vec::new(), + model: None, + usage: Some(Usage::default()), + } + } +} + +impl DescribedChunk { + /// Adds the answer to one request about the chunk. + pub(super) fn add(&mut self, response: Response>) { + self.actions.extend(response.output); + self.model = self.model.take().or(response.model); + self.usage = super::add_usage(self.usage, response.usage); + } +} + +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct SavedChunk { + #[serde(default, skip_serializing_if = "Option::is_none")] + model: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + usage: Option, + actions: Vec, +} + +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct SavedAction { + offset_seconds: f64, + description: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + object: Option, + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + parameters: BTreeMap, +} + +/// Writes the checkpoint through a temporary file, so a crash never leaves a partial one. +pub(super) fn write(path: &Utf8Path, chunk: &DescribedChunk) -> anyhow::Result<()> { + let partial = path.with_extension("partial"); + + let saved = SavedChunk { + model: chunk.model.clone(), + usage: chunk.usage.map(TokenUsage::from), + actions: chunk + .actions + .iter() + .map(|action| SavedAction { + offset_seconds: action.offset.as_secs_f64(), + description: action.description.clone(), + object: action.object.clone(), + parameters: action.parameters.clone(), + }) + .collect(), + }; + + let mut out = BufWriter::new(std::fs::File::create(&partial).with_context(|| format!("create {partial}"))?); + serde_json::to_writer(&mut out, &saved)?; + out.into_inner()?.sync_all()?; + std::fs::rename(&partial, path).with_context(|| format!("rename {partial}"))?; + + Ok(()) +} + +pub(super) fn read(path: &Utf8Path) -> anyhow::Result { + let file = BufReader::new(std::fs::File::open(path).with_context(|| format!("open {path}"))?); + let saved: SavedChunk = serde_json::from_reader(file).with_context(|| format!("read {path}"))?; + + let actions = saved + .actions + .into_iter() + .map(|saved| { + Ok(Action { + offset: Duration::try_from_secs_f64(saved.offset_seconds)?, + description: saved.description, + object: saved.object, + parameters: saved.parameters, + }) + }) + .collect::>()?; + + Ok(DescribedChunk { + actions, + model: saved.model, + usage: saved.usage.map(Usage::from), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn described_chunk_round_trips() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = path(Utf8Path::from_path(dir.path()).expect("UTF-8"), 3); + let chunk = DescribedChunk { + actions: vec![ + Action { + offset: Duration::from_millis(1500), + description: "Listed files".to_owned(), + object: Some("/var/log".to_owned()), + parameters: BTreeMap::from([("Command".to_owned(), "ls".to_owned())]), + }, + Action { + offset: Duration::from_secs(62), + description: "Closed the shell".to_owned(), + object: None, + parameters: BTreeMap::new(), + }, + ], + model: Some("gpt-test-2026-09-30".to_owned()), + usage: Some(Usage { + input_tokens: 10, + output_tokens: 20, + }), + }; + + write(&path, &chunk).expect("written"); + + assert!(path.as_str().ends_with("chunk-0003.json")); + assert_eq!(read(&path).expect("read"), chunk); + assert!(!path.with_extension("partial").exists()); + } + + #[test] + fn unreported_model_and_usage_stay_unknown() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = path(Utf8Path::from_path(dir.path()).expect("UTF-8"), 0); + let chunk = DescribedChunk { + actions: Vec::new(), + model: None, + usage: None, + }; + + write(&path, &chunk).expect("written"); + + assert_eq!(std::fs::read_to_string(&path).expect("checkpoint"), r#"{"actions":[]}"#); + assert_eq!(read(&path).expect("read"), chunk); + } +} diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs new file mode 100644 index 000000000..b0ba21126 --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -0,0 +1,1020 @@ +//! Purpose: describe what the user did in a recorded session, with [`AiClient::analyze_recording`]. +//! +//! A session is recorded as a terminal or as a video, never both, see [`Recordings`]: +//! - terminal recordings (asciicast and TRP) become a timed transcript, see [`transcript`], described with +//! [`session_actions`](crate::session_actions); +//! - video recordings (WebM) become the screenshots where the screen changed, see [`video`], described with +//! [`screen_actions`](crate::screen_actions). +//! +//! The model cites the transcript line or the screenshot that shows each action, and the action gets its time: the +//! model never writes times. +//! +//! The input is cut into chunks small enough for one request, and the actions found in every chunk are merged into a +//! Session Recording Log (`.slog`). The work keeps its files next to the manifest, in `.ai-analysis/`: +//! 1. `analysis.json`: the model, the prompt versions and the chunk sizes the other files were made with; +//! 2. the transcript chunks `transcript-NNNN.txt` and the screenshots `FNNNNN.png`, then `plan.json`, the list of +//! chunks, once they are all written; +//! 3. what the AI answered for each chunk, `chunk-NNNN.json`, see [`checkpoint`], so that analysing the recording again +//! after a failure only asks about the chunks left; +//! 4. the log merged from them, see [`slog`]. +//! +//! A successful analysis moves the log next to the manifest, as `.ai-analysis.slog`, and removes the rest. + +mod checkpoint; +mod recordings; +mod slog; +mod transcript; +mod video; + +use std::collections::VecDeque; +use std::fmt; +use std::fs::File; +use std::io::{self, BufWriter}; +use std::ops::Range; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use camino::{Utf8Path, Utf8PathBuf}; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tracing::{debug, info, warn}; + +use self::checkpoint::DescribedChunk; +pub use self::recordings::{ + RecordingFile, Recordings, RecordingsError, TerminalFormat, TerminalRecording, VideoRecording, WrongRecordingKind, +}; +use self::transcript::TranscriptError; +use self::video::VideoError; +use crate::screen_actions::Screenshot; +use crate::{AiClient, Error, Usage}; + +/// Input tokens sent in one AI request, estimated at 4 characters per token. +const MAX_INPUT_TOKENS_PER_REQUEST: usize = 100_000; + +const CHARS_PER_TOKEN: usize = 4; + +/// Longest transcript chunk, in bytes; it bounds the memory an analysis uses. +const MAX_CHUNK_LEN: usize = MAX_INPUT_TOKENS_PER_REQUEST * CHARS_PER_TOKEN; + +/// A truncated AI answer for a transcript part shorter than this fails the analysis instead of splitting the part again. +const MIN_SPLIT_LEN: usize = 2_000; + +/// Screenshots sent in one AI request. +const MAX_SCREENSHOTS_PER_REQUEST: usize = 40; + +/// Bytes of screenshots sent in one AI request; providers limit the size of a request, such as 32 MB for Anthropic. +const MAX_SCREENSHOT_BYTES_PER_REQUEST: usize = 16 * 1024 * 1024; + +/// Folder of the working files, next to the manifest. +const WORKSPACE_DIR: &str = ".ai-analysis"; + +/// Name of the finished log, next to the manifest. +const LOG_FILE: &str = ".ai-analysis.slog"; + +const SETTINGS_FILE: &str = "analysis.json"; + +/// Written once every chunk file is complete; lists the chunks. +const PLAN_FILE: &str = "plan.json"; + +const WORKSPACE_LOG_FILE: &str = "log.slog"; + +/// What [`AiClient::analyze_recording`] needs from the manifest of a recorded session, `recording.json`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RecordingManifest { + /// Start of the session, in Unix seconds. + pub start_time: i64, + /// Length of the session, in seconds. + pub duration: i64, + /// Recordings of the session, in the order of the manifest, such as from [`Recordings::from_files`]. + pub recordings: Recordings, +} +/// How far an analysis got, in [`RecordingAnalysisEvent::Progress`]. +/// +/// The stages come in order: [`Reading`](Self::Reading), [`Describing`](Self::Describing), then +/// [`Writing`](Self::Writing). An analysis that resumes from its working files skips reading. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum Progress { + /// Reading the recordings, before the AI is asked anything, such as decoding a video and keeping its screenshots: + /// bytes of the recording files read so far, out of their total size. + /// + /// Sent every half second while it lasts. + Reading { read_bytes: u64, total_bytes: u64 }, + /// Asking the AI: chunks already described, out of all chunks. + /// + /// Sent before each chunk, then once all are described. + Describing { + described_chunks: usize, + total_chunks: usize, + }, + /// Writing the log from the answers. + Writing, +} + +impl Progress { + /// Name of the stage, `reading`, `describing` or `writing`, stable enough to be stored. + pub fn stage(&self) -> &'static str { + match self { + Self::Reading { .. } => "reading", + Self::Describing { .. } => "describing", + Self::Writing => "writing", + } + } + + /// Work done in the stage and the work it has in all, in bytes while reading and in chunks while describing; + /// `None` while writing. + pub fn done_and_total(&self) -> Option<(u64, u64)> { + match *self { + Self::Reading { + read_bytes, + total_bytes, + } => Some((read_bytes, total_bytes)), + Self::Describing { + described_chunks, + total_chunks, + } => Some(( + u64::try_from(described_chunks).unwrap_or(u64::MAX), + u64::try_from(total_chunks).unwrap_or(u64::MAX), + )), + Self::Writing => None, + } + } + + /// Part of the stage done, from 0 to 100. + pub fn percent(&self) -> u8 { + match self.done_and_total() { + Some((done, total)) if total > 0 => { + u8::try_from(done.min(total).saturating_mul(100) / total).unwrap_or(100) + } + Some(_) => 0, + None => 100, + } + } +} + +/// A short sentence for people following the analysis, such as `asking the AI: 2 of 5 chunks described`. +impl fmt::Display for Progress { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match *self { + Self::Reading { + read_bytes, + total_bytes, + } => write!( + f, + "reading the recordings: {:.1} of {:.1} MB ({}%)", + megabytes(read_bytes), + megabytes(total_bytes), + self.percent() + ), + Self::Describing { + described_chunks, + total_chunks, + } => write!( + f, + "asking the AI: {described_chunks} of {total_chunks} chunks described" + ), + Self::Writing => f.write_str("writing the log"), + } + } +} + +#[expect(clippy::cast_precision_loss, reason = "shown with one decimal")] +fn megabytes(bytes: u64) -> f64 { + bytes as f64 / 1_048_576.0 +} +/// Time between two [`Progress::Reading`] events. +const PREPARING_PROGRESS_INTERVAL: Duration = Duration::from_millis(500); + +/// A finished analysis. +#[derive(Debug, Clone, PartialEq)] +pub struct RecordingAnalysis { + /// The log, next to the manifest: move it where it belongs, or remove it with [`discard`]. + pub log_path: Utf8PathBuf, + /// Actions in the log. + pub actions: usize, + /// Model the log names: the first one the provider reported, or the requested one when it reported none. + pub model: String, + /// Version of the prompt used, such as `session-actions-3`; the log names it too. + pub prompt_version: String, + /// Tokens the provider counted for the whole analysis, answers cut and asked again included. + /// + /// `None` when the provider did not report them for every request. + pub usage: Option, +} + +/// Error returned when an analysis fails. +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum RecordingAnalysisError { + /// The manifest lists no recording. + #[error("session has no terminal or video recording")] + NoRecording, + /// The session has video recordings, but XMF, which decodes them, is not loaded. + #[error("video decoding is not available: XMF is not loaded")] + VideoDecoderUnavailable, + #[error("failed to read {file_name}: {reason}")] + InvalidRecording { file_name: String, reason: String }, + /// A request to the AI failed. + #[error(transparent)] + Ai(#[from] Error), + /// The answer was cut short even for a part of the input too small to split again. + #[error("AI answer was cut at the output token limit, even for a small part of the recording")] + Truncated, + /// The working files could not be written or read. + #[error("AI analysis files: {0}")] + Files(String), + /// The analysis stopped because its cancellation token was cancelled. + #[error("AI analysis cancelled")] + Cancelled, +} + +impl RecordingAnalysisError { + /// Returns `true` when analysing the recording again later may succeed, such as after a rate limit. + /// + /// The chunks already described are not asked about again. + pub fn is_transient(&self) -> bool { + matches!(self, Self::Ai(error) if error.is_transient()) + } +} + +/// Removes what [`AiClient::analyze_recording`] left next to `manifest_path`: its working files and a log nobody moved. +/// +/// Call it once the analysis of the recording is over for good; a successful analysis already removed its working +/// files. +pub async fn discard(manifest_path: &Utf8Path) -> io::Result<()> { + let dir = recording_dir(manifest_path); + + match tokio::fs::remove_dir_all(dir.join(WORKSPACE_DIR)).await { + Err(error) if error.kind() != io::ErrorKind::NotFound => return Err(error), + _ => {} + } + + match tokio::fs::remove_file(dir.join(LOG_FILE)).await { + Err(error) if error.kind() != io::ErrorKind::NotFound => Err(error), + _ => Ok(()), + } +} + +impl AiClient { + /// Describes what the user did in a recorded session, as a Session Recording Log. + /// + /// `manifest_path` is the path of the session manifest, `recording.json`, and `manifest` what it holds. + /// The recordings are read from the folder of the manifest, and the working files are kept there too. + /// Video recordings need XMF, loaded with `cadeau::xmf::init`. + /// Analysing a recording again after a failure resumes from the chunks already described, as long as the model and + /// the prompts are the same; otherwise it starts over. + /// Only one analysis of a recording may run at a time. + pub fn analyze_recording<'a>( + &'a self, + manifest: &'a RecordingManifest, + manifest_path: &'a Utf8Path, + ) -> AnalyzeRecording<'a> { + AnalyzeRecording { + client: self, + manifest, + manifest_path, + max_output_tokens: None, + cancel: CancellationToken::new(), + } + } +} + +/// Request built by [`AiClient::analyze_recording`]. +#[must_use = "the analysis runs only by `start`"] +pub struct AnalyzeRecording<'a> { + client: &'a AiClient, + manifest: &'a RecordingManifest, + manifest_path: &'a Utf8Path, + max_output_tokens: Option, + cancel: CancellationToken, +} + +impl fmt::Debug for AnalyzeRecording<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AnalyzeRecording") + .field("client", self.client) + .field("manifest_path", &self.manifest_path) + .field("max_output_tokens", &self.max_output_tokens) + .finish_non_exhaustive() + } +} + +impl AnalyzeRecording<'_> { + /// Upper bound of tokens in each AI answer, see + /// [`DescribeSessionActions::max_output_tokens`](crate::session_actions::DescribeSessionActions::max_output_tokens). + pub fn max_output_tokens(mut self, max_output_tokens: u32) -> Self { + self.max_output_tokens = Some(max_output_tokens); + self + } + + /// Stops the analysis soon after `cancel` is cancelled, with [`RecordingAnalysisError::Cancelled`]. + /// + /// Dropping the receiver of the events cancels the analysis too. + pub fn cancellation(mut self, cancel: CancellationToken) -> Self { + self.cancel = cancel; + self + } + + /// Starts the analysis on the Tokio runtime and returns its events. + /// + /// The events are [`RecordingAnalysisEvent::Progress`] while it runs, then [`RecordingAnalysisEvent::Complete`] + /// or [`RecordingAnalysisEvent::Failure`], which is always the last one. Progress events are dropped while the + /// channel is full, so a slow reader only misses some of them; the last event is never dropped. + /// The channel ending without a last event means the analysis panicked. + /// + /// Dropping the receiver cancels the analysis. + pub fn start(self) -> mpsc::Receiver { + let (events, receiver) = mpsc::channel(EVENT_BUFFER); + let finished = CancellationToken::new(); + + // A receiver that went away means nobody waits for the result anymore. + tokio::spawn({ + let events = events.clone(); + let cancel = self.cancel.clone(); + let finished = finished.clone(); + async move { + tokio::select! { + () = events.closed() => cancel.cancel(), + () = finished.cancelled() => {} + } + } + }); + + let run = Run { + client: self.client.clone(), + manifest: self.manifest.clone(), + manifest_path: self.manifest_path.to_owned(), + max_output_tokens: self.max_output_tokens, + cancel: self.cancel, + events, + }; + + tokio::spawn(async move { + let last = match run.run().await { + Ok(analysis) => RecordingAnalysisEvent::Complete(analysis), + Err(error) => RecordingAnalysisEvent::Failure(error), + }; + finished.cancel(); + let _ = run.events.send(last).await; + }); + + receiver + } +} + +/// An event of an analysis started with [`AnalyzeRecording::start`]. +#[derive(Debug)] +#[non_exhaustive] +pub enum RecordingAnalysisEvent { + /// The analysis goes on. + Progress(Progress), + /// The analysis succeeded; it is the last event. + Complete(RecordingAnalysis), + /// The analysis failed; it is the last event. + Failure(RecordingAnalysisError), +} + +/// Events waiting for the reader before progress events are dropped. +const EVENT_BUFFER: usize = 16; + +/// A started analysis. +struct Run { + client: AiClient, + manifest: RecordingManifest, + manifest_path: Utf8PathBuf, + max_output_tokens: Option, + cancel: CancellationToken, + events: mpsc::Sender, +} + +impl Run { + async fn run(&self) -> Result { + let dir = recording_dir(&self.manifest_path).to_owned(); + let workspace = dir.join(WORKSPACE_DIR); + let cancel = self.cancel.clone(); + + let settings = WorkspaceSettings { + prompt_version: self.manifest.recordings.prompt_version().to_owned(), + model: self.client.model().to_owned(), + chunk_len: MAX_CHUNK_LEN, + }; + + let total_bytes = self.manifest.recordings.total_bytes(&dir); + let read_bytes = Arc::new(AtomicU64::new(0)); + + let mut preparing = tokio::task::spawn_blocking({ + let manifest = self.manifest.clone(); + let dir = dir.clone(); + let workspace = workspace.clone(); + let cancel = cancel.clone(); + let read_bytes = Arc::clone(&read_bytes); + move || prepare(&manifest, &dir, &workspace, &settings, &read_bytes, &cancel) + }); + let mut ticks = tokio::time::interval(PREPARING_PROGRESS_INTERVAL); + + // Decoding a long video takes a while, so the bytes read so far are reported until it is done. + let plan = loop { + tokio::select! { + prepared = &mut preparing => { + break prepared + .map_err(|_| RecordingAnalysisError::Files("an AI analysis step panicked".to_owned()))??; + } + _ = ticks.tick() => { + self.report(Progress::Reading { + read_bytes: read_bytes.load(Ordering::Relaxed), + total_bytes, + }); + } + } + }; + let total = plan.len(); + info!(chunks = total, "Describe the session actions"); + + for (index, chunk) in plan.iter().enumerate() { + if cancel.is_cancelled() { + return Err(RecordingAnalysisError::Cancelled); + } + + let checkpoint = checkpoint::path(&workspace, index); + + if checkpoint.exists() { + debug!(index, "Chunk already described"); + continue; + } + + self.report(Progress::Describing { + described_chunks: index, + total_chunks: total, + }); + + let described = match chunk { + Chunk::Transcript { file } => { + let transcript = tokio::fs::read_to_string(workspace.join(file)) + .await + .map_err(|error| files_error(&error))?; + + tokio::select! { + () = cancel.cancelled() => return Err(RecordingAnalysisError::Cancelled), + described = describe_transcript(&self.client, self.max_output_tokens, &transcript) => described?, + } + } + Chunk::Screenshots { frames, context } => { + let frames = read_screenshots(&workspace, frames).await?; + let context = match context { + Some(context) => read_screenshots(&workspace, std::slice::from_ref(context)).await?.pop(), + None => None, + }; + + tokio::select! { + () = cancel.cancelled() => return Err(RecordingAnalysisError::Cancelled), + described = describe_screenshots(&self.client, self.max_output_tokens, &frames, context.as_ref()) => described?, + } + } + }; + + debug!( + index, + actions = described.actions.len(), + model = ?described.model, + usage = ?described.usage, + "Chunk described" + ); + + blocking(move || checkpoint::write(&checkpoint, &described).map_err(|error| files_error(&error))).await?; + } + + self.report(Progress::Describing { + described_chunks: total, + total_chunks: total, + }); + self.report(Progress::Writing); + + let prompt_version = self.manifest.recordings.prompt_version().to_owned(); + let merged = blocking({ + let workspace = workspace.clone(); + let start_time = self.manifest.start_time; + let duration = self.manifest.duration; + let requested_model = self.client.model().to_owned(); + let prompt_version = prompt_version.clone(); + let cancel = cancel.clone(); + move || { + merge_checkpoints( + &workspace, + total, + start_time, + duration, + &requested_model, + &prompt_version, + &cancel, + ) + } + }) + .await?; + + // Once the log leaves the working files, the analysis is done even if cancelled, so this is the last check. + if cancel.is_cancelled() { + return Err(RecordingAnalysisError::Cancelled); + } + + let log_path = dir.join(LOG_FILE); + tokio::fs::rename(workspace.join(WORKSPACE_LOG_FILE), &log_path) + .await + .map_err(|error| files_error(&error))?; + + if let Err(error) = tokio::fs::remove_dir_all(&workspace).await { + warn!(%error, path = %workspace, "Failed to remove the AI analysis working files"); + } + + Ok(RecordingAnalysis { + log_path, + actions: merged.actions, + model: merged.model, + prompt_version, + usage: merged.usage, + }) + } + + /// Sends a progress event, unless the reader is behind; a reader that went away cancels the analysis. + fn report(&self, progress: Progress) { + if let Err(mpsc::error::TrySendError::Closed(_)) = + self.events.try_send(RecordingAnalysisEvent::Progress(progress)) + { + self.cancel.cancel(); + } + } +} +/// Counts the bytes read from a recording, for [`Progress::Reading`]. +struct CountingReader { + inner: R, + read_bytes: Arc, +} + +impl io::Read for CountingReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + let read = self.inner.read(buf)?; + self.read_bytes + .fetch_add(u64::try_from(read).unwrap_or(u64::MAX), Ordering::Relaxed); + Ok(read) + } +} + +/// Opens a recording file, counting the bytes read from it. +fn open_counted(path: &Utf8Path, read_bytes: &Arc) -> io::Result> { + Ok(CountingReader { + inner: File::open(path)?, + read_bytes: Arc::clone(read_bytes), + }) +} +/// Asks the AI about one transcript chunk; a truncated answer is asked again as two halves. +async fn describe_transcript( + client: &AiClient, + max_output_tokens: Option, + chunk: &str, +) -> Result { + let mut parts = VecDeque::from([chunk]); + let mut described = DescribedChunk::default(); + + while let Some(part) = parts.pop_front() { + let mut request = client.describe_session_actions(part); + + if let Some(max_output_tokens) = max_output_tokens { + request = request.max_output_tokens(max_output_tokens); + } + + match request.send().await { + Ok(response) => described.add(response), + Err(Error::Truncated { usage }) => { + described.usage = add_usage(described.usage, usage); + + let Some((first, second)) = split_in_half(part) else { + return Err(RecordingAnalysisError::Truncated); + }; + + debug!(part_len = part.len(), "AI answer truncated; asking again in two halves"); + parts.push_front(second); + parts.push_front(first); + } + Err(error) => return Err(error.into()), + } + } + + Ok(described) +} + +/// A screenshot read back from the working files. +struct LoadedScreenshot { + frame: PlannedFrame, + png: Vec, +} + +impl LoadedScreenshot { + fn as_screenshot(&self, context: bool) -> Screenshot<'_> { + Screenshot { + id: &self.frame.id, + offset: Duration::from_millis(self.frame.offset_ms), + png: &self.png, + context, + } + } +} + +async fn read_screenshots( + workspace: &Utf8Path, + frames: &[PlannedFrame], +) -> Result, RecordingAnalysisError> { + let mut loaded = Vec::with_capacity(frames.len()); + + for frame in frames { + let png = tokio::fs::read(video::frame_path(workspace, &frame.id)) + .await + .map_err(|error| files_error(&error))?; + loaded.push(LoadedScreenshot { + frame: frame.clone(), + png, + }); + } + + Ok(loaded) +} + +/// Asks the AI about one chunk of screenshots, in parts that fit in a request; a truncated answer is asked again as +/// two halves. +/// +/// Each part starts with the screenshot before it as context: `context` for the first part. +async fn describe_screenshots( + client: &AiClient, + max_output_tokens: Option, + frames: &[LoadedScreenshot], + context: Option<&LoadedScreenshot>, +) -> Result { + let mut parts = VecDeque::from(split_by_size(frames)); + let mut described = DescribedChunk::default(); + + while let Some(part) = parts.pop_front() { + let before = match part.start { + 0 => context, + start => frames.get(start - 1), + }; + + let screenshots = before + .map(|before| before.as_screenshot(true)) + .into_iter() + .chain(frames[part.clone()].iter().map(|frame| frame.as_screenshot(false))) + .collect::>(); + + let mut request = client.describe_screen_actions(&screenshots); + + if let Some(max_output_tokens) = max_output_tokens { + request = request.max_output_tokens(max_output_tokens); + } + + match request.send().await { + Ok(response) => described.add(response), + Err(Error::Truncated { usage }) => { + described.usage = add_usage(described.usage, usage); + + if part.len() < 2 { + return Err(RecordingAnalysisError::Truncated); + } + + let middle = part.start + part.len() / 2; + debug!( + screenshots = part.len(), + "AI answer truncated; asking again in two halves" + ); + parts.push_front(middle..part.end); + parts.push_front(part.start..middle); + } + Err(error) => return Err(error.into()), + } + } + + Ok(described) +} + +/// Cuts `frames` into consecutive parts of at most [`MAX_SCREENSHOT_BYTES_PER_REQUEST`] bytes, each with one +/// screenshot at least. +fn split_by_size(frames: &[LoadedScreenshot]) -> Vec> { + let mut parts = Vec::new(); + let mut start = 0; + let mut bytes = 0; + + for (index, frame) in frames.iter().enumerate() { + if index > start && bytes + frame.png.len() > MAX_SCREENSHOT_BYTES_PER_REQUEST { + parts.push(start..index); + start = index; + bytes = 0; + } + bytes += frame.png.len(); + } + + if start < frames.len() { + parts.push(start..frames.len()); + } + + parts +} + +/// What the working files were made with; files made with other settings are never reused. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct WorkspaceSettings { + prompt_version: String, + model: String, + chunk_len: usize, +} + +/// A part of the input the AI is asked about in one go. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase", rename_all_fields = "camelCase")] +enum Chunk { + /// A transcript chunk file. + Transcript { file: String }, + /// Screenshots, with the one before them for context. + Screenshots { + frames: Vec, + context: Option, + }, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct PlannedFrame { + id: String, + /// Elapsed time since the start of the session. + offset_ms: u64, +} + +fn recording_dir(manifest_path: &Utf8Path) -> &Utf8Path { + match manifest_path.parent() { + Some(dir) if !dir.as_str().is_empty() => dir, + _ => Utf8Path::new("."), + } +} + +async fn blocking( + work: impl FnOnce() -> Result + Send + 'static, +) -> Result { + tokio::task::spawn_blocking(work) + .await + .map_err(|_| RecordingAnalysisError::Files("an AI analysis step panicked".to_owned()))? +} + +fn files_error(error: &dyn fmt::Display) -> RecordingAnalysisError { + RecordingAnalysisError::Files(format!("{error:#}")) +} + +/// Gets the working files ready and returns the chunks: the transcript or the screenshots are made once, and kept for +/// another analysis with the same settings. +fn prepare( + manifest: &RecordingManifest, + dir: &Utf8Path, + workspace: &Utf8Path, + settings: &WorkspaceSettings, + read_bytes: &Arc, + cancel: &CancellationToken, +) -> Result, RecordingAnalysisError> { + if manifest.recordings.is_empty() { + return Err(RecordingAnalysisError::NoRecording); + } + + let saved = std::fs::read(workspace.join(SETTINGS_FILE)) + .ok() + .and_then(|saved| serde_json::from_slice::(&saved).ok()); + + if saved.as_ref() == Some(settings) { + let plan = std::fs::read(workspace.join(PLAN_FILE)) + .ok() + .and_then(|plan| serde_json::from_slice::>(&plan).ok()); + + if let Some(plan) = plan { + return Ok(plan); + } + } else if workspace.exists() { + info!(path = %workspace, "Discard the working files of another AI analysis"); + std::fs::remove_dir_all(workspace).map_err(|error| files_error(&error))?; + } + + std::fs::create_dir_all(workspace).map_err(|error| files_error(&error))?; + + let settings = serde_json::to_vec(settings).map_err(|error| files_error(&error))?; + std::fs::write(workspace.join(SETTINGS_FILE), settings).map_err(|error| files_error(&error))?; + + let plan = match &manifest.recordings { + Recordings::Terminal(recordings) => { + plan_transcript(recordings, manifest.start_time, dir, workspace, read_bytes, cancel)? + } + Recordings::Video(recordings) => { + plan_screenshots(recordings, manifest.start_time, dir, workspace, read_bytes, cancel)? + } + }; + let saved_plan = serde_json::to_vec(&plan).map_err(|error| files_error(&error))?; + std::fs::write(workspace.join(PLAN_FILE), saved_plan).map_err(|error| files_error(&error))?; + + Ok(plan) +} + +/// Writes the transcript chunks of terminal recordings. +fn plan_transcript( + recordings: &[TerminalRecording], + session_start: i64, + dir: &Utf8Path, + workspace: &Utf8Path, + read_bytes: &Arc, + cancel: &CancellationToken, +) -> Result, RecordingAnalysisError> { + let chunks = transcript::write_chunks( + &in_time_order(recordings, TerminalRecording::start_time), + session_start, + dir, + workspace, + MAX_CHUNK_LEN, + read_bytes, + cancel, + ) + .map_err(|error| match error { + TranscriptError::Invalid { file_name, reason } => { + RecordingAnalysisError::InvalidRecording { file_name, reason } + } + TranscriptError::Write(error) => files_error(&error), + TranscriptError::Cancelled => RecordingAnalysisError::Cancelled, + })?; + + Ok((0..chunks) + .map(|index| Chunk::Transcript { + file: transcript::chunk_path(workspace, index) + .file_name() + .unwrap_or_default() + .to_owned(), + }) + .collect()) +} + +/// Keeps the screenshots of video recordings and groups them by request, each group after the screenshot before it. +fn plan_screenshots( + recordings: &[VideoRecording], + session_start: i64, + dir: &Utf8Path, + workspace: &Utf8Path, + read_bytes: &Arc, + cancel: &CancellationToken, +) -> Result, RecordingAnalysisError> { + let mut frames = Vec::new(); + let mut next_id = 0; + + for recording in in_time_order(recordings, VideoRecording::start_time) { + let offset = + Duration::from_secs(u64::try_from(recording.start_time().saturating_sub(session_start)).unwrap_or(0)); + + let kept = video::keep_changed_frames( + &dir.join(recording.file_name()), + offset, + workspace, + &mut next_id, + read_bytes, + cancel, + ) + .map_err(|error| match error { + VideoError::DecoderUnavailable => RecordingAnalysisError::VideoDecoderUnavailable, + VideoError::Invalid(reason) => RecordingAnalysisError::InvalidRecording { + file_name: recording.file_name().to_owned(), + reason, + }, + VideoError::Write(error) => files_error(&error), + VideoError::Cancelled => RecordingAnalysisError::Cancelled, + })?; + + info!(file = %recording.file_name(), screenshots = kept.len(), "Screenshots kept from a video"); + + frames.extend(kept.into_iter().map(|frame| PlannedFrame { + id: frame.id, + offset_ms: u64::try_from(frame.offset.as_millis()).unwrap_or(u64::MAX), + })); + } + + Ok(group_screenshots(frames)) +} + +/// Groups the screenshots of a session, in time order, into requests of at most [`MAX_SCREENSHOTS_PER_REQUEST`], each +/// after the last screenshot of the group before it; a session recorded in several videos is grouped as one. +fn group_screenshots(frames: Vec) -> Vec { + let mut groups = Vec::new(); + let mut previous = None; + + for frames in frames.chunks(MAX_SCREENSHOTS_PER_REQUEST) { + groups.push(Chunk::Screenshots { + frames: frames.to_vec(), + context: previous.clone(), + }); + previous = frames.last().cloned(); + } + + groups +} + +/// The recordings of a session sorted by start time, since a manifest may list them in another order. +fn in_time_order(recordings: &[R], start_time: impl Fn(&R) -> i64) -> Vec { + let mut sorted = recordings.to_vec(); + sorted.sort_by_key(|recording| start_time(recording)); + sorted +} + +/// Sums token counts; the total is unknown as soon as one count is. +fn add_usage(total: Option, more: Option) -> Option { + Some(total? + more?) +} + +/// Cuts `part` on the line boundary closest to its middle, unless it is too short to split. +fn split_in_half(part: &str) -> Option<(&str, &str)> { + if part.len() < MIN_SPLIT_LEN { + return None; + } + + let bytes = part.as_bytes(); + let middle = part.len() / 2; + + let before = bytes[..middle] + .iter() + .rposition(|&byte| byte == b'\n') + .map(|end| end + 1); + let after = bytes[middle..] + .iter() + .position(|&byte| byte == b'\n') + .map(|end| middle + end + 1); + + let cut = match (before, after) { + (Some(before), Some(after)) if middle - before <= after - middle => before, + (_, Some(after)) if after < part.len() => after, + (Some(before), _) => before, + _ => return None, + }; + + Some(part.split_at(cut)) +} + +/// What [`merge_checkpoints`] wrote. +struct MergedLog { + actions: usize, + model: String, + usage: Option, +} + +/// Writes the log from the checkpoints, into the working files, with every action in time order. +/// +/// The log names the first model reported for a chunk, or `requested_model` when the provider reported none. +fn merge_checkpoints( + workspace: &Utf8Path, + total: usize, + start_time: i64, + duration: i64, + requested_model: &str, + prompt_version: &str, + cancel: &CancellationToken, +) -> Result { + let log_error = |error: anyhow::Error| files_error(&format!("failed to write the log: {error:#}")); + + let mut actions = Vec::new(); + let mut model = None; + let mut usage = Some(Usage::default()); + + for index in 0..total { + if cancel.is_cancelled() { + return Err(RecordingAnalysisError::Cancelled); + } + + let described = checkpoint::read(&checkpoint::path(workspace, index)).map_err(|error| files_error(&error))?; + model = model.or(described.model); + usage = add_usage(usage, described.usage); + actions.extend(described.actions); + } + + // A model may list the actions of a chunk slightly out of order; the sort is stable, so equal times keep their order. + actions.sort_by_key(|action| action.offset); + let model = model.unwrap_or_else(|| requested_model.to_owned()); + + let out = File::create(workspace.join(WORKSPACE_LOG_FILE)).map_err(|error| files_error(&error))?; + let mut log = + slog::SlogWriter::start(BufWriter::new(out), start_time, &model, prompt_version).map_err(log_error)?; + + for action in &actions { + log.action(action).map_err(log_error)?; + } + + log.finish(duration) + .map_err(log_error)? + .into_inner() + .map_err(|error| files_error(&error.into_error()))? + .sync_all() + .map_err(|error| files_error(&error))?; + + Ok(MergedLog { + actions: actions.len(), + model, + usage, + }) +} + +#[cfg(test)] +mod tests; diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/recordings.rs b/crates/devolutions-gateway-ai/src/recording_analysis/recordings.rs new file mode 100644 index 000000000..144668921 --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/recordings.rs @@ -0,0 +1,250 @@ +//! The recordings of a session, sorted by kind when they are accepted. +//! +//! A session is recorded either as a terminal or as a video, never both, and the types hold that: [`Recordings`] is +//! all terminal or all video, and each file is accepted only with the extension of its kind. + +use camino::Utf8Path; + +/// A file listed in a session manifest, `recording.json`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RecordingFile { + /// Name of the file, in the folder of the manifest. + pub file_name: String, + /// Start of the recording, in Unix seconds. + pub start_time: i64, +} + +/// The recordings of one session: all terminal, or all video. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Recordings { + Terminal(Vec), + Video(Vec), +} + +/// Why the files of a session manifest are not one kind of recording. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum RecordingsError { + #[error("session has no terminal or video recording")] + NoRecording, + #[error("session mixes terminal and video recordings")] + Mixed, +} + +impl Recordings { + /// Sorts the files of a session manifest by kind; files that are not recordings, such as session logs, are + /// skipped. + pub fn from_files(files: impl IntoIterator) -> Result { + let mut terminal = Vec::new(); + let mut video = Vec::new(); + + for file in files { + match Kind::of(&file.file_name) { + Some(Kind::Terminal(format)) => terminal.push(TerminalRecording { + file_name: file.file_name, + start_time: file.start_time, + format, + }), + Some(Kind::Video) => video.push(VideoRecording { + file_name: file.file_name, + start_time: file.start_time, + }), + None => {} + } + } + + match (terminal.is_empty(), video.is_empty()) { + (false, true) => Ok(Self::Terminal(terminal)), + (true, false) => Ok(Self::Video(video)), + (true, true) => Err(RecordingsError::NoRecording), + (false, false) => Err(RecordingsError::Mixed), + } + } + + /// Version of the prompt that describes this kind of recording. + pub(super) fn prompt_version(&self) -> &'static str { + match self { + Self::Terminal(_) => crate::session_actions::PROMPT_VERSION, + Self::Video(_) => crate::screen_actions::PROMPT_VERSION, + } + } + + /// Total size of the recording files in `dir`; a file that cannot be read counts as empty. + pub(super) fn total_bytes(&self, dir: &Utf8Path) -> u64 { + let size = |file_name: &str| std::fs::metadata(dir.join(file_name)).map_or(0, |metadata| metadata.len()); + + match self { + Self::Terminal(recordings) => recordings.iter().map(|recording| size(recording.file_name())).sum(), + Self::Video(recordings) => recordings.iter().map(|recording| size(recording.file_name())).sum(), + } + } + + pub(super) fn is_empty(&self) -> bool { + match self { + Self::Terminal(recordings) => recordings.is_empty(), + Self::Video(recordings) => recordings.is_empty(), + } + } +} + +/// A file that is not a recording of the expected kind. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("{file_name} is not a {expected} recording")] +pub struct WrongRecordingKind { + pub file_name: String, + pub expected: &'static str, +} + +/// Format of a terminal recording. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TerminalFormat { + /// Asciicast v2 or v3, `.cast`. + Asciicast, + /// Terminal Playback Recording, `.trp`. + Trp, +} + +/// A terminal recording, `.cast` or `.trp`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TerminalRecording { + file_name: String, + start_time: i64, + format: TerminalFormat, +} + +impl TerminalRecording { + /// Accepts `file_name` only with the `.cast` or `.trp` extension. + pub fn new(file_name: impl Into, start_time: i64) -> Result { + let file_name = file_name.into(); + + match Kind::of(&file_name) { + Some(Kind::Terminal(format)) => Ok(Self { + file_name, + start_time, + format, + }), + _ => Err(WrongRecordingKind { + file_name, + expected: "terminal", + }), + } + } + + /// Name of the file, in the folder of the manifest. + pub fn file_name(&self) -> &str { + &self.file_name + } + + /// Start of the recording, in Unix seconds. + pub fn start_time(&self) -> i64 { + self.start_time + } + + pub fn format(&self) -> TerminalFormat { + self.format + } +} + +/// A video recording, `.webm`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct VideoRecording { + file_name: String, + start_time: i64, +} + +impl VideoRecording { + /// Accepts `file_name` only with the `.webm` extension. + pub fn new(file_name: impl Into, start_time: i64) -> Result { + let file_name = file_name.into(); + + match Kind::of(&file_name) { + Some(Kind::Video) => Ok(Self { file_name, start_time }), + _ => Err(WrongRecordingKind { + file_name, + expected: "video", + }), + } + } + + /// Name of the file, in the folder of the manifest. + pub fn file_name(&self) -> &str { + &self.file_name + } + + /// Start of the recording, in Unix seconds. + pub fn start_time(&self) -> i64 { + self.start_time + } +} + +enum Kind { + Terminal(TerminalFormat), + Video, +} + +impl Kind { + fn of(file_name: &str) -> Option { + match Utf8Path::new(file_name).extension()? { + "cast" => Some(Self::Terminal(TerminalFormat::Asciicast)), + "trp" => Some(Self::Terminal(TerminalFormat::Trp)), + "webm" => Some(Self::Video), + _ => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn file(file_name: &str) -> RecordingFile { + RecordingFile { + file_name: file_name.to_owned(), + start_time: 1, + } + } + + #[test] + fn files_are_sorted_by_kind_and_other_files_skipped() { + let terminal = Recordings::from_files([file("recording-0.cast"), file("log.slog"), file("recording-1.trp")]); + let Ok(Recordings::Terminal(terminal)) = terminal else { + panic!("terminal session: {terminal:?}"); + }; + assert_eq!( + terminal.iter().map(TerminalRecording::format).collect::>(), + [TerminalFormat::Asciicast, TerminalFormat::Trp] + ); + + let video = Recordings::from_files([file("recording-0.webm"), file("recording-1.webm")]); + assert!( + matches!(video, Ok(Recordings::Video(ref video)) if video.len() == 2), + "{video:?}" + ); + } + + #[test] + fn a_session_with_both_kinds_or_none_is_refused() { + assert_eq!( + Recordings::from_files([file("recording-0.cast"), file("recording-1.webm")]), + Err(RecordingsError::Mixed) + ); + assert_eq!( + Recordings::from_files([file("log.slog"), file("notes.txt")]), + Err(RecordingsError::NoRecording) + ); + assert_eq!(Recordings::from_files([]), Err(RecordingsError::NoRecording)); + } + + #[test] + fn a_recording_is_accepted_only_with_the_extension_of_its_kind() { + assert!(TerminalRecording::new("recording-0.cast", 1).is_ok()); + assert!(TerminalRecording::new("recording-0.trp", 1).is_ok()); + assert!(VideoRecording::new("recording-0.webm", 1).is_ok()); + + for name in ["recording-0.webm", "log.slog", "recording-0", "recording-0.CAST"] { + assert!(TerminalRecording::new(name, 1).is_err(), "{name}"); + } + for name in ["recording-0.cast", "recording-0.trp", "recording-0.mp4"] { + assert!(VideoRecording::new(name, 1).is_err(), "{name}"); + } + } +} diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/slog.rs b/crates/devolutions-gateway-ai/src/recording_analysis/slog.rs new file mode 100644 index 000000000..7349125f4 --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/slog.rs @@ -0,0 +1,174 @@ +//! Writes AI actions as a Session Recording Log (`.slog`): one JSON object per line, in the shipped schema. + +use std::collections::BTreeMap; +use std::time::Duration; + +use serde::Serialize; + +use crate::session_actions::Action; + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct Entry<'a> { + timestamp: String, + seq: usize, + event: &'static str, + description: &'a str, + #[serde(skip_serializing_if = "Option::is_none")] + object: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + parameters: Option<&'a BTreeMap>, + #[serde(skip_serializing_if = "Option::is_none")] + source: Option<&'static str>, + #[serde(skip_serializing_if = "Option::is_none")] + model: Option<&'a str>, + #[serde(skip_serializing_if = "Option::is_none")] + prompt_version: Option<&'a str>, +} + +impl<'a> Entry<'a> { + fn new(seq: usize, timestamp: String, event: &'static str, description: &'a str) -> Self { + Self { + timestamp, + seq, + event, + description, + object: None, + parameters: None, + source: None, + model: None, + prompt_version: None, + } + } +} + +/// Writes the log of a session entry by entry, so the actions never have to be all in memory. +pub(super) struct SlogWriter { + out: W, + start_time: i64, + seq: usize, +} + +impl SlogWriter { + /// Writes `session.start` for a session that started at `start_time` (unix seconds), naming the model and the + /// versions of the prompts that found the actions. + pub(super) fn start(out: W, start_time: i64, model: &str, prompt_version: &str) -> anyhow::Result { + let mut writer = Self { + out, + start_time, + seq: 0, + }; + + let timestamp = timestamp(start_time, Duration::ZERO)?; + writer.write(Entry { + source: Some("ai"), + model: Some(model), + prompt_version: Some(prompt_version), + ..Entry::new(0, timestamp, "session.start", "Session started") + })?; + + Ok(writer) + } + + /// Actions must come in offset order. + pub(super) fn action(&mut self, action: &Action) -> anyhow::Result<()> { + let timestamp = timestamp(self.start_time, action.offset)?; + self.write(Entry { + object: action.object.as_deref(), + parameters: Some(&action.parameters).filter(|parameters| !parameters.is_empty()), + ..Entry::new(self.seq, timestamp, "session.action", &action.description) + }) + } + + /// Writes `session.end` for a session that lasted `duration` seconds. + pub(super) fn finish(mut self, duration: i64) -> anyhow::Result { + let end_offset = Duration::from_secs(u64::try_from(duration).unwrap_or(0)); + let timestamp = timestamp(self.start_time, end_offset)?; + self.write(Entry::new(self.seq, timestamp, "session.end", "Session ended"))?; + self.out.flush()?; + Ok(self.out) + } + + fn write(&mut self, entry: Entry<'_>) -> anyhow::Result<()> { + serde_json::to_writer(&mut self.out, &entry)?; + self.out.write_all(b"\n")?; + self.seq += 1; + Ok(()) + } +} + +/// ISO 8601 UTC time with milliseconds, like the other `.slog` writers. +fn timestamp(start_time: i64, offset: Duration) -> anyhow::Result { + let time = time::OffsetDateTime::from_unix_timestamp(start_time)? + .checked_add(time::Duration::try_from(offset)?) + .ok_or_else(|| anyhow::anyhow!("timestamp out of range"))?; + + Ok(format!( + "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:03}Z", + time.year(), + u8::from(time.month()), + time.day(), + time.hour(), + time.minute(), + time.second(), + time.millisecond(), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn write(start_time: i64, duration: i64, actions: &[Action]) -> String { + let mut writer = SlogWriter::start(Vec::new(), start_time, "gpt-test", "session-actions-3").expect("started"); + for action in actions { + writer.action(action).expect("written"); + } + String::from_utf8(writer.finish(duration).expect("finished")).expect("UTF-8") + } + + fn action(offset_ms: u64, description: &str, object: Option<&str>, parameters: &[(&str, &str)]) -> Action { + Action { + offset: Duration::from_millis(offset_ms), + description: description.to_owned(), + object: object.map(str::to_owned), + parameters: parameters + .iter() + .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) + .collect(), + } + } + + #[test] + fn writes_the_shipped_schema_with_ai_fields_on_session_start() { + let actions = [ + action(1500, "Listed files", Some("/var/log"), &[("Command", "ls /var/log")]), + action(62_250, "Closed the shell", None, &[]), + ]; + + let slog = write(1_787_255_035, 90, &actions); + + let expected = [ + r#"{"timestamp":"2026-08-20T19:43:55.000Z","seq":0,"event":"session.start","description":"Session started","source":"ai","model":"gpt-test","promptVersion":"session-actions-3"}"#, + r#"{"timestamp":"2026-08-20T19:43:56.500Z","seq":1,"event":"session.action","description":"Listed files","object":"/var/log","parameters":{"Command":"ls /var/log"}}"#, + r#"{"timestamp":"2026-08-20T19:44:57.250Z","seq":2,"event":"session.action","description":"Closed the shell"}"#, + r#"{"timestamp":"2026-08-20T19:45:25.000Z","seq":3,"event":"session.end","description":"Session ended"}"#, + ]; + assert_eq!(slog, expected.map(|line| format!("{line}\n")).concat()); + } + + #[test] + fn a_session_without_actions_still_starts_and_ends() { + let slog = write(0, 5, &[]); + + let events: Vec = slog + .lines() + .map(|line| serde_json::from_str(line).expect("JSON line")) + .collect(); + assert_eq!(events.len(), 2); + assert_eq!(events[0]["event"], "session.start"); + assert_eq!(events[1]["event"], "session.end"); + assert_eq!(events[1]["seq"], 1); + assert_eq!(events[1]["timestamp"], "1970-01-01T00:00:05.000Z"); + } +} diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs new file mode 100644 index 000000000..064e30c5a --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs @@ -0,0 +1,501 @@ +#![allow(clippy::unwrap_used, reason = "test code can panic on errors")] + +use std::sync::{Arc, Mutex}; + +use tokio::io::{AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt as _, BufReader}; +use tokio::net::TcpListener; +use tokio::sync::mpsc; + +use super::*; +use crate::Provider; + +const MODEL: &str = "gpt-test"; +const REPORTED_MODEL: &str = "gpt-test-2026-09-30"; + +/// A mock OpenAI-compatible provider: `status` picks the answer of the n-th request, 0 being a valid answer. +/// +/// Returns a client for it and the bodies of the requests it got. +async fn spawn_provider(status: fn(usize) -> u16) -> (AiClient, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + + tokio::spawn({ + let requests = Arc::clone(&requests); + async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + let requests = Arc::clone(&requests); + + tokio::spawn(async move { + let mut stream = BufReader::new(stream); + let mut content_length = 0; + + loop { + let mut line = String::new(); + stream.read_line(&mut line).await.unwrap(); + let line = line.trim_end(); + + if line.is_empty() { + break; + } + + if let Some((name, value)) = line.split_once(':') + && name.eq_ignore_ascii_case("content-length") + { + content_length = value.trim().parse().unwrap(); + } + } + + let mut body = vec![0; content_length]; + stream.read_exact(&mut body).await.unwrap(); + + let index = { + let mut requests = requests.lock().unwrap(); + requests.push(String::from_utf8(body).unwrap()); + requests.len() - 1 + }; + + let (status, answer) = match status(index) { + 0 => ( + 200, + serde_json::json!({ + "model": REPORTED_MODEL, + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": format!("{{\"line\":\"L1\",\"description\":\"Step {index}\"}}"), + }, + "finish_reason": "stop", + }], + "usage": { "prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30 }, + }), + ), + status => ( + status, + serde_json::json!({ "error": { "message": "scripted failure" } }), + ), + }; + + let answer = answer.to_string(); + let response = format!( + "HTTP/1.1 {status} Scripted\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{answer}", + answer.len() + ); + stream.get_mut().write_all(response.as_bytes()).await.unwrap(); + }); + } + } + }); + + let client = AiClient::builder() + .provider(Provider::OpenAiCompatible) + .base_url(format!("http://{addr}/v1/").parse().unwrap()) + .model(MODEL) + .api_key("sk-test-secret") + .http_client(reqwest::Client::new()) + .build() + .unwrap(); + + (client, requests) +} + +/// A finished session with one terminal recording of `lines` lines of about 100 bytes each. +struct Session { + _dir: tempfile::TempDir, + manifest: RecordingManifest, + manifest_path: Utf8PathBuf, +} + +impl Session { + fn new(lines: usize) -> Self { + let dir = tempfile::tempdir().unwrap(); + let path = Utf8Path::from_path(dir.path()).unwrap(); + + let mut cast = String::from("{\"version\": 2}\n"); + for line in 0..lines { + cast.push_str(&format!("[{line}.0,\"o\",\"line {line} {}\\r\\n\"]\n", "x".repeat(80))); + } + std::fs::write(path.join("recording-0.cast"), cast).unwrap(); + + Self { + manifest: RecordingManifest { + start_time: 1_787_255_035, + duration: i64::try_from(lines).unwrap(), + recordings: Recordings::Terminal(vec![ + TerminalRecording::new("recording-0.cast", 1_787_255_035).unwrap(), + ]), + }, + manifest_path: path.join("recording.json"), + _dir: dir, + } + } + + fn workspace(&self) -> Utf8PathBuf { + recording_dir(&self.manifest_path).join(WORKSPACE_DIR) + } + + async fn analyze(&self, client: &AiClient) -> Result { + finish(client.analyze_recording(&self.manifest, &self.manifest_path).start()) + .await + .1 + } +} + +/// Reads the events of an analysis to the last one, and returns its progress and its result. +async fn finish( + mut events: mpsc::Receiver, +) -> (Vec, Result) { + let mut progress = Vec::new(); + + while let Some(event) = events.recv().await { + match event { + RecordingAnalysisEvent::Progress(step) => progress.push(step), + RecordingAnalysisEvent::Complete(analysis) => { + assert!(events.recv().await.is_none(), "Complete is the last event"); + return (progress, Ok(analysis)); + } + RecordingAnalysisEvent::Failure(error) => { + assert!(events.recv().await.is_none(), "Failure is the last event"); + return (progress, Err(error)); + } + } + } + + panic!("the events ended without a last event"); +} + +fn descriptions(log: &Utf8Path) -> Vec { + std::fs::read_to_string(log) + .unwrap() + .lines() + .map(|line| { + serde_json::from_str::(line).unwrap()["description"] + .as_str() + .unwrap() + .to_owned() + }) + .collect() +} + +#[tokio::test] +async fn the_log_is_left_next_to_the_manifest_and_the_working_files_removed() { + let session = Session::new(10); + let (client, _requests) = spawn_provider(|_| 0).await; + let (progress, analysis) = finish( + client + .analyze_recording(&session.manifest, &session.manifest_path) + .start(), + ) + .await; + let analysis = analysis.unwrap(); + + assert_eq!(analysis.log_path, recording_dir(&session.manifest_path).join(LOG_FILE)); + assert_eq!(analysis.actions, 1); + assert_eq!(analysis.model, REPORTED_MODEL); + assert_eq!( + analysis.usage, + Some(Usage { + input_tokens: 10, + output_tokens: 20 + }) + ); + assert_eq!( + descriptions(&analysis.log_path), + ["Session started", "Step 0", "Session ended"] + ); + assert!(!session.workspace().exists()); + let recording_size = std::fs::metadata(recording_dir(&session.manifest_path).join("recording-0.cast")) + .unwrap() + .len(); + let Some(Progress::Reading { + read_bytes, + total_bytes, + }) = progress.first().copied() + else { + panic!("reading the recording is reported first: {progress:?}"); + }; + assert_eq!(total_bytes, recording_size); + assert!(read_bytes <= total_bytes); + assert_eq!( + progress + .iter() + .filter(|step| matches!(step, Progress::Describing { .. })) + .copied() + .collect::>(), + [ + Progress::Describing { + described_chunks: 0, + total_chunks: 1 + }, + Progress::Describing { + described_chunks: 1, + total_chunks: 1 + }, + ] + ); + assert_eq!(progress.last(), Some(&Progress::Writing)); +} + +#[tokio::test] +async fn analysing_again_only_asks_about_the_chunks_left() { + // About 1 MB of transcript: three chunks. + let session = Session::new(10_000); + let (client, requests) = spawn_provider(|index| if index == 1 { 503 } else { 0 }).await; + + let error = session.analyze(&client).await.unwrap_err(); + assert!(error.is_transient(), "{error:?}"); + assert!(session.workspace().join("chunk-0000.json").exists()); + assert!(!session.workspace().join("chunk-0001.json").exists()); + + let analysis = session.analyze(&client).await.unwrap(); + + assert_eq!(requests.lock().unwrap().len(), 4, "only chunks 1 and 2 are sent again"); + assert_eq!( + descriptions(&analysis.log_path), + ["Session started", "Step 0", "Step 2", "Step 3", "Session ended"] + ); + assert_eq!( + analysis.usage, + Some(Usage { + input_tokens: 30, + output_tokens: 60 + }), + "chunk 0 counts from its checkpoint, the failed request not at all" + ); +} + +#[tokio::test] +async fn working_files_of_another_analysis_are_not_reused() { + let session = Session::new(10); + let (client, requests) = spawn_provider(|index| if index == 0 { 503 } else { 0 }).await; + session.analyze(&client).await.unwrap_err(); + + let settings = session.workspace().join(SETTINGS_FILE); + let saved = std::fs::read_to_string(&settings).unwrap(); + std::fs::write(&settings, saved.replace(MODEL, "another-model")).unwrap(); + checkpoint::write( + &checkpoint::path(&session.workspace(), 0), + &DescribedChunk { + actions: Vec::new(), + model: None, + usage: None, + }, + ) + .unwrap(); + + let analysis = session.analyze(&client).await.unwrap(); + + assert_eq!(requests.lock().unwrap().len(), 2, "the chunk is asked about again"); + assert_eq!( + descriptions(&analysis.log_path), + ["Session started", "Step 1", "Session ended"] + ); +} + +#[tokio::test] +async fn a_manifest_without_recording_fails_for_good() { + let mut session = Session::new(1); + session.manifest.recordings = Recordings::Terminal(Vec::new()); + let (client, requests) = spawn_provider(|_| 0).await; + + let error = session.analyze(&client).await.unwrap_err(); + + assert!(matches!(error, RecordingAnalysisError::NoRecording), "{error:?}"); + assert!(!error.is_transient()); + assert!(requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn a_cancelled_analysis_stops() { + let session = Session::new(10); + let (client, requests) = spawn_provider(|_| 0).await; + let cancel = CancellationToken::new(); + cancel.cancel(); + + let (_, result) = finish( + client + .analyze_recording(&session.manifest, &session.manifest_path) + .cancellation(cancel) + .start(), + ) + .await; + let error = result.unwrap_err(); + + assert!(matches!(error, RecordingAnalysisError::Cancelled), "{error:?}"); + assert!(requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn discard_removes_the_working_files_and_the_log() { + let session = Session::new(10); + let (client, _requests) = spawn_provider(|index| if index == 0 { 503 } else { 0 }).await; + session.analyze(&client).await.unwrap_err(); + let log = recording_dir(&session.manifest_path).join(LOG_FILE); + std::fs::write(&log, "left behind").unwrap(); + + discard(&session.manifest_path).await.unwrap(); + + assert!(!session.workspace().exists()); + assert!(!log.exists()); + assert!(recording_dir(&session.manifest_path).join("recording-0.cast").exists()); + discard(&session.manifest_path).await.unwrap(); +} + +#[tokio::test] +async fn dropping_the_events_cancels_the_analysis() { + // About 1 MB of transcript: three chunks, each answered after a while. + let session = Session::new(10_000); + let (client, requests) = spawn_provider(|_| 0).await; + + let mut events = client + .analyze_recording(&session.manifest, &session.manifest_path) + .start(); + while let Some(event) = events.recv().await { + if matches!(event, RecordingAnalysisEvent::Progress(Progress::Describing { .. })) { + break; + } + } + drop(events); + + // The analysis stops at its next step: it never gets to the log. + tokio::time::sleep(Duration::from_secs(1)).await; + assert!(requests.lock().unwrap().len() < 3, "{}", requests.lock().unwrap().len()); + assert!(!recording_dir(&session.manifest_path).join(LOG_FILE).exists()); +} + +#[test] +fn progress_tells_its_stage_and_how_far_it_is() { + let reading = Progress::Reading { + read_bytes: 1_572_864, + total_bytes: 2_097_152, + }; + assert_eq!(reading.stage(), "reading"); + assert_eq!(reading.done_and_total(), Some((1_572_864, 2_097_152))); + assert_eq!(reading.percent(), 75); + assert_eq!(reading.to_string(), "reading the recordings: 1.5 of 2.0 MB (75%)"); + + let describing = Progress::Describing { + described_chunks: 2, + total_chunks: 5, + }; + assert_eq!(describing.stage(), "describing"); + assert_eq!(describing.percent(), 40); + assert_eq!(describing.to_string(), "asking the AI: 2 of 5 chunks described"); + + assert_eq!(Progress::Writing.stage(), "writing"); + assert_eq!(Progress::Writing.done_and_total(), None); + assert_eq!(Progress::Writing.to_string(), "writing the log"); + + let nothing = Progress::Reading { + read_bytes: 0, + total_bytes: 0, + }; + assert_eq!(nothing.percent(), 0); +} +#[test] +fn screenshots_of_several_videos_are_grouped_as_one_session() { + // 25 screenshots from a first video, then 45 from a second one. + let frames = (0..70) + .map(|index| PlannedFrame { + id: format!("F{index:05}"), + offset_ms: if index < 25 { index * 500 } else { 60_000 + index * 500 }, + }) + .collect::>(); + + let groups = group_screenshots(frames.clone()); + + assert_eq!( + groups, + [ + Chunk::Screenshots { + frames: frames[..40].to_vec(), + context: None, + }, + Chunk::Screenshots { + frames: frames[40..].to_vec(), + context: Some(frames[39].clone()), + }, + ], + "the second video does not start a request of its own" + ); +} + +#[test] +fn recordings_listed_out_of_order_are_read_in_time_order() { + let dir = tempfile::tempdir().unwrap(); + let path = Utf8Path::from_path(dir.path()).unwrap(); + std::fs::write( + path.join("recording-0.cast"), + "{\"version\": 2}\n[1.0,\"o\",\"first\\r\\n\"]\n", + ) + .unwrap(); + std::fs::write( + path.join("recording-1.cast"), + "{\"version\": 2}\n[1.0,\"o\",\"second\\r\\n\"]\n", + ) + .unwrap(); + + // A manifest that lists the later recording first. + let manifest = RecordingManifest { + start_time: 100, + duration: 60, + recordings: Recordings::Terminal(vec![ + TerminalRecording::new("recording-1.cast", 130).unwrap(), + TerminalRecording::new("recording-0.cast", 100).unwrap(), + ]), + }; + let workspace = path.join(WORKSPACE_DIR); + let settings = WorkspaceSettings { + prompt_version: "test".to_owned(), + model: MODEL.to_owned(), + chunk_len: MAX_CHUNK_LEN, + }; + + let plan = prepare( + &manifest, + path, + &workspace, + &settings, + &Arc::new(AtomicU64::new(0)), + &CancellationToken::new(), + ) + .unwrap(); + + assert_eq!(plan.len(), 1); + assert_eq!( + std::fs::read_to_string(transcript::chunk_path(&workspace, 0)).unwrap(), + "[1.0] first\n[31.0] second\n" + ); +} +#[test] +fn halves_are_cut_on_the_line_boundary_closest_to_the_middle() { + let line = format!("[1.0] {}\n", "a".repeat(94)); + let part = line.repeat(30); + + let (first, second) = split_in_half(&part).unwrap(); + + assert_eq!(first.len(), 15 * line.len()); + assert_eq!(second.len(), 15 * line.len()); + + let uneven = format!("{line}{}", "b".repeat(1900)); + let (first, second) = split_in_half(&uneven).unwrap(); + assert_eq!(first, line); + assert_eq!(second, "b".repeat(1900)); + + assert!(split_in_half(&line.repeat(3)).is_none(), "too short"); + assert!(split_in_half(&"c".repeat(MIN_SPLIT_LEN * 2)).is_none(), "one line"); +} + +#[test] +fn usage_total_is_unknown_once_a_count_is() { + let usage = |input_tokens, output_tokens| Usage { + input_tokens, + output_tokens, + }; + + assert_eq!(add_usage(Some(usage(1, 2)), Some(usage(10, 20))), Some(usage(11, 22))); + assert_eq!(add_usage(Some(usage(1, 2)), None), None); + assert_eq!(add_usage(None, Some(usage(10, 20))), None); +} diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/transcript.rs b/crates/devolutions-gateway-ai/src/recording_analysis/transcript.rs new file mode 100644 index 000000000..d67385600 --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/transcript.rs @@ -0,0 +1,665 @@ +//! Turns the terminal recordings of a session into transcript chunk files sent to the AI. +//! +//! Each transcript line reads `[] `. +//! Recordings are read as a stream, so memory use depends on the chunk size, not on the recording size. + +use std::fs::File; +use std::io::{self, BufRead, BufReader, BufWriter, Write as _}; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; + +use camino::{Utf8Path, Utf8PathBuf}; +use tokio_util::sync::CancellationToken; +use tracing::warn; + +use super::{TerminalFormat, TerminalRecording}; + +/// Longest text kept from one terminal line, in characters. +const MAX_LINE_CHARS: usize = 500; + +/// Longest asciicast event read; a longer one is skipped. +const MAX_CAST_EVENT_LEN: usize = 1024 * 1024; + +#[derive(Debug, thiserror::Error)] +pub(super) enum TranscriptError { + #[error("failed to read {file_name}: {reason}")] + Invalid { file_name: String, reason: String }, + #[error("failed to write the transcript: {0}")] + Write(io::Error), + #[error("transcript cancelled before its end")] + Cancelled, +} + +pub(super) fn chunk_path(workspace: &Utf8Path, index: usize) -> Utf8PathBuf { + workspace.join(format!("transcript-{index:04}.txt")) +} + +/// Writes the transcript of `recordings`, read from `dir`, into chunk files of at most `max_len` bytes in `workspace`, +/// cut on line boundaries, and returns the number of chunks. +/// +/// Times are counted from `session_start`, in Unix seconds. The bytes read from the recordings are added to +/// `read_bytes`. +/// Returns [`TranscriptError::Cancelled`] soon after `cancel` is cancelled, leaving the chunk files incomplete. +pub(super) fn write_chunks( + recordings: &[TerminalRecording], + session_start: i64, + dir: &Utf8Path, + workspace: &Utf8Path, + max_len: usize, + read_bytes: &Arc, + cancel: &CancellationToken, +) -> Result { + let mut transcript = Transcript::new(ChunkWriter::new(workspace, max_len), cancel); + + for recording in recordings { + let to_transcript_error = |error: CastError| match error { + CastError::Read(reason) => TranscriptError::Invalid { + file_name: recording.file_name().to_owned(), + reason, + }, + CastError::Write(error) => TranscriptError::Write(error), + CastError::Cancelled => TranscriptError::Cancelled, + }; + + let offset = u32::try_from(recording.start_time().saturating_sub(session_start).max(0)).unwrap_or(u32::MAX); + let offset = f64::from(offset); + + let file = super::open_counted(&dir.join(recording.file_name()), read_bytes) + .map(BufReader::new) + .map_err(|error| to_transcript_error(CastError::Read(error.to_string())))?; + + match recording.format() { + TerminalFormat::Asciicast => transcript.add_cast(file, offset), + TerminalFormat::Trp => transcript.add_trp(file, offset), + } + .map_err(to_transcript_error)?; + } + + transcript.chunks.finish().map_err(TranscriptError::Write) +} +enum CastError { + Read(String), + Write(io::Error), + Cancelled, +} + +/// Writes lines into numbered chunk files, starting a new file when the next line would not fit. +struct ChunkWriter { + workspace: Utf8PathBuf, + max_len: usize, + count: usize, + current: Option>, + current_len: usize, +} + +impl ChunkWriter { + fn new(workspace: &Utf8Path, max_len: usize) -> Self { + Self { + workspace: workspace.to_owned(), + max_len, + count: 0, + current: None, + current_len: 0, + } + } + + fn write_line(&mut self, line: &str) -> io::Result<()> { + if self.current.is_some() && self.current_len + line.len() > self.max_len { + self.close_current()?; + } + + let current = match &mut self.current { + Some(current) => current, + None => { + let file = File::create(chunk_path(&self.workspace, self.count))?; + self.count += 1; + self.current_len = 0; + self.current.insert(BufWriter::new(file)) + } + }; + + current.write_all(line.as_bytes())?; + self.current_len += line.len(); + + Ok(()) + } + + fn close_current(&mut self) -> io::Result<()> { + if let Some(current) = self.current.take() { + current + .into_inner() + .map_err(io::IntoInnerError::into_error)? + .sync_all()?; + } + + Ok(()) + } + + fn finish(mut self) -> io::Result { + self.close_current()?; + Ok(self.count) + } +} + +struct Transcript<'a> { + chunks: ChunkWriter, + last_line: String, + cancel: &'a CancellationToken, +} + +impl<'a> Transcript<'a> { + fn new(chunks: ChunkWriter, cancel: &'a CancellationToken) -> Self { + Self { + chunks, + last_line: String::new(), + cancel, + } + } + + /// Adds the terminal output of an asciicast v2 or v3 recording that started `offset` seconds into the session. + fn add_cast(&mut self, mut cast: impl BufRead, offset: f64) -> Result<(), CastError> { + let mut line = Vec::new(); + + let events = core::iter::from_fn(|| match read_event(&mut cast, &mut line) { + Ok(true) => Some(Ok(core::mem::take(&mut line))), + Ok(false) => None, + Err(error) => Some(Err(CastError::Read(error.to_string()))), + }); + + self.add_cast_events(events, offset) + } + + /// Adds the terminal output of a TRP recording that started `offset` seconds into the session. + fn add_trp(&mut self, trp: impl io::Read, offset: f64) -> Result<(), CastError> { + let events = terminal_streamer::trp_decoder::AsciicastLines::new(trp).map(|line| { + line.map(String::into_bytes) + .map_err(|error| CastError::Read(format!("{error:#}"))) + }); + + self.add_cast_events(events, offset) + } + + /// Adds the terminal output of asciicast lines, the header first. + fn add_cast_events( + &mut self, + mut lines: impl Iterator, CastError>>, + offset: f64, + ) -> Result<(), CastError> { + let header = loop { + let line = lines + .next() + .ok_or_else(|| CastError::Read("empty asciicast".to_owned()))??; + + if !line.trim_ascii().is_empty() { + break serde_json::from_slice::(&line) + .map_err(|error| CastError::Read(format!("invalid asciicast header: {error}")))?; + } + }; + + let relative_times = header["version"].as_u64() == Some(3); + + let mut terminal = TerminalText::default(); + let mut time = 0.0; + + for line in lines { + if self.cancel.is_cancelled() { + return Err(CastError::Cancelled); + } + + // Only the output is used: typed passwords are usually not echoed, but printed secrets still reach the AI. + let Ok((event_time, code, data)) = serde_json::from_slice::<(f64, String, String)>(&line?) else { + continue; + }; + + time = if relative_times { time + event_time } else { event_time }; + + if code == "o" { + terminal.feed(offset + time, &data, self).map_err(CastError::Write)?; + } + } + + terminal.end_line(self).map_err(CastError::Write) + } + + fn push_line(&mut self, time: f64, text: &str) -> io::Result<()> { + let text = text.trim(); + + if text.is_empty() || text == self.last_line { + return Ok(()); + } + + let text = match text.char_indices().nth(MAX_LINE_CHARS) { + Some((cut, _)) => &text[..cut], + None => text, + }; + + self.chunks.write_line(&format!("[{time:.1}] {text}\n"))?; + text.clone_into(&mut self.last_line); + + Ok(()) + } +} + +/// Reads the next asciicast event into `line`, skipping events longer than [`MAX_CAST_EVENT_LEN`]. +fn read_event(reader: &mut impl BufRead, line: &mut Vec) -> io::Result { + let limit = u64::try_from(MAX_CAST_EVENT_LEN).unwrap_or(u64::MAX); + + loop { + line.clear(); + + if io::Read::take(&mut *reader, limit).read_until(b'\n', line)? == 0 { + return Ok(false); + } + + if line.len() < MAX_CAST_EVENT_LEN || line.ends_with(b"\n") { + return Ok(true); + } + + warn!(max_len = MAX_CAST_EVENT_LEN, "Skipped an oversized asciicast event"); + skip_line(reader)?; + } +} + +fn skip_line(reader: &mut impl BufRead) -> io::Result<()> { + loop { + let buffer = reader.fill_buf()?; + + if buffer.is_empty() { + return Ok(()); + } + + if let Some(end) = buffer.iter().position(|&byte| byte == b'\n') { + reader.consume(end + 1); + return Ok(()); + } + + let len = buffer.len(); + reader.consume(len); + } +} + +#[derive(Default, Clone, Copy)] +enum EscapeState { + #[default] + Text, + Escape, + EscapeArgument, + ControlSequence, + OperatingSystemCommand, + OperatingSystemCommandEscape, +} + +/// Rebuilds the text lines of a terminal output stream, without escape sequences. +#[derive(Default)] +struct TerminalText { + state: EscapeState, + line: String, + line_start: Option, + carriage_return: bool, +} + +impl TerminalText { + fn feed(&mut self, time: f64, data: &str, transcript: &mut Transcript<'_>) -> io::Result<()> { + for c in data.chars() { + self.state = match (self.state, c) { + (EscapeState::Text, '\x1b') => EscapeState::Escape, + (EscapeState::Text, '\n') => { + self.end_line(transcript)?; + EscapeState::Text + } + (EscapeState::Text, '\r') => { + self.carriage_return = true; + EscapeState::Text + } + (EscapeState::Text, '\x08') => { + self.line.pop(); + EscapeState::Text + } + (EscapeState::Text, '\t') => { + self.put(' ', time); + EscapeState::Text + } + (EscapeState::Text, c) => { + if !c.is_control() { + self.put(c, time); + } + EscapeState::Text + } + (EscapeState::Escape, '[') => EscapeState::ControlSequence, + (EscapeState::Escape, ']') => EscapeState::OperatingSystemCommand, + (EscapeState::Escape, '(' | ')' | '*' | '+' | '#' | '%') => EscapeState::EscapeArgument, + (EscapeState::Escape | EscapeState::EscapeArgument | EscapeState::OperatingSystemCommandEscape, _) => { + EscapeState::Text + } + (EscapeState::ControlSequence, '\x40'..='\x7e') => EscapeState::Text, + (EscapeState::ControlSequence, _) => EscapeState::ControlSequence, + (EscapeState::OperatingSystemCommand, '\x07') => EscapeState::Text, + (EscapeState::OperatingSystemCommand, '\x1b') => EscapeState::OperatingSystemCommandEscape, + (EscapeState::OperatingSystemCommand, _) => EscapeState::OperatingSystemCommand, + }; + } + + Ok(()) + } + + fn put(&mut self, c: char, time: f64) { + // A carriage return not followed by a line feed means the line is redrawn. + if self.carriage_return { + self.carriage_return = false; + self.line.clear(); + self.line_start = None; + } + + self.line_start.get_or_insert(time); + + // Text past the kept length is dropped anyway, so a line that never ends cannot grow without bound. + if self.line.len() < MAX_LINE_CHARS * 4 { + self.line.push(c); + } + } + + fn end_line(&mut self, transcript: &mut Transcript<'_>) -> io::Result<()> { + self.carriage_return = false; + + if let Some(start) = self.line_start.take() { + transcript.push_line(start, &self.line)?; + } + + self.line.clear(); + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + fn read_chunks(dir: &Utf8Path, count: usize) -> Vec { + (0..count) + .map(|index| std::fs::read_to_string(chunk_path(dir, index)).expect("chunk file")) + .collect() + } + + fn transcript_of(cast: &str, offset: f64) -> String { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); + let cancel = CancellationToken::new(); + let mut transcript = Transcript::new(ChunkWriter::new(dir, usize::MAX), &cancel); + assert!(transcript.add_cast(cast.as_bytes(), offset).is_ok()); + let count = transcript.chunks.finish().expect("written"); + read_chunks(dir, count).concat() + } + + const CAST: &str = r#"{"version": 2, "width": 80, "height": 24} +[0.3,"o","\u001b]0;user@host: ~\u0007\u001b[01;32muser@host\u001b[00m:~$ "] +[1.0,"i","l"] +[1.1,"o","l"] +[1.2,"o","s\r\n"] +[1.5,"o","file-a file-b\r\n"] +[2.0,"o","user@host:~$ sudo passwd david\r\n[sudo] password for user: "] +[3.0,"i","hunter2\r"] +[3.5,"o","\r\n"] +[4.0,"o","progress 10%\rprogress 100%\r\n"] +[5.0,"o","same\r\nsame\r\n"] +[6.0,"o","typo\b\b\u001b[Kps\r\n"] +[7.0,"r","100x30"] +"#; + + #[test] + fn cast_output_becomes_timed_plain_lines() { + assert_eq!( + transcript_of(CAST, 10.0), + concat!( + "[10.3] user@host:~$ ls\n", + "[11.5] file-a file-b\n", + "[12.0] user@host:~$ sudo passwd david\n", + "[12.0] [sudo] password for user:\n", + "[14.0] progress 100%\n", + "[15.0] same\n", + "[16.0] typs\n", + ) + ); + } + + #[test] + fn asciicast_v3_times_are_intervals() { + let cast = "{\"version\": 3, \"term\": {\"cols\": 80, \"rows\": 24}}\n[1.0,\"o\",\"a\\r\\n\"]\n[0.5,\"o\",\"b\\r\\n\"]\n"; + + assert_eq!(transcript_of(cast, 0.0), "[1.0] a\n[1.5] b\n"); + } + + #[test] + fn long_lines_are_cut() { + let cast = format!( + "{{\"version\": 2}}\n[0,\"o\",\"{}\"]\n", + "é".repeat(MAX_LINE_CHARS + 10) + ); + + let transcript = transcript_of(&cast, 0.0); + + assert_eq!(transcript.trim_end().chars().count(), "[0.0] ".len() + MAX_LINE_CHARS); + } + + #[test] + fn oversized_cast_events_are_skipped() { + let cast = format!( + "{{\"version\": 2}}\n[0,\"o\",\"{}\"]\n[1,\"o\",\"kept\\r\\n\"]\n", + "x".repeat(MAX_CAST_EVENT_LEN) + ); + + assert_eq!(transcript_of(&cast, 0.0), "[1.0] kept\n"); + } + + #[test] + fn chunks_keep_whole_lines_within_the_budget() { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); + let mut chunks = ChunkWriter::new(dir, 22); + for line in ["[1.0] aaaa\n", "[2.0] bbbb\n", "[3.0] cccc\n"] { + chunks.write_line(line).expect("written"); + } + + let count = chunks.finish().expect("written"); + + assert_eq!(read_chunks(dir, count), ["[1.0] aaaa\n[2.0] bbbb\n", "[3.0] cccc\n"]); + assert_eq!(ChunkWriter::new(dir, 22).finish().expect("nothing written"), 0); + } + + /// Generates a large recording lazily and records how many events were read when the first chunk file was complete. + struct LazyRecording { + next_event: usize, + events: usize, + event: fn(usize) -> Vec, + pending: Vec, + second_chunk: Utf8PathBuf, + read_when_first_chunk_done: Arc, + } + + impl LazyRecording { + fn new(dir: &Utf8Path, events: usize, event: fn(usize) -> Vec) -> (Self, Arc) { + let read_when_first_chunk_done = Arc::new(AtomicUsize::new(0)); + let recording = Self { + next_event: 0, + events, + event, + pending: Vec::new(), + second_chunk: chunk_path(dir, 1), + read_when_first_chunk_done: Arc::clone(&read_when_first_chunk_done), + }; + (recording, read_when_first_chunk_done) + } + } + + impl io::Read for LazyRecording { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + if self.pending.is_empty() { + if self.next_event == self.events { + return Ok(0); + } + + if self.second_chunk.exists() { + let _ = self.read_when_first_chunk_done.compare_exchange( + 0, + self.next_event, + Ordering::Relaxed, + Ordering::Relaxed, + ); + } + + self.pending = (self.event)(self.next_event); + self.next_event += 1; + } + + let n = buf.len().min(self.pending.len()); + buf[..n].copy_from_slice(&self.pending[..n]); + self.pending.drain(..n); + Ok(n) + } + } + + fn cast_event(event: usize) -> Vec { + if event == 0 { + b"{\"version\": 2}\n".to_vec() + } else { + format!("[{event}.0,\"o\",\"line {event} of a long session\\r\\n\"]\n").into_bytes() + } + } + + fn trp_packet(time_delta: u32, event_type: u16, payload: &[u8]) -> Vec { + let mut packet = Vec::new(); + packet.extend_from_slice(&time_delta.to_le_bytes()); + packet.extend_from_slice(&event_type.to_le_bytes()); + packet.extend_from_slice(&u16::try_from(payload.len()).expect("small").to_le_bytes()); + packet.extend_from_slice(payload); + packet + } + + fn trp_event(event: usize) -> Vec { + if event == 0 { + trp_packet(0, 4, b"") + } else { + trp_packet(1000, 0, format!("line {event} of a long session\r\n").as_bytes()) + } + } + + fn assert_streamed_into_bounded_chunk_files(event: fn(usize) -> Vec, trp: bool) { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); + let events = 20_000; + let (recording, read_when_first_chunk_done) = LazyRecording::new(dir, events, event); + + let cancel = CancellationToken::new(); + let mut transcript = Transcript::new(ChunkWriter::new(dir, 4096), &cancel); + let recording = BufReader::with_capacity(64, recording); + let added = if trp { + transcript.add_trp(recording, 0.0) + } else { + transcript.add_cast(recording, 0.0) + }; + assert!(added.is_ok()); + let count = transcript.chunks.finish().expect("written"); + + assert!(count > 100, "{count}"); + for chunk in read_chunks(dir, count) { + assert!(chunk.len() <= 4096); + assert!(chunk.ends_with('\n')); + } + + let read = read_when_first_chunk_done.load(Ordering::Relaxed); + assert!( + read > 0 && read < events / 50, + "first chunk done after {read} of {events} events" + ); + } + + #[test] + fn large_casts_are_streamed_into_bounded_chunk_files() { + assert_streamed_into_bounded_chunk_files(cast_event, false); + } + + #[test] + fn large_trp_recordings_are_streamed_into_bounded_chunk_files() { + assert_streamed_into_bounded_chunk_files(trp_event, true); + } + + #[test] + fn transcript_stops_once_cancelled() { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); + let cancel = CancellationToken::new(); + cancel.cancel(); + let (recording, _) = LazyRecording::new(dir, 1000, cast_event); + + let mut transcript = Transcript::new(ChunkWriter::new(dir, 4096), &cancel); + let result = transcript.add_cast(BufReader::new(recording), 0.0); + + assert!(matches!(result, Err(CastError::Cancelled))); + } + + fn session_dir(files: &[(&str, &[u8])]) -> tempfile::TempDir { + let dir = tempfile::tempdir().expect("temp dir"); + for (name, contents) in files { + std::fs::write(dir.path().join(name), contents).expect("write recording"); + } + dir + } + + fn recordings(files: &[(&str, i64)]) -> Vec { + files + .iter() + .map(|&(file_name, start_time)| TerminalRecording::new(file_name, start_time).expect("terminal recording")) + .collect() + } + + fn read( + recordings: &[TerminalRecording], + session_start: i64, + dir: &tempfile::TempDir, + ) -> Result { + let path = Utf8Path::from_path(dir.path()).expect("utf8 path"); + let workspace = path.join("workspace"); + std::fs::create_dir(&workspace).expect("workspace"); + let count = write_chunks( + recordings, + session_start, + path, + &workspace, + usize::MAX, + &Arc::new(AtomicU64::new(0)), + &CancellationToken::new(), + )?; + Ok(read_chunks(&workspace, count).concat()) + } + #[test] + fn recordings_are_offset_by_their_start_time() { + let first: &[u8] = b"{\"version\": 2}\n[1.5,\"o\",\"whoami\\r\\n\"]\n"; + let second: &[u8] = b"{\"version\": 2}\n[2.0,\"o\",\"exit\\r\\n\"]\n"; + let dir = session_dir(&[("recording-0.cast", first), ("recording-1.cast", second)]); + + assert_eq!( + read( + &recordings(&[("recording-0.cast", 100), ("recording-1.cast", 160)]), + 100, + &dir + ) + .expect("transcript"), + "[1.5] whoami\n[62.0] exit\n" + ); + } + + #[test] + fn trp_recordings_are_decoded() { + let mut trp = trp_packet(0, 4, b""); + trp.extend(trp_packet(2500, 0, b"uptime\r\n")); + let dir = session_dir(&[("recording-0.trp", &trp)]); + + assert_eq!( + read(&recordings(&[("recording-0.trp", 100)]), 100, &dir).expect("transcript"), + "[2.5] uptime\n" + ); + } +} diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs new file mode 100644 index 000000000..5ca93be10 --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs @@ -0,0 +1,697 @@ +//! Picks the screenshots of a video recording worth showing to the AI: the frames where the screen changed. +//! +//! The video is decoded with XMF and sampled at [`SAMPLES_PER_SECOND`]. A sample is compared with the previous one in +//! 16x16 tiles of a half-resolution gray image. Tiles that keep changing, such as a clock or a spinner, count as +//! animated and are ignored. A frame is kept at the start, once the screen settles after a change, and regularly while +//! a change goes on. +//! Only the current and the previous samples are held in memory, whatever the length of the video. + +use std::collections::VecDeque; +use std::fs::File; +use std::io::{self, BufReader, BufWriter}; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::time::Duration; + +use cadeau::xmf::vpx::{VpxCodec, VpxDecoder, VpxImage}; +use camino::{Utf8Path, Utf8PathBuf}; +use tokio_util::sync::CancellationToken; +use webm_iterable::WebmIterator; +use webm_iterable::errors::TagIteratorError; +use webm_iterable::matroska_spec::{Block, Master, MatroskaSpec, SimpleBlock}; + +/// Decoded samples per second of video. +const SAMPLES_PER_SECOND: u32 = 2; + +const SAMPLE_INTERVAL_MS: u64 = 1000 / SAMPLES_PER_SECOND as u64; + +/// Side of a tile, in pixels of the half-resolution gray image. +const TILE: usize = 16; + +/// Mean absolute gray difference over which a tile counts as changed. +const TILE_DIFF: u32 = 3; + +const TILE_PIXELS: u32 = 16 * 16; + +/// Samples of history used to find animated tiles (10 s). +const ANIMATION_WINDOW: usize = 20; + +/// A tile that changed in at least this many samples of the window is animated. +const ANIMATION_HITS: usize = 3; + +/// A change ends after this many quiet samples (1 s). +const SETTLE_SAMPLES: usize = 2; + +/// During a long change, one frame is kept every this many samples (3 s). +const KEEP_DURING_CHANGE: usize = 6; + +/// Largest screenshot sent to the AI. +const MAX_WIDTH: usize = 1280; +const MAX_HEIGHT: usize = 720; + +#[derive(Debug, thiserror::Error)] +pub(super) enum VideoError { + #[error("video decoding is not available: XMF is not loaded")] + DecoderUnavailable, + #[error("{0}")] + Invalid(String), + #[error("failed to write a screenshot: {0}")] + Write(io::Error), + #[error("cancelled")] + Cancelled, +} + +/// A screenshot kept from a video, saved as `.png` in the working files. +#[derive(Debug, Clone, PartialEq)] +pub(super) struct KeptFrame { + pub(super) id: String, + /// Elapsed time since the start of the session. + pub(super) offset: Duration, +} + +pub(super) fn frame_path(workspace: &Utf8Path, id: &str) -> Utf8PathBuf { + workspace.join(format!("{id}.png")) +} + +/// Keeps the frames of the video at `path` where the screen changed, as PNG files in `workspace`. +/// +/// `offset` is when the video starts in the session. `next_id` numbers the frames across the videos of a session. +/// The bytes read from the video are added to `read_bytes`. +pub(super) fn keep_changed_frames( + path: &Utf8Path, + offset: Duration, + workspace: &Utf8Path, + next_id: &mut usize, + read_bytes: &Arc, + cancel: &CancellationToken, +) -> Result, VideoError> { + if !cadeau::xmf::is_init() { + return Err(VideoError::DecoderUnavailable); + } + + let file = super::open_counted(path, read_bytes) + .map_err(|error| VideoError::Invalid(format!("failed to open the video: {error}")))?; + let mut tags = WebmIterator::new(BufReader::new(file), &[MatroskaSpec::BlockGroup(Master::Start)]); + tags.emit_master_end_when_eof(false); + + let mut keeper = Keeper { + selector: FrameSelector::default(), + offset, + workspace, + next_id, + kept: Vec::new(), + current: None, + }; + let mut track = TrackInfo::default(); + let mut decoder: Option = None; + let mut cluster_time = None; + let mut timestamp_scale_ns = 1_000_000u64; + let mut next_sample_ms = 0u64; + let mut previous: Option<(u64, Picture)> = None; + + for tag in &mut tags { + if cancel.is_cancelled() { + return Err(VideoError::Cancelled); + } + + let tag = match tag { + Ok(tag) => tag, + // A recording cut short, such as after a crash, ends without its last element. + Err(TagIteratorError::UnexpectedEOF { .. }) => break, + Err(error) => return Err(VideoError::Invalid(format!("invalid WebM: {error}"))), + }; + + let block = match tag { + MatroskaSpec::TimestampScale(scale) => { + timestamp_scale_ns = scale.max(1); + continue; + } + MatroskaSpec::CodecID(id) => { + track.codec = Some(match id.as_str() { + "V_VP8" => VpxCodec::VP8, + "V_VP9" => VpxCodec::VP9, + other => return Err(VideoError::Invalid(format!("unsupported video codec {other}"))), + }); + continue; + } + MatroskaSpec::PixelWidth(width) => { + track.width = u32::try_from(width).ok(); + continue; + } + MatroskaSpec::PixelHeight(height) => { + track.height = u32::try_from(height).ok(); + continue; + } + MatroskaSpec::Timestamp(time) => { + cluster_time = Some(time); + continue; + } + MatroskaSpec::SimpleBlock(data) => SimpleBlock::try_from(&data) + .map_err(|error| VideoError::Invalid(format!("invalid block: {error}"))) + .and_then(|block| frame_of(block.timestamp, block.read_frame_data()))?, + MatroskaSpec::BlockGroup(Master::Full(children)) => { + let Some(raw) = children.iter().find_map(|child| match child { + MatroskaSpec::Block(raw) => Some(raw), + _ => None, + }) else { + continue; + }; + Block::try_from(raw) + .map_err(|error| VideoError::Invalid(format!("invalid block: {error}"))) + .and_then(|block| frame_of(block.timestamp, block.read_frame_data()))? + } + _ => continue, + }; + + let Some(cluster_time) = cluster_time else { + return Err(VideoError::Invalid("block before its cluster time".to_owned())); + }; + + let ticks = i64::try_from(cluster_time) + .unwrap_or(i64::MAX) + .saturating_add(i64::from(block.0)); + let pts_ms = u64::try_from(ticks).unwrap_or(0).saturating_mul(timestamp_scale_ns) / 1_000_000; + + let decoder = match &mut decoder { + Some(decoder) => decoder, + None => decoder.insert(track.decoder()?), + }; + + decoder + .decode(&block.1) + .map_err(|error| VideoError::Invalid(format!("failed to decode a frame: {error:?}")))?; + + // Exactly one call per decoded frame: XMF would hand back the same image again. + let Ok(image) = decoder.next_frame() else { + continue; + }; + let picture = Picture::of(&image)?; + + // Every sample time before this frame shows the previous one. + if let Some((_, shown)) = &previous { + while next_sample_ms < pts_ms { + keeper.sample(next_sample_ms, shown)?; + next_sample_ms += SAMPLE_INTERVAL_MS; + } + } else { + next_sample_ms = pts_ms; + } + + previous = Some((pts_ms, picture)); + } + + if let Some((pts_ms, shown)) = &previous { + // The last frame stays on screen until the end: it gets at least one sample. + loop { + keeper.sample(next_sample_ms, shown)?; + if *pts_ms < next_sample_ms + SAMPLE_INTERVAL_MS { + break; + } + next_sample_ms += SAMPLE_INTERVAL_MS; + } + + keeper.finish(shown)?; + } + + Ok(keeper.kept) +} + +/// The time of a block relative to its cluster, and its frame. +fn frame_of( + timestamp: i16, + frames: Result>, webm_iterable::errors::WebmCoercionError>, +) -> Result<(i16, Vec), VideoError> { + let frames = frames.map_err(|error| VideoError::Invalid(format!("invalid block: {error}")))?; + + match frames.as_slice() { + [frame] => Ok((timestamp, frame.data.to_vec())), + _ => Err(VideoError::Invalid("laced blocks are not supported".to_owned())), + } +} + +#[derive(Default)] +struct TrackInfo { + codec: Option, + width: Option, + height: Option, +} + +impl TrackInfo { + fn decoder(&self) -> Result { + let (Some(codec), Some(width), Some(height)) = (self.codec, self.width, self.height) else { + return Err(VideoError::Invalid("video track without codec or size".to_owned())); + }; + + VpxDecoder::builder() + .threads(2) + .width(width) + .height(height) + .codec(codec) + .build() + .map_err(|error| VideoError::Invalid(format!("failed to start the video decoder: {error:?}"))) + } +} + +/// A decoded frame, with its I420 planes copied out of the decoder. +struct Picture { + width: usize, + height: usize, + y: Vec, + u: Vec, + v: Vec, +} + +impl Picture { + fn of(image: &VpxImage<'_>) -> Result { + let planes = image + .i420_planes() + .ok_or_else(|| VideoError::Invalid("decoded frame is not 8-bit I420".to_owned()))?; + + let copy = |plane: cadeau::xmf::vpx::VpxPlane<'_>| { + let mut out = Vec::with_capacity(plane.width() * plane.height()); + for row in plane.rows() { + out.extend_from_slice(&row[..plane.width()]); + } + out + }; + + Ok(Self { + width: planes.y.width(), + height: planes.y.height(), + y: copy(planes.y), + u: copy(planes.u), + v: copy(planes.v), + }) + } + + /// Every other luma pixel of every other row. + fn half_gray(&self) -> Gray { + let width = self.width / 2; + let height = self.height / 2; + let mut pixels = Vec::with_capacity(width * height); + + for row in self.y.chunks_exact(self.width).step_by(2).take(height) { + pixels.extend(row.iter().step_by(2).take(width)); + } + + Gray { width, height, pixels } + } + + /// Writes the picture as a PNG no larger than [`MAX_WIDTH`]x[`MAX_HEIGHT`], averaging the pixels it shrinks. + /// + /// Colors are read as BT.709 full range, which is what Gateway and RDM record. + fn write_png(&self, path: &Utf8Path) -> io::Result<()> { + // Shrink to fit both limits, keeping the aspect ratio. + let (out_width, out_height) = if self.width * MAX_HEIGHT >= self.height * MAX_WIDTH { + let out_width = self.width.min(MAX_WIDTH); + (out_width, (self.height * out_width / self.width).max(1)) + } else { + let out_height = self.height.min(MAX_HEIGHT); + ((self.width * out_height / self.height).max(1), out_height) + }; + let chroma_width = self.width.div_ceil(2); + + let mut rgb = Vec::with_capacity(out_width * out_height * 3); + + for out_y in 0..out_height { + let y0 = out_y * self.height / out_height; + let y1 = ((out_y + 1) * self.height / out_height).max(y0 + 1); + + for out_x in 0..out_width { + let x0 = out_x * self.width / out_width; + let x1 = ((out_x + 1) * self.width / out_width).max(x0 + 1); + + let (mut luma, mut cb, mut cr, mut count) = (0u32, 0u32, 0u32, 0u32); + for y in y0..y1 { + for x in x0..x1 { + let chroma = (y / 2) * chroma_width + x / 2; + luma += u32::from(self.y[y * self.width + x]); + cb += u32::from(self.u[chroma]); + cr += u32::from(self.v[chroma]); + count += 1; + } + } + + let luma = f64::from(luma) / f64::from(count); + let cb = f64::from(cb) / f64::from(count) - 128.0; + let cr = f64::from(cr) / f64::from(count) - 128.0; + + rgb.push(to_byte(luma + 1.5748 * cr)); + rgb.push(to_byte(luma - 0.1873 * cb - 0.4681 * cr)); + rgb.push(to_byte(luma + 1.8556 * cb)); + } + } + + let out = BufWriter::new(File::create(path)?); + let size = |value: usize| u32::try_from(value).map_err(io::Error::other); + let mut encoder = png::Encoder::new(out, size(out_width)?, size(out_height)?); + encoder.set_color(png::ColorType::Rgb); + encoder.set_depth(png::BitDepth::Eight); + encoder.set_compression(png::Compression::Fast); + encoder + .write_header() + .and_then(|mut writer| writer.write_image_data(&rgb)) + .map_err(io::Error::other) + } +} + +#[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "the value is clamped to 0..=255 first" +)] +fn to_byte(value: f64) -> u8 { + value.round().clamp(0.0, 255.0) as u8 +} + +/// Writes the kept frames, numbering them across the videos of a session. +struct Keeper<'a> { + selector: FrameSelector, + offset: Duration, + workspace: &'a Utf8Path, + next_id: &'a mut usize, + kept: Vec, + /// The last sample, kept if a change is still going on when the video ends. + current: Option<(u64, String)>, +} + +impl Keeper<'_> { + fn sample(&mut self, time_ms: u64, picture: &Picture) -> Result<(), VideoError> { + let id = format!("F{:05}", *self.next_id); + *self.next_id += 1; + + if self.selector.sample(picture.half_gray()) { + self.keep(time_ms, &id, picture)?; + } + + self.current = Some((time_ms, id)); + Ok(()) + } + + fn keep(&mut self, time_ms: u64, id: &str, picture: &Picture) -> Result<(), VideoError> { + if self.kept.last().is_some_and(|last| last.id == id) { + return Ok(()); + } + + picture + .write_png(&frame_path(self.workspace, id)) + .map_err(VideoError::Write)?; + self.kept.push(KeptFrame { + id: id.to_owned(), + offset: self.offset + Duration::from_millis(time_ms), + }); + Ok(()) + } + + /// Keeps the last sample, which shows last, if the screen was still changing when the video ended. + fn finish(&mut self, last: &Picture) -> Result<(), VideoError> { + match self.current.take() { + Some((time_ms, id)) if self.selector.in_change() => self.keep(time_ms, &id, last), + _ => Ok(()), + } + } +} + +/// Half-resolution gray image of a sample. +struct Gray { + width: usize, + height: usize, + pixels: Vec, +} + +/// Decides which samples of a video to keep, from their gray images, in order. +#[derive(Default)] +struct FrameSelector { + previous: Option, + /// Changed tiles of the last samples, to find animated ones. + history: VecDeque>, + in_change: bool, + change_samples: usize, + quiet_samples: usize, +} + +impl FrameSelector { + /// Returns `true` when the sample is worth keeping. + fn sample(&mut self, gray: Gray) -> bool { + let Some(previous) = self.previous.replace(gray) else { + // The first sample shows where the video starts. + return true; + }; + let gray = self.previous.as_ref().expect("just set"); + + let active = if previous.width == gray.width && previous.height == gray.height { + let changed = changed_tiles(&previous, gray); + let animated = self.animated_tiles(changed.len(), gray.width / TILE); + + if self.history.len() == ANIMATION_WINDOW { + self.history.pop_front(); + } + let active = changed + .iter() + .zip(&animated) + .any(|(&changed, &animated)| changed && !animated); + self.history.push_back(changed); + active + } else { + // A new screen size changes everything, and the old history no longer lines up. + self.history.clear(); + true + }; + + if active { + if !self.in_change { + self.in_change = true; + self.change_samples = 0; + } + self.change_samples += 1; + self.quiet_samples = 0; + return self.change_samples.is_multiple_of(KEEP_DURING_CHANGE); + } + + if self.in_change { + self.quiet_samples += 1; + if self.quiet_samples >= SETTLE_SAMPLES { + self.in_change = false; + return true; + } + } + + false + } + + /// Whether a change is still going on, so that the last sample should be kept. + fn in_change(&self) -> bool { + self.in_change + } + + /// Tiles that changed in at least [`ANIMATION_HITS`] samples of the history, grown by one tile in every direction. + fn animated_tiles(&self, tiles: usize, columns: usize) -> Vec { + let mut hits = vec![0usize; tiles]; + for changed in self.history.iter().filter(|changed| changed.len() == tiles) { + for (hits, &changed) in hits.iter_mut().zip(changed) { + *hits += usize::from(changed); + } + } + + let animated: Vec = hits.iter().map(|&hits| hits >= ANIMATION_HITS).collect(); + let mut grown = animated.clone(); + + if columns == 0 { + return grown; + } + + for (index, _) in animated.iter().enumerate().filter(|(_, animated)| **animated) { + let (row, column) = (index / columns, index % columns); + let rows = tiles / columns; + + for (neighbor_row, neighbor_column) in [ + (row.wrapping_sub(1), column), + (row + 1, column), + (row, column.wrapping_sub(1)), + (row, column + 1), + ] { + if neighbor_row < rows && neighbor_column < columns { + grown[neighbor_row * columns + neighbor_column] = true; + } + } + } + + grown + } +} + +/// Which 16x16 tiles changed between two gray images of the same size, row by row; partial tiles are ignored. +fn changed_tiles(before: &Gray, after: &Gray) -> Vec { + let columns = after.width / TILE; + let rows = after.height / TILE; + let mut changed = Vec::with_capacity(columns * rows); + + for row in 0..rows { + for column in 0..columns { + let mut sum = 0u32; + for y in row * TILE..(row + 1) * TILE { + let start = y * after.width + column * TILE; + for x in start..start + TILE { + sum += u32::from(before.pixels[x].abs_diff(after.pixels[x])); + } + } + changed.push(sum > TILE_DIFF * TILE_PIXELS); + } + } + + changed +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 16 tiles wide and 2 tiles high. + const WIDTH: usize = 256; + const HEIGHT: usize = 32; + + /// A gray image with the 16x16 tiles in `bright` set to white, row by row. + fn screen(bright: &[usize]) -> Gray { + let mut pixels = vec![0; WIDTH * HEIGHT]; + let columns = WIDTH / TILE; + + for &tile in bright { + let (row, column) = (tile / columns, tile % columns); + for y in row * TILE..(row + 1) * TILE { + for x in column * TILE..(column + 1) * TILE { + pixels[y * WIDTH + x] = 255; + } + } + } + + Gray { + width: WIDTH, + height: HEIGHT, + pixels, + } + } + + /// Indexes of the kept samples. + fn kept(samples: impl IntoIterator) -> Vec { + let mut selector = FrameSelector::default(); + samples + .into_iter() + .enumerate() + .filter_map(|(index, sample)| selector.sample(sample).then_some(index)) + .collect() + } + + #[test] + fn a_still_screen_keeps_only_its_first_frame() { + assert_eq!(kept((0..20).map(|_| screen(&[]))), [0]); + } + + #[test] + fn a_change_keeps_the_frame_once_the_screen_settles() { + let samples = [ + screen(&[]), + screen(&[]), + screen(&[1]), + screen(&[1]), + screen(&[1]), + screen(&[1]), + ]; + + // The change is at 2; the screen is quiet at 3 and 4, so 4 is kept. + assert_eq!(kept(samples), [0, 4]); + } + + #[test] + fn a_long_change_keeps_a_frame_every_three_seconds() { + // Each sample lights a different tile, so every sample changes, and no tile changes often enough to count as + // animated. + let samples = (0..14).map(|index| screen(&[index])); + + assert_eq!(kept(samples), [0, 6, 12]); + } + + #[test] + fn a_blinking_tile_stops_counting_as_a_change() { + // Tile 5 blinks every sample: after three changes it is animated and ignored. + let samples = (0..12).map(|index| screen(if index % 2 == 0 { &[] } else { &[5] })); + let kept = kept(samples); + + assert!(kept.len() <= 3, "{kept:?}"); + assert_eq!(kept[0], 0); + } + + #[test] + fn a_new_screen_size_is_a_change() { + let small = Gray { + width: 32, + height: 32, + pixels: vec![0; 32 * 32], + }; + + // Two size changes in a row, then the screen settles after two quiet samples. + assert_eq!( + kept([screen(&[]), small, screen(&[]), screen(&[]), screen(&[])]), + [0, 4] + ); + } + + #[test] + fn small_differences_are_not_changes() { + let mut nearly = screen(&[]); + nearly.pixels.iter_mut().for_each(|pixel| *pixel = 2); + + assert_eq!(changed_tiles(&screen(&[]), &nearly), vec![false; 32]); + assert_eq!( + changed_tiles(&screen(&[]), &screen(&[3])) + .iter() + .filter(|changed| **changed) + .count(), + 1 + ); + } + + fn picture(width: usize, height: usize) -> Picture { + Picture { + width, + height, + y: vec![128; width * height], + u: vec![128; width.div_ceil(2) * height.div_ceil(2)], + v: vec![128; width.div_ceil(2) * height.div_ceil(2)], + } + } + + fn png_size(png: &Utf8Path) -> (u32, u32) { + let decoder = png::Decoder::new(BufReader::new(File::open(png).expect("png"))); + let reader = decoder.read_info().expect("PNG header"); + let info = reader.info(); + (info.width, info.height) + } + + #[test] + fn screenshots_shrink_to_fit_1280_by_720() { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); + + for ((width, height), expected) in [ + ((1920, 1080), (1280, 720)), + ((3840, 2400), (1152, 720)), + ((800, 600), (800, 600)), + ] { + let path = dir.join(format!("{width}x{height}.png")); + picture(width, height).write_png(&path).expect("written"); + + assert_eq!(png_size(&path), expected, "{width}x{height}"); + } + } + + #[test] + fn half_gray_takes_every_other_pixel() { + let mut picture = picture(4, 2); + picture.y = vec![10, 20, 30, 40, 50, 60, 70, 80]; + + let gray = picture.half_gray(); + + assert_eq!((gray.width, gray.height), (2, 1)); + assert_eq!(gray.pixels, [10, 30]); + } +} From b5653d50504f917b2e0fb25ff8ad01e878538421 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 7 Oct 2026 20:34:03 -0400 Subject: [PATCH 04/18] docs(dgw): describe recording analysis in the AI crate intent --- crates/devolutions-gateway-ai/INTENT.md | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/crates/devolutions-gateway-ai/INTENT.md b/crates/devolutions-gateway-ai/INTENT.md index d0a74ce18..e65555498 100644 --- a/crates/devolutions-gateway-ai/INTENT.md +++ b/crates/devolutions-gateway-ai/INTENT.md @@ -1,3 +1,19 @@ +# Background + The AI crate serves as an HTTP client and a runner for tasks that are related to AI. We do not want to serve it as a generic LLM client only. -The crate provides methods for specific purposes, for example, to summarize a terminal session, or to generate AI analysis for a given recording video file. \ No newline at end of file +The crate provides methods for specific purposes, for example, to summarize a terminal session, or to generate AI analysis for a given recording video file. + +## Tasks + +### Analyze Recording + +Analyze recording takes a parsed manifest file and its path to a recording folder. +It returns a stream of events that contains the progress of the analysis, and final results of the analysis and the potential errors that happened during the analysis. + +The final result is a series of JSON objects with timestamps and a natural language description of the events that happened during the recording. + +Analyze recording is trivial for terminal recordings. We just send the recording files to the AI service and get the results back. + +Analyze video recordings breaks down the video into frames, and we apply motion detection altorithms to filter out idle or close enough frames. Then we send the the frames in chronological order to the AI services in batches. + From 07158420388e6f274a475b938fc26323c1c24769 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 11:59:12 -0400 Subject: [PATCH 05/18] fix(dgw): read only the video track of a WebM recording Every CodecID was taken as the video codec, so a WebM with an audio track, such as VP8 video with Opus audio, was refused as an unsupported video codec, and blocks were decoded whatever their track. The first video track is now picked from the track entries, and only its blocks are decoded; audio and other tracks are skipped. Checked locally with XMF on a real session remuxed with an Opus track before and after the video: the same 32 screenshots as the original. --- .../src/recording_analysis/video.rs | 144 +++++++++++++++--- 1 file changed, 119 insertions(+), 25 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs index 5ca93be10..a2e502695 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs @@ -102,7 +102,8 @@ pub(super) fn keep_changed_frames( kept: Vec::new(), current: None, }; - let mut track = TrackInfo::default(); + let mut entry: Option = None; + let mut video: Option = None; let mut decoder: Option = None; let mut cluster_time = None; let mut timestamp_scale_ns = 1_000_000u64; @@ -126,20 +127,48 @@ pub(super) fn keep_changed_frames( timestamp_scale_ns = scale.max(1); continue; } + MatroskaSpec::TrackEntry(Master::Start) => { + entry = Some(TrackEntry::default()); + continue; + } + MatroskaSpec::TrackNumber(number) => { + if let Some(entry) = &mut entry { + entry.number = Some(number); + } + continue; + } + MatroskaSpec::TrackType(kind) => { + if let Some(entry) = &mut entry { + entry.kind = Some(kind); + } + continue; + } MatroskaSpec::CodecID(id) => { - track.codec = Some(match id.as_str() { - "V_VP8" => VpxCodec::VP8, - "V_VP9" => VpxCodec::VP9, - other => return Err(VideoError::Invalid(format!("unsupported video codec {other}"))), - }); + if let Some(entry) = &mut entry { + entry.codec_id = Some(id); + } continue; } MatroskaSpec::PixelWidth(width) => { - track.width = u32::try_from(width).ok(); + if let Some(entry) = &mut entry { + entry.width = u32::try_from(width).ok(); + } continue; } MatroskaSpec::PixelHeight(height) => { - track.height = u32::try_from(height).ok(); + if let Some(entry) = &mut entry { + entry.height = u32::try_from(height).ok(); + } + continue; + } + MatroskaSpec::TrackEntry(Master::End) => { + // Only the first video track is read; audio, subtitles and any other video track are skipped. + if let Some(entry) = entry.take() + && video.is_none() + && entry.is_video() + { + video = Some(VideoTrack::try_from(entry)?); + } continue; } MatroskaSpec::Timestamp(time) => { @@ -148,7 +177,7 @@ pub(super) fn keep_changed_frames( } MatroskaSpec::SimpleBlock(data) => SimpleBlock::try_from(&data) .map_err(|error| VideoError::Invalid(format!("invalid block: {error}"))) - .and_then(|block| frame_of(block.timestamp, block.read_frame_data()))?, + .and_then(|block| frame_of(block.track, block.timestamp, block.read_frame_data()))?, MatroskaSpec::BlockGroup(Master::Full(children)) => { let Some(raw) = children.iter().find_map(|child| match child { MatroskaSpec::Block(raw) => Some(raw), @@ -158,27 +187,35 @@ pub(super) fn keep_changed_frames( }; Block::try_from(raw) .map_err(|error| VideoError::Invalid(format!("invalid block: {error}"))) - .and_then(|block| frame_of(block.timestamp, block.read_frame_data()))? + .and_then(|block| frame_of(block.track, block.timestamp, block.read_frame_data()))? } _ => continue, }; + let Some(video) = &video else { + return Err(VideoError::Invalid("block before any video track".to_owned())); + }; + + if block.track != video.number { + continue; + } + let Some(cluster_time) = cluster_time else { return Err(VideoError::Invalid("block before its cluster time".to_owned())); }; let ticks = i64::try_from(cluster_time) .unwrap_or(i64::MAX) - .saturating_add(i64::from(block.0)); + .saturating_add(i64::from(block.timestamp)); let pts_ms = u64::try_from(ticks).unwrap_or(0).saturating_mul(timestamp_scale_ns) / 1_000_000; let decoder = match &mut decoder { Some(decoder) => decoder, - None => decoder.insert(track.decoder()?), + None => decoder.insert(video.decoder()?), }; decoder - .decode(&block.1) + .decode(&block.data) .map_err(|error| VideoError::Invalid(format!("failed to decode a frame: {error:?}")))?; // Exactly one call per decoded frame: XMF would hand back the same image again. @@ -216,37 +253,94 @@ pub(super) fn keep_changed_frames( Ok(keeper.kept) } -/// The time of a block relative to its cluster, and its frame. +/// A frame of a block, with the track it belongs to and its time relative to its cluster. +struct BlockFrame { + track: u64, + timestamp: i16, + data: Vec, +} + fn frame_of( + track: u64, timestamp: i16, frames: Result>, webm_iterable::errors::WebmCoercionError>, -) -> Result<(i16, Vec), VideoError> { +) -> Result { let frames = frames.map_err(|error| VideoError::Invalid(format!("invalid block: {error}")))?; match frames.as_slice() { - [frame] => Ok((timestamp, frame.data.to_vec())), + [frame] => Ok(BlockFrame { + track, + timestamp, + data: frame.data.to_vec(), + }), _ => Err(VideoError::Invalid("laced blocks are not supported".to_owned())), } } +/// Matroska `TrackType` of a video track. +const VIDEO_TRACK_TYPE: u64 = 1; + +/// What a `TrackEntry` says about its track, of any kind. #[derive(Default)] -struct TrackInfo { - codec: Option, +struct TrackEntry { + number: Option, + kind: Option, + codec_id: Option, width: Option, height: Option, } -impl TrackInfo { - fn decoder(&self) -> Result { - let (Some(codec), Some(width), Some(height)) = (self.codec, self.width, self.height) else { - return Err(VideoError::Invalid("video track without codec or size".to_owned())); +impl TrackEntry { + fn is_video(&self) -> bool { + match self.kind { + Some(kind) => kind == VIDEO_TRACK_TYPE, + None => self.codec_id.as_deref().is_some_and(|id| id.starts_with("V_")), + } + } +} + +/// The video track the screenshots are taken from. +struct VideoTrack { + number: u64, + codec: VpxCodec, + width: u32, + height: u32, +} + +impl TryFrom for VideoTrack { + type Error = VideoError; + + fn try_from(entry: TrackEntry) -> Result { + let (Some(number), Some(codec_id), Some(width), Some(height)) = + (entry.number, entry.codec_id, entry.width, entry.height) + else { + return Err(VideoError::Invalid( + "video track without number, codec or size".to_owned(), + )); }; + let codec = match codec_id.as_str() { + "V_VP8" => VpxCodec::VP8, + "V_VP9" => VpxCodec::VP9, + other => return Err(VideoError::Invalid(format!("unsupported video codec {other}"))), + }; + + Ok(Self { + number, + codec, + width, + height, + }) + } +} + +impl VideoTrack { + fn decoder(&self) -> Result { VpxDecoder::builder() .threads(2) - .width(width) - .height(height) - .codec(codec) + .width(self.width) + .height(self.height) + .codec(self.codec) .build() .map_err(|error| VideoError::Invalid(format!("failed to start the video decoder: {error:?}"))) } From 32dcab8c025a2845cd6f52a8cd73aa66b0adb976 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 11:59:12 -0400 Subject: [PATCH 06/18] fix(dgw): always report the reading progress first The first reading progress came from an interval tick racing the end of the reading, so a small recording could be read before any progress was sent. The start of the reading is now reported before the files are read. --- .../src/recording_analysis/mod.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs index b0ba21126..1d7d7eb34 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -412,7 +412,15 @@ impl Run { let read_bytes = Arc::clone(&read_bytes); move || prepare(&manifest, &dir, &workspace, &settings, &read_bytes, &cancel) }); - let mut ticks = tokio::time::interval(PREPARING_PROGRESS_INTERVAL); + // The first event of an analysis is always the start of the reading, however quickly the files are read. + self.report(Progress::Reading { + read_bytes: 0, + total_bytes, + }); + let mut ticks = tokio::time::interval_at( + tokio::time::Instant::now() + PREPARING_PROGRESS_INTERVAL, + PREPARING_PROGRESS_INTERVAL, + ); // Decoding a long video takes a while, so the bytes read so far are reported until it is done. let plan = loop { From 57e71898f8ce911a17658f3f79bbbec7a4fa1ce7 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 11:59:12 -0400 Subject: [PATCH 07/18] docs(dgw): remove a control character from a doc comment --- crates/devolutions-gateway-ai/src/wire/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/devolutions-gateway-ai/src/wire/mod.rs b/crates/devolutions-gateway-ai/src/wire/mod.rs index 7887e176e..32b1aa6ac 100644 --- a/crates/devolutions-gateway-ai/src/wire/mod.rs +++ b/crates/devolutions-gateway-ai/src/wire/mod.rs @@ -56,7 +56,7 @@ enum Content<'a> { } impl<'a> Content<'a> { - /// Builds the content of input, writing each part with lock. + /// Builds the content of input, writing each part with `block`. fn of(input: &[Input<'a>], block: impl Fn(&Input<'a>) -> Value) -> Self { match input { [Input::Text(text)] => Self::Text(text), From 56b01c1b4b515200f33c12299258f721436ac296 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 12:03:54 -0400 Subject: [PATCH 08/18] docs(dgw): describe how terminal recordings are prepared in the AI crate intent --- crates/devolutions-gateway-ai/INTENT.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/devolutions-gateway-ai/INTENT.md b/crates/devolutions-gateway-ai/INTENT.md index e65555498..512705dc3 100644 --- a/crates/devolutions-gateway-ai/INTENT.md +++ b/crates/devolutions-gateway-ai/INTENT.md @@ -13,7 +13,7 @@ It returns a stream of events that contains the progress of the analysis, and fi The final result is a series of JSON objects with timestamps and a natural language description of the events that happened during the recording. -Analyze recording is trivial for terminal recordings. We just send the recording files to the AI service and get the results back. +Analyze recording is trivial for terminal recordings. We pre-process the .trp file to .cast, and strip out some noises, then send as text to the AI services. Analyze video recordings breaks down the video into frames, and we apply motion detection altorithms to filter out idle or close enough frames. Then we send the the frames in chronological order to the AI services in batches. From f38a8d5a285139768bb43761dadc730077435d8a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 12:51:07 -0400 Subject: [PATCH 09/18] fix(dgw): keep the provider URL out of AI transport errors reqwest errors repeat the URL of the request, so a transport error named the provider's endpoint, which callers store in Task results and logs. Transport errors now drop the URL, and the request debug log no longer names the base URL. --- crates/devolutions-gateway-ai/src/client.rs | 5 ++--- crates/devolutions-gateway-ai/src/error.rs | 10 +++++++-- .../devolutions-gateway-ai/tests/providers.rs | 22 ++++++++++++++++--- 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/client.rs b/crates/devolutions-gateway-ai/src/client.rs index 7316d2fe0..c9fff4489 100644 --- a/crates/devolutions-gateway-ai/src/client.rs +++ b/crates/devolutions-gateway-ai/src/client.rs @@ -243,7 +243,6 @@ impl AiClient { debug!( provider = ?self.provider, model = %self.model, - base_url = %self.base_url, input_len = prompt .input .iter() @@ -273,13 +272,13 @@ impl AiClient { .timeout(self.request_timeout) .send() .await - .map_err(|error| error::transport(&error, api_key))?; + .map_err(|error| error::transport(error, api_key))?; let status = response.status(); let retry_after = error::retry_after(response.headers()); let body = read_body(response) .await - .map_err(|error| error::transport(&error, api_key))?; + .map_err(|error| error::transport(error, api_key))?; if !status.is_success() { // An error body over the limit is not parsed, so the error holds the reason phrase of the status. diff --git a/crates/devolutions-gateway-ai/src/error.rs b/crates/devolutions-gateway-ai/src/error.rs index a918b59dd..f01b67661 100644 --- a/crates/devolutions-gateway-ai/src/error.rs +++ b/crates/devolutions-gateway-ai/src/error.rs @@ -77,7 +77,13 @@ const SPENT_QUOTA_CODES: [&str; 6] = [ "enforced_spend_limit_reached", ]; -pub(crate) fn transport(error: &dyn std::error::Error, api_key: &str) -> Error { +/// Builds a transport error that names neither the API key nor the URL of the request. +pub(crate) fn transport(error: reqwest::Error, api_key: &str) -> Error { + // The URL names the endpoint of the provider, which callers keep out of their logs and records. + redacted_transport(&error.without_url(), api_key) +} + +fn redacted_transport(error: &dyn std::error::Error, api_key: &str) -> Error { let mut message = error.to_string(); let mut source = error.source(); while let Some(cause) = source { @@ -131,7 +137,7 @@ mod tests { fn transport_error_redacts_api_key() { let error = std::io::Error::other(format!("invalid key {API_KEY} provided")); - let error = transport(&error, API_KEY); + let error = redacted_transport(&error, API_KEY); assert!(!error.to_string().contains(API_KEY)); assert!(!format!("{error:?}").contains(API_KEY)); diff --git a/crates/devolutions-gateway-ai/tests/providers.rs b/crates/devolutions-gateway-ai/tests/providers.rs index 1c0926111..790744e9e 100644 --- a/crates/devolutions-gateway-ai/tests/providers.rs +++ b/crates/devolutions-gateway-ai/tests/providers.rs @@ -855,6 +855,19 @@ async fn error_answer_over_the_size_limit_keeps_the_status() { } } +/// A base URL whose host, port and path a transport error must not repeat. +fn private_url(addr: std::net::SocketAddr) -> Url { + Url::parse(&format!("http://{addr}/private-endpoint/v1/")).unwrap() +} + +// Errors end up in logs and in stored Task results, where the endpoint of the provider must not appear. +fn assert_url_hidden(error: &Error, addr: std::net::SocketAddr) { + for shown in [error.to_string(), format!("{error:?}")] { + assert!(!shown.contains("private-endpoint"), "{shown}"); + assert!(!shown.contains(&addr.to_string()), "{shown}"); + } +} + #[tokio::test] async fn answer_cut_by_a_closed_connection_is_transient() { let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap(); @@ -869,7 +882,7 @@ async fn answer_cut_by_a_closed_connection_is_transient() { .unwrap(); }); - let error = client(Provider::OpenAi, Url::parse(&format!("http://{addr}/v1/")).unwrap()) + let error = client(Provider::OpenAi, private_url(addr)) .describe_session_actions("[0] whoami") .send() .await @@ -877,6 +890,7 @@ async fn answer_cut_by_a_closed_connection_is_transient() { assert!(matches!(error, Error::Transport { .. }), "{error:?}"); assert!(error.is_transient()); + assert_url_hidden(&error, addr); } #[tokio::test] @@ -885,7 +899,7 @@ async fn unreachable_provider_is_transient() { let addr = listener.local_addr().unwrap(); drop(listener); - let error = client(Provider::OpenAi, Url::parse(&format!("http://{addr}/v1/")).unwrap()) + let error = client(Provider::OpenAi, private_url(addr)) .describe_session_actions("[0] whoami") .send() .await @@ -893,6 +907,7 @@ async fn unreachable_provider_is_transient() { assert!(matches!(error, Error::Transport { .. }), "{error:?}"); assert!(error.is_transient()); + assert_url_hidden(&error, addr); } #[tokio::test] @@ -915,7 +930,7 @@ async fn silent_provider_times_out_as_transient() { .provider(Provider::OpenAiCompatible) .model(MODEL) .api_key(API_KEY) - .base_url(Url::parse(&format!("http://{addr}/v1/")).unwrap()) + .base_url(private_url(addr)) .http_client(reqwest::Client::builder().no_proxy().build().unwrap()) .request_timeout(Duration::from_millis(200)) .build() @@ -927,6 +942,7 @@ async fn silent_provider_times_out_as_transient() { assert!(matches!(error, Error::Transport { .. }), "{error:?}"); assert!(error.is_transient()); + assert_url_hidden(&error, addr); assert!(started.elapsed() < Duration::from_secs(5), "{:?}", started.elapsed()); } From 00560de127c6fb7e036a45ad95560efa6afba91d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:57:37 -0400 Subject: [PATCH 10/18] fix(dgw): bound WebM tag and frame sizes in the AI crate The WebM parser zero-fills the size an element claims before reading it, up to 4 GB by default, so a small hostile recording could claim gigabytes. Elements larger than the file are now refused. The cap is the file length, not a fixed constant, because remuxed recordings have a known Segment size as large as the file and the cap applies to every element. Decoded frames larger than 8192x8192 are refused before they are copied. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/recording_analysis/video.rs | 55 +++++++++++++++++-- 1 file changed, 50 insertions(+), 5 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs index a2e502695..1e4411166 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs @@ -8,7 +8,7 @@ use std::collections::VecDeque; use std::fs::File; -use std::io::{self, BufReader, BufWriter}; +use std::io::{self, BufReader, BufWriter, Read}; use std::sync::Arc; use std::sync::atomic::AtomicU64; use std::time::Duration; @@ -49,6 +49,9 @@ const KEEP_DURING_CHANGE: usize = 6; const MAX_WIDTH: usize = 1280; const MAX_HEIGHT: usize = 720; +/// Largest side of a decoded frame, in pixels; it bounds the memory a frame takes. +const MAX_FRAME_SIDE: usize = 8192; + #[derive(Debug, thiserror::Error)] pub(super) enum VideoError { #[error("video decoding is not available: XMF is not loaded")] @@ -91,8 +94,12 @@ pub(super) fn keep_changed_frames( let file = super::open_counted(path, read_bytes) .map_err(|error| VideoError::Invalid(format!("failed to open the video: {error}")))?; - let mut tags = WebmIterator::new(BufReader::new(file), &[MatroskaSpec::BlockGroup(Master::Start)]); - tags.emit_master_end_when_eof(false); + let len = file + .inner + .metadata() + .map_err(|error| VideoError::Invalid(format!("failed to open the video: {error}")))? + .len(); + let mut tags = webm_tags(BufReader::new(file), len); let mut keeper = Keeper { selector: FrameSelector::default(), @@ -253,6 +260,17 @@ pub(super) fn keep_changed_frames( Ok(keeper.kept) } +/// Reads the WebM elements of `source`, a file of `len` bytes. +/// +/// Every element starts with the size it claims, and the parser sets that much memory aside before reading it. A broken +/// or hostile file could claim gigabytes, so any element bigger than the whole file is refused up front. +fn webm_tags(source: R, len: u64) -> WebmIterator { + let mut tags = WebmIterator::new(source, &[MatroskaSpec::BlockGroup(Master::Start)]); + tags.emit_master_end_when_eof(false); + tags.set_max_allowable_tag_size(Some(usize::try_from(len).unwrap_or(usize::MAX))); + tags +} + /// A frame of a block, with the track it belongs to and its time relative to its cluster. struct BlockFrame { track: u64, @@ -361,6 +379,13 @@ impl Picture { .i420_planes() .ok_or_else(|| VideoError::Invalid("decoded frame is not 8-bit I420".to_owned()))?; + let (width, height) = (planes.y.width(), planes.y.height()); + if width > MAX_FRAME_SIDE || height > MAX_FRAME_SIDE { + return Err(VideoError::Invalid(format!( + "decoded frame of {width}x{height} is larger than {MAX_FRAME_SIDE}x{MAX_FRAME_SIDE}" + ))); + } + let copy = |plane: cadeau::xmf::vpx::VpxPlane<'_>| { let mut out = Vec::with_capacity(plane.width() * plane.height()); for row in plane.rows() { @@ -370,8 +395,8 @@ impl Picture { }; Ok(Self { - width: planes.y.width(), - height: planes.y.height(), + width, + height, y: copy(planes.y), u: copy(planes.u), v: copy(planes.v), @@ -788,4 +813,24 @@ mod tests { assert_eq!((gray.width, gray.height), (2, 1)); assert_eq!(gray.pixels, [10, 30]); } + + #[test] + fn an_element_larger_than_its_file_is_refused() { + let webm = [ + // Cluster of unknown size. + &[0x1F, 0x43, 0xB6, 0x75, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF][..], + // SimpleBlock claiming 1 MiB. + &[0xA3, 0x10, 0x10, 0x00, 0x00], + &[0x81, 0x00, 0x00, 0x80], + ] + .concat(); + let mut tags = webm_tags(webm.as_slice(), u64::try_from(webm.len()).expect("small")); + + assert!(matches!(tags.next(), Some(Ok(MatroskaSpec::Cluster(Master::Start))))); + let next = tags.next(); + assert!( + matches!(next, Some(Err(TagIteratorError::CorruptedFileData(_)))), + "{next:?}" + ); + } } From 7c977bff120b529cb5a494240a1f18ef48fc3282 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:58:45 -0400 Subject: [PATCH 11/18] fix(dgw): check cancellation while sampling a long video frame A frame that stays on screen for a long time, as in an idle session, is sampled many times. Each sample now checks the cancellation token, and the half-resolution gray image of a frame is computed on its first sample only, instead of on every sample. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/recording_analysis/video.rs | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs index 1e4411166..b7c875df5 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs @@ -9,6 +9,7 @@ use std::collections::VecDeque; use std::fs::File; use std::io::{self, BufReader, BufWriter, Read}; +use std::rc::Rc; use std::sync::Arc; use std::sync::atomic::AtomicU64; use std::time::Duration; @@ -233,8 +234,14 @@ pub(super) fn keep_changed_frames( // Every sample time before this frame shows the previous one. if let Some((_, shown)) = &previous { + let mut gray = None; while next_sample_ms < pts_ms { - keeper.sample(next_sample_ms, shown)?; + if cancel.is_cancelled() { + return Err(VideoError::Cancelled); + } + + let gray = gray.get_or_insert_with(|| Rc::new(shown.half_gray())); + keeper.sample(next_sample_ms, shown, gray)?; next_sample_ms += SAMPLE_INTERVAL_MS; } } else { @@ -245,9 +252,11 @@ pub(super) fn keep_changed_frames( } if let Some((pts_ms, shown)) = &previous { + let gray = Rc::new(shown.half_gray()); + // The last frame stays on screen until the end: it gets at least one sample. loop { - keeper.sample(next_sample_ms, shown)?; + keeper.sample(next_sample_ms, shown, &gray)?; if *pts_ms < next_sample_ms + SAMPLE_INTERVAL_MS { break; } @@ -495,11 +504,11 @@ struct Keeper<'a> { } impl Keeper<'_> { - fn sample(&mut self, time_ms: u64, picture: &Picture) -> Result<(), VideoError> { + fn sample(&mut self, time_ms: u64, picture: &Picture, gray: &Rc) -> Result<(), VideoError> { let id = format!("F{:05}", *self.next_id); *self.next_id += 1; - if self.selector.sample(picture.half_gray()) { + if self.selector.sample(Rc::clone(gray)) { self.keep(time_ms, &id, picture)?; } @@ -541,7 +550,7 @@ struct Gray { /// Decides which samples of a video to keep, from their gray images, in order. #[derive(Default)] struct FrameSelector { - previous: Option, + previous: Option>, /// Changed tiles of the last samples, to find animated ones. history: VecDeque>, in_change: bool, @@ -551,7 +560,7 @@ struct FrameSelector { impl FrameSelector { /// Returns `true` when the sample is worth keeping. - fn sample(&mut self, gray: Gray) -> bool { + fn sample(&mut self, gray: Rc) -> bool { let Some(previous) = self.previous.replace(gray) else { // The first sample shows where the video starts. return true; @@ -696,7 +705,7 @@ mod tests { samples .into_iter() .enumerate() - .filter_map(|(index, sample)| selector.sample(sample).then_some(index)) + .filter_map(|(index, sample)| selector.sample(Rc::new(sample)).then_some(index)) .collect() } From 182d84158975c122fe795b1bd293ab86fdf2e9e4 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:59:53 -0400 Subject: [PATCH 12/18] fix(terminal-streamer): count TRP packet headers toward the setup cap Events held until the terminal setup were counted by payload size only, so a flood of empty packets was held whole in memory. Each packet now counts its 8-byte header too. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/terminal-streamer/src/trp_decoder.rs | 36 +++++++++++---------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/crates/terminal-streamer/src/trp_decoder.rs b/crates/terminal-streamer/src/trp_decoder.rs index 77c9d09ed..b9342e707 100644 --- a/crates/terminal-streamer/src/trp_decoder.rs +++ b/crates/terminal-streamer/src/trp_decoder.rs @@ -175,7 +175,7 @@ impl PacketHeader { } } -/// Most payload bytes held while waiting for the terminal setup; past it, the header is sent with the size known so far. +/// Most packet bytes held while waiting for the terminal setup; past it, the header is sent with the size known so far. const MAX_BEFORE_SETUP_LEN: usize = 1024 * 1024; /// Turns TRP packets into asciicast lines, holding the events that come before the terminal setup. @@ -216,7 +216,7 @@ impl TrpDecoder { match self.before_setup_cache { Some(ref mut cache) => { cache.push(event); - self.before_setup_len += packet_len; + self.before_setup_len += PACKET_HEADER_SIZE + packet_len; // A recording that never sends its terminal setup must not be held whole in memory. if self.before_setup_len > MAX_BEFORE_SETUP_LEN { @@ -416,23 +416,25 @@ mod tests { #[test] fn events_waiting_for_the_terminal_setup_are_bounded() { - let payload = vec![b'x'; 60_000]; - let packets = MAX_BEFORE_SETUP_LEN / payload.len() + 1; - let mut trp = Vec::new(); - for _ in 0..packets { - trp.extend(packet(0, 0, &payload)); - } + // Empty packets cost their header, or a flood of them would be held whole. + for payload in [vec![b'x'; 60_000], Vec::new()] { + let packets = MAX_BEFORE_SETUP_LEN / (PACKET_HEADER_SIZE + payload.len()) + 1; + let mut trp = Vec::new(); + for _ in 0..packets { + trp.extend(packet(0, 0, &payload)); + } - // The setup never comes: the rest of the recording fails to read. - let lines = AsciicastLines::new(Read::chain(trp.as_slice(), FailingReader)).collect::>(); + // The setup never comes: the rest of the recording fails to read. + let lines = AsciicastLines::new(Read::chain(trp.as_slice(), FailingReader)).collect::>(); - assert_eq!(lines.len(), 1 + packets + 1, "header, every event, then the read error"); - assert_eq!( - lines[0].as_ref().expect("valid header"), - "{\"version\": 2, \"width\": 80, \"height\": 24}" - ); - assert!(lines[1..=packets].iter().all(Result::is_ok)); - assert!(lines[packets + 1].is_err()); + assert_eq!(lines.len(), 1 + packets + 1, "header, every event, then the read error"); + assert_eq!( + lines[0].as_ref().expect("valid header"), + "{\"version\": 2, \"width\": 80, \"height\": 24}" + ); + assert!(lines[1..=packets].iter().all(Result::is_ok)); + assert!(lines[packets + 1].is_err()); + } } #[test] From 5c57c3f7ee5095dbcdc044282c6d3e7043f3f2bd Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 21:00:18 -0400 Subject: [PATCH 13/18] refactor(dgw): drop the casts from the video sampling constants Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/recording_analysis/video.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs index b7c875df5..bbd5fa6c7 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs @@ -22,17 +22,17 @@ use webm_iterable::errors::TagIteratorError; use webm_iterable::matroska_spec::{Block, Master, MatroskaSpec, SimpleBlock}; /// Decoded samples per second of video. -const SAMPLES_PER_SECOND: u32 = 2; +const SAMPLES_PER_SECOND: u64 = 2; -const SAMPLE_INTERVAL_MS: u64 = 1000 / SAMPLES_PER_SECOND as u64; +const SAMPLE_INTERVAL_MS: u64 = 1000 / SAMPLES_PER_SECOND; /// Side of a tile, in pixels of the half-resolution gray image. const TILE: usize = 16; /// Mean absolute gray difference over which a tile counts as changed. -const TILE_DIFF: u32 = 3; +const TILE_DIFF: usize = 3; -const TILE_PIXELS: u32 = 16 * 16; +const TILE_PIXELS: usize = TILE * TILE; /// Samples of history used to find animated tiles (10 s). const ANIMATION_WINDOW: usize = 20; @@ -656,11 +656,11 @@ fn changed_tiles(before: &Gray, after: &Gray) -> Vec { for row in 0..rows { for column in 0..columns { - let mut sum = 0u32; + let mut sum = 0usize; for y in row * TILE..(row + 1) * TILE { let start = y * after.width + column * TILE; for x in start..start + TILE { - sum += u32::from(before.pixels[x].abs_diff(after.pixels[x])); + sum += usize::from(before.pixels[x].abs_diff(after.pixels[x])); } } changed.push(sum > TILE_DIFF * TILE_PIXELS); From 134dffa9faba0eda52054bc162a642b9cb7ab7be Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 21:00:38 -0400 Subject: [PATCH 14/18] refactor(dgw): show video decoder errors with their cause VpxError displays only "decoder error"; the libvpx code is in its source. The messages now show the Display chain instead of the Debug output. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/recording_analysis/video.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs index bbd5fa6c7..c86f5aac7 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs @@ -222,9 +222,9 @@ pub(super) fn keep_changed_frames( None => decoder.insert(video.decoder()?), }; - decoder - .decode(&block.data) - .map_err(|error| VideoError::Invalid(format!("failed to decode a frame: {error:?}")))?; + decoder.decode(&block.data).map_err(|error| { + VideoError::Invalid(format!("failed to decode a frame: {:#}", anyhow::Error::new(error))) + })?; // Exactly one call per decoded frame: XMF would hand back the same image again. let Ok(image) = decoder.next_frame() else { @@ -369,7 +369,12 @@ impl VideoTrack { .height(self.height) .codec(self.codec) .build() - .map_err(|error| VideoError::Invalid(format!("failed to start the video decoder: {error:?}"))) + .map_err(|error| { + VideoError::Invalid(format!( + "failed to start the video decoder: {:#}", + anyhow::Error::new(error) + )) + }) } } From a016fc8a6fb9971e5d79c9e8c11814c3524b349f Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 21:01:02 -0400 Subject: [PATCH 15/18] refactor(dgw): log AI analysis steps at debug level The library logs its steps at debug; the caller decides what is worth an info line. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../devolutions-gateway-ai/src/recording_analysis/mod.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs index 1d7d7eb34..f360a50b5 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -39,7 +39,7 @@ use camino::{Utf8Path, Utf8PathBuf}; use serde::{Deserialize, Serialize}; use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; -use tracing::{debug, info, warn}; +use tracing::{debug, warn}; use self::checkpoint::DescribedChunk; pub use self::recordings::{ @@ -438,7 +438,7 @@ impl Run { } }; let total = plan.len(); - info!(chunks = total, "Describe the session actions"); + debug!(chunks = total, "Describe the session actions"); for (index, chunk) in plan.iter().enumerate() { if cancel.is_cancelled() { @@ -801,7 +801,7 @@ fn prepare( return Ok(plan); } } else if workspace.exists() { - info!(path = %workspace, "Discard the working files of another AI analysis"); + debug!(path = %workspace, "Discard the working files of another AI analysis"); std::fs::remove_dir_all(workspace).map_err(|error| files_error(&error))?; } @@ -894,7 +894,7 @@ fn plan_screenshots( VideoError::Cancelled => RecordingAnalysisError::Cancelled, })?; - info!(file = %recording.file_name(), screenshots = kept.len(), "Screenshots kept from a video"); + debug!(file = %recording.file_name(), screenshots = kept.len(), "Screenshots kept from a video"); frames.extend(kept.into_iter().map(|frame| PlannedFrame { id: frame.id, From 74f78f8e75183d06b6713a2c5529173e71505dea Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 21:01:27 -0400 Subject: [PATCH 16/18] refactor(dgw): keep blocking file checks off the async AI analysis task The size of the recordings is read on a blocking thread, and the check for a chunk already described uses tokio::fs. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../devolutions-gateway-ai/src/recording_analysis/mod.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs index f360a50b5..66b1dd542 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -401,7 +401,12 @@ impl Run { chunk_len: MAX_CHUNK_LEN, }; - let total_bytes = self.manifest.recordings.total_bytes(&dir); + let total_bytes = blocking({ + let recordings = self.manifest.recordings.clone(); + let dir = dir.clone(); + move || Ok(recordings.total_bytes(&dir)) + }) + .await?; let read_bytes = Arc::new(AtomicU64::new(0)); let mut preparing = tokio::task::spawn_blocking({ @@ -447,7 +452,7 @@ impl Run { let checkpoint = checkpoint::path(&workspace, index); - if checkpoint.exists() { + if tokio::fs::try_exists(&checkpoint).await.unwrap_or(false) { debug!(index, "Chunk already described"); continue; } From 0079cc5ab326caaf8beb12910c264663e4274562 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:06:31 -0400 Subject: [PATCH 17/18] fix(dgw): close the AI analysis events on a panic and bound long video pauses - A panic in the analysis left the watcher holding a sender, so the events channel never ended and the caller waited forever. The watcher now stops through a drop guard, so the channel ends as documented. - Dropping the receiver cancelled the caller's own token; it now cancels a child token, so a shared token is never cancelled for other work. - A pause between two video frames was sampled every half second, so a corrupt timestamp far ahead spun until the attempt timeout. Past the animation window one picture no longer changes the frame selector, so the rest of the pause is skipped and counted, with the same frame IDs. - PNG frames are finished explicitly, so a write or flush error is reported instead of dropped. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/recording_analysis/mod.rs | 9 ++- .../src/recording_analysis/video.rs | 56 ++++++++++++++++++- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs index 66b1dd542..7f65c614e 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -328,11 +328,13 @@ impl AnalyzeRecording<'_> { pub fn start(self) -> mpsc::Receiver { let (events, receiver) = mpsc::channel(EVENT_BUFFER); let finished = CancellationToken::new(); + // A child, so a receiver that goes away never cancels a token the caller shares with other work. + let cancel = self.cancel.child_token(); // A receiver that went away means nobody waits for the result anymore. tokio::spawn({ let events = events.clone(); - let cancel = self.cancel.clone(); + let cancel = cancel.clone(); let finished = finished.clone(); async move { tokio::select! { @@ -347,16 +349,17 @@ impl AnalyzeRecording<'_> { manifest: self.manifest.clone(), manifest_path: self.manifest_path.to_owned(), max_output_tokens: self.max_output_tokens, - cancel: self.cancel, + cancel, events, }; tokio::spawn(async move { + // Also stops the watcher when the analysis panics, so the channel closes instead of staying open forever. + let _finished = finished.drop_guard(); let last = match run.run().await { Ok(analysis) => RecordingAnalysisEvent::Complete(analysis), Err(error) => RecordingAnalysisEvent::Failure(error), }; - finished.cancel(); let _ = run.events.send(last).await; }); diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs index c86f5aac7..ec0ec8ae8 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/video.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs @@ -235,14 +235,25 @@ pub(super) fn keep_changed_frames( // Every sample time before this frame shows the previous one. if let Some((_, shown)) = &previous { let mut gray = None; + let mut repeats = 0; while next_sample_ms < pts_ms { if cancel.is_cancelled() { return Err(VideoError::Cancelled); } + // Past this many samples of one picture the selector no longer changes, so a long pause, or a + // corrupt time far ahead, costs no more work. + if repeats > ANIMATION_WINDOW { + let skipped = (pts_ms - next_sample_ms).div_ceil(SAMPLE_INTERVAL_MS); + next_sample_ms += skipped * SAMPLE_INTERVAL_MS; + keeper.skip(skipped, next_sample_ms - SAMPLE_INTERVAL_MS); + break; + } + let gray = gray.get_or_insert_with(|| Rc::new(shown.half_gray())); keeper.sample(next_sample_ms, shown, gray)?; next_sample_ms += SAMPLE_INTERVAL_MS; + repeats += 1; } } else { next_sample_ms = pts_ms; @@ -483,7 +494,10 @@ impl Picture { encoder.set_compression(png::Compression::Fast); encoder .write_header() - .and_then(|mut writer| writer.write_image_data(&rgb)) + .and_then(|mut writer| { + writer.write_image_data(&rgb)?; + writer.finish() + }) .map_err(io::Error::other) } } @@ -521,6 +535,13 @@ impl Keeper<'_> { Ok(()) } + /// Counts `samples` that would change nothing, the last one at `last_ms`, as if each was taken. + fn skip(&mut self, samples: u64, last_ms: u64) { + let samples = usize::try_from(samples).unwrap_or(usize::MAX); + *self.next_id = self.next_id.saturating_add(samples); + self.current = Some((last_ms, format!("F{:05}", *self.next_id - 1))); + } + fn keep(&mut self, time_ms: u64, id: &str, picture: &Picture) -> Result<(), VideoError> { if self.kept.last().is_some_and(|last| last.id == id) { return Ok(()); @@ -753,6 +774,39 @@ mod tests { assert_eq!(kept[0], 0); } + #[test] + fn one_picture_sampled_past_the_animation_window_changes_nothing_more() { + let mut selector = FrameSelector::default(); + for index in 0..4 { + selector.sample(Rc::new(screen(&[index, 8]))); + } + + // What the sampling loop skips over: one picture, sampled again and again. + let still = Rc::new(screen(&[30])); + for _ in 0..=ANIMATION_WINDOW { + selector.sample(Rc::clone(&still)); + } + let state = ( + selector.history.clone(), + selector.in_change, + selector.change_samples, + selector.quiet_samples, + ); + + for _ in 0..3 { + assert!(!selector.sample(Rc::clone(&still))); + assert_eq!( + ( + selector.history.clone(), + selector.in_change, + selector.change_samples, + selector.quiet_samples + ), + state + ); + } + } + #[test] fn a_new_screen_size_is_a_change() { let small = Gray { From cfa0e8cc7a33234531f232a0fa451d209c5d2545 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:37:34 -0400 Subject: [PATCH 18/18] feat(dgw): give the AI analysis working files an owner The working files next to a recording were reused by any later analysis with the same model and prompt version, and `discard` removed them whoever made them. A caller that runs several analyses of one recording, such as one Task after another, could resume from another one's checkpoints or delete the checkpoints of one that is still retrying. `AnalyzeRecording::owner` records whose the files are: only an analysis of the same owner resumes from them, and another owner starts fresh. `discard(manifest_path, Some(owner))` leaves another owner's files in place; `discard(manifest_path, None)` removes them whoever made them, for a caller that starts over from scratch. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/recording_analysis/mod.rs | 40 +++++++++++++++---- .../src/recording_analysis/tests.rs | 38 +++++++++++++++++- 2 files changed, 68 insertions(+), 10 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs index 7f65c614e..cda11cdab 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -239,16 +239,26 @@ impl RecordingAnalysisError { } } -/// Removes what [`AiClient::analyze_recording`] left next to `manifest_path`: its working files and a log nobody moved. +/// Removes what [`AiClient::analyze_recording`] left next to `manifest_path`: the working files of `owner`, or of +/// nobody known, and a log nobody moved. With no `owner`, the working files go whoever made them. /// -/// Call it once the analysis of the recording is over for good; a successful analysis already removed its working -/// files. -pub async fn discard(manifest_path: &Utf8Path) -> io::Result<()> { +/// Call it once the analysis of the recording is over for good, or before a new analysis that must start from scratch; +/// a successful analysis already removed its working files. Working files of another owner stay, since that analysis +/// may still resume from them. +pub async fn discard(manifest_path: &Utf8Path, owner: Option<&str>) -> io::Result<()> { let dir = recording_dir(manifest_path); + let workspace = dir.join(WORKSPACE_DIR); - match tokio::fs::remove_dir_all(dir.join(WORKSPACE_DIR)).await { - Err(error) if error.kind() != io::ErrorKind::NotFound => return Err(error), - _ => {} + let saved = tokio::fs::read(workspace.join(SETTINGS_FILE)) + .await + .ok() + .and_then(|saved| serde_json::from_slice::(&saved).ok()); + + if owner.is_none() || saved.is_none_or(|saved| saved.owner.as_deref() == owner) { + match tokio::fs::remove_dir_all(&workspace).await { + Err(error) if error.kind() != io::ErrorKind::NotFound => return Err(error), + _ => {} + } } match tokio::fs::remove_file(dir.join(LOG_FILE)).await { @@ -275,6 +285,7 @@ impl AiClient { client: self, manifest, manifest_path, + owner: None, max_output_tokens: None, cancel: CancellationToken::new(), } @@ -287,6 +298,7 @@ pub struct AnalyzeRecording<'a> { client: &'a AiClient, manifest: &'a RecordingManifest, manifest_path: &'a Utf8Path, + owner: Option, max_output_tokens: Option, cancel: CancellationToken, } @@ -296,12 +308,20 @@ impl fmt::Debug for AnalyzeRecording<'_> { f.debug_struct("AnalyzeRecording") .field("client", self.client) .field("manifest_path", &self.manifest_path) + .field("owner", &self.owner) .field("max_output_tokens", &self.max_output_tokens) .finish_non_exhaustive() } } impl AnalyzeRecording<'_> { + /// Who the working files belong to, such as a Task ID: only an analysis of the same owner resumes from them, and + /// another owner starts fresh. + pub fn owner(mut self, owner: impl Into) -> Self { + self.owner = Some(owner.into()); + self + } + /// Upper bound of tokens in each AI answer, see /// [`DescribeSessionActions::max_output_tokens`](crate::session_actions::DescribeSessionActions::max_output_tokens). pub fn max_output_tokens(mut self, max_output_tokens: u32) -> Self { @@ -348,6 +368,7 @@ impl AnalyzeRecording<'_> { client: self.client.clone(), manifest: self.manifest.clone(), manifest_path: self.manifest_path.to_owned(), + owner: self.owner, max_output_tokens: self.max_output_tokens, cancel, events, @@ -387,6 +408,7 @@ struct Run { client: AiClient, manifest: RecordingManifest, manifest_path: Utf8PathBuf, + owner: Option, max_output_tokens: Option, cancel: CancellationToken, events: mpsc::Sender, @@ -399,6 +421,7 @@ impl Run { let cancel = self.cancel.clone(); let settings = WorkspaceSettings { + owner: self.owner.clone(), prompt_version: self.manifest.recordings.prompt_version().to_owned(), model: self.client.model().to_owned(), chunk_len: MAX_CHUNK_LEN, @@ -733,10 +756,11 @@ fn split_by_size(frames: &[LoadedScreenshot]) -> Vec> { parts } -/// What the working files were made with; files made with other settings are never reused. +/// Whose the working files are and what they were made with; files of another owner or settings are never reused. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] struct WorkspaceSettings { + owner: Option, prompt_version: String, model: String, chunk_len: usize, diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs index 064e30c5a..12d7ebcc9 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs @@ -335,12 +335,45 @@ async fn discard_removes_the_working_files_and_the_log() { let log = recording_dir(&session.manifest_path).join(LOG_FILE); std::fs::write(&log, "left behind").unwrap(); - discard(&session.manifest_path).await.unwrap(); + discard(&session.manifest_path, None).await.unwrap(); assert!(!session.workspace().exists()); assert!(!log.exists()); assert!(recording_dir(&session.manifest_path).join("recording-0.cast").exists()); - discard(&session.manifest_path).await.unwrap(); + discard(&session.manifest_path, None).await.unwrap(); +} + +#[tokio::test] +async fn working_files_belong_to_their_owner() { + // About 1 MB of transcript: three chunks. + let session = Session::new(10_000); + let (client, requests) = spawn_provider(|index| if index == 1 { 503 } else { 0 }).await; + let analyze = |owner: &'static str| { + finish( + client + .analyze_recording(&session.manifest, &session.manifest_path) + .owner(owner) + .start(), + ) + }; + + analyze("task-a").await.1.unwrap_err(); + assert!(session.workspace().join("chunk-0000.json").exists()); + + discard(&session.manifest_path, Some("task-b")).await.unwrap(); + assert!(session.workspace().exists(), "another owner's cleanup leaves the files"); + discard(&session.manifest_path, None).await.unwrap(); + assert!( + !session.workspace().exists(), + "a cleanup for nobody in particular removes them" + ); + + analyze("task-b").await.1.unwrap(); + assert_eq!( + requests.lock().unwrap().len(), + 5, + "another owner asks about every chunk again" + ); } #[tokio::test] @@ -448,6 +481,7 @@ fn recordings_listed_out_of_order_are_read_in_time_order() { }; let workspace = path.join(WORKSPACE_DIR); let settings = WorkspaceSettings { + owner: None, prompt_version: "test".to_owned(), model: MODEL.to_owned(), chunk_len: MAX_CHUNK_LEN,