diff --git a/Cargo.lock b/Cargo.lock index 7c9a28980..feaaa934e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1847,6 +1847,7 @@ dependencies = [ "cfg-if", "chacha20poly1305", "devolutions-agent-shared", + "devolutions-gateway-ai", "devolutions-gateway-generators", "devolutions-gateway-task", "devolutions-log", @@ -1885,6 +1886,8 @@ dependencies = [ "picky-krb", "pin-project-lite 0.2.17", "proptest", + "provisioner-task", + "provisioner-task-libsql", "rand 0.10.2", "reqwest", "rstest", @@ -1913,6 +1916,7 @@ dependencies = [ "tokio-rustls", "tokio-test", "tokio-tungstenite", + "tokio-util", "tower 0.5.3", "tower-http 0.5.2", "tracing", @@ -3993,6 +3997,7 @@ dependencies = [ "libsql", "serde", "time", + "tokio 1.52.3", "tracing", "typed-builder", "ulid", @@ -7748,6 +7753,7 @@ dependencies = [ "sysevent-winevent", "tempfile", "test-utils", + "time", "tokio 1.52.3", "tokio-rustls", "tokio-tungstenite", diff --git a/crates/devolutions-gateway-ai/examples/analyze_recording.rs b/crates/devolutions-gateway-ai/examples/analyze_recording.rs index 94bf8b1d1..d390b7c5d 100644 --- a/crates/devolutions-gateway-ai/examples/analyze_recording.rs +++ b/crates/devolutions-gateway-ai/examples/analyze_recording.rs @@ -246,10 +246,16 @@ async fn main() -> anyhow::Result<()> { 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), - ); + let usage = match analysis.usage.unreported_requests { + 0 => format!( + "{} in, {} out", + analysis.usage.input_tokens, analysis.usage.output_tokens + ), + unreported => format!( + "{} in, {} out, and {unreported} requests not reported", + analysis.usage.input_tokens, analysis.usage.output_tokens + ), + }; println!(" {} {}", style(DIM, "model "), analysis.model); println!(" {} {}", style(DIM, "prompt "), analysis.prompt_version); println!(" {} {usage}", style(DIM, "tokens ")); diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/checkpoint.rs b/crates/devolutions-gateway-ai/src/recording_analysis/checkpoint.rs index e5a3a4c4e..a0b81ece2 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/checkpoint.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/checkpoint.rs @@ -8,60 +8,19 @@ use anyhow::Context as _; use camino::{Utf8Path, Utf8PathBuf}; use serde::{Deserialize, Serialize}; +use crate::Response; 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)] +/// Actions found in one chunk, with the model that described them. +#[derive(Debug, Clone, Default, 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 { @@ -69,7 +28,6 @@ impl DescribedChunk { 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); } } @@ -78,8 +36,6 @@ impl DescribedChunk { struct SavedChunk { #[serde(default, skip_serializing_if = "Option::is_none")] model: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - usage: Option, actions: Vec, } @@ -100,7 +56,6 @@ pub(super) fn write(path: &Utf8Path, chunk: &DescribedChunk) -> anyhow::Result<( let saved = SavedChunk { model: chunk.model.clone(), - usage: chunk.usage.map(TokenUsage::from), actions: chunk .actions .iter() @@ -141,7 +96,6 @@ pub(super) fn read(path: &Utf8Path) -> anyhow::Result { Ok(DescribedChunk { actions, model: saved.model, - usage: saved.usage.map(Usage::from), }) } @@ -169,10 +123,6 @@ mod tests { }, ], model: Some("gpt-test-2026-09-30".to_owned()), - usage: Some(Usage { - input_tokens: 10, - output_tokens: 20, - }), }; write(&path, &chunk).expect("written"); @@ -183,18 +133,17 @@ mod tests { } #[test] - fn unreported_model_and_usage_stay_unknown() { + fn an_unreported_model_stays_unknown() { + // oracle: property round-trip 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 index e598d9833..1a5c3b6cb 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -33,8 +33,8 @@ 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::sync::{Arc, Mutex, PoisonError}; use std::time::Duration; use camino::{Utf8Path, Utf8PathBuf}; @@ -51,7 +51,7 @@ use self::transcript::TranscriptError; use self::video::VideoError; use crate::screen_actions::Screenshot; use crate::session_actions::Action; -use crate::{AiClient, Error, Language, Usage}; +use crate::{AiClient, Error, Language, Response, Usage}; /// Input tokens sent in one AI request, estimated at 4 characters per token. const MAX_INPUT_TOKENS_PER_REQUEST: usize = 100_000; @@ -105,13 +105,14 @@ pub enum Progress { /// /// Sent every half second while it lasts. Reading { read_bytes: u64, total_bytes: u64 }, - /// Asking the AI: chunks already described, out of all chunks. + /// Asking the AI: chunks already described, out of all chunks, and the tokens of the requests sent so far. /// - /// Sent before each chunk, then once all are described. That last event also covers the summary request, which - /// follows when the chunks have actions. + /// Sent before and after each request to the AI, and once all chunks are described. + /// A request counts as unreported until its answer comes, so usage cut off with a request is never shown as complete. Describing { described_chunks: usize, total_chunks: usize, + usage: UsageSoFar, }, /// Writing the log from the answers. Writing, @@ -138,6 +139,7 @@ impl Progress { 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), @@ -175,6 +177,7 @@ impl fmt::Display for Progress { Self::Describing { described_chunks, total_chunks, + .. } => write!( f, "asking the AI: {described_chunks} of {total_chunks} chunks described" @@ -184,6 +187,32 @@ impl fmt::Display for Progress { } } +/// Tokens the provider counted for the requests of an analysis, and the requests it did not count. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct UsageSoFar { + /// Tokens of the prompts and the inputs. + pub input_tokens: u64, + /// Tokens of the answers. + pub output_tokens: u64, + /// Requests that may have cost tokens the provider did not report, such as an answer without a token count, a request + /// that got no answer, or one still waiting for it; while it is not zero, the two counts are short of what the + /// provider billed. + pub unreported_requests: u64, +} + +impl UsageSoFar { + /// Counts the tokens of one more request, `None` when the provider did not report them. + fn add(&mut self, usage: Option) { + match usage { + Some(usage) => { + self.input_tokens = self.input_tokens.saturating_add(usage.input_tokens); + self.output_tokens = self.output_tokens.saturating_add(usage.output_tokens); + } + None => self.unreported_requests = self.unreported_requests.saturating_add(1), + } + } +} + #[expect(clippy::cast_precision_loss, reason = "shown with one decimal")] fn megabytes(bytes: u64) -> f64 { bytes as f64 / 1_048_576.0 @@ -203,10 +232,10 @@ pub struct RecordingAnalysis { /// Versions of the prompts that wrote the log, such as /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. pub prompt_version: String, - /// Tokens the provider counted for the whole analysis, answers cut and asked again included. + /// Tokens of the requests this analysis sent, answers cut and asked again included. /// - /// `None` when the provider did not report them for every request. - pub usage: Option, + /// An analysis that resumed from the working files of an earlier one does not count the requests of that one. + pub usage: UsageSoFar, } /// Error returned when an analysis fails. @@ -354,8 +383,9 @@ impl AnalyzeRecording<'_> { /// 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. + /// or [`RecordingAnalysisEvent::Failure`], which is always the last one. + /// [`Progress::Reading`] events are dropped while the channel is full, so a slow reader only misses some of them; + /// the other events never are, so the last [`Progress::Describing`] counts every request sent before a failure. /// The channel ending without a last event means the analysis panicked. /// /// Dropping the receiver cancels the analysis. @@ -387,6 +417,7 @@ impl AnalyzeRecording<'_> { max_output_tokens: self.max_output_tokens, cancel, events, + usage: Mutex::new(UsageSoFar::default()), }; tokio::spawn(async move { @@ -428,6 +459,7 @@ struct Run { max_output_tokens: Option, cancel: CancellationToken, events: mpsc::Sender, + usage: Mutex, } impl Run { @@ -461,7 +493,7 @@ impl Run { 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 { + self.report_reading(Progress::Reading { read_bytes: 0, total_bytes, }); @@ -478,7 +510,7 @@ impl Run { .map_err(|_| RecordingAnalysisError::Files("an AI analysis step panicked".to_owned()))??; } _ = ticks.tick() => { - self.report(Progress::Reading { + self.report_reading(Progress::Reading { read_bytes: read_bytes.load(Ordering::Relaxed), total_bytes, }); @@ -500,19 +532,13 @@ impl Run { 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))?; - let describe = - describe_transcript(&self.client, self.max_output_tokens, &self.language, &transcript); + let describe = self.describe_transcript(index, total, &transcript); tokio::select! { () = cancel.cancelled() => return Err(RecordingAnalysisError::Cancelled), @@ -526,13 +552,7 @@ impl Run { None => None, }; - let describe = describe_screenshots( - &self.client, - self.max_output_tokens, - &self.language, - &frames, - context.as_ref(), - ); + let describe = self.describe_screenshots(index, total, &frames, context.as_ref()); tokio::select! { () = cancel.cancelled() => return Err(RecordingAnalysisError::Cancelled), @@ -545,19 +565,15 @@ impl Run { 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_describing(total, total).await; - let DescribedChunk { actions, model, usage } = blocking({ + let DescribedChunk { actions, model } = blocking({ let workspace = workspace.clone(); let cancel = cancel.clone(); move || read_checkpoints(&workspace, total, &cancel) @@ -575,17 +591,15 @@ impl Run { let summary = tokio::select! { () = cancel.cancelled() => return Err(RecordingAnalysisError::Cancelled), - summary = request.send() => summary, + summary = self.ask(request.send(), total, total) => summary, }; Some(summary?) }; - let (model, usage) = match &summary { - Some(summary) => (model.or_else(|| summary.model.clone()), add_usage(usage, summary.usage)), - None => (model, usage), - }; - let model = model.unwrap_or_else(|| self.client.model().to_owned()); + let model = model + .or_else(|| summary.as_ref().and_then(|summary| summary.model.clone())) + .unwrap_or_else(|| self.client.model().to_owned()); let prompt_version = match &summary { Some(_) => format!( "{}+{}", @@ -596,7 +610,7 @@ impl Run { }; let action_count = actions.len(); - self.report(Progress::Writing); + self.report(Progress::Writing).await; blocking({ let workspace = workspace.clone(); @@ -642,12 +656,170 @@ impl Run { actions: action_count, model, prompt_version, - usage, + usage: self.usage(), + }) + } + + /// Asks the AI about one transcript chunk, the chunk `index` of `total`; a truncated answer is asked again as two + /// halves. + async fn describe_transcript( + &self, + index: usize, + total: usize, + chunk: &str, + ) -> Result { + let mut parts = VecDeque::from([chunk]); + let mut described = DescribedChunk::default(); + + while let Some(part) = parts.pop_front() { + let mut request = self + .client + .describe_session_actions(part) + .language(self.language.clone()); + + if let Some(max_output_tokens) = self.max_output_tokens { + request = request.max_output_tokens(max_output_tokens); + } + + let answer = self.ask(request.send(), index, total).await; + + match answer { + Ok(response) => described.add(response), + Err(Error::Truncated { .. }) => { + 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) + } + + /// Asks the AI about one chunk of screenshots, the chunk `index` of `total`, 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( + &self, + index: usize, + total: usize, + 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 = self + .client + .describe_screen_actions(&screenshots) + .language(self.language.clone()); + + if let Some(max_output_tokens) = self.max_output_tokens { + request = request.max_output_tokens(max_output_tokens); + } + + let answer = self.ask(request.send(), index, total).await; + + match answer { + Ok(response) => described.add(response), + Err(Error::Truncated { .. }) => { + 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) + } + + /// Sends one request to the AI and counts its tokens, reporting them with the chunks described so far before the + /// request, as one unreported request, and again once it is answered. + async fn ask( + &self, + request: impl Future, Error>>, + described_chunks: usize, + total_chunks: usize, + ) -> Result, Error> { + self.usage.lock().unwrap_or_else(PoisonError::into_inner).add(None); + self.report_describing(described_chunks, total_chunks).await; + + let answer = request.await; + + { + let mut usage = self.usage.lock().unwrap_or_else(PoisonError::into_inner); + usage.unreported_requests -= 1; + + match &answer { + Ok(response) => usage.add(response.usage), + Err(Error::Truncated { usage: reported }) => usage.add(*reported), + // An error status means the provider turned the request down instead of answering it. + Err(Error::Status { .. }) => {} + // The answer, if any, was lost or unreadable, so the provider may have counted tokens we never saw. + Err(_) => usage.add(None), + } + } + + self.report_describing(described_chunks, total_chunks).await; + + answer + } + + fn usage(&self) -> UsageSoFar { + *self.usage.lock().unwrap_or_else(PoisonError::into_inner) + } + + async fn report_describing(&self, described_chunks: usize, total_chunks: usize) { + self.report(Progress::Describing { + described_chunks, + total_chunks, + usage: self.usage(), }) + .await; + } + + /// Sends a progress event, waiting while the reader is behind; a reader that went away cancels the analysis. + async fn report(&self, progress: Progress) { + if self + .events + .send(RecordingAnalysisEvent::Progress(progress)) + .await + .is_err() + { + self.cancel.cancel(); + } } - /// Sends a progress event, unless the reader is behind; a reader that went away cancels the analysis. - fn report(&self, progress: Progress) { + /// Sends a reading event, unless the reader is behind, since another one follows soon; a reader that went away + /// cancels the analysis. + fn report_reading(&self, progress: Progress) { if let Err(mpsc::error::TrySendError::Closed(_)) = self.events.try_send(RecordingAnalysisEvent::Progress(progress)) { @@ -677,43 +849,6 @@ fn open_counted(path: &Utf8Path, read_bytes: &Arc) -> io::Result, - language: &Language, - 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).language(language.clone()); - - 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, @@ -750,62 +885,6 @@ async fn read_screenshots( 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, - language: &Language, - 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).language(language.clone()); - - 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> { @@ -1035,11 +1114,6 @@ fn in_time_order(recordings: &[R], start_time: impl Fn(&R) -> i64) -> 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 { @@ -1068,8 +1142,7 @@ fn split_in_half(part: &str) -> Option<(&str, &str)> { Some(part.split_at(cut)) } -/// Reads the actions of every checkpoint, in time order, with the first model reported for a chunk and the tokens of -/// all of them. +/// Reads the actions of every checkpoint, in time order, with the first model reported for a chunk. fn read_checkpoints( workspace: &Utf8Path, total: usize, @@ -1084,7 +1157,6 @@ fn read_checkpoints( let described = checkpoint::read(&checkpoint::path(workspace, index)).map_err(|error| files_error(&error))?; all.model = all.model.or(described.model); - all.usage = add_usage(all.usage, described.usage); all.actions.extend(described.actions); } diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs index 82b18b771..b14e508f5 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs @@ -15,7 +15,8 @@ const REPORTED_MODEL: &str = "gpt-test-2026-09-30"; const SUMMARY_PROMPT_START: &str = "You summarize"; /// A mock OpenAI-compatible provider: `status` picks the answer of the n-th request, 0 being a valid answer, which is -/// the action `Step n`, or `Summary n` for a summary request. +/// the action `Step n`, or `Summary n` for a summary request, counted as 10 input and 20 output tokens; 1 is the same +/// answer without a token count. /// /// Returns a client for it and the bodies of the requests it got. async fn spawn_provider(status: fn(usize) -> u16) -> (AiClient, Arc>>) { @@ -73,9 +74,8 @@ async fn spawn_provider_answering( }; let (status, answer) = match status(index) { - 0 => ( - 200, - serde_json::json!({ + scripted @ (0 | 1) => { + let mut answer = serde_json::json!({ "model": REPORTED_MODEL, "choices": [{ "index": 0, @@ -85,9 +85,12 @@ async fn spawn_provider_answering( }, "finish_reason": "stop", }], - "usage": { "prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30 }, - }), - ), + }); + if scripted == 0 { + answer["usage"] = serde_json::json!({ "prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30 }); + } + (200, answer) + } status => ( status, serde_json::json!({ "error": { "message": "scripted failure" } }), @@ -182,6 +185,14 @@ async fn finish( panic!("the events ended without a last event"); } +fn tokens(input_tokens: u64, output_tokens: u64) -> UsageSoFar { + UsageSoFar { + input_tokens, + output_tokens, + unreported_requests: 0, + } +} + fn entries(log: &Utf8Path) -> Vec { std::fs::read_to_string(log) .unwrap() @@ -222,10 +233,11 @@ async fn the_log_is_left_next_to_the_manifest_and_the_working_files_removed() { ); assert_eq!( analysis.usage, - Some(Usage { + UsageSoFar { input_tokens: 20, - output_tokens: 40 - }), + output_tokens: 40, + unreported_requests: 0, + }, "the summary request counts too" ); assert_eq!( @@ -254,6 +266,7 @@ async fn the_log_is_left_next_to_the_manifest_and_the_working_files_removed() { }; assert_eq!(total_bytes, recording_size); assert!(read_bytes <= total_bytes); + // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis assert_eq!( progress .iter() @@ -263,13 +276,37 @@ async fn the_log_is_left_next_to_the_manifest_and_the_working_files_removed() { [ Progress::Describing { described_chunks: 0, - total_chunks: 1 + total_chunks: 1, + usage: UsageSoFar { + unreported_requests: 1, + ..tokens(0, 0) + }, + }, + Progress::Describing { + described_chunks: 0, + total_chunks: 1, + usage: tokens(10, 20), }, Progress::Describing { described_chunks: 1, - total_chunks: 1 + total_chunks: 1, + usage: tokens(10, 20), }, - ] + Progress::Describing { + described_chunks: 1, + total_chunks: 1, + usage: UsageSoFar { + unreported_requests: 1, + ..tokens(10, 20) + }, + }, + Progress::Describing { + described_chunks: 1, + total_chunks: 1, + usage: tokens(20, 40), + }, + ], + "a request counts as unreported until its answer comes" ); assert_eq!(progress.last(), Some(&Progress::Writing)); } @@ -298,14 +335,52 @@ async fn analysing_again_only_asks_about_the_chunks_left() { ); assert_eq!( analysis.usage, - Some(Usage { - input_tokens: 40, - output_tokens: 80 - }), - "chunk 0 counts from its checkpoint, the failed request not at all" + tokens(30, 60), + "only the requests of this analysis count, so chunk 0 belongs to the first one" + ); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn requests_without_a_token_count_are_counted_apart() { + let session = Session::new(10); + // The chunk is request 0, answered without a token count, and the summary request 1. + let (client, _requests) = spawn_provider(|index| if index == 0 { 1 } else { 0 }).await; + + let analysis = session.analyze(&client).await.unwrap(); + + assert_eq!( + analysis.usage, + UsageSoFar { + input_tokens: 10, + output_tokens: 20, + unreported_requests: 1, + } ); } +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn a_failed_analysis_reports_the_tokens_it_spent() { + // About 1 MB of transcript: three chunks. + let session = Session::new(10_000); + // Chunk 0 is answered, then the provider fails chunk 1 with a server error, which costs no token. + let (client, _requests) = spawn_provider(|index| if index == 1 { 503 } else { 0 }).await; + + let (progress, result) = finish( + client + .analyze_recording(&session.manifest, &session.manifest_path) + .start(), + ) + .await; + + result.unwrap_err(); + let Some(Progress::Describing { usage, .. }) = progress.last() else { + panic!("the tokens are reported before the failure: {progress:?}"); + }; + assert_eq!(*usage, tokens(10, 20)); +} + #[tokio::test] async fn working_files_of_another_analysis_are_not_reused() { let session = Session::new(10); @@ -320,7 +395,6 @@ async fn working_files_of_another_analysis_are_not_reused() { &DescribedChunk { actions: Vec::new(), model: None, - usage: None, }, ) .unwrap(); @@ -541,6 +615,7 @@ fn progress_tells_its_stage_and_how_far_it_is() { let describing = Progress::Describing { described_chunks: 2, total_chunks: 5, + usage: UsageSoFar::default(), }; assert_eq!(describing.stage(), "describing"); assert_eq!(describing.percent(), 40); @@ -651,15 +726,3 @@ fn halves_are_cut_on_the_line_boundary_closest_to_the_middle() { 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/job-queue-libsql/Cargo.toml b/crates/job-queue-libsql/Cargo.toml index efdf7da01..a4a4934a3 100644 --- a/crates/job-queue-libsql/Cargo.toml +++ b/crates/job-queue-libsql/Cargo.toml @@ -19,3 +19,6 @@ tracing = "0.1" typed-builder = "0.21" ulid = { version = "1.2", features = ["uuid"] } uuid = "1.23" + +[dev-dependencies] +tokio = { version = "1.52", features = ["macros", "rt"] } diff --git a/crates/job-queue-libsql/src/lib.rs b/crates/job-queue-libsql/src/lib.rs index 0efbfd23f..f8224ba38 100644 --- a/crates/job-queue-libsql/src/lib.rs +++ b/crates/job-queue-libsql/src/lib.rs @@ -402,6 +402,27 @@ impl JobQueue for LibSqlJobQueue { Ok(Some(scheduled_for)) } + + async fn job_definitions(&self, name: &str) -> anyhow::Result> { + let sql_query = "SELECT json(def) FROM job_queue WHERE name = :name AND failed_attempts < :max_attempts"; + + let params = ((":name", name), (":max_attempts", self.max_attempts)); + + trace!(%sql_query, ?params, "Listing job definitions"); + + let mut rows = self + .conn + .query(sql_query, params) + .await + .context("failed to execute SQL query")?; + + let mut definitions = Vec::new(); + while let Some(row) = rows.next().await.context("failed to read the row")? { + definitions.push(row.get::(0).context("failed to read the job definition")?); + } + + Ok(definitions) + } } // Typically, migrations should not be modified once released, and we should only be appending to this list. @@ -425,3 +446,56 @@ const MIGRATIONS: &[&str] = &[ CREATE INDEX idx_scheduled_for ON job_queue(scheduled_for);", ]; + +#[cfg(test)] +mod tests { + use job_queue::{DynJob, Job, JobQueue as _, RunnerWaker}; + + use super::*; + + struct Named { + name: &'static str, + def: &'static str, + } + + #[async_trait] + impl Job for Named { + fn name(&self) -> &str { + self.name + } + + fn write_json(&self) -> anyhow::Result { + Ok(self.def.to_owned()) + } + + async fn run(&mut self) -> anyhow::Result<()> { + Ok(()) + } + } + + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + #[tokio::test] + async fn job_definitions_lists_the_stored_jobs_of_one_name() { + let conn = libsql::Builder::new_local(":memory:") + .build() + .await + .expect("build database") + .connect() + .expect("connect"); + let queue = LibSqlJobQueue::builder() + .runner_waker(RunnerWaker::new(|| {})) + .conn(conn) + .build(); + queue.setup().await.expect("setup"); + + for (name, def) in [("a", r#"{"n":1}"#), ("b", r#"{"n":2}"#), ("a", r#"{"n":3}"#)] { + let job: DynJob = Box::new(Named { name, def }); + queue.push_job(&job, None).await.expect("push"); + } + + let mut definitions = queue.job_definitions("a").await.expect("list"); + definitions.sort(); + + assert_eq!(definitions, [r#"{"n":1}"#, r#"{"n":3}"#]); + } +} diff --git a/crates/job-queue/src/lib.rs b/crates/job-queue/src/lib.rs index 93232bd86..5521da528 100644 --- a/crates/job-queue/src/lib.rs +++ b/crates/job-queue/src/lib.rs @@ -62,6 +62,9 @@ pub trait JobQueue: Send + Sync { /// Retrieves the closest future scheduled date async fn next_scheduled_date(&self) -> anyhow::Result>; + + /// Lists the definitions of the stored jobs named `name` that may still run, in no particular order. + async fn job_definitions(&self, name: &str) -> anyhow::Result>; } pub struct JobCtx { diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 3a0f03972..1b1ff2f75 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -45,7 +45,7 @@ impl LibSqlProvisionerTaskStore { /// Ends the Task as `state` if it is currently in one of `from`. /// - /// A Task past its deadline is failed as timed out first, and then not ended again. + /// A Task past its deadline is failed for it first, and then not ended again. async fn finish( &self, id: Uuid, @@ -131,11 +131,10 @@ async fn fail_if_overdue( return Ok(task); } - let payload = serde_json::json!({ - "reason": "timed out", - "details": { "lastPayload": task.payload }, - "attempts": task.attempts, - }); + let mut payload = serde_json::json!({ "error": { "kind": "deadlinePassed" } }); + if let Some(usage) = task.payload.get("usage") { + payload["usage"] = usage.clone(); + } conn.execute( "UPDATE provisioner_task_records SET state = 'failed', payload = ?2, finished_at = ?3 @@ -145,7 +144,7 @@ async fn fail_if_overdue( .await .context("fail overdue task")?; - warn!(task.id = %task.id, task.kind = %task.kind, attempts = task.attempts, "Task timed out"); + warn!(task.id = %task.id, task.kind = %task.kind, attempts = task.attempts, "Task failed past its deadline"); find(conn, task.id).await?.context("overdue task is missing") } @@ -226,6 +225,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { id: Uuid, job_token: Uuid, payload: serde_json::Value, + deadline_at: OffsetDateTime, now: OffsetDateTime, ) -> anyhow::Result { let conn = self.conn.lock().await; @@ -253,9 +253,15 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { tx.execute( "UPDATE provisioner_task_records - SET state = 'running', payload = ?2, attempts = attempts + 1, started_at = COALESCE(started_at, ?3) + SET state = 'running', payload = ?2, attempts = attempts + 1, started_at = COALESCE(started_at, ?3), + deadline_at = ?4 WHERE id = ?1", - params![id.to_string(), payload.to_string(), now.unix_timestamp()], + params![ + id.to_string(), + payload.to_string(), + now.unix_timestamp(), + deadline_at.unix_timestamp() + ], ) .await .context("start task attempt")?; @@ -331,7 +337,13 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(unfinished) } - async fn update_running(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + async fn update_running( + &self, + id: Uuid, + payload: serde_json::Value, + deadline_at: OffsetDateTime, + now: OffsetDateTime, + ) -> anyhow::Result { let conn = self.conn.lock().await; let tx = conn .transaction_with_behavior(TransactionBehavior::Immediate) @@ -348,8 +360,8 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { if running { tx.execute( - "UPDATE provisioner_task_records SET payload = ?2 WHERE id = ?1", - params![id.to_string(), payload.to_string()], + "UPDATE provisioner_task_records SET payload = ?2, deadline_at = ?3 WHERE id = ?1", + params![id.to_string(), payload.to_string(), deadline_at.unix_timestamp()], ) .await .context("update running task")?; @@ -671,7 +683,8 @@ mod tests { .start_attempt( overdue.id, overdue.job_token, - serde_json::json!({ "done": 3, "total": 10 }), + serde_json::json!({ "done": 3, "total": 10, "usage": { "inputTokens": 1200 } }), + now() - Duration::minutes(1), now() - Duration::minutes(2), ) .await @@ -688,9 +701,10 @@ mod tests { .expect("get task") .expect("task exists"); assert_eq!(failed.state, ProvisionerTaskState::Failed); - assert_eq!(failed.payload["reason"], "timed out"); - assert_eq!(failed.payload["attempts"], 1); - assert_eq!(failed.payload["details"]["lastPayload"]["done"], 3); + assert_eq!( + failed.payload, + serde_json::json!({ "error": { "kind": "deadlinePassed" }, "usage": { "inputTokens": 1200 } }) + ); } #[tokio::test] @@ -701,20 +715,37 @@ mod tests { assert!(matches!( store - .start_attempt(task.id, task.job_token, serde_json::json!({ "done": 0 }), now()) + .start_attempt( + task.id, + task.job_token, + serde_json::json!({ "done": 0 }), + now() + Duration::hours(2), + now() + ) .await .expect("start"), AttemptStart::Started(_) )); assert!( store - .update_running(task.id, serde_json::json!({ "done": 1 }), now()) + .update_running( + task.id, + serde_json::json!({ "done": 1 }), + now() + Duration::hours(2), + now() + ) .await .expect("update") ); assert!(matches!( store - .start_attempt(task.id, task.job_token, serde_json::json!({ "done": 1 }), now()) + .start_attempt( + task.id, + task.job_token, + serde_json::json!({ "done": 1 }), + now() + Duration::hours(2), + now() + ) .await .expect("retry"), AttemptStart::Started(_) @@ -739,14 +770,20 @@ mod tests { ); assert!(matches!( store - .start_attempt(task.id, task.job_token, serde_json::json!({}), now()) + .start_attempt( + task.id, + task.job_token, + serde_json::json!({}), + now() + Duration::hours(2), + now() + ) .await .expect("start after finish"), AttemptStart::Finished(_) )); assert!( !store - .update_running(task.id, serde_json::json!({}), now()) + .update_running(task.id, serde_json::json!({}), now() + Duration::hours(2), now()) .await .expect("update after finish") ); @@ -763,7 +800,13 @@ mod tests { let task = new_task("session-1"); create(&store, task.clone()).await; store - .start_attempt(task.id, task.job_token, serde_json::json!({ "done": 0 }), now()) + .start_attempt( + task.id, + task.job_token, + serde_json::json!({ "done": 0 }), + now() + Duration::hours(2), + now(), + ) .await .expect("start"); @@ -787,14 +830,25 @@ mod tests { overdue.deadline_at = now() + Duration::minutes(1); create(&store, overdue.clone()).await; store - .start_attempt(overdue.id, overdue.job_token, serde_json::json!({ "done": 3 }), now()) + .start_attempt( + overdue.id, + overdue.job_token, + serde_json::json!({ "done": 3 }), + now() + Duration::minutes(1), + now(), + ) .await .expect("start"); let late = now() + Duration::minutes(2); assert!( !store - .update_running(overdue.id, serde_json::json!({ "done": 4 }), late) + .update_running( + overdue.id, + serde_json::json!({ "done": 4 }), + late + Duration::minutes(30), + late, + ) .await .expect("update late") ); @@ -817,8 +871,77 @@ mod tests { .expect("get task") .expect("task exists"); assert_eq!(failed.state, ProvisionerTaskState::Failed); - assert_eq!(failed.payload["reason"], "timed out"); - assert_eq!(failed.payload["details"]["lastPayload"]["done"], 3); + assert_eq!( + failed.payload, + serde_json::json!({ "error": { "kind": "deadlinePassed" } }) + ); + } + + #[tokio::test] + async fn an_attempt_start_sets_the_deadline() { + // oracle: intent crates/provisioner-task/INTENT.md §invariants + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + store + .start_attempt( + task.id, + task.job_token, + serde_json::json!({}), + now() + Duration::minutes(30), + now(), + ) + .await + .expect("start"); + + // Still within the deadline of the queued Task, but past the one its attempt set. + let task = store + .get(task.id, now() + Duration::minutes(40)) + .await + .expect("get task") + .expect("task exists"); + assert_eq!(task.state, ProvisionerTaskState::Failed); + } + + #[tokio::test] + async fn progress_moves_the_deadline_forward() { + // oracle: intent crates/provisioner-task/INTENT.md §invariants + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + store + .start_attempt( + task.id, + task.job_token, + serde_json::json!({}), + now() + Duration::minutes(30), + now(), + ) + .await + .expect("start"); + let moved = now() + Duration::minutes(50); + assert!( + store + .update_running( + task.id, + serde_json::json!({ "done": 1 }), + moved, + now() + Duration::minutes(20) + ) + .await + .expect("report progress") + ); + + let running = store + .get(task.id, now() + Duration::minutes(40)) + .await + .expect("get task") + .expect("task exists"); + assert_eq!(running.state, ProvisionerTaskState::Running); + assert_eq!(running.deadline_at, moved); + + let failed = store.get(task.id, moved).await.expect("get task").expect("task exists"); + assert_eq!(failed.state, ProvisionerTaskState::Failed); } #[tokio::test] @@ -830,7 +953,13 @@ mod tests { assert!(matches!( store - .start_attempt(overdue.id, overdue.job_token, serde_json::json!({}), now()) + .start_attempt( + overdue.id, + overdue.job_token, + serde_json::json!({}), + now() + Duration::hours(2), + now() + ) .await .expect("start attempt"), AttemptStart::Finished(_) @@ -841,7 +970,10 @@ mod tests { .expect("get task") .expect("task exists"); assert_eq!(read.state, ProvisionerTaskState::Failed); - assert_eq!(read.payload["reason"], "timed out"); + assert_eq!( + read.payload, + serde_json::json!({ "error": { "kind": "deadlinePassed" } }) + ); let CreateOutcome::Existing(again) = create(&store, overdue).await else { panic!("same request must return the existing task"); @@ -857,7 +989,13 @@ mod tests { assert_eq!( store - .start_attempt(task.id, Uuid::new_v4(), serde_json::json!({}), now()) + .start_attempt( + task.id, + Uuid::new_v4(), + serde_json::json!({}), + now() + Duration::hours(2), + now() + ) .await .expect("start with another token"), AttemptStart::Superseded @@ -882,14 +1020,20 @@ mod tests { ); assert_eq!( store - .start_attempt(task.id, task.job_token, serde_json::json!({}), now()) + .start_attempt( + task.id, + task.job_token, + serde_json::json!({}), + now() + Duration::hours(2), + now() + ) .await .expect("start with the old token"), AttemptStart::Superseded ); assert!(matches!( store - .start_attempt(task.id, new_token, serde_json::json!({}), now()) + .start_attempt(task.id, new_token, serde_json::json!({}), now() + Duration::hours(2), now()) .await .expect("start with the new token"), AttemptStart::Started(started) if started.job_token == new_token && started.attempts == 1 @@ -920,7 +1064,13 @@ mod tests { create(&store, task.clone()).await; } store - .start_attempt(running.id, running.job_token, serde_json::json!({}), now()) + .start_attempt( + running.id, + running.job_token, + serde_json::json!({}), + now() + Duration::hours(2), + now(), + ) .await .expect("start"); store @@ -942,8 +1092,8 @@ mod tests { .await .expect("get task") .expect("task") - .payload["reason"], - "timed out" + .payload["error"]["kind"], + "deadlinePassed" ); } @@ -954,7 +1104,13 @@ mod tests { assert_eq!(store.get(Uuid::new_v4(), now()).await.expect("get task"), None); assert_eq!( store - .start_attempt(Uuid::new_v4(), Uuid::new_v4(), serde_json::json!({}), now()) + .start_attempt( + Uuid::new_v4(), + Uuid::new_v4(), + serde_json::json!({}), + now() + Duration::hours(2), + now() + ) .await .expect("start attempt"), AttemptStart::Unknown @@ -966,7 +1122,9 @@ mod tests { .expect("replace job token") ); for finished in [ - store.update_running(Uuid::new_v4(), serde_json::json!({}), now()).await, + store + .update_running(Uuid::new_v4(), serde_json::json!({}), now() + Duration::hours(2), now()) + .await, store.succeed(Uuid::new_v4(), serde_json::json!({}), now()).await, store.fail(Uuid::new_v4(), serde_json::json!({}), now()).await, ] { diff --git a/crates/provisioner-task/INTENT.md b/crates/provisioner-task/INTENT.md index 40d2a64d2..642b6271b 100644 --- a/crates/provisioner-task/INTENT.md +++ b/crates/provisioner-task/INTENT.md @@ -13,7 +13,7 @@ Only these transitions exist: - `Running` -> `Succeeded` - `Running` -> `Failed` -Each Task has a deadline counted from when Gateway receives it. +Each Task has a deadline. While the Task runs, its deadline may move forward each time it makes progress. An unfinished Task past its deadline is failed, so a lost job can't block its target forever. # Tasks and jobs diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index be3d8431d..587a4b93c 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -46,6 +46,8 @@ pub struct ProvisionerTaskRecord { pub started_at: Option, pub finished_at: Option, /// After this, an unfinished Task is failed, so a lost job can't block its target forever. + /// + /// It moves forward while the Task runs, each time an attempt starts or reports progress. pub deadline_at: OffsetDateTime, /// The job that may run this Task: a job with another token is stale or a duplicate, and must not run it. pub job_token: Uuid, @@ -87,6 +89,10 @@ pub enum AttemptStart { Unknown, } +/// Where Tasks are kept. +/// +/// A Task past its deadline fails with the payload `{ "error": { "kind": "deadlinePassed" } }`, plus the `usage` of +/// its last payload when it has one, so what it spent stays known. #[async_trait] pub trait ProvisionerTaskStore: Send + Sync { /// Records a new `queued` Task, unless the ID is known or its target is busy. @@ -97,7 +103,8 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Reads a Task, after failing it if it is past its deadline. async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result>; - /// Marks the start of one more attempt by the job holding `job_token`: `running`, with one more attempt counted. + /// Marks the start of one more attempt by the job holding `job_token`: `running`, with one more attempt counted and + /// its deadline moved to `deadline_at`. /// /// Only [`AttemptStart::Started`] lets the attempt run. A job that no longer holds the token gets /// [`AttemptStart::Superseded`], and nothing changes. @@ -106,6 +113,7 @@ pub trait ProvisionerTaskStore: Send + Sync { id: Uuid, job_token: Uuid, payload: serde_json::Value, + deadline_at: OffsetDateTime, now: OffsetDateTime, ) -> anyhow::Result; @@ -118,9 +126,15 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Every unfinished Task, after failing those past their deadline. async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; - /// Replaces the payload of a `running` Task, e.g. to report progress. Returns `false` if it is not running, such - /// as when it was just failed for being past its deadline. - async fn update_running(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; + /// Replaces the payload of a `running` Task and moves its deadline to `deadline_at`, e.g. to report progress. + /// Returns `false` if it is not running, such as when it was just failed for being past its deadline. + async fn update_running( + &self, + id: Uuid, + payload: serde_json::Value, + deadline_at: OffsetDateTime, + now: OffsetDateTime, + ) -> anyhow::Result; /// Ends a `running` Task as `succeeded`. Returns `false` if it is not running, such as when it was just failed for /// being past its deadline. diff --git a/devolutions-gateway/Cargo.toml b/devolutions-gateway/Cargo.toml index ee26e0adb..b486331c5 100644 --- a/devolutions-gateway/Cargo.toml +++ b/devolutions-gateway/Cargo.toml @@ -24,6 +24,7 @@ transport.path = "../crates/transport" jmux-proxy.path = "../crates/jmux-proxy" devolutions-agent-shared.path = "../crates/devolutions-agent-shared" devolutions-gateway-task.path = "../crates/devolutions-gateway-task" +devolutions-gateway-ai.path = "../crates/devolutions-gateway-ai" devolutions-log.path = "../crates/devolutions-log" job-queue.path = "../crates/job-queue" job-queue-libsql.path = "../crates/job-queue-libsql" @@ -38,6 +39,8 @@ sysevent-codes.path = "../crates/sysevent-codes" agent-tunnel = { path = "../crates/agent-tunnel" } agent-tunnel-libsql.path = "../crates/agent-tunnel-libsql" gateway-db.path = "../crates/gateway-db" +provisioner-task.path = "../crates/provisioner-task" +provisioner-task-libsql.path = "../crates/provisioner-task-libsql" ironrdp-pdu = { version = "0.9", features = ["std"] } ironrdp-core = { version = "0.2", features = ["std"] } ironrdp-rdcleanpath = "0.2" @@ -91,6 +94,7 @@ tracing = "0.1" # Async, futures… tokio = { version = "1.52", features = ["signal", "net", "io-util", "time", "rt", "rt-multi-thread", "sync", "macros", "parking_lot", "fs"] } +tokio-util = "0.7" tokio-rustls = { version = "0.26", default-features = false, features = ["logging", "tls12", "ring"] } socket2 = "0.5" reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-native-roots", "json", "stream", "http2", "socks"] } diff --git a/devolutions-gateway/src/api/jrec.intent.md b/devolutions-gateway/src/api/jrec.intent.md new file mode 100644 index 000000000..ccd9fd423 --- /dev/null +++ b/devolutions-gateway/src/api/jrec.intent.md @@ -0,0 +1,14 @@ +# Background + +Jrec is the endpoint to push, streaming, read recordings and its associated files. + + +# AI analysis + +An AI analysis is a task initiated by the provisioner. +It requires the scope token with access `gateway.tasks.recording-ai-analysis.start` to access the url `POST /jet/jrec/{session_id}/ai-analysis`; Naturally in order to read the task's state, it requires the scope token with access `gateway.tasks.recording-ai-analysis.read`. +The analyzed output belongs to the recording reading scope, however, the associated metadata, including but not limited to error messages, token usages etc. belong to the task reading scope. +The caller chooses the language of the analysis in the request, under the rules of `crates/devolutions-gateway-ai/INTENT.md`. The language is part of the Task like every other parameter, so the same Task ID sent again with another language, model or provider is a conflict. +A request sent again is compared on its parameters only: the API key and the base URL are never stored, so another key or base URL gets the same Task, which keeps the first ones. +A new Task answers 202 with its path in `Location`, and the same request again answers 200 with the same Task. A conflict answers 409 with a `code`: `task_id_conflict`, `recording_active`, or `analysis_in_progress` with the running `taskId`. A session without recording, or an unknown Task, answers 404. An invalid request answers 400 and records no Task. Unknown fields are ignored, so a newer provisioner works with an older Gateway. The provider is one of `openai`, `anthropic`, `mistral`, `gemini` or `openai-compatible`. +Every Task reports the tokens it used so far, whether it is running, succeeded or failed, so the provisioner can add up what all analyses cost, failed ones included. Requests whose usage the provider didn't report are counted apart, so a total is never shown as complete when it isn't. A succeeded Task also shows the log's file name, the model the provider reported, and the prompt versions. diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index 32f4b7d14..396d15d00 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -12,27 +12,35 @@ use anyhow::Context as _; use axum::body::Body; use axum::extract::ws::WebSocket; use axum::extract::{self, ConnectInfo, Query, State, WebSocketUpgrade}; -use axum::http::header::{CONTENT_DISPOSITION, CONTENT_TYPE, HeaderValue}; -use axum::response::Response; +use axum::http::header::{CONTENT_DISPOSITION, CONTENT_TYPE, HeaderValue, LOCATION}; +use axum::response::{IntoResponse as _, Response}; use axum::routing::{delete, get, post}; use axum::{Json, Router}; use bytes::Bytes; use cadeau::xmf; use camino::{Utf8Path, Utf8PathBuf}; +use devolutions_gateway_ai::Language; use devolutions_gateway_task::ShutdownSignal; use futures::stream; use hyper::StatusCode; +use secrecy::SecretString; use tokio::io::AsyncReadExt as _; use tracing::Instrument as _; +use url::Url; use uuid::Uuid; use zip::CompressionMethod; use zip::write::SimpleFileOptions; use crate::DgwState; use crate::api::heartbeat::recording_storage_health; +use crate::api::tasks::TaskInfo; use crate::artifacts::JrecArtifacts; -use crate::extract::{JrecToken, RecordingDeleteScope, RecordingsReadScope}; +use crate::extract::{ + JrecToken, ProvisionerTaskReadScope, ProvisionerTaskStartScope, RecordingDeleteScope, RecordingsReadScope, +}; use crate::http::{HttpError, HttpErrorBuilder}; +use crate::provisioner_tasks::ai::{AiProvider, AiSettings}; +use crate::provisioner_tasks::recording_ai_analysis::{self, StartError, StartOutcome, StartRequest}; use crate::recording::{PushOutcome, RecordingMessageSender}; use crate::token::{JrecTokenClaims, RecordingFileType, RecordingOperation}; @@ -62,6 +70,8 @@ pub fn make_router(state: DgwState) -> Router { .route("/play/{*path}", get(get_player)) .route("/shadow/{id}", get(shadow_recording)) .route("/search", post(search::search_recordings)) + .route("/{session_id}/ai-analysis", post(start_ai_analysis)) + .route("/{session_id}/ai-analysis/{task_id}", get(get_ai_analysis)) .with_state(state) } @@ -429,6 +439,211 @@ async fn delete_recording(recording_path: &Utf8Path) -> anyhow::Result<()> { Ok(()) } +/// Settings of an AI analysis +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisRequest { + /// Task ID, chosen by the caller so it can send the same request again safely. + task_id: Uuid, + provider: AiProvider, + /// Model identifier, passed to the provider as is. + model: String, + /// Endpoint of the provider API, kept in memory only; required for `openai-compatible`. + #[cfg_attr(feature = "openapi", schema(value_type = Option))] + #[serde(default)] + base_url: Option, + /// Upper bound of tokens in each AI answer. + #[cfg_attr(feature = "openapi", schema(value_type = Option, minimum = 0, maximum = 4_294_967_295_u64))] + #[serde(default)] + max_output_tokens: Option, + /// Language the AI writes in, any text of 1 to 10 characters such as `fr-CA` or `français`; English when absent. + #[cfg_attr(feature = "openapi", schema(min_length = 1, max_length = 10))] + #[serde(default)] + language: Option, + /// API key of the provider, kept in memory only until the Task ends. + #[cfg_attr(feature = "openapi", schema(value_type = String, format = Password))] + api_key: SecretString, +} + +/// Why an AI analysis cannot start now +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisConflict { + code: AiAnalysisConflictCode, + /// The AI analysis of this session that is not finished yet; set for `analysis_in_progress`. + #[serde(skip_serializing_if = "Option::is_none")] + task_id: Option, +} + +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum AiAnalysisConflictCode { + /// The Task ID is already used by a Task with other parameters. + TaskIdConflict, + /// The session is still recording. + RecordingActive, + /// Another AI analysis of this session is not finished yet. + AnalysisInProgress, +} + +fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> Response { + debug!(?code, task.id = ?task_id, "AI analysis refused"); + (StatusCode::CONFLICT, Json(AiAnalysisConflict { code, task_id })).into_response() +} + +/// Starts an AI analysis of a recorded session +/// +/// Gateway asks the AI provider what the user did in the session recording, and puts the answer in the session as its +/// log, listed under `artifacts.ai-analysis` in its manifest. +/// The answer is in the `language` of the request, English without one. +/// The work runs in the background as a Task; follow it with `GET /jet/jrec/{session_id}/ai-analysis/{task_id}`, +/// which the `Location` header of a 202 answer names. +/// +/// Sending the same request again returns the same Task and starts nothing new. +/// A new Task for the session starts from scratch, and replaces the earlier AI analysis of the session once its success +/// is recorded, so a session holds at most one; a Task that fails leaves the earlier analysis in place and removes +/// everything it made. +/// A Task fails if it does not start within 2 hours, or, once it runs, after 30 minutes without progress. +/// The API key and the base URL are kept in memory only, until the Task ends; a Gateway restart fails the unfinished +/// Tasks. +#[cfg_attr(feature = "openapi", utoipa::path( + post, + operation_id = "StartAiAnalysis", + tag = "Jrec", + path = "/jet/jrec/{session_id}/ai-analysis", + params( + ("session_id" = Uuid, Path, description = "Recorded session ID"), + ), + request_body(content = AiAnalysisRequest, description = "AI analysis settings", content_type = "application/json"), + responses( + (status = 200, description = "The same request came in before; nothing new starts", body = TaskInfo), + (status = 202, description = "The AI analysis Task is created and runs in the background", body = TaskInfo, headers( + ("Location" = String, description = "Path of the new Task, `/jet/jrec/{session_id}/ai-analysis/{task_id}`"), + )), + (status = 400, description = "Invalid request body, AI settings or language"), + (status = 401, description = "Invalid or missing authorization token"), + (status = 403, description = "Insufficient permissions"), + (status = 404, description = "The session has no recording"), + (status = 409, description = "The AI analysis cannot start now", body = AiAnalysisConflict), + (status = 500, description = "Unexpected server error"), + ), + security(("scope_token" = ["gateway.tasks.recording-ai-analysis.start"])), +))] +pub(crate) async fn start_ai_analysis( + State(state): State, + _scope: ProvisionerTaskStartScope, + extract::Path(session_id): extract::Path, + body: Bytes, +) -> Result { + let request: AiAnalysisRequest = serde_json::from_slice(&body).map_err(|error| { + // The serde message may quote a value of the body, such as a misplaced API key, so only its place is logged. + debug!( + category = ?error.classify(), + line = error.line(), + column = error.column(), + "Invalid AI analysis request body" + ); + HttpError::bad_request().msg("invalid AI analysis request body") + })?; + + let AiAnalysisRequest { + task_id, + provider, + model, + base_url, + max_output_tokens, + language, + api_key, + } = request; + + let language = language + .map(Language::new) + .transpose() + .map_err(HttpError::bad_request().with_msg("invalid language").err())?; + + let outcome = recording_ai_analysis::start( + &state, + StartRequest { + task_id, + session_id, + settings: AiSettings { + provider, + model, + base_url, + max_output_tokens, + }, + language, + api_key, + }, + ) + .await; + + match outcome { + Ok(StartOutcome::Created(task)) => { + let location = HeaderValue::try_from(format!("/jet/jrec/{session_id}/ai-analysis/{}", task.id)) + .map_err(HttpError::internal().with_msg("invalid Location header").err())?; + Ok((StatusCode::ACCEPTED, [(LOCATION, location)], Json(TaskInfo::from(task))).into_response()) + } + Ok(StartOutcome::Existing(task)) => Ok((StatusCode::OK, Json(TaskInfo::from(task))).into_response()), + Ok(StartOutcome::IdConflict(_)) => Ok(ai_analysis_conflict(AiAnalysisConflictCode::TaskIdConflict, None)), + Ok(StartOutcome::TargetBusy { active_task_id }) => Ok(ai_analysis_conflict( + AiAnalysisConflictCode::AnalysisInProgress, + Some(active_task_id), + )), + Err(StartError::RecordingActive) => Ok(ai_analysis_conflict(AiAnalysisConflictCode::RecordingActive, None)), + Err(StartError::InvalidSettings(error)) => { + Err(HttpError::bad_request().with_msg("invalid AI settings").build(error)) + } + Err(StartError::RecordingNotFound) => Err(HttpError::not_found().msg("session has no recording")), + Err(StartError::Internal(error)) => Err(HttpError::internal() + .with_msg("failed to start the AI analysis") + .build(error)), + } +} + +/// Gets an AI analysis of a recorded session +/// +/// Returns the Task of the analysis: its state, its progress while it runs, its result once it is over, and the tokens +/// it used in every state but `queued`. +/// Tasks are kept forever, across Gateway restarts; an unfinished Task past its deadline is reported as failed. +#[cfg_attr(feature = "openapi", utoipa::path( + get, + operation_id = "GetAiAnalysis", + tag = "Jrec", + path = "/jet/jrec/{session_id}/ai-analysis/{task_id}", + params( + ("session_id" = Uuid, Path, description = "Recorded session ID"), + ("task_id" = Uuid, Path, description = "Task ID of the AI analysis"), + ), + responses( + (status = 200, description = "The AI analysis Task", body = TaskInfo), + (status = 400, description = "Bad request"), + (status = 401, description = "Invalid or missing authorization token"), + (status = 403, description = "Insufficient permissions"), + (status = 404, description = "The session has no AI analysis with this Task ID"), + (status = 500, description = "Unexpected server error"), + ), + security(("scope_token" = ["gateway.tasks.recording-ai-analysis.read"])), +))] +pub(crate) async fn get_ai_analysis( + State(DgwState { provisioner_tasks, .. }): State, + _scope: ProvisionerTaskReadScope, + extract::Path((session_id, task_id)): extract::Path<(Uuid, Uuid)>, +) -> Result, HttpError> { + let task = provisioner_tasks + .get(task_id, time::OffsetDateTime::now_utc()) + .await + .map_err(HttpError::internal().with_msg("failed to read the AI analysis").err())? + // Another kind of Task, or the analysis of another session, is not found under this path. + .filter(|task| task.kind == recording_ai_analysis::KIND && task.target == session_id.to_string()) + .ok_or_else(|| HttpError::not_found().msg("AI analysis not found"))?; + + Ok(Json(TaskInfo::from(task))) +} + /// Lists all recordings stored on this instance #[cfg_attr(feature = "openapi", utoipa::path( get, @@ -1012,6 +1227,53 @@ mod tests { use super::*; + // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis + #[tokio::test] + async fn ai_analysis_of_a_session_still_recording_is_a_conflict() { + use axum::extract::connect_info::MockConnectInfo; + use base64::Engine as _; + use tower::ServiceExt as _; + + let dir = tempfile::tempdir().expect("temp dir"); + let config = serde_json::json!({ + "ProvisionerPublicKeyData": { + "Value": "mMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4vuqLOkl1pWobt6su1XO9VskgCAwevEGs6kkNjJQBwkGnPKYLmNF1E/af1yCocfVn/OnPf9e4x+lXVyZ6LMDJxFxu+axdgOq3Ld392J1iAEbfvwlyRFnEXFOJNyylqg3bY6LvnWHL/XZczVdMD9xYfq2sO9bg3xjRW4s7r9EEYOFjqVT3VFznH9iWJVtcSEKukmS/3uKoO6lGhacvu0HhjXXdgq0R8zvR4XRJ9Fcnf0f9Ypoc+i6L80NVjrRCeVOH+Ld/2fA9bocpfLarcVqG3RjS+qgOtpyCc0jWVFF4zaGQ7LUDFkEIYILkICeMMn2ll29hmZNzsJzZJ9s6NocgQIDAQAB" + }, + "Listeners": [{ "InternalUrl": "http://*:7171", "ExternalUrl": "https://*:7171" }], + "RecordingPath": dir.path().to_str().expect("UTF-8"), + "__debug__": { "disable_token_validation": true }, + }); + let (state, _handles) = DgwState::mock(&config.to_string()).await.expect("mock state"); + let session_id = Uuid::new_v4(); + state.recordings.active_recordings.insert(session_id); + let app = crate::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 3000)))); + + let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD; + let claims = serde_json::json!({ "jti": Uuid::new_v4(), "exp": i64::from(u32::MAX), "scope": "gateway.tasks.recording-ai-analysis.start" }); + let token = format!( + "{}.{}.{}", + engine.encode(r#"{"alg":"RS256","cty":"SCOPE"}"#), + engine.encode(claims.to_string()), + engine.encode(b"signature"), + ); + let body = serde_json::json!({ "taskId": Uuid::new_v4(), "provider": "openai", "model": "gpt-test", "apiKey": "sk-test" }); + let request = axum::http::Request::builder() + .method("POST") + .uri(format!("/jet/jrec/{session_id}/ai-analysis")) + .header(axum::http::header::AUTHORIZATION, format!("Bearer {token}")) + .body(Body::from(body.to_string())) + .expect("request"); + + let response = app.oneshot(request).await.expect("response"); + + assert_eq!(response.status(), StatusCode::CONFLICT); + let body = response.into_body().collect().await.expect("body").to_bytes(); + assert_eq!( + serde_json::from_slice::(&body).expect("JSON"), + serde_json::json!({ "code": "recording_active" }) + ); + } + #[test] fn rejects_unsafe_recording_file_names() { assert!(is_safe_recording_file_name("recording-0.webm")); diff --git a/devolutions-gateway/src/api/mod.rs b/devolutions-gateway/src/api/mod.rs index c7b222ebc..f31a32d15 100644 --- a/devolutions-gateway/src/api/mod.rs +++ b/devolutions-gateway/src/api/mod.rs @@ -13,6 +13,7 @@ pub mod preflight; pub mod rdp; pub mod session; pub mod sessions; +pub mod tasks; pub mod traffic; pub mod tunnel; pub mod update; diff --git a/devolutions-gateway/src/api/tasks.rs b/devolutions-gateway/src/api/tasks.rs new file mode 100644 index 000000000..dfba09fb7 --- /dev/null +++ b/devolutions-gateway/src/api/tasks.rs @@ -0,0 +1,86 @@ +//! The Task JSON returned by the endpoints of provisioner tasks, such as the AI analysis of a recording. + +use provisioner_task::{ProvisionerTaskRecord, ProvisionerTaskState}; +use time::OffsetDateTime; +use uuid::Uuid; + +/// A Task: the permanent record of one piece of work the provisioner asked Gateway to do +/// +/// The work runs in the background; `state` and `payload` tell how it goes. +/// The shape of `params` and `payload` belongs to the kind. +/// For `recording-ai-analysis`, `params` is an `AiAnalysisParams`, and `payload` is an `AiAnalysisRunningPayload` while +/// `running`, an `AiAnalysisSucceededPayload` once `succeeded`, and an `AiAnalysisFailedPayload` once `failed`. +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct TaskInfo { + /// Task ID, chosen by the provisioner. + id: Uuid, + /// Kind of work, such as `recording-ai-analysis`. + kind: String, + /// What the work is about, such as a session ID. + target: String, + /// What was asked, without secrets. + #[cfg_attr(feature = "openapi", schema(value_type = Object))] + params: serde_json::Value, + state: TaskState, + /// Details of the current state; null while `queued`. + #[cfg_attr(feature = "openapi", schema(value_type = Option))] + payload: serde_json::Value, + /// Number of attempts started so far. + attempts: u32, + /// When the Task was created. + #[serde(with = "time::serde::rfc3339")] + created_at: OffsetDateTime, + /// When its first attempt started. + #[serde(with = "time::serde::rfc3339::option")] + started_at: Option, + /// When it succeeded or failed. + #[serde(with = "time::serde::rfc3339::option")] + finished_at: Option, + /// Time at which a Task still unfinished fails; while the Task runs, it moves ahead each time an attempt starts or + /// the Task makes progress. + #[serde(with = "time::serde::rfc3339")] + deadline_at: OffsetDateTime, +} + +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum TaskState { + /// Waiting for its first attempt. + Queued, + /// An attempt runs, or the next one waits for a retry. + Running, + Succeeded, + Failed, +} + +impl From for TaskState { + fn from(state: ProvisionerTaskState) -> Self { + match state { + ProvisionerTaskState::Queued => Self::Queued, + ProvisionerTaskState::Running => Self::Running, + ProvisionerTaskState::Succeeded => Self::Succeeded, + ProvisionerTaskState::Failed => Self::Failed, + } + } +} + +impl From for TaskInfo { + fn from(task: ProvisionerTaskRecord) -> Self { + Self { + id: task.id, + kind: task.kind, + target: task.target, + params: task.params, + state: TaskState::from(task.state), + payload: task.payload, + attempts: task.attempts, + created_at: task.created_at, + started_at: task.started_at, + finished_at: task.finished_at, + deadline_at: task.deadline_at, + } + } +} diff --git a/devolutions-gateway/src/artifacts.rs b/devolutions-gateway/src/artifacts.rs index a2ecd5fa8..d62f5d803 100644 --- a/devolutions-gateway/src/artifacts.rs +++ b/devolutions-gateway/src/artifacts.rs @@ -2,8 +2,8 @@ use serde::{Deserialize, Serialize}; use crate::token::RecordingFileType; -/// Non-recording artifacts, one list per [`ArtifactKind`]. Each list is append-only, like `files`: names -/// are derived from positions. +/// Non-recording artifacts, one list per [`ArtifactKind`]. Each list only grows, like `files`, until it is cleared +/// as a whole: names are derived from positions. #[derive(Debug, Clone, Default, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] pub(crate) struct JrecArtifacts { diff --git a/devolutions-gateway/src/extract.rs b/devolutions-gateway/src/extract.rs index a000e46cf..224c9382a 100644 --- a/devolutions-gateway/src/extract.rs +++ b/devolutions-gateway/src/extract.rs @@ -6,6 +6,7 @@ use axum::http::request::Parts; use crate::DgwState; use crate::http::HttpError; +use crate::provisioner_tasks::runner::ProvisionerTask; use crate::token::{ AccessScope, AccessTokenClaims, AssociationTokenClaims, BridgeTokenClaims, EnrollmentTokenClaims, JmuxTokenClaims, JrecTokenClaims, JrlTokenClaims, KdcTokenClaims, ScopeTokenClaims, WebAppTokenClaims, @@ -321,6 +322,48 @@ where } } +/// A scope token allowed to read the records of the provisioner task `T`: `*`, or the read scope of `T`. +pub struct ProvisionerTaskReadScope(core::marker::PhantomData); + +impl FromRequestParts for ProvisionerTaskReadScope +where + S: Send + Sync, + T: ProvisionerTask, +{ + type Rejection = HttpError; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + let scope = ScopeToken::from_request_parts(parts, state).await?.0.scope; + + if scope == AccessScope::Wildcard || scope == T::READ_SCOPE { + Ok(Self(core::marker::PhantomData)) + } else { + Err(HttpError::forbidden().msg("invalid scope for route")) + } + } +} + +/// A scope token allowed to start the provisioner task `T`: `*`, or the start scope of `T`. +pub struct ProvisionerTaskStartScope(core::marker::PhantomData); + +impl FromRequestParts for ProvisionerTaskStartScope +where + S: Send + Sync, + T: ProvisionerTask, +{ + type Rejection = HttpError; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + let scope = ScopeToken::from_request_parts(parts, state).await?.0.scope; + + if scope == AccessScope::Wildcard || scope == T::START_SCOPE { + Ok(Self(core::marker::PhantomData)) + } else { + Err(HttpError::forbidden().msg("invalid scope for route")) + } + } +} + #[derive(Clone, Copy)] pub struct RecordingsSearchScope; diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 8985f3184..656e9887b 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -6,11 +6,16 @@ use std::time::Duration; use anyhow::Context as _; use async_trait::async_trait; use devolutions_gateway_task::{ChildTask, ShutdownSignal, Task}; -use job_queue::{DynJobQueue, Job, JobCtx, JobQueue, JobReader, JobRunner, RunnerWaker}; +use job_queue::{DynJob, DynJobQueue, Job, JobCtx, JobQueue, JobReader, JobRunner, RunnerWaker}; use job_queue_libsql::libsql; use time::OffsetDateTime; use tokio::sync::{Notify, mpsc}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; + +/// Runs of a job before the job queue gives up on it. +pub const MAX_ATTEMPTS: u32 = 5; + pub struct JobQueueCtx { notify_runner: Arc, runner_waker: RunnerWaker, @@ -38,6 +43,7 @@ pub struct JobRunnerTask { notify_runner: Arc, runner_waker: RunnerWaker, queue: DynJobQueue, + tasks: ProvisionerTaskRunner, } impl JobQueueCtx { @@ -59,6 +65,7 @@ impl JobQueueCtx { let queue = job_queue_libsql::LibSqlJobQueue::builder() .runner_waker(runner_waker.clone()) .conn(conn) + .max_attempts(MAX_ATTEMPTS) .build(); let queue = Arc::new(queue); @@ -82,6 +89,27 @@ impl JobQueueCtx { job_queue_handle: handle, }) } + + /// The queue itself, for jobs whose callers must know they are stored. + pub fn durable(&self) -> DurableJobQueue { + DurableJobQueue(Arc::clone(&self.queue)) + } +} + +/// The job queue database: a job pushed here is stored before `push` returns, unlike with [`JobQueueHandle`]. +#[derive(Clone)] +pub struct DurableJobQueue(DynJobQueue); + +#[async_trait] +impl ProvisionerTaskJobQueue for DurableJobQueue { + async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()> { + let job: DynJob = Box::new(job); + self.0.push_job(&job, schedule_for).await + } + + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + self.0.job_definitions(name).await + } } impl JobQueueHandle { @@ -197,11 +225,12 @@ async fn job_queue_task(ctx: JobQueueTask, mut shutdown_signal: ShutdownSignal) } impl JobRunnerTask { - pub fn new(ctx: &JobQueueCtx) -> Self { + pub fn new(ctx: &JobQueueCtx, tasks: ProvisionerTaskRunner) -> Self { Self { notify_runner: Arc::clone(&ctx.notify_runner), runner_waker: RunnerWaker::clone(&ctx.runner_waker), queue: Arc::clone(&ctx.queue), + tasks, } } } @@ -225,9 +254,10 @@ async fn job_runner_task(ctx: JobRunnerTask, mut shutdown_signal: ShutdownSignal notify_runner, runner_waker, queue, + tasks, } = ctx; - let reader = DgwJobReader; + let reader = DgwJobReader { tasks }; let spawn = |mut ctx: JobCtx, callback: job_queue::SpawnCallback| { tokio::spawn(async move { @@ -279,10 +309,12 @@ async fn job_runner_task(ctx: JobRunnerTask, mut shutdown_signal: ShutdownSignal Ok(()) } -struct DgwJobReader; +struct DgwJobReader { + tasks: ProvisionerTaskRunner, +} impl JobReader for DgwJobReader { - fn read_json(&self, name: &str, json: &str) -> anyhow::Result { + fn read_json(&self, name: &str, json: &str) -> anyhow::Result { use crate::api::jrec::DeleteRecordingsJob; use crate::recording::RemuxJob; @@ -296,6 +328,10 @@ impl JobReader for DgwJobReader { serde_json::from_str(json).context("failed to deserialize DeleteRecordingsJob")?; Ok(Box::new(job)) } + ProvisionerTaskJob::NAME => { + let job = ProvisionerTaskJob::read_json(json, self.tasks.clone())?; + Ok(Box::new(job)) + } _ => anyhow::bail!("unknown job name: {name}"), } } diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index a163bfa40..c20eabc64 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -30,6 +30,7 @@ pub mod log; pub mod middleware; pub mod ngrok; pub mod plugin_manager; +pub mod provisioner_tasks; pub mod provisioning; pub mod proxy; pub mod rd_clean_path; @@ -69,6 +70,8 @@ pub struct DgwState { pub monitoring_state: Arc, pub traffic_audit_handle: traffic_audit::TrafficAuditHandle, pub agent_tunnel_handle: Option>, + pub provisioner_tasks: provisioner_tasks::runner::ProvisionerTaskRunner, + pub secrets: secret_store::SecretStore, } #[doc(hidden)] @@ -83,7 +86,7 @@ pub struct MockHandles { impl DgwState { #[doc(hidden)] - pub fn mock(json_config: &str) -> anyhow::Result<(Self, MockHandles)> { + pub async fn mock(json_config: &str) -> anyhow::Result<(Self, MockHandles)> { let conf_handle = config::ConfHandle::mock(json_config)?; let token_cache = Arc::new(token::new_token_cache()); let jrl = Arc::new(parking_lot::Mutex::new(token::JrlTokenClaims::default())); @@ -96,6 +99,16 @@ impl DgwState { let provisioning = provisioning::ProvisioningStore::new(); let synthetic_kdc_registry = credential_injection::SyntheticKdcRegistry::new(); let monitoring_state = Arc::new(network_monitor::State::new(Arc::new(MockMonitorsCache))?); + let secrets = secret_store::SecretStore::default(); + let task_store = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?; + let task_store = Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store).await?); + let provisioner_tasks = provisioner_tasks::task_runner( + task_store, + Arc::new(job_queue_handle.clone()), + conf_handle.clone(), + recording_manager_handle.clone(), + secrets.clone(), + ); let state = Self { conf_handle, @@ -111,6 +124,8 @@ impl DgwState { synthetic_kdc_registry, monitoring_state, agent_tunnel_handle: None, + provisioner_tasks, + secrets, }; let handles = MockHandles { diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index e023be143..e105dbd25 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -5,6 +5,8 @@ use uuid::Uuid; use crate::api::preflight::PreflightAlertStatus; use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; +use crate::provisioner_tasks::ai::AiProvider; +use crate::provisioner_tasks::{BadAnswer, Permanent, Recoverable}; #[derive(OpenApi)] #[openapi( @@ -25,6 +27,8 @@ use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; crate::api::jrec::pull_recording_session, crate::api::jrec::pull_recording_file, crate::api::jrec::search::search_recordings, + crate::api::jrec::start_ai_analysis, + crate::api::jrec::get_ai_analysis, crate::api::webapp::sign_app_token, crate::api::webapp::sign_session_token, crate::api::update::get_update_products, @@ -64,6 +68,22 @@ use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; crate::api::jrec::search::StartTimeRange, crate::api::jrec::search::RecordingSearchResponse, crate::api::jrec::search::RecordingSearchHit, + crate::api::jrec::AiAnalysisRequest, + crate::api::jrec::AiAnalysisConflict, + crate::api::jrec::AiAnalysisConflictCode, + AiProvider, + crate::api::tasks::TaskInfo, + crate::api::tasks::TaskState, + AiAnalysisParams, + AiAnalysisRunningPayload, + AiAnalysisSucceededPayload, + AiAnalysisUsage, + AiAnalysisFailedPayload, + AiAnalysisError, + AiAnalysisErrorKind, + AiAnalysisBadAnswer, + AiAnalysisRecoverableError, + AiAnalysisRecoverableErrorKind, crate::token::AccessScope, crate::api::webapp::AppTokenSignRequest, crate::api::webapp::AppTokenContentType, @@ -142,6 +162,180 @@ struct SessionInfo { destination_host: Option, } +/// Params of an AI analysis Task: what was asked, without the API key and the base URL +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisParams { + provider: AiProvider, + /// Model identifier, passed to the provider as is. + model: String, + /// Upper bound of tokens in each AI answer; null when not asked. + #[schema(value_type = Option, minimum = 0, maximum = 4_294_967_295_u64)] + max_output_tokens: Option, + /// Language the AI writes in; absent when not asked, so the AI writes in English. + #[serde(skip_serializing_if = "Option::is_none")] + language: Option, +} + +/// Payload of a `running` AI analysis Task: its progress, and the tokens it used so far +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisRunningPayload { + /// Current step: `preparing` (the attempt starts), `reading` (the recordings are read), `describing` (the AI is + /// asked about the recording, part by part), then `writing` (the log is written). + step: String, + /// Work done in the step: bytes read while `reading`, parts described while `describing`; absent in the other steps. + done: Option, + /// Work the step has in all, in the same unit as `done`; absent in the other steps. + total: Option, + usage: AiAnalysisUsage, +} + +/// Payload of a `succeeded` AI analysis Task: the log it put in the session +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisSucceededPayload { + /// Name of the log in the session manifest, such as `ai-analysis-0.slog`. + file_name: String, + /// Model that wrote the log, as reported by the AI provider; the requested model when it reports none. + model: String, + usage: AiAnalysisUsage, + /// Versions of the prompts that wrote the log, such as + /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. + prompt_version: String, +} + +/// Tokens the AI provider counted for an AI analysis Task, over all its attempts +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisUsage { + input_tokens: u64, + output_tokens: u64, + /// Requests that may have cost tokens the provider did not report, such as a request that got no answer, or one + /// still waiting for it; while it is not 0, `inputTokens` and `outputTokens` are short of what the provider billed. + unreported_requests: u64, +} + +/// Payload of a `failed` AI analysis Task: why it failed, and the tokens it used +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisFailedPayload { + error: AiAnalysisError, + /// Absent only for a Task that never ran before its deadline passed, since it used no token. + usage: Option, +} + +/// Why an AI analysis Task failed: a kind, with the fields of that kind; there is no text for people, so the caller +/// words it +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisError { + kind: AiAnalysisErrorKind, + /// HTTP status the AI provider answered, for `aiRejected`. + status: Option, + /// What was wrong with the answer, for `aiBadAnswer`. + why: Option, + /// Recording file that cannot be read, for `unreadableRecording`. + file: Option, + /// Failure of the last attempt, for `attemptsExhausted`. + last: Option, +} + +/// Kind of failure of an AI analysis Task +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +enum AiAnalysisErrorKind { + /// The AI provider refused the request for good, such as for a wrong API key (401) or a spent quota. + AiRejected, + /// The AI provider answered, but not with something usable. + AiBadAnswer, + /// The API key is gone, such as after a Gateway restart. + KeyLost, + /// The session has no recording, or it was deleted. + NoRecording, + /// The session mixes terminal and video recordings. + MixedRecording, + /// A recording file of the session cannot be read. + UnreadableRecording, + /// The Task itself is not valid, such as with parameters this Gateway cannot read. + InvalidTask, + /// The attempt panicked. + Panicked, + /// Something went wrong inside Gateway, such as a file it could not write. + Internal, + /// Every attempt failed with a failure worth another attempt. + AttemptsExhausted, + /// The Task did not start in time, or went too long without progress. + DeadlinePassed, + /// The job of the Task could not be queued. + NotQueued, +} + +/// What was wrong with an answer of the AI provider +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +enum AiAnalysisBadAnswer { + /// The answer is not in the format of the provider API. + Invalid, + /// The answer was cut short by the output token limit or the context window. + Truncated, + /// The provider refused to answer, such as when its content filter blocked the request. + Refused, + /// The model did not answer in the format it was asked for. + NotExpectedJson, +} + +/// A failure of an attempt that was worth another attempt +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisRecoverableError { + kind: AiAnalysisRecoverableErrorKind, + /// HTTP status the AI provider answered, for `aiBusy`. + status: Option, +} + +/// Kind of a failure worth another attempt +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +enum AiAnalysisRecoverableErrorKind { + /// The AI provider could not be reached, or did not answer in time. + AiUnreachable, + /// The AI provider is busy or failing for now, such as with a rate limit (429) or a server error (5xx). + AiBusy, + /// Gateway could not read its own job queue. + Storage, +} + +// The error kinds are written out above for their descriptions, which utoipa drops from the variants without fields, so +// these matches stop the build until a new kind or field is written out too. +const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| match error { + Permanent::AiRejected { status: _ } => AiAnalysisErrorKind::AiRejected, + Permanent::AiBadAnswer { why: _ } => AiAnalysisErrorKind::AiBadAnswer, + Permanent::KeyLost => AiAnalysisErrorKind::KeyLost, + Permanent::NoRecording => AiAnalysisErrorKind::NoRecording, + Permanent::MixedRecording => AiAnalysisErrorKind::MixedRecording, + Permanent::UnreadableRecording { file: _ } => AiAnalysisErrorKind::UnreadableRecording, + Permanent::InvalidTask => AiAnalysisErrorKind::InvalidTask, + Permanent::Panicked => AiAnalysisErrorKind::Panicked, + Permanent::Internal => AiAnalysisErrorKind::Internal, + Permanent::AttemptsExhausted { last: _ } => AiAnalysisErrorKind::AttemptsExhausted, + Permanent::DeadlinePassed => AiAnalysisErrorKind::DeadlinePassed, + Permanent::NotQueued => AiAnalysisErrorKind::NotQueued, +}; + +const _: fn(&Recoverable) -> AiAnalysisRecoverableErrorKind = |error| match error { + Recoverable::AiUnreachable => AiAnalysisRecoverableErrorKind::AiUnreachable, + Recoverable::AiBusy { status: _ } => AiAnalysisRecoverableErrorKind::AiBusy, + Recoverable::Storage => AiAnalysisRecoverableErrorKind::Storage, +}; + +const _: fn(&BadAnswer) -> AiAnalysisBadAnswer = |why| match why { + BadAnswer::Invalid => AiAnalysisBadAnswer::Invalid, + BadAnswer::Truncated => AiAnalysisBadAnswer::Truncated, + BadAnswer::Refused => AiAnalysisBadAnswer::Refused, + BadAnswer::NotExpectedJson => AiAnalysisBadAnswer::NotExpectedJson, +}; + #[allow(unused)] #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "kebab-case")] diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md new file mode 100644 index 000000000..ca2b0010b --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -0,0 +1,41 @@ +# Background +Provisioner tasks are tasks issued for Gateway to execute. Tasks's state are pesisted in `gateway.db`, see more at `../../../crates/provisioner-task-libsql/`. + + +# Invariants + +## Management +`ProvisionerTaskRunner` MUST be the only place touches both the Task records and the job queue. It is also the only place that manages how tasks are executed in relation to the job queue. + +A provisioner task only defines it's work: what one attempt does, and what to drop once its task ends. The coordination belongs to the `ProvisionerTaskRunner`. + +The task record is the source of truth; the `ProvisionerTaskJob` only wakes it up and holds no state. + +One task at the same time only has one job to execute task. + +Jobs doesn't have state, tasks do. + +An attempt error is recoverable or permanent: a recoverable one is retried with the job up to a limit, and a permanent one fails the Task at once. A failed Task's payload is that permanent error, giving up included, as a closed set of kinds without text for people, because Gateway does not translate: the caller words it. + +When Gateway starts, it gives a new job to every unfinished Task that has none. This can happen if Gateway stopped after saving a Task but before saving its job. Gateway starts the job without checking whether the Task can still succeed: each kind of Task decides that for itself. For example, an AI analysis fails because its key is gone. If the new job can't be saved, the Task fails. + + +## Naming +`ProvisionerTask` is the work itself; `ProvisionerTaskRecord` is what is stored about it. + +## Secrets +Secrets a Task needs, such as an API key or a base URL that may carry credentials, never reach the disk: not `gateway.db`, not the job, not the logs. +They are kept encrypted in memory, in the `SecretStore`, until the Task ends or reaches its deadline. +After a Gateway restart they are gone, and the Tasks that needed them fail; this is accepted. + +## AI analysis +Only a finished recording is analyzed. +A session that mixes terminal and video recordings cannot be analyzed. +Two analyses of the same session never run at the same time. +An analysis keeps working files next to the recording, so a retry can pick up where it stopped. When a Task ends, it deletes its own working files, whether it succeeded or failed. A new Task first deletes any working files an earlier analysis left behind. A Task never deletes another Task's working files. +A session keeps one AI analysis: a new Task replaces the earlier one atomically once it succeeds, and a Task that fails leaves the earlier one in place. Keeping one analysis per language is a future extension that must stay possible. +A Task succeeds only once its success is recorded in `gateway.db`. Until then, its new log sits in the session unlisted and the earlier analysis stays. After that, the manifest lists the new log and drops the earlier one; if this last step is cut short, Gateway does it again when the Task ends and when Gateway starts. +A recoverable failure, such as a rate limit (429), is retried with the job, and each retry continues from the progress saved so far; only a new Task starts over. A Task that fails keeps its record, but removes everything it produced, leaving no file on disk, even if that wastes the money already spent, because an analysis missing a large part would mislead whoever searches it, more than the misses of a nondeterministic AI already do. A Gateway restart loses the API key, so an unfinished Task fails with `keyLost` at its next attempt, without a retry. +A rate limit, a server error, a lost connection or a request timeout is recoverable; a refused API key (401), a spent quota, or an answer the model cut short or refused is permanent. +An AI analysis fails if it hasn't started within 2 hours. Once it runs, it fails only after 30 minutes without progress, so a long recording that keeps making progress is never cut off. Each request to the AI times out after 10 minutes and is retried like any other recoverable failure. +An analysis waits until Gateway has finished remuxing the session's videos, because remuxing replaces the file the analysis reads. Waiting is not progress, since remuxing should not take that long. diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs new file mode 100644 index 000000000..294cf18bb --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -0,0 +1,303 @@ +//! AI settings shared by the provisioner tasks that call an AI provider. +//! +//! The API key and the base URL never reach `gateway.db`: they stay in memory, encrypted, as the secret of their Task +//! in the [`SecretStore`](crate::secret_store::SecretStore), until the Task ends. + +use core::fmt; + +use anyhow::Context as _; +use devolutions_gateway_ai::recording_analysis::RecordingAnalysisError; +use devolutions_gateway_ai::{AiClient, BuildError, Error, Provider}; +use secrecy::{ExposeSecret as _, SecretString}; +use url::Url; + +use super::{AttemptError, BadAnswer, Permanent, Recoverable}; +use crate::config::Conf; + +/// AI provider to ask; `openai-compatible` needs a base URL +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum AiProvider { + #[serde(rename = "openai")] + OpenAi, + #[serde(rename = "anthropic")] + Anthropic, + #[serde(rename = "mistral")] + Mistral, + #[serde(rename = "gemini")] + Gemini, + #[serde(rename = "openai-compatible")] + OpenAiCompatible, +} + +impl From for Provider { + fn from(provider: AiProvider) -> Self { + match provider { + AiProvider::OpenAi => Provider::OpenAi, + AiProvider::Anthropic => Provider::Anthropic, + AiProvider::Mistral => Provider::Mistral, + AiProvider::Gemini => Provider::Gemini, + AiProvider::OpenAiCompatible => Provider::OpenAiCompatible, + } + } +} + +/// How to reach an AI provider, except the API key. +#[derive(Clone, PartialEq, Eq)] +pub struct AiSettings { + pub provider: AiProvider, + pub model: String, + /// Required for [`AiProvider::OpenAiCompatible`]; the others have a default. + pub base_url: Option, + /// Upper bound of tokens in each AI answer. + pub max_output_tokens: Option, +} + +impl fmt::Debug for AiSettings { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AiSettings") + .field("provider", &self.provider) + .field("model", &self.model) + .field("base_url", &self.base_url.as_ref().map(|_| "")) + .field("max_output_tokens", &self.max_output_tokens) + .finish() + } +} + +/// Why AI settings were refused before a Task was recorded. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum AiSettingsError { + #[error("AI model is missing")] + MissingModel, + #[error("API key is missing")] + MissingApiKey, + #[error("base URL is missing for this AI provider")] + MissingBaseUrl, + #[error("invalid AI settings: {0}")] + Invalid(String), +} + +impl From<&BuildError> for AiSettingsError { + fn from(error: &BuildError) -> Self { + match error { + BuildError::MissingModel => Self::MissingModel, + BuildError::MissingApiKey(_) => Self::MissingApiKey, + BuildError::MissingBaseUrl(_) => Self::MissingBaseUrl, + _ => Self::Invalid(error.to_string()), + } + } +} + +pub(crate) enum ClientError { + Build(BuildError), + HttpClient(reqwest::Error), +} + +impl fmt::Display for ClientError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ClientError::Build(error) => fmt::Display::fmt(error, f), + ClientError::HttpClient(error) => write!(f, "failed to build the HTTP client: {error}"), + } + } +} + +/// What a Task needs to reach its AI provider and must never write to disk. +#[derive(Clone)] +pub(crate) struct AiAccess { + pub(crate) api_key: SecretString, + pub(crate) base_url: Option, +} + +impl fmt::Debug for AiAccess { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AiAccess") + .field("api_key", &self.api_key) + .field("base_url", &self.base_url.as_ref().map(|_| "")) + .finish() + } +} + +/// [`AiAccess`] as written into its task secret. +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct AiAccessSecretRef<'a> { + api_key: &'a str, + base_url: Option<&'a Url>, +} + +/// [`AiAccess`] as read back from its task secret. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct AiAccessSecret { + api_key: String, + base_url: Option, +} + +impl AiAccess { + /// The task secret holding this access. + pub(crate) fn to_secret(&self) -> anyhow::Result { + let json = serde_json::to_string(&AiAccessSecretRef { + api_key: self.api_key.expose_secret(), + base_url: self.base_url.as_ref(), + }) + .context("serialize the AI access")?; + + Ok(SecretString::from(json)) + } + + /// Reads back an access from its task secret. + pub(crate) fn from_secret(secret: &SecretString) -> anyhow::Result { + // The parse error is dropped: it may quote the API key. + let access: AiAccessSecret = + serde_json::from_str(secret.expose_secret()).map_err(|_| anyhow::anyhow!("invalid AI access secret"))?; + + Ok(Self { + api_key: SecretString::from(access.api_key), + base_url: access.base_url, + }) + } +} + +impl AiSettings { + /// Builds a client going through the proxy configured for Gateway. + pub(crate) fn client(&self, conf: &Conf, api_key: &SecretString) -> Result { + let provider = Provider::from(self.provider); + + let mut builder = AiClient::builder() + .provider(provider) + .model(self.model.clone()) + .api_key(api_key.clone()); + + let endpoint = self.base_url.clone().or_else(|| provider.default_base_url()); + + if let Some(base_url) = self.base_url.clone() { + builder = builder.base_url(base_url); + } + + // Without an endpoint, `build` reports the missing base URL before it needs the HTTP client. + if let Some(endpoint) = endpoint { + let proxy_config = conf.proxy.to_proxy_config(); + + let http_client = + http_client_proxy::get_or_create_cached_client(reqwest::Client::builder(), &endpoint, &proxy_config) + .map_err(|error| ClientError::HttpClient(error.without_url()))?; + + builder = builder.http_client(http_client); + } + + builder.build().map_err(ClientError::Build) + } +} + +/// Retries only the failures that may pass later, such as a rate limit or a network error. +impl From for AttemptError { + fn from(error: RecordingAnalysisError) -> Self { + match error { + RecordingAnalysisError::Ai(error) => AttemptError::from(error), + RecordingAnalysisError::NoRecording => AttemptError::Permanent(Permanent::NoRecording), + RecordingAnalysisError::InvalidRecording { file_name, .. } => { + AttemptError::Permanent(Permanent::UnreadableRecording { file: file_name }) + } + RecordingAnalysisError::Truncated => AttemptError::Permanent(Permanent::AiBadAnswer { + why: BadAnswer::Truncated, + }), + // The runner cancels an attempt only once its Task is past its deadline. + RecordingAnalysisError::Cancelled => AttemptError::Permanent(Permanent::DeadlinePassed), + _ => AttemptError::Permanent(Permanent::Internal), + } + } +} + +impl From for AttemptError { + fn from(error: Error) -> Self { + let bad_answer = |why| AttemptError::Permanent(Permanent::AiBadAnswer { why }); + + match &error { + Error::Transport { .. } => AttemptError::Recoverable(Recoverable::AiUnreachable), + Error::Status { status, .. } if error.is_transient() => { + AttemptError::Recoverable(Recoverable::AiBusy { status: *status }) + } + Error::Status { status, .. } => AttemptError::Permanent(Permanent::AiRejected { status: *status }), + Error::InvalidResponse { .. } => bad_answer(BadAnswer::Invalid), + Error::Truncated { .. } => bad_answer(BadAnswer::Truncated), + Error::Refused { .. } => bad_answer(BadAnswer::Refused), + Error::InvalidOutput { .. } => bad_answer(BadAnswer::NotExpectedJson), + _ => AttemptError::Permanent(Permanent::Internal), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + #[test] + fn transient_ai_errors_are_recoverable_and_the_others_permanent() { + let attempt = |error: Error| AttemptError::from(RecordingAnalysisError::Ai(error)); + let status = |status, code: Option<&str>| Error::Status { + status, + message: "failed".to_owned(), + code: code.map(str::to_owned), + retry_after: None, + }; + + let network = Error::Transport { + message: "connection refused".to_owned(), + }; + assert_eq!(attempt(network), AttemptError::Recoverable(Recoverable::AiUnreachable)); + + for code in [429, 500, 503] { + assert_eq!( + attempt(status(code, None)), + AttemptError::Recoverable(Recoverable::AiBusy { status: code }) + ); + } + + for code in [400, 401, 403, 404] { + assert_eq!( + attempt(status(code, None)), + AttemptError::Permanent(Permanent::AiRejected { status: code }) + ); + } + + assert_eq!( + attempt(status(429, Some("insufficient_quota"))), + AttemptError::Permanent(Permanent::AiRejected { status: 429 }), + "a spent quota" + ); + assert_eq!( + AttemptError::from(RecordingAnalysisError::Truncated), + AttemptError::Permanent(Permanent::AiBadAnswer { + why: BadAnswer::Truncated + }) + ); + assert_eq!( + attempt(Error::Refused { + reason: "content filter".to_owned() + }), + AttemptError::Permanent(Permanent::AiBadAnswer { + why: BadAnswer::Refused + }) + ); + } + + // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis + #[test] + fn providers_keep_their_wire_names() { + for (provider, name) in [ + (AiProvider::OpenAi, "openai"), + (AiProvider::Anthropic, "anthropic"), + (AiProvider::Mistral, "mistral"), + (AiProvider::Gemini, "gemini"), + (AiProvider::OpenAiCompatible, "openai-compatible"), + ] { + assert_eq!(serde_json::to_value(provider).expect("serializable"), name); + assert_eq!( + serde_json::from_value::(serde_json::Value::from(name)).expect("known name"), + provider + ); + } + } +} diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs new file mode 100644 index 000000000..2c49d0330 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -0,0 +1,97 @@ +//! Work the provisioner asks Gateway to do: each request is a Task recorded in `gateway.db`, run as a job on the job +//! queue by the [`runner::ProvisionerTaskRunner`]. + +pub mod ai; +pub mod recording_ai_analysis; +pub mod runner; + +use std::sync::Arc; + +use provisioner_task::DynProvisionerTaskStore; + +use self::recording_ai_analysis::RecordingAiAnalysis; +use self::runner::{DynProvisionerTaskJobQueue, ProvisionerTaskRunner}; +use crate::config::ConfHandle; +use crate::recording::RecordingMessageSender; +use crate::secret_store::SecretStore; + +/// The provisioner task runner, with every provisioner task Gateway knows. +pub fn task_runner( + store: DynProvisionerTaskStore, + job_queue: DynProvisionerTaskJobQueue, + conf_handle: ConfHandle, + recordings: RecordingMessageSender, + secrets: SecretStore, +) -> ProvisionerTaskRunner { + ProvisionerTaskRunner::builder(store, Arc::clone(&job_queue)) + .register(RecordingAiAnalysis::new(conf_handle, recordings, secrets, job_queue)) + .build() +} + +/// Why one attempt of a Task failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AttemptError { + /// Worth another attempt later, such as a rate limit or a network error. + Recoverable(Recoverable), + /// Another attempt would fail the same way, so the Task fails at once. + Permanent(Permanent), +} + +/// A failure worth another attempt. +/// +/// It reaches the caller only as the `last` failure of [`Permanent::AttemptsExhausted`]. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(tag = "kind", rename_all = "camelCase", rename_all_fields = "camelCase")] +pub enum Recoverable { + /// The AI provider could not be reached, or did not answer in time. + AiUnreachable, + /// The AI provider is busy or failing for now, such as with a rate limit (429) or a server error (5xx). + AiBusy { status: u16 }, + /// Gateway could not read its own job queue. + Storage, +} + +/// Why a Task failed, as the `error` of its payload: a closed set of kinds without text for people, since the caller +/// words them. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(tag = "kind", rename_all = "camelCase", rename_all_fields = "camelCase")] +pub enum Permanent { + /// The AI provider refused the request for good, such as for a wrong API key (401) or a spent quota. + AiRejected { status: u16 }, + /// The AI provider answered, but not with something usable. + AiBadAnswer { why: BadAnswer }, + /// The API key is gone, such as after a Gateway restart. + KeyLost, + /// The session has no recording, or it was deleted. + NoRecording, + /// The session mixes terminal and video recordings. + MixedRecording, + /// A recording file of the session cannot be read. + UnreadableRecording { file: String }, + /// The Task itself is not valid, such as with parameters this Gateway cannot read. + InvalidTask, + /// The attempt panicked. + Panicked, + /// Something went wrong inside Gateway, such as a file it could not write. + Internal, + /// Every attempt failed with a recoverable failure, the last one being `last`. + AttemptsExhausted { last: Recoverable }, + /// The Task did not start in time, or went too long without progress. + DeadlinePassed, + /// The job of the Task could not be queued. + NotQueued, +} + +/// What was wrong with an answer of the AI provider. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub enum BadAnswer { + /// The answer is not in the format of the provider API. + Invalid, + /// The answer was cut short by the output token limit or the context window. + Truncated, + /// The provider refused to answer, such as when its content filter blocked the request. + Refused, + /// The model did not answer in the format it was asked for. + NotExpectedJson, +} diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs new file mode 100644 index 000000000..33a5ca427 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -0,0 +1,685 @@ +//! AI analysis (`recording-ai-analysis`): describes what the user did in one recorded session, and puts the result in +//! the session as its log, in place of the earlier one. +//! +//! [`start`] checks the request and records the Task through the +//! [`ProvisionerTaskRunner`](super::runner::ProvisionerTaskRunner). +//! Each attempt runs the analysis of the AI crate, +//! [`AiClient::analyze_recording`](devolutions_gateway_ai::AiClient::analyze_recording), which keeps its working files +//! next to the session manifest so that a retry resumes where the last attempt stopped, then leaves the log in the +//! session folder, unlisted, under a name of its own Task. +//! Once the runner records the success, the end of the Task lists that log in the manifest in place of the earlier one; +//! a Task that fails removes it instead. +//! The API key and the working files go away when the Task ends. + +#[cfg(test)] +mod tests; + +use std::collections::HashMap; +use std::sync::{Arc, LazyLock, Weak}; + +use anyhow::Context as _; +use async_trait::async_trait; +use camino::{Utf8Path, Utf8PathBuf}; +use devolutions_gateway_ai::Language; +use devolutions_gateway_ai::recording_analysis::{ + self, Progress, RecordingAnalysisEvent, RecordingFile, RecordingManifest, Recordings, RecordingsError, UsageSoFar, +}; +use provisioner_task::{CreateOutcome, ProvisionerTaskRecord, ProvisionerTaskState}; +use secrecy::SecretString; +use serde_json::json; +use time::OffsetDateTime; +use uuid::Uuid; + +use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; +use super::runner::{Attempt, DynProvisionerTaskJobQueue, ProvisionerTask}; +use super::{AttemptError, Permanent, Recoverable}; +use crate::DgwState; +use crate::artifacts::ArtifactKind; +use crate::config::ConfHandle; +use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender, RemuxJob, remux_pending}; +use crate::secret_store::{SecretId, SecretStore}; +use crate::token::AccessScope; + +pub const KIND: &str = "recording-ai-analysis"; + +/// Time from the creation of an AI analysis until it must start. +pub const START_WINDOW: time::Duration = time::Duration::hours(2); + +/// Time a running AI analysis may go without progress before it fails. +pub const PROGRESS_TIMEOUT: time::Duration = time::Duration::minutes(30); + +/// Time between two checks for a remux of the session. +const REMUX_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1); + +/// Tokens the AI provider counted for a Task, over all its attempts, as kept in its payloads. +/// +/// It mirrors [`UsageSoFar`], so a change in the AI crate never changes what we store. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct TokenUsage { + input_tokens: u64, + output_tokens: u64, + /// Requests whose tokens the provider did not report, so the two counts are short of what it billed. + unreported_requests: u64, +} + +impl TokenUsage { + /// The usage in a payload of the Task; none before its first attempt. + fn of(payload: &serde_json::Value) -> Self { + payload + .get("usage") + .and_then(|usage| serde_json::from_value(usage.clone()).ok()) + .unwrap_or_default() + } + + /// This usage, with that of the attempt running now. + fn with(self, attempt: UsageSoFar) -> Self { + Self { + input_tokens: self.input_tokens.saturating_add(attempt.input_tokens), + output_tokens: self.output_tokens.saturating_add(attempt.output_tokens), + unreported_requests: self.unreported_requests.saturating_add(attempt.unreported_requests), + } + } +} + +/// Payload of a running AI analysis Task: how far it is, and the tokens it used so far. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct Running { + /// `preparing` until the analysis reports its stage: `reading`, `describing` or `writing`. + step: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + done: Option, + #[serde(skip_serializing_if = "Option::is_none")] + total: Option, + usage: TokenUsage, +} + +/// What a caller asks for when it starts an AI analysis. +#[derive(Debug)] +pub struct StartRequest { + /// Chosen by the caller, so it can ask again safely. + pub task_id: Uuid, + pub session_id: Uuid, + pub settings: AiSettings, + /// Language of the analysis; English when absent. + pub language: Option, + pub api_key: SecretString, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum StartOutcome { + /// A new Task is recorded and its job is queued. + Created(ProvisionerTaskRecord), + /// The same request came in before; nothing new runs. + Existing(ProvisionerTaskRecord), + /// The Task ID is already used for something else. + IdConflict(ProvisionerTaskRecord), + /// Another AI analysis of this session is not finished yet. + TargetBusy { active_task_id: Uuid }, +} + +#[derive(Debug, thiserror::Error)] +pub enum StartError { + #[error(transparent)] + InvalidSettings(AiSettingsError), + /// The session is still recording, so there is nothing complete to analyse yet. + #[error("session is still recording")] + RecordingActive, + /// The session has no recording to analyse. + #[error("session has no recording")] + RecordingNotFound, + #[error(transparent)] + Internal(anyhow::Error), +} + +/// What an AI analysis Task records as its parameters: the AI settings without the API key and the base URL, and the +/// language. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Params { + provider: AiProvider, + model: String, + max_output_tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + language: Option, +} + +/// Result payload of a successful AI analysis. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct Analysis { + /// Name of the new log in the session manifest, such as `ai-analysis-0.slog`. + file_name: String, + /// Model that wrote the log, as reported by the AI provider; the requested model when it reports none. + model: String, + /// Tokens the AI provider counted for the Task, over all its attempts. + usage: TokenUsage, + /// Versions of the prompts that wrote the log, such as + /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. + prompt_version: String, +} + +/// Records an AI analysis Task, whose job the [`ProvisionerTaskRunner`](super::runner::ProvisionerTaskRunner) queues. +/// +/// Only a new Task keeps the API key; asking again for a known Task changes nothing, even while its session is +/// recording again. +pub async fn start(state: &DgwState, request: StartRequest) -> Result { + let StartRequest { + task_id, + session_id, + settings, + language, + api_key, + } = request; + + let conf = state.conf_handle.get_conf(); + + match settings.client(&conf, &api_key) { + Ok(_) => {} + Err(ClientError::Build(error)) => { + debug!(%error, "Invalid AI settings"); + return Err(StartError::InvalidSettings(AiSettingsError::from(&error))); + } + Err(ClientError::HttpClient(error)) => { + return Err(StartError::Internal( + anyhow::Error::new(error).context("build the HTTP client for the AI provider"), + )); + } + } + + let params = Params { + provider: settings.provider, + model: settings.model, + max_output_tokens: settings.max_output_tokens, + language, + }; + let target = session_id.to_string(); + + // A repeated request gets its Task back whatever its session does now, so the ID is looked up first. + let known = state + .provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .map_err(StartError::Internal)?; + + if let Some(record) = known { + let params = serde_json::to_value(¶ms) + .context("serialize the AI analysis parameters") + .map_err(StartError::Internal)?; + let same_request = record.kind == KIND && record.target == target && record.params == params; + return Ok(if same_request { + StartOutcome::Existing(record) + } else { + StartOutcome::IdConflict(record) + }); + } + + if state.recordings.active_recordings.contains(session_id) { + return Err(StartError::RecordingActive); + } + + match state.recordings.get_finished(session_id).await { + Ok(_) => {} + Err(FinishedRecordingError::Recording) => return Err(StartError::RecordingActive), + Err(FinishedRecordingError::NotFound) => return Err(StartError::RecordingNotFound), + Err(FinishedRecordingError::Other(error)) => { + return Err(StartError::Internal(error.context("read the recording"))); + } + } + + let access = AiAccess { + api_key, + base_url: settings.base_url, + }; + + let outcome = state + .provisioner_tasks + .create::(task_id, target, ¶ms, |record| { + // The job may run as soon as it is queued, so the key must be in place first. + state + .secrets + .insert(SecretId(record.id), &access.to_secret()?, record.deadline_at) + }) + .await + .map_err(StartError::Internal)?; + + Ok(match outcome { + CreateOutcome::Created(task) => { + debug!(task.id = %task_id, session.id = %session_id, "AI analysis queued"); + StartOutcome::Created(task) + } + CreateOutcome::Existing(task) => StartOutcome::Existing(task), + CreateOutcome::IdConflict(task) => StartOutcome::IdConflict(task), + CreateOutcome::TargetBusy(active) => StartOutcome::TargetBusy { + active_task_id: active.id, + }, + }) +} + +/// The `recording-ai-analysis` provisioner task. +pub struct RecordingAiAnalysis { + conf_handle: ConfHandle, + recordings: RecordingMessageSender, + secrets: SecretStore, + job_queue: DynProvisionerTaskJobQueue, + progress_timeout: time::Duration, +} + +impl RecordingAiAnalysis { + pub fn new( + conf_handle: ConfHandle, + recordings: RecordingMessageSender, + secrets: SecretStore, + job_queue: DynProvisionerTaskJobQueue, + ) -> Self { + Self { + conf_handle, + recordings, + secrets, + job_queue, + progress_timeout: PROGRESS_TIMEOUT, + } + } + + fn manifest_path(&self, session_id: Uuid) -> Utf8PathBuf { + self.conf_handle + .get_conf() + .recording_path + .join(session_id.to_string()) + .join("recording.json") + } +} + +const UNLISTED_LOG_PREFIX: &str = ".ai-analysis-"; + +const UNLISTED_LOG_SUFFIX: &str = ".slog"; + +/// Where a Task leaves its log until its success is recorded: in the session folder, but not in the manifest, under a +/// name of its own, so the end of another Task leaves it alone. +fn unlisted_log_path(manifest_path: &Utf8Path, task_id: Uuid) -> Utf8PathBuf { + manifest_path.with_file_name(format!("{UNLISTED_LOG_PREFIX}{task_id}{UNLISTED_LOG_SUFFIX}")) +} + +/// The logs that Tasks of the session left unlisted, each with the ID of its Task. +async fn unlisted_logs(manifest_path: &Utf8Path) -> std::io::Result> { + let Some(dir) = manifest_path.parent() else { + return Ok(Vec::new()); + }; + let mut entries = match tokio::fs::read_dir(dir).await { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) => return Err(error), + }; + let mut logs = Vec::new(); + + while let Some(entry) = entries.next_entry().await? { + let name = entry.file_name(); + let Some(name) = name.to_str() else { + continue; + }; + let task_id = name + .strip_prefix(UNLISTED_LOG_PREFIX) + .and_then(|name| name.strip_suffix(UNLISTED_LOG_SUFFIX)) + .and_then(|task_id| Uuid::parse_str(task_id).ok()); + + if let Some(task_id) = task_id { + logs.push((task_id, dir.join(name))); + } + } + + Ok(logs) +} + +/// Removes `path`, unless it is already gone. +async fn remove_if_exists(path: &Utf8Path) -> std::io::Result<()> { + match tokio::fs::remove_file(path).await { + Err(error) if error.kind() != std::io::ErrorKind::NotFound => Err(error), + _ => Ok(()), + } +} + +#[async_trait] +impl ProvisionerTask for RecordingAiAnalysis { + const KIND: &'static str = KIND; + + const START_SCOPE: AccessScope = AccessScope::RecordingAiAnalysisStart; + + const READ_SCOPE: AccessScope = AccessScope::RecordingAiAnalysisRead; + + type Params = Params; + + type Output = Analysis; + + fn start_window(&self) -> time::Duration { + START_WINDOW + } + + fn progress_timeout(&self) -> time::Duration { + self.progress_timeout + } + + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value { + json!(Running { + step: "preparing", + done: None, + total: None, + usage: TokenUsage::of(last), + }) + } + + fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value { + json!({ "error": error, "usage": TokenUsage::of(last) }) + } + + async fn run(&self, attempt: &Attempt, params: Params) -> Result { + let session_id = Uuid::parse_str(&attempt.target).map_err(|error| { + warn!(task.id = %attempt.task_id, %error, "Invalid session ID"); + AttemptError::Permanent(Permanent::InvalidTask) + })?; + + // UNTESTED: the key lives until the deadline, which this attempt start moved, so a retry hours after the Task was + // created still finds it; every move of the deadline must move the key's expiry too. + self.secrets.set_expiry(SecretId(attempt.task_id), attempt.deadline()); + + // Waiting for an earlier analysis of the session is no progress, so the deadline cancels it too. + let _session = tokio::select! { + session = lock_session(session_id) => session, + () = attempt.cancel.cancelled() => return Err(AttemptError::Permanent(Permanent::DeadlinePassed)), + }; + + // A new Task never resumes from what an earlier analysis left, while its retries resume from their own files. + if attempt.attempts == 1 { + let manifest_path = self.manifest_path(session_id); + + recording_analysis::discard(&manifest_path, None) + .await + .map_err(|error| { + error!(task.id = %attempt.task_id, %error, "Failed to remove the earlier AI analysis working files"); + AttemptError::Permanent(Permanent::Internal) + })?; + + self.end_unlisted_logs(attempt, session_id, &manifest_path) + .await + .map_err(|error| { + error!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to end the logs earlier AI analyses left unlisted"); + AttemptError::Permanent(Permanent::Internal) + })?; + } + + let secret = self.secrets.get(SecretId(attempt.task_id)).map_err(|error| { + error!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to read the API key"); + AttemptError::Permanent(Permanent::Internal) + })?; + + let Some(secret) = secret else { + warn!(task.id = %attempt.task_id, attempts = attempt.attempts, "AI analysis lost its API key"); + return Err(AttemptError::Permanent(Permanent::KeyLost)); + }; + + let access = AiAccess::from_secret(&secret).map_err(|error| { + error!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to read the API key"); + AttemptError::Permanent(Permanent::Internal) + })?; + + let settings = AiSettings { + provider: params.provider, + model: params.model, + base_url: access.base_url, + max_output_tokens: params.max_output_tokens, + }; + + self.analyze(attempt, session_id, settings, params.language, access.api_key) + .await + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { + self.secrets.remove(SecretId(record.id)); + + let Ok(session_id) = Uuid::parse_str(&record.target) else { + return Ok(()); + }; + + let manifest_path = self.manifest_path(session_id); + let log_path = unlisted_log_path(&manifest_path, record.id); + let _session = lock_session(session_id).await; + + self.end_log(session_id, &log_path, record.state == ProvisionerTaskState::Succeeded) + .await?; + + recording_analysis::discard(&manifest_path, Some(&record.id.to_string())) + .await + .context("remove the AI analysis working files") + } +} + +impl RecordingAiAnalysis { + /// Runs the analysis while the caller holds the lock of the session. + async fn analyze( + &self, + attempt: &Attempt, + session_id: Uuid, + settings: AiSettings, + language: Option, + api_key: SecretString, + ) -> Result { + let conf = self.conf_handle.get_conf(); + + let client = settings.client(&conf, &api_key).map_err(|error| { + error!(task.id = %attempt.task_id, %error, "Failed to build the AI client"); + AttemptError::Permanent(Permanent::Internal) + })?; + + let recording = self + .recordings + .get_finished(session_id) + .await + .map_err(|error| match error { + FinishedRecordingError::NotFound => AttemptError::Permanent(Permanent::NoRecording), + error => { + error!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to read the recording"); + AttemptError::Permanent(Permanent::Internal) + } + })?; + + let manifest = analysis_manifest(&recording.manifest).map_err(|error| match error { + RecordingsError::Mixed => AttemptError::Permanent(Permanent::MixedRecording), + _ => AttemptError::Permanent(Permanent::NoRecording), + })?; + let manifest_path = recording.manifest_path(); + + self.wait_for_remux(attempt, &recording.dir).await?; + + let mut request = client + .analyze_recording(&manifest, &manifest_path) + .owner(attempt.task_id.to_string()) + .cancellation(attempt.cancel.clone()); + + if let Some(max_output_tokens) = settings.max_output_tokens { + request = request.max_output_tokens(max_output_tokens); + } + + if let Some(language) = language { + request = request.language(language); + } + + let earlier = TokenUsage::of(&attempt.payload); + let mut usage = earlier; + let mut events = request.start(); + let mut recorded = None; + + let analysis = loop { + match events.recv().await { + Some(RecordingAnalysisEvent::Progress(progress)) => { + // Reading sends an event every half second, so while reading the Task is updated only when the percent + // changes. + let done = match progress { + Progress::Reading { .. } => u64::from(progress.percent()), + _ => progress.done_and_total().map_or(0, |(done, _)| done), + }; + let attempt_usage = match progress { + Progress::Describing { usage, .. } => Some(usage), + _ => None, + }; + let shown = (progress.stage(), done, attempt_usage); + + if recorded != Some(shown) { + recorded = Some(shown); + + if let Some(attempt_usage) = attempt_usage { + usage = earlier.with(attempt_usage); + } + + attempt.report_progress(progress_payload(&progress, usage)).await; + self.secrets.set_expiry(SecretId(attempt.task_id), attempt.deadline()); + } + } + Some(RecordingAnalysisEvent::Complete(analysis)) => break analysis, + Some(RecordingAnalysisEvent::Failure(error)) => { + warn!(task.id = %attempt.task_id, %error, "AI analysis attempt failed"); + return Err(error.into()); + } + Some(_) => {} + None => { + error!(task.id = %attempt.task_id, "The AI analysis stopped unexpectedly"); + return Err(AttemptError::Permanent(Permanent::Panicked)); + } + } + }; + + // The log stays out of the manifest until the runner records the success, so a Task that fails never replaces + // the earlier analysis. + let log_path = unlisted_log_path(&manifest_path, attempt.task_id); + tokio::fs::rename(&analysis.log_path, &log_path) + .await + .map_err(|error| { + error!(task.id = %attempt.task_id, %error, "Failed to keep the AI analysis log"); + AttemptError::Permanent(Permanent::Internal) + })?; + + debug!(task.id = %attempt.task_id, actions = analysis.actions, "AI analysis log written"); + + Ok(Analysis { + file_name: ArtifactKind::AiAnalysis.file_name(0), + model: analysis.model, + usage: earlier.with(analysis.usage), + prompt_version: analysis.prompt_version, + }) + } +} + +impl RecordingAiAnalysis { + /// What the end of a Task does to the log it left unlisted: the log of a Task whose success is recorded takes the + /// place of the earlier analysis, and any other log is removed. + async fn end_log(&self, session_id: Uuid, log_path: &Utf8Path, succeeded: bool) -> anyhow::Result<()> { + if succeeded { + self.recordings + .replace_artifacts(session_id, ArtifactKind::AiAnalysis, log_path) + .await + .with_context(|| format!("list {log_path}")) + } else { + remove_if_exists(log_path) + .await + .with_context(|| format!("remove {log_path}")) + } + } + + /// Does what the end of an earlier Task does to the log it left unlisted, since that end may still be due, so a + /// newer Task never loses an analysis whose success is recorded: the log of a succeeded Task is listed, so it stays + /// until a newer analysis succeeds, and the log of a failed or unknown one is removed. + async fn end_unlisted_logs( + &self, + attempt: &Attempt, + session_id: Uuid, + manifest_path: &Utf8Path, + ) -> anyhow::Result<()> { + for (task_id, log_path) in unlisted_logs(manifest_path).await.context("list the unlisted logs")? { + match attempt + .state_of(task_id) + .await + .context("read the Task of an unlisted log")? + { + Some(ProvisionerTaskState::Queued | ProvisionerTaskState::Running) => {} + state => { + self.end_log(session_id, &log_path, state == Some(ProvisionerTaskState::Succeeded)) + .await? + } + } + } + + Ok(()) + } + + /// Waits until no remux job of the session is stored, since a remux replaces a video that would be read. + async fn wait_for_remux(&self, attempt: &Attempt, dir: &Utf8Path) -> Result<(), AttemptError> { + loop { + let jobs = self.job_queue.stored_jobs(RemuxJob::NAME).await.map_err(|error| { + warn!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to list the remux jobs"); + AttemptError::Recoverable(Recoverable::Storage) + })?; + + if !remux_pending(&jobs, dir) { + return Ok(()); + } + + debug!(task.id = %attempt.task_id, %dir, "Waiting for the remux of the session"); + + tokio::select! { + () = tokio::time::sleep(REMUX_POLL_INTERVAL) => {} + () = attempt.cancel.cancelled() => return Err(AttemptError::Permanent(Permanent::DeadlinePassed)), + } + } + } +} + +/// The payload of a running AI analysis Task: its progress, for whoever asks about it, and the tokens it used so far. +fn progress_payload(progress: &Progress, usage: TokenUsage) -> serde_json::Value { + let (done, total) = progress.done_and_total().unzip(); + + json!(Running { + step: progress.stage(), + done, + total, + usage, + }) +} + +/// Waits until no other AI analysis of the session runs. +/// +/// The working files belong to the session, and a Task that ended past its deadline may still run an attempt while +/// another Task of the session starts. +async fn lock_session(session_id: Uuid) -> tokio::sync::OwnedMutexGuard<()> { + static LOCKS: LazyLock>>>> = + LazyLock::new(Default::default); + + let lock = { + let mut locks = LOCKS.lock(); + locks.retain(|_, lock| lock.strong_count() > 0); + + match locks.get(&session_id).and_then(Weak::upgrade) { + Some(lock) => lock, + None => { + let lock = Arc::new(tokio::sync::Mutex::new(())); + locks.insert(session_id, Arc::downgrade(&lock)); + lock + } + } + }; + + lock.lock_owned().await +} + +/// What the AI crate needs from the session manifest; a session without recordings, or with both terminal and video ones, +/// cannot be analyzed. +fn analysis_manifest(manifest: &JrecManifest) -> Result { + let recordings = Recordings::from_files(manifest.files().iter().map(|file| RecordingFile { + file_name: file.file_name().to_owned(), + start_time: file.start_time(), + }))?; + + Ok(RecordingManifest { + start_time: manifest.start_time(), + duration: manifest.duration(), + recordings, + }) +} diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs new file mode 100644 index 000000000..42e7990ce --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -0,0 +1,1278 @@ +use std::sync::Arc; +use std::time::Duration; + +use axum::http::StatusCode; +use camino::Utf8Path; +use devolutions_gateway_ai::session_actions::PROMPT_VERSION; +use devolutions_gateway_ai::{Language, session_summary}; +use devolutions_gateway_task::ShutdownHandle; +use job_queue::Job as _; +use parking_lot::Mutex; +use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; + +use super::*; +use crate::job_queue::{JobMessage, JobQueueHandle, JobQueueReceiver, MAX_ATTEMPTS}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; +use crate::recording::RecordingManagerTask; +use crate::{DgwState, MockHandles}; + +const API_KEY: &str = "sk-recording-ai-analysis-test-secret"; + +fn config(recording_path: &Utf8Path) -> String { + json!({ + "ProvisionerPublicKeyData": { + "Value": "mMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4vuqLOkl1pWobt6su1XO9VskgCAwevEGs6kkNjJQBwkGnPKYLmNF1E/af1yCocfVn/OnPf9e4x+lXVyZ6LMDJxFxu+axdgOq3Ld392J1iAEbfvwlyRFnEXFOJNyylqg3bY6LvnWHL/XZczVdMD9xYfq2sO9bg3xjRW4s7r9EEYOFjqVT3VFznH9iWJVtcSEKukmS/3uKoO6lGhacvu0HhjXXdgq0R8zvR4XRJ9Fcnf0f9Ypoc+i6L80NVjrRCeVOH+Ld/2fA9bocpfLarcVqG3RjS+qgOtpyCc0jWVFF4zaGQ7LUDFkEIYILkICeMMn2ll29hmZNzsJzZJ9s6NocgQIDAQAB" + }, + "Listeners": [{ "InternalUrl": "http://*:7171", "ExternalUrl": "https://*:7171" }], + "Proxy": { "Mode": "Off" }, + "RecordingPath": recording_path, + }) + .to_string() +} + +fn openai(model: &str) -> AiSettings { + AiSettings { + provider: AiProvider::OpenAi, + model: model.to_owned(), + base_url: None, + max_output_tokens: None, + } +} + +struct Harness { + state: DgwState, + job_queue_rx: JobQueueReceiver, + recording_path: Utf8PathBuf, + _keep: Box, + _dir: tempfile::TempDir, +} + +impl Harness { + async fn new() -> Self { + let dir = tempfile::tempdir().expect("temp dir"); + let recording_path = Utf8PathBuf::from_path_buf(dir.path().join("recordings")).expect("UTF-8"); + + let (state, handles) = DgwState::mock(&config(&recording_path)).await.expect("mock state"); + let MockHandles { + session_manager_rx, + recording_manager_rx, + subscriber_rx, + job_queue_rx, + traffic_audit_rx, + shutdown_handle, + } = handles; + + let recording_manager = RecordingManagerTask::new( + recording_manager_rx, + recording_path.clone(), + state.sessions.clone(), + state.job_queue_handle.clone(), + ); + let (recording_shutdown_handle, recording_shutdown_signal) = ShutdownHandle::new(); + tokio::spawn(devolutions_gateway_task::Task::run( + recording_manager, + recording_shutdown_signal, + )); + + Self { + state, + job_queue_rx, + recording_path, + _keep: Box::new(( + session_manager_rx, + subscriber_rx, + traffic_audit_rx, + shutdown_handle, + recording_shutdown_handle, + )), + _dir: dir, + } + } + + async fn start(&self, task_id: Uuid, session_id: Uuid, settings: AiSettings) -> Result { + self.start_in(task_id, session_id, settings, None).await + } + + /// Same as [`Harness::start`], asking for the analysis in `language`. + async fn start_in( + &self, + task_id: Uuid, + session_id: Uuid, + settings: AiSettings, + language: Option, + ) -> Result { + start( + &self.state, + StartRequest { + task_id, + session_id, + settings, + language, + api_key: SecretString::from(API_KEY), + }, + ) + .await + } + + /// Starts an AI analysis of `session_id` and returns its Task ID. + async fn start_created(&self, session_id: Uuid, settings: AiSettings) -> Uuid { + let task_id = Uuid::new_v4(); + let outcome = self.start(task_id, session_id, settings).await.expect("started"); + assert!(matches!(outcome, StartOutcome::Created(_)), "{outcome:?}"); + task_id + } + + /// The job holding the current token of the Task, as the job queue would run it. + async fn job(&self, task_id: Uuid) -> ProvisionerTaskJob { + // Read as of long ago, so reading never fails the Task in place of the job. + let job_token = self + .state + .provisioner_tasks + .store() + .get(task_id, OffsetDateTime::UNIX_EPOCH) + .await + .expect("read task") + .expect("task exists") + .job_token; + let json = json!({ "taskId": task_id, "jobToken": job_token }).to_string(); + + ProvisionerTaskJob::read_json(&json, self.state.provisioner_tasks.clone()).expect("valid job") + } + + async fn run(&self, task_id: Uuid) -> anyhow::Result<()> { + self.job(task_id).await.run().await + } + + /// Records a Task directly, as if it was created at `created_at`. + async fn create_task( + &self, + task_id: Uuid, + session_id: Uuid, + provider: &str, + created_at: OffsetDateTime, + deadline_at: OffsetDateTime, + ) { + self.state + .provisioner_tasks + .store() + .create( + NewProvisionerTaskRecord { + id: task_id, + kind: KIND.to_owned(), + target: session_id.to_string(), + params: json!({ "provider": provider, "model": "gpt-test", "maxOutputTokens": null }), + deadline_at, + job_token: Uuid::new_v4(), + }, + created_at, + ) + .await + .expect("create"); + } + + async fn task(&self, task_id: Uuid) -> ProvisionerTaskRecord { + self.state + .provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .expect("read task") + .expect("task exists") + } + + /// Runs the Tasks with an AI analysis that fails after `progress_timeout` without progress. + fn with_progress_timeout(&mut self, progress_timeout: time::Duration) { + let mut task = RecordingAiAnalysis::new( + self.state.conf_handle.clone(), + self.state.recordings.clone(), + self.state.secrets.clone(), + Arc::new(self.state.job_queue_handle.clone()), + ); + task.progress_timeout = progress_timeout; + self.state.provisioner_tasks = ProvisionerTaskRunner::builder( + Arc::clone(self.state.provisioner_tasks.store()), + Arc::new(self.state.job_queue_handle.clone()), + ) + .register(task) + .build(); + } + + /// Working files of the AI analyses of session_id. + fn workspace(&self, session_id: Uuid) -> Utf8PathBuf { + self.recording_path.join(session_id.to_string()).join(".ai-analysis") + } + + fn access(&self, task_id: Uuid) -> Option { + self.state + .secrets + .get(SecretId(task_id)) + .expect("read the secret") + .map(|secret| AiAccess::from_secret(&secret).expect("AI access")) + } + + fn has_key(&self, task_id: Uuid) -> bool { + self.access(task_id).is_some() + } + + /// Hands `task_id` an API key directly, as `start` does. + fn insert_key(&self, task_id: Uuid, base_url: Option, expires_at: OffsetDateTime) { + let access = AiAccess { + api_key: SecretString::from(API_KEY), + base_url, + }; + self.state + .secrets + .insert(SecretId(task_id), &access.to_secret().expect("secret"), expires_at) + .expect("insert the secret"); + } + + /// Writes a finished session with a short terminal recording. + fn write_session(&self) -> Uuid { + self.write_long_session(10) + } + + /// Writes a finished session whose terminal output is `lines` lines of about 100 bytes each. + fn write_long_session(&self, lines: usize) -> Uuid { + let session_id = Uuid::new_v4(); + let dir = self.recording_path.join(session_id.to_string()); + std::fs::create_dir_all(&dir).expect("session dir"); + + let manifest = json!({ + "sessionId": session_id, + "startTime": 1_787_255_035, + "duration": lines, + "files": [{ "fileName": "recording-0.cast", "startTime": 1_787_255_035, "duration": lines }], + }); + std::fs::write(dir.join("recording.json"), manifest.to_string()).expect("manifest"); + + 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(dir.join("recording-0.cast"), cast).expect("cast"); + + session_id + } +} + +/// A mock OpenAI-compatible provider: `status` picks the answer of the n-th request, 0 being a valid answer. +async fn spawn_ai_provider( + status: impl Fn(usize) -> u16 + Send + Sync + 'static, +) -> (AiSettings, Arc>>) { + spawn_ai_provider_after(Duration::ZERO, status).await +} + +/// Like [`spawn_ai_provider`], answering each request after `delay`. +async fn spawn_ai_provider_after( + delay: Duration, + status: impl Fn(usize) -> u16 + Send + Sync + 'static, +) -> (AiSettings, Arc>>) { + let requests = Arc::new(Mutex::new(Vec::new())); + let status = Arc::new(status); + + let app = axum::Router::new().fallback(axum::routing::post({ + let requests = Arc::clone(&requests); + move |body: String| { + let requests = Arc::clone(&requests); + let status = Arc::clone(&status); + async move { + tokio::time::sleep(delay).await; + + let index = { + let mut requests = requests.lock(); + requests.push(body); + requests.len() - 1 + }; + + let answer = json!({ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-test-2026-09-30", + "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 } + }); + + match status(index) { + 0 => (StatusCode::OK, axum::Json(answer)), + code => ( + StatusCode::from_u16(code).expect("valid status"), + axum::Json(json!({ "error": { "message": "scripted failure" } })), + ), + } + } + } + })); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("address"); + tokio::spawn(async move { axum::serve(listener, app).await.expect("serve") }); + + let settings = AiSettings { + provider: AiProvider::OpenAiCompatible, + model: "gpt-test".to_owned(), + base_url: Some(format!("http://{addr}/v1/").parse().expect("URL")), + max_output_tokens: None, + }; + + (settings, requests) +} + +fn first_line(path: &Utf8Path) -> String { + std::fs::read_to_string(path) + .expect("chunk") + .lines() + .next() + .expect("line") + .to_owned() +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets +#[tokio::test] +async fn start_records_the_task_keeps_the_key_and_queues_a_job() { + let mut harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = harness.write_session(); + let mut settings = openai("gpt-test"); + settings.base_url = Some("https://ai.example/v1/".parse().expect("URL")); + + let StartOutcome::Created(task) = harness.start(task_id, session_id, settings).await.expect("started") else { + panic!("task must be created"); + }; + + assert_eq!(task.id, task_id); + assert_eq!(task.kind, KIND); + assert_eq!(task.target, session_id.to_string()); + assert_eq!(task.state, ProvisionerTaskState::Queued); + assert_eq!( + task.params, + json!({ "provider": "openai", "model": "gpt-test", "maxOutputTokens": null }), + "neither the key nor the base URL is recorded" + ); + assert_eq!(task.deadline_at - task.created_at, START_WINDOW); + + let access = harness.access(task_id).expect("key is kept"); + assert_eq!(access.base_url.expect("base URL").as_str(), "https://ai.example/v1/"); + + let JobMessage { job, schedule_for } = harness.job_queue_rx.try_recv().expect("job queued"); + assert_eq!(job.name(), ProvisionerTaskJob::NAME); + assert_eq!( + job.write_json().expect("job JSON"), + format!(r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, task.job_token) + ); + assert!(schedule_for.is_none()); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn asking_again_creates_nothing_new() { + let mut harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = harness.write_session(); + harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect("started"); + harness.job_queue_rx.try_recv().expect("first job"); + + let same = harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect("same request"); + assert!( + matches!(same, StartOutcome::Existing(ref task) if task.id == task_id), + "{same:?}" + ); + + let other_params = harness + .start(task_id, session_id, openai("gpt-other")) + .await + .expect("other parameters"); + assert!(matches!(other_params, StartOutcome::IdConflict(_)), "{other_params:?}"); + + let busy = harness + .start(Uuid::new_v4(), session_id, openai("gpt-test")) + .await + .expect("other task"); + assert_eq!( + busy, + StartOutcome::TargetBusy { + active_task_id: task_id + } + ); + + assert!(harness.job_queue_rx.try_recv().is_err(), "no other job is queued"); + assert!(harness.has_key(task_id), "the first key stays"); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn asking_again_returns_the_task_even_while_the_session_records_again() { + let mut harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = harness.write_session(); + harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect("started"); + harness.job_queue_rx.try_recv().expect("first job"); + harness.state.recordings.active_recordings.insert(session_id); + + let same = harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect("same request"); + assert!(matches!(same, StartOutcome::Existing(_)), "{same:?}"); + + let other_params = harness + .start(task_id, session_id, openai("gpt-other")) + .await + .expect("other parameters"); + assert!(matches!(other_params, StartOutcome::IdConflict(_)), "{other_params:?}"); + + let new_task = harness.start(Uuid::new_v4(), session_id, openai("gpt-test")).await; + assert!(matches!(new_task, Err(StartError::RecordingActive)), "{new_task:?}"); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn refuses_a_session_without_recording() { + let harness = Harness::new().await; + let task_id = Uuid::new_v4(); + + let error = harness + .start(task_id, Uuid::new_v4(), openai("gpt-test")) + .await + .expect_err("no recording"); + + assert!(matches!(error, StartError::RecordingNotFound), "{error:?}"); + assert!( + harness + .state + .provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .expect("read") + .is_none() + ); + assert!(!harness.has_key(task_id)); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn refuses_a_session_that_is_still_recording() { + let harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = Uuid::new_v4(); + harness.state.recordings.active_recordings.insert(session_id); + + let error = harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect_err("still recording"); + + assert!(matches!(error, StartError::RecordingActive), "{error:?}"); + assert!( + harness + .state + .provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .expect("read") + .is_none() + ); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn refuses_invalid_ai_settings() { + let harness = Harness::new().await; + let task_id = Uuid::new_v4(); + + let error = harness + .start(task_id, Uuid::new_v4(), openai(" ")) + .await + .expect_err("empty model"); + assert!( + matches!(error, StartError::InvalidSettings(AiSettingsError::MissingModel)), + "{error:?}" + ); + + let mut no_base_url = openai("gpt-test"); + no_base_url.provider = AiProvider::OpenAiCompatible; + let error = harness + .start(task_id, Uuid::new_v4(), no_base_url) + .await + .expect_err("no base URL"); + assert!( + matches!(error, StartError::InvalidSettings(AiSettingsError::MissingBaseUrl)), + "{error:?}" + ); + + assert!(!harness.has_key(task_id)); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { + let harness = Harness::new().await; + // About 1 MB of transcript: three chunks. + let session_id = harness.write_long_session(10_000); + let (settings, requests) = spawn_ai_provider(|index| if index == 1 { 503 } else { 0 }).await; + let task_id = harness.start_created(session_id, settings).await; + let workspace = harness.workspace(session_id); + + assert!(harness.run(task_id).await.is_err(), "a provider outage is retried"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Running); + assert_eq!( + task.payload, + json!({ + "step": "describing", + "done": 1, + "total": 3, + "usage": { "inputTokens": 10, "outputTokens": 20, "unreportedRequests": 0 }, + }) + ); + assert!(harness.has_key(task_id), "the key is kept for the retry"); + let plan: Vec = + serde_json::from_slice(&std::fs::read(workspace.join("plan.json")).expect("plan")).expect("plan JSON"); + assert_eq!(plan.len(), 3, "three transcript chunks"); + assert!(workspace.join("chunk-0000.json").exists()); + assert!(!workspace.join("chunk-0001.json").exists()); + let chunk_starts = (0..3) + .map(|index| first_line(&workspace.join(format!("transcript-{index:04}.txt")))) + .collect::>(); + assert_eq!(requests.lock().len(), 2); + + harness.run(task_id).await.expect("second attempt succeeds"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert_eq!(task.attempts, 2); + assert_eq!( + task.payload, + json!({ + "fileName": "ai-analysis-0.slog", + "model": "gpt-test-2026-09-30", + "usage": { "inputTokens": 40, "outputTokens": 80, "unreportedRequests": 0 }, + "promptVersion": format!("{PROMPT_VERSION}+{}", session_summary::PROMPT_VERSION), + }), + "every attempt counts, and the refused request costs nothing" + ); + + let requests = requests.lock().clone(); + assert_eq!( + requests.len(), + 5, + "only chunks 1 and 2 are sent again, then the summary" + ); + for (request, chunk) in requests.iter().zip([0, 1, 1, 2]) { + assert!(request.contains(&chunk_starts[chunk]), "request for chunk {chunk}"); + } + assert!(requests.iter().all(|request| !request.contains(API_KEY))); + + let log = std::fs::read_to_string( + harness + .recording_path + .join(session_id.to_string()) + .join("ai-analysis-0.slog"), + ) + .expect("log"); + let descriptions = log + .lines() + .map(|line| serde_json::from_str::(line).expect("JSON")["description"].clone()) + .collect::>(); + assert_eq!( + descriptions, + ["Session started", "Step 0", "Step 2", "Step 3", "Session ended"].map(serde_json::Value::from) + ); + + assert!(!workspace.exists(), "the workspace is deleted on success"); + assert!(!harness.has_key(task_id), "the key is dropped on success"); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn permanent_error_fails_without_retry() { + let harness = Harness::new().await; + let session_id = harness.write_long_session(10); + let (settings, requests) = spawn_ai_provider(|_| 401).await; + let task_id = harness.start_created(session_id, settings).await; + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!( + task.payload, + json!({ + "error": { "kind": "aiRejected", "status": 401 }, + "usage": { "inputTokens": 0, "outputTokens": 0, "unreportedRequests": 0 }, + }) + ); + assert_eq!(task.attempts, 1); + assert_eq!(requests.lock().len(), 1); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_session_mixing_terminal_and_video_fails() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let manifest_path = harness + .recording_path + .join(session_id.to_string()) + .join("recording.json"); + let mut manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(&manifest_path).expect("manifest")).expect("manifest JSON"); + manifest["files"] + .as_array_mut() + .expect("files") + .push(json!({ "fileName": "recording-1.webm", "startTime": 1_787_255_035, "duration": 1 })); + std::fs::write(&manifest_path, manifest.to_string()).expect("manifest"); + let task_id = harness.start_created(session_id, openai("gpt-test")).await; + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "mixedRecording" })); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn recording_deleted_after_start_fails() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = harness.start_created(session_id, openai("gpt-test")).await; + std::fs::remove_dir_all(harness.recording_path.join(session_id.to_string())).expect("delete the recording"); + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "noRecording" })); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { + let harness = Harness::new().await; + let session_id = harness.write_long_session(10); + let (settings, requests) = spawn_ai_provider(|_| 503).await; + let task_id = harness.start_created(session_id, settings).await; + + for attempt in 1..MAX_ATTEMPTS { + assert!(harness.run(task_id).await.is_err(), "attempt {attempt} is retried"); + assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Running); + } + + harness.run(task_id).await.expect("the last attempt is not retried"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.attempts, MAX_ATTEMPTS); + assert_eq!( + task.payload["error"], + json!({ "kind": "attemptsExhausted", "last": { "kind": "aiBusy", "status": 503 } }) + ); + assert_eq!(requests.lock().len(), 5); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn an_analysis_without_progress_fails_at_its_deadline() { + let mut harness = Harness::new().await; + let session_id = harness.write_long_session(10); + + // A provider that never answers stops the progress of the analysis. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("address"); + tokio::spawn(async move { + while let Ok((connection, _)) = listener.accept().await { + tokio::spawn(async move { + let _connection = connection; + std::future::pending::<()>().await; + }); + } + }); + let settings = AiSettings { + provider: AiProvider::OpenAiCompatible, + model: "gpt-test".to_owned(), + base_url: Some(format!("http://{addr}/v1/").parse().expect("URL")), + max_output_tokens: None, + }; + let task_id = harness.start_created(session_id, settings).await; + harness.with_progress_timeout(time::Duration::seconds(2)); + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis + assert_eq!( + task.payload, + json!({ + "error": { "kind": "deadlinePassed" }, + "usage": { "inputTokens": 0, "outputTokens": 0, "unreportedRequests": 1 }, + }), + "the request cut off at the deadline may have cost tokens" + ); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn lost_key_fails_the_task() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = harness.start_created(session_id, openai("gpt-test")).await; + std::fs::create_dir_all(harness.workspace(session_id)).expect("workspace"); + + // A restart keeps `gateway.db` and the job queue, but not the keys held in memory. + harness.state.secrets.remove(SecretId(task_id)); + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!( + task.payload, + json!({ + "error": { "kind": "keyLost" }, + "usage": { "inputTokens": 0, "outputTokens": 0, "unreportedRequests": 0 }, + }) + ); + assert!(!harness.workspace(session_id).exists()); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn overdue_task_does_not_run() { + let harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = Uuid::new_v4(); + let now = OffsetDateTime::now_utc(); + harness + .create_task( + task_id, + session_id, + "openai", + now - time::Duration::hours(2), + now - time::Duration::seconds(1), + ) + .await; + harness.insert_key(task_id, None, now + time::Duration::hours(1)); + std::fs::create_dir_all(harness.workspace(session_id)).expect("workspace"); + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!( + task.payload, + json!({ "error": { "kind": "deadlinePassed" } }), + "a Task that never ran spent nothing" + ); + assert_eq!(task.attempts, 0); + assert!(!harness.has_key(task_id)); + assert!(!harness.workspace(session_id).exists()); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn finished_task_does_not_run_again() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 401).await; + let task_id = harness.start_created(session_id, settings).await; + + harness.run(task_id).await.expect("first run"); + std::fs::create_dir_all(harness.workspace(session_id)).expect("workspace left behind"); + harness.run(task_id).await.expect("second run"); + + assert_eq!(harness.task(task_id).await.attempts, 1); + assert!( + !harness.workspace(session_id).exists(), + "a job that finds its Task finished deletes the workspace" + ); +} + +/// The job queue channel, plus remux jobs stored as the job queue database would list them. +struct WithRemuxJobs { + queue: JobQueueHandle, + remux_jobs: Arc>>, +} + +#[async_trait] +impl ProvisionerTaskJobQueue for WithRemuxJobs { + async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()> { + self.queue.push(job, schedule_for).await + } + + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + Ok(if name == RemuxJob::NAME { + self.remux_jobs.lock().clone() + } else { + Vec::new() + }) + } +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn analysis_waits_for_the_remux_of_the_session() { + let mut harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, requests) = spawn_ai_provider(|_| 0).await; + let task_id = harness.start_created(session_id, settings).await; + + let dir = harness.recording_path.join(session_id.to_string()); + let other_dir = harness.recording_path.join(Uuid::new_v4().to_string()); + let remux_jobs = Arc::new(Mutex::new(vec![ + json!({ "input_path": dir.join("recording-1.webm") }).to_string(), + json!({ "input_path": other_dir.join("recording-0.webm") }).to_string(), + ])); + let queue: DynProvisionerTaskJobQueue = Arc::new(WithRemuxJobs { + queue: harness.state.job_queue_handle.clone(), + remux_jobs: Arc::clone(&remux_jobs), + }); + harness.state.provisioner_tasks = + ProvisionerTaskRunner::builder(Arc::clone(harness.state.provisioner_tasks.store()), Arc::clone(&queue)) + .register(RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.state.secrets.clone(), + queue, + )) + .build(); + + let mut job = harness.job(task_id).await; + let job = tokio::spawn(async move { job.run().await }); + + tokio::time::sleep(REMUX_POLL_INTERVAL * 2).await; + assert!(!job.is_finished(), "the analysis waits"); + assert!(requests.lock().is_empty(), "the AI is not asked yet"); + + remux_jobs.lock().remove(0); + job.await.expect("job task").expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert!(!requests.lock().is_empty()); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn waiting_for_a_remux_is_no_progress() { + let mut harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, requests) = spawn_ai_provider(|_| 0).await; + let task_id = harness.start_created(session_id, settings).await; + + let dir = harness.recording_path.join(session_id.to_string()); + let queue: DynProvisionerTaskJobQueue = Arc::new(WithRemuxJobs { + queue: harness.state.job_queue_handle.clone(), + remux_jobs: Arc::new(Mutex::new(vec![ + json!({ "input_path": dir.join("recording-1.webm") }).to_string(), + ])), + }); + let mut task = RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.state.secrets.clone(), + Arc::clone(&queue), + ); + task.progress_timeout = time::Duration::seconds(2); + harness.state.provisioner_tasks = + ProvisionerTaskRunner::builder(Arc::clone(harness.state.provisioner_tasks.store()), queue) + .register(task) + .build(); + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "deadlinePassed" })); + assert!(requests.lock().is_empty(), "the AI is never asked"); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn analyses_of_one_session_run_one_at_a_time() { + let session_id = Uuid::new_v4(); + let first = lock_session(session_id).await; + + let other_session = tokio::time::timeout(Duration::from_secs(1), lock_session(Uuid::new_v4())).await; + assert!(other_session.is_ok(), "another session is not blocked"); + + let second = tokio::spawn(lock_session(session_id)); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(!second.is_finished(), "the second analysis waits"); + + drop(first); + tokio::time::timeout(Duration::from_secs(1), second) + .await + .expect("the second analysis goes on") + .expect("lock task"); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn working_files_left_by_an_earlier_task_go_with_the_task_that_ends() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 401).await; + let task_id = harness.start_created(session_id, settings).await; + std::fs::create_dir_all(harness.workspace(session_id)).expect("working files of a crashed analysis"); + let log = harness + .recording_path + .join(session_id.to_string()) + .join(".ai-analysis.slog"); + std::fs::write(&log, "never moved").expect("log of a crashed analysis"); + + harness.run(task_id).await.expect("no retry"); + + assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Failed); + assert!(!harness.workspace(session_id).exists()); + assert!(!log.exists()); +} +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_new_task_replaces_the_result_of_the_earlier_one() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 0).await; + let first = harness.start_created(session_id, settings.clone()).await; + harness.run(first).await.expect("first analysis"); + + let second = harness.start_created(session_id, settings).await; + harness.run(second).await.expect("second analysis"); + + let task = harness.task(second).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert_eq!(task.payload["fileName"], "ai-analysis-0.slog"); + let dir = harness.recording_path.join(session_id.to_string()); + let manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(dir.join("recording.json")).expect("manifest")).expect("manifest JSON"); + assert_eq!( + manifest["artifacts"]["ai-analysis"], + json!([{ "fileName": "ai-analysis-0.slog" }]), + "one analysis per recording" + ); + assert!(!dir.join("ai-analysis-1.slog").exists()); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_new_task_that_fails_keeps_the_earlier_analysis() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 0).await; + let first = harness.start_created(session_id, settings.clone()).await; + harness.run(first).await.expect("first analysis"); + + let second = harness.start_created(session_id, settings).await; + harness.state.secrets.remove(SecretId(second)); + harness.run(second).await.expect("no retry"); + + let task = harness.task(second).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "keyLost" })); + let dir = harness.recording_path.join(session_id.to_string()); + assert_eq!(ai_analyses(&dir), json!([{ "fileName": "ai-analysis-0.slog" }])); + assert!(dir.join("ai-analysis-0.slog").exists()); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_failed_task_reports_the_tokens_it_spent() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + // The chunk is answered, then the provider refuses the summary request for good. + let (settings, _requests) = spawn_ai_provider(|index| if index == 1 { 401 } else { 0 }).await; + let task_id = harness.start_created(session_id, settings).await; + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!( + task.payload["usage"], + json!({ "inputTokens": 10, "outputTokens": 20, "unreportedRequests": 0 }) + ); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_task_that_fails_removes_its_unlisted_log() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = harness.start_created(session_id, openai("gpt-test")).await; + let dir = harness.recording_path.join(session_id.to_string()); + let log = dir.join(format!(".ai-analysis-{task_id}.slog")); + std::fs::write(&log, "an analysis whose success was never recorded").expect("unlisted log"); + harness + .state + .provisioner_tasks + .store() + .fail(task_id, json!({}), OffsetDateTime::now_utc()) + .await + .expect("fail the Task"); + + harness.run(task_id).await.expect("the job ends the Task"); + + assert!(!log.exists()); + assert_eq!(ai_analyses(&dir), json!([])); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_success_cut_short_before_the_log_is_listed_is_finished_by_the_job() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 0).await; + let task_id = harness.start_created(session_id, settings).await; + harness.run(task_id).await.expect("analysis"); + + // As if Gateway stopped right after recording the success, before the manifest listed the log. + let dir = harness.recording_path.join(session_id.to_string()); + let log = dir.join(format!(".ai-analysis-{task_id}.slog")); + std::fs::rename(dir.join("ai-analysis-0.slog"), &log).expect("unlist the log"); + let manifest_path = dir.join("recording.json"); + let mut manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(&manifest_path).expect("manifest")).expect("manifest JSON"); + manifest["artifacts"] = json!({}); + std::fs::write(&manifest_path, manifest.to_string()).expect("manifest"); + + // The job of the Task stays stored until it returns, so it runs again when Gateway starts. + harness.run(task_id).await.expect("the job ends the Task again"); + + assert_eq!(ai_analyses(&dir), json!([{ "fileName": "ai-analysis-0.slog" }])); + assert!(dir.join("ai-analysis-0.slog").exists()); + assert!(!log.exists()); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_late_end_of_an_earlier_task_never_lists_its_log_over_a_newer_analysis() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 0).await; + let first = harness.start_created(session_id, settings.clone()).await; + harness.run(first).await.expect("first analysis"); + + // As if the end of the first Task failed before the manifest listed its log, so it is done again later. + let dir = harness.recording_path.join(session_id.to_string()); + let first_log = dir.join(format!(".ai-analysis-{first}.slog")); + std::fs::remove_file(dir.join("ai-analysis-0.slog")).expect("unlist the log"); + std::fs::write(&first_log, "the analysis of the first Task").expect("unlisted log"); + let manifest_path = dir.join("recording.json"); + let mut manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(&manifest_path).expect("manifest")).expect("manifest JSON"); + manifest["artifacts"] = json!({}); + std::fs::write(&manifest_path, manifest.to_string()).expect("manifest"); + + let second = harness.start_created(session_id, settings).await; + harness.run(second).await.expect("second analysis"); + let newer = std::fs::read_to_string(dir.join("ai-analysis-0.slog")).expect("newer analysis"); + + harness.run(first).await.expect("the late end of the first Task"); + + assert_eq!(ai_analyses(&dir), json!([{ "fileName": "ai-analysis-0.slog" }])); + assert_eq!( + std::fs::read_to_string(dir.join("ai-analysis-0.slog")).expect("analysis"), + newer + ); + assert!(!first_log.exists()); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_new_task_that_fails_keeps_an_earlier_success_whose_log_was_not_listed_yet() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 0).await; + let first = harness.start_created(session_id, settings.clone()).await; + harness.run(first).await.expect("first analysis"); + + // As if the end of the first Task failed before the manifest listed its log, so it is done again later. + let dir = harness.recording_path.join(session_id.to_string()); + std::fs::remove_file(dir.join("ai-analysis-0.slog")).expect("unlist the log"); + std::fs::write( + dir.join(format!(".ai-analysis-{first}.slog")), + "the analysis of the first Task", + ) + .expect("unlisted log"); + let manifest_path = dir.join("recording.json"); + let mut manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(&manifest_path).expect("manifest")).expect("manifest JSON"); + manifest["artifacts"] = json!({}); + std::fs::write(&manifest_path, manifest.to_string()).expect("manifest"); + + let second = harness.start_created(session_id, settings).await; + harness.state.secrets.remove(SecretId(second)); + harness.run(second).await.expect("no retry"); + assert_eq!(harness.task(second).await.state, ProvisionerTaskState::Failed); + + harness.run(first).await.expect("the late end of the first Task"); + + assert_eq!(ai_analyses(&dir), json!([{ "fileName": "ai-analysis-0.slog" }])); + assert_eq!( + std::fs::read_to_string(dir.join("ai-analysis-0.slog")).expect("analysis"), + "the analysis of the first Task" + ); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_new_task_removes_a_log_left_unlisted_by_an_unknown_task() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let dir = harness.recording_path.join(session_id.to_string()); + let unknown_log = dir.join(format!(".ai-analysis-{}.slog", Uuid::new_v4())); + std::fs::write(&unknown_log, "an analysis whose success was never recorded").expect("unlisted log"); + let (settings, _requests) = spawn_ai_provider(|_| 401).await; + let task_id = harness.start_created(session_id, settings).await; + + harness.run(task_id).await.expect("no retry"); + + assert!(!unknown_log.exists()); + assert_eq!(ai_analyses(&dir), json!([])); +} + +/// The AI analyses the manifest in `dir` lists. +fn ai_analyses(dir: &Utf8Path) -> serde_json::Value { + let manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(dir.join("recording.json")).expect("manifest")).expect("manifest JSON"); + manifest["artifacts"] + .get("ai-analysis") + .cloned() + .unwrap_or_else(|| json!([])) +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn the_analysis_is_in_the_language_of_the_request() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, requests) = spawn_ai_provider(|_| 0).await; + let task_id = Uuid::new_v4(); + let language = Language::new("fr-CA").expect("language"); + + let outcome = harness + .start_in(task_id, session_id, settings.clone(), Some(language.clone())) + .await + .expect("started"); + assert!(matches!(outcome, StartOutcome::Created(_)), "{outcome:?}"); + harness.run(task_id).await.expect("analysis"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert_eq!(task.params["language"], "fr-CA"); + assert!(requests.lock().iter().all(|request| request.contains("`fr-CA`"))); + + let other = harness + .start_in( + task_id, + session_id, + settings, + Some(Language::new("de-DE").expect("language")), + ) + .await + .expect("started"); + assert!( + matches!(other, StartOutcome::IdConflict(_)), + "the same Task ID with another language is another request: {other:?}" + ); +} + +// oracle: intent crates/devolutions-gateway-ai/INTENT.md §Analyze Recording +#[tokio::test] +async fn the_analysis_is_in_english_without_a_language() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, requests) = spawn_ai_provider(|_| 0).await; + let task_id = harness.start_created(session_id, settings).await; + + harness.run(task_id).await.expect("analysis"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert!(task.params.get("language").is_none(), "{}", task.params); + assert!(requests.lock().iter().all(|request| request.contains("`English`"))); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn a_new_task_never_resumes_from_an_earlier_task() { + let harness = Harness::new().await; + // About 1 MB of transcript: three chunks. + let session_id = harness.write_long_session(10_000); + let (settings, requests) = spawn_ai_provider(|index| if index == 1 { 503 } else { 0 }).await; + let first = harness.start_created(session_id, settings.clone()).await; + assert!(harness.run(first).await.is_err(), "the outage is retried"); + assert!(harness.workspace(session_id).join("chunk-0000.json").exists()); + harness + .state + .provisioner_tasks + .store() + .fail(first, json!({}), OffsetDateTime::now_utc()) + .await + .expect("end the first Task"); + + let second = harness.start_created(session_id, settings).await; + harness.run(second).await.expect("second analysis"); + + assert_eq!(harness.task(second).await.state, ProvisionerTaskState::Succeeded); + assert_eq!( + requests.lock().len(), + 2 + 3 + 1, + "the new Task asks about every chunk again, then for the summary" + ); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn an_ended_task_leaves_the_working_files_of_another_task() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 0).await; + let task_id = harness.start_created(session_id, settings).await; + let workspace = harness.workspace(session_id); + std::fs::create_dir_all(&workspace).expect("working files of a newer Task"); + std::fs::write( + workspace.join("analysis.json"), + json!({ + "owner": Uuid::new_v4().to_string(), + "promptVersion": "p", + "model": "m", + "language": "English", + "chunkLen": 1, + }) + .to_string(), + ) + .expect("settings of a newer Task"); + let record = harness.task(task_id).await; + let task = RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.state.secrets.clone(), + Arc::new(harness.state.job_queue_handle.clone()), + ); + + task.on_end(&record).await.expect("end"); + + assert!(workspace.exists(), "the newer Task still resumes from them"); + assert!(!harness.has_key(task_id)); +} diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs new file mode 100644 index 000000000..a8278960e --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -0,0 +1,719 @@ +//! Runs provisioner tasks: the one place that keeps a task's record and the job running it in step. +//! +//! Each kind of work is a [`ProvisionerTask`]: what one attempt does, what its payloads hold, and what to do once it has +//! ended. +//! [`ProvisionerTaskRunner`] does the rest for every task. It creates the record, queues the job, starts each attempt, +//! retries recoverable failures, moves the deadline ahead with each progress, cancels an attempt once the deadline passes, +//! and records how the task ended. +//! +//! The record is the source of truth and the job only wakes it up. Each record names the one job allowed to run it, by +//! its token, so a stale or duplicate job stops without running anything. At startup, before any job runs, +//! [`ProvisionerTaskRunner::resume`] hands every unfinished record to a new job, which also covers a record created just +//! before a crash, before its job was queued. +//! A record whose end fails is handed to a new job that ends it again later, since the runs of its job are shared with +//! its attempts and may be used up; that job is stored, so a restart runs it too. + +#[cfg(test)] +mod tests; + +use std::collections::HashMap; +use std::sync::Arc; + +use anyhow::Context as _; +use async_trait::async_trait; +use parking_lot::Mutex; +use provisioner_task::{ + AttemptStart, CreateOutcome, DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskRecord, + ProvisionerTaskState, +}; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::json; +use time::OffsetDateTime; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use super::{AttemptError, Permanent}; +use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; +use crate::token::AccessScope; + +/// Time before a record whose end failed is ended again. +const END_RETRY_DELAY: time::Duration = time::Duration::minutes(5); + +/// Where the runner queues the jobs of its records. +#[async_trait] +pub trait ProvisionerTaskJobQueue: Send + Sync { + /// Queues `job` to run at `schedule_for`, or at once without it, and returns once it is stored, so a record never + /// waits for a job that was lost. + async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()>; + + /// Lists the JSON definitions of the stored jobs named `name` that may still run. + async fn stored_jobs(&self, name: &str) -> anyhow::Result>; +} + +pub type DynProvisionerTaskJobQueue = Arc; + +/// Hands the job to the job queue task, which stores it later: for tests, which take the jobs from its channel. +#[async_trait] +impl ProvisionerTaskJobQueue for JobQueueHandle { + async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()> { + match schedule_for { + Some(schedule_for) => self.schedule(job, schedule_for).await, + None => self.enqueue(job).await, + } + } + + /// The channel cannot tell what is stored: every unfinished record looks like it has no job. + async fn stored_jobs(&self, _name: &str) -> anyhow::Result> { + Ok(Vec::new()) + } +} + +/// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. +/// +/// One value serves every record of its kind, so it keeps only what they share; what belongs to one record comes in +/// through [`Attempt`] and the parameters. +#[async_trait] +pub trait ProvisionerTask: Send + Sync + 'static { + /// The kind written in the records of this task. + const KIND: &'static str; + + /// The scope a token needs to start this task, `gateway.tasks..start`. + const START_SCOPE: AccessScope; + + /// The scope a token needs to read the records of this task, `gateway.tasks..read`. + const READ_SCOPE: AccessScope; + + /// What the record keeps as asked; never secrets, since records are kept and returned to callers. + type Params: Serialize + DeserializeOwned + Send; + + /// The payload of a `succeeded` record. + type Output: Serialize + Send; + + /// Time from the creation of a record until its first attempt must start. + fn start_window(&self) -> time::Duration; + + /// Time a running record may go without progress: each attempt start and each progress report moves its deadline + /// this far ahead. + fn progress_timeout(&self) -> time::Duration; + + /// The payload of the running record as an attempt starts, from `last`, its payload until then: `null` before its + /// first attempt. + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value; + + /// The payload of a record that failed with `error`, from `last`, its payload until then: `null` if it never ran. + fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value; + + /// Runs one attempt. + /// + /// The attempt should stop soon after [`Attempt::cancel`] is cancelled, which happens once the record is past its + /// deadline. + async fn run(&self, attempt: &Attempt, params: Self::Params) -> Result; + + /// Finishes what the record leaves to do once it has ended, however it ended, such as dropping what was kept only + /// while it ran. + /// + /// It may be called again for a record that already ended, so it must be idempotent; an error has it called again + /// later, until it succeeds. + async fn on_end(&self, _record: &ProvisionerTaskRecord) -> anyhow::Result<()> { + Ok(()) + } +} + +/// One attempt of a record, as its task sees it. +pub struct Attempt { + pub task_id: Uuid, + pub target: String, + /// Attempts started so far, this one included. + pub attempts: u32, + /// Payload of the record as this attempt started, from [`ProvisionerTask::starting_payload`]. + pub payload: serde_json::Value, + /// Cancelled once the record is past its deadline or no longer running. + pub cancel: CancellationToken, + progress_timeout: time::Duration, + store: DynProvisionerTaskStore, + latest: Arc>, +} + +/// What the runner knows of the running record: the last payload of the attempt, and the deadline it stored. +struct Latest { + payload: serde_json::Value, + deadline_at: OffsetDateTime, +} + +impl Attempt { + /// Deadline of the record, after which the attempt is cancelled. + pub fn deadline(&self) -> OffsetDateTime { + self.latest.lock().deadline_at + } + + /// State of another record of the same target, such as an earlier Task; `None` for an unknown one or one of another + /// target. One past its deadline is failed first. + pub async fn state_of(&self, task_id: Uuid) -> anyhow::Result> { + let record = self.store.get(task_id, OffsetDateTime::now_utc()).await?; + Ok(record + .filter(|record| record.target == self.target) + .map(|record| record.state)) + } + + /// Replaces the payload of the running record and moves its deadline ahead, so whoever asks about it sees how the + /// attempt goes. + /// + /// A record that is no longer running, such as one past its deadline, cancels the attempt. + pub async fn report_progress(&self, payload: serde_json::Value) { + let now = OffsetDateTime::now_utc(); + let deadline_at = deadline(now, self.progress_timeout); + self.latest.lock().payload = payload.clone(); + + match self.store.update_running(self.task_id, payload, deadline_at, now).await { + Ok(true) => self.latest.lock().deadline_at = deadline_at, + Ok(false) => { + debug!(task.id = %self.task_id, "Task is no longer running, so its attempt stops"); + self.cancel.cancel(); + } + Err(error) => { + warn!(task.id = %self.task_id, error = format!("{error:#}"), "Failed to record the Task progress"); + } + } + } +} + +/// [`ProvisionerTask`] with its types erased, so the runner can hold every task together. +#[async_trait] +trait AnyProvisionerTask: Send + Sync { + fn start_window(&self) -> time::Duration; + + fn progress_timeout(&self) -> time::Duration; + + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value; + + fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value; + + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result; + + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()>; +} + +#[async_trait] +impl AnyProvisionerTask for T { + fn start_window(&self) -> time::Duration { + ProvisionerTask::start_window(self) + } + + fn progress_timeout(&self) -> time::Duration { + ProvisionerTask::progress_timeout(self) + } + + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value { + ProvisionerTask::starting_payload(self, last) + } + + fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value { + ProvisionerTask::failed_payload(self, error, last) + } + + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { + let params = serde_json::from_value(params).map_err(|error| { + warn!(task.id = %attempt.task_id, %error, "Invalid Task parameters"); + AttemptError::Permanent(Permanent::InvalidTask) + })?; + + let output = ProvisionerTask::run(self, attempt, params).await?; + + serde_json::to_value(output).map_err(|error| { + error!(task.id = %attempt.task_id, %error, "Failed to serialize the Task result"); + AttemptError::Permanent(Permanent::Internal) + }) + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { + ProvisionerTask::on_end(self, record).await + } +} + +/// Creates provisioner task records and runs them as jobs on the job queue. +#[derive(Clone)] +pub struct ProvisionerTaskRunner { + inner: Arc, +} + +struct Inner { + store: DynProvisionerTaskStore, + job_queue: DynProvisionerTaskJobQueue, + /// Kind → the task that runs records of this kind. + tasks: HashMap<&'static str, Arc>, +} + +pub struct ProvisionerTaskRunnerBuilder { + store: DynProvisionerTaskStore, + job_queue: DynProvisionerTaskJobQueue, + tasks: HashMap<&'static str, Arc>, +} + +impl ProvisionerTaskRunnerBuilder { + /// Registers `task`, so records of its kind can be created and run. + #[must_use] + pub fn register(mut self, task: T) -> Self { + self.tasks.insert(T::KIND, Arc::new(task)); + self + } + + pub fn build(self) -> ProvisionerTaskRunner { + ProvisionerTaskRunner { + inner: Arc::new(Inner { + store: self.store, + job_queue: self.job_queue, + tasks: self.tasks, + }), + } + } +} + +impl ProvisionerTaskRunner { + pub fn builder( + store: DynProvisionerTaskStore, + job_queue: DynProvisionerTaskJobQueue, + ) -> ProvisionerTaskRunnerBuilder { + ProvisionerTaskRunnerBuilder { + store, + job_queue, + tasks: HashMap::new(), + } + } + + /// Returns the record with this ID; one past its deadline is failed first. + pub async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result> { + self.inner.store.get(id, now).await + } + + /// Creates a record for the task `T` and queues the job that runs it. + /// + /// `on_created` runs only for a new record, before its job is queued, so what the job needs is in place before it + /// can run; if it fails, the record fails and no job is queued. A known ID or a busy target creates nothing and + /// queues nothing. + pub async fn create( + &self, + id: Uuid, + target: String, + params: &T::Params, + on_created: impl FnOnce(&ProvisionerTaskRecord) -> anyhow::Result<()> + Send, + ) -> anyhow::Result { + let task = self.task(T::KIND).context("unregistered provisioner task")?; + let params = serde_json::to_value(params).context("serialize the task parameters")?; + let now = OffsetDateTime::now_utc(); + let job_token = Uuid::new_v4(); + + let outcome = self + .inner + .store + .create( + NewProvisionerTaskRecord { + id, + kind: T::KIND.to_owned(), + target, + params, + deadline_at: deadline(now, task.start_window()), + job_token, + }, + now, + ) + .await?; + + let CreateOutcome::Created(record) = &outcome else { + return Ok(outcome); + }; + + let queued = match on_created(record) { + Ok(()) => self.enqueue(id, job_token, None).await, + Err(error) => Err(error.context("prepare the Task")), + }; + + if let Err(error) = queued { + error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to queue the Task"); + self.fail_not_queued(task.as_ref(), id, job_token, &serde_json::Value::Null) + .await; + return Err(error); + } + + info!(task.id = %id, task.kind = T::KIND, "Task queued"); + + Ok(outcome) + } + + /// Hands every unfinished record that has no job left to a new job, such as a record created just before a crash. + /// + /// A record whose job is still stored keeps it. Call it at startup, before the job runner claims any job, so no job + /// runs an attempt meanwhile. A record whose new job cannot be queued is failed, since no job is left to run it. + pub async fn resume(&self) -> anyhow::Result<()> { + let now = OffsetDateTime::now_utc(); + let records = self.inner.store.list_unfinished(now).await?; + + let stored = self + .inner + .job_queue + .stored_jobs(ProvisionerTaskJob::NAME) + .await + .context("list the stored provisioner task jobs")? + .iter() + .filter_map(|json| serde_json::from_str::(json).ok()) + .map(|def| (def.task_id, def.job_token)) + .collect::>(); + + for record in records { + if stored.contains(&(record.id, record.job_token)) { + debug!(task.id = %record.id, task.kind = %record.kind, "Task still has its job"); + continue; + } + + let job_token = Uuid::new_v4(); + + match self.inner.store.replace_job_token(record.id, job_token, now).await { + Ok(true) => {} + Ok(false) => continue, + Err(error) => { + error!( + task.id = %record.id, + task.kind = %record.kind, + error = format!("{error:#}"), + "Failed to hand the Task to a new job" + ); + continue; + } + } + + match self.enqueue(record.id, job_token, None).await { + Ok(()) => info!(task.id = %record.id, task.kind = %record.kind, "Task queued again"), + Err(error) => { + error!( + task.id = %record.id, + task.kind = %record.kind, + error = format!("{error:#}"), + "Failed to queue the Task again" + ); + if let Some(task) = self.task(&record.kind) { + self.fail_not_queued(task.as_ref(), record.id, job_token, &record.payload) + .await; + } + } + } + } + + Ok(()) + } + + #[cfg(test)] + pub(crate) fn store(&self) -> &DynProvisionerTaskStore { + &self.inner.store + } + + /// The task that runs records of `kind`. + fn task(&self, kind: &str) -> Option> { + self.inner.tasks.get(kind).cloned() + } + + async fn enqueue( + &self, + task_id: Uuid, + job_token: Uuid, + schedule_for: Option, + ) -> anyhow::Result<()> { + self.inner + .job_queue + .push( + ProvisionerTaskJob { + task_id, + job_token, + runner: self.clone(), + }, + schedule_for, + ) + .await + } + + /// Runs one attempt for the job holding `job_token`; returns an error only to ask the job queue for a retry. + async fn run_attempt(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { + let store = &self.inner.store; + let now = OffsetDateTime::now_utc(); + + // The starting payload carries over what earlier attempts reported, so the record is read first. + let Some(record) = store.get(task_id, now).await? else { + debug!(task.id = %task_id, "Task is unknown, so this attempt does not run"); + return Ok(()); + }; + + let Some(task) = self.task(&record.kind) else { + if record.job_token == job_token && !record.state.is_finished() { + error!(task.id = %task_id, task.kind = %record.kind, "Task of an unknown kind"); + // Without its kind, nothing else of the payload is known, so it holds the error alone. + store + .fail(task_id, json!({ "error": Permanent::InvalidTask }), now) + .await + .context("record the Task failure")?; + } + return Ok(()); + }; + + let starting_payload = task.starting_payload(&record.payload); + let deadline_at = deadline(now, task.progress_timeout()); + + let record = match store + .start_attempt(task_id, job_token, starting_payload.clone(), deadline_at, now) + .await? + { + AttemptStart::Started(record) => record, + AttemptStart::Finished(record) => { + debug!(task.id = %task_id, "Task is over, so this attempt does not run"); + return self.end_record(task.as_ref(), &record, job_token).await; + } + AttemptStart::Superseded => { + debug!(task.id = %task_id, "Another job runs this Task, so this one stops"); + return Ok(()); + } + AttemptStart::Unknown => { + debug!(task.id = %task_id, "Task is unknown, so this attempt does not run"); + return Ok(()); + } + }; + + let attempts = record.attempts; + let cancel = CancellationToken::new(); + let latest = Arc::new(Mutex::new(Latest { + payload: starting_payload.clone(), + deadline_at, + })); + let attempt = Attempt { + task_id, + target: record.target.clone(), + attempts, + payload: starting_payload, + cancel: cancel.clone(), + progress_timeout: task.progress_timeout(), + store: Arc::clone(store), + latest: Arc::clone(&latest), + }; + + info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task attempt started"); + + // The attempt has its own Tokio task, so a panic fails the record instead of leaving it running. + let mut handle = tokio::spawn({ + let task = Arc::clone(&task); + let params = record.params.clone(); + async move { task.run(&attempt, params).await } + }); + + let mut past_deadline = false; + let joined = loop { + let deadline_at = latest.lock().deadline_at; + let wait = std::time::Duration::try_from(deadline_at - OffsetDateTime::now_utc()).unwrap_or_default(); + + tokio::select! { + joined = &mut handle => break joined, + () = tokio::time::sleep(wait) => { + // Progress may have moved the deadline while the timer ran. + if OffsetDateTime::now_utc() < latest.lock().deadline_at { + continue; + } + + warn!(task.id = %task_id, task.kind = %record.kind, attempts, "Task went past its deadline; its attempt is cancelled"); + past_deadline = true; + // Aborting cannot stop blocking steps, so the attempt is cancelled and awaited instead: what runs + // once the Task ends never runs alongside it. + cancel.cancel(); + break handle.await; + } + } + }; + + let outcome = match joined { + // The attempt stopped because it was cancelled, whatever error it reports for it. + Ok(Err(_)) if past_deadline => Err(AttemptError::Permanent(Permanent::DeadlinePassed)), + Ok(outcome) => outcome, + Err(_) => Err(AttemptError::Permanent(Permanent::Panicked)), + }; + + let last = latest.lock().payload.clone(); + + match outcome { + Ok(output) => match store.succeed(task_id, output, OffsetDateTime::now_utc()).await { + Ok(ended) => { + if ended { + info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"); + } else { + warn!( + task.id = %task_id, + task.kind = %record.kind, + "Task attempt succeeded after the Task had ended, such as past its deadline; the Task keeps its earlier end" + ); + } + + self.end(task.as_ref(), task_id, job_token).await?; + } + // UNTESTED: a success that cannot be recorded fails the Task, which drops what the attempt made, + // because a retry would do the work again, such as paying for a second analysis; no error may leave + // this branch before the failure is recorded. + Err(error) => { + error!( + task.id = %task_id, + error = format!("{error:#}"), + "Failed to record the Task success; the Task fails instead" + ); + + let payload = task.failed_payload(&Permanent::Internal, &last); + + match store.fail(task_id, payload, OffsetDateTime::now_utc()).await { + // The job retries only when the end could not be queued again, which then only ends the Task. + Ok(_) => self.end(task.as_ref(), task_id, job_token).await?, + Err(error) => { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to fail the Task; it fails at its deadline"); + + // The store fails the Task once its deadline passes, so a job that runs after it only ends it. + let after_deadline = latest.lock().deadline_at + time::Duration::SECOND; + + if let Err(error) = self.enqueue(task_id, job_token, Some(after_deadline)).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to queue the end of the Task"); + } + } + } + } + }, + Err(AttemptError::Recoverable(error)) if attempts < MAX_ATTEMPTS => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, "Task attempt failed; it will be retried"); + anyhow::bail!("Task attempt failed: {error:?}"); + } + Err(AttemptError::Recoverable(error)) => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, "Task failed after its last attempt"); + self.fail( + task.as_ref(), + task_id, + job_token, + &Permanent::AttemptsExhausted { last: error }, + &last, + ) + .await?; + } + Err(AttemptError::Permanent(error)) => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, "Task failed"); + self.fail(task.as_ref(), task_id, job_token, &error, &last).await?; + } + } + + Ok(()) + } + + /// Records the failure, then ends the record. + /// + /// If the failure cannot be recorded, nothing is dropped, since the record may still run, and the error is returned so + /// the job queue tries again. + async fn fail( + &self, + task: &dyn AnyProvisionerTask, + task_id: Uuid, + job_token: Uuid, + error: &Permanent, + last: &serde_json::Value, + ) -> anyhow::Result<()> { + self.inner + .store + .fail(task_id, task.failed_payload(error, last), OffsetDateTime::now_utc()) + .await + .context("record the Task failure")?; + + self.end(task, task_id, job_token).await + } + + /// Fails a record no job will run. + async fn fail_not_queued( + &self, + task: &dyn AnyProvisionerTask, + task_id: Uuid, + job_token: Uuid, + last: &serde_json::Value, + ) { + if let Err(error) = self.fail(task, task_id, job_token, &Permanent::NotQueued, last).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to end the Task that has no job"); + } + } + + /// Ends the record, read again so its task sees how it ended. + async fn end(&self, task: &dyn AnyProvisionerTask, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { + match self.inner.store.get(task_id, OffsetDateTime::now_utc()).await { + Ok(Some(record)) => self.end_record(task, &record, job_token).await, + Ok(None) => Ok(()), + Err(error) => self.end_later(task_id, job_token, error).await, + } + } + + async fn end_record( + &self, + task: &dyn AnyProvisionerTask, + record: &ProvisionerTaskRecord, + job_token: Uuid, + ) -> anyhow::Result<()> { + match task.on_end(record).await { + Ok(()) => Ok(()), + Err(error) => self.end_later(record.id, job_token, error).await, + } + } + + /// Hands the record to a new job that ends it again later; an error means that job could not be queued. + async fn end_later(&self, task_id: Uuid, job_token: Uuid, error: anyhow::Error) -> anyhow::Result<()> { + warn!(task.id = %task_id, error = format!("{error:#}"), "Failed to end the Task; it is ended again later"); + + self.enqueue(task_id, job_token, Some(OffsetDateTime::now_utc() + END_RETRY_DELAY)) + .await + .context("queue the end of the Task again") + } +} + +/// The deadline `timeout` after `now`, in whole seconds as the store keeps it, so the runner and the store agree on when +/// it passes. +fn deadline(now: OffsetDateTime, timeout: time::Duration) -> OffsetDateTime { + (now + timeout).truncate_to_second() +} + +#[derive(Serialize, serde::Deserialize)] +#[serde(rename_all = "camelCase")] +struct JobDef { + task_id: Uuid, + job_token: Uuid, +} + +/// The job of a provisioner task record: it only wakes the record up, and the record decides what runs. +pub struct ProvisionerTaskJob { + task_id: Uuid, + job_token: Uuid, + runner: ProvisionerTaskRunner, +} + +impl ProvisionerTaskJob { + pub const NAME: &'static str = "provisioner-task"; + + pub fn read_json(json: &str, runner: ProvisionerTaskRunner) -> anyhow::Result { + let def: JobDef = serde_json::from_str(json).context("failed to deserialize ProvisionerTaskJob")?; + + Ok(Self { + task_id: def.task_id, + job_token: def.job_token, + runner, + }) + } +} + +#[async_trait] +impl job_queue::Job for ProvisionerTaskJob { + fn name(&self) -> &str { + Self::NAME + } + + fn write_json(&self) -> anyhow::Result { + serde_json::to_string(&JobDef { + task_id: self.task_id, + job_token: self.job_token, + }) + .context("failed to serialize ProvisionerTaskJob") + } + + async fn run(&mut self) -> anyhow::Result<()> { + self.runner.run_attempt(self.task_id, self.job_token).await + } +} diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs new file mode 100644 index 000000000..b3338e923 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -0,0 +1,515 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use parking_lot::Mutex; +use provisioner_task::ProvisionerTaskState; + +use super::*; +use crate::job_queue::{JobMessage, JobQueueReceiver}; +use crate::provisioner_tasks::Recoverable; +use crate::token::AccessScope; + +const KIND: &str = "test.echo"; + +/// A value whose attempt waits until it is cancelled. +const HANG: u32 = u32::MAX; + +/// A value whose attempt reports progress five times, half a second apart, then succeeds. +const PROGRESS: u32 = u32::MAX - 1; + +#[derive(Debug, Serialize, serde::Deserialize)] +struct EchoParams { + value: u32, +} + +/// A task whose attempts return `outcome`, counting its runs and the records it ended. +struct Echo { + outcome: fn(u32) -> Result, + progress_timeout: time::Duration, + runs: Arc, + ended: Arc>>, + /// Ends that fail before one succeeds. + end_failures: Arc, +} + +#[async_trait] +impl ProvisionerTask for Echo { + const KIND: &'static str = KIND; + + const START_SCOPE: AccessScope = AccessScope::Wildcard; + + const READ_SCOPE: AccessScope = AccessScope::Wildcard; + + type Params = EchoParams; + + type Output = serde_json::Value; + + fn start_window(&self) -> time::Duration { + time::Duration::hours(1) + } + + fn progress_timeout(&self) -> time::Duration { + self.progress_timeout + } + + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value { + json!({ "after": last }) + } + + fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value { + json!({ "error": error, "last": last }) + } + + async fn run(&self, attempt: &Attempt, params: EchoParams) -> Result { + self.runs.fetch_add(1, Ordering::SeqCst); + + match params.value { + HANG => { + attempt.cancel.cancelled().await; + Err(AttemptError::Recoverable(Recoverable::Storage)) + } + PROGRESS => { + for step in 0..5 { + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + attempt.report_progress(json!({ "step": step })).await; + } + Ok(json!("done")) + } + value => (self.outcome)(value), + } + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { + if self + .end_failures + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |left| left.checked_sub(1)) + .is_ok() + { + anyhow::bail!("the end failed"); + } + + self.ended.lock().push(record.id); + Ok(()) + } +} + +struct Harness { + runner: ProvisionerTaskRunner, + job_queue_rx: JobQueueReceiver, + runs: Arc, + ended: Arc>>, + end_failures: Arc, +} + +impl Harness { + async fn new(outcome: fn(u32) -> Result) -> Self { + Self::with_progress_timeout(outcome, time::Duration::hours(1)).await + } + + async fn with_progress_timeout( + outcome: fn(u32) -> Result, + progress_timeout: time::Duration, + ) -> Self { + let conn = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open gateway database") + .connect() + .await + .expect("connect"); + let store = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) + .await + .expect("open task store"), + ); + let (job_queue, job_queue_rx) = JobQueueHandle::new(); + let runs = Arc::new(AtomicUsize::new(0)); + let ended = Arc::new(Mutex::new(Vec::new())); + let end_failures = Arc::new(AtomicUsize::new(0)); + + let runner = ProvisionerTaskRunner::builder(store, Arc::new(job_queue)) + .register(Echo { + outcome, + progress_timeout, + runs: Arc::clone(&runs), + ended: Arc::clone(&ended), + end_failures: Arc::clone(&end_failures), + }) + .build(); + + Self { + runner, + job_queue_rx, + runs, + ended, + end_failures, + } + } + + async fn create(&self, value: u32) -> Uuid { + let id = Uuid::new_v4(); + let outcome = self + .runner + .create::(id, format!("target-{id}"), &EchoParams { value }, |_| Ok(())) + .await + .expect("create"); + assert!(matches!(outcome, CreateOutcome::Created(_)), "{outcome:?}"); + id + } + + fn queued_job(&mut self) -> Box { + let JobMessage { job, .. } = self.job_queue_rx.try_recv().expect("a queued job"); + job + } + + async fn task(&self, id: Uuid) -> ProvisionerTaskRecord { + self.runner + .get(id, OffsetDateTime::now_utc()) + .await + .expect("read task") + .expect("task exists") + } +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn only_a_new_task_runs_on_created_and_gets_a_job() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = Uuid::new_v4(); + let created = AtomicUsize::new(0); + let create = || { + harness + .runner + .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + created.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + }; + + assert!(matches!(create().await.expect("create"), CreateOutcome::Created(_))); + assert!(matches!(create().await.expect("again"), CreateOutcome::Existing(_))); + + assert_eq!(created.load(Ordering::SeqCst), 1); + harness.queued_job(); + assert!(harness.job_queue_rx.try_recv().is_err(), "one job only"); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn a_task_whose_job_cannot_be_queued_fails() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + harness.job_queue_rx.close(); + let id = Uuid::new_v4(); + + harness + .runner + .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| Ok(())) + .await + .expect_err("the job queue is gone"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload, json!({ "error": { "kind": "notQueued" }, "last": null })); + assert_eq!(*harness.ended.lock(), [id], "what was set up for the Task is dropped"); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn a_task_that_cannot_be_prepared_fails_without_a_job() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = Uuid::new_v4(); + + harness + .runner + .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + anyhow::bail!("no room for the secret") + }) + .await + .expect_err("preparing fails"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "notQueued" })); + assert!(harness.job_queue_rx.try_recv().is_err(), "no job is queued"); + assert_eq!(*harness.ended.lock(), [id]); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn a_superseded_job_runs_nothing_once_the_task_is_resumed() { + let mut harness = Harness::new(|value| Ok(json!({ "echo": value }))).await; + let id = harness.create(7).await; + let mut stale = harness.queued_job(); + + harness.runner.resume().await.expect("resume"); + let mut resumed = harness.queued_job(); + assert!( + harness.job_queue_rx.try_recv().is_err(), + "the unfinished Task gets one new job" + ); + + stale.run().await.expect("the stale job stops"); + assert_eq!(harness.runs.load(Ordering::SeqCst), 0); + assert_eq!(harness.task(id).await.attempts, 0, "a stale job counts no attempt"); + assert!(harness.ended.lock().is_empty(), "a stale job ends nothing"); + + resumed.run().await.expect("the new job runs"); + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded); + assert_eq!(task.payload, json!({ "echo": 7 })); + assert_eq!(*harness.ended.lock(), [id]); + + harness.runner.resume().await.expect("resume again"); + assert!( + harness.job_queue_rx.try_recv().is_err(), + "a finished Task gets no new job" + ); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn a_task_that_cannot_be_queued_again_fails() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + harness.queued_job(); + harness.job_queue_rx.close(); + + harness.runner.resume().await.expect("resume"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed, "no job is left to run it"); + assert_eq!(task.payload["error"], json!({ "kind": "notQueued" })); + assert_eq!(*harness.ended.lock(), [id]); +} + +/// A job queue that remembers the jobs it stores, as the job queue database does. +#[derive(Default)] +struct StoredJobs(Mutex>); + +#[async_trait] +impl ProvisionerTaskJobQueue for StoredJobs { + async fn push(&self, job: ProvisionerTaskJob, _schedule_for: Option) -> anyhow::Result<()> { + self.0.lock().push(job_queue::Job::write_json(&job)?); + Ok(()) + } + + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + assert_eq!(name, ProvisionerTaskJob::NAME); + Ok(self.0.lock().clone()) + } +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn resume_queues_only_the_tasks_left_without_a_job() { + let conn = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open gateway database") + .connect() + .await + .expect("connect"); + let store = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) + .await + .expect("open task store"), + ); + let queue = Arc::new(StoredJobs::default()); + let runner = ProvisionerTaskRunner::builder(store, Arc::clone(&queue) as DynProvisionerTaskJobQueue) + .register(Echo { + outcome: |value| Ok(json!(value)), + progress_timeout: time::Duration::hours(1), + runs: Arc::default(), + ended: Arc::default(), + end_failures: Arc::default(), + }) + .build(); + let create = |id: Uuid| runner.create::(id, format!("target-{id}"), &EchoParams { value: 1 }, |_| Ok(())); + let with_job = Uuid::new_v4(); + let without_job = Uuid::new_v4(); + create(with_job).await.expect("create"); + create(without_job).await.expect("create"); + let token_with_job = runner + .get(with_job, OffsetDateTime::now_utc()) + .await + .expect("read") + .expect("record") + .job_token; + + // The job of the second record is lost, as when Gateway stops right after creating the record. + queue.0.lock().retain(|json| json.contains(&with_job.to_string())); + + runner.resume().await.expect("resume"); + + let stored = queue.0.lock().clone(); + assert_eq!(stored.len(), 2, "one new job, for the record without one: {stored:?}"); + assert!(stored[1].contains(&without_job.to_string())); + assert_eq!( + runner + .get(with_job, OffsetDateTime::now_utc()) + .await + .expect("read") + .expect("record") + .job_token, + token_with_job, + "a record with its job keeps it" + ); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn a_panicking_attempt_fails_the_task() { + let mut harness = Harness::new(|_| panic!("the attempt panics")).await; + let id = harness.create(1).await; + + harness.queued_job().run().await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!( + task.payload, + json!({ "error": { "kind": "panicked" }, "last": { "after": null } }), + "the failure keeps what the attempt started with" + ); + assert_eq!(*harness.ended.lock(), [id]); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn a_permanent_failure_fails_the_task_at_once() { + let mut harness = Harness::new(|_| Err(AttemptError::Permanent(Permanent::KeyLost))).await; + let id = harness.create(1).await; + + harness.queued_job().run().await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.attempts, 1); + assert_eq!(task.payload["error"], json!({ "kind": "keyLost" })); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { + let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(Recoverable::AiBusy { status: 429 }))).await; + let id = harness.create(1).await; + let mut job = harness.queued_job(); + + for attempt in 1..MAX_ATTEMPTS { + job.run().await.expect_err("the job queue retries"); + assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Running, "{attempt}"); + } + job.run().await.expect("the last attempt"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.attempts, MAX_ATTEMPTS); + assert_eq!( + task.payload["error"], + json!({ "kind": "attemptsExhausted", "last": { "kind": "aiBusy", "status": 429 } }) + ); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn invalid_parameters_and_unknown_kinds_fail_the_task() { + let harness = Harness::new(|value| Ok(json!(value))).await; + let store = harness.runner.store(); + let now = OffsetDateTime::now_utc(); + + for (kind, params) in [(KIND, json!({ "value": "not a number" })), ("test.unknown", json!({}))] { + let id = Uuid::new_v4(); + let job_token = Uuid::new_v4(); + store + .create( + NewProvisionerTaskRecord { + id, + kind: kind.to_owned(), + target: id.to_string(), + params, + deadline_at: now + time::Duration::hours(1), + job_token, + }, + now, + ) + .await + .expect("create"); + + harness.runner.run_attempt(id, job_token).await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed, "{kind}"); + assert_eq!(task.payload["error"], json!({ "kind": "invalidTask" }), "{kind}"); + } + + assert_eq!(harness.runs.load(Ordering::SeqCst), 0); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn an_end_that_fails_is_done_again_by_a_new_job_until_it_succeeds() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + let mut job = harness.queued_job(); + harness.end_failures.store(2, Ordering::SeqCst); + + for _ in 0..2 { + job.run().await.expect("the end is left to a new job"); + + let JobMessage { + job: next, + schedule_for, + } = harness.job_queue_rx.try_recv().expect("a new job"); + assert!( + schedule_for.is_some_and(|at| OffsetDateTime::now_utc() < at), + "it runs later" + ); + assert!(harness.ended.lock().is_empty()); + job = next; + } + + job.run().await.expect("the end succeeds"); + + assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Succeeded); + assert_eq!(*harness.ended.lock(), [id]); + assert!(harness.job_queue_rx.try_recv().is_err(), "no job is left"); +} + +// oracle: intent crates/provisioner-task/INTENT.md §Invariants +#[tokio::test] +async fn an_attempt_without_progress_is_cancelled_at_the_deadline() { + let mut harness = Harness::with_progress_timeout(|value| Ok(json!(value)), time::Duration::seconds(2)).await; + let id = harness.create(HANG).await; + + harness.queued_job().run().await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "deadlinePassed" })); + assert_eq!(*harness.ended.lock(), [id]); +} + +// oracle: intent crates/provisioner-task/INTENT.md §Invariants +#[tokio::test] +async fn progress_keeps_an_attempt_running_past_its_first_deadline() { + let mut harness = Harness::with_progress_timeout(|value| Ok(json!(value)), time::Duration::seconds(3)).await; + let id = harness.create(PROGRESS).await; + + harness.queued_job().run().await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert_eq!(task.payload, json!("done")); +} + +// oracle: property round-trip +#[tokio::test] +async fn the_job_is_read_back_from_its_json() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + let job = harness.queued_job(); + + let json = job.write_json().expect("job JSON"); + let read = ProvisionerTaskJob::read_json(&json, harness.runner.clone()).expect("valid job"); + + assert_eq!(read.task_id, id); + assert_eq!(read.job_token, harness.task(id).await.job_token); + assert_eq!(job.name(), ProvisionerTaskJob::NAME); +} diff --git a/devolutions-gateway/src/recording.intent.md b/devolutions-gateway/src/recording.intent.md index 14bb0371e..e86a51ee9 100644 --- a/devolutions-gateway/src/recording.intent.md +++ b/devolutions-gateway/src/recording.intent.md @@ -7,4 +7,5 @@ When connection is established with request of recordings for the session, if re # Artifacts Artifacts are files that are not recordings, currently only have `ai-analysis` with combination to `slog` file type. -Currently, we do not have an api endpoint to push artifacts, but we will have one in the future. \ No newline at end of file +Currently, we do not have an api endpoint to push artifacts, but we will have one in the future. +Gateway replaces the manifest in one step and never rewrites it in place, so neither a reader nor a crash ever leaves half of it. diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 8ece9b9b6..ef0fd0f67 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -8,7 +8,7 @@ use std::time::Duration; use anyhow::Context as _; use async_trait::async_trait; -use camino::Utf8PathBuf; +use camino::{Utf8Path, Utf8PathBuf}; use devolutions_gateway_task::{ShutdownSignal, Task}; use futures::future::Either; use parking_lot::Mutex; @@ -30,15 +30,26 @@ const BUFFER_WRITER_SIZE: usize = 64 * 1024; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -struct JrecFile { +pub(crate) struct JrecFile { file_name: String, start_time: i64, duration: i64, } +impl JrecFile { + pub(crate) fn file_name(&self) -> &str { + &self.file_name + } + + /// Unix seconds. + pub(crate) fn start_time(&self) -> i64 { + self.start_time + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -struct JrecManifest { +pub(crate) struct JrecManifest { session_id: Uuid, start_time: i64, duration: i64, @@ -48,15 +59,48 @@ struct JrecManifest { } impl JrecManifest { + /// Unix seconds. + pub(crate) fn start_time(&self) -> i64 { + self.start_time + } + + /// Seconds. + pub(crate) fn duration(&self) -> i64 { + self.duration + } + + pub(crate) fn files(&self) -> &[JrecFile] { + &self.files + } + fn read_from_file(path: impl AsRef) -> anyhow::Result { let json = std::fs::read(path)?; let manifest = serde_json::from_slice(&json)?; Ok(manifest) } + /// Writes the manifest to a file of its own first, then renames it over `path`, so that neither a reader nor a crash + /// ever finds half of it. fn save_to_file(&self, path: impl AsRef) -> anyhow::Result<()> { + use std::io::Write as _; + + let path = path.as_ref(); + // jrec never serves a name that starts with a dot, so a file a crash left behind is never downloaded. + let mut partial_name = std::ffi::OsString::from("."); + partial_name.push(path.file_name().context("manifest path without a file name")?); + partial_name.push(".partial"); + let partial = path.with_file_name(partial_name); let json = serde_json::to_string_pretty(&self)?; - std::fs::write(path, json)?; + + let mut file = std::fs::File::create(&partial)?; + file.write_all(json.as_bytes())?; + // UNTESTED: a crash leaves the old manifest or the new one, never half of it, only while the file is synced + // before the rename. + file.sync_all()?; + drop(file); + + std::fs::rename(&partial, path)?; + Ok(()) } } @@ -316,7 +360,7 @@ impl ActiveRecordings { self.0.lock().clone() } - fn insert(&self, id: Uuid) -> usize { + pub(crate) fn insert(&self, id: Uuid) -> usize { let mut guard = self.0.lock(); guard.insert(id); guard.len() @@ -357,6 +401,16 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, + ReplaceArtifacts { + id: Uuid, + kind: ArtifactKind, + file: Utf8PathBuf, + channel: oneshot::Sender>, + }, + GetFinished { + id: Uuid, + channel: oneshot::Sender>, + }, Disconnect { id: Uuid, }, @@ -377,6 +431,29 @@ enum RecordingManagerMessage { }, } +/// A session that is not recording anymore and has at least one Recording. +#[derive(Debug)] +pub(crate) struct FinishedRecording { + pub(crate) dir: Utf8PathBuf, + pub(crate) manifest: JrecManifest, +} + +impl FinishedRecording { + pub(crate) fn manifest_path(&self) -> Utf8PathBuf { + self.dir.join("recording.json") + } +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum FinishedRecordingError { + #[error("session has no recording")] + NotFound, + #[error("session is still recording")] + Recording, + #[error(transparent)] + Other(anyhow::Error), +} + impl fmt::Debug for RecordingManagerMessage { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { @@ -396,6 +473,20 @@ impl fmt::Debug for RecordingManagerMessage { .field("id", id) .field("kind", kind) .finish_non_exhaustive(), + RecordingManagerMessage::ReplaceArtifacts { + id, + kind, + file, + channel: _, + } => f + .debug_struct("ReplaceArtifacts") + .field("id", id) + .field("kind", kind) + .field("file", file) + .finish_non_exhaustive(), + RecordingManagerMessage::GetFinished { id, channel: _ } => { + f.debug_struct("GetFinished").field("id", id).finish_non_exhaustive() + } RecordingManagerMessage::Disconnect { id } => f.debug_struct("Disconnect").field("id", id).finish(), RecordingManagerMessage::GetState { id, channel: _ } => { f.debug_struct("GetState").field("id", id).finish_non_exhaustive() @@ -459,6 +550,58 @@ impl RecordingMessageSender { rx.await.context("couldn't receive AddArtifact result")? } + /// Moves `file` into the session as its only artifact of `kind`, named `kind.file_name(0)`, in place of the earlier + /// ones. + /// + /// `file` must be in the session folder, so that a reader finds the earlier artifact or the new one, never none or a + /// partial one. + /// Calling it again after it was cut short finishes the work, a `file` already gone lists nothing new, and a session + /// whose recording was deleted meanwhile only loses `file`. + pub(crate) async fn replace_artifacts(&self, id: Uuid, kind: ArtifactKind, file: &Utf8Path) -> anyhow::Result<()> { + let (tx, rx) = oneshot::channel(); + self.channel + .send(RecordingManagerMessage::ReplaceArtifacts { + id, + kind, + file: file.to_owned(), + channel: tx, + }) + .await + .ok() + .context("couldn't send ReplaceArtifacts message")?; + rx.await.context("couldn't receive ReplaceArtifacts result")? + } + + pub(crate) async fn get_finished(&self, id: Uuid) -> Result { + let (tx, rx) = oneshot::channel(); + self.channel + .send(RecordingManagerMessage::GetFinished { id, channel: tx }) + .await + .ok() + .context("couldn't send GetFinished message") + .map_err(FinishedRecordingError::Other)?; + let dir = rx + .await + .context("couldn't receive GetFinished result") + .map_err(FinishedRecordingError::Other)??; + + // The manifest is read here, so the recording manager task never waits on the disk for it. + let manifest = match fs::read(dir.join("recording.json")).await { + Ok(manifest) => manifest, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Err(FinishedRecordingError::NotFound), + Err(error) => { + return Err(FinishedRecordingError::Other( + anyhow::Error::new(error).context("read manifest from disk"), + )); + } + }; + let manifest = serde_json::from_slice(&manifest) + .context("parse manifest") + .map_err(FinishedRecordingError::Other)?; + + Ok(FinishedRecording { dir, manifest }) + } + async fn disconnect(&self, id: Uuid) -> anyhow::Result<()> { self.channel .send(RecordingManagerMessage::Disconnect { id }) @@ -819,6 +962,71 @@ impl RecordingManagerTask { Ok(artifact_path) } + async fn handle_replace_artifacts(&self, id: Uuid, kind: ArtifactKind, file: &Utf8Path) -> anyhow::Result<()> { + let recording_path = self.recordings_path.join(id.to_string()); + let manifest_path = recording_path.join("recording.json"); + + // A recording deleted while its artifact was being generated must stay deleted. + if !manifest_path.exists() { + match fs::remove_file(file).await { + Err(error) if error.kind() != io::ErrorKind::NotFound => { + return Err(anyhow::Error::new(error).context(format!("remove {file}"))); + } + _ => return Ok(()), + } + } + + let mut manifest = JrecManifest::read_from_file(&manifest_path).context("read manifest from disk")?; + + let file_name = kind.file_name(0); + let artifact_path = recording_path.join(&file_name); + + // The rename swaps the file in at once and the manifest follows, so it never lists a missing or partial file. + match fs::rename(file, &artifact_path).await { + Ok(()) => {} + // Already moved, by an earlier call cut short before the manifest was written, or by a newer Task that listed + // it for the Task that made it; what is in place is listed. + Err(error) if error.kind() == io::ErrorKind::NotFound && artifact_path.exists() => {} + Err(error) if error.kind() == io::ErrorKind::NotFound => { + warn!(%file, "The artifact to list is gone"); + return Ok(()); + } + Err(error) => return Err(anyhow::Error::new(error).context(format!("move {file} to {artifact_path}"))), + } + + let earlier = core::mem::replace( + manifest.artifacts.of_kind_mut(kind), + vec![JrecArtifact { + file_name: file_name.clone(), + }], + ); + + manifest + .save_to_file(&manifest_path) + .context("write manifest to disk")?; + + for artifact in earlier.into_iter().filter(|artifact| artifact.file_name != file_name) { + let artifact_path = recording_path.join(&artifact.file_name); + match fs::remove_file(&artifact_path).await { + Err(error) if error.kind() != io::ErrorKind::NotFound => { + warn!(%error, %artifact_path, "Failed to remove an earlier artifact"); + } + _ => {} + } + } + + Ok(()) + } + + fn handle_get_finished(&self, id: Uuid) -> Result { + // The transcript is built from the whole recording, so one still being pushed is not ready. + if self.ongoing_recordings.contains_key(&id) { + return Err(FinishedRecordingError::Recording); + } + + Ok(self.recordings_path.join(id.to_string())) + } + fn handle_remove(&mut self, id: Uuid) { if let Some(ongoing) = self.ongoing_recordings.get(&id) { match ongoing.state { @@ -937,6 +1145,12 @@ async fn recording_manager_task( RecordingManagerMessage::AddArtifact { id, kind, channel } => { let _ = channel.send(manager.handle_add_artifact(id, kind).await); }, + RecordingManagerMessage::ReplaceArtifacts { id, kind, file, channel } => { + let _ = channel.send(manager.handle_replace_artifacts(id, kind, &file).await); + }, + RecordingManagerMessage::GetFinished { id, channel } => { + let _ = channel.send(manager.handle_get_finished(id)); + } RecordingManagerMessage::Disconnect { id } => { if let Err(e) = manager.handle_disconnect(id).await { error!(error = format!("{e:#}"), "handle_disconnect"); @@ -1030,6 +1244,14 @@ impl RemuxJob { pub const NAME: &'static str = "remux"; } +/// Whether one of the remux job `definitions` is for a recording file in `dir`. +pub(crate) fn remux_pending(definitions: &[String], dir: &Utf8Path) -> bool { + definitions + .iter() + .filter_map(|json| serde_json::from_str::(json).ok()) + .any(|job| job.input_path.parent() == Some(dir)) +} + #[async_trait] impl job_queue::Job for RemuxJob { fn name(&self) -> &str { @@ -1421,6 +1643,26 @@ mod tests { assert_eq!(harness.sender.get_count().await.expect("count"), 1); } + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + #[tokio::test] + async fn finished_recordings_are_read_through_the_manager() { + let harness = Harness::start(); + let id = Uuid::new_v4(); + harness.write_manifest(id, MASTER_MANIFEST); + + let finished = harness.sender.get_finished(id).await.expect("finished"); + assert_eq!(finished.dir, harness.recordings_path.join(id.to_string())); + assert_eq!(finished.manifest.files()[0].file_name(), "recording-0.webm"); + + let ongoing = Uuid::new_v4(); + harness.connect(ongoing, WEBM).await; + let result = harness.sender.get_finished(ongoing).await; + assert!(matches!(result, Err(FinishedRecordingError::Recording)), "{result:?}"); + + let missing = harness.sender.get_finished(Uuid::new_v4()).await; + assert!(matches!(missing, Err(FinishedRecordingError::NotFound)), "{missing:?}"); + } + #[test] fn manifest_without_artifacts_round_trips_byte_for_byte() { let dir = tempfile::tempdir().expect("temp dir"); diff --git a/devolutions-gateway/src/secret_store.rs b/devolutions-gateway/src/secret_store.rs index 650316046..ea8d9179f 100644 --- a/devolutions-gateway/src/secret_store.rs +++ b/devolutions-gateway/src/secret_store.rs @@ -72,6 +72,22 @@ impl SecretStore { .map(Some) } + /// Moves the expiry of the secret of `id` to `expires_at`, unless it is already removed or expired. + pub fn set_expiry(&self, id: SecretId, expires_at: time::OffsetDateTime) { + self.set_expiry_at(id, expires_at, time::OffsetDateTime::now_utc()); + } + + fn set_expiry_at(&self, id: SecretId, expires_at: time::OffsetDateTime, now: time::OffsetDateTime) { + if let Some(entry) = self.entries.lock().get_mut(&id) + && now < entry.expires_at + { + entry.expires_at = expires_at; + // UNTESTED: the cleanup task wakes up for the new expiry, so a secret moved to an earlier expiry is still + // removed on time; any change of an expiry must notify it. + self.cleanup_notify.notify_one(); + } + } + /// Removes the secret of `id`. pub fn remove(&self, id: SecretId) { self.entries.lock().remove(&id); @@ -195,4 +211,32 @@ mod tests { store.remove_expired(now + time::Duration::minutes(1)); assert!(store.entries.lock().is_empty(), "the cleanup pass removes it"); } + + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets + #[test] + fn a_moved_expiry_keeps_the_secret_until_then() { + let store = SecretStore::default(); + let now = time::OffsetDateTime::now_utc(); + let moved = SecretId(Uuid::new_v4()); + let expired = SecretId(Uuid::new_v4()); + + store + .insert(moved, &SecretString::from("a"), now + time::Duration::minutes(1)) + .expect("insert"); + store.insert(expired, &SecretString::from("b"), now).expect("insert"); + + store.set_expiry_at(moved, now + time::Duration::hours(1), now); + store.set_expiry_at(expired, now + time::Duration::hours(1), now); + + assert!( + store + .get_at(moved, now + time::Duration::minutes(30)) + .expect("read") + .is_some() + ); + assert!( + store.get_at(expired, now).expect("read").is_none(), + "an expired secret stays gone" + ); + } } diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index 258e15d23..0c6666733 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -340,6 +340,35 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { None }; + let provisioner_tasks_conn = gateway_db + .connect() + .await + .context("failed to connect the provisioner tasks to the gateway database")?; + let task_store: provisioner_task::DynProvisionerTaskStore = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(provisioner_tasks_conn) + .await + .context("failed to initialize the provisioner task store")?, + ); + + let secrets = devolutions_gateway::secret_store::SecretStore::default(); + + let provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( + task_store, + Arc::new(job_queue_ctx.durable()), + conf_handle.clone(), + recording_manager_handle.clone(), + secrets.clone(), + ); + + // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. + // Unfinished Tasks still fail at their deadline, so this must not keep the other features from starting. + if let Err(error) = provisioner_tasks.resume().await { + error!( + error = format!("{error:#}"), + "Failed to queue the unfinished provisioner tasks again" + ); + } + let state = DgwState { conf_handle: conf_handle.clone(), token_cache: Arc::clone(&token_cache), @@ -354,6 +383,8 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { monitoring_state, traffic_audit_handle: traffic_audit_task.handle(), agent_tunnel_handle, + provisioner_tasks: provisioner_tasks.clone(), + secrets: secrets.clone(), }; for listener in &conf.listeners { @@ -387,6 +418,8 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { tasks.register(devolutions_gateway::provisioning::CleanupTask { handle: provisioning }); + tasks.register(devolutions_gateway::secret_store::CleanupTask { handle: secrets }); + tasks.register(devolutions_gateway::credential_injection::CleanupTask { handle: synthetic_kdc_registry, }); @@ -418,7 +451,10 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { job_queue_ctx.job_queue_handle.clone(), )); - tasks.register(devolutions_gateway::job_queue::JobRunnerTask::new(&job_queue_ctx)); + tasks.register(devolutions_gateway::job_queue::JobRunnerTask::new( + &job_queue_ctx, + provisioner_tasks, + )); tasks.register(devolutions_gateway::job_queue::JobQueueTask::new(job_queue_ctx)); tasks.register(traffic_audit_task); diff --git a/devolutions-gateway/src/token.rs b/devolutions-gateway/src/token.rs index c922869a5..d645c2e09 100644 --- a/devolutions-gateway/src/token.rs +++ b/devolutions-gateway/src/token.rs @@ -513,6 +513,10 @@ pub enum AccessScope { AgentRead, #[serde(rename = "gateway.recordings.search")] RecordingsSearch, + #[serde(rename = "gateway.tasks.recording-ai-analysis.start")] + RecordingAiAnalysisStart, + #[serde(rename = "gateway.tasks.recording-ai-analysis.read")] + RecordingAiAnalysisRead, } #[derive(Clone, Serialize, Deserialize)] diff --git a/devolutions-gateway/tests/preflight.rs b/devolutions-gateway/tests/preflight.rs index 1cbcf8480..71a312c62 100644 --- a/devolutions-gateway/tests/preflight.rs +++ b/devolutions-gateway/tests/preflight.rs @@ -58,15 +58,15 @@ fn unsigned_jws(payload: serde_json::Value) -> anyhow::Result { Ok(format!("{header}.{payload}.{signature}")) } -fn make_router() -> anyhow::Result<(Router, DgwState, MockHandles)> { - let (state, handles) = DgwState::mock(CONFIG)?; +async fn make_router() -> anyhow::Result<(Router, DgwState, MockHandles)> { + let (state, handles) = DgwState::mock(CONFIG).await?; let app = devolutions_gateway::make_http_service(state.clone()) .layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); Ok((app, state, handles)) } -fn make_router_with_config(config: &str) -> anyhow::Result<(Router, DgwState, MockHandles)> { - let (state, handles) = DgwState::mock(config)?; +async fn make_router_with_config(config: &str) -> anyhow::Result<(Router, DgwState, MockHandles)> { + let (state, handles) = DgwState::mock(config).await?; let app = devolutions_gateway::make_http_service(state.clone()) .layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); Ok((app, state, handles)) @@ -83,7 +83,7 @@ fn init_logger() -> tracing::subscriber::DefaultGuard { async fn test_provision_credentials_success() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let jti = Uuid::new_v4(); let token = unsigned_jws(json!({ @@ -124,7 +124,7 @@ async fn test_provision_credentials_success_when_unstable_disabled() -> anyhow:: // `provision-credentials` is protocol-neutral: NTLM credential injection relies on this // path even when the unstable feature flag is off. let config = CONFIG.replace("\"enable_unstable\": true", "\"enable_unstable\": false"); - let (app, _state, _handles) = make_router_with_config(&config)?; + let (app, _state, _handles) = make_router_with_config(&config).await?; let jti = Uuid::new_v4(); let token = unsigned_jws(json!({ @@ -161,7 +161,7 @@ async fn test_provision_credentials_success_when_unstable_disabled() -> anyhow:: async fn test_provision_credentials_rejects_alternate_targets() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let token = unsigned_jws(json!({ "jti": "5e3e833f-84c7-4541-b676-acc3299e39b8", @@ -205,7 +205,7 @@ async fn test_provision_credentials_rejects_alternate_targets() -> anyhow::Resul async fn test_provision_credentials_rejects_missing_target_hostname() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let token = unsigned_jws(json!({ "jti": "5e3e833f-84c7-4541-b676-acc3299e39b8" @@ -247,7 +247,7 @@ async fn test_provision_credentials_rejects_missing_target_hostname() -> anyhow: async fn test_provision_token_overwrite_alert() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; // Same JTI twice: second provision-token replaces the stored token-only entry (master behavior). let token = unsigned_jws(json!({ @@ -288,7 +288,7 @@ async fn test_provision_token_overwrite_alert() -> anyhow::Result<()> { async fn test_provision_connection_options_success() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let jti = Uuid::new_v4(); let token = unsigned_jws(json!({ @@ -319,7 +319,7 @@ async fn test_provision_connection_options_success() -> anyhow::Result<()> { async fn test_provision_credentials_and_connection_options_fold() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let jti = Uuid::new_v4(); let token = unsigned_jws(json!({ @@ -361,7 +361,7 @@ async fn test_provision_credentials_and_connection_options_fold() -> anyhow::Res async fn test_provision_invalid_params() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJqdGkiOiI1ZTNlODMzZi04NGM3LTQ1NDEtYjY3Ni1hY2MzMjk5ZTM5YjgifQ.1qECGlrW7y9HWFArc6GPHLGTOY7PhAvzKJ5XMRBg4k4"; diff --git a/devolutions-gateway/tests/traffic.rs b/devolutions-gateway/tests/traffic.rs index 7175d3bb5..3c9e4a63c 100644 --- a/devolutions-gateway/tests/traffic.rs +++ b/devolutions-gateway/tests/traffic.rs @@ -71,7 +71,7 @@ impl Drop for HandlesGuard { /// Build a Router with a live TrafficAuditManagerTask and real handle, /// similar to tests/preflight.rs harness. async fn make_router() -> anyhow::Result<(Router, DgwState, HandlesGuard)> { - let (mut state, handles) = DgwState::mock(CONFIG)?; + let (mut state, handles) = DgwState::mock(CONFIG).await?; // Start the manager with an in-memory DB. let manager = TrafficAuditManagerTask::init(":memory:").await?; diff --git a/devolutions-gateway/tests/webapp.rs b/devolutions-gateway/tests/webapp.rs index 8e090e9d1..25755d7e0 100644 --- a/devolutions-gateway/tests/webapp.rs +++ b/devolutions-gateway/tests/webapp.rs @@ -60,7 +60,7 @@ fn initialize_conf() { async fn custom_authentication_flow() -> anyhow::Result<()> { let (cov, _guard) = init_cov_mark(); initialize_conf(); - let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG)?; + let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG).await?; let mut app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); @@ -191,7 +191,7 @@ async fn static_assets_are_served(#[case] subdir: &str, #[case] uri_prefix: &str config["WebApp"]["StaticRootPath"] = json!(static_root); let config = config.to_string(); - let (state, _handle) = devolutions_gateway::DgwState::mock(&config)?; + let (state, _handle) = devolutions_gateway::DgwState::mock(&config).await?; let mut app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); @@ -229,7 +229,7 @@ async fn static_assets_are_served(#[case] subdir: &str, #[case] uri_prefix: &str async fn sign_app_token_bad_password() -> anyhow::Result<()> { let (cov, _guard) = init_cov_mark(); initialize_conf(); - let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG)?; + let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG).await?; let app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); @@ -267,7 +267,7 @@ async fn sign_app_token_bad_password() -> anyhow::Result<()> { async fn sign_app_token_username_mismatch() -> anyhow::Result<()> { let (cov, _guard) = init_cov_mark(); initialize_conf(); - let (state, _handles) = devolutions_gateway::DgwState::mock(CONFIG)?; + let (state, _handles) = devolutions_gateway::DgwState::mock(CONFIG).await?; let app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); diff --git a/testsuite/Cargo.toml b/testsuite/Cargo.toml index aa8fa4fc7..5b3ecb7b3 100644 --- a/testsuite/Cargo.toml +++ b/testsuite/Cargo.toml @@ -62,6 +62,7 @@ sysevent.path = "../crates/sysevent" sysevent-codes.path = "../crates/sysevent-codes" tempfile = "3" test-utils.path = "../crates/test-utils" +time = { version = "0.3", default-features = false, features = ["std", "parsing"] } tokio-rustls = { version = "0.26", features = ["ring"] } uuid = { version = "1", features = ["v4"] } diff --git a/testsuite/tests/cli/dgw/ai_analysis.rs b/testsuite/tests/cli/dgw/ai_analysis.rs new file mode 100644 index 000000000..592efd58e --- /dev/null +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -0,0 +1,551 @@ +//! `/jet/jrec/{sessionId}/ai-analysis` tests, against a running Gateway and a mock AI provider. + +use std::path::PathBuf; +use std::time::{Duration, Instant}; + +use base64::Engine as _; +use reqwest::StatusCode; +use serde_json::{Value, json}; +use testsuite::cli::start_dgw; +use testsuite::dgw_config::{DgwConfig, DgwConfigHandle, ProxyMode, VerbosityProfile}; +use time::OffsetDateTime; +use time::format_description::well_known::Rfc3339; +use tokio::io::{AsyncBufReadExt as _, AsyncReadExt as _, AsyncWriteExt as _, BufReader}; +use tokio::net::TcpListener; +use uuid::Uuid; + +const API_KEY: &str = "sk-ai-analysis-endpoint-test-secret"; +const START_SCOPE: &str = "gateway.tasks.recording-ai-analysis.start"; +const READ_SCOPE: &str = "gateway.tasks.recording-ai-analysis.read"; + +/// An unsigned scope token; the Gateway under test runs with token validation disabled. +fn scope_token(scope: &str) -> String { + let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD; + let header = engine.encode(r#"{"alg":"RS256","typ":"JWT","cty":"SCOPE"}"#); + let payload = engine.encode( + json!({ + "jti": Uuid::new_v4(), + "exp": OffsetDateTime::now_utc().unix_timestamp() + 3600, + "scope": scope, + }) + .to_string(), + ); + let signature = engine.encode(b"signature"); + format!("{header}.{payload}.{signature}") +} + +struct Gateway { + config: DgwConfigHandle, + process: tokio::process::Child, + client: reqwest::Client, +} + +impl Drop for Gateway { + // The config folder is deleted right after, and Windows can't delete the files Gateway still holds open. + fn drop(&mut self) { + let _ = self.process.start_kill(); + let deadline = Instant::now() + Duration::from_secs(10); + while matches!(self.process.try_wait(), Ok(None)) && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(50)); + } + } +} + +impl Gateway { + async fn start() -> Self { + let config = DgwConfig::builder() + .disable_token_validation(true) + .verbosity_profile(VerbosityProfile::ALL) + .proxy_mode(ProxyMode::Off) + .build() + .init() + .unwrap(); + let process = start_dgw(&config).await.unwrap(); + + Self { + config, + process, + client: reqwest::Client::new(), + } + } + + fn url(&self, path: &str) -> String { + format!("http://127.0.0.1:{}{path}", self.config.http_port()) + } + + /// Writes a finished session with one terminal recording in the default recording folder. + fn write_session(&self) -> Uuid { + let session_id = Uuid::new_v4(); + let dir = self.config.config_dir().join("recordings").join(session_id.to_string()); + std::fs::create_dir_all(&dir).unwrap(); + + let manifest = json!({ + "sessionId": session_id, + "startTime": 1_787_255_035, + "duration": 3, + "files": [{ "fileName": "recording-0.cast", "startTime": 1_787_255_035, "duration": 3 }], + }); + std::fs::write(dir.join("recording.json"), manifest.to_string()).unwrap(); + std::fs::write( + dir.join("recording-0.cast"), + "{\"version\": 2}\n[0.5,\"o\",\"$ ls\\r\\n\"]\n[1.0,\"o\",\"file.txt\\r\\n\"]\n", + ) + .unwrap(); + + session_id + } + + async fn post_raw(&self, session_id: Uuid, scope: &str, body: String) -> reqwest::Response { + self.client + .post(self.url(&format!("/jet/jrec/{session_id}/ai-analysis"))) + .header("content-type", "application/json") + .bearer_auth(scope_token(scope)) + .body(body) + .send() + .await + .unwrap() + } + + async fn post(&self, session_id: Uuid, body: &Value) -> (StatusCode, Value) { + read(self.post_raw(session_id, START_SCOPE, body.to_string()).await).await + } + + async fn get_analysis(&self, session_id: Uuid, task_id: Uuid, scope: &str) -> (StatusCode, Value) { + let response = self + .client + .get(self.url(&format!("/jet/jrec/{session_id}/ai-analysis/{task_id}"))) + .bearer_auth(scope_token(scope)) + .send() + .await + .unwrap(); + read(response).await + } + + /// Asks for the analysis until it has succeeded or failed. + async fn finished_analysis(&self, session_id: Uuid, task_id: Uuid) -> Value { + let deadline = Instant::now() + Duration::from_secs(60); + + loop { + let (status, task) = self.get_analysis(session_id, task_id, READ_SCOPE).await; + assert_eq!(status, StatusCode::OK, "{task}"); + + if task["state"] == "succeeded" || task["state"] == "failed" { + return task; + } + + assert!(Instant::now() < deadline, "the analysis did not finish: {task}"); + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + /// Reads the log files until they contain `text`, as the Gateway writes them in the background. + async fn logs_containing(&self, text: &str) -> String { + let deadline = Instant::now() + Duration::from_secs(10); + + loop { + let logs = self + .files(|name| name.starts_with("gateway") && name.ends_with(".log")) + .map(|path| std::fs::read_to_string(path).unwrap()) + .collect::(); + + if logs.contains(text) { + return logs; + } + + assert!(Instant::now() < deadline, "no {text:?} in the logs"); + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + /// `gateway.db` and its write-ahead log, where the Tasks are recorded. + fn database_bytes(&self) -> Vec { + self.files(|name| name == "gateway.db" || name == "gateway.db-wal") + .flat_map(|path| std::fs::read(path).unwrap()) + .collect() + } + + fn files(&self, keep: impl Fn(&str) -> bool) -> impl Iterator { + std::fs::read_dir(self.config.config_dir()) + .unwrap() + .map(|entry| entry.unwrap().path()) + .filter(move |path| path.file_name().and_then(|name| name.to_str()).is_some_and(&keep)) + } +} + +async fn read(response: reqwest::Response) -> (StatusCode, Value) { + let status = response.status(); + let body = response.bytes().await.unwrap(); + let body = if body.is_empty() { + Value::Null + } else { + serde_json::from_slice(&body).unwrap() + }; + (status, body) +} + +fn request(task_id: Uuid, model: &str, base_url: &str) -> Value { + json!({ + "taskId": task_id, + "provider": "openai-compatible", + "model": model, + "baseUrl": base_url, + "apiKey": API_KEY, + }) +} + +/// A mock AI provider that never answers, so an analysis started against it stays in progress. +async fn spawn_silent_ai_provider() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + tokio::spawn(async move { + let _stream = stream; + std::future::pending::<()>().await; + }); + } + }); + + format!("http://{addr}/v1/") +} + +/// A mock OpenAI-compatible provider answering every request with one action, or with `status` when it is not 200. +async fn spawn_ai_provider(status: u16) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + + 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 answer = if status == 200 { + json!({ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-test-2026-09-30", + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": "{\"line\":\"L1\",\"description\":\"Listed files\"}" }, + "finish_reason": "stop" + }], + "usage": { "prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30 } + }) + } else { + json!({ "error": { "message": "scripted failure" } }) + } + .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(); + }); + } + }); + + format!("http://{addr}/v1/") +} + +fn assert_rfc3339(value: &Value) { + let text = value.as_str().unwrap_or_else(|| panic!("not a string: {value}")); + OffsetDateTime::parse(text, &Rfc3339).unwrap_or_else(|error| panic!("{text}: {error}")); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn start_returns_202_then_200_for_the_same_request() { + let gateway = Gateway::start().await; + let session_id = gateway.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_silent_ai_provider().await; + + let (status, created) = gateway.post(session_id, &request(task_id, "gpt-test", &base_url)).await; + assert_eq!(status, StatusCode::ACCEPTED, "{created}"); + assert_eq!(created["id"], task_id.to_string()); + assert_eq!(created["kind"], "recording-ai-analysis"); + assert_eq!(created["target"], session_id.to_string()); + assert_eq!( + created["params"], + json!({ "provider": "openai-compatible", "model": "gpt-test", "maxOutputTokens": null }) + ); + assert_eq!(created["state"], "queued"); + assert_eq!(created["payload"], Value::Null); + assert_eq!(created["attempts"], 0); + assert_rfc3339(&created["createdAt"]); + assert_rfc3339(&created["deadlineAt"]); + assert_eq!(created["startedAt"], Value::Null); + assert_eq!(created["finishedAt"], Value::Null); + + let (status, again) = gateway.post(session_id, &request(task_id, "gpt-test", &base_url)).await; + assert_eq!(status, StatusCode::OK, "{again}"); + // The job may have started meanwhile, so only the fields a run leaves alone are compared. + for field in ["id", "kind", "target", "params", "createdAt"] { + assert_eq!(again[field], created[field], "{field}"); + } +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn start_conflicts_return_409_with_a_code() { + let gateway = Gateway::start().await; + let session_id = gateway.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_silent_ai_provider().await; + gateway.post(session_id, &request(task_id, "gpt-test", &base_url)).await; + + let (status, body) = gateway + .post(session_id, &request(task_id, "gpt-other", &base_url)) + .await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(body, json!({ "code": "task_id_conflict" })); + + let (status, body) = gateway + .post(session_id, &request(Uuid::new_v4(), "gpt-test", &base_url)) + .await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(body, json!({ "code": "analysis_in_progress", "taskId": task_id })); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn start_without_recording_returns_404() { + let gateway = Gateway::start().await; + let base_url = spawn_silent_ai_provider().await; + + let (status, _) = gateway + .post(Uuid::new_v4(), &request(Uuid::new_v4(), "gpt-test", &base_url)) + .await; + + assert_eq!(status, StatusCode::NOT_FOUND); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn start_ignores_unknown_fields() { + let gateway = Gateway::start().await; + let session_id = gateway.write_session(); + let base_url = spawn_silent_ai_provider().await; + let mut body = request(Uuid::new_v4(), "gpt-test", &base_url); + body["temperature"] = json!(0.5); + + let (status, created) = gateway.post(session_id, &body).await; + + assert_eq!(status, StatusCode::ACCEPTED, "{created}"); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn invalid_bodies_and_settings_return_400() { + let gateway = Gateway::start().await; + let session_id = gateway.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_silent_ai_provider().await; + let valid = request(task_id, "gpt-test", &base_url); + + let mut no_key = valid.clone(); + no_key.as_object_mut().unwrap().remove("apiKey"); + + let mut empty_key = valid.clone(); + empty_key["apiKey"] = json!(""); + + let mut no_base_url = valid.clone(); + no_base_url.as_object_mut().unwrap().remove("baseUrl"); + + let mut bad_base_url = valid.clone(); + bad_base_url["baseUrl"] = json!("ftp://ai.example/v1/"); + + let mut unknown_provider = valid.clone(); + unknown_provider["provider"] = json!("skynet"); + + let mut long_language = valid.clone(); + long_language["language"] = json!("zh-Hant-TW-x"); + + let mut empty_language = valid.clone(); + empty_language["language"] = json!(""); + + for (case, body) in [ + ("missing API key", no_key.to_string()), + ("empty API key", empty_key.to_string()), + ("empty model", request(task_id, " ", &base_url).to_string()), + ("missing base URL", no_base_url.to_string()), + ("unsupported base URL", bad_base_url.to_string()), + ("unknown provider", unknown_provider.to_string()), + ("language longer than 10 characters", long_language.to_string()), + ("empty language", empty_language.to_string()), + ("not JSON", "{".to_owned()), + ] { + let response = gateway.post_raw(session_id, START_SCOPE, body).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{case}"); + } + + let (status, _) = gateway.get_analysis(session_id, task_id, READ_SCOPE).await; + assert_eq!(status, StatusCode::NOT_FOUND, "no Task is recorded"); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn endpoints_require_their_scope() { + let gateway = Gateway::start().await; + let session_id = gateway.write_session(); + let base_url = spawn_silent_ai_provider().await; + let body = request(Uuid::new_v4(), "gpt-test", &base_url).to_string(); + + for scope in [READ_SCOPE, "gateway.recording.delete", "gateway.recordings.read"] { + let response = gateway.post_raw(session_id, scope, body.clone()).await; + assert_eq!(response.status(), StatusCode::FORBIDDEN, "{scope}"); + } + + for scope in [START_SCOPE, "gateway.sessions.read"] { + let (status, _) = gateway.get_analysis(session_id, Uuid::new_v4(), scope).await; + assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); + } + + let response = gateway.post_raw(session_id, "*", body).await; + assert_eq!( + response.status(), + StatusCode::ACCEPTED, + "the wildcard scope is accepted" + ); + + let (status, _) = gateway.get_analysis(session_id, Uuid::new_v4(), "*").await; + assert_eq!(status, StatusCode::NOT_FOUND, "the wildcard scope is accepted"); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn an_analysis_is_found_at_its_location_only() { + let gateway = Gateway::start().await; + let session_id = gateway.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_silent_ai_provider().await; + + let response = gateway + .post_raw( + session_id, + START_SCOPE, + request(task_id, "gpt-test", &base_url).to_string(), + ) + .await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + let location = response.headers()["location"].to_str().unwrap().to_owned(); + assert_eq!(location, format!("/jet/jrec/{session_id}/ai-analysis/{task_id}")); + + let (status, task) = gateway.get_analysis(session_id, task_id, READ_SCOPE).await; + assert_eq!(status, StatusCode::OK, "{task}"); + assert_eq!(task["id"], task_id.to_string()); + + let (status, _) = gateway.get_analysis(Uuid::new_v4(), task_id, READ_SCOPE).await; + assert_eq!(status, StatusCode::NOT_FOUND, "not under another session"); +} + +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +#[tokio::test] +async fn get_analysis_returns_404_for_an_unknown_id() { + let gateway = Gateway::start().await; + + let (status, _) = gateway.get_analysis(Uuid::new_v4(), Uuid::new_v4(), READ_SCOPE).await; + + assert_eq!(status, StatusCode::NOT_FOUND); +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets +#[tokio::test] +async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { + let gateway = Gateway::start().await; + let session_id = gateway.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_ai_provider(200).await; + let mut body = request(task_id, "gpt-test", &base_url); + body["maxOutputTokens"] = json!(1000); + + let (status, created) = gateway.post(session_id, &body).await; + assert_eq!(status, StatusCode::ACCEPTED, "{created}"); + + let task = gateway.finished_analysis(session_id, task_id).await; + assert_eq!(task["state"], "succeeded", "{task}"); + assert_eq!(task["attempts"], 1); + assert_eq!(task["payload"]["fileName"], "ai-analysis-0.slog"); + assert_eq!(task["payload"]["model"], "gpt-test-2026-09-30"); + // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis + assert_eq!( + task["payload"]["usage"], + json!({ "inputTokens": 20, "outputTokens": 40, "unreportedRequests": 0 }), + "one request for the actions, one for the summary" + ); + assert_rfc3339(&task["startedAt"]); + assert_rfc3339(&task["finishedAt"]); + + let (status, again) = gateway.post(session_id, &body).await; + assert_eq!(status, StatusCode::OK, "{again}"); + + let responses = [created, task, again].map(|body| body.to_string()).join("\n"); + let logs = gateway.logs_containing("Task succeeded").await; + let database = String::from_utf8_lossy(&gateway.database_bytes()).into_owned(); + assert!(!database.is_empty(), "gateway.db is read"); + + for secret in [API_KEY, base_url.as_str()] { + assert!(!responses.contains(secret), "{secret} in a response"); + assert!(!logs.contains(secret), "{secret} in the logs"); + assert!(!database.contains(secret), "{secret} in gateway.db"); + } +} + +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +#[tokio::test] +async fn provider_refusing_the_key_fails_the_task() { + let gateway = Gateway::start().await; + let session_id = gateway.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_ai_provider(401).await; + + let (status, _) = gateway.post(session_id, &request(task_id, "gpt-test", &base_url)).await; + assert_eq!(status, StatusCode::ACCEPTED); + + let task = gateway.finished_analysis(session_id, task_id).await; + assert_eq!(task["state"], "failed", "{task}"); + assert_eq!(task["attempts"], 1); + assert_eq!( + task["payload"], + json!({ + "error": { "kind": "aiRejected", "status": 401 }, + "usage": { "inputTokens": 0, "outputTokens": 0, "unreportedRequests": 0 }, + }) + ); + + let (status, _) = gateway + .post(session_id, &request(Uuid::new_v4(), "gpt-test", &base_url)) + .await; + assert_eq!( + status, + StatusCode::ACCEPTED, + "a failed analysis no longer blocks the session" + ); +} diff --git a/testsuite/tests/cli/dgw/mod.rs b/testsuite/tests/cli/dgw/mod.rs index 5bd4e0c49..0e29f568d 100644 --- a/testsuite/tests/cli/dgw/mod.rs +++ b/testsuite/tests/cli/dgw/mod.rs @@ -1,3 +1,4 @@ +mod ai_analysis; mod benign_disconnect; mod cli_args; mod heartbeat;