From d0bab560038b0e0ff8d5ff84bf8c5d8f4e350277 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 6 Oct 2026 15:10:51 -0400 Subject: [PATCH 01/53] refactor(dgw): give the gateway state the provisioner task store The Gateway state now holds the provisioner task store, on its own gateway.db connection, so the work it records can reach it. Building the mock state opens an in-memory database, so `DgwState::mock` is now async. The job queue attempt limit is now the named constant `MAX_ATTEMPTS`, still 5, so a job can refer to it. No behavior change. Co-Authored-By: Claude Opus 5.5 (1M context) --- Cargo.lock | 2 ++ devolutions-gateway/Cargo.toml | 2 ++ devolutions-gateway/src/job_queue.rs | 4 ++++ devolutions-gateway/src/lib.rs | 7 ++++++- devolutions-gateway/src/service.rs | 11 +++++++++++ devolutions-gateway/tests/preflight.rs | 24 ++++++++++++------------ devolutions-gateway/tests/traffic.rs | 2 +- devolutions-gateway/tests/webapp.rs | 8 ++++---- 8 files changed, 42 insertions(+), 18 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7c9a28980..fd7667d18 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1885,6 +1885,8 @@ dependencies = [ "picky-krb", "pin-project-lite 0.2.17", "proptest", + "provisioner-task", + "provisioner-task-libsql", "rand 0.10.2", "reqwest", "rstest", diff --git a/devolutions-gateway/Cargo.toml b/devolutions-gateway/Cargo.toml index ee26e0adb..ea8fefd64 100644 --- a/devolutions-gateway/Cargo.toml +++ b/devolutions-gateway/Cargo.toml @@ -38,6 +38,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" diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 8985f3184..615bbfb63 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -11,6 +11,9 @@ use job_queue_libsql::libsql; use time::OffsetDateTime; use tokio::sync::{Notify, mpsc}; +/// 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, @@ -59,6 +62,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); diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index a163bfa40..f6e20c20f 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -69,6 +69,7 @@ pub struct DgwState { pub monitoring_state: Arc, pub traffic_audit_handle: traffic_audit::TrafficAuditHandle, pub agent_tunnel_handle: Option>, + pub provisioner_tasks: provisioner_task::DynProvisionerTaskStore, } #[doc(hidden)] @@ -83,7 +84,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 +97,9 @@ 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 provisioner_tasks = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?; + let provisioner_tasks = + Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(provisioner_tasks).await?); let state = Self { conf_handle, @@ -111,6 +115,7 @@ impl DgwState { synthetic_kdc_registry, monitoring_state, agent_tunnel_handle: None, + provisioner_tasks, }; let handles = MockHandles { diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index 258e15d23..32abd87a7 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -340,6 +340,16 @@ 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 provisioner_tasks = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(provisioner_tasks_conn) + .await + .context("failed to initialize the provisioner task store")?, + ); + let state = DgwState { conf_handle: conf_handle.clone(), token_cache: Arc::clone(&token_cache), @@ -354,6 +364,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { monitoring_state, traffic_audit_handle: traffic_audit_task.handle(), agent_tunnel_handle, + provisioner_tasks, }; for listener in &conf.listeners { 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)))); From f398d0c64a061e4a72329484e575c38f91c5e141 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 7 Oct 2026 17:03:00 -0400 Subject: [PATCH 02/53] feat(dgw): run AI analysis as a provisioner task Gateway can now describe what the user did in a recorded session with an AI provider, and add the result to the session as a new log. The analysis itself is the AI crate's `analyze_recording`; Gateway gives it the session manifest once the recording is finished, then adds the log it returns to the session through the recording manager. Provisioner tasks run through a ProvisionerTaskRunner, the one place that keeps a task's record in gateway.db and the job running it on the job queue in step. Each kind of work is a ProvisionerTask: what one attempt does, and what to drop once its record ends. The runner records the Task and queues its job, starts each attempt, cancels and awaits an attempt past its timeout, retries transient failures until the job queue's attempts are used up, and records how the Task ends. The Task is the source of truth and its job only wakes it up: each Task names the one job allowed to run it, so a stale or duplicate job stops without running anything. At startup, before the job runner claims anything, every unfinished Task is handed to a new job, which also covers a Task recorded just before a crash, before its job was queued. AI analysis (`recording-ai-analysis`) is the first kind. A repeated request with the same Task ID gets its Task back, even while the session records again. A new request is refused while the session is recording, or when it has no recording. The API key and the base URL stay in memory only, encrypted, as the Task's secret in the provisioning store, until the Task ends or reaches its deadline. After a restart the attempt finds no key and fails its Task. Neither the key nor the URL appears in errors, logs or stored results: the AI client (#2043) leaves the URL out of its transport errors. Analyses of one session run one at a time, so two attempts never share the working files the AI crate keeps next to the manifest. Those files are removed when the Task ends. Starting an AI analysis is an internal function for now; the HTTP endpoint comes next. Co-Authored-By: Claude Opus 5.5 (1M context) --- Cargo.lock | 2 + devolutions-gateway/Cargo.toml | 2 + devolutions-gateway/src/job_queue.rs | 29 +- devolutions-gateway/src/lib.rs | 15 +- .../src/provisioner_tasks/ai.rs | 262 ++++++ .../src/provisioner_tasks/mod.rs | 41 + .../recording_ai_analysis/mod.rs | 452 +++++++++ .../recording_ai_analysis/tests.rs | 864 ++++++++++++++++++ .../src/provisioner_tasks/runner/mod.rs | 483 ++++++++++ .../src/provisioner_tasks/runner/tests.rs | 301 ++++++ devolutions-gateway/src/recording.rs | 122 ++- devolutions-gateway/src/service.rs | 23 +- 12 files changed, 2581 insertions(+), 15 deletions(-) create mode 100644 devolutions-gateway/src/provisioner_tasks/ai.rs create mode 100644 devolutions-gateway/src/provisioner_tasks/mod.rs create mode 100644 devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs create mode 100644 devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs create mode 100644 devolutions-gateway/src/provisioner_tasks/runner/mod.rs create mode 100644 devolutions-gateway/src/provisioner_tasks/runner/tests.rs diff --git a/Cargo.lock b/Cargo.lock index fd7667d18..dcebbe13a 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", @@ -1915,6 +1916,7 @@ dependencies = [ "tokio-rustls", "tokio-test", "tokio-tungstenite", + "tokio-util", "tower 0.5.3", "tower-http 0.5.2", "tracing", diff --git a/devolutions-gateway/Cargo.toml b/devolutions-gateway/Cargo.toml index ea8fefd64..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" @@ -93,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/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 615bbfb63..3e4b611fc 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -6,11 +6,13 @@ 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, ProvisionerTaskRunner}; + /// Runs of a job before the job queue gives up on it. pub const MAX_ATTEMPTS: u32 = 5; @@ -41,6 +43,7 @@ pub struct JobRunnerTask { notify_runner: Arc, runner_waker: RunnerWaker, queue: DynJobQueue, + tasks: ProvisionerTaskRunner, } impl JobQueueCtx { @@ -86,6 +89,14 @@ impl JobQueueCtx { job_queue_handle: handle, }) } + + /// Writes a job to the queue right away, without going through [`JobQueueHandle`]. + /// + /// For startup, before the job runner claims anything. + pub async fn push(&self, job: T) -> anyhow::Result<()> { + let job: DynJob = Box::new(job); + self.queue.push_job(&job, None).await + } } impl JobQueueHandle { @@ -201,11 +212,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, } } } @@ -229,9 +241,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 { @@ -283,10 +296,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; @@ -300,6 +315,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 f6e20c20f..6a6ca3783 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,7 +70,7 @@ pub struct DgwState { pub monitoring_state: Arc, pub traffic_audit_handle: traffic_audit::TrafficAuditHandle, pub agent_tunnel_handle: Option>, - pub provisioner_tasks: provisioner_task::DynProvisionerTaskStore, + pub provisioner_tasks: provisioner_tasks::runner::ProvisionerTaskRunner, } #[doc(hidden)] @@ -97,9 +98,15 @@ 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 provisioner_tasks = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?; - let provisioner_tasks = - Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(provisioner_tasks).await?); + 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, + job_queue_handle.clone(), + conf_handle.clone(), + recording_manager_handle.clone(), + provisioning.clone(), + ); let state = Self { conf_handle, diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs new file mode 100644 index 000000000..8b17d3746 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -0,0 +1,262 @@ +//! 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 task secret of their +//! Task in the [`ProvisioningStore`](crate::provisioning::ProvisioningStore), until the Task ends. + +use core::fmt; + +use anyhow::Context as _; +use devolutions_gateway_ai::recording_analysis::RecordingAnalysisError; +use devolutions_gateway_ai::{AiClient, BuildError, Provider}; +use secrecy::{ExposeSecret as _, SecretString}; +use url::Url; + +use super::AttemptError; +use crate::config::Conf; + +#[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(Debug, 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, +} + +/// 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(Debug, Clone)] +pub(crate) struct AiAccess { + pub(crate) api_key: SecretString, + pub(crate) base_url: Option, +} + +/// [`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::Cancelled => AttemptError::Cancelled, + error if error.is_transient() => AttemptError::Transient(error.to_string()), + error => AttemptError::Permanent(error.to_string()), + } + } +} + +#[cfg(test)] +mod tests { + use devolutions_gateway_ai::Error; + + use super::*; + + #[test] + fn build_errors_map_to_settings_errors() { + assert_eq!( + AiSettingsError::from(&BuildError::MissingModel), + AiSettingsError::MissingModel + ); + assert_eq!( + AiSettingsError::from(&BuildError::MissingApiKey(Provider::OpenAi)), + AiSettingsError::MissingApiKey + ); + assert_eq!( + AiSettingsError::from(&BuildError::MissingBaseUrl(Provider::OpenAiCompatible)), + AiSettingsError::MissingBaseUrl + ); + assert!(matches!( + AiSettingsError::from(&BuildError::InvalidApiKey), + AiSettingsError::Invalid(_) + )); + } + + #[test] + fn rate_limits_server_and_network_errors_are_transient() { + let attempt = |error: Error| AttemptError::from(RecordingAnalysisError::Ai(error)); + let status = |status| Error::Status { + status, + message: "failed".to_owned(), + code: None, + retry_after: None, + }; + + let network = Error::Transport { + message: "connection refused".to_owned(), + }; + assert!(matches!(attempt(network), AttemptError::Transient(_))); + + for code in [429, 500, 503] { + assert!(matches!(attempt(status(code)), AttemptError::Transient(_)), "{code}"); + } + + for code in [400, 401, 403, 404] { + assert!(matches!(attempt(status(code)), AttemptError::Permanent(_)), "{code}"); + } + + let invalid = Error::InvalidOutput { + reason: "no valid action line".to_owned(), + }; + assert!(matches!(attempt(invalid), AttemptError::Permanent(_))); + + assert!(matches!( + AttemptError::from(RecordingAnalysisError::Truncated), + AttemptError::Permanent(_) + )); + assert_eq!( + AttemptError::from(RecordingAnalysisError::Cancelled), + AttemptError::Cancelled + ); + } + + #[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..b8bc845eb --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -0,0 +1,41 @@ +//! 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 provisioner_task::DynProvisionerTaskStore; + +use self::recording_ai_analysis::RecordingAiAnalysis; +use self::runner::ProvisionerTaskRunner; +use crate::config::ConfHandle; +use crate::job_queue::JobQueueHandle; +use crate::provisioning::ProvisioningStore; +use crate::recording::RecordingMessageSender; + +/// The provisioner task runner, with every provisioner task Gateway knows. +pub fn task_runner( + store: DynProvisionerTaskStore, + job_queue: JobQueueHandle, + conf_handle: ConfHandle, + recordings: RecordingMessageSender, + provisioning: ProvisioningStore, +) -> ProvisionerTaskRunner { + ProvisionerTaskRunner::builder(store, job_queue) + .register(RecordingAiAnalysis::new(conf_handle, recordings, provisioning)) + .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. + Transient(String), + /// Another attempt would fail the same way. + Permanent(String), + /// Fails the Task with a reason of its own, such as `key lost`, instead of `permanent error`. + Failed { reason: &'static str, details: String }, + /// The attempt was cancelled before its end, such as when it ran out of time. + Cancelled, +} 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..25cf3a8f0 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -0,0 +1,452 @@ +//! AI analysis (`recording-ai-analysis`): describes what the user did in one recorded session, and adds the result to +//! the session as a new log. +//! +//! [`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 adds the log to the +//! session. 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::Usage; +use devolutions_gateway_ai::recording_analysis::{ + self, Progress, RecordingAnalysisEvent, RecordingFile, RecordingManifest, Recordings, RecordingsError, +}; +use provisioner_task::{CreateOutcome, ProvisionerTaskRecord}; +use secrecy::SecretString; +use serde_json::json; +use time::OffsetDateTime; +use uuid::Uuid; + +use super::AttemptError; +use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; +use super::runner::{Attempt, ProvisionerTask}; +use crate::DgwState; +use crate::artifacts::ArtifactKind; +use crate::config::ConfHandle; +use crate::provisioning::ProvisioningStore; +use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; + +pub const KIND: &str = "recording-ai-analysis"; + +/// Time an AI analysis gets from its creation to its end, waiting and retries included. +pub const TASK_DEADLINE: time::Duration = time::Duration::hours(2); + +/// Time one attempt of an AI analysis gets. +pub const ATTEMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60); + +const KEY_LOST: &str = "the API key is no longer available, such as after a Gateway restart"; + +/// Tokens counted by the AI provider, as kept in Task results. +/// +/// It mirrors [`Usage`], so a change in the AI crate never changes what we store. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[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, + } + } +} + +/// 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, + 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. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Params { + provider: AiProvider, + model: String, + max_output_tokens: 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 log, answers cut and asked again included. + /// + /// Absent when the provider did not report them for every request. + #[serde(skip_serializing_if = "Option::is_none")] + usage: Option, + /// Versions of the prompts that wrote the log, such as `session-actions-3`. + 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, + 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, + }; + 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 + .provisioning + .insert_task_secret(record.id, &access.to_secret()?, record.deadline_at) + }) + .await + .map_err(StartError::Internal)?; + + Ok(match outcome { + CreateOutcome::Created(task) => { + info!(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, + provisioning: ProvisioningStore, + attempt_timeout: std::time::Duration, +} + +impl RecordingAiAnalysis { + pub fn new(conf_handle: ConfHandle, recordings: RecordingMessageSender, provisioning: ProvisioningStore) -> Self { + Self { + conf_handle, + recordings, + provisioning, + attempt_timeout: ATTEMPT_TIMEOUT, + } + } + + fn manifest_path(&self, session_id: Uuid) -> Utf8PathBuf { + self.conf_handle + .get_conf() + .recording_path + .join(session_id.to_string()) + .join("recording.json") + } +} + +#[async_trait] +impl ProvisionerTask for RecordingAiAnalysis { + const KIND: &'static str = KIND; + + type Params = Params; + + type Output = Analysis; + + fn deadline(&self) -> time::Duration { + TASK_DEADLINE + } + + fn attempt_timeout(&self) -> std::time::Duration { + self.attempt_timeout + } + + async fn run(&self, attempt: &Attempt, params: Params) -> Result { + let session_id = Uuid::parse_str(&attempt.target).map_err(|error| AttemptError::Failed { + reason: "invalid task", + details: format!("invalid session ID: {error}"), + })?; + + let secret = self + .provisioning + .task_secret(attempt.task_id) + .map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + + let Some(secret) = secret else { + warn!(task.id = %attempt.task_id, attempts = attempt.attempts, "AI analysis lost its API key"); + return Err(AttemptError::Failed { + reason: "key lost", + details: KEY_LOST.to_owned(), + }); + }; + + let access = AiAccess::from_secret(&secret).map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + + 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, access.api_key).await + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) { + self.provisioning.remove_task_secret(record.id); + + let Ok(session_id) = Uuid::parse_str(&record.target) else { + return; + }; + + let manifest_path = self.manifest_path(session_id); + let _session = lock_session(session_id).await; + + if let Err(error) = recording_analysis::discard(&manifest_path).await { + warn!(task.id = %record.id, session.id = %session_id, %error, "Failed to remove the AI analysis working files"); + } + } +} + +impl RecordingAiAnalysis { + async fn analyze( + &self, + attempt: &Attempt, + session_id: Uuid, + settings: AiSettings, + api_key: SecretString, + ) -> Result { + let conf = self.conf_handle.get_conf(); + + let client = settings + .client(&conf, &api_key) + .map_err(|error| AttemptError::Permanent(error.to_string()))?; + + let _session = lock_session(session_id).await; + + let recording = self + .recordings + .get_finished(session_id) + .await + .map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + + let manifest = + analysis_manifest(&recording.manifest).map_err(|error| AttemptError::Permanent(error.to_string()))?; + let manifest_path = recording.manifest_path(); + + let mut request = client + .analyze_recording(&manifest, &manifest_path) + .cancellation(attempt.cancel.clone()); + + if let Some(max_output_tokens) = settings.max_output_tokens { + request = request.max_output_tokens(max_output_tokens); + } + + let mut events = request.start(); + let mut recorded: Option<(&'static str, u8)> = None; + + let analysis = loop { + match events.recv().await { + Some(RecordingAnalysisEvent::Progress(progress)) => { + // Reading sends an event every half second: the Task is only updated when the numbers it shows change. + let shown = (progress.stage(), progress.percent()); + if recorded != Some(shown) { + recorded = Some(shown); + attempt.report_progress(progress_payload(&progress)).await; + } + } + Some(RecordingAnalysisEvent::Complete(analysis)) => break analysis, + Some(RecordingAnalysisEvent::Failure(error)) => return Err(error.into()), + Some(_) => {} + None => { + return Err(AttemptError::Permanent( + "the AI analysis stopped unexpectedly".to_owned(), + )); + } + } + }; + + let file_name = add_log(&self.recordings, session_id, &analysis.log_path) + .await + .map_err(|error| AttemptError::Permanent(format!("failed to add the log: {error:#}")))?; + + debug!(task.id = %attempt.task_id, file = %file_name, actions = analysis.actions, "Session log added"); + + Ok(Analysis { + file_name, + model: analysis.model, + usage: analysis.usage.map(TokenUsage::from), + prompt_version: analysis.prompt_version, + }) + } +} + +/// The payload of a running AI analysis Task: its progress, for whoever asks about it. +fn progress_payload(progress: &Progress) -> serde_json::Value { + let (done, total) = progress.done_and_total().unzip(); + + json!({ + "step": progress.stage(), + "done": done, + "total": total, + "percent": progress.percent(), + "message": progress.to_string(), + }) +} + +/// 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, + }) +} + +/// Adds the complete log the analysis left next to the manifest to the session. +async fn add_log(recordings: &RecordingMessageSender, session_id: Uuid, log_path: &Utf8Path) -> anyhow::Result { + let artifact_path = recordings.add_artifact(session_id, ArtifactKind::AiAnalysis).await?; + + // The log is in the session folder, so a rename moves it in at once and the session never holds a partial one. + tokio::fs::rename(log_path, &artifact_path) + .await + .with_context(|| format!("move the log to {artifact_path}"))?; + + Ok(artifact_path + .file_name() + .context("artifact path has no file name")? + .to_owned()) +} 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..3b07f8bbd --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -0,0 +1,864 @@ +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_task::ShutdownHandle; +use job_queue::Job as _; +use parking_lot::Mutex; +use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; + +use super::*; +use crate::job_queue::{JobMessage, JobQueueReceiver, MAX_ATTEMPTS}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, 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 { + start( + &self.state, + StartRequest { + task_id, + session_id, + settings, + 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 + .expect("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") + } + + /// 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 + .provisioning + .task_secret(task_id) + .expect("read the task 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 + .provisioning + .insert_task_secret(task_id, &access.to_secret().expect("secret"), expires_at) + .expect("insert the task 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() +} + +#[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, TASK_DEADLINE); + + 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.expect("job token") + ) + ); + assert!(schedule_for.is_none()); +} + +#[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"); +} + +#[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:?}"); +} + +#[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)); +} + +#[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() + ); +} + +#[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)); +} + +#[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, + "percent": 33, + "message": "asking the AI: 1 of 3 chunks described", + }) + ); + 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": 30, "outputTokens": 60 }, + "promptVersion": PROMPT_VERSION, + }), + "chunk 0 counts from its checkpoint, the failed request not at all" + ); + + let requests = requests.lock().clone(); + assert_eq!(requests.len(), 4, "only chunks 1 and 2 are sent again"); + 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"); +} + +#[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["reason"], "permanent error"); + assert!( + task.payload["details"].as_str().expect("details").contains("401"), + "{}", + task.payload + ); + assert_eq!(task.payload["attempts"], 1); + assert_eq!(requests.lock().len(), 1); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +#[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["reason"], "permanent error"); + assert_eq!(task.payload["details"], "session mixes terminal and video recordings"); +} +#[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["details"], "session has no recording"); +} + +#[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["reason"], "attempts exhausted"); + assert!( + task.payload["details"].as_str().expect("details").contains("503"), + "{}", + task.payload + ); + assert_eq!(requests.lock().len(), 5); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +#[tokio::test] +async fn slow_attempt_is_retried() { + let mut harness = Harness::new().await; + let session_id = harness.write_long_session(10); + + // A provider that never answers makes the attempt outlast its timeout. + 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; + + let mut task = RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.state.provisioning.clone(), + ); + task.attempt_timeout = Duration::from_millis(500); + harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( + Arc::clone(harness.state.provisioner_tasks.store()), + harness.state.job_queue_handle.clone(), + ) + .register(task) + .build(); + + let error = harness.run(task_id).await.expect_err("timeout is retried"); + + assert!(error.to_string().contains("took longer than"), "{error:#}"); + assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Running); + assert!(harness.has_key(task_id)); + assert!( + harness.workspace(session_id).join("plan.json").exists(), + "the chunks are kept for the retry" + ); +} + +#[tokio::test] +async fn log_added_after_the_deadline_leaves_the_task_failed() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider_after(Duration::from_secs(4), |_| 0).await; + let task_id = Uuid::new_v4(); + let now = OffsetDateTime::now_utc(); + harness + .create_task( + task_id, + session_id, + "openai-compatible", + now, + now + time::Duration::seconds(2), + ) + .await; + harness.insert_key(task_id, settings.base_url, now + time::Duration::hours(1)); + + let mut job = harness.job(task_id).await; + let job = tokio::spawn(async move { job.run().await }); + + // Reading the Task past its deadline, as a status request does, ends it while the AI is still answering. + tokio::time::sleep(Duration::from_secs(3)).await; + assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Failed); + + job.await.expect("job task").expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "timed out"); + assert!( + harness + .recording_path + .join(session_id.to_string()) + .join("ai-analysis-0.slog") + .exists(), + "the log stays in the session" + ); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +#[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.provisioning.remove_task_secret(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!({ "reason": "key lost", "details": KEY_LOST, "attempts": 1 }) + ); + assert!(!harness.workspace(session_id).exists()); +} + +#[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["reason"], "timed out"); + assert_eq!(task.attempts, 0); + assert!(!harness.has_key(task_id)); + assert!(!harness.workspace(session_id).exists()); +} + +#[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" + ); +} + +#[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"); +} + +#[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()); +} +#[test] +fn unknown_usage_is_left_out_of_the_result() { + let analysis = Analysis { + file_name: "ai-analysis-0.slog".to_owned(), + model: "gpt-test-2026-09-30".to_owned(), + usage: None, + prompt_version: PROMPT_VERSION.to_owned(), + }; + + assert_eq!( + json!(analysis), + json!({ + "fileName": "ai-analysis-0.slog", + "model": "gpt-test-2026-09-30", + "promptVersion": PROMPT_VERSION, + }) + ); +} 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..308ada5e6 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -0,0 +1,483 @@ +//! 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, and what to clean up once it has ended. +//! [`ProvisionerTaskRunner`] does the rest for every task. It creates the record, queues the job, starts each attempt, +//! retries transient failures, enforces the attempt timeout 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. + +#[cfg(test)] +mod tests; + +use std::collections::HashMap; +use std::future::Future; +use std::sync::Arc; + +use anyhow::Context as _; +use async_trait::async_trait; +use provisioner_task::{ + AttemptStart, CreateOutcome, DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskRecord, +}; +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; +use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; + +/// 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; + + /// 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 to its end, waiting and retries included. + fn deadline(&self) -> time::Duration; + + /// Time one attempt gets before it is cancelled and retried. + fn attempt_timeout(&self) -> std::time::Duration; + + /// Runs one attempt. + /// + /// The attempt should stop soon after [`Attempt::cancel`] is cancelled, and return [`AttemptError::Cancelled`]. + async fn run(&self, attempt: &Attempt, params: Self::Params) -> Result; + + /// Drops what is kept only while the record runs, once it has ended, however it ended. + /// + /// It may be called again for a record that already ended, so it must be idempotent. + async fn on_end(&self, _record: &ProvisionerTaskRecord) {} +} + +/// 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, + pub cancel: CancellationToken, + store: DynProvisionerTaskStore, +} + +impl Attempt { + /// Replaces the payload of the running record, so whoever asks about it sees how the attempt goes. + pub async fn report_progress(&self, payload: serde_json::Value) { + if let Err(error) = self.store.update_running(self.task_id, payload).await { + 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 deadline(&self) -> time::Duration; + + fn attempt_timeout(&self) -> std::time::Duration; + + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result; + + async fn on_end(&self, record: &ProvisionerTaskRecord); +} + +#[async_trait] +impl AnyProvisionerTask for T { + fn deadline(&self) -> time::Duration { + ProvisionerTask::deadline(self) + } + + fn attempt_timeout(&self) -> std::time::Duration { + ProvisionerTask::attempt_timeout(self) + } + + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { + let params = serde_json::from_value(params).map_err(|error| AttemptError::Failed { + reason: "invalid task", + details: format!("invalid parameters: {error}"), + })?; + + let output = ProvisionerTask::run(self, attempt, params).await?; + + serde_json::to_value(output).map_err(|error| AttemptError::Permanent(format!("invalid result: {error}"))) + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) { + 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: JobQueueHandle, + /// Kind → the task that runs records of this kind. + tasks: HashMap<&'static str, Arc>, +} + +pub struct ProvisionerTaskRunnerBuilder { + store: DynProvisionerTaskStore, + job_queue: JobQueueHandle, + 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, + }), + } + } +} + +/// Initial payload of a running record, until its task reports progress. +fn starting_payload() -> serde_json::Value { + json!({ "step": "preparing" }) +} + +impl ProvisionerTaskRunner { + pub fn builder(store: DynProvisionerTaskStore, job_queue: JobQueueHandle) -> 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: now + task.deadline(), + job_token, + }, + now, + ) + .await?; + + let CreateOutcome::Created(record) = &outcome else { + return Ok(outcome); + }; + + if let Err(error) = on_created(record) { + error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to prepare the Task"); + self.fail(task.as_ref(), id, failure("not queued", &format!("{error:#}"), 0)) + .await; + return Err(error); + } + + if let Err(error) = self.enqueue(id, job_token).await { + error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to queue the Task"); + self.fail(task.as_ref(), id, failure("not queued", &format!("{error:#}"), 0)) + .await; + return Err(error); + } + + info!(task.id = %id, task.kind = T::KIND, "Task queued"); + + Ok(outcome) + } + + /// Hands every unfinished record to a new job queued with `push`: the jobs queued before stop at their next attempt. + /// + /// Call it at startup, before the job runner claims any job, so no earlier job is running an attempt meanwhile. + pub async fn resume(&self, push: F) -> anyhow::Result<()> + where + F: Fn(ProvisionerTaskJob) -> Fut, + Fut: Future>, + { + let now = OffsetDateTime::now_utc(); + let records = self.inner.store.list_unfinished(now).await?; + + for record in records { + let job_token = Uuid::new_v4(); + + if !self.inner.store.replace_job_token(record.id, job_token, now).await? { + continue; + } + + let job = ProvisionerTaskJob { + task_id: record.id, + job_token, + runner: self.clone(), + }; + + match push(job).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; it fails at its deadline" + ), + } + } + + 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) -> anyhow::Result<()> { + self.inner + .job_queue + .enqueue(ProvisionerTaskJob { + task_id, + job_token, + runner: self.clone(), + }) + .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(); + + let record = match store.start_attempt(task_id, job_token, starting_payload(), now).await? { + AttemptStart::Started(record) => record, + AttemptStart::Finished(record) => { + debug!(task.id = %task_id, "Task is over, so this attempt does not run"); + if let Some(task) = self.task(&record.kind) { + task.on_end(&record).await; + } + return Ok(()); + } + 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 Some(task) = self.task(&record.kind) else { + error!(task.id = %task_id, task.kind = %record.kind, "Task of an unknown kind"); + let details = format!("unknown task kind {}", record.kind); + if let Err(error) = store + .fail(task_id, failure("invalid task", &details, record.attempts), now) + .await + { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure"); + } + return Ok(()); + }; + + let attempts = record.attempts; + let attempt_timeout = task.attempt_timeout(); + let cancel = CancellationToken::new(); + let attempt = Attempt { + task_id, + target: record.target.clone(), + attempts, + cancel: cancel.clone(), + store: Arc::clone(store), + }; + + 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 joined = match tokio::time::timeout(attempt_timeout, &mut handle).await { + Ok(joined) => joined, + Err(_) => { + // Aborting cannot stop blocking steps, so the attempt is cancelled and awaited instead: the next + // attempt never runs alongside it. + cancel.cancel(); + handle.await + } + }; + + let outcome = match joined { + Ok(Err(AttemptError::Cancelled)) => Err(AttemptError::Transient(format!( + "the attempt took longer than {attempt_timeout:?}" + ))), + Ok(outcome) => outcome, + Err(_) => Err(AttemptError::Permanent("the attempt panicked".to_owned())), + }; + + match outcome { + Ok(output) => { + // An error here must not ask for a retry: the work is done. + match store.succeed(task_id, output, OffsetDateTime::now_utc()).await { + Ok(true) => info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"), + Ok(false) => 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" + ), + Err(error) => error!( + task.id = %task_id, + error = format!("{error:#}"), + "Failed to record the Task result" + ), + } + + self.end(task.as_ref(), task_id).await; + } + Err(AttemptError::Transient(details)) if attempts < MAX_ATTEMPTS => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task attempt failed; it will be retried"); + anyhow::bail!("Task attempt failed: {details}"); + } + Err(AttemptError::Transient(details)) => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task failed after its last attempt"); + self.fail( + task.as_ref(), + task_id, + failure("attempts exhausted", &details, attempts), + ) + .await; + } + Err(AttemptError::Permanent(details)) => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task failed"); + self.fail(task.as_ref(), task_id, failure("permanent error", &details, attempts)) + .await; + } + Err(AttemptError::Failed { reason, details }) => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, reason, %details, "Task failed"); + self.fail(task.as_ref(), task_id, failure(reason, &details, attempts)) + .await; + } + Err(AttemptError::Cancelled) => unreachable!("a cancelled attempt is reported as a timeout above"), + } + + Ok(()) + } + + async fn fail(&self, task: &dyn AnyProvisionerTask, task_id: Uuid, payload: serde_json::Value) { + if let Err(error) = self.inner.store.fail(task_id, payload, OffsetDateTime::now_utc()).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure"); + } + + self.end(task, task_id).await; + } + + async fn end(&self, task: &dyn AnyProvisionerTask, task_id: Uuid) { + match self.inner.store.get(task_id, OffsetDateTime::now_utc()).await { + Ok(Some(record)) => task.on_end(&record).await, + Ok(None) => {} + Err(error) => error!(task.id = %task_id, error = format!("{error:#}"), "Failed to read the ended Task"), + } + } +} + +fn failure(reason: &str, details: &str, attempts: u32) -> serde_json::Value { + json!({ "reason": reason, "details": details, "attempts": attempts }) +} + +#[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..a40619393 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -0,0 +1,301 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use job_queue::Job as _; +use parking_lot::Mutex; +use provisioner_task::ProvisionerTaskState; + +use super::*; +use crate::job_queue::{JobMessage, JobQueueReceiver}; + +const KIND: &str = "test.echo"; + +#[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, + runs: Arc, + ended: Arc>>, +} + +#[async_trait] +impl ProvisionerTask for Echo { + const KIND: &'static str = KIND; + + type Params = EchoParams; + + type Output = serde_json::Value; + + fn deadline(&self) -> time::Duration { + time::Duration::hours(1) + } + + fn attempt_timeout(&self) -> std::time::Duration { + std::time::Duration::from_secs(5) + } + + async fn run(&self, _attempt: &Attempt, params: EchoParams) -> Result { + self.runs.fetch_add(1, Ordering::SeqCst); + (self.outcome)(params.value) + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) { + self.ended.lock().push(record.id); + } +} + +struct Harness { + runner: ProvisionerTaskRunner, + job_queue_rx: JobQueueReceiver, + runs: Arc, + ended: Arc>>, +} + +impl Harness { + async fn new(outcome: fn(u32) -> Result) -> 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 runner = ProvisionerTaskRunner::builder(store, job_queue) + .register(Echo { + outcome, + runs: Arc::clone(&runs), + ended: Arc::clone(&ended), + }) + .build(); + + Self { + runner, + job_queue_rx, + runs, + ended, + } + } + + 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") + } +} + +#[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"); +} + +#[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["reason"], "not queued"); + assert_eq!(*harness.ended.lock(), [id], "what was set up for the Task is dropped"); +} + +#[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["reason"], "not queued"); + assert!(harness.job_queue_rx.try_recv().is_err(), "no job is queued"); + assert_eq!(*harness.ended.lock(), [id]); +} + +#[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(); + + let resumed = Mutex::new(Vec::new()); + harness + .runner + .resume(|job| { + resumed.lock().push(job); + async { Ok(()) } + }) + .await + .expect("resume"); + let mut resumed = resumed.into_inner(); + assert_eq!(resumed.len(), 1, "the unfinished Task gets a 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[0].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]); + + let queued_again = AtomicUsize::new(0); + harness + .runner + .resume(|_| { + queued_again.fetch_add(1, Ordering::SeqCst); + async { Ok(()) } + }) + .await + .expect("resume again"); + assert_eq!( + queued_again.load(Ordering::SeqCst), + 0, + "a finished Task gets no new job" + ); +} + +#[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!({ "reason": "permanent error", "details": "the attempt panicked", "attempts": 1 }) + ); + assert_eq!(*harness.ended.lock(), [id]); +} + +#[tokio::test] +async fn a_failure_keeps_the_reason_of_its_kind() { + let mut harness = Harness::new(|_| { + Err(AttemptError::Failed { + reason: "key lost", + details: "gone".to_owned(), + }) + }) + .await; + let id = harness.create(1).await; + + harness.queued_job().run().await.expect("no retry"); + + assert_eq!( + harness.task(id).await.payload, + json!({ "reason": "key lost", "details": "gone", "attempts": 1 }) + ); +} + +#[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["reason"], "invalid task", "{kind}"); + } + + assert_eq!(harness.runs.load(Ordering::SeqCst), 0); +} + +#[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!(Some(read.job_token), harness.task(id).await.job_token); + assert_eq!(job.name(), ProvisionerTaskJob::NAME); +} diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 8ece9b9b6..5c5814230 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -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,6 +59,20 @@ 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)?; @@ -316,7 +341,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 +382,10 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, + GetFinished { + id: Uuid, + channel: oneshot::Sender>, + }, Disconnect { id: Uuid, }, @@ -377,6 +406,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 +448,9 @@ impl fmt::Debug for RecordingManagerMessage { .field("id", id) .field("kind", kind) .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 +514,36 @@ impl RecordingMessageSender { rx.await.context("couldn't receive AddArtifact 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 +904,15 @@ impl RecordingManagerTask { Ok(artifact_path) } + 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 +1031,9 @@ async fn recording_manager_task( RecordingManagerMessage::AddArtifact { id, kind, channel } => { let _ = channel.send(manager.handle_add_artifact(id, kind).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"); @@ -1421,6 +1518,25 @@ mod tests { assert_eq!(harness.sender.get_count().await.expect("count"), 1); } + #[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/service.rs b/devolutions-gateway/src/service.rs index 32abd87a7..eea186dca 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -344,12 +344,26 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { .connect() .await .context("failed to connect the provisioner tasks to the gateway database")?; - let provisioner_tasks = Arc::new( + 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 provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( + task_store, + job_queue_ctx.job_queue_handle.clone(), + conf_handle.clone(), + recording_manager_handle.clone(), + provisioning.clone(), + ); + + // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. + provisioner_tasks + .resume(|job| job_queue_ctx.push(job)) + .await + .context("failed to queue the unfinished provisioner tasks again")?; + let state = DgwState { conf_handle: conf_handle.clone(), token_cache: Arc::clone(&token_cache), @@ -364,7 +378,7 @@ 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: provisioner_tasks.clone(), }; for listener in &conf.listeners { @@ -429,7 +443,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); From e66ab4dc018a13b348dfb530a9bf0910404b9562 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 6 Oct 2026 16:21:34 -0400 Subject: [PATCH 03/53] feat(dgw): add endpoints to start and read an AI analysis The provisioner can now start an AI analysis of a recorded session with `POST /jet/jrec/{session_id}/ai-analysis`, and follow its Task with `GET /jet/jrec/{session_id}/ai-analysis/{task_id}`, which the Location header of a 202 answer names. A Task of another kind, or the analysis of another session, is not found under that path. Each provisioner task declares its own scopes, named after its kind: `gateway.tasks..start` and `gateway.tasks..read`. AI analysis uses `gateway.tasks.recording-ai-analysis.start` and `gateway.tasks.recording-ai-analysis.read`. The request body holds the Task ID chosen by the caller, the AI provider and model, an optional base URL and output token limit, and the API key. A new Task answers 202, the same request again answers 200 with the same Task. Gateway answers 409 with a code when the Task ID is used for other parameters, when the session still records, or when another AI analysis of the session is not finished; the last one names the Task ID. A session without recording answers 404, invalid settings or an unknown field answer 400. The API key and the base URL stay in memory only: they are never logged, returned or written to gateway.db. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/jrec/mod.rs | 247 +++++++- devolutions-gateway/src/api/mod.rs | 1 + devolutions-gateway/src/api/tasks.rs | 85 +++ devolutions-gateway/src/extract.rs | 41 ++ devolutions-gateway/src/openapi.rs | 67 +++ .../src/provisioner_tasks/ai.rs | 2 + .../recording_ai_analysis/mod.rs | 5 + .../src/provisioner_tasks/runner/mod.rs | 7 + .../src/provisioner_tasks/runner/tests.rs | 5 + devolutions-gateway/src/token.rs | 4 + devolutions-gateway/tests/ai_analysis.rs | 566 ++++++++++++++++++ .../tests/dvls_compatibility.rs | 6 + 12 files changed, 1033 insertions(+), 3 deletions(-) create mode 100644 devolutions-gateway/src/api/tasks.rs create mode 100644 devolutions-gateway/tests/ai_analysis.rs diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index 32f4b7d14..bc9c86cd9 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -12,8 +12,8 @@ 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; @@ -22,17 +22,24 @@ use camino::{Utf8Path, Utf8PathBuf}; 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 +69,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 +438,192 @@ 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", deny_unknown_fields)] +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. + #[serde(default)] + max_output_tokens: 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 adds the answer to the session as a +/// new log, listed under `artifacts.ai-analysis` in its manifest. +/// 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. +/// 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 or AI settings"), + (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, + api_key, + } = request; + + let outcome = recording_ai_analysis::start( + &state, + StartRequest { + task_id, + session_id, + settings: AiSettings { + provider, + model, + base_url, + max_output_tokens, + }, + 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, and its result once it is over. +/// 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 +1207,52 @@ mod tests { use super::*; + #[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..c82e99076 --- /dev/null +++ b/devolutions-gateway/src/api/tasks.rs @@ -0,0 +1,85 @@ +//! 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`, `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, + /// When an unfinished Task becomes failed, waiting and retries included. + #[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/extract.rs b/devolutions-gateway/src/extract.rs index a000e46cf..0212a79d1 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,46 @@ 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/openapi.rs b/devolutions-gateway/src/openapi.rs index e023be143..6ef5e0f66 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -25,6 +25,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 +66,16 @@ 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, + crate::provisioner_tasks::ai::AiProvider, + crate::api::tasks::TaskInfo, + crate::api::tasks::TaskState, + AiAnalysisRunningPayload, + AiAnalysisSucceededPayload, + AiAnalysisUsage, + AiAnalysisFailedPayload, crate::token::AccessScope, crate::api::webapp::AppTokenSignRequest, crate::api::webapp::AppTokenContentType, @@ -142,6 +154,61 @@ struct SessionInfo { destination_host: Option, } +/// Payload of a `running` AI analysis Task: its progress +#[allow(dead_code)] +#[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`. + done: Option, + /// Work the step has in all, in the same unit as `done`. + total: Option, + /// Part of the step done, from 0 to 100. + percent: Option, + /// Short sentence for people following the analysis, such as `asking the AI: 2 of 5 chunks described`. + message: Option, +} + +/// Payload of a `succeeded` AI analysis Task: the log it added to the session +#[allow(dead_code)] +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisSucceededPayload { + /// 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 counted by the AI provider; absent when it did not report them for every request. + usage: Option, + /// Version of the instructions sent to the AI provider. + prompt_version: String, +} + +/// Tokens counted by an AI provider +#[allow(dead_code)] +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisUsage { + input_tokens: u64, + output_tokens: u64, +} + +/// Payload of a `failed` AI analysis Task: why it failed +#[allow(dead_code)] +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisFailedPayload { + /// Kind of failure, such as `permanent error`, `attempts exhausted`, `key lost` or `timed out`. + reason: String, + /// What went wrong: a text, except for `timed out`, where it is `{ "lastPayload": }`. + details: serde_json::Value, + /// Number of attempts started. + attempts: u32, +} + #[allow(unused)] #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "kebab-case")] diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index 8b17d3746..6ae4bf015 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -14,6 +14,8 @@ use url::Url; use super::AttemptError; 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")] diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 25cf3a8f0..48b07b0b4 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -33,6 +33,7 @@ use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; use crate::provisioning::ProvisioningStore; use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; +use crate::token::AccessScope; pub const KIND: &str = "recording-ai-analysis"; @@ -251,6 +252,10 @@ impl RecordingAiAnalysis { 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; diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 308ada5e6..4e43724f2 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -30,6 +30,7 @@ use uuid::Uuid; use super::AttemptError; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; +use crate::token::AccessScope; /// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. /// @@ -40,6 +41,12 @@ 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; diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index a40619393..f063a771f 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -6,6 +6,7 @@ use provisioner_task::ProvisionerTaskState; use super::*; use crate::job_queue::{JobMessage, JobQueueReceiver}; +use crate::token::AccessScope; const KIND: &str = "test.echo"; @@ -25,6 +26,10 @@ struct Echo { 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; 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/ai_analysis.rs b/devolutions-gateway/tests/ai_analysis.rs new file mode 100644 index 000000000..d69889bf2 --- /dev/null +++ b/devolutions-gateway/tests/ai_analysis.rs @@ -0,0 +1,566 @@ +#![allow(unused_crate_dependencies)] +#![allow(clippy::unwrap_used)] + +use std::io::Write; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; + +use axum::Router; +use axum::body::Body; +use axum::extract::connect_info::MockConnectInfo; +use axum::http::{Request, StatusCode}; +use base64::Engine as _; +use camino::{Utf8Path, Utf8PathBuf}; +use devolutions_gateway::job_queue::{JobMessage, JobQueueReceiver}; +use devolutions_gateway::recording::RecordingManagerTask; +use devolutions_gateway::{DgwState, MockHandles}; +use devolutions_gateway_task::ShutdownHandle; +use http_body_util::BodyExt as _; +use serde_json::{Value, json}; +use time::OffsetDateTime; +use time::format_description::well_known::Rfc3339; +use tower::ServiceExt as _; +use tracing_subscriber::util::SubscriberInitExt as _; +use uuid::Uuid; + +const API_KEY: &str = "sk-ai-analysis-endpoint-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, + "__debug__": { "disable_token_validation": true }, + }) + .to_string() +} + +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}") +} + +/// Collects every log line, so a test can check that no secret reaches the logs. +#[derive(Clone, Default)] +struct LogBuffer(Arc>>); + +impl Write for LogBuffer { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +impl LogBuffer { + fn text(&self) -> String { + String::from_utf8_lossy(&self.0.lock().unwrap()).into_owned() + } +} + +struct Harness { + app: Router, + job_queue_rx: JobQueueReceiver, + recording_path: Utf8PathBuf, + db_path: Utf8PathBuf, + task_store: provisioner_task::DynProvisionerTaskStore, + _keep: Box, + _dir: tempfile::TempDir, +} + +impl Harness { + async fn new() -> Self { + let dir = tempfile::tempdir().unwrap(); + let root = Utf8PathBuf::from_path_buf(dir.path().to_path_buf()).unwrap(); + let recording_path = root.join("recordings"); + let db_path = root.join("gateway.db"); + + let (mut state, handles) = DgwState::mock(&config(&recording_path)).await.unwrap(); + + // A file, unlike the in-memory default, can be searched for secrets afterwards. + let connection = gateway_db::GatewayDb::open_path(db_path.as_str()) + .await + .unwrap() + .connect() + .await + .unwrap(); + let task_store: provisioner_task::DynProvisionerTaskStore = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(connection) + .await + .unwrap(), + ); + state.provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( + Arc::clone(&task_store), + state.job_queue_handle.clone(), + state.conf_handle.clone(), + state.recordings.clone(), + state.provisioning.clone(), + ); + + 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, + )); + + let app = devolutions_gateway::make_http_service(state) + .layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 3000)))); + + Self { + app, + job_queue_rx, + recording_path, + db_path, + task_store, + _keep: Box::new(( + session_manager_rx, + subscriber_rx, + traffic_audit_rx, + shutdown_handle, + recording_shutdown_handle, + )), + _dir: dir, + } + } + + fn write_session(&self) -> Uuid { + let session_id = Uuid::new_v4(); + let dir = self.recording_path.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 send(&self, request: Request) -> (StatusCode, Value) { + let response = self.app.clone().oneshot(request).await.unwrap(); + let status = response.status(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let body = if body.is_empty() { + Value::Null + } else { + serde_json::from_slice(&body).unwrap() + }; + (status, body) + } + + async fn post_raw(&self, session_id: Uuid, scope: &str, body: String) -> (StatusCode, Value) { + let request = Request::builder() + .method("POST") + .uri(format!("/jet/jrec/{session_id}/ai-analysis")) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {}", scope_token(scope))) + .body(Body::from(body)) + .unwrap(); + self.send(request).await + } + + async fn post(&self, session_id: Uuid, body: &Value) -> (StatusCode, Value) { + self.post_raw( + session_id, + "gateway.tasks.recording-ai-analysis.start", + body.to_string(), + ) + .await + } + + async fn get_analysis(&self, session_id: Uuid, task_id: Uuid, scope: &str) -> (StatusCode, Value) { + let request = Request::builder() + .uri(format!("/jet/jrec/{session_id}/ai-analysis/{task_id}")) + .header("authorization", format!("Bearer {}", scope_token(scope))) + .body(Body::empty()) + .unwrap(); + self.send(request).await + } + + async fn run_queued_job(&mut self) { + let JobMessage { mut job, .. } = self.job_queue_rx.try_recv().expect("a queued job"); + job.run().await.unwrap(); + } + + fn database_bytes(&self) -> Vec { + let mut bytes = Vec::new(); + for suffix in ["", "-wal"] { + if let Ok(content) = std::fs::read(format!("{}{suffix}", self.db_path)) { + bytes.extend(content); + } + } + bytes + } +} + +fn request(task_id: Uuid, model: &str) -> Value { + json!({ "taskId": task_id, "provider": "openai", "model": model, "apiKey": API_KEY }) +} + +/// 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 app = Router::new().fallback(axum::routing::post(move || async move { + if status != 200 { + return ( + StatusCode::from_u16(status).unwrap(), + axum::Json(json!({ "error": { "message": "scripted failure" } })), + ); + } + + let answer = 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 } + }); + (StatusCode::OK, axum::Json(answer)) + })); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, app).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}")); +} + +#[tokio::test] +async fn start_returns_202_then_200_for_the_same_request() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + + let (status, created) = harness.post(session_id, &request(task_id, "gpt-test")).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", "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) = harness.post(session_id, &request(task_id, "gpt-test")).await; + assert_eq!(status, StatusCode::OK, "{again}"); + assert_eq!(again, created); +} + +#[tokio::test] +async fn start_conflicts_return_409_with_a_code() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + harness.post(session_id, &request(task_id, "gpt-test")).await; + + let (status, body) = harness.post(session_id, &request(task_id, "gpt-other")).await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(body, json!({ "code": "task_id_conflict" })); + + let (status, body) = harness.post(session_id, &request(Uuid::new_v4(), "gpt-test")).await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(body, json!({ "code": "analysis_in_progress", "taskId": task_id })); +} + +#[tokio::test] +async fn start_without_recording_returns_404() { + let harness = Harness::new().await; + + let (status, _) = harness.post(Uuid::new_v4(), &request(Uuid::new_v4(), "gpt-test")).await; + + assert_eq!(status, StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn invalid_bodies_and_settings_return_400() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + + let mut unknown_field = request(task_id, "gpt-test"); + unknown_field["temperature"] = json!(0.5); + + let mut no_key = request(task_id, "gpt-test"); + no_key.as_object_mut().unwrap().remove("apiKey"); + + let mut empty_key = request(task_id, "gpt-test"); + empty_key["apiKey"] = json!(""); + + let mut no_base_url = request(task_id, "gpt-test"); + no_base_url["provider"] = json!("openai-compatible"); + + let mut bad_base_url = request(task_id, "gpt-test"); + bad_base_url["baseUrl"] = json!("ftp://ai.example/v1/"); + + let mut unknown_provider = request(task_id, "gpt-test"); + unknown_provider["provider"] = json!("skynet"); + + for (case, body) in [ + ("unknown field", unknown_field.to_string()), + ("missing API key", no_key.to_string()), + ("empty API key", empty_key.to_string()), + ("empty model", request(task_id, " ").to_string()), + ("missing base URL", no_base_url.to_string()), + ("unsupported base URL", bad_base_url.to_string()), + ("unknown provider", unknown_provider.to_string()), + ("not JSON", "{".to_owned()), + ] { + let (status, _) = harness + .post_raw(session_id, "gateway.tasks.recording-ai-analysis.start", body) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{case}"); + } + + let (status, _) = harness + .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::NOT_FOUND, "no Task is recorded"); +} + +#[tokio::test] +async fn endpoints_require_their_scope() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let body = request(Uuid::new_v4(), "gpt-test").to_string(); + + for scope in [ + "gateway.tasks.recording-ai-analysis.read", + "gateway.recording.delete", + "gateway.recordings.read", + ] { + let (status, _) = harness.post_raw(session_id, scope, body.clone()).await; + assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); + } + + for scope in ["gateway.tasks.recording-ai-analysis.start", "gateway.sessions.read"] { + let (status, _) = harness.get_analysis(session_id, Uuid::new_v4(), scope).await; + assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); + } + + let (status, _) = harness.post_raw(session_id, "*", body).await; + assert_eq!(status, StatusCode::ACCEPTED, "the wildcard scope is accepted"); + + let (status, _) = harness.get_analysis(session_id, Uuid::new_v4(), "*").await; + assert_eq!(status, StatusCode::NOT_FOUND, "the wildcard scope is accepted"); +} + +#[tokio::test] +async fn an_analysis_is_found_at_its_location_only() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + let request = Request::builder() + .method("POST") + .uri(format!("/jet/jrec/{session_id}/ai-analysis")) + .header("content-type", "application/json") + .header( + "authorization", + format!("Bearer {}", scope_token("gateway.tasks.recording-ai-analysis.start")), + ) + .body(Body::from(request(task_id, "gpt-test").to_string())) + .unwrap(); + + let response = harness.app.clone().oneshot(request).await.unwrap(); + + 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) = harness + .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::OK, "{task}"); + assert_eq!(task["id"], task_id.to_string()); + + let (status, _) = harness + .get_analysis(Uuid::new_v4(), task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::NOT_FOUND, "not under another session"); + + // A Task of another kind is not an AI analysis, whatever the scope. + let now = OffsetDateTime::now_utc(); + let other_kind = Uuid::new_v4(); + harness + .task_store + .create( + provisioner_task::NewProvisionerTaskRecord { + id: other_kind, + kind: "test.other-kind".to_owned(), + target: session_id.to_string(), + params: json!({}), + deadline_at: now + time::Duration::hours(1), + job_token: Uuid::new_v4(), + }, + now, + ) + .await + .unwrap(); + let (status, _) = harness.get_analysis(session_id, other_kind, "*").await; + assert_eq!(status, StatusCode::NOT_FOUND); +} +#[tokio::test] +async fn get_analysis_returns_404_for_an_unknown_id() { + let harness = Harness::new().await; + + let (status, _) = harness + .get_analysis( + Uuid::new_v4(), + Uuid::new_v4(), + "gateway.tasks.recording-ai-analysis.read", + ) + .await; + + assert_eq!(status, StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { + let logs = LogBuffer::default(); + let _guard = tracing_subscriber::fmt() + .with_writer({ + let logs = logs.clone(); + move || logs.clone() + }) + .with_ansi(false) + .with_max_level(tracing::Level::TRACE) + .set_default(); + + let mut harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_ai_provider(200).await; + + let body = json!({ + "taskId": task_id, + "provider": "openai-compatible", + "model": "gpt-test", + "baseUrl": base_url, + "maxOutputTokens": 1000, + "apiKey": API_KEY, + }); + + let (status, created) = harness.post(session_id, &body).await; + assert_eq!(status, StatusCode::ACCEPTED, "{created}"); + + harness.run_queued_job().await; + + let (status, task) = harness + .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::OK); + 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"); + assert_eq!( + task["payload"]["usage"], + json!({ "inputTokens": 10, "outputTokens": 20 }) + ); + assert_rfc3339(&task["startedAt"]); + assert_rfc3339(&task["finishedAt"]); + + let (status, again) = harness.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 = logs.text(); + let database = String::from_utf8_lossy(&harness.database_bytes()).into_owned(); + assert!(!database.is_empty(), "gateway.db is read"); + assert!(logs.contains("Task succeeded"), "logs are captured"); + + 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"); + } +} + +#[tokio::test] +async fn provider_refusing_the_key_fails_the_task() { + let mut harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_ai_provider(401).await; + + let body = json!({ + "taskId": task_id, + "provider": "openai-compatible", + "model": "gpt-test", + "baseUrl": base_url, + "apiKey": API_KEY, + }); + let (status, _) = harness.post(session_id, &body).await; + assert_eq!(status, StatusCode::ACCEPTED); + + harness.run_queued_job().await; + + let (status, task) = harness + .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(task["state"], "failed", "{task}"); + assert_eq!(task["payload"]["reason"], "permanent error"); + assert_eq!(task["payload"]["attempts"], 1); + + let (status, _) = harness.post(session_id, &request(Uuid::new_v4(), "gpt-test")).await; + assert_eq!( + status, + StatusCode::ACCEPTED, + "a failed analysis no longer blocks the session" + ); +} diff --git a/devolutions-gateway/tests/dvls_compatibility.rs b/devolutions-gateway/tests/dvls_compatibility.rs index e727c8e01..fa0d080cc 100644 --- a/devolutions-gateway/tests/dvls_compatibility.rs +++ b/devolutions-gateway/tests/dvls_compatibility.rs @@ -301,6 +301,10 @@ mod as_of_v2022_3_0_0 { JrlRead, #[serde(rename = "gateway.config.write")] ConfigWrite, + #[serde(rename = "gateway.tasks.recording-ai-analysis.start")] + RecordingAiAnalysisStart, + #[serde(rename = "gateway.tasks.recording-ai-analysis.read")] + RecordingAiAnalysisRead, } fn dvls_access_scope() -> impl Strategy { @@ -312,6 +316,8 @@ mod as_of_v2022_3_0_0 { Just(DvlsAccessScope::JrlRead), Just(DvlsAccessScope::ConfigWrite), Just(DvlsAccessScope::ConfigWrite), + Just(DvlsAccessScope::RecordingAiAnalysisStart), + Just(DvlsAccessScope::RecordingAiAnalysisRead), ] } From 3b58f94bff6290d6deddc36b8354abe15fcbc007 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 16:41:35 -0400 Subject: [PATCH 04/53] fix(dgw): store provisioner task jobs before reporting them queued A Task was reported created once its job reached the channel of the job queue task, which gives up after a few failed writes without telling anyone, so a Task could be created without a job. The runner now writes jobs to the job queue database itself, through ProvisionerTaskJobQueue, and only reports a Task created once its job is stored. At startup, a Task whose new job cannot be stored is failed as not queued, instead of waiting for its deadline with no job. A Task is only cleaned up once its end is recorded: when recording a failure fails, the job queue tries the attempt again; when recording a success fails, the Task keeps its key and working files and fails at its deadline, since running the work again would add a second log. Waiting for an earlier AI analysis of the same session now stops when the attempt is cancelled, so the attempt timeout also covers the wait. --- devolutions-gateway/src/job_queue.rs | 21 ++- devolutions-gateway/src/lib.rs | 2 +- .../src/provisioner_tasks/mod.rs | 5 +- .../recording_ai_analysis/mod.rs | 6 +- .../recording_ai_analysis/tests.rs | 2 +- .../src/provisioner_tasks/runner/mod.rs | 161 +++++++++++------- .../src/provisioner_tasks/runner/tests.rs | 52 +++--- devolutions-gateway/src/service.rs | 4 +- devolutions-gateway/tests/ai_analysis.rs | 2 +- 9 files changed, 156 insertions(+), 99 deletions(-) diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 3e4b611fc..60db0a367 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -11,7 +11,7 @@ use job_queue_libsql::libsql; use time::OffsetDateTime; use tokio::sync::{Notify, mpsc}; -use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskRunner}; +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; @@ -90,12 +90,21 @@ impl JobQueueCtx { }) } - /// Writes a job to the queue right away, without going through [`JobQueueHandle`]. - /// - /// For startup, before the job runner claims anything. - pub async fn push(&self, job: T) -> anyhow::Result<()> { + /// 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) -> anyhow::Result<()> { let job: DynJob = Box::new(job); - self.queue.push_job(&job, None).await + self.0.push_job(&job, None).await } } diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index 6a6ca3783..c39576552 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -102,7 +102,7 @@ impl DgwState { let task_store = Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store).await?); let provisioner_tasks = provisioner_tasks::task_runner( task_store, - job_queue_handle.clone(), + Arc::new(job_queue_handle.clone()), conf_handle.clone(), recording_manager_handle.clone(), provisioning.clone(), diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index b8bc845eb..ccb1d4938 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -8,16 +8,15 @@ pub mod runner; use provisioner_task::DynProvisionerTaskStore; use self::recording_ai_analysis::RecordingAiAnalysis; -use self::runner::ProvisionerTaskRunner; +use self::runner::{DynProvisionerTaskJobQueue, ProvisionerTaskRunner}; use crate::config::ConfHandle; -use crate::job_queue::JobQueueHandle; use crate::provisioning::ProvisioningStore; use crate::recording::RecordingMessageSender; /// The provisioner task runner, with every provisioner task Gateway knows. pub fn task_runner( store: DynProvisionerTaskStore, - job_queue: JobQueueHandle, + job_queue: DynProvisionerTaskJobQueue, conf_handle: ConfHandle, recordings: RecordingMessageSender, provisioning: ProvisioningStore, diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 48b07b0b4..dea325249 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -329,7 +329,11 @@ impl RecordingAiAnalysis { .client(&conf, &api_key) .map_err(|error| AttemptError::Permanent(error.to_string()))?; - let _session = lock_session(session_id).await; + // Waiting for an earlier analysis of the session counts in the attempt time, so a timeout cancels it too. + let _session = tokio::select! { + session = lock_session(session_id) => session, + () = attempt.cancel.cancelled() => return Err(AttemptError::Cancelled), + }; let recording = self .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 index 3b07f8bbd..fdcee3db4 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -681,7 +681,7 @@ async fn slow_attempt_is_retried() { task.attempt_timeout = Duration::from_millis(500); harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( Arc::clone(harness.state.provisioner_tasks.store()), - harness.state.job_queue_handle.clone(), + Arc::new(harness.state.job_queue_handle.clone()), ) .register(task) .build(); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 4e43724f2..c3b848d39 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -13,7 +13,6 @@ mod tests; use std::collections::HashMap; -use std::future::Future; use std::sync::Arc; use anyhow::Context as _; @@ -30,6 +29,23 @@ use uuid::Uuid; use super::AttemptError; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; + +/// Where the runner queues the jobs of its records. +#[async_trait] +pub trait ProvisionerTaskJobQueue: Send + Sync { + /// Queues `job`, and returns once it is stored, so a record never waits for a job that was lost. + async fn push(&self, job: ProvisionerTaskJob) -> 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) -> anyhow::Result<()> { + self.enqueue(job).await + } +} use crate::token::AccessScope; /// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. @@ -135,14 +151,14 @@ pub struct ProvisionerTaskRunner { struct Inner { store: DynProvisionerTaskStore, - job_queue: JobQueueHandle, + job_queue: DynProvisionerTaskJobQueue, /// Kind → the task that runs records of this kind. tasks: HashMap<&'static str, Arc>, } pub struct ProvisionerTaskRunnerBuilder { store: DynProvisionerTaskStore, - job_queue: JobQueueHandle, + job_queue: DynProvisionerTaskJobQueue, tasks: HashMap<&'static str, Arc>, } @@ -171,7 +187,10 @@ fn starting_payload() -> serde_json::Value { } impl ProvisionerTaskRunner { - pub fn builder(store: DynProvisionerTaskStore, job_queue: JobQueueHandle) -> ProvisionerTaskRunnerBuilder { + pub fn builder( + store: DynProvisionerTaskStore, + job_queue: DynProvisionerTaskJobQueue, + ) -> ProvisionerTaskRunnerBuilder { ProvisionerTaskRunnerBuilder { store, job_queue, @@ -221,17 +240,14 @@ impl ProvisionerTaskRunner { return Ok(outcome); }; - if let Err(error) = on_created(record) { - error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to prepare the Task"); - self.fail(task.as_ref(), id, failure("not queued", &format!("{error:#}"), 0)) - .await; - return Err(error); - } + let queued = match on_created(record) { + Ok(()) => self.enqueue(id, job_token).await, + Err(error) => Err(error.context("prepare the Task")), + }; - if let Err(error) = self.enqueue(id, job_token).await { + if let Err(error) = queued { error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to queue the Task"); - self.fail(task.as_ref(), id, failure("not queued", &format!("{error:#}"), 0)) - .await; + self.fail_not_queued(task.as_ref(), id, &error, 0).await; return Err(error); } @@ -240,14 +256,11 @@ impl ProvisionerTaskRunner { Ok(outcome) } - /// Hands every unfinished record to a new job queued with `push`: the jobs queued before stop at their next attempt. + /// Hands every unfinished record to a new job: the jobs queued before stop at their next attempt. /// - /// Call it at startup, before the job runner claims any job, so no earlier job is running an attempt meanwhile. - pub async fn resume(&self, push: F) -> anyhow::Result<()> - where - F: Fn(ProvisionerTaskJob) -> Fut, - Fut: Future>, - { + /// Call it at startup, before the job runner claims any job, so no earlier job is running 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?; @@ -258,20 +271,20 @@ impl ProvisionerTaskRunner { continue; } - let job = ProvisionerTaskJob { - task_id: record.id, - job_token, - runner: self.clone(), - }; - - match push(job).await { + match self.enqueue(record.id, job_token).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; it fails at its deadline" - ), + 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, &error, record.attempts) + .await; + } + } } } @@ -291,7 +304,7 @@ impl ProvisionerTaskRunner { async fn enqueue(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { self.inner .job_queue - .enqueue(ProvisionerTaskJob { + .push(ProvisionerTaskJob { task_id, job_token, runner: self.clone(), @@ -374,24 +387,28 @@ impl ProvisionerTaskRunner { }; match outcome { - Ok(output) => { - // An error here must not ask for a retry: the work is done. - match store.succeed(task_id, output, OffsetDateTime::now_utc()).await { - Ok(true) => info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"), - Ok(false) => 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" - ), - Err(error) => error!( - task.id = %task_id, - error = format!("{error:#}"), - "Failed to record the Task result" - ), + 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).await; } - - self.end(task.as_ref(), task_id).await; - } + // No retry: the work is done, and doing it again would repeat it, such as adding a second log. Nothing is + // dropped either, since the record still runs; it fails at its deadline. + Err(error) => error!( + task.id = %task_id, + error = format!("{error:#}"), + "Failed to record the Task result; the Task fails at its deadline" + ), + }, Err(AttemptError::Transient(details)) if attempts < MAX_ATTEMPTS => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task attempt failed; it will be retried"); anyhow::bail!("Task attempt failed: {details}"); @@ -403,17 +420,17 @@ impl ProvisionerTaskRunner { task_id, failure("attempts exhausted", &details, attempts), ) - .await; + .await?; } Err(AttemptError::Permanent(details)) => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task failed"); self.fail(task.as_ref(), task_id, failure("permanent error", &details, attempts)) - .await; + .await?; } Err(AttemptError::Failed { reason, details }) => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, reason, %details, "Task failed"); self.fail(task.as_ref(), task_id, failure(reason, &details, attempts)) - .await; + .await?; } Err(AttemptError::Cancelled) => unreachable!("a cancelled attempt is reported as a timeout above"), } @@ -421,12 +438,40 @@ impl ProvisionerTaskRunner { Ok(()) } - async fn fail(&self, task: &dyn AnyProvisionerTask, task_id: Uuid, payload: serde_json::Value) { - if let Err(error) = self.inner.store.fail(task_id, payload, OffsetDateTime::now_utc()).await { - error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure"); - } + /// Records the failure, then drops what the record kept while running. + /// + /// 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, + payload: serde_json::Value, + ) -> anyhow::Result<()> { + self.inner + .store + .fail(task_id, payload, OffsetDateTime::now_utc()) + .await + .context("record the Task failure")?; self.end(task, task_id).await; + + Ok(()) + } + + /// Fails a record no job will run. + async fn fail_not_queued( + &self, + task: &dyn AnyProvisionerTask, + task_id: Uuid, + error: &anyhow::Error, + attempts: u32, + ) { + let payload = failure("not queued", &format!("{error:#}"), attempts); + + if let Err(error) = self.fail(task, task_id, payload).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure; it fails at its deadline"); + } } async fn end(&self, task: &dyn AnyProvisionerTask, task_id: Uuid) { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index f063a771f..6cbd047ae 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -1,6 +1,5 @@ use std::sync::atomic::{AtomicUsize, Ordering}; -use job_queue::Job as _; use parking_lot::Mutex; use provisioner_task::ProvisionerTaskState; @@ -76,7 +75,7 @@ impl Harness { let runs = Arc::new(AtomicUsize::new(0)); let ended = Arc::new(Mutex::new(Vec::new())); - let runner = ProvisionerTaskRunner::builder(store, job_queue) + let runner = ProvisionerTaskRunner::builder(store, Arc::new(job_queue)) .register(Echo { outcome, runs: Arc::clone(&runs), @@ -183,45 +182,46 @@ async fn a_superseded_job_runs_nothing_once_the_task_is_resumed() { let id = harness.create(7).await; let mut stale = harness.queued_job(); - let resumed = Mutex::new(Vec::new()); - harness - .runner - .resume(|job| { - resumed.lock().push(job); - async { Ok(()) } - }) - .await - .expect("resume"); - let mut resumed = resumed.into_inner(); - assert_eq!(resumed.len(), 1, "the unfinished Task gets a new 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[0].run().await.expect("the new job runs"); + 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]); - let queued_again = AtomicUsize::new(0); - harness - .runner - .resume(|_| { - queued_again.fetch_add(1, Ordering::SeqCst); - async { Ok(()) } - }) - .await - .expect("resume again"); - assert_eq!( - queued_again.load(Ordering::SeqCst), - 0, + harness.runner.resume().await.expect("resume again"); + assert!( + harness.job_queue_rx.try_recv().is_err(), "a finished Task gets no new job" ); } +#[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["reason"], "not queued"); + assert_eq!(*harness.ended.lock(), [id]); +} + #[tokio::test] async fn a_panicking_attempt_fails_the_task() { let mut harness = Harness::new(|_| panic!("the attempt panics")).await; diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index eea186dca..efc0ea707 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -352,7 +352,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { let provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( task_store, - job_queue_ctx.job_queue_handle.clone(), + Arc::new(job_queue_ctx.durable()), conf_handle.clone(), recording_manager_handle.clone(), provisioning.clone(), @@ -360,7 +360,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. provisioner_tasks - .resume(|job| job_queue_ctx.push(job)) + .resume() .await .context("failed to queue the unfinished provisioner tasks again")?; diff --git a/devolutions-gateway/tests/ai_analysis.rs b/devolutions-gateway/tests/ai_analysis.rs index d69889bf2..79e28edb6 100644 --- a/devolutions-gateway/tests/ai_analysis.rs +++ b/devolutions-gateway/tests/ai_analysis.rs @@ -107,7 +107,7 @@ impl Harness { ); state.provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( Arc::clone(&task_store), - state.job_queue_handle.clone(), + Arc::new(state.job_queue_handle.clone()), state.conf_handle.clone(), state.recordings.clone(), state.provisioning.clone(), From 3cb3ca36ae0e1a8e3ad44241178e60681b96903d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:03:51 -0400 Subject: [PATCH 05/53] fix(dgw): keep the stored job of a provisioner task on restart At startup, every unfinished Task got a new job, even when the job it had was still stored in the job queue and would run anyway, leaving a stale job behind. The runner now reads the provisioner task jobs the job queue still holds, and only gives a new job to a Task left without one, such as a Task created just before a crash or whose job was dropped. The job queue gains `job_definitions` to list the stored jobs of a name that may still run. --- Cargo.lock | 1 + crates/job-queue-libsql/Cargo.toml | 3 + crates/job-queue-libsql/src/lib.rs | 73 +++++++++++++++++++ crates/job-queue/src/lib.rs | 3 + devolutions-gateway/src/job_queue.rs | 4 + .../src/provisioner_tasks/runner/mod.rs | 33 ++++++++- .../src/provisioner_tasks/runner/tests.rs | 69 ++++++++++++++++++ 7 files changed, 183 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dcebbe13a..020845f37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3997,6 +3997,7 @@ dependencies = [ "libsql", "serde", "time", + "tokio 1.52.3", "tracing", "typed-builder", "ulid", 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..04f76b63d 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,55 @@ 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(()) + } + } + + #[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/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 60db0a367..42e4f3c65 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -106,6 +106,10 @@ impl ProvisionerTaskJobQueue for DurableJobQueue { let job: DynJob = Box::new(job); self.0.push_job(&job, None).await } + + async fn stored_jobs(&self) -> anyhow::Result> { + self.0.job_definitions(ProvisionerTaskJob::NAME).await + } } impl JobQueueHandle { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index c3b848d39..9e7b08968 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -35,6 +35,9 @@ use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; pub trait ProvisionerTaskJobQueue: Send + Sync { /// Queues `job`, and returns once it is stored, so a record never waits for a job that was lost. async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()>; + + /// Lists the JSON definitions of the provisioner task jobs already stored that may still run. + async fn stored_jobs(&self) -> anyhow::Result>; } pub type DynProvisionerTaskJobQueue = Arc; @@ -45,6 +48,11 @@ impl ProvisionerTaskJobQueue for JobQueueHandle { async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { 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) -> anyhow::Result> { + Ok(Vec::new()) + } } use crate::token::AccessScope; @@ -256,15 +264,34 @@ impl ProvisionerTaskRunner { Ok(outcome) } - /// Hands every unfinished record to a new job: the jobs queued before stop at their next attempt. + /// Hands every unfinished record that has no job left to a new job, such as a record created just before a crash. /// - /// Call it at startup, before the job runner claims any job, so no earlier job is running an attempt meanwhile. A - /// record whose new job cannot be queued is failed, since no job is left to run it. + /// 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() + .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 record + .job_token + .is_some_and(|job_token| stored.contains(&(record.id, job_token))) + { + debug!(task.id = %record.id, task.kind = %record.kind, "Task still has its job"); + continue; + } + let job_token = Uuid::new_v4(); if !self.inner.store.replace_job_token(record.id, job_token, now).await? { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 6cbd047ae..67e80b2d1 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -222,6 +222,75 @@ async fn a_task_that_cannot_be_queued_again_fails() { 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) -> anyhow::Result<()> { + self.0.lock().push(job_queue::Job::write_json(&job)?); + Ok(()) + } + + async fn stored_jobs(&self) -> anyhow::Result> { + Ok(self.0.lock().clone()) + } +} + +#[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)), + runs: Arc::default(), + ended: 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" + ); +} + #[tokio::test] async fn a_panicking_attempt_fails_the_task() { let mut harness = Harness::new(|_| panic!("the attempt panics")).await; From 4400ea0656fc38d7199219c41e0547728cc572bf Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:31:15 -0400 Subject: [PATCH 06/53] docs(dgw): add the provisioner tasks intent --- .../src/provisioner_tasks/INTENT.md | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 devolutions-gateway/src/provisioner_tasks/INTENT.md diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md new file mode 100644 index 000000000..e96647acd --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -0,0 +1,25 @@ +# 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/README.md`. + + +# Invariants + +## Management +`ProvisionerTaskRunner` MUST be the only place touches both the Task records and the job queue. +A provisioner task only defines it's work: what one attempt does, and what to drop once its task ends. +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. + + +## Naming +`ProvisionerTask` is the work itself; `ProvisionerTaskRecord` is what is stored about it. + +## Secrets +Secrets a Task needs, such as an API key, never reach the disk: not `gateway.db`, not the job, not the logs. +They are kept encrypted in memory, in the provisioning store, 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. \ No newline at end of file From 18572ef2c212e6bda1cb35ebc2e1f37d4a42ff1d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:40:42 -0400 Subject: [PATCH 07/53] docs(dgw): point the provisioner tasks intent to the task store crate --- devolutions-gateway/src/provisioner_tasks/INTENT.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index e96647acd..b7e8f1245 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -1,5 +1,5 @@ # 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/README.md`. +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 From d2c11a05757f020c83ea967f2b81ad93ec26364b Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:46:09 -0400 Subject: [PATCH 08/53] docs(dgw): add the jrec intent --- devolutions-gateway/src/api/jrec.intent.md | 10 ++++++++++ 1 file changed, 10 insertions(+) create mode 100644 devolutions-gateway/src/api/jrec.intent.md diff --git a/devolutions-gateway/src/api/jrec.intent.md b/devolutions-gateway/src/api/jrec.intent.md new file mode 100644 index 000000000..e65cc4955 --- /dev/null +++ b/devolutions-gateway/src/api/jrec.intent.md @@ -0,0 +1,10 @@ +# 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. \ No newline at end of file From 3cc6fe9c086e2a790794074ad3642cfab0a3b5e5 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:57:41 -0400 Subject: [PATCH 09/53] docs(dgw): expand the provisioner tasks intent --- devolutions-gateway/src/provisioner_tasks/INTENT.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index b7e8f1245..346de3c4b 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -5,11 +5,16 @@ Provisioner tasks are tasks issued for Gateway to execute. Tasks's state are pes # Invariants ## Management -`ProvisionerTaskRunner` MUST be the only place touches both the Task records and the job queue. -A provisioner task only defines it's work: what one attempt does, and what to drop once its task ends. +`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. + ## Naming `ProvisionerTask` is the work itself; `ProvisionerTaskRecord` is what is stored about it. From fc1207fa0180c0dc6893bc891464acc97b24256c Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:56 -0400 Subject: [PATCH 10/53] test(dgw): keep the DVLS 2022.3 scope snapshot frozen `as_of_v2022_3_0_0` lists what DVLS 2022.3 sends; later scopes such as `gateway.recording.delete` were never added there. The AI analysis scopes are covered by the `ai_analysis` tests. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/tests/dvls_compatibility.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/devolutions-gateway/tests/dvls_compatibility.rs b/devolutions-gateway/tests/dvls_compatibility.rs index fa0d080cc..e727c8e01 100644 --- a/devolutions-gateway/tests/dvls_compatibility.rs +++ b/devolutions-gateway/tests/dvls_compatibility.rs @@ -301,10 +301,6 @@ mod as_of_v2022_3_0_0 { JrlRead, #[serde(rename = "gateway.config.write")] ConfigWrite, - #[serde(rename = "gateway.tasks.recording-ai-analysis.start")] - RecordingAiAnalysisStart, - #[serde(rename = "gateway.tasks.recording-ai-analysis.read")] - RecordingAiAnalysisRead, } fn dvls_access_scope() -> impl Strategy { @@ -316,8 +312,6 @@ mod as_of_v2022_3_0_0 { Just(DvlsAccessScope::JrlRead), Just(DvlsAccessScope::ConfigWrite), Just(DvlsAccessScope::ConfigWrite), - Just(DvlsAccessScope::RecordingAiAnalysisStart), - Just(DvlsAccessScope::RecordingAiAnalysisRead), ] } From 232d4661cce0600ce9987969432ec1a8b7a99c86 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:56 -0400 Subject: [PATCH 11/53] fix(dgw): keep the AI provider URL out of Debug output `AiSettings` and `AiAccess` (and so `StartRequest`) derived `Debug`, which prints the base URL that must never reach the logs. Nothing logs them today; this keeps it that way. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/provisioner_tasks/ai.rs | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index 6ae4bf015..089805ee6 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -43,7 +43,7 @@ impl From for Provider { } /// How to reach an AI provider, except the API key. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Clone, PartialEq, Eq)] pub struct AiSettings { pub provider: AiProvider, pub model: String, @@ -53,6 +53,17 @@ pub struct AiSettings { 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 { @@ -92,12 +103,21 @@ impl fmt::Display for ClientError { } /// What a Task needs to reach its AI provider and must never write to disk. -#[derive(Debug, Clone)] +#[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")] From 576a537b6172bca494015904cc335bc2f0270efd Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 12/53] fix(dgw): bound the AI analysis progress percent in the OpenAPI schema The server only ever sends 0 to 100, as documented. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/openapi.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 6ef5e0f66..7163baad4 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -167,6 +167,7 @@ struct AiAnalysisRunningPayload { /// Work the step has in all, in the same unit as `done`. total: Option, /// Part of the step done, from 0 to 100. + #[schema(maximum = 100)] percent: Option, /// Short sentence for people following the analysis, such as `asking the AI: 2 of 5 chunks described`. message: Option, From 34728dc2eddce1fdbfe4dfd3e1184549426278f5 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 13/53] fix(dgw): keep Gateway starting when provisioner tasks cannot be resumed A failure to list or hand over the unfinished Tasks stopped the whole service. Unfinished Tasks still fail at their deadline, so the error is logged and startup goes on; one Task that cannot get a new job no longer stops the others. Also follows the provisioner task store changes below this PR: `update_running` takes `now` and `job_token` is no longer optional. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../recording_ai_analysis/tests.rs | 8 ++---- .../src/provisioner_tasks/runner/mod.rs | 27 +++++++++++++------ .../src/provisioner_tasks/runner/tests.rs | 2 +- devolutions-gateway/src/service.rs | 11 +++++--- 4 files changed, 29 insertions(+), 19 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index fdcee3db4..b4157ac7e 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -120,8 +120,7 @@ impl Harness { .await .expect("read task") .expect("task exists") - .job_token - .expect("job token"); + .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") @@ -331,10 +330,7 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { assert_eq!(job.name(), ProvisionerTaskJob::NAME); assert_eq!( job.write_json().expect("job JSON"), - format!( - r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, - task.job_token.expect("job token") - ) + format!(r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, task.job_token) ); assert!(schedule_for.is_none()); } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 9e7b08968..44476192c 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -29,6 +29,7 @@ use uuid::Uuid; use super::AttemptError; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; +use crate::token::AccessScope; /// Where the runner queues the jobs of its records. #[async_trait] @@ -54,7 +55,6 @@ impl ProvisionerTaskJobQueue for JobQueueHandle { Ok(Vec::new()) } } -use crate::token::AccessScope; /// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. /// @@ -107,7 +107,11 @@ pub struct Attempt { impl Attempt { /// Replaces the payload of the running record, so whoever asks about it sees how the attempt goes. pub async fn report_progress(&self, payload: serde_json::Value) { - if let Err(error) = self.store.update_running(self.task_id, payload).await { + if let Err(error) = self + .store + .update_running(self.task_id, payload, OffsetDateTime::now_utc()) + .await + { warn!(task.id = %self.task_id, error = format!("{error:#}"), "Failed to record the Task progress"); } } @@ -284,18 +288,25 @@ impl ProvisionerTaskRunner { .collect::>(); for record in records { - if record - .job_token - .is_some_and(|job_token| stored.contains(&(record.id, job_token))) - { + 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(); - if !self.inner.store.replace_job_token(record.id, job_token, now).await? { - continue; + 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).await { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 67e80b2d1..4c4fc8d33 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -370,6 +370,6 @@ async fn the_job_is_read_back_from_its_json() { let read = ProvisionerTaskJob::read_json(&json, harness.runner.clone()).expect("valid job"); assert_eq!(read.task_id, id); - assert_eq!(Some(read.job_token), harness.task(id).await.job_token); + assert_eq!(read.job_token, harness.task(id).await.job_token); assert_eq!(job.name(), ProvisionerTaskJob::NAME); } diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index efc0ea707..6a33137f7 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -359,10 +359,13 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { ); // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. - provisioner_tasks - .resume() - .await - .context("failed to queue the unfinished provisioner tasks again")?; + // 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(), From 79f65b076058320f84ac511e710239bb37066142 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 14/53] style(dgw): separate items and group imports in the AI analysis code Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/jrec/mod.rs | 1 + devolutions-gateway/src/extract.rs | 2 ++ 2 files changed, 3 insertions(+) diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index bc9c86cd9..728ece393 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -624,6 +624,7 @@ pub(crate) async fn get_ai_analysis( Ok(Json(TaskInfo::from(task))) } + /// Lists all recordings stored on this instance #[cfg_attr(feature = "openapi", utoipa::path( get, diff --git a/devolutions-gateway/src/extract.rs b/devolutions-gateway/src/extract.rs index 0212a79d1..224c9382a 100644 --- a/devolutions-gateway/src/extract.rs +++ b/devolutions-gateway/src/extract.rs @@ -342,6 +342,7 @@ where } } } + /// A scope token allowed to start the provisioner task `T`: `*`, or the start scope of `T`. pub struct ProvisionerTaskStartScope(core::marker::PhantomData); @@ -362,6 +363,7 @@ where } } } + #[derive(Clone, Copy)] pub struct RecordingsSearchScope; From 369a08667433ce750284e9c1f2942270dc3b731b Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 09:48:45 -0400 Subject: [PATCH 15/53] fix(dgw): ignore unknown fields in the AI analysis request A newer DVLS may send a new optional field to an older Gateway. With deny_unknown_fields the older Gateway refused the whole request with 400, so the analysis never ran. Every other Gateway endpoint ignores unknown fields; this one now does too. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/jrec/mod.rs | 2 +- devolutions-gateway/tests/ai_analysis.rs | 16 ++++++++++++---- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index 728ece393..d85b9b05e 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -441,7 +441,7 @@ async fn delete_recording(recording_path: &Utf8Path) -> anyhow::Result<()> { /// Settings of an AI analysis #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] #[derive(Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] +#[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, diff --git a/devolutions-gateway/tests/ai_analysis.rs b/devolutions-gateway/tests/ai_analysis.rs index 79e28edb6..d471a42fb 100644 --- a/devolutions-gateway/tests/ai_analysis.rs +++ b/devolutions-gateway/tests/ai_analysis.rs @@ -326,15 +326,24 @@ async fn start_without_recording_returns_404() { assert_eq!(status, StatusCode::NOT_FOUND); } +#[tokio::test] +async fn start_ignores_unknown_fields() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let mut body = request(Uuid::new_v4(), "gpt-test"); + body["temperature"] = json!(0.5); + + let (status, created) = harness.post(session_id, &body).await; + + assert_eq!(status, StatusCode::ACCEPTED, "{created}"); +} + #[tokio::test] async fn invalid_bodies_and_settings_return_400() { let harness = Harness::new().await; let session_id = harness.write_session(); let task_id = Uuid::new_v4(); - let mut unknown_field = request(task_id, "gpt-test"); - unknown_field["temperature"] = json!(0.5); - let mut no_key = request(task_id, "gpt-test"); no_key.as_object_mut().unwrap().remove("apiKey"); @@ -351,7 +360,6 @@ async fn invalid_bodies_and_settings_return_400() { unknown_provider["provider"] = json!("skynet"); for (case, body) in [ - ("unknown field", unknown_field.to_string()), ("missing API key", no_key.to_string()), ("empty API key", empty_key.to_string()), ("empty model", request(task_id, " ").to_string()), From b5bebdd50e02fe365e1dbc73324e7479df1c0684 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:02:58 -0400 Subject: [PATCH 16/53] fix(dgw): start no AI analysis attempt that cannot finish before the deadline An attempt started near the 2-hour deadline could keep calling the AI for 30 more minutes, add its log to the recording, and still leave the Task failed as timed out. Attempts now start only during the first 2 hours, and the deadline is one attempt timeout after that, so an attempt started in time can finish and its result counts. The runner fails a Task that has no time left for a full attempt as timed out, without running it. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../recording_ai_analysis/mod.rs | 8 ++-- .../recording_ai_analysis/tests.rs | 31 +++++--------- .../src/provisioner_tasks/runner/mod.rs | 30 +++++++++++--- .../src/provisioner_tasks/runner/tests.rs | 41 ++++++++++++++++++- 4 files changed, 78 insertions(+), 32 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index dea325249..f29828ca6 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -37,8 +37,8 @@ use crate::token::AccessScope; pub const KIND: &str = "recording-ai-analysis"; -/// Time an AI analysis gets from its creation to its end, waiting and retries included. -pub const TASK_DEADLINE: time::Duration = time::Duration::hours(2); +/// Time from the creation of an AI analysis during which its attempts may start, waiting and retries included. +pub const ATTEMPT_START_WINDOW: time::Duration = time::Duration::hours(2); /// Time one attempt of an AI analysis gets. pub const ATTEMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60); @@ -260,8 +260,8 @@ impl ProvisionerTask for RecordingAiAnalysis { type Output = Analysis; - fn deadline(&self) -> time::Duration { - TASK_DEADLINE + fn attempt_start_window(&self) -> time::Duration { + ATTEMPT_START_WINDOW } fn attempt_timeout(&self) -> std::time::Duration { diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index b4157ac7e..f0916e583 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -321,7 +321,10 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { 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, TASK_DEADLINE); + assert_eq!( + task.deadline_at - task.created_at, + ATTEMPT_START_WINDOW + ATTEMPT_TIMEOUT + ); let access = harness.access(task_id).expect("key is kept"); assert_eq!(access.base_url.expect("base URL").as_str(), "https://ai.example/v1/"); @@ -694,10 +697,10 @@ async fn slow_attempt_is_retried() { } #[tokio::test] -async fn log_added_after_the_deadline_leaves_the_task_failed() { +async fn task_without_time_for_an_attempt_fails_without_calling_the_ai() { let harness = Harness::new().await; let session_id = harness.write_session(); - let (settings, _requests) = spawn_ai_provider_after(Duration::from_secs(4), |_| 0).await; + let (settings, requests) = spawn_ai_provider(|_| 0).await; let task_id = Uuid::new_v4(); let now = OffsetDateTime::now_utc(); harness @@ -706,32 +709,18 @@ async fn log_added_after_the_deadline_leaves_the_task_failed() { session_id, "openai-compatible", now, - now + time::Duration::seconds(2), + now + ATTEMPT_TIMEOUT - time::Duration::minutes(1), ) .await; harness.insert_key(task_id, settings.base_url, now + time::Duration::hours(1)); - let mut job = harness.job(task_id).await; - let job = tokio::spawn(async move { job.run().await }); - - // Reading the Task past its deadline, as a status request does, ends it while the AI is still answering. - tokio::time::sleep(Duration::from_secs(3)).await; - assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Failed); - - job.await.expect("job task").expect("no retry"); + harness.job(task_id).await.run().await.expect("no retry"); let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); assert_eq!(task.payload["reason"], "timed out"); - assert!( - harness - .recording_path - .join(session_id.to_string()) - .join("ai-analysis-0.slog") - .exists(), - "the log stays in the session" - ); - assert!(!harness.workspace(session_id).exists()); + assert_eq!(task.attempts, 0); + assert!(requests.lock().is_empty()); 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 index 44476192c..a5121cc59 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -77,8 +77,10 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// The payload of a `succeeded` record. type Output: Serialize + Send; - /// Time from the creation of a record to its end, waiting and retries included. - fn deadline(&self) -> time::Duration; + /// Time from the creation of a record during which its attempts may start, the first one and retries alike. + /// + /// The record's deadline comes one [`Self::attempt_timeout`] later, so an attempt started in time can finish. + fn attempt_start_window(&self) -> time::Duration; /// Time one attempt gets before it is cancelled and retried. fn attempt_timeout(&self) -> std::time::Duration; @@ -120,7 +122,7 @@ impl Attempt { /// [`ProvisionerTask`] with its types erased, so the runner can hold every task together. #[async_trait] trait AnyProvisionerTask: Send + Sync { - fn deadline(&self) -> time::Duration; + fn attempt_start_window(&self) -> time::Duration; fn attempt_timeout(&self) -> std::time::Duration; @@ -131,8 +133,8 @@ trait AnyProvisionerTask: Send + Sync { #[async_trait] impl AnyProvisionerTask for T { - fn deadline(&self) -> time::Duration { - ProvisionerTask::deadline(self) + fn attempt_start_window(&self) -> time::Duration { + ProvisionerTask::attempt_start_window(self) } fn attempt_timeout(&self) -> std::time::Duration { @@ -241,7 +243,7 @@ impl ProvisionerTaskRunner { kind: T::KIND.to_owned(), target, params, - deadline_at: now + task.deadline(), + deadline_at: now + task.attempt_start_window() + task.attempt_timeout(), job_token, }, now, @@ -355,6 +357,22 @@ impl ProvisionerTaskRunner { let store = &self.inner.store; let now = OffsetDateTime::now_utc(); + // An attempt that cannot finish before the deadline would do work the record never keeps, such as paid AI calls. + if let Some(record) = store.get(task_id, now).await? + && record.job_token == job_token + && !record.state.is_finished() + && let Some(task) = self.task(&record.kind) + && now + task.attempt_timeout() > record.deadline_at + { + warn!(task.id = %task_id, task.kind = %record.kind, attempts = record.attempts, "Task timed out before its next attempt"); + let payload = json!({ + "reason": "timed out", + "details": { "lastPayload": record.payload }, + "attempts": record.attempts, + }); + return self.fail(task.as_ref(), task_id, payload).await; + } + let record = match store.start_attempt(task_id, job_token, starting_payload(), now).await? { AttemptStart::Started(record) => record, AttemptStart::Finished(record) => { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 4c4fc8d33..92f4666ac 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -33,7 +33,7 @@ impl ProvisionerTask for Echo { type Output = serde_json::Value; - fn deadline(&self) -> time::Duration { + fn attempt_start_window(&self) -> time::Duration { time::Duration::hours(1) } @@ -360,6 +360,45 @@ async fn invalid_parameters_and_unknown_kinds_fail_the_task() { assert_eq!(harness.runs.load(Ordering::SeqCst), 0); } +#[tokio::test] +async fn an_attempt_starts_only_if_it_can_finish_before_the_deadline() { + let harness = Harness::new(|value| Ok(json!(value))).await; + let store = harness.runner.store(); + let now = OffsetDateTime::now_utc(); + + // The echo attempt timeout is 5 seconds. + for (left, state) in [(6, ProvisionerTaskState::Succeeded), (4, ProvisionerTaskState::Failed)] { + let id = Uuid::new_v4(); + let job_token = Uuid::new_v4(); + store + .create( + NewProvisionerTaskRecord { + id, + kind: KIND.to_owned(), + target: id.to_string(), + params: json!({ "value": 1 }), + deadline_at: now + time::Duration::seconds(left), + 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, state, "{left} seconds left"); + if state == ProvisionerTaskState::Failed { + assert_eq!(task.payload["reason"], "timed out"); + assert_eq!(task.attempts, 0); + assert!(harness.ended.lock().contains(&id)); + } + } + + assert_eq!(harness.runs.load(Ordering::SeqCst), 1); +} + #[tokio::test] async fn the_job_is_read_back_from_its_json() { let mut harness = Harness::new(|value| Ok(json!(value))).await; From c5f72c84595e7b882cf1ed823e12813d769bb968 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:23:39 -0400 Subject: [PATCH 17/53] refactor(dgw): keep only step, done and total in the AI analysis progress `percent` is `done / total`, and `message` is English that clients translate anyway, so neither belongs in the API. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/openapi.rs | 5 ----- .../src/provisioner_tasks/recording_ai_analysis/mod.rs | 4 +--- .../src/provisioner_tasks/recording_ai_analysis/tests.rs | 2 -- 3 files changed, 1 insertion(+), 10 deletions(-) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 7163baad4..65128cca5 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -166,11 +166,6 @@ struct AiAnalysisRunningPayload { done: Option, /// Work the step has in all, in the same unit as `done`. total: Option, - /// Part of the step done, from 0 to 100. - #[schema(maximum = 100)] - percent: Option, - /// Short sentence for people following the analysis, such as `asking the AI: 2 of 5 chunks described`. - message: Option, } /// Payload of a `succeeded` AI analysis Task: the log it added to the session diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index f29828ca6..39ca2daf3 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -359,7 +359,7 @@ impl RecordingAiAnalysis { let analysis = loop { match events.recv().await { Some(RecordingAnalysisEvent::Progress(progress)) => { - // Reading sends an event every half second: the Task is only updated when the numbers it shows change. + // Reading sends an event every half second, so the Task is updated only when the stage's percent changes. let shown = (progress.stage(), progress.percent()); if recorded != Some(shown) { recorded = Some(shown); @@ -400,8 +400,6 @@ fn progress_payload(progress: &Progress) -> serde_json::Value { "step": progress.stage(), "done": done, "total": total, - "percent": progress.percent(), - "message": progress.to_string(), }) } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index f0916e583..d1845a5ab 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -501,8 +501,6 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { "step": "describing", "done": 1, "total": 3, - "percent": 33, - "message": "asking the AI: 1 of 3 chunks described", }) ); assert!(harness.has_key(task_id), "the key is kept for the retry"); From 645a294b02039d947e13c12c6e9cbd0d5e702905 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:40:54 -0400 Subject: [PATCH 18/53] fix(dgw): wait for the remux of the session before the AI analysis reads it A recording is remuxed 60 seconds after it ends, and the remux replaces the video file. An analysis started right after the session ended read the video before it was remuxed. The analysis now waits until the job queue stores no remux job for a file of the session, checking every second, before it reads the recordings. The job queue lists remux jobs that are scheduled or running, including those stored before a Gateway restart. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/job_queue.rs | 4 +- .../src/provisioner_tasks/mod.rs | 11 ++- .../recording_ai_analysis/mod.rs | 42 +++++++++++- .../recording_ai_analysis/tests.rs | 67 ++++++++++++++++++- .../src/provisioner_tasks/runner/mod.rs | 8 +-- .../src/provisioner_tasks/runner/tests.rs | 3 +- devolutions-gateway/src/recording.rs | 10 ++- 7 files changed, 131 insertions(+), 14 deletions(-) diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 42e4f3c65..9d5d38017 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -107,8 +107,8 @@ impl ProvisionerTaskJobQueue for DurableJobQueue { self.0.push_job(&job, None).await } - async fn stored_jobs(&self) -> anyhow::Result> { - self.0.job_definitions(ProvisionerTaskJob::NAME).await + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + self.0.job_definitions(name).await } } diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index ccb1d4938..0278157fd 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -5,6 +5,8 @@ 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; @@ -21,8 +23,13 @@ pub fn task_runner( recordings: RecordingMessageSender, provisioning: ProvisioningStore, ) -> ProvisionerTaskRunner { - ProvisionerTaskRunner::builder(store, job_queue) - .register(RecordingAiAnalysis::new(conf_handle, recordings, provisioning)) + ProvisionerTaskRunner::builder(store, Arc::clone(&job_queue)) + .register(RecordingAiAnalysis::new( + conf_handle, + recordings, + provisioning, + job_queue, + )) .build() } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 39ca2daf3..c3b65141d 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -27,12 +27,12 @@ use uuid::Uuid; use super::AttemptError; use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; -use super::runner::{Attempt, ProvisionerTask}; +use super::runner::{Attempt, DynProvisionerTaskJobQueue, ProvisionerTask}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; use crate::provisioning::ProvisioningStore; -use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; +use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender, RemuxJob, remux_pending}; use crate::token::AccessScope; pub const KIND: &str = "recording-ai-analysis"; @@ -43,6 +43,9 @@ pub const ATTEMPT_START_WINDOW: time::Duration = time::Duration::hours(2); /// Time one attempt of an AI analysis gets. pub const ATTEMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60); +/// Time between two checks for a remux of the session. +const REMUX_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1); + const KEY_LOST: &str = "the API key is no longer available, such as after a Gateway restart"; /// Tokens counted by the AI provider, as kept in Task results. @@ -226,15 +229,22 @@ pub struct RecordingAiAnalysis { conf_handle: ConfHandle, recordings: RecordingMessageSender, provisioning: ProvisioningStore, + job_queue: DynProvisionerTaskJobQueue, attempt_timeout: std::time::Duration, } impl RecordingAiAnalysis { - pub fn new(conf_handle: ConfHandle, recordings: RecordingMessageSender, provisioning: ProvisioningStore) -> Self { + pub fn new( + conf_handle: ConfHandle, + recordings: RecordingMessageSender, + provisioning: ProvisioningStore, + job_queue: DynProvisionerTaskJobQueue, + ) -> Self { Self { conf_handle, recordings, provisioning, + job_queue, attempt_timeout: ATTEMPT_TIMEOUT, } } @@ -345,6 +355,8 @@ impl RecordingAiAnalysis { analysis_manifest(&recording.manifest).map_err(|error| AttemptError::Permanent(error.to_string()))?; let manifest_path = recording.manifest_path(); + self.wait_for_remux(attempt, &recording.dir).await?; + let mut request = client .analyze_recording(&manifest, &manifest_path) .cancellation(attempt.cancel.clone()); @@ -392,6 +404,30 @@ impl RecordingAiAnalysis { } } +impl RecordingAiAnalysis { + /// 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| AttemptError::Transient(format!("failed to list the remux jobs: {error:#}")))?; + + 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::Cancelled), + } + } + } +} + /// The payload of a running AI analysis Task: its progress, for whoever asks about it. fn progress_payload(progress: &Progress) -> serde_json::Value { let (done, total) = progress.done_and_total().unzip(); diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index d1845a5ab..702989741 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -10,8 +10,9 @@ use parking_lot::Mutex; use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; +use crate::job_queue::JobQueueHandle; use crate::job_queue::{JobMessage, JobQueueReceiver, MAX_ATTEMPTS}; -use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskRunner}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles}; @@ -674,6 +675,7 @@ async fn slow_attempt_is_retried() { harness.state.conf_handle.clone(), harness.state.recordings.clone(), harness.state.provisioning.clone(), + Arc::new(harness.state.job_queue_handle.clone()), ); task.attempt_timeout = Duration::from_millis(500); harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( @@ -789,6 +791,69 @@ async fn finished_task_does_not_run_again() { ); } +/// 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) -> anyhow::Result<()> { + self.queue.push(job).await + } + + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + Ok(if name == RemuxJob::NAME { + self.remux_jobs.lock().clone() + } else { + Vec::new() + }) + } +} + +#[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.provisioning.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()); +} + #[tokio::test] async fn analyses_of_one_session_run_one_at_a_time() { let session_id = Uuid::new_v4(); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index a5121cc59..f1135f90f 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -37,8 +37,8 @@ pub trait ProvisionerTaskJobQueue: Send + Sync { /// Queues `job`, and returns once it is stored, so a record never waits for a job that was lost. async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()>; - /// Lists the JSON definitions of the provisioner task jobs already stored that may still run. - async fn stored_jobs(&self) -> 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; @@ -51,7 +51,7 @@ impl ProvisionerTaskJobQueue for JobQueueHandle { } /// The channel cannot tell what is stored: every unfinished record looks like it has no job. - async fn stored_jobs(&self) -> anyhow::Result> { + async fn stored_jobs(&self, _name: &str) -> anyhow::Result> { Ok(Vec::new()) } } @@ -281,7 +281,7 @@ impl ProvisionerTaskRunner { let stored = self .inner .job_queue - .stored_jobs() + .stored_jobs(ProvisionerTaskJob::NAME) .await .context("list the stored provisioner task jobs")? .iter() diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 92f4666ac..e962b99a6 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -233,7 +233,8 @@ impl ProvisionerTaskJobQueue for StoredJobs { Ok(()) } - async fn stored_jobs(&self) -> anyhow::Result> { + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + assert_eq!(name, ProvisionerTaskJob::NAME); Ok(self.0.lock().clone()) } } diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 5c5814230..172ab0880 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; @@ -1127,6 +1127,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 { From 49c9c4cb8b20a2ef645620f65d7ac4e196c87e13 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:13:21 -0400 Subject: [PATCH 19/53] fix(dgw): tie the AI analysis working files to their Task - The working files of a session were reused by any later Task with the same model, and the late job of a Task that ended without it, such as one failed past its deadline on a status read, could delete the files of a newer Task, which then paid for every chunk again. The analysis now names its Task as the owner of the files: the attempts of one Task resume from them, a new Task starts fresh, and a Task's cleanup leaves another Task's files in place. - An attempt starts only with one minute to spare past its timeout: it ends a little later while it stops, and stored deadlines lose their fraction of a second. This also keeps the runner test clear of rounding. - While describing, progress is written whenever another chunk is done, not when the percent changes, since the payload no longer has a percent. - `deadlineAt` says a Task can time out earlier, and the failure reasons are listed in full. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/tasks.rs | 3 ++- devolutions-gateway/src/openapi.rs | 2 +- .../recording_ai_analysis/mod.rs | 13 ++++++--- .../recording_ai_analysis/tests.rs | 27 +++++++++++++++++++ .../src/provisioner_tasks/runner/mod.rs | 6 ++++- .../src/provisioner_tasks/runner/tests.rs | 10 ++++--- 6 files changed, 51 insertions(+), 10 deletions(-) diff --git a/devolutions-gateway/src/api/tasks.rs b/devolutions-gateway/src/api/tasks.rs index c82e99076..d0f2b9504 100644 --- a/devolutions-gateway/src/api/tasks.rs +++ b/devolutions-gateway/src/api/tasks.rs @@ -38,7 +38,8 @@ pub(crate) struct TaskInfo { /// When it succeeded or failed. #[serde(with = "time::serde::rfc3339::option")] finished_at: Option, - /// When an unfinished Task becomes failed, waiting and retries included. + /// Latest time an unfinished Task becomes failed, waiting and retries included; it fails as `timed out` earlier + /// once too little time is left for another attempt. #[serde(with = "time::serde::rfc3339")] deadline_at: OffsetDateTime, } diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 65128cca5..c98d21e96 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -197,7 +197,7 @@ struct AiAnalysisUsage { #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "camelCase")] struct AiAnalysisFailedPayload { - /// Kind of failure, such as `permanent error`, `attempts exhausted`, `key lost` or `timed out`. + /// Kind of failure: `permanent error`, `attempts exhausted`, `key lost`, `timed out`, `invalid task` or `not queued`. reason: String, /// What went wrong: a text, except for `timed out`, where it is `{ "lastPayload": }`. details: serde_json::Value, diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index c3b65141d..f536510ff 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -319,7 +319,7 @@ impl ProvisionerTask for RecordingAiAnalysis { let manifest_path = self.manifest_path(session_id); let _session = lock_session(session_id).await; - if let Err(error) = recording_analysis::discard(&manifest_path).await { + if let Err(error) = recording_analysis::discard(&manifest_path, Some(&record.id.to_string())).await { warn!(task.id = %record.id, session.id = %session_id, %error, "Failed to remove the AI analysis working files"); } } @@ -359,6 +359,7 @@ impl RecordingAiAnalysis { 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 { @@ -366,13 +367,17 @@ impl RecordingAiAnalysis { } let mut events = request.start(); - let mut recorded: Option<(&'static str, u8)> = None; + let mut recorded: Option<(&'static str, u64)> = None; let analysis = loop { match events.recv().await { Some(RecordingAnalysisEvent::Progress(progress)) => { - // Reading sends an event every half second, so the Task is updated only when the stage's percent changes. - let shown = (progress.stage(), progress.percent()); + // Reading sends an event every half second, so while reading the Task is updated only when the percent + // changes. + let shown = match progress { + Progress::Reading { .. } => (progress.stage(), u64::from(progress.percent())), + _ => (progress.stage(), progress.done_and_total().map_or(0, |(done, _)| done)), + }; if recorded != Some(shown) { recorded = Some(shown); attempt.report_progress(progress_payload(&progress)).await; diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 702989741..85fa41649 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -892,6 +892,33 @@ async fn working_files_left_by_an_earlier_task_go_with_the_task_that_ends() { assert!(!harness.workspace(session_id).exists()); assert!(!log.exists()); } +#[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", "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.provisioning.clone(), + Arc::new(harness.state.job_queue_handle.clone()), + ); + + task.on_end(&record).await; + + assert!(workspace.exists(), "the newer Task still resumes from them"); + assert!(!harness.has_key(task_id)); +} + #[test] fn unknown_usage_is_left_out_of_the_result() { let analysis = Analysis { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index f1135f90f..2b7e45429 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -31,6 +31,10 @@ use super::AttemptError; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; use crate::token::AccessScope; +/// Time an attempt may need past its timeout: it ends a little later while it stops, and a stored deadline loses its +/// fraction of a second. +pub(crate) const ATTEMPT_END_GRACE: std::time::Duration = std::time::Duration::from_secs(60); + /// Where the runner queues the jobs of its records. #[async_trait] pub trait ProvisionerTaskJobQueue: Send + Sync { @@ -362,7 +366,7 @@ impl ProvisionerTaskRunner { && record.job_token == job_token && !record.state.is_finished() && let Some(task) = self.task(&record.kind) - && now + task.attempt_timeout() > record.deadline_at + && now + task.attempt_timeout() + ATTEMPT_END_GRACE > record.deadline_at { warn!(task.id = %task_id, task.kind = %record.kind, attempts = record.attempts, "Task timed out before its next attempt"); let payload = json!({ diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index e962b99a6..a69558afd 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -367,8 +367,12 @@ async fn an_attempt_starts_only_if_it_can_finish_before_the_deadline() { let store = harness.runner.store(); let now = OffsetDateTime::now_utc(); - // The echo attempt timeout is 5 seconds. - for (left, state) in [(6, ProvisionerTaskState::Succeeded), (4, ProvisionerTaskState::Failed)] { + // The echo attempt timeout is 5 seconds; 2 more or 2 less keep the test clear of whole-second rounding. + let needed = time::Duration::seconds(5) + time::Duration::try_from(ATTEMPT_END_GRACE).expect("grace"); + for (left, state) in [ + (needed + time::Duration::seconds(2), ProvisionerTaskState::Succeeded), + (needed - time::Duration::seconds(2), ProvisionerTaskState::Failed), + ] { let id = Uuid::new_v4(); let job_token = Uuid::new_v4(); store @@ -378,7 +382,7 @@ async fn an_attempt_starts_only_if_it_can_finish_before_the_deadline() { kind: KIND.to_owned(), target: id.to_string(), params: json!({ "value": 1 }), - deadline_at: now + time::Duration::seconds(left), + deadline_at: now + left, job_token, }, now, From 8e3c8dd0b6d87c444c73900557cd6e2afc724277 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:56:31 -0400 Subject: [PATCH 20/53] feat(dgw): start a new AI analysis Task of a recording from scratch Analysing a recording again kept the earlier results: each success added another `ai-analysis-N.slog` to the manifest, so one recording could hold several analyses. A new Task is now a hard restart. Its first attempt removes every earlier AI analysis result of the session, the file and its manifest entry, even a succeeded one, and the working files whoever left them. Its later attempts are job retries and resume from the working files. A recording holds at most one AI analysis. The recording manager gains `ClearArtifacts`, which drops the manifest entries first, so the manifest never lists a file that is gone. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/jrec/mod.rs | 3 +- devolutions-gateway/src/artifacts.rs | 4 +- .../recording_ai_analysis/mod.rs | 15 +++++ .../recording_ai_analysis/tests.rs | 53 +++++++++++++++++ devolutions-gateway/src/recording.rs | 59 +++++++++++++++++++ 5 files changed, 131 insertions(+), 3 deletions(-) diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index d85b9b05e..7b4fe64d6 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -495,7 +495,8 @@ fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> /// 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. +/// Sending the same request again returns the same Task and starts nothing new. A new Task for the session starts from +/// scratch: it removes the earlier AI analysis of the session, even a succeeded one, so a session holds at most one. /// 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( 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/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index f536510ff..acc421195 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -355,6 +355,21 @@ impl RecordingAiAnalysis { analysis_manifest(&recording.manifest).map_err(|error| AttemptError::Permanent(error.to_string()))?; let manifest_path = recording.manifest_path(); + // A new Task is a hard restart: nothing an earlier analysis of the session left stays, not even its result. + // Its later attempts are job retries, which resume from the working files. + if attempt.attempts == 1 { + let cleared = async { + self.recordings + .clear_artifacts(session_id, ArtifactKind::AiAnalysis) + .await?; + recording_analysis::discard(&manifest_path, None).await?; + anyhow::Ok(()) + }; + cleared.await.map_err(|error| { + AttemptError::Permanent(format!("failed to remove the earlier AI analysis: {error:#}")) + })?; + } + self.wait_for_remux(attempt, &recording.dir).await?; let mut request = client diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 85fa41649..c098fb712 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -892,6 +892,59 @@ async fn working_files_left_by_an_earlier_task_go_with_the_task_that_ends() { assert!(!harness.workspace(session_id).exists()); assert!(!log.exists()); } +#[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()); +} + +#[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, + "the new Task asks about every chunk again" + ); +} + #[tokio::test] async fn an_ended_task_leaves_the_working_files_of_another_task() { let harness = Harness::new().await; diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 172ab0880..1207f1af6 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -382,6 +382,11 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, + ClearArtifacts { + id: Uuid, + kind: ArtifactKind, + channel: oneshot::Sender>, + }, GetFinished { id: Uuid, channel: oneshot::Sender>, @@ -448,6 +453,11 @@ impl fmt::Debug for RecordingManagerMessage { .field("id", id) .field("kind", kind) .finish_non_exhaustive(), + RecordingManagerMessage::ClearArtifacts { id, kind, channel: _ } => f + .debug_struct("ClearArtifacts") + .field("id", id) + .field("kind", kind) + .finish_non_exhaustive(), RecordingManagerMessage::GetFinished { id, channel: _ } => { f.debug_struct("GetFinished").field("id", id).finish_non_exhaustive() } @@ -514,6 +524,19 @@ impl RecordingMessageSender { rx.await.context("couldn't receive AddArtifact result")? } + /// Removes every artifact of `kind` from the session, both its file and its manifest entry. + /// + /// Fails when the session has no recording. + pub(crate) async fn clear_artifacts(&self, id: Uuid, kind: ArtifactKind) -> anyhow::Result<()> { + let (tx, rx) = oneshot::channel(); + self.channel + .send(RecordingManagerMessage::ClearArtifacts { id, kind, channel: tx }) + .await + .ok() + .context("couldn't send ClearArtifacts message")?; + rx.await.context("couldn't receive ClearArtifacts result")? + } + pub(crate) async fn get_finished(&self, id: Uuid) -> Result { let (tx, rx) = oneshot::channel(); self.channel @@ -904,6 +927,39 @@ impl RecordingManagerTask { Ok(artifact_path) } + async fn handle_clear_artifacts(&self, id: Uuid, kind: ArtifactKind) -> anyhow::Result<()> { + let recording_path = self.recordings_path.join(id.to_string()); + let manifest_path = recording_path.join("recording.json"); + + if !manifest_path.exists() { + anyhow::bail!("session {id} has no recording"); + } + + let mut manifest = JrecManifest::read_from_file(&manifest_path).context("read manifest from disk")?; + let artifacts = core::mem::take(manifest.artifacts.of_kind_mut(kind)); + + if artifacts.is_empty() { + return Ok(()); + } + + // The manifest goes first, so it never lists a file that is gone. + manifest + .save_to_file(&manifest_path) + .context("write manifest to disk")?; + + for artifact in artifacts { + let artifact_path = recording_path.join(&artifact.file_name); + match fs::remove_file(&artifact_path).await { + Err(error) if error.kind() != io::ErrorKind::NotFound => { + return Err(anyhow::Error::new(error).context(format!("remove {artifact_path}"))); + } + _ => {} + } + } + + 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) { @@ -1031,6 +1087,9 @@ async fn recording_manager_task( RecordingManagerMessage::AddArtifact { id, kind, channel } => { let _ = channel.send(manager.handle_add_artifact(id, kind).await); }, + RecordingManagerMessage::ClearArtifacts { id, kind, channel } => { + let _ = channel.send(manager.handle_clear_artifacts(id, kind).await); + }, RecordingManagerMessage::GetFinished { id, channel } => { let _ = channel.send(manager.handle_get_finished(id)); } From 74d0153f9e5011440509fe9d546df1b9c9836606 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 14:19:27 -0400 Subject: [PATCH 21/53] style(dgw): format the AI analysis task tests --- .../src/provisioner_tasks/recording_ai_analysis/tests.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index c098fb712..3c91ad0d5 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -10,8 +10,7 @@ use parking_lot::Mutex; use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; -use crate::job_queue::JobQueueHandle; -use crate::job_queue::{JobMessage, JobQueueReceiver, MAX_ATTEMPTS}; +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}; From f3a9d67359a59b3a593d641fde197d8bddaa1e55 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:19:39 -0400 Subject: [PATCH 22/53] feat(dgw): write the AI analysis in the language of the request MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The caller of `POST /jet/jrec/{session_id}/ai-analysis` now picks the language in an optional `language` field: any text of 1 to 10 characters, such as `fr-CA` or `français`; anything else is a 400. Without it, the analysis is in English. The language is part of the Task parameters, so the same Task ID with another language is a `task_id_conflict`. Gateway no longer reads anything the producer pushed: there is no events log to take a `locale` from. - The tests count the summary request the AI crate now sends. - "AI analysis queued" is a debug line, since the runner already logs the queued Task. - The jrec and provisioner tasks intents state the language, how many analyses a session keeps, and how retries and restarts behave. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/jrec.intent.md | 3 +- devolutions-gateway/src/api/jrec/mod.rs | 20 +++++- devolutions-gateway/src/openapi.rs | 3 +- .../src/provisioner_tasks/INTENT.md | 4 +- .../recording_ai_analysis/mod.rs | 42 ++++++++--- .../recording_ai_analysis/tests.rs | 72 +++++++++++++++++-- devolutions-gateway/tests/ai_analysis.rs | 11 ++- 7 files changed, 134 insertions(+), 21 deletions(-) diff --git a/devolutions-gateway/src/api/jrec.intent.md b/devolutions-gateway/src/api/jrec.intent.md index e65cc4955..8bcee2600 100644 --- a/devolutions-gateway/src/api/jrec.intent.md +++ b/devolutions-gateway/src/api/jrec.intent.md @@ -7,4 +7,5 @@ Jrec is the endpoint to push, streaming, read recordings and its associated file 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. \ No newline at end of file +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; a language longer than 10 characters is refused, and the language is part of the Task, so the same Task ID sent again with another language is a conflict. \ No newline at end of file diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index 7b4fe64d6..13a1bf5ce 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -19,6 +19,7 @@ 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; @@ -455,6 +456,10 @@ pub(crate) struct AiAnalysisRequest { /// Upper bound of tokens in each AI answer. #[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, @@ -492,11 +497,13 @@ fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> /// /// Gateway asks the AI provider what the user did in the session recording, and adds the answer to the session as a /// new 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: it removes the earlier AI analysis of the session, even a succeeded one, so a session holds at most one. +/// Sending the same request again returns the same Task and starts nothing new. +/// A new Task for the session starts from scratch: it removes the earlier AI analysis of the session, even a succeeded +/// one, so a session holds at most one. /// 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( @@ -513,7 +520,7 @@ fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> (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 or AI settings"), + (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"), @@ -545,9 +552,15 @@ pub(crate) async fn start_ai_analysis( 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 { @@ -559,6 +572,7 @@ pub(crate) async fn start_ai_analysis( base_url, max_output_tokens, }, + language, api_key, }, ) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index c98d21e96..dcd42fe21 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -179,7 +179,8 @@ struct AiAnalysisSucceededPayload { model: String, /// Tokens counted by the AI provider; absent when it did not report them for every request. usage: Option, - /// Version of the instructions sent to the AI provider. + /// Versions of the prompts that wrote the log, such as + /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. prompt_version: String, } diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index 346de3c4b..b8746192f 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -27,4 +27,6 @@ After a Gateway restart they are gone, and the Tasks that needed them fail; this ## 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. \ No newline at end of file +Two analyses of the same session never run at the same time. +A session keeps one AI analysis: a new Task removes the earlier one when it starts, and a Task that fails leaves nothing, so the session then has none. Keeping one analysis per language is a future extension that must stay possible. +A transient failure such as a rate limit (429) is retried with the job, and each retry resumes where the last attempt stopped; only a new Task starts over. A Gateway restart loses the API key, so an unfinished Task fails as "key lost" at its next attempt, without a retry, and leaves nothing. \ No newline at end of file diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index acc421195..a0aa6beca 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -1,10 +1,12 @@ //! AI analysis (`recording-ai-analysis`): describes what the user did in one recorded session, and adds the result to //! the session as a new log. //! -//! [`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 adds the log to the -//! session. The API key and the working files go away when the Task ends. +//! [`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 adds the log to the session. +//! The API key and the working files go away when the Task ends. #[cfg(test)] mod tests; @@ -15,7 +17,7 @@ use std::sync::{Arc, LazyLock, Weak}; use anyhow::Context as _; use async_trait::async_trait; use camino::{Utf8Path, Utf8PathBuf}; -use devolutions_gateway_ai::Usage; +use devolutions_gateway_ai::{Language, Usage}; use devolutions_gateway_ai::recording_analysis::{ self, Progress, RecordingAnalysisEvent, RecordingFile, RecordingManifest, Recordings, RecordingsError, }; @@ -74,6 +76,8 @@ pub struct StartRequest { 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, } @@ -103,13 +107,16 @@ pub enum StartError { Internal(anyhow::Error), } -/// What an AI analysis Task records as its parameters: the AI settings without the API key and the base URL. +/// 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(skip_serializing_if = "Option::is_none")] + language: Option, } /// Result payload of a successful AI analysis. @@ -125,7 +132,8 @@ pub struct Analysis { /// Absent when the provider did not report them for every request. #[serde(skip_serializing_if = "Option::is_none")] usage: Option, - /// Versions of the prompts that wrote the log, such as `session-actions-3`. + /// Versions of the prompts that wrote the log, such as + /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. prompt_version: String, } @@ -138,6 +146,7 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result Result Result { - info!(task.id = %task_id, session.id = %session_id, "AI analysis queued"); + debug!(task.id = %task_id, session.id = %session_id, "AI analysis queued"); StartOutcome::Created(task) } CreateOutcome::Existing(task) => StartOutcome::Existing(task), @@ -299,6 +309,15 @@ impl ProvisionerTask for RecordingAiAnalysis { let access = AiAccess::from_secret(&secret).map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + let language = params + .language + .map(Language::new) + .transpose() + .map_err(|error| AttemptError::Failed { + reason: "invalid task", + details: format!("invalid language: {error}"), + })?; + let settings = AiSettings { provider: params.provider, model: params.model, @@ -306,7 +325,7 @@ impl ProvisionerTask for RecordingAiAnalysis { max_output_tokens: params.max_output_tokens, }; - self.analyze(attempt, session_id, settings, access.api_key).await + self.analyze(attempt, session_id, settings, language, access.api_key).await } async fn on_end(&self, record: &ProvisionerTaskRecord) { @@ -331,6 +350,7 @@ impl RecordingAiAnalysis { attempt: &Attempt, session_id: Uuid, settings: AiSettings, + language: Option, api_key: SecretString, ) -> Result { let conf = self.conf_handle.get_conf(); @@ -381,6 +401,10 @@ impl RecordingAiAnalysis { request = request.max_output_tokens(max_output_tokens); } + if let Some(language) = language { + request = request.language(language); + } + let mut events = request.start(); let mut recorded: Option<(&'static str, u64)> = None; diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 3c91ad0d5..1fc8ebf24 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -4,6 +4,7 @@ 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; @@ -89,12 +90,24 @@ impl Harness { } 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), }, ) @@ -524,14 +537,18 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { json!({ "fileName": "ai-analysis-0.slog", "model": "gpt-test-2026-09-30", - "usage": { "inputTokens": 30, "outputTokens": 60 }, - "promptVersion": PROMPT_VERSION, + "usage": { "inputTokens": 40, "outputTokens": 80 }, + "promptVersion": format!("{PROMPT_VERSION}+{}", session_summary::PROMPT_VERSION), }), "chunk 0 counts from its checkpoint, the failed request not at all" ); let requests = requests.lock().clone(); - assert_eq!(requests.len(), 4, "only chunks 1 and 2 are sent again"); + 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}"); } @@ -916,6 +933,51 @@ async fn a_new_task_replaces_the_result_of_the_earlier_one() { assert!(!dir.join("ai-analysis-1.slog").exists()); } +#[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:?}" + ); +} + +#[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`"))); +} + #[tokio::test] async fn a_new_task_never_resumes_from_an_earlier_task() { let harness = Harness::new().await; @@ -939,8 +1001,8 @@ async fn a_new_task_never_resumes_from_an_earlier_task() { assert_eq!(harness.task(second).await.state, ProvisionerTaskState::Succeeded); assert_eq!( requests.lock().len(), - 2 + 3, - "the new Task asks about every chunk again" + 2 + 3 + 1, + "the new Task asks about every chunk again, then for the summary" ); } diff --git a/devolutions-gateway/tests/ai_analysis.rs b/devolutions-gateway/tests/ai_analysis.rs index d471a42fb..e1cc0d944 100644 --- a/devolutions-gateway/tests/ai_analysis.rs +++ b/devolutions-gateway/tests/ai_analysis.rs @@ -359,6 +359,12 @@ async fn invalid_bodies_and_settings_return_400() { let mut unknown_provider = request(task_id, "gpt-test"); unknown_provider["provider"] = json!("skynet"); + let mut long_language = request(task_id, "gpt-test"); + long_language["language"] = json!("zh-Hant-TW-x"); + + let mut empty_language = request(task_id, "gpt-test"); + empty_language["language"] = json!(""); + for (case, body) in [ ("missing API key", no_key.to_string()), ("empty API key", empty_key.to_string()), @@ -366,6 +372,8 @@ async fn invalid_bodies_and_settings_return_400() { ("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 (status, _) = harness @@ -517,7 +525,8 @@ async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { assert_eq!(task["payload"]["model"], "gpt-test-2026-09-30"); assert_eq!( task["payload"]["usage"], - json!({ "inputTokens": 10, "outputTokens": 20 }) + json!({ "inputTokens": 20, "outputTokens": 40 }), + "one request for the actions, one for the summary" ); assert_rfc3339(&task["startedAt"]); assert_rfc3339(&task["finishedAt"]); From 634a41f6a3fad2c0b800a0441313aa3ce1f40066 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:41:47 -0400 Subject: [PATCH 23/53] test(dgw): give the fixture of another task its language The working files of an analysis now always name their language, so a fixture without one is not a valid workspace. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/provisioner_tasks/recording_ai_analysis/tests.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 1fc8ebf24..145480cfb 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -1016,7 +1016,14 @@ async fn an_ended_task_leaves_the_working_files_of_another_task() { 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", "chunkLen": 1 }).to_string(), + 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; From 2873ccab204e643190d97c05917343fa0542cbad Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:54:55 -0400 Subject: [PATCH 24/53] fix(dgw): clear the earlier AI analysis on every attempt of a new Task Clearing on the first attempt only missed two paths. A Task that failed before its first attempt cleared, such as on a lost key, kept the earlier analysis. A first attempt cancelled while waiting for the session retried as attempt 2 and never cleared, so the session ended with two analyses. Each attempt now takes the session lock, then clears, before anything else can fail. The Task params keep the language as a `Language`, so it's checked once, when the request comes in. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../recording_ai_analysis/mod.rs | 72 ++++++++-------- .../recording_ai_analysis/tests.rs | 83 ++++++++++++++++++- 2 files changed, 118 insertions(+), 37 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index a0aa6beca..a7a456744 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -17,10 +17,10 @@ use std::sync::{Arc, LazyLock, Weak}; use anyhow::Context as _; use async_trait::async_trait; use camino::{Utf8Path, Utf8PathBuf}; -use devolutions_gateway_ai::{Language, Usage}; use devolutions_gateway_ai::recording_analysis::{ self, Progress, RecordingAnalysisEvent, RecordingFile, RecordingManifest, Recordings, RecordingsError, }; +use devolutions_gateway_ai::{Language, Usage}; use provisioner_task::{CreateOutcome, ProvisionerTaskRecord}; use secrecy::SecretString; use serde_json::json; @@ -115,8 +115,8 @@ pub struct Params { provider: AiProvider, model: String, max_output_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - language: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + language: Option, } /// Result payload of a successful AI analysis. @@ -169,7 +169,7 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result session, + () = attempt.cancel.cancelled() => return Err(AttemptError::Cancelled), + }; + + // A new Task is a hard restart, so the session loses its earlier analysis even when the Task then fails, such as + // on a lost key. Every attempt clears, since an earlier one may have ended before it could. + self.clear_earlier_analysis(attempt, session_id).await?; + let secret = self .provisioning .task_secret(attempt.task_id) @@ -309,15 +319,6 @@ impl ProvisionerTask for RecordingAiAnalysis { let access = AiAccess::from_secret(&secret).map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; - let language = params - .language - .map(Language::new) - .transpose() - .map_err(|error| AttemptError::Failed { - reason: "invalid task", - details: format!("invalid language: {error}"), - })?; - let settings = AiSettings { provider: params.provider, model: params.model, @@ -325,7 +326,8 @@ impl ProvisionerTask for RecordingAiAnalysis { max_output_tokens: params.max_output_tokens, }; - self.analyze(attempt, session_id, settings, language, access.api_key).await + self.analyze(attempt, session_id, settings, params.language, access.api_key) + .await } async fn on_end(&self, record: &ProvisionerTaskRecord) { @@ -345,6 +347,27 @@ impl ProvisionerTask for RecordingAiAnalysis { } impl RecordingAiAnalysis { + /// Removes the analysis an earlier Task added to the session, and on the first attempt the working files any + /// analysis left; later attempts are job retries, which resume from their own working files. + async fn clear_earlier_analysis(&self, attempt: &Attempt, session_id: Uuid) -> Result<(), AttemptError> { + let cleared = async { + self.recordings + .clear_artifacts(session_id, ArtifactKind::AiAnalysis) + .await?; + + if attempt.attempts == 1 { + recording_analysis::discard(&self.manifest_path(session_id), None).await?; + } + + anyhow::Ok(()) + }; + + cleared + .await + .map_err(|error| AttemptError::Permanent(format!("failed to remove the earlier AI analysis: {error:#}"))) + } + + /// Runs the analysis while the caller holds the lock of the session. async fn analyze( &self, attempt: &Attempt, @@ -359,12 +382,6 @@ impl RecordingAiAnalysis { .client(&conf, &api_key) .map_err(|error| AttemptError::Permanent(error.to_string()))?; - // Waiting for an earlier analysis of the session counts in the attempt time, so a timeout cancels it too. - let _session = tokio::select! { - session = lock_session(session_id) => session, - () = attempt.cancel.cancelled() => return Err(AttemptError::Cancelled), - }; - let recording = self .recordings .get_finished(session_id) @@ -375,21 +392,6 @@ impl RecordingAiAnalysis { analysis_manifest(&recording.manifest).map_err(|error| AttemptError::Permanent(error.to_string()))?; let manifest_path = recording.manifest_path(); - // A new Task is a hard restart: nothing an earlier analysis of the session left stays, not even its result. - // Its later attempts are job retries, which resume from the working files. - if attempt.attempts == 1 { - let cleared = async { - self.recordings - .clear_artifacts(session_id, ArtifactKind::AiAnalysis) - .await?; - recording_analysis::discard(&manifest_path, None).await?; - anyhow::Ok(()) - }; - cleared.await.map_err(|error| { - AttemptError::Permanent(format!("failed to remove the earlier AI analysis: {error:#}")) - })?; - } - self.wait_for_remux(attempt, &recording.dir).await?; let mut request = client diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 145480cfb..de482567d 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -632,7 +632,9 @@ async fn recording_deleted_after_start_fails() { let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["details"], "session has no recording"); + assert_eq!(task.payload["reason"], "permanent error"); + let details = task.payload["details"].as_str().expect("details"); + assert!(details.ends_with("has no recording"), "{details}"); } #[tokio::test] @@ -933,6 +935,78 @@ async fn a_new_task_replaces_the_result_of_the_earlier_one() { assert!(!dir.join("ai-analysis-1.slog").exists()); } +#[tokio::test] +async fn a_new_task_that_fails_still_removes_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.provisioning.remove_task_secret(second); + harness.run(second).await.expect("no retry"); + + let task = harness.task(second).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "key lost"); + let dir = harness.recording_path.join(session_id.to_string()); + assert_eq!(ai_analyses(&dir), json!([])); + assert!(!dir.join("ai-analysis-0.slog").exists()); +} + +#[tokio::test] +async fn a_retry_of_an_attempt_that_never_got_the_session_removes_the_earlier_analysis() { + let mut 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; + + let mut task = RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.state.provisioning.clone(), + Arc::new(harness.state.job_queue_handle.clone()), + ); + task.attempt_timeout = Duration::from_millis(500); + harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( + Arc::clone(harness.state.provisioner_tasks.store()), + Arc::new(harness.state.job_queue_handle.clone()), + ) + .register(task) + .build(); + + // Another analysis of the session holds it past the attempt timeout. + let session = lock_session(session_id).await; + let error = harness.run(second).await.expect_err("timeout is retried"); + assert!(error.to_string().contains("took longer than"), "{error:#}"); + drop(session); + + harness.run(second).await.expect("second attempt"); + + let task = harness.task(second).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert_eq!(task.attempts, 2); + let dir = harness.recording_path.join(session_id.to_string()); + assert_eq!( + ai_analyses(&dir), + json!([{ "fileName": "ai-analysis-0.slog" }]), + "one analysis per recording" + ); +} + +/// 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!([])) +} + #[tokio::test] async fn the_analysis_is_in_the_language_of_the_request() { let harness = Harness::new().await; @@ -954,7 +1028,12 @@ async fn the_analysis_is_in_the_language_of_the_request() { 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"))) + .start_in( + task_id, + session_id, + settings, + Some(Language::new("de-DE").expect("language")), + ) .await .expect("started"); assert!( From b1e3d327220cf96440f34573b553fbc4694aaf47 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:54:56 -0400 Subject: [PATCH 25/53] docs(dgw): describe the params of an AI analysis Task in the OpenAPI document `params` of a `recording-ai-analysis` Task now has the `AiAnalysisParams` schema, like its payloads. The lint attributes on these schema structs are gone, since the lint never fires on them. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/tasks.rs | 4 ++-- devolutions-gateway/src/openapi.rs | 19 +++++++++++++++---- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/devolutions-gateway/src/api/tasks.rs b/devolutions-gateway/src/api/tasks.rs index d0f2b9504..2431febb6 100644 --- a/devolutions-gateway/src/api/tasks.rs +++ b/devolutions-gateway/src/api/tasks.rs @@ -8,8 +8,8 @@ use uuid::Uuid; /// /// 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`, `payload` is an `AiAnalysisRunningPayload` while `running`, an -/// `AiAnalysisSucceededPayload` once `succeeded`, and an `AiAnalysisFailedPayload` once `failed`. +/// 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")] diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index dcd42fe21..2354f3a6e 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -72,6 +72,7 @@ use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; crate::provisioner_tasks::ai::AiProvider, crate::api::tasks::TaskInfo, crate::api::tasks::TaskState, + AiAnalysisParams, AiAnalysisRunningPayload, AiAnalysisSucceededPayload, AiAnalysisUsage, @@ -154,8 +155,21 @@ 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: crate::provisioner_tasks::ai::AiProvider, + /// Model identifier, passed to the provider as is. + model: String, + /// Upper bound of tokens in each AI answer; null when not asked. + 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 -#[allow(dead_code)] #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "camelCase")] struct AiAnalysisRunningPayload { @@ -169,7 +183,6 @@ struct AiAnalysisRunningPayload { } /// Payload of a `succeeded` AI analysis Task: the log it added to the session -#[allow(dead_code)] #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "camelCase")] struct AiAnalysisSucceededPayload { @@ -185,7 +198,6 @@ struct AiAnalysisSucceededPayload { } /// Tokens counted by an AI provider -#[allow(dead_code)] #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "camelCase")] struct AiAnalysisUsage { @@ -194,7 +206,6 @@ struct AiAnalysisUsage { } /// Payload of a `failed` AI analysis Task: why it failed -#[allow(dead_code)] #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "camelCase")] struct AiAnalysisFailedPayload { From 815aeeb4590563b4da0977b2faeef52f1056d07e Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:54:57 -0400 Subject: [PATCH 26/53] test(dgw): run the AI analysis endpoint tests against the Gateway binary The tests move from the crate's own `tests/` to the testsuite, as AGENTS.md asks, and talk to a running Gateway and a mock AI provider instead of linking Gateway internals. They read the key leak checks from the log files and `gateway.db` in the config folder. The check that a Task of another kind is not found at this endpoint is gone: the binary only knows one kind. Co-Authored-By: Claude Opus 5.5 (1M context) --- Cargo.lock | 1 + devolutions-gateway/tests/ai_analysis.rs | 583 ----------------------- testsuite/Cargo.toml | 1 + testsuite/tests/cli/dgw/ai_analysis.rs | 523 ++++++++++++++++++++ testsuite/tests/cli/dgw/mod.rs | 1 + 5 files changed, 526 insertions(+), 583 deletions(-) delete mode 100644 devolutions-gateway/tests/ai_analysis.rs create mode 100644 testsuite/tests/cli/dgw/ai_analysis.rs diff --git a/Cargo.lock b/Cargo.lock index 020845f37..feaaa934e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7753,6 +7753,7 @@ dependencies = [ "sysevent-winevent", "tempfile", "test-utils", + "time", "tokio 1.52.3", "tokio-rustls", "tokio-tungstenite", diff --git a/devolutions-gateway/tests/ai_analysis.rs b/devolutions-gateway/tests/ai_analysis.rs deleted file mode 100644 index e1cc0d944..000000000 --- a/devolutions-gateway/tests/ai_analysis.rs +++ /dev/null @@ -1,583 +0,0 @@ -#![allow(unused_crate_dependencies)] -#![allow(clippy::unwrap_used)] - -use std::io::Write; -use std::net::SocketAddr; -use std::sync::{Arc, Mutex}; - -use axum::Router; -use axum::body::Body; -use axum::extract::connect_info::MockConnectInfo; -use axum::http::{Request, StatusCode}; -use base64::Engine as _; -use camino::{Utf8Path, Utf8PathBuf}; -use devolutions_gateway::job_queue::{JobMessage, JobQueueReceiver}; -use devolutions_gateway::recording::RecordingManagerTask; -use devolutions_gateway::{DgwState, MockHandles}; -use devolutions_gateway_task::ShutdownHandle; -use http_body_util::BodyExt as _; -use serde_json::{Value, json}; -use time::OffsetDateTime; -use time::format_description::well_known::Rfc3339; -use tower::ServiceExt as _; -use tracing_subscriber::util::SubscriberInitExt as _; -use uuid::Uuid; - -const API_KEY: &str = "sk-ai-analysis-endpoint-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, - "__debug__": { "disable_token_validation": true }, - }) - .to_string() -} - -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}") -} - -/// Collects every log line, so a test can check that no secret reaches the logs. -#[derive(Clone, Default)] -struct LogBuffer(Arc>>); - -impl Write for LogBuffer { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.0.lock().unwrap().extend_from_slice(buf); - Ok(buf.len()) - } - - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } -} - -impl LogBuffer { - fn text(&self) -> String { - String::from_utf8_lossy(&self.0.lock().unwrap()).into_owned() - } -} - -struct Harness { - app: Router, - job_queue_rx: JobQueueReceiver, - recording_path: Utf8PathBuf, - db_path: Utf8PathBuf, - task_store: provisioner_task::DynProvisionerTaskStore, - _keep: Box, - _dir: tempfile::TempDir, -} - -impl Harness { - async fn new() -> Self { - let dir = tempfile::tempdir().unwrap(); - let root = Utf8PathBuf::from_path_buf(dir.path().to_path_buf()).unwrap(); - let recording_path = root.join("recordings"); - let db_path = root.join("gateway.db"); - - let (mut state, handles) = DgwState::mock(&config(&recording_path)).await.unwrap(); - - // A file, unlike the in-memory default, can be searched for secrets afterwards. - let connection = gateway_db::GatewayDb::open_path(db_path.as_str()) - .await - .unwrap() - .connect() - .await - .unwrap(); - let task_store: provisioner_task::DynProvisionerTaskStore = Arc::new( - provisioner_task_libsql::LibSqlProvisionerTaskStore::open(connection) - .await - .unwrap(), - ); - state.provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( - Arc::clone(&task_store), - Arc::new(state.job_queue_handle.clone()), - state.conf_handle.clone(), - state.recordings.clone(), - state.provisioning.clone(), - ); - - 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, - )); - - let app = devolutions_gateway::make_http_service(state) - .layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 3000)))); - - Self { - app, - job_queue_rx, - recording_path, - db_path, - task_store, - _keep: Box::new(( - session_manager_rx, - subscriber_rx, - traffic_audit_rx, - shutdown_handle, - recording_shutdown_handle, - )), - _dir: dir, - } - } - - fn write_session(&self) -> Uuid { - let session_id = Uuid::new_v4(); - let dir = self.recording_path.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 send(&self, request: Request) -> (StatusCode, Value) { - let response = self.app.clone().oneshot(request).await.unwrap(); - let status = response.status(); - let body = response.into_body().collect().await.unwrap().to_bytes(); - let body = if body.is_empty() { - Value::Null - } else { - serde_json::from_slice(&body).unwrap() - }; - (status, body) - } - - async fn post_raw(&self, session_id: Uuid, scope: &str, body: String) -> (StatusCode, Value) { - let request = Request::builder() - .method("POST") - .uri(format!("/jet/jrec/{session_id}/ai-analysis")) - .header("content-type", "application/json") - .header("authorization", format!("Bearer {}", scope_token(scope))) - .body(Body::from(body)) - .unwrap(); - self.send(request).await - } - - async fn post(&self, session_id: Uuid, body: &Value) -> (StatusCode, Value) { - self.post_raw( - session_id, - "gateway.tasks.recording-ai-analysis.start", - body.to_string(), - ) - .await - } - - async fn get_analysis(&self, session_id: Uuid, task_id: Uuid, scope: &str) -> (StatusCode, Value) { - let request = Request::builder() - .uri(format!("/jet/jrec/{session_id}/ai-analysis/{task_id}")) - .header("authorization", format!("Bearer {}", scope_token(scope))) - .body(Body::empty()) - .unwrap(); - self.send(request).await - } - - async fn run_queued_job(&mut self) { - let JobMessage { mut job, .. } = self.job_queue_rx.try_recv().expect("a queued job"); - job.run().await.unwrap(); - } - - fn database_bytes(&self) -> Vec { - let mut bytes = Vec::new(); - for suffix in ["", "-wal"] { - if let Ok(content) = std::fs::read(format!("{}{suffix}", self.db_path)) { - bytes.extend(content); - } - } - bytes - } -} - -fn request(task_id: Uuid, model: &str) -> Value { - json!({ "taskId": task_id, "provider": "openai", "model": model, "apiKey": API_KEY }) -} - -/// 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 app = Router::new().fallback(axum::routing::post(move || async move { - if status != 200 { - return ( - StatusCode::from_u16(status).unwrap(), - axum::Json(json!({ "error": { "message": "scripted failure" } })), - ); - } - - let answer = 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 } - }); - (StatusCode::OK, axum::Json(answer)) - })); - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - tokio::spawn(async move { axum::serve(listener, app).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}")); -} - -#[tokio::test] -async fn start_returns_202_then_200_for_the_same_request() { - let harness = Harness::new().await; - let session_id = harness.write_session(); - let task_id = Uuid::new_v4(); - - let (status, created) = harness.post(session_id, &request(task_id, "gpt-test")).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", "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) = harness.post(session_id, &request(task_id, "gpt-test")).await; - assert_eq!(status, StatusCode::OK, "{again}"); - assert_eq!(again, created); -} - -#[tokio::test] -async fn start_conflicts_return_409_with_a_code() { - let harness = Harness::new().await; - let session_id = harness.write_session(); - let task_id = Uuid::new_v4(); - harness.post(session_id, &request(task_id, "gpt-test")).await; - - let (status, body) = harness.post(session_id, &request(task_id, "gpt-other")).await; - assert_eq!(status, StatusCode::CONFLICT); - assert_eq!(body, json!({ "code": "task_id_conflict" })); - - let (status, body) = harness.post(session_id, &request(Uuid::new_v4(), "gpt-test")).await; - assert_eq!(status, StatusCode::CONFLICT); - assert_eq!(body, json!({ "code": "analysis_in_progress", "taskId": task_id })); -} - -#[tokio::test] -async fn start_without_recording_returns_404() { - let harness = Harness::new().await; - - let (status, _) = harness.post(Uuid::new_v4(), &request(Uuid::new_v4(), "gpt-test")).await; - - assert_eq!(status, StatusCode::NOT_FOUND); -} - -#[tokio::test] -async fn start_ignores_unknown_fields() { - let harness = Harness::new().await; - let session_id = harness.write_session(); - let mut body = request(Uuid::new_v4(), "gpt-test"); - body["temperature"] = json!(0.5); - - let (status, created) = harness.post(session_id, &body).await; - - assert_eq!(status, StatusCode::ACCEPTED, "{created}"); -} - -#[tokio::test] -async fn invalid_bodies_and_settings_return_400() { - let harness = Harness::new().await; - let session_id = harness.write_session(); - let task_id = Uuid::new_v4(); - - let mut no_key = request(task_id, "gpt-test"); - no_key.as_object_mut().unwrap().remove("apiKey"); - - let mut empty_key = request(task_id, "gpt-test"); - empty_key["apiKey"] = json!(""); - - let mut no_base_url = request(task_id, "gpt-test"); - no_base_url["provider"] = json!("openai-compatible"); - - let mut bad_base_url = request(task_id, "gpt-test"); - bad_base_url["baseUrl"] = json!("ftp://ai.example/v1/"); - - let mut unknown_provider = request(task_id, "gpt-test"); - unknown_provider["provider"] = json!("skynet"); - - let mut long_language = request(task_id, "gpt-test"); - long_language["language"] = json!("zh-Hant-TW-x"); - - let mut empty_language = request(task_id, "gpt-test"); - 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, " ").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 (status, _) = harness - .post_raw(session_id, "gateway.tasks.recording-ai-analysis.start", body) - .await; - assert_eq!(status, StatusCode::BAD_REQUEST, "{case}"); - } - - let (status, _) = harness - .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") - .await; - assert_eq!(status, StatusCode::NOT_FOUND, "no Task is recorded"); -} - -#[tokio::test] -async fn endpoints_require_their_scope() { - let harness = Harness::new().await; - let session_id = harness.write_session(); - let body = request(Uuid::new_v4(), "gpt-test").to_string(); - - for scope in [ - "gateway.tasks.recording-ai-analysis.read", - "gateway.recording.delete", - "gateway.recordings.read", - ] { - let (status, _) = harness.post_raw(session_id, scope, body.clone()).await; - assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); - } - - for scope in ["gateway.tasks.recording-ai-analysis.start", "gateway.sessions.read"] { - let (status, _) = harness.get_analysis(session_id, Uuid::new_v4(), scope).await; - assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); - } - - let (status, _) = harness.post_raw(session_id, "*", body).await; - assert_eq!(status, StatusCode::ACCEPTED, "the wildcard scope is accepted"); - - let (status, _) = harness.get_analysis(session_id, Uuid::new_v4(), "*").await; - assert_eq!(status, StatusCode::NOT_FOUND, "the wildcard scope is accepted"); -} - -#[tokio::test] -async fn an_analysis_is_found_at_its_location_only() { - let harness = Harness::new().await; - let session_id = harness.write_session(); - let task_id = Uuid::new_v4(); - let request = Request::builder() - .method("POST") - .uri(format!("/jet/jrec/{session_id}/ai-analysis")) - .header("content-type", "application/json") - .header( - "authorization", - format!("Bearer {}", scope_token("gateway.tasks.recording-ai-analysis.start")), - ) - .body(Body::from(request(task_id, "gpt-test").to_string())) - .unwrap(); - - let response = harness.app.clone().oneshot(request).await.unwrap(); - - 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) = harness - .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") - .await; - assert_eq!(status, StatusCode::OK, "{task}"); - assert_eq!(task["id"], task_id.to_string()); - - let (status, _) = harness - .get_analysis(Uuid::new_v4(), task_id, "gateway.tasks.recording-ai-analysis.read") - .await; - assert_eq!(status, StatusCode::NOT_FOUND, "not under another session"); - - // A Task of another kind is not an AI analysis, whatever the scope. - let now = OffsetDateTime::now_utc(); - let other_kind = Uuid::new_v4(); - harness - .task_store - .create( - provisioner_task::NewProvisionerTaskRecord { - id: other_kind, - kind: "test.other-kind".to_owned(), - target: session_id.to_string(), - params: json!({}), - deadline_at: now + time::Duration::hours(1), - job_token: Uuid::new_v4(), - }, - now, - ) - .await - .unwrap(); - let (status, _) = harness.get_analysis(session_id, other_kind, "*").await; - assert_eq!(status, StatusCode::NOT_FOUND); -} -#[tokio::test] -async fn get_analysis_returns_404_for_an_unknown_id() { - let harness = Harness::new().await; - - let (status, _) = harness - .get_analysis( - Uuid::new_v4(), - Uuid::new_v4(), - "gateway.tasks.recording-ai-analysis.read", - ) - .await; - - assert_eq!(status, StatusCode::NOT_FOUND); -} - -#[tokio::test] -async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { - let logs = LogBuffer::default(); - let _guard = tracing_subscriber::fmt() - .with_writer({ - let logs = logs.clone(); - move || logs.clone() - }) - .with_ansi(false) - .with_max_level(tracing::Level::TRACE) - .set_default(); - - let mut harness = Harness::new().await; - let session_id = harness.write_session(); - let task_id = Uuid::new_v4(); - let base_url = spawn_ai_provider(200).await; - - let body = json!({ - "taskId": task_id, - "provider": "openai-compatible", - "model": "gpt-test", - "baseUrl": base_url, - "maxOutputTokens": 1000, - "apiKey": API_KEY, - }); - - let (status, created) = harness.post(session_id, &body).await; - assert_eq!(status, StatusCode::ACCEPTED, "{created}"); - - harness.run_queued_job().await; - - let (status, task) = harness - .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") - .await; - assert_eq!(status, StatusCode::OK); - 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"); - assert_eq!( - task["payload"]["usage"], - json!({ "inputTokens": 20, "outputTokens": 40 }), - "one request for the actions, one for the summary" - ); - assert_rfc3339(&task["startedAt"]); - assert_rfc3339(&task["finishedAt"]); - - let (status, again) = harness.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 = logs.text(); - let database = String::from_utf8_lossy(&harness.database_bytes()).into_owned(); - assert!(!database.is_empty(), "gateway.db is read"); - assert!(logs.contains("Task succeeded"), "logs are captured"); - - 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"); - } -} - -#[tokio::test] -async fn provider_refusing_the_key_fails_the_task() { - let mut harness = Harness::new().await; - let session_id = harness.write_session(); - let task_id = Uuid::new_v4(); - let base_url = spawn_ai_provider(401).await; - - let body = json!({ - "taskId": task_id, - "provider": "openai-compatible", - "model": "gpt-test", - "baseUrl": base_url, - "apiKey": API_KEY, - }); - let (status, _) = harness.post(session_id, &body).await; - assert_eq!(status, StatusCode::ACCEPTED); - - harness.run_queued_job().await; - - let (status, task) = harness - .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") - .await; - assert_eq!(status, StatusCode::OK); - assert_eq!(task["state"], "failed", "{task}"); - assert_eq!(task["payload"]["reason"], "permanent error"); - assert_eq!(task["payload"]["attempts"], 1); - - let (status, _) = harness.post(session_id, &request(Uuid::new_v4(), "gpt-test")).await; - assert_eq!( - status, - StatusCode::ACCEPTED, - "a failed analysis no longer blocks the session" - ); -} 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..df30628d7 --- /dev/null +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -0,0 +1,523 @@ +//! `/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 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: 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}")); +} + +#[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", "deadlineAt"] { + assert_eq!(again[field], created[field], "{field}"); + } +} + +#[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 })); +} + +#[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); +} + +#[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}"); +} + +#[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"); +} + +#[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"); +} + +#[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"); +} + +#[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); +} + +#[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"); + assert_eq!( + task["payload"]["usage"], + json!({ "inputTokens": 20, "outputTokens": 40 }), + "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"); + } +} + +#[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["payload"]["reason"], "permanent error"); + assert_eq!(task["payload"]["attempts"], 1); + + 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; From e14b2836e1f025028b5ecd920b0a486570c9c3e9 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:57:53 -0400 Subject: [PATCH 27/53] fix(dgw): name the AI provider schema of the Task params as the request does The full path made the OpenAPI document reference a schema that does not exist. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/openapi.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 2354f3a6e..4978ef675 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -5,6 +5,7 @@ use uuid::Uuid; use crate::api::preflight::PreflightAlertStatus; use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; +use crate::provisioner_tasks::ai::AiProvider; #[derive(OpenApi)] #[openapi( @@ -159,7 +160,7 @@ struct SessionInfo { #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "camelCase")] struct AiAnalysisParams { - provider: crate::provisioner_tasks::ai::AiProvider, + provider: AiProvider, /// Model identifier, passed to the provider as is. model: String, /// Upper bound of tokens in each AI answer; null when not asked. From de9e7c785b84a66ecd0e9a1d82ef5daa7e2269f2 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:16:24 -0400 Subject: [PATCH 28/53] fix(dgw): fail an analysis of a deleted recording as missing, not as a failed removal A session without recording has no AI analysis to remove, so clearing it succeeds and the attempt then reports "session has no recording". Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/provisioner_tasks/recording_ai_analysis/tests.rs | 4 +--- devolutions-gateway/src/recording.rs | 4 ++-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index de482567d..9919f41cf 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -632,9 +632,7 @@ async fn recording_deleted_after_start_fails() { let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["reason"], "permanent error"); - let details = task.payload["details"].as_str().expect("details"); - assert!(details.ends_with("has no recording"), "{details}"); + assert_eq!(task.payload["details"], "session has no recording"); } #[tokio::test] diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 1207f1af6..47147c778 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -526,7 +526,7 @@ impl RecordingMessageSender { /// Removes every artifact of `kind` from the session, both its file and its manifest entry. /// - /// Fails when the session has no recording. + /// A session without recording has nothing to remove. pub(crate) async fn clear_artifacts(&self, id: Uuid, kind: ArtifactKind) -> anyhow::Result<()> { let (tx, rx) = oneshot::channel(); self.channel @@ -932,7 +932,7 @@ impl RecordingManagerTask { let manifest_path = recording_path.join("recording.json"); if !manifest_path.exists() { - anyhow::bail!("session {id} has no recording"); + return Ok(()); } let mut manifest = JrecManifest::read_from_file(&manifest_path).context("read manifest from disk")?; From 4dab45a0d663d85d3d989d6254208b37aa2bec6f Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:16:24 -0400 Subject: [PATCH 29/53] test(dgw): stop Gateway before its config folder goes Windows can't delete the database and log files a running Gateway holds open, so each test left its folder behind. Co-Authored-By: Claude Opus 5.5 (1M context) --- testsuite/tests/cli/dgw/ai_analysis.rs | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/testsuite/tests/cli/dgw/ai_analysis.rs b/testsuite/tests/cli/dgw/ai_analysis.rs index df30628d7..6786cf58e 100644 --- a/testsuite/tests/cli/dgw/ai_analysis.rs +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -36,10 +36,21 @@ fn scope_token(scope: &str) -> String { struct Gateway { config: DgwConfigHandle, - _process: tokio::process::Child, + 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() @@ -53,7 +64,7 @@ impl Gateway { Self { config, - _process: process, + process, client: reqwest::Client::new(), } } From b4e8a3c612a82bafa2143e0eb2dd8bc0ffc667eb Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:16:25 -0400 Subject: [PATCH 30/53] style(dgw): use the imported AI provider name in the OpenAPI components Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/openapi.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 4978ef675..6a414bd7f 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -70,7 +70,7 @@ use crate::provisioner_tasks::ai::AiProvider; crate::api::jrec::AiAnalysisRequest, crate::api::jrec::AiAnalysisConflict, crate::api::jrec::AiAnalysisConflictCode, - crate::provisioner_tasks::ai::AiProvider, + AiProvider, crate::api::tasks::TaskInfo, crate::api::tasks::TaskState, AiAnalysisParams, From e536bb487d90b98f3b25bd2a4a5bdd75d1aade27 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:19:48 -0400 Subject: [PATCH 31/53] test(dgw): check the reason of an analysis of a deleted recording Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/provisioner_tasks/recording_ai_analysis/tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 9919f41cf..25176806b 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -632,6 +632,7 @@ async fn recording_deleted_after_start_fails() { let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "permanent error"); assert_eq!(task.payload["details"], "session has no recording"); } From d306d458a317e850cc495ff3edfd12894b6a6693 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:37:44 -0400 Subject: [PATCH 32/53] docs(dgw): describe the maximum output tokens with the range the server accepts The server reads a `u32`, which an `int32` schema can't hold, so the clients now get a 64-bit integer from 0 to 4294967295. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/jrec/mod.rs | 1 + devolutions-gateway/src/openapi.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index 13a1bf5ce..fb50d30b9 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -454,6 +454,7 @@ pub(crate) struct AiAnalysisRequest { #[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. diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 6a414bd7f..cbb9d5036 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -164,6 +164,7 @@ struct AiAnalysisParams { /// 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")] From 366ccba67a643a89aaf8b1de2612605943addafe Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 21 Oct 2026 15:48:02 -0400 Subject: [PATCH 33/53] refactor(dgw): keep the AI access in the secret store The AI access of a Task now lives in the standalone `SecretStore`, keyed by the Task ID, instead of the provisioning store, which is about sessions only. Gateway holds one store in `DgwState`, hands it to the provisioner tasks, and runs its expiry cleanup. Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/lib.rs | 5 +++- .../src/provisioner_tasks/ai.rs | 4 ++-- .../src/provisioner_tasks/mod.rs | 11 +++------ .../recording_ai_analysis/mod.rs | 18 +++++++------- .../recording_ai_analysis/tests.rs | 24 +++++++++---------- devolutions-gateway/src/service.rs | 7 +++++- 6 files changed, 36 insertions(+), 33 deletions(-) diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index c39576552..c20eabc64 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -71,6 +71,7 @@ pub struct DgwState { 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)] @@ -98,6 +99,7 @@ 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( @@ -105,7 +107,7 @@ impl DgwState { Arc::new(job_queue_handle.clone()), conf_handle.clone(), recording_manager_handle.clone(), - provisioning.clone(), + secrets.clone(), ); let state = Self { @@ -123,6 +125,7 @@ impl DgwState { monitoring_state, agent_tunnel_handle: None, provisioner_tasks, + secrets, }; let handles = MockHandles { diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index 089805ee6..c923e5877 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -1,7 +1,7 @@ //! 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 task secret of their -//! Task in the [`ProvisioningStore`](crate::provisioning::ProvisioningStore), until the Task ends. +//! 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; diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index 0278157fd..7e900f13c 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -12,8 +12,8 @@ use provisioner_task::DynProvisionerTaskStore; use self::recording_ai_analysis::RecordingAiAnalysis; use self::runner::{DynProvisionerTaskJobQueue, ProvisionerTaskRunner}; use crate::config::ConfHandle; -use crate::provisioning::ProvisioningStore; use crate::recording::RecordingMessageSender; +use crate::secret_store::SecretStore; /// The provisioner task runner, with every provisioner task Gateway knows. pub fn task_runner( @@ -21,15 +21,10 @@ pub fn task_runner( job_queue: DynProvisionerTaskJobQueue, conf_handle: ConfHandle, recordings: RecordingMessageSender, - provisioning: ProvisioningStore, + secrets: SecretStore, ) -> ProvisionerTaskRunner { ProvisionerTaskRunner::builder(store, Arc::clone(&job_queue)) - .register(RecordingAiAnalysis::new( - conf_handle, - recordings, - provisioning, - job_queue, - )) + .register(RecordingAiAnalysis::new(conf_handle, recordings, secrets, job_queue)) .build() } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index a7a456744..4b192d1e9 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -33,8 +33,8 @@ use super::runner::{Attempt, DynProvisionerTaskJobQueue, ProvisionerTask}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; -use crate::provisioning::ProvisioningStore; 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"; @@ -215,8 +215,8 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result(task_id, target, ¶ms, |record| { // The job may run as soon as it is queued, so the key must be in place first. state - .provisioning - .insert_task_secret(record.id, &access.to_secret()?, record.deadline_at) + .secrets + .insert(SecretId(record.id), &access.to_secret()?, record.deadline_at) }) .await .map_err(StartError::Internal)?; @@ -238,7 +238,7 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result Self { Self { conf_handle, recordings, - provisioning, + secrets, job_queue, attempt_timeout: ATTEMPT_TIMEOUT, } @@ -305,8 +305,8 @@ impl ProvisionerTask for RecordingAiAnalysis { self.clear_earlier_analysis(attempt, session_id).await?; let secret = self - .provisioning - .task_secret(attempt.task_id) + .secrets + .get(SecretId(attempt.task_id)) .map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; let Some(secret) = secret else { @@ -331,7 +331,7 @@ impl ProvisionerTask for RecordingAiAnalysis { } async fn on_end(&self, record: &ProvisionerTaskRecord) { - self.provisioning.remove_task_secret(record.id); + self.secrets.remove(SecretId(record.id)); let Ok(session_id) = Uuid::parse_str(&record.target) else { return; diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 25176806b..538d50c64 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -186,9 +186,9 @@ impl Harness { fn access(&self, task_id: Uuid) -> Option { self.state - .provisioning - .task_secret(task_id) - .expect("read the task secret") + .secrets + .get(SecretId(task_id)) + .expect("read the secret") .map(|secret| AiAccess::from_secret(&secret).expect("AI access")) } @@ -203,9 +203,9 @@ impl Harness { base_url, }; self.state - .provisioning - .insert_task_secret(task_id, &access.to_secret().expect("secret"), expires_at) - .expect("insert the task secret"); + .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. @@ -691,7 +691,7 @@ async fn slow_attempt_is_retried() { let mut task = RecordingAiAnalysis::new( harness.state.conf_handle.clone(), harness.state.recordings.clone(), - harness.state.provisioning.clone(), + harness.state.secrets.clone(), Arc::new(harness.state.job_queue_handle.clone()), ); task.attempt_timeout = Duration::from_millis(500); @@ -749,7 +749,7 @@ async fn lost_key_fails_the_task() { 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.provisioning.remove_task_secret(task_id); + harness.state.secrets.remove(SecretId(task_id)); harness.run(task_id).await.expect("no retry"); @@ -851,7 +851,7 @@ async fn analysis_waits_for_the_remux_of_the_session() { .register(RecordingAiAnalysis::new( harness.state.conf_handle.clone(), harness.state.recordings.clone(), - harness.state.provisioning.clone(), + harness.state.secrets.clone(), queue, )) .build(); @@ -943,7 +943,7 @@ async fn a_new_task_that_fails_still_removes_the_earlier_analysis() { harness.run(first).await.expect("first analysis"); let second = harness.start_created(session_id, settings).await; - harness.state.provisioning.remove_task_secret(second); + harness.state.secrets.remove(SecretId(second)); harness.run(second).await.expect("no retry"); let task = harness.task(second).await; @@ -966,7 +966,7 @@ async fn a_retry_of_an_attempt_that_never_got_the_session_removes_the_earlier_an let mut task = RecordingAiAnalysis::new( harness.state.conf_handle.clone(), harness.state.recordings.clone(), - harness.state.provisioning.clone(), + harness.state.secrets.clone(), Arc::new(harness.state.job_queue_handle.clone()), ); task.attempt_timeout = Duration::from_millis(500); @@ -1108,7 +1108,7 @@ async fn an_ended_task_leaves_the_working_files_of_another_task() { let task = RecordingAiAnalysis::new( harness.state.conf_handle.clone(), harness.state.recordings.clone(), - harness.state.provisioning.clone(), + harness.state.secrets.clone(), Arc::new(harness.state.job_queue_handle.clone()), ); diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index 6a33137f7..0c6666733 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -350,12 +350,14 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { .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(), - provisioning.clone(), + secrets.clone(), ); // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. @@ -382,6 +384,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { traffic_audit_handle: traffic_audit_task.handle(), agent_tunnel_handle, provisioner_tasks: provisioner_tasks.clone(), + secrets: secrets.clone(), }; for listener in &conf.listeners { @@ -415,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, }); From 8b1637035308f84e12cdbee4f30847de48d288a9 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 21 Oct 2026 17:18:57 -0400 Subject: [PATCH 34/53] docs(dgw): name the SecretStore in the tasks intent Task secrets live in the standalone `SecretStore`, not in the provisioning store, which is about sessions only. Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/provisioner_tasks/INTENT.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index b8746192f..742db342d 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -21,7 +21,7 @@ Jobs doesn't have state, tasks do. ## Secrets Secrets a Task needs, such as an API key, never reach the disk: not `gateway.db`, not the job, not the logs. -They are kept encrypted in memory, in the provisioning store, until the Task ends or reaches its deadline. +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 From f231c4802aa3f2876f5b5e4bafa0f2fdeaf59d69 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 15:26:34 -0400 Subject: [PATCH 35/53] docs(dgw): settle the AI analysis intent Record the owner's rulings from the intent review of the AI analysis Task: a new Task replaces the earlier analysis only once its success is recorded, a failed Task leaves no file behind, attempt errors are recoverable or permanent, token usage is always reported, a running Task fails after 30 minutes without progress, and the base URL is a secret like the API key. Co-Authored-By: Claude Opus 5.5 --- crates/provisioner-task/INTENT.md | 2 +- devolutions-gateway/src/api/jrec.intent.md | 5 ++++- .../src/provisioner_tasks/INTENT.md | 15 ++++++++++++--- devolutions-gateway/src/recording.intent.md | 3 ++- 4 files changed, 19 insertions(+), 6 deletions(-) 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/devolutions-gateway/src/api/jrec.intent.md b/devolutions-gateway/src/api/jrec.intent.md index 8bcee2600..ccd9fd423 100644 --- a/devolutions-gateway/src/api/jrec.intent.md +++ b/devolutions-gateway/src/api/jrec.intent.md @@ -8,4 +8,7 @@ Jrec is the endpoint to push, streaming, read recordings and its associated file 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; a language longer than 10 characters is refused, and the language is part of the Task, so the same Task ID sent again with another language is a conflict. \ No newline at end of file +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/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index 742db342d..ca2b0010b 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -15,12 +15,16 @@ 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, never reach the disk: not `gateway.db`, not the job, not the logs. +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. @@ -28,5 +32,10 @@ After a Gateway restart they are gone, and the Tasks that needed them fail; this 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. -A session keeps one AI analysis: a new Task removes the earlier one when it starts, and a Task that fails leaves nothing, so the session then has none. Keeping one analysis per language is a future extension that must stay possible. -A transient failure such as a rate limit (429) is retried with the job, and each retry resumes where the last attempt stopped; only a new Task starts over. A Gateway restart loses the API key, so an unfinished Task fails as "key lost" at its next attempt, without a retry, and leaves nothing. \ No newline at end of file +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/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. From 04718526e913495bb70d9dde369d10e9af569ec3 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 16:53:29 -0400 Subject: [PATCH 36/53] feat(dgw): let the expiry of a stored secret move A Task keeps its API key until its deadline, and that deadline now moves while the Task makes progress, so the key must follow it. An expired or removed secret stays gone. Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/secret_store.rs | 42 +++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/devolutions-gateway/src/secret_store.rs b/devolutions-gateway/src/secret_store.rs index 650316046..c8a75e022 100644 --- a/devolutions-gateway/src/secret_store.rs +++ b/devolutions-gateway/src/secret_store.rs @@ -72,6 +72,20 @@ 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; + self.cleanup_notify.notify_one(); + } + } + /// Removes the secret of `id`. pub fn remove(&self, id: SecretId) { self.entries.lock().remove(&id); @@ -195,4 +209,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" + ); + } } From 178f802a383bc11a7249c2ab60f9d5f5eb5c2e3c Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 16:53:30 -0400 Subject: [PATCH 37/53] feat(dgw): run the AI analysis Task as its intent says Errors: an attempt fails with a closed kind and no text, recoverable or permanent. A failed Task's payload is `{ error, usage }`, so the caller words the error. Usage: every payload but `queued` carries the tokens used so far, over all attempts. Requests whose usage the provider did not report are counted apart in `unreportedRequests`. Time: a Task fails if it does not start within 2 hours. Once it runs, each attempt start and each progress report moves its deadline 30 minutes ahead, and the runner cancels an attempt past it. Waiting for a remux is not progress. The store fails an overdue Task with `deadlinePassed` and keeps the usage of its last payload. Success: the new log sits in the session unlisted until the success is recorded. Ending the Task then lists it and drops the earlier analysis. If Gateway stops in between, the job runs again at startup and redoes that step. A failed Task removes everything it made and leaves the earlier analysis in place. The session manifest is written to a file of its own, then renamed over the old one, so a crash never leaves half of it. Co-Authored-By: Claude Opus 5.5 --- .../examples/analyze_recording.rs | 14 +- .../src/recording_analysis/checkpoint.rs | 61 +--- .../src/recording_analysis/mod.rs | 344 +++++++++++------- .../src/recording_analysis/tests.rs | 110 ++++-- crates/job-queue-libsql/src/lib.rs | 1 + crates/provisioner-task-libsql/src/lib.rs | 230 ++++++++++-- crates/provisioner-task/src/lib.rs | 22 +- devolutions-gateway/src/api/jrec/mod.rs | 14 +- devolutions-gateway/src/api/tasks.rs | 4 +- devolutions-gateway/src/openapi.rs | 116 +++++- .../src/provisioner_tasks/ai.rs | 105 +++--- .../src/provisioner_tasks/mod.rs | 69 +++- .../recording_ai_analysis/mod.rs | 301 +++++++++------ .../recording_ai_analysis/tests.rs | 306 +++++++++------- .../src/provisioner_tasks/runner/mod.rs | 324 ++++++++++------- .../src/provisioner_tasks/runner/tests.rs | 165 ++++++--- devolutions-gateway/src/recording.rs | 88 ++++- testsuite/tests/cli/dgw/ai_analysis.rs | 25 +- 18 files changed, 1490 insertions(+), 809 deletions(-) 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..b033c8608 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 each chunk, after each request to the AI, and once all chunks are described. The summary request, + /// which follows when the chunks have actions, is counted by the event sent after it. 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,31 @@ 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 or a + /// request that got no answer; 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 +231,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 +382,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 +416,7 @@ impl AnalyzeRecording<'_> { max_output_tokens: self.max_output_tokens, cancel, events, + usage: Mutex::new(UsageSoFar::default()), }; tokio::spawn(async move { @@ -428,6 +458,7 @@ struct Run { max_output_tokens: Option, cancel: CancellationToken, events: mpsc::Sender, + usage: Mutex, } impl Run { @@ -461,7 +492,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 +509,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,10 +531,7 @@ impl Run { continue; } - self.report(Progress::Describing { - described_chunks: index, - total_chunks: total, - }); + self.report_describing(index, total).await; let described = match chunk { Chunk::Transcript { file } => { @@ -511,8 +539,7 @@ impl Run { .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 +553,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 +566,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) @@ -577,15 +594,14 @@ impl Run { () = cancel.cancelled() => return Err(RecordingAnalysisError::Cancelled), summary = request.send() => summary, }; + self.count(&summary, total, total).await; 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 +612,7 @@ impl Run { }; let action_count = actions.len(); - self.report(Progress::Writing); + self.report(Progress::Writing).await; blocking({ let workspace = workspace.clone(); @@ -642,12 +658,158 @@ 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 = request.send().await; + self.count(&answer, 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 = request.send().await; + self.count(&answer, 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) + } + + /// Counts the tokens of one request to the AI, then reports them with the chunks described so far. + async fn count(&self, answer: &Result, Error>, described_chunks: usize, total_chunks: usize) { + { + let mut usage = self.usage.lock().unwrap_or_else(PoisonError::into_inner); + + 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 { .. }) => return, + // 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; + } + + 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 +839,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 +875,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 +1104,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 +1132,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 +1147,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..7baff5efd 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!( @@ -263,11 +275,23 @@ 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::default(), + }, + 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: tokens(20, 40), }, ] ); @@ -298,14 +322,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 +382,6 @@ async fn working_files_of_another_analysis_are_not_reused() { &DescribedChunk { actions: Vec::new(), model: None, - usage: None, }, ) .unwrap(); @@ -541,6 +602,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 +713,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/src/lib.rs b/crates/job-queue-libsql/src/lib.rs index 04f76b63d..f8224ba38 100644 --- a/crates/job-queue-libsql/src/lib.rs +++ b/crates/job-queue-libsql/src/lib.rs @@ -473,6 +473,7 @@ mod tests { } } + // 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:") 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/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/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index fb50d30b9..396d15d00 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -496,15 +496,17 @@ fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> /// Starts an AI analysis of a recorded session /// -/// Gateway asks the AI provider what the user did in the session recording, and adds the answer to the session as a -/// new log, listed under `artifacts.ai-analysis` in its manifest. +/// 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: it removes the earlier AI analysis of the session, even a succeeded -/// one, so a session holds at most one. +/// 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( @@ -604,7 +606,8 @@ pub(crate) async fn start_ai_analysis( /// Gets an AI analysis of a recorded session /// -/// Returns the Task of the analysis: its state, its progress while it runs, and its result once it is over. +/// 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, @@ -1224,6 +1227,7 @@ 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; diff --git a/devolutions-gateway/src/api/tasks.rs b/devolutions-gateway/src/api/tasks.rs index 2431febb6..dfba09fb7 100644 --- a/devolutions-gateway/src/api/tasks.rs +++ b/devolutions-gateway/src/api/tasks.rs @@ -38,8 +38,8 @@ pub(crate) struct TaskInfo { /// When it succeeded or failed. #[serde(with = "time::serde::rfc3339::option")] finished_at: Option, - /// Latest time an unfinished Task becomes failed, waiting and retries included; it fails as `timed out` earlier - /// once too little time is left for another attempt. + /// 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, } diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index cbb9d5036..f00639d11 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -78,6 +78,11 @@ use crate::provisioner_tasks::ai::AiProvider; AiAnalysisSucceededPayload, AiAnalysisUsage, AiAnalysisFailedPayload, + AiAnalysisError, + AiAnalysisErrorKind, + AiAnalysisBadAnswer, + AiAnalysisRecoverableError, + AiAnalysisRecoverableErrorKind, crate::token::AccessScope, crate::api::webapp::AppTokenSignRequest, crate::api::webapp::AppTokenContentType, @@ -171,52 +176,133 @@ struct AiAnalysisParams { language: Option, } -/// Payload of a `running` AI analysis Task: its progress +/// 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`. + /// 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`. + /// 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 added to the session +/// 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 new log in the session manifest, such as `ai-analysis-0.slog`. + /// 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, - /// Tokens counted by the AI provider; absent when it did not report them for every request. - usage: Option, + usage: AiAnalysisUsage, /// Versions of the prompts that wrote the log, such as /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. prompt_version: String, } -/// Tokens counted by an AI provider +/// 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; 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 +/// 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 { - /// Kind of failure: `permanent error`, `attempts exhausted`, `key lost`, `timed out`, `invalid task` or `not queued`. - reason: String, - /// What went wrong: a text, except for `timed out`, where it is `{ "lastPayload": }`. - details: serde_json::Value, - /// Number of attempts started. - attempts: u32, + 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 within 2 hours, or went 30 minutes 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, } #[allow(unused)] diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index c923e5877..bf87e91a3 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -7,11 +7,11 @@ use core::fmt; use anyhow::Context as _; use devolutions_gateway_ai::recording_analysis::RecordingAnalysisError; -use devolutions_gateway_ai::{AiClient, BuildError, Provider}; +use devolutions_gateway_ai::{AiClient, BuildError, Error, Provider}; use secrecy::{ExposeSecret as _, SecretString}; use url::Url; -use super::AttemptError; +use super::{AttemptError, BadAnswer, Permanent, Recoverable}; use crate::config::Conf; /// AI provider to ask; `openai-compatible` needs a base URL @@ -194,77 +194,96 @@ impl AiSettings { impl From for AttemptError { fn from(error: RecordingAnalysisError) -> Self { match error { - RecordingAnalysisError::Cancelled => AttemptError::Cancelled, - error if error.is_transient() => AttemptError::Transient(error.to_string()), - error => AttemptError::Permanent(error.to_string()), + 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 devolutions_gateway_ai::Error; - use super::*; + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[test] - fn build_errors_map_to_settings_errors() { - assert_eq!( - AiSettingsError::from(&BuildError::MissingModel), - AiSettingsError::MissingModel - ); - assert_eq!( - AiSettingsError::from(&BuildError::MissingApiKey(Provider::OpenAi)), - AiSettingsError::MissingApiKey - ); - assert_eq!( - AiSettingsError::from(&BuildError::MissingBaseUrl(Provider::OpenAiCompatible)), - AiSettingsError::MissingBaseUrl - ); - assert!(matches!( - AiSettingsError::from(&BuildError::InvalidApiKey), - AiSettingsError::Invalid(_) - )); - } - - #[test] - fn rate_limits_server_and_network_errors_are_transient() { + fn rate_limits_server_and_network_errors_are_recoverable() { let attempt = |error: Error| AttemptError::from(RecordingAnalysisError::Ai(error)); - let status = |status| Error::Status { + let status = |status, code: Option<&str>| Error::Status { status, message: "failed".to_owned(), - code: None, + code: code.map(str::to_owned), retry_after: None, }; let network = Error::Transport { message: "connection refused".to_owned(), }; - assert!(matches!(attempt(network), AttemptError::Transient(_))); + assert_eq!(attempt(network), AttemptError::Recoverable(Recoverable::AiUnreachable)); for code in [429, 500, 503] { - assert!(matches!(attempt(status(code)), AttemptError::Transient(_)), "{code}"); + assert_eq!( + attempt(status(code, None)), + AttemptError::Recoverable(Recoverable::AiBusy { status: code }) + ); } for code in [400, 401, 403, 404] { - assert!(matches!(attempt(status(code)), AttemptError::Permanent(_)), "{code}"); + assert_eq!( + attempt(status(code, None)), + AttemptError::Permanent(Permanent::AiRejected { status: code }) + ); } - let invalid = Error::InvalidOutput { - reason: "no valid action line".to_owned(), - }; - assert!(matches!(attempt(invalid), AttemptError::Permanent(_))); - - assert!(matches!( + 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(_) - )); + AttemptError::Permanent(Permanent::AiBadAnswer { + why: BadAnswer::Truncated + }) + ); assert_eq!( - AttemptError::from(RecordingAnalysisError::Cancelled), - AttemptError::Cancelled + 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 [ diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index 7e900f13c..2c49d0330 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -32,11 +32,66 @@ pub fn task_runner( #[derive(Debug, Clone, PartialEq, Eq)] pub enum AttemptError { /// Worth another attempt later, such as a rate limit or a network error. - Transient(String), - /// Another attempt would fail the same way. - Permanent(String), - /// Fails the Task with a reason of its own, such as `key lost`, instead of `permanent error`. - Failed { reason: &'static str, details: String }, - /// The attempt was cancelled before its end, such as when it ran out of time. - Cancelled, + 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 index 4b192d1e9..d45cba0d0 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -1,11 +1,14 @@ -//! AI analysis (`recording-ai-analysis`): describes what the user did in one recorded session, and adds the result to -//! the session as a new log. +//! 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 adds the log to the session. +//! 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)] @@ -17,19 +20,19 @@ 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, + self, Progress, RecordingAnalysisEvent, RecordingFile, RecordingManifest, Recordings, RecordingsError, UsageSoFar, }; -use devolutions_gateway_ai::{Language, Usage}; -use provisioner_task::{CreateOutcome, ProvisionerTaskRecord}; +use provisioner_task::{CreateOutcome, ProvisionerTaskRecord, ProvisionerTaskState}; use secrecy::SecretString; use serde_json::json; use time::OffsetDateTime; use uuid::Uuid; -use super::AttemptError; 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; @@ -39,36 +42,59 @@ use crate::token::AccessScope; pub const KIND: &str = "recording-ai-analysis"; -/// Time from the creation of an AI analysis during which its attempts may start, waiting and retries included. -pub const ATTEMPT_START_WINDOW: time::Duration = time::Duration::hours(2); +/// Time from the creation of an AI analysis until it must start. +pub const START_WINDOW: time::Duration = time::Duration::hours(2); -/// Time one attempt of an AI analysis gets. -pub const ATTEMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60); +/// 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); -const KEY_LOST: &str = "the API key is no longer available, such as after a Gateway restart"; - -/// Tokens counted by the AI provider, as kept in Task results. +/// Tokens the AI provider counted for a Task, over all its attempts, as kept in its payloads. /// -/// It mirrors [`Usage`], so a change in the AI crate never changes what we store. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +/// 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 From for TokenUsage { - fn from(usage: Usage) -> Self { +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: usage.input_tokens, - output_tokens: usage.output_tokens, + 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 { @@ -127,11 +153,8 @@ pub struct Analysis { 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 log, answers cut and asked again included. - /// - /// Absent when the provider did not report them for every request. - #[serde(skip_serializing_if = "Option::is_none")] - usage: Option, + /// 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, @@ -240,7 +263,7 @@ pub struct RecordingAiAnalysis { recordings: RecordingMessageSender, secrets: SecretStore, job_queue: DynProvisionerTaskJobQueue, - attempt_timeout: std::time::Duration, + progress_timeout: time::Duration, } impl RecordingAiAnalysis { @@ -255,7 +278,7 @@ impl RecordingAiAnalysis { recordings, secrets, job_queue, - attempt_timeout: ATTEMPT_TIMEOUT, + progress_timeout: PROGRESS_TIMEOUT, } } @@ -268,6 +291,12 @@ impl RecordingAiAnalysis { } } +/// Where a Task leaves its log until its success is recorded: in the session folder, but not in the manifest, under a +/// name that only this Task uses, so that removing the working files of another analysis never removes it. +fn unlisted_log_path(manifest_path: &Utf8Path, task_id: Uuid) -> Utf8PathBuf { + manifest_path.with_file_name(format!(".ai-analysis-{task_id}.slog")) +} + #[async_trait] impl ProvisionerTask for RecordingAiAnalysis { const KIND: &'static str = KIND; @@ -280,44 +309,66 @@ impl ProvisionerTask for RecordingAiAnalysis { type Output = Analysis; - fn attempt_start_window(&self) -> time::Duration { - ATTEMPT_START_WINDOW + fn start_window(&self) -> time::Duration { + START_WINDOW } - fn attempt_timeout(&self) -> std::time::Duration { - self.attempt_timeout + 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| AttemptError::Failed { - reason: "invalid task", - details: format!("invalid session ID: {error}"), + 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) })?; - // Waiting for an earlier analysis of the session counts in the attempt time, so a timeout cancels it too. + // The API key is kept until the deadline, which the start of this attempt moved. + 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::Cancelled), + () = attempt.cancel.cancelled() => return Err(AttemptError::Permanent(Permanent::DeadlinePassed)), }; - // A new Task is a hard restart, so the session loses its earlier analysis even when the Task then fails, such as - // on a lost key. Every attempt clears, since an earlier one may have ended before it could. - self.clear_earlier_analysis(attempt, session_id).await?; + // A new Task never resumes from what an earlier analysis left, while its retries resume from their own files. + if attempt.attempts == 1 { + recording_analysis::discard(&self.manifest_path(session_id), 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) + })?; + } - let secret = self - .secrets - .get(SecretId(attempt.task_id)) - .map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + 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::Failed { - reason: "key lost", - details: KEY_LOST.to_owned(), - }); + return Err(AttemptError::Permanent(Permanent::KeyLost)); }; - let access = AiAccess::from_secret(&secret).map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + 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, @@ -330,43 +381,39 @@ impl ProvisionerTask for RecordingAiAnalysis { .await } - async fn on_end(&self, record: &ProvisionerTaskRecord) { + 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; + 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; - if let Err(error) = recording_analysis::discard(&manifest_path, Some(&record.id.to_string())).await { - warn!(task.id = %record.id, session.id = %session_id, %error, "Failed to remove the AI analysis working files"); - } - } -} - -impl RecordingAiAnalysis { - /// Removes the analysis an earlier Task added to the session, and on the first attempt the working files any - /// analysis left; later attempts are job retries, which resume from their own working files. - async fn clear_earlier_analysis(&self, attempt: &Attempt, session_id: Uuid) -> Result<(), AttemptError> { - let cleared = async { + if record.state == ProvisionerTaskState::Succeeded { + // The success is recorded, so the log takes the place of the earlier analysis. self.recordings - .clear_artifacts(session_id, ArtifactKind::AiAnalysis) - .await?; - - if attempt.attempts == 1 { - recording_analysis::discard(&self.manifest_path(session_id), None).await?; + .replace_artifacts(session_id, ArtifactKind::AiAnalysis, &log_path) + .await + .context("list the AI analysis log in the session")?; + } else { + match tokio::fs::remove_file(&log_path).await { + Err(error) if error.kind() != std::io::ErrorKind::NotFound => { + return Err(anyhow::Error::new(error).context(format!("remove {log_path}"))); + } + _ => {} } + } - anyhow::Ok(()) - }; - - cleared + recording_analysis::discard(&manifest_path, Some(&record.id.to_string())) .await - .map_err(|error| AttemptError::Permanent(format!("failed to remove the earlier AI analysis: {error:#}"))) + .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, @@ -378,18 +425,27 @@ impl RecordingAiAnalysis { ) -> Result { let conf = self.conf_handle.get_conf(); - let client = settings - .client(&conf, &api_key) - .map_err(|error| AttemptError::Permanent(error.to_string()))?; + 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| AttemptError::Permanent(format!("{error:#}")))?; + .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| AttemptError::Permanent(error.to_string()))?; + 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?; @@ -407,8 +463,10 @@ impl RecordingAiAnalysis { request = request.language(language); } + let earlier = TokenUsage::of(&attempt.payload); + let mut usage = earlier; let mut events = request.start(); - let mut recorded: Option<(&'static str, u64)> = None; + let mut recorded = None; let analysis = loop { match events.recv().await { @@ -416,35 +474,59 @@ impl RecordingAiAnalysis { // Reading sends an event every half second, so while reading the Task is updated only when the percent // changes. let shown = match progress { - Progress::Reading { .. } => (progress.stage(), u64::from(progress.percent())), - _ => (progress.stage(), progress.done_and_total().map_or(0, |(done, _)| done)), + Progress::Reading { .. } => (progress.stage(), u64::from(progress.percent()), None), + Progress::Describing { usage, .. } => ( + progress.stage(), + progress.done_and_total().map_or(0, |(done, _)| done), + Some(usage), + ), + _ => ( + progress.stage(), + progress.done_and_total().map_or(0, |(done, _)| done), + None, + ), }; + if recorded != Some(shown) { recorded = Some(shown); - attempt.report_progress(progress_payload(&progress)).await; + + if let Some(attempt_usage) = shown.2 { + 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)) => return Err(error.into()), + Some(RecordingAnalysisEvent::Failure(error)) => { + warn!(task.id = %attempt.task_id, %error, "AI analysis attempt failed"); + return Err(error.into()); + } Some(_) => {} None => { - return Err(AttemptError::Permanent( - "the AI analysis stopped unexpectedly".to_owned(), - )); + error!(task.id = %attempt.task_id, "The AI analysis stopped unexpectedly"); + return Err(AttemptError::Permanent(Permanent::Panicked)); } } }; - let file_name = add_log(&self.recordings, session_id, &analysis.log_path) + // 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| AttemptError::Permanent(format!("failed to add the log: {error:#}")))?; + .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, file = %file_name, actions = analysis.actions, "Session log added"); + debug!(task.id = %attempt.task_id, actions = analysis.actions, "AI analysis log written"); Ok(Analysis { - file_name, + file_name: ArtifactKind::AiAnalysis.file_name(0), model: analysis.model, - usage: analysis.usage.map(TokenUsage::from), + usage: earlier.with(analysis.usage), prompt_version: analysis.prompt_version, }) } @@ -454,11 +536,10 @@ impl RecordingAiAnalysis { /// 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| AttemptError::Transient(format!("failed to list the remux jobs: {error:#}")))?; + 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(()); @@ -468,20 +549,21 @@ impl RecordingAiAnalysis { tokio::select! { () = tokio::time::sleep(REMUX_POLL_INTERVAL) => {} - () = attempt.cancel.cancelled() => return Err(AttemptError::Cancelled), + () = attempt.cancel.cancelled() => return Err(AttemptError::Permanent(Permanent::DeadlinePassed)), } } } } -/// The payload of a running AI analysis Task: its progress, for whoever asks about it. -fn progress_payload(progress: &Progress) -> serde_json::Value { +/// 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!({ - "step": progress.stage(), - "done": done, - "total": total, + json!(Running { + step: progress.stage(), + done, + total, + usage, }) } @@ -524,18 +606,3 @@ fn analysis_manifest(manifest: &JrecManifest) -> Result anyhow::Result { - let artifact_path = recordings.add_artifact(session_id, ArtifactKind::AiAnalysis).await?; - - // The log is in the session folder, so a rename moves it in at once and the session never holds a partial one. - tokio::fs::rename(log_path, &artifact_path) - .await - .with_context(|| format!("move the log to {artifact_path}"))?; - - Ok(artifact_path - .file_name() - .context("artifact path has no file name")? - .to_owned()) -} diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 538d50c64..dae881fdb 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -179,6 +179,23 @@ impl Harness { .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") @@ -313,6 +330,7 @@ fn first_line(path: &Utf8Path) -> String { .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; @@ -334,10 +352,7 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { 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, - ATTEMPT_START_WINDOW + ATTEMPT_TIMEOUT - ); + 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/"); @@ -351,6 +366,7 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { 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; @@ -392,6 +408,7 @@ async fn asking_again_creates_nothing_new() { 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; @@ -420,6 +437,7 @@ async fn asking_again_returns_the_task_even_while_the_session_records_again() { 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; @@ -443,6 +461,7 @@ async fn refuses_a_session_without_recording() { 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; @@ -467,6 +486,7 @@ async fn refuses_a_session_that_is_still_recording() { ); } +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis #[tokio::test] async fn refuses_invalid_ai_settings() { let harness = Harness::new().await; @@ -495,6 +515,7 @@ async fn refuses_invalid_ai_settings() { 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; @@ -514,6 +535,7 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { "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"); @@ -537,10 +559,10 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { json!({ "fileName": "ai-analysis-0.slog", "model": "gpt-test-2026-09-30", - "usage": { "inputTokens": 40, "outputTokens": 80 }, + "usage": { "inputTokens": 40, "outputTokens": 80, "unreportedRequests": 0 }, "promptVersion": format!("{PROMPT_VERSION}+{}", session_summary::PROMPT_VERSION), }), - "chunk 0 counts from its checkpoint, the failed request not at all" + "every attempt counts, and the refused request costs nothing" ); let requests = requests.lock().clone(); @@ -574,6 +596,7 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { 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; @@ -585,18 +608,20 @@ async fn permanent_error_fails_without_retry() { let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["reason"], "permanent error"); - assert!( - task.payload["details"].as_str().expect("details").contains("401"), - "{}", - task.payload + assert_eq!( + task.payload, + json!({ + "error": { "kind": "aiRejected", "status": 401 }, + "usage": { "inputTokens": 0, "outputTokens": 0, "unreportedRequests": 0 }, + }) ); - assert_eq!(task.payload["attempts"], 1); + 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; @@ -618,9 +643,10 @@ async fn a_session_mixing_terminal_and_video_fails() { let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["reason"], "permanent error"); - assert_eq!(task.payload["details"], "session mixes terminal and video recordings"); + 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; @@ -632,10 +658,10 @@ async fn recording_deleted_after_start_fails() { let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["reason"], "permanent error"); - assert_eq!(task.payload["details"], "session has no recording"); + 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; @@ -653,23 +679,22 @@ async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); assert_eq!(task.attempts, MAX_ATTEMPTS); - assert_eq!(task.payload["reason"], "attempts exhausted"); - assert!( - task.payload["details"].as_str().expect("details").contains("503"), - "{}", - task.payload + 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 slow_attempt_is_retried() { +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 makes the attempt outlast its timeout. + // 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 { @@ -687,60 +712,24 @@ async fn slow_attempt_is_retried() { max_output_tokens: None, }; let task_id = harness.start_created(session_id, settings).await; + harness.with_progress_timeout(time::Duration::seconds(2)); - let mut 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.attempt_timeout = Duration::from_millis(500); - harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( - Arc::clone(harness.state.provisioner_tasks.store()), - Arc::new(harness.state.job_queue_handle.clone()), - ) - .register(task) - .build(); - - let error = harness.run(task_id).await.expect_err("timeout is retried"); - - assert!(error.to_string().contains("took longer than"), "{error:#}"); - assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Running); - assert!(harness.has_key(task_id)); - assert!( - harness.workspace(session_id).join("plan.json").exists(), - "the chunks are kept for the retry" - ); -} - -#[tokio::test] -async fn task_without_time_for_an_attempt_fails_without_calling_the_ai() { - 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 now = OffsetDateTime::now_utc(); - harness - .create_task( - task_id, - session_id, - "openai-compatible", - now, - now + ATTEMPT_TIMEOUT - time::Duration::minutes(1), - ) - .await; - harness.insert_key(task_id, settings.base_url, now + time::Duration::hours(1)); - - harness.job(task_id).await.run().await.expect("no retry"); + 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["reason"], "timed out"); - assert_eq!(task.attempts, 0); - assert!(requests.lock().is_empty()); + assert_eq!( + task.payload, + json!({ + "error": { "kind": "deadlinePassed" }, + "usage": { "inputTokens": 0, "outputTokens": 0, "unreportedRequests": 0 }, + }) + ); + 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; @@ -757,11 +746,15 @@ async fn lost_key_fails_the_task() { assert_eq!(task.state, ProvisionerTaskState::Failed); assert_eq!( task.payload, - json!({ "reason": "key lost", "details": KEY_LOST, "attempts": 1 }) + 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; @@ -784,12 +777,17 @@ async fn overdue_task_does_not_run() { let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["reason"], "timed out"); + 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; @@ -829,6 +827,7 @@ impl ProvisionerTaskJobQueue for WithRemuxJobs { } } +// 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; @@ -871,6 +870,42 @@ async fn analysis_waits_for_the_remux_of_the_session() { 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(); @@ -890,6 +925,7 @@ async fn analyses_of_one_session_run_one_at_a_time() { .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; @@ -909,6 +945,7 @@ async fn working_files_left_by_an_earlier_task_go_with_the_task_that_ends() { 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; @@ -934,8 +971,9 @@ async fn a_new_task_replaces_the_result_of_the_earlier_one() { 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_still_removes_the_earlier_analysis() { +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; @@ -948,52 +986,79 @@ async fn a_new_task_that_fails_still_removes_the_earlier_analysis() { let task = harness.task(second).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["reason"], "key lost"); + assert_eq!(task.payload["error"], json!({ "kind": "keyLost" })); let dir = harness.recording_path.join(session_id.to_string()); - assert_eq!(ai_analyses(&dir), json!([])); - assert!(!dir.join("ai-analysis-0.slog").exists()); + 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_retry_of_an_attempt_that_never_got_the_session_removes_the_earlier_analysis() { - let mut harness = Harness::new().await; +async fn a_failed_task_reports_the_tokens_it_spent() { + 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; + // 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; - let mut task = RecordingAiAnalysis::new( - harness.state.conf_handle.clone(), - harness.state.recordings.clone(), - harness.state.secrets.clone(), - Arc::new(harness.state.job_queue_handle.clone()), + 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 }) ); - task.attempt_timeout = Duration::from_millis(500); - harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( - Arc::clone(harness.state.provisioner_tasks.store()), - Arc::new(harness.state.job_queue_handle.clone()), - ) - .register(task) - .build(); +} - // Another analysis of the session holds it past the attempt timeout. - let session = lock_session(session_id).await; - let error = harness.run(second).await.expect_err("timeout is retried"); - assert!(error.to_string().contains("took longer than"), "{error:#}"); - drop(session); +// 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(second).await.expect("second attempt"); + harness.run(task_id).await.expect("the job ends the Task"); - let task = harness.task(second).await; - assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); - assert_eq!(task.attempts, 2); + 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()); - assert_eq!( - ai_analyses(&dir), - json!([{ "fileName": "ai-analysis-0.slog" }]), - "one analysis per recording" - ); + 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()); } /// The AI analyses the manifest in `dir` lists. @@ -1006,6 +1071,7 @@ fn ai_analyses(dir: &Utf8Path) -> serde_json::Value { .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; @@ -1041,6 +1107,7 @@ async fn the_analysis_is_in_the_language_of_the_request() { ); } +// 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; @@ -1056,6 +1123,7 @@ async fn the_analysis_is_in_english_without_a_language() { 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; @@ -1084,6 +1152,7 @@ async fn a_new_task_never_resumes_from_an_earlier_task() { ); } +// 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; @@ -1112,27 +1181,8 @@ async fn an_ended_task_leaves_the_working_files_of_another_task() { Arc::new(harness.state.job_queue_handle.clone()), ); - task.on_end(&record).await; + task.on_end(&record).await.expect("end"); assert!(workspace.exists(), "the newer Task still resumes from them"); assert!(!harness.has_key(task_id)); } - -#[test] -fn unknown_usage_is_left_out_of_the_result() { - let analysis = Analysis { - file_name: "ai-analysis-0.slog".to_owned(), - model: "gpt-test-2026-09-30".to_owned(), - usage: None, - prompt_version: PROMPT_VERSION.to_owned(), - }; - - assert_eq!( - json!(analysis), - json!({ - "fileName": "ai-analysis-0.slog", - "model": "gpt-test-2026-09-30", - "promptVersion": PROMPT_VERSION, - }) - ); -} diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 2b7e45429..51e287436 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -1,8 +1,10 @@ //! 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, and what to clean up once it has ended. +//! 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 transient failures, enforces the attempt timeout and records how the task ended. +//! 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, @@ -17,6 +19,7 @@ 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, }; @@ -27,14 +30,10 @@ use time::OffsetDateTime; use tokio_util::sync::CancellationToken; use uuid::Uuid; -use super::AttemptError; +use super::{AttemptError, Permanent}; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; use crate::token::AccessScope; -/// Time an attempt may need past its timeout: it ends a little later while it stops, and a stored deadline loses its -/// fraction of a second. -pub(crate) const ATTEMPT_END_GRACE: std::time::Duration = std::time::Duration::from_secs(60); - /// Where the runner queues the jobs of its records. #[async_trait] pub trait ProvisionerTaskJobQueue: Send + Sync { @@ -81,23 +80,34 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// The payload of a `succeeded` record. type Output: Serialize + Send; - /// Time from the creation of a record during which its attempts may start, the first one and retries alike. - /// - /// The record's deadline comes one [`Self::attempt_timeout`] later, so an attempt started in time can finish. - fn attempt_start_window(&self) -> time::Duration; + /// 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; - /// Time one attempt gets before it is cancelled and retried. - fn attempt_timeout(&self) -> std::time::Duration; + /// 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, and return [`AttemptError::Cancelled`]. + /// 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; - /// Drops what is kept only while the record runs, once it has ended, however it ended. + /// 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. - async fn on_end(&self, _record: &ProvisionerTaskRecord) {} + /// It may be called again for a record that already ended, so it must be idempotent; an error has the job call it + /// again later. + async fn on_end(&self, _record: &ProvisionerTaskRecord) -> anyhow::Result<()> { + Ok(()) + } } /// One attempt of a record, as its task sees it. @@ -106,19 +116,45 @@ pub struct Attempt { 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 { - /// Replaces the payload of the running record, so whoever asks about it sees how the attempt goes. + /// Deadline of the record, after which the attempt is cancelled. + pub fn deadline(&self) -> OffsetDateTime { + self.latest.lock().deadline_at + } + + /// 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) { - if let Err(error) = self - .store - .update_running(self.task_id, payload, OffsetDateTime::now_utc()) - .await - { - warn!(task.id = %self.task_id, error = format!("{error:#}"), "Failed to record the Task progress"); + let now = OffsetDateTime::now_utc(); + let deadline_at = 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"); + } } } } @@ -126,38 +162,53 @@ impl Attempt { /// [`ProvisionerTask`] with its types erased, so the runner can hold every task together. #[async_trait] trait AnyProvisionerTask: Send + Sync { - fn attempt_start_window(&self) -> time::Duration; + fn start_window(&self) -> time::Duration; + + fn progress_timeout(&self) -> time::Duration; - fn attempt_timeout(&self) -> std::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); + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()>; } #[async_trait] impl AnyProvisionerTask for T { - fn attempt_start_window(&self) -> time::Duration { - ProvisionerTask::attempt_start_window(self) + 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 attempt_timeout(&self) -> std::time::Duration { - ProvisionerTask::attempt_timeout(self) + 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| AttemptError::Failed { - reason: "invalid task", - details: format!("invalid parameters: {error}"), + 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| AttemptError::Permanent(format!("invalid result: {error}"))) + 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) { - ProvisionerTask::on_end(self, record).await; + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { + ProvisionerTask::on_end(self, record).await } } @@ -199,11 +250,6 @@ impl ProvisionerTaskRunnerBuilder { } } -/// Initial payload of a running record, until its task reports progress. -fn starting_payload() -> serde_json::Value { - json!({ "step": "preparing" }) -} - impl ProvisionerTaskRunner { pub fn builder( store: DynProvisionerTaskStore, @@ -247,7 +293,7 @@ impl ProvisionerTaskRunner { kind: T::KIND.to_owned(), target, params, - deadline_at: now + task.attempt_start_window() + task.attempt_timeout(), + deadline_at: now + task.start_window(), job_token, }, now, @@ -265,7 +311,7 @@ impl ProvisionerTaskRunner { 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, &error, 0).await; + self.fail_not_queued(task.as_ref(), id, &serde_json::Value::Null).await; return Err(error); } @@ -325,8 +371,7 @@ impl ProvisionerTaskRunner { "Failed to queue the Task again" ); if let Some(task) = self.task(&record.kind) { - self.fail_not_queued(task.as_ref(), record.id, &error, record.attempts) - .await; + self.fail_not_queued(task.as_ref(), record.id, &record.payload).await; } } } @@ -361,30 +406,35 @@ impl ProvisionerTaskRunner { let store = &self.inner.store; let now = OffsetDateTime::now_utc(); - // An attempt that cannot finish before the deadline would do work the record never keeps, such as paid AI calls. - if let Some(record) = store.get(task_id, now).await? - && record.job_token == job_token - && !record.state.is_finished() - && let Some(task) = self.task(&record.kind) - && now + task.attempt_timeout() + ATTEMPT_END_GRACE > record.deadline_at - { - warn!(task.id = %task_id, task.kind = %record.kind, attempts = record.attempts, "Task timed out before its next attempt"); - let payload = json!({ - "reason": "timed out", - "details": { "lastPayload": record.payload }, - "attempts": record.attempts, - }); - return self.fail(task.as_ref(), task_id, payload).await; - } + // 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 = now + task.progress_timeout(); - let record = match store.start_attempt(task_id, job_token, starting_payload(), now).await? { + 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"); - if let Some(task) = self.task(&record.kind) { - task.on_end(&record).await; - } - return Ok(()); + return task.on_end(&record).await.context("end the Task"); } AttemptStart::Superseded => { debug!(task.id = %task_id, "Another job runs this Task, so this one stops"); @@ -396,27 +446,21 @@ impl ProvisionerTaskRunner { } }; - let Some(task) = self.task(&record.kind) else { - error!(task.id = %task_id, task.kind = %record.kind, "Task of an unknown kind"); - let details = format!("unknown task kind {}", record.kind); - if let Err(error) = store - .fail(task_id, failure("invalid task", &details, record.attempts), now) - .await - { - error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure"); - } - return Ok(()); - }; - let attempts = record.attempts; - let attempt_timeout = task.attempt_timeout(); 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"); @@ -428,24 +472,38 @@ impl ProvisionerTaskRunner { async move { task.run(&attempt, params).await } }); - let joined = match tokio::time::timeout(attempt_timeout, &mut handle).await { - Ok(joined) => joined, - Err(_) => { - // Aborting cannot stop blocking steps, so the attempt is cancelled and awaited instead: the next - // attempt never runs alongside it. - cancel.cancel(); - handle.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 { - Ok(Err(AttemptError::Cancelled)) => Err(AttemptError::Transient(format!( - "the attempt took longer than {attempt_timeout:?}" - ))), + // 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("the attempt panicked".to_owned())), + 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) => { @@ -459,46 +517,46 @@ impl ProvisionerTaskRunner { ); } - self.end(task.as_ref(), task_id).await; + self.end(task.as_ref(), task_id).await?; + } + // A retry would do the work again, such as paying for a second analysis, so the Task fails instead, + // which drops what the attempt made. + Err(error) => { + error!( + task.id = %task_id, + error = format!("{error:#}"), + "Failed to record the Task success; the Task fails instead" + ); + + if let Err(error) = self.fail(task.as_ref(), task_id, &Permanent::Internal, &last).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to fail the Task; it fails at its deadline"); + } } - // No retry: the work is done, and doing it again would repeat it, such as adding a second log. Nothing is - // dropped either, since the record still runs; it fails at its deadline. - Err(error) => error!( - task.id = %task_id, - error = format!("{error:#}"), - "Failed to record the Task result; the Task fails at its deadline" - ), }, - Err(AttemptError::Transient(details)) if attempts < MAX_ATTEMPTS => { - warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task attempt failed; it will be retried"); - anyhow::bail!("Task attempt failed: {details}"); + 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::Transient(details)) => { - warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task failed after its last attempt"); + 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, - failure("attempts exhausted", &details, attempts), + &Permanent::AttemptsExhausted { last: error }, + &last, ) .await?; } - Err(AttemptError::Permanent(details)) => { - warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task failed"); - self.fail(task.as_ref(), task_id, failure("permanent error", &details, attempts)) - .await?; - } - Err(AttemptError::Failed { reason, details }) => { - warn!(task.id = %task_id, task.kind = %record.kind, attempts, reason, %details, "Task failed"); - self.fail(task.as_ref(), task_id, failure(reason, &details, attempts)) - .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, &error, &last).await?; } - Err(AttemptError::Cancelled) => unreachable!("a cancelled attempt is reported as a timeout above"), } Ok(()) } - /// Records the failure, then drops what the record kept while running. + /// 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. @@ -506,47 +564,33 @@ impl ProvisionerTaskRunner { &self, task: &dyn AnyProvisionerTask, task_id: Uuid, - payload: serde_json::Value, + error: &Permanent, + last: &serde_json::Value, ) -> anyhow::Result<()> { self.inner .store - .fail(task_id, payload, OffsetDateTime::now_utc()) + .fail(task_id, task.failed_payload(error, last), OffsetDateTime::now_utc()) .await .context("record the Task failure")?; - self.end(task, task_id).await; - - Ok(()) + self.end(task, task_id).await } /// Fails a record no job will run. - async fn fail_not_queued( - &self, - task: &dyn AnyProvisionerTask, - task_id: Uuid, - error: &anyhow::Error, - attempts: u32, - ) { - let payload = failure("not queued", &format!("{error:#}"), attempts); - - if let Err(error) = self.fail(task, task_id, payload).await { - error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure; it fails at its deadline"); + async fn fail_not_queued(&self, task: &dyn AnyProvisionerTask, task_id: Uuid, last: &serde_json::Value) { + if let Err(error) = self.fail(task, task_id, &Permanent::NotQueued, last).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to end the Task that has no job"); } } - async fn end(&self, task: &dyn AnyProvisionerTask, task_id: Uuid) { - match self.inner.store.get(task_id, OffsetDateTime::now_utc()).await { - Ok(Some(record)) => task.on_end(&record).await, - Ok(None) => {} - Err(error) => error!(task.id = %task_id, error = format!("{error:#}"), "Failed to read the ended Task"), + async fn end(&self, task: &dyn AnyProvisionerTask, task_id: Uuid) -> anyhow::Result<()> { + match self.inner.store.get(task_id, OffsetDateTime::now_utc()).await? { + Some(record) => task.on_end(&record).await.context("end the Task"), + None => Ok(()), } } } -fn failure(reason: &str, details: &str, attempts: u32) -> serde_json::Value { - json!({ "reason": reason, "details": details, "attempts": attempts }) -} - #[derive(Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] struct JobDef { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index a69558afd..2db72f6b4 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -5,10 +5,17 @@ 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, @@ -17,6 +24,7 @@ struct EchoParams { /// 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>>, } @@ -33,21 +41,44 @@ impl ProvisionerTask for Echo { type Output = serde_json::Value; - fn attempt_start_window(&self) -> time::Duration { + fn start_window(&self) -> time::Duration { time::Duration::hours(1) } - fn attempt_timeout(&self) -> std::time::Duration { - std::time::Duration::from_secs(5) + fn progress_timeout(&self) -> time::Duration { + self.progress_timeout + } + + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value { + json!({ "after": last }) } - async fn run(&self, _attempt: &Attempt, params: EchoParams) -> Result { + 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); - (self.outcome)(params.value) + + 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) { + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { self.ended.lock().push(record.id); + Ok(()) } } @@ -60,6 +91,13 @@ struct Harness { 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") @@ -78,6 +116,7 @@ impl Harness { let runner = ProvisionerTaskRunner::builder(store, Arc::new(job_queue)) .register(Echo { outcome, + progress_timeout, runs: Arc::clone(&runs), ended: Arc::clone(&ended), }) @@ -116,6 +155,7 @@ impl Harness { } } +// 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; @@ -138,6 +178,7 @@ async fn only_a_new_task_runs_on_created_and_gets_a_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; @@ -152,10 +193,11 @@ async fn a_task_whose_job_cannot_be_queued_fails() { let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["reason"], "not queued"); + 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; @@ -171,11 +213,12 @@ async fn a_task_that_cannot_be_prepared_fails_without_a_job() { let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["reason"], "not queued"); + 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; @@ -207,6 +250,7 @@ async fn a_superseded_job_runs_nothing_once_the_task_is_resumed() { ); } +// 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; @@ -218,7 +262,7 @@ async fn a_task_that_cannot_be_queued_again_fails() { let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed, "no job is left to run it"); - assert_eq!(task.payload["reason"], "not queued"); + assert_eq!(task.payload["error"], json!({ "kind": "notQueued" })); assert_eq!(*harness.ended.lock(), [id]); } @@ -239,6 +283,7 @@ impl ProvisionerTaskJobQueue for StoredJobs { } } +// 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:") @@ -256,6 +301,7 @@ async fn resume_queues_only_the_tasks_left_without_a_job() { 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(), }) @@ -292,6 +338,7 @@ async fn resume_queues_only_the_tasks_left_without_a_job() { ); } +// 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; @@ -303,30 +350,49 @@ async fn a_panicking_attempt_fails_the_task() { assert_eq!(task.state, ProvisionerTaskState::Failed); assert_eq!( task.payload, - json!({ "reason": "permanent error", "details": "the attempt panicked", "attempts": 1 }) + 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_failure_keeps_the_reason_of_its_kind() { - let mut harness = Harness::new(|_| { - Err(AttemptError::Failed { - reason: "key lost", - details: "gone".to_owned(), - }) - }) - .await; +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!( - harness.task(id).await.payload, - json!({ "reason": "key lost", "details": "gone", "attempts": 1 }) + 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; @@ -355,55 +421,40 @@ async fn invalid_parameters_and_unknown_kinds_fail_the_task() { let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed, "{kind}"); - assert_eq!(task.payload["reason"], "invalid task", "{kind}"); + assert_eq!(task.payload["error"], json!({ "kind": "invalidTask" }), "{kind}"); } assert_eq!(harness.runs.load(Ordering::SeqCst), 0); } +// oracle: intent crates/provisioner-task/INTENT.md §Invariants #[tokio::test] -async fn an_attempt_starts_only_if_it_can_finish_before_the_deadline() { - let harness = Harness::new(|value| Ok(json!(value))).await; - let store = harness.runner.store(); - let now = OffsetDateTime::now_utc(); +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; - // The echo attempt timeout is 5 seconds; 2 more or 2 less keep the test clear of whole-second rounding. - let needed = time::Duration::seconds(5) + time::Duration::try_from(ATTEMPT_END_GRACE).expect("grace"); - for (left, state) in [ - (needed + time::Duration::seconds(2), ProvisionerTaskState::Succeeded), - (needed - time::Duration::seconds(2), ProvisionerTaskState::Failed), - ] { - let id = Uuid::new_v4(); - let job_token = Uuid::new_v4(); - store - .create( - NewProvisionerTaskRecord { - id, - kind: KIND.to_owned(), - target: id.to_string(), - params: json!({ "value": 1 }), - deadline_at: now + left, - job_token, - }, - now, - ) - .await - .expect("create"); + harness.queued_job().run().await.expect("no retry"); - harness.runner.run_attempt(id, job_token).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]); +} - let task = harness.task(id).await; - assert_eq!(task.state, state, "{left} seconds left"); - if state == ProvisionerTaskState::Failed { - assert_eq!(task.payload["reason"], "timed out"); - assert_eq!(task.attempts, 0); - assert!(harness.ended.lock().contains(&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(2)).await; + let id = harness.create(PROGRESS).await; - assert_eq!(harness.runs.load(Ordering::SeqCst), 1); + 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; diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 47147c778..8666ee9de 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -79,9 +79,22 @@ impl JrecManifest { 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(); + let partial = path.with_extension("json.partial"); 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())?; + file.sync_all()?; + drop(file); + + std::fs::rename(&partial, path)?; + Ok(()) } } @@ -382,9 +395,10 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, - ClearArtifacts { + ReplaceArtifacts { id: Uuid, kind: ArtifactKind, + file: Utf8PathBuf, channel: oneshot::Sender>, }, GetFinished { @@ -453,10 +467,16 @@ impl fmt::Debug for RecordingManagerMessage { .field("id", id) .field("kind", kind) .finish_non_exhaustive(), - RecordingManagerMessage::ClearArtifacts { id, kind, channel: _ } => f - .debug_struct("ClearArtifacts") + 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() @@ -524,17 +544,26 @@ impl RecordingMessageSender { rx.await.context("couldn't receive AddArtifact result")? } - /// Removes every artifact of `kind` from the session, both its file and its manifest entry. + /// Moves `file` into the session as its only artifact of `kind`, named `kind.file_name(0)`, in place of the earlier + /// ones. /// - /// A session without recording has nothing to remove. - pub(crate) async fn clear_artifacts(&self, id: Uuid, kind: ArtifactKind) -> anyhow::Result<()> { + /// `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, 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::ClearArtifacts { id, kind, channel: tx }) + .send(RecordingManagerMessage::ReplaceArtifacts { + id, + kind, + file: file.to_owned(), + channel: tx, + }) .await .ok() - .context("couldn't send ClearArtifacts message")?; - rx.await.context("couldn't receive ClearArtifacts result")? + .context("couldn't send ReplaceArtifacts message")?; + rx.await.context("couldn't receive ReplaceArtifacts result")? } pub(crate) async fn get_finished(&self, id: Uuid) -> Result { @@ -927,31 +956,49 @@ impl RecordingManagerTask { Ok(artifact_path) } - async fn handle_clear_artifacts(&self, id: Uuid, kind: ArtifactKind) -> anyhow::Result<()> { + 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() { - return Ok(()); + 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 artifacts = core::mem::take(manifest.artifacts.of_kind_mut(kind)); - if artifacts.is_empty() { - return Ok(()); + 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(()) => {} + // An earlier call moved it, then stopped before the manifest was written. + Err(error) if error.kind() == io::ErrorKind::NotFound && artifact_path.exists() => {} + Err(error) => return Err(anyhow::Error::new(error).context(format!("move {file} to {artifact_path}"))), } - // The manifest goes first, so it never lists a file that is gone. + 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 artifacts { + 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 => { - return Err(anyhow::Error::new(error).context(format!("remove {artifact_path}"))); + warn!(%error, %artifact_path, "Failed to remove an earlier artifact"); } _ => {} } @@ -1087,8 +1134,8 @@ async fn recording_manager_task( RecordingManagerMessage::AddArtifact { id, kind, channel } => { let _ = channel.send(manager.handle_add_artifact(id, kind).await); }, - RecordingManagerMessage::ClearArtifacts { id, kind, channel } => { - let _ = channel.send(manager.handle_clear_artifacts(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)); @@ -1585,6 +1632,7 @@ 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(); diff --git a/testsuite/tests/cli/dgw/ai_analysis.rs b/testsuite/tests/cli/dgw/ai_analysis.rs index 6786cf58e..592efd58e 100644 --- a/testsuite/tests/cli/dgw/ai_analysis.rs +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -278,6 +278,7 @@ fn assert_rfc3339(value: &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; @@ -305,11 +306,12 @@ async fn start_returns_202_then_200_for_the_same_request() { 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", "deadlineAt"] { + 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; @@ -331,6 +333,7 @@ async fn start_conflicts_return_409_with_a_code() { 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; @@ -343,6 +346,7 @@ async fn start_without_recording_returns_404() { 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; @@ -356,6 +360,7 @@ async fn start_ignores_unknown_fields() { 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; @@ -404,6 +409,7 @@ async fn invalid_bodies_and_settings_return_400() { 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; @@ -432,6 +438,7 @@ async fn endpoints_require_their_scope() { 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; @@ -459,6 +466,7 @@ async fn an_analysis_is_found_at_its_location_only() { 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; @@ -468,6 +476,7 @@ async fn get_analysis_returns_404_for_an_unknown_id() { 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; @@ -485,9 +494,10 @@ async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { 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 }), + json!({ "inputTokens": 20, "outputTokens": 40, "unreportedRequests": 0 }), "one request for the actions, one for the summary" ); assert_rfc3339(&task["startedAt"]); @@ -508,6 +518,7 @@ async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { } } +// 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; @@ -520,8 +531,14 @@ async fn provider_refusing_the_key_fails_the_task() { let task = gateway.finished_analysis(session_id, task_id).await; assert_eq!(task["state"], "failed", "{task}"); - assert_eq!(task["payload"]["reason"], "permanent error"); - assert_eq!(task["payload"]["attempts"], 1); + 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)) From 9cbcde936589e2015fa74fb863df7086359e3839 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:04 -0400 Subject: [PATCH 38/53] fix(dgw): retry the end of a provisioner Task until it succeeds The end of a Task, such as listing its AI analysis log, ran from the Task's own job, whose runs are shared with the attempts. Once those runs were used up, a failed end was never done again, not even at the next start. A failed end now hands the Task to a new job that ends it again five minutes later. The job is stored, so a restart runs it too. Deadlines are also cut to whole seconds, as the store keeps them, so the runner and the store agree on when one passes. Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/job_queue.rs | 4 +- .../recording_ai_analysis/tests.rs | 4 +- .../src/provisioner_tasks/runner/mod.rs | 125 +++++++++++++----- .../src/provisioner_tasks/runner/tests.rs | 49 ++++++- 4 files changed, 143 insertions(+), 39 deletions(-) diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 9d5d38017..656e9887b 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -102,9 +102,9 @@ pub struct DurableJobQueue(DynJobQueue); #[async_trait] impl ProvisionerTaskJobQueue for DurableJobQueue { - async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { + async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()> { let job: DynJob = Box::new(job); - self.0.push_job(&job, None).await + self.0.push_job(&job, schedule_for).await } async fn stored_jobs(&self, name: &str) -> anyhow::Result> { diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index dae881fdb..337b74eac 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -814,8 +814,8 @@ struct WithRemuxJobs { #[async_trait] impl ProvisionerTaskJobQueue for WithRemuxJobs { - async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { - self.queue.push(job).await + 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> { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 51e287436..83a3b9997 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -10,6 +10,8 @@ //! 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; @@ -34,11 +36,15 @@ 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`, and returns once it is stored, so a record never waits for a job that was lost. - async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()>; + /// 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>; @@ -49,8 +55,11 @@ 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) -> anyhow::Result<()> { - self.enqueue(job).await + 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. @@ -103,8 +112,8 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// 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 the job call it - /// again later. + /// 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(()) } @@ -143,7 +152,7 @@ impl Attempt { /// 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 = now + self.progress_timeout; + 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 { @@ -293,7 +302,7 @@ impl ProvisionerTaskRunner { kind: T::KIND.to_owned(), target, params, - deadline_at: now + task.start_window(), + deadline_at: deadline(now, task.start_window()), job_token, }, now, @@ -305,13 +314,14 @@ impl ProvisionerTaskRunner { }; let queued = match on_created(record) { - Ok(()) => self.enqueue(id, job_token).await, + 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, &serde_json::Value::Null).await; + self.fail_not_queued(task.as_ref(), id, job_token, &serde_json::Value::Null) + .await; return Err(error); } @@ -361,7 +371,7 @@ impl ProvisionerTaskRunner { } } - match self.enqueue(record.id, job_token).await { + 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!( @@ -371,7 +381,8 @@ impl ProvisionerTaskRunner { "Failed to queue the Task again" ); if let Some(task) = self.task(&record.kind) { - self.fail_not_queued(task.as_ref(), record.id, &record.payload).await; + self.fail_not_queued(task.as_ref(), record.id, job_token, &record.payload) + .await; } } } @@ -390,14 +401,22 @@ impl ProvisionerTaskRunner { self.inner.tasks.get(kind).cloned() } - async fn enqueue(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { + 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(), - }) + .push( + ProvisionerTaskJob { + task_id, + job_token, + runner: self.clone(), + }, + schedule_for, + ) .await } @@ -425,7 +444,7 @@ impl ProvisionerTaskRunner { }; let starting_payload = task.starting_payload(&record.payload); - let deadline_at = now + task.progress_timeout(); + let deadline_at = deadline(now, task.progress_timeout()); let record = match store .start_attempt(task_id, job_token, starting_payload.clone(), deadline_at, now) @@ -434,7 +453,7 @@ impl ProvisionerTaskRunner { AttemptStart::Started(record) => record, AttemptStart::Finished(record) => { debug!(task.id = %task_id, "Task is over, so this attempt does not run"); - return task.on_end(&record).await.context("end the Task"); + 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"); @@ -517,10 +536,10 @@ impl ProvisionerTaskRunner { ); } - self.end(task.as_ref(), task_id).await?; + self.end(task.as_ref(), task_id, job_token).await?; } - // A retry would do the work again, such as paying for a second analysis, so the Task fails instead, - // which drops what the attempt made. + // 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. Err(error) => { error!( task.id = %task_id, @@ -528,8 +547,11 @@ impl ProvisionerTaskRunner { "Failed to record the Task success; the Task fails instead" ); - if let Err(error) = self.fail(task.as_ref(), task_id, &Permanent::Internal, &last).await { - error!(task.id = %task_id, error = format!("{error:#}"), "Failed to fail the Task; it fails at its deadline"); + if let Err(error) = self + .fail(task.as_ref(), task_id, job_token, &Permanent::Internal, &last) + .await + { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to fail or end the Task"); } } }, @@ -542,6 +564,7 @@ impl ProvisionerTaskRunner { self.fail( task.as_ref(), task_id, + job_token, &Permanent::AttemptsExhausted { last: error }, &last, ) @@ -549,7 +572,7 @@ impl ProvisionerTaskRunner { } Err(AttemptError::Permanent(error)) => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, "Task failed"); - self.fail(task.as_ref(), task_id, &error, &last).await?; + self.fail(task.as_ref(), task_id, job_token, &error, &last).await?; } } @@ -564,6 +587,7 @@ impl ProvisionerTaskRunner { &self, task: &dyn AnyProvisionerTask, task_id: Uuid, + job_token: Uuid, error: &Permanent, last: &serde_json::Value, ) -> anyhow::Result<()> { @@ -573,22 +597,57 @@ impl ProvisionerTaskRunner { .await .context("record the Task failure")?; - self.end(task, task_id).await + 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, last: &serde_json::Value) { - if let Err(error) = self.fail(task, task_id, &Permanent::NotQueued, last).await { + 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"); } } - async fn end(&self, task: &dyn AnyProvisionerTask, task_id: Uuid) -> anyhow::Result<()> { - match self.inner.store.get(task_id, OffsetDateTime::now_utc()).await? { - Some(record) => task.on_end(&record).await.context("end the Task"), - None => Ok(()), + /// 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)] diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 2db72f6b4..b3338e923 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -27,6 +27,8 @@ struct Echo { progress_timeout: time::Duration, runs: Arc, ended: Arc>>, + /// Ends that fail before one succeeds. + end_failures: Arc, } #[async_trait] @@ -77,6 +79,14 @@ impl ProvisionerTask for Echo { } 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(()) } @@ -87,6 +97,7 @@ struct Harness { job_queue_rx: JobQueueReceiver, runs: Arc, ended: Arc>>, + end_failures: Arc, } impl Harness { @@ -112,6 +123,7 @@ impl Harness { 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 { @@ -119,6 +131,7 @@ impl Harness { progress_timeout, runs: Arc::clone(&runs), ended: Arc::clone(&ended), + end_failures: Arc::clone(&end_failures), }) .build(); @@ -127,6 +140,7 @@ impl Harness { job_queue_rx, runs, ended, + end_failures, } } @@ -272,7 +286,7 @@ struct StoredJobs(Mutex>); #[async_trait] impl ProvisionerTaskJobQueue for StoredJobs { - async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { + async fn push(&self, job: ProvisionerTaskJob, _schedule_for: Option) -> anyhow::Result<()> { self.0.lock().push(job_queue::Job::write_json(&job)?); Ok(()) } @@ -304,6 +318,7 @@ async fn resume_queues_only_the_tasks_left_without_a_job() { 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(())); @@ -427,6 +442,36 @@ async fn invalid_parameters_and_unknown_kinds_fail_the_task() { 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() { @@ -444,7 +489,7 @@ async fn an_attempt_without_progress_is_cancelled_at_the_deadline() { // 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(2)).await; + 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"); From 14f325154b34cb166fc15a3216e6af12a84b4c84 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:05 -0400 Subject: [PATCH 39/53] fix(dgw): count an AI request as unreported until it is answered A request was counted only once its answer came. One cut off by the deadline or a restart was never counted, so the usage looked complete when it was not. Each request now counts as unreported from the moment it is sent, and its reported tokens replace that count on answer. Co-Authored-By: Claude Opus 5.5 --- .../src/recording_analysis/mod.rs | 48 +++++++++++-------- .../src/recording_analysis/tests.rs | 17 ++++++- devolutions-gateway/src/openapi.rs | 4 +- .../recording_ai_analysis/tests.rs | 6 ++- 4 files changed, 50 insertions(+), 25 deletions(-) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs index b033c8608..1a5c3b6cb 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/mod.rs @@ -107,8 +107,8 @@ pub enum Progress { Reading { read_bytes: u64, total_bytes: u64 }, /// Asking the AI: chunks already described, out of all chunks, and the tokens of the requests sent so far. /// - /// Sent before each chunk, after each request to the AI, and once all chunks are described. The summary request, - /// which follows when the chunks have actions, is counted by the event sent after it. + /// 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, @@ -194,8 +194,9 @@ pub struct UsageSoFar { 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 or a - /// request that got no answer; while it is not zero, the two counts are short of what the provider billed. + /// 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, } @@ -382,9 +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::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. + /// 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. @@ -531,8 +532,6 @@ impl Run { continue; } - self.report_describing(index, total).await; - let described = match chunk { Chunk::Transcript { file } => { let transcript = tokio::fs::read_to_string(workspace.join(file)) @@ -592,9 +591,8 @@ impl Run { let summary = tokio::select! { () = cancel.cancelled() => return Err(RecordingAnalysisError::Cancelled), - summary = request.send() => summary, + summary = self.ask(request.send(), total, total) => summary, }; - self.count(&summary, total, total).await; Some(summary?) }; @@ -683,8 +681,7 @@ impl Run { request = request.max_output_tokens(max_output_tokens); } - let answer = request.send().await; - self.count(&answer, index, total).await; + let answer = self.ask(request.send(), index, total).await; match answer { Ok(response) => described.add(response), @@ -739,8 +736,7 @@ impl Run { request = request.max_output_tokens(max_output_tokens); } - let answer = request.send().await; - self.count(&answer, index, total).await; + let answer = self.ask(request.send(), index, total).await; match answer { Ok(response) => described.add(response), @@ -764,22 +760,36 @@ impl Run { Ok(described) } - /// Counts the tokens of one request to the AI, then reports them with the chunks described so far. - async fn count(&self, answer: &Result, Error>, described_chunks: usize, total_chunks: usize) { + /// 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 { + 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 { .. }) => return, + 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 { diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs index 7baff5efd..b14e508f5 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs @@ -266,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() @@ -276,7 +277,10 @@ async fn the_log_is_left_next_to_the_manifest_and_the_working_files_removed() { Progress::Describing { described_chunks: 0, total_chunks: 1, - usage: UsageSoFar::default(), + usage: UsageSoFar { + unreported_requests: 1, + ..tokens(0, 0) + }, }, Progress::Describing { described_chunks: 0, @@ -288,12 +292,21 @@ async fn the_log_is_left_next_to_the_manifest_and_the_working_files_removed() { 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)); } diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index f00639d11..a1a62354b 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -210,8 +210,8 @@ struct AiAnalysisSucceededPayload { 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; while it - /// is not 0, `inputTokens` and `outputTokens` are short of what the provider billed. + /// 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, } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 337b74eac..4b391b2d6 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -718,12 +718,14 @@ async fn an_analysis_without_progress_fails_at_its_deadline() { 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": 0 }, - }) + "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)); From 2a79588b8c9068246283f784dddab6c6f887b502 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:06 -0400 Subject: [PATCH 40/53] fix(dgw): never list a late AI analysis log over a newer one A Task whose end failed could list its log later, over the analysis of a newer Task of the session. A new Task now removes the logs earlier Tasks left unlisted, so such a late end finds nothing to list. Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/mod.rs | 42 ++++++++++++++++++- .../recording_ai_analysis/tests.rs | 34 +++++++++++++++ devolutions-gateway/src/recording.rs | 13 ++++-- 3 files changed, 83 insertions(+), 6 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index d45cba0d0..6640b308e 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -292,11 +292,42 @@ impl RecordingAiAnalysis { } /// Where a Task leaves its log until its success is recorded: in the session folder, but not in the manifest, under a -/// name that only this Task uses, so that removing the working files of another analysis never removes it. +/// 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!(".ai-analysis-{task_id}.slog")) } +/// Removes the logs that earlier Tasks of the session left unlisted, so the late end of an earlier Task finds nothing to +/// list over a newer analysis. +async fn remove_unlisted_logs(manifest_path: &Utf8Path) -> std::io::Result<()> { + let Some(dir) = manifest_path.parent() else { + return Ok(()); + }; + let mut entries = match tokio::fs::read_dir(dir).await { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(error), + }; + + while let Some(entry) = entries.next_entry().await? { + let name = entry.file_name(); + let is_unlisted_log = name + .to_str() + .and_then(|name| name.strip_prefix(".ai-analysis-")) + .and_then(|name| name.strip_suffix(".slog")) + .is_some_and(|task_id| Uuid::parse_str(task_id).is_ok()); + + if is_unlisted_log { + match tokio::fs::remove_file(entry.path()).await { + Err(error) if error.kind() != std::io::ErrorKind::NotFound => return Err(error), + _ => {} + } + } + } + + Ok(()) +} + #[async_trait] impl ProvisionerTask for RecordingAiAnalysis { const KIND: &'static str = KIND; @@ -347,12 +378,19 @@ impl ProvisionerTask for RecordingAiAnalysis { // A new Task never resumes from what an earlier analysis left, while its retries resume from their own files. if attempt.attempts == 1 { - recording_analysis::discard(&self.manifest_path(session_id), None) + 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) })?; + + remove_unlisted_logs(&manifest_path).await.map_err(|error| { + error!(task.id = %attempt.task_id, %error, "Failed to remove the earlier unlisted AI analysis logs"); + AttemptError::Permanent(Permanent::Internal) + })?; } let secret = self.secrets.get(SecretId(attempt.task_id)).map_err(|error| { diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 4b391b2d6..22f5cf37c 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -1063,6 +1063,40 @@ async fn a_success_cut_short_before_the_log_is_listed_is_finished_by_the_job() { 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()); +} + /// The AI analyses the manifest in `dir` lists. fn ai_analyses(dir: &Utf8Path) -> serde_json::Value { let manifest: serde_json::Value = diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 8666ee9de..79219be9d 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -549,8 +549,8 @@ impl RecordingMessageSender { /// /// `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, and a session whose recording was deleted meanwhile - /// only loses `file`. + /// 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 @@ -978,8 +978,13 @@ impl RecordingManagerTask { // 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(()) => {} - // An earlier call moved it, then stopped before the manifest was written. - Err(error) if error.kind() == io::ErrorKind::NotFound && artifact_path.exists() => {} + // An earlier call moved it and stopped before the manifest was written, or it was removed for a newer + // artifact; either way, what is in place is what to list. + Err(error) if error.kind() == io::ErrorKind::NotFound => { + if !artifact_path.exists() { + return Ok(()); + } + } Err(error) => return Err(anyhow::Error::new(error).context(format!("move {file} to {artifact_path}"))), } From 573164982ccc8adc329509625d0814f2bf6105c0 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:07 -0400 Subject: [PATCH 41/53] fix(dgw): keep the manifest's temporary file out of jrec downloads The manifest is written to a temporary file first. Its name now starts with a dot, which jrec never serves, so a file a crash leaves behind is never downloaded. Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/recording.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 79219be9d..b870d147b 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -85,11 +85,14 @@ impl JrecManifest { use std::io::Write as _; let path = path.as_ref(); - let partial = path.with_extension("json.partial"); + // jrec never serves a name that starts with a dot, so a file a crash left behind is never downloaded. + let partial = path.with_file_name(".recording.json.partial"); let json = serde_json::to_string_pretty(&self)?; 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); From 850236cb4876c211a6d052d09176dca63621ff5d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:07 -0400 Subject: [PATCH 42/53] fix(dgw): keep the documented AI analysis error kinds in step The OpenAPI document writes the error kinds out by hand, since utoipa drops the descriptions of variants without fields. An exhaustive match now stops the build until a new kind is documented too. Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/openapi.rs | 33 +++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index a1a62354b..e8a0982f5 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -6,6 +6,7 @@ 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( @@ -264,7 +265,7 @@ enum AiAnalysisErrorKind { Internal, /// Every attempt failed with a failure worth another attempt. AttemptsExhausted, - /// The Task did not start within 2 hours, or went 30 minutes without progress. + /// The Task did not start in time, or went too long without progress. DeadlinePassed, /// The job of the Task could not be queued. NotQueued, @@ -305,6 +306,36 @@ enum AiAnalysisRecoverableErrorKind { 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 is written out too. +const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| match error { + Permanent::AiRejected { .. } => AiAnalysisErrorKind::AiRejected, + Permanent::AiBadAnswer { .. } => AiAnalysisErrorKind::AiBadAnswer, + Permanent::KeyLost => AiAnalysisErrorKind::KeyLost, + Permanent::NoRecording => AiAnalysisErrorKind::NoRecording, + Permanent::MixedRecording => AiAnalysisErrorKind::MixedRecording, + Permanent::UnreadableRecording { .. } => AiAnalysisErrorKind::UnreadableRecording, + Permanent::InvalidTask => AiAnalysisErrorKind::InvalidTask, + Permanent::Panicked => AiAnalysisErrorKind::Panicked, + Permanent::Internal => AiAnalysisErrorKind::Internal, + Permanent::AttemptsExhausted { .. } => AiAnalysisErrorKind::AttemptsExhausted, + Permanent::DeadlinePassed => AiAnalysisErrorKind::DeadlinePassed, + Permanent::NotQueued => AiAnalysisErrorKind::NotQueued, +}; + +const _: fn(&Recoverable) -> AiAnalysisRecoverableErrorKind = |error| match error { + Recoverable::AiUnreachable => AiAnalysisRecoverableErrorKind::AiUnreachable, + Recoverable::AiBusy { .. } => 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")] From 7fbfc20f42724f85ebbf41fd644b87c126a9a9dc Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:08 -0400 Subject: [PATCH 43/53] docs(dgw): mark the AI analysis key expiry as untested Co-Authored-By: Claude Opus 5.5 --- .../src/provisioner_tasks/recording_ai_analysis/mod.rs | 3 ++- devolutions-gateway/src/secret_store.rs | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 6640b308e..4eb29a8de 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -367,7 +367,8 @@ impl ProvisionerTask for RecordingAiAnalysis { AttemptError::Permanent(Permanent::InvalidTask) })?; - // The API key is kept until the deadline, which the start of this attempt moved. + // 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. diff --git a/devolutions-gateway/src/secret_store.rs b/devolutions-gateway/src/secret_store.rs index c8a75e022..a1af793c6 100644 --- a/devolutions-gateway/src/secret_store.rs +++ b/devolutions-gateway/src/secret_store.rs @@ -82,6 +82,8 @@ impl SecretStore { && 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. self.cleanup_notify.notify_one(); } } From cd9493b41a237c4fa7a74cb333c9c6b593b24335 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:08 -0400 Subject: [PATCH 44/53] refactor(dgw): name the progress values the AI analysis records Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/mod.rs | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 4eb29a8de..d07408fac 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -512,24 +512,20 @@ impl RecordingAiAnalysis { Some(RecordingAnalysisEvent::Progress(progress)) => { // Reading sends an event every half second, so while reading the Task is updated only when the percent // changes. - let shown = match progress { - Progress::Reading { .. } => (progress.stage(), u64::from(progress.percent()), None), - Progress::Describing { usage, .. } => ( - progress.stage(), - progress.done_and_total().map_or(0, |(done, _)| done), - Some(usage), - ), - _ => ( - progress.stage(), - progress.done_and_total().map_or(0, |(done, _)| done), - None, - ), + 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) = shown.2 { + if let Some(attempt_usage) = attempt_usage { usage = earlier.with(attempt_usage); } From aaee72bc7b85264afc2da910fcb95e710fbac910 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:09 -0400 Subject: [PATCH 45/53] test(dgw): name the AI error mapping test after what it checks Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/provisioner_tasks/ai.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index bf87e91a3..294cf18bb 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -234,7 +234,7 @@ mod tests { // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[test] - fn rate_limits_server_and_network_errors_are_recoverable() { + 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, From 481a10407f45e35ea6e3539c6089864638b9db34 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:49 -0400 Subject: [PATCH 46/53] fix(dgw): list the log of an earlier success instead of dropping it A new Task removed every log earlier Tasks left unlisted. When the end of a succeeded Task was still due, the new Task dropped an analysis whose success was recorded, and a new Task that then failed left the session without it. A new Task now does what the end of the earlier Task would: it lists the log of a succeeded Task, and removes the log of a failed or unknown one. The runner gives an attempt the state of another record for this. Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/mod.rs | 96 +++++++++++++------ .../recording_ai_analysis/tests.rs | 54 +++++++++++ .../src/provisioner_tasks/runner/mod.rs | 7 ++ devolutions-gateway/src/recording.rs | 9 +- 4 files changed, 133 insertions(+), 33 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index d07408fac..1d00fe562 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -291,41 +291,52 @@ impl RecordingAiAnalysis { } } +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!(".ai-analysis-{task_id}.slog")) + manifest_path.with_file_name(format!("{UNLISTED_LOG_PREFIX}{task_id}{UNLISTED_LOG_SUFFIX}")) } -/// Removes the logs that earlier Tasks of the session left unlisted, so the late end of an earlier Task finds nothing to -/// list over a newer analysis. -async fn remove_unlisted_logs(manifest_path: &Utf8Path) -> std::io::Result<()> { +/// 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(()); + 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(()), + 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 is_unlisted_log = name - .to_str() - .and_then(|name| name.strip_prefix(".ai-analysis-")) - .and_then(|name| name.strip_suffix(".slog")) - .is_some_and(|task_id| Uuid::parse_str(task_id).is_ok()); - - if is_unlisted_log { - match tokio::fs::remove_file(entry.path()).await { - Err(error) if error.kind() != std::io::ErrorKind::NotFound => return Err(error), - _ => {} - } + 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(()) + 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] @@ -388,10 +399,12 @@ impl ProvisionerTask for RecordingAiAnalysis { AttemptError::Permanent(Permanent::Internal) })?; - remove_unlisted_logs(&manifest_path).await.map_err(|error| { - error!(task.id = %attempt.task_id, %error, "Failed to remove the earlier unlisted AI analysis logs"); - 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| { @@ -438,12 +451,9 @@ impl ProvisionerTask for RecordingAiAnalysis { .await .context("list the AI analysis log in the session")?; } else { - match tokio::fs::remove_file(&log_path).await { - Err(error) if error.kind() != std::io::ErrorKind::NotFound => { - return Err(anyhow::Error::new(error).context(format!("remove {log_path}"))); - } - _ => {} - } + remove_if_exists(&log_path) + .await + .with_context(|| format!("remove {log_path}"))?; } recording_analysis::discard(&manifest_path, Some(&record.id.to_string())) @@ -568,6 +578,36 @@ impl RecordingAiAnalysis { } impl RecordingAiAnalysis { + /// 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::Succeeded) => self + .recordings + .replace_artifacts(session_id, ArtifactKind::AiAnalysis, &log_path) + .await + .with_context(|| format!("list {log_path}"))?, + Some(ProvisionerTaskState::Failed) | None => remove_if_exists(&log_path) + .await + .with_context(|| format!("remove {log_path}"))?, + Some(ProvisionerTaskState::Queued | ProvisionerTaskState::Running) => {} + } + } + + 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 { diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 22f5cf37c..200d1ddc7 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -1097,6 +1097,60 @@ async fn a_late_end_of_an_earlier_task_never_lists_its_log_over_a_newer_analysis 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_a_task_that_did_not_succeed() { + 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 = diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 83a3b9997..691c6393e 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -24,6 +24,7 @@ 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; @@ -146,6 +147,12 @@ impl Attempt { self.latest.lock().deadline_at } + /// State of another record, such as an earlier Task of the same target; `None` for an unknown one. + pub async fn state_of(&self, task_id: Uuid) -> anyhow::Result> { + let record = self.store.get(task_id, OffsetDateTime::now_utc()).await?; + Ok(record.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. /// diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index b870d147b..c28b61a20 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -981,12 +981,11 @@ impl RecordingManagerTask { // 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(()) => {} - // An earlier call moved it and stopped before the manifest was written, or it was removed for a newer - // artifact; either way, what is in place is what to list. + // An earlier call moved it, then stopped before the manifest was written. + Err(error) if error.kind() == io::ErrorKind::NotFound && artifact_path.exists() => {} Err(error) if error.kind() == io::ErrorKind::NotFound => { - if !artifact_path.exists() { - return Ok(()); - } + 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}"))), } From 9d38acaa411fa952d6dd92a5118a0a0223a7cd91 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:50 -0400 Subject: [PATCH 47/53] fix(dgw): retry the end when a failed success cannot queue it When the success of a Task could not be recorded, any error that came after the failure was recorded was only logged. An end that could not be queued again was then lost. That error now reaches the job queue, whose retry finds the Task finished and only ends it. Co-Authored-By: Claude Opus 5.5 --- .../src/provisioner_tasks/runner/mod.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 691c6393e..fab0b03b8 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -546,7 +546,8 @@ impl ProvisionerTaskRunner { 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. + // 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, @@ -554,11 +555,14 @@ impl ProvisionerTaskRunner { "Failed to record the Task success; the Task fails instead" ); - if let Err(error) = self - .fail(task.as_ref(), task_id, job_token, &Permanent::Internal, &last) - .await - { - error!(task.id = %task_id, error = format!("{error:#}"), "Failed to fail or end the Task"); + let payload = task.failed_payload(&Permanent::Internal, &last); + + match self.inner.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"); + } } } }, From f94532ddf442bd45ffd6d881552b6492a26f11d2 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:51 -0400 Subject: [PATCH 48/53] fix(dgw): stop the build on an undocumented AI analysis error field Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/openapi.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index e8a0982f5..e105dbd25 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -307,25 +307,25 @@ enum AiAnalysisRecoverableErrorKind { } // 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 is written out too. +// these matches stop the build until a new kind or field is written out too. const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| match error { - Permanent::AiRejected { .. } => AiAnalysisErrorKind::AiRejected, - Permanent::AiBadAnswer { .. } => AiAnalysisErrorKind::AiBadAnswer, + Permanent::AiRejected { status: _ } => AiAnalysisErrorKind::AiRejected, + Permanent::AiBadAnswer { why: _ } => AiAnalysisErrorKind::AiBadAnswer, Permanent::KeyLost => AiAnalysisErrorKind::KeyLost, Permanent::NoRecording => AiAnalysisErrorKind::NoRecording, Permanent::MixedRecording => AiAnalysisErrorKind::MixedRecording, - Permanent::UnreadableRecording { .. } => AiAnalysisErrorKind::UnreadableRecording, + Permanent::UnreadableRecording { file: _ } => AiAnalysisErrorKind::UnreadableRecording, Permanent::InvalidTask => AiAnalysisErrorKind::InvalidTask, Permanent::Panicked => AiAnalysisErrorKind::Panicked, Permanent::Internal => AiAnalysisErrorKind::Internal, - Permanent::AttemptsExhausted { .. } => AiAnalysisErrorKind::AttemptsExhausted, + 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 { .. } => AiAnalysisRecoverableErrorKind::AiBusy, + Recoverable::AiBusy { status: _ } => AiAnalysisRecoverableErrorKind::AiBusy, Recoverable::Storage => AiAnalysisRecoverableErrorKind::Storage, }; From 90b21f8bc7b347e9458695526995817903f16d40 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:53 -0400 Subject: [PATCH 49/53] refactor(dgw): name the manifest's temporary file after the manifest Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/recording.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index c28b61a20..55f3cc73b 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -86,7 +86,10 @@ impl JrecManifest { 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 partial = path.with_file_name(".recording.json.partial"); + 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)?; let mut file = std::fs::File::create(&partial)?; From ee826cccdcc9b12fe0c95f8f6773a722aa44886a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:53 -0400 Subject: [PATCH 50/53] docs(dgw): say what keeps a moved secret expiry on time Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/secret_store.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devolutions-gateway/src/secret_store.rs b/devolutions-gateway/src/secret_store.rs index a1af793c6..ea8d9179f 100644 --- a/devolutions-gateway/src/secret_store.rs +++ b/devolutions-gateway/src/secret_store.rs @@ -83,7 +83,7 @@ impl SecretStore { { 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. + // removed on time; any change of an expiry must notify it. self.cleanup_notify.notify_one(); } } From 4762f95a763590dafdd297a1bdc4e1ad980f76c5 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:44:23 -0400 Subject: [PATCH 51/53] fix(dgw): end a Task whose success and failure both went unrecorded When neither the success nor the failure of an attempt could be recorded, the job ended and nothing ended the Task, so its log stayed on disk. A job now runs just after the deadline, when the store has failed the Task, and only ends it. Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/provisioner_tasks/runner/mod.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index fab0b03b8..786ffe4c3 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -557,11 +557,18 @@ impl ProvisionerTaskRunner { let payload = task.failed_payload(&Permanent::Internal, &last); - match self.inner.store.fail(task_id, payload, OffsetDateTime::now_utc()).await { + 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"); + } } } } From a4caae62a1b837a78a2f39c33044564a430e920e Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:44:24 -0400 Subject: [PATCH 52/53] fix(dgw): read only Tasks of the same session for an unlisted log The ID of an unlisted log comes from a file name, so a Task of another session now reads as unknown, and its log is removed instead of listed. Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/provisioner_tasks/runner/mod.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 786ffe4c3..a8278960e 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -147,10 +147,13 @@ impl Attempt { self.latest.lock().deadline_at } - /// State of another record, such as an earlier Task of the same target; `None` for an unknown one. + /// 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.map(|record| record.state)) + 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 From a6e3fa864896989c9deaf7f29aae23075390c819 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:44:25 -0400 Subject: [PATCH 53/53] refactor(dgw): keep one rule for what the end does to an unlisted log Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/mod.rs | 40 ++++++++++--------- .../recording_ai_analysis/tests.rs | 2 +- devolutions-gateway/src/recording.rs | 3 +- 3 files changed, 24 insertions(+), 21 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 1d00fe562..33a5ca427 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -444,17 +444,8 @@ impl ProvisionerTask for RecordingAiAnalysis { let log_path = unlisted_log_path(&manifest_path, record.id); let _session = lock_session(session_id).await; - if record.state == ProvisionerTaskState::Succeeded { - // The success is recorded, so the log takes the place of the earlier analysis. - self.recordings - .replace_artifacts(session_id, ArtifactKind::AiAnalysis, &log_path) - .await - .context("list the AI analysis log in the session")?; - } else { - remove_if_exists(&log_path) - .await - .with_context(|| format!("remove {log_path}"))?; - } + self.end_log(session_id, &log_path, record.state == ProvisionerTaskState::Succeeded) + .await?; recording_analysis::discard(&manifest_path, Some(&record.id.to_string())) .await @@ -578,6 +569,21 @@ impl RecordingAiAnalysis { } 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. @@ -593,15 +599,11 @@ impl RecordingAiAnalysis { .await .context("read the Task of an unlisted log")? { - Some(ProvisionerTaskState::Succeeded) => self - .recordings - .replace_artifacts(session_id, ArtifactKind::AiAnalysis, &log_path) - .await - .with_context(|| format!("list {log_path}"))?, - Some(ProvisionerTaskState::Failed) | None => remove_if_exists(&log_path) - .await - .with_context(|| format!("remove {log_path}"))?, Some(ProvisionerTaskState::Queued | ProvisionerTaskState::Running) => {} + state => { + self.end_log(session_id, &log_path, state == Some(ProvisionerTaskState::Succeeded)) + .await? + } } } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 200d1ddc7..42e7990ce 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -1136,7 +1136,7 @@ async fn a_new_task_that_fails_keeps_an_earlier_success_whose_log_was_not_listed // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[tokio::test] -async fn a_new_task_removes_a_log_left_unlisted_by_a_task_that_did_not_succeed() { +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()); diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 55f3cc73b..ef0fd0f67 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -984,7 +984,8 @@ impl RecordingManagerTask { // 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(()) => {} - // An earlier call moved it, then stopped before the manifest was written. + // 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");