From 5f2bdf8f16a3c1679291da04aef5973443da19af Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 6 Oct 2026 15:10:51 -0400 Subject: [PATCH 01/80] 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 688d542242..ee53d09ec8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2083,6 +2083,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 d42ea03f45..95ced16348 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 8985f3184f..615bbfb631 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 a163bfa403..f6e20c20f9 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 258e15d230..32abd87a72 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 1cbcf84808..71a312c624 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 7175d3bb59..3c9e4a63cb 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 8e090e9d10..25755d7e0f 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 67f134f6b0755b131cae9020a2e3a0b5db9fae26 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 7 Oct 2026 17:03:00 -0400 Subject: [PATCH 02/80] 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 ee53d09ec8..c9f61194c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2045,6 +2045,7 @@ dependencies = [ "cfg-if", "chacha20poly1305", "devolutions-agent-shared", + "devolutions-gateway-ai", "devolutions-gateway-generators", "devolutions-gateway-task", "devolutions-log", @@ -2113,6 +2114,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 95ced16348..a2c18215b2 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 615bbfb631..3e4b611fc6 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 f6e20c20f9..6a6ca3783c 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 0000000000..8b17d3746b --- /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 0000000000..b8bc845ebc --- /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 0000000000..25cf3a8f06 --- /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 0000000000..3b07f8bbd4 --- /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 0000000000..308ada5e68 --- /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 0000000000..a40619393e --- /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 0ebc4c730d..1efc1df059 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -31,15 +31,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, @@ -49,6 +60,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)?; @@ -335,7 +360,7 @@ impl ActiveRecordings { self.0.lock().clone() } - fn insert(&self, id: Uuid) -> usize { + pub(crate) fn insert(&self, id: Uuid) -> usize { let mut guard = self.0.lock(); guard.insert(id); guard.len() @@ -376,6 +401,10 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, + GetFinished { + id: Uuid, + channel: oneshot::Sender>, + }, Disconnect { id: Uuid, end: PushEnd, @@ -397,6 +426,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 { @@ -416,6 +468,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, end } => { f.debug_struct("Disconnect").field("id", id).field("end", end).finish() } @@ -481,6 +536,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, end: PushEnd) -> anyhow::Result<()> { self.channel .send(RecordingManagerMessage::Disconnect { id, end }) @@ -846,6 +931,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 { @@ -964,6 +1058,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, end } => { if let Err(e) = manager.handle_disconnect(id, end).await { error!(error = format!("{e:#}"), "handle_disconnect"); @@ -1487,6 +1584,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 32abd87a72..eea186dca9 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 c379b23b58f546880dd51fca74eeaca365d2591c Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 6 Oct 2026 16:21:34 -0400 Subject: [PATCH 03/80] 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 4dfa03784c..337c09d9c3 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::{PushEnd, 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) } @@ -444,6 +453,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, @@ -1027,6 +1222,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 c7b222ebc1..f31a32d158 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 0000000000..c82e990761 --- /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 a000e46cf1..0212a79d1e 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 e023be1436..6ef5e0f669 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 8b17d3746b..6ae4bf015d 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 25cf3a8f06..48b07b0b41 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 308ada5e68..4e43724f2c 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 a40619393e..f063a771f8 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 c922869a58..d645c2e095 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 0000000000..d69889bf2c --- /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 e727c8e015..fa0d080ccc 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 cba330ca3842549cd7e4e8aabde75a8e6f6029eb Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 16:41:35 -0400 Subject: [PATCH 04/80] 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 3e4b611fc6..60db0a367f 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 6a6ca3783c..c39576552d 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 b8bc845ebc..ccb1d4938c 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 48b07b0b41..dea3252499 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 3b07f8bbd4..fdcee3db4f 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 4e43724f2c..c3b848d397 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 f063a771f8..6cbd047ae6 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 eea186dca9..efc0ea7070 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 d69889bf2c..79e28edb6c 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 0d940009d3cffc04064147d0b1cae0bb052d37c8 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:03:51 -0400 Subject: [PATCH 05/80] 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 c9f61194c7..97351453b6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4321,6 +4321,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 efdf7da010..a4a4934a3f 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 0efbfd23f0..04f76b63d3 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 93232bd868..5521da5281 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 60db0a367f..42e4f3c652 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 c3b848d397..9e7b089689 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 6cbd047ae6..67e80b2d11 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 eb4547ca23b132d1df53e3acbd2ac67b5637c361 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:31:15 -0400 Subject: [PATCH 06/80] 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 0000000000..e96647acdd --- /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 b6ef90ebf4b68c1ba097a5701001904d91449662 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:40:42 -0400 Subject: [PATCH 07/80] 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 e96647acdd..b7e8f12459 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 54b61cbdb9c1d85984d3837f04509d7e2c09d1d6 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:46:09 -0400 Subject: [PATCH 08/80] 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 0000000000..e65cc49550 --- /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 f78dc41ea18eb7c47f1f66a0dc7777e96ae74ce3 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:57:41 -0400 Subject: [PATCH 09/80] 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 b7e8f12459..346de3c4bd 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 2fa28a332c5d4ad4baf88268d78ecca8737d9b11 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:56 -0400 Subject: [PATCH 10/80] 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 fa0d080ccc..e727c8e015 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 a556dd679712dc2ada317c73fcbf8bf74e92730d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:56 -0400 Subject: [PATCH 11/80] 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 6ae4bf015d..089805ee6c 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 ae9b2522e56ca656e4ca57b94e2bcdb3f10542c1 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 12/80] 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 6ef5e0f669..7163baad4d 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 d77eaa045586a0010ddeec912eb4e9dd34bb2d98 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 13/80] 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 fdcee3db4f..b4157ac7e1 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 9e7b089689..44476192c8 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 67e80b2d11..4c4fc8d331 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 efc0ea7070..6a33137f7d 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 0055e44b801a60b59e004b273556e77467e0641d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 14/80] 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 337c09d9c3..3030975aec 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -639,6 +639,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 0212a79d1e..224c9382a6 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 ba8341d899d47fe4dcf840ff4352f149f97cf6b1 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 09:48:45 -0400 Subject: [PATCH 15/80] 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 3030975aec..218e66c84f 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -456,7 +456,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 79e28edb6c..d471a42fb5 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 2ced8c7924e2a2376b5c0166115c3274ff8c28f4 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:02:58 -0400 Subject: [PATCH 16/80] 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 dea3252499..f29828ca6d 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 b4157ac7e1..f0916e583f 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 44476192c8..a5121cc59f 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 4c4fc8d331..92f4666ac0 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 e946651683a14e8add7786afb6288e7f807e4fea Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:23:39 -0400 Subject: [PATCH 17/80] 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 7163baad4d..65128cca56 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 f29828ca6d..39ca2daf3a 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 f0916e583f..d1845a5ab5 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 567229c122a7da361a747105fcd8fc4932df880a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:40:54 -0400 Subject: [PATCH 18/80] 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 42e4f3c652..9d5d38017e 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 ccb1d4938c..0278157fd3 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 39ca2daf3a..c3b65141db 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 d1845a5ab5..7029897415 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 a5121cc59f..f1135f90f5 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 92f4666ac0..e962b99a6d 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 1efc1df059..e7eedafc8e 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; @@ -1154,6 +1154,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 839e3b355d1a5b4e5aa94461089ff84cef891dcf Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:13:21 -0400 Subject: [PATCH 19/80] 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 c82e990761..d0f2b95040 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 65128cca56..c98d21e961 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 c3b65141db..f536510ff0 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 7029897415..85fa416490 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 f1135f90f5..2b7e454292 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 e962b99a6d..a69558afda 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 cc765f4169d297bafa96f152c60963fabf7ee082 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:56:31 -0400 Subject: [PATCH 20/80] 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 218e66c84f..b22dd986b9 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -510,7 +510,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 a2ecd5fa89..d62f5d8037 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 f536510ff0..acc4211953 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 85fa416490..c098fb7126 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 e7eedafc8e..f38ac6d6ef 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -401,6 +401,11 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, + ClearArtifacts { + id: Uuid, + kind: ArtifactKind, + channel: oneshot::Sender>, + }, GetFinished { id: Uuid, channel: oneshot::Sender>, @@ -468,6 +473,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() } @@ -536,6 +546,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 @@ -931,6 +954,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) { @@ -1058,6 +1114,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 be1cef5faf7ec78c3c34b072ce678f1b8798c4b8 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 14:19:27 -0400 Subject: [PATCH 21/80] 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 c098fb7126..3c91ad0d54 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 626ebcb5e74976155f218de44bedf68942d60e8a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:19:39 -0400 Subject: [PATCH 22/80] 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 e65cc49550..8bcee2600b 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 b22dd986b9..e60d1f8c8d 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; @@ -470,6 +471,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, @@ -507,11 +512,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( @@ -528,7 +535,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"), @@ -560,9 +567,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 { @@ -574,6 +587,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 c98d21e961..dcd42fe21f 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 346de3c4bd..b8746192f5 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 acc4211953..a0aa6becac 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 3c91ad0d54..1fc8ebf249 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 d471a42fb5..e1cc0d944f 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 50e1b1dc649397f663f4dff54c2b349af7a3559b Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:41:47 -0400 Subject: [PATCH 23/80] 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 1fc8ebf249..145480cfbc 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 68d7619f008339eed76cbe8b6928cbe1f22cc57d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:54:55 -0400 Subject: [PATCH 24/80] 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 a0aa6becac..a7a456744c 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 145480cfbc..de482567d9 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 72c1b284b3c1d8150a9f6d4fee60a7adbf054a8a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:54:56 -0400 Subject: [PATCH 25/80] 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 d0f2b95040..2431febb62 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 dcd42fe21f..2354f3a6e5 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 c4150e0ae6bdbe0b3f820d5eb51312228df56973 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:54:57 -0400 Subject: [PATCH 26/80] 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 97351453b6..2a2d350493 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8201,6 +8201,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 e1cc0d944f..0000000000 --- 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 cbd210e7be..1f47e506db 100644 --- a/testsuite/Cargo.toml +++ b/testsuite/Cargo.toml @@ -64,6 +64,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 0000000000..df30628d73 --- /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 23da8ba6c8..85e2b02bcd 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 d7e6ff2625cdeaafb17a887916d1368b8ed00fe3 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 14:57:53 -0400 Subject: [PATCH 27/80] 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 2354f3a6e5..4978ef675a 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 2d99d33df721e4e845c6e9a171c6404a6715b897 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:16:24 -0400 Subject: [PATCH 28/80] 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 de482567d9..9919f41cf7 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 f38ac6d6ef..21bbf78e02 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -548,7 +548,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 @@ -959,7 +959,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 72d155971395cda667873adf3bba17bbe66ee707 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:16:24 -0400 Subject: [PATCH 29/80] 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 df30628d73..6786cf58e8 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 5e23d9201d6665012b8eae71a2e32555bf58e11a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:16:25 -0400 Subject: [PATCH 30/80] 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 4978ef675a..6a414bd7fb 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 8c72c4383d43bbeb841e6c1f729b4e9921ac6ee4 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:19:48 -0400 Subject: [PATCH 31/80] 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 9919f41cf7..25176806b5 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 333c87a1375cf685c4aac766e3b13888cbd4ba92 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 20 Oct 2026 15:37:44 -0400 Subject: [PATCH 32/80] 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 e60d1f8c8d..f3e8b55550 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -469,6 +469,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 6a414bd7fb..cbb9d5036c 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 938b89a7de2c0eb0590cb4b31c50900a14085961 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 21 Oct 2026 15:48:02 -0400 Subject: [PATCH 33/80] 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 c39576552d..c20eabc64d 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 089805ee6c..c923e5877a 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 0278157fd3..7e900f13c8 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 a7a456744c..4b192d1e98 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 25176806b5..538d50c64f 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 6a33137f7d..0c66667332 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 a285daafa74c6f1671e3fc9dc1d04898dfae9b1e Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 21 Oct 2026 17:18:57 -0400 Subject: [PATCH 34/80] 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 b8746192f5..742db342d9 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 a338b621f008dbd55eb9ce9ac8c2c2f092d881a8 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 15:26:34 -0400 Subject: [PATCH 35/80] 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 | 1 + 4 files changed, 18 insertions(+), 5 deletions(-) diff --git a/crates/provisioner-task/INTENT.md b/crates/provisioner-task/INTENT.md index 40d2a64d25..642b6271b4 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 8bcee2600b..ccd9fd423f 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 742db342d9..ca2b0010b6 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 0a6b6ca5e5..6c56dd1b2f 100644 --- a/devolutions-gateway/src/recording.intent.md +++ b/devolutions-gateway/src/recording.intent.md @@ -9,6 +9,7 @@ When a session whose token requires recording starts and no recording starts for # Artifacts Artifacts are files that are not recordings; the only kind today is `ai-analysis`, stored with the `slog` file type. An endpoint to push artifacts is wanted later. +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. # Ending a recording From 69991abaf36af8bd3eb349e9922d0a528d905a4f Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 16:53:29 -0400 Subject: [PATCH 36/80] 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 6503160464..c8a75e0220 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 330f505ff92903a86e1ba42d1554dd1d1e0c7c38 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 16:53:30 -0400 Subject: [PATCH 37/80] 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 94bf8b1d13..d390b7c5d4 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 e5a3a4c4e5..a0b81ece2d 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 e598d98339..b033c8608c 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 82b18b7712..7baff5efd0 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 04f76b63d3..f8224ba385 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 3a0f03972d..1b1ff2f756 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 be3d8431dc..587a4b93c5 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 f3e8b55550..a643d40c00 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -511,15 +511,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( @@ -619,7 +621,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, @@ -1239,6 +1242,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 2431febb62..dfba09fb7d 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 cbb9d5036c..f00639d11a 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 c923e5877a..bf87e91a3f 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 7e900f13c8..2c49d0330e 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 4b192d1e98..d45cba0d0f 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 538d50c64f..dae881fdb7 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 2b7e454292..51e2874362 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 a69558afda..2db72f6b49 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 21bbf78e02..b48df2616c 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -80,9 +80,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(()) } } @@ -401,9 +414,10 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, - ClearArtifacts { + ReplaceArtifacts { id: Uuid, kind: ArtifactKind, + file: Utf8PathBuf, channel: oneshot::Sender>, }, GetFinished { @@ -473,10 +487,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() @@ -546,17 +566,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 { @@ -954,31 +983,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"); } _ => {} } @@ -1114,8 +1161,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)); @@ -1651,6 +1698,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 6786cf58e8..592efd58e6 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 6b9c9c89154e9df3992fd989d0b690cef98c58ba Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:04 -0400 Subject: [PATCH 38/80] 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 9d5d38017e..656e9887bc 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 dae881fdb7..337b74eac5 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 51e2874362..83a3b99973 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 2db72f6b49..b3338e9232 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 08beff0230e44b5ff242de2105e31db05d965fd8 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:05 -0400 Subject: [PATCH 39/80] 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 b033c8608c..1a5c3b6cb4 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 7baff5efd0..b14e508f55 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 f00639d11a..a1a62354be 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 337b74eac5..4b391b2d68 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 3386fab1e9df3f91c31adfdf8812fdad1e66f7c7 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:06 -0400 Subject: [PATCH 40/80] 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 d45cba0d0f..6640b308e8 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 4b391b2d68..22f5cf37c4 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 b48df2616c..0df6be2dce 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -571,8 +571,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 @@ -1005,8 +1005,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 b11c0ff0a411a8d8152a4e173f579ab58f24e662 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:07 -0400 Subject: [PATCH 41/80] 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 0df6be2dce..f07c0769e8 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -86,11 +86,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 57e05dba58fa5fe7666ef19135485499b89bb78d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:07 -0400 Subject: [PATCH 42/80] 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 a1a62354be..e8a0982f53 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 5f31e70c76503ef527934a960ba884ee3b1414ec Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:08 -0400 Subject: [PATCH 43/80] 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 6640b308e8..4eb29a8dec 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 c8a75e0220..a1af793c64 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 64fe1bd10d86cab216966101e76f181c0ee8d04b Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:08 -0400 Subject: [PATCH 44/80] 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 4eb29a8dec..d07408fac6 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 70a50981a10628e340a9b17518b726f0602726f5 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 17:36:09 -0400 Subject: [PATCH 45/80] 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 bf87e91a3f..294cf18bb7 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 3a951452d59b500181883d8d85b0cb1ad9419053 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:49 -0400 Subject: [PATCH 46/80] 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 d07408fac6..1d00fe5620 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 22f5cf37c4..200d1ddc75 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 83a3b99973..691c6393ef 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 f07c0769e8..38f4153b08 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -1008,12 +1008,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 ea03a4cc2feddcd102d4852849891d2d1ccee4c0 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:50 -0400 Subject: [PATCH 47/80] 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 691c6393ef..fab0b03b86 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 a2b3751e084d71068bbe07cf2d0ada940922fb18 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:51 -0400 Subject: [PATCH 48/80] 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 e8a0982f53..e105dbd25e 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 4caaadcb1c21872c85428eda8193a665a0ce80eb Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:53 -0400 Subject: [PATCH 49/80] 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 38f4153b08..1cfee6486d 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -87,7 +87,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 b932a976bc40c57c2290a09d60dcc33ab7929f5b Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:06:53 -0400 Subject: [PATCH 50/80] 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 a1af793c64..ea8d9179f7 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 7100e2f77b6fa712d6f260b37c01929242cec69a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:44:23 -0400 Subject: [PATCH 51/80] 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 fab0b03b86..786ffe4c3c 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 61f5b57d3ea7111bf159dbb577db397c61b0e6b0 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:44:24 -0400 Subject: [PATCH 52/80] 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 786ffe4c3c..a8278960e6 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 840f4727e90f710e05fb265a7b65dfe78e51c71f Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 22 Oct 2026 18:44:25 -0400 Subject: [PATCH 53/80] 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 1d00fe5620..33a5ca4270 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 200d1ddc75..42e7990ce7 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 1cfee6486d..afdb165f74 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -1011,7 +1011,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"); From c4247507e18c78de4950c901ec4330557301931f Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 27 Oct 2026 10:08:49 -0400 Subject: [PATCH 54/80] docs(dgw): say why a running Task's deadline moves Co-Authored-By: Claude Opus 5.5 --- crates/provisioner-task/INTENT.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/provisioner-task/INTENT.md b/crates/provisioner-task/INTENT.md index 642b6271b4..6fefec608c 100644 --- a/crates/provisioner-task/INTENT.md +++ b/crates/provisioner-task/INTENT.md @@ -13,7 +13,10 @@ Only these transitions exist: - `Running` -> `Succeeded` - `Running` -> `Failed` -Each Task has a deadline. While the Task runs, its deadline may move forward each time it makes progress. +Each Task has a deadline. +While the Task runs, its deadline may move forward each time it makes progress. +Some work has no known length: a recording can be arbitrarily long, so any fixed deadline would leave some recordings impossible to analyze. +With a moving deadline, a Task fails only once it stops making progress, however long it runs. An unfinished Task past its deadline is failed, so a lost job can't block its target forever. # Tasks and jobs From de4010d5a6fa3c741c52ccbef25c5a86a41ea09a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 27 Oct 2026 12:10:01 -0400 Subject: [PATCH 55/80] test(dgw): mark tests kept for the faults they catch No intent sentence decides what these seven tests expect, but each one alone catches a fault, so their markers now name that fault instead of an intent section that does not decide the value. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- crates/job-queue-libsql/src/lib.rs | 2 +- .../src/provisioner_tasks/recording_ai_analysis/tests.rs | 4 ++-- devolutions-gateway/src/provisioner_tasks/runner/tests.rs | 8 ++++---- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/job-queue-libsql/src/lib.rs b/crates/job-queue-libsql/src/lib.rs index f8224ba385..10737ce59d 100644 --- a/crates/job-queue-libsql/src/lib.rs +++ b/crates/job-queue-libsql/src/lib.rs @@ -473,7 +473,7 @@ mod tests { } } - // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + // oracle: mutant the name filter is joined with OR, so the jobs of every name are listed #[tokio::test] async fn job_definitions_lists_the_stored_jobs_of_one_name() { let conn = libsql::Builder::new_local(":memory:") 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 42e7990ce7..cb21a5a23f 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -646,7 +646,7 @@ async fn a_session_mixing_terminal_and_video_fails() { assert_eq!(task.payload["error"], json!({ "kind": "mixedRecording" })); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +// oracle: mutant a recording deleted after the Task starts fails it as internal instead of noRecording #[tokio::test] async fn recording_deleted_after_start_fails() { let harness = Harness::new().await; @@ -1134,7 +1134,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 +// oracle: mutant the log left unlisted by a Task that has no record is kept #[tokio::test] async fn a_new_task_removes_a_log_left_unlisted_by_an_unknown_task() { let harness = Harness::new().await; diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index b3338e9232..d37333b6be 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -192,7 +192,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 +// oracle: mutant a job that cannot be queued is ignored, so the Task stays queued and its creation succeeds #[tokio::test] async fn a_task_whose_job_cannot_be_queued_fails() { let mut harness = Harness::new(|value| Ok(json!(value))).await; @@ -211,7 +211,7 @@ async fn a_task_whose_job_cannot_be_queued_fails() { 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 +// oracle: mutant the job is queued before the Task is prepared, so a Task that cannot be prepared still gets a job #[tokio::test] async fn a_task_that_cannot_be_prepared_fails_without_a_job() { let mut harness = Harness::new(|value| Ok(json!(value))).await; @@ -353,7 +353,7 @@ async fn resume_queues_only_the_tasks_left_without_a_job() { ); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +// oracle: mutant a panicking attempt fails the Task as internal instead of panicked #[tokio::test] async fn a_panicking_attempt_fails_the_task() { let mut harness = Harness::new(|_| panic!("the attempt panics")).await; @@ -407,7 +407,7 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { ); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +// oracle: mutant a queued Task of an unknown kind is never failed #[tokio::test] async fn invalid_parameters_and_unknown_kinds_fail_the_task() { let harness = Harness::new(|value| Ok(json!(value))).await; From 473658033e2e2038ab531783b2973ab1ec942e96 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 27 Oct 2026 12:13:19 -0400 Subject: [PATCH 56/80] test(dgw): run the progress test past its first deadline The work of the progress test took 2.5 s against a first deadline 2 to 3 s away, and it ignored the cancellation, so the test passed even when progress never moved the deadline. It now works for 4 s and stops once cancelled, so only a moved deadline lets it succeed. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/provisioner_tasks/runner/tests.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index d37333b6be..9f0cdfe76b 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -68,8 +68,12 @@ impl ProvisionerTask for Echo { Err(AttemptError::Recoverable(Recoverable::Storage)) } PROGRESS => { - for step in 0..5 { - tokio::time::sleep(std::time::Duration::from_millis(500)).await; + // 4 s of work outlasts the first deadline, 2 to 3 s away, so only the progress keeps the attempt running. + for step in 0..8 { + tokio::select! { + () = tokio::time::sleep(std::time::Duration::from_millis(500)) => {} + () = attempt.cancel.cancelled() => return Err(AttemptError::Recoverable(Recoverable::Storage)), + } attempt.report_progress(json!({ "step": step })).await; } Ok(json!("done")) From fbb05df3f4241537a30bbe09e7bac57b96cb8199 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 27 Oct 2026 12:17:08 -0400 Subject: [PATCH 57/80] test(dgw): cite the intent that decides two tests The tokens of a failed Task are decided by the jrec intent, and the deadline an attempt start sets by the 30 minutes without progress of the provisioner tasks intent, not by the sections these tests cited. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- crates/provisioner-task-libsql/src/lib.rs | 2 +- .../src/provisioner_tasks/recording_ai_analysis/tests.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 1b1ff2f756..deee0b5546 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -879,7 +879,7 @@ mod tests { #[tokio::test] async fn an_attempt_start_sets_the_deadline() { - // oracle: intent crates/provisioner-task/INTENT.md §invariants + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis let store = store().await; let task = new_task("session-1"); create(&store, task.clone()).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 cb21a5a23f..2fd9835ec4 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -994,7 +994,7 @@ async fn a_new_task_that_fails_keeps_the_earlier_analysis() { assert!(dir.join("ai-analysis-0.slog").exists()); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis +// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis #[tokio::test] async fn a_failed_task_reports_the_tokens_it_spent() { let harness = Harness::new().await; From b5fb2e2a36ba8a2c0c009f7fce801ca3438368b6 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 27 Oct 2026 12:17:38 -0400 Subject: [PATCH 58/80] test(dgw): refuse the recordings read scope on a Task read The metadata of a Task belongs to the task reading scope, so reading it with only `gateway.recordings.read` must be refused; no test sent that scope. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- testsuite/tests/cli/dgw/ai_analysis.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testsuite/tests/cli/dgw/ai_analysis.rs b/testsuite/tests/cli/dgw/ai_analysis.rs index 592efd58e6..72689d832a 100644 --- a/testsuite/tests/cli/dgw/ai_analysis.rs +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -422,7 +422,7 @@ async fn endpoints_require_their_scope() { assert_eq!(response.status(), StatusCode::FORBIDDEN, "{scope}"); } - for scope in [START_SCOPE, "gateway.sessions.read"] { + for scope in [START_SCOPE, "gateway.sessions.read", "gateway.recordings.read"] { let (status, _) = gateway.get_analysis(session_id, Uuid::new_v4(), scope).await; assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); } From 1060659a5ba60c7aac167e07db86cf6183fbffe6 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 13:23:36 -0400 Subject: [PATCH 59/80] refactor(dgw): serve the AI analysis under /jet/tasks/recording-ai-analysis Each kind of task gets its own typed route under `/jet/tasks//{id}`, a public convention the next kinds can follow. `PUT /jet/tasks/recording-ai-analysis/{taskId}` starts an analysis, with the session ID in the body, and a GET on the same path reads it. The answer carries `sessionId` and leaves out `kind`, `target` and `attempts`, which only matter inside Gateway. Its params and payload are documented as the AI analysis schemas. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/api/jrec/mod.rs | 266 +------------- devolutions-gateway/src/api/mod.rs | 1 + devolutions-gateway/src/api/tasks.rs | 75 +--- .../src/api/tasks/recording_ai_analysis.rs | 341 ++++++++++++++++++ devolutions-gateway/src/openapi.rs | 24 +- testsuite/tests/cli/dgw/ai_analysis.rs | 109 +++--- 6 files changed, 433 insertions(+), 383 deletions(-) create mode 100644 devolutions-gateway/src/api/tasks/recording_ai_analysis.rs diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index a643d40c00..0af099cedb 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -12,35 +12,27 @@ 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, LOCATION}; +use axum::http::header::{CONTENT_DISPOSITION, CONTENT_TYPE, HeaderValue}; use axum::response::{IntoResponse as _, Response}; use axum::routing::{delete, get, post}; use axum::{Json, Router}; use bytes::Bytes; use cadeau::xmf; use camino::{Utf8Path, Utf8PathBuf}; -use devolutions_gateway_ai::Language; use devolutions_gateway_task::ShutdownSignal; use futures::stream; use hyper::StatusCode; -use secrecy::SecretString; use tokio::io::AsyncReadExt as _; use tracing::Instrument as _; -use url::Url; use uuid::Uuid; use zip::CompressionMethod; use zip::write::SimpleFileOptions; use crate::DgwState; use crate::api::heartbeat::recording_storage_health; -use crate::api::tasks::TaskInfo; use crate::artifacts::JrecArtifacts; -use crate::extract::{ - JrecToken, ProvisionerTaskReadScope, ProvisionerTaskStartScope, RecordingDeleteScope, RecordingsReadScope, -}; +use crate::extract::{JrecToken, 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::{PushEnd, PushOutcome, RecordingMessageSender}; use crate::token::{JrecTokenClaims, RecordingFileType, RecordingOperation}; @@ -70,8 +62,6 @@ 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) } @@ -454,211 +444,6 @@ async fn delete_recording(recording_path: &Utf8Path) -> anyhow::Result<()> { Ok(()) } -/// Settings of an AI analysis -#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] -#[derive(Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AiAnalysisRequest { - /// Task ID, chosen by the caller so it can send the same request again safely. - task_id: Uuid, - provider: AiProvider, - /// Model identifier, passed to the provider as is. - model: String, - /// Endpoint of the provider API, kept in memory only; required for `openai-compatible`. - #[cfg_attr(feature = "openapi", schema(value_type = Option))] - #[serde(default)] - base_url: Option, - /// Upper bound of tokens in each AI answer. - #[cfg_attr(feature = "openapi", schema(value_type = Option, minimum = 0, maximum = 4_294_967_295_u64))] - #[serde(default)] - max_output_tokens: Option, - /// Language the AI writes in, any text of 1 to 10 characters such as `fr-CA` or `français`; English when absent. - #[cfg_attr(feature = "openapi", schema(min_length = 1, max_length = 10))] - #[serde(default)] - language: Option, - /// API key of the provider, kept in memory only until the Task ends. - #[cfg_attr(feature = "openapi", schema(value_type = String, format = Password))] - api_key: SecretString, -} - -/// Why an AI analysis cannot start now -#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AiAnalysisConflict { - code: AiAnalysisConflictCode, - /// The AI analysis of this session that is not finished yet; set for `analysis_in_progress`. - #[serde(skip_serializing_if = "Option::is_none")] - task_id: Option, -} - -#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] -#[serde(rename_all = "snake_case")] -pub(crate) enum AiAnalysisConflictCode { - /// The Task ID is already used by a Task with other parameters. - TaskIdConflict, - /// The session is still recording. - RecordingActive, - /// Another AI analysis of this session is not finished yet. - AnalysisInProgress, -} - -fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> Response { - debug!(?code, task.id = ?task_id, "AI analysis refused"); - (StatusCode::CONFLICT, Json(AiAnalysisConflict { code, task_id })).into_response() -} - -/// Starts an AI analysis of a recorded session -/// -/// Gateway asks the AI provider what the user did in the session recording, and puts the answer in the session as its -/// log, listed under `artifacts.ai-analysis` in its manifest. -/// The answer is in the `language` of the request, English without one. -/// The work runs in the background as a Task; follow it with `GET /jet/jrec/{session_id}/ai-analysis/{task_id}`, -/// which the `Location` header of a 202 answer names. -/// -/// Sending the same request again returns the same Task and starts nothing new. -/// A new Task for the session starts from scratch, and replaces the earlier AI analysis of the session once its success -/// is recorded, so a session holds at most one; a Task that fails leaves the earlier analysis in place and removes -/// everything it made. -/// A Task fails if it does not start within 2 hours, or, once it runs, after 30 minutes without progress. -/// The API key and the base URL are kept in memory only, until the Task ends; a Gateway restart fails the unfinished -/// Tasks. -#[cfg_attr(feature = "openapi", utoipa::path( - post, - operation_id = "StartAiAnalysis", - tag = "Jrec", - path = "/jet/jrec/{session_id}/ai-analysis", - params( - ("session_id" = Uuid, Path, description = "Recorded session ID"), - ), - request_body(content = AiAnalysisRequest, description = "AI analysis settings", content_type = "application/json"), - responses( - (status = 200, description = "The same request came in before; nothing new starts", body = TaskInfo), - (status = 202, description = "The AI analysis Task is created and runs in the background", body = TaskInfo, headers( - ("Location" = String, description = "Path of the new Task, `/jet/jrec/{session_id}/ai-analysis/{task_id}`"), - )), - (status = 400, description = "Invalid request body, AI settings or language"), - (status = 401, description = "Invalid or missing authorization token"), - (status = 403, description = "Insufficient permissions"), - (status = 404, description = "The session has no recording"), - (status = 409, description = "The AI analysis cannot start now", body = AiAnalysisConflict), - (status = 500, description = "Unexpected server error"), - ), - security(("scope_token" = ["gateway.tasks.recording-ai-analysis.start"])), -))] -pub(crate) async fn start_ai_analysis( - State(state): State, - _scope: ProvisionerTaskStartScope, - extract::Path(session_id): extract::Path, - body: Bytes, -) -> Result { - let request: AiAnalysisRequest = serde_json::from_slice(&body).map_err(|error| { - // The serde message may quote a value of the body, such as a misplaced API key, so only its place is logged. - debug!( - category = ?error.classify(), - line = error.line(), - column = error.column(), - "Invalid AI analysis request body" - ); - HttpError::bad_request().msg("invalid AI analysis request body") - })?; - - let AiAnalysisRequest { - task_id, - provider, - model, - base_url, - max_output_tokens, - language, - api_key, - } = request; - - let language = language - .map(Language::new) - .transpose() - .map_err(HttpError::bad_request().with_msg("invalid language").err())?; - - let outcome = recording_ai_analysis::start( - &state, - StartRequest { - task_id, - session_id, - settings: AiSettings { - provider, - model, - base_url, - max_output_tokens, - }, - language, - api_key, - }, - ) - .await; - - match outcome { - Ok(StartOutcome::Created(task)) => { - let location = HeaderValue::try_from(format!("/jet/jrec/{session_id}/ai-analysis/{}", task.id)) - .map_err(HttpError::internal().with_msg("invalid Location header").err())?; - Ok((StatusCode::ACCEPTED, [(LOCATION, location)], Json(TaskInfo::from(task))).into_response()) - } - Ok(StartOutcome::Existing(task)) => Ok((StatusCode::OK, Json(TaskInfo::from(task))).into_response()), - Ok(StartOutcome::IdConflict(_)) => Ok(ai_analysis_conflict(AiAnalysisConflictCode::TaskIdConflict, None)), - Ok(StartOutcome::TargetBusy { active_task_id }) => Ok(ai_analysis_conflict( - AiAnalysisConflictCode::AnalysisInProgress, - Some(active_task_id), - )), - Err(StartError::RecordingActive) => Ok(ai_analysis_conflict(AiAnalysisConflictCode::RecordingActive, None)), - Err(StartError::InvalidSettings(error)) => { - Err(HttpError::bad_request().with_msg("invalid AI settings").build(error)) - } - Err(StartError::RecordingNotFound) => Err(HttpError::not_found().msg("session has no recording")), - Err(StartError::Internal(error)) => Err(HttpError::internal() - .with_msg("failed to start the AI analysis") - .build(error)), - } -} - -/// Gets an AI analysis of a recorded session -/// -/// Returns the Task of the analysis: its state, its progress while it runs, its result once it is over, and the tokens -/// it used in every state but `queued`. -/// Tasks are kept forever, across Gateway restarts; an unfinished Task past its deadline is reported as failed. -#[cfg_attr(feature = "openapi", utoipa::path( - get, - operation_id = "GetAiAnalysis", - tag = "Jrec", - path = "/jet/jrec/{session_id}/ai-analysis/{task_id}", - params( - ("session_id" = Uuid, Path, description = "Recorded session ID"), - ("task_id" = Uuid, Path, description = "Task ID of the AI analysis"), - ), - responses( - (status = 200, description = "The AI analysis Task", body = TaskInfo), - (status = 400, description = "Bad request"), - (status = 401, description = "Invalid or missing authorization token"), - (status = 403, description = "Insufficient permissions"), - (status = 404, description = "The session has no AI analysis with this Task ID"), - (status = 500, description = "Unexpected server error"), - ), - security(("scope_token" = ["gateway.tasks.recording-ai-analysis.read"])), -))] -pub(crate) async fn get_ai_analysis( - State(DgwState { provisioner_tasks, .. }): State, - _scope: ProvisionerTaskReadScope, - extract::Path((session_id, task_id)): extract::Path<(Uuid, Uuid)>, -) -> Result, HttpError> { - let task = provisioner_tasks - .get(task_id, time::OffsetDateTime::now_utc()) - .await - .map_err(HttpError::internal().with_msg("failed to read the AI analysis").err())? - // Another kind of Task, or the analysis of another session, is not found under this path. - .filter(|task| task.kind == recording_ai_analysis::KIND && task.target == session_id.to_string()) - .ok_or_else(|| HttpError::not_found().msg("AI analysis not found"))?; - - Ok(Json(TaskInfo::from(task))) -} - /// Lists all recordings stored on this instance #[cfg_attr(feature = "openapi", utoipa::path( get, @@ -1242,53 +1027,6 @@ mod tests { use super::*; - // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis - #[tokio::test] - async fn ai_analysis_of_a_session_still_recording_is_a_conflict() { - use axum::extract::connect_info::MockConnectInfo; - use base64::Engine as _; - use tower::ServiceExt as _; - - let dir = tempfile::tempdir().expect("temp dir"); - let config = serde_json::json!({ - "ProvisionerPublicKeyData": { - "Value": "mMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4vuqLOkl1pWobt6su1XO9VskgCAwevEGs6kkNjJQBwkGnPKYLmNF1E/af1yCocfVn/OnPf9e4x+lXVyZ6LMDJxFxu+axdgOq3Ld392J1iAEbfvwlyRFnEXFOJNyylqg3bY6LvnWHL/XZczVdMD9xYfq2sO9bg3xjRW4s7r9EEYOFjqVT3VFznH9iWJVtcSEKukmS/3uKoO6lGhacvu0HhjXXdgq0R8zvR4XRJ9Fcnf0f9Ypoc+i6L80NVjrRCeVOH+Ld/2fA9bocpfLarcVqG3RjS+qgOtpyCc0jWVFF4zaGQ7LUDFkEIYILkICeMMn2ll29hmZNzsJzZJ9s6NocgQIDAQAB" - }, - "Listeners": [{ "InternalUrl": "http://*:7171", "ExternalUrl": "https://*:7171" }], - "RecordingPath": dir.path().to_str().expect("UTF-8"), - "__debug__": { "disable_token_validation": true }, - }); - let (state, _handles) = DgwState::mock(&config.to_string()).await.expect("mock state"); - let session_id = Uuid::new_v4(); - state.recordings.active_recordings.insert(session_id); - let app = crate::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 3000)))); - - let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD; - let claims = serde_json::json!({ "jti": Uuid::new_v4(), "exp": i64::from(u32::MAX), "scope": "gateway.tasks.recording-ai-analysis.start" }); - let token = format!( - "{}.{}.{}", - engine.encode(r#"{"alg":"RS256","cty":"SCOPE"}"#), - engine.encode(claims.to_string()), - engine.encode(b"signature"), - ); - let body = serde_json::json!({ "taskId": Uuid::new_v4(), "provider": "openai", "model": "gpt-test", "apiKey": "sk-test" }); - let request = axum::http::Request::builder() - .method("POST") - .uri(format!("/jet/jrec/{session_id}/ai-analysis")) - .header(axum::http::header::AUTHORIZATION, format!("Bearer {token}")) - .body(Body::from(body.to_string())) - .expect("request"); - - let response = app.oneshot(request).await.expect("response"); - - assert_eq!(response.status(), StatusCode::CONFLICT); - let body = response.into_body().collect().await.expect("body").to_bytes(); - assert_eq!( - serde_json::from_slice::(&body).expect("JSON"), - serde_json::json!({ "code": "recording_active" }) - ); - } - #[test] fn rejects_unsafe_recording_file_names() { assert!(is_safe_recording_file_name("recording-0.webm")); diff --git a/devolutions-gateway/src/api/mod.rs b/devolutions-gateway/src/api/mod.rs index f31a32d158..241b446016 100644 --- a/devolutions-gateway/src/api/mod.rs +++ b/devolutions-gateway/src/api/mod.rs @@ -26,6 +26,7 @@ pub fn make_router(state: crate::DgwState) -> axum::Router { .route("/jet/preflight", axum::routing::post(preflight::post_preflight)) .nest("/jet/jrl", jrl::make_router(state.clone())) .nest("/jet/jrec", jrec::make_router(state.clone())) + .nest("/jet/tasks", tasks::make_router(state.clone())) .nest("/jet/config", config::make_router(state.clone())) .nest("/jet/session", session::make_router(state.clone())) .nest("/jet/sessions", sessions::make_router(state.clone())) diff --git a/devolutions-gateway/src/api/tasks.rs b/devolutions-gateway/src/api/tasks.rs index dfba09fb7d..fe05b0580b 100644 --- a/devolutions-gateway/src/api/tasks.rs +++ b/devolutions-gateway/src/api/tasks.rs @@ -1,47 +1,20 @@ -//! The Task JSON returned by the endpoints of provisioner tasks, such as the AI analysis of a recording. +//! Endpoints of provisioner tasks: each kind of Task is started and read under `/jet/tasks/{kind}/{task_id}`. -use provisioner_task::{ProvisionerTaskRecord, ProvisionerTaskState}; -use time::OffsetDateTime; -use uuid::Uuid; +pub(crate) mod recording_ai_analysis; -/// A Task: the permanent record of one piece of work the provisioner asked Gateway to do -/// -/// The work runs in the background; `state` and `payload` tell how it goes. -/// The shape of `params` and `payload` belongs to the kind. -/// For `recording-ai-analysis`, `params` is an `AiAnalysisParams`, and `payload` is an `AiAnalysisRunningPayload` while -/// `running`, an `AiAnalysisSucceededPayload` once `succeeded`, and an `AiAnalysisFailedPayload` once `failed`. -#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct TaskInfo { - /// Task ID, chosen by the provisioner. - id: Uuid, - /// Kind of work, such as `recording-ai-analysis`. - kind: String, - /// What the work is about, such as a session ID. - target: String, - /// What was asked, without secrets. - #[cfg_attr(feature = "openapi", schema(value_type = Object))] - params: serde_json::Value, - state: TaskState, - /// Details of the current state; null while `queued`. - #[cfg_attr(feature = "openapi", schema(value_type = Option))] - payload: serde_json::Value, - /// Number of attempts started so far. - attempts: u32, - /// When the Task was created. - #[serde(with = "time::serde::rfc3339")] - created_at: OffsetDateTime, - /// When its first attempt started. - #[serde(with = "time::serde::rfc3339::option")] - started_at: Option, - /// When it succeeded or failed. - #[serde(with = "time::serde::rfc3339::option")] - finished_at: Option, - /// Time at which a Task still unfinished fails; while the Task runs, it moves ahead each time an attempt starts or - /// the Task makes progress. - #[serde(with = "time::serde::rfc3339")] - deadline_at: OffsetDateTime, +use axum::Router; +use axum::routing::put; +use provisioner_task::ProvisionerTaskState; + +use crate::DgwState; + +pub fn make_router(state: DgwState) -> Router { + Router::new() + .route( + "/recording-ai-analysis/{task_id}", + put(recording_ai_analysis::start_ai_analysis).get(recording_ai_analysis::get_ai_analysis), + ) + .with_state(state) } #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] @@ -66,21 +39,3 @@ impl From for TaskState { } } } - -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/api/tasks/recording_ai_analysis.rs b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs new file mode 100644 index 0000000000..5cd16bf0ea --- /dev/null +++ b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs @@ -0,0 +1,341 @@ +//! `/jet/tasks/recording-ai-analysis/{task_id}`: start and read the AI analysis of a recorded session. + +use axum::Json; +use axum::body::Bytes; +use axum::extract::{self, State}; +use axum::http::header::{HeaderValue, LOCATION}; +use axum::response::{IntoResponse as _, Response}; +use devolutions_gateway_ai::Language; +use hyper::StatusCode; +use provisioner_task::ProvisionerTaskRecord; +use secrecy::SecretString; +use time::OffsetDateTime; +use url::Url; +use uuid::Uuid; + +use super::TaskState; +use crate::DgwState; +use crate::extract::{ProvisionerTaskReadScope, ProvisionerTaskStartScope}; +use crate::http::HttpError; +use crate::provisioner_tasks::ai::{AiProvider, AiSettings}; +use crate::provisioner_tasks::recording_ai_analysis::{ + self as analysis, RecordingAiAnalysis, StartError, StartOutcome, StartRequest, +}; + +/// Settings of an AI analysis +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisRequest { + /// Recorded session to analyze. + session_id: Uuid, + provider: AiProvider, + /// Model identifier, passed to the provider as is. + model: String, + /// Endpoint of the provider API, kept in memory only; required for `openai-compatible`. + #[cfg_attr(feature = "openapi", schema(value_type = Option))] + #[serde(default)] + base_url: Option, + /// Upper bound of tokens in each AI answer. + #[cfg_attr(feature = "openapi", schema(value_type = Option, minimum = 0, maximum = 4_294_967_295_u64))] + #[serde(default)] + max_output_tokens: Option, + /// Language the AI writes in, any text of 1 to 10 characters such as `fr-CA` or `français`; English when absent. + #[cfg_attr(feature = "openapi", schema(min_length = 1, max_length = 10))] + #[serde(default)] + language: Option, + /// API key of the provider, kept in memory only until the Task ends. + #[cfg_attr(feature = "openapi", schema(value_type = String, format = Password))] + api_key: SecretString, +} + +/// Why an AI analysis cannot start now +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisConflict { + code: AiAnalysisConflictCode, + /// The AI analysis of this session that is not finished yet; set for `analysis_in_progress`. + #[serde(skip_serializing_if = "Option::is_none")] + task_id: Option, +} + +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum AiAnalysisConflictCode { + /// The Task ID is already used by a Task with other parameters. + TaskIdConflict, + /// The session is still recording. + RecordingActive, + /// Another AI analysis of this session is not finished yet. + AnalysisInProgress, +} + +/// An AI analysis Task: the permanent record of one analysis the provisioner asked for +/// +/// The work runs in the background; `state` and `payload` tell how it goes. +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisTask { + /// Task ID, chosen by the provisioner. + id: Uuid, + /// Recorded session the analysis is about. + session_id: Uuid, + /// What was asked, without the API key and the base URL. + #[cfg_attr(feature = "openapi", schema(value_type = crate::openapi::AiAnalysisParams))] + params: serde_json::Value, + state: TaskState, + /// Details of the state: progress while `running`, the log once `succeeded`, the reason once `failed`; null while + /// `queued`. + #[cfg_attr(feature = "openapi", schema(value_type = Option))] + payload: serde_json::Value, + /// When the Task was created. + #[serde(with = "time::serde::rfc3339")] + created_at: OffsetDateTime, + /// When its first attempt started. + #[serde(with = "time::serde::rfc3339::option")] + started_at: Option, + /// When it succeeded or failed. + #[serde(with = "time::serde::rfc3339::option")] + finished_at: Option, + /// Time at which a Task still unfinished fails; while the Task runs, it moves ahead each time an attempt starts or + /// the Task makes progress. + #[serde(with = "time::serde::rfc3339")] + deadline_at: OffsetDateTime, +} + +impl TryFrom for AiAnalysisTask { + type Error = uuid::Error; + + fn try_from(task: ProvisionerTaskRecord) -> Result { + Ok(Self { + id: task.id, + session_id: Uuid::parse_str(&task.target)?, + params: task.params, + state: TaskState::from(task.state), + payload: task.payload, + created_at: task.created_at, + started_at: task.started_at, + finished_at: task.finished_at, + deadline_at: task.deadline_at, + }) + } +} + +fn task_json(task: ProvisionerTaskRecord) -> Result, HttpError> { + AiAnalysisTask::try_from(task) + .map(Json) + .map_err(HttpError::internal().with_msg("invalid session ID in the AI analysis").err()) +} + +fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> Response { + debug!(?code, task.id = ?task_id, "AI analysis refused"); + (StatusCode::CONFLICT, Json(AiAnalysisConflict { code, task_id })).into_response() +} + +/// Starts an AI analysis of a recorded session +/// +/// Gateway asks the AI provider what the user did in the session recording, and puts the answer in the session as its +/// log, listed under `artifacts.ai-analysis` in its manifest. +/// The answer is in the `language` of the request, English without one. +/// The work runs in the background as a Task; follow it with `GET /jet/tasks/recording-ai-analysis/{task_id}`. +/// +/// Sending the same request again returns the same Task and starts nothing new. +/// A new Task for the session starts from scratch, and replaces the earlier AI analysis of the session once its success +/// is recorded, so a session holds at most one; a Task that fails leaves the earlier analysis in place and removes +/// everything it made. +/// A Task fails if it does not start within 2 hours, or, once it runs, after 30 minutes without progress. +/// The API key and the base URL are kept in memory only, until the Task ends; a Gateway restart fails the unfinished +/// Tasks. +#[cfg_attr(feature = "openapi", utoipa::path( + put, + operation_id = "StartAiAnalysis", + tag = "Tasks", + path = "/jet/tasks/recording-ai-analysis/{task_id}", + params( + ("task_id" = Uuid, Path, description = "Task ID, chosen by the caller so it can send the same request again safely"), + ), + 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 = AiAnalysisTask), + (status = 202, description = "The AI analysis Task is created and runs in the background", body = AiAnalysisTask, headers( + ("Location" = String, description = "Path of the new Task, `/jet/tasks/recording-ai-analysis/{task_id}`"), + )), + (status = 400, description = "Invalid request body, AI settings or language"), + (status = 401, description = "Invalid or missing authorization token"), + (status = 403, description = "Insufficient permissions"), + (status = 404, description = "The session has no recording"), + (status = 409, description = "The AI analysis cannot start now", body = AiAnalysisConflict), + (status = 500, description = "Unexpected server error"), + ), + security(("scope_token" = ["gateway.tasks.recording-ai-analysis.start"])), +))] +pub(crate) async fn start_ai_analysis( + State(state): State, + _scope: ProvisionerTaskStartScope, + extract::Path(task_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 { + session_id, + provider, + model, + base_url, + max_output_tokens, + language, + api_key, + } = request; + + let language = language + .map(Language::new) + .transpose() + .map_err(HttpError::bad_request().with_msg("invalid language").err())?; + + let outcome = analysis::start( + &state, + StartRequest { + task_id, + session_id, + settings: AiSettings { + provider, + model, + base_url, + max_output_tokens, + }, + language, + api_key, + }, + ) + .await; + + match outcome { + Ok(StartOutcome::Created(task)) => { + let location = HeaderValue::try_from(format!("/jet/tasks/recording-ai-analysis/{}", task.id)) + .map_err(HttpError::internal().with_msg("invalid Location header").err())?; + Ok((StatusCode::ACCEPTED, [(LOCATION, location)], task_json(task)?).into_response()) + } + Ok(StartOutcome::Existing(task)) => Ok((StatusCode::OK, task_json(task)?).into_response()), + Ok(StartOutcome::IdConflict(_)) => Ok(ai_analysis_conflict(AiAnalysisConflictCode::TaskIdConflict, None)), + Ok(StartOutcome::TargetBusy { active_task_id }) => Ok(ai_analysis_conflict( + AiAnalysisConflictCode::AnalysisInProgress, + Some(active_task_id), + )), + Err(StartError::RecordingActive) => Ok(ai_analysis_conflict(AiAnalysisConflictCode::RecordingActive, None)), + Err(StartError::InvalidSettings(error)) => { + Err(HttpError::bad_request().with_msg("invalid AI settings").build(error)) + } + Err(StartError::RecordingNotFound) => Err(HttpError::not_found().msg("session has no recording")), + Err(StartError::Internal(error)) => Err(HttpError::internal() + .with_msg("failed to start the AI analysis") + .build(error)), + } +} + +/// Gets an AI analysis of a recorded session +/// +/// Returns the Task of the analysis: its state, its progress while it runs, its result once it is over, and the tokens +/// it used in every state but `queued`. +/// Tasks are kept forever, across Gateway restarts; an unfinished Task past its deadline is reported as failed. +#[cfg_attr(feature = "openapi", utoipa::path( + get, + operation_id = "GetAiAnalysis", + tag = "Tasks", + path = "/jet/tasks/recording-ai-analysis/{task_id}", + params( + ("task_id" = Uuid, Path, description = "Task ID of the AI analysis"), + ), + responses( + (status = 200, description = "The AI analysis Task", body = AiAnalysisTask), + (status = 400, description = "Bad request"), + (status = 401, description = "Invalid or missing authorization token"), + (status = 403, description = "Insufficient permissions"), + (status = 404, description = "No AI analysis has 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(task_id): extract::Path, +) -> Result, HttpError> { + let task = provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .map_err(HttpError::internal().with_msg("failed to read the AI analysis").err())? + // A Task of another kind is not found under this path. + .filter(|task| task.kind == analysis::KIND) + .ok_or_else(|| HttpError::not_found().msg("AI analysis not found"))?; + + task_json(task) +} + +#[cfg(test)] +mod tests { + use std::net::SocketAddr; + + use axum::body::Body; + use http_body_util::BodyExt as _; + + use super::*; + + // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis + #[tokio::test] + async fn ai_analysis_of_a_session_still_recording_is_a_conflict() { + use axum::extract::connect_info::MockConnectInfo; + use base64::Engine as _; + use tower::ServiceExt as _; + + let dir = tempfile::tempdir().expect("temp dir"); + let config = serde_json::json!({ + "ProvisionerPublicKeyData": { + "Value": "mMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4vuqLOkl1pWobt6su1XO9VskgCAwevEGs6kkNjJQBwkGnPKYLmNF1E/af1yCocfVn/OnPf9e4x+lXVyZ6LMDJxFxu+axdgOq3Ld392J1iAEbfvwlyRFnEXFOJNyylqg3bY6LvnWHL/XZczVdMD9xYfq2sO9bg3xjRW4s7r9EEYOFjqVT3VFznH9iWJVtcSEKukmS/3uKoO6lGhacvu0HhjXXdgq0R8zvR4XRJ9Fcnf0f9Ypoc+i6L80NVjrRCeVOH+Ld/2fA9bocpfLarcVqG3RjS+qgOtpyCc0jWVFF4zaGQ7LUDFkEIYILkICeMMn2ll29hmZNzsJzZJ9s6NocgQIDAQAB" + }, + "Listeners": [{ "InternalUrl": "http://*:7171", "ExternalUrl": "https://*:7171" }], + "RecordingPath": dir.path().to_str().expect("UTF-8"), + "__debug__": { "disable_token_validation": true }, + }); + let (state, _handles) = DgwState::mock(&config.to_string()).await.expect("mock state"); + let session_id = Uuid::new_v4(); + state.recordings.active_recordings.insert(session_id); + let app = crate::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 3000)))); + + let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD; + let claims = serde_json::json!({ "jti": Uuid::new_v4(), "exp": i64::from(u32::MAX), "scope": "gateway.tasks.recording-ai-analysis.start" }); + let token = format!( + "{}.{}.{}", + engine.encode(r#"{"alg":"RS256","cty":"SCOPE"}"#), + engine.encode(claims.to_string()), + engine.encode(b"signature"), + ); + let body = serde_json::json!({ "sessionId": session_id, "provider": "openai", "model": "gpt-test", "apiKey": "sk-test" }); + let request = axum::http::Request::builder() + .method("PUT") + .uri(format!("/jet/tasks/recording-ai-analysis/{}", Uuid::new_v4())) + .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" }) + ); + } +} diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index e105dbd25e..d2162d8fab 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -27,8 +27,8 @@ use crate::provisioner_tasks::{BadAnswer, Permanent, Recoverable}; 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::tasks::recording_ai_analysis::start_ai_analysis, + crate::api::tasks::recording_ai_analysis::get_ai_analysis, crate::api::webapp::sign_app_token, crate::api::webapp::sign_session_token, crate::api::update::get_update_products, @@ -68,13 +68,14 @@ use crate::provisioner_tasks::{BadAnswer, Permanent, Recoverable}; 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::api::tasks::recording_ai_analysis::AiAnalysisRequest, + crate::api::tasks::recording_ai_analysis::AiAnalysisConflict, + crate::api::tasks::recording_ai_analysis::AiAnalysisConflictCode, AiProvider, - crate::api::tasks::TaskInfo, + crate::api::tasks::recording_ai_analysis::AiAnalysisTask, crate::api::tasks::TaskState, AiAnalysisParams, + AiAnalysisPayload, AiAnalysisRunningPayload, AiAnalysisSucceededPayload, AiAnalysisUsage, @@ -165,7 +166,7 @@ struct SessionInfo { /// 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 { +pub(crate) struct AiAnalysisParams { provider: AiProvider, /// Model identifier, passed to the provider as is. model: String, @@ -177,6 +178,15 @@ struct AiAnalysisParams { language: Option, } +/// Details of the state of an AI analysis Task: `running`, `succeeded` and `failed` each have their own +#[derive(Serialize, utoipa::ToSchema)] +#[serde(untagged)] +pub(crate) enum AiAnalysisPayload { + Running(AiAnalysisRunningPayload), + Succeeded(AiAnalysisSucceededPayload), + Failed(AiAnalysisFailedPayload), +} + /// Payload of a `running` AI analysis Task: its progress, and the tokens it used so far #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "camelCase")] diff --git a/testsuite/tests/cli/dgw/ai_analysis.rs b/testsuite/tests/cli/dgw/ai_analysis.rs index 72689d832a..bfcdd37331 100644 --- a/testsuite/tests/cli/dgw/ai_analysis.rs +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -1,4 +1,4 @@ -//! `/jet/jrec/{sessionId}/ai-analysis` tests, against a running Gateway and a mock AI provider. +//! `/jet/tasks/recording-ai-analysis/{taskId}` tests, against a running Gateway and a mock AI provider. use std::path::PathBuf; use std::time::{Duration, Instant}; @@ -95,9 +95,9 @@ impl Gateway { session_id } - async fn post_raw(&self, session_id: Uuid, scope: &str, body: String) -> reqwest::Response { + async fn put_raw(&self, task_id: Uuid, scope: &str, body: String) -> reqwest::Response { self.client - .post(self.url(&format!("/jet/jrec/{session_id}/ai-analysis"))) + .put(self.url(&format!("/jet/tasks/recording-ai-analysis/{task_id}"))) .header("content-type", "application/json") .bearer_auth(scope_token(scope)) .body(body) @@ -106,14 +106,14 @@ impl Gateway { .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 put(&self, task_id: Uuid, body: &Value) -> (StatusCode, Value) { + read(self.put_raw(task_id, START_SCOPE, body.to_string()).await).await } - async fn get_analysis(&self, session_id: Uuid, task_id: Uuid, scope: &str) -> (StatusCode, Value) { + async fn get_analysis(&self, task_id: Uuid, scope: &str) -> (StatusCode, Value) { let response = self .client - .get(self.url(&format!("/jet/jrec/{session_id}/ai-analysis/{task_id}"))) + .get(self.url(&format!("/jet/tasks/recording-ai-analysis/{task_id}"))) .bearer_auth(scope_token(scope)) .send() .await @@ -122,11 +122,11 @@ impl Gateway { } /// Asks for the analysis until it has succeeded or failed. - async fn finished_analysis(&self, session_id: Uuid, task_id: Uuid) -> Value { + async fn finished_analysis(&self, 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; + let (status, task) = self.get_analysis(task_id, READ_SCOPE).await; assert_eq!(status, StatusCode::OK, "{task}"); if task["state"] == "succeeded" || task["state"] == "failed" { @@ -183,9 +183,9 @@ async fn read(response: reqwest::Response) -> (StatusCode, Value) { (status, body) } -fn request(task_id: Uuid, model: &str, base_url: &str) -> Value { +fn request(session_id: Uuid, model: &str, base_url: &str) -> Value { json!({ - "taskId": task_id, + "sessionId": session_id, "provider": "openai-compatible", "model": model, "baseUrl": base_url, @@ -286,27 +286,28 @@ async fn start_returns_202_then_200_for_the_same_request() { 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; + let (status, created) = gateway.put(task_id, &request(session_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["sessionId"], 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); + for internal in ["kind", "target", "attempts"] { + assert!(created.get(internal).is_none(), "{internal} is not returned"); + } - let (status, again) = gateway.post(session_id, &request(task_id, "gpt-test", &base_url)).await; + let (status, again) = gateway.put(task_id, &request(session_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"] { + for field in ["id", "sessionId", "params", "createdAt"] { assert_eq!(again[field], created[field], "{field}"); } } @@ -318,16 +319,20 @@ async fn start_conflicts_return_409_with_a_code() { 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; + gateway.put(task_id, &request(session_id, "gpt-test", &base_url)).await; + + let (status, body) = gateway.put(task_id, &request(session_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(task_id, "gpt-other", &base_url)) + .put(task_id, &request(gateway.write_session(), "gpt-test", &base_url)) .await; - assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(status, StatusCode::CONFLICT, "the same Task ID for another session"); assert_eq!(body, json!({ "code": "task_id_conflict" })); let (status, body) = gateway - .post(session_id, &request(Uuid::new_v4(), "gpt-test", &base_url)) + .put(Uuid::new_v4(), &request(session_id, "gpt-test", &base_url)) .await; assert_eq!(status, StatusCode::CONFLICT); assert_eq!(body, json!({ "code": "analysis_in_progress", "taskId": task_id })); @@ -340,7 +345,7 @@ async fn start_without_recording_returns_404() { let base_url = spawn_silent_ai_provider().await; let (status, _) = gateway - .post(Uuid::new_v4(), &request(Uuid::new_v4(), "gpt-test", &base_url)) + .put(Uuid::new_v4(), &request(Uuid::new_v4(), "gpt-test", &base_url)) .await; assert_eq!(status, StatusCode::NOT_FOUND); @@ -352,10 +357,10 @@ 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); + let mut body = request(session_id, "gpt-test", &base_url); body["temperature"] = json!(0.5); - let (status, created) = gateway.post(session_id, &body).await; + let (status, created) = gateway.put(Uuid::new_v4(), &body).await; assert_eq!(status, StatusCode::ACCEPTED, "{created}"); } @@ -367,7 +372,10 @@ async fn invalid_bodies_and_settings_return_400() { 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 valid = request(session_id, "gpt-test", &base_url); + + let mut no_session = valid.clone(); + no_session.as_object_mut().unwrap().remove("sessionId"); let mut no_key = valid.clone(); no_key.as_object_mut().unwrap().remove("apiKey"); @@ -391,9 +399,10 @@ async fn invalid_bodies_and_settings_return_400() { empty_language["language"] = json!(""); for (case, body) in [ + ("missing session ID", no_session.to_string()), ("missing API key", no_key.to_string()), ("empty API key", empty_key.to_string()), - ("empty model", request(task_id, " ", &base_url).to_string()), + ("empty model", request(session_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()), @@ -401,11 +410,11 @@ async fn invalid_bodies_and_settings_return_400() { ("empty language", empty_language.to_string()), ("not JSON", "{".to_owned()), ] { - let response = gateway.post_raw(session_id, START_SCOPE, body).await; + let response = gateway.put_raw(task_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; + let (status, _) = gateway.get_analysis(task_id, READ_SCOPE).await; assert_eq!(status, StatusCode::NOT_FOUND, "no Task is recorded"); } @@ -415,55 +424,53 @@ 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(); + let body = request(session_id, "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; + let response = gateway.put_raw(Uuid::new_v4(), scope, body.clone()).await; assert_eq!(response.status(), StatusCode::FORBIDDEN, "{scope}"); } for scope in [START_SCOPE, "gateway.sessions.read", "gateway.recordings.read"] { - let (status, _) = gateway.get_analysis(session_id, Uuid::new_v4(), scope).await; + let (status, _) = gateway.get_analysis(Uuid::new_v4(), scope).await; assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); } - let response = gateway.post_raw(session_id, "*", body).await; + let response = gateway.put_raw(Uuid::new_v4(), "*", body).await; assert_eq!( response.status(), StatusCode::ACCEPTED, "the wildcard scope is accepted" ); - let (status, _) = gateway.get_analysis(session_id, Uuid::new_v4(), "*").await; + let (status, _) = gateway.get_analysis(Uuid::new_v4(), "*").await; assert_eq!(status, StatusCode::NOT_FOUND, "the wildcard scope is accepted"); } // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis #[tokio::test] -async fn an_analysis_is_found_at_its_location_only() { +async fn an_analysis_is_found_at_its_location() { 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, + .put_raw( + task_id, START_SCOPE, - request(task_id, "gpt-test", &base_url).to_string(), + request(session_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}")); + assert_eq!(location, format!("/jet/tasks/recording-ai-analysis/{task_id}")); - let (status, task) = gateway.get_analysis(session_id, task_id, READ_SCOPE).await; + let (status, task) = gateway.get_analysis(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"); + assert_eq!(task["sessionId"], session_id.to_string()); } // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis @@ -471,7 +478,7 @@ async fn an_analysis_is_found_at_its_location_only() { 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; + let (status, _) = gateway.get_analysis(Uuid::new_v4(), READ_SCOPE).await; assert_eq!(status, StatusCode::NOT_FOUND); } @@ -483,15 +490,14 @@ async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { 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); + let mut body = request(session_id, "gpt-test", &base_url); body["maxOutputTokens"] = json!(1000); - let (status, created) = gateway.post(session_id, &body).await; + let (status, created) = gateway.put(task_id, &body).await; assert_eq!(status, StatusCode::ACCEPTED, "{created}"); - let task = gateway.finished_analysis(session_id, task_id).await; + let task = gateway.finished_analysis(task_id).await; assert_eq!(task["state"], "succeeded", "{task}"); - assert_eq!(task["attempts"], 1); assert_eq!(task["payload"]["fileName"], "ai-analysis-0.slog"); assert_eq!(task["payload"]["model"], "gpt-test-2026-09-30"); // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis @@ -503,7 +509,7 @@ async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { assert_rfc3339(&task["startedAt"]); assert_rfc3339(&task["finishedAt"]); - let (status, again) = gateway.post(session_id, &body).await; + let (status, again) = gateway.put(task_id, &body).await; assert_eq!(status, StatusCode::OK, "{again}"); let responses = [created, task, again].map(|body| body.to_string()).join("\n"); @@ -526,12 +532,11 @@ async fn provider_refusing_the_key_fails_the_task() { 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; + let (status, _) = gateway.put(task_id, &request(session_id, "gpt-test", &base_url)).await; assert_eq!(status, StatusCode::ACCEPTED); - let task = gateway.finished_analysis(session_id, task_id).await; + let task = gateway.finished_analysis(task_id).await; assert_eq!(task["state"], "failed", "{task}"); - assert_eq!(task["attempts"], 1); assert_eq!( task["payload"], json!({ @@ -541,7 +546,7 @@ async fn provider_refusing_the_key_fails_the_task() { ); let (status, _) = gateway - .post(session_id, &request(Uuid::new_v4(), "gpt-test", &base_url)) + .put(Uuid::new_v4(), &request(session_id, "gpt-test", &base_url)) .await; assert_eq!( status, From f9e45d25fbb3462468a455a0c66bdf4ceda9ab06 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 13:27:28 -0400 Subject: [PATCH 60/80] refactor(dgw): mark pending remuxes on disk for the AI analysis An AI analysis must not read a video that its remux is about to replace, and only the recording module should know how a remux is tracked. The recording manager writes `.remuxing` before it schedules the remux, the remux job removes it, and the AI analysis asks the recording module whether a session has such a mark. A mark older than an hour no longer counts, so a crash between the mark and its job cannot block the analyses of the session for good. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/openapi.rs | 2 +- .../src/provisioner_tasks/mod.rs | 8 +- .../recording_ai_analysis/mod.rs | 21 ++--- .../recording_ai_analysis/tests.rs | 73 +++------------ devolutions-gateway/src/recording.rs | 93 +++++++++++++++++-- 5 files changed, 108 insertions(+), 89 deletions(-) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index d2162d8fab..31919c0c76 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -312,7 +312,7 @@ enum AiAnalysisRecoverableErrorKind { 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. + /// Gateway could not read its own storage, such as the folder of the session. Storage, } diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index 2c49d0330e..3eb5c940a8 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -5,8 +5,6 @@ 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; @@ -23,8 +21,8 @@ pub fn task_runner( recordings: RecordingMessageSender, secrets: SecretStore, ) -> ProvisionerTaskRunner { - ProvisionerTaskRunner::builder(store, Arc::clone(&job_queue)) - .register(RecordingAiAnalysis::new(conf_handle, recordings, secrets, job_queue)) + ProvisionerTaskRunner::builder(store, job_queue) + .register(RecordingAiAnalysis::new(conf_handle, recordings, secrets)) .build() } @@ -47,7 +45,7 @@ pub enum Recoverable { 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. + /// Gateway could not read its own storage, such as the folder of the session. Storage, } 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 33a5ca4270..f63c0d78b0 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -31,12 +31,12 @@ use time::OffsetDateTime; use uuid::Uuid; use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; -use super::runner::{Attempt, DynProvisionerTaskJobQueue, ProvisionerTask}; +use super::runner::{Attempt, ProvisionerTask}; use super::{AttemptError, Permanent, Recoverable}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; -use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender, RemuxJob, remux_pending}; +use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender, remux_pending}; use crate::secret_store::{SecretId, SecretStore}; use crate::token::AccessScope; @@ -262,22 +262,15 @@ pub struct RecordingAiAnalysis { conf_handle: ConfHandle, recordings: RecordingMessageSender, secrets: SecretStore, - job_queue: DynProvisionerTaskJobQueue, progress_timeout: time::Duration, } impl RecordingAiAnalysis { - pub fn new( - conf_handle: ConfHandle, - recordings: RecordingMessageSender, - secrets: SecretStore, - job_queue: DynProvisionerTaskJobQueue, - ) -> Self { + pub fn new(conf_handle: ConfHandle, recordings: RecordingMessageSender, secrets: SecretStore) -> Self { Self { conf_handle, recordings, secrets, - job_queue, progress_timeout: PROGRESS_TIMEOUT, } } @@ -610,15 +603,15 @@ impl RecordingAiAnalysis { Ok(()) } - /// Waits until no remux job of the session is stored, since a remux replaces a video that would be read. + /// Waits until no recording of the session waits for its remux, since a remux replaces a video that would be read. async fn wait_for_remux(&self, attempt: &Attempt, dir: &Utf8Path) -> Result<(), AttemptError> { loop { - let jobs = self.job_queue.stored_jobs(RemuxJob::NAME).await.map_err(|error| { - warn!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to list the remux jobs"); + let pending = remux_pending(dir).await.map_err(|error| { + warn!(task.id = %attempt.task_id, %error, %dir, "Failed to look for a pending remux"); AttemptError::Recoverable(Recoverable::Storage) })?; - if !remux_pending(&jobs, dir) { + if !pending { return Ok(()); } 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 2fd9835ec4..228a9246b2 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -12,7 +12,7 @@ use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; use crate::job_queue::{JobMessage, JobQueueHandle, JobQueueReceiver, MAX_ATTEMPTS}; -use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskRunner}; use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles}; @@ -185,7 +185,6 @@ impl Harness { 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( @@ -808,54 +807,22 @@ 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, schedule_for: Option) -> anyhow::Result<()> { - self.queue.push(job, schedule_for).await - } - - async fn stored_jobs(&self, name: &str) -> anyhow::Result> { - Ok(if name == RemuxJob::NAME { - self.remux_jobs.lock().clone() - } else { - Vec::new() - }) - } -} - // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[tokio::test] async fn analysis_waits_for_the_remux_of_the_session() { - let mut harness = Harness::new().await; + let 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 marker = harness + .recording_path + .join(session_id.to_string()) + .join("recording-1.webm.remuxing"); + std::fs::write(&marker, b"").expect("remux mark"); let other_dir = harness.recording_path.join(Uuid::new_v4().to_string()); - let remux_jobs = Arc::new(Mutex::new(vec![ - json!({ "input_path": dir.join("recording-1.webm") }).to_string(), - json!({ "input_path": other_dir.join("recording-0.webm") }).to_string(), - ])); - let queue: DynProvisionerTaskJobQueue = Arc::new(WithRemuxJobs { - queue: harness.state.job_queue_handle.clone(), - remux_jobs: Arc::clone(&remux_jobs), - }); - harness.state.provisioner_tasks = - ProvisionerTaskRunner::builder(Arc::clone(harness.state.provisioner_tasks.store()), Arc::clone(&queue)) - .register(RecordingAiAnalysis::new( - harness.state.conf_handle.clone(), - harness.state.recordings.clone(), - harness.state.secrets.clone(), - queue, - )) - .build(); + std::fs::create_dir_all(&other_dir).expect("other session"); + std::fs::write(other_dir.join("recording-0.webm.remuxing"), b"").expect("remux mark of another session"); let mut job = harness.job(task_id).await; let job = tokio::spawn(async move { job.run().await }); @@ -864,7 +831,7 @@ async fn analysis_waits_for_the_remux_of_the_session() { assert!(!job.is_finished(), "the analysis waits"); assert!(requests.lock().is_empty(), "the AI is not asked yet"); - remux_jobs.lock().remove(0); + std::fs::remove_file(&marker).expect("remux done"); job.await.expect("job task").expect("no retry"); let task = harness.task(task_id).await; @@ -881,23 +848,8 @@ async fn waiting_for_a_remux_is_no_progress() { 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(); + std::fs::write(dir.join("recording-1.webm.remuxing"), b"").expect("remux mark"); + harness.with_progress_timeout(time::Duration::seconds(2)); harness.run(task_id).await.expect("no retry"); @@ -1268,7 +1220,6 @@ async fn an_ended_task_leaves_the_working_files_of_another_task() { harness.state.conf_handle.clone(), harness.state.recordings.clone(), harness.state.secrets.clone(), - Arc::new(harness.state.job_queue_handle.clone()), ); task.on_end(&record).await.expect("end"); diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index afdb165f74..55fb28db4e 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -943,8 +943,15 @@ impl RecordingManagerTask { if cadeau::xmf::is_init() { debug!(%recording_file_path, "Enqueue video remuxing operation"); + let marker = remux_marker(&recording_file_path); + + // An AI analysis that finds no mark reads the file even while it is remuxed. + if let Err(error) = fs::write(&marker, b"").await { + warn!(%marker, %error, "Failed to mark the recording as waiting for its remux"); + } + // Schedule 60 seconds to wait for the streamers to release the file. - let _ = self + let scheduled = self .job_queue_handle .schedule( RemuxJob { @@ -953,6 +960,10 @@ impl RecordingManagerTask { time::OffsetDateTime::now_utc() + time::Duration::seconds(60), ) .await; + + if scheduled.is_err() { + remove_remux_marker(&marker).await; + } } else { debug!("Video remuxing was skipped because XMF native library is not loaded"); } @@ -1271,12 +1282,49 @@ 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)) +const REMUX_MARKER_SUFFIX: &str = ".remuxing"; + +/// Age past which a remux mark no longer counts. +/// +/// A crash between the mark and the storing of its job leaves a mark that no remux removes; without this limit, it +/// would make every AI analysis of the session wait until it fails. No remux takes this long. +const REMUX_MARKER_STALE_AFTER: std::time::Duration = std::time::Duration::from_secs(60 * 60); + +/// The file that marks `recording_file` as waiting for its remux, from the time the remux is scheduled until it ends. +fn remux_marker(recording_file: &Utf8Path) -> Utf8PathBuf { + Utf8PathBuf::from(format!("{recording_file}{REMUX_MARKER_SUFFIX}")) +} + +async fn remove_remux_marker(marker: &Utf8Path) { + match fs::remove_file(marker).await { + Err(error) if error.kind() != io::ErrorKind::NotFound => { + warn!(%marker, %error, "Failed to remove the remux mark"); + } + _ => {} + } +} + +/// Whether a recording file of the session folder `dir` waits for its remux, which would replace the file. +pub(crate) async fn remux_pending(dir: &Utf8Path) -> io::Result { + let mut entries = match fs::read_dir(dir).await { + Ok(entries) => entries, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(false), + Err(error) => return Err(error), + }; + + while let Some(entry) = entries.next_entry().await? { + if !entry.file_name().to_string_lossy().ends_with(REMUX_MARKER_SUFFIX) { + continue; + } + + let age = entry.metadata().await?.modified()?.elapsed().unwrap_or_default(); + + if age < REMUX_MARKER_STALE_AFTER { + return Ok(true); + } + } + + Ok(false) } #[async_trait] @@ -1290,7 +1338,10 @@ impl job_queue::Job for RemuxJob { } async fn run(&mut self) -> anyhow::Result<()> { - remux(core::mem::take(&mut self.input_path)).await; + let input_path = core::mem::take(&mut self.input_path); + let marker = remux_marker(&input_path); + remux(input_path).await; + remove_remux_marker(&marker).await; Ok(()) } } @@ -1433,6 +1484,32 @@ mod tests { } } + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + #[tokio::test] + async fn a_fresh_remux_mark_is_pending_and_a_stale_one_is_not() { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); + let marker = remux_marker(&dir.join("recording-0.webm")); + + assert!(!remux_pending(&dir.join("gone")).await.expect("missing folder")); + assert!(!remux_pending(dir).await.expect("no mark")); + + std::fs::write(&marker, b"").expect("mark"); + assert!(remux_pending(dir).await.expect("fresh mark")); + + let stale = std::time::SystemTime::now() - REMUX_MARKER_STALE_AFTER - std::time::Duration::from_secs(60); + std::fs::File::options() + .write(true) + .open(&marker) + .and_then(|file| file.set_modified(stale)) + .expect("age the mark"); + assert!(!remux_pending(dir).await.expect("stale mark")); + + remove_remux_marker(&marker).await; + remove_remux_marker(&marker).await; + assert!(!marker.exists(), "removing an absent mark is no error"); + } + #[tokio::test] async fn slog_recording_stays_in_files() { let harness = Harness::start(); From 9963cb0f3fdd515af3bcbe7f7b83addb2e9db773 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 13:30:18 -0400 Subject: [PATCH 61/80] refactor(dgw): count Task attempts on the record only, check scopes in the routes The record's `attempts` is the one count of attempts, so the job queue and the runner never disagree on when a Task gives up. A recoverable failure queues the next attempt with the same token, 30 seconds later and doubling, and the job returns Ok. The job queue then retries a job only when what it had to record or queue could not be stored. The extractors of the AI analysis routes check the start and read scopes, since authorization belongs to the API layer. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../src/api/tasks/recording_ai_analysis.rs | 10 ++--- devolutions-gateway/src/extract.rs | 37 ++++++++----------- devolutions-gateway/src/job_queue.rs | 4 -- .../recording_ai_analysis/mod.rs | 5 --- .../recording_ai_analysis/tests.rs | 20 +++++++--- .../src/provisioner_tasks/runner/mod.rs | 29 +++++++++------ .../src/provisioner_tasks/runner/tests.rs | 17 ++++++--- 7 files changed, 63 insertions(+), 59 deletions(-) diff --git a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs index 5cd16bf0ea..c4c3e37d99 100644 --- a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs +++ b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs @@ -15,12 +15,10 @@ use uuid::Uuid; use super::TaskState; use crate::DgwState; -use crate::extract::{ProvisionerTaskReadScope, ProvisionerTaskStartScope}; +use crate::extract::{RecordingAiAnalysisReadScope, RecordingAiAnalysisStartScope}; use crate::http::HttpError; use crate::provisioner_tasks::ai::{AiProvider, AiSettings}; -use crate::provisioner_tasks::recording_ai_analysis::{ - self as analysis, RecordingAiAnalysis, StartError, StartOutcome, StartRequest, -}; +use crate::provisioner_tasks::recording_ai_analysis::{self as analysis, StartError, StartOutcome, StartRequest}; /// Settings of an AI analysis #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] @@ -174,7 +172,7 @@ fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> ))] pub(crate) async fn start_ai_analysis( State(state): State, - _scope: ProvisionerTaskStartScope, + _scope: RecordingAiAnalysisStartScope, extract::Path(task_id): extract::Path, body: Bytes, ) -> Result { @@ -269,7 +267,7 @@ pub(crate) async fn start_ai_analysis( ))] pub(crate) async fn get_ai_analysis( State(DgwState { provisioner_tasks, .. }): State, - _scope: ProvisionerTaskReadScope, + _scope: RecordingAiAnalysisReadScope, extract::Path(task_id): extract::Path, ) -> Result, HttpError> { let task = provisioner_tasks diff --git a/devolutions-gateway/src/extract.rs b/devolutions-gateway/src/extract.rs index 224c9382a6..4f5f1549ae 100644 --- a/devolutions-gateway/src/extract.rs +++ b/devolutions-gateway/src/extract.rs @@ -6,7 +6,6 @@ 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, @@ -322,44 +321,40 @@ 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); +/// A scope token allowed to start an AI analysis: `*`, or `gateway.tasks.recording-ai-analysis.start`. +#[derive(Clone, Copy)] +pub struct RecordingAiAnalysisStartScope; -impl FromRequestParts for ProvisionerTaskReadScope +impl FromRequestParts for RecordingAiAnalysisStartScope 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")) + match ScopeToken::from_request_parts(parts, state).await?.0.scope { + AccessScope::Wildcard => Ok(Self), + AccessScope::RecordingAiAnalysisStart => Ok(Self), + _ => 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); +/// A scope token allowed to read AI analyses: `*`, or `gateway.tasks.recording-ai-analysis.read`. +#[derive(Clone, Copy)] +pub struct RecordingAiAnalysisReadScope; -impl FromRequestParts for ProvisionerTaskStartScope +impl FromRequestParts for RecordingAiAnalysisReadScope 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")) + match ScopeToken::from_request_parts(parts, state).await?.0.scope { + AccessScope::Wildcard => Ok(Self), + AccessScope::RecordingAiAnalysisRead => Ok(Self), + _ => Err(HttpError::forbidden().msg("invalid scope for route")), } } } diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 656e9887bc..079efbc757 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -13,9 +13,6 @@ use tokio::sync::{Notify, mpsc}; use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; -/// Runs of a job before the job queue gives up on it. -pub const MAX_ATTEMPTS: u32 = 5; - pub struct JobQueueCtx { notify_runner: Arc, runner_waker: RunnerWaker, @@ -65,7 +62,6 @@ 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/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index f63c0d78b0..d3b4cb2c3c 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -38,7 +38,6 @@ use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender, remux_pending}; use crate::secret_store::{SecretId, SecretStore}; -use crate::token::AccessScope; pub const KIND: &str = "recording-ai-analysis"; @@ -336,10 +335,6 @@ async fn remove_if_exists(path: &Utf8Path) -> std::io::Result<()> { 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/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 228a9246b2..2656b3e211 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -11,8 +11,8 @@ use parking_lot::Mutex; use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; -use crate::job_queue::{JobMessage, JobQueueHandle, JobQueueReceiver, MAX_ATTEMPTS}; -use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskRunner}; +use crate::job_queue::{JobMessage, JobQueueHandle, JobQueueReceiver}; +use crate::provisioner_tasks::runner::{MAX_ATTEMPTS, ProvisionerTaskJob, ProvisionerTaskRunner}; use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles}; @@ -143,6 +143,11 @@ impl Harness { self.job(task_id).await.run().await } + /// Whether a job waits to run a later attempt, as after a recoverable failure; it takes every queued job. + fn later_attempt_queued(&mut self) -> bool { + std::iter::from_fn(|| self.job_queue_rx.try_recv().ok()).any(|message| message.schedule_for.is_some()) + } + /// Records a Task directly, as if it was created at `created_at`. async fn create_task( &self, @@ -517,14 +522,15 @@ async fn refuses_invalid_ai_settings() { // 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; + let mut 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"); + harness.run(task_id).await.expect("the job queues the next attempt itself"); + assert!(harness.later_attempt_queued(), "a provider outage is retried"); let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Running); @@ -663,17 +669,19 @@ async fn recording_deleted_after_start_fails() { // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management #[tokio::test] async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { - let harness = Harness::new().await; + let mut 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"); + harness.run(task_id).await.expect("the job queues the next attempt itself"); + assert!(harness.later_attempt_queued(), "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"); + assert!(!harness.later_attempt_queued(), "no attempt after the last"); let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index a8278960e6..103dc4c718 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -34,12 +34,19 @@ use tokio_util::sync::CancellationToken; use uuid::Uuid; use super::{AttemptError, Permanent}; -use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; -use crate::token::AccessScope; +use crate::job_queue::JobQueueHandle; /// Time before a record whose end failed is ended again. const END_RETRY_DELAY: time::Duration = time::Duration::minutes(5); +/// Attempts a record gets: a recoverable failure of the last one fails the record. +pub const MAX_ATTEMPTS: u32 = 5; + +/// Time before the attempt that follows attempt number `attempts`: 30 seconds, doubled after each attempt. +fn retry_delay(attempts: u32) -> time::Duration { + time::Duration::seconds(30) * (1_i32 << attempts.saturating_sub(1).min(MAX_ATTEMPTS)) +} + /// Where the runner queues the jobs of its records. #[async_trait] pub trait ProvisionerTaskJobQueue: Send + Sync { @@ -78,12 +85,6 @@ 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; @@ -430,7 +431,10 @@ impl ProvisionerTaskRunner { .await } - /// Runs one attempt for the job holding `job_token`; returns an error only to ask the job queue for a retry. + /// Runs one attempt for the job holding `job_token`. + /// + /// The record counts its attempts, and the job queues the next one itself, so the job queue retries a job only when + /// what it had to record or queue could not be stored. async fn run_attempt(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { let store = &self.inner.store; let now = OffsetDateTime::now_utc(); @@ -577,8 +581,11 @@ impl ProvisionerTaskRunner { } }, 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:?}"); + let retry_at = OffsetDateTime::now_utc() + retry_delay(attempts); + warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, %retry_at, "Task attempt failed; it will be retried"); + self.enqueue(task_id, job_token, Some(retry_at)) + .await + .context("queue the next attempt")?; } Err(AttemptError::Recoverable(error)) => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, "Task failed after its last attempt"); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 9f0cdfe76b..d971a19823 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -6,7 +6,6 @@ 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"; @@ -35,10 +34,6 @@ 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; @@ -397,10 +392,20 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { let mut job = harness.queued_job(); for attempt in 1..MAX_ATTEMPTS { - job.run().await.expect_err("the job queue retries"); + let failed_at = OffsetDateTime::now_utc(); + job.run().await.expect("the job queues the next attempt itself"); assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Running, "{attempt}"); + + let JobMessage { + job: next, + schedule_for, + } = harness.job_queue_rx.try_recv().expect("the next attempt is queued"); + let retry_at = schedule_for.expect("the next attempt waits"); + assert!(retry_at >= failed_at + retry_delay(attempt), "{attempt}: {retry_at}"); + job = next; } job.run().await.expect("the last attempt"); + assert!(harness.job_queue_rx.try_recv().is_err(), "no attempt after the last"); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); From 2f1aaa17938f7a61506008554a14cc9df0ace1de Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 13:34:00 -0400 Subject: [PATCH 62/80] refactor(dgw): give every unfinished Task a new job at startup Every unfinished Task gets a new job token and a new job at startup, so a job left from before stops as superseded and the runner never needs to know which jobs are stored. A kind that cannot outlive a restart says so with `after_restart`: an AI analysis, whose key lived in memory only, fails with `keyLost` at once, and its new job only ends it. `JobQueue::job_definitions` is removed, since nothing uses it. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 - crates/job-queue-libsql/Cargo.toml | 3 - crates/job-queue-libsql/src/lib.rs | 74 ----------- crates/job-queue/src/lib.rs | 3 - devolutions-gateway/src/job_queue.rs | 4 - .../recording_ai_analysis/mod.rs | 5 + .../recording_ai_analysis/tests.rs | 21 ++++ .../src/provisioner_tasks/runner/mod.rs | 73 ++++++----- .../src/provisioner_tasks/runner/tests.rs | 116 ++++++++---------- 9 files changed, 124 insertions(+), 176 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2a2d350493..ec3bdc383b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4321,7 +4321,6 @@ 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 a4a4934a3f..efdf7da010 100644 --- a/crates/job-queue-libsql/Cargo.toml +++ b/crates/job-queue-libsql/Cargo.toml @@ -19,6 +19,3 @@ 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 10737ce59d..0efbfd23f0 100644 --- a/crates/job-queue-libsql/src/lib.rs +++ b/crates/job-queue-libsql/src/lib.rs @@ -402,27 +402,6 @@ 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. @@ -446,56 +425,3 @@ const MIGRATIONS: &[&str] = &[ CREATE INDEX idx_scheduled_for ON job_queue(scheduled_for);", ]; - -#[cfg(test)] -mod tests { - use job_queue::{DynJob, Job, JobQueue as _, RunnerWaker}; - - use super::*; - - struct Named { - name: &'static str, - def: &'static str, - } - - #[async_trait] - impl Job for Named { - fn name(&self) -> &str { - self.name - } - - fn write_json(&self) -> anyhow::Result { - Ok(self.def.to_owned()) - } - - async fn run(&mut self) -> anyhow::Result<()> { - Ok(()) - } - } - - // oracle: mutant the name filter is joined with OR, so the jobs of every name are listed - #[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 5521da5281..93232bd868 100644 --- a/crates/job-queue/src/lib.rs +++ b/crates/job-queue/src/lib.rs @@ -62,9 +62,6 @@ 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 079efbc757..d51a623ddd 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -102,10 +102,6 @@ impl ProvisionerTaskJobQueue for DurableJobQueue { let job: DynJob = Box::new(job); self.0.push_job(&job, schedule_for).await } - - async fn stored_jobs(&self, name: &str) -> anyhow::Result> { - self.0.job_definitions(name).await - } } impl JobQueueHandle { 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 d3b4cb2c3c..c4336be601 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -421,6 +421,11 @@ impl ProvisionerTask for RecordingAiAnalysis { .await } + /// The API key was kept in memory only, so no attempt can run after a restart. + fn after_restart(&self) -> Option { + Some(Permanent::KeyLost) + } + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { self.secrets.remove(SecretId(record.id)); 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 2656b3e211..2c4652c184 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -954,6 +954,27 @@ async fn a_new_task_that_fails_keeps_the_earlier_analysis() { assert!(dir.join("ai-analysis-0.slog").exists()); } +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets +#[tokio::test] +async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { + 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; + std::fs::create_dir_all(harness.workspace(session_id)).expect("working files of an attempt"); + + harness.state.provisioner_tasks.resume().await.expect("resume"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed, "no job ran yet"); + assert_eq!(task.payload["error"], json!({ "kind": "keyLost" })); + + harness.run(task_id).await.expect("the new job only ends the Task"); + assert!(requests.lock().is_empty(), "the AI is never asked"); + assert!(!harness.workspace(session_id).exists(), "the working files are removed"); + assert!(!harness.has_key(task_id)); +} + // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis #[tokio::test] async fn a_failed_task_reports_the_tokens_it_spent() { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 103dc4c718..2bf86b7d54 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -53,9 +53,6 @@ pub trait ProvisionerTaskJobQueue: Send + Sync { /// Queues `job` to run at `schedule_for`, or at once without it, and returns once it is stored, so a record never /// waits for a job that was lost. async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()>; - - /// Lists the JSON definitions of the stored jobs named `name` that may still run. - async fn stored_jobs(&self, name: &str) -> anyhow::Result>; } pub type DynProvisionerTaskJobQueue = Arc; @@ -69,11 +66,6 @@ impl ProvisionerTaskJobQueue for JobQueueHandle { None => self.enqueue(job).await, } } - - /// The channel cannot tell what is stored: every unfinished record looks like it has no job. - async fn stored_jobs(&self, _name: &str) -> anyhow::Result> { - Ok(Vec::new()) - } } /// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. @@ -105,6 +97,12 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// 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; + /// Why an unfinished record fails when Gateway starts again, for a task that cannot outlive a restart, such as one + /// whose secrets were kept in memory only; `None` for a task that goes on. + fn after_restart(&self) -> Option { + None + } + /// Runs one attempt. /// /// The attempt should stop soon after [`Attempt::cancel`] is cancelled, which happens once the record is past its @@ -190,6 +188,8 @@ trait AnyProvisionerTask: Send + Sync { fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value; + fn after_restart(&self) -> Option; + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result; async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()>; @@ -213,6 +213,10 @@ impl AnyProvisionerTask for T { ProvisionerTask::failed_payload(self, error, last) } + fn after_restart(&self) -> Option { + ProvisionerTask::after_restart(self) + } + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { let params = serde_json::from_value(params).map_err(|error| { warn!(task.id = %attempt.task_id, %error, "Invalid Task parameters"); @@ -341,31 +345,17 @@ impl ProvisionerTaskRunner { Ok(outcome) } - /// Hands every unfinished record that has no job left to a new job, such as a record created just before a crash. + /// Hands every unfinished record to a new job, since a record created just before Gateway stopped may have none. /// - /// 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. + /// Each record gets a new `job_token`, so a job left from before finds itself superseded and stops, and the job queue + /// is never read. A record whose task cannot outlive a restart fails first, and its new job only ends it. Call it at + /// startup, before the job runner claims any job, so no job runs an attempt meanwhile. A record whose new job cannot + /// be queued is failed, since no job is left to run it. pub async fn resume(&self) -> anyhow::Result<()> { let now = OffsetDateTime::now_utc(); let records = self.inner.store.list_unfinished(now).await?; - let stored = self - .inner - .job_queue - .stored_jobs(ProvisionerTaskJob::NAME) - .await - .context("list the stored provisioner task jobs")? - .iter() - .filter_map(|json| serde_json::from_str::(json).ok()) - .map(|def| (def.task_id, def.job_token)) - .collect::>(); - for record in records { - if stored.contains(&(record.id, record.job_token)) { - debug!(task.id = %record.id, task.kind = %record.kind, "Task still has its job"); - continue; - } - let job_token = Uuid::new_v4(); match self.inner.store.replace_job_token(record.id, job_token, now).await { @@ -382,6 +372,33 @@ impl ProvisionerTaskRunner { } } + let task = self.task(&record.kind); + + if let Some(task) = &task + && let Some(failure) = task.after_restart() + { + match self + .inner + .store + .fail(record.id, task.failed_payload(&failure, &record.payload), now) + .await + { + Ok(_) => info!( + task.id = %record.id, + task.kind = %record.kind, + ?failure, + "Task failed, since it cannot outlive a restart" + ), + // Its new job runs the attempt then, which fails the same way. + Err(error) => error!( + task.id = %record.id, + task.kind = %record.kind, + error = format!("{error:#}"), + "Failed to fail the Task that cannot outlive a restart" + ), + } + } + match self.enqueue(record.id, job_token, None).await { Ok(()) => info!(task.id = %record.id, task.kind = %record.kind, "Task queued again"), Err(error) => { @@ -391,7 +408,7 @@ impl ProvisionerTaskRunner { error = format!("{error:#}"), "Failed to queue the Task again" ); - if let Some(task) = self.task(&record.kind) { + if let Some(task) = &task { self.fail_not_queued(task.as_ref(), record.id, job_token, &record.payload) .await; } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index d971a19823..d62a88a37e 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -28,6 +28,8 @@ struct Echo { ended: Arc>>, /// Ends that fail before one succeeds. end_failures: Arc, + /// Why its unfinished records fail when Gateway starts again. + restart_failure: Option, } #[async_trait] @@ -54,6 +56,10 @@ impl ProvisionerTask for Echo { json!({ "error": error, "last": last }) } + fn after_restart(&self) -> Option { + self.restart_failure.clone() + } + async fn run(&self, attempt: &Attempt, params: EchoParams) -> Result { self.runs.fetch_add(1, Ordering::SeqCst); @@ -107,6 +113,19 @@ impl Harness { async fn with_progress_timeout( outcome: fn(u32) -> Result, progress_timeout: time::Duration, + ) -> Self { + Self::build(outcome, progress_timeout, None).await + } + + /// A harness whose task cannot outlive a restart, like one that keeps its secrets in memory only. + async fn forgetting_at_restart(outcome: fn(u32) -> Result) -> Self { + Self::build(outcome, time::Duration::hours(1), Some(Permanent::KeyLost)).await + } + + async fn build( + outcome: fn(u32) -> Result, + progress_timeout: time::Duration, + restart_failure: Option, ) -> Self { let conn = gateway_db::GatewayDb::open_path(":memory:") .await @@ -131,6 +150,7 @@ impl Harness { runs: Arc::clone(&runs), ended: Arc::clone(&ended), end_failures: Arc::clone(&end_failures), + restart_failure, }) .build(); @@ -279,77 +299,47 @@ 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>); +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +#[tokio::test] +async fn resume_gives_every_unfinished_task_a_new_job() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let first = harness.create(1).await; + let second = harness.create(2).await; + let tokens = [harness.task(first).await.job_token, harness.task(second).await.job_token]; + // The jobs of the creation may be stored or lost; the runner never asks which. + while harness.job_queue_rx.try_recv().is_ok() {} -#[async_trait] -impl ProvisionerTaskJobQueue for StoredJobs { - async fn push(&self, job: ProvisionerTaskJob, _schedule_for: Option) -> anyhow::Result<()> { - self.0.lock().push(job_queue::Job::write_json(&job)?); - Ok(()) - } + harness.runner.resume().await.expect("resume"); - async fn stored_jobs(&self, name: &str) -> anyhow::Result> { - assert_eq!(name, ProvisionerTaskJob::NAME); - Ok(self.0.lock().clone()) - } + assert_eq!( + std::iter::from_fn(|| harness.job_queue_rx.try_recv().ok()).count(), + 2, + "one new job for each unfinished Task" + ); + assert_ne!(harness.task(first).await.job_token, tokens[0]); + assert_ne!(harness.task(second).await.job_token, tokens[1]); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets #[tokio::test] -async fn resume_queues_only_the_tasks_left_without_a_job() { - let conn = gateway_db::GatewayDb::open_path(":memory:") - .await - .expect("open gateway database") - .connect() - .await - .expect("connect"); - let store = Arc::new( - provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) - .await - .expect("open task store"), - ); - let queue = Arc::new(StoredJobs::default()); - let runner = ProvisionerTaskRunner::builder(store, Arc::clone(&queue) as DynProvisionerTaskJobQueue) - .register(Echo { - outcome: |value| Ok(json!(value)), - progress_timeout: time::Duration::hours(1), - runs: Arc::default(), - ended: Arc::default(), - end_failures: Arc::default(), - }) - .build(); - let create = |id: Uuid| runner.create::(id, format!("target-{id}"), &EchoParams { value: 1 }, |_| Ok(())); - let with_job = Uuid::new_v4(); - let without_job = Uuid::new_v4(); - create(with_job).await.expect("create"); - create(without_job).await.expect("create"); - let token_with_job = runner - .get(with_job, OffsetDateTime::now_utc()) - .await - .expect("read") - .expect("record") - .job_token; +async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { + let mut harness = Harness::forgetting_at_restart(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + let mut stale = harness.queued_job(); - // 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())); + harness.runner.resume().await.expect("resume"); - runner.resume().await.expect("resume"); + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed, "it fails before any job runs"); + assert_eq!(task.payload, json!({ "error": { "kind": "keyLost" }, "last": null })); - 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" - ); + let mut ending = harness.queued_job(); + stale.run().await.expect("the stale job stops"); + assert!(harness.ended.lock().is_empty(), "a stale job ends nothing"); + + ending.run().await.expect("the new job ends the Task"); + assert_eq!(harness.runs.load(Ordering::SeqCst), 0, "no attempt runs"); + assert_eq!(*harness.ended.lock(), [id]); } // oracle: mutant a panicking attempt fails the Task as internal instead of panicked From fb531c4df6dfd7431c40da0be698d7797d13ec89 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 13:56:48 -0400 Subject: [PATCH 63/80] fix(dgw): name the AI analysis schemas the Task schema refers to utoipa turns a `value_type` path into the schema name as written, so `crate::openapi::AiAnalysisParams` became the reference `crate.openapi.AiAnalysisParams`, which names no schema. A bare name refers to the registered schema, and since utoipa never uses the Rust type behind it, the schema types stay private to `openapi`. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/api/tasks/recording_ai_analysis.rs | 4 ++-- devolutions-gateway/src/openapi.rs | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs index c4c3e37d99..c30ef59f6f 100644 --- a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs +++ b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs @@ -82,12 +82,12 @@ pub(crate) struct AiAnalysisTask { /// Recorded session the analysis is about. session_id: Uuid, /// What was asked, without the API key and the base URL. - #[cfg_attr(feature = "openapi", schema(value_type = crate::openapi::AiAnalysisParams))] + #[cfg_attr(feature = "openapi", schema(value_type = AiAnalysisParams))] params: serde_json::Value, state: TaskState, /// Details of the state: progress while `running`, the log once `succeeded`, the reason once `failed`; null while /// `queued`. - #[cfg_attr(feature = "openapi", schema(value_type = Option))] + #[cfg_attr(feature = "openapi", schema(value_type = Option))] payload: serde_json::Value, /// When the Task was created. #[serde(with = "time::serde::rfc3339")] diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 31919c0c76..fa79b71c2f 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -166,7 +166,7 @@ struct SessionInfo { /// 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")] -pub(crate) struct AiAnalysisParams { +struct AiAnalysisParams { provider: AiProvider, /// Model identifier, passed to the provider as is. model: String, @@ -181,7 +181,7 @@ pub(crate) struct AiAnalysisParams { /// Details of the state of an AI analysis Task: `running`, `succeeded` and `failed` each have their own #[derive(Serialize, utoipa::ToSchema)] #[serde(untagged)] -pub(crate) enum AiAnalysisPayload { +enum AiAnalysisPayload { Running(AiAnalysisRunningPayload), Succeeded(AiAnalysisSucceededPayload), Failed(AiAnalysisFailedPayload), From 8a2b8cdcd1f49335ea67bc8a70694a0759a4a596 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 13:58:45 -0400 Subject: [PATCH 64/80] fix(dgw): let a forgotten remux mark expire before an AI analysis times out Waiting for a remux is not progress, so an AI analysis waiting on a mark that no remux will remove fails once 30 minutes pass without progress. A mark that stops counting after 10 minutes lets the analysis go on in time, and no remux takes that long. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/recording.rs | 46 ++++++++++++++++++++-------- 1 file changed, 34 insertions(+), 12 deletions(-) diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 55fb28db4e..59ec4f1727 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -1286,9 +1286,10 @@ const REMUX_MARKER_SUFFIX: &str = ".remuxing"; /// Age past which a remux mark no longer counts. /// -/// A crash between the mark and the storing of its job leaves a mark that no remux removes; without this limit, it -/// would make every AI analysis of the session wait until it fails. No remux takes this long. -const REMUX_MARKER_STALE_AFTER: std::time::Duration = std::time::Duration::from_secs(60 * 60); +/// A crash between the mark and the storing of its job leaves a mark that no remux removes. +/// The limit lets an AI analysis waiting on such a mark go on well before it fails for lack of progress, and no remux +/// takes this long. +const REMUX_MARKER_STALE_AFTER: Duration = Duration::from_secs(10 * 60); /// The file that marks `recording_file` as waiting for its remux, from the time the remux is scheduled until it ends. fn remux_marker(recording_file: &Utf8Path) -> Utf8PathBuf { @@ -1486,28 +1487,49 @@ mod tests { // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[tokio::test] - async fn a_fresh_remux_mark_is_pending_and_a_stale_one_is_not() { + async fn a_recording_waits_for_its_remux_while_its_mark_is_fresh() { let dir = tempfile::tempdir().expect("temp dir"); let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); - let marker = remux_marker(&dir.join("recording-0.webm")); + std::fs::write(remux_marker(&dir.join("recording-0.webm")), b"").expect("mark"); - assert!(!remux_pending(&dir.join("gone")).await.expect("missing folder")); - assert!(!remux_pending(dir).await.expect("no mark")); + assert!(remux_pending(dir).await.expect("fresh mark")); + } + // oracle: mutant a mark left by a crash before its remux job was stored keeps every AI analysis waiting + #[tokio::test] + async fn a_remux_mark_older_than_any_remux_no_longer_counts() { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); + let marker = remux_marker(&dir.join("recording-0.webm")); std::fs::write(&marker, b"").expect("mark"); - assert!(remux_pending(dir).await.expect("fresh mark")); - let stale = std::time::SystemTime::now() - REMUX_MARKER_STALE_AFTER - std::time::Duration::from_secs(60); + let stale = std::time::SystemTime::now() - REMUX_MARKER_STALE_AFTER - Duration::from_secs(60); std::fs::File::options() .write(true) .open(&marker) .and_then(|file| file.set_modified(stale)) .expect("age the mark"); + assert!(!remux_pending(dir).await.expect("stale mark")); + } - remove_remux_marker(&marker).await; - remove_remux_marker(&marker).await; - assert!(!marker.exists(), "removing an absent mark is no error"); + // oracle: mutant a session folder without a mark, or without a folder yet, fails the AI analysis as storage + #[tokio::test] + async fn a_session_without_remux_marks_has_no_pending_remux() { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); + std::fs::write(dir.join("recording-0.webm"), b"").expect("recording"); + + assert!(!remux_pending(dir).await.expect("no mark")); + assert!(!remux_pending(&dir.join("gone")).await.expect("missing folder")); + } + + // oracle: mutant a remux mark left by a crash outlasts the progress deadline, so AI analyses of the session fail + #[test] + fn a_forgotten_remux_mark_stops_counting_before_an_ai_analysis_times_out() { + let progress_timeout = crate::provisioner_tasks::recording_ai_analysis::PROGRESS_TIMEOUT.unsigned_abs(); + + assert!(REMUX_MARKER_STALE_AFTER < progress_timeout); } #[tokio::test] From 453c13229c2cadc806d86f48c48becb58e4a2509 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:04:29 -0400 Subject: [PATCH 65/80] refactor(dgw): have every provisioner task say whether a restart fails it A record that cannot be failed at startup still gets a new job, so a task that cannot outlive a restart must fail the same way when that job runs; a required `restart_failure` makes each kind decide it. A success that cannot be recorded fails the Task instead of being retried, since a retry would do the paid work again; a store that fails to record a success now checks it. The retry delay doubles with `saturating_pow`, so no attempt count can overflow the shift. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/mod.rs | 4 +- .../src/provisioner_tasks/runner/mod.rs | 46 +++--- .../src/provisioner_tasks/runner/tests.rs | 133 +++++++++++++++--- 3 files changed, 137 insertions(+), 46 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 c4336be601..9912b37744 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -421,8 +421,8 @@ impl ProvisionerTask for RecordingAiAnalysis { .await } - /// The API key was kept in memory only, so no attempt can run after a restart. - fn after_restart(&self) -> Option { + /// The API key is kept in memory only, so no attempt can run after a restart. + fn restart_failure(&self) -> Option { Some(Permanent::KeyLost) } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 2bf86b7d54..f2950e8b25 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -6,12 +6,13 @@ //! retries recoverable failures, moves the deadline ahead with each progress, cancels an attempt once the deadline passes, //! and records how the task ended. //! -//! The record is the source of truth and the job only wakes it up. Each record names the one job allowed to run it, by -//! its token, so a stale or duplicate job stops without running anything. At startup, before any job runs, -//! [`ProvisionerTaskRunner::resume`] hands every unfinished record to a new job, which also covers a record created just -//! before a crash, before its job was queued. -//! A record whose end fails is handed to a new job that ends it again later, since the runs of its job are shared with -//! its attempts and may be used up; that job is stored, so a restart runs it too. +//! The record is the source of truth and the job only wakes it up. +//! Each record names the one job allowed to run it, by its token, so a stale or duplicate job stops without running +//! anything. +//! At startup, before any job runs, [`ProvisionerTaskRunner::resume`] hands every unfinished record to a new job, which +//! also covers a record created just before a crash, before its job was queued. +//! A record whose end fails is handed to a new job that ends it again later, since the job queue gives up on a job +//! after a few failed runs; that job is stored, so a restart runs it too. #[cfg(test)] mod tests; @@ -44,7 +45,7 @@ pub const MAX_ATTEMPTS: u32 = 5; /// Time before the attempt that follows attempt number `attempts`: 30 seconds, doubled after each attempt. fn retry_delay(attempts: u32) -> time::Duration { - time::Duration::seconds(30) * (1_i32 << attempts.saturating_sub(1).min(MAX_ATTEMPTS)) + time::Duration::seconds(30) * 2_i32.saturating_pow(attempts.saturating_sub(1)) } /// Where the runner queues the jobs of its records. @@ -98,10 +99,11 @@ pub trait ProvisionerTask: Send + Sync + 'static { fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value; /// Why an unfinished record fails when Gateway starts again, for a task that cannot outlive a restart, such as one - /// whose secrets were kept in memory only; `None` for a task that goes on. - fn after_restart(&self) -> Option { - None - } + /// whose secrets are kept in memory only; `None` for a task that goes on. + /// + /// A record that cannot be failed at startup is still handed to a new job, so [`ProvisionerTask::run`] must fail + /// the same way after a restart. + fn restart_failure(&self) -> Option; /// Runs one attempt. /// @@ -188,7 +190,7 @@ trait AnyProvisionerTask: Send + Sync { fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value; - fn after_restart(&self) -> Option; + fn restart_failure(&self) -> Option; async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result; @@ -213,8 +215,8 @@ impl AnyProvisionerTask for T { ProvisionerTask::failed_payload(self, error, last) } - fn after_restart(&self) -> Option { - ProvisionerTask::after_restart(self) + fn restart_failure(&self) -> Option { + ProvisionerTask::restart_failure(self) } async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { @@ -347,10 +349,11 @@ impl ProvisionerTaskRunner { /// Hands every unfinished record to a new job, since a record created just before Gateway stopped may have none. /// - /// Each record gets a new `job_token`, so a job left from before finds itself superseded and stops, and the job queue - /// is never read. A record whose task cannot outlive a restart fails first, and its new job only ends 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. + /// Each record gets a new `job_token`, so a job left from before finds itself superseded and stops, and the job + /// queue is never read. + /// A record whose task cannot outlive a restart fails first, and its new job only ends 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?; @@ -375,7 +378,7 @@ impl ProvisionerTaskRunner { let task = self.task(&record.kind); if let Some(task) = &task - && let Some(failure) = task.after_restart() + && let Some(failure) = task.restart_failure() { match self .inner @@ -569,9 +572,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; no error may leave - // this branch before the failure is recorded. + // A success that cannot be recorded fails the Task instead of being retried, since a retry would do + // the work again, such as paying for a second analysis. Err(error) => { error!( task.id = %task_id, diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index d62a88a37e..eceae76e9f 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -1,7 +1,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use parking_lot::Mutex; -use provisioner_task::ProvisionerTaskState; +use provisioner_task::{ProvisionerTaskState, ProvisionerTaskStore}; use super::*; use crate::job_queue::{JobMessage, JobQueueReceiver}; @@ -56,7 +56,7 @@ impl ProvisionerTask for Echo { json!({ "error": error, "last": last }) } - fn after_restart(&self) -> Option { + fn restart_failure(&self) -> Option { self.restart_failure.clone() } @@ -97,12 +97,76 @@ impl ProvisionerTask for Echo { } } +/// The task store, failing to record a success when asked, as a database that is briefly unavailable does. +struct Store { + inner: provisioner_task_libsql::LibSqlProvisionerTaskStore, + /// Successes that fail to be recorded before one is. + success_failures: Arc, +} + +#[async_trait] +impl ProvisionerTaskStore for Store { + async fn create(&self, task: NewProvisionerTaskRecord, now: OffsetDateTime) -> anyhow::Result { + self.inner.create(task, now).await + } + + async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result> { + self.inner.get(id, now).await + } + + async fn start_attempt( + &self, + id: Uuid, + job_token: Uuid, + payload: serde_json::Value, + deadline_at: OffsetDateTime, + now: OffsetDateTime, + ) -> anyhow::Result { + self.inner.start_attempt(id, job_token, payload, deadline_at, now).await + } + + async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result { + self.inner.replace_job_token(id, job_token, now).await + } + + async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result> { + self.inner.list_unfinished(now).await + } + + async fn update_running( + &self, + id: Uuid, + payload: serde_json::Value, + deadline_at: OffsetDateTime, + now: OffsetDateTime, + ) -> anyhow::Result { + self.inner.update_running(id, payload, deadline_at, now).await + } + + async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + if self + .success_failures + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |left| left.checked_sub(1)) + .is_ok() + { + anyhow::bail!("the success could not be recorded"); + } + + self.inner.succeed(id, payload, now).await + } + + async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + self.inner.fail(id, payload, now).await + } +} + struct Harness { runner: ProvisionerTaskRunner, job_queue_rx: JobQueueReceiver, runs: Arc, ended: Arc>>, end_failures: Arc, + success_failures: Arc, } impl Harness { @@ -117,11 +181,6 @@ impl Harness { Self::build(outcome, progress_timeout, None).await } - /// A harness whose task cannot outlive a restart, like one that keeps its secrets in memory only. - async fn forgetting_at_restart(outcome: fn(u32) -> Result) -> Self { - Self::build(outcome, time::Duration::hours(1), Some(Permanent::KeyLost)).await - } - async fn build( outcome: fn(u32) -> Result, progress_timeout: time::Duration, @@ -133,11 +192,13 @@ impl Harness { .connect() .await .expect("connect"); - let store = Arc::new( - provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) + let success_failures = Arc::new(AtomicUsize::new(0)); + let store = Arc::new(Store { + inner: provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) .await .expect("open task store"), - ); + success_failures: Arc::clone(&success_failures), + }); let (job_queue, job_queue_rx) = JobQueueHandle::new(); let runs = Arc::new(AtomicUsize::new(0)); let ended = Arc::new(Mutex::new(Vec::new())); @@ -160,6 +221,7 @@ impl Harness { runs, ended, end_failures, + success_failures, } } @@ -179,6 +241,11 @@ impl Harness { job } + /// Takes every queued job. + fn queued_jobs(&mut self) -> Vec { + std::iter::from_fn(|| self.job_queue_rx.try_recv().ok()).collect() + } + async fn task(&self, id: Uuid) -> ProvisionerTaskRecord { self.runner .get(id, OffsetDateTime::now_utc()) @@ -299,7 +366,7 @@ async fn a_task_that_cannot_be_queued_again_fails() { assert_eq!(*harness.ended.lock(), [id]); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management +// oracle: mutant resume skips a Task whose job may still be stored, so a Task created just before a crash never runs #[tokio::test] async fn resume_gives_every_unfinished_task_a_new_job() { let mut harness = Harness::new(|value| Ok(json!(value))).await; @@ -307,23 +374,24 @@ async fn resume_gives_every_unfinished_task_a_new_job() { let second = harness.create(2).await; let tokens = [harness.task(first).await.job_token, harness.task(second).await.job_token]; // The jobs of the creation may be stored or lost; the runner never asks which. - while harness.job_queue_rx.try_recv().is_ok() {} + harness.queued_jobs(); harness.runner.resume().await.expect("resume"); - assert_eq!( - std::iter::from_fn(|| harness.job_queue_rx.try_recv().ok()).count(), - 2, - "one new job for each unfinished Task" - ); + assert_eq!(harness.queued_jobs().len(), 2, "one new job for each unfinished Task"); assert_ne!(harness.task(first).await.job_token, tokens[0]); assert_ne!(harness.task(second).await.job_token, tokens[1]); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets +// oracle: mutant a Task that cannot outlive a restart stays unfinished until a job runs it #[tokio::test] async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { - let mut harness = Harness::forgetting_at_restart(|value| Ok(json!(value))).await; + let mut harness = Harness::build( + |value| Ok(json!(value)), + time::Duration::hours(1), + Some(Permanent::KeyLost), + ) + .await; let id = harness.create(1).await; let mut stale = harness.queued_job(); @@ -382,7 +450,6 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { let mut job = harness.queued_job(); for attempt in 1..MAX_ATTEMPTS { - let failed_at = OffsetDateTime::now_utc(); job.run().await.expect("the job queues the next attempt itself"); assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Running, "{attempt}"); @@ -390,8 +457,10 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { job: next, schedule_for, } = harness.job_queue_rx.try_recv().expect("the next attempt is queued"); - let retry_at = schedule_for.expect("the next attempt waits"); - assert!(retry_at >= failed_at + retry_delay(attempt), "{attempt}: {retry_at}"); + assert!( + schedule_for.is_some_and(|at| OffsetDateTime::now_utc() < at), + "attempt {attempt} is retried later" + ); job = next; } job.run().await.expect("the last attempt"); @@ -406,6 +475,26 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { ); } +// oracle: mutant a success that cannot be recorded is retried, so the work is done and paid for twice +#[tokio::test] +async fn a_success_that_cannot_be_recorded_fails_the_task_and_ends_it() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + harness.success_failures.store(1, Ordering::SeqCst); + + harness.queued_job().run().await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!( + task.payload, + json!({ "error": { "kind": "internal" }, "last": { "after": null } }) + ); + assert_eq!(harness.runs.load(Ordering::SeqCst), 1, "the work is not done again"); + assert_eq!(*harness.ended.lock(), [id], "what the attempt made is dropped"); + assert!(harness.queued_jobs().is_empty(), "no attempt follows"); +} + // oracle: mutant a queued Task of an unknown kind is never failed #[tokio::test] async fn invalid_parameters_and_unknown_kinds_fail_the_task() { From 25d5de363d9be7551982970822d1e6cf8f947109 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:04:30 -0400 Subject: [PATCH 66/80] test(dgw): check the AI analysis across a restart A restart after the log was written but before the success was recorded fails the Task with `keyLost` and removes the unlisted log. A restart after the success was recorded but before the manifest listed the log keeps the job of the Task, which lists the log. The remux tests name the mark through `remux_marker`, so they follow its file name. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/tests.rs | 46 ++++++++++++------- devolutions-gateway/src/recording.rs | 2 +- 2 files changed, 30 insertions(+), 18 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 2c4652c184..874d7ac002 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -11,9 +11,9 @@ use parking_lot::Mutex; use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; -use crate::job_queue::{JobMessage, JobQueueHandle, JobQueueReceiver}; +use crate::job_queue::{JobMessage, JobQueueReceiver}; use crate::provisioner_tasks::runner::{MAX_ATTEMPTS, ProvisionerTaskJob, ProvisionerTaskRunner}; -use crate::recording::RecordingManagerTask; +use crate::recording::{RecordingManagerTask, remux_marker}; use crate::{DgwState, MockHandles}; const API_KEY: &str = "sk-recording-ai-analysis-test-secret"; @@ -143,9 +143,9 @@ impl Harness { self.job(task_id).await.run().await } - /// Whether a job waits to run a later attempt, as after a recoverable failure; it takes every queued job. - fn later_attempt_queued(&mut self) -> bool { - std::iter::from_fn(|| self.job_queue_rx.try_recv().ok()).any(|message| message.schedule_for.is_some()) + /// Takes every queued job. + fn queued_jobs(&mut self) -> Vec { + std::iter::from_fn(|| self.job_queue_rx.try_recv().ok()).collect() } /// Records a Task directly, as if it was created at `created_at`. @@ -530,7 +530,10 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { let workspace = harness.workspace(session_id); harness.run(task_id).await.expect("the job queues the next attempt itself"); - assert!(harness.later_attempt_queued(), "a provider outage is retried"); + assert!( + harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), + "a provider outage is retried" + ); let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Running); @@ -676,12 +679,15 @@ async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { for attempt in 1..MAX_ATTEMPTS { harness.run(task_id).await.expect("the job queues the next attempt itself"); - assert!(harness.later_attempt_queued(), "attempt {attempt} is retried"); + assert!( + harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), + "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"); - assert!(!harness.later_attempt_queued(), "no attempt after the last"); + assert!(harness.queued_jobs().is_empty(), "no attempt after the last"); let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); @@ -823,14 +829,11 @@ async fn analysis_waits_for_the_remux_of_the_session() { let (settings, requests) = spawn_ai_provider(|_| 0).await; let task_id = harness.start_created(session_id, settings).await; - let marker = harness - .recording_path - .join(session_id.to_string()) - .join("recording-1.webm.remuxing"); + let marker = remux_marker(&harness.recording_path.join(session_id.to_string()).join("recording-1.webm")); std::fs::write(&marker, b"").expect("remux mark"); let other_dir = harness.recording_path.join(Uuid::new_v4().to_string()); std::fs::create_dir_all(&other_dir).expect("other session"); - std::fs::write(other_dir.join("recording-0.webm.remuxing"), b"").expect("remux mark of another session"); + std::fs::write(remux_marker(&other_dir.join("recording-0.webm")), b"").expect("remux mark of another session"); let mut job = harness.job(task_id).await; let job = tokio::spawn(async move { job.run().await }); @@ -856,7 +859,7 @@ async fn waiting_for_a_remux_is_no_progress() { let task_id = harness.start_created(session_id, settings).await; let dir = harness.recording_path.join(session_id.to_string()); - std::fs::write(dir.join("recording-1.webm.remuxing"), b"").expect("remux mark"); + std::fs::write(remux_marker(&dir.join("recording-1.webm")), b"").expect("remux mark"); harness.with_progress_timeout(time::Duration::seconds(2)); harness.run(task_id).await.expect("no retry"); @@ -954,7 +957,7 @@ async fn a_new_task_that_fails_keeps_the_earlier_analysis() { assert!(dir.join("ai-analysis-0.slog").exists()); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets +// oracle: mutant an AI analysis that cannot run after a restart stays unfinished until its next attempt #[tokio::test] async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { let harness = Harness::new().await; @@ -962,6 +965,10 @@ async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { let (settings, requests) = spawn_ai_provider(|_| 0).await; let task_id = harness.start_created(session_id, settings).await; std::fs::create_dir_all(harness.workspace(session_id)).expect("working files of an attempt"); + // As if Gateway stopped after the attempt wrote its log, before its success was recorded. + 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.resume().await.expect("resume"); @@ -972,6 +979,8 @@ async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { harness.run(task_id).await.expect("the new job only ends the Task"); assert!(requests.lock().is_empty(), "the AI is never asked"); assert!(!harness.workspace(session_id).exists(), "the working files are removed"); + assert!(!log.exists(), "the unlisted log is removed"); + assert_eq!(ai_analyses(&dir), json!([])); assert!(!harness.has_key(task_id)); } @@ -1036,8 +1045,11 @@ async fn a_success_cut_short_before_the_log_is_listed_is_finished_by_the_job() { 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"); + // The job of the Task stays stored until it returns, so it runs again when Gateway starts, and it still holds the + // Task since a finished Task keeps its job. + let mut job = harness.job(task_id).await; + harness.state.provisioner_tasks.resume().await.expect("resume"); + job.run().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()); diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 59ec4f1727..ba8d49c3dd 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -1292,7 +1292,7 @@ const REMUX_MARKER_SUFFIX: &str = ".remuxing"; const REMUX_MARKER_STALE_AFTER: Duration = Duration::from_secs(10 * 60); /// The file that marks `recording_file` as waiting for its remux, from the time the remux is scheduled until it ends. -fn remux_marker(recording_file: &Utf8Path) -> Utf8PathBuf { +pub(crate) fn remux_marker(recording_file: &Utf8Path) -> Utf8PathBuf { Utf8PathBuf::from(format!("{recording_file}{REMUX_MARKER_SUFFIX}")) } From 1b49e4a984abfde8fa2de4d63083c73d79c830c6 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:17:09 -0400 Subject: [PATCH 67/80] fix(dgw): age a remux mark from the start of its remux A mark written when the remux is scheduled also ages while the job waits to run, so a slow queue could make it stale before the remux replaced the file. The remux job marks the file again as it starts, and a mark stops counting after 20 minutes, which a compile-time check keeps under the 30 minutes an AI analysis may go without progress. A mark that disappears while the session folder is read means its remux ended, not that the storage failed. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/mod.rs | 7 +++- devolutions-gateway/src/recording.rs | 37 +++++++++++++------ 2 files changed, 31 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 9912b37744..8e2521411f 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -36,7 +36,9 @@ use super::{AttemptError, Permanent, Recoverable}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; -use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender, remux_pending}; +use crate::recording::{ + FinishedRecordingError, JrecManifest, REMUX_MARKER_STALE_AFTER, RecordingMessageSender, remux_pending, +}; use crate::secret_store::{SecretId, SecretStore}; pub const KIND: &str = "recording-ai-analysis"; @@ -47,6 +49,9 @@ pub const START_WINDOW: time::Duration = time::Duration::hours(2); /// Time a running AI analysis may go without progress before it fails. pub const PROGRESS_TIMEOUT: time::Duration = time::Duration::minutes(30); +// Waiting for a remux is no progress, so a mark that no remux will remove must stop counting before the wait times out. +const _: () = assert!(REMUX_MARKER_STALE_AFTER.as_secs() < PROGRESS_TIMEOUT.whole_seconds().unsigned_abs()); + /// Time between two checks for a remux of the session. const REMUX_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1); diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index ba8d49c3dd..3e00eb609a 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -1287,9 +1287,10 @@ const REMUX_MARKER_SUFFIX: &str = ".remuxing"; /// Age past which a remux mark no longer counts. /// /// A crash between the mark and the storing of its job leaves a mark that no remux removes. -/// The limit lets an AI analysis waiting on such a mark go on well before it fails for lack of progress, and no remux -/// takes this long. -const REMUX_MARKER_STALE_AFTER: Duration = Duration::from_secs(10 * 60); +/// The limit lets an AI analysis waiting on such a mark go on before it fails for lack of progress. +/// The remux job marks the file again as it starts, so only the remux itself, which should take far less, counts +/// against it. +pub(crate) const REMUX_MARKER_STALE_AFTER: Duration = Duration::from_secs(20 * 60); /// The file that marks `recording_file` as waiting for its remux, from the time the remux is scheduled until it ends. pub(crate) fn remux_marker(recording_file: &Utf8Path) -> Utf8PathBuf { @@ -1318,7 +1319,13 @@ pub(crate) async fn remux_pending(dir: &Utf8Path) -> io::Result { continue; } - let age = entry.metadata().await?.modified()?.elapsed().unwrap_or_default(); + let metadata = match entry.metadata().await { + Ok(metadata) => metadata, + // The remux ended between the listing and this look. + Err(error) if error.kind() == io::ErrorKind::NotFound => continue, + Err(error) => return Err(error), + }; + let age = metadata.modified()?.elapsed().unwrap_or_default(); if age < REMUX_MARKER_STALE_AFTER { return Ok(true); @@ -1341,6 +1348,12 @@ impl job_queue::Job for RemuxJob { async fn run(&mut self) -> anyhow::Result<()> { let input_path = core::mem::take(&mut self.input_path); let marker = remux_marker(&input_path); + + // The mark ages from now, so neither the wait for this job nor a run again after a restart makes it stale. + if let Err(error) = fs::write(&marker, b"").await { + warn!(%marker, %error, "Failed to mark the recording as waiting for its remux"); + } + remux(input_path).await; remove_remux_marker(&marker).await; Ok(()) @@ -1485,7 +1498,7 @@ mod tests { } } - // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + // oracle: mutant a fresh remux mark is ignored, so an AI analysis reads a video its remux is about to replace #[tokio::test] async fn a_recording_waits_for_its_remux_while_its_mark_is_fresh() { let dir = tempfile::tempdir().expect("temp dir"); @@ -1513,7 +1526,7 @@ mod tests { assert!(!remux_pending(dir).await.expect("stale mark")); } - // oracle: mutant a session folder without a mark, or without a folder yet, fails the AI analysis as storage + // oracle: mutant a recording file is taken for a remux mark, so every AI analysis of a video session waits #[tokio::test] async fn a_session_without_remux_marks_has_no_pending_remux() { let dir = tempfile::tempdir().expect("temp dir"); @@ -1521,15 +1534,15 @@ mod tests { std::fs::write(dir.join("recording-0.webm"), b"").expect("recording"); assert!(!remux_pending(dir).await.expect("no mark")); - assert!(!remux_pending(&dir.join("gone")).await.expect("missing folder")); } - // oracle: mutant a remux mark left by a crash outlasts the progress deadline, so AI analyses of the session fail - #[test] - fn a_forgotten_remux_mark_stops_counting_before_an_ai_analysis_times_out() { - let progress_timeout = crate::provisioner_tasks::recording_ai_analysis::PROGRESS_TIMEOUT.unsigned_abs(); + // oracle: mutant a session folder that does not exist fails the AI analysis as storage + #[tokio::test] + async fn a_missing_session_folder_has_no_pending_remux() { + let dir = tempfile::tempdir().expect("temp dir"); + let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); - assert!(REMUX_MARKER_STALE_AFTER < progress_timeout); + assert!(!remux_pending(&dir.join("gone")).await.expect("missing folder")); } #[tokio::test] From 48f15f4d060c57fe49cb930a85b4f971032d8d73 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:17:09 -0400 Subject: [PATCH 68/80] test(dgw): check provisioner tasks when their end cannot be recorded A Task whose success and failure both fail to be recorded is left to a job that runs past its deadline, when the store fails it, so the paid work is never redone. A Task whose failure at startup cannot be recorded still gets a new job, and a Task keeps the attempts it made across a restart. A Task that was waiting to retry runs at once after a restart, since retry timing is no promise to the caller. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../src/provisioner_tasks/runner/mod.rs | 10 ++- .../src/provisioner_tasks/runner/tests.rs | 89 +++++++++++++++++-- 2 files changed, 89 insertions(+), 10 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index f2950e8b25..37a49cfdfa 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -352,6 +352,7 @@ impl ProvisionerTaskRunner { /// Each record gets a new `job_token`, so a job left from before finds itself superseded and stops, and the job /// queue is never read. /// A record whose task cannot outlive a restart fails first, and its new job only ends it. + /// A record that was waiting to retry runs at once, since retry timing is no promise to the caller. /// 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<()> { @@ -601,7 +602,14 @@ impl ProvisionerTaskRunner { }, Err(AttemptError::Recoverable(error)) if attempts < MAX_ATTEMPTS => { let retry_at = OffsetDateTime::now_utc() + retry_delay(attempts); - warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, %retry_at, "Task attempt failed; it will be retried"); + warn!( + task.id = %task_id, + task.kind = %record.kind, + attempts, + ?error, + %retry_at, + "Task attempt failed; it will be retried" + ); self.enqueue(task_id, job_token, Some(retry_at)) .await .context("queue the next attempt")?; diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index eceae76e9f..5930ac3e39 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -97,15 +97,24 @@ impl ProvisionerTask for Echo { } } -/// The task store, failing to record a success when asked, as a database that is briefly unavailable does. -struct Store { +/// The task store, failing to record an end when asked, as a database that is briefly unavailable does. +struct FailingStore { inner: provisioner_task_libsql::LibSqlProvisionerTaskStore, /// Successes that fail to be recorded before one is. success_failures: Arc, + /// Failures that fail to be recorded before one is. + failure_failures: Arc, +} + +/// Whether one more of the `failures` left happens, counting it. +fn fails(failures: &AtomicUsize) -> bool { + failures + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |left| left.checked_sub(1)) + .is_ok() } #[async_trait] -impl ProvisionerTaskStore for Store { +impl ProvisionerTaskStore for FailingStore { async fn create(&self, task: NewProvisionerTaskRecord, now: OffsetDateTime) -> anyhow::Result { self.inner.create(task, now).await } @@ -144,11 +153,7 @@ impl ProvisionerTaskStore for Store { } async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { - if self - .success_failures - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |left| left.checked_sub(1)) - .is_ok() - { + if fails(&self.success_failures) { anyhow::bail!("the success could not be recorded"); } @@ -156,6 +161,10 @@ impl ProvisionerTaskStore for Store { } async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + if fails(&self.failure_failures) { + anyhow::bail!("the failure could not be recorded"); + } + self.inner.fail(id, payload, now).await } } @@ -167,6 +176,7 @@ struct Harness { ended: Arc>>, end_failures: Arc, success_failures: Arc, + failure_failures: Arc, } impl Harness { @@ -193,11 +203,13 @@ impl Harness { .await .expect("connect"); let success_failures = Arc::new(AtomicUsize::new(0)); - let store = Arc::new(Store { + let failure_failures = Arc::new(AtomicUsize::new(0)); + let store = Arc::new(FailingStore { inner: provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) .await .expect("open task store"), success_failures: Arc::clone(&success_failures), + failure_failures: Arc::clone(&failure_failures), }); let (job_queue, job_queue_rx) = JobQueueHandle::new(); let runs = Arc::new(AtomicUsize::new(0)); @@ -222,6 +234,7 @@ impl Harness { ended, end_failures, success_failures, + failure_failures, } } @@ -410,6 +423,42 @@ async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { assert_eq!(*harness.ended.lock(), [id]); } +// oracle: mutant a Task whose failure at startup cannot be recorded gets no new job, so nothing ends it +#[tokio::test] +async fn a_task_whose_restart_failure_cannot_be_recorded_still_gets_a_new_job() { + let mut harness = Harness::build( + |value| Ok(json!(value)), + time::Duration::hours(1), + Some(Permanent::KeyLost), + ) + .await; + let id = harness.create(1).await; + let token = harness.task(id).await.job_token; + harness.queued_jobs(); + harness.failure_failures.store(1, Ordering::SeqCst); + + harness.runner.resume().await.expect("resume"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Queued, "the failure was not recorded"); + assert_ne!(task.job_token, token); + assert_eq!(harness.queued_jobs().len(), 1, "its new job runs it"); +} + +// oracle: mutant resume resets the attempts of a Task, so a Task that fails at each restart is retried forever +#[tokio::test] +async fn resume_keeps_the_attempts_a_task_made() { + let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(Recoverable::AiBusy { status: 429 }))).await; + let id = harness.create(1).await; + harness.queued_job().run().await.expect("the job queues the next attempt itself"); + harness.queued_jobs(); + + harness.runner.resume().await.expect("resume"); + harness.queued_job().run().await.expect("the job queues the next attempt itself"); + + assert_eq!(harness.task(id).await.attempts, 2); +} + // oracle: mutant a panicking attempt fails the Task as internal instead of panicked #[tokio::test] async fn a_panicking_attempt_fails_the_task() { @@ -457,6 +506,7 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { job: next, schedule_for, } = harness.job_queue_rx.try_recv().expect("the next attempt is queued"); + // oracle: mutant a recoverable failure is retried at once, without the delay a busy provider asks for assert!( schedule_for.is_some_and(|at| OffsetDateTime::now_utc() < at), "attempt {attempt} is retried later" @@ -495,6 +545,27 @@ async fn a_success_that_cannot_be_recorded_fails_the_task_and_ends_it() { assert!(harness.queued_jobs().is_empty(), "no attempt follows"); } +// oracle: mutant an end that can be recorded neither as a success nor as a failure is retried, so the work is redone +#[tokio::test] +async fn a_task_whose_end_cannot_be_recorded_is_ended_after_its_deadline() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + harness.success_failures.store(1, Ordering::SeqCst); + harness.failure_failures.store(1, Ordering::SeqCst); + + harness.queued_job().run().await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Running, "nothing could be recorded"); + let JobMessage { schedule_for, .. } = harness.job_queue_rx.try_recv().expect("a job ends it later"); + assert!( + schedule_for.is_some_and(|at| task.deadline_at < at), + "the job runs once the store fails the Task" + ); + assert_eq!(harness.runs.load(Ordering::SeqCst), 1, "the work is not done again"); + assert!(harness.ended.lock().is_empty()); +} + // oracle: mutant a queued Task of an unknown kind is never failed #[tokio::test] async fn invalid_parameters_and_unknown_kinds_fail_the_task() { From e40467d88d5cc86dc556e3c0ce675523a9dc1b2d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:17:09 -0400 Subject: [PATCH 69/80] test(dgw): mark what the AI analysis tests check beyond the intent The answer's `sessionId`, the fields it leaves out, retries and the failure at startup are not decided by an intent file, so those assertions carry their own markers. The location test follows the `Location` header instead of comparing it with a path. A restart after an attempt that spent tokens keeps them in the failed Task. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/tests.rs | 21 +++++++++++++------ testsuite/tests/cli/dgw/ai_analysis.rs | 12 +++++++---- 2 files changed, 23 insertions(+), 10 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 874d7ac002..783add0d31 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -530,6 +530,7 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { let workspace = harness.workspace(session_id); harness.run(task_id).await.expect("the job queues the next attempt itself"); + // oracle: mutant a recoverable failure is not retried assert!( harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), "a provider outage is retried" @@ -679,6 +680,7 @@ async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { for attempt in 1..MAX_ATTEMPTS { harness.run(task_id).await.expect("the job queues the next attempt itself"); + // oracle: mutant a recoverable failure is not retried assert!( harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), "attempt {attempt} is retried" @@ -957,15 +959,16 @@ async fn a_new_task_that_fails_keeps_the_earlier_analysis() { assert!(dir.join("ai-analysis-0.slog").exists()); } -// oracle: mutant an AI analysis that cannot run after a restart stays unfinished until its next attempt +// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[tokio::test] async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { let harness = Harness::new().await; - let session_id = harness.write_session(); - let (settings, requests) = spawn_ai_provider(|_| 0).await; + // About 1 MB of transcript: three chunks, so the first attempt stops with the tokens of one. + 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; - std::fs::create_dir_all(harness.workspace(session_id)).expect("working files of an attempt"); - // As if Gateway stopped after the attempt wrote its log, before its success was recorded. + harness.run(task_id).await.expect("the job queues the next attempt itself"); + // As if Gateway stopped after an attempt wrote its log, before its success was recorded. 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"); @@ -973,11 +976,17 @@ async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { harness.state.provisioner_tasks.resume().await.expect("resume"); let task = harness.task(task_id).await; + // oracle: mutant an AI analysis that cannot run after a restart stays unfinished until its next attempt assert_eq!(task.state, ProvisionerTaskState::Failed, "no job ran yet"); assert_eq!(task.payload["error"], json!({ "kind": "keyLost" })); + // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis + assert_eq!( + task.payload["usage"], + json!({ "inputTokens": 10, "outputTokens": 20, "unreportedRequests": 0 }) + ); harness.run(task_id).await.expect("the new job only ends the Task"); - assert!(requests.lock().is_empty(), "the AI is never asked"); + assert_eq!(requests.lock().len(), 2, "the AI is not asked again"); assert!(!harness.workspace(session_id).exists(), "the working files are removed"); assert!(!log.exists(), "the unlisted log is removed"); assert_eq!(ai_analyses(&dir), json!([])); diff --git a/testsuite/tests/cli/dgw/ai_analysis.rs b/testsuite/tests/cli/dgw/ai_analysis.rs index bfcdd37331..a873d5d5a5 100644 --- a/testsuite/tests/cli/dgw/ai_analysis.rs +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -111,9 +111,13 @@ impl Gateway { } async fn get_analysis(&self, task_id: Uuid, scope: &str) -> (StatusCode, Value) { + self.get(&format!("/jet/tasks/recording-ai-analysis/{task_id}"), scope).await + } + + async fn get(&self, path: &str, scope: &str) -> (StatusCode, Value) { let response = self .client - .get(self.url(&format!("/jet/tasks/recording-ai-analysis/{task_id}"))) + .get(self.url(path)) .bearer_auth(scope_token(scope)) .send() .await @@ -289,6 +293,7 @@ async fn start_returns_202_then_200_for_the_same_request() { let (status, created) = gateway.put(task_id, &request(session_id, "gpt-test", &base_url)).await; assert_eq!(status, StatusCode::ACCEPTED, "{created}"); assert_eq!(created["id"], task_id.to_string()); + // oracle: mutant the answer leaves out the session the analysis is about assert_eq!(created["sessionId"], session_id.to_string()); assert_eq!( created["params"], @@ -301,6 +306,7 @@ async fn start_returns_202_then_200_for_the_same_request() { assert_eq!(created["startedAt"], Value::Null); assert_eq!(created["finishedAt"], Value::Null); for internal in ["kind", "target", "attempts"] { + // oracle: mutant the answer shows `kind`, `target` and `attempts`, which only matter inside Gateway assert!(created.get(internal).is_none(), "{internal} is not returned"); } @@ -465,12 +471,10 @@ async fn an_analysis_is_found_at_its_location() { assert_eq!(response.status(), StatusCode::ACCEPTED); let location = response.headers()["location"].to_str().unwrap().to_owned(); - assert_eq!(location, format!("/jet/tasks/recording-ai-analysis/{task_id}")); - let (status, task) = gateway.get_analysis(task_id, READ_SCOPE).await; + let (status, task) = gateway.get(&location, READ_SCOPE).await; assert_eq!(status, StatusCode::OK, "{task}"); assert_eq!(task["id"], task_id.to_string()); - assert_eq!(task["sessionId"], session_id.to_string()); } // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis From b3295d9884f146461bba8ff43109f41066fa2ee3 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:20:05 -0400 Subject: [PATCH 70/80] fix(dgw): say what a remux mark covers before its job starts The remux job marks the file again as it starts, so the 20 minutes cover the remux itself, but until then the mark still ages from the scheduling; the doc now says so instead of promising more. Both places that write the mark share one helper, as removing it already does. A Task whose end can be neither recorded nor queued must not make the job fail, or the job queue would redo the paid work; a test now checks it. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/mod.rs | 3 ++- .../recording_ai_analysis/tests.rs | 4 ++-- .../src/provisioner_tasks/runner/tests.rs | 20 ++++++++++++++++++ devolutions-gateway/src/recording.rs | 21 +++++++++++-------- 4 files changed, 36 insertions(+), 12 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 8e2521411f..030278ebe9 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -49,7 +49,8 @@ pub const START_WINDOW: time::Duration = time::Duration::hours(2); /// Time a running AI analysis may go without progress before it fails. pub const PROGRESS_TIMEOUT: time::Duration = time::Duration::minutes(30); -// Waiting for a remux is no progress, so a mark that no remux will remove must stop counting before the wait times out. +// INVARIANT: a mark that no remux will remove stops counting before an analysis waiting on it times out, since waiting +// for a remux is no progress. const _: () = assert!(REMUX_MARKER_STALE_AFTER.as_secs() < PROGRESS_TIMEOUT.whole_seconds().unsigned_abs()); /// Time between two checks for a remux of the session. 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 783add0d31..a949d20cb4 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -530,7 +530,7 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { let workspace = harness.workspace(session_id); harness.run(task_id).await.expect("the job queues the next attempt itself"); - // oracle: mutant a recoverable failure is not retried + // oracle: mutant a recoverable failure is retried at once, without the delay a busy provider asks for assert!( harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), "a provider outage is retried" @@ -680,7 +680,7 @@ async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { for attempt in 1..MAX_ATTEMPTS { harness.run(task_id).await.expect("the job queues the next attempt itself"); - // oracle: mutant a recoverable failure is not retried + // oracle: mutant a recoverable failure is retried at once, without the delay a busy provider asks for assert!( harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), "attempt {attempt} is retried" diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 5930ac3e39..2bc5e38bb6 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -566,6 +566,26 @@ async fn a_task_whose_end_cannot_be_recorded_is_ended_after_its_deadline() { assert!(harness.ended.lock().is_empty()); } +// oracle: mutant an end that can be neither recorded nor queued returns an error, so the job queue redoes the paid work +#[tokio::test] +async fn a_task_whose_end_cannot_be_recorded_or_queued_is_not_retried() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + let mut job = harness.queued_job(); + harness.job_queue_rx.close(); + harness.success_failures.store(1, Ordering::SeqCst); + harness.failure_failures.store(1, Ordering::SeqCst); + + job.run().await.expect("the job queue does not run the attempt again"); + + assert_eq!(harness.runs.load(Ordering::SeqCst), 1); + assert_eq!( + harness.task(id).await.state, + ProvisionerTaskState::Running, + "the store fails it at its deadline" + ); +} + // oracle: mutant a queued Task of an unknown kind is never failed #[tokio::test] async fn invalid_parameters_and_unknown_kinds_fail_the_task() { diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 3e00eb609a..c1a7846f5d 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -946,9 +946,7 @@ impl RecordingManagerTask { let marker = remux_marker(&recording_file_path); // An AI analysis that finds no mark reads the file even while it is remuxed. - if let Err(error) = fs::write(&marker, b"").await { - warn!(%marker, %error, "Failed to mark the recording as waiting for its remux"); - } + write_remux_marker(&marker).await; // Schedule 60 seconds to wait for the streamers to release the file. let scheduled = self @@ -1288,8 +1286,9 @@ const REMUX_MARKER_SUFFIX: &str = ".remuxing"; /// /// A crash between the mark and the storing of its job leaves a mark that no remux removes. /// The limit lets an AI analysis waiting on such a mark go on before it fails for lack of progress. -/// The remux job marks the file again as it starts, so only the remux itself, which should take far less, counts -/// against it. +/// The remux job marks the file again as it starts, so the limit covers the remux itself. +/// Until then the mark ages from the scheduling, so a job that waits longer than the limit to start lets an AI analysis +/// read the file the remux is about to replace. pub(crate) const REMUX_MARKER_STALE_AFTER: Duration = Duration::from_secs(20 * 60); /// The file that marks `recording_file` as waiting for its remux, from the time the remux is scheduled until it ends. @@ -1297,6 +1296,12 @@ pub(crate) fn remux_marker(recording_file: &Utf8Path) -> Utf8PathBuf { Utf8PathBuf::from(format!("{recording_file}{REMUX_MARKER_SUFFIX}")) } +async fn write_remux_marker(marker: &Utf8Path) { + if let Err(error) = fs::write(marker, b"").await { + warn!(%marker, %error, "Failed to mark the recording as waiting for its remux"); + } +} + async fn remove_remux_marker(marker: &Utf8Path) { match fs::remove_file(marker).await { Err(error) if error.kind() != io::ErrorKind::NotFound => { @@ -1349,10 +1354,8 @@ impl job_queue::Job for RemuxJob { let input_path = core::mem::take(&mut self.input_path); let marker = remux_marker(&input_path); - // The mark ages from now, so neither the wait for this job nor a run again after a restart makes it stale. - if let Err(error) = fs::write(&marker, b"").await { - warn!(%marker, %error, "Failed to mark the recording as waiting for its remux"); - } + // From here the mark ages with the remux itself, also when the job runs again after a restart. + write_remux_marker(&marker).await; remux(input_path).await; remove_remux_marker(&marker).await; From 549b2cd87c1dfd73cb7c80bc1bf41005c86dabc0 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:49:25 -0400 Subject: [PATCH 71/80] refactor(dgw): keep AI analysis failures with their kind, run an attempt in three steps `Permanent` keeps only the failures any Task may have, such as a passed deadline, and holds those of the Task's own kind in an untagged variant, so a new kind of Task brings its failures without adding to the shared ones. The AI analysis defines its own failures (`aiRejected`, `keyLost`, `noRecording` and the others) next to its work, and the JSON of a failed Task stays the same. An attempt now runs as a prelude that starts it on the record, a body that runs the work, and a completion that records how it ended, so each step can be read on its own. `AiAnalysisPayload` only describes the payload in the OpenAPI document, so its never-built variants are expected. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../src/api/{tasks.rs => tasks/mod.rs} | 0 devolutions-gateway/src/openapi.rs | 18 +- .../src/provisioner_tasks/ai.rs | 93 +--------- .../src/provisioner_tasks/mod.rs | 61 ++---- .../recording_ai_analysis/failure.rs | 149 +++++++++++++++ .../recording_ai_analysis/mod.rs | 54 +++--- .../src/provisioner_tasks/runner/mod.rs | 175 +++++++++++++----- .../src/provisioner_tasks/runner/tests.rs | 77 +++++--- 8 files changed, 393 insertions(+), 234 deletions(-) rename devolutions-gateway/src/api/{tasks.rs => tasks/mod.rs} (100%) create mode 100644 devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs diff --git a/devolutions-gateway/src/api/tasks.rs b/devolutions-gateway/src/api/tasks/mod.rs similarity index 100% rename from devolutions-gateway/src/api/tasks.rs rename to devolutions-gateway/src/api/tasks/mod.rs diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index fa79b71c2f..4b157fcb35 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -5,8 +5,9 @@ use uuid::Uuid; use crate::api::preflight::PreflightAlertStatus; use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; +use crate::provisioner_tasks::Permanent; use crate::provisioner_tasks::ai::AiProvider; -use crate::provisioner_tasks::{BadAnswer, Permanent, Recoverable}; +use crate::provisioner_tasks::recording_ai_analysis::{BadAnswer, Failure, Recoverable}; #[derive(OpenApi)] #[openapi( @@ -181,6 +182,7 @@ struct AiAnalysisParams { /// Details of the state of an AI analysis Task: `running`, `succeeded` and `failed` each have their own #[derive(Serialize, utoipa::ToSchema)] #[serde(untagged)] +#[expect(dead_code, reason = "it only describes the payload, which the Task keeps as JSON")] enum AiAnalysisPayload { Running(AiAnalysisRunningPayload), Succeeded(AiAnalysisSucceededPayload), @@ -318,13 +320,13 @@ 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 or field is written out too. -const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| match error { - Permanent::AiRejected { status: _ } => AiAnalysisErrorKind::AiRejected, - Permanent::AiBadAnswer { why: _ } => AiAnalysisErrorKind::AiBadAnswer, - Permanent::KeyLost => AiAnalysisErrorKind::KeyLost, - Permanent::NoRecording => AiAnalysisErrorKind::NoRecording, - Permanent::MixedRecording => AiAnalysisErrorKind::MixedRecording, - Permanent::UnreadableRecording { file: _ } => AiAnalysisErrorKind::UnreadableRecording, +const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| match error { + Permanent::Task(Failure::AiRejected { status: _ }) => AiAnalysisErrorKind::AiRejected, + Permanent::Task(Failure::AiBadAnswer { why: _ }) => AiAnalysisErrorKind::AiBadAnswer, + Permanent::Task(Failure::KeyLost) => AiAnalysisErrorKind::KeyLost, + Permanent::Task(Failure::NoRecording) => AiAnalysisErrorKind::NoRecording, + Permanent::Task(Failure::MixedRecording) => AiAnalysisErrorKind::MixedRecording, + Permanent::Task(Failure::UnreadableRecording { file: _ }) => AiAnalysisErrorKind::UnreadableRecording, Permanent::InvalidTask => AiAnalysisErrorKind::InvalidTask, Permanent::Panicked => AiAnalysisErrorKind::Panicked, Permanent::Internal => AiAnalysisErrorKind::Internal, diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index 294cf18bb7..ee6a29c6d7 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -6,12 +6,10 @@ use core::fmt; use anyhow::Context as _; -use devolutions_gateway_ai::recording_analysis::RecordingAnalysisError; -use devolutions_gateway_ai::{AiClient, BuildError, Error, Provider}; +use devolutions_gateway_ai::{AiClient, BuildError, Provider}; use secrecy::{ExposeSecret as _, SecretString}; use url::Url; -use super::{AttemptError, BadAnswer, Permanent, Recoverable}; use crate::config::Conf; /// AI provider to ask; `openai-compatible` needs a base URL @@ -190,99 +188,10 @@ impl AiSettings { } } -/// Retries only the failures that may pass later, such as a rate limit or a network error. -impl From for AttemptError { - fn from(error: RecordingAnalysisError) -> Self { - match error { - RecordingAnalysisError::Ai(error) => AttemptError::from(error), - RecordingAnalysisError::NoRecording => AttemptError::Permanent(Permanent::NoRecording), - RecordingAnalysisError::InvalidRecording { file_name, .. } => { - AttemptError::Permanent(Permanent::UnreadableRecording { file: file_name }) - } - RecordingAnalysisError::Truncated => AttemptError::Permanent(Permanent::AiBadAnswer { - why: BadAnswer::Truncated, - }), - // The runner cancels an attempt only once its Task is past its deadline. - RecordingAnalysisError::Cancelled => AttemptError::Permanent(Permanent::DeadlinePassed), - _ => AttemptError::Permanent(Permanent::Internal), - } - } -} - -impl From for AttemptError { - fn from(error: Error) -> Self { - let bad_answer = |why| AttemptError::Permanent(Permanent::AiBadAnswer { why }); - - match &error { - Error::Transport { .. } => AttemptError::Recoverable(Recoverable::AiUnreachable), - Error::Status { status, .. } if error.is_transient() => { - AttemptError::Recoverable(Recoverable::AiBusy { status: *status }) - } - Error::Status { status, .. } => AttemptError::Permanent(Permanent::AiRejected { status: *status }), - Error::InvalidResponse { .. } => bad_answer(BadAnswer::Invalid), - Error::Truncated { .. } => bad_answer(BadAnswer::Truncated), - Error::Refused { .. } => bad_answer(BadAnswer::Refused), - Error::InvalidOutput { .. } => bad_answer(BadAnswer::NotExpectedJson), - _ => AttemptError::Permanent(Permanent::Internal), - } - } -} - #[cfg(test)] mod tests { use super::*; - // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis - #[test] - fn transient_ai_errors_are_recoverable_and_the_others_permanent() { - let attempt = |error: Error| AttemptError::from(RecordingAnalysisError::Ai(error)); - let status = |status, code: Option<&str>| Error::Status { - status, - message: "failed".to_owned(), - code: code.map(str::to_owned), - retry_after: None, - }; - - let network = Error::Transport { - message: "connection refused".to_owned(), - }; - assert_eq!(attempt(network), AttemptError::Recoverable(Recoverable::AiUnreachable)); - - for code in [429, 500, 503] { - assert_eq!( - attempt(status(code, None)), - AttemptError::Recoverable(Recoverable::AiBusy { status: code }) - ); - } - - for code in [400, 401, 403, 404] { - assert_eq!( - attempt(status(code, None)), - AttemptError::Permanent(Permanent::AiRejected { status: code }) - ); - } - - assert_eq!( - attempt(status(429, Some("insufficient_quota"))), - AttemptError::Permanent(Permanent::AiRejected { status: 429 }), - "a spent quota" - ); - assert_eq!( - AttemptError::from(RecordingAnalysisError::Truncated), - AttemptError::Permanent(Permanent::AiBadAnswer { - why: BadAnswer::Truncated - }) - ); - assert_eq!( - attempt(Error::Refused { - reason: "content filter".to_owned() - }), - AttemptError::Permanent(Permanent::AiBadAnswer { - why: BadAnswer::Refused - }) - ); - } - // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis #[test] fn providers_keep_their_wire_names() { diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index 3eb5c940a8..e5f853972f 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -26,46 +26,30 @@ pub fn task_runner( .build() } -/// Why one attempt of a Task failed. +/// Why one attempt of a Task failed, with `F` and `R` the failures of its own kind of work. #[derive(Debug, Clone, PartialEq, Eq)] -pub enum AttemptError { +pub enum AttemptError { /// Worth another attempt later, such as a rate limit or a network error. - Recoverable(Recoverable), + Recoverable(R), /// Another attempt would fail the same way, so the Task fails at once. - Permanent(Permanent), + 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 storage, such as the folder of the session. - Storage, +impl AttemptError { + /// A permanent failure of the Task's own kind of work. + pub fn task(failure: F) -> Self { + Self::Permanent(Permanent::Task(failure)) + } } /// Why a Task failed, as the `error` of its payload: a closed set of kinds without text for people, since the caller /// words them. +/// +/// The kinds below can happen to any Task, while [`Permanent::Task`] holds those of its own kind of work, so a new kind +/// of Task brings its failures without adding to these. #[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 }, +pub enum Permanent { /// The Task itself is not valid, such as with parameters this Gateway cannot read. InvalidTask, /// The attempt panicked. @@ -73,23 +57,12 @@ pub enum Permanent { /// 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 }, + AttemptsExhausted { last: R }, /// 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, + /// A failure of the Task's own kind of work, which names its own `kind`. + #[serde(untagged)] + Task(F), } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs new file mode 100644 index 0000000000..f0997c5f99 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs @@ -0,0 +1,149 @@ +//! Why an AI analysis fails, next to the failures any Task may have. + +use devolutions_gateway_ai::Error; +use devolutions_gateway_ai::recording_analysis::RecordingAnalysisError; + +use crate::provisioner_tasks::{AttemptError, Permanent}; + +/// Why one attempt of an AI analysis failed. +pub type AnalysisError = AttemptError; + +/// Why an AI analysis failed for good, as the `error` of its payload, next to the failures any Task may have. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(tag = "kind", rename_all = "camelCase", rename_all_fields = "camelCase")] +pub enum Failure { + /// 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 }, +} + +/// A failure of an AI analysis attempt 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 storage, such as the folder of the session. + Storage, +} + +/// 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, +} + +/// Retries only the failures that may pass later, such as a rate limit or a network error. +impl From for AnalysisError { + fn from(error: RecordingAnalysisError) -> Self { + match error { + RecordingAnalysisError::Ai(error) => Self::from(error), + RecordingAnalysisError::NoRecording => Self::task(Failure::NoRecording), + RecordingAnalysisError::InvalidRecording { file_name, .. } => { + Self::task(Failure::UnreadableRecording { file: file_name }) + } + RecordingAnalysisError::Truncated => Self::task(Failure::AiBadAnswer { + why: BadAnswer::Truncated, + }), + // The runner cancels an attempt only once its Task is past its deadline. + RecordingAnalysisError::Cancelled => Self::Permanent(Permanent::DeadlinePassed), + _ => Self::Permanent(Permanent::Internal), + } + } +} + +impl From for AnalysisError { + fn from(error: Error) -> Self { + let bad_answer = |why| Self::task(Failure::AiBadAnswer { why }); + + match &error { + Error::Transport { .. } => Self::Recoverable(Recoverable::AiUnreachable), + Error::Status { status, .. } if error.is_transient() => { + Self::Recoverable(Recoverable::AiBusy { status: *status }) + } + Error::Status { status, .. } => Self::task(Failure::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), + _ => Self::Permanent(Permanent::Internal), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + #[test] + fn transient_ai_errors_are_recoverable_and_the_others_permanent() { + let attempt = |error: Error| AnalysisError::from(RecordingAnalysisError::Ai(error)); + let status = |status, code: Option<&str>| Error::Status { + status, + message: "failed".to_owned(), + code: code.map(str::to_owned), + retry_after: None, + }; + + let network = Error::Transport { + message: "connection refused".to_owned(), + }; + assert_eq!(attempt(network), AnalysisError::Recoverable(Recoverable::AiUnreachable)); + + for code in [429, 500, 503] { + assert_eq!( + attempt(status(code, None)), + AnalysisError::Recoverable(Recoverable::AiBusy { status: code }) + ); + } + + for code in [400, 401, 403, 404] { + assert_eq!( + attempt(status(code, None)), + AnalysisError::task(Failure::AiRejected { status: code }) + ); + } + + assert_eq!( + attempt(status(429, Some("insufficient_quota"))), + AnalysisError::task(Failure::AiRejected { status: 429 }), + "a spent quota" + ); + assert_eq!( + AnalysisError::from(RecordingAnalysisError::Truncated), + AnalysisError::task(Failure::AiBadAnswer { + why: BadAnswer::Truncated + }) + ); + assert_eq!( + attempt(Error::Refused { + reason: "content filter".to_owned() + }), + AnalysisError::task(Failure::AiBadAnswer { + why: BadAnswer::Refused + }) + ); + } +} 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 030278ebe9..3c1efcbd5b 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -11,6 +11,7 @@ //! a Task that fails removes it instead. //! The API key and the working files go away when the Task ends. +mod failure; #[cfg(test)] mod tests; @@ -30,9 +31,10 @@ use serde_json::json; use time::OffsetDateTime; use uuid::Uuid; +pub use self::failure::{AnalysisError, BadAnswer, Failure, Recoverable}; +use super::Permanent; use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; -use super::runner::{Attempt, ProvisionerTask}; -use super::{AttemptError, Permanent, Recoverable}; +use super::runner::{AnyPermanent, Attempt, ProvisionerTask}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; @@ -345,6 +347,10 @@ impl ProvisionerTask for RecordingAiAnalysis { type Output = Analysis; + type Failure = Failure; + + type Recoverable = Recoverable; + fn start_window(&self) -> time::Duration { START_WINDOW } @@ -362,14 +368,14 @@ impl ProvisionerTask for RecordingAiAnalysis { }) } - fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value { + fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value { json!({ "error": error, "usage": TokenUsage::of(last) }) } - async fn run(&self, attempt: &Attempt, params: Params) -> Result { + async fn run(&self, attempt: &Attempt, params: Params) -> Result { let session_id = Uuid::parse_str(&attempt.target).map_err(|error| { warn!(task.id = %attempt.task_id, %error, "Invalid session ID"); - AttemptError::Permanent(Permanent::InvalidTask) + AnalysisError::Permanent(Permanent::InvalidTask) })?; // UNTESTED: the key lives until the deadline, which this attempt start moved, so a retry hours after the Task was @@ -379,7 +385,7 @@ impl ProvisionerTask for RecordingAiAnalysis { // Waiting for an earlier analysis of the session is no progress, so the deadline cancels it too. let _session = tokio::select! { session = lock_session(session_id) => session, - () = attempt.cancel.cancelled() => return Err(AttemptError::Permanent(Permanent::DeadlinePassed)), + () = attempt.cancel.cancelled() => return Err(AnalysisError::Permanent(Permanent::DeadlinePassed)), }; // A new Task never resumes from what an earlier analysis left, while its retries resume from their own files. @@ -390,30 +396,30 @@ impl ProvisionerTask for RecordingAiAnalysis { .await .map_err(|error| { error!(task.id = %attempt.task_id, %error, "Failed to remove the earlier AI analysis working files"); - AttemptError::Permanent(Permanent::Internal) + AnalysisError::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) + AnalysisError::Permanent(Permanent::Internal) })?; } let secret = self.secrets.get(SecretId(attempt.task_id)).map_err(|error| { error!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to read the API key"); - AttemptError::Permanent(Permanent::Internal) + AnalysisError::Permanent(Permanent::Internal) })?; let Some(secret) = secret else { warn!(task.id = %attempt.task_id, attempts = attempt.attempts, "AI analysis lost its API key"); - return Err(AttemptError::Permanent(Permanent::KeyLost)); + return Err(AnalysisError::task(Failure::KeyLost)); }; let access = AiAccess::from_secret(&secret).map_err(|error| { error!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to read the API key"); - AttemptError::Permanent(Permanent::Internal) + AnalysisError::Permanent(Permanent::Internal) })?; let settings = AiSettings { @@ -428,8 +434,8 @@ impl ProvisionerTask for RecordingAiAnalysis { } /// The API key is kept in memory only, so no attempt can run after a restart. - fn restart_failure(&self) -> Option { - Some(Permanent::KeyLost) + fn restart_failure(&self) -> Option { + Some(Failure::KeyLost) } async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { @@ -461,12 +467,12 @@ impl RecordingAiAnalysis { settings: AiSettings, language: Option, api_key: SecretString, - ) -> Result { + ) -> Result { let conf = self.conf_handle.get_conf(); let client = settings.client(&conf, &api_key).map_err(|error| { error!(task.id = %attempt.task_id, %error, "Failed to build the AI client"); - AttemptError::Permanent(Permanent::Internal) + AnalysisError::Permanent(Permanent::Internal) })?; let recording = self @@ -474,16 +480,16 @@ impl RecordingAiAnalysis { .get_finished(session_id) .await .map_err(|error| match error { - FinishedRecordingError::NotFound => AttemptError::Permanent(Permanent::NoRecording), + FinishedRecordingError::NotFound => AnalysisError::task(Failure::NoRecording), error => { error!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to read the recording"); - AttemptError::Permanent(Permanent::Internal) + AnalysisError::Permanent(Permanent::Internal) } })?; let manifest = analysis_manifest(&recording.manifest).map_err(|error| match error { - RecordingsError::Mixed => AttemptError::Permanent(Permanent::MixedRecording), - _ => AttemptError::Permanent(Permanent::NoRecording), + RecordingsError::Mixed => AnalysisError::task(Failure::MixedRecording), + _ => AnalysisError::task(Failure::NoRecording), })?; let manifest_path = recording.manifest_path(); @@ -541,7 +547,7 @@ impl RecordingAiAnalysis { Some(_) => {} None => { error!(task.id = %attempt.task_id, "The AI analysis stopped unexpectedly"); - return Err(AttemptError::Permanent(Permanent::Panicked)); + return Err(AnalysisError::Permanent(Permanent::Panicked)); } } }; @@ -553,7 +559,7 @@ impl RecordingAiAnalysis { .await .map_err(|error| { error!(task.id = %attempt.task_id, %error, "Failed to keep the AI analysis log"); - AttemptError::Permanent(Permanent::Internal) + AnalysisError::Permanent(Permanent::Internal) })?; debug!(task.id = %attempt.task_id, actions = analysis.actions, "AI analysis log written"); @@ -610,11 +616,11 @@ impl RecordingAiAnalysis { } /// Waits until no recording of the session waits for its remux, since a remux replaces a video that would be read. - async fn wait_for_remux(&self, attempt: &Attempt, dir: &Utf8Path) -> Result<(), AttemptError> { + async fn wait_for_remux(&self, attempt: &Attempt, dir: &Utf8Path) -> Result<(), AnalysisError> { loop { let pending = remux_pending(dir).await.map_err(|error| { warn!(task.id = %attempt.task_id, %error, %dir, "Failed to look for a pending remux"); - AttemptError::Recoverable(Recoverable::Storage) + AnalysisError::Recoverable(Recoverable::Storage) })?; if !pending { @@ -625,7 +631,7 @@ impl RecordingAiAnalysis { tokio::select! { () = tokio::time::sleep(REMUX_POLL_INTERVAL) => {} - () = attempt.cancel.cancelled() => return Err(AttemptError::Permanent(Permanent::DeadlinePassed)), + () = attempt.cancel.cancelled() => return Err(AnalysisError::Permanent(Permanent::DeadlinePassed)), } } } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 37a49cfdfa..a474e64276 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -45,7 +45,7 @@ pub const MAX_ATTEMPTS: u32 = 5; /// Time before the attempt that follows attempt number `attempts`: 30 seconds, doubled after each attempt. fn retry_delay(attempts: u32) -> time::Duration { - time::Duration::seconds(30) * 2_i32.saturating_pow(attempts.saturating_sub(1)) + time::Duration::seconds(30) * 2i32.saturating_pow(attempts.saturating_sub(1)) } /// Where the runner queues the jobs of its records. @@ -84,6 +84,12 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// The payload of a `succeeded` record. type Output: Serialize + Send; + /// Why an attempt of this task fails for good, next to the failures any task may have. + type Failure: Serialize + Send; + + /// Why an attempt of this task fails for now, so a later attempt may succeed. + type Recoverable: Serialize + Send; + /// Time from the creation of a record until its first attempt must start. fn start_window(&self) -> time::Duration; @@ -96,20 +102,24 @@ pub trait ProvisionerTask: Send + Sync + 'static { fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value; /// The payload of a record that failed with `error`, from `last`, its payload until then: `null` if it never ran. - fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value; + fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value; /// Why an unfinished record fails when Gateway starts again, for a task that cannot outlive a restart, such as one /// whose secrets are kept in memory only; `None` for a task that goes on. /// /// A record that cannot be failed at startup is still handed to a new job, so [`ProvisionerTask::run`] must fail /// the same way after a restart. - fn restart_failure(&self) -> Option; + fn restart_failure(&self) -> Option; /// Runs one attempt. /// /// The attempt should stop soon after [`Attempt::cancel`] is cancelled, which happens once the record is past its /// deadline. - async fn run(&self, attempt: &Attempt, params: Self::Params) -> Result; + async fn run( + &self, + attempt: &Attempt, + params: Self::Params, + ) -> Result>; /// Finishes what the record leaves to do once it has ended, however it ended, such as dropping what was kept only /// while it ran. @@ -136,6 +146,12 @@ pub struct Attempt { latest: Arc>, } +/// An attempt the prelude started: the task that runs it, and the record as the attempt started. +struct Started { + task: Arc, + record: ProvisionerTaskRecord, +} + /// 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, @@ -179,6 +195,27 @@ impl Attempt { } } +/// [`Permanent`] with the failures of its kind of task as JSON, as the runner holds every kind together. +pub type AnyPermanent = Permanent; + +/// [`AttemptError`] with the failures of its kind of task as JSON. +type AnyAttemptError = AttemptError; + +fn erase_attempt_error(error: AttemptError) -> AnyAttemptError { + match error { + AttemptError::Recoverable(error) => AttemptError::Recoverable(json!(error)), + AttemptError::Permanent(error) => AttemptError::Permanent(match error { + Permanent::InvalidTask => Permanent::InvalidTask, + Permanent::Panicked => Permanent::Panicked, + Permanent::Internal => Permanent::Internal, + Permanent::AttemptsExhausted { last } => Permanent::AttemptsExhausted { last: json!(last) }, + Permanent::DeadlinePassed => Permanent::DeadlinePassed, + Permanent::NotQueued => Permanent::NotQueued, + Permanent::Task(failure) => Permanent::Task(json!(failure)), + }), + } +} + /// [`ProvisionerTask`] with its types erased, so the runner can hold every task together. #[async_trait] trait AnyProvisionerTask: Send + Sync { @@ -188,11 +225,11 @@ trait AnyProvisionerTask: Send + Sync { fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value; - fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value; + fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value; - fn restart_failure(&self) -> Option; + fn restart_failure(&self) -> Option; - async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result; + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result; async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()>; } @@ -211,25 +248,27 @@ impl AnyProvisionerTask for T { ProvisionerTask::starting_payload(self, last) } - fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value { + fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value { ProvisionerTask::failed_payload(self, error, last) } - fn restart_failure(&self) -> Option { - ProvisionerTask::restart_failure(self) + fn restart_failure(&self) -> Option { + ProvisionerTask::restart_failure(self).map(|failure| Permanent::Task(json!(failure))) } - async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { let params = serde_json::from_value(params).map_err(|error| { warn!(task.id = %attempt.task_id, %error, "Invalid Task parameters"); - AttemptError::Permanent(Permanent::InvalidTask) + AnyAttemptError::Permanent(Permanent::InvalidTask) })?; - let output = ProvisionerTask::run(self, attempt, params).await?; + let output = ProvisionerTask::run(self, attempt, params) + .await + .map_err(erase_attempt_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) + AnyAttemptError::Permanent(Permanent::Internal) }) } @@ -452,18 +491,33 @@ impl ProvisionerTaskRunner { .await } - /// Runs one attempt for the job holding `job_token`. + /// Runs one attempt for the job holding `job_token`, in three steps. /// + /// The prelude marks the start of the attempt on the record, the body runs the work of the task, and the completion + /// records how it ended, then ends the record or queues its next attempt. /// The record counts its attempts, and the job queues the next one itself, so the job queue retries a job only when /// what it had to record or queue could not be stored. async fn run_attempt(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { + let Some(started) = self.begin_attempt(task_id, job_token).await? else { + return Ok(()); + }; + + let (outcome, latest) = self.execute(&started).await; + + self.complete(&started, job_token, outcome, latest).await + } + + /// The prelude: marks the start of one more attempt on the record, unless another job holds it or it has ended. + /// + /// A record that has ended is ended again instead, since its end may have been cut short. + async fn begin_attempt(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result> { let store = &self.inner.store; let now = OffsetDateTime::now_utc(); // The starting payload carries over what earlier attempts reported, so the record is read first. let Some(record) = store.get(task_id, now).await? else { debug!(task.id = %task_id, "Task is unknown, so this attempt does not run"); - return Ok(()); + return Ok(None); }; let Some(task) = self.task(&record.kind) else { @@ -471,57 +525,63 @@ impl ProvisionerTaskRunner { 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) + .fail(task_id, json!({ "error": AnyPermanent::InvalidTask }), now) .await .context("record the Task failure")?; } - return Ok(()); + return Ok(None); }; let starting_payload = task.starting_payload(&record.payload); let deadline_at = deadline(now, task.progress_timeout()); - let record = match store - .start_attempt(task_id, job_token, starting_payload.clone(), deadline_at, now) + match store + .start_attempt(task_id, job_token, starting_payload, deadline_at, now) .await? { - AttemptStart::Started(record) => record, + AttemptStart::Started(record) => Ok(Some(Started { task, record })), AttemptStart::Finished(record) => { debug!(task.id = %task_id, "Task is over, so this attempt does not run"); - return self.end_record(task.as_ref(), &record, job_token).await; + self.end_record(task.as_ref(), &record, job_token).await?; + Ok(None) } AttemptStart::Superseded => { debug!(task.id = %task_id, "Another job runs this Task, so this one stops"); - return Ok(()); + Ok(None) } AttemptStart::Unknown => { debug!(task.id = %task_id, "Task is unknown, so this attempt does not run"); - return Ok(()); + Ok(None) } - }; + } + } - let attempts = record.attempts; + /// The body: runs the attempt of the task, and cancels it once the record is past its deadline. + /// + /// Returns how the attempt ended, with the last payload it reported and the deadline it last set. + async fn execute(&self, started: &Started) -> (Result, Latest) { + let Started { task, record } = started; let cancel = CancellationToken::new(); let latest = Arc::new(Mutex::new(Latest { - payload: starting_payload.clone(), - deadline_at, + payload: record.payload.clone(), + deadline_at: record.deadline_at, })); let attempt = Attempt { - task_id, + task_id: record.id, target: record.target.clone(), - attempts, - payload: starting_payload, + attempts: record.attempts, + payload: record.payload.clone(), cancel: cancel.clone(), progress_timeout: task.progress_timeout(), - store: Arc::clone(store), + store: Arc::clone(&self.inner.store), latest: Arc::clone(&latest), }; - info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task attempt started"); + info!(task.id = %record.id, task.kind = %record.kind, attempts = record.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 task = Arc::clone(task); let params = record.params.clone(); async move { task.run(&attempt, params).await } }); @@ -539,7 +599,12 @@ impl ProvisionerTaskRunner { continue; } - warn!(task.id = %task_id, task.kind = %record.kind, attempts, "Task went past its deadline; its attempt is cancelled"); + warn!( + task.id = %record.id, + task.kind = %record.kind, + attempts = record.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. @@ -556,7 +621,31 @@ impl ProvisionerTaskRunner { Err(_) => Err(AttemptError::Permanent(Permanent::Panicked)), }; - let last = latest.lock().payload.clone(); + let latest = latest.lock(); + + ( + outcome, + Latest { + payload: latest.payload.clone(), + deadline_at: latest.deadline_at, + }, + ) + } + + /// The completion: records how the attempt ended, then ends the record or queues its next attempt. + /// + /// Recording the end and dropping what the record kept are separate steps, so a failed drop never undoes the end. + async fn complete( + &self, + started: &Started, + job_token: Uuid, + outcome: Result, + latest: Latest, + ) -> anyhow::Result<()> { + let Started { task, record } = started; + let (task_id, attempts) = (record.id, record.attempts); + let store = &self.inner.store; + let last = &latest.payload; match outcome { Ok(output) => match store.succeed(task_id, output, OffsetDateTime::now_utc()).await { @@ -582,7 +671,7 @@ impl ProvisionerTaskRunner { "Failed to record the Task success; the Task fails instead" ); - let payload = task.failed_payload(&Permanent::Internal, &last); + let payload = task.failed_payload(&Permanent::Internal, last); match store.fail(task_id, payload, OffsetDateTime::now_utc()).await { // The job retries only when the end could not be queued again, which then only ends the Task. @@ -591,7 +680,7 @@ impl ProvisionerTaskRunner { 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; + let after_deadline = latest.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"); @@ -606,7 +695,7 @@ impl ProvisionerTaskRunner { task.id = %task_id, task.kind = %record.kind, attempts, - ?error, + %error, %retry_at, "Task attempt failed; it will be retried" ); @@ -615,19 +704,19 @@ impl ProvisionerTaskRunner { .context("queue the next attempt")?; } Err(AttemptError::Recoverable(error)) => { - warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, "Task failed after its last attempt"); + warn!(task.id = %task_id, task.kind = %record.kind, attempts, %error, "Task failed after its last attempt"); self.fail( task.as_ref(), task_id, job_token, &Permanent::AttemptsExhausted { last: error }, - &last, + last, ) .await?; } Err(AttemptError::Permanent(error)) => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, "Task failed"); - self.fail(task.as_ref(), task_id, job_token, &error, &last).await?; + self.fail(task.as_ref(), task_id, job_token, &error, last).await?; } } @@ -643,7 +732,7 @@ impl ProvisionerTaskRunner { task: &dyn AnyProvisionerTask, task_id: Uuid, job_token: Uuid, - error: &Permanent, + error: &AnyPermanent, last: &serde_json::Value, ) -> anyhow::Result<()> { self.inner diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 2bc5e38bb6..d99e09ce24 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -5,7 +5,6 @@ use provisioner_task::{ProvisionerTaskState, ProvisionerTaskStore}; use super::*; use crate::job_queue::{JobMessage, JobQueueReceiver}; -use crate::provisioner_tasks::Recoverable; const KIND: &str = "test.echo"; @@ -15,6 +14,23 @@ 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; +/// Why an Echo attempt fails for good. +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +enum EchoFailure { + /// What the attempt needed is gone, as a secret kept in memory is after a restart. + Gone, +} + +/// Why an Echo attempt fails for now. +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +enum EchoRecoverable { + Busy, +} + +type EchoError = AttemptError; + #[derive(Debug, Serialize, serde::Deserialize)] struct EchoParams { value: u32, @@ -22,14 +38,14 @@ struct EchoParams { /// A task whose attempts return `outcome`, counting its runs and the records it ended. struct Echo { - outcome: fn(u32) -> Result, + outcome: fn(u32) -> Result, progress_timeout: time::Duration, runs: Arc, ended: Arc>>, /// Ends that fail before one succeeds. end_failures: Arc, /// Why its unfinished records fail when Gateway starts again. - restart_failure: Option, + restart_failure: Option, } #[async_trait] @@ -40,6 +56,10 @@ impl ProvisionerTask for Echo { type Output = serde_json::Value; + type Failure = EchoFailure; + + type Recoverable = EchoRecoverable; + fn start_window(&self) -> time::Duration { time::Duration::hours(1) } @@ -52,28 +72,28 @@ impl ProvisionerTask for Echo { json!({ "after": last }) } - fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value { + fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value { json!({ "error": error, "last": last }) } - fn restart_failure(&self) -> Option { + fn restart_failure(&self) -> Option { self.restart_failure.clone() } - async fn run(&self, attempt: &Attempt, params: EchoParams) -> Result { + async fn run(&self, attempt: &Attempt, params: EchoParams) -> Result { self.runs.fetch_add(1, Ordering::SeqCst); match params.value { HANG => { attempt.cancel.cancelled().await; - Err(AttemptError::Recoverable(Recoverable::Storage)) + Err(AttemptError::Recoverable(EchoRecoverable::Busy)) } PROGRESS => { // 4 s of work outlasts the first deadline, 2 to 3 s away, so only the progress keeps the attempt running. for step in 0..8 { tokio::select! { () = tokio::time::sleep(std::time::Duration::from_millis(500)) => {} - () = attempt.cancel.cancelled() => return Err(AttemptError::Recoverable(Recoverable::Storage)), + () = attempt.cancel.cancelled() => return Err(AttemptError::Recoverable(EchoRecoverable::Busy)), } attempt.report_progress(json!({ "step": step })).await; } @@ -180,21 +200,21 @@ struct Harness { } impl Harness { - async fn new(outcome: fn(u32) -> Result) -> Self { + 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, + outcome: fn(u32) -> Result, progress_timeout: time::Duration, ) -> Self { Self::build(outcome, progress_timeout, None).await } async fn build( - outcome: fn(u32) -> Result, + outcome: fn(u32) -> Result, progress_timeout: time::Duration, - restart_failure: Option, + restart_failure: Option, ) -> Self { let conn = gateway_db::GatewayDb::open_path(":memory:") .await @@ -385,7 +405,10 @@ async fn resume_gives_every_unfinished_task_a_new_job() { let mut harness = Harness::new(|value| Ok(json!(value))).await; let first = harness.create(1).await; let second = harness.create(2).await; - let tokens = [harness.task(first).await.job_token, harness.task(second).await.job_token]; + let tokens = [ + harness.task(first).await.job_token, + harness.task(second).await.job_token, + ]; // The jobs of the creation may be stored or lost; the runner never asks which. harness.queued_jobs(); @@ -402,7 +425,7 @@ async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { let mut harness = Harness::build( |value| Ok(json!(value)), time::Duration::hours(1), - Some(Permanent::KeyLost), + Some(EchoFailure::Gone), ) .await; let id = harness.create(1).await; @@ -412,7 +435,7 @@ async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed, "it fails before any job runs"); - assert_eq!(task.payload, json!({ "error": { "kind": "keyLost" }, "last": null })); + assert_eq!(task.payload, json!({ "error": { "kind": "gone" }, "last": null })); let mut ending = harness.queued_job(); stale.run().await.expect("the stale job stops"); @@ -429,7 +452,7 @@ async fn a_task_whose_restart_failure_cannot_be_recorded_still_gets_a_new_job() let mut harness = Harness::build( |value| Ok(json!(value)), time::Duration::hours(1), - Some(Permanent::KeyLost), + Some(EchoFailure::Gone), ) .await; let id = harness.create(1).await; @@ -448,13 +471,21 @@ async fn a_task_whose_restart_failure_cannot_be_recorded_still_gets_a_new_job() // oracle: mutant resume resets the attempts of a Task, so a Task that fails at each restart is retried forever #[tokio::test] async fn resume_keeps_the_attempts_a_task_made() { - let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(Recoverable::AiBusy { status: 429 }))).await; + let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).await; let id = harness.create(1).await; - harness.queued_job().run().await.expect("the job queues the next attempt itself"); + harness + .queued_job() + .run() + .await + .expect("the job queues the next attempt itself"); harness.queued_jobs(); harness.runner.resume().await.expect("resume"); - harness.queued_job().run().await.expect("the job queues the next attempt itself"); + harness + .queued_job() + .run() + .await + .expect("the job queues the next attempt itself"); assert_eq!(harness.task(id).await.attempts, 2); } @@ -480,7 +511,7 @@ async fn a_panicking_attempt_fails_the_task() { // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management #[tokio::test] async fn a_permanent_failure_fails_the_task_at_once() { - let mut harness = Harness::new(|_| Err(AttemptError::Permanent(Permanent::KeyLost))).await; + let mut harness = Harness::new(|_| Err(AttemptError::task(EchoFailure::Gone))).await; let id = harness.create(1).await; harness.queued_job().run().await.expect("no retry"); @@ -488,13 +519,13 @@ async fn a_permanent_failure_fails_the_task_at_once() { 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" })); + assert_eq!(task.payload["error"], json!({ "kind": "gone" })); } // 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 mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).await; let id = harness.create(1).await; let mut job = harness.queued_job(); @@ -521,7 +552,7 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { assert_eq!(task.attempts, MAX_ATTEMPTS); assert_eq!( task.payload["error"], - json!({ "kind": "attemptsExhausted", "last": { "kind": "aiBusy", "status": 429 } }) + json!({ "kind": "attemptsExhausted", "last": { "kind": "busy" } }) ); } From 0f7813425455a1cb1b24d3ca5356f938a721bac0 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:49:25 -0400 Subject: [PATCH 72/80] fix(dgw): pass clippy, rustfmt and the retry test of the task routes The tasks route module has a child module, so it lives in `api/tasks/mod.rs`, as the lint on module files requires. An attempt that fails for now queues the next one itself and returns Ok, so the test of a new Task after an outage expects Ok, as the other retry tests do. The jrec routes do not use `IntoResponse`. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/api/jrec/mod.rs | 2 +- .../src/api/tasks/recording_ai_analysis.rs | 8 ++++--- .../recording_ai_analysis/tests.rs | 24 +++++++++++++++---- testsuite/tests/cli/dgw/ai_analysis.rs | 3 ++- 4 files changed, 27 insertions(+), 10 deletions(-) diff --git a/devolutions-gateway/src/api/jrec/mod.rs b/devolutions-gateway/src/api/jrec/mod.rs index 0af099cedb..4dfa03784c 100644 --- a/devolutions-gateway/src/api/jrec/mod.rs +++ b/devolutions-gateway/src/api/jrec/mod.rs @@ -13,7 +13,7 @@ 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::{IntoResponse as _, Response}; +use axum::response::Response; use axum::routing::{delete, get, post}; use axum::{Json, Router}; use bytes::Bytes; diff --git a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs index c30ef59f6f..09382dc729 100644 --- a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs +++ b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs @@ -123,9 +123,11 @@ impl TryFrom for AiAnalysisTask { } fn task_json(task: ProvisionerTaskRecord) -> Result, HttpError> { - AiAnalysisTask::try_from(task) - .map(Json) - .map_err(HttpError::internal().with_msg("invalid session ID in the AI analysis").err()) + AiAnalysisTask::try_from(task).map(Json).map_err( + HttpError::internal() + .with_msg("invalid session ID in the AI analysis") + .err(), + ) } fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> Response { 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 a949d20cb4..aa16100fdb 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -529,7 +529,10 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { let task_id = harness.start_created(session_id, settings).await; let workspace = harness.workspace(session_id); - harness.run(task_id).await.expect("the job queues the next attempt itself"); + harness + .run(task_id) + .await + .expect("the job queues the next attempt itself"); // oracle: mutant a recoverable failure is retried at once, without the delay a busy provider asks for assert!( harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), @@ -679,7 +682,10 @@ async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { let task_id = harness.start_created(session_id, settings).await; for attempt in 1..MAX_ATTEMPTS { - harness.run(task_id).await.expect("the job queues the next attempt itself"); + harness + .run(task_id) + .await + .expect("the job queues the next attempt itself"); // oracle: mutant a recoverable failure is retried at once, without the delay a busy provider asks for assert!( harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), @@ -831,7 +837,12 @@ async fn analysis_waits_for_the_remux_of_the_session() { let (settings, requests) = spawn_ai_provider(|_| 0).await; let task_id = harness.start_created(session_id, settings).await; - let marker = remux_marker(&harness.recording_path.join(session_id.to_string()).join("recording-1.webm")); + let marker = remux_marker( + &harness + .recording_path + .join(session_id.to_string()) + .join("recording-1.webm"), + ); std::fs::write(&marker, b"").expect("remux mark"); let other_dir = harness.recording_path.join(Uuid::new_v4().to_string()); std::fs::create_dir_all(&other_dir).expect("other session"); @@ -967,7 +978,10 @@ async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { 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; - harness.run(task_id).await.expect("the job queues the next attempt itself"); + harness + .run(task_id) + .await + .expect("the job queues the next attempt itself"); // As if Gateway stopped after an attempt wrote its log, before its success was recorded. let dir = harness.recording_path.join(session_id.to_string()); let log = dir.join(format!(".ai-analysis-{task_id}.slog")); @@ -1223,7 +1237,7 @@ async fn a_new_task_never_resumes_from_an_earlier_task() { 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"); + harness.run(first).await.expect("the job queues the next attempt itself"); assert!(harness.workspace(session_id).join("chunk-0000.json").exists()); harness .state diff --git a/testsuite/tests/cli/dgw/ai_analysis.rs b/testsuite/tests/cli/dgw/ai_analysis.rs index a873d5d5a5..6909714e9e 100644 --- a/testsuite/tests/cli/dgw/ai_analysis.rs +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -111,7 +111,8 @@ impl Gateway { } async fn get_analysis(&self, task_id: Uuid, scope: &str) -> (StatusCode, Value) { - self.get(&format!("/jet/tasks/recording-ai-analysis/{task_id}"), scope).await + self.get(&format!("/jet/tasks/recording-ai-analysis/{task_id}"), scope) + .await } async fn get(&self, path: &str, scope: &str) -> (StatusCode, Value) { From af358209536be9fc772ceec0e7b3a7b8935988bc Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 14:58:03 -0400 Subject: [PATCH 73/80] refactor(dgw): name a Task's own failures by kind and convert them without unwrapping `Permanent::Kind` and `AttemptError::of_kind` say that the failure belongs to the Task's kind of work; `Task` read as if it built a Task. `Permanent::try_map` keeps the list of shared failures next to the enum, and a failure that cannot be written as JSON becomes `internal` instead of a panic at startup or in an attempt. A test keeps the AI analysis failure names apart from the shared ones, since the caller tells failures apart by name only. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/openapi.rs | 12 ++--- .../src/provisioner_tasks/mod.rs | 27 ++++++++-- .../recording_ai_analysis/failure.rs | 52 +++++++++++++++---- .../recording_ai_analysis/mod.rs | 8 +-- .../recording_ai_analysis/tests.rs | 5 +- .../src/provisioner_tasks/runner/mod.rs | 47 ++++++++++++----- .../src/provisioner_tasks/runner/tests.rs | 2 +- 7 files changed, 115 insertions(+), 38 deletions(-) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 4b157fcb35..17ffa5e1aa 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -321,12 +321,12 @@ 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 or field is written out too. const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| match error { - Permanent::Task(Failure::AiRejected { status: _ }) => AiAnalysisErrorKind::AiRejected, - Permanent::Task(Failure::AiBadAnswer { why: _ }) => AiAnalysisErrorKind::AiBadAnswer, - Permanent::Task(Failure::KeyLost) => AiAnalysisErrorKind::KeyLost, - Permanent::Task(Failure::NoRecording) => AiAnalysisErrorKind::NoRecording, - Permanent::Task(Failure::MixedRecording) => AiAnalysisErrorKind::MixedRecording, - Permanent::Task(Failure::UnreadableRecording { file: _ }) => AiAnalysisErrorKind::UnreadableRecording, + Permanent::Kind(Failure::AiRejected { status: _ }) => AiAnalysisErrorKind::AiRejected, + Permanent::Kind(Failure::AiBadAnswer { why: _ }) => AiAnalysisErrorKind::AiBadAnswer, + Permanent::Kind(Failure::KeyLost) => AiAnalysisErrorKind::KeyLost, + Permanent::Kind(Failure::NoRecording) => AiAnalysisErrorKind::NoRecording, + Permanent::Kind(Failure::MixedRecording) => AiAnalysisErrorKind::MixedRecording, + Permanent::Kind(Failure::UnreadableRecording { file: _ }) => AiAnalysisErrorKind::UnreadableRecording, Permanent::InvalidTask => AiAnalysisErrorKind::InvalidTask, Permanent::Panicked => AiAnalysisErrorKind::Panicked, Permanent::Internal => AiAnalysisErrorKind::Internal, diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index e5f853972f..c99b2d7488 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -37,15 +37,15 @@ pub enum AttemptError { impl AttemptError { /// A permanent failure of the Task's own kind of work. - pub fn task(failure: F) -> Self { - Self::Permanent(Permanent::Task(failure)) + pub fn of_kind(failure: F) -> Self { + Self::Permanent(Permanent::Kind(failure)) } } /// Why a Task failed, as the `error` of its payload: a closed set of kinds without text for people, since the caller /// words them. /// -/// The kinds below can happen to any Task, while [`Permanent::Task`] holds those of its own kind of work, so a new kind +/// The kinds below can happen to any Task, while [`Permanent::Kind`] holds those of its own kind of work, so a new kind /// of Task brings its failures without adding to these. #[derive(Debug, Clone, PartialEq, Eq, Serialize)] #[serde(tag = "kind", rename_all = "camelCase", rename_all_fields = "camelCase")] @@ -64,5 +64,24 @@ pub enum Permanent { NotQueued, /// A failure of the Task's own kind of work, which names its own `kind`. #[serde(untagged)] - Task(F), + Kind(F), +} + +impl Permanent { + /// The same failure, with the failures of its kind converted by `kind` and `last`. + pub fn try_map( + self, + kind: impl FnOnce(F) -> Result, + last: impl FnOnce(R) -> Result, + ) -> Result, E> { + Ok(match self { + Self::InvalidTask => Permanent::InvalidTask, + Self::Panicked => Permanent::Panicked, + Self::Internal => Permanent::Internal, + Self::AttemptsExhausted { last: failure } => Permanent::AttemptsExhausted { last: last(failure)? }, + Self::DeadlinePassed => Permanent::DeadlinePassed, + Self::NotQueued => Permanent::NotQueued, + Self::Kind(failure) => Permanent::Kind(kind(failure)?), + }) + } } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs index f0997c5f99..1aed23d9e4 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs @@ -59,11 +59,11 @@ impl From for AnalysisError { fn from(error: RecordingAnalysisError) -> Self { match error { RecordingAnalysisError::Ai(error) => Self::from(error), - RecordingAnalysisError::NoRecording => Self::task(Failure::NoRecording), + RecordingAnalysisError::NoRecording => Self::of_kind(Failure::NoRecording), RecordingAnalysisError::InvalidRecording { file_name, .. } => { - Self::task(Failure::UnreadableRecording { file: file_name }) + Self::of_kind(Failure::UnreadableRecording { file: file_name }) } - RecordingAnalysisError::Truncated => Self::task(Failure::AiBadAnswer { + RecordingAnalysisError::Truncated => Self::of_kind(Failure::AiBadAnswer { why: BadAnswer::Truncated, }), // The runner cancels an attempt only once its Task is past its deadline. @@ -75,14 +75,14 @@ impl From for AnalysisError { impl From for AnalysisError { fn from(error: Error) -> Self { - let bad_answer = |why| Self::task(Failure::AiBadAnswer { why }); + let bad_answer = |why| Self::of_kind(Failure::AiBadAnswer { why }); match &error { Error::Transport { .. } => Self::Recoverable(Recoverable::AiUnreachable), Error::Status { status, .. } if error.is_transient() => { Self::Recoverable(Recoverable::AiBusy { status: *status }) } - Error::Status { status, .. } => Self::task(Failure::AiRejected { status: *status }), + Error::Status { status, .. } => Self::of_kind(Failure::AiRejected { status: *status }), Error::InvalidResponse { .. } => bad_answer(BadAnswer::Invalid), Error::Truncated { .. } => bad_answer(BadAnswer::Truncated), Error::Refused { .. } => bad_answer(BadAnswer::Refused), @@ -96,6 +96,40 @@ impl From for AnalysisError { mod tests { use super::*; + // oracle: mutant a failure of the AI analysis takes the name of one any Task may have, so callers mix them up + #[test] + fn ai_analysis_failures_have_names_of_their_own() { + fn kind(failure: impl serde::Serialize) -> serde_json::Value { + serde_json::to_value(failure).expect("JSON")["kind"].clone() + } + + let shared = [ + Permanent::<(), ()>::InvalidTask, + Permanent::Panicked, + Permanent::Internal, + Permanent::AttemptsExhausted { last: () }, + Permanent::DeadlinePassed, + Permanent::NotQueued, + ] + .map(kind); + + for failure in [ + Failure::AiRejected { status: 401 }, + Failure::AiBadAnswer { + why: BadAnswer::Invalid, + }, + Failure::KeyLost, + Failure::NoRecording, + Failure::MixedRecording, + Failure::UnreadableRecording { + file: "recording-0.webm".to_owned(), + }, + ] { + let name = kind(&failure); + assert!(!shared.contains(&name), "{name} is also a failure any Task may have"); + } + } + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[test] fn transient_ai_errors_are_recoverable_and_the_others_permanent() { @@ -122,18 +156,18 @@ mod tests { for code in [400, 401, 403, 404] { assert_eq!( attempt(status(code, None)), - AnalysisError::task(Failure::AiRejected { status: code }) + AnalysisError::of_kind(Failure::AiRejected { status: code }) ); } assert_eq!( attempt(status(429, Some("insufficient_quota"))), - AnalysisError::task(Failure::AiRejected { status: 429 }), + AnalysisError::of_kind(Failure::AiRejected { status: 429 }), "a spent quota" ); assert_eq!( AnalysisError::from(RecordingAnalysisError::Truncated), - AnalysisError::task(Failure::AiBadAnswer { + AnalysisError::of_kind(Failure::AiBadAnswer { why: BadAnswer::Truncated }) ); @@ -141,7 +175,7 @@ mod tests { attempt(Error::Refused { reason: "content filter".to_owned() }), - AnalysisError::task(Failure::AiBadAnswer { + AnalysisError::of_kind(Failure::AiBadAnswer { why: BadAnswer::Refused }) ); 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 3c1efcbd5b..3a32d2b22e 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -414,7 +414,7 @@ impl ProvisionerTask for RecordingAiAnalysis { let Some(secret) = secret else { warn!(task.id = %attempt.task_id, attempts = attempt.attempts, "AI analysis lost its API key"); - return Err(AnalysisError::task(Failure::KeyLost)); + return Err(AnalysisError::of_kind(Failure::KeyLost)); }; let access = AiAccess::from_secret(&secret).map_err(|error| { @@ -480,7 +480,7 @@ impl RecordingAiAnalysis { .get_finished(session_id) .await .map_err(|error| match error { - FinishedRecordingError::NotFound => AnalysisError::task(Failure::NoRecording), + FinishedRecordingError::NotFound => AnalysisError::of_kind(Failure::NoRecording), error => { error!(task.id = %attempt.task_id, error = format!("{error:#}"), "Failed to read the recording"); AnalysisError::Permanent(Permanent::Internal) @@ -488,8 +488,8 @@ impl RecordingAiAnalysis { })?; let manifest = analysis_manifest(&recording.manifest).map_err(|error| match error { - RecordingsError::Mixed => AnalysisError::task(Failure::MixedRecording), - _ => AnalysisError::task(Failure::NoRecording), + RecordingsError::Mixed => AnalysisError::of_kind(Failure::MixedRecording), + _ => AnalysisError::of_kind(Failure::NoRecording), })?; let manifest_path = recording.manifest_path(); 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 aa16100fdb..1b2f0617d3 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -1237,7 +1237,10 @@ async fn a_new_task_never_resumes_from_an_earlier_task() { 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; - harness.run(first).await.expect("the job queues the next attempt itself"); + harness + .run(first) + .await + .expect("the job queues the next attempt itself"); assert!(harness.workspace(session_id).join("chunk-0000.json").exists()); harness .state diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index a474e64276..d2cc80c3f0 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -201,18 +201,26 @@ pub type AnyPermanent = Permanent; /// [`AttemptError`] with the failures of its kind of task as JSON. type AnyAttemptError = AttemptError; +/// `error` with the failures of its kind as JSON; one that cannot be written as JSON is an internal failure. +fn erase_permanent(error: Permanent) -> AnyPermanent { + error + .try_map(serde_json::to_value, serde_json::to_value) + .unwrap_or_else(|error| { + error!(%error, "Failed to write a Task failure as JSON"); + Permanent::Internal + }) +} + fn erase_attempt_error(error: AttemptError) -> AnyAttemptError { match error { - AttemptError::Recoverable(error) => AttemptError::Recoverable(json!(error)), - AttemptError::Permanent(error) => AttemptError::Permanent(match error { - Permanent::InvalidTask => Permanent::InvalidTask, - Permanent::Panicked => Permanent::Panicked, - Permanent::Internal => Permanent::Internal, - Permanent::AttemptsExhausted { last } => Permanent::AttemptsExhausted { last: json!(last) }, - Permanent::DeadlinePassed => Permanent::DeadlinePassed, - Permanent::NotQueued => Permanent::NotQueued, - Permanent::Task(failure) => Permanent::Task(json!(failure)), - }), + AttemptError::Recoverable(error) => match serde_json::to_value(error) { + Ok(error) => AttemptError::Recoverable(error), + Err(error) => { + error!(%error, "Failed to write a Task failure as JSON"); + AttemptError::Permanent(Permanent::Internal) + } + }, + AttemptError::Permanent(error) => AttemptError::Permanent(erase_permanent(error)), } } @@ -253,7 +261,8 @@ impl AnyProvisionerTask for T { } fn restart_failure(&self) -> Option { - ProvisionerTask::restart_failure(self).map(|failure| Permanent::Task(json!(failure))) + ProvisionerTask::restart_failure(self) + .map(|failure| erase_permanent(Permanent::<_, T::Recoverable>::Kind(failure))) } async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { @@ -704,7 +713,13 @@ impl ProvisionerTaskRunner { .context("queue the next attempt")?; } Err(AttemptError::Recoverable(error)) => { - warn!(task.id = %task_id, task.kind = %record.kind, attempts, %error, "Task failed after its last attempt"); + warn!( + task.id = %task_id, + task.kind = %record.kind, + attempts, + %error, + "Task failed after its last attempt" + ); self.fail( task.as_ref(), task_id, @@ -715,7 +730,13 @@ impl ProvisionerTaskRunner { .await?; } Err(AttemptError::Permanent(error)) => { - warn!(task.id = %task_id, task.kind = %record.kind, attempts, ?error, "Task failed"); + warn!( + task.id = %task_id, + task.kind = %record.kind, + attempts, + error = %serde_json::to_value(&error).unwrap_or_default(), + "Task failed" + ); self.fail(task.as_ref(), task_id, job_token, &error, last).await?; } } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index d99e09ce24..13283babe4 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -511,7 +511,7 @@ async fn a_panicking_attempt_fails_the_task() { // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management #[tokio::test] async fn a_permanent_failure_fails_the_task_at_once() { - let mut harness = Harness::new(|_| Err(AttemptError::task(EchoFailure::Gone))).await; + let mut harness = Harness::new(|_| Err(AttemptError::of_kind(EchoFailure::Gone))).await; let id = harness.create(1).await; harness.queued_job().run().await.expect("no retry"); From a14aee14da1a208954254cfcf5452860a3ba9ee2 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 15:40:25 -0400 Subject: [PATCH 74/80] refactor(dgw): read a recording for the AI analysis without waiting for its remux The analysis needs nothing a remux adds: it reads the tracks, the cluster timestamps and the blocks of a WebM in order, and stops at a cut-short end. A remux replaces only the video file, by renaming its output over it, and an analysis that has the file open keeps reading the bytes it opened. So the remux mark, its 20-minute cap and the wait go, and the remux job is back to what master has. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../recording_ai_analysis/mod.rs | 34 +---- .../recording_ai_analysis/tests.rs | 56 +------- devolutions-gateway/src/recording.rs | 127 +----------------- 3 files changed, 4 insertions(+), 213 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 3a32d2b22e..e40c9af5d0 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -38,9 +38,7 @@ use super::runner::{AnyPermanent, Attempt, ProvisionerTask}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; -use crate::recording::{ - FinishedRecordingError, JrecManifest, REMUX_MARKER_STALE_AFTER, RecordingMessageSender, remux_pending, -}; +use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; use crate::secret_store::{SecretId, SecretStore}; pub const KIND: &str = "recording-ai-analysis"; @@ -51,13 +49,6 @@ pub const START_WINDOW: time::Duration = time::Duration::hours(2); /// Time a running AI analysis may go without progress before it fails. pub const PROGRESS_TIMEOUT: time::Duration = time::Duration::minutes(30); -// INVARIANT: a mark that no remux will remove stops counting before an analysis waiting on it times out, since waiting -// for a remux is no progress. -const _: () = assert!(REMUX_MARKER_STALE_AFTER.as_secs() < PROGRESS_TIMEOUT.whole_seconds().unsigned_abs()); - -/// Time between two checks for a remux of the session. -const REMUX_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1); - /// Tokens the AI provider counted for a Task, over all its attempts, as kept in its payloads. /// /// It mirrors [`UsageSoFar`], so a change in the AI crate never changes what we store. @@ -493,8 +484,6 @@ impl RecordingAiAnalysis { })?; let manifest_path = recording.manifest_path(); - self.wait_for_remux(attempt, &recording.dir).await?; - let mut request = client .analyze_recording(&manifest, &manifest_path) .owner(attempt.task_id.to_string()) @@ -614,27 +603,6 @@ impl RecordingAiAnalysis { Ok(()) } - - /// Waits until no recording of the session waits for its remux, since a remux replaces a video that would be read. - async fn wait_for_remux(&self, attempt: &Attempt, dir: &Utf8Path) -> Result<(), AnalysisError> { - loop { - let pending = remux_pending(dir).await.map_err(|error| { - warn!(task.id = %attempt.task_id, %error, %dir, "Failed to look for a pending remux"); - AnalysisError::Recoverable(Recoverable::Storage) - })?; - - if !pending { - 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(AnalysisError::Permanent(Permanent::DeadlinePassed)), - } - } - } } /// The payload of a running AI analysis Task: its progress, for whoever asks about it, and the tokens it used so far. 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 1b2f0617d3..f35d82e5cc 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -13,7 +13,7 @@ use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; use crate::job_queue::{JobMessage, JobQueueReceiver}; use crate::provisioner_tasks::runner::{MAX_ATTEMPTS, ProvisionerTaskJob, ProvisionerTaskRunner}; -use crate::recording::{RecordingManagerTask, remux_marker}; +use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles}; const API_KEY: &str = "sk-recording-ai-analysis-test-secret"; @@ -829,60 +829,6 @@ async fn finished_task_does_not_run_again() { ); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis -#[tokio::test] -async fn analysis_waits_for_the_remux_of_the_session() { - 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 marker = remux_marker( - &harness - .recording_path - .join(session_id.to_string()) - .join("recording-1.webm"), - ); - std::fs::write(&marker, b"").expect("remux mark"); - let other_dir = harness.recording_path.join(Uuid::new_v4().to_string()); - std::fs::create_dir_all(&other_dir).expect("other session"); - std::fs::write(remux_marker(&other_dir.join("recording-0.webm")), b"").expect("remux mark of another session"); - - 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"); - - std::fs::remove_file(&marker).expect("remux done"); - job.await.expect("job task").expect("no retry"); - - let task = harness.task(task_id).await; - assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); - assert!(!requests.lock().is_empty()); -} - -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis -#[tokio::test] -async fn waiting_for_a_remux_is_no_progress() { - let mut harness = Harness::new().await; - let session_id = harness.write_session(); - let (settings, requests) = spawn_ai_provider(|_| 0).await; - let task_id = harness.start_created(session_id, settings).await; - - let dir = harness.recording_path.join(session_id.to_string()); - std::fs::write(remux_marker(&dir.join("recording-1.webm")), b"").expect("remux mark"); - harness.with_progress_timeout(time::Duration::seconds(2)); - - harness.run(task_id).await.expect("no retry"); - - let task = harness.task(task_id).await; - assert_eq!(task.state, ProvisionerTaskState::Failed); - 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() { diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index c1a7846f5d..5fc4125555 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -943,13 +943,8 @@ impl RecordingManagerTask { if cadeau::xmf::is_init() { debug!(%recording_file_path, "Enqueue video remuxing operation"); - let marker = remux_marker(&recording_file_path); - - // An AI analysis that finds no mark reads the file even while it is remuxed. - write_remux_marker(&marker).await; - // Schedule 60 seconds to wait for the streamers to release the file. - let scheduled = self + let _ = self .job_queue_handle .schedule( RemuxJob { @@ -958,10 +953,6 @@ impl RecordingManagerTask { time::OffsetDateTime::now_utc() + time::Duration::seconds(60), ) .await; - - if scheduled.is_err() { - remove_remux_marker(&marker).await; - } } else { debug!("Video remuxing was skipped because XMF native library is not loaded"); } @@ -1280,66 +1271,6 @@ impl RemuxJob { pub const NAME: &'static str = "remux"; } -const REMUX_MARKER_SUFFIX: &str = ".remuxing"; - -/// Age past which a remux mark no longer counts. -/// -/// A crash between the mark and the storing of its job leaves a mark that no remux removes. -/// The limit lets an AI analysis waiting on such a mark go on before it fails for lack of progress. -/// The remux job marks the file again as it starts, so the limit covers the remux itself. -/// Until then the mark ages from the scheduling, so a job that waits longer than the limit to start lets an AI analysis -/// read the file the remux is about to replace. -pub(crate) const REMUX_MARKER_STALE_AFTER: Duration = Duration::from_secs(20 * 60); - -/// The file that marks `recording_file` as waiting for its remux, from the time the remux is scheduled until it ends. -pub(crate) fn remux_marker(recording_file: &Utf8Path) -> Utf8PathBuf { - Utf8PathBuf::from(format!("{recording_file}{REMUX_MARKER_SUFFIX}")) -} - -async fn write_remux_marker(marker: &Utf8Path) { - if let Err(error) = fs::write(marker, b"").await { - warn!(%marker, %error, "Failed to mark the recording as waiting for its remux"); - } -} - -async fn remove_remux_marker(marker: &Utf8Path) { - match fs::remove_file(marker).await { - Err(error) if error.kind() != io::ErrorKind::NotFound => { - warn!(%marker, %error, "Failed to remove the remux mark"); - } - _ => {} - } -} - -/// Whether a recording file of the session folder `dir` waits for its remux, which would replace the file. -pub(crate) async fn remux_pending(dir: &Utf8Path) -> io::Result { - let mut entries = match fs::read_dir(dir).await { - Ok(entries) => entries, - Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(false), - Err(error) => return Err(error), - }; - - while let Some(entry) = entries.next_entry().await? { - if !entry.file_name().to_string_lossy().ends_with(REMUX_MARKER_SUFFIX) { - continue; - } - - let metadata = match entry.metadata().await { - Ok(metadata) => metadata, - // The remux ended between the listing and this look. - Err(error) if error.kind() == io::ErrorKind::NotFound => continue, - Err(error) => return Err(error), - }; - let age = metadata.modified()?.elapsed().unwrap_or_default(); - - if age < REMUX_MARKER_STALE_AFTER { - return Ok(true); - } - } - - Ok(false) -} - #[async_trait] impl job_queue::Job for RemuxJob { fn name(&self) -> &str { @@ -1351,14 +1282,7 @@ impl job_queue::Job for RemuxJob { } async fn run(&mut self) -> anyhow::Result<()> { - let input_path = core::mem::take(&mut self.input_path); - let marker = remux_marker(&input_path); - - // From here the mark ages with the remux itself, also when the job runs again after a restart. - write_remux_marker(&marker).await; - - remux(input_path).await; - remove_remux_marker(&marker).await; + remux(core::mem::take(&mut self.input_path)).await; Ok(()) } } @@ -1501,53 +1425,6 @@ mod tests { } } - // oracle: mutant a fresh remux mark is ignored, so an AI analysis reads a video its remux is about to replace - #[tokio::test] - async fn a_recording_waits_for_its_remux_while_its_mark_is_fresh() { - let dir = tempfile::tempdir().expect("temp dir"); - let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); - std::fs::write(remux_marker(&dir.join("recording-0.webm")), b"").expect("mark"); - - assert!(remux_pending(dir).await.expect("fresh mark")); - } - - // oracle: mutant a mark left by a crash before its remux job was stored keeps every AI analysis waiting - #[tokio::test] - async fn a_remux_mark_older_than_any_remux_no_longer_counts() { - let dir = tempfile::tempdir().expect("temp dir"); - let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); - let marker = remux_marker(&dir.join("recording-0.webm")); - std::fs::write(&marker, b"").expect("mark"); - - let stale = std::time::SystemTime::now() - REMUX_MARKER_STALE_AFTER - Duration::from_secs(60); - std::fs::File::options() - .write(true) - .open(&marker) - .and_then(|file| file.set_modified(stale)) - .expect("age the mark"); - - assert!(!remux_pending(dir).await.expect("stale mark")); - } - - // oracle: mutant a recording file is taken for a remux mark, so every AI analysis of a video session waits - #[tokio::test] - async fn a_session_without_remux_marks_has_no_pending_remux() { - let dir = tempfile::tempdir().expect("temp dir"); - let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); - std::fs::write(dir.join("recording-0.webm"), b"").expect("recording"); - - assert!(!remux_pending(dir).await.expect("no mark")); - } - - // oracle: mutant a session folder that does not exist fails the AI analysis as storage - #[tokio::test] - async fn a_missing_session_folder_has_no_pending_remux() { - let dir = tempfile::tempdir().expect("temp dir"); - let dir = Utf8Path::from_path(dir.path()).expect("UTF-8"); - - assert!(!remux_pending(&dir.join("gone")).await.expect("missing folder")); - } - #[tokio::test] async fn slog_recording_stays_in_files() { let harness = Harness::start(); From f860456789152c5be01e755bd55603568156ba40 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 15:45:18 -0400 Subject: [PATCH 75/80] refactor(dgw): drop the storage failure of the AI analysis Only the wait for the remux produced it, so no attempt can fail with it now, and a caller would handle a kind that never comes. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/openapi.rs | 3 --- .../src/provisioner_tasks/recording_ai_analysis/failure.rs | 2 -- 2 files changed, 5 deletions(-) diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 17ffa5e1aa..d229136eeb 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -314,8 +314,6 @@ enum AiAnalysisRecoverableErrorKind { 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 storage, such as the folder of the session. - Storage, } // The error kinds are written out above for their descriptions, which utoipa drops from the variants without fields, so @@ -338,7 +336,6 @@ const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| m const _: fn(&Recoverable) -> AiAnalysisRecoverableErrorKind = |error| match error { Recoverable::AiUnreachable => AiAnalysisRecoverableErrorKind::AiUnreachable, Recoverable::AiBusy { status: _ } => AiAnalysisRecoverableErrorKind::AiBusy, - Recoverable::Storage => AiAnalysisRecoverableErrorKind::Storage, }; const _: fn(&BadAnswer) -> AiAnalysisBadAnswer = |why| match why { diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs index 1aed23d9e4..e38206827d 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs @@ -36,8 +36,6 @@ pub enum Recoverable { 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 storage, such as the folder of the session. - Storage, } /// What was wrong with an answer of the AI provider. From 6a57963446049a66bf578deb84a438665d57df15 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 28 Oct 2026 17:05:59 -0400 Subject: [PATCH 76/80] fix(dgw): keep the end a Task owes on its record and never rerun a started attempt A review against Benoit's Evaluation found two ways recovery got worse. A crash inside `resume`, between failing a Task with `keyLost` and queuing its new job, left the Task's working files behind for good. And when a terminal write failed, the job queue ran the job again, which ran the paid AI attempt again, even past the 5-attempt limit. The record now keeps whether its end is done (`ended_at`, migration 2), and `resume` queues the end of every finished Task that has none, so an end no longer depends on a job surviving. Each job names the attempt it starts, and `start_attempt` starts it only when it is the next one, so a rerun or duplicate job never runs the attempt twice; the Task fails at its deadline and is ended then. A success that can't be recorded now ends that way too: `deadlinePassed` at the deadline, instead of trying to record `internal` right away. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../migrations/02_provisioner_task_end.sql | 4 + crates/provisioner-task-libsql/src/lib.rs | 164 ++++++++++++- crates/provisioner-task/src/lib.rs | 21 +- .../recording_ai_analysis/tests.rs | 84 ++++--- .../src/provisioner_tasks/runner/mod.rs | 228 +++++++++++------- .../src/provisioner_tasks/runner/tests.rs | 115 ++++++--- 6 files changed, 452 insertions(+), 164 deletions(-) create mode 100644 crates/provisioner-task-libsql/migrations/02_provisioner_task_end.sql diff --git a/crates/provisioner-task-libsql/migrations/02_provisioner_task_end.sql b/crates/provisioner-task-libsql/migrations/02_provisioner_task_end.sql new file mode 100644 index 0000000000..70d9ad4bee --- /dev/null +++ b/crates/provisioner-task-libsql/migrations/02_provisioner_task_end.sql @@ -0,0 +1,4 @@ +ALTER TABLE provisioner_task_records ADD COLUMN ended_at INTEGER CHECK (ended_at IS NULL OR finished_at IS NOT NULL); + +CREATE INDEX provisioner_task_records_unended ON provisioner_task_records (finished_at) + WHERE finished_at IS NOT NULL AND ended_at IS NULL; diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index deee0b5546..0af13acd4f 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -19,12 +19,15 @@ use tokio::sync::Mutex; use uuid::Uuid; /// Every schema change ever made to the `provisioner_task_*` tables, oldest first. Only ever append to it. -const MIGRATIONS: &[&str] = &[include_str!("../migrations/01_provisioner_task.sql")]; +const MIGRATIONS: &[&str] = &[ + include_str!("../migrations/01_provisioner_task.sql"), + include_str!("../migrations/02_provisioner_task_end.sql"), +]; const SCHEMA_VERSION_KEY: &str = "schema_version"; -const COLUMNS: &str = - "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at, job_token"; +const COLUMNS: &str = "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, \ + deadline_at, job_token, ended_at"; /// Tasks kept in the `provisioner_task_records` table of `gateway.db`. pub struct LibSqlProvisionerTaskStore { @@ -224,6 +227,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { &self, id: Uuid, job_token: Uuid, + attempt: u32, payload: serde_json::Value, deadline_at: OffsetDateTime, now: OffsetDateTime, @@ -251,6 +255,11 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { return Ok(AttemptStart::Finished(task)); } + if task.attempts.checked_add(1) != Some(attempt) { + tx.rollback().await.context("end task attempt")?; + return Ok(AttemptStart::AlreadyStarted(task)); + } + tx.execute( "UPDATE provisioner_task_records SET state = 'running', payload = ?2, attempts = attempts + 1, started_at = COALESCE(started_at, ?3), @@ -337,6 +346,27 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(unfinished) } + async fn list_unended(&self) -> anyhow::Result> { + let conn = self.conn.lock().await; + let mut rows = conn + .query( + &format!( + "SELECT {COLUMNS} FROM provisioner_task_records + WHERE finished_at IS NOT NULL AND ended_at IS NULL ORDER BY finished_at" + ), + (), + ) + .await + .context("query unended tasks")?; + + let mut tasks = Vec::new(); + while let Some(row) = rows.next().await.context("read unended task")? { + tasks.push(task_from_row(&row)?); + } + + Ok(tasks) + } + async fn update_running( &self, id: Uuid, @@ -393,6 +423,22 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { ) .await } + + async fn mark_ended(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result { + let changed = self + .conn + .lock() + .await + .execute( + "UPDATE provisioner_task_records SET ended_at = ?2 + WHERE id = ?1 AND finished_at IS NOT NULL AND ended_at IS NULL", + params![id.to_string(), now.unix_timestamp()], + ) + .await + .context("mark task ended")?; + + Ok(changed == 1) + } } fn state_to_sql(state: ProvisionerTaskState) -> &'static str { @@ -453,6 +499,7 @@ fn task_from_row(row: &Row) -> anyhow::Result { .context("decode task job token")? .parse() .context("parse task job token")?, + ended_at: optional_timestamp(12, "end time")?, }) } @@ -683,6 +730,7 @@ mod tests { .start_attempt( overdue.id, overdue.job_token, + 1, serde_json::json!({ "done": 3, "total": 10, "usage": { "inputTokens": 1200 } }), now() - Duration::minutes(1), now() - Duration::minutes(2), @@ -718,6 +766,7 @@ mod tests { .start_attempt( task.id, task.job_token, + 1, serde_json::json!({ "done": 0 }), now() + Duration::hours(2), now() @@ -742,6 +791,7 @@ mod tests { .start_attempt( task.id, task.job_token, + 2, serde_json::json!({ "done": 1 }), now() + Duration::hours(2), now() @@ -773,6 +823,7 @@ mod tests { .start_attempt( task.id, task.job_token, + 3, serde_json::json!({}), now() + Duration::hours(2), now() @@ -803,6 +854,7 @@ mod tests { .start_attempt( task.id, task.job_token, + 1, serde_json::json!({ "done": 0 }), now() + Duration::hours(2), now(), @@ -833,6 +885,7 @@ mod tests { .start_attempt( overdue.id, overdue.job_token, + 1, serde_json::json!({ "done": 3 }), now() + Duration::minutes(1), now(), @@ -887,6 +940,7 @@ mod tests { .start_attempt( task.id, task.job_token, + 1, serde_json::json!({}), now() + Duration::minutes(30), now(), @@ -913,6 +967,7 @@ mod tests { .start_attempt( task.id, task.job_token, + 1, serde_json::json!({}), now() + Duration::minutes(30), now(), @@ -956,6 +1011,7 @@ mod tests { .start_attempt( overdue.id, overdue.job_token, + 1, serde_json::json!({}), now() + Duration::hours(2), now() @@ -992,6 +1048,7 @@ mod tests { .start_attempt( task.id, Uuid::new_v4(), + 1, serde_json::json!({}), now() + Duration::hours(2), now() @@ -1023,6 +1080,7 @@ mod tests { .start_attempt( task.id, task.job_token, + 1, serde_json::json!({}), now() + Duration::hours(2), now() @@ -1033,7 +1091,7 @@ mod tests { ); assert!(matches!( store - .start_attempt(task.id, new_token, serde_json::json!({}), now() + Duration::hours(2), now()) + .start_attempt(task.id, new_token, 1, 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 @@ -1067,6 +1125,7 @@ mod tests { .start_attempt( running.id, running.job_token, + 1, serde_json::json!({}), now() + Duration::hours(2), now(), @@ -1107,6 +1166,7 @@ mod tests { .start_attempt( Uuid::new_v4(), Uuid::new_v4(), + 1, serde_json::json!({}), now() + Duration::hours(2), now() @@ -1132,6 +1192,102 @@ mod tests { } } + #[tokio::test] + async fn an_attempt_starts_only_once() { + // oracle: mutant start_attempt ignores the attempt number, so a job run again runs its attempt again + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + let start = |attempt| { + store.start_attempt( + task.id, + task.job_token, + attempt, + serde_json::json!({}), + now() + Duration::hours(2), + now(), + ) + }; + + assert!(matches!(start(1).await.expect("start"), AttemptStart::Started(_))); + for attempt in [1, 3] { + assert!( + matches!( + start(attempt).await.expect("start again"), + AttemptStart::AlreadyStarted(started) if started.attempts == 1 + ), + "{attempt}" + ); + } + assert!(matches!( + start(2).await.expect("start the next"), + AttemptStart::Started(started) if started.attempts == 2 + )); + } + + #[tokio::test] + async fn a_finished_task_is_listed_until_marked_ended() { + let store = store().await; + let unfinished = new_task("session-1"); + let failed = new_task("session-2"); + let ended = new_task("session-3"); + for task in [&unfinished, &failed, &ended] { + create(&store, task.clone()).await; + } + for task in [&failed, &ended] { + store.fail(task.id, serde_json::json!({}), now()).await.expect("fail"); + } + + assert!(store.mark_ended(ended.id, now()).await.expect("mark ended")); + assert!(!store.mark_ended(ended.id, now()).await.expect("mark ended again")); + assert!(!store.mark_ended(unfinished.id, now()).await.expect("mark unfinished")); + assert!(!store.mark_ended(Uuid::new_v4(), now()).await.expect("mark unknown")); + + let unended = store.list_unended().await.expect("list"); + assert_eq!(unended.iter().map(|task| task.id).collect::>(), [failed.id]); + assert_eq!(unended[0].ended_at, None); + assert_eq!( + store + .get(ended.id, now()) + .await + .expect("get task") + .expect("task") + .ended_at, + Some(now()) + ); + } + + #[tokio::test] + async fn tasks_finished_before_the_end_was_recorded_are_listed_unended() { + let conn = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open gateway database") + .connect() + .await + .expect("connect"); + conn.execute_batch(MIGRATIONS[0]).await.expect("first schema"); + conn.execute( + "INSERT INTO provisioner_task_metadata (key, value) VALUES (?1, 1)", + params![SCHEMA_VERSION_KEY], + ) + .await + .expect("record the first schema version"); + let id = Uuid::new_v4(); + conn.execute( + "INSERT INTO provisioner_task_records + (id, kind, target, params, state, payload, created_at, finished_at, deadline_at, job_token) + VALUES (?1, 'recording-ai-analysis', 'session-1', '{}', 'failed', '{}', 1, 2, 3, ?2)", + params![id.to_string(), Uuid::new_v4().to_string()], + ) + .await + .expect("insert a finished task"); + + let store = LibSqlProvisionerTaskStore::open(conn).await.expect("migrate"); + + let unended = store.list_unended().await.expect("list"); + assert_eq!(unended.iter().map(|task| task.id).collect::>(), [id]); + } + #[tokio::test] async fn new_tables_record_their_schema_version_in_metadata() { let store = store().await; diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index 587a4b93c5..65ba179203 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -45,6 +45,9 @@ pub struct ProvisionerTaskRecord { pub created_at: OffsetDateTime, pub started_at: Option, pub finished_at: Option, + /// When what the Task left to do once finished, such as dropping what it kept while it ran, was done. `None` until + /// then, so Gateway still does it after a restart. + pub ended_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. @@ -83,6 +86,9 @@ pub enum AttemptStart { Started(ProvisionerTaskRecord), /// The Task is finished, possibly just now because it is past its deadline. Finished(ProvisionerTaskRecord), + /// The attempt the job comes to start is not the next one: it already started, such as when the job runs again or + /// twice, so it must not run again. + AlreadyStarted(ProvisionerTaskRecord), /// Another job now runs the Task, so this one is stale or a duplicate. Superseded, /// No Task has this ID. @@ -103,15 +109,17 @@ 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 and - /// its deadline moved to `deadline_at`. + /// Marks the start of attempt number `attempt`, counting from 1, 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. + /// [`AttemptStart::Superseded`], one whose attempt is not the next gets [`AttemptStart::AlreadyStarted`], and + /// nothing changes for either. async fn start_attempt( &self, id: Uuid, job_token: Uuid, + attempt: u32, payload: serde_json::Value, deadline_at: OffsetDateTime, now: OffsetDateTime, @@ -126,6 +134,9 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Every unfinished Task, after failing those past their deadline. async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; + /// Every finished Task whose end is not recorded yet by [`ProvisionerTaskStore::mark_ended`]. + async fn list_unended(&self) -> 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( @@ -143,4 +154,8 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Ends a `queued` or `running` Task as `failed`. Returns `false` if it was already finished, such as when it was /// just failed for being past its deadline. async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; + + /// Records that what a finished Task left to do is done, so a restart no longer does it. Returns `false` if the + /// Task is unknown, unfinished or already ended. + async fn mark_ended(&self, id: Uuid, now: OffsetDateTime) -> 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 f35d82e5cc..343d9ebcc2 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -8,7 +8,7 @@ use devolutions_gateway_ai::{Language, session_summary}; use devolutions_gateway_task::ShutdownHandle; use job_queue::Job as _; use parking_lot::Mutex; -use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; +use provisioner_task::{AttemptStart, NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; use crate::job_queue::{JobMessage, JobQueueReceiver}; @@ -122,21 +122,49 @@ impl Harness { task_id } - /// The job holding the current token of the Task, as the job queue would run it. + /// The job holding the current token of the Task, for its next attempt, 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 + let record = self .state .provisioner_tasks .store() .get(task_id, OffsetDateTime::UNIX_EPOCH) .await .expect("read task") + .expect("task exists"); + let json = json!({ + "taskId": task_id, + "jobToken": record.job_token, + "step": { "attempt": record.attempts + 1 }, + }) + .to_string(); + + ProvisionerTaskJob::read_json(&json, self.state.provisioner_tasks.clone()).expect("valid job") + } + + /// Records the success of the first attempt of the Task without ending it, as if Gateway stopped right after. + async fn record_success(&self, task_id: Uuid) { + let store = self.state.provisioner_tasks.store(); + let now = OffsetDateTime::now_utc(); + let job_token = store + .get(task_id, now) + .await + .expect("read task") .expect("task exists") .job_token; - let json = json!({ "taskId": task_id, "jobToken": job_token }).to_string(); - ProvisionerTaskJob::read_json(&json, self.state.provisioner_tasks.clone()).expect("valid job") + let started = store + .start_attempt(task_id, job_token, 1, json!({}), now + time::Duration::hours(1), now) + .await + .expect("start the attempt"); + assert!(matches!(started, AttemptStart::Started(_)), "{started:?}"); + assert!( + store + .succeed(task_id, json!({}), now) + .await + .expect("record the success") + ); } async fn run(&self, task_id: Uuid) -> anyhow::Result<()> { @@ -365,7 +393,10 @@ 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) + format!( + r#"{{"taskId":"{task_id}","jobToken":"{}","step":{{"attempt":1}}}}"#, + task.job_token + ) ); assert!(schedule_for.is_none()); } @@ -815,18 +846,15 @@ async fn overdue_task_does_not_run() { 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 (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"); + let asked = requests.lock().len(); 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" - ); + assert_eq!(requests.lock().len(), asked, "the AI is not asked again"); } // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis @@ -997,28 +1025,22 @@ async fn a_task_that_fails_removes_its_unlisted_log() { // 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; +async fn a_success_cut_short_before_the_log_is_listed_is_finished_at_startup() { + 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; - harness.run(task_id).await.expect("analysis"); + let task_id = harness.start_created(session_id, openai("gpt-test")).await; + harness.queued_jobs(); // As if Gateway stopped right after recording the success, before the manifest listed the log. let dir = harness.recording_path.join(session_id.to_string()); let log = dir.join(format!(".ai-analysis-{task_id}.slog")); - std::fs::rename(dir.join("ai-analysis-0.slog"), &log).expect("unlist the log"); - let manifest_path = dir.join("recording.json"); - let mut manifest: serde_json::Value = - serde_json::from_slice(&std::fs::read(&manifest_path).expect("manifest")).expect("manifest JSON"); - manifest["artifacts"] = json!({}); - std::fs::write(&manifest_path, manifest.to_string()).expect("manifest"); + std::fs::write(&log, "an analysis whose log was never listed").expect("unlisted log"); + harness.record_success(task_id).await; - // The job of the Task stays stored until it returns, so it runs again when Gateway starts, and it still holds the - // Task since a finished Task keeps its job. - let mut job = harness.job(task_id).await; harness.state.provisioner_tasks.resume().await.expect("resume"); - job.run().await.expect("the job ends the Task again"); + for JobMessage { mut job, .. } in harness.queued_jobs() { + job.run().await.expect("the job ends the Task"); + } assert_eq!(ai_analyses(&dir), json!([{ "fileName": "ai-analysis-0.slog" }])); assert!(dir.join("ai-analysis-0.slog").exists()); @@ -1032,18 +1054,12 @@ async fn a_late_end_of_an_earlier_task_never_lists_its_log_over_a_newer_analysis 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"); + harness.record_success(first).await; let second = harness.start_created(session_id, settings).await; harness.run(second).await.expect("second analysis"); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index d2cc80c3f0..d673bb7bec 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -7,12 +7,12 @@ //! 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. +//! Each record names the one job allowed to run it, by its token, and each job names the attempt it starts, so a stale, +//! duplicate or rerun 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 job queue gives up on a job -//! after a few failed runs; that job is stored, so a restart runs it too. +//! A finished record keeps whether its end is done, so the same startup ends every record whose end was cut short, +//! whether or not a job for it survived. #[cfg(test)] mod tests; @@ -379,7 +379,7 @@ impl ProvisionerTaskRunner { }; let queued = match on_created(record) { - Ok(()) => self.enqueue(id, job_token, None).await, + Ok(()) => self.enqueue(id, job_token, Step::Attempt(1), None).await, Err(error) => Err(error.context("prepare the Task")), }; @@ -395,22 +395,23 @@ impl ProvisionerTaskRunner { Ok(outcome) } - /// Hands every unfinished record to a new job, since a record created just before Gateway stopped may have none. + /// Hands every unfinished record to a new job, since a record created just before Gateway stopped may have none, + /// then queues the end of every finished record whose end is not done, since Gateway may have stopped before it. /// - /// Each record gets a new `job_token`, so a job left from before finds itself superseded and stops, and the job - /// queue is never read. - /// A record whose task cannot outlive a restart fails first, and its new job only ends it. + /// Each unfinished record gets a new `job_token`, so a job left from before finds itself superseded and stops, and + /// the job queue is never read. + /// A record whose task cannot outlive a restart fails instead, and is ended with the other finished records. /// A record that was waiting to retry runs at once, since retry timing is no promise to the caller. /// 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 store = &self.inner.store; let now = OffsetDateTime::now_utc(); - let records = self.inner.store.list_unfinished(now).await?; - for record in records { + for record in store.list_unfinished(now).await? { let job_token = Uuid::new_v4(); - match self.inner.store.replace_job_token(record.id, job_token, now).await { + match store.replace_job_token(record.id, job_token, now).await { Ok(true) => {} Ok(false) => continue, Err(error) => { @@ -429,18 +430,19 @@ impl ProvisionerTaskRunner { if let Some(task) = &task && let Some(failure) = task.restart_failure() { - match self - .inner - .store + match store .fail(record.id, task.failed_payload(&failure, &record.payload), now) .await { - Ok(_) => info!( - task.id = %record.id, - task.kind = %record.kind, - ?failure, - "Task failed, since it cannot outlive a restart" - ), + Ok(_) => { + info!( + task.id = %record.id, + task.kind = %record.kind, + ?failure, + "Task failed, since it cannot outlive a restart" + ); + continue; + } // Its new job runs the attempt then, which fails the same way. Err(error) => error!( task.id = %record.id, @@ -451,7 +453,9 @@ impl ProvisionerTaskRunner { } } - match self.enqueue(record.id, job_token, None).await { + let next = Step::Attempt(record.attempts.saturating_add(1)); + + match self.enqueue(record.id, job_token, next, None).await { Ok(()) => info!(task.id = %record.id, task.kind = %record.kind, "Task queued again"), Err(error) => { error!( @@ -468,6 +472,18 @@ impl ProvisionerTaskRunner { } } + for record in store.list_unended().await? { + match self.enqueue(record.id, record.job_token, Step::End, None).await { + Ok(()) => info!(task.id = %record.id, task.kind = %record.kind, "Task end queued again"), + Err(error) => error!( + task.id = %record.id, + task.kind = %record.kind, + error = format!("{error:#}"), + "Failed to queue the end of the Task again" + ), + } + } + Ok(()) } @@ -485,6 +501,7 @@ impl ProvisionerTaskRunner { &self, task_id: Uuid, job_token: Uuid, + step: Step, schedule_for: Option, ) -> anyhow::Result<()> { self.inner @@ -493,6 +510,7 @@ impl ProvisionerTaskRunner { ProvisionerTaskJob { task_id, job_token, + step, runner: self.clone(), }, schedule_for, @@ -500,14 +518,14 @@ impl ProvisionerTaskRunner { .await } - /// Runs one attempt for the job holding `job_token`, in three steps. + /// Runs attempt number `attempt` for the job holding `job_token`, in three steps. /// /// The prelude marks the start of the attempt on the record, the body runs the work of the task, and the completion /// records how it ended, then ends the record or queues its next attempt. - /// The record counts its attempts, and the job queues the next one itself, so the job queue retries a job only when - /// what it had to record or queue could not be stored. - async fn run_attempt(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { - let Some(started) = self.begin_attempt(task_id, job_token).await? else { + /// The record counts its attempts, and the job queues the next one itself, so the job queue runs a job again only + /// when what it had to record or queue could not be stored; the prelude then finds its attempt already started. + async fn run_attempt(&self, task_id: Uuid, job_token: Uuid, attempt: u32) -> anyhow::Result<()> { + let Some(started) = self.begin_attempt(task_id, job_token, attempt).await? else { return Ok(()); }; @@ -516,10 +534,13 @@ impl ProvisionerTaskRunner { self.complete(&started, job_token, outcome, latest).await } - /// The prelude: marks the start of one more attempt on the record, unless another job holds it or it has ended. + /// The prelude: marks the start of attempt number `attempt` on the record, unless another job holds it, it has + /// ended, or the attempt already started. /// - /// A record that has ended is ended again instead, since its end may have been cut short. - async fn begin_attempt(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result> { + /// A record that has ended is ended instead, if its end is not done yet. + /// An attempt that already started is never run again, since that could do paid work twice: its record is ended + /// once its deadline passes, unless the attempt is still running and ends it first. + async fn begin_attempt(&self, task_id: Uuid, job_token: Uuid, attempt: u32) -> anyhow::Result> { let store = &self.inner.store; let now = OffsetDateTime::now_utc(); @@ -545,13 +566,30 @@ impl ProvisionerTaskRunner { let deadline_at = deadline(now, task.progress_timeout()); match store - .start_attempt(task_id, job_token, starting_payload, deadline_at, now) + .start_attempt(task_id, job_token, attempt, starting_payload, deadline_at, now) .await? { AttemptStart::Started(record) => Ok(Some(Started { task, record })), AttemptStart::Finished(record) => { debug!(task.id = %task_id, "Task is over, so this attempt does not run"); - self.end_record(task.as_ref(), &record, job_token).await?; + self.end_record(&record, job_token).await?; + Ok(None) + } + AttemptStart::AlreadyStarted(record) => { + warn!( + task.id = %task_id, + task.kind = %record.kind, + attempt, + attempts = record.attempts, + "Task attempt already started, so it does not run again" + ); + + // The store fails the record once its deadline passes, so a job that runs after it only ends it. + let after_deadline = record.deadline_at + time::Duration::SECOND; + + self.enqueue(task_id, job_token, Step::End, Some(after_deadline)) + .await + .context("queue the end of the Task")?; Ok(None) } AttemptStart::Superseded => { @@ -643,6 +681,8 @@ impl ProvisionerTaskRunner { /// The completion: records how the attempt ended, then ends the record or queues its next attempt. /// + /// If that cannot be stored, the error is returned so the job queue runs the job again, and the record fails at its + /// deadline, see [`ProvisionerTaskRunner::begin_attempt`]. /// Recording the end and dropping what the record kept are separate steps, so a failed drop never undoes the end. async fn complete( &self, @@ -653,51 +693,29 @@ impl ProvisionerTaskRunner { ) -> anyhow::Result<()> { let Started { task, record } = started; let (task_id, attempts) = (record.id, record.attempts); - let store = &self.inner.store; let last = &latest.payload; match outcome { - Ok(output) => match store.succeed(task_id, output, OffsetDateTime::now_utc()).await { - Ok(ended) => { - if ended { - info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"); - } else { - warn!( - task.id = %task_id, - task.kind = %record.kind, - "Task attempt succeeded after the Task had ended, such as past its deadline; the Task keeps its earlier end" - ); - } + Ok(output) => { + let ended = self + .inner + .store + .succeed(task_id, output, OffsetDateTime::now_utc()) + .await + .context("record the Task success")?; - self.end(task.as_ref(), task_id, job_token).await?; - } - // A success that cannot be recorded fails the Task instead of being retried, since a retry would do - // the work again, such as paying for a second analysis. - Err(error) => { - error!( + if ended { + info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"); + } else { + warn!( task.id = %task_id, - error = format!("{error:#}"), - "Failed to record the Task success; the Task fails instead" + task.kind = %record.kind, + "Task attempt succeeded after the Task had ended, such as past its deadline; the Task keeps its earlier end" ); - - let payload = task.failed_payload(&Permanent::Internal, last); - - match store.fail(task_id, payload, OffsetDateTime::now_utc()).await { - // The job retries only when the end could not be queued again, which then only ends the Task. - Ok(_) => self.end(task.as_ref(), task_id, job_token).await?, - Err(error) => { - error!(task.id = %task_id, error = format!("{error:#}"), "Failed to fail the Task; it fails at its deadline"); - - // The store fails the Task once its deadline passes, so a job that runs after it only ends it. - let after_deadline = latest.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"); - } - } - } } - }, + + self.end(task_id, job_token).await?; + } Err(AttemptError::Recoverable(error)) if attempts < MAX_ATTEMPTS => { let retry_at = OffsetDateTime::now_utc() + retry_delay(attempts); warn!( @@ -708,7 +726,7 @@ impl ProvisionerTaskRunner { %retry_at, "Task attempt failed; it will be retried" ); - self.enqueue(task_id, job_token, Some(retry_at)) + self.enqueue(task_id, job_token, Step::Attempt(attempts + 1), Some(retry_at)) .await .context("queue the next attempt")?; } @@ -746,8 +764,7 @@ impl ProvisionerTaskRunner { /// 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. + /// If the failure cannot be recorded, nothing is dropped, since the record may still run, and the error is returned. async fn fail( &self, task: &dyn AnyProvisionerTask, @@ -762,7 +779,7 @@ impl ProvisionerTaskRunner { .await .context("record the Task failure")?; - self.end(task, task_id, job_token).await + self.end(task_id, job_token).await } /// Fails a record no job will run. @@ -779,33 +796,51 @@ impl ProvisionerTaskRunner { } /// 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<()> { + async fn end(&self, 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(Some(record)) => self.end_record(&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(()), + /// Does what the finished record leaves to do, then records that it is done. + /// + /// A record still unfinished is left to its attempt, and one of a kind this Gateway does not know is left for one + /// that does. + async fn end_record(&self, record: &ProvisionerTaskRecord, job_token: Uuid) -> anyhow::Result<()> { + if !record.state.is_finished() || record.ended_at.is_some() { + return Ok(()); + } + + let Some(task) = self.task(&record.kind) else { + return Ok(()); + }; + + if let Err(error) = task.on_end(record).await { + return self.end_later(record.id, job_token, error).await; + } + + match self.inner.store.mark_ended(record.id, OffsetDateTime::now_utc()).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. + /// + /// The end stays due on the record, so the next startup ends it if that job is lost too. 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") + self.enqueue( + task_id, + job_token, + Step::End, + Some(OffsetDateTime::now_utc() + END_RETRY_DELAY), + ) + .await + .context("queue the end of the Task again") } } @@ -815,17 +850,29 @@ fn deadline(now: OffsetDateTime, timeout: time::Duration) -> OffsetDateTime { (now + timeout).truncate_to_second() } +/// What a job comes to do for its record. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] +#[serde(rename_all = "camelCase")] +enum Step { + /// Start the attempt with this number, counting from 1, unless it already started. + Attempt(u32), + /// End the finished record, unless that is already done. + End, +} + #[derive(Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] struct JobDef { task_id: Uuid, job_token: Uuid, + step: Step, } /// 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, + step: Step, runner: ProvisionerTaskRunner, } @@ -838,6 +885,7 @@ impl ProvisionerTaskJob { Ok(Self { task_id: def.task_id, job_token: def.job_token, + step: def.step, runner, }) } @@ -853,11 +901,15 @@ impl job_queue::Job for ProvisionerTaskJob { serde_json::to_string(&JobDef { task_id: self.task_id, job_token: self.job_token, + step: self.step, }) .context("failed to serialize ProvisionerTaskJob") } async fn run(&mut self) -> anyhow::Result<()> { - self.runner.run_attempt(self.task_id, self.job_token).await + match self.step { + Step::Attempt(attempt) => self.runner.run_attempt(self.task_id, self.job_token, attempt).await, + Step::End => self.runner.end(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 index 13283babe4..e62689bc9b 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -1,5 +1,6 @@ use std::sync::atomic::{AtomicUsize, Ordering}; +use job_queue::Job as _; use parking_lot::Mutex; use provisioner_task::{ProvisionerTaskState, ProvisionerTaskStore}; @@ -147,11 +148,14 @@ impl ProvisionerTaskStore for FailingStore { &self, id: Uuid, job_token: Uuid, + attempt: u32, payload: serde_json::Value, deadline_at: OffsetDateTime, now: OffsetDateTime, ) -> anyhow::Result { - self.inner.start_attempt(id, job_token, payload, deadline_at, now).await + self.inner + .start_attempt(id, job_token, attempt, payload, deadline_at, now) + .await } async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result { @@ -162,6 +166,10 @@ impl ProvisionerTaskStore for FailingStore { self.inner.list_unfinished(now).await } + async fn list_unended(&self) -> anyhow::Result> { + self.inner.list_unended().await + } + async fn update_running( &self, id: Uuid, @@ -187,6 +195,10 @@ impl ProvisionerTaskStore for FailingStore { self.inner.fail(id, payload, now).await } + + async fn mark_ended(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result { + self.inner.mark_ended(id, now).await + } } struct Harness { @@ -446,6 +458,31 @@ async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { assert_eq!(*harness.ended.lock(), [id]); } +// oracle: mutant resume looks at unfinished Tasks only, so an end cut short by a crash is never done +#[tokio::test] +async fn resume_ends_a_finished_task_whose_end_was_cut_short() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + harness.queued_jobs(); + // As if Gateway stopped right after the failure was recorded, before the Task was ended. + harness + .runner + .store() + .fail(id, json!({ "error": { "kind": "gone" } }), OffsetDateTime::now_utc()) + .await + .expect("fail"); + + harness.runner.resume().await.expect("resume"); + let mut ending = harness.queued_job(); + assert!(harness.job_queue_rx.try_recv().is_err(), "one job ends it"); + ending.run().await.expect("the job ends the Task"); + + assert_eq!(harness.runs.load(Ordering::SeqCst), 0, "no attempt runs"); + assert_eq!(*harness.ended.lock(), [id]); + harness.runner.resume().await.expect("resume again"); + assert!(harness.job_queue_rx.try_recv().is_err(), "an ended Task gets no job"); +} + // oracle: mutant a Task whose failure at startup cannot be recorded gets no new job, so nothing ends it #[tokio::test] async fn a_task_whose_restart_failure_cannot_be_recorded_still_gets_a_new_job() { @@ -558,63 +595,70 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { // oracle: mutant a success that cannot be recorded is retried, so the work is done and paid for twice #[tokio::test] -async fn a_success_that_cannot_be_recorded_fails_the_task_and_ends_it() { +async fn a_success_that_cannot_be_recorded_is_not_done_again_and_ends_after_the_deadline() { let mut harness = Harness::new(|value| Ok(json!(value))).await; let id = harness.create(1).await; + let mut job = harness.queued_job(); harness.success_failures.store(1, Ordering::SeqCst); - harness.queued_job().run().await.expect("no retry"); + job.run().await.expect_err("the job queue runs the job again"); + job.run().await.expect("its attempt already started"); let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!( - task.payload, - json!({ "error": { "kind": "internal" }, "last": { "after": null } }) - ); + assert_eq!(task.state, ProvisionerTaskState::Running, "nothing could be recorded"); assert_eq!(harness.runs.load(Ordering::SeqCst), 1, "the work is not done again"); - assert_eq!(*harness.ended.lock(), [id], "what the attempt made is dropped"); - assert!(harness.queued_jobs().is_empty(), "no attempt follows"); + assert!( + harness.ended.lock().is_empty(), + "nothing is dropped while the Task may still run" + ); + let JobMessage { schedule_for, .. } = harness.job_queue_rx.try_recv().expect("a job ends it later"); + assert!( + schedule_for.is_some_and(|at| task.deadline_at < at), + "the job runs once the store fails the Task" + ); } -// oracle: mutant an end that can be recorded neither as a success nor as a failure is retried, so the work is redone +// oracle: mutant a job run again after its last failure could not be recorded runs one more attempt past the limit #[tokio::test] -async fn a_task_whose_end_cannot_be_recorded_is_ended_after_its_deadline() { - let mut harness = Harness::new(|value| Ok(json!(value))).await; +async fn a_last_attempt_whose_failure_cannot_be_recorded_is_not_run_again() { + let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).await; let id = harness.create(1).await; - harness.success_failures.store(1, Ordering::SeqCst); + let mut job = harness.queued_job(); + for _ in 1..MAX_ATTEMPTS { + job.run().await.expect("the job queues the next attempt itself"); + job = harness.queued_job(); + } harness.failure_failures.store(1, Ordering::SeqCst); - harness.queued_job().run().await.expect("no retry"); + job.run().await.expect_err("the job queue runs the job again"); + job.run().await.expect("its attempt already started"); let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Running, "nothing could be recorded"); - let JobMessage { schedule_for, .. } = harness.job_queue_rx.try_recv().expect("a job ends it later"); - assert!( - schedule_for.is_some_and(|at| task.deadline_at < at), - "the job runs once the store fails the Task" + assert_eq!(task.attempts, MAX_ATTEMPTS); + assert_eq!( + harness.runs.load(Ordering::SeqCst), + usize::try_from(MAX_ATTEMPTS).expect("a few attempts"), + "no attempt past the last" ); - assert_eq!(harness.runs.load(Ordering::SeqCst), 1, "the work is not done again"); - assert!(harness.ended.lock().is_empty()); + let JobMessage { schedule_for, .. } = harness.job_queue_rx.try_recv().expect("a job ends it later"); + assert!(schedule_for.is_some_and(|at| task.deadline_at < at)); } -// oracle: mutant an end that can be neither recorded nor queued returns an error, so the job queue redoes the paid work +// oracle: mutant two jobs holding the same token both run the attempt #[tokio::test] -async fn a_task_whose_end_cannot_be_recorded_or_queued_is_not_retried() { - let mut harness = Harness::new(|value| Ok(json!(value))).await; +async fn a_duplicate_job_does_not_run_its_attempt_again() { + let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).await; let id = harness.create(1).await; let mut job = harness.queued_job(); - harness.job_queue_rx.close(); - harness.success_failures.store(1, Ordering::SeqCst); - harness.failure_failures.store(1, Ordering::SeqCst); + let mut duplicate = + ProvisionerTaskJob::read_json(&job.write_json().expect("job JSON"), harness.runner.clone()).expect("valid job"); - job.run().await.expect("the job queue does not run the attempt again"); + job.run().await.expect("the job queues the next attempt itself"); + harness.queued_job(); + duplicate.run().await.expect("its attempt already started"); assert_eq!(harness.runs.load(Ordering::SeqCst), 1); - assert_eq!( - harness.task(id).await.state, - ProvisionerTaskState::Running, - "the store fails it at its deadline" - ); + assert_eq!(harness.task(id).await.attempts, 1); } // oracle: mutant a queued Task of an unknown kind is never failed @@ -642,7 +686,7 @@ async fn invalid_parameters_and_unknown_kinds_fail_the_task() { .await .expect("create"); - harness.runner.run_attempt(id, job_token).await.expect("no retry"); + harness.runner.run_attempt(id, job_token, 1).await.expect("no retry"); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed, "{kind}"); @@ -721,5 +765,6 @@ async fn the_job_is_read_back_from_its_json() { assert_eq!(read.task_id, id); assert_eq!(read.job_token, harness.task(id).await.job_token); + assert_eq!(read.step, Step::Attempt(1)); assert_eq!(job.name(), ProvisionerTaskJob::NAME); } From f949771a422f5b7ed446870285d227856715ffd3 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 29 Oct 2026 14:03:20 -0400 Subject: [PATCH 77/80] refactor(dgw): run each Task attempt as a TaskJob hydrated from job dependencies `TaskJob` now runs one attempt in three steps: the prelude records its start, the body runs the task, and the completion records how it ended. A recoverable failure returns an error, so the job queue runs the same job again. Migration 2 is gone: `resume` gives each unfinished Task a new job token, stores its new job, then fails a Task that can't outlive a restart. `ProvisionerTaskRunner` becomes `ProvisionerTaskManager`, since it doesn't run anything: it creates, reads and resumes Tasks, and holds a push-only `DurableJobQueue`. The job reader names each Task job explicitly (`TaskJob::::NAME`), parses it into a `TaskJobUnhydrated`, and hydrates it with the `InMemoryJobDependencies` set once on the job queue handle. `DgwState` no longer carries `provisioner_tasks` or `secrets`. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- .../migrations/02_provisioner_task_end.sql | 4 - crates/provisioner-task-libsql/src/lib.rs | 164 +--- crates/provisioner-task/src/lib.rs | 21 +- .../src/api/tasks/recording_ai_analysis.rs | 7 +- devolutions-gateway/src/job_queue.rs | 189 +++- devolutions-gateway/src/lib.rs | 16 +- .../src/provisioner_tasks/INTENT.md | 6 +- .../src/provisioner_tasks/mod.rs | 53 +- .../recording_ai_analysis/failure.rs | 2 +- .../recording_ai_analysis/mod.rs | 61 +- .../recording_ai_analysis/tests.rs | 169 ++-- .../src/provisioner_tasks/runner/mod.rs | 829 ++++++++---------- .../src/provisioner_tasks/runner/tests.rs | 328 +++---- devolutions-gateway/src/service.rs | 23 +- devolutions-gateway/tests/preflight.rs | 24 +- devolutions-gateway/tests/traffic.rs | 2 +- devolutions-gateway/tests/webapp.rs | 8 +- 17 files changed, 876 insertions(+), 1030 deletions(-) delete mode 100644 crates/provisioner-task-libsql/migrations/02_provisioner_task_end.sql diff --git a/crates/provisioner-task-libsql/migrations/02_provisioner_task_end.sql b/crates/provisioner-task-libsql/migrations/02_provisioner_task_end.sql deleted file mode 100644 index 70d9ad4bee..0000000000 --- a/crates/provisioner-task-libsql/migrations/02_provisioner_task_end.sql +++ /dev/null @@ -1,4 +0,0 @@ -ALTER TABLE provisioner_task_records ADD COLUMN ended_at INTEGER CHECK (ended_at IS NULL OR finished_at IS NOT NULL); - -CREATE INDEX provisioner_task_records_unended ON provisioner_task_records (finished_at) - WHERE finished_at IS NOT NULL AND ended_at IS NULL; diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 0af13acd4f..deee0b5546 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -19,15 +19,12 @@ use tokio::sync::Mutex; use uuid::Uuid; /// Every schema change ever made to the `provisioner_task_*` tables, oldest first. Only ever append to it. -const MIGRATIONS: &[&str] = &[ - include_str!("../migrations/01_provisioner_task.sql"), - include_str!("../migrations/02_provisioner_task_end.sql"), -]; +const MIGRATIONS: &[&str] = &[include_str!("../migrations/01_provisioner_task.sql")]; const SCHEMA_VERSION_KEY: &str = "schema_version"; -const COLUMNS: &str = "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, \ - deadline_at, job_token, ended_at"; +const COLUMNS: &str = + "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at, job_token"; /// Tasks kept in the `provisioner_task_records` table of `gateway.db`. pub struct LibSqlProvisionerTaskStore { @@ -227,7 +224,6 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { &self, id: Uuid, job_token: Uuid, - attempt: u32, payload: serde_json::Value, deadline_at: OffsetDateTime, now: OffsetDateTime, @@ -255,11 +251,6 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { return Ok(AttemptStart::Finished(task)); } - if task.attempts.checked_add(1) != Some(attempt) { - tx.rollback().await.context("end task attempt")?; - return Ok(AttemptStart::AlreadyStarted(task)); - } - tx.execute( "UPDATE provisioner_task_records SET state = 'running', payload = ?2, attempts = attempts + 1, started_at = COALESCE(started_at, ?3), @@ -346,27 +337,6 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(unfinished) } - async fn list_unended(&self) -> anyhow::Result> { - let conn = self.conn.lock().await; - let mut rows = conn - .query( - &format!( - "SELECT {COLUMNS} FROM provisioner_task_records - WHERE finished_at IS NOT NULL AND ended_at IS NULL ORDER BY finished_at" - ), - (), - ) - .await - .context("query unended tasks")?; - - let mut tasks = Vec::new(); - while let Some(row) = rows.next().await.context("read unended task")? { - tasks.push(task_from_row(&row)?); - } - - Ok(tasks) - } - async fn update_running( &self, id: Uuid, @@ -423,22 +393,6 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { ) .await } - - async fn mark_ended(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result { - let changed = self - .conn - .lock() - .await - .execute( - "UPDATE provisioner_task_records SET ended_at = ?2 - WHERE id = ?1 AND finished_at IS NOT NULL AND ended_at IS NULL", - params![id.to_string(), now.unix_timestamp()], - ) - .await - .context("mark task ended")?; - - Ok(changed == 1) - } } fn state_to_sql(state: ProvisionerTaskState) -> &'static str { @@ -499,7 +453,6 @@ fn task_from_row(row: &Row) -> anyhow::Result { .context("decode task job token")? .parse() .context("parse task job token")?, - ended_at: optional_timestamp(12, "end time")?, }) } @@ -730,7 +683,6 @@ mod tests { .start_attempt( overdue.id, overdue.job_token, - 1, serde_json::json!({ "done": 3, "total": 10, "usage": { "inputTokens": 1200 } }), now() - Duration::minutes(1), now() - Duration::minutes(2), @@ -766,7 +718,6 @@ mod tests { .start_attempt( task.id, task.job_token, - 1, serde_json::json!({ "done": 0 }), now() + Duration::hours(2), now() @@ -791,7 +742,6 @@ mod tests { .start_attempt( task.id, task.job_token, - 2, serde_json::json!({ "done": 1 }), now() + Duration::hours(2), now() @@ -823,7 +773,6 @@ mod tests { .start_attempt( task.id, task.job_token, - 3, serde_json::json!({}), now() + Duration::hours(2), now() @@ -854,7 +803,6 @@ mod tests { .start_attempt( task.id, task.job_token, - 1, serde_json::json!({ "done": 0 }), now() + Duration::hours(2), now(), @@ -885,7 +833,6 @@ mod tests { .start_attempt( overdue.id, overdue.job_token, - 1, serde_json::json!({ "done": 3 }), now() + Duration::minutes(1), now(), @@ -940,7 +887,6 @@ mod tests { .start_attempt( task.id, task.job_token, - 1, serde_json::json!({}), now() + Duration::minutes(30), now(), @@ -967,7 +913,6 @@ mod tests { .start_attempt( task.id, task.job_token, - 1, serde_json::json!({}), now() + Duration::minutes(30), now(), @@ -1011,7 +956,6 @@ mod tests { .start_attempt( overdue.id, overdue.job_token, - 1, serde_json::json!({}), now() + Duration::hours(2), now() @@ -1048,7 +992,6 @@ mod tests { .start_attempt( task.id, Uuid::new_v4(), - 1, serde_json::json!({}), now() + Duration::hours(2), now() @@ -1080,7 +1023,6 @@ mod tests { .start_attempt( task.id, task.job_token, - 1, serde_json::json!({}), now() + Duration::hours(2), now() @@ -1091,7 +1033,7 @@ mod tests { ); assert!(matches!( store - .start_attempt(task.id, new_token, 1, serde_json::json!({}), now() + Duration::hours(2), 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 @@ -1125,7 +1067,6 @@ mod tests { .start_attempt( running.id, running.job_token, - 1, serde_json::json!({}), now() + Duration::hours(2), now(), @@ -1166,7 +1107,6 @@ mod tests { .start_attempt( Uuid::new_v4(), Uuid::new_v4(), - 1, serde_json::json!({}), now() + Duration::hours(2), now() @@ -1192,102 +1132,6 @@ mod tests { } } - #[tokio::test] - async fn an_attempt_starts_only_once() { - // oracle: mutant start_attempt ignores the attempt number, so a job run again runs its attempt again - let store = store().await; - let task = new_task("session-1"); - create(&store, task.clone()).await; - let start = |attempt| { - store.start_attempt( - task.id, - task.job_token, - attempt, - serde_json::json!({}), - now() + Duration::hours(2), - now(), - ) - }; - - assert!(matches!(start(1).await.expect("start"), AttemptStart::Started(_))); - for attempt in [1, 3] { - assert!( - matches!( - start(attempt).await.expect("start again"), - AttemptStart::AlreadyStarted(started) if started.attempts == 1 - ), - "{attempt}" - ); - } - assert!(matches!( - start(2).await.expect("start the next"), - AttemptStart::Started(started) if started.attempts == 2 - )); - } - - #[tokio::test] - async fn a_finished_task_is_listed_until_marked_ended() { - let store = store().await; - let unfinished = new_task("session-1"); - let failed = new_task("session-2"); - let ended = new_task("session-3"); - for task in [&unfinished, &failed, &ended] { - create(&store, task.clone()).await; - } - for task in [&failed, &ended] { - store.fail(task.id, serde_json::json!({}), now()).await.expect("fail"); - } - - assert!(store.mark_ended(ended.id, now()).await.expect("mark ended")); - assert!(!store.mark_ended(ended.id, now()).await.expect("mark ended again")); - assert!(!store.mark_ended(unfinished.id, now()).await.expect("mark unfinished")); - assert!(!store.mark_ended(Uuid::new_v4(), now()).await.expect("mark unknown")); - - let unended = store.list_unended().await.expect("list"); - assert_eq!(unended.iter().map(|task| task.id).collect::>(), [failed.id]); - assert_eq!(unended[0].ended_at, None); - assert_eq!( - store - .get(ended.id, now()) - .await - .expect("get task") - .expect("task") - .ended_at, - Some(now()) - ); - } - - #[tokio::test] - async fn tasks_finished_before_the_end_was_recorded_are_listed_unended() { - let conn = gateway_db::GatewayDb::open_path(":memory:") - .await - .expect("open gateway database") - .connect() - .await - .expect("connect"); - conn.execute_batch(MIGRATIONS[0]).await.expect("first schema"); - conn.execute( - "INSERT INTO provisioner_task_metadata (key, value) VALUES (?1, 1)", - params![SCHEMA_VERSION_KEY], - ) - .await - .expect("record the first schema version"); - let id = Uuid::new_v4(); - conn.execute( - "INSERT INTO provisioner_task_records - (id, kind, target, params, state, payload, created_at, finished_at, deadline_at, job_token) - VALUES (?1, 'recording-ai-analysis', 'session-1', '{}', 'failed', '{}', 1, 2, 3, ?2)", - params![id.to_string(), Uuid::new_v4().to_string()], - ) - .await - .expect("insert a finished task"); - - let store = LibSqlProvisionerTaskStore::open(conn).await.expect("migrate"); - - let unended = store.list_unended().await.expect("list"); - assert_eq!(unended.iter().map(|task| task.id).collect::>(), [id]); - } - #[tokio::test] async fn new_tables_record_their_schema_version_in_metadata() { let store = store().await; diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index 65ba179203..587a4b93c5 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -45,9 +45,6 @@ pub struct ProvisionerTaskRecord { pub created_at: OffsetDateTime, pub started_at: Option, pub finished_at: Option, - /// When what the Task left to do once finished, such as dropping what it kept while it ran, was done. `None` until - /// then, so Gateway still does it after a restart. - pub ended_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. @@ -86,9 +83,6 @@ pub enum AttemptStart { Started(ProvisionerTaskRecord), /// The Task is finished, possibly just now because it is past its deadline. Finished(ProvisionerTaskRecord), - /// The attempt the job comes to start is not the next one: it already started, such as when the job runs again or - /// twice, so it must not run again. - AlreadyStarted(ProvisionerTaskRecord), /// Another job now runs the Task, so this one is stale or a duplicate. Superseded, /// No Task has this ID. @@ -109,17 +103,15 @@ 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 attempt number `attempt`, counting from 1, by the job holding `job_token`: `running`, with one - /// more attempt counted and its deadline moved to `deadline_at`. + /// 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`], one whose attempt is not the next gets [`AttemptStart::AlreadyStarted`], and - /// nothing changes for either. + /// [`AttemptStart::Superseded`], and nothing changes. async fn start_attempt( &self, id: Uuid, job_token: Uuid, - attempt: u32, payload: serde_json::Value, deadline_at: OffsetDateTime, now: OffsetDateTime, @@ -134,9 +126,6 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Every unfinished Task, after failing those past their deadline. async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; - /// Every finished Task whose end is not recorded yet by [`ProvisionerTaskStore::mark_ended`]. - async fn list_unended(&self) -> 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( @@ -154,8 +143,4 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Ends a `queued` or `running` Task as `failed`. Returns `false` if it was already finished, such as when it was /// just failed for being past its deadline. async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; - - /// Records that what a finished Task left to do is done, so a restart no longer does it. Returns `false` if the - /// Task is unknown, unfinished or already ended. - async fn mark_ended(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result; } diff --git a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs index 09382dc729..743b651d69 100644 --- a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs +++ b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs @@ -268,12 +268,11 @@ pub(crate) async fn start_ai_analysis( security(("scope_token" = ["gateway.tasks.recording-ai-analysis.read"])), ))] pub(crate) async fn get_ai_analysis( - State(DgwState { provisioner_tasks, .. }): State, + State(state): State, _scope: RecordingAiAnalysisReadScope, extract::Path(task_id): extract::Path, ) -> Result, HttpError> { - let task = provisioner_tasks - .get(task_id, OffsetDateTime::now_utc()) + let task = analysis::get(&state, task_id, OffsetDateTime::now_utc()) .await .map_err(HttpError::internal().with_msg("failed to read the AI analysis").err())? // A Task of another kind is not found under this path. @@ -308,7 +307,7 @@ mod tests { "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 (state, _handles) = DgwState::mock(&config.to_string()).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)))); diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index d51a623ddd..16b2b83153 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -1,6 +1,6 @@ use std::future::Future; use std::path::Path; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::time::Duration; use anyhow::Context as _; @@ -8,10 +8,15 @@ use async_trait::async_trait; use devolutions_gateway_task::{ChildTask, ShutdownSignal, Task}; use job_queue::{DynJob, DynJobQueue, Job, JobCtx, JobQueue, JobReader, JobRunner, RunnerWaker}; use job_queue_libsql::libsql; +use provisioner_task::DynProvisionerTaskStore; use time::OffsetDateTime; use tokio::sync::{Notify, mpsc}; -use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; +use crate::config::ConfHandle; +use crate::provisioner_tasks::recording_ai_analysis::RecordingAiAnalysis; +use crate::provisioner_tasks::runner::TaskJob; +use crate::recording::RecordingMessageSender; +use crate::secret_store::SecretStore; pub struct JobQueueCtx { notify_runner: Arc, @@ -27,10 +32,28 @@ pub struct JobMessage { } #[derive(Clone)] -pub struct JobQueueHandle(mpsc::Sender); +pub struct JobQueueHandle { + sender: mpsc::Sender, + in_memory_job_dependencies: Arc>, +} pub type JobQueueReceiver = mpsc::Receiver; +/// What jobs need to run beyond their stored JSON, kept in memory only, such as the secrets of a Task. +#[derive(Clone)] +pub struct InMemoryJobDependencies { + pub conf_handle: ConfHandle, + pub recordings: RecordingMessageSender, + pub secrets: SecretStore, + pub task_store: DynProvisionerTaskStore, + pub job_queue: DurableJobQueue, +} + +/// A job read back from its JSON that still needs the in-memory dependencies to run. +pub trait UnhydratedJob { + fn hydrate(self, dependencies: &InMemoryJobDependencies) -> DynJob; +} + pub struct JobQueueTask { queue: DynJobQueue, job_queue_rx: JobQueueReceiver, @@ -40,7 +63,7 @@ pub struct JobRunnerTask { notify_runner: Arc, runner_waker: RunnerWaker, queue: DynJobQueue, - tasks: ProvisionerTaskRunner, + in_memory_job_dependencies: Arc>, } impl JobQueueCtx { @@ -96,9 +119,9 @@ impl JobQueueCtx { #[derive(Clone)] pub struct DurableJobQueue(DynJobQueue); -#[async_trait] -impl ProvisionerTaskJobQueue for DurableJobQueue { - async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()> { +impl DurableJobQueue { + /// Stores `job` to run at `schedule_for`, or at once without it. + pub async fn push(&self, job: T, schedule_for: Option) -> anyhow::Result<()> { let job: DynJob = Box::new(job); self.0.push_job(&job, schedule_for).await } @@ -107,11 +130,26 @@ impl ProvisionerTaskJobQueue for DurableJobQueue { impl JobQueueHandle { pub fn new() -> (Self, JobQueueReceiver) { let (tx, rx) = mpsc::channel(512); - (Self(tx), rx) + let handle = Self { + sender: tx, + in_memory_job_dependencies: Arc::default(), + }; + (handle, rx) + } + + /// Sets the dependencies the job runner hydrates each job with; once only. + pub fn set_in_memory_job_dependencies(&self, dependencies: InMemoryJobDependencies) -> anyhow::Result<()> { + self.in_memory_job_dependencies + .set(dependencies) + .map_err(|_| anyhow::anyhow!("the in-memory job dependencies are already set")) + } + + pub fn in_memory_job_dependencies(&self) -> Option<&InMemoryJobDependencies> { + self.in_memory_job_dependencies.get() } pub fn blocking_enqueue(&self, job: T) -> anyhow::Result<()> { - self.0 + self.sender .blocking_send(JobMessage { job: Box::new(job), schedule_for: None, @@ -120,7 +158,7 @@ impl JobQueueHandle { } pub async fn enqueue(&self, job: T) -> anyhow::Result<()> { - self.0 + self.sender .send(JobMessage { job: Box::new(job), schedule_for: None, @@ -134,7 +172,7 @@ impl JobQueueHandle { job: T, schedule_for: OffsetDateTime, ) -> anyhow::Result<()> { - self.0 + self.sender .blocking_send(JobMessage { job: Box::new(job), schedule_for: Some(schedule_for), @@ -143,7 +181,7 @@ impl JobQueueHandle { } pub async fn schedule(&self, job: T, schedule_for: OffsetDateTime) -> anyhow::Result<()> { - self.0 + self.sender .send(JobMessage { job: Box::new(job), schedule_for: Some(schedule_for), @@ -217,12 +255,12 @@ async fn job_queue_task(ctx: JobQueueTask, mut shutdown_signal: ShutdownSignal) } impl JobRunnerTask { - pub fn new(ctx: &JobQueueCtx, tasks: ProvisionerTaskRunner) -> Self { + pub fn new(ctx: &JobQueueCtx) -> Self { Self { notify_runner: Arc::clone(&ctx.notify_runner), runner_waker: RunnerWaker::clone(&ctx.runner_waker), queue: Arc::clone(&ctx.queue), - tasks, + in_memory_job_dependencies: Arc::clone(&ctx.job_queue_handle.in_memory_job_dependencies), } } } @@ -246,10 +284,12 @@ async fn job_runner_task(ctx: JobRunnerTask, mut shutdown_signal: ShutdownSignal notify_runner, runner_waker, queue, - tasks, + in_memory_job_dependencies, } = ctx; - let reader = DgwJobReader { tasks }; + let reader = DgwJobReader { + in_memory_job_dependencies, + }; let spawn = |mut ctx: JobCtx, callback: job_queue::SpawnCallback| { tokio::spawn(async move { @@ -302,7 +342,17 @@ async fn job_runner_task(ctx: JobRunnerTask, mut shutdown_signal: ShutdownSignal } struct DgwJobReader { - tasks: ProvisionerTaskRunner, + in_memory_job_dependencies: Arc>, +} + +impl DgwJobReader { + fn hydrate(&self, job: impl UnhydratedJob) -> anyhow::Result { + let dependencies = self + .in_memory_job_dependencies + .get() + .context("the in-memory job dependencies are not set")?; + Ok(job.hydrate(dependencies)) + } } impl JobReader for DgwJobReader { @@ -320,11 +370,108 @@ 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)) - } + TaskJob::::NAME => self.hydrate(TaskJob::::read_json(json)?), _ => anyhow::bail!("unknown job name: {name}"), } } } + +/// A job the job queue of tests stored, as the job runner reads it back. +#[cfg(test)] +pub(crate) struct StoredJob { + pub(crate) name: String, + pub(crate) json: String, + pub(crate) schedule_for: Option, +} + +/// The job queue of tests: the queue Gateway runs, in memory, which also hands each job it stores to the test, and +/// refuses every job once closed, as a database that can't be written does. +#[cfg(test)] +pub(crate) struct TestJobQueue { + inner: job_queue_libsql::LibSqlJobQueue, + stored: parking_lot::Mutex>, + closed: std::sync::atomic::AtomicBool, +} + +#[cfg(test)] +impl TestJobQueue { + pub(crate) async fn new() -> Arc { + let conn = libsql::Builder::new_local(":memory:") + .build() + .await + .expect("build database") + .connect() + .expect("open database connection"); + let inner = job_queue_libsql::LibSqlJobQueue::builder() + .runner_waker(RunnerWaker::new(|| {})) + .conn(conn) + .build(); + inner.setup().await.expect("queue setup"); + + Arc::new(Self { + inner, + stored: parking_lot::Mutex::default(), + closed: std::sync::atomic::AtomicBool::new(false), + }) + } + + /// Takes the oldest job stored and not taken yet. + pub(crate) fn take(&self) -> Option { + self.stored.lock().pop_front() + } + + pub(crate) fn close(&self) { + self.closed.store(true, std::sync::atomic::Ordering::SeqCst); + } + + /// This queue, as [`JobQueueCtx::durable`] hands out the queue Gateway runs. + pub(crate) fn durable(self: &Arc) -> DurableJobQueue { + DurableJobQueue(Arc::::clone(self)) + } +} + +#[cfg(test)] +#[async_trait] +impl JobQueue for TestJobQueue { + async fn setup(&self) -> anyhow::Result<()> { + self.inner.setup().await + } + + async fn reset_claimed_jobs(&self) -> anyhow::Result<()> { + self.inner.reset_claimed_jobs().await + } + + async fn push_job(&self, job: &DynJob, schedule_for: Option) -> anyhow::Result<()> { + anyhow::ensure!( + !self.closed.load(std::sync::atomic::Ordering::SeqCst), + "the job queue is closed" + ); + self.inner.push_job(job, schedule_for).await?; + self.stored.lock().push_back(StoredJob { + name: job.name().to_owned(), + json: job.write_json()?, + schedule_for, + }); + Ok(()) + } + + async fn claim_jobs(&self, reader: &dyn JobReader, number_of_jobs: usize) -> anyhow::Result> { + self.inner.claim_jobs(reader, number_of_jobs).await + } + + async fn delete_job(&self, job_id: uuid::Uuid) -> anyhow::Result<()> { + self.inner.delete_job(job_id).await + } + + async fn fail_job(&self, job_id: uuid::Uuid, schedule_for: OffsetDateTime) -> anyhow::Result<()> { + self.inner.fail_job(job_id, schedule_for).await + } + + async fn clear_failed(&self) -> anyhow::Result<()> { + self.inner.clear_failed().await + } + + async fn next_scheduled_date(&self) -> anyhow::Result> { + self.inner.next_scheduled_date().await + } +} diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index c20eabc64d..3480c5c21f 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -70,8 +70,6 @@ pub struct DgwState { pub monitoring_state: Arc, pub traffic_audit_handle: traffic_audit::TrafficAuditHandle, pub agent_tunnel_handle: Option>, - pub provisioner_tasks: provisioner_tasks::runner::ProvisionerTaskRunner, - pub secrets: secret_store::SecretStore, } #[doc(hidden)] @@ -86,7 +84,7 @@ pub struct MockHandles { impl DgwState { #[doc(hidden)] - pub async fn mock(json_config: &str) -> anyhow::Result<(Self, MockHandles)> { + pub 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())); @@ -99,16 +97,6 @@ impl DgwState { let provisioning = provisioning::ProvisioningStore::new(); let synthetic_kdc_registry = credential_injection::SyntheticKdcRegistry::new(); let monitoring_state = Arc::new(network_monitor::State::new(Arc::new(MockMonitorsCache))?); - let secrets = secret_store::SecretStore::default(); - let task_store = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?; - let task_store = Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store).await?); - let provisioner_tasks = provisioner_tasks::task_runner( - task_store, - Arc::new(job_queue_handle.clone()), - conf_handle.clone(), - recording_manager_handle.clone(), - secrets.clone(), - ); let state = Self { conf_handle, @@ -124,8 +112,6 @@ impl DgwState { synthetic_kdc_registry, monitoring_state, agent_tunnel_handle: None, - provisioner_tasks, - secrets, }; let handles = MockHandles { diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index ca2b0010b6..8e54b089d9 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -5,11 +5,11 @@ 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. It is also the only place that manages how tasks are executed in relation to the job queue. +`ProvisionerTaskManager` 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`. +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 `ProvisionerTaskManager`. -The task record is the source of truth; the `ProvisionerTaskJob` only wakes it up and holds no state. +The task record is the source of truth; the `TaskJob` only wakes it up and holds no state. One task at the same time only has one job to execute task. diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index c99b2d7488..a71dbf5b0b 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -1,29 +1,27 @@ -//! 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`]. +//! Work the provisioner asks Gateway to do: each request is a Task recorded in `gateway.db`, queued as a job on the job +//! queue by the [`runner::ProvisionerTaskManager`]. pub mod ai; pub mod recording_ai_analysis; pub mod runner; -use provisioner_task::DynProvisionerTaskStore; +use std::sync::Arc; use self::recording_ai_analysis::RecordingAiAnalysis; -use self::runner::{DynProvisionerTaskJobQueue, ProvisionerTaskRunner}; -use crate::config::ConfHandle; -use crate::recording::RecordingMessageSender; -use crate::secret_store::SecretStore; +use self::runner::ProvisionerTaskManager; +use crate::job_queue::InMemoryJobDependencies; -/// The provisioner task runner, with every provisioner task Gateway knows. -pub fn task_runner( - store: DynProvisionerTaskStore, - job_queue: DynProvisionerTaskJobQueue, - conf_handle: ConfHandle, - recordings: RecordingMessageSender, - secrets: SecretStore, -) -> ProvisionerTaskRunner { - ProvisionerTaskRunner::builder(store, job_queue) - .register(RecordingAiAnalysis::new(conf_handle, recordings, secrets)) - .build() +/// The manager of the AI analysis Tasks. +pub fn task_manager(dependencies: &InMemoryJobDependencies) -> ProvisionerTaskManager { + ProvisionerTaskManager::new( + RecordingAiAnalysis::new( + dependencies.conf_handle.clone(), + dependencies.recordings.clone(), + dependencies.secrets.clone(), + ), + Arc::clone(&dependencies.task_store), + dependencies.job_queue.clone(), + ) } /// Why one attempt of a Task failed, with `F` and `R` the failures of its own kind of work. @@ -66,22 +64,3 @@ pub enum Permanent { #[serde(untagged)] Kind(F), } - -impl Permanent { - /// The same failure, with the failures of its kind converted by `kind` and `last`. - pub fn try_map( - self, - kind: impl FnOnce(F) -> Result, - last: impl FnOnce(R) -> Result, - ) -> Result, E> { - Ok(match self { - Self::InvalidTask => Permanent::InvalidTask, - Self::Panicked => Permanent::Panicked, - Self::Internal => Permanent::Internal, - Self::AttemptsExhausted { last: failure } => Permanent::AttemptsExhausted { last: last(failure)? }, - Self::DeadlinePassed => Permanent::DeadlinePassed, - Self::NotQueued => Permanent::NotQueued, - Self::Kind(failure) => Permanent::Kind(kind(failure)?), - }) - } -} diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs index e38206827d..7be4132b65 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs @@ -64,7 +64,7 @@ impl From for AnalysisError { RecordingAnalysisError::Truncated => Self::of_kind(Failure::AiBadAnswer { why: BadAnswer::Truncated, }), - // The runner cancels an attempt only once its Task is past its deadline. + // The job cancels an attempt only once its Task is past its deadline. RecordingAnalysisError::Cancelled => Self::Permanent(Permanent::DeadlinePassed), _ => Self::Permanent(Permanent::Internal), } 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 e40c9af5d0..226ba52ee0 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -2,12 +2,12 @@ //! 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). +//! [`ProvisionerTaskManager`](super::runner::ProvisionerTaskManager). //! Each attempt runs the analysis of the AI crate, //! [`AiClient::analyze_recording`](devolutions_gateway_ai::AiClient::analyze_recording), which keeps its working files //! next to the session manifest so that a retry resumes where the last attempt stopped, then leaves the log in the //! session folder, unlisted, under a name of its own Task. -//! Once the runner records the success, the end of the Task lists that log in the manifest in place of the earlier one; +//! Once the job 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. @@ -34,10 +34,11 @@ use uuid::Uuid; pub use self::failure::{AnalysisError, BadAnswer, Failure, Recoverable}; use super::Permanent; use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; -use super::runner::{AnyPermanent, Attempt, ProvisionerTask}; +use super::runner::{Attempt, ProvisionerTask, TaskJobUnhydrated}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; +use crate::job_queue::{InMemoryJobDependencies, UnhydratedJob}; use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; use crate::secret_store::{SecretId, SecretStore}; @@ -158,7 +159,7 @@ pub struct Analysis { prompt_version: String, } -/// Records an AI analysis Task, whose job the [`ProvisionerTaskRunner`](super::runner::ProvisionerTaskRunner) queues. +/// Records an AI analysis Task, whose job the [`ProvisionerTaskManager`](super::runner::ProvisionerTaskManager) queues. /// /// Only a new Task keeps the API key; asking again for a known Task changes nothing, even while its session is /// recording again. @@ -193,15 +194,18 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result Result {} Err(FinishedRecordingError::Recording) => return Err(StartError::RecordingActive), @@ -231,11 +237,11 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result(task_id, target, ¶ms, |record| { + let outcome = 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 + tasks + .task() .secrets .insert(SecretId(record.id), &access.to_secret()?, record.deadline_at) }) @@ -255,6 +261,27 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result anyhow::Result> { + state + .job_queue_handle + .in_memory_job_dependencies() + .map(super::task_manager) + .context("the in-memory job dependencies are not set")? + .get(task_id, now) + .await +} + +impl UnhydratedJob for TaskJobUnhydrated { + fn hydrate(self, dependencies: &InMemoryJobDependencies) -> job_queue::DynJob { + Box::new(self.into_job(super::task_manager(dependencies))) + } +} + /// The `recording-ai-analysis` provisioner task. pub struct RecordingAiAnalysis { conf_handle: ConfHandle, @@ -359,7 +386,7 @@ impl ProvisionerTask for RecordingAiAnalysis { }) } - fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value { + fn failed_payload(&self, error: &Permanent, last: &serde_json::Value) -> serde_json::Value { json!({ "error": error, "usage": TokenUsage::of(last) }) } @@ -541,7 +568,7 @@ impl RecordingAiAnalysis { } }; - // The log stays out of the manifest until the runner records the success, so a Task that fails never replaces + // The log stays out of the manifest until the job 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) 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 343d9ebcc2..115a5aa86a 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -11,8 +11,8 @@ use parking_lot::Mutex; use provisioner_task::{AttemptStart, NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; -use crate::job_queue::{JobMessage, JobQueueReceiver}; -use crate::provisioner_tasks::runner::{MAX_ATTEMPTS, ProvisionerTaskJob, ProvisionerTaskRunner}; +use crate::job_queue::{InMemoryJobDependencies, StoredJob, TestJobQueue}; +use crate::provisioner_tasks::runner::{MAX_ATTEMPTS, ProvisionerTaskManager, TaskJob}; use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles}; @@ -41,7 +41,9 @@ fn openai(model: &str) -> AiSettings { struct Harness { state: DgwState, - job_queue_rx: JobQueueReceiver, + tasks: ProvisionerTaskManager, + secrets: SecretStore, + job_queue: Arc, recording_path: Utf8PathBuf, _keep: Box, _dir: tempfile::TempDir, @@ -52,7 +54,7 @@ impl Harness { 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 (state, handles) = DgwState::mock(&config(&recording_path)).expect("mock state"); let MockHandles { session_manager_rx, recording_manager_rx, @@ -74,12 +76,39 @@ impl Harness { recording_shutdown_signal, )); + let task_store = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open the gateway database") + .connect() + .await + .expect("connect to the gateway database"); + let task_store = provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store) + .await + .expect("open the task store"); + let secrets = SecretStore::default(); + let job_queue = TestJobQueue::new().await; + let dependencies = InMemoryJobDependencies { + conf_handle: state.conf_handle.clone(), + recordings: state.recordings.clone(), + secrets: secrets.clone(), + task_store: Arc::new(task_store), + job_queue: job_queue.durable(), + }; + let tasks = crate::provisioner_tasks::task_manager(&dependencies); + state + .job_queue_handle + .set_in_memory_job_dependencies(dependencies) + .expect("set the job dependencies"); + Self { state, - job_queue_rx, + tasks, + secrets, + job_queue, recording_path, _keep: Box::new(( session_manager_rx, + job_queue_rx, subscriber_rx, traffic_audit_rx, shutdown_handle, @@ -122,30 +151,26 @@ impl Harness { task_id } - /// The job holding the current token of the Task, for its next attempt, as the job queue would run it. - async fn job(&self, task_id: Uuid) -> ProvisionerTaskJob { + /// The job holding the current token of the Task, as the job queue would run it. + async fn job(&self, task_id: Uuid) -> TaskJob { // Read as of long ago, so reading never fails the Task in place of the job. let record = self - .state - .provisioner_tasks + .tasks .store() .get(task_id, OffsetDateTime::UNIX_EPOCH) .await .expect("read task") .expect("task exists"); - let json = json!({ - "taskId": task_id, - "jobToken": record.job_token, - "step": { "attempt": record.attempts + 1 }, - }) - .to_string(); + let json = json!({ "taskId": task_id, "jobToken": record.job_token }).to_string(); - ProvisionerTaskJob::read_json(&json, self.state.provisioner_tasks.clone()).expect("valid job") + TaskJob::::read_json(&json) + .expect("valid job") + .into_job(self.tasks.clone()) } /// Records the success of the first attempt of the Task without ending it, as if Gateway stopped right after. async fn record_success(&self, task_id: Uuid) { - let store = self.state.provisioner_tasks.store(); + let store = self.tasks.store(); let now = OffsetDateTime::now_utc(); let job_token = store .get(task_id, now) @@ -155,7 +180,7 @@ impl Harness { .job_token; let started = store - .start_attempt(task_id, job_token, 1, json!({}), now + time::Duration::hours(1), now) + .start_attempt(task_id, job_token, json!({}), now + time::Duration::hours(1), now) .await .expect("start the attempt"); assert!(matches!(started, AttemptStart::Started(_)), "{started:?}"); @@ -172,8 +197,8 @@ impl Harness { } /// Takes every queued job. - fn queued_jobs(&mut self) -> Vec { - std::iter::from_fn(|| self.job_queue_rx.try_recv().ok()).collect() + fn queued_jobs(&mut self) -> Vec { + std::iter::from_fn(|| self.job_queue.take()).collect() } /// Records a Task directly, as if it was created at `created_at`. @@ -185,8 +210,7 @@ impl Harness { created_at: OffsetDateTime, deadline_at: OffsetDateTime, ) { - self.state - .provisioner_tasks + self.tasks .store() .create( NewProvisionerTaskRecord { @@ -204,8 +228,7 @@ impl Harness { } async fn task(&self, task_id: Uuid) -> ProvisionerTaskRecord { - self.state - .provisioner_tasks + self.tasks .get(task_id, OffsetDateTime::now_utc()) .await .expect("read task") @@ -217,15 +240,10 @@ impl Harness { let mut task = RecordingAiAnalysis::new( self.state.conf_handle.clone(), self.state.recordings.clone(), - self.state.secrets.clone(), + self.secrets.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(); + self.tasks = ProvisionerTaskManager::new(task, Arc::clone(self.tasks.store()), self.job_queue.durable()); } /// Working files of the AI analyses of session_id. @@ -234,8 +252,7 @@ impl Harness { } fn access(&self, task_id: Uuid) -> Option { - self.state - .secrets + self.secrets .get(SecretId(task_id)) .expect("read the secret") .map(|secret| AiAccess::from_secret(&secret).expect("AI access")) @@ -251,8 +268,7 @@ impl Harness { api_key: SecretString::from(API_KEY), base_url, }; - self.state - .secrets + self.secrets .insert(SecretId(task_id), &access.to_secret().expect("secret"), expires_at) .expect("insert the secret"); } @@ -365,7 +381,7 @@ fn first_line(path: &Utf8Path) -> String { // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets #[tokio::test] async fn start_records_the_task_keeps_the_key_and_queues_a_job() { - let mut harness = Harness::new().await; + let harness = Harness::new().await; let task_id = Uuid::new_v4(); let session_id = harness.write_session(); let mut settings = openai("gpt-test"); @@ -389,14 +405,15 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { 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); + let StoredJob { + name, + json, + schedule_for, + } = harness.job_queue.take().expect("job queued"); + assert_eq!(name, KIND); assert_eq!( - job.write_json().expect("job JSON"), - format!( - r#"{{"taskId":"{task_id}","jobToken":"{}","step":{{"attempt":1}}}}"#, - task.job_token - ) + json, + format!(r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, task.job_token) ); assert!(schedule_for.is_none()); } @@ -404,14 +421,14 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis #[tokio::test] async fn asking_again_creates_nothing_new() { - let mut harness = Harness::new().await; + let 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.job_queue.take().expect("first job"); let same = harness .start(task_id, session_id, openai("gpt-test")) @@ -439,21 +456,21 @@ async fn asking_again_creates_nothing_new() { } ); - assert!(harness.job_queue_rx.try_recv().is_err(), "no other job is queued"); + assert!(harness.job_queue.take().is_none(), "no other job is queued"); assert!(harness.has_key(task_id), "the first key stays"); } // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis #[tokio::test] async fn asking_again_returns_the_task_even_while_the_session_records_again() { - let mut harness = Harness::new().await; + let 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.job_queue.take().expect("first job"); harness.state.recordings.active_recordings.insert(session_id); let same = harness @@ -486,8 +503,7 @@ async fn refuses_a_session_without_recording() { assert!(matches!(error, StartError::RecordingNotFound), "{error:?}"); assert!( harness - .state - .provisioner_tasks + .tasks .get(task_id, OffsetDateTime::now_utc()) .await .expect("read") @@ -512,8 +528,7 @@ async fn refuses_a_session_that_is_still_recording() { assert!(matches!(error, StartError::RecordingActive), "{error:?}"); assert!( harness - .state - .provisioner_tasks + .tasks .get(task_id, OffsetDateTime::now_utc()) .await .expect("read") @@ -558,16 +573,17 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { 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; + // The job of the creation, so only a job queued by the retry is left to find. + harness.queued_jobs(); let workspace = harness.workspace(session_id); harness .run(task_id) .await - .expect("the job queues the next attempt itself"); - // oracle: mutant a recoverable failure is retried at once, without the delay a busy provider asks for + .expect_err("the job queue runs the job again"); assert!( - harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), - "a provider outage is retried" + harness.queued_jobs().is_empty(), + "the job queue runs the same job again" ); let task = harness.task(task_id).await; @@ -711,16 +727,17 @@ async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { 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; + // The job of the creation, so only a job queued by a retry is left to find. + harness.queued_jobs(); for attempt in 1..MAX_ATTEMPTS { harness .run(task_id) .await - .expect("the job queues the next attempt itself"); - // oracle: mutant a recoverable failure is retried at once, without the delay a busy provider asks for + .expect_err("the job queue runs the job again"); assert!( - harness.queued_jobs().iter().any(|job| job.schedule_for.is_some()), - "attempt {attempt} is retried" + harness.queued_jobs().is_empty(), + "the job queue runs the same job again after attempt {attempt}" ); assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Running); } @@ -792,7 +809,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.secrets.remove(SecretId(task_id)); + harness.secrets.remove(SecretId(task_id)); harness.run(task_id).await.expect("no retry"); @@ -933,7 +950,7 @@ async fn a_new_task_that_fails_keeps_the_earlier_analysis() { harness.run(first).await.expect("first analysis"); let second = harness.start_created(session_id, settings).await; - harness.state.secrets.remove(SecretId(second)); + harness.secrets.remove(SecretId(second)); harness.run(second).await.expect("no retry"); let task = harness.task(second).await; @@ -955,13 +972,13 @@ async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { harness .run(task_id) .await - .expect("the job queues the next attempt itself"); + .expect_err("the job queue runs the job again"); // As if Gateway stopped after an attempt wrote its log, before its success was recorded. 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.resume().await.expect("resume"); + harness.tasks.resume().await.expect("resume"); let task = harness.task(task_id).await; // oracle: mutant an AI analysis that cannot run after a restart stays unfinished until its next attempt @@ -1010,8 +1027,7 @@ async fn a_task_that_fails_removes_its_unlisted_log() { 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 + .tasks .store() .fail(task_id, json!({}), OffsetDateTime::now_utc()) .await @@ -1029,7 +1045,6 @@ async fn a_success_cut_short_before_the_log_is_listed_is_finished_at_startup() { let mut harness = Harness::new().await; let session_id = harness.write_session(); let task_id = harness.start_created(session_id, openai("gpt-test")).await; - harness.queued_jobs(); // 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()); @@ -1037,8 +1052,12 @@ async fn a_success_cut_short_before_the_log_is_listed_is_finished_at_startup() { std::fs::write(&log, "an analysis whose log was never listed").expect("unlisted log"); harness.record_success(task_id).await; - harness.state.provisioner_tasks.resume().await.expect("resume"); - for JobMessage { mut job, .. } in harness.queued_jobs() { + harness.tasks.resume().await.expect("resume"); + // The job of its creation never returned, so it stays stored and runs again when Gateway starts. + for StoredJob { json, .. } in harness.queued_jobs() { + let mut job = TaskJob::::read_json(&json) + .expect("valid job") + .into_job(harness.tasks.clone()); job.run().await.expect("the job ends the Task"); } @@ -1099,7 +1118,7 @@ async fn a_new_task_that_fails_keeps_an_earlier_success_whose_log_was_not_listed 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.secrets.remove(SecretId(second)); harness.run(second).await.expect("no retry"); assert_eq!(harness.task(second).await.state, ProvisionerTaskState::Failed); @@ -1199,14 +1218,10 @@ async fn a_new_task_never_resumes_from_an_earlier_task() { 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; - harness - .run(first) - .await - .expect("the job queues the next attempt itself"); + harness.run(first).await.expect_err("the job queue runs the job again"); assert!(harness.workspace(session_id).join("chunk-0000.json").exists()); harness - .state - .provisioner_tasks + .tasks .store() .fail(first, json!({}), OffsetDateTime::now_utc()) .await @@ -1248,7 +1263,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.secrets.clone(), + harness.secrets.clone(), ); task.on_end(&record).await.expect("end"); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index d673bb7bec..72ac1a76c6 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -2,22 +2,20 @@ //! //! Each kind of work is a [`ProvisionerTask`]: what one attempt does, what its payloads hold, and what to do once it has //! ended. -//! [`ProvisionerTaskRunner`] does the rest for every task. It creates the record, queues the job, starts each attempt, -//! retries recoverable failures, moves the deadline ahead with each progress, cancels an attempt once the deadline passes, -//! and records how the task ended. +//! Its job is a [`TaskJob`], a job like any other on the job queue, which wraps each attempt in a prelude and a +//! completion: the prelude records the start of the attempt, and the completion records how it ended and ends the record. +//! A recoverable failure goes back to the job queue, which runs the same job again later for the next attempt. +//! [`ProvisionerTaskManager`] creates the records, reads them, and recovers them at startup. //! //! 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, and each job names the attempt it starts, so a stale, -//! duplicate or rerun 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 finished record keeps whether its end is done, so the same startup ends every record whose end was cut short, -//! whether or not a job for it survived. +//! Each record names the one job allowed to run it, by its token, so a stale job stops without running anything. +//! At startup, before any job runs, [`ProvisionerTaskManager::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::marker::PhantomData; use std::sync::Arc; use anyhow::Context as _; @@ -29,13 +27,12 @@ use provisioner_task::{ }; use serde::Serialize; use serde::de::DeserializeOwned; -use serde_json::json; use time::OffsetDateTime; use tokio_util::sync::CancellationToken; use uuid::Uuid; use super::{AttemptError, Permanent}; -use crate::job_queue::JobQueueHandle; +use crate::job_queue::DurableJobQueue; /// Time before a record whose end failed is ended again. const END_RETRY_DELAY: time::Duration = time::Duration::minutes(5); @@ -43,52 +40,26 @@ const END_RETRY_DELAY: time::Duration = time::Duration::minutes(5); /// Attempts a record gets: a recoverable failure of the last one fails the record. pub const MAX_ATTEMPTS: u32 = 5; -/// Time before the attempt that follows attempt number `attempts`: 30 seconds, doubled after each attempt. -fn retry_delay(attempts: u32) -> time::Duration { - time::Duration::seconds(30) * 2i32.saturating_pow(attempts.saturating_sub(1)) -} - -/// Where the runner queues the jobs of its records. -#[async_trait] -pub trait ProvisionerTaskJobQueue: Send + Sync { - /// Queues `job` to run at `schedule_for`, or at once without it, and returns once it is stored, so a record never - /// waits for a job that was lost. - async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()>; -} - -pub type DynProvisionerTaskJobQueue = Arc; - -/// Hands the job to the job queue task, which stores it later: for tests, which take the jobs from its channel. -#[async_trait] -impl ProvisionerTaskJobQueue for JobQueueHandle { - async fn push(&self, job: ProvisionerTaskJob, schedule_for: Option) -> anyhow::Result<()> { - match schedule_for { - Some(schedule_for) => self.schedule(job, schedule_for).await, - None => self.enqueue(job).await, - } - } -} - /// 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. + /// The kind written in the records of this task, and the name of its jobs. 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; + type Params: Serialize + DeserializeOwned + Send + 'static; /// The payload of a `succeeded` record. - type Output: Serialize + Send; + type Output: Serialize + Send + 'static; /// Why an attempt of this task fails for good, next to the failures any task may have. - type Failure: Serialize + Send; + type Failure: Serialize + Send + Sync + 'static; /// Why an attempt of this task fails for now, so a later attempt may succeed. - type Recoverable: Serialize + Send; + type Recoverable: Serialize + Send + Sync + 'static; /// Time from the creation of a record until its first attempt must start. fn start_window(&self) -> time::Duration; @@ -102,7 +73,11 @@ pub trait ProvisionerTask: Send + Sync + 'static { fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value; /// The payload of a record that failed with `error`, from `last`, its payload until then: `null` if it never ran. - fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value; + fn failed_payload( + &self, + error: &Permanent, + last: &serde_json::Value, + ) -> serde_json::Value; /// Why an unfinished record fails when Gateway starts again, for a task that cannot outlive a restart, such as one /// whose secrets are kept in memory only; `None` for a task that goes on. @@ -146,13 +121,7 @@ pub struct Attempt { latest: Arc>, } -/// An attempt the prelude started: the task that runs it, and the record as the attempt started. -struct Started { - task: Arc, - record: ProvisionerTaskRecord, -} - -/// What the runner knows of the running record: the last payload of the attempt, and the deadline it stored. +/// What the job 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, @@ -195,145 +164,39 @@ impl Attempt { } } -/// [`Permanent`] with the failures of its kind of task as JSON, as the runner holds every kind together. -pub type AnyPermanent = Permanent; - -/// [`AttemptError`] with the failures of its kind of task as JSON. -type AnyAttemptError = AttemptError; - -/// `error` with the failures of its kind as JSON; one that cannot be written as JSON is an internal failure. -fn erase_permanent(error: Permanent) -> AnyPermanent { - error - .try_map(serde_json::to_value, serde_json::to_value) - .unwrap_or_else(|error| { - error!(%error, "Failed to write a Task failure as JSON"); - Permanent::Internal - }) -} - -fn erase_attempt_error(error: AttemptError) -> AnyAttemptError { - match error { - AttemptError::Recoverable(error) => match serde_json::to_value(error) { - Ok(error) => AttemptError::Recoverable(error), - Err(error) => { - error!(%error, "Failed to write a Task failure as JSON"); - AttemptError::Permanent(Permanent::Internal) - } - }, - AttemptError::Permanent(error) => AttemptError::Permanent(erase_permanent(error)), - } +/// `value` as JSON for the logs, empty if it cannot be written. +fn log_json(value: &impl Serialize) -> serde_json::Value { + serde_json::to_value(value).unwrap_or_default() } -/// [`ProvisionerTask`] with its types erased, so the runner can hold every task together. -#[async_trait] -trait AnyProvisionerTask: Send + Sync { - fn start_window(&self) -> time::Duration; - - fn progress_timeout(&self) -> time::Duration; - - fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value; - - fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value; - - fn restart_failure(&self) -> Option; - - async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result; - - async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()>; +/// Creates, reads and recovers the records of the task `T`, whose jobs are [`TaskJob`]. +pub struct ProvisionerTaskManager { + inner: Arc>, } -#[async_trait] -impl AnyProvisionerTask for T { - fn start_window(&self) -> time::Duration { - ProvisionerTask::start_window(self) - } - - fn progress_timeout(&self) -> time::Duration { - ProvisionerTask::progress_timeout(self) - } - - fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value { - ProvisionerTask::starting_payload(self, last) - } - - fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value { - ProvisionerTask::failed_payload(self, error, last) - } - - fn restart_failure(&self) -> Option { - ProvisionerTask::restart_failure(self) - .map(|failure| erase_permanent(Permanent::<_, T::Recoverable>::Kind(failure))) - } - - async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { - let params = serde_json::from_value(params).map_err(|error| { - warn!(task.id = %attempt.task_id, %error, "Invalid Task parameters"); - AnyAttemptError::Permanent(Permanent::InvalidTask) - })?; - - let output = ProvisionerTask::run(self, attempt, params) - .await - .map_err(erase_attempt_error)?; - - serde_json::to_value(output).map_err(|error| { - error!(task.id = %attempt.task_id, %error, "Failed to serialize the Task result"); - AnyAttemptError::Permanent(Permanent::Internal) - }) - } - - async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { - ProvisionerTask::on_end(self, record).await +impl Clone for ProvisionerTaskManager { + fn clone(&self) -> Self { + Self { + inner: Arc::clone(&self.inner), + } } } -/// Creates provisioner task records and runs them as jobs on the job queue. -#[derive(Clone)] -pub struct ProvisionerTaskRunner { - inner: Arc, -} - -struct Inner { +struct Inner { + task: T, store: DynProvisionerTaskStore, - job_queue: DynProvisionerTaskJobQueue, - /// Kind → the task that runs records of this kind. - tasks: HashMap<&'static str, Arc>, + job_queue: DurableJobQueue, } -pub struct ProvisionerTaskRunnerBuilder { - store: DynProvisionerTaskStore, - job_queue: DynProvisionerTaskJobQueue, - tasks: HashMap<&'static str, Arc>, -} - -impl ProvisionerTaskRunnerBuilder { - /// Registers `task`, so records of its kind can be created and run. - #[must_use] - pub fn register(mut self, task: T) -> Self { - self.tasks.insert(T::KIND, Arc::new(task)); - self - } - - pub fn build(self) -> ProvisionerTaskRunner { - ProvisionerTaskRunner { - inner: Arc::new(Inner { - store: self.store, - job_queue: self.job_queue, - tasks: self.tasks, - }), +impl ProvisionerTaskManager { + pub fn new(task: T, store: DynProvisionerTaskStore, job_queue: DurableJobQueue) -> Self { + Self { + inner: Arc::new(Inner { task, store, job_queue }), } } -} -impl ProvisionerTaskRunner { - pub fn builder( - store: DynProvisionerTaskStore, - job_queue: DynProvisionerTaskJobQueue, - ) -> ProvisionerTaskRunnerBuilder { - ProvisionerTaskRunnerBuilder { - store, - job_queue, - tasks: HashMap::new(), - } + pub(crate) fn task(&self) -> &T { + &self.inner.task } /// Returns the record with this ID; one past its deadline is failed first. @@ -341,19 +204,18 @@ impl ProvisionerTaskRunner { self.inner.store.get(id, now).await } - /// Creates a record for the task `T` and queues the job that runs it. + /// Creates a record 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( + 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(); @@ -367,7 +229,7 @@ impl ProvisionerTaskRunner { kind: T::KIND.to_owned(), target, params, - deadline_at: deadline(now, task.start_window()), + deadline_at: deadline(now, self.inner.task.start_window()), job_token, }, now, @@ -379,14 +241,13 @@ impl ProvisionerTaskRunner { }; let queued = match on_created(record) { - Ok(()) => self.enqueue(id, job_token, Step::Attempt(1), None).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, job_token, &serde_json::Value::Null) - .await; + self.fail_not_queued(id, job_token, &serde_json::Value::Null).await; return Err(error); } @@ -395,20 +256,25 @@ impl ProvisionerTaskRunner { Ok(outcome) } - /// Hands every unfinished record to a new job, since a record created just before Gateway stopped may have none, - /// then queues the end of every finished record whose end is not done, since Gateway may have stopped before it. + /// Hands every unfinished record of `T` to a new job, since a record created just before Gateway stopped may have + /// none. /// - /// Each unfinished record gets a new `job_token`, so a job left from before finds itself superseded and stops, and - /// the job queue is never read. - /// A record whose task cannot outlive a restart fails instead, and is ended with the other finished records. + /// Each record gets a new `job_token`, so a job left from before finds itself superseded and stops, and the job + /// queue is never read. + /// The new job is stored before a record whose task cannot outlive a restart is failed, so a crash in between never + /// leaves a finished record without a job to end it. /// A record that was waiting to retry runs at once, since retry timing is no promise to the caller. + /// A record whose new job cannot be stored is failed, since no job is left to run 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 store = &self.inner.store; + let Inner { task, store, .. } = self.inner.as_ref(); let now = OffsetDateTime::now_utc(); for record in store.list_unfinished(now).await? { + if record.kind != T::KIND { + continue; + } + let job_token = Uuid::new_v4(); match store.replace_job_token(record.id, job_token, now).await { @@ -425,61 +291,40 @@ impl ProvisionerTaskRunner { } } - let task = self.task(&record.kind); - - if let Some(task) = &task - && let Some(failure) = task.restart_failure() - { - match store - .fail(record.id, task.failed_payload(&failure, &record.payload), now) - .await - { - Ok(_) => { - info!( - task.id = %record.id, - task.kind = %record.kind, - ?failure, - "Task failed, since it cannot outlive a restart" - ); - continue; - } - // Its new job runs the attempt then, which fails the same way. - Err(error) => error!( - task.id = %record.id, - task.kind = %record.kind, - error = format!("{error:#}"), - "Failed to fail the Task that cannot outlive a restart" - ), - } + if let Err(error) = self.enqueue(record.id, job_token, None).await { + error!( + task.id = %record.id, + task.kind = %record.kind, + error = format!("{error:#}"), + "Failed to queue the Task again" + ); + self.fail_not_queued(record.id, job_token, &record.payload).await; + continue; } - let next = Step::Attempt(record.attempts.saturating_add(1)); + info!(task.id = %record.id, task.kind = %record.kind, "Task queued again"); - match self.enqueue(record.id, job_token, next, None).await { - Ok(()) => info!(task.id = %record.id, task.kind = %record.kind, "Task queued again"), - Err(error) => { - error!( - task.id = %record.id, - task.kind = %record.kind, - error = format!("{error:#}"), - "Failed to queue the Task again" - ); - if let Some(task) = &task { - self.fail_not_queued(task.as_ref(), record.id, job_token, &record.payload) - .await; - } - } - } - } + let Some(failure) = task.restart_failure() else { + continue; + }; + let failure = Permanent::Kind(failure); - for record in store.list_unended().await? { - match self.enqueue(record.id, record.job_token, Step::End, None).await { - Ok(()) => info!(task.id = %record.id, task.kind = %record.kind, "Task end queued again"), + match store + .fail(record.id, task.failed_payload(&failure, &record.payload), now) + .await + { + Ok(_) => info!( + task.id = %record.id, + task.kind = %record.kind, + failure = %log_json(&failure), + "Task failed, since it cannot outlive a restart" + ), + // Its new job runs the attempt then, which fails the same way. Err(error) => error!( task.id = %record.id, task.kind = %record.kind, error = format!("{error:#}"), - "Failed to queue the end of the Task again" + "Failed to fail the Task that cannot outlive a restart" ), } } @@ -492,135 +337,244 @@ impl ProvisionerTaskRunner { &self.inner.store } - /// The task that runs records of `kind`. - fn task(&self, kind: &str) -> Option> { - self.inner.tasks.get(kind).cloned() - } - + /// Stores the job holding `job_token`, which wakes the record up at `schedule_for`, or at once without it. async fn enqueue( &self, task_id: Uuid, job_token: Uuid, - step: Step, schedule_for: Option, + ) -> anyhow::Result<()> { + let job = TaskJob { + task_id, + job_token, + manager: self.clone(), + }; + + self.inner.job_queue.push(job, schedule_for).await + } + + /// 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. + async fn fail( + &self, + task_id: Uuid, + job_token: Uuid, + error: &Permanent, + last: &serde_json::Value, ) -> anyhow::Result<()> { self.inner - .job_queue - .push( - ProvisionerTaskJob { - task_id, - job_token, - step, - runner: self.clone(), - }, - schedule_for, + .store + .fail( + task_id, + self.inner.task.failed_payload(error, last), + OffsetDateTime::now_utc(), ) .await + .context("record the Task failure")?; + + self.end(task_id, job_token).await } - /// Runs attempt number `attempt` for the job holding `job_token`, in three steps. - /// - /// The prelude marks the start of the attempt on the record, the body runs the work of the task, and the completion - /// records how it ended, then ends the record or queues its next attempt. - /// The record counts its attempts, and the job queues the next one itself, so the job queue runs a job again only - /// when what it had to record or queue could not be stored; the prelude then finds its attempt already started. - async fn run_attempt(&self, task_id: Uuid, job_token: Uuid, attempt: u32) -> anyhow::Result<()> { - let Some(started) = self.begin_attempt(task_id, job_token, attempt).await? else { + /// Fails a record no job will run. + async fn fail_not_queued(&self, task_id: Uuid, job_token: Uuid, last: &serde_json::Value) { + if let Err(error) = self.fail(task_id, job_token, &Permanent::NotQueued, last).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to end the Task that has no job"); + } + } + + /// Ends the record, read again so its task sees how it ended. + async fn end(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { + match self.inner.store.get(task_id, OffsetDateTime::now_utc()).await { + Ok(Some(record)) => self.end_record(&record, job_token).await, + Ok(None) => Ok(()), + Err(error) => self.end_later(task_id, job_token, error).await, + } + } + + /// Does what the finished record leaves to do; a record still unfinished is left to its attempt. + async fn end_record(&self, record: &ProvisionerTaskRecord, job_token: Uuid) -> anyhow::Result<()> { + if !record.state.is_finished() { return Ok(()); - }; + } + + match self.inner.task.on_end(record).await { + Ok(()) => Ok(()), + Err(error) => self.end_later(record.id, job_token, error).await, + } + } + + /// Wakes the record up again later, to end it then; an error means that job could not be stored. + 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 manager, the job and the store +/// agree on when it passes. +fn deadline(now: OffsetDateTime, timeout: time::Duration) -> OffsetDateTime { + (now + timeout).truncate_to_second() +} + +#[derive(Serialize, serde::Deserialize)] +#[serde(rename_all = "camelCase")] +struct JobDef { + task_id: Uuid, + job_token: Uuid, +} + +/// The job of a record of the task `T`: it wakes the record up, and the record decides what runs. +/// +/// It runs one attempt of `T` in three steps: the prelude records its start, the body runs [`ProvisionerTask::run`], +/// and the completion records how it ended. A record that has ended is ended instead. +pub struct TaskJob { + task_id: Uuid, + job_token: Uuid, + manager: ProvisionerTaskManager, +} - let (outcome, latest) = self.execute(&started).await; +/// A [`TaskJob`] read back from its JSON, before it gets the manager of `T`. +pub struct TaskJobUnhydrated { + task_id: Uuid, + job_token: Uuid, + task: PhantomData, +} - self.complete(&started, job_token, outcome, latest).await +impl TaskJobUnhydrated { + pub fn into_job(self, manager: ProvisionerTaskManager) -> TaskJob { + TaskJob { + task_id: self.task_id, + job_token: self.job_token, + manager, + } + } +} + +impl TaskJob { + /// The name its jobs are stored under. + pub const NAME: &'static str = T::KIND; + + pub fn read_json(json: &str) -> anyhow::Result> { + let JobDef { task_id, job_token }: JobDef = + serde_json::from_str(json).context("failed to deserialize TaskJob")?; + + Ok(TaskJobUnhydrated { + task_id, + job_token, + task: PhantomData, + }) } - /// The prelude: marks the start of attempt number `attempt` on the record, unless another job holds it, it has - /// ended, or the attempt already started. + /// The prelude: marks the start of one more attempt on the record, unless another job holds it, it has ended, or + /// it has no attempt left. /// - /// A record that has ended is ended instead, if its end is not done yet. - /// An attempt that already started is never run again, since that could do paid work twice: its record is ended - /// once its deadline passes, unless the attempt is still running and ends it first. - async fn begin_attempt(&self, task_id: Uuid, job_token: Uuid, attempt: u32) -> anyhow::Result> { - let store = &self.inner.store; + /// A record that has ended is ended instead. + async fn prelude(&self) -> anyhow::Result> { + let Inner { task, store, .. } = self.manager.inner.as_ref(); let now = OffsetDateTime::now_utc(); // The starting payload carries over what earlier attempts reported, so the record is read first. - let Some(record) = store.get(task_id, now).await? else { - debug!(task.id = %task_id, "Task is unknown, so this attempt does not run"); + let Some(record) = store.get(self.task_id, now).await? else { + debug!(task.id = %self.task_id, "Task is unknown, so this attempt does not run"); return Ok(None); }; - 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": AnyPermanent::InvalidTask }), now) - .await - .context("record the Task failure")?; - } + if record.kind != T::KIND { + debug!(task.id = %self.task_id, task.kind = %record.kind, "Task of another kind, so this attempt does not run"); return Ok(None); - }; + } + + if record.job_token != self.job_token { + debug!(task.id = %self.task_id, "Another job runs this Task, so this one stops"); + return Ok(None); + } + + if record.state.is_finished() { + debug!(task.id = %self.task_id, "Task is over, so this attempt does not run"); + self.manager.end_record(&record, self.job_token).await?; + return Ok(None); + } + + // A job run again after its last attempt could not record how it ended would otherwise start one past the last. + if MAX_ATTEMPTS <= record.attempts { + warn!( + task.id = %self.task_id, + task.kind = %record.kind, + attempts = record.attempts, + "Task has no attempt left, so it fails without running another" + ); + self.fail_terminal(&Permanent::Internal, &record.payload, record.deadline_at) + .await?; + return Ok(None); + } let starting_payload = task.starting_payload(&record.payload); let deadline_at = deadline(now, task.progress_timeout()); match store - .start_attempt(task_id, job_token, attempt, starting_payload, deadline_at, now) + .start_attempt(self.task_id, self.job_token, starting_payload, deadline_at, now) .await? { - AttemptStart::Started(record) => Ok(Some(Started { task, record })), + AttemptStart::Started(record) => Ok(Some(record)), AttemptStart::Finished(record) => { - debug!(task.id = %task_id, "Task is over, so this attempt does not run"); - self.end_record(&record, job_token).await?; - Ok(None) - } - AttemptStart::AlreadyStarted(record) => { - warn!( - task.id = %task_id, - task.kind = %record.kind, - attempt, - attempts = record.attempts, - "Task attempt already started, so it does not run again" - ); - - // The store fails the record once its deadline passes, so a job that runs after it only ends it. - let after_deadline = record.deadline_at + time::Duration::SECOND; - - self.enqueue(task_id, job_token, Step::End, Some(after_deadline)) - .await - .context("queue the end of the Task")?; + debug!(task.id = %self.task_id, "Task is over, so this attempt does not run"); + self.manager.end_record(&record, self.job_token).await?; Ok(None) } AttemptStart::Superseded => { - debug!(task.id = %task_id, "Another job runs this Task, so this one stops"); + debug!(task.id = %self.task_id, "Another job runs this Task, so this one stops"); Ok(None) } AttemptStart::Unknown => { - debug!(task.id = %task_id, "Task is unknown, so this attempt does not run"); + debug!(task.id = %self.task_id, "Task is unknown, so this attempt does not run"); Ok(None) } } } - /// The body: runs the attempt of the task, and cancels it once the record is past its deadline. + /// The body: runs the attempt of `T`, and cancels it once the record is past its deadline. /// /// Returns how the attempt ended, with the last payload it reported and the deadline it last set. - async fn execute(&self, started: &Started) -> (Result, Latest) { - let Started { task, record } = started; - let cancel = CancellationToken::new(); + async fn body( + &self, + record: &ProvisionerTaskRecord, + ) -> (Result>, Latest) { let latest = Arc::new(Mutex::new(Latest { payload: record.payload.clone(), deadline_at: record.deadline_at, })); + let take_latest = |latest: &Mutex| { + let latest = latest.lock(); + Latest { + payload: latest.payload.clone(), + deadline_at: latest.deadline_at, + } + }; + + let params = match serde_json::from_value::(record.params.clone()) { + Ok(params) => params, + Err(error) => { + warn!(task.id = %record.id, %error, "Invalid Task parameters"); + return ( + Err(AttemptError::Permanent(Permanent::InvalidTask)), + take_latest(&latest), + ); + } + }; + + let cancel = CancellationToken::new(); let attempt = Attempt { task_id: record.id, target: record.target.clone(), attempts: record.attempts, payload: record.payload.clone(), cancel: cancel.clone(), - progress_timeout: task.progress_timeout(), - store: Arc::clone(&self.inner.store), + progress_timeout: self.manager.inner.task.progress_timeout(), + store: Arc::clone(&self.manager.inner.store), latest: Arc::clone(&latest), }; @@ -628,9 +582,8 @@ impl ProvisionerTaskRunner { // 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 manager = self.manager.clone(); + async move { manager.inner.task.run(&attempt, params).await } }); let mut past_deadline = false; @@ -668,231 +621,137 @@ impl ProvisionerTaskRunner { Err(_) => Err(AttemptError::Permanent(Permanent::Panicked)), }; - let latest = latest.lock(); - - ( - outcome, - Latest { - payload: latest.payload.clone(), - deadline_at: latest.deadline_at, - }, - ) + (outcome, take_latest(&latest)) } - /// The completion: records how the attempt ended, then ends the record or queues its next attempt. + /// The completion: records how the attempt ended, then ends the record, or returns a recoverable failure to the job + /// queue, which runs this job again later for the next attempt. /// - /// If that cannot be stored, the error is returned so the job queue runs the job again, and the record fails at its - /// deadline, see [`ProvisionerTaskRunner::begin_attempt`]. + /// An end that cannot be recorded never goes back to the job queue, since that would run the attempt again, see + /// [`TaskJob::fail_terminal`]. /// Recording the end and dropping what the record kept are separate steps, so a failed drop never undoes the end. - async fn complete( + async fn completion( &self, - started: &Started, - job_token: Uuid, - outcome: Result, + record: &ProvisionerTaskRecord, + outcome: Result>, latest: Latest, ) -> anyhow::Result<()> { - let Started { task, record } = started; + let store = &self.manager.inner.store; let (task_id, attempts) = (record.id, record.attempts); let last = &latest.payload; match outcome { Ok(output) => { - let ended = self - .inner - .store - .succeed(task_id, output, OffsetDateTime::now_utc()) - .await - .context("record the Task success")?; + let output = match serde_json::to_value(output) { + Ok(output) => output, + Err(error) => { + error!(task.id = %task_id, %error, "Failed to serialize the Task result"); + return self.fail_terminal(&Permanent::Internal, last, latest.deadline_at).await; + } + }; - if ended { - info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"); - } else { - warn!( + 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" - ); + ), + // A retry would do the work again, such as paying for a second analysis. + Err(error) => { + error!( + task.id = %task_id, + error = format!("{error:#}"), + "Failed to record the Task success; the Task fails instead" + ); + return self.fail_terminal(&Permanent::Internal, last, latest.deadline_at).await; + } } - self.end(task_id, job_token).await?; + self.manager.end(task_id, self.job_token).await } Err(AttemptError::Recoverable(error)) if attempts < MAX_ATTEMPTS => { - let retry_at = OffsetDateTime::now_utc() + retry_delay(attempts); + let error = log_json(&error); warn!( task.id = %task_id, task.kind = %record.kind, attempts, %error, - %retry_at, - "Task attempt failed; it will be retried" + "Task attempt failed; the job queue runs it again later" ); - self.enqueue(task_id, job_token, Step::Attempt(attempts + 1), Some(retry_at)) - .await - .context("queue the next attempt")?; + anyhow::bail!("Task attempt {attempts} failed with {error}, to be retried") } Err(AttemptError::Recoverable(error)) => { warn!( task.id = %task_id, task.kind = %record.kind, attempts, - %error, + error = %log_json(&error), "Task failed after its last attempt" ); - self.fail( - task.as_ref(), - task_id, - job_token, - &Permanent::AttemptsExhausted { last: error }, - last, - ) - .await?; + self.fail_terminal(&Permanent::AttemptsExhausted { last: error }, last, latest.deadline_at) + .await } Err(AttemptError::Permanent(error)) => { warn!( task.id = %task_id, task.kind = %record.kind, attempts, - error = %serde_json::to_value(&error).unwrap_or_default(), + error = %log_json(&error), "Task failed" ); - self.fail(task.as_ref(), task_id, job_token, &error, last).await?; + self.fail_terminal(&error, last, latest.deadline_at).await } } - - Ok(()) } - /// Records the failure, then ends the record. + /// Records that the record failed with `failure`, then ends it. /// - /// If the failure cannot be recorded, nothing is dropped, since the record may still run, and the error is returned. - async fn fail( + /// If the failure cannot be recorded, no error goes back to the job queue, since it would run the attempt again + /// before the deadline: a job is stored to wake the record up past `deadline_at` instead, once the store has failed + /// it and only its end is left. + async fn fail_terminal( &self, - task: &dyn AnyProvisionerTask, - task_id: Uuid, - job_token: Uuid, - error: &AnyPermanent, + failure: &Permanent, last: &serde_json::Value, + deadline_at: OffsetDateTime, ) -> anyhow::Result<()> { - self.inner - .store - .fail(task_id, task.failed_payload(error, last), OffsetDateTime::now_utc()) - .await - .context("record the Task failure")?; + let Inner { task, store, .. } = self.manager.inner.as_ref(); - self.end(task_id, job_token).await - } + let recorded = store + .fail( + self.task_id, + task.failed_payload(failure, last), + OffsetDateTime::now_utc(), + ) + .await; - /// Fails a record no job will run. - async fn fail_not_queued( - &self, - task: &dyn AnyProvisionerTask, - task_id: Uuid, - job_token: Uuid, - last: &serde_json::Value, - ) { - if let Err(error) = self.fail(task, task_id, job_token, &Permanent::NotQueued, last).await { - error!(task.id = %task_id, error = format!("{error:#}"), "Failed to end the Task that has no job"); - } - } + if let Err(error) = recorded { + error!( + task.id = %self.task_id, + error = format!("{error:#}"), + "Failed to fail the Task; it fails at its deadline" + ); - /// Ends the record, read again so its task sees how it ended. - async fn end(&self, 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(&record, job_token).await, - Ok(None) => Ok(()), - Err(error) => self.end_later(task_id, job_token, error).await, - } - } + let after_deadline = deadline_at + time::Duration::SECOND; - /// Does what the finished record leaves to do, then records that it is done. - /// - /// A record still unfinished is left to its attempt, and one of a kind this Gateway does not know is left for one - /// that does. - async fn end_record(&self, record: &ProvisionerTaskRecord, job_token: Uuid) -> anyhow::Result<()> { - if !record.state.is_finished() || record.ended_at.is_some() { - return Ok(()); - } + if let Err(error) = self + .manager + .enqueue(self.task_id, self.job_token, Some(after_deadline)) + .await + { + error!(task.id = %self.task_id, error = format!("{error:#}"), "Failed to queue the end of the Task"); + } - let Some(task) = self.task(&record.kind) else { return Ok(()); - }; - - if let Err(error) = task.on_end(record).await { - return self.end_later(record.id, job_token, error).await; - } - - match self.inner.store.mark_ended(record.id, OffsetDateTime::now_utc()).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. - /// - /// The end stays due on the record, so the next startup ends it if that job is lost too. - 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, - Step::End, - 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() -} - -/// What a job comes to do for its record. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, serde::Deserialize)] -#[serde(rename_all = "camelCase")] -enum Step { - /// Start the attempt with this number, counting from 1, unless it already started. - Attempt(u32), - /// End the finished record, unless that is already done. - End, -} -#[derive(Serialize, serde::Deserialize)] -#[serde(rename_all = "camelCase")] -struct JobDef { - task_id: Uuid, - job_token: Uuid, - step: Step, -} - -/// 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, - step: Step, - 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, - step: def.step, - runner, - }) + self.manager.end(self.task_id, self.job_token).await } } #[async_trait] -impl job_queue::Job for ProvisionerTaskJob { +impl job_queue::Job for TaskJob { fn name(&self) -> &str { Self::NAME } @@ -901,15 +760,17 @@ impl job_queue::Job for ProvisionerTaskJob { serde_json::to_string(&JobDef { task_id: self.task_id, job_token: self.job_token, - step: self.step, }) - .context("failed to serialize ProvisionerTaskJob") + .context("failed to serialize TaskJob") } async fn run(&mut self) -> anyhow::Result<()> { - match self.step { - Step::Attempt(attempt) => self.runner.run_attempt(self.task_id, self.job_token, attempt).await, - Step::End => self.runner.end(self.task_id, self.job_token).await, - } + let Some(record) = self.prelude().await? else { + return Ok(()); + }; + + let (outcome, latest) = self.body(&record).await; + + self.completion(&record, outcome, latest).await } } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index e62689bc9b..c2c6c863f0 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -3,9 +3,10 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use job_queue::Job as _; use parking_lot::Mutex; use provisioner_task::{ProvisionerTaskState, ProvisionerTaskStore}; +use serde_json::json; use super::*; -use crate::job_queue::{JobMessage, JobQueueReceiver}; +use crate::job_queue::{StoredJob, TestJobQueue}; const KIND: &str = "test.echo"; @@ -73,7 +74,11 @@ impl ProvisionerTask for Echo { json!({ "after": last }) } - fn failed_payload(&self, error: &AnyPermanent, last: &serde_json::Value) -> serde_json::Value { + fn failed_payload( + &self, + error: &Permanent, + last: &serde_json::Value, + ) -> serde_json::Value { json!({ "error": error, "last": last }) } @@ -148,14 +153,11 @@ impl ProvisionerTaskStore for FailingStore { &self, id: Uuid, job_token: Uuid, - attempt: u32, payload: serde_json::Value, deadline_at: OffsetDateTime, now: OffsetDateTime, ) -> anyhow::Result { - self.inner - .start_attempt(id, job_token, attempt, payload, deadline_at, now) - .await + self.inner.start_attempt(id, job_token, payload, deadline_at, now).await } async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result { @@ -166,10 +168,6 @@ impl ProvisionerTaskStore for FailingStore { self.inner.list_unfinished(now).await } - async fn list_unended(&self) -> anyhow::Result> { - self.inner.list_unended().await - } - async fn update_running( &self, id: Uuid, @@ -195,15 +193,11 @@ impl ProvisionerTaskStore for FailingStore { self.inner.fail(id, payload, now).await } - - async fn mark_ended(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result { - self.inner.mark_ended(id, now).await - } } struct Harness { - runner: ProvisionerTaskRunner, - job_queue_rx: JobQueueReceiver, + manager: ProvisionerTaskManager, + job_queue: Arc, runs: Arc, ended: Arc>>, end_failures: Arc, @@ -243,25 +237,27 @@ impl Harness { success_failures: Arc::clone(&success_failures), failure_failures: Arc::clone(&failure_failures), }); - let (job_queue, job_queue_rx) = JobQueueHandle::new(); + let job_queue = TestJobQueue::new().await; 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 { + let manager = ProvisionerTaskManager::new( + Echo { outcome, progress_timeout, runs: Arc::clone(&runs), ended: Arc::clone(&ended), end_failures: Arc::clone(&end_failures), restart_failure, - }) - .build(); + }, + store, + job_queue.durable(), + ); Self { - runner, - job_queue_rx, + manager, + job_queue, runs, ended, end_failures, @@ -273,26 +269,41 @@ impl Harness { async fn create(&self, value: u32) -> Uuid { let id = Uuid::new_v4(); let outcome = self - .runner - .create::(id, format!("target-{id}"), &EchoParams { value }, |_| Ok(())) + .manager + .create(id, format!("target-{id}"), &EchoParams { value }, |_| Ok(())) .await .expect("create"); assert!(matches!(outcome, CreateOutcome::Created(_)), "{outcome:?}"); id } + /// Takes the oldest queued job, read back as the job runner reads it, with when it runs. + fn queued(&mut self) -> Option<(Box, Option)> { + let StoredJob { + name, + json, + schedule_for, + } = self.job_queue.take()?; + assert_eq!(name, KIND); + let job: Box = Box::new( + TaskJob::::read_json(&json) + .expect("valid job") + .into_job(self.manager.clone()), + ); + Some((job, schedule_for)) + } + fn queued_job(&mut self) -> Box { - let JobMessage { job, .. } = self.job_queue_rx.try_recv().expect("a queued job"); - job + self.queued().expect("a queued job").0 } /// Takes every queued job. - fn queued_jobs(&mut self) -> Vec { - std::iter::from_fn(|| self.job_queue_rx.try_recv().ok()).collect() + fn queued_jobs(&mut self) -> Vec { + std::iter::from_fn(|| self.job_queue.take()).collect() } async fn task(&self, id: Uuid) -> ProvisionerTaskRecord { - self.runner + self.manager .get(id, OffsetDateTime::now_utc()) .await .expect("read task") @@ -308,8 +319,8 @@ async fn only_a_new_task_runs_on_created_and_gets_a_job() { let created = AtomicUsize::new(0); let create = || { harness - .runner - .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + .manager + .create(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { created.fetch_add(1, Ordering::SeqCst); Ok(()) }) @@ -320,19 +331,19 @@ async fn only_a_new_task_runs_on_created_and_gets_a_job() { assert_eq!(created.load(Ordering::SeqCst), 1); harness.queued_job(); - assert!(harness.job_queue_rx.try_recv().is_err(), "one job only"); + assert!(harness.job_queue.take().is_none(), "one job only"); } // oracle: mutant a job that cannot be queued is ignored, so the Task stays queued and its creation succeeds #[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 harness = Harness::new(|value| Ok(json!(value))).await; + harness.job_queue.close(); let id = Uuid::new_v4(); harness - .runner - .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| Ok(())) + .manager + .create(id, "target".to_owned(), &EchoParams { value: 1 }, |_| Ok(())) .await .expect_err("the job queue is gone"); @@ -345,12 +356,12 @@ async fn a_task_whose_job_cannot_be_queued_fails() { // oracle: mutant the job is queued before the Task is prepared, so a Task that cannot be prepared still gets a job #[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 harness = Harness::new(|value| Ok(json!(value))).await; let id = Uuid::new_v4(); harness - .runner - .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + .manager + .create(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { anyhow::bail!("no room for the secret") }) .await @@ -359,7 +370,7 @@ 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["error"], json!({ "kind": "notQueued" })); - assert!(harness.job_queue_rx.try_recv().is_err(), "no job is queued"); + assert!(harness.job_queue.take().is_none(), "no job is queued"); assert_eq!(*harness.ended.lock(), [id]); } @@ -370,10 +381,10 @@ async fn a_superseded_job_runs_nothing_once_the_task_is_resumed() { let id = harness.create(7).await; let mut stale = harness.queued_job(); - harness.runner.resume().await.expect("resume"); + harness.manager.resume().await.expect("resume"); let mut resumed = harness.queued_job(); assert!( - harness.job_queue_rx.try_recv().is_err(), + harness.job_queue.take().is_none(), "the unfinished Task gets one new job" ); @@ -388,11 +399,8 @@ async fn a_superseded_job_runs_nothing_once_the_task_is_resumed() { assert_eq!(task.payload, json!({ "echo": 7 })); assert_eq!(*harness.ended.lock(), [id]); - harness.runner.resume().await.expect("resume again"); - assert!( - harness.job_queue_rx.try_recv().is_err(), - "a finished Task gets no new job" - ); + harness.manager.resume().await.expect("resume again"); + assert!(harness.job_queue.take().is_none(), "a finished Task gets no new job"); } // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management @@ -401,9 +409,9 @@ 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.job_queue.close(); - harness.runner.resume().await.expect("resume"); + harness.manager.resume().await.expect("resume"); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed, "no job is left to run it"); @@ -421,10 +429,10 @@ async fn resume_gives_every_unfinished_task_a_new_job() { harness.task(first).await.job_token, harness.task(second).await.job_token, ]; - // The jobs of the creation may be stored or lost; the runner never asks which. + // The jobs of the creation may be stored or lost; the manager never asks which. harness.queued_jobs(); - harness.runner.resume().await.expect("resume"); + harness.manager.resume().await.expect("resume"); assert_eq!(harness.queued_jobs().len(), 2, "one new job for each unfinished Task"); assert_ne!(harness.task(first).await.job_token, tokens[0]); @@ -443,7 +451,7 @@ async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { let id = harness.create(1).await; let mut stale = harness.queued_job(); - harness.runner.resume().await.expect("resume"); + harness.manager.resume().await.expect("resume"); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed, "it fails before any job runs"); @@ -458,31 +466,6 @@ async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { assert_eq!(*harness.ended.lock(), [id]); } -// oracle: mutant resume looks at unfinished Tasks only, so an end cut short by a crash is never done -#[tokio::test] -async fn resume_ends_a_finished_task_whose_end_was_cut_short() { - let mut harness = Harness::new(|value| Ok(json!(value))).await; - let id = harness.create(1).await; - harness.queued_jobs(); - // As if Gateway stopped right after the failure was recorded, before the Task was ended. - harness - .runner - .store() - .fail(id, json!({ "error": { "kind": "gone" } }), OffsetDateTime::now_utc()) - .await - .expect("fail"); - - harness.runner.resume().await.expect("resume"); - let mut ending = harness.queued_job(); - assert!(harness.job_queue_rx.try_recv().is_err(), "one job ends it"); - ending.run().await.expect("the job ends the Task"); - - assert_eq!(harness.runs.load(Ordering::SeqCst), 0, "no attempt runs"); - assert_eq!(*harness.ended.lock(), [id]); - harness.runner.resume().await.expect("resume again"); - assert!(harness.job_queue_rx.try_recv().is_err(), "an ended Task gets no job"); -} - // oracle: mutant a Task whose failure at startup cannot be recorded gets no new job, so nothing ends it #[tokio::test] async fn a_task_whose_restart_failure_cannot_be_recorded_still_gets_a_new_job() { @@ -497,7 +480,7 @@ async fn a_task_whose_restart_failure_cannot_be_recorded_still_gets_a_new_job() harness.queued_jobs(); harness.failure_failures.store(1, Ordering::SeqCst); - harness.runner.resume().await.expect("resume"); + harness.manager.resume().await.expect("resume"); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Queued, "the failure was not recorded"); @@ -514,15 +497,14 @@ async fn resume_keeps_the_attempts_a_task_made() { .queued_job() .run() .await - .expect("the job queues the next attempt itself"); - harness.queued_jobs(); + .expect_err("the job queue runs the job again"); - harness.runner.resume().await.expect("resume"); + harness.manager.resume().await.expect("resume"); harness .queued_job() .run() .await - .expect("the job queues the next attempt itself"); + .expect_err("the job queue runs the job again"); assert_eq!(harness.task(id).await.attempts, 2); } @@ -567,22 +549,17 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { let mut job = harness.queued_job(); for attempt in 1..MAX_ATTEMPTS { - job.run().await.expect("the job queues the next attempt itself"); - assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Running, "{attempt}"); + job.run().await.expect_err("the job queue runs the job again"); - let JobMessage { - job: next, - schedule_for, - } = harness.job_queue_rx.try_recv().expect("the next attempt is queued"); - // oracle: mutant a recoverable failure is retried at once, without the delay a busy provider asks for + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Running, "{attempt}"); assert!( - schedule_for.is_some_and(|at| OffsetDateTime::now_utc() < at), - "attempt {attempt} is retried later" + harness.job_queue.take().is_none(), + "the job queue runs the same job again" ); - job = next; } job.run().await.expect("the last attempt"); - assert!(harness.job_queue_rx.try_recv().is_err(), "no attempt after the last"); + assert!(harness.job_queue.take().is_none(), "no attempt after the last"); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); @@ -595,104 +572,137 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { // oracle: mutant a success that cannot be recorded is retried, so the work is done and paid for twice #[tokio::test] -async fn a_success_that_cannot_be_recorded_is_not_done_again_and_ends_after_the_deadline() { +async fn a_success_that_cannot_be_recorded_fails_the_task_instead_of_running_again() { let mut harness = Harness::new(|value| Ok(json!(value))).await; let id = harness.create(1).await; - let mut job = harness.queued_job(); harness.success_failures.store(1, Ordering::SeqCst); - job.run().await.expect_err("the job queue runs the job again"); - job.run().await.expect("its attempt already started"); + harness.queued_job().run().await.expect("no retry"); let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Running, "nothing could be recorded"); + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "internal" })); assert_eq!(harness.runs.load(Ordering::SeqCst), 1, "the work is not done again"); + assert_eq!(*harness.ended.lock(), [id]); + assert!(harness.job_queue.take().is_none(), "no job is left"); +} + +// oracle: mutant an end that cannot be recorded goes back to the job queue, which runs the attempt again +#[tokio::test] +async fn an_end_that_cannot_be_recorded_is_left_to_a_job_past_the_deadline() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + harness.success_failures.store(1, Ordering::SeqCst); + harness.failure_failures.store(1, Ordering::SeqCst); + + harness + .queued_job() + .run() + .await + .expect("nothing goes back to the job queue"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Running, "nothing could be recorded"); + assert_eq!(harness.runs.load(Ordering::SeqCst), 1); assert!( harness.ended.lock().is_empty(), "nothing is dropped while the Task may still run" ); - let JobMessage { schedule_for, .. } = harness.job_queue_rx.try_recv().expect("a job ends it later"); + let (_, schedule_for) = harness.queued().expect("a job ends it later"); assert!( schedule_for.is_some_and(|at| task.deadline_at < at), "the job runs once the store fails the Task" ); } -// oracle: mutant a job run again after its last failure could not be recorded runs one more attempt past the limit +// oracle: mutant an end that can be neither recorded nor queued returns an error, so the job queue redoes the paid work #[tokio::test] -async fn a_last_attempt_whose_failure_cannot_be_recorded_is_not_run_again() { - let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).await; +async fn a_task_whose_end_cannot_be_recorded_or_queued_is_not_retried() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; let id = harness.create(1).await; let mut job = harness.queued_job(); - for _ in 1..MAX_ATTEMPTS { - job.run().await.expect("the job queues the next attempt itself"); - job = harness.queued_job(); - } + harness.job_queue.close(); + harness.success_failures.store(1, Ordering::SeqCst); harness.failure_failures.store(1, Ordering::SeqCst); - job.run().await.expect_err("the job queue runs the job again"); - job.run().await.expect("its attempt already started"); + job.run().await.expect("the job queue does not run the attempt again"); - let task = harness.task(id).await; - assert_eq!(task.attempts, MAX_ATTEMPTS); + assert_eq!(harness.runs.load(Ordering::SeqCst), 1); assert_eq!( - harness.runs.load(Ordering::SeqCst), - usize::try_from(MAX_ATTEMPTS).expect("a few attempts"), - "no attempt past the last" + harness.task(id).await.state, + ProvisionerTaskState::Running, + "the store fails it at its deadline" ); - let JobMessage { schedule_for, .. } = harness.job_queue_rx.try_recv().expect("a job ends it later"); - assert!(schedule_for.is_some_and(|at| task.deadline_at < at)); } -// oracle: mutant two jobs holding the same token both run the attempt +// oracle: mutant the prelude starts an attempt whatever the attempts made, so a Task runs one attempt past the last #[tokio::test] -async fn a_duplicate_job_does_not_run_its_attempt_again() { +async fn a_last_attempt_whose_failure_cannot_be_recorded_is_not_run_again() { let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).await; let id = harness.create(1).await; let mut job = harness.queued_job(); - let mut duplicate = - ProvisionerTaskJob::read_json(&job.write_json().expect("job JSON"), harness.runner.clone()).expect("valid job"); + for _ in 1..MAX_ATTEMPTS { + job.run().await.expect_err("the job queue runs the job again"); + } + harness.failure_failures.store(1, Ordering::SeqCst); - job.run().await.expect("the job queues the next attempt itself"); - harness.queued_job(); - duplicate.run().await.expect("its attempt already started"); + job.run().await.expect("its end is left to a job past the deadline"); + let task = harness.task(id).await; + assert_eq!( + task.state, + ProvisionerTaskState::Running, + "the failure was not recorded" + ); + let (_, schedule_for) = harness.queued().expect("a job ends it later"); + assert!(schedule_for.is_some_and(|at| task.deadline_at < at)); - assert_eq!(harness.runs.load(Ordering::SeqCst), 1); - assert_eq!(harness.task(id).await.attempts, 1); + // As if Gateway stopped before the deadline, so a new job finds the Task still running. + harness.manager.resume().await.expect("resume"); + harness.queued_job().run().await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!( + harness.runs.load(Ordering::SeqCst), + usize::try_from(MAX_ATTEMPTS).expect("a few attempts"), + "no attempt past the last" + ); + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "internal" })); } -// oracle: mutant a queued Task of an unknown kind is never failed +// oracle: mutant a Task whose parameters cannot be read runs its attempt, or stays unfinished #[tokio::test] -async fn invalid_parameters_and_unknown_kinds_fail_the_task() { +async fn invalid_parameters_fail_the_task() { let harness = Harness::new(|value| Ok(json!(value))).await; - let store = harness.runner.store(); + let id = Uuid::new_v4(); + let job_token = Uuid::new_v4(); let now = OffsetDateTime::now_utc(); + harness + .manager + .store() + .create( + NewProvisionerTaskRecord { + id, + kind: KIND.to_owned(), + target: id.to_string(), + params: json!({ "value": "not a number" }), + deadline_at: now + time::Duration::hours(1), + job_token, + }, + now, + ) + .await + .expect("create"); + let json = json!({ "taskId": id, "jobToken": job_token }).to_string(); + let mut job = TaskJob::::read_json(&json) + .expect("valid job") + .into_job(harness.manager.clone()); - 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, 1).await.expect("no retry"); - - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Failed, "{kind}"); - assert_eq!(task.payload["error"], json!({ "kind": "invalidTask" }), "{kind}"); - } + job.run().await.expect("no retry"); + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["error"], json!({ "kind": "invalidTask" })); assert_eq!(harness.runs.load(Ordering::SeqCst), 0); } @@ -707,10 +717,7 @@ async fn an_end_that_fails_is_done_again_by_a_new_job_until_it_succeeds() { 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"); + let (next, schedule_for) = harness.queued().expect("a new job"); assert!( schedule_for.is_some_and(|at| OffsetDateTime::now_utc() < at), "it runs later" @@ -723,7 +730,7 @@ async fn an_end_that_fails_is_done_again_by_a_new_job_until_it_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"); + assert!(harness.job_queue.take().is_none(), "no job is left"); } // oracle: intent crates/provisioner-task/INTENT.md §Invariants @@ -761,10 +768,11 @@ async fn the_job_is_read_back_from_its_json() { 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"); + let read = TaskJob::::read_json(&json) + .expect("valid job") + .into_job(harness.manager.clone()); assert_eq!(read.task_id, id); assert_eq!(read.job_token, harness.task(id).await.job_token); - assert_eq!(read.step, Step::Attempt(1)); - assert_eq!(job.name(), ProvisionerTaskJob::NAME); + assert_eq!(job.name(), KIND); } diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index 0c66667332..79a7f94f36 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -352,13 +352,17 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { let secrets = devolutions_gateway::secret_store::SecretStore::default(); - let provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( + let in_memory_job_dependencies = devolutions_gateway::job_queue::InMemoryJobDependencies { + conf_handle: conf_handle.clone(), + recordings: recording_manager_handle.clone(), + secrets: secrets.clone(), task_store, - Arc::new(job_queue_ctx.durable()), - conf_handle.clone(), - recording_manager_handle.clone(), - secrets.clone(), - ); + job_queue: job_queue_ctx.durable(), + }; + let provisioner_tasks = devolutions_gateway::provisioner_tasks::task_manager(&in_memory_job_dependencies); + job_queue_ctx + .job_queue_handle + .set_in_memory_job_dependencies(in_memory_job_dependencies)?; // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. // Unfinished Tasks still fail at their deadline, so this must not keep the other features from starting. @@ -383,8 +387,6 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { monitoring_state, traffic_audit_handle: traffic_audit_task.handle(), agent_tunnel_handle, - provisioner_tasks: provisioner_tasks.clone(), - secrets: secrets.clone(), }; for listener in &conf.listeners { @@ -451,10 +453,7 @@ 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, - provisioner_tasks, - )); + tasks.register(devolutions_gateway::job_queue::JobRunnerTask::new(&job_queue_ctx)); tasks.register(devolutions_gateway::job_queue::JobQueueTask::new(job_queue_ctx)); tasks.register(traffic_audit_task); diff --git a/devolutions-gateway/tests/preflight.rs b/devolutions-gateway/tests/preflight.rs index 71a312c624..1cbcf84808 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}")) } -async fn make_router() -> anyhow::Result<(Router, DgwState, MockHandles)> { - let (state, handles) = DgwState::mock(CONFIG).await?; +fn make_router() -> anyhow::Result<(Router, DgwState, MockHandles)> { + let (state, handles) = DgwState::mock(CONFIG)?; let app = devolutions_gateway::make_http_service(state.clone()) .layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); Ok((app, state, handles)) } -async fn make_router_with_config(config: &str) -> anyhow::Result<(Router, DgwState, MockHandles)> { - let (state, handles) = DgwState::mock(config).await?; +fn make_router_with_config(config: &str) -> anyhow::Result<(Router, DgwState, MockHandles)> { + let (state, handles) = DgwState::mock(config)?; 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().await?; + let (app, _state, _handles) = make_router()?; 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).await?; + let (app, _state, _handles) = make_router_with_config(&config)?; 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().await?; + let (app, _state, _handles) = make_router()?; 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().await?; + let (app, _state, _handles) = make_router()?; 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().await?; + let (app, _state, _handles) = make_router()?; // 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().await?; + let (app, _state, _handles) = make_router()?; 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().await?; + let (app, _state, _handles) = make_router()?; 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().await?; + let (app, _state, _handles) = make_router()?; let token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJqdGkiOiI1ZTNlODMzZi04NGM3LTQ1NDEtYjY3Ni1hY2MzMjk5ZTM5YjgifQ.1qECGlrW7y9HWFArc6GPHLGTOY7PhAvzKJ5XMRBg4k4"; diff --git a/devolutions-gateway/tests/traffic.rs b/devolutions-gateway/tests/traffic.rs index 3c9e4a63cb..7175d3bb59 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).await?; + let (mut state, handles) = DgwState::mock(CONFIG)?; // 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 25755d7e0f..8e090e9d10 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).await?; + let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG)?; 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).await?; + let (state, _handle) = devolutions_gateway::DgwState::mock(&config)?; 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).await?; + let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG)?; 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).await?; + let (state, _handles) = devolutions_gateway::DgwState::mock(CONFIG)?; let app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); From 62fba8f9bdf559da5db94e12177d7e30e26fd968 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 29 Oct 2026 15:41:24 -0400 Subject: [PATCH 78/80] refactor(dgw): retry Task attempts inside the job and drop the job token `create` only writes the two records: the Task in gateway.db and its job in job_queue.db, with the API key kept in the secret store in between. The job queue delivers each job at least once; the record keeps where the Task left off, so after a restart the job picks it up from there, and the prelude fails a Task whose attempts are spent. `TaskJob::run` waits `retry_delay` and runs the next attempt itself after a recoverable failure, instead of returning an error for the job queue to retry. An error from `run` now only means a write or an end failed. There is no startup resume anymore. `requeue_lost_jobs` stores the job of a queued Task that has none, before the job runner starts. With a single job per Task, the job token has nothing left to guard, so it is removed, and migration 2 drops its column. `DgwState` holds the task store, the secret store and a `DurableJobQueue`, shared with the job dependencies, which only the job runner gets. A job is pushed in its stored form (`StorableJob`) and hydrated from it (`UnhydratedJob`), so a runnable `TaskJob` exists only in the runner. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- crates/job-queue-libsql/src/lib.rs | 29 +- crates/job-queue/src/lib.rs | 5 +- .../migrations/02_drop_job_token.sql | 1 + crates/provisioner-task-libsql/src/lib.rs | 236 ++----- crates/provisioner-task/src/lib.rs | 20 +- .../src/api/tasks/recording_ai_analysis.rs | 2 +- devolutions-gateway/src/job_queue.rs | 144 ++--- devolutions-gateway/src/lib.rs | 11 +- devolutions-gateway/src/openapi.rs | 6 +- .../src/provisioner_tasks/mod.rs | 28 +- .../recording_ai_analysis/failure.rs | 1 - .../recording_ai_analysis/mod.rs | 86 +-- .../recording_ai_analysis/tests.rs | 250 ++------ .../src/provisioner_tasks/runner/mod.rs | 581 ++++++------------ .../src/provisioner_tasks/runner/tests.rs | 525 ++++------------ devolutions-gateway/src/service.rs | 39 +- devolutions-gateway/tests/preflight.rs | 24 +- devolutions-gateway/tests/traffic.rs | 2 +- devolutions-gateway/tests/webapp.rs | 8 +- 19 files changed, 616 insertions(+), 1382 deletions(-) create mode 100644 crates/provisioner-task-libsql/migrations/02_drop_job_token.sql diff --git a/crates/job-queue-libsql/src/lib.rs b/crates/job-queue-libsql/src/lib.rs index 0efbfd23f0..60fae7daf1 100644 --- a/crates/job-queue-libsql/src/lib.rs +++ b/crates/job-queue-libsql/src/lib.rs @@ -3,7 +3,7 @@ extern crate tracing; use anyhow::Context as _; use async_trait::async_trait; -use job_queue::{DynJob, JobCtx, JobQueue, JobReader, RunnerWaker}; +use job_queue::{JobCtx, JobQueue, JobReader, RunnerWaker}; use libsql::Connection; use time::OffsetDateTime; use ulid::Ulid; @@ -199,7 +199,7 @@ impl JobQueue for LibSqlJobQueue { Ok(()) } - async fn push_job(&self, job: &DynJob, schedule_for: Option) -> anyhow::Result<()> { + async fn push_job(&self, name: &str, def: &str, schedule_for: Option) -> anyhow::Result<()> { let sql_query = "INSERT INTO job_queue (id, scheduled_for, failed_attempts, status, name, def) VALUES (:id, :scheduled_for, :failed_attempts, :status, :name, jsonb(:def))"; @@ -216,8 +216,8 @@ impl JobQueue for LibSqlJobQueue { (":scheduled_for", schedule_for.unix_timestamp()), (":failed_attempts", 0), (":status", JobStatus::Queued as u32), - (":name", job.name()), - (":def", job.write_json()?), + (":name", name), + (":def", def), ); trace!(%sql_query, ?params, "Pushing a new job"); @@ -402,6 +402,27 @@ impl JobQueue for LibSqlJobQueue { Ok(Some(scheduled_for)) } + + async fn job_defs(&self, name: &str) -> anyhow::Result> { + let sql_query = "SELECT json(def) FROM job_queue WHERE name = $1"; + + let params = [name]; + + trace!(%sql_query, ?params, "Fetching the jobs stored under a name"); + + let mut rows = self + .conn + .query(sql_query, params) + .await + .context("failed to execute SQL query")?; + + let mut defs = Vec::new(); + while let Some(row) = rows.next().await.context("failed to read the row")? { + defs.push(row.get::(0).context("failed to read def value")?); + } + + Ok(defs) + } } // Typically, migrations should not be modified once released, and we should only be appending to this list. diff --git a/crates/job-queue/src/lib.rs b/crates/job-queue/src/lib.rs index 93232bd868..0c1d8c4d03 100644 --- a/crates/job-queue/src/lib.rs +++ b/crates/job-queue/src/lib.rs @@ -44,7 +44,7 @@ pub trait JobQueue: Send + Sync { /// Pushes a new job into the queue /// /// This function should ideally call `RunnerWaker::wake()` once the job is enqueued. - async fn push_job(&self, job: &DynJob, schedule_for: Option) -> anyhow::Result<()>; + async fn push_job(&self, name: &str, def: &str, schedule_for: Option) -> anyhow::Result<()>; /// Fetches at most `number_of_jobs` from the queue async fn claim_jobs(&self, reader: &dyn JobReader, number_of_jobs: usize) -> anyhow::Result>; @@ -62,6 +62,9 @@ pub trait JobQueue: Send + Sync { /// Retrieves the closest future scheduled date async fn next_scheduled_date(&self) -> anyhow::Result>; + + /// Returns the JSON of every job stored under `name` + async fn job_defs(&self, name: &str) -> anyhow::Result>; } pub struct JobCtx { diff --git a/crates/provisioner-task-libsql/migrations/02_drop_job_token.sql b/crates/provisioner-task-libsql/migrations/02_drop_job_token.sql new file mode 100644 index 0000000000..dbfbaf096d --- /dev/null +++ b/crates/provisioner-task-libsql/migrations/02_drop_job_token.sql @@ -0,0 +1 @@ +ALTER TABLE provisioner_task_records DROP COLUMN job_token; diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index deee0b5546..7a786e76bf 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -19,12 +19,15 @@ use tokio::sync::Mutex; use uuid::Uuid; /// Every schema change ever made to the `provisioner_task_*` tables, oldest first. Only ever append to it. -const MIGRATIONS: &[&str] = &[include_str!("../migrations/01_provisioner_task.sql")]; +const MIGRATIONS: &[&str] = &[ + include_str!("../migrations/01_provisioner_task.sql"), + include_str!("../migrations/02_drop_job_token.sql"), +]; const SCHEMA_VERSION_KEY: &str = "schema_version"; const COLUMNS: &str = - "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at, job_token"; + "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at"; /// Tasks kept in the `provisioner_task_records` table of `gateway.db`. pub struct LibSqlProvisionerTaskStore { @@ -182,8 +185,8 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { tx.execute( "INSERT INTO provisioner_task_records - (id, kind, target, params, state, payload, created_at, deadline_at, job_token) - VALUES (?1, ?2, ?3, ?4, 'queued', 'null', ?5, ?6, ?7)", + (id, kind, target, params, state, payload, created_at, deadline_at) + VALUES (?1, ?2, ?3, ?4, 'queued', 'null', ?5, ?6)", params![ task.id.to_string(), task.kind.clone(), @@ -191,7 +194,6 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { task.params.to_string(), now.unix_timestamp(), task.deadline_at.unix_timestamp(), - task.job_token.to_string(), ], ) .await @@ -223,7 +225,6 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { async fn start_attempt( &self, id: Uuid, - job_token: Uuid, payload: serde_json::Value, deadline_at: OffsetDateTime, now: OffsetDateTime, @@ -239,11 +240,6 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { return Ok(AttemptStart::Unknown); }; - if task.job_token != job_token { - tx.rollback().await.context("end task attempt")?; - return Ok(AttemptStart::Superseded); - } - let task = fail_if_overdue(&tx, task, now).await?; if task.state.is_finished() { @@ -272,35 +268,6 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(AttemptStart::Started(started)) } - async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result { - let conn = self.conn.lock().await; - let tx = conn - .transaction_with_behavior(TransactionBehavior::Immediate) - .await - .context("begin job token replacement")?; - - let Some(task) = find(&tx, id).await? else { - tx.rollback().await.context("end job token replacement")?; - return Ok(false); - }; - let task = fail_if_overdue(&tx, task, now).await?; - - let replaced = !task.state.is_finished(); - - if replaced { - tx.execute( - "UPDATE provisioner_task_records SET job_token = ?2 WHERE id = ?1", - params![id.to_string(), job_token.to_string()], - ) - .await - .context("replace task job token")?; - } - - tx.commit().await.context("commit job token replacement")?; - - Ok(replaced) - } - async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result> { let conn = self.conn.lock().await; let tx = conn @@ -448,11 +415,6 @@ fn task_from_row(row: &Row) -> anyhow::Result { started_at: optional_timestamp(8, "start time")?, finished_at: optional_timestamp(9, "finish time")?, deadline_at: timestamp(10, "deadline")?, - job_token: row - .get::(11) - .context("decode task job token")? - .parse() - .context("parse task job token")?, }) } @@ -558,7 +520,6 @@ mod tests { target: String::from(target), params: serde_json::json!({ "provider": "openai", "model": "gpt-6-luna" }), deadline_at: now() + Duration::hours(2), - job_token: Uuid::new_v4(), } } @@ -587,15 +548,10 @@ mod tests { let task = new_task("session-1"); create(&store, task.clone()).await; - let mut again = task.clone(); - again.job_token = Uuid::new_v4(); - let CreateOutcome::Existing(existing) = create(&store, again).await else { + let CreateOutcome::Existing(existing) = create(&store, task.clone()).await else { panic!("the same request must return the existing task"); }; - assert_eq!( - existing.job_token, task.job_token, - "the job of the first request keeps the task" - ); + assert_eq!(existing.id, task.id); } #[tokio::test] @@ -682,7 +638,6 @@ mod tests { store .start_attempt( overdue.id, - overdue.job_token, serde_json::json!({ "done": 3, "total": 10, "usage": { "inputTokens": 1200 } }), now() - Duration::minutes(1), now() - Duration::minutes(2), @@ -717,7 +672,6 @@ mod tests { store .start_attempt( task.id, - task.job_token, serde_json::json!({ "done": 0 }), now() + Duration::hours(2), now() @@ -741,7 +695,6 @@ mod tests { store .start_attempt( task.id, - task.job_token, serde_json::json!({ "done": 1 }), now() + Duration::hours(2), now() @@ -770,13 +723,7 @@ mod tests { ); assert!(matches!( store - .start_attempt( - task.id, - task.job_token, - serde_json::json!({}), - now() + Duration::hours(2), - now() - ) + .start_attempt(task.id, serde_json::json!({}), now() + Duration::hours(2), now()) .await .expect("start after finish"), AttemptStart::Finished(_) @@ -802,7 +749,6 @@ mod tests { store .start_attempt( task.id, - task.job_token, serde_json::json!({ "done": 0 }), now() + Duration::hours(2), now(), @@ -832,7 +778,6 @@ mod tests { store .start_attempt( overdue.id, - overdue.job_token, serde_json::json!({ "done": 3 }), now() + Duration::minutes(1), now(), @@ -884,13 +829,7 @@ mod tests { 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(), - ) + .start_attempt(task.id, serde_json::json!({}), now() + Duration::minutes(30), now()) .await .expect("start"); @@ -910,13 +849,7 @@ mod tests { 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(), - ) + .start_attempt(task.id, serde_json::json!({}), now() + Duration::minutes(30), now()) .await .expect("start"); let moved = now() + Duration::minutes(50); @@ -953,13 +886,7 @@ mod tests { assert!(matches!( store - .start_attempt( - overdue.id, - overdue.job_token, - serde_json::json!({}), - now() + Duration::hours(2), - now() - ) + .start_attempt(overdue.id, serde_json::json!({}), now() + Duration::hours(2), now()) .await .expect("start attempt"), AttemptStart::Finished(_) @@ -981,77 +908,6 @@ mod tests { assert_eq!(again.state, ProvisionerTaskState::Failed); } - #[tokio::test] - async fn only_the_job_holding_the_token_runs_the_task() { - let store = store().await; - let task = new_task("session-1"); - create(&store, task.clone()).await; - - assert_eq!( - store - .start_attempt( - task.id, - Uuid::new_v4(), - serde_json::json!({}), - now() + Duration::hours(2), - now() - ) - .await - .expect("start with another token"), - AttemptStart::Superseded - ); - assert_eq!( - store - .get(task.id, now()) - .await - .expect("get task") - .expect("task") - .attempts, - 0, - "a superseded job counts no attempt" - ); - - let new_token = Uuid::new_v4(); - assert!( - store - .replace_job_token(task.id, new_token, now()) - .await - .expect("replace token") - ); - assert_eq!( - store - .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() + Duration::hours(2), now()) - .await - .expect("start with the new token"), - AttemptStart::Started(started) if started.job_token == new_token && started.attempts == 1 - )); - - store - .fail(task.id, serde_json::json!({ "reason": "boom" }), now()) - .await - .expect("fail"); - assert!( - !store - .replace_job_token(task.id, Uuid::new_v4(), now()) - .await - .expect("replace token of a finished task"), - "a finished Task gets no new job" - ); - } - #[tokio::test] async fn unfinished_tasks_are_listed_and_overdue_ones_failed() { let store = store().await; @@ -1064,13 +920,7 @@ mod tests { create(&store, task.clone()).await; } store - .start_attempt( - running.id, - running.job_token, - serde_json::json!({}), - now() + Duration::hours(2), - now(), - ) + .start_attempt(running.id, serde_json::json!({}), now() + Duration::hours(2), now()) .await .expect("start"); store @@ -1104,23 +954,11 @@ 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() + Duration::hours(2), - now() - ) + .start_attempt(Uuid::new_v4(), serde_json::json!({}), now() + Duration::hours(2), now()) .await .expect("start attempt"), AttemptStart::Unknown ); - assert!( - !store - .replace_job_token(Uuid::new_v4(), Uuid::new_v4(), now()) - .await - .expect("replace job token") - ); for finished in [ store .update_running(Uuid::new_v4(), serde_json::json!({}), now() + Duration::hours(2), now()) @@ -1169,4 +1007,48 @@ mod tests { .expect("newer schema must be rejected"); assert!(error.to_string().contains("newer than supported")); } + + #[tokio::test] + async fn a_task_recorded_with_a_job_token_is_read_after_the_upgrade() { + let temp_dir = tempfile::tempdir().expect("create temporary directory"); + let path = temp_dir.path().join("gateway.db"); + let db = gateway_db::GatewayDb::open_path(path.to_str().expect("temporary path is UTF-8")) + .await + .expect("open gateway database"); + let conn = db.connect().await.expect("connect"); + conn.execute_batch(MIGRATIONS[0]).await.expect("apply migration 1"); + conn.execute( + "INSERT INTO provisioner_task_metadata (key, value) VALUES (?1, 1)", + params![SCHEMA_VERSION_KEY], + ) + .await + .expect("record schema version 1"); + let id = Uuid::new_v4(); + conn.execute( + "INSERT INTO provisioner_task_records + (id, kind, target, params, state, payload, created_at, deadline_at, job_token) + VALUES (?1, 'recording-ai-analysis', 'session-1', '{}', 'queued', 'null', ?2, ?3, ?4)", + params![ + id.to_string(), + now().unix_timestamp(), + (now() + Duration::hours(2)).unix_timestamp(), + Uuid::new_v4().to_string() + ], + ) + .await + .expect("insert a task of schema version 1"); + + let store = LibSqlProvisionerTaskStore::open(db.connect().await.expect("connect")) + .await + .expect("upgrade the provisioner task store"); + + let task = store.get(id, now()).await.expect("get task").expect("task"); + assert_eq!(task.state, ProvisionerTaskState::Queued); + assert_eq!( + schema_version(&*store.conn.lock().await) + .await + .expect("read schema version"), + MIGRATIONS.len() + ); + } } diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index 587a4b93c5..31c7c4105f 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -49,8 +49,6 @@ pub struct ProvisionerTaskRecord { /// /// 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, } /// The record of a Task about to be created. @@ -61,8 +59,6 @@ pub struct NewProvisionerTaskRecord { pub target: String, pub params: serde_json::Value, pub deadline_at: OffsetDateTime, - /// Token of the job that will run the new Task. - pub job_token: Uuid, } #[derive(Debug, Clone, PartialEq)] @@ -83,8 +79,6 @@ pub enum AttemptStart { Started(ProvisionerTaskRecord), /// The Task is finished, possibly just now because it is past its deadline. Finished(ProvisionerTaskRecord), - /// Another job now runs the Task, so this one is stale or a duplicate. - Superseded, /// No Task has this ID. Unknown, } @@ -103,26 +97,18 @@ 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 and - /// its deadline moved to `deadline_at`. + /// Marks the start of one more attempt: `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. + /// Only [`AttemptStart::Started`] lets the attempt run. async fn start_attempt( &self, id: Uuid, - job_token: Uuid, payload: serde_json::Value, deadline_at: OffsetDateTime, now: OffsetDateTime, ) -> anyhow::Result; - /// Hands an unfinished Task to a new job, so the job that had it stops at its next attempt. - /// - /// Only the start of an attempt checks the token, so call it only while the old job runs no attempt, such as at - /// startup before any job runs. Returns `false` if the Task is unknown or finished. - async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result; - /// Every unfinished Task, after failing those past their deadline. async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; diff --git a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs index 743b651d69..6e895b7241 100644 --- a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs +++ b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs @@ -307,7 +307,7 @@ mod tests { "RecordingPath": dir.path().to_str().expect("UTF-8"), "__debug__": { "disable_token_validation": true }, }); - let (state, _handles) = DgwState::mock(&config.to_string()).expect("mock state"); + 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)))); diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 16b2b83153..45f64f392b 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -1,6 +1,6 @@ use std::future::Future; use std::path::Path; -use std::sync::{Arc, OnceLock}; +use std::sync::Arc; use std::time::Duration; use anyhow::Context as _; @@ -14,7 +14,7 @@ use tokio::sync::{Notify, mpsc}; use crate::config::ConfHandle; use crate::provisioner_tasks::recording_ai_analysis::RecordingAiAnalysis; -use crate::provisioner_tasks::runner::TaskJob; +use crate::provisioner_tasks::runner::{TaskJob, TaskJobUnhydrated}; use crate::recording::RecordingMessageSender; use crate::secret_store::SecretStore; @@ -31,25 +31,33 @@ pub struct JobMessage { pub schedule_for: Option, } +/// Sends jobs to be stored by the job queue task, which tries a few times then drops the job, so the caller never +/// knows whether it was stored; see [`DurableJobQueue`] for a caller that must know. #[derive(Clone)] pub struct JobQueueHandle { sender: mpsc::Sender, - in_memory_job_dependencies: Arc>, } pub type JobQueueReceiver = mpsc::Receiver; /// What jobs need to run beyond their stored JSON, kept in memory only, such as the secrets of a Task. -#[derive(Clone)] +/// +/// Only the job runner holds it, to hydrate the jobs it reads back. pub struct InMemoryJobDependencies { pub conf_handle: ConfHandle, pub recordings: RecordingMessageSender, pub secrets: SecretStore, pub task_store: DynProvisionerTaskStore, - pub job_queue: DurableJobQueue, } -/// A job read back from its JSON that still needs the in-memory dependencies to run. +/// A job as the job queue stores it: a name and its JSON. +pub trait StorableJob { + fn name(&self) -> &str; + + fn write_json(&self) -> anyhow::Result; +} + +/// A job read back from its JSON that needs the in-memory dependencies to run. pub trait UnhydratedJob { fn hydrate(self, dependencies: &InMemoryJobDependencies) -> DynJob; } @@ -63,7 +71,7 @@ pub struct JobRunnerTask { notify_runner: Arc, runner_waker: RunnerWaker, queue: DynJobQueue, - in_memory_job_dependencies: Arc>, + in_memory_job_dependencies: InMemoryJobDependencies, } impl JobQueueCtx { @@ -113,41 +121,55 @@ impl JobQueueCtx { pub fn durable(&self) -> DurableJobQueue { DurableJobQueue(Arc::clone(&self.queue)) } + + /// The JSON of every job stored under `name`, for startup, before the job runner claims any job. + pub async fn job_defs(&self, name: &str) -> anyhow::Result> { + self.queue.job_defs(name).await + } +} + +/// A job queue database in memory. +async fn in_memory_job_queue() -> anyhow::Result { + let conn = libsql::Builder::new_local(":memory:") + .build() + .await + .context("build database")? + .connect() + .context("open database connection")?; + let queue = job_queue_libsql::LibSqlJobQueue::builder() + .runner_waker(RunnerWaker::new(|| {})) + .conn(conn) + .build(); + queue.setup().await.context("queue setup")?; + + Ok(queue) } -/// The job queue database: a job pushed here is stored before `push` returns, unlike with [`JobQueueHandle`]. +/// Stores jobs in the job queue database: a job pushed here is stored before `push` returns, unlike with +/// [`JobQueueHandle`]. #[derive(Clone)] pub struct DurableJobQueue(DynJobQueue); impl DurableJobQueue { + /// A job queue in memory, for [`crate::DgwState::mock`]. + #[doc(hidden)] + pub async fn mock() -> anyhow::Result { + Ok(Self(Arc::new(in_memory_job_queue().await?))) + } + /// Stores `job` to run at `schedule_for`, or at once without it. - pub async fn push(&self, job: T, schedule_for: Option) -> anyhow::Result<()> { - let job: DynJob = Box::new(job); - self.0.push_job(&job, schedule_for).await + pub async fn push(&self, job: &impl StorableJob, schedule_for: Option) -> anyhow::Result<()> { + self.0.push_job(job.name(), &job.write_json()?, schedule_for).await } } impl JobQueueHandle { pub fn new() -> (Self, JobQueueReceiver) { let (tx, rx) = mpsc::channel(512); - let handle = Self { - sender: tx, - in_memory_job_dependencies: Arc::default(), - }; + let handle = Self { sender: tx }; (handle, rx) } - /// Sets the dependencies the job runner hydrates each job with; once only. - pub fn set_in_memory_job_dependencies(&self, dependencies: InMemoryJobDependencies) -> anyhow::Result<()> { - self.in_memory_job_dependencies - .set(dependencies) - .map_err(|_| anyhow::anyhow!("the in-memory job dependencies are already set")) - } - - pub fn in_memory_job_dependencies(&self) -> Option<&InMemoryJobDependencies> { - self.in_memory_job_dependencies.get() - } - pub fn blocking_enqueue(&self, job: T) -> anyhow::Result<()> { self.sender .blocking_send(JobMessage { @@ -232,8 +254,16 @@ async fn job_queue_task(ctx: JobQueueTask, mut shutdown_signal: ShutdownSignal) let queue = Arc::clone(&queue); async move { + let def = match msg.job.write_json() { + Ok(def) => def, + Err(e) => { + error!(error = format!("{e:#}"), "Failed to serialize job"); + return; + } + }; + for _ in 0..5 { - match queue.push_job(&msg.job, msg.schedule_for).await { + match queue.push_job(msg.job.name(), &def, msg.schedule_for).await { Ok(()) => break, Err(e) => { warn!(error = format!("{e:#}"), "Failed to push job"); @@ -255,12 +285,13 @@ async fn job_queue_task(ctx: JobQueueTask, mut shutdown_signal: ShutdownSignal) } impl JobRunnerTask { - pub fn new(ctx: &JobQueueCtx) -> Self { + /// The job runner, which hydrates each job it reads back with `in_memory_job_dependencies`. + pub fn new(ctx: &JobQueueCtx, in_memory_job_dependencies: InMemoryJobDependencies) -> Self { Self { notify_runner: Arc::clone(&ctx.notify_runner), runner_waker: RunnerWaker::clone(&ctx.runner_waker), queue: Arc::clone(&ctx.queue), - in_memory_job_dependencies: Arc::clone(&ctx.job_queue_handle.in_memory_job_dependencies), + in_memory_job_dependencies, } } } @@ -342,17 +373,7 @@ async fn job_runner_task(ctx: JobRunnerTask, mut shutdown_signal: ShutdownSignal } struct DgwJobReader { - in_memory_job_dependencies: Arc>, -} - -impl DgwJobReader { - fn hydrate(&self, job: impl UnhydratedJob) -> anyhow::Result { - let dependencies = self - .in_memory_job_dependencies - .get() - .context("the in-memory job dependencies are not set")?; - Ok(job.hydrate(dependencies)) - } + in_memory_job_dependencies: InMemoryJobDependencies, } impl JobReader for DgwJobReader { @@ -370,7 +391,9 @@ impl JobReader for DgwJobReader { serde_json::from_str(json).context("failed to deserialize DeleteRecordingsJob")?; Ok(Box::new(job)) } - TaskJob::::NAME => self.hydrate(TaskJob::::read_json(json)?), + TaskJob::::NAME => Ok( + TaskJobUnhydrated::::read_json(json)?.hydrate(&self.in_memory_job_dependencies) + ), _ => anyhow::bail!("unknown job name: {name}"), } } @@ -381,37 +404,23 @@ impl JobReader for DgwJobReader { pub(crate) struct StoredJob { pub(crate) name: String, pub(crate) json: String, - pub(crate) schedule_for: Option, } -/// The job queue of tests: the queue Gateway runs, in memory, which also hands each job it stores to the test, and -/// refuses every job once closed, as a database that can't be written does. +/// The job queue of tests: the queue Gateway runs, in memory, which also hands each job it stores to the test. #[cfg(test)] pub(crate) struct TestJobQueue { inner: job_queue_libsql::LibSqlJobQueue, stored: parking_lot::Mutex>, - closed: std::sync::atomic::AtomicBool, } #[cfg(test)] impl TestJobQueue { pub(crate) async fn new() -> Arc { - let conn = libsql::Builder::new_local(":memory:") - .build() - .await - .expect("build database") - .connect() - .expect("open database connection"); - let inner = job_queue_libsql::LibSqlJobQueue::builder() - .runner_waker(RunnerWaker::new(|| {})) - .conn(conn) - .build(); - inner.setup().await.expect("queue setup"); + let inner = in_memory_job_queue().await.expect("in-memory job queue"); Arc::new(Self { inner, stored: parking_lot::Mutex::default(), - closed: std::sync::atomic::AtomicBool::new(false), }) } @@ -420,10 +429,6 @@ impl TestJobQueue { self.stored.lock().pop_front() } - pub(crate) fn close(&self) { - self.closed.store(true, std::sync::atomic::Ordering::SeqCst); - } - /// This queue, as [`JobQueueCtx::durable`] hands out the queue Gateway runs. pub(crate) fn durable(self: &Arc) -> DurableJobQueue { DurableJobQueue(Arc::::clone(self)) @@ -441,16 +446,11 @@ impl JobQueue for TestJobQueue { self.inner.reset_claimed_jobs().await } - async fn push_job(&self, job: &DynJob, schedule_for: Option) -> anyhow::Result<()> { - anyhow::ensure!( - !self.closed.load(std::sync::atomic::Ordering::SeqCst), - "the job queue is closed" - ); - self.inner.push_job(job, schedule_for).await?; + async fn push_job(&self, name: &str, def: &str, schedule_for: Option) -> anyhow::Result<()> { + self.inner.push_job(name, def, schedule_for).await?; self.stored.lock().push_back(StoredJob { - name: job.name().to_owned(), - json: job.write_json()?, - schedule_for, + name: name.to_owned(), + json: def.to_owned(), }); Ok(()) } @@ -474,4 +474,8 @@ impl JobQueue for TestJobQueue { async fn next_scheduled_date(&self) -> anyhow::Result> { self.inner.next_scheduled_date().await } + + async fn job_defs(&self, name: &str) -> anyhow::Result> { + self.inner.job_defs(name).await + } } diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index 3480c5c21f..9cee7088e4 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -70,6 +70,9 @@ pub struct DgwState { pub monitoring_state: Arc, pub traffic_audit_handle: traffic_audit::TrafficAuditHandle, pub agent_tunnel_handle: Option>, + pub task_store: provisioner_task::DynProvisionerTaskStore, + pub secrets: secret_store::SecretStore, + pub durable_job_queue: job_queue::DurableJobQueue, } #[doc(hidden)] @@ -84,7 +87,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())); @@ -97,6 +100,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 task_store_conn = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?; + let task_store = Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store_conn).await?); + let durable_job_queue = job_queue::DurableJobQueue::mock().await?; let state = Self { conf_handle, @@ -112,6 +118,9 @@ impl DgwState { synthetic_kdc_registry, monitoring_state, agent_tunnel_handle: None, + task_store, + secrets: secret_store::SecretStore::default(), + durable_job_queue, }; let handles = MockHandles { diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index d229136eeb..923d54ea80 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -273,14 +273,13 @@ enum AiAnalysisErrorKind { InvalidTask, /// The attempt panicked. Panicked, - /// Something went wrong inside Gateway, such as a file it could not write. + /// Something went wrong inside Gateway, such as a file it could not write, or a restart that cut the last attempt + /// short. Internal, /// Every attempt failed with a failure worth another attempt. AttemptsExhausted, /// The Task did not start in time, or went too long without progress. DeadlinePassed, - /// The job of the Task could not be queued. - NotQueued, } /// What was wrong with an answer of the AI provider @@ -330,7 +329,6 @@ const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| m Permanent::Internal => AiAnalysisErrorKind::Internal, Permanent::AttemptsExhausted { last: _ } => AiAnalysisErrorKind::AttemptsExhausted, Permanent::DeadlinePassed => AiAnalysisErrorKind::DeadlinePassed, - Permanent::NotQueued => AiAnalysisErrorKind::NotQueued, }; const _: fn(&Recoverable) -> AiAnalysisRecoverableErrorKind = |error| match error { diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index a71dbf5b0b..fe25851377 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -1,29 +1,10 @@ -//! Work the provisioner asks Gateway to do: each request is a Task recorded in `gateway.db`, queued as a job on the job -//! queue by the [`runner::ProvisionerTaskManager`]. +//! Work the provisioner asks Gateway to do: each request is a Task recorded in `gateway.db`, run on the job queue as a +//! [`runner::TaskJob`]. pub mod ai; pub mod recording_ai_analysis; pub mod runner; -use std::sync::Arc; - -use self::recording_ai_analysis::RecordingAiAnalysis; -use self::runner::ProvisionerTaskManager; -use crate::job_queue::InMemoryJobDependencies; - -/// The manager of the AI analysis Tasks. -pub fn task_manager(dependencies: &InMemoryJobDependencies) -> ProvisionerTaskManager { - ProvisionerTaskManager::new( - RecordingAiAnalysis::new( - dependencies.conf_handle.clone(), - dependencies.recordings.clone(), - dependencies.secrets.clone(), - ), - Arc::clone(&dependencies.task_store), - dependencies.job_queue.clone(), - ) -} - /// Why one attempt of a Task failed, with `F` and `R` the failures of its own kind of work. #[derive(Debug, Clone, PartialEq, Eq)] pub enum AttemptError { @@ -52,14 +33,13 @@ pub enum Permanent { InvalidTask, /// The attempt panicked. Panicked, - /// Something went wrong inside Gateway, such as a file it could not write. + /// Something went wrong inside Gateway, such as a file it could not write, or a restart that cut the last attempt + /// short. Internal, /// Every attempt failed with a recoverable failure, the last one being `last`. AttemptsExhausted { last: R }, /// The Task did not start in time, or went too long without progress. DeadlinePassed, - /// The job of the Task could not be queued. - NotQueued, /// A failure of the Task's own kind of work, which names its own `kind`. #[serde(untagged)] Kind(F), diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs index 7be4132b65..582a4c246b 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/failure.rs @@ -107,7 +107,6 @@ mod tests { Permanent::Internal, Permanent::AttemptsExhausted { last: () }, Permanent::DeadlinePassed, - Permanent::NotQueued, ] .map(kind); 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 226ba52ee0..a6c3c72474 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -1,8 +1,7 @@ //! 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 -//! [`ProvisionerTaskManager`](super::runner::ProvisionerTaskManager). +//! [`start`] checks the request and records the Task with [`create`](super::runner::create). //! Each attempt runs the analysis of the AI crate, //! [`AiClient::analyze_recording`](devolutions_gateway_ai::AiClient::analyze_recording), which keeps its working files //! next to the session manifest so that a retry resumes where the last attempt stopped, then leaves the log in the @@ -50,6 +49,9 @@ pub const START_WINDOW: time::Duration = time::Duration::hours(2); /// Time a running AI analysis may go without progress before it fails. pub const PROGRESS_TIMEOUT: time::Duration = time::Duration::minutes(30); +/// Time to wait after the first recoverable failure, doubled after each next one. +pub const FIRST_RETRY_DELAY: time::Duration = time::Duration::seconds(30); + /// Tokens the AI provider counted for a Task, over all its attempts, as kept in its payloads. /// /// It mirrors [`UsageSoFar`], so a change in the AI crate never changes what we store. @@ -159,7 +161,7 @@ pub struct Analysis { prompt_version: String, } -/// Records an AI analysis Task, whose job the [`ProvisionerTaskManager`](super::runner::ProvisionerTaskManager) queues. +/// Records an AI analysis Task and stores its job. /// /// Only a new Task keeps the API key; asking again for a known Task changes nothing, even while its session is /// recording again. @@ -194,18 +196,15 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result Result {} Err(FinishedRecordingError::Recording) => return Err(StartError::RecordingActive), @@ -237,16 +234,21 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result( + &state.task_store, + &state.durable_job_queue, + task_id, + target, + ¶ms, + |record| { // The job may run as soon as it is queued, so the key must be in place first. - tasks - .task() + state .secrets .insert(SecretId(record.id), &access.to_secret()?, record.deadline_at) - }) - .await - .map_err(StartError::Internal)?; + }, + ) + .await + .map_err(StartError::Internal)?; Ok(match outcome { CreateOutcome::Created(task) => { @@ -267,18 +269,15 @@ pub async fn get( task_id: Uuid, now: OffsetDateTime, ) -> anyhow::Result> { - state - .job_queue_handle - .in_memory_job_dependencies() - .map(super::task_manager) - .context("the in-memory job dependencies are not set")? - .get(task_id, now) - .await + state.task_store.get(task_id, now).await } impl UnhydratedJob for TaskJobUnhydrated { fn hydrate(self, dependencies: &InMemoryJobDependencies) -> job_queue::DynJob { - Box::new(self.into_job(super::task_manager(dependencies))) + Box::new(self.into_job( + Arc::new(RecordingAiAnalysis::from_dependencies(dependencies)), + Arc::clone(&dependencies.task_store), + )) } } @@ -288,18 +287,28 @@ pub struct RecordingAiAnalysis { recordings: RecordingMessageSender, secrets: SecretStore, progress_timeout: time::Duration, + first_retry_delay: time::Duration, } impl RecordingAiAnalysis { - pub fn new(conf_handle: ConfHandle, recordings: RecordingMessageSender, secrets: SecretStore) -> Self { + fn new(conf_handle: ConfHandle, recordings: RecordingMessageSender, secrets: SecretStore) -> Self { Self { conf_handle, recordings, secrets, progress_timeout: PROGRESS_TIMEOUT, + first_retry_delay: FIRST_RETRY_DELAY, } } + fn from_dependencies(dependencies: &InMemoryJobDependencies) -> Self { + Self::new( + dependencies.conf_handle.clone(), + dependencies.recordings.clone(), + dependencies.secrets.clone(), + ) + } + fn manifest_path(&self, session_id: Uuid) -> Utf8PathBuf { self.conf_handle .get_conf() @@ -361,6 +370,8 @@ async fn remove_if_exists(path: &Utf8Path) -> std::io::Result<()> { impl ProvisionerTask for RecordingAiAnalysis { const KIND: &'static str = KIND; + const START_WINDOW: time::Duration = START_WINDOW; + type Params = Params; type Output = Analysis; @@ -369,14 +380,14 @@ impl ProvisionerTask for RecordingAiAnalysis { type Recoverable = Recoverable; - fn start_window(&self) -> time::Duration { - START_WINDOW - } - fn progress_timeout(&self) -> time::Duration { self.progress_timeout } + fn retry_delay(&self, attempts: u32) -> time::Duration { + self.first_retry_delay * 2u32.pow(attempts.saturating_sub(1)) + } + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value { json!(Running { step: "preparing", @@ -451,11 +462,6 @@ impl ProvisionerTask for RecordingAiAnalysis { .await } - /// The API key is kept in memory only, so no attempt can run after a restart. - fn restart_failure(&self) -> Option { - Some(Failure::KeyLost) - } - async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { self.secrets.remove(SecretId(record.id)); 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 115a5aa86a..0d1a25cb28 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -8,11 +8,11 @@ use devolutions_gateway_ai::{Language, session_summary}; use devolutions_gateway_task::ShutdownHandle; use job_queue::Job as _; use parking_lot::Mutex; -use provisioner_task::{AttemptStart, NewProvisionerTaskRecord, ProvisionerTaskState}; +use provisioner_task::{AttemptStart, DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; -use crate::job_queue::{InMemoryJobDependencies, StoredJob, TestJobQueue}; -use crate::provisioner_tasks::runner::{MAX_ATTEMPTS, ProvisionerTaskManager, TaskJob}; +use crate::job_queue::{StoredJob, TestJobQueue}; +use crate::provisioner_tasks::runner::{MAX_ATTEMPTS, TaskJob, TaskJobUnhydrated}; use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles}; @@ -41,7 +41,8 @@ fn openai(model: &str) -> AiSettings { struct Harness { state: DgwState, - tasks: ProvisionerTaskManager, + task: Arc, + store: DynProvisionerTaskStore, secrets: SecretStore, job_queue: Arc, recording_path: Utf8PathBuf, @@ -54,7 +55,7 @@ impl Harness { 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)).expect("mock state"); + let (mut state, handles) = DgwState::mock(&config(&recording_path)).await.expect("mock state"); let MockHandles { session_manager_rx, recording_manager_rx, @@ -76,33 +77,19 @@ impl Harness { recording_shutdown_signal, )); - let task_store = gateway_db::GatewayDb::open_path(":memory:") - .await - .expect("open the gateway database") - .connect() - .await - .expect("connect to the gateway database"); - let task_store = provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store) - .await - .expect("open the task store"); - let secrets = SecretStore::default(); + // `start` stores its jobs where the test can take them. let job_queue = TestJobQueue::new().await; - let dependencies = InMemoryJobDependencies { - conf_handle: state.conf_handle.clone(), - recordings: state.recordings.clone(), - secrets: secrets.clone(), - task_store: Arc::new(task_store), - job_queue: job_queue.durable(), - }; - let tasks = crate::provisioner_tasks::task_manager(&dependencies); - state - .job_queue_handle - .set_in_memory_job_dependencies(dependencies) - .expect("set the job dependencies"); + state.durable_job_queue = job_queue.durable(); + let secrets = state.secrets.clone(); + let mut task = RecordingAiAnalysis::new(state.conf_handle.clone(), state.recordings.clone(), secrets.clone()); + task.first_retry_delay = time::Duration::ZERO; + let task = Arc::new(task); + let store = Arc::clone(&state.task_store); Self { state, - tasks, + task, + store, secrets, job_queue, recording_path, @@ -151,36 +138,22 @@ impl Harness { 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) -> TaskJob { - // Read as of long ago, so reading never fails the Task in place of the job. - let record = self - .tasks - .store() - .get(task_id, OffsetDateTime::UNIX_EPOCH) - .await - .expect("read task") - .expect("task exists"); - let json = json!({ "taskId": task_id, "jobToken": record.job_token }).to_string(); + /// The job of the Task, as the job queue would run it. + fn job(&self, task_id: Uuid) -> TaskJob { + let json = json!({ "taskId": task_id }).to_string(); - TaskJob::::read_json(&json) + TaskJobUnhydrated::::read_json(&json) .expect("valid job") - .into_job(self.tasks.clone()) + .into_job(Arc::clone(&self.task), Arc::clone(&self.store)) } /// Records the success of the first attempt of the Task without ending it, as if Gateway stopped right after. async fn record_success(&self, task_id: Uuid) { - let store = self.tasks.store(); + let store = &self.store; let now = OffsetDateTime::now_utc(); - let job_token = store - .get(task_id, now) - .await - .expect("read task") - .expect("task exists") - .job_token; let started = store - .start_attempt(task_id, job_token, json!({}), now + time::Duration::hours(1), now) + .start_attempt(task_id, json!({}), now + time::Duration::hours(1), now) .await .expect("start the attempt"); assert!(matches!(started, AttemptStart::Started(_)), "{started:?}"); @@ -193,7 +166,7 @@ impl Harness { } async fn run(&self, task_id: Uuid) -> anyhow::Result<()> { - self.job(task_id).await.run().await + self.job(task_id).run().await } /// Takes every queued job. @@ -210,8 +183,7 @@ impl Harness { created_at: OffsetDateTime, deadline_at: OffsetDateTime, ) { - self.tasks - .store() + self.store .create( NewProvisionerTaskRecord { id: task_id, @@ -219,7 +191,6 @@ impl Harness { target: session_id.to_string(), params: json!({ "provider": provider, "model": "gpt-test", "maxOutputTokens": null }), deadline_at, - job_token: Uuid::new_v4(), }, created_at, ) @@ -228,7 +199,7 @@ impl Harness { } async fn task(&self, task_id: Uuid) -> ProvisionerTaskRecord { - self.tasks + self.store .get(task_id, OffsetDateTime::now_utc()) .await .expect("read task") @@ -243,7 +214,8 @@ impl Harness { self.secrets.clone(), ); task.progress_timeout = progress_timeout; - self.tasks = ProvisionerTaskManager::new(task, Arc::clone(self.tasks.store()), self.job_queue.durable()); + task.first_retry_delay = time::Duration::ZERO; + self.task = Arc::new(task); } /// Working files of the AI analyses of session_id. @@ -369,15 +341,6 @@ async fn spawn_ai_provider_after( (settings, requests) } -fn first_line(path: &Utf8Path) -> String { - std::fs::read_to_string(path) - .expect("chunk") - .lines() - .next() - .expect("line") - .to_owned() -} - // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Secrets #[tokio::test] async fn start_records_the_task_keeps_the_key_and_queues_a_job() { @@ -405,17 +368,9 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { 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 StoredJob { - name, - json, - schedule_for, - } = harness.job_queue.take().expect("job queued"); + let StoredJob { name, json } = harness.job_queue.take().expect("job queued"); assert_eq!(name, KIND); - assert_eq!( - json, - format!(r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, task.job_token) - ); - assert!(schedule_for.is_none()); + assert_eq!(json, format!(r#"{{"taskId":"{task_id}"}}"#)); } // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis @@ -503,7 +458,7 @@ async fn refuses_a_session_without_recording() { assert!(matches!(error, StartError::RecordingNotFound), "{error:?}"); assert!( harness - .tasks + .store .get(task_id, OffsetDateTime::now_utc()) .await .expect("read") @@ -528,7 +483,7 @@ async fn refuses_a_session_that_is_still_recording() { assert!(matches!(error, StartError::RecordingActive), "{error:?}"); assert!( harness - .tasks + .store .get(task_id, OffsetDateTime::now_utc()) .await .expect("read") @@ -573,42 +528,12 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { 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; - // The job of the creation, so only a job queued by the retry is left to find. + // The job of the creation, so a job queued by a retry would be left to find. harness.queued_jobs(); let workspace = harness.workspace(session_id); - harness - .run(task_id) - .await - .expect_err("the job queue runs the job again"); - assert!( - harness.queued_jobs().is_empty(), - "the job queue runs the same job again" - ); - - let task = harness.task(task_id).await; - assert_eq!(task.state, ProvisionerTaskState::Running); - assert_eq!( - task.payload, - json!({ - "step": "describing", - "done": 1, - "total": 3, - "usage": { "inputTokens": 10, "outputTokens": 20, "unreportedRequests": 0 }, - }) - ); - assert!(harness.has_key(task_id), "the key is kept for the retry"); - let plan: Vec = - serde_json::from_slice(&std::fs::read(workspace.join("plan.json")).expect("plan")).expect("plan JSON"); - assert_eq!(plan.len(), 3, "three transcript chunks"); - assert!(workspace.join("chunk-0000.json").exists()); - assert!(!workspace.join("chunk-0001.json").exists()); - let chunk_starts = (0..3) - .map(|index| first_line(&workspace.join(format!("transcript-{index:04}.txt")))) - .collect::>(); - assert_eq!(requests.lock().len(), 2); - - harness.run(task_id).await.expect("second attempt succeeds"); + harness.run(task_id).await.expect("the second attempt succeeds"); + assert!(harness.queued_jobs().is_empty(), "the job runs the next attempt itself"); let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); @@ -630,9 +555,9 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { 5, "only chunks 1 and 2 are sent again, then the summary" ); - for (request, chunk) in requests.iter().zip([0, 1, 1, 2]) { - assert!(request.contains(&chunk_starts[chunk]), "request for chunk {chunk}"); - } + assert_eq!(requests[2], requests[1], "the refused chunk is sent again"); + assert_ne!(requests[1], requests[0]); + assert_ne!(requests[3], requests[1]); assert!(requests.iter().all(|request| !request.contains(API_KEY))); let log = std::fs::read_to_string( @@ -730,20 +655,8 @@ async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { // The job of the creation, so only a job queued by a retry is left to find. harness.queued_jobs(); - for attempt in 1..MAX_ATTEMPTS { - harness - .run(task_id) - .await - .expect_err("the job queue runs the job again"); - assert!( - harness.queued_jobs().is_empty(), - "the job queue runs the same job again after attempt {attempt}" - ); - assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Running); - } - - harness.run(task_id).await.expect("the last attempt is not retried"); - assert!(harness.queued_jobs().is_empty(), "no attempt after the last"); + harness.run(task_id).await.expect("the job runs every attempt itself"); + assert!(harness.queued_jobs().is_empty(), "no other job is queued"); let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); @@ -961,43 +874,6 @@ async fn a_new_task_that_fails_keeps_the_earlier_analysis() { assert!(dir.join("ai-analysis-0.slog").exists()); } -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis -#[tokio::test] -async fn a_restart_fails_an_unfinished_analysis_at_once_and_ends_it() { - let harness = Harness::new().await; - // About 1 MB of transcript: three chunks, so the first attempt stops with the tokens of one. - 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; - harness - .run(task_id) - .await - .expect_err("the job queue runs the job again"); - // As if Gateway stopped after an attempt wrote its log, before its success was recorded. - 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.tasks.resume().await.expect("resume"); - - let task = harness.task(task_id).await; - // oracle: mutant an AI analysis that cannot run after a restart stays unfinished until its next attempt - assert_eq!(task.state, ProvisionerTaskState::Failed, "no job ran yet"); - assert_eq!(task.payload["error"], json!({ "kind": "keyLost" })); - // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis - assert_eq!( - task.payload["usage"], - json!({ "inputTokens": 10, "outputTokens": 20, "unreportedRequests": 0 }) - ); - - harness.run(task_id).await.expect("the new job only ends the Task"); - assert_eq!(requests.lock().len(), 2, "the AI is not asked again"); - assert!(!harness.workspace(session_id).exists(), "the working files are removed"); - assert!(!log.exists(), "the unlisted log is removed"); - assert_eq!(ai_analyses(&dir), json!([])); - assert!(!harness.has_key(task_id)); -} - // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis #[tokio::test] async fn a_failed_task_reports_the_tokens_it_spent() { @@ -1027,8 +903,7 @@ async fn a_task_that_fails_removes_its_unlisted_log() { 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 - .tasks - .store() + .store .fail(task_id, json!({}), OffsetDateTime::now_utc()) .await .expect("fail the Task"); @@ -1039,33 +914,6 @@ async fn a_task_that_fails_removes_its_unlisted_log() { 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_at_startup() { - let mut harness = Harness::new().await; - let session_id = harness.write_session(); - let task_id = harness.start_created(session_id, openai("gpt-test")).await; - - // As if Gateway stopped right after recording the success, before the manifest listed the log. - let dir = harness.recording_path.join(session_id.to_string()); - let log = dir.join(format!(".ai-analysis-{task_id}.slog")); - std::fs::write(&log, "an analysis whose log was never listed").expect("unlisted log"); - harness.record_success(task_id).await; - - harness.tasks.resume().await.expect("resume"); - // The job of its creation never returned, so it stays stored and runs again when Gateway starts. - for StoredJob { json, .. } in harness.queued_jobs() { - let mut job = TaskJob::::read_json(&json) - .expect("valid job") - .into_job(harness.tasks.clone()); - job.run().await.expect("the job ends the Task"); - } - - assert_eq!(ai_analyses(&dir), json!([{ "fileName": "ai-analysis-0.slog" }])); - assert!(dir.join("ai-analysis-0.slog").exists()); - assert!(!log.exists()); -} - // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[tokio::test] async fn a_late_end_of_an_earlier_task_never_lists_its_log_over_a_newer_analysis() { @@ -1213,16 +1061,26 @@ async fn the_analysis_is_in_english_without_a_language() { // 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; + let mut harness = Harness::new().await; + harness.task = Arc::new(RecordingAiAnalysis { + first_retry_delay: time::Duration::hours(1), + ..RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.secrets.clone(), + ) + }); // 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; - harness.run(first).await.expect_err("the job queue runs the job again"); + // As if Gateway stopped while the first Task waited for its next attempt. + tokio::time::timeout(Duration::from_secs(3), harness.run(first)) + .await + .expect_err("the first Task waits for its next attempt"); assert!(harness.workspace(session_id).join("chunk-0000.json").exists()); harness - .tasks - .store() + .store .fail(first, json!({}), OffsetDateTime::now_utc()) .await .expect("end the first Task"); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 72ac1a76c6..22aea2b6ab 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -4,17 +4,18 @@ //! ended. //! Its job is a [`TaskJob`], a job like any other on the job queue, which wraps each attempt in a prelude and a //! completion: the prelude records the start of the attempt, and the completion records how it ended and ends the record. -//! A recoverable failure goes back to the job queue, which runs the same job again later for the next attempt. -//! [`ProvisionerTaskManager`] creates the records, reads them, and recovers them at startup. +//! After a recoverable failure, the job waits [`ProvisionerTask::retry_delay`] and runs the next attempt itself, up to +//! [`MAX_ATTEMPTS`]. +//! [`create`] records a Task and stores its job. //! //! 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 job stops without running anything. -//! At startup, before any job runs, [`ProvisionerTaskManager::resume`] hands every unfinished record to a new job, -//! which also covers a record created just before a crash, before its job was queued. +//! The job queue delivers the job at least once, so after a restart the job picks the record up where it left off, and +//! [`requeue_lost_jobs`] stores the job of a record Gateway stopped before storing it. #[cfg(test)] mod tests; +use std::collections::HashSet; use std::marker::PhantomData; use std::sync::Arc; @@ -32,10 +33,7 @@ use tokio_util::sync::CancellationToken; use uuid::Uuid; use super::{AttemptError, Permanent}; -use crate::job_queue::DurableJobQueue; - -/// Time before a record whose end failed is ended again. -const END_RETRY_DELAY: time::Duration = time::Duration::minutes(5); +use crate::job_queue::{DurableJobQueue, StorableJob}; /// Attempts a record gets: a recoverable failure of the last one fails the record. pub const MAX_ATTEMPTS: u32 = 5; @@ -49,6 +47,9 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// The kind written in the records of this task, and the name of its jobs. const KIND: &'static str; + /// Time from the creation of a record until its first attempt must start. + const START_WINDOW: time::Duration; + /// What the record keeps as asked; never secrets, since records are kept and returned to callers. type Params: Serialize + DeserializeOwned + Send + 'static; @@ -61,13 +62,13 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// Why an attempt of this task fails for now, so a later attempt may succeed. type Recoverable: Serialize + Send + Sync + 'static; - /// 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; + /// Time to wait after attempt `attempts` failed with a recoverable failure, before the next attempt starts. + fn retry_delay(&self, attempts: u32) -> 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; @@ -79,13 +80,6 @@ pub trait ProvisionerTask: Send + Sync + 'static { last: &serde_json::Value, ) -> serde_json::Value; - /// Why an unfinished record fails when Gateway starts again, for a task that cannot outlive a restart, such as one - /// whose secrets are kept in memory only; `None` for a task that goes on. - /// - /// A record that cannot be failed at startup is still handed to a new job, so [`ProvisionerTask::run`] must fail - /// the same way after a restart. - fn restart_failure(&self) -> Option; - /// Runs one attempt. /// /// The attempt should stop soon after [`Attempt::cancel`] is cancelled, which happens once the record is past its @@ -99,8 +93,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 it called again - /// later, until it succeeds. + /// It may be called again for a record that already ended, so it must be idempotent; an error returns from the job, + /// and the job queue runs the job again, which calls it again. async fn on_end(&self, _record: &ProvisionerTaskRecord) -> anyhow::Result<()> { Ok(()) } @@ -169,337 +163,191 @@ fn log_json(value: &impl Serialize) -> serde_json::Value { serde_json::to_value(value).unwrap_or_default() } -/// Creates, reads and recovers the records of the task `T`, whose jobs are [`TaskJob`]. -pub struct ProvisionerTaskManager { - inner: Arc>, -} - -impl Clone for ProvisionerTaskManager { - fn clone(&self) -> Self { - Self { - inner: Arc::clone(&self.inner), - } - } -} - -struct Inner { - task: T, - store: DynProvisionerTaskStore, - job_queue: DurableJobQueue, +/// Records a Task of `T` and stores the job that runs it. +/// +/// `on_created` runs only for a new record, before its job is stored, so what the job needs is in place before it can +/// run; if it fails, no job is stored. A known ID or a busy target creates nothing and stores nothing. +pub async fn create( + store: &DynProvisionerTaskStore, + job_queue: &DurableJobQueue, + id: Uuid, + target: String, + params: &T::Params, + on_created: impl FnOnce(&ProvisionerTaskRecord) -> anyhow::Result<()> + Send, +) -> anyhow::Result { + let params = serde_json::to_value(params).context("serialize the task parameters")?; + let now = OffsetDateTime::now_utc(); + let job = TaskJobUnhydrated:: { + task_id: id, + task: PhantomData, + }; + + let outcome = store + .create( + NewProvisionerTaskRecord { + id, + kind: T::KIND.to_owned(), + target, + params, + deadline_at: deadline(now, T::START_WINDOW), + }, + now, + ) + .await?; + + let CreateOutcome::Created(record) = &outcome else { + return Ok(outcome); + }; + + on_created(record).context("prepare the Task")?; + + job_queue.push(&job, None).await?; + + info!(task.id = %id, task.kind = T::KIND, "Task queued"); + + Ok(outcome) } -impl ProvisionerTaskManager { - pub fn new(task: T, store: DynProvisionerTaskStore, job_queue: DurableJobQueue) -> Self { - Self { - inner: Arc::new(Inner { task, store, job_queue }), +/// Stores again the job of every queued record of `T` that has none among `job_defs`, the JSON of the jobs stored under +/// [`TaskJob::NAME`], as when Gateway stopped between recording a Task and storing its job. +/// +/// Call it at startup, before the job runner claims any job, so a job that runs meanwhile is never taken for lost. +pub async fn requeue_lost_jobs( + store: &DynProvisionerTaskStore, + job_queue: &DurableJobQueue, + job_defs: &[String], +) -> anyhow::Result<()> { + let stored = job_defs + .iter() + .map(|json| TaskJobUnhydrated::::read_json(json).map(|job| job.task_id)) + .collect::>>()?; + + for record in store.list_unfinished(OffsetDateTime::now_utc()).await? { + if record.kind != T::KIND || record.state != ProvisionerTaskState::Queued || stored.contains(&record.id) { + continue; } - } - - pub(crate) fn task(&self) -> &T { - &self.inner.task - } - - /// 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 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 params = serde_json::to_value(params).context("serialize the task parameters")?; - let now = OffsetDateTime::now_utc(); - let job_token = Uuid::new_v4(); - - let outcome = self - .inner - .store - .create( - NewProvisionerTaskRecord { - id, - kind: T::KIND.to_owned(), - target, - params, - deadline_at: deadline(now, self.inner.task.start_window()), - job_token, - }, - now, - ) - .await?; - - let CreateOutcome::Created(record) = &outcome else { - return Ok(outcome); - }; - let queued = match on_created(record) { - Ok(()) => self.enqueue(id, job_token, None).await, - Err(error) => Err(error.context("prepare the Task")), + let job = TaskJobUnhydrated:: { + task_id: record.id, + task: PhantomData, }; + job_queue.push(&job, None).await?; - 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(id, job_token, &serde_json::Value::Null).await; - return Err(error); - } - - info!(task.id = %id, task.kind = T::KIND, "Task queued"); - - Ok(outcome) + info!(task.id = %record.id, task.kind = %record.kind, "Task job queued again"); } - /// Hands every unfinished record of `T` to a new job, since a record created just before Gateway stopped may have - /// none. - /// - /// Each record gets a new `job_token`, so a job left from before finds itself superseded and stops, and the job - /// queue is never read. - /// The new job is stored before a record whose task cannot outlive a restart is failed, so a crash in between never - /// leaves a finished record without a job to end it. - /// A record that was waiting to retry runs at once, since retry timing is no promise to the caller. - /// A record whose new job cannot be stored is failed, since no job is left to run it. - /// Call it at startup, before the job runner claims any job, so no job runs an attempt meanwhile. - pub async fn resume(&self) -> anyhow::Result<()> { - let Inner { task, store, .. } = self.inner.as_ref(); - let now = OffsetDateTime::now_utc(); - - for record in store.list_unfinished(now).await? { - if record.kind != T::KIND { - continue; - } + Ok(()) +} - let job_token = Uuid::new_v4(); +/// The deadline `timeout` after `now`, in whole seconds as the store keeps it, so Gateway and the store agree on when +/// it passes. +fn deadline(now: OffsetDateTime, timeout: time::Duration) -> OffsetDateTime { + (now + timeout).truncate_to_second() +} - match 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; - } - } +/// The job of a record of the task `T`: it wakes the record up, and the record decides what runs. +/// +/// It runs the attempts of `T` until the record ends, each in three steps: the prelude records its start, the body runs +/// [`ProvisionerTask::run`], and the completion records how the last one ended. A record that has ended is ended +/// instead. +pub struct TaskJob { + task_id: Uuid, + task: Arc, + store: DynProvisionerTaskStore, +} - if let Err(error) = self.enqueue(record.id, job_token, None).await { - error!( - task.id = %record.id, - task.kind = %record.kind, - error = format!("{error:#}"), - "Failed to queue the Task again" - ); - self.fail_not_queued(record.id, job_token, &record.payload).await; - continue; - } +/// A [`TaskJob`] as the job queue stores it, before it gets its task and the task store. +#[derive(Serialize, serde::Deserialize)] +#[serde(rename_all = "camelCase", bound = "")] +pub struct TaskJobUnhydrated { + task_id: Uuid, + #[serde(skip)] + task: PhantomData, +} - info!(task.id = %record.id, task.kind = %record.kind, "Task queued again"); +impl TaskJobUnhydrated { + pub fn read_json(json: &str) -> anyhow::Result { + serde_json::from_str(json).context("failed to deserialize TaskJob") + } - let Some(failure) = task.restart_failure() else { - continue; - }; - let failure = Permanent::Kind(failure); - - match store - .fail(record.id, task.failed_payload(&failure, &record.payload), now) - .await - { - Ok(_) => info!( - task.id = %record.id, - task.kind = %record.kind, - failure = %log_json(&failure), - "Task failed, since it cannot outlive a restart" - ), - // Its new job runs the attempt then, which fails the same way. - Err(error) => error!( - task.id = %record.id, - task.kind = %record.kind, - error = format!("{error:#}"), - "Failed to fail the Task that cannot outlive a restart" - ), - } + pub fn into_job(self, task: Arc, store: DynProvisionerTaskStore) -> TaskJob { + TaskJob { + task_id: self.task_id, + task, + store, } - - Ok(()) } +} - #[cfg(test)] - pub(crate) fn store(&self) -> &DynProvisionerTaskStore { - &self.inner.store +impl StorableJob for TaskJobUnhydrated { + fn name(&self) -> &str { + TaskJob::::NAME } - /// Stores the job holding `job_token`, which wakes the record up at `schedule_for`, or at once without it. - async fn enqueue( - &self, - task_id: Uuid, - job_token: Uuid, - schedule_for: Option, - ) -> anyhow::Result<()> { - let job = TaskJob { - task_id, - job_token, - manager: self.clone(), - }; - - self.inner.job_queue.push(job, schedule_for).await + fn write_json(&self) -> anyhow::Result { + serde_json::to_string(self).context("failed to serialize TaskJob") } +} + +impl TaskJob { + /// The name its jobs are stored under. + pub const NAME: &'static str = T::KIND; /// 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. async fn fail( &self, - task_id: Uuid, - job_token: Uuid, error: &Permanent, last: &serde_json::Value, ) -> anyhow::Result<()> { - self.inner - .store + self.store .fail( - task_id, - self.inner.task.failed_payload(error, last), + self.task_id, + self.task.failed_payload(error, last), OffsetDateTime::now_utc(), ) .await .context("record the Task failure")?; - self.end(task_id, job_token).await - } - - /// Fails a record no job will run. - async fn fail_not_queued(&self, task_id: Uuid, job_token: Uuid, last: &serde_json::Value) { - if let Err(error) = self.fail(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"); - } + self.end().await } /// Ends the record, read again so its task sees how it ended. - async fn end(&self, 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(&record, job_token).await, - Ok(None) => Ok(()), - Err(error) => self.end_later(task_id, job_token, error).await, + async fn end(&self) -> anyhow::Result<()> { + match self.store.get(self.task_id, OffsetDateTime::now_utc()).await? { + Some(record) => self.end_record(&record).await, + None => Ok(()), } } /// Does what the finished record leaves to do; a record still unfinished is left to its attempt. - async fn end_record(&self, record: &ProvisionerTaskRecord, job_token: Uuid) -> anyhow::Result<()> { + async fn end_record(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { if !record.state.is_finished() { return Ok(()); } - match self.inner.task.on_end(record).await { - Ok(()) => Ok(()), - Err(error) => self.end_later(record.id, job_token, error).await, - } + self.task.on_end(record).await.context("end the Task") } - /// Wakes the record up again later, to end it then; an error means that job could not be stored. - 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 manager, the job and the store -/// agree on when it passes. -fn deadline(now: OffsetDateTime, timeout: time::Duration) -> OffsetDateTime { - (now + timeout).truncate_to_second() -} - -#[derive(Serialize, serde::Deserialize)] -#[serde(rename_all = "camelCase")] -struct JobDef { - task_id: Uuid, - job_token: Uuid, -} - -/// The job of a record of the task `T`: it wakes the record up, and the record decides what runs. -/// -/// It runs one attempt of `T` in three steps: the prelude records its start, the body runs [`ProvisionerTask::run`], -/// and the completion records how it ended. A record that has ended is ended instead. -pub struct TaskJob { - task_id: Uuid, - job_token: Uuid, - manager: ProvisionerTaskManager, -} - -/// A [`TaskJob`] read back from its JSON, before it gets the manager of `T`. -pub struct TaskJobUnhydrated { - task_id: Uuid, - job_token: Uuid, - task: PhantomData, -} - -impl TaskJobUnhydrated { - pub fn into_job(self, manager: ProvisionerTaskManager) -> TaskJob { - TaskJob { - task_id: self.task_id, - job_token: self.job_token, - manager, - } - } -} - -impl TaskJob { - /// The name its jobs are stored under. - pub const NAME: &'static str = T::KIND; - - pub fn read_json(json: &str) -> anyhow::Result> { - let JobDef { task_id, job_token }: JobDef = - serde_json::from_str(json).context("failed to deserialize TaskJob")?; - - Ok(TaskJobUnhydrated { - task_id, - job_token, - task: PhantomData, - }) - } - - /// The prelude: marks the start of one more attempt on the record, unless another job holds it, it has ended, or - /// it has no attempt left. + /// The prelude: marks the start of one more attempt on the record, unless it has ended. /// /// A record that has ended is ended instead. async fn prelude(&self) -> anyhow::Result> { - let Inner { task, store, .. } = self.manager.inner.as_ref(); + let Self { task, store, .. } = self; let now = OffsetDateTime::now_utc(); // The starting payload carries over what earlier attempts reported, so the record is read first. let Some(record) = store.get(self.task_id, now).await? else { - debug!(task.id = %self.task_id, "Task is unknown, so this attempt does not run"); + error!(task.id = %self.task_id, "Task is unknown, so this attempt does not run"); return Ok(None); }; - if record.kind != T::KIND { - debug!(task.id = %self.task_id, task.kind = %record.kind, "Task of another kind, so this attempt does not run"); - return Ok(None); - } - - if record.job_token != self.job_token { - debug!(task.id = %self.task_id, "Another job runs this Task, so this one stops"); - return Ok(None); - } - if record.state.is_finished() { debug!(task.id = %self.task_id, "Task is over, so this attempt does not run"); - self.manager.end_record(&record, self.job_token).await?; + self.end_record(&record).await?; return Ok(None); } - // A job run again after its last attempt could not record how it ended would otherwise start one past the last. + // A restart in the middle of the last attempt delivers the job again. if MAX_ATTEMPTS <= record.attempts { warn!( task.id = %self.task_id, @@ -507,8 +355,7 @@ impl TaskJob { attempts = record.attempts, "Task has no attempt left, so it fails without running another" ); - self.fail_terminal(&Permanent::Internal, &record.payload, record.deadline_at) - .await?; + self.fail(&Permanent::Internal, &record.payload).await?; return Ok(None); } @@ -516,21 +363,17 @@ impl TaskJob { let deadline_at = deadline(now, task.progress_timeout()); match store - .start_attempt(self.task_id, self.job_token, starting_payload, deadline_at, now) + .start_attempt(self.task_id, starting_payload, deadline_at, now) .await? { AttemptStart::Started(record) => Ok(Some(record)), AttemptStart::Finished(record) => { debug!(task.id = %self.task_id, "Task is over, so this attempt does not run"); - self.manager.end_record(&record, self.job_token).await?; - Ok(None) - } - AttemptStart::Superseded => { - debug!(task.id = %self.task_id, "Another job runs this Task, so this one stops"); + self.end_record(&record).await?; Ok(None) } AttemptStart::Unknown => { - debug!(task.id = %self.task_id, "Task is unknown, so this attempt does not run"); + error!(task.id = %self.task_id, "Task is unknown, so this attempt does not run"); Ok(None) } } @@ -573,8 +416,8 @@ impl TaskJob { attempts: record.attempts, payload: record.payload.clone(), cancel: cancel.clone(), - progress_timeout: self.manager.inner.task.progress_timeout(), - store: Arc::clone(&self.manager.inner.store), + progress_timeout: self.task.progress_timeout(), + store: Arc::clone(&self.store), latest: Arc::clone(&latest), }; @@ -582,8 +425,8 @@ impl TaskJob { // The attempt has its own Tokio task, so a panic fails the record instead of leaving it running. let mut handle = tokio::spawn({ - let manager = self.manager.clone(); - async move { manager.inner.task.run(&attempt, params).await } + let task = Arc::clone(&self.task); + async move { task.run(&attempt, params).await } }); let mut past_deadline = false; @@ -624,11 +467,8 @@ impl TaskJob { (outcome, take_latest(&latest)) } - /// The completion: records how the attempt ended, then ends the record, or returns a recoverable failure to the job - /// queue, which runs this job again later for the next attempt. + /// The completion: records how the last attempt ended, then ends the record. /// - /// An end that cannot be recorded never goes back to the job queue, since that would run the attempt again, see - /// [`TaskJob::fail_terminal`]. /// Recording the end and dropping what the record kept are separate steps, so a failed drop never undoes the end. async fn completion( &self, @@ -636,7 +476,7 @@ impl TaskJob { outcome: Result>, latest: Latest, ) -> anyhow::Result<()> { - let store = &self.manager.inner.store; + let store = &self.store; let (task_id, attempts) = (record.id, record.attempts); let last = &latest.payload; @@ -646,40 +486,26 @@ impl TaskJob { Ok(output) => output, Err(error) => { error!(task.id = %task_id, %error, "Failed to serialize the Task result"); - return self.fail_terminal(&Permanent::Internal, last, latest.deadline_at).await; + return self.fail(&Permanent::Internal, last).await; } }; - 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!( + let succeeded = store + .succeed(task_id, output, OffsetDateTime::now_utc()) + .await + .context("record the Task success")?; + + if succeeded { + 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" - ), - // A retry would do the work again, such as paying for a second analysis. - Err(error) => { - error!( - task.id = %task_id, - error = format!("{error:#}"), - "Failed to record the Task success; the Task fails instead" - ); - return self.fail_terminal(&Permanent::Internal, last, latest.deadline_at).await; - } + ); } - self.manager.end(task_id, self.job_token).await - } - Err(AttemptError::Recoverable(error)) if attempts < MAX_ATTEMPTS => { - let error = log_json(&error); - warn!( - task.id = %task_id, - task.kind = %record.kind, - attempts, - %error, - "Task attempt failed; the job queue runs it again later" - ); - anyhow::bail!("Task attempt {attempts} failed with {error}, to be retried") + self.end().await } Err(AttemptError::Recoverable(error)) => { warn!( @@ -689,8 +515,7 @@ impl TaskJob { error = %log_json(&error), "Task failed after its last attempt" ); - self.fail_terminal(&Permanent::AttemptsExhausted { last: error }, last, latest.deadline_at) - .await + self.fail(&Permanent::AttemptsExhausted { last: error }, last).await } Err(AttemptError::Permanent(error)) => { warn!( @@ -700,54 +525,10 @@ impl TaskJob { error = %log_json(&error), "Task failed" ); - self.fail_terminal(&error, last, latest.deadline_at).await + self.fail(&error, last).await } } } - - /// Records that the record failed with `failure`, then ends it. - /// - /// If the failure cannot be recorded, no error goes back to the job queue, since it would run the attempt again - /// before the deadline: a job is stored to wake the record up past `deadline_at` instead, once the store has failed - /// it and only its end is left. - async fn fail_terminal( - &self, - failure: &Permanent, - last: &serde_json::Value, - deadline_at: OffsetDateTime, - ) -> anyhow::Result<()> { - let Inner { task, store, .. } = self.manager.inner.as_ref(); - - let recorded = store - .fail( - self.task_id, - task.failed_payload(failure, last), - OffsetDateTime::now_utc(), - ) - .await; - - if let Err(error) = recorded { - error!( - task.id = %self.task_id, - error = format!("{error:#}"), - "Failed to fail the Task; it fails at its deadline" - ); - - let after_deadline = deadline_at + time::Duration::SECOND; - - if let Err(error) = self - .manager - .enqueue(self.task_id, self.job_token, Some(after_deadline)) - .await - { - error!(task.id = %self.task_id, error = format!("{error:#}"), "Failed to queue the end of the Task"); - } - - return Ok(()); - } - - self.manager.end(self.task_id, self.job_token).await - } } #[async_trait] @@ -757,20 +538,36 @@ impl job_queue::Job for TaskJob { } fn write_json(&self) -> anyhow::Result { - serde_json::to_string(&JobDef { + TaskJobUnhydrated:: { task_id: self.task_id, - job_token: self.job_token, - }) - .context("failed to serialize TaskJob") + task: PhantomData, + } + .write_json() } async fn run(&mut self) -> anyhow::Result<()> { - let Some(record) = self.prelude().await? else { - return Ok(()); - }; + loop { + let Some(record) = self.prelude().await? else { + return Ok(()); + }; - let (outcome, latest) = self.body(&record).await; + let (outcome, latest) = self.body(&record).await; - self.completion(&record, outcome, latest).await + match outcome { + Err(AttemptError::Recoverable(error)) if record.attempts < MAX_ATTEMPTS => { + let delay = self.task.retry_delay(record.attempts); + warn!( + task.id = %record.id, + task.kind = %record.kind, + attempts = record.attempts, + error = %log_json(&error), + %delay, + "Task attempt failed; the next one runs after a delay" + ); + tokio::time::sleep(std::time::Duration::try_from(delay).unwrap_or_default()).await; + } + outcome => return self.completion(&record, outcome, latest).await, + } + } } } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index c2c6c863f0..a2a493df4b 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -1,8 +1,8 @@ use std::sync::atomic::{AtomicUsize, Ordering}; -use job_queue::Job as _; +use job_queue::{Job as _, JobQueue as _}; use parking_lot::Mutex; -use provisioner_task::{ProvisionerTaskState, ProvisionerTaskStore}; +use provisioner_task::ProvisionerTaskState; use serde_json::json; use super::*; @@ -20,7 +20,7 @@ const PROGRESS: u32 = u32::MAX - 1; #[derive(Debug, Clone, Serialize)] #[serde(tag = "kind", rename_all = "camelCase")] enum EchoFailure { - /// What the attempt needed is gone, as a secret kept in memory is after a restart. + /// What the attempt needed is gone. Gone, } @@ -46,14 +46,14 @@ struct Echo { ended: Arc>>, /// Ends that fail before one succeeds. end_failures: Arc, - /// Why its unfinished records fail when Gateway starts again. - restart_failure: Option, } #[async_trait] impl ProvisionerTask for Echo { const KIND: &'static str = KIND; + const START_WINDOW: time::Duration = time::Duration::hours(1); + type Params = EchoParams; type Output = serde_json::Value; @@ -62,14 +62,14 @@ impl ProvisionerTask for Echo { type Recoverable = EchoRecoverable; - fn start_window(&self) -> time::Duration { - time::Duration::hours(1) - } - fn progress_timeout(&self) -> time::Duration { self.progress_timeout } + fn retry_delay(&self, _: u32) -> time::Duration { + time::Duration::ZERO + } + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value { json!({ "after": last }) } @@ -82,10 +82,6 @@ impl ProvisionerTask for Echo { json!({ "error": error, "last": last }) } - fn restart_failure(&self) -> Option { - self.restart_failure.clone() - } - async fn run(&self, attempt: &Attempt, params: EchoParams) -> Result { self.runs.fetch_add(1, Ordering::SeqCst); @@ -123,86 +119,13 @@ impl ProvisionerTask for Echo { } } -/// The task store, failing to record an end when asked, as a database that is briefly unavailable does. -struct FailingStore { - inner: provisioner_task_libsql::LibSqlProvisionerTaskStore, - /// Successes that fail to be recorded before one is. - success_failures: Arc, - /// Failures that fail to be recorded before one is. - failure_failures: Arc, -} - -/// Whether one more of the `failures` left happens, counting it. -fn fails(failures: &AtomicUsize) -> bool { - failures - .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |left| left.checked_sub(1)) - .is_ok() -} - -#[async_trait] -impl ProvisionerTaskStore for FailingStore { - async fn create(&self, task: NewProvisionerTaskRecord, now: OffsetDateTime) -> anyhow::Result { - self.inner.create(task, now).await - } - - async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result> { - self.inner.get(id, now).await - } - - async fn start_attempt( - &self, - id: Uuid, - job_token: Uuid, - payload: serde_json::Value, - deadline_at: OffsetDateTime, - now: OffsetDateTime, - ) -> anyhow::Result { - self.inner.start_attempt(id, job_token, payload, deadline_at, now).await - } - - async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result { - self.inner.replace_job_token(id, job_token, now).await - } - - async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result> { - self.inner.list_unfinished(now).await - } - - async fn update_running( - &self, - id: Uuid, - payload: serde_json::Value, - deadline_at: OffsetDateTime, - now: OffsetDateTime, - ) -> anyhow::Result { - self.inner.update_running(id, payload, deadline_at, now).await - } - - async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { - if fails(&self.success_failures) { - anyhow::bail!("the success could not be recorded"); - } - - self.inner.succeed(id, payload, now).await - } - - async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { - if fails(&self.failure_failures) { - anyhow::bail!("the failure could not be recorded"); - } - - self.inner.fail(id, payload, now).await - } -} - struct Harness { - manager: ProvisionerTaskManager, + task: Arc, + store: DynProvisionerTaskStore, job_queue: Arc, runs: Arc, ended: Arc>>, end_failures: Arc, - success_failures: Arc, - failure_failures: Arc, } impl Harness { @@ -213,14 +136,6 @@ impl Harness { async fn with_progress_timeout( outcome: fn(u32) -> Result, progress_timeout: time::Duration, - ) -> Self { - Self::build(outcome, progress_timeout, None).await - } - - async fn build( - outcome: fn(u32) -> Result, - progress_timeout: time::Duration, - restart_failure: Option, ) -> Self { let conn = gateway_db::GatewayDb::open_path(":memory:") .await @@ -228,82 +143,67 @@ impl Harness { .connect() .await .expect("connect"); - let success_failures = Arc::new(AtomicUsize::new(0)); - let failure_failures = Arc::new(AtomicUsize::new(0)); - let store = Arc::new(FailingStore { - inner: provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) + let store: DynProvisionerTaskStore = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) .await .expect("open task store"), - success_failures: Arc::clone(&success_failures), - failure_failures: Arc::clone(&failure_failures), - }); + ); let job_queue = TestJobQueue::new().await; let runs = Arc::new(AtomicUsize::new(0)); let ended = Arc::new(Mutex::new(Vec::new())); let end_failures = Arc::new(AtomicUsize::new(0)); - let manager = ProvisionerTaskManager::new( - Echo { - outcome, - progress_timeout, - runs: Arc::clone(&runs), - ended: Arc::clone(&ended), - end_failures: Arc::clone(&end_failures), - restart_failure, - }, - store, - job_queue.durable(), - ); + let task = Arc::new(Echo { + outcome, + progress_timeout, + runs: Arc::clone(&runs), + ended: Arc::clone(&ended), + end_failures: Arc::clone(&end_failures), + }); Self { - manager, + task, + store, job_queue, runs, ended, end_failures, - success_failures, - failure_failures, } } async fn create(&self, value: u32) -> Uuid { let id = Uuid::new_v4(); let outcome = self - .manager - .create(id, format!("target-{id}"), &EchoParams { value }, |_| Ok(())) + .create_task(id, format!("target-{id}"), &EchoParams { value }, |_| Ok(())) .await .expect("create"); assert!(matches!(outcome, CreateOutcome::Created(_)), "{outcome:?}"); id } - /// Takes the oldest queued job, read back as the job runner reads it, with when it runs. - fn queued(&mut self) -> Option<(Box, Option)> { - let StoredJob { - name, - json, - schedule_for, - } = self.job_queue.take()?; - assert_eq!(name, KIND); - let job: Box = Box::new( - TaskJob::::read_json(&json) - .expect("valid job") - .into_job(self.manager.clone()), - ); - Some((job, schedule_for)) + async fn create_task( + &self, + id: Uuid, + target: String, + params: &EchoParams, + on_created: impl FnOnce(&ProvisionerTaskRecord) -> anyhow::Result<()> + Send, + ) -> anyhow::Result { + create::(&self.store, &self.job_queue.durable(), id, target, params, on_created).await } + /// Takes the oldest queued job, read back as the job runner reads it. fn queued_job(&mut self) -> Box { - self.queued().expect("a queued job").0 - } - - /// Takes every queued job. - fn queued_jobs(&mut self) -> Vec { - std::iter::from_fn(|| self.job_queue.take()).collect() + let StoredJob { name, json, .. } = self.job_queue.take().expect("a queued job"); + assert_eq!(name, KIND); + Box::new( + TaskJobUnhydrated::::read_json(&json) + .expect("valid job") + .into_job(Arc::clone(&self.task), Arc::clone(&self.store)), + ) } async fn task(&self, id: Uuid) -> ProvisionerTaskRecord { - self.manager + self.store .get(id, OffsetDateTime::now_utc()) .await .expect("read task") @@ -318,12 +218,10 @@ async fn only_a_new_task_runs_on_created_and_gets_a_job() { let id = Uuid::new_v4(); let created = AtomicUsize::new(0); let create = || { - harness - .manager - .create(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { - created.fetch_add(1, Ordering::SeqCst); - Ok(()) - }) + harness.create_task(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + created.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) }; assert!(matches!(create().await.expect("create"), CreateOutcome::Created(_))); @@ -334,179 +232,20 @@ async fn only_a_new_task_runs_on_created_and_gets_a_job() { assert!(harness.job_queue.take().is_none(), "one job only"); } -// oracle: mutant a job that cannot be queued is ignored, so the Task stays queued and its creation succeeds -#[tokio::test] -async fn a_task_whose_job_cannot_be_queued_fails() { - let harness = Harness::new(|value| Ok(json!(value))).await; - harness.job_queue.close(); - let id = Uuid::new_v4(); - - harness - .manager - .create(id, "target".to_owned(), &EchoParams { value: 1 }, |_| Ok(())) - .await - .expect_err("the job queue is gone"); - - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload, json!({ "error": { "kind": "notQueued" }, "last": null })); - assert_eq!(*harness.ended.lock(), [id], "what was set up for the Task is dropped"); -} - // oracle: mutant the job is queued before the Task is prepared, so a Task that cannot be prepared still gets a job #[tokio::test] -async fn a_task_that_cannot_be_prepared_fails_without_a_job() { +async fn a_task_that_cannot_be_prepared_gets_no_job() { let harness = Harness::new(|value| Ok(json!(value))).await; let id = Uuid::new_v4(); harness - .manager - .create(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + .create_task(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { anyhow::bail!("no room for the secret") }) .await .expect_err("preparing fails"); - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["error"], json!({ "kind": "notQueued" })); assert!(harness.job_queue.take().is_none(), "no job is queued"); - assert_eq!(*harness.ended.lock(), [id]); -} - -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management -#[tokio::test] -async fn a_superseded_job_runs_nothing_once_the_task_is_resumed() { - let mut harness = Harness::new(|value| Ok(json!({ "echo": value }))).await; - let id = harness.create(7).await; - let mut stale = harness.queued_job(); - - harness.manager.resume().await.expect("resume"); - let mut resumed = harness.queued_job(); - assert!( - harness.job_queue.take().is_none(), - "the unfinished Task gets one new job" - ); - - stale.run().await.expect("the stale job stops"); - assert_eq!(harness.runs.load(Ordering::SeqCst), 0); - assert_eq!(harness.task(id).await.attempts, 0, "a stale job counts no attempt"); - assert!(harness.ended.lock().is_empty(), "a stale job ends nothing"); - - resumed.run().await.expect("the new job runs"); - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Succeeded); - assert_eq!(task.payload, json!({ "echo": 7 })); - assert_eq!(*harness.ended.lock(), [id]); - - harness.manager.resume().await.expect("resume again"); - assert!(harness.job_queue.take().is_none(), "a finished Task gets no new job"); -} - -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management -#[tokio::test] -async fn a_task_that_cannot_be_queued_again_fails() { - let mut harness = Harness::new(|value| Ok(json!(value))).await; - let id = harness.create(1).await; - harness.queued_job(); - harness.job_queue.close(); - - harness.manager.resume().await.expect("resume"); - - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Failed, "no job is left to run it"); - assert_eq!(task.payload["error"], json!({ "kind": "notQueued" })); - assert_eq!(*harness.ended.lock(), [id]); -} - -// oracle: mutant resume skips a Task whose job may still be stored, so a Task created just before a crash never runs -#[tokio::test] -async fn resume_gives_every_unfinished_task_a_new_job() { - let mut harness = Harness::new(|value| Ok(json!(value))).await; - let first = harness.create(1).await; - let second = harness.create(2).await; - let tokens = [ - harness.task(first).await.job_token, - harness.task(second).await.job_token, - ]; - // The jobs of the creation may be stored or lost; the manager never asks which. - harness.queued_jobs(); - - harness.manager.resume().await.expect("resume"); - - assert_eq!(harness.queued_jobs().len(), 2, "one new job for each unfinished Task"); - assert_ne!(harness.task(first).await.job_token, tokens[0]); - assert_ne!(harness.task(second).await.job_token, tokens[1]); -} - -// oracle: mutant a Task that cannot outlive a restart stays unfinished until a job runs it -#[tokio::test] -async fn a_task_that_cannot_outlive_a_restart_fails_at_startup_and_is_ended() { - let mut harness = Harness::build( - |value| Ok(json!(value)), - time::Duration::hours(1), - Some(EchoFailure::Gone), - ) - .await; - let id = harness.create(1).await; - let mut stale = harness.queued_job(); - - harness.manager.resume().await.expect("resume"); - - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Failed, "it fails before any job runs"); - assert_eq!(task.payload, json!({ "error": { "kind": "gone" }, "last": null })); - - let mut ending = harness.queued_job(); - stale.run().await.expect("the stale job stops"); - assert!(harness.ended.lock().is_empty(), "a stale job ends nothing"); - - ending.run().await.expect("the new job ends the Task"); - assert_eq!(harness.runs.load(Ordering::SeqCst), 0, "no attempt runs"); - assert_eq!(*harness.ended.lock(), [id]); -} - -// oracle: mutant a Task whose failure at startup cannot be recorded gets no new job, so nothing ends it -#[tokio::test] -async fn a_task_whose_restart_failure_cannot_be_recorded_still_gets_a_new_job() { - let mut harness = Harness::build( - |value| Ok(json!(value)), - time::Duration::hours(1), - Some(EchoFailure::Gone), - ) - .await; - let id = harness.create(1).await; - let token = harness.task(id).await.job_token; - harness.queued_jobs(); - harness.failure_failures.store(1, Ordering::SeqCst); - - harness.manager.resume().await.expect("resume"); - - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Queued, "the failure was not recorded"); - assert_ne!(task.job_token, token); - assert_eq!(harness.queued_jobs().len(), 1, "its new job runs it"); -} - -// oracle: mutant resume resets the attempts of a Task, so a Task that fails at each restart is retried forever -#[tokio::test] -async fn resume_keeps_the_attempts_a_task_made() { - let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).await; - let id = harness.create(1).await; - harness - .queued_job() - .run() - .await - .expect_err("the job queue runs the job again"); - - harness.manager.resume().await.expect("resume"); - harness - .queued_job() - .run() - .await - .expect_err("the job queue runs the job again"); - - assert_eq!(harness.task(id).await.attempts, 2); } // oracle: mutant a panicking attempt fails the Task as internal instead of panicked @@ -546,20 +285,16 @@ async fn a_permanent_failure_fails_the_task_at_once() { async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).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 runs the job again"); - - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Running, "{attempt}"); - assert!( - harness.job_queue.take().is_none(), - "the job queue runs the same job again" - ); - } - job.run().await.expect("the last attempt"); - assert!(harness.job_queue.take().is_none(), "no attempt after the last"); + harness + .queued_job() + .run() + .await + .expect("the job runs every attempt itself"); + assert!(harness.job_queue.take().is_none(), "no other job is queued"); + assert_eq!( + harness.runs.load(Ordering::SeqCst), + usize::try_from(MAX_ATTEMPTS).expect("a few attempts") + ); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); @@ -570,104 +305,64 @@ async fn a_recoverable_failure_is_retried_up_to_the_last_attempt() { ); } -// oracle: mutant a success that cannot be recorded is retried, so the work is done and paid for twice +// oracle: mutant the prelude starts an attempt whatever the attempts made, so a restart during the last attempt runs one past it #[tokio::test] -async fn a_success_that_cannot_be_recorded_fails_the_task_instead_of_running_again() { +async fn a_job_delivered_again_after_the_last_attempt_fails_the_task() { let mut harness = Harness::new(|value| Ok(json!(value))).await; let id = harness.create(1).await; - harness.success_failures.store(1, Ordering::SeqCst); + let now = OffsetDateTime::now_utc(); + for _ in 0..MAX_ATTEMPTS { + harness + .store + .start_attempt(id, json!(null), now + time::Duration::hours(1), now) + .await + .expect("start an attempt"); + } harness.queued_job().run().await.expect("no retry"); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); assert_eq!(task.payload["error"], json!({ "kind": "internal" })); - assert_eq!(harness.runs.load(Ordering::SeqCst), 1, "the work is not done again"); + assert_eq!(harness.runs.load(Ordering::SeqCst), 0, "no attempt past the last"); assert_eq!(*harness.ended.lock(), [id]); - assert!(harness.job_queue.take().is_none(), "no job is left"); } -// oracle: mutant an end that cannot be recorded goes back to the job queue, which runs the attempt again +// oracle: mutant a queued Task whose job was never stored stays queued until its deadline after a restart #[tokio::test] -async fn an_end_that_cannot_be_recorded_is_left_to_a_job_past_the_deadline() { +async fn a_queued_task_without_a_job_gets_one_at_startup() { let mut harness = Harness::new(|value| Ok(json!(value))).await; - let id = harness.create(1).await; - harness.success_failures.store(1, Ordering::SeqCst); - harness.failure_failures.store(1, Ordering::SeqCst); - + harness.create(1).await; + // The job stays in the job queue; only the copy handed to the test is taken. + harness.queued_job(); + let id = Uuid::new_v4(); + let now = OffsetDateTime::now_utc(); harness - .queued_job() - .run() + .store + .create( + NewProvisionerTaskRecord { + id, + kind: KIND.to_owned(), + target: id.to_string(), + params: json!({ "value": 2 }), + deadline_at: now + time::Duration::hours(1), + }, + now, + ) .await - .expect("nothing goes back to the job queue"); - - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Running, "nothing could be recorded"); - assert_eq!(harness.runs.load(Ordering::SeqCst), 1); - assert!( - harness.ended.lock().is_empty(), - "nothing is dropped while the Task may still run" - ); - let (_, schedule_for) = harness.queued().expect("a job ends it later"); - assert!( - schedule_for.is_some_and(|at| task.deadline_at < at), - "the job runs once the store fails the Task" - ); -} - -// oracle: mutant an end that can be neither recorded nor queued returns an error, so the job queue redoes the paid work -#[tokio::test] -async fn a_task_whose_end_cannot_be_recorded_or_queued_is_not_retried() { - let mut harness = Harness::new(|value| Ok(json!(value))).await; - let id = harness.create(1).await; - let mut job = harness.queued_job(); - harness.job_queue.close(); - harness.success_failures.store(1, Ordering::SeqCst); - harness.failure_failures.store(1, Ordering::SeqCst); - - job.run().await.expect("the job queue does not run the attempt again"); - - assert_eq!(harness.runs.load(Ordering::SeqCst), 1); - assert_eq!( - harness.task(id).await.state, - ProvisionerTaskState::Running, - "the store fails it at its deadline" - ); -} - -// oracle: mutant the prelude starts an attempt whatever the attempts made, so a Task runs one attempt past the last -#[tokio::test] -async fn a_last_attempt_whose_failure_cannot_be_recorded_is_not_run_again() { - let mut harness = Harness::new(|_| Err(AttemptError::Recoverable(EchoRecoverable::Busy))).await; - let id = harness.create(1).await; - let mut job = harness.queued_job(); - for _ in 1..MAX_ATTEMPTS { - job.run().await.expect_err("the job queue runs the job again"); - } - harness.failure_failures.store(1, Ordering::SeqCst); - - job.run().await.expect("its end is left to a job past the deadline"); - let task = harness.task(id).await; - assert_eq!( - task.state, - ProvisionerTaskState::Running, - "the failure was not recorded" - ); - let (_, schedule_for) = harness.queued().expect("a job ends it later"); - assert!(schedule_for.is_some_and(|at| task.deadline_at < at)); + .expect("create"); - // As if Gateway stopped before the deadline, so a new job finds the Task still running. - harness.manager.resume().await.expect("resume"); - harness.queued_job().run().await.expect("no retry"); + let job_defs = harness.job_queue.job_defs(KIND).await.expect("stored jobs"); + requeue_lost_jobs::(&harness.store, &harness.job_queue.durable(), &job_defs) + .await + .expect("requeue"); - let task = harness.task(id).await; - assert_eq!( - harness.runs.load(Ordering::SeqCst), - usize::try_from(MAX_ATTEMPTS).expect("a few attempts"), - "no attempt past the last" + harness.queued_job().run().await.expect("the job runs the Task"); + assert!( + harness.job_queue.take().is_none(), + "the Task that has a job gets no other" ); - assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!(task.payload["error"], json!({ "kind": "internal" })); + assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Succeeded); } // oracle: mutant a Task whose parameters cannot be read runs its attempt, or stays unfinished @@ -675,11 +370,9 @@ async fn a_last_attempt_whose_failure_cannot_be_recorded_is_not_run_again() { async fn invalid_parameters_fail_the_task() { let harness = Harness::new(|value| Ok(json!(value))).await; let id = Uuid::new_v4(); - let job_token = Uuid::new_v4(); let now = OffsetDateTime::now_utc(); harness - .manager - .store() + .store .create( NewProvisionerTaskRecord { id, @@ -687,16 +380,15 @@ async fn invalid_parameters_fail_the_task() { target: id.to_string(), params: json!({ "value": "not a number" }), deadline_at: now + time::Duration::hours(1), - job_token, }, now, ) .await .expect("create"); - let json = json!({ "taskId": id, "jobToken": job_token }).to_string(); - let mut job = TaskJob::::read_json(&json) + let json = json!({ "taskId": id }).to_string(); + let mut job = TaskJobUnhydrated::::read_json(&json) .expect("valid job") - .into_job(harness.manager.clone()); + .into_job(Arc::clone(&harness.task), Arc::clone(&harness.store)); job.run().await.expect("no retry"); @@ -708,27 +400,21 @@ async fn invalid_parameters_fail_the_task() { // 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() { +async fn an_end_that_fails_is_done_again_when_the_job_queue_runs_the_job_again() { 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 (next, schedule_for) = harness.queued().expect("a new job"); - assert!( - schedule_for.is_some_and(|at| OffsetDateTime::now_utc() < at), - "it runs later" - ); + job.run().await.expect_err("the job queue runs the job again"); 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.runs.load(Ordering::SeqCst), 1, "the work is not done again"); assert_eq!(*harness.ended.lock(), [id]); assert!(harness.job_queue.take().is_none(), "no job is left"); } @@ -768,11 +454,10 @@ async fn the_job_is_read_back_from_its_json() { let job = harness.queued_job(); let json = job.write_json().expect("job JSON"); - let read = TaskJob::::read_json(&json) + let read = TaskJobUnhydrated::::read_json(&json) .expect("valid job") - .into_job(harness.manager.clone()); + .into_job(Arc::clone(&harness.task), Arc::clone(&harness.store)); assert_eq!(read.task_id, id); - assert_eq!(read.job_token, harness.task(id).await.job_token); assert_eq!(job.name(), KIND); } diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index 79a7f94f36..0c9666b85c 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -3,8 +3,11 @@ use std::time::Duration; use anyhow::Context as _; use devolutions_gateway::config::{Conf, ConfHandle}; +use devolutions_gateway::job_queue::{InMemoryJobDependencies, JobQueueTask, JobRunnerTask}; use devolutions_gateway::listener::{GatewayListener, ListenerUrls}; use devolutions_gateway::log::GatewayLog; +use devolutions_gateway::provisioner_tasks::recording_ai_analysis::RecordingAiAnalysis; +use devolutions_gateway::provisioner_tasks::runner::{TaskJob, requeue_lost_jobs}; use devolutions_gateway::recording::recording_message_channel; use devolutions_gateway::session::session_manager_channel; use devolutions_gateway::subscriber::subscriber_channel; @@ -351,28 +354,27 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { ); let secrets = devolutions_gateway::secret_store::SecretStore::default(); + let durable_job_queue = job_queue_ctx.durable(); - let in_memory_job_dependencies = devolutions_gateway::job_queue::InMemoryJobDependencies { - conf_handle: conf_handle.clone(), - recordings: recording_manager_handle.clone(), - secrets: secrets.clone(), - task_store, - job_queue: job_queue_ctx.durable(), + // Before the job runner starts, so a job that runs meanwhile is never taken for lost. + let requeued = async { + let job_defs = job_queue_ctx.job_defs(TaskJob::::NAME).await?; + requeue_lost_jobs::(&task_store, &durable_job_queue, &job_defs).await }; - let provisioner_tasks = devolutions_gateway::provisioner_tasks::task_manager(&in_memory_job_dependencies); - job_queue_ctx - .job_queue_handle - .set_in_memory_job_dependencies(in_memory_job_dependencies)?; - - // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. - // Unfinished Tasks still fail at their deadline, so this must not keep the other features from starting. - if let Err(error) = provisioner_tasks.resume().await { + if let Err(error) = requeued.await { error!( error = format!("{error:#}"), - "Failed to queue the unfinished provisioner tasks again" + "Failed to queue again the jobs of the provisioner tasks" ); } + let in_memory_job_dependencies = InMemoryJobDependencies { + conf_handle: conf_handle.clone(), + recordings: recording_manager_handle.clone(), + secrets: secrets.clone(), + task_store: Arc::clone(&task_store), + }; + let state = DgwState { conf_handle: conf_handle.clone(), token_cache: Arc::clone(&token_cache), @@ -387,6 +389,9 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { monitoring_state, traffic_audit_handle: traffic_audit_task.handle(), agent_tunnel_handle, + task_store, + secrets: secrets.clone(), + durable_job_queue, }; for listener in &conf.listeners { @@ -453,8 +458,8 @@ 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::JobQueueTask::new(job_queue_ctx)); + tasks.register(JobRunnerTask::new(&job_queue_ctx, in_memory_job_dependencies)); + tasks.register(JobQueueTask::new(job_queue_ctx)); tasks.register(traffic_audit_task); diff --git a/devolutions-gateway/tests/preflight.rs b/devolutions-gateway/tests/preflight.rs index 1cbcf84808..71a312c624 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 7175d3bb59..3c9e4a63cb 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 8e090e9d10..25755d7e0f 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 6d21bff5314df80ec35dda7108ee26a1e860bf6e Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 29 Oct 2026 20:00:46 -0400 Subject: [PATCH 79/80] refactor(dgw): store a Task's job with its record and split the task modules The provisioner task service, the one task holding the task store and the job queue, stores the job inside the transaction that records the Task, so a Task is never recorded without its job; the handlers only ask it, as they ask the recording manager. Nothing is left to find at startup, so `requeue_lost_jobs` and the job queue's `job_defs` go. A job stored for a Task whose transaction rolled back stays in the job queue, so the job token comes back: a job runs only the Task recorded with its token. Migration 2 is gone, and the table keeps the schema of #2041. A success is published in the write that records it (`publish`), so a succeeded Task always has its output listed; the unlisted log a Task kept until its end, and the end that listed it, are gone with it. `provisioner_tasks::runner` is split into `attempt`, `job` and `service`, with the `ProvisionerTask` trait in `provisioner_tasks` itself. The service pushes a `TaskJob` like any other job, so the job queue crates are as on master. A `TaskJob` owns its task and runs the attempt inline: a panic ends the job, and the record fails at its deadline. The API answers with typed DTOs: the handler reads `params` and the payload of a Task into the shapes the OpenAPI document declares, so a record of another shape is an error on Gateway's side and never reaches the caller. The intent states the guarantee a Task gives, how recovery is split between the job queue and the job, that only a success has a guarantee on its result, and that the secret store is shared. The AI analysis API intent moves next to `api/tasks/`, with its routes. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 Co-Authored-By: Claude Fable 5.1 --- .../src/recording_analysis/tests.rs | 6 +- crates/job-queue-libsql/src/lib.rs | 29 +- crates/job-queue/src/lib.rs | 5 +- .../migrations/02_drop_job_token.sql | 1 - crates/provisioner-task-libsql/src/lib.rs | 218 +++++-- crates/provisioner-task/INTENT.md | 4 +- crates/provisioner-task/src/lib.rs | 32 +- devolutions-gateway/src/api/jrec.intent.md | 11 - .../api/tasks/recording_ai_analysis.intent.md | 9 + .../src/api/tasks/recording_ai_analysis.rs | 272 ++++++++- devolutions-gateway/src/job_queue.rs | 128 ++-- devolutions-gateway/src/lib.rs | 12 +- devolutions-gateway/src/openapi.rs | 174 +----- .../src/provisioner_tasks/INTENT.md | 41 +- .../src/provisioner_tasks/ai.rs | 2 +- .../src/provisioner_tasks/attempt.rs | 52 ++ .../src/provisioner_tasks/job.rs | 323 ++++++++++ .../src/provisioner_tasks/mod.rs | 110 +++- .../recording_ai_analysis/mod.rs | 208 ++----- .../recording_ai_analysis/tests.rs | 220 ++----- .../src/provisioner_tasks/runner/mod.rs | 573 ------------------ .../src/provisioner_tasks/service.rs | 184 ++++++ .../provisioner_tasks/{runner => }/tests.rs | 241 +++++--- devolutions-gateway/src/recording.rs | 20 +- .../src/secret_store.intent.md | 5 + devolutions-gateway/src/service.rs | 26 +- testsuite/tests/cli/dgw/ai_analysis.rs | 18 +- 27 files changed, 1493 insertions(+), 1431 deletions(-) delete mode 100644 crates/provisioner-task-libsql/migrations/02_drop_job_token.sql create mode 100644 devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md create mode 100644 devolutions-gateway/src/provisioner_tasks/attempt.rs create mode 100644 devolutions-gateway/src/provisioner_tasks/job.rs delete mode 100644 devolutions-gateway/src/provisioner_tasks/runner/mod.rs create mode 100644 devolutions-gateway/src/provisioner_tasks/service.rs rename devolutions-gateway/src/provisioner_tasks/{runner => }/tests.rs (67%) diff --git a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs index b14e508f55..702381e8fe 100644 --- a/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs +++ b/crates/devolutions-gateway-ai/src/recording_analysis/tests.rs @@ -266,7 +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 + // oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis assert_eq!( progress .iter() @@ -340,7 +340,7 @@ async fn analysing_again_only_asks_about_the_chunks_left() { ); } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn requests_without_a_token_count_are_counted_apart() { let session = Session::new(10); @@ -359,7 +359,7 @@ async fn requests_without_a_token_count_are_counted_apart() { ); } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn a_failed_analysis_reports_the_tokens_it_spent() { // About 1 MB of transcript: three chunks. diff --git a/crates/job-queue-libsql/src/lib.rs b/crates/job-queue-libsql/src/lib.rs index 60fae7daf1..0efbfd23f0 100644 --- a/crates/job-queue-libsql/src/lib.rs +++ b/crates/job-queue-libsql/src/lib.rs @@ -3,7 +3,7 @@ extern crate tracing; use anyhow::Context as _; use async_trait::async_trait; -use job_queue::{JobCtx, JobQueue, JobReader, RunnerWaker}; +use job_queue::{DynJob, JobCtx, JobQueue, JobReader, RunnerWaker}; use libsql::Connection; use time::OffsetDateTime; use ulid::Ulid; @@ -199,7 +199,7 @@ impl JobQueue for LibSqlJobQueue { Ok(()) } - async fn push_job(&self, name: &str, def: &str, schedule_for: Option) -> anyhow::Result<()> { + async fn push_job(&self, job: &DynJob, schedule_for: Option) -> anyhow::Result<()> { let sql_query = "INSERT INTO job_queue (id, scheduled_for, failed_attempts, status, name, def) VALUES (:id, :scheduled_for, :failed_attempts, :status, :name, jsonb(:def))"; @@ -216,8 +216,8 @@ impl JobQueue for LibSqlJobQueue { (":scheduled_for", schedule_for.unix_timestamp()), (":failed_attempts", 0), (":status", JobStatus::Queued as u32), - (":name", name), - (":def", def), + (":name", job.name()), + (":def", job.write_json()?), ); trace!(%sql_query, ?params, "Pushing a new job"); @@ -402,27 +402,6 @@ impl JobQueue for LibSqlJobQueue { Ok(Some(scheduled_for)) } - - async fn job_defs(&self, name: &str) -> anyhow::Result> { - let sql_query = "SELECT json(def) FROM job_queue WHERE name = $1"; - - let params = [name]; - - trace!(%sql_query, ?params, "Fetching the jobs stored under a name"); - - let mut rows = self - .conn - .query(sql_query, params) - .await - .context("failed to execute SQL query")?; - - let mut defs = Vec::new(); - while let Some(row) = rows.next().await.context("failed to read the row")? { - defs.push(row.get::(0).context("failed to read def value")?); - } - - Ok(defs) - } } // Typically, migrations should not be modified once released, and we should only be appending to this list. diff --git a/crates/job-queue/src/lib.rs b/crates/job-queue/src/lib.rs index 0c1d8c4d03..93232bd868 100644 --- a/crates/job-queue/src/lib.rs +++ b/crates/job-queue/src/lib.rs @@ -44,7 +44,7 @@ pub trait JobQueue: Send + Sync { /// Pushes a new job into the queue /// /// This function should ideally call `RunnerWaker::wake()` once the job is enqueued. - async fn push_job(&self, name: &str, def: &str, schedule_for: Option) -> anyhow::Result<()>; + async fn push_job(&self, job: &DynJob, schedule_for: Option) -> anyhow::Result<()>; /// Fetches at most `number_of_jobs` from the queue async fn claim_jobs(&self, reader: &dyn JobReader, number_of_jobs: usize) -> anyhow::Result>; @@ -62,9 +62,6 @@ pub trait JobQueue: Send + Sync { /// Retrieves the closest future scheduled date async fn next_scheduled_date(&self) -> anyhow::Result>; - - /// Returns the JSON of every job stored under `name` - async fn job_defs(&self, name: &str) -> anyhow::Result>; } pub struct JobCtx { diff --git a/crates/provisioner-task-libsql/migrations/02_drop_job_token.sql b/crates/provisioner-task-libsql/migrations/02_drop_job_token.sql deleted file mode 100644 index dbfbaf096d..0000000000 --- a/crates/provisioner-task-libsql/migrations/02_drop_job_token.sql +++ /dev/null @@ -1 +0,0 @@ -ALTER TABLE provisioner_task_records DROP COLUMN job_token; diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 7a786e76bf..2a66cae6de 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -11,7 +11,7 @@ use anyhow::{Context as _, bail}; use async_trait::async_trait; use libsql::{Connection, Row, TransactionBehavior, params}; use provisioner_task::{ - AttemptStart, CreateOutcome, NewProvisionerTaskRecord, ProvisionerTaskRecord, ProvisionerTaskState, + AttemptStart, BeforeCommit, CreateOutcome, NewProvisionerTaskRecord, ProvisionerTaskRecord, ProvisionerTaskState, ProvisionerTaskStore, }; use time::OffsetDateTime; @@ -19,15 +19,12 @@ use tokio::sync::Mutex; use uuid::Uuid; /// Every schema change ever made to the `provisioner_task_*` tables, oldest first. Only ever append to it. -const MIGRATIONS: &[&str] = &[ - include_str!("../migrations/01_provisioner_task.sql"), - include_str!("../migrations/02_drop_job_token.sql"), -]; +const MIGRATIONS: &[&str] = &[include_str!("../migrations/01_provisioner_task.sql")]; const SCHEMA_VERSION_KEY: &str = "schema_version"; const COLUMNS: &str = - "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at"; + "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at, job_token"; /// Tasks kept in the `provisioner_task_records` table of `gateway.db`. pub struct LibSqlProvisionerTaskStore { @@ -46,7 +43,7 @@ impl LibSqlProvisionerTaskStore { Ok(Self { conn: Mutex::new(conn) }) } - /// Ends the Task as `state` if it is currently in one of `from`. + /// Ends the Task as `state` if it is currently in one of `from`, once `before_commit` has run. /// /// A Task past its deadline is failed for it first, and then not ended again. async fn finish( @@ -56,6 +53,7 @@ impl LibSqlProvisionerTaskStore { state: ProvisionerTaskState, payload: serde_json::Value, now: OffsetDateTime, + before_commit: Option>, ) -> anyhow::Result { let conn = self.conn.lock().await; let tx = conn @@ -74,6 +72,13 @@ impl LibSqlProvisionerTaskStore { return Ok(false); } + if let Some(before_commit) = before_commit + && let Err(error) = before_commit.await + { + tx.rollback().await.context("end task finish")?; + return Err(error); + } + tx.execute( "UPDATE provisioner_task_records SET state = ?2, payload = ?3, finished_at = ?4 WHERE id = ?1", params![ @@ -154,7 +159,12 @@ async fn fail_if_overdue( #[async_trait] impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { - async fn create(&self, task: NewProvisionerTaskRecord, now: OffsetDateTime) -> anyhow::Result { + async fn create( + &self, + task: NewProvisionerTaskRecord, + now: OffsetDateTime, + store_job: BeforeCommit<'_>, + ) -> anyhow::Result { let conn = self.conn.lock().await; let tx = conn .transaction_with_behavior(TransactionBehavior::Immediate) @@ -185,8 +195,8 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { tx.execute( "INSERT INTO provisioner_task_records - (id, kind, target, params, state, payload, created_at, deadline_at) - VALUES (?1, ?2, ?3, ?4, 'queued', 'null', ?5, ?6)", + (id, kind, target, params, state, payload, created_at, deadline_at, job_token) + VALUES (?1, ?2, ?3, ?4, 'queued', 'null', ?5, ?6, ?7)", params![ task.id.to_string(), task.kind.clone(), @@ -194,11 +204,17 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { task.params.to_string(), now.unix_timestamp(), task.deadline_at.unix_timestamp(), + task.job_token.to_string(), ], ) .await .context("insert task")?; + if let Err(error) = store_job.await { + tx.rollback().await.context("end task creation")?; + return Err(error.context("store the task job")); + } + let created = find(&tx, task.id).await?.context("created task is missing")?; tx.commit().await.context("commit task creation")?; @@ -339,13 +355,20 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(running) } - async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + async fn succeed( + &self, + id: Uuid, + payload: serde_json::Value, + now: OffsetDateTime, + publish: BeforeCommit<'_>, + ) -> anyhow::Result { self.finish( id, &[ProvisionerTaskState::Running], ProvisionerTaskState::Succeeded, payload, now, + Some(publish), ) .await } @@ -357,6 +380,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { ProvisionerTaskState::Failed, payload, now, + None, ) .await } @@ -415,6 +439,11 @@ fn task_from_row(row: &Row) -> anyhow::Result { started_at: optional_timestamp(8, "start time")?, finished_at: optional_timestamp(9, "finish time")?, deadline_at: timestamp(10, "deadline")?, + job_token: row + .get::(11) + .context("decode task job token")? + .parse() + .context("parse task job token")?, }) } @@ -513,6 +542,11 @@ mod tests { OffsetDateTime::from_unix_timestamp(1_800_000_000).expect("valid timestamp") } + /// Nothing to do before the commit. + fn nothing() -> BeforeCommit<'static> { + Box::pin(async { Ok(()) }) + } + fn new_task(target: &str) -> NewProvisionerTaskRecord { NewProvisionerTaskRecord { id: Uuid::new_v4(), @@ -520,11 +554,12 @@ mod tests { target: String::from(target), params: serde_json::json!({ "provider": "openai", "model": "gpt-6-luna" }), deadline_at: now() + Duration::hours(2), + job_token: Uuid::new_v4(), } } async fn create(store: &LibSqlProvisionerTaskStore, task: NewProvisionerTaskRecord) -> CreateOutcome { - store.create(task, now()).await.expect("create task") + store.create(task, now(), nothing()).await.expect("create task") } #[tokio::test] @@ -539,19 +574,45 @@ mod tests { assert_eq!(created.state, ProvisionerTaskState::Queued); assert_eq!(created.attempts, 0); assert_eq!(created.params, task.params); + assert_eq!(created.job_token, task.job_token); assert_eq!(store.get(task.id, now()).await.expect("get task"), Some(created)); } + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management + #[tokio::test] + async fn a_task_whose_job_is_not_stored_is_not_recorded() { + let store = store().await; + let task = new_task("session-1"); + + let error = store + .create( + task.clone(), + now(), + Box::pin(async { anyhow::bail!("the job queue is gone") }), + ) + .await + .expect_err("the job is not stored"); + + assert!(error.to_string().contains("store the task job"), "{error:#}"); + assert_eq!(store.get(task.id, now()).await.expect("get task"), None); + assert!(matches!(create(&store, task).await, CreateOutcome::Created(_))); + } + #[tokio::test] async fn same_request_again_returns_the_existing_task() { let store = store().await; let task = new_task("session-1"); create(&store, task.clone()).await; - let CreateOutcome::Existing(existing) = create(&store, task.clone()).await else { + let mut again = task.clone(); + again.job_token = Uuid::new_v4(); + let CreateOutcome::Existing(existing) = create(&store, again).await else { panic!("the same request must return the existing task"); }; - assert_eq!(existing.id, task.id); + assert_eq!( + existing.job_token, task.job_token, + "the job of the first request keeps the task" + ); } #[tokio::test] @@ -608,7 +669,7 @@ mod tests { assert!( !store - .succeed(task.id, serde_json::json!({}), now()) + .succeed(task.id, serde_json::json!({}), now(), nothing()) .await .expect("succeed queued task") ); @@ -629,6 +690,74 @@ mod tests { ); } + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + #[tokio::test] + async fn a_task_whose_result_is_not_published_stays_running() { + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + store + .start_attempt(task.id, serde_json::json!({}), now() + Duration::hours(2), now()) + .await + .expect("start"); + let published = std::sync::atomic::AtomicBool::new(false); + + let error = store + .succeed( + task.id, + serde_json::json!({ "log": "ai-analysis-0.slog" }), + now(), + Box::pin(async { anyhow::bail!("the manifest cannot be written") }), + ) + .await + .expect_err("the success is not recorded"); + assert!(error.to_string().contains("manifest"), "{error:#}"); + assert_eq!( + store.get(task.id, now()).await.expect("get task").expect("task").state, + ProvisionerTaskState::Running + ); + + assert!( + store + .succeed( + task.id, + serde_json::json!({ "log": "ai-analysis-0.slog" }), + now(), + Box::pin(async { + published.store(true, std::sync::atomic::Ordering::SeqCst); + Ok(()) + }), + ) + .await + .expect("succeed") + ); + assert!(published.load(std::sync::atomic::Ordering::SeqCst)); + assert_eq!( + store.get(task.id, now()).await.expect("get task").expect("task").state, + ProvisionerTaskState::Succeeded + ); + } + + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis + #[tokio::test] + async fn a_task_that_is_not_running_does_not_publish() { + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + + assert!( + !store + .succeed( + task.id, + serde_json::json!({}), + now(), + Box::pin(async { panic!("a queued Task publishes nothing") }), + ) + .await + .expect("succeed queued task") + ); + } + #[tokio::test] async fn overdue_task_is_failed_and_frees_its_target() { let store = store().await; @@ -717,7 +846,7 @@ mod tests { ); assert!( !store - .succeed(task.id, serde_json::json!({}), now()) + .succeed(task.id, serde_json::json!({}), now(), nothing()) .await .expect("finish again") ); @@ -758,7 +887,12 @@ mod tests { assert!( store - .succeed(task.id, serde_json::json!({ "log": "ai-analysis-0.slog" }), now()) + .succeed( + task.id, + serde_json::json!({ "log": "ai-analysis-0.slog" }), + now(), + nothing() + ) .await .expect("succeed") ); @@ -799,7 +933,7 @@ mod tests { ); assert!( !store - .succeed(overdue.id, serde_json::json!({ "log": "late" }), late) + .succeed(overdue.id, serde_json::json!({ "log": "late" }), late, nothing()) .await .expect("succeed late") ); @@ -963,7 +1097,9 @@ mod tests { 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 + .succeed(Uuid::new_v4(), serde_json::json!({}), now(), nothing()) + .await, store.fail(Uuid::new_v4(), serde_json::json!({}), now()).await, ] { assert!(!finished.expect("update unknown task")); @@ -1007,48 +1143,4 @@ mod tests { .expect("newer schema must be rejected"); assert!(error.to_string().contains("newer than supported")); } - - #[tokio::test] - async fn a_task_recorded_with_a_job_token_is_read_after_the_upgrade() { - let temp_dir = tempfile::tempdir().expect("create temporary directory"); - let path = temp_dir.path().join("gateway.db"); - let db = gateway_db::GatewayDb::open_path(path.to_str().expect("temporary path is UTF-8")) - .await - .expect("open gateway database"); - let conn = db.connect().await.expect("connect"); - conn.execute_batch(MIGRATIONS[0]).await.expect("apply migration 1"); - conn.execute( - "INSERT INTO provisioner_task_metadata (key, value) VALUES (?1, 1)", - params![SCHEMA_VERSION_KEY], - ) - .await - .expect("record schema version 1"); - let id = Uuid::new_v4(); - conn.execute( - "INSERT INTO provisioner_task_records - (id, kind, target, params, state, payload, created_at, deadline_at, job_token) - VALUES (?1, 'recording-ai-analysis', 'session-1', '{}', 'queued', 'null', ?2, ?3, ?4)", - params![ - id.to_string(), - now().unix_timestamp(), - (now() + Duration::hours(2)).unix_timestamp(), - Uuid::new_v4().to_string() - ], - ) - .await - .expect("insert a task of schema version 1"); - - let store = LibSqlProvisionerTaskStore::open(db.connect().await.expect("connect")) - .await - .expect("upgrade the provisioner task store"); - - let task = store.get(id, now()).await.expect("get task").expect("task"); - assert_eq!(task.state, ProvisionerTaskState::Queued); - assert_eq!( - schema_version(&*store.conn.lock().await) - .await - .expect("read schema version"), - MIGRATIONS.len() - ); - } } diff --git a/crates/provisioner-task/INTENT.md b/crates/provisioner-task/INTENT.md index 6fefec608c..f8dcc766a4 100644 --- a/crates/provisioner-task/INTENT.md +++ b/crates/provisioner-task/INTENT.md @@ -1,5 +1,5 @@ # Background -Provisioner Tasks are work that a provisioner (DVLS, Hub) assigns to Gateway. +Provisioner Tasks (hereafter "Tasks") are work that a provisioner (DVLS, Hub) assigns to Gateway. # Invariants The provisioner chooses each Task's ID, so it can ask again without knowing whether its first request went through. @@ -21,4 +21,4 @@ An unfinished Task past its deadline is failed, so a lost job can't block its ta # Tasks and jobs A Task is not a job. -A Task keeps what was asked, how it went and how it ended, and stays for auditing after its job is gone; the job queue only schedules and retries the work. +A Task keeps what was asked, how it went and how it ended, and stays for auditing after its job is gone; the job queue only delivers the job, and the job runs and retries the attempts itself. diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index 31c7c4105f..33c002c0ed 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -3,6 +3,8 @@ //! A Task remembers what was asked, how it is going and how it ended, and it stays around for auditing. The work //! itself runs as a job on the job queue; that job is gone once it finishes, but its Task is not. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use async_trait::async_trait; @@ -49,6 +51,9 @@ pub struct ProvisionerTaskRecord { /// /// It moves forward while the Task runs, each time an attempt starts or reports progress. pub deadline_at: OffsetDateTime, + /// The job that runs this Task: a job with another token was stored for a Task that was never recorded, and must + /// not run this one. + pub job_token: Uuid, } /// The record of a Task about to be created. @@ -59,8 +64,14 @@ pub struct NewProvisionerTaskRecord { pub target: String, pub params: serde_json::Value, pub deadline_at: OffsetDateTime, + /// Token of the job that will run the new Task. + pub job_token: Uuid, } +/// Work the store does inside the transaction of a write, before it commits; an error rolls the write back and is +/// returned. See [`ProvisionerTaskStore::create`] and [`ProvisionerTaskStore::succeed`]. +pub type BeforeCommit<'a> = Pin> + Send + 'a>>; + #[derive(Debug, Clone, PartialEq)] pub enum CreateOutcome { Created(ProvisionerTaskRecord), @@ -91,8 +102,16 @@ pub enum AttemptStart { pub trait ProvisionerTaskStore: Send + Sync { /// Records a new `queued` Task, unless the ID is known or its target is busy. /// + /// `store_job` runs once the record is written and before it is committed, so a Task is recorded only with its + /// job: if it fails, nothing is recorded and its error is returned. + /// /// Like every method that takes `now`, it first fails the Tasks it looks at that are past their deadline. - async fn create(&self, task: NewProvisionerTaskRecord, now: OffsetDateTime) -> anyhow::Result; + async fn create( + &self, + task: NewProvisionerTaskRecord, + now: OffsetDateTime, + store_job: BeforeCommit<'_>, + ) -> anyhow::Result; /// Reads a Task, after failing it if it is past its deadline. async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result>; @@ -124,7 +143,16 @@ pub trait ProvisionerTaskStore: Send + Sync { /// 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. - async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; + /// + /// `publish` runs once the Task is found running and before its success is written, so a succeeded Task always + /// has what `publish` does: if it fails, the Task stays running and its error is returned. + async fn succeed( + &self, + id: Uuid, + payload: serde_json::Value, + now: OffsetDateTime, + publish: BeforeCommit<'_>, + ) -> anyhow::Result; /// Ends a `queued` or `running` Task as `failed`. Returns `false` if it was already finished, such as when it was /// just failed for being past its deadline. diff --git a/devolutions-gateway/src/api/jrec.intent.md b/devolutions-gateway/src/api/jrec.intent.md index ccd9fd423f..b181fc5c15 100644 --- a/devolutions-gateway/src/api/jrec.intent.md +++ b/devolutions-gateway/src/api/jrec.intent.md @@ -1,14 +1,3 @@ # Background Jrec is the endpoint to push, streaming, read recordings and its associated files. - - -# AI analysis - -An AI analysis is a task initiated by the provisioner. -It requires the scope token with access `gateway.tasks.recording-ai-analysis.start` to access the url `POST /jet/jrec/{session_id}/ai-analysis`; Naturally in order to read the task's state, it requires the scope token with access `gateway.tasks.recording-ai-analysis.read`. -The analyzed output belongs to the recording reading scope, however, the associated metadata, including but not limited to error messages, token usages etc. belong to the task reading scope. -The caller chooses the language of the analysis in the request, under the rules of `crates/devolutions-gateway-ai/INTENT.md`. The language is part of the Task like every other parameter, so the same Task ID sent again with another language, model or provider is a conflict. -A request sent again is compared on its parameters only: the API key and the base URL are never stored, so another key or base URL gets the same Task, which keeps the first ones. -A new Task answers 202 with its path in `Location`, and the same request again answers 200 with the same Task. A conflict answers 409 with a `code`: `task_id_conflict`, `recording_active`, or `analysis_in_progress` with the running `taskId`. A session without recording, or an unknown Task, answers 404. An invalid request answers 400 and records no Task. Unknown fields are ignored, so a newer provisioner works with an older Gateway. The provider is one of `openai`, `anthropic`, `mistral`, `gemini` or `openai-compatible`. -Every Task reports the tokens it used so far, whether it is running, succeeded or failed, so the provisioner can add up what all analyses cost, failed ones included. Requests whose usage the provider didn't report are counted apart, so a total is never shown as complete when it isn't. A succeeded Task also shows the log's file name, the model the provider reported, and the prompt versions. diff --git a/devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md b/devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md new file mode 100644 index 0000000000..38d59475d1 --- /dev/null +++ b/devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md @@ -0,0 +1,9 @@ +# 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 start one with `PUT /jet/tasks/recording-ai-analysis/{task_id}`, and the scope token with access `gateway.tasks.recording-ai-analysis.read` to read its state with `GET /jet/tasks/recording-ai-analysis/{task_id}`. +The analyzed output belongs to the recording reading scope, however, the associated metadata, including but not limited to error messages, token usages etc. belong to the task reading scope. +The caller chooses the language of the analysis in the request, under the rules of `crates/devolutions-gateway-ai/INTENT.md`. The language is part of the Task like every other parameter, so the same Task ID sent again with another language, model or provider is a conflict. +A request sent again is compared on its parameters only: the API key and the base URL are never stored, so another key or base URL gets the same Task, which keeps the first ones. +A new Task answers 202 with its path in `Location`, and the same request again answers 200 with the same Task. A conflict answers 409 with a `code`: `task_id_conflict`, `recording_active`, or `analysis_in_progress` with the running `taskId`. A session without recording, or an unknown Task, answers 404. An invalid request answers 400 and records no Task. Unknown fields are ignored, so a newer provisioner works with an older Gateway. The provider is one of `openai`, `anthropic`, `mistral`, `gemini` or `openai-compatible`. +Every Task reports the tokens it used so far, whether it is running, succeeded or failed, so the provisioner can add up what all analyses cost, failed ones included. Requests whose usage the provider didn't report are counted apart, so a total is never shown as complete when it isn't. A succeeded Task also shows the log's file name, the model the provider reported, and the prompt versions. diff --git a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs index 6e895b7241..00a9fcd963 100644 --- a/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs +++ b/devolutions-gateway/src/api/tasks/recording_ai_analysis.rs @@ -1,5 +1,6 @@ //! `/jet/tasks/recording-ai-analysis/{task_id}`: start and read the AI analysis of a recorded session. +use anyhow::Context as _; use axum::Json; use axum::body::Bytes; use axum::extract::{self, State}; @@ -82,13 +83,11 @@ pub(crate) struct AiAnalysisTask { /// Recorded session the analysis is about. session_id: Uuid, /// What was asked, without the API key and the base URL. - #[cfg_attr(feature = "openapi", schema(value_type = AiAnalysisParams))] - params: serde_json::Value, + params: AiAnalysisParams, state: TaskState, /// Details of the state: progress while `running`, the log once `succeeded`, the reason once `failed`; null while /// `queued`. - #[cfg_attr(feature = "openapi", schema(value_type = Option))] - payload: serde_json::Value, + payload: Option, /// When the Task was created. #[serde(with = "time::serde::rfc3339")] created_at: OffsetDateTime, @@ -104,16 +103,201 @@ pub(crate) struct AiAnalysisTask { deadline_at: OffsetDateTime, } +/// Params of an AI analysis Task: what was asked, without the API key and the base URL +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisParams { + provider: AiProvider, + /// Model identifier, passed to the provider as is. + model: String, + /// Upper bound of tokens in each AI answer; null when not asked. + #[cfg_attr(feature = "openapi", 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(default, skip_serializing_if = "Option::is_none")] + language: Option, +} + +/// Details of the state of an AI analysis Task: `running`, `succeeded` and `failed` each have their own +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize)] +#[serde(untagged)] +pub(crate) enum AiAnalysisPayload { + Running(AiAnalysisRunningPayload), + Succeeded(AiAnalysisSucceededPayload), + Failed(AiAnalysisFailedPayload), +} + +/// Payload of a `running` AI analysis Task: its progress, and the tokens it used so far +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisRunningPayload { + /// Current step: `preparing` (the attempt starts), `reading` (the recordings are read), `describing` (the AI is + /// asked about the recording, part by part), then `writing` (the log is written). + step: String, + /// Work done in the step: bytes read while `reading`, parts described while `describing`; absent in the other steps. + #[serde(default, skip_serializing_if = "Option::is_none")] + done: Option, + /// Work the step has in all, in the same unit as `done`; absent in the other steps. + #[serde(default, skip_serializing_if = "Option::is_none")] + total: Option, + usage: AiAnalysisUsage, +} + +/// Payload of a `succeeded` AI analysis Task: the log it put in the session +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisSucceededPayload { + /// Name of the log in the session manifest, such as `ai-analysis-0.slog`. + file_name: String, + /// Model that wrote the log, as reported by the AI provider; the requested model when it reports none. + model: String, + usage: AiAnalysisUsage, + /// Versions of the prompts that wrote the log, such as + /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. + prompt_version: String, +} + +/// Tokens the AI provider counted for an AI analysis Task, over all its attempts +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisUsage { + input_tokens: u64, + output_tokens: u64, + /// Requests that may have cost tokens the provider did not report, such as a request that got no answer, or one + /// still waiting for it; while it is not 0, `inputTokens` and `outputTokens` are short of what the provider billed. + unreported_requests: u64, +} + +/// Payload of a `failed` AI analysis Task: why it failed, and the tokens it used +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisFailedPayload { + error: AiAnalysisError, + /// Absent only for a Task that never ran before its deadline passed, since it used no token. + #[serde(default, skip_serializing_if = "Option::is_none")] + 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 +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisError { + kind: AiAnalysisErrorKind, + /// HTTP status the AI provider answered, for `aiRejected`. + #[serde(default, skip_serializing_if = "Option::is_none")] + status: Option, + /// What was wrong with the answer, for `aiBadAnswer`. + #[serde(default, skip_serializing_if = "Option::is_none")] + why: Option, + /// Recording file that cannot be read, for `unreadableRecording`. + #[serde(default, skip_serializing_if = "Option::is_none")] + file: Option, + /// Failure of the last attempt, for `attemptsExhausted`. + #[serde(default, skip_serializing_if = "Option::is_none")] + last: Option, +} + +/// Kind of failure of an AI analysis Task +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) 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, or a restart that cut the last attempt + /// short. + Internal, + /// Every attempt failed with a failure worth another attempt. + AttemptsExhausted, + /// The Task did not start in time, or went too long without progress. + DeadlinePassed, +} + +/// What was wrong with an answer of the AI provider +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) 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 +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AiAnalysisRecoverableError { + kind: AiAnalysisRecoverableErrorKind, + /// HTTP status the AI provider answered, for `aiBusy`. + #[serde(default, skip_serializing_if = "Option::is_none")] + status: Option, +} + +/// Kind of a failure worth another attempt +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) 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, +} + impl TryFrom for AiAnalysisTask { - type Error = uuid::Error; + type Error = anyhow::Error; + /// The record as the API shows it; a record whose params or payload are not of this shape is an error, so a + /// Gateway that stored something else never passes it on. fn try_from(task: ProvisionerTaskRecord) -> Result { + let state = TaskState::from(task.state); + let payload = match state { + TaskState::Queued => None, + TaskState::Running => Some(AiAnalysisPayload::Running( + serde_json::from_value(task.payload).context("running payload")?, + )), + TaskState::Succeeded => Some(AiAnalysisPayload::Succeeded( + serde_json::from_value(task.payload).context("succeeded payload")?, + )), + TaskState::Failed => Some(AiAnalysisPayload::Failed( + serde_json::from_value(task.payload).context("failed payload")?, + )), + }; + Ok(Self { id: task.id, - session_id: Uuid::parse_str(&task.target)?, - params: task.params, - state: TaskState::from(task.state), - payload: task.payload, + session_id: Uuid::parse_str(&task.target).context("session ID")?, + params: serde_json::from_value(task.params).context("params")?, + state, + payload, created_at: task.created_at, started_at: task.started_at, finished_at: task.finished_at, @@ -125,7 +309,7 @@ impl TryFrom for AiAnalysisTask { fn task_json(task: ProvisionerTaskRecord) -> Result, HttpError> { AiAnalysisTask::try_from(task).map(Json).map_err( HttpError::internal() - .with_msg("invalid session ID in the AI analysis") + .with_msg("the AI analysis Task is not of the expected shape") .err(), ) } @@ -272,7 +456,7 @@ pub(crate) async fn get_ai_analysis( _scope: RecordingAiAnalysisReadScope, extract::Path(task_id): extract::Path, ) -> Result, HttpError> { - let task = analysis::get(&state, task_id, OffsetDateTime::now_utc()) + let task = analysis::get(&state, task_id) .await .map_err(HttpError::internal().with_msg("failed to read the AI analysis").err())? // A Task of another kind is not found under this path. @@ -291,7 +475,52 @@ mod tests { use super::*; - // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis + fn record(state: provisioner_task::ProvisionerTaskState, payload: serde_json::Value) -> ProvisionerTaskRecord { + let now = OffsetDateTime::now_utc(); + ProvisionerTaskRecord { + id: Uuid::new_v4(), + kind: analysis::KIND.to_owned(), + target: Uuid::new_v4().to_string(), + params: serde_json::json!({ "provider": "openai", "model": "gpt-test", "maxOutputTokens": null }), + state, + payload, + attempts: 0, + created_at: now, + started_at: None, + finished_at: None, + deadline_at: now, + job_token: Uuid::new_v4(), + } + } + + // oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis + #[test] + fn a_task_failed_before_it_ran_is_shown_without_usage() { + let record = record( + provisioner_task::ProvisionerTaskState::Failed, + serde_json::json!({ "error": { "kind": "deadlinePassed" } }), + ); + + let task = AiAnalysisTask::try_from(record).expect("a failed Task without usage"); + + let payload = serde_json::to_value(task.payload).expect("JSON"); + assert_eq!(payload, serde_json::json!({ "error": { "kind": "deadlinePassed" } })); + } + + // oracle: mutant a payload of another shape is passed on to the caller as it is + #[test] + fn a_task_of_another_shape_is_refused() { + let record = record( + provisioner_task::ProvisionerTaskState::Succeeded, + serde_json::json!({ "error": { "kind": "internal" } }), + ); + + let error = AiAnalysisTask::try_from(record).expect_err("a succeeded Task with a failed payload"); + + assert!(error.to_string().contains("succeeded payload"), "{error:#}"); + } + + // oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn ai_analysis_of_a_session_still_recording_is_a_conflict() { use axum::extract::connect_info::MockConnectInfo; @@ -307,7 +536,24 @@ mod tests { "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 (state, handles) = DgwState::mock(&config.to_string()).await.expect("mock state"); + // `start` asks the provisioner task service whether it knows the Task before it looks at the session. + let conn = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open gateway database") + .connect() + .await + .expect("connect"); + let store = provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) + .await + .expect("open task store"); + let service = crate::provisioner_tasks::ProvisionerTaskService::new( + std::sync::Arc::new(store), + crate::job_queue::TestJobQueue::new().await, + handles.provisioner_task_rx, + ); + let (_service_shutdown, service_shutdown_signal) = devolutions_gateway_task::ShutdownHandle::new(); + tokio::spawn(devolutions_gateway_task::Task::run(service, service_shutdown_signal)); 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)))); diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 45f64f392b..832d68ff67 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -13,8 +13,8 @@ use time::OffsetDateTime; use tokio::sync::{Notify, mpsc}; use crate::config::ConfHandle; +use crate::provisioner_tasks::TaskJob; use crate::provisioner_tasks::recording_ai_analysis::RecordingAiAnalysis; -use crate::provisioner_tasks::runner::{TaskJob, TaskJobUnhydrated}; use crate::recording::RecordingMessageSender; use crate::secret_store::SecretStore; @@ -31,12 +31,8 @@ pub struct JobMessage { pub schedule_for: Option, } -/// Sends jobs to be stored by the job queue task, which tries a few times then drops the job, so the caller never -/// knows whether it was stored; see [`DurableJobQueue`] for a caller that must know. #[derive(Clone)] -pub struct JobQueueHandle { - sender: mpsc::Sender, -} +pub struct JobQueueHandle(mpsc::Sender); pub type JobQueueReceiver = mpsc::Receiver; @@ -50,18 +46,6 @@ pub struct InMemoryJobDependencies { pub task_store: DynProvisionerTaskStore, } -/// A job as the job queue stores it: a name and its JSON. -pub trait StorableJob { - fn name(&self) -> &str; - - fn write_json(&self) -> anyhow::Result; -} - -/// A job read back from its JSON that needs the in-memory dependencies to run. -pub trait UnhydratedJob { - fn hydrate(self, dependencies: &InMemoryJobDependencies) -> DynJob; -} - pub struct JobQueueTask { queue: DynJobQueue, job_queue_rx: JobQueueReceiver, @@ -117,61 +101,20 @@ impl JobQueueCtx { }) } - /// The queue itself, for jobs whose callers must know they are stored. - pub fn durable(&self) -> DurableJobQueue { - DurableJobQueue(Arc::clone(&self.queue)) - } - - /// The JSON of every job stored under `name`, for startup, before the job runner claims any job. - pub async fn job_defs(&self, name: &str) -> anyhow::Result> { - self.queue.job_defs(name).await - } -} - -/// A job queue database in memory. -async fn in_memory_job_queue() -> anyhow::Result { - let conn = libsql::Builder::new_local(":memory:") - .build() - .await - .context("build database")? - .connect() - .context("open database connection")?; - let queue = job_queue_libsql::LibSqlJobQueue::builder() - .runner_waker(RunnerWaker::new(|| {})) - .conn(conn) - .build(); - queue.setup().await.context("queue setup")?; - - Ok(queue) -} - -/// Stores jobs in the job queue database: a job pushed here is stored before `push` returns, unlike with -/// [`JobQueueHandle`]. -#[derive(Clone)] -pub struct DurableJobQueue(DynJobQueue); - -impl DurableJobQueue { - /// A job queue in memory, for [`crate::DgwState::mock`]. - #[doc(hidden)] - pub async fn mock() -> anyhow::Result { - Ok(Self(Arc::new(in_memory_job_queue().await?))) - } - - /// Stores `job` to run at `schedule_for`, or at once without it. - pub async fn push(&self, job: &impl StorableJob, schedule_for: Option) -> anyhow::Result<()> { - self.0.push_job(job.name(), &job.write_json()?, schedule_for).await + /// The queue itself, for the one service whose callers must know their job is stored. + pub fn queue(&self) -> DynJobQueue { + Arc::clone(&self.queue) } } impl JobQueueHandle { pub fn new() -> (Self, JobQueueReceiver) { let (tx, rx) = mpsc::channel(512); - let handle = Self { sender: tx }; - (handle, rx) + (Self(tx), rx) } pub fn blocking_enqueue(&self, job: T) -> anyhow::Result<()> { - self.sender + self.0 .blocking_send(JobMessage { job: Box::new(job), schedule_for: None, @@ -180,7 +123,7 @@ impl JobQueueHandle { } pub async fn enqueue(&self, job: T) -> anyhow::Result<()> { - self.sender + self.0 .send(JobMessage { job: Box::new(job), schedule_for: None, @@ -194,7 +137,7 @@ impl JobQueueHandle { job: T, schedule_for: OffsetDateTime, ) -> anyhow::Result<()> { - self.sender + self.0 .blocking_send(JobMessage { job: Box::new(job), schedule_for: Some(schedule_for), @@ -203,7 +146,7 @@ impl JobQueueHandle { } pub async fn schedule(&self, job: T, schedule_for: OffsetDateTime) -> anyhow::Result<()> { - self.sender + self.0 .send(JobMessage { job: Box::new(job), schedule_for: Some(schedule_for), @@ -254,16 +197,8 @@ async fn job_queue_task(ctx: JobQueueTask, mut shutdown_signal: ShutdownSignal) let queue = Arc::clone(&queue); async move { - let def = match msg.job.write_json() { - Ok(def) => def, - Err(e) => { - error!(error = format!("{e:#}"), "Failed to serialize job"); - return; - } - }; - for _ in 0..5 { - match queue.push_job(msg.job.name(), &def, msg.schedule_for).await { + match queue.push_job(&msg.job, msg.schedule_for).await { Ok(()) => break, Err(e) => { warn!(error = format!("{e:#}"), "Failed to push job"); @@ -391,9 +326,15 @@ impl JobReader for DgwJobReader { serde_json::from_str(json).context("failed to deserialize DeleteRecordingsJob")?; Ok(Box::new(job)) } - TaskJob::::NAME => Ok( - TaskJobUnhydrated::::read_json(json)?.hydrate(&self.in_memory_job_dependencies) - ), + TaskJob::::NAME => { + let dependencies = &self.in_memory_job_dependencies; + let job = TaskJob::read_json( + json, + RecordingAiAnalysis::from_dependencies(dependencies), + Arc::clone(&dependencies.task_store), + )?; + Ok(Box::new(job)) + } _ => anyhow::bail!("unknown job name: {name}"), } } @@ -416,7 +357,17 @@ pub(crate) struct TestJobQueue { #[cfg(test)] impl TestJobQueue { pub(crate) async fn new() -> Arc { - let inner = in_memory_job_queue().await.expect("in-memory job queue"); + let conn = libsql::Builder::new_local(":memory:") + .build() + .await + .expect("build database") + .connect() + .expect("open database connection"); + let inner = job_queue_libsql::LibSqlJobQueue::builder() + .runner_waker(RunnerWaker::new(|| {})) + .conn(conn) + .build(); + inner.setup().await.expect("queue setup"); Arc::new(Self { inner, @@ -428,11 +379,6 @@ impl TestJobQueue { pub(crate) fn take(&self) -> Option { self.stored.lock().pop_front() } - - /// This queue, as [`JobQueueCtx::durable`] hands out the queue Gateway runs. - pub(crate) fn durable(self: &Arc) -> DurableJobQueue { - DurableJobQueue(Arc::::clone(self)) - } } #[cfg(test)] @@ -446,11 +392,11 @@ impl JobQueue for TestJobQueue { self.inner.reset_claimed_jobs().await } - async fn push_job(&self, name: &str, def: &str, schedule_for: Option) -> anyhow::Result<()> { - self.inner.push_job(name, def, schedule_for).await?; + async fn push_job(&self, job: &DynJob, schedule_for: Option) -> anyhow::Result<()> { + self.inner.push_job(job, schedule_for).await?; self.stored.lock().push_back(StoredJob { - name: name.to_owned(), - json: def.to_owned(), + name: job.name().to_owned(), + json: job.write_json()?, }); Ok(()) } @@ -474,8 +420,4 @@ impl JobQueue for TestJobQueue { async fn next_scheduled_date(&self) -> anyhow::Result> { self.inner.next_scheduled_date().await } - - async fn job_defs(&self, name: &str) -> anyhow::Result> { - self.inner.job_defs(name).await - } } diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index 9cee7088e4..56dca423a6 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -70,9 +70,8 @@ pub struct DgwState { pub monitoring_state: Arc, pub traffic_audit_handle: traffic_audit::TrafficAuditHandle, pub agent_tunnel_handle: Option>, - pub task_store: provisioner_task::DynProvisionerTaskStore, + pub provisioner_tasks: provisioner_tasks::ProvisionerTaskHandle, pub secrets: secret_store::SecretStore, - pub durable_job_queue: job_queue::DurableJobQueue, } #[doc(hidden)] @@ -82,6 +81,7 @@ pub struct MockHandles { pub subscriber_rx: subscriber::SubscriberReceiver, pub job_queue_rx: job_queue::JobQueueReceiver, pub traffic_audit_rx: traffic_audit::TrafficAuditReceiver, + pub provisioner_task_rx: provisioner_tasks::ProvisionerTaskReceiver, pub shutdown_handle: devolutions_gateway_task::ShutdownHandle, } @@ -100,9 +100,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 task_store_conn = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?; - let task_store = Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store_conn).await?); - let durable_job_queue = job_queue::DurableJobQueue::mock().await?; + let (provisioner_tasks, provisioner_task_rx) = provisioner_tasks::provisioner_task_channel(); let state = Self { conf_handle, @@ -118,9 +116,8 @@ impl DgwState { synthetic_kdc_registry, monitoring_state, agent_tunnel_handle: None, - task_store, + provisioner_tasks, secrets: secret_store::SecretStore::default(), - durable_job_queue, }; let handles = MockHandles { @@ -129,6 +126,7 @@ impl DgwState { subscriber_rx, job_queue_rx, traffic_audit_rx, + provisioner_task_rx, shutdown_handle, }; diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 923d54ea80..e1166a9396 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -4,6 +4,9 @@ use utoipa::{Modify, OpenApi}; use uuid::Uuid; use crate::api::preflight::PreflightAlertStatus; +use crate::api::tasks::recording_ai_analysis::{ + AiAnalysisBadAnswer, AiAnalysisErrorKind, AiAnalysisRecoverableErrorKind, +}; use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; use crate::provisioner_tasks::Permanent; use crate::provisioner_tasks::ai::AiProvider; @@ -75,16 +78,16 @@ use crate::provisioner_tasks::recording_ai_analysis::{BadAnswer, Failure, Recove AiProvider, crate::api::tasks::recording_ai_analysis::AiAnalysisTask, crate::api::tasks::TaskState, - AiAnalysisParams, - AiAnalysisPayload, - AiAnalysisRunningPayload, - AiAnalysisSucceededPayload, - AiAnalysisUsage, - AiAnalysisFailedPayload, - AiAnalysisError, + crate::api::tasks::recording_ai_analysis::AiAnalysisParams, + crate::api::tasks::recording_ai_analysis::AiAnalysisPayload, + crate::api::tasks::recording_ai_analysis::AiAnalysisRunningPayload, + crate::api::tasks::recording_ai_analysis::AiAnalysisSucceededPayload, + crate::api::tasks::recording_ai_analysis::AiAnalysisUsage, + crate::api::tasks::recording_ai_analysis::AiAnalysisFailedPayload, + crate::api::tasks::recording_ai_analysis::AiAnalysisError, AiAnalysisErrorKind, AiAnalysisBadAnswer, - AiAnalysisRecoverableError, + crate::api::tasks::recording_ai_analysis::AiAnalysisRecoverableError, AiAnalysisRecoverableErrorKind, crate::token::AccessScope, crate::api::webapp::AppTokenSignRequest, @@ -164,159 +167,8 @@ struct SessionInfo { destination_host: Option, } -/// Params of an AI analysis Task: what was asked, without the API key and the base URL -#[derive(Serialize, utoipa::ToSchema)] -#[serde(rename_all = "camelCase")] -struct AiAnalysisParams { - provider: AiProvider, - /// Model identifier, passed to the provider as is. - model: String, - /// Upper bound of tokens in each AI answer; null when not asked. - #[schema(value_type = Option, minimum = 0, maximum = 4_294_967_295_u64)] - max_output_tokens: Option, - /// Language the AI writes in; absent when not asked, so the AI writes in English. - #[serde(skip_serializing_if = "Option::is_none")] - language: Option, -} - -/// Details of the state of an AI analysis Task: `running`, `succeeded` and `failed` each have their own -#[derive(Serialize, utoipa::ToSchema)] -#[serde(untagged)] -#[expect(dead_code, reason = "it only describes the payload, which the Task keeps as JSON")] -enum AiAnalysisPayload { - Running(AiAnalysisRunningPayload), - Succeeded(AiAnalysisSucceededPayload), - Failed(AiAnalysisFailedPayload), -} - -/// Payload of a `running` AI analysis Task: its progress, and the tokens it used so far -#[derive(Serialize, utoipa::ToSchema)] -#[serde(rename_all = "camelCase")] -struct AiAnalysisRunningPayload { - /// Current step: `preparing` (the attempt starts), `reading` (the recordings are read), `describing` (the AI is - /// asked about the recording, part by part), then `writing` (the log is written). - step: String, - /// Work done in the step: bytes read while `reading`, parts described while `describing`; absent in the other steps. - done: Option, - /// Work the step has in all, in the same unit as `done`; absent in the other steps. - total: Option, - usage: AiAnalysisUsage, -} - -/// Payload of a `succeeded` AI analysis Task: the log it put in the session -#[derive(Serialize, utoipa::ToSchema)] -#[serde(rename_all = "camelCase")] -struct AiAnalysisSucceededPayload { - /// Name of the log in the session manifest, such as `ai-analysis-0.slog`. - file_name: String, - /// Model that wrote the log, as reported by the AI provider; the requested model when it reports none. - model: String, - usage: AiAnalysisUsage, - /// Versions of the prompts that wrote the log, such as - /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. - prompt_version: String, -} - -/// Tokens the AI provider counted for an AI analysis Task, over all its attempts -#[derive(Serialize, utoipa::ToSchema)] -#[serde(rename_all = "camelCase")] -struct AiAnalysisUsage { - input_tokens: u64, - output_tokens: u64, - /// Requests that may have cost tokens the provider did not report, such as a request that got no answer, or one - /// still waiting for it; while it is not 0, `inputTokens` and `outputTokens` are short of what the provider billed. - unreported_requests: u64, -} - -/// Payload of a `failed` AI analysis Task: why it failed, and the tokens it used -#[derive(Serialize, utoipa::ToSchema)] -#[serde(rename_all = "camelCase")] -struct AiAnalysisFailedPayload { - error: AiAnalysisError, - /// Absent only for a Task that never ran before its deadline passed, since it used no token. - usage: Option, -} - -/// Why an AI analysis Task failed: a kind, with the fields of that kind; there is no text for people, so the caller -/// words it -#[derive(Serialize, utoipa::ToSchema)] -#[serde(rename_all = "camelCase")] -struct AiAnalysisError { - kind: AiAnalysisErrorKind, - /// HTTP status the AI provider answered, for `aiRejected`. - status: Option, - /// What was wrong with the answer, for `aiBadAnswer`. - why: Option, - /// Recording file that cannot be read, for `unreadableRecording`. - file: Option, - /// Failure of the last attempt, for `attemptsExhausted`. - last: Option, -} - -/// Kind of failure of an AI analysis Task -#[derive(Serialize, utoipa::ToSchema)] -#[serde(rename_all = "camelCase")] -enum AiAnalysisErrorKind { - /// The AI provider refused the request for good, such as for a wrong API key (401) or a spent quota. - AiRejected, - /// The AI provider answered, but not with something usable. - AiBadAnswer, - /// The API key is gone, such as after a Gateway restart. - KeyLost, - /// The session has no recording, or it was deleted. - NoRecording, - /// The session mixes terminal and video recordings. - MixedRecording, - /// A recording file of the session cannot be read. - UnreadableRecording, - /// The Task itself is not valid, such as with parameters this Gateway cannot read. - InvalidTask, - /// The attempt panicked. - Panicked, - /// Something went wrong inside Gateway, such as a file it could not write, or a restart that cut the last attempt - /// short. - Internal, - /// Every attempt failed with a failure worth another attempt. - AttemptsExhausted, - /// The Task did not start in time, or went too long without progress. - DeadlinePassed, -} - -/// 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, -} - -// The error kinds are written out above for their descriptions, which utoipa drops from the variants without fields, so -// these matches stop the build until a new kind or field is written out too. +// The API writes the error kinds out as its own enums, so these matches stop the build until a new kind or field of +// the Task is written out there too. const _: fn(&Permanent) -> AiAnalysisErrorKind = |error| match error { Permanent::Kind(Failure::AiRejected { status: _ }) => AiAnalysisErrorKind::AiRejected, Permanent::Kind(Failure::AiBadAnswer { why: _ }) => AiAnalysisErrorKind::AiBadAnswer, diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index 8e54b089d9..28bbed3d4f 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -1,27 +1,30 @@ # 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/`. +Provisioner Tasks (hereafter "Tasks") are tasks issued for Gateway to execute. Tasks's state are pesisted in `gateway.db`, see more at `../../../crates/provisioner-task-libsql/`. # Invariants ## Management -`ProvisionerTaskManager` 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 `ProvisionerTaskManager`. - -The task record is the source of truth; the `TaskJob` only wakes it up and holds no state. - -One task at the same time only has one job to execute task. - -Jobs doesn't have state, tasks do. - -An attempt error is recoverable or permanent: a recoverable one is retried with the job up to a limit, and a permanent one fails the Task at once. A failed Task's payload is that permanent error, giving up included, as a closed set of kinds without text for people, because Gateway does not translate: the caller words it. - -When Gateway starts, it gives a new job to every unfinished Task that has none. This can happen if Gateway stopped after saving a Task but before saving its job. Gateway starts the job without checking whether the Task can still succeed: each kind of Task decides that for itself. For example, an AI analysis fails because its key is gone. If the new job can't be saved, the Task fails. +Once a provisioner gets a successful answer to its start request, the Task is never lost and never stuck, even if Gateway restarts: it ends in success or failure, within its limit of attempts and its deadline. +It may end in failure, such as `keyLost`. +Gateway does not handle a database failure, such as a failed write to `gateway.db` or to the job queue, since it cannot be handled reasonably. +Provisioner tasks run inside a job. +Recovery from a crashed Gateway is managed by the job queue: when Gateway starts, it delivers again every job it had not finished. +The job keeps only the ID of its Provisioner Task and the job token recorded with it, and runs the Provisioner Task only if the two tokens match. +The record keeps the attempts, the deadline, the payload and the end of a Provisioner Task, and each kind keeps its own progress, such as the working files of an AI analysis, so a job delivered again picks up where the last one left off. +A job delivered again in the middle of an attempt starts a new one, which counts toward the limit of attempts; with none left, the Provisioner Task fails as `internal`. +Recoverable errors are recovered within the task job itself, which waits and runs the next attempt up to the limit, and are never handed to the job queue. +The job returns an error to the job queue only when it cannot read or write the record, cannot publish the result of a success, or cannot finish what a Provisioner Task leaves to do once it ends; the job queue then runs the job again. +The handlers only record the Provisioner Task with its job and keep the secrets it needs, and never take part in its execution. +Reading a Provisioner Task past its deadline fails it. + +## Failures +The error in the payload of a failed Provisioner Task is its permanent error, giving up included, as a closed set of kinds without text for people, because Gateway does not translate: the caller words it. ## Naming -`ProvisionerTask` is the work itself; `ProvisionerTaskRecord` is what is stored about it. +A Provisioner Task is a long-running task a provisioner asks Gateway to do. +Its record, in `gateway.db`, and its job, on the job queue, are only the means to carry it out. ## Secrets Secrets a Task needs, such as an API key or a base URL that may carry credentials, never reach the disk: not `gateway.db`, not the job, not the logs. @@ -33,9 +36,9 @@ Only a finished recording is analyzed. A session that mixes terminal and video recordings cannot be analyzed. Two analyses of the same session never run at the same time. An analysis keeps working files next to the recording, so a retry can pick up where it stopped. When a Task ends, it deletes its own working files, whether it succeeded or failed. A new Task first deletes any working files an earlier analysis left behind. A Task never deletes another Task's working files. -A session keeps one AI analysis: a new Task replaces the earlier one atomically once it succeeds, and a Task that fails leaves the earlier one in place. Keeping one analysis per language is a future extension that must stay possible. -A Task succeeds only once its success is recorded in `gateway.db`. Until then, its new log sits in the session unlisted and the earlier analysis stays. After that, the manifest lists the new log and drops the earlier one; if this last step is cut short, Gateway does it again when the Task ends and when Gateway starts. -A recoverable failure, such as a rate limit (429), is retried with the job, and each retry continues from the progress saved so far; only a new Task starts over. A Task that fails keeps its record, but removes everything it produced, leaving no file on disk, even if that wastes the money already spent, because an analysis missing a large part would mislead whoever searches it, more than the misses of a nondeterministic AI already do. A Gateway restart loses the API key, so an unfinished Task fails with `keyLost` at its next attempt, without a retry. +A session keeps one AI analysis: a new Task replaces the earlier one atomically once it succeeds. Keeping one analysis per language is a future extension that must stay possible. +The only guarantee on results is that a succeeded Task always has its result persisted: its log is listed in the manifest before `succeeded` is written. Any other outcome may leave anything behind, and the provisioner runs the analysis again. +A recoverable failure, such as a rate limit (429), is retried within the job, and each retry continues from the progress saved so far; only a new Task starts over. A Gateway restart loses the API key, so an unfinished Task fails with `keyLost` at its next attempt, without a retry; this is enforced for security, since the key never reaches the disk. 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. +An analysis does not wait for Gateway to remux the session's videos: it reads the tracks, the cluster timestamps and the blocks of a WebM in order, which a remux does not change, and a remux replaces the file by renaming its output over it, so an analysis that has the file open keeps reading what it opened. diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index ee6a29c6d7..4f103c475b 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -192,7 +192,7 @@ impl AiSettings { mod tests { use super::*; - // oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis + // oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[test] fn providers_keep_their_wire_names() { for (provider, name) in [ diff --git a/devolutions-gateway/src/provisioner_tasks/attempt.rs b/devolutions-gateway/src/provisioner_tasks/attempt.rs new file mode 100644 index 0000000000..b89072e495 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/attempt.rs @@ -0,0 +1,52 @@ +use std::sync::Arc; + +use parking_lot::Mutex; +use provisioner_task::DynProvisionerTaskStore; +use time::OffsetDateTime; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use super::deadline; + +/// One attempt of a record, as its task sees it. +pub struct Attempt { + pub task_id: Uuid, + pub target: String, + /// Attempts started so far, this one included. + pub attempts: u32, + /// Payload of the record as this attempt started, from [`ProvisionerTask::starting_payload`]. + pub payload: serde_json::Value, + /// Cancelled once the record is past its deadline or no longer running. + pub cancel: CancellationToken, + pub(super) progress_timeout: time::Duration, + pub(super) store: DynProvisionerTaskStore, + /// The deadline the record has, shared with the job so it cancels the attempt once the deadline passes. + pub(super) deadline_at: Arc>, +} + +impl Attempt { + /// Deadline of the record, after which the attempt is cancelled. + pub fn deadline(&self) -> OffsetDateTime { + *self.deadline_at.lock() + } + + /// Replaces the payload of the running record and moves its deadline ahead, so whoever asks about it sees how the + /// attempt goes. + /// + /// A record that is no longer running, such as one past its deadline, cancels the attempt. + pub async fn report_progress(&self, payload: serde_json::Value) { + let now = OffsetDateTime::now_utc(); + let deadline_at = deadline(now, self.progress_timeout); + + match self.store.update_running(self.task_id, payload, deadline_at, now).await { + Ok(true) => *self.deadline_at.lock() = 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"); + } + } + } +} diff --git a/devolutions-gateway/src/provisioner_tasks/job.rs b/devolutions-gateway/src/provisioner_tasks/job.rs new file mode 100644 index 0000000000..3071810af5 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/job.rs @@ -0,0 +1,323 @@ +use std::marker::PhantomData; +use std::sync::Arc; + +use anyhow::Context as _; +use async_trait::async_trait; +use parking_lot::Mutex; +use provisioner_task::{AttemptStart, DynProvisionerTaskStore, ProvisionerTaskRecord}; +use serde::Serialize; +use time::OffsetDateTime; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use super::{Attempt, AttemptError, MAX_ATTEMPTS, Permanent, ProvisionerTask, deadline}; + +/// `value` as JSON for the logs, empty if it cannot be written. +fn log_json(value: &impl Serialize) -> serde_json::Value { + serde_json::to_value(value).unwrap_or_default() +} + +/// The job of a record of the task `T`: it wakes the record up, and the record decides what runs. +/// +/// It runs the attempts of `T` until the record ends, each in three steps: the prelude records its start, the body runs +/// [`ProvisionerTask::run`], and the completion records how the last one ended. A record that has ended is ended +/// instead. +pub struct TaskJob { + pub(super) task_id: Uuid, + job_token: Uuid, + task: T, + store: DynProvisionerTaskStore, +} + +/// What the job queue stores of a [`TaskJob`]: the ID of its record and the token recorded with it. +#[derive(Serialize, serde::Deserialize)] +#[serde(rename_all = "camelCase", bound = "")] +struct TaskJobUnhydrated { + task_id: Uuid, + /// Matches [`ProvisionerTaskRecord::job_token`] of the record this job was stored with. + job_token: Uuid, + #[serde(skip)] + task: PhantomData, +} + +impl TaskJob { + /// The name its jobs are stored under. + pub const NAME: &'static str = T::KIND; + + /// The job of the record `task_id`, recorded with `job_token`. + pub(super) fn new(task_id: Uuid, job_token: Uuid, task: T, store: DynProvisionerTaskStore) -> Self { + Self { + task_id, + job_token, + task, + store, + } + } + + /// The job read back from what the job queue stored of it, with its task and the task store. + pub fn read_json(json: &str, task: T, store: DynProvisionerTaskStore) -> anyhow::Result { + let TaskJobUnhydrated:: { task_id, job_token, .. } = + serde_json::from_str(json).context("failed to deserialize TaskJob")?; + + Ok(Self::new(task_id, job_token, task, store)) + } + + /// Records the failure, from the payload the record has, then ends the record. + async fn fail(&self, error: &Permanent) -> anyhow::Result<()> { + let now = OffsetDateTime::now_utc(); + let last = self + .store + .get(self.task_id, now) + .await + .context("read the Task before its failure")? + .map_or(serde_json::Value::Null, |record| record.payload); + + self.store + .fail(self.task_id, self.task.failed_payload(error, &last), now) + .await + .context("record the Task failure")?; + + self.end().await + } + + /// Ends the record, read again so its task sees how it ended. + async fn end(&self) -> anyhow::Result<()> { + match self.store.get(self.task_id, OffsetDateTime::now_utc()).await? { + Some(record) => self.end_record(&record).await, + None => Ok(()), + } + } + + /// Does what the finished record leaves to do; a record still unfinished is left to its attempt. + async fn end_record(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { + if !record.state.is_finished() { + return Ok(()); + } + + self.task.on_end(record).await.context("end the Task") + } + + /// The prelude: marks the start of one more attempt on the record, unless it has ended. + /// + /// A record that has ended is ended instead. A record not there yet is an error, so the job queue runs the job + /// again once it is; a record with another job token belongs to another job, so this one does nothing. + async fn prelude(&self) -> anyhow::Result> { + let Self { task, store, .. } = self; + let now = OffsetDateTime::now_utc(); + + // The starting payload carries over what earlier attempts reported, so the record is read first. + let Some(record) = store.get(self.task_id, now).await? else { + anyhow::bail!("Task {} is not recorded yet", self.task_id); + }; + + if record.job_token != self.job_token { + debug!(task.id = %self.task_id, "Task has another job, so this attempt does not run"); + return Ok(None); + } + + if record.state.is_finished() { + debug!(task.id = %self.task_id, "Task is over, so this attempt does not run"); + self.end_record(&record).await?; + return Ok(None); + } + + // A restart in the middle of the last attempt delivers the job again. + if MAX_ATTEMPTS <= record.attempts { + warn!( + task.id = %self.task_id, + task.kind = %record.kind, + attempts = record.attempts, + "Task has no attempt left, so it fails without running another" + ); + self.fail(&Permanent::Internal).await?; + return Ok(None); + } + + let starting_payload = task.starting_payload(&record.payload); + let deadline_at = deadline(now, task.progress_timeout()); + + match store + .start_attempt(self.task_id, starting_payload, deadline_at, now) + .await? + { + AttemptStart::Started(record) => Ok(Some(record)), + AttemptStart::Finished(record) => { + debug!(task.id = %self.task_id, "Task is over, so this attempt does not run"); + self.end_record(&record).await?; + Ok(None) + } + AttemptStart::Unknown => anyhow::bail!("Task {} is not recorded anymore", self.task_id), + } + } + + /// The body: runs the attempt of `T`, and cancels it once the record is past its deadline. + async fn body( + &self, + record: &ProvisionerTaskRecord, + ) -> Result> { + let params = match serde_json::from_value::(record.params.clone()) { + Ok(params) => params, + Err(error) => { + warn!(task.id = %record.id, %error, "Invalid Task parameters"); + return Err(AttemptError::Permanent(Permanent::InvalidTask)); + } + }; + + // Progress moves the deadline, so the attempt shares it with the loop below that enforces it. + let deadline_at = Arc::new(Mutex::new(record.deadline_at)); + let cancel = CancellationToken::new(); + let attempt = Attempt { + task_id: record.id, + target: record.target.clone(), + attempts: record.attempts, + payload: record.payload.clone(), + cancel: cancel.clone(), + progress_timeout: self.task.progress_timeout(), + store: Arc::clone(&self.store), + deadline_at: Arc::clone(&deadline_at), + }; + + info!(task.id = %record.id, task.kind = %record.kind, attempts = record.attempts, "Task attempt started"); + + let run = self.task.run(&attempt, params); + tokio::pin!(run); + + let mut past_deadline = false; + let outcome = loop { + let wait = + std::time::Duration::try_from(*deadline_at.lock() - OffsetDateTime::now_utc()).unwrap_or_default(); + + tokio::select! { + outcome = &mut run => break outcome, + () = tokio::time::sleep(wait) => { + // Progress may have moved the deadline while the timer ran. + if OffsetDateTime::now_utc() < *deadline_at.lock() { + continue; + } + + warn!( + task.id = %record.id, + task.kind = %record.kind, + attempts = record.attempts, + "Task went past its deadline; its attempt is cancelled" + ); + past_deadline = true; + // Dropping 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 run.await; + } + } + }; + + match outcome { + // The attempt stopped because it was cancelled, whatever error it reports for it. + Err(_) if past_deadline => Err(AttemptError::Permanent(Permanent::DeadlinePassed)), + outcome => outcome, + } + } + + /// The completion: publishes and records how the last attempt ended, then ends the record. + /// + /// A success is published in the step that records it, so a `succeeded` record always has its output published. + /// Recording the end and dropping what the record kept are separate steps, so a failed drop never undoes the end. + async fn completion( + &self, + record: &ProvisionerTaskRecord, + outcome: Result>, + ) -> anyhow::Result<()> { + let store = &self.store; + let (task_id, attempts) = (record.id, record.attempts); + + match outcome { + Ok(output) => { + let payload = match serde_json::to_value(&output) { + Ok(payload) => payload, + Err(error) => { + error!(task.id = %task_id, %error, "Failed to serialize the Task result"); + return self.fail(&Permanent::Internal).await; + } + }; + + let publish = Box::pin(async { + self.task + .publish(record, &output) + .await + .context("publish the Task result") + }); + let succeeded = store + .succeed(task_id, payload, OffsetDateTime::now_utc(), publish) + .await + .context("record the Task success")?; + + if succeeded { + 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().await + } + Err(error) => { + warn!( + task.id = %task_id, + task.kind = %record.kind, + attempts, + error = %log_json(&error), + "Task failed" + ); + self.fail(&error).await + } + } + } +} + +#[async_trait] +impl job_queue::Job for TaskJob { + fn name(&self) -> &str { + Self::NAME + } + + fn write_json(&self) -> anyhow::Result { + serde_json::to_string(&TaskJobUnhydrated:: { + task_id: self.task_id, + job_token: self.job_token, + task: PhantomData, + }) + .context("failed to serialize TaskJob") + } + + async fn run(&mut self) -> anyhow::Result<()> { + loop { + let Some(record) = self.prelude().await? else { + return Ok(()); + }; + + // The body ends the attempt; the completion ends the record, so only what ends it reaches the completion. + let outcome = match self.body(&record).await { + Err(AttemptError::Recoverable(error)) if record.attempts < MAX_ATTEMPTS => { + let delay = self.task.retry_delay(record.attempts); + warn!( + task.id = %record.id, + task.kind = %record.kind, + attempts = record.attempts, + error = %log_json(&error), + %delay, + "Task attempt failed; the next one runs after a delay" + ); + tokio::time::sleep(std::time::Duration::try_from(delay).unwrap_or_default()).await; + continue; + } + Err(AttemptError::Recoverable(last)) => Err(Permanent::AttemptsExhausted { last }), + Err(AttemptError::Permanent(error)) => Err(error), + Ok(output) => Ok(output), + }; + + return self.completion(&record, outcome).await; + } + } +} diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index fe25851377..22152dadca 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -1,9 +1,115 @@ //! Work the provisioner asks Gateway to do: each request is a Task recorded in `gateway.db`, run on the job queue as a -//! [`runner::TaskJob`]. +//! [`TaskJob`]. This module is the one place that keeps a task's record and the job running it in step. +//! +//! Each kind of work is a [`ProvisionerTask`]: what one attempt does, what its payloads hold, and what to do once it has +//! ended. +//! Its job is a [`TaskJob`], a job like any other on the job queue, which wraps each attempt in a prelude and a +//! completion: the prelude records the start of the attempt, and the completion publishes and records how it ended, +//! then ends the record. +//! After a recoverable failure, the job waits [`ProvisionerTask::retry_delay`] and runs the next attempt itself, up to +//! [`MAX_ATTEMPTS`]. +//! [`ProvisionerTaskService`] stores the job while it records the Task, so a record never exists without its job. +//! +//! The record is the source of truth and the job only wakes it up. +//! The job queue delivers the job at least once, so after a restart the job picks the record up where it left off. +//! A job whose record is not there yet returns an error, so the job queue delivers it again later; the record of a job +//! whose creation was cut short never comes, and the job queue drops the job once its retries are spent. pub mod ai; pub mod recording_ai_analysis; -pub mod runner; + +mod attempt; +mod job; +mod service; +#[cfg(test)] +mod tests; + +use async_trait::async_trait; +use provisioner_task::ProvisionerTaskRecord; +use serde::Serialize; +use serde::de::DeserializeOwned; +use time::OffsetDateTime; + +pub use self::attempt::Attempt; +pub use self::job::TaskJob; +pub use self::service::{ + ProvisionerTaskHandle, ProvisionerTaskReceiver, ProvisionerTaskService, provisioner_task_channel, +}; + +/// Attempts a record gets: a recoverable failure of the last one fails the record. +pub const MAX_ATTEMPTS: u32 = 5; + +/// 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, and the name of its jobs. + const KIND: &'static str; + + /// Time from the creation of a record until its first attempt must start. + const START_WINDOW: time::Duration; + + /// What the record keeps as asked; never secrets, since records are kept and returned to callers. + type Params: Serialize + DeserializeOwned + Send + 'static; + + /// The payload of a `succeeded` record. + type Output: Serialize + Send + Sync + 'static; + + /// Why an attempt of this task fails for good, next to the failures any task may have. + type Failure: Serialize + Send + Sync + 'static; + + /// Why an attempt of this task fails for now, so a later attempt may succeed. + type Recoverable: Serialize + Send + Sync + 'static; + + /// 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; + + /// Time to wait after attempt `attempts` failed with a recoverable failure, before the next attempt starts. + fn retry_delay(&self, attempts: u32) -> time::Duration; + + /// The payload of the running record as an attempt starts, from `last`, its payload until then: `null` before its + /// first attempt. + fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value; + + /// The payload of a record that failed with `error`, from `last`, its payload until then: `null` if it never ran. + fn failed_payload( + &self, + error: &Permanent, + last: &serde_json::Value, + ) -> serde_json::Value; + + /// Runs one attempt. + /// + /// The attempt should stop soon after [`Attempt::cancel`] is cancelled, which happens once the record is past its + /// deadline. + async fn run( + &self, + attempt: &Attempt, + params: Self::Params, + ) -> Result>; + + /// Publishes what a successful attempt produced, such as listing its file where callers find it. + /// + /// It runs in the step that records the success, before the success is written, so a `succeeded` record always + /// has its output published; an error leaves the record running, and the job queue runs the job again. + async fn publish(&self, record: &ProvisionerTaskRecord, output: &Self::Output) -> anyhow::Result<()>; + + /// Finishes what the record leaves to do once it has ended, however it ended, such as dropping what was kept only + /// while it ran. + /// + /// It may be called again for a record that already ended, so it must be idempotent; an error returns from the job, + /// and the job queue runs the job again, which calls it again. + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()>; +} + +/// The deadline `timeout` after `now`, in whole seconds as the store keeps it, so Gateway and the store agree on when +/// it passes. +fn deadline(now: OffsetDateTime, timeout: time::Duration) -> OffsetDateTime { + (now + timeout).truncate_to_second() +} /// Why one attempt of a Task failed, with `F` and `R` the failures of its own kind of work. #[derive(Debug, Clone, PartialEq, Eq)] 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 a6c3c72474..99afc6b334 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -1,13 +1,13 @@ //! 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 with [`create`](super::runner::create). +//! [`start`] checks the request and records the Task with [`create`](super::create). //! Each attempt runs the analysis of the AI crate, //! [`AiClient::analyze_recording`](devolutions_gateway_ai::AiClient::analyze_recording), which keeps its working files //! next to the session manifest so that a retry resumes where the last attempt stopped, then leaves the log in the -//! session folder, unlisted, under a name of its own Task. -//! Once the job 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. +//! session folder. +//! The step that records the success lists that log in the manifest in place of the earlier one, so a succeeded Task +//! always has its log listed; a Task that fails removes it instead. //! The API key and the working files go away when the Task ends. mod failure; @@ -19,25 +19,23 @@ use std::sync::{Arc, LazyLock, Weak}; use anyhow::Context as _; use async_trait::async_trait; -use camino::{Utf8Path, Utf8PathBuf}; +use camino::Utf8PathBuf; use devolutions_gateway_ai::Language; use devolutions_gateway_ai::recording_analysis::{ self, Progress, RecordingAnalysisEvent, RecordingFile, RecordingManifest, Recordings, RecordingsError, UsageSoFar, }; -use provisioner_task::{CreateOutcome, ProvisionerTaskRecord, ProvisionerTaskState}; +use provisioner_task::{CreateOutcome, ProvisionerTaskRecord}; use secrecy::SecretString; use serde_json::json; -use time::OffsetDateTime; use uuid::Uuid; pub use self::failure::{AnalysisError, BadAnswer, Failure, Recoverable}; -use super::Permanent; use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; -use super::runner::{Attempt, ProvisionerTask, TaskJobUnhydrated}; +use super::{Attempt, Permanent, ProvisionerTask}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; -use crate::job_queue::{InMemoryJobDependencies, UnhydratedJob}; +use crate::job_queue::InMemoryJobDependencies; use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; use crate::secret_store::{SecretId, SecretStore}; @@ -159,6 +157,9 @@ pub struct Analysis { /// Versions of the prompts that wrote the log, such as /// `session-actions-3f9a1c2b7e4d0a51+session-summary-8b20e9a4c1f70d36`. prompt_version: String, + /// Where the analysis left the log, until the success lists it under `file_name`. + #[serde(skip)] + log_path: Utf8PathBuf, } /// Records an AI analysis Task and stores its job. @@ -199,8 +200,8 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result Result( - &state.task_store, - &state.durable_job_queue, - task_id, - target, - ¶ms, - |record| { - // The job may run as soon as it is queued, so the key must be in place first. - state - .secrets - .insert(SecretId(record.id), &access.to_secret()?, record.deadline_at) - }, - ) - .await - .map_err(StartError::Internal)?; + let secrets = state.secrets.clone(); + let outcome = state + .provisioner_tasks + .create( + RecordingAiAnalysis::new( + state.conf_handle.clone(), + state.recordings.clone(), + state.secrets.clone(), + ), + task_id, + target, + ¶ms, + move |record| { + // The job may run as soon as the Task is recorded, so the key must be in place first. + secrets.insert(SecretId(record.id), &access.to_secret()?, record.deadline_at) + }, + ) + .await + .map_err(StartError::Internal)?; Ok(match outcome { CreateOutcome::Created(task) => { @@ -264,21 +269,8 @@ pub async fn start(state: &DgwState, request: StartRequest) -> Result anyhow::Result> { - state.task_store.get(task_id, now).await -} - -impl UnhydratedJob for TaskJobUnhydrated { - fn hydrate(self, dependencies: &InMemoryJobDependencies) -> job_queue::DynJob { - Box::new(self.into_job( - Arc::new(RecordingAiAnalysis::from_dependencies(dependencies)), - Arc::clone(&dependencies.task_store), - )) - } +pub async fn get(state: &DgwState, task_id: Uuid) -> anyhow::Result> { + state.provisioner_tasks.get(task_id).await } /// The `recording-ai-analysis` provisioner task. @@ -301,7 +293,7 @@ impl RecordingAiAnalysis { } } - fn from_dependencies(dependencies: &InMemoryJobDependencies) -> Self { + pub(crate) fn from_dependencies(dependencies: &InMemoryJobDependencies) -> Self { Self::new( dependencies.conf_handle.clone(), dependencies.recordings.clone(), @@ -318,54 +310,6 @@ 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!("{UNLISTED_LOG_PREFIX}{task_id}{UNLISTED_LOG_SUFFIX}")) -} - -/// The logs that Tasks of the session left unlisted, each with the ID of its Task. -async fn unlisted_logs(manifest_path: &Utf8Path) -> std::io::Result> { - let Some(dir) = manifest_path.parent() else { - return Ok(Vec::new()); - }; - let mut entries = match tokio::fs::read_dir(dir).await { - Ok(entries) => entries, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), - Err(error) => return Err(error), - }; - let mut logs = Vec::new(); - - while let Some(entry) = entries.next_entry().await? { - let name = entry.file_name(); - let Some(name) = name.to_str() else { - continue; - }; - let task_id = name - .strip_prefix(UNLISTED_LOG_PREFIX) - .and_then(|name| name.strip_suffix(UNLISTED_LOG_SUFFIX)) - .and_then(|task_id| Uuid::parse_str(task_id).ok()); - - if let Some(task_id) = task_id { - logs.push((task_id, dir.join(name))); - } - } - - Ok(logs) -} - -/// Removes `path`, unless it is already gone. -async fn remove_if_exists(path: &Utf8Path) -> std::io::Result<()> { - match tokio::fs::remove_file(path).await { - Err(error) if error.kind() != std::io::ErrorKind::NotFound => Err(error), - _ => Ok(()), - } -} - #[async_trait] impl ProvisionerTask for RecordingAiAnalysis { const KIND: &'static str = KIND; @@ -419,21 +363,12 @@ 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 { - let manifest_path = self.manifest_path(session_id); - - recording_analysis::discard(&manifest_path, None) + 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"); AnalysisError::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"); - AnalysisError::Permanent(Permanent::Internal) - })?; } let secret = self.secrets.get(SecretId(attempt.task_id)).map_err(|error| { @@ -462,6 +397,16 @@ impl ProvisionerTask for RecordingAiAnalysis { .await } + async fn publish(&self, record: &ProvisionerTaskRecord, output: &Analysis) -> anyhow::Result<()> { + let session_id = Uuid::parse_str(&record.target).context("parse the session ID")?; + + // Nothing else runs on the session meanwhile: a newer Task of the session cannot start before this one ends. + self.recordings + .replace_artifacts(session_id, ArtifactKind::AiAnalysis, &output.log_path) + .await + .with_context(|| format!("list {}", output.log_path)) + } + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { self.secrets.remove(SecretId(record.id)); @@ -469,14 +414,9 @@ impl ProvisionerTask for RecordingAiAnalysis { return Ok(()); }; - let manifest_path = self.manifest_path(session_id); - let log_path = unlisted_log_path(&manifest_path, record.id); let _session = lock_session(session_id).await; - self.end_log(session_id, &log_path, record.state == ProvisionerTaskState::Succeeded) - .await?; - - recording_analysis::discard(&manifest_path, Some(&record.id.to_string())) + recording_analysis::discard(&self.manifest_path(session_id), Some(&record.id.to_string())) .await .context("remove the AI analysis working files") } @@ -574,16 +514,6 @@ impl RecordingAiAnalysis { } }; - // The log stays out of the manifest until the job records the success, so a Task that fails never replaces - // the earlier analysis. - let log_path = unlisted_log_path(&manifest_path, attempt.task_id); - tokio::fs::rename(&analysis.log_path, &log_path) - .await - .map_err(|error| { - error!(task.id = %attempt.task_id, %error, "Failed to keep the AI analysis log"); - AnalysisError::Permanent(Permanent::Internal) - })?; - debug!(task.id = %attempt.task_id, actions = analysis.actions, "AI analysis log written"); Ok(Analysis { @@ -591,53 +521,11 @@ impl RecordingAiAnalysis { model: analysis.model, usage: earlier.with(analysis.usage), prompt_version: analysis.prompt_version, + log_path: analysis.log_path, }) } } -impl RecordingAiAnalysis { - /// What the end of a Task does to the log it left unlisted: the log of a Task whose success is recorded takes the - /// place of the earlier analysis, and any other log is removed. - async fn end_log(&self, session_id: Uuid, log_path: &Utf8Path, succeeded: bool) -> anyhow::Result<()> { - if succeeded { - self.recordings - .replace_artifacts(session_id, ArtifactKind::AiAnalysis, log_path) - .await - .with_context(|| format!("list {log_path}")) - } else { - remove_if_exists(log_path) - .await - .with_context(|| format!("remove {log_path}")) - } - } - - /// Does what the end of an earlier Task does to the log it left unlisted, since that end may still be due, so a - /// newer Task never loses an analysis whose success is recorded: the log of a succeeded Task is listed, so it stays - /// until a newer analysis succeeds, and the log of a failed or unknown one is removed. - async fn end_unlisted_logs( - &self, - attempt: &Attempt, - session_id: Uuid, - manifest_path: &Utf8Path, - ) -> anyhow::Result<()> { - for (task_id, log_path) in unlisted_logs(manifest_path).await.context("list the unlisted logs")? { - match attempt - .state_of(task_id) - .await - .context("read the Task of an unlisted log")? - { - Some(ProvisionerTaskState::Queued | ProvisionerTaskState::Running) => {} - state => { - self.end_log(session_id, &log_path, state == Some(ProvisionerTaskState::Succeeded)) - .await? - } - } - } - - Ok(()) - } -} - /// 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(); 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 0d1a25cb28..7a5182b731 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -8,11 +8,12 @@ use devolutions_gateway_ai::{Language, session_summary}; use devolutions_gateway_task::ShutdownHandle; use job_queue::Job as _; use parking_lot::Mutex; -use provisioner_task::{AttemptStart, DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskState}; +use provisioner_task::{DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskState}; +use time::OffsetDateTime; use super::*; use crate::job_queue::{StoredJob, TestJobQueue}; -use crate::provisioner_tasks::runner::{MAX_ATTEMPTS, TaskJob, TaskJobUnhydrated}; +use crate::provisioner_tasks::{MAX_ATTEMPTS, ProvisionerTaskService, TaskJob}; use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles}; @@ -41,7 +42,8 @@ fn openai(model: &str) -> AiSettings { struct Harness { state: DgwState, - task: Arc, + progress_timeout: time::Duration, + first_retry_delay: time::Duration, store: DynProvisionerTaskStore, secrets: SecretStore, job_queue: Arc, @@ -55,13 +57,14 @@ impl Harness { let dir = tempfile::tempdir().expect("temp dir"); let recording_path = Utf8PathBuf::from_path_buf(dir.path().join("recordings")).expect("UTF-8"); - let (mut state, handles) = DgwState::mock(&config(&recording_path)).await.expect("mock state"); + 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, + provisioner_task_rx, shutdown_handle, } = handles; @@ -79,16 +82,30 @@ impl Harness { // `start` stores its jobs where the test can take them. let job_queue = TestJobQueue::new().await; - state.durable_job_queue = job_queue.durable(); + let conn = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open gateway database") + .connect() + .await + .expect("connect"); + let store: DynProvisionerTaskStore = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) + .await + .expect("open task store"), + ); + let service = ProvisionerTaskService::new( + Arc::clone(&store), + Arc::::clone(&job_queue), + provisioner_task_rx, + ); + let (service_shutdown_handle, service_shutdown_signal) = ShutdownHandle::new(); + tokio::spawn(devolutions_gateway_task::Task::run(service, service_shutdown_signal)); let secrets = state.secrets.clone(); - let mut task = RecordingAiAnalysis::new(state.conf_handle.clone(), state.recordings.clone(), secrets.clone()); - task.first_retry_delay = time::Duration::ZERO; - let task = Arc::new(task); - let store = Arc::clone(&state.task_store); Self { state, - task, + progress_timeout: PROGRESS_TIMEOUT, + first_retry_delay: time::Duration::ZERO, store, secrets, job_queue, @@ -100,6 +117,7 @@ impl Harness { traffic_audit_rx, shutdown_handle, recording_shutdown_handle, + service_shutdown_handle, )), _dir: dir, } @@ -138,35 +156,28 @@ impl Harness { task_id } - /// The job of the Task, as the job queue would run it. - fn job(&self, task_id: Uuid) -> TaskJob { - let json = json!({ "taskId": task_id }).to_string(); - - TaskJobUnhydrated::::read_json(&json) - .expect("valid job") - .into_job(Arc::clone(&self.task), Arc::clone(&self.store)) + /// The task every job of this harness runs, with its timings. + fn analysis(&self) -> RecordingAiAnalysis { + RecordingAiAnalysis { + progress_timeout: self.progress_timeout, + first_retry_delay: self.first_retry_delay, + ..RecordingAiAnalysis::new( + self.state.conf_handle.clone(), + self.state.recordings.clone(), + self.secrets.clone(), + ) + } } - /// Records the success of the first attempt of the Task without ending it, as if Gateway stopped right after. - async fn record_success(&self, task_id: Uuid) { - let store = &self.store; - let now = OffsetDateTime::now_utc(); + /// The job of the Task, as the job queue would run it. + async fn job(&self, task_id: Uuid) -> TaskJob { + let json = json!({ "taskId": task_id, "jobToken": self.task(task_id).await.job_token }).to_string(); - let started = store - .start_attempt(task_id, json!({}), now + time::Duration::hours(1), now) - .await - .expect("start the attempt"); - assert!(matches!(started, AttemptStart::Started(_)), "{started:?}"); - assert!( - store - .succeed(task_id, json!({}), now) - .await - .expect("record the success") - ); + TaskJob::read_json(&json, self.analysis(), Arc::clone(&self.store)).expect("valid job") } async fn run(&self, task_id: Uuid) -> anyhow::Result<()> { - self.job(task_id).run().await + self.job(task_id).await.run().await } /// Takes every queued job. @@ -191,8 +202,10 @@ impl Harness { target: session_id.to_string(), params: json!({ "provider": provider, "model": "gpt-test", "maxOutputTokens": null }), deadline_at, + job_token: Uuid::new_v4(), }, created_at, + Box::pin(async { Ok(()) }), ) .await .expect("create"); @@ -208,14 +221,7 @@ impl Harness { /// 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.secrets.clone(), - ); - task.progress_timeout = progress_timeout; - task.first_retry_delay = time::Duration::ZERO; - self.task = Arc::new(task); + self.progress_timeout = progress_timeout; } /// Working files of the AI analyses of session_id. @@ -370,10 +376,14 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { let StoredJob { name, json } = harness.job_queue.take().expect("job queued"); assert_eq!(name, KIND); - assert_eq!(json, format!(r#"{{"taskId":"{task_id}"}}"#)); + assert_eq!( + json, + format!(r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, task.job_token), + "the job is stored with the token of the record" + ); } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn asking_again_creates_nothing_new() { let harness = Harness::new().await; @@ -415,7 +425,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 +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn asking_again_returns_the_task_even_while_the_session_records_again() { let harness = Harness::new().await; @@ -444,7 +454,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 +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn refuses_a_session_without_recording() { let harness = Harness::new().await; @@ -491,7 +501,7 @@ async fn refuses_a_session_that_is_still_recording() { ); } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn refuses_invalid_ai_settings() { let harness = Harness::new().await; @@ -700,7 +710,7 @@ 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 + // oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis assert_eq!( task.payload, json!({ @@ -874,7 +884,7 @@ async fn a_new_task_that_fails_keeps_the_earlier_analysis() { assert!(dir.join("ai-analysis-0.slog").exists()); } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn a_failed_task_reports_the_tokens_it_spent() { let harness = Harness::new().await; @@ -893,109 +903,6 @@ async fn a_failed_task_reports_the_tokens_it_spent() { ); } -// 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 - .store - .fail(task_id, json!({}), OffsetDateTime::now_utc()) - .await - .expect("fail the Task"); - - harness.run(task_id).await.expect("the job ends the Task"); - - assert!(!log.exists()); - assert_eq!(ai_analyses(&dir), json!([])); -} - -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis -#[tokio::test] -async fn a_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; - - // 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::write(&first_log, "the analysis of the first Task").expect("unlisted log"); - harness.record_success(first).await; - - let second = harness.start_created(session_id, settings).await; - harness.run(second).await.expect("second analysis"); - let newer = std::fs::read_to_string(dir.join("ai-analysis-0.slog")).expect("newer analysis"); - - harness.run(first).await.expect("the late end of the first Task"); - - assert_eq!(ai_analyses(&dir), json!([{ "fileName": "ai-analysis-0.slog" }])); - assert_eq!( - std::fs::read_to_string(dir.join("ai-analysis-0.slog")).expect("analysis"), - newer - ); - assert!(!first_log.exists()); -} - -// oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis -#[tokio::test] -async fn a_new_task_that_fails_keeps_an_earlier_success_whose_log_was_not_listed_yet() { - let harness = Harness::new().await; - let session_id = harness.write_session(); - let (settings, _requests) = spawn_ai_provider(|_| 0).await; - let first = harness.start_created(session_id, settings.clone()).await; - harness.run(first).await.expect("first analysis"); - - // As if the end of the first Task failed before the manifest listed its log, so it is done again later. - let dir = harness.recording_path.join(session_id.to_string()); - std::fs::remove_file(dir.join("ai-analysis-0.slog")).expect("unlist the log"); - std::fs::write( - dir.join(format!(".ai-analysis-{first}.slog")), - "the analysis of the first Task", - ) - .expect("unlisted log"); - let manifest_path = dir.join("recording.json"); - let mut manifest: serde_json::Value = - serde_json::from_slice(&std::fs::read(&manifest_path).expect("manifest")).expect("manifest JSON"); - manifest["artifacts"] = json!({}); - std::fs::write(&manifest_path, manifest.to_string()).expect("manifest"); - - let second = harness.start_created(session_id, settings).await; - harness.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: mutant the log left unlisted by a Task that has no record is kept -#[tokio::test] -async fn a_new_task_removes_a_log_left_unlisted_by_an_unknown_task() { - let harness = Harness::new().await; - let session_id = harness.write_session(); - let dir = harness.recording_path.join(session_id.to_string()); - let unknown_log = dir.join(format!(".ai-analysis-{}.slog", Uuid::new_v4())); - std::fs::write(&unknown_log, "an analysis whose success was never recorded").expect("unlisted log"); - let (settings, _requests) = spawn_ai_provider(|_| 401).await; - let task_id = harness.start_created(session_id, settings).await; - - harness.run(task_id).await.expect("no retry"); - - assert!(!unknown_log.exists()); - assert_eq!(ai_analyses(&dir), json!([])); -} - /// The AI analyses the manifest in `dir` lists. fn ai_analyses(dir: &Utf8Path) -> serde_json::Value { let manifest: serde_json::Value = @@ -1006,7 +913,7 @@ fn ai_analyses(dir: &Utf8Path) -> serde_json::Value { .unwrap_or_else(|| json!([])) } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn the_analysis_is_in_the_language_of_the_request() { let harness = Harness::new().await; @@ -1062,14 +969,7 @@ async fn the_analysis_is_in_english_without_a_language() { #[tokio::test] async fn a_new_task_never_resumes_from_an_earlier_task() { let mut harness = Harness::new().await; - harness.task = Arc::new(RecordingAiAnalysis { - first_retry_delay: time::Duration::hours(1), - ..RecordingAiAnalysis::new( - harness.state.conf_handle.clone(), - harness.state.recordings.clone(), - harness.secrets.clone(), - ) - }); + harness.first_retry_delay = time::Duration::hours(1); // 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; diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs deleted file mode 100644 index 22aea2b6ab..0000000000 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ /dev/null @@ -1,573 +0,0 @@ -//! Runs provisioner tasks: the one place that keeps a task's record and the job running it in step. -//! -//! Each kind of work is a [`ProvisionerTask`]: what one attempt does, what its payloads hold, and what to do once it has -//! ended. -//! Its job is a [`TaskJob`], a job like any other on the job queue, which wraps each attempt in a prelude and a -//! completion: the prelude records the start of the attempt, and the completion records how it ended and ends the record. -//! After a recoverable failure, the job waits [`ProvisionerTask::retry_delay`] and runs the next attempt itself, up to -//! [`MAX_ATTEMPTS`]. -//! [`create`] records a Task and stores its job. -//! -//! The record is the source of truth and the job only wakes it up. -//! The job queue delivers the job at least once, so after a restart the job picks the record up where it left off, and -//! [`requeue_lost_jobs`] stores the job of a record Gateway stopped before storing it. - -#[cfg(test)] -mod tests; - -use std::collections::HashSet; -use std::marker::PhantomData; -use std::sync::Arc; - -use anyhow::Context as _; -use async_trait::async_trait; -use parking_lot::Mutex; -use provisioner_task::{ - AttemptStart, CreateOutcome, DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskRecord, - ProvisionerTaskState, -}; -use serde::Serialize; -use serde::de::DeserializeOwned; -use time::OffsetDateTime; -use tokio_util::sync::CancellationToken; -use uuid::Uuid; - -use super::{AttemptError, Permanent}; -use crate::job_queue::{DurableJobQueue, StorableJob}; - -/// Attempts a record gets: a recoverable failure of the last one fails the record. -pub const MAX_ATTEMPTS: u32 = 5; - -/// 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, and the name of its jobs. - const KIND: &'static str; - - /// Time from the creation of a record until its first attempt must start. - const START_WINDOW: time::Duration; - - /// What the record keeps as asked; never secrets, since records are kept and returned to callers. - type Params: Serialize + DeserializeOwned + Send + 'static; - - /// The payload of a `succeeded` record. - type Output: Serialize + Send + 'static; - - /// Why an attempt of this task fails for good, next to the failures any task may have. - type Failure: Serialize + Send + Sync + 'static; - - /// Why an attempt of this task fails for now, so a later attempt may succeed. - type Recoverable: Serialize + Send + Sync + 'static; - - /// 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; - - /// Time to wait after attempt `attempts` failed with a recoverable failure, before the next attempt starts. - fn retry_delay(&self, attempts: u32) -> time::Duration; - - /// The payload of the running record as an attempt starts, from `last`, its payload until then: `null` before its - /// first attempt. - fn starting_payload(&self, last: &serde_json::Value) -> serde_json::Value; - - /// The payload of a record that failed with `error`, from `last`, its payload until then: `null` if it never ran. - fn failed_payload( - &self, - error: &Permanent, - last: &serde_json::Value, - ) -> serde_json::Value; - - /// Runs one attempt. - /// - /// The attempt should stop soon after [`Attempt::cancel`] is cancelled, which happens once the record is past its - /// deadline. - async fn run( - &self, - attempt: &Attempt, - params: Self::Params, - ) -> Result>; - - /// Finishes what the record leaves to do once it has ended, however it ended, such as dropping what was kept only - /// while it ran. - /// - /// It may be called again for a record that already ended, so it must be idempotent; an error returns from the job, - /// and the job queue runs the job again, which calls it again. - async fn on_end(&self, _record: &ProvisionerTaskRecord) -> anyhow::Result<()> { - Ok(()) - } -} - -/// One attempt of a record, as its task sees it. -pub struct Attempt { - pub task_id: Uuid, - pub target: String, - /// Attempts started so far, this one included. - pub attempts: u32, - /// Payload of the record as this attempt started, from [`ProvisionerTask::starting_payload`]. - pub payload: serde_json::Value, - /// Cancelled once the record is past its deadline or no longer running. - pub cancel: CancellationToken, - progress_timeout: time::Duration, - store: DynProvisionerTaskStore, - latest: Arc>, -} - -/// What the job knows of the running record: the last payload of the attempt, and the deadline it stored. -struct Latest { - payload: serde_json::Value, - deadline_at: OffsetDateTime, -} - -impl Attempt { - /// Deadline of the record, after which the attempt is cancelled. - pub fn deadline(&self) -> OffsetDateTime { - self.latest.lock().deadline_at - } - - /// State of another record of the same target, such as an earlier Task; `None` for an unknown one or one of another - /// target. One past its deadline is failed first. - pub async fn state_of(&self, task_id: Uuid) -> anyhow::Result> { - let record = self.store.get(task_id, OffsetDateTime::now_utc()).await?; - Ok(record - .filter(|record| record.target == self.target) - .map(|record| record.state)) - } - - /// Replaces the payload of the running record and moves its deadline ahead, so whoever asks about it sees how the - /// attempt goes. - /// - /// A record that is no longer running, such as one past its deadline, cancels the attempt. - pub async fn report_progress(&self, payload: serde_json::Value) { - let now = OffsetDateTime::now_utc(); - let deadline_at = deadline(now, self.progress_timeout); - self.latest.lock().payload = payload.clone(); - - match self.store.update_running(self.task_id, payload, deadline_at, now).await { - Ok(true) => self.latest.lock().deadline_at = deadline_at, - Ok(false) => { - debug!(task.id = %self.task_id, "Task is no longer running, so its attempt stops"); - self.cancel.cancel(); - } - Err(error) => { - warn!(task.id = %self.task_id, error = format!("{error:#}"), "Failed to record the Task progress"); - } - } - } -} - -/// `value` as JSON for the logs, empty if it cannot be written. -fn log_json(value: &impl Serialize) -> serde_json::Value { - serde_json::to_value(value).unwrap_or_default() -} - -/// Records a Task of `T` and stores the job that runs it. -/// -/// `on_created` runs only for a new record, before its job is stored, so what the job needs is in place before it can -/// run; if it fails, no job is stored. A known ID or a busy target creates nothing and stores nothing. -pub async fn create( - store: &DynProvisionerTaskStore, - job_queue: &DurableJobQueue, - id: Uuid, - target: String, - params: &T::Params, - on_created: impl FnOnce(&ProvisionerTaskRecord) -> anyhow::Result<()> + Send, -) -> anyhow::Result { - let params = serde_json::to_value(params).context("serialize the task parameters")?; - let now = OffsetDateTime::now_utc(); - let job = TaskJobUnhydrated:: { - task_id: id, - task: PhantomData, - }; - - let outcome = store - .create( - NewProvisionerTaskRecord { - id, - kind: T::KIND.to_owned(), - target, - params, - deadline_at: deadline(now, T::START_WINDOW), - }, - now, - ) - .await?; - - let CreateOutcome::Created(record) = &outcome else { - return Ok(outcome); - }; - - on_created(record).context("prepare the Task")?; - - job_queue.push(&job, None).await?; - - info!(task.id = %id, task.kind = T::KIND, "Task queued"); - - Ok(outcome) -} - -/// Stores again the job of every queued record of `T` that has none among `job_defs`, the JSON of the jobs stored under -/// [`TaskJob::NAME`], as when Gateway stopped between recording a Task and storing its job. -/// -/// Call it at startup, before the job runner claims any job, so a job that runs meanwhile is never taken for lost. -pub async fn requeue_lost_jobs( - store: &DynProvisionerTaskStore, - job_queue: &DurableJobQueue, - job_defs: &[String], -) -> anyhow::Result<()> { - let stored = job_defs - .iter() - .map(|json| TaskJobUnhydrated::::read_json(json).map(|job| job.task_id)) - .collect::>>()?; - - for record in store.list_unfinished(OffsetDateTime::now_utc()).await? { - if record.kind != T::KIND || record.state != ProvisionerTaskState::Queued || stored.contains(&record.id) { - continue; - } - - let job = TaskJobUnhydrated:: { - task_id: record.id, - task: PhantomData, - }; - job_queue.push(&job, None).await?; - - info!(task.id = %record.id, task.kind = %record.kind, "Task job queued again"); - } - - Ok(()) -} - -/// The deadline `timeout` after `now`, in whole seconds as the store keeps it, so Gateway and the store agree on when -/// it passes. -fn deadline(now: OffsetDateTime, timeout: time::Duration) -> OffsetDateTime { - (now + timeout).truncate_to_second() -} - -/// The job of a record of the task `T`: it wakes the record up, and the record decides what runs. -/// -/// It runs the attempts of `T` until the record ends, each in three steps: the prelude records its start, the body runs -/// [`ProvisionerTask::run`], and the completion records how the last one ended. A record that has ended is ended -/// instead. -pub struct TaskJob { - task_id: Uuid, - task: Arc, - store: DynProvisionerTaskStore, -} - -/// A [`TaskJob`] as the job queue stores it, before it gets its task and the task store. -#[derive(Serialize, serde::Deserialize)] -#[serde(rename_all = "camelCase", bound = "")] -pub struct TaskJobUnhydrated { - task_id: Uuid, - #[serde(skip)] - task: PhantomData, -} - -impl TaskJobUnhydrated { - pub fn read_json(json: &str) -> anyhow::Result { - serde_json::from_str(json).context("failed to deserialize TaskJob") - } - - pub fn into_job(self, task: Arc, store: DynProvisionerTaskStore) -> TaskJob { - TaskJob { - task_id: self.task_id, - task, - store, - } - } -} - -impl StorableJob for TaskJobUnhydrated { - fn name(&self) -> &str { - TaskJob::::NAME - } - - fn write_json(&self) -> anyhow::Result { - serde_json::to_string(self).context("failed to serialize TaskJob") - } -} - -impl TaskJob { - /// The name its jobs are stored under. - pub const NAME: &'static str = T::KIND; - - /// Records the failure, then ends the record. - async fn fail( - &self, - error: &Permanent, - last: &serde_json::Value, - ) -> anyhow::Result<()> { - self.store - .fail( - self.task_id, - self.task.failed_payload(error, last), - OffsetDateTime::now_utc(), - ) - .await - .context("record the Task failure")?; - - self.end().await - } - - /// Ends the record, read again so its task sees how it ended. - async fn end(&self) -> anyhow::Result<()> { - match self.store.get(self.task_id, OffsetDateTime::now_utc()).await? { - Some(record) => self.end_record(&record).await, - None => Ok(()), - } - } - - /// Does what the finished record leaves to do; a record still unfinished is left to its attempt. - async fn end_record(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { - if !record.state.is_finished() { - return Ok(()); - } - - self.task.on_end(record).await.context("end the Task") - } - - /// The prelude: marks the start of one more attempt on the record, unless it has ended. - /// - /// A record that has ended is ended instead. - async fn prelude(&self) -> anyhow::Result> { - let Self { task, store, .. } = self; - let now = OffsetDateTime::now_utc(); - - // The starting payload carries over what earlier attempts reported, so the record is read first. - let Some(record) = store.get(self.task_id, now).await? else { - error!(task.id = %self.task_id, "Task is unknown, so this attempt does not run"); - return Ok(None); - }; - - if record.state.is_finished() { - debug!(task.id = %self.task_id, "Task is over, so this attempt does not run"); - self.end_record(&record).await?; - return Ok(None); - } - - // A restart in the middle of the last attempt delivers the job again. - if MAX_ATTEMPTS <= record.attempts { - warn!( - task.id = %self.task_id, - task.kind = %record.kind, - attempts = record.attempts, - "Task has no attempt left, so it fails without running another" - ); - self.fail(&Permanent::Internal, &record.payload).await?; - return Ok(None); - } - - let starting_payload = task.starting_payload(&record.payload); - let deadline_at = deadline(now, task.progress_timeout()); - - match store - .start_attempt(self.task_id, starting_payload, deadline_at, now) - .await? - { - AttemptStart::Started(record) => Ok(Some(record)), - AttemptStart::Finished(record) => { - debug!(task.id = %self.task_id, "Task is over, so this attempt does not run"); - self.end_record(&record).await?; - Ok(None) - } - AttemptStart::Unknown => { - error!(task.id = %self.task_id, "Task is unknown, so this attempt does not run"); - Ok(None) - } - } - } - - /// The body: runs the attempt of `T`, and cancels it once the record is past its deadline. - /// - /// Returns how the attempt ended, with the last payload it reported and the deadline it last set. - async fn body( - &self, - record: &ProvisionerTaskRecord, - ) -> (Result>, Latest) { - let latest = Arc::new(Mutex::new(Latest { - payload: record.payload.clone(), - deadline_at: record.deadline_at, - })); - let take_latest = |latest: &Mutex| { - let latest = latest.lock(); - Latest { - payload: latest.payload.clone(), - deadline_at: latest.deadline_at, - } - }; - - let params = match serde_json::from_value::(record.params.clone()) { - Ok(params) => params, - Err(error) => { - warn!(task.id = %record.id, %error, "Invalid Task parameters"); - return ( - Err(AttemptError::Permanent(Permanent::InvalidTask)), - take_latest(&latest), - ); - } - }; - - let cancel = CancellationToken::new(); - let attempt = Attempt { - task_id: record.id, - target: record.target.clone(), - attempts: record.attempts, - payload: record.payload.clone(), - cancel: cancel.clone(), - progress_timeout: self.task.progress_timeout(), - store: Arc::clone(&self.store), - latest: Arc::clone(&latest), - }; - - info!(task.id = %record.id, task.kind = %record.kind, attempts = record.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(&self.task); - async move { task.run(&attempt, params).await } - }); - - let mut past_deadline = false; - let joined = loop { - let deadline_at = latest.lock().deadline_at; - let wait = std::time::Duration::try_from(deadline_at - OffsetDateTime::now_utc()).unwrap_or_default(); - - tokio::select! { - joined = &mut handle => break joined, - () = tokio::time::sleep(wait) => { - // Progress may have moved the deadline while the timer ran. - if OffsetDateTime::now_utc() < latest.lock().deadline_at { - continue; - } - - warn!( - task.id = %record.id, - task.kind = %record.kind, - attempts = record.attempts, - "Task went past its deadline; its attempt is cancelled" - ); - past_deadline = true; - // Aborting cannot stop blocking steps, so the attempt is cancelled and awaited instead: what runs - // once the Task ends never runs alongside it. - cancel.cancel(); - break handle.await; - } - } - }; - - let outcome = match joined { - // The attempt stopped because it was cancelled, whatever error it reports for it. - Ok(Err(_)) if past_deadline => Err(AttemptError::Permanent(Permanent::DeadlinePassed)), - Ok(outcome) => outcome, - Err(_) => Err(AttemptError::Permanent(Permanent::Panicked)), - }; - - (outcome, take_latest(&latest)) - } - - /// The completion: records how the last attempt ended, then ends the record. - /// - /// Recording the end and dropping what the record kept are separate steps, so a failed drop never undoes the end. - async fn completion( - &self, - record: &ProvisionerTaskRecord, - outcome: Result>, - latest: Latest, - ) -> anyhow::Result<()> { - let store = &self.store; - let (task_id, attempts) = (record.id, record.attempts); - let last = &latest.payload; - - match outcome { - Ok(output) => { - let output = match serde_json::to_value(output) { - Ok(output) => output, - Err(error) => { - error!(task.id = %task_id, %error, "Failed to serialize the Task result"); - return self.fail(&Permanent::Internal, last).await; - } - }; - - let succeeded = store - .succeed(task_id, output, OffsetDateTime::now_utc()) - .await - .context("record the Task success")?; - - if succeeded { - 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().await - } - Err(AttemptError::Recoverable(error)) => { - warn!( - task.id = %task_id, - task.kind = %record.kind, - attempts, - error = %log_json(&error), - "Task failed after its last attempt" - ); - self.fail(&Permanent::AttemptsExhausted { last: error }, last).await - } - Err(AttemptError::Permanent(error)) => { - warn!( - task.id = %task_id, - task.kind = %record.kind, - attempts, - error = %log_json(&error), - "Task failed" - ); - self.fail(&error, last).await - } - } - } -} - -#[async_trait] -impl job_queue::Job for TaskJob { - fn name(&self) -> &str { - Self::NAME - } - - fn write_json(&self) -> anyhow::Result { - TaskJobUnhydrated:: { - task_id: self.task_id, - task: PhantomData, - } - .write_json() - } - - async fn run(&mut self) -> anyhow::Result<()> { - loop { - let Some(record) = self.prelude().await? else { - return Ok(()); - }; - - let (outcome, latest) = self.body(&record).await; - - match outcome { - Err(AttemptError::Recoverable(error)) if record.attempts < MAX_ATTEMPTS => { - let delay = self.task.retry_delay(record.attempts); - warn!( - task.id = %record.id, - task.kind = %record.kind, - attempts = record.attempts, - error = %log_json(&error), - %delay, - "Task attempt failed; the next one runs after a delay" - ); - tokio::time::sleep(std::time::Duration::try_from(delay).unwrap_or_default()).await; - } - outcome => return self.completion(&record, outcome, latest).await, - } - } - } -} diff --git a/devolutions-gateway/src/provisioner_tasks/service.rs b/devolutions-gateway/src/provisioner_tasks/service.rs new file mode 100644 index 0000000000..2509e037a5 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/service.rs @@ -0,0 +1,184 @@ +//! The one task that records Tasks: it holds the task store and the job queue, so a Task and its job are written +//! together or not at all, and the handlers only ask it. + +use std::sync::Arc; + +use anyhow::Context as _; +use async_trait::async_trait; +use devolutions_gateway_task::{ShutdownSignal, Task}; +use job_queue::{DynJob, DynJobQueue}; +use provisioner_task::{CreateOutcome, DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskRecord}; +use time::OffsetDateTime; +use tokio::sync::{mpsc, oneshot}; +use uuid::Uuid; + +use super::{ProvisionerTask, TaskJob, deadline}; + +/// How far ahead the job of a new Task is scheduled, so the job queue claims it after the record is committed. +/// +/// The job is stored before the record is committed, and storing it wakes the job queue; without the delay, the job +/// would be claimed first and come back only after the job queue's retry delay. The job queue keeps schedules in +/// whole seconds, so a one-second delay could be over at once. +const JOB_DELAY: time::Duration = time::Duration::seconds(2); + +/// What is done for a new Task once its record is written and before it is committed, such as keeping its secrets. +type OnCreated = Box anyhow::Result<()> + Send>; + +/// Builds the job of a new Task from the task store, which only the service holds. +type JobOf = Box DynJob + Send>; + +enum Message { + Create { + record: NewProvisionerTaskRecord, + job_of: JobOf, + on_created: OnCreated, + reply: oneshot::Sender>, + }, + Get { + id: Uuid, + reply: oneshot::Sender>>, + }, +} + +/// Asks the provisioner task service to record or read Tasks. +#[derive(Clone)] +pub struct ProvisionerTaskHandle(mpsc::Sender); + +pub struct ProvisionerTaskReceiver(mpsc::Receiver); + +pub fn provisioner_task_channel() -> (ProvisionerTaskHandle, ProvisionerTaskReceiver) { + let (tx, rx) = mpsc::channel(64); + (ProvisionerTaskHandle(tx), ProvisionerTaskReceiver(rx)) +} + +impl ProvisionerTaskHandle { + /// Records a Task of `task` and stores the job that runs it, both or neither. + /// + /// `on_created` runs only for a new record, before its job is stored, so what the job needs is in place before it + /// can run; if it or storing the job fails, nothing is recorded. A known ID or a busy target creates nothing and + /// stores nothing. + pub async fn create( + &self, + task: T, + id: Uuid, + target: String, + params: &T::Params, + on_created: impl FnOnce(&NewProvisionerTaskRecord) -> anyhow::Result<()> + Send + 'static, + ) -> anyhow::Result { + let params = serde_json::to_value(params).context("serialize the task parameters")?; + let now = OffsetDateTime::now_utc(); + let job_token = Uuid::new_v4(); + let record = NewProvisionerTaskRecord { + id, + kind: T::KIND.to_owned(), + target, + params, + deadline_at: deadline(now, T::START_WINDOW), + job_token, + }; + let job_of: JobOf = Box::new(move |store| Box::new(TaskJob::new(id, job_token, task, store))); + + let (reply, outcome) = oneshot::channel(); + self.0 + .send(Message::Create { + record, + job_of, + on_created: Box::new(on_created), + reply, + }) + .await + .ok() + .context("couldn't send Create message")?; + outcome.await.context("couldn't receive Create result")? + } + + /// Reads a Task, after failing it if it is past its deadline. + pub async fn get(&self, id: Uuid) -> anyhow::Result> { + let (reply, record) = oneshot::channel(); + self.0 + .send(Message::Get { id, reply }) + .await + .ok() + .context("couldn't send Get message")?; + record.await.context("couldn't receive Get result")? + } +} + +/// The provisioner task service: the only one to write Tasks and their jobs from a request. +pub struct ProvisionerTaskService { + store: DynProvisionerTaskStore, + job_queue: DynJobQueue, + rx: ProvisionerTaskReceiver, +} + +impl ProvisionerTaskService { + pub fn new(store: DynProvisionerTaskStore, job_queue: DynJobQueue, rx: ProvisionerTaskReceiver) -> Self { + Self { store, job_queue, rx } + } + + async fn create( + &self, + record: NewProvisionerTaskRecord, + job_of: JobOf, + on_created: OnCreated, + ) -> anyhow::Result { + let now = OffsetDateTime::now_utc(); + let job = job_of(Arc::clone(&self.store)); + + let store_job = async { + on_created(&record).context("prepare the Task")?; + self.job_queue.push_job(&job, Some(now + JOB_DELAY)).await + }; + let outcome = self.store.create(record.clone(), now, Box::pin(store_job)).await?; + + if matches!(outcome, CreateOutcome::Created(_)) { + info!(task.id = %record.id, task.kind = %record.kind, "Task queued"); + } + + Ok(outcome) + } +} + +#[async_trait] +impl Task for ProvisionerTaskService { + type Output = anyhow::Result<()>; + + const NAME: &'static str = "provisioner task service"; + + async fn run(self, shutdown_signal: ShutdownSignal) -> Self::Output { + provisioner_task_service(self, shutdown_signal).await + } +} + +#[instrument(skip_all)] +async fn provisioner_task_service( + mut service: ProvisionerTaskService, + mut shutdown_signal: ShutdownSignal, +) -> anyhow::Result<()> { + debug!("Task started"); + + loop { + tokio::select! { + message = service.rx.0.recv() => { + let Some(message) = message else { + debug!("All senders are dead"); + break; + }; + + match message { + Message::Create { record, job_of, on_created, reply } => { + let _ = reply.send(service.create(record, job_of, on_created).await); + } + Message::Get { id, reply } => { + let _ = reply.send(service.store.get(id, OffsetDateTime::now_utc()).await); + } + } + } + () = shutdown_signal.wait() => break, + } + } + + debug!("Task terminated"); + + Ok(()) +} diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/tests.rs similarity index 67% rename from devolutions-gateway/src/provisioner_tasks/runner/tests.rs rename to devolutions-gateway/src/provisioner_tasks/tests.rs index a2a493df4b..64501ce258 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/tests.rs @@ -1,9 +1,12 @@ +use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use job_queue::{Job as _, JobQueue as _}; +use devolutions_gateway_task::ShutdownHandle; +use job_queue::Job as _; use parking_lot::Mutex; -use provisioner_task::ProvisionerTaskState; +use provisioner_task::{CreateOutcome, DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskState}; use serde_json::json; +use uuid::Uuid; use super::*; use crate::job_queue::{StoredJob, TestJobQueue}; @@ -46,6 +49,10 @@ struct Echo { ended: Arc>>, /// Ends that fail before one succeeds. end_failures: Arc, + /// Publications that fail before one succeeds. + publish_failures: Arc, + /// The records whose output was published. + published: Arc>>, } #[async_trait] @@ -105,6 +112,19 @@ impl ProvisionerTask for Echo { } } + async fn publish(&self, record: &ProvisionerTaskRecord, _: &serde_json::Value) -> anyhow::Result<()> { + if self + .publish_failures + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |left| left.checked_sub(1)) + .is_ok() + { + anyhow::bail!("the publication failed"); + } + + self.published.lock().push(record.id); + Ok(()) + } + async fn on_end(&self, record: &ProvisionerTaskRecord) -> anyhow::Result<()> { if self .end_failures @@ -120,12 +140,18 @@ impl ProvisionerTask for Echo { } struct Harness { - task: Arc, + outcome: fn(u32) -> Result, + progress_timeout: time::Duration, store: DynProvisionerTaskStore, job_queue: Arc, + tasks: ProvisionerTaskHandle, + /// Keeps the provisioner task service running. + _service_shutdown: ShutdownHandle, runs: Arc, ended: Arc>>, end_failures: Arc, + publish_failures: Arc, + published: Arc>>, } impl Harness { @@ -149,25 +175,41 @@ impl Harness { .expect("open task store"), ); let job_queue = TestJobQueue::new().await; + let (tasks, rx) = provisioner_task_channel(); + let (service_shutdown, shutdown_signal) = ShutdownHandle::new(); + let service = ProvisionerTaskService::new(Arc::clone(&store), Arc::::clone(&job_queue), rx); + tokio::spawn(devolutions_gateway_task::Task::run(service, shutdown_signal)); let runs = Arc::new(AtomicUsize::new(0)); let ended = Arc::new(Mutex::new(Vec::new())); let end_failures = Arc::new(AtomicUsize::new(0)); + let publish_failures = Arc::new(AtomicUsize::new(0)); + let published = Arc::new(Mutex::new(Vec::new())); - let task = Arc::new(Echo { + Self { outcome, progress_timeout, - runs: Arc::clone(&runs), - ended: Arc::clone(&ended), - end_failures: Arc::clone(&end_failures), - }); - - Self { - task, store, job_queue, + tasks, + _service_shutdown: service_shutdown, runs, ended, end_failures, + publish_failures, + published, + } + } + + /// The task of this harness: every one shares its counters. + fn echo(&self) -> Echo { + Echo { + outcome: self.outcome, + progress_timeout: self.progress_timeout, + runs: Arc::clone(&self.runs), + ended: Arc::clone(&self.ended), + end_failures: Arc::clone(&self.end_failures), + publish_failures: Arc::clone(&self.publish_failures), + published: Arc::clone(&self.published), } } @@ -186,20 +228,42 @@ impl Harness { id: Uuid, target: String, params: &EchoParams, - on_created: impl FnOnce(&ProvisionerTaskRecord) -> anyhow::Result<()> + Send, + on_created: impl FnOnce(&NewProvisionerTaskRecord) -> anyhow::Result<()> + Send + 'static, ) -> anyhow::Result { - create::(&self.store, &self.job_queue.durable(), id, target, params, on_created).await + self.tasks.create(self.echo(), id, target, params, on_created).await + } + + /// Records a Task directly, with the job token `job_token` and no job. + async fn record_task(&self, id: Uuid, params: serde_json::Value, job_token: Uuid) { + let now = OffsetDateTime::now_utc(); + self.store + .create( + NewProvisionerTaskRecord { + id, + kind: KIND.to_owned(), + target: id.to_string(), + params, + deadline_at: now + time::Duration::hours(1), + job_token, + }, + now, + Box::pin(async { Ok(()) }), + ) + .await + .expect("create"); + } + + /// The job of Task `id` with the job token `job_token`, as the job runner reads it. + fn job(&self, id: Uuid, job_token: Uuid) -> TaskJob { + let json = json!({ "taskId": id, "jobToken": job_token }).to_string(); + TaskJob::read_json(&json, self.echo(), Arc::clone(&self.store)).expect("valid job") } /// Takes the oldest queued job, read back as the job runner reads it. fn queued_job(&mut self) -> Box { let StoredJob { name, json, .. } = self.job_queue.take().expect("a queued job"); assert_eq!(name, KIND); - Box::new( - TaskJobUnhydrated::::read_json(&json) - .expect("valid job") - .into_job(Arc::clone(&self.task), Arc::clone(&self.store)), - ) + Box::new(TaskJob::read_json(&json, self.echo(), Arc::clone(&self.store)).expect("valid job")) } async fn task(&self, id: Uuid) -> ProvisionerTaskRecord { @@ -216,9 +280,10 @@ impl Harness { 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 created = Arc::new(AtomicUsize::new(0)); let create = || { - harness.create_task(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + let created = Arc::clone(&created); + harness.create_task(id, "target".to_owned(), &EchoParams { value: 1 }, move |_| { created.fetch_add(1, Ordering::SeqCst); Ok(()) }) @@ -246,24 +311,49 @@ async fn a_task_that_cannot_be_prepared_gets_no_job() { .expect_err("preparing fails"); assert!(harness.job_queue.take().is_none(), "no job is queued"); + assert!( + harness + .store + .get(id, OffsetDateTime::now_utc()) + .await + .expect("read task") + .is_none(), + "no Task is recorded" + ); +} + +// oracle: mutant a job claimed before its Task is committed is dropped as done, and the Task never runs +#[tokio::test] +async fn a_job_whose_task_is_not_recorded_yet_is_run_again_later() { + let harness = Harness::new(|value| Ok(json!(value))).await; + let id = Uuid::new_v4(); + let job_token = Uuid::new_v4(); + let mut job = harness.job(id, job_token); + + job.run().await.expect_err("the job queue runs the job again"); + assert_eq!(harness.runs.load(Ordering::SeqCst), 0); + + harness.record_task(id, json!({ "value": 1 }), job_token).await; + job.run().await.expect("the Task is recorded now"); + + assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Succeeded); + assert_eq!(harness.runs.load(Ordering::SeqCst), 1); } -// oracle: mutant a panicking attempt fails the Task as internal instead of panicked +// oracle: mutant the job stored for a Task that was never committed runs the Task recorded since under the same ID #[tokio::test] -async fn a_panicking_attempt_fails_the_task() { - let mut harness = Harness::new(|_| panic!("the attempt panics")).await; +async fn only_the_job_recorded_with_the_task_runs_it() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; let id = harness.create(1).await; + let mut other_job = harness.job(id, Uuid::new_v4()); - harness.queued_job().run().await.expect("no retry"); + other_job.run().await.expect("the other job is done for good"); + assert_eq!(harness.runs.load(Ordering::SeqCst), 0); + assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Queued); - let task = harness.task(id).await; - assert_eq!(task.state, ProvisionerTaskState::Failed); - assert_eq!( - task.payload, - json!({ "error": { "kind": "panicked" }, "last": { "after": null } }), - "the failure keeps what the attempt started with" - ); - assert_eq!(*harness.ended.lock(), [id]); + harness.queued_job().run().await.expect("the job of the Task runs it"); + assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Succeeded); + assert_eq!(harness.runs.load(Ordering::SeqCst), 1); } // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §Management @@ -328,67 +418,16 @@ async fn a_job_delivered_again_after_the_last_attempt_fails_the_task() { assert_eq!(*harness.ended.lock(), [id]); } -// oracle: mutant a queued Task whose job was never stored stays queued until its deadline after a restart -#[tokio::test] -async fn a_queued_task_without_a_job_gets_one_at_startup() { - let mut harness = Harness::new(|value| Ok(json!(value))).await; - harness.create(1).await; - // The job stays in the job queue; only the copy handed to the test is taken. - harness.queued_job(); - let id = Uuid::new_v4(); - let now = OffsetDateTime::now_utc(); - harness - .store - .create( - NewProvisionerTaskRecord { - id, - kind: KIND.to_owned(), - target: id.to_string(), - params: json!({ "value": 2 }), - deadline_at: now + time::Duration::hours(1), - }, - now, - ) - .await - .expect("create"); - - let job_defs = harness.job_queue.job_defs(KIND).await.expect("stored jobs"); - requeue_lost_jobs::(&harness.store, &harness.job_queue.durable(), &job_defs) - .await - .expect("requeue"); - - harness.queued_job().run().await.expect("the job runs the Task"); - assert!( - harness.job_queue.take().is_none(), - "the Task that has a job gets no other" - ); - assert_eq!(harness.task(id).await.state, ProvisionerTaskState::Succeeded); -} - // oracle: mutant a Task whose parameters cannot be read runs its attempt, or stays unfinished #[tokio::test] async fn invalid_parameters_fail_the_task() { let harness = Harness::new(|value| Ok(json!(value))).await; let id = Uuid::new_v4(); - let now = OffsetDateTime::now_utc(); + let job_token = Uuid::new_v4(); harness - .store - .create( - NewProvisionerTaskRecord { - id, - kind: KIND.to_owned(), - target: id.to_string(), - params: json!({ "value": "not a number" }), - deadline_at: now + time::Duration::hours(1), - }, - now, - ) - .await - .expect("create"); - let json = json!({ "taskId": id }).to_string(); - let mut job = TaskJobUnhydrated::::read_json(&json) - .expect("valid job") - .into_job(Arc::clone(&harness.task), Arc::clone(&harness.store)); + .record_task(id, json!({ "value": "not a number" }), job_token) + .await; + let mut job = harness.job(id, job_token); job.run().await.expect("no retry"); @@ -398,6 +437,30 @@ async fn invalid_parameters_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 a_success_is_recorded_only_once_its_output_is_published() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + let mut job = harness.queued_job(); + harness.publish_failures.store(1, Ordering::SeqCst); + + job.run().await.expect_err("the job queue runs the job again"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Running, "{:?}", task.payload); + assert!(harness.published.lock().is_empty()); + assert!(harness.ended.lock().is_empty()); + + job.run().await.expect("the publication succeeds"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert_eq!(task.attempts, 2, "the output comes from a new attempt"); + assert_eq!(*harness.published.lock(), [id]); + assert_eq!(*harness.ended.lock(), [id]); +} + // oracle: intent devolutions-gateway/src/provisioner_tasks/INTENT.md §AI analysis #[tokio::test] async fn an_end_that_fails_is_done_again_when_the_job_queue_runs_the_job_again() { @@ -454,9 +517,7 @@ async fn the_job_is_read_back_from_its_json() { let job = harness.queued_job(); let json = job.write_json().expect("job JSON"); - let read = TaskJobUnhydrated::::read_json(&json) - .expect("valid job") - .into_job(Arc::clone(&harness.task), Arc::clone(&harness.store)); + let read = TaskJob::read_json(&json, harness.echo(), Arc::clone(&harness.store)).expect("valid job"); assert_eq!(read.task_id, id); assert_eq!(job.name(), KIND); diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 5fc4125555..f362a69b60 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -577,8 +577,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, a `file` already gone lists nothing new, and a session - /// whose recording was deleted meanwhile only loses `file`. + /// A `file` that is gone is an error, since only its maker moves it, and so is a session whose recording was deleted + /// meanwhile, which 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 @@ -999,7 +999,7 @@ impl RecordingManagerTask { Err(error) if error.kind() != io::ErrorKind::NotFound => { return Err(anyhow::Error::new(error).context(format!("remove {file}"))); } - _ => return Ok(()), + _ => anyhow::bail!("recording {id} was deleted"), } } @@ -1009,17 +1009,9 @@ impl RecordingManagerTask { let artifact_path = recording_path.join(&file_name); // The rename swaps the file in at once and the manifest follows, so it never lists a missing or partial file. - match fs::rename(file, &artifact_path).await { - Ok(()) => {} - // Already moved, by an earlier call cut short before the manifest was written, or by a newer Task that listed - // it for the Task that made it; what is in place is listed. - Err(error) if error.kind() == io::ErrorKind::NotFound && artifact_path.exists() => {} - Err(error) if error.kind() == io::ErrorKind::NotFound => { - warn!(%file, "The artifact to list is gone"); - return Ok(()); - } - Err(error) => return Err(anyhow::Error::new(error).context(format!("move {file} to {artifact_path}"))), - } + fs::rename(file, &artifact_path) + .await + .with_context(|| format!("move {file} to {artifact_path}"))?; let earlier = core::mem::replace( manifest.artifacts.of_kind_mut(kind), diff --git a/devolutions-gateway/src/secret_store.intent.md b/devolutions-gateway/src/secret_store.intent.md index d24e08089d..9e65642ca9 100644 --- a/devolutions-gateway/src/secret_store.intent.md +++ b/devolutions-gateway/src/secret_store.intent.md @@ -13,3 +13,8 @@ All of this is best effort, as if someone has root access and we still try our b ## Replacement Storing a secret under an ID that already has one replaces the old one. + + +## Sharing + +Gateway has one secret store, shared by the handlers and the background services, so a secret a handler stores is found by the background service that uses it, such as the API key of an AI analysis. diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index 0c9666b85c..828e1b3743 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -6,8 +6,7 @@ use devolutions_gateway::config::{Conf, ConfHandle}; use devolutions_gateway::job_queue::{InMemoryJobDependencies, JobQueueTask, JobRunnerTask}; use devolutions_gateway::listener::{GatewayListener, ListenerUrls}; use devolutions_gateway::log::GatewayLog; -use devolutions_gateway::provisioner_tasks::recording_ai_analysis::RecordingAiAnalysis; -use devolutions_gateway::provisioner_tasks::runner::{TaskJob, requeue_lost_jobs}; +use devolutions_gateway::provisioner_tasks::{ProvisionerTaskService, provisioner_task_channel}; use devolutions_gateway::recording::recording_message_channel; use devolutions_gateway::session::session_manager_channel; use devolutions_gateway::subscriber::subscriber_channel; @@ -354,25 +353,15 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { ); let secrets = devolutions_gateway::secret_store::SecretStore::default(); - let durable_job_queue = job_queue_ctx.durable(); - - // Before the job runner starts, so a job that runs meanwhile is never taken for lost. - let requeued = async { - let job_defs = job_queue_ctx.job_defs(TaskJob::::NAME).await?; - requeue_lost_jobs::(&task_store, &durable_job_queue, &job_defs).await - }; - if let Err(error) = requeued.await { - error!( - error = format!("{error:#}"), - "Failed to queue again the jobs of the provisioner tasks" - ); - } + let (provisioner_tasks, provisioner_task_rx) = provisioner_task_channel(); + let provisioner_task_service = + ProvisionerTaskService::new(Arc::clone(&task_store), job_queue_ctx.queue(), provisioner_task_rx); let in_memory_job_dependencies = InMemoryJobDependencies { conf_handle: conf_handle.clone(), recordings: recording_manager_handle.clone(), secrets: secrets.clone(), - task_store: Arc::clone(&task_store), + task_store, }; let state = DgwState { @@ -389,9 +378,8 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { monitoring_state, traffic_audit_handle: traffic_audit_task.handle(), agent_tunnel_handle, - task_store, + provisioner_tasks, secrets: secrets.clone(), - durable_job_queue, }; for listener in &conf.listeners { @@ -427,6 +415,8 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { tasks.register(devolutions_gateway::secret_store::CleanupTask { handle: secrets }); + tasks.register(provisioner_task_service); + tasks.register(devolutions_gateway::credential_injection::CleanupTask { handle: synthetic_kdc_registry, }); diff --git a/testsuite/tests/cli/dgw/ai_analysis.rs b/testsuite/tests/cli/dgw/ai_analysis.rs index 6909714e9e..ad5b676898 100644 --- a/testsuite/tests/cli/dgw/ai_analysis.rs +++ b/testsuite/tests/cli/dgw/ai_analysis.rs @@ -283,7 +283,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 +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn start_returns_202_then_200_for_the_same_request() { let gateway = Gateway::start().await; @@ -319,7 +319,7 @@ async fn start_returns_202_then_200_for_the_same_request() { } } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn start_conflicts_return_409_with_a_code() { let gateway = Gateway::start().await; @@ -345,7 +345,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 +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn start_without_recording_returns_404() { let gateway = Gateway::start().await; @@ -358,7 +358,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 +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn start_ignores_unknown_fields() { let gateway = Gateway::start().await; @@ -372,7 +372,7 @@ async fn start_ignores_unknown_fields() { assert_eq!(status, StatusCode::ACCEPTED, "{created}"); } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn invalid_bodies_and_settings_return_400() { let gateway = Gateway::start().await; @@ -425,7 +425,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 +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn endpoints_require_their_scope() { let gateway = Gateway::start().await; @@ -454,7 +454,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 +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn an_analysis_is_found_at_its_location() { let gateway = Gateway::start().await; @@ -478,7 +478,7 @@ async fn an_analysis_is_found_at_its_location() { assert_eq!(task["id"], task_id.to_string()); } -// oracle: intent devolutions-gateway/src/api/jrec.intent.md §AI analysis +// oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis #[tokio::test] async fn get_analysis_returns_404_for_an_unknown_id() { let gateway = Gateway::start().await; @@ -505,7 +505,7 @@ async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { assert_eq!(task["state"], "succeeded", "{task}"); 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 + // oracle: intent devolutions-gateway/src/api/tasks/recording_ai_analysis.intent.md §AI analysis assert_eq!( task["payload"]["usage"], json!({ "inputTokens": 20, "outputTokens": 40, "unreportedRequests": 0 }), From d82fd46b0fbad5b5211fbaf5a1c774740680324d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Fri, 30 Oct 2026 12:37:36 -0400 Subject: [PATCH 80/80] refactor(dgw): store a new Task's job without a delay The job of a new Task was scheduled 2 s ahead so the job queue would claim it after the record is committed. That was never needed: the job's first read of its record waits on the same store lock `create` holds until the commit, so an early claim reads the committed record anyway. Issue: DGW-462 Co-Authored-By: Claude Opus 5.5 --- devolutions-gateway/src/provisioner_tasks/service.rs | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/service.rs b/devolutions-gateway/src/provisioner_tasks/service.rs index 2509e037a5..7577ade0a2 100644 --- a/devolutions-gateway/src/provisioner_tasks/service.rs +++ b/devolutions-gateway/src/provisioner_tasks/service.rs @@ -14,13 +14,6 @@ use uuid::Uuid; use super::{ProvisionerTask, TaskJob, deadline}; -/// How far ahead the job of a new Task is scheduled, so the job queue claims it after the record is committed. -/// -/// The job is stored before the record is committed, and storing it wakes the job queue; without the delay, the job -/// would be claimed first and come back only after the job queue's retry delay. The job queue keeps schedules in -/// whole seconds, so a one-second delay could be over at once. -const JOB_DELAY: time::Duration = time::Duration::seconds(2); - /// What is done for a new Task once its record is written and before it is committed, such as keeping its secrets. type OnCreated = Box anyhow::Result<()> + Send>; @@ -127,7 +120,7 @@ impl ProvisionerTaskService { let store_job = async { on_created(&record).context("prepare the Task")?; - self.job_queue.push_job(&job, Some(now + JOB_DELAY)).await + self.job_queue.push_job(&job, None).await }; let outcome = self.store.create(record.clone(), now, Box::pin(store_job)).await?;