diff --git a/Cargo.lock b/Cargo.lock index 33d22d9a8..f4f7b5e88 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1937,14 +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]] @@ -2453,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" @@ -5616,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 b6c6f2178..77f7fb1da 100644 --- a/crates/devolutions-gateway-ai/Cargo.toml +++ b/crates/devolutions-gateway-ai/Cargo.toml @@ -10,13 +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/INTENT.md b/crates/devolutions-gateway-ai/INTENT.md index d0a74ce18..8d0c25f8a 100644 --- a/crates/devolutions-gateway-ai/INTENT.md +++ b/crates/devolutions-gateway-ai/INTENT.md @@ -1,3 +1,21 @@ -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 +# 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. + +## 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 pre-process the .trp file to .cast, strip out some noise, then send it as text to the AI services. + +Analyze video recordings breaks down the video into frames, and we apply motion detection algorithms to filter out idle or close enough frames. +Then we send the frames in chronological order to the AI services in batches. + 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 ca15a1717..c9fff4489 100644 --- a/crates/devolutions-gateway-ai/src/client.rs +++ b/crates/devolutions-gateway-ai/src/client.rs @@ -208,16 +208,29 @@ 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() } + /// 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 @@ -230,8 +243,15 @@ impl AiClient { debug!( 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" ); @@ -252,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/src/lib.rs b/crates/devolutions-gateway-ai/src/lib.rs index 99bf42f89..7ccbf9a13 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. @@ -27,10 +27,16 @@ //! //! 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; mod wire; 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..cda11cdab --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -0,0 +1,1060 @@ +//! 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, 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`: 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, 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); + + 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 { + 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, + owner: None, + 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, + owner: Option, + 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("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 { + 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 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 = 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(), + owner: self.owner, + max_output_tokens: self.max_output_tokens, + 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), + }; + 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, + owner: Option, + 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 { + owner: self.owner.clone(), + prompt_version: self.manifest.recordings.prompt_version().to_owned(), + model: self.client.model().to_owned(), + chunk_len: MAX_CHUNK_LEN, + }; + + 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({ + 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) + }); + // 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 { + 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(); + debug!(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 tokio::fs::try_exists(&checkpoint).await.unwrap_or(false) { + 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 +} + +/// 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, +} + +/// 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() { + debug!(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, + })?; + + debug!(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..12d7ebcc9 --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs @@ -0,0 +1,535 @@ +#![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, 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, 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] +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 { + owner: None, + 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..ec0ec8ae8 --- /dev/null +++ b/crates/devolutions-gateway-ai/src/recording_analysis/video.rs @@ -0,0 +1,904 @@ +//! 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, Read}; +use std::rc::Rc; +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: u64 = 2; + +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: usize = 3; + +const TILE_PIXELS: usize = TILE * TILE; + +/// 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; + +/// 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")] + 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 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(), + offset, + workspace, + next_id, + kept: Vec::new(), + current: None, + }; + 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; + 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::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) => { + if let Some(entry) = &mut entry { + entry.codec_id = Some(id); + } + continue; + } + MatroskaSpec::PixelWidth(width) => { + if let Some(entry) = &mut entry { + entry.width = u32::try_from(width).ok(); + } + continue; + } + MatroskaSpec::PixelHeight(height) => { + 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) => { + 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.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), + _ => None, + }) else { + continue; + }; + Block::try_from(raw) + .map_err(|error| VideoError::Invalid(format!("invalid block: {error}"))) + .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.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(video.decoder()?), + }; + + 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 { + continue; + }; + let picture = Picture::of(&image)?; + + // 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; + } + + previous = Some((pts_ms, picture)); + } + + 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, &gray)?; + if *pts_ms < next_sample_ms + SAMPLE_INTERVAL_MS { + break; + } + next_sample_ms += SAMPLE_INTERVAL_MS; + } + + keeper.finish(shown)?; + } + + 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, + timestamp: i16, + data: Vec, +} + +fn frame_of( + track: u64, + timestamp: i16, + frames: Result>, webm_iterable::errors::WebmCoercionError>, +) -> Result { + let frames = frames.map_err(|error| VideoError::Invalid(format!("invalid block: {error}")))?; + + match frames.as_slice() { + [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 TrackEntry { + number: Option, + kind: Option, + codec_id: Option, + width: Option, + height: Option, +} + +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(self.width) + .height(self.height) + .codec(self.codec) + .build() + .map_err(|error| { + VideoError::Invalid(format!( + "failed to start the video decoder: {:#}", + anyhow::Error::new(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 (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() { + out.extend_from_slice(&row[..plane.width()]); + } + out + }; + + Ok(Self { + width, + 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)?; + writer.finish() + }) + .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, gray: &Rc) -> Result<(), VideoError> { + let id = format!("F{:05}", *self.next_id); + *self.next_id += 1; + + if self.selector.sample(Rc::clone(gray)) { + self.keep(time_ms, &id, picture)?; + } + + self.current = Some((time_ms, id)); + 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(()); + } + + 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: Rc) -> 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 = 0usize; + for y in row * TILE..(row + 1) * TILE { + let start = y * after.width + column * TILE; + for x in start..start + TILE { + sum += usize::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(Rc::new(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 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 { + 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]); + } + + #[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:?}" + ); + } +} 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..32b1aa6ac 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 `block`. + 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..790744e9e 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(); @@ -846,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(); @@ -860,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 @@ -868,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] @@ -876,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 @@ -884,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] @@ -906,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() @@ -918,5 +942,77 @@ 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()); } + +#[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] + ); + } + } + } +} diff --git a/crates/terminal-streamer/src/trp_decoder.rs b/crates/terminal-streamer/src/trp_decoder.rs index 751ec3d8c..b9342e707 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 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. +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_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 { + 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,153 @@ 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() { + // 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::>(); + + 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" + ); + } +}