From 02286371c57efb5a73b73c1321024b75a438d6ac Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 6 Oct 2026 15:10:51 -0400 Subject: [PATCH 01/21] 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 73f9702c8..1e7cb4056 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1885,6 +1885,8 @@ dependencies = [ "picky-krb", "pin-project-lite 0.2.17", "proptest", + "provisioner-task", + "provisioner-task-libsql", "rand 0.10.2", "reqwest", "rstest", diff --git a/devolutions-gateway/Cargo.toml b/devolutions-gateway/Cargo.toml index bc65fd8de..5b49988d7 100644 --- a/devolutions-gateway/Cargo.toml +++ b/devolutions-gateway/Cargo.toml @@ -38,6 +38,8 @@ sysevent-codes.path = "../crates/sysevent-codes" agent-tunnel = { path = "../crates/agent-tunnel" } agent-tunnel-libsql.path = "../crates/agent-tunnel-libsql" gateway-db.path = "../crates/gateway-db" +provisioner-task.path = "../crates/provisioner-task" +provisioner-task-libsql.path = "../crates/provisioner-task-libsql" ironrdp-pdu = { version = "0.9", features = ["std"] } ironrdp-core = { version = "0.2", features = ["std"] } ironrdp-rdcleanpath = "0.2" diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 8985f3184..615bbfb63 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -11,6 +11,9 @@ use job_queue_libsql::libsql; use time::OffsetDateTime; use tokio::sync::{Notify, mpsc}; +/// Runs of a job before the job queue gives up on it. +pub const MAX_ATTEMPTS: u32 = 5; + pub struct JobQueueCtx { notify_runner: Arc, runner_waker: RunnerWaker, @@ -59,6 +62,7 @@ impl JobQueueCtx { let queue = job_queue_libsql::LibSqlJobQueue::builder() .runner_waker(runner_waker.clone()) .conn(conn) + .max_attempts(MAX_ATTEMPTS) .build(); let queue = Arc::new(queue); diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index cd53900b2..a928aff20 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -68,6 +68,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)] @@ -82,7 +83,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())); @@ -95,6 +96,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, @@ -110,6 +114,7 @@ impl DgwState { synthetic_kdc_registry, monitoring_state, agent_tunnel_handle: None, + provisioner_tasks, }; let handles = MockHandles { diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index 258e15d23..32abd87a7 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -340,6 +340,16 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { None }; + let provisioner_tasks_conn = gateway_db + .connect() + .await + .context("failed to connect the provisioner tasks to the gateway database")?; + let provisioner_tasks = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(provisioner_tasks_conn) + .await + .context("failed to initialize the provisioner task store")?, + ); + let state = DgwState { conf_handle: conf_handle.clone(), token_cache: Arc::clone(&token_cache), @@ -354,6 +364,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { monitoring_state, traffic_audit_handle: traffic_audit_task.handle(), agent_tunnel_handle, + provisioner_tasks, }; for listener in &conf.listeners { diff --git a/devolutions-gateway/tests/preflight.rs b/devolutions-gateway/tests/preflight.rs index 1cbcf8480..71a312c62 100644 --- a/devolutions-gateway/tests/preflight.rs +++ b/devolutions-gateway/tests/preflight.rs @@ -58,15 +58,15 @@ fn unsigned_jws(payload: serde_json::Value) -> anyhow::Result { Ok(format!("{header}.{payload}.{signature}")) } -fn make_router() -> anyhow::Result<(Router, DgwState, MockHandles)> { - let (state, handles) = DgwState::mock(CONFIG)?; +async fn make_router() -> anyhow::Result<(Router, DgwState, MockHandles)> { + let (state, handles) = DgwState::mock(CONFIG).await?; let app = devolutions_gateway::make_http_service(state.clone()) .layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); Ok((app, state, handles)) } -fn make_router_with_config(config: &str) -> anyhow::Result<(Router, DgwState, MockHandles)> { - let (state, handles) = DgwState::mock(config)?; +async fn make_router_with_config(config: &str) -> anyhow::Result<(Router, DgwState, MockHandles)> { + let (state, handles) = DgwState::mock(config).await?; let app = devolutions_gateway::make_http_service(state.clone()) .layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); Ok((app, state, handles)) @@ -83,7 +83,7 @@ fn init_logger() -> tracing::subscriber::DefaultGuard { async fn test_provision_credentials_success() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let jti = Uuid::new_v4(); let token = unsigned_jws(json!({ @@ -124,7 +124,7 @@ async fn test_provision_credentials_success_when_unstable_disabled() -> anyhow:: // `provision-credentials` is protocol-neutral: NTLM credential injection relies on this // path even when the unstable feature flag is off. let config = CONFIG.replace("\"enable_unstable\": true", "\"enable_unstable\": false"); - let (app, _state, _handles) = make_router_with_config(&config)?; + let (app, _state, _handles) = make_router_with_config(&config).await?; let jti = Uuid::new_v4(); let token = unsigned_jws(json!({ @@ -161,7 +161,7 @@ async fn test_provision_credentials_success_when_unstable_disabled() -> anyhow:: async fn test_provision_credentials_rejects_alternate_targets() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let token = unsigned_jws(json!({ "jti": "5e3e833f-84c7-4541-b676-acc3299e39b8", @@ -205,7 +205,7 @@ async fn test_provision_credentials_rejects_alternate_targets() -> anyhow::Resul async fn test_provision_credentials_rejects_missing_target_hostname() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let token = unsigned_jws(json!({ "jti": "5e3e833f-84c7-4541-b676-acc3299e39b8" @@ -247,7 +247,7 @@ async fn test_provision_credentials_rejects_missing_target_hostname() -> anyhow: async fn test_provision_token_overwrite_alert() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; // Same JTI twice: second provision-token replaces the stored token-only entry (master behavior). let token = unsigned_jws(json!({ @@ -288,7 +288,7 @@ async fn test_provision_token_overwrite_alert() -> anyhow::Result<()> { async fn test_provision_connection_options_success() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let jti = Uuid::new_v4(); let token = unsigned_jws(json!({ @@ -319,7 +319,7 @@ async fn test_provision_connection_options_success() -> anyhow::Result<()> { async fn test_provision_credentials_and_connection_options_fold() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let jti = Uuid::new_v4(); let token = unsigned_jws(json!({ @@ -361,7 +361,7 @@ async fn test_provision_credentials_and_connection_options_fold() -> anyhow::Res async fn test_provision_invalid_params() -> anyhow::Result<()> { let _guard = init_logger(); - let (app, _state, _handles) = make_router()?; + let (app, _state, _handles) = make_router().await?; let token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJqdGkiOiI1ZTNlODMzZi04NGM3LTQ1NDEtYjY3Ni1hY2MzMjk5ZTM5YjgifQ.1qECGlrW7y9HWFArc6GPHLGTOY7PhAvzKJ5XMRBg4k4"; diff --git a/devolutions-gateway/tests/traffic.rs b/devolutions-gateway/tests/traffic.rs index 7175d3bb5..3c9e4a63c 100644 --- a/devolutions-gateway/tests/traffic.rs +++ b/devolutions-gateway/tests/traffic.rs @@ -71,7 +71,7 @@ impl Drop for HandlesGuard { /// Build a Router with a live TrafficAuditManagerTask and real handle, /// similar to tests/preflight.rs harness. async fn make_router() -> anyhow::Result<(Router, DgwState, HandlesGuard)> { - let (mut state, handles) = DgwState::mock(CONFIG)?; + let (mut state, handles) = DgwState::mock(CONFIG).await?; // Start the manager with an in-memory DB. let manager = TrafficAuditManagerTask::init(":memory:").await?; diff --git a/devolutions-gateway/tests/webapp.rs b/devolutions-gateway/tests/webapp.rs index 8e090e9d1..25755d7e0 100644 --- a/devolutions-gateway/tests/webapp.rs +++ b/devolutions-gateway/tests/webapp.rs @@ -60,7 +60,7 @@ fn initialize_conf() { async fn custom_authentication_flow() -> anyhow::Result<()> { let (cov, _guard) = init_cov_mark(); initialize_conf(); - let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG)?; + let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG).await?; let mut app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); @@ -191,7 +191,7 @@ async fn static_assets_are_served(#[case] subdir: &str, #[case] uri_prefix: &str config["WebApp"]["StaticRootPath"] = json!(static_root); let config = config.to_string(); - let (state, _handle) = devolutions_gateway::DgwState::mock(&config)?; + let (state, _handle) = devolutions_gateway::DgwState::mock(&config).await?; let mut app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); @@ -229,7 +229,7 @@ async fn static_assets_are_served(#[case] subdir: &str, #[case] uri_prefix: &str async fn sign_app_token_bad_password() -> anyhow::Result<()> { let (cov, _guard) = init_cov_mark(); initialize_conf(); - let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG)?; + let (state, _handle) = devolutions_gateway::DgwState::mock(CONFIG).await?; let app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); @@ -267,7 +267,7 @@ async fn sign_app_token_bad_password() -> anyhow::Result<()> { async fn sign_app_token_username_mismatch() -> anyhow::Result<()> { let (cov, _guard) = init_cov_mark(); initialize_conf(); - let (state, _handles) = devolutions_gateway::DgwState::mock(CONFIG)?; + let (state, _handles) = devolutions_gateway::DgwState::mock(CONFIG).await?; let app = devolutions_gateway::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([0, 0, 0, 0], 3000)))); From 28e7a5545214fb53b94aa32500c5b87bef844f0c Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 7 Oct 2026 17:03:00 -0400 Subject: [PATCH 02/21] 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 1e7cb4056..81bf14fc0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1847,6 +1847,7 @@ dependencies = [ "cfg-if", "chacha20poly1305", "devolutions-agent-shared", + "devolutions-gateway-ai", "devolutions-gateway-generators", "devolutions-gateway-task", "devolutions-log", @@ -1914,6 +1915,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 5b49988d7..317d40694 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"] } reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-native-roots", "json", "stream", "http2", "socks"] } http-client-proxy = { path = "../crates/http-client-proxy" } diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 615bbfb63..3e4b611fc 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -6,11 +6,13 @@ use std::time::Duration; use anyhow::Context as _; use async_trait::async_trait; use devolutions_gateway_task::{ChildTask, ShutdownSignal, Task}; -use job_queue::{DynJobQueue, Job, JobCtx, JobQueue, JobReader, JobRunner, RunnerWaker}; +use job_queue::{DynJob, DynJobQueue, Job, JobCtx, JobQueue, JobReader, JobRunner, RunnerWaker}; use job_queue_libsql::libsql; use time::OffsetDateTime; use tokio::sync::{Notify, mpsc}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskRunner}; + /// Runs of a job before the job queue gives up on it. pub const MAX_ATTEMPTS: u32 = 5; @@ -41,6 +43,7 @@ pub struct JobRunnerTask { notify_runner: Arc, runner_waker: RunnerWaker, queue: DynJobQueue, + tasks: ProvisionerTaskRunner, } impl JobQueueCtx { @@ -86,6 +89,14 @@ impl JobQueueCtx { job_queue_handle: handle, }) } + + /// Writes a job to the queue right away, without going through [`JobQueueHandle`]. + /// + /// For startup, before the job runner claims anything. + pub async fn push(&self, job: T) -> anyhow::Result<()> { + let job: DynJob = Box::new(job); + self.queue.push_job(&job, None).await + } } impl JobQueueHandle { @@ -201,11 +212,12 @@ async fn job_queue_task(ctx: JobQueueTask, mut shutdown_signal: ShutdownSignal) } impl JobRunnerTask { - pub fn new(ctx: &JobQueueCtx) -> Self { + pub fn new(ctx: &JobQueueCtx, tasks: ProvisionerTaskRunner) -> Self { Self { notify_runner: Arc::clone(&ctx.notify_runner), runner_waker: RunnerWaker::clone(&ctx.runner_waker), queue: Arc::clone(&ctx.queue), + tasks, } } } @@ -229,9 +241,10 @@ async fn job_runner_task(ctx: JobRunnerTask, mut shutdown_signal: ShutdownSignal notify_runner, runner_waker, queue, + tasks, } = ctx; - let reader = DgwJobReader; + let reader = DgwJobReader { tasks }; let spawn = |mut ctx: JobCtx, callback: job_queue::SpawnCallback| { tokio::spawn(async move { @@ -283,10 +296,12 @@ async fn job_runner_task(ctx: JobRunnerTask, mut shutdown_signal: ShutdownSignal Ok(()) } -struct DgwJobReader; +struct DgwJobReader { + tasks: ProvisionerTaskRunner, +} impl JobReader for DgwJobReader { - fn read_json(&self, name: &str, json: &str) -> anyhow::Result { + fn read_json(&self, name: &str, json: &str) -> anyhow::Result { use crate::api::jrec::DeleteRecordingsJob; use crate::recording::RemuxJob; @@ -300,6 +315,10 @@ impl JobReader for DgwJobReader { serde_json::from_str(json).context("failed to deserialize DeleteRecordingsJob")?; Ok(Box::new(job)) } + ProvisionerTaskJob::NAME => { + let job = ProvisionerTaskJob::read_json(json, self.tasks.clone())?; + Ok(Box::new(job)) + } _ => anyhow::bail!("unknown job name: {name}"), } } diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index a928aff20..328d015df 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; @@ -68,7 +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, + pub provisioner_tasks: provisioner_tasks::runner::ProvisionerTaskRunner, } #[doc(hidden)] @@ -96,9 +97,15 @@ impl DgwState { let provisioning = provisioning::ProvisioningStore::new(); let synthetic_kdc_registry = credential_injection::SyntheticKdcRegistry::new(); let monitoring_state = Arc::new(network_monitor::State::new(Arc::new(MockMonitorsCache))?); - let provisioner_tasks = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?; - let provisioner_tasks = - Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(provisioner_tasks).await?); + let task_store = gateway_db::GatewayDb::open_path(":memory:").await?.connect().await?; + let task_store = Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store).await?); + let provisioner_tasks = provisioner_tasks::task_runner( + task_store, + job_queue_handle.clone(), + conf_handle.clone(), + recording_manager_handle.clone(), + provisioning.clone(), + ); let state = Self { conf_handle, diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs new file mode 100644 index 000000000..8b17d3746 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -0,0 +1,262 @@ +//! AI settings shared by the provisioner tasks that call an AI provider. +//! +//! The API key and the base URL never reach `gateway.db`: they stay in memory, encrypted, as the task secret of their +//! Task in the [`ProvisioningStore`](crate::provisioning::ProvisioningStore), until the Task ends. + +use core::fmt; + +use anyhow::Context as _; +use devolutions_gateway_ai::recording_analysis::RecordingAnalysisError; +use devolutions_gateway_ai::{AiClient, BuildError, Provider}; +use secrecy::{ExposeSecret as _, SecretString}; +use url::Url; + +use super::AttemptError; +use crate::config::Conf; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum AiProvider { + #[serde(rename = "openai")] + OpenAi, + #[serde(rename = "anthropic")] + Anthropic, + #[serde(rename = "mistral")] + Mistral, + #[serde(rename = "gemini")] + Gemini, + #[serde(rename = "openai-compatible")] + OpenAiCompatible, +} + +impl From for Provider { + fn from(provider: AiProvider) -> Self { + match provider { + AiProvider::OpenAi => Provider::OpenAi, + AiProvider::Anthropic => Provider::Anthropic, + AiProvider::Mistral => Provider::Mistral, + AiProvider::Gemini => Provider::Gemini, + AiProvider::OpenAiCompatible => Provider::OpenAiCompatible, + } + } +} + +/// How to reach an AI provider, except the API key. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AiSettings { + pub provider: AiProvider, + pub model: String, + /// Required for [`AiProvider::OpenAiCompatible`]; the others have a default. + pub base_url: Option, + /// Upper bound of tokens in each AI answer. + pub max_output_tokens: Option, +} + +/// Why AI settings were refused before a Task was recorded. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum AiSettingsError { + #[error("AI model is missing")] + MissingModel, + #[error("API key is missing")] + MissingApiKey, + #[error("base URL is missing for this AI provider")] + MissingBaseUrl, + #[error("invalid AI settings: {0}")] + Invalid(String), +} + +impl From<&BuildError> for AiSettingsError { + fn from(error: &BuildError) -> Self { + match error { + BuildError::MissingModel => Self::MissingModel, + BuildError::MissingApiKey(_) => Self::MissingApiKey, + BuildError::MissingBaseUrl(_) => Self::MissingBaseUrl, + _ => Self::Invalid(error.to_string()), + } + } +} + +pub(crate) enum ClientError { + Build(BuildError), + HttpClient(reqwest::Error), +} + +impl fmt::Display for ClientError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ClientError::Build(error) => fmt::Display::fmt(error, f), + ClientError::HttpClient(error) => write!(f, "failed to build the HTTP client: {error}"), + } + } +} + +/// What a Task needs to reach its AI provider and must never write to disk. +#[derive(Debug, Clone)] +pub(crate) struct AiAccess { + pub(crate) api_key: SecretString, + pub(crate) base_url: Option, +} + +/// [`AiAccess`] as written into its task secret. +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct AiAccessSecretRef<'a> { + api_key: &'a str, + base_url: Option<&'a Url>, +} + +/// [`AiAccess`] as read back from its task secret. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct AiAccessSecret { + api_key: String, + base_url: Option, +} + +impl AiAccess { + /// The task secret holding this access. + pub(crate) fn to_secret(&self) -> anyhow::Result { + let json = serde_json::to_string(&AiAccessSecretRef { + api_key: self.api_key.expose_secret(), + base_url: self.base_url.as_ref(), + }) + .context("serialize the AI access")?; + + Ok(SecretString::from(json)) + } + + /// Reads back an access from its task secret. + pub(crate) fn from_secret(secret: &SecretString) -> anyhow::Result { + // The parse error is dropped: it may quote the API key. + let access: AiAccessSecret = + serde_json::from_str(secret.expose_secret()).map_err(|_| anyhow::anyhow!("invalid AI access secret"))?; + + Ok(Self { + api_key: SecretString::from(access.api_key), + base_url: access.base_url, + }) + } +} + +impl AiSettings { + /// Builds a client going through the proxy configured for Gateway. + pub(crate) fn client(&self, conf: &Conf, api_key: &SecretString) -> Result { + let provider = Provider::from(self.provider); + + let mut builder = AiClient::builder() + .provider(provider) + .model(self.model.clone()) + .api_key(api_key.clone()); + + let endpoint = self.base_url.clone().or_else(|| provider.default_base_url()); + + if let Some(base_url) = self.base_url.clone() { + builder = builder.base_url(base_url); + } + + // Without an endpoint, `build` reports the missing base URL before it needs the HTTP client. + if let Some(endpoint) = endpoint { + let proxy_config = conf.proxy.to_proxy_config(); + + let http_client = + http_client_proxy::get_or_create_cached_client(reqwest::Client::builder(), &endpoint, &proxy_config) + .map_err(|error| ClientError::HttpClient(error.without_url()))?; + + builder = builder.http_client(http_client); + } + + builder.build().map_err(ClientError::Build) + } +} + +/// Retries only the failures that may pass later, such as a rate limit or a network error. +impl From for AttemptError { + fn from(error: RecordingAnalysisError) -> Self { + match error { + RecordingAnalysisError::Cancelled => AttemptError::Cancelled, + error if error.is_transient() => AttemptError::Transient(error.to_string()), + error => AttemptError::Permanent(error.to_string()), + } + } +} + +#[cfg(test)] +mod tests { + use devolutions_gateway_ai::Error; + + use super::*; + + #[test] + fn build_errors_map_to_settings_errors() { + assert_eq!( + AiSettingsError::from(&BuildError::MissingModel), + AiSettingsError::MissingModel + ); + assert_eq!( + AiSettingsError::from(&BuildError::MissingApiKey(Provider::OpenAi)), + AiSettingsError::MissingApiKey + ); + assert_eq!( + AiSettingsError::from(&BuildError::MissingBaseUrl(Provider::OpenAiCompatible)), + AiSettingsError::MissingBaseUrl + ); + assert!(matches!( + AiSettingsError::from(&BuildError::InvalidApiKey), + AiSettingsError::Invalid(_) + )); + } + + #[test] + fn rate_limits_server_and_network_errors_are_transient() { + let attempt = |error: Error| AttemptError::from(RecordingAnalysisError::Ai(error)); + let status = |status| Error::Status { + status, + message: "failed".to_owned(), + code: None, + retry_after: None, + }; + + let network = Error::Transport { + message: "connection refused".to_owned(), + }; + assert!(matches!(attempt(network), AttemptError::Transient(_))); + + for code in [429, 500, 503] { + assert!(matches!(attempt(status(code)), AttemptError::Transient(_)), "{code}"); + } + + for code in [400, 401, 403, 404] { + assert!(matches!(attempt(status(code)), AttemptError::Permanent(_)), "{code}"); + } + + let invalid = Error::InvalidOutput { + reason: "no valid action line".to_owned(), + }; + assert!(matches!(attempt(invalid), AttemptError::Permanent(_))); + + assert!(matches!( + AttemptError::from(RecordingAnalysisError::Truncated), + AttemptError::Permanent(_) + )); + assert_eq!( + AttemptError::from(RecordingAnalysisError::Cancelled), + AttemptError::Cancelled + ); + } + + #[test] + fn providers_keep_their_wire_names() { + for (provider, name) in [ + (AiProvider::OpenAi, "openai"), + (AiProvider::Anthropic, "anthropic"), + (AiProvider::Mistral, "mistral"), + (AiProvider::Gemini, "gemini"), + (AiProvider::OpenAiCompatible, "openai-compatible"), + ] { + assert_eq!(serde_json::to_value(provider).expect("serializable"), name); + assert_eq!( + serde_json::from_value::(serde_json::Value::from(name)).expect("known name"), + provider + ); + } + } +} diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs new file mode 100644 index 000000000..b8bc845eb --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -0,0 +1,41 @@ +//! Work the provisioner asks Gateway to do: each request is a Task recorded in `gateway.db`, run as a job on the job +//! queue by the [`runner::ProvisionerTaskRunner`]. + +pub mod ai; +pub mod recording_ai_analysis; +pub mod runner; + +use provisioner_task::DynProvisionerTaskStore; + +use self::recording_ai_analysis::RecordingAiAnalysis; +use self::runner::ProvisionerTaskRunner; +use crate::config::ConfHandle; +use crate::job_queue::JobQueueHandle; +use crate::provisioning::ProvisioningStore; +use crate::recording::RecordingMessageSender; + +/// The provisioner task runner, with every provisioner task Gateway knows. +pub fn task_runner( + store: DynProvisionerTaskStore, + job_queue: JobQueueHandle, + conf_handle: ConfHandle, + recordings: RecordingMessageSender, + provisioning: ProvisioningStore, +) -> ProvisionerTaskRunner { + ProvisionerTaskRunner::builder(store, job_queue) + .register(RecordingAiAnalysis::new(conf_handle, recordings, provisioning)) + .build() +} + +/// Why one attempt of a Task failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AttemptError { + /// Worth another attempt later, such as a rate limit or a network error. + Transient(String), + /// Another attempt would fail the same way. + Permanent(String), + /// Fails the Task with a reason of its own, such as `key lost`, instead of `permanent error`. + Failed { reason: &'static str, details: String }, + /// The attempt was cancelled before its end, such as when it ran out of time. + Cancelled, +} diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs new file mode 100644 index 000000000..25cf3a8f0 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -0,0 +1,452 @@ +//! AI analysis (`recording-ai-analysis`): describes what the user did in one recorded session, and adds the result to +//! the session as a new log. +//! +//! [`start`] checks the request and records the Task through the [`ProvisionerTaskRunner`](super::runner::ProvisionerTaskRunner). Each attempt runs the analysis of the AI +//! crate, [`AiClient::analyze_recording`](devolutions_gateway_ai::AiClient::analyze_recording), which keeps its working +//! files next to the session manifest so that a retry resumes where the last attempt stopped, then adds the log to the +//! session. The API key and the working files go away when the Task ends. + +#[cfg(test)] +mod tests; + +use std::collections::HashMap; +use std::sync::{Arc, LazyLock, Weak}; + +use anyhow::Context as _; +use async_trait::async_trait; +use camino::{Utf8Path, Utf8PathBuf}; +use devolutions_gateway_ai::Usage; +use devolutions_gateway_ai::recording_analysis::{ + self, Progress, RecordingAnalysisEvent, RecordingFile, RecordingManifest, Recordings, RecordingsError, +}; +use provisioner_task::{CreateOutcome, ProvisionerTaskRecord}; +use secrecy::SecretString; +use serde_json::json; +use time::OffsetDateTime; +use uuid::Uuid; + +use super::AttemptError; +use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; +use super::runner::{Attempt, ProvisionerTask}; +use crate::DgwState; +use crate::artifacts::ArtifactKind; +use crate::config::ConfHandle; +use crate::provisioning::ProvisioningStore; +use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; + +pub const KIND: &str = "recording-ai-analysis"; + +/// Time an AI analysis gets from its creation to its end, waiting and retries included. +pub const TASK_DEADLINE: time::Duration = time::Duration::hours(2); + +/// Time one attempt of an AI analysis gets. +pub const ATTEMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60); + +const KEY_LOST: &str = "the API key is no longer available, such as after a Gateway restart"; + +/// Tokens counted by the AI provider, as kept in Task results. +/// +/// It mirrors [`Usage`], so a change in the AI crate never changes what we store. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +struct TokenUsage { + input_tokens: u64, + output_tokens: u64, +} + +impl From for TokenUsage { + fn from(usage: Usage) -> Self { + Self { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + } + } +} + +/// What a caller asks for when it starts an AI analysis. +#[derive(Debug)] +pub struct StartRequest { + /// Chosen by the caller, so it can ask again safely. + pub task_id: Uuid, + pub session_id: Uuid, + pub settings: AiSettings, + pub api_key: SecretString, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum StartOutcome { + /// A new Task is recorded and its job is queued. + Created(ProvisionerTaskRecord), + /// The same request came in before; nothing new runs. + Existing(ProvisionerTaskRecord), + /// The Task ID is already used for something else. + IdConflict(ProvisionerTaskRecord), + /// Another AI analysis of this session is not finished yet. + TargetBusy { active_task_id: Uuid }, +} + +#[derive(Debug, thiserror::Error)] +pub enum StartError { + #[error(transparent)] + InvalidSettings(AiSettingsError), + /// The session is still recording, so there is nothing complete to analyse yet. + #[error("session is still recording")] + RecordingActive, + /// The session has no recording to analyse. + #[error("session has no recording")] + RecordingNotFound, + #[error(transparent)] + Internal(anyhow::Error), +} + +/// What an AI analysis Task records as its parameters: the AI settings without the API key and the base URL. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Params { + provider: AiProvider, + model: String, + max_output_tokens: Option, +} + +/// Result payload of a successful AI analysis. +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct Analysis { + /// Name of the new log in the session manifest, such as `ai-analysis-0.slog`. + file_name: String, + /// Model that wrote the log, as reported by the AI provider; the requested model when it reports none. + model: String, + /// Tokens the AI provider counted for the log, answers cut and asked again included. + /// + /// Absent when the provider did not report them for every request. + #[serde(skip_serializing_if = "Option::is_none")] + usage: Option, + /// Versions of the prompts that wrote the log, such as `session-actions-3`. + prompt_version: String, +} + +/// Records an AI analysis Task, whose job the [`ProvisionerTaskRunner`](super::runner::ProvisionerTaskRunner) queues. +/// +/// Only a new Task keeps the API key; asking again for a known Task changes nothing, even while its session is +/// recording again. +pub async fn start(state: &DgwState, request: StartRequest) -> Result { + let StartRequest { + task_id, + session_id, + settings, + api_key, + } = request; + + let conf = state.conf_handle.get_conf(); + + match settings.client(&conf, &api_key) { + Ok(_) => {} + Err(ClientError::Build(error)) => { + debug!(%error, "Invalid AI settings"); + return Err(StartError::InvalidSettings(AiSettingsError::from(&error))); + } + Err(ClientError::HttpClient(error)) => { + return Err(StartError::Internal( + anyhow::Error::new(error).context("build the HTTP client for the AI provider"), + )); + } + } + + let params = Params { + provider: settings.provider, + model: settings.model, + max_output_tokens: settings.max_output_tokens, + }; + let target = session_id.to_string(); + + // A repeated request gets its Task back whatever its session does now, so the ID is looked up first. + let known = state + .provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .map_err(StartError::Internal)?; + + if let Some(record) = known { + let params = serde_json::to_value(¶ms) + .context("serialize the AI analysis parameters") + .map_err(StartError::Internal)?; + let same_request = record.kind == KIND && record.target == target && record.params == params; + return Ok(if same_request { + StartOutcome::Existing(record) + } else { + StartOutcome::IdConflict(record) + }); + } + + if state.recordings.active_recordings.contains(session_id) { + return Err(StartError::RecordingActive); + } + + match state.recordings.get_finished(session_id).await { + Ok(_) => {} + Err(FinishedRecordingError::Recording) => return Err(StartError::RecordingActive), + Err(FinishedRecordingError::NotFound) => return Err(StartError::RecordingNotFound), + Err(FinishedRecordingError::Other(error)) => { + return Err(StartError::Internal(error.context("read the recording"))); + } + } + + let access = AiAccess { + api_key, + base_url: settings.base_url, + }; + + let outcome = state + .provisioner_tasks + .create::(task_id, target, ¶ms, |record| { + // The job may run as soon as it is queued, so the key must be in place first. + state + .provisioning + .insert_task_secret(record.id, &access.to_secret()?, record.deadline_at) + }) + .await + .map_err(StartError::Internal)?; + + Ok(match outcome { + CreateOutcome::Created(task) => { + info!(task.id = %task_id, session.id = %session_id, "AI analysis queued"); + StartOutcome::Created(task) + } + CreateOutcome::Existing(task) => StartOutcome::Existing(task), + CreateOutcome::IdConflict(task) => StartOutcome::IdConflict(task), + CreateOutcome::TargetBusy(active) => StartOutcome::TargetBusy { + active_task_id: active.id, + }, + }) +} + +/// The `recording-ai-analysis` provisioner task. +pub struct RecordingAiAnalysis { + conf_handle: ConfHandle, + recordings: RecordingMessageSender, + provisioning: ProvisioningStore, + attempt_timeout: std::time::Duration, +} + +impl RecordingAiAnalysis { + pub fn new(conf_handle: ConfHandle, recordings: RecordingMessageSender, provisioning: ProvisioningStore) -> Self { + Self { + conf_handle, + recordings, + provisioning, + attempt_timeout: ATTEMPT_TIMEOUT, + } + } + + fn manifest_path(&self, session_id: Uuid) -> Utf8PathBuf { + self.conf_handle + .get_conf() + .recording_path + .join(session_id.to_string()) + .join("recording.json") + } +} + +#[async_trait] +impl ProvisionerTask for RecordingAiAnalysis { + const KIND: &'static str = KIND; + + type Params = Params; + + type Output = Analysis; + + fn deadline(&self) -> time::Duration { + TASK_DEADLINE + } + + fn attempt_timeout(&self) -> std::time::Duration { + self.attempt_timeout + } + + async fn run(&self, attempt: &Attempt, params: Params) -> Result { + let session_id = Uuid::parse_str(&attempt.target).map_err(|error| AttemptError::Failed { + reason: "invalid task", + details: format!("invalid session ID: {error}"), + })?; + + let secret = self + .provisioning + .task_secret(attempt.task_id) + .map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + + let Some(secret) = secret else { + warn!(task.id = %attempt.task_id, attempts = attempt.attempts, "AI analysis lost its API key"); + return Err(AttemptError::Failed { + reason: "key lost", + details: KEY_LOST.to_owned(), + }); + }; + + let access = AiAccess::from_secret(&secret).map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + + let settings = AiSettings { + provider: params.provider, + model: params.model, + base_url: access.base_url, + max_output_tokens: params.max_output_tokens, + }; + + self.analyze(attempt, session_id, settings, access.api_key).await + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) { + self.provisioning.remove_task_secret(record.id); + + let Ok(session_id) = Uuid::parse_str(&record.target) else { + return; + }; + + let manifest_path = self.manifest_path(session_id); + let _session = lock_session(session_id).await; + + if let Err(error) = recording_analysis::discard(&manifest_path).await { + warn!(task.id = %record.id, session.id = %session_id, %error, "Failed to remove the AI analysis working files"); + } + } +} + +impl RecordingAiAnalysis { + async fn analyze( + &self, + attempt: &Attempt, + session_id: Uuid, + settings: AiSettings, + api_key: SecretString, + ) -> Result { + let conf = self.conf_handle.get_conf(); + + let client = settings + .client(&conf, &api_key) + .map_err(|error| AttemptError::Permanent(error.to_string()))?; + + let _session = lock_session(session_id).await; + + let recording = self + .recordings + .get_finished(session_id) + .await + .map_err(|error| AttemptError::Permanent(format!("{error:#}")))?; + + let manifest = + analysis_manifest(&recording.manifest).map_err(|error| AttemptError::Permanent(error.to_string()))?; + let manifest_path = recording.manifest_path(); + + let mut request = client + .analyze_recording(&manifest, &manifest_path) + .cancellation(attempt.cancel.clone()); + + if let Some(max_output_tokens) = settings.max_output_tokens { + request = request.max_output_tokens(max_output_tokens); + } + + let mut events = request.start(); + let mut recorded: Option<(&'static str, u8)> = None; + + let analysis = loop { + match events.recv().await { + Some(RecordingAnalysisEvent::Progress(progress)) => { + // Reading sends an event every half second: the Task is only updated when the numbers it shows change. + let shown = (progress.stage(), progress.percent()); + if recorded != Some(shown) { + recorded = Some(shown); + attempt.report_progress(progress_payload(&progress)).await; + } + } + Some(RecordingAnalysisEvent::Complete(analysis)) => break analysis, + Some(RecordingAnalysisEvent::Failure(error)) => return Err(error.into()), + Some(_) => {} + None => { + return Err(AttemptError::Permanent( + "the AI analysis stopped unexpectedly".to_owned(), + )); + } + } + }; + + let file_name = add_log(&self.recordings, session_id, &analysis.log_path) + .await + .map_err(|error| AttemptError::Permanent(format!("failed to add the log: {error:#}")))?; + + debug!(task.id = %attempt.task_id, file = %file_name, actions = analysis.actions, "Session log added"); + + Ok(Analysis { + file_name, + model: analysis.model, + usage: analysis.usage.map(TokenUsage::from), + prompt_version: analysis.prompt_version, + }) + } +} + +/// The payload of a running AI analysis Task: its progress, for whoever asks about it. +fn progress_payload(progress: &Progress) -> serde_json::Value { + let (done, total) = progress.done_and_total().unzip(); + + json!({ + "step": progress.stage(), + "done": done, + "total": total, + "percent": progress.percent(), + "message": progress.to_string(), + }) +} + +/// Waits until no other AI analysis of the session runs. +/// +/// The working files belong to the session, and a Task that ended past its deadline may still run an attempt while +/// another Task of the session starts. +async fn lock_session(session_id: Uuid) -> tokio::sync::OwnedMutexGuard<()> { + static LOCKS: LazyLock>>>> = + LazyLock::new(Default::default); + + let lock = { + let mut locks = LOCKS.lock(); + locks.retain(|_, lock| lock.strong_count() > 0); + + match locks.get(&session_id).and_then(Weak::upgrade) { + Some(lock) => lock, + None => { + let lock = Arc::new(tokio::sync::Mutex::new(())); + locks.insert(session_id, Arc::downgrade(&lock)); + lock + } + } + }; + + lock.lock_owned().await +} + +/// What the AI crate needs from the session manifest; a session without recordings, or with both terminal and video ones, +/// cannot be analyzed. +fn analysis_manifest(manifest: &JrecManifest) -> Result { + let recordings = Recordings::from_files(manifest.files().iter().map(|file| RecordingFile { + file_name: file.file_name().to_owned(), + start_time: file.start_time(), + }))?; + + Ok(RecordingManifest { + start_time: manifest.start_time(), + duration: manifest.duration(), + recordings, + }) +} + +/// Adds the complete log the analysis left next to the manifest to the session. +async fn add_log(recordings: &RecordingMessageSender, session_id: Uuid, log_path: &Utf8Path) -> anyhow::Result { + let artifact_path = recordings.add_artifact(session_id, ArtifactKind::AiAnalysis).await?; + + // The log is in the session folder, so a rename moves it in at once and the session never holds a partial one. + tokio::fs::rename(log_path, &artifact_path) + .await + .with_context(|| format!("move the log to {artifact_path}"))?; + + Ok(artifact_path + .file_name() + .context("artifact path has no file name")? + .to_owned()) +} diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs new file mode 100644 index 000000000..3b07f8bbd --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -0,0 +1,864 @@ +use std::sync::Arc; +use std::time::Duration; + +use axum::http::StatusCode; +use camino::Utf8Path; +use devolutions_gateway_ai::session_actions::PROMPT_VERSION; +use devolutions_gateway_task::ShutdownHandle; +use job_queue::Job as _; +use parking_lot::Mutex; +use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; + +use super::*; +use crate::job_queue::{JobMessage, JobQueueReceiver, MAX_ATTEMPTS}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskRunner}; +use crate::recording::RecordingManagerTask; +use crate::{DgwState, MockHandles}; + +const API_KEY: &str = "sk-recording-ai-analysis-test-secret"; + +fn config(recording_path: &Utf8Path) -> String { + json!({ + "ProvisionerPublicKeyData": { + "Value": "mMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4vuqLOkl1pWobt6su1XO9VskgCAwevEGs6kkNjJQBwkGnPKYLmNF1E/af1yCocfVn/OnPf9e4x+lXVyZ6LMDJxFxu+axdgOq3Ld392J1iAEbfvwlyRFnEXFOJNyylqg3bY6LvnWHL/XZczVdMD9xYfq2sO9bg3xjRW4s7r9EEYOFjqVT3VFznH9iWJVtcSEKukmS/3uKoO6lGhacvu0HhjXXdgq0R8zvR4XRJ9Fcnf0f9Ypoc+i6L80NVjrRCeVOH+Ld/2fA9bocpfLarcVqG3RjS+qgOtpyCc0jWVFF4zaGQ7LUDFkEIYILkICeMMn2ll29hmZNzsJzZJ9s6NocgQIDAQAB" + }, + "Listeners": [{ "InternalUrl": "http://*:7171", "ExternalUrl": "https://*:7171" }], + "Proxy": { "Mode": "Off" }, + "RecordingPath": recording_path, + }) + .to_string() +} + +fn openai(model: &str) -> AiSettings { + AiSettings { + provider: AiProvider::OpenAi, + model: model.to_owned(), + base_url: None, + max_output_tokens: None, + } +} + +struct Harness { + state: DgwState, + job_queue_rx: JobQueueReceiver, + recording_path: Utf8PathBuf, + _keep: Box, + _dir: tempfile::TempDir, +} + +impl Harness { + async fn new() -> Self { + let dir = tempfile::tempdir().expect("temp dir"); + let recording_path = Utf8PathBuf::from_path_buf(dir.path().join("recordings")).expect("UTF-8"); + + let (state, handles) = DgwState::mock(&config(&recording_path)).await.expect("mock state"); + let MockHandles { + session_manager_rx, + recording_manager_rx, + subscriber_rx, + job_queue_rx, + traffic_audit_rx, + shutdown_handle, + } = handles; + + let recording_manager = RecordingManagerTask::new( + recording_manager_rx, + recording_path.clone(), + state.sessions.clone(), + state.job_queue_handle.clone(), + ); + let (recording_shutdown_handle, recording_shutdown_signal) = ShutdownHandle::new(); + tokio::spawn(devolutions_gateway_task::Task::run( + recording_manager, + recording_shutdown_signal, + )); + + Self { + state, + job_queue_rx, + recording_path, + _keep: Box::new(( + session_manager_rx, + subscriber_rx, + traffic_audit_rx, + shutdown_handle, + recording_shutdown_handle, + )), + _dir: dir, + } + } + + async fn start(&self, task_id: Uuid, session_id: Uuid, settings: AiSettings) -> Result { + start( + &self.state, + StartRequest { + task_id, + session_id, + settings, + api_key: SecretString::from(API_KEY), + }, + ) + .await + } + + /// Starts an AI analysis of `session_id` and returns its Task ID. + async fn start_created(&self, session_id: Uuid, settings: AiSettings) -> Uuid { + let task_id = Uuid::new_v4(); + let outcome = self.start(task_id, session_id, settings).await.expect("started"); + assert!(matches!(outcome, StartOutcome::Created(_)), "{outcome:?}"); + task_id + } + + /// The job holding the current token of the Task, as the job queue would run it. + async fn job(&self, task_id: Uuid) -> ProvisionerTaskJob { + // Read as of long ago, so reading never fails the Task in place of the job. + let job_token = self + .state + .provisioner_tasks + .store() + .get(task_id, OffsetDateTime::UNIX_EPOCH) + .await + .expect("read task") + .expect("task exists") + .job_token + .expect("job token"); + let json = json!({ "taskId": task_id, "jobToken": job_token }).to_string(); + + ProvisionerTaskJob::read_json(&json, self.state.provisioner_tasks.clone()).expect("valid job") + } + + async fn run(&self, task_id: Uuid) -> anyhow::Result<()> { + self.job(task_id).await.run().await + } + + /// Records a Task directly, as if it was created at `created_at`. + async fn create_task( + &self, + task_id: Uuid, + session_id: Uuid, + provider: &str, + created_at: OffsetDateTime, + deadline_at: OffsetDateTime, + ) { + self.state + .provisioner_tasks + .store() + .create( + NewProvisionerTaskRecord { + id: task_id, + kind: KIND.to_owned(), + target: session_id.to_string(), + params: json!({ "provider": provider, "model": "gpt-test", "maxOutputTokens": null }), + deadline_at, + job_token: Uuid::new_v4(), + }, + created_at, + ) + .await + .expect("create"); + } + + async fn task(&self, task_id: Uuid) -> ProvisionerTaskRecord { + self.state + .provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .expect("read task") + .expect("task exists") + } + + /// Working files of the AI analyses of session_id. + fn workspace(&self, session_id: Uuid) -> Utf8PathBuf { + self.recording_path.join(session_id.to_string()).join(".ai-analysis") + } + + fn access(&self, task_id: Uuid) -> Option { + self.state + .provisioning + .task_secret(task_id) + .expect("read the task secret") + .map(|secret| AiAccess::from_secret(&secret).expect("AI access")) + } + + fn has_key(&self, task_id: Uuid) -> bool { + self.access(task_id).is_some() + } + + /// Hands `task_id` an API key directly, as `start` does. + fn insert_key(&self, task_id: Uuid, base_url: Option, expires_at: OffsetDateTime) { + let access = AiAccess { + api_key: SecretString::from(API_KEY), + base_url, + }; + self.state + .provisioning + .insert_task_secret(task_id, &access.to_secret().expect("secret"), expires_at) + .expect("insert the task secret"); + } + + /// Writes a finished session with a short terminal recording. + fn write_session(&self) -> Uuid { + self.write_long_session(10) + } + + /// Writes a finished session whose terminal output is `lines` lines of about 100 bytes each. + fn write_long_session(&self, lines: usize) -> Uuid { + let session_id = Uuid::new_v4(); + let dir = self.recording_path.join(session_id.to_string()); + std::fs::create_dir_all(&dir).expect("session dir"); + + let manifest = json!({ + "sessionId": session_id, + "startTime": 1_787_255_035, + "duration": lines, + "files": [{ "fileName": "recording-0.cast", "startTime": 1_787_255_035, "duration": lines }], + }); + std::fs::write(dir.join("recording.json"), manifest.to_string()).expect("manifest"); + + let mut cast = String::from("{\"version\": 2}\n"); + for line in 0..lines { + cast.push_str(&format!("[{line}.0,\"o\",\"line {line} {}\\r\\n\"]\n", "x".repeat(80))); + } + std::fs::write(dir.join("recording-0.cast"), cast).expect("cast"); + + session_id + } +} + +/// A mock OpenAI-compatible provider: `status` picks the answer of the n-th request, 0 being a valid answer. +async fn spawn_ai_provider( + status: impl Fn(usize) -> u16 + Send + Sync + 'static, +) -> (AiSettings, Arc>>) { + spawn_ai_provider_after(Duration::ZERO, status).await +} + +/// Like [`spawn_ai_provider`], answering each request after `delay`. +async fn spawn_ai_provider_after( + delay: Duration, + status: impl Fn(usize) -> u16 + Send + Sync + 'static, +) -> (AiSettings, Arc>>) { + let requests = Arc::new(Mutex::new(Vec::new())); + let status = Arc::new(status); + + let app = axum::Router::new().fallback(axum::routing::post({ + let requests = Arc::clone(&requests); + move |body: String| { + let requests = Arc::clone(&requests); + let status = Arc::clone(&status); + async move { + tokio::time::sleep(delay).await; + + let index = { + let mut requests = requests.lock(); + requests.push(body); + requests.len() - 1 + }; + + let answer = json!({ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-test-2026-09-30", + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": format!("{{\"line\":\"L1\",\"description\":\"Step {index}\"}}") }, + "finish_reason": "stop" + }], + "usage": { "prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30 } + }); + + match status(index) { + 0 => (StatusCode::OK, axum::Json(answer)), + code => ( + StatusCode::from_u16(code).expect("valid status"), + axum::Json(json!({ "error": { "message": "scripted failure" } })), + ), + } + } + } + })); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("address"); + tokio::spawn(async move { axum::serve(listener, app).await.expect("serve") }); + + let settings = AiSettings { + provider: AiProvider::OpenAiCompatible, + model: "gpt-test".to_owned(), + base_url: Some(format!("http://{addr}/v1/").parse().expect("URL")), + max_output_tokens: None, + }; + + (settings, requests) +} + +fn first_line(path: &Utf8Path) -> String { + std::fs::read_to_string(path) + .expect("chunk") + .lines() + .next() + .expect("line") + .to_owned() +} + +#[tokio::test] +async fn start_records_the_task_keeps_the_key_and_queues_a_job() { + let mut harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = harness.write_session(); + let mut settings = openai("gpt-test"); + settings.base_url = Some("https://ai.example/v1/".parse().expect("URL")); + + let StartOutcome::Created(task) = harness.start(task_id, session_id, settings).await.expect("started") else { + panic!("task must be created"); + }; + + assert_eq!(task.id, task_id); + assert_eq!(task.kind, KIND); + assert_eq!(task.target, session_id.to_string()); + assert_eq!(task.state, ProvisionerTaskState::Queued); + assert_eq!( + task.params, + json!({ "provider": "openai", "model": "gpt-test", "maxOutputTokens": null }), + "neither the key nor the base URL is recorded" + ); + assert_eq!(task.deadline_at - task.created_at, TASK_DEADLINE); + + let access = harness.access(task_id).expect("key is kept"); + assert_eq!(access.base_url.expect("base URL").as_str(), "https://ai.example/v1/"); + + let JobMessage { job, schedule_for } = harness.job_queue_rx.try_recv().expect("job queued"); + assert_eq!(job.name(), ProvisionerTaskJob::NAME); + assert_eq!( + job.write_json().expect("job JSON"), + format!( + r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, + task.job_token.expect("job token") + ) + ); + assert!(schedule_for.is_none()); +} + +#[tokio::test] +async fn asking_again_creates_nothing_new() { + let mut harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = harness.write_session(); + harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect("started"); + harness.job_queue_rx.try_recv().expect("first job"); + + let same = harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect("same request"); + assert!( + matches!(same, StartOutcome::Existing(ref task) if task.id == task_id), + "{same:?}" + ); + + let other_params = harness + .start(task_id, session_id, openai("gpt-other")) + .await + .expect("other parameters"); + assert!(matches!(other_params, StartOutcome::IdConflict(_)), "{other_params:?}"); + + let busy = harness + .start(Uuid::new_v4(), session_id, openai("gpt-test")) + .await + .expect("other task"); + assert_eq!( + busy, + StartOutcome::TargetBusy { + active_task_id: task_id + } + ); + + assert!(harness.job_queue_rx.try_recv().is_err(), "no other job is queued"); + assert!(harness.has_key(task_id), "the first key stays"); +} + +#[tokio::test] +async fn asking_again_returns_the_task_even_while_the_session_records_again() { + let mut harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = harness.write_session(); + harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect("started"); + harness.job_queue_rx.try_recv().expect("first job"); + harness.state.recordings.active_recordings.insert(session_id); + + let same = harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect("same request"); + assert!(matches!(same, StartOutcome::Existing(_)), "{same:?}"); + + let other_params = harness + .start(task_id, session_id, openai("gpt-other")) + .await + .expect("other parameters"); + assert!(matches!(other_params, StartOutcome::IdConflict(_)), "{other_params:?}"); + + let new_task = harness.start(Uuid::new_v4(), session_id, openai("gpt-test")).await; + assert!(matches!(new_task, Err(StartError::RecordingActive)), "{new_task:?}"); +} + +#[tokio::test] +async fn refuses_a_session_without_recording() { + let harness = Harness::new().await; + let task_id = Uuid::new_v4(); + + let error = harness + .start(task_id, Uuid::new_v4(), openai("gpt-test")) + .await + .expect_err("no recording"); + + assert!(matches!(error, StartError::RecordingNotFound), "{error:?}"); + assert!( + harness + .state + .provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .expect("read") + .is_none() + ); + assert!(!harness.has_key(task_id)); +} + +#[tokio::test] +async fn refuses_a_session_that_is_still_recording() { + let harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = Uuid::new_v4(); + harness.state.recordings.active_recordings.insert(session_id); + + let error = harness + .start(task_id, session_id, openai("gpt-test")) + .await + .expect_err("still recording"); + + assert!(matches!(error, StartError::RecordingActive), "{error:?}"); + assert!( + harness + .state + .provisioner_tasks + .get(task_id, OffsetDateTime::now_utc()) + .await + .expect("read") + .is_none() + ); +} + +#[tokio::test] +async fn refuses_invalid_ai_settings() { + let harness = Harness::new().await; + let task_id = Uuid::new_v4(); + + let error = harness + .start(task_id, Uuid::new_v4(), openai(" ")) + .await + .expect_err("empty model"); + assert!( + matches!(error, StartError::InvalidSettings(AiSettingsError::MissingModel)), + "{error:?}" + ); + + let mut no_base_url = openai("gpt-test"); + no_base_url.provider = AiProvider::OpenAiCompatible; + let error = harness + .start(task_id, Uuid::new_v4(), no_base_url) + .await + .expect_err("no base URL"); + assert!( + matches!(error, StartError::InvalidSettings(AiSettingsError::MissingBaseUrl)), + "{error:?}" + ); + + assert!(!harness.has_key(task_id)); +} + +#[tokio::test] +async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { + let harness = Harness::new().await; + // About 1 MB of transcript: three chunks. + let session_id = harness.write_long_session(10_000); + let (settings, requests) = spawn_ai_provider(|index| if index == 1 { 503 } else { 0 }).await; + let task_id = harness.start_created(session_id, settings).await; + let workspace = harness.workspace(session_id); + + assert!(harness.run(task_id).await.is_err(), "a provider outage is retried"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Running); + assert_eq!( + task.payload, + json!({ + "step": "describing", + "done": 1, + "total": 3, + "percent": 33, + "message": "asking the AI: 1 of 3 chunks described", + }) + ); + assert!(harness.has_key(task_id), "the key is kept for the retry"); + let plan: Vec = + serde_json::from_slice(&std::fs::read(workspace.join("plan.json")).expect("plan")).expect("plan JSON"); + assert_eq!(plan.len(), 3, "three transcript chunks"); + assert!(workspace.join("chunk-0000.json").exists()); + assert!(!workspace.join("chunk-0001.json").exists()); + let chunk_starts = (0..3) + .map(|index| first_line(&workspace.join(format!("transcript-{index:04}.txt")))) + .collect::>(); + assert_eq!(requests.lock().len(), 2); + + harness.run(task_id).await.expect("second attempt succeeds"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert_eq!(task.attempts, 2); + assert_eq!( + task.payload, + json!({ + "fileName": "ai-analysis-0.slog", + "model": "gpt-test-2026-09-30", + "usage": { "inputTokens": 30, "outputTokens": 60 }, + "promptVersion": PROMPT_VERSION, + }), + "chunk 0 counts from its checkpoint, the failed request not at all" + ); + + let requests = requests.lock().clone(); + assert_eq!(requests.len(), 4, "only chunks 1 and 2 are sent again"); + for (request, chunk) in requests.iter().zip([0, 1, 1, 2]) { + assert!(request.contains(&chunk_starts[chunk]), "request for chunk {chunk}"); + } + assert!(requests.iter().all(|request| !request.contains(API_KEY))); + + let log = std::fs::read_to_string( + harness + .recording_path + .join(session_id.to_string()) + .join("ai-analysis-0.slog"), + ) + .expect("log"); + let descriptions = log + .lines() + .map(|line| serde_json::from_str::(line).expect("JSON")["description"].clone()) + .collect::>(); + assert_eq!( + descriptions, + ["Session started", "Step 0", "Step 2", "Step 3", "Session ended"].map(serde_json::Value::from) + ); + + assert!(!workspace.exists(), "the workspace is deleted on success"); + assert!(!harness.has_key(task_id), "the key is dropped on success"); +} + +#[tokio::test] +async fn permanent_error_fails_without_retry() { + let harness = Harness::new().await; + let session_id = harness.write_long_session(10); + let (settings, requests) = spawn_ai_provider(|_| 401).await; + let task_id = harness.start_created(session_id, settings).await; + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "permanent error"); + assert!( + task.payload["details"].as_str().expect("details").contains("401"), + "{}", + task.payload + ); + assert_eq!(task.payload["attempts"], 1); + assert_eq!(requests.lock().len(), 1); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +#[tokio::test] +async fn a_session_mixing_terminal_and_video_fails() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let manifest_path = harness + .recording_path + .join(session_id.to_string()) + .join("recording.json"); + let mut manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(&manifest_path).expect("manifest")).expect("manifest JSON"); + manifest["files"] + .as_array_mut() + .expect("files") + .push(json!({ "fileName": "recording-1.webm", "startTime": 1_787_255_035, "duration": 1 })); + std::fs::write(&manifest_path, manifest.to_string()).expect("manifest"); + let task_id = harness.start_created(session_id, openai("gpt-test")).await; + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "permanent error"); + assert_eq!(task.payload["details"], "session mixes terminal and video recordings"); +} +#[tokio::test] +async fn recording_deleted_after_start_fails() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = harness.start_created(session_id, openai("gpt-test")).await; + std::fs::remove_dir_all(harness.recording_path.join(session_id.to_string())).expect("delete the recording"); + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["details"], "session has no recording"); +} + +#[tokio::test] +async fn transient_errors_fail_the_task_once_attempts_are_exhausted() { + let harness = Harness::new().await; + let session_id = harness.write_long_session(10); + let (settings, requests) = spawn_ai_provider(|_| 503).await; + let task_id = harness.start_created(session_id, settings).await; + + for attempt in 1..MAX_ATTEMPTS { + assert!(harness.run(task_id).await.is_err(), "attempt {attempt} is retried"); + assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Running); + } + + harness.run(task_id).await.expect("the last attempt is not retried"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.attempts, MAX_ATTEMPTS); + assert_eq!(task.payload["reason"], "attempts exhausted"); + assert!( + task.payload["details"].as_str().expect("details").contains("503"), + "{}", + task.payload + ); + assert_eq!(requests.lock().len(), 5); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +#[tokio::test] +async fn slow_attempt_is_retried() { + let mut harness = Harness::new().await; + let session_id = harness.write_long_session(10); + + // A provider that never answers makes the attempt outlast its timeout. + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("address"); + tokio::spawn(async move { + while let Ok((connection, _)) = listener.accept().await { + tokio::spawn(async move { + let _connection = connection; + std::future::pending::<()>().await; + }); + } + }); + let settings = AiSettings { + provider: AiProvider::OpenAiCompatible, + model: "gpt-test".to_owned(), + base_url: Some(format!("http://{addr}/v1/").parse().expect("URL")), + max_output_tokens: None, + }; + let task_id = harness.start_created(session_id, settings).await; + + let mut task = RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.state.provisioning.clone(), + ); + task.attempt_timeout = Duration::from_millis(500); + harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( + Arc::clone(harness.state.provisioner_tasks.store()), + harness.state.job_queue_handle.clone(), + ) + .register(task) + .build(); + + let error = harness.run(task_id).await.expect_err("timeout is retried"); + + assert!(error.to_string().contains("took longer than"), "{error:#}"); + assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Running); + assert!(harness.has_key(task_id)); + assert!( + harness.workspace(session_id).join("plan.json").exists(), + "the chunks are kept for the retry" + ); +} + +#[tokio::test] +async fn log_added_after_the_deadline_leaves_the_task_failed() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider_after(Duration::from_secs(4), |_| 0).await; + let task_id = Uuid::new_v4(); + let now = OffsetDateTime::now_utc(); + harness + .create_task( + task_id, + session_id, + "openai-compatible", + now, + now + time::Duration::seconds(2), + ) + .await; + harness.insert_key(task_id, settings.base_url, now + time::Duration::hours(1)); + + let mut job = harness.job(task_id).await; + let job = tokio::spawn(async move { job.run().await }); + + // Reading the Task past its deadline, as a status request does, ends it while the AI is still answering. + tokio::time::sleep(Duration::from_secs(3)).await; + assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Failed); + + job.await.expect("job task").expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "timed out"); + assert!( + harness + .recording_path + .join(session_id.to_string()) + .join("ai-analysis-0.slog") + .exists(), + "the log stays in the session" + ); + assert!(!harness.workspace(session_id).exists()); + assert!(!harness.has_key(task_id)); +} + +#[tokio::test] +async fn lost_key_fails_the_task() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = harness.start_created(session_id, openai("gpt-test")).await; + std::fs::create_dir_all(harness.workspace(session_id)).expect("workspace"); + + // A restart keeps `gateway.db` and the job queue, but not the keys held in memory. + harness.state.provisioning.remove_task_secret(task_id); + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!( + task.payload, + json!({ "reason": "key lost", "details": KEY_LOST, "attempts": 1 }) + ); + assert!(!harness.workspace(session_id).exists()); +} + +#[tokio::test] +async fn overdue_task_does_not_run() { + let harness = Harness::new().await; + let task_id = Uuid::new_v4(); + let session_id = Uuid::new_v4(); + let now = OffsetDateTime::now_utc(); + harness + .create_task( + task_id, + session_id, + "openai", + now - time::Duration::hours(2), + now - time::Duration::seconds(1), + ) + .await; + harness.insert_key(task_id, None, now + time::Duration::hours(1)); + std::fs::create_dir_all(harness.workspace(session_id)).expect("workspace"); + + harness.run(task_id).await.expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "timed out"); + assert_eq!(task.attempts, 0); + assert!(!harness.has_key(task_id)); + assert!(!harness.workspace(session_id).exists()); +} + +#[tokio::test] +async fn finished_task_does_not_run_again() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 401).await; + let task_id = harness.start_created(session_id, settings).await; + + harness.run(task_id).await.expect("first run"); + std::fs::create_dir_all(harness.workspace(session_id)).expect("workspace left behind"); + harness.run(task_id).await.expect("second run"); + + assert_eq!(harness.task(task_id).await.attempts, 1); + assert!( + !harness.workspace(session_id).exists(), + "a job that finds its Task finished deletes the workspace" + ); +} + +#[tokio::test] +async fn analyses_of_one_session_run_one_at_a_time() { + let session_id = Uuid::new_v4(); + let first = lock_session(session_id).await; + + let other_session = tokio::time::timeout(Duration::from_secs(1), lock_session(Uuid::new_v4())).await; + assert!(other_session.is_ok(), "another session is not blocked"); + + let second = tokio::spawn(lock_session(session_id)); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(!second.is_finished(), "the second analysis waits"); + + drop(first); + tokio::time::timeout(Duration::from_secs(1), second) + .await + .expect("the second analysis goes on") + .expect("lock task"); +} + +#[tokio::test] +async fn working_files_left_by_an_earlier_task_go_with_the_task_that_ends() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 401).await; + let task_id = harness.start_created(session_id, settings).await; + std::fs::create_dir_all(harness.workspace(session_id)).expect("working files of a crashed analysis"); + let log = harness + .recording_path + .join(session_id.to_string()) + .join(".ai-analysis.slog"); + std::fs::write(&log, "never moved").expect("log of a crashed analysis"); + + harness.run(task_id).await.expect("no retry"); + + assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Failed); + assert!(!harness.workspace(session_id).exists()); + assert!(!log.exists()); +} +#[test] +fn unknown_usage_is_left_out_of_the_result() { + let analysis = Analysis { + file_name: "ai-analysis-0.slog".to_owned(), + model: "gpt-test-2026-09-30".to_owned(), + usage: None, + prompt_version: PROMPT_VERSION.to_owned(), + }; + + assert_eq!( + json!(analysis), + json!({ + "fileName": "ai-analysis-0.slog", + "model": "gpt-test-2026-09-30", + "promptVersion": PROMPT_VERSION, + }) + ); +} diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs new file mode 100644 index 000000000..308ada5e6 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -0,0 +1,483 @@ +//! Runs provisioner tasks: the one place that keeps a task's record and the job running it in step. +//! +//! Each kind of work is a [`ProvisionerTask`]: what one attempt does, and what to clean up once it has ended. +//! [`ProvisionerTaskRunner`] does the rest for every task. It creates the record, queues the job, starts each attempt, +//! retries transient failures, enforces the attempt timeout and records how the task ended. +//! +//! The record is the source of truth and the job only wakes it up. Each record names the one job allowed to run it, by +//! its token, so a stale or duplicate job stops without running anything. At startup, before any job runs, +//! [`ProvisionerTaskRunner::resume`] hands every unfinished record to a new job, which also covers a record created just +//! before a crash, before its job was queued. + +#[cfg(test)] +mod tests; + +use std::collections::HashMap; +use std::future::Future; +use std::sync::Arc; + +use anyhow::Context as _; +use async_trait::async_trait; +use provisioner_task::{ + AttemptStart, CreateOutcome, DynProvisionerTaskStore, NewProvisionerTaskRecord, ProvisionerTaskRecord, +}; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::json; +use time::OffsetDateTime; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use super::AttemptError; +use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; + +/// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. +/// +/// One value serves every record of its kind, so it keeps only what they share; what belongs to one record comes in +/// through [`Attempt`] and the parameters. +#[async_trait] +pub trait ProvisionerTask: Send + Sync + 'static { + /// The kind written in the records of this task. + const KIND: &'static str; + + /// What the record keeps as asked; never secrets, since records are kept and returned to callers. + type Params: Serialize + DeserializeOwned + Send; + + /// The payload of a `succeeded` record. + type Output: Serialize + Send; + + /// Time from the creation of a record to its end, waiting and retries included. + fn deadline(&self) -> time::Duration; + + /// Time one attempt gets before it is cancelled and retried. + fn attempt_timeout(&self) -> std::time::Duration; + + /// Runs one attempt. + /// + /// The attempt should stop soon after [`Attempt::cancel`] is cancelled, and return [`AttemptError::Cancelled`]. + async fn run(&self, attempt: &Attempt, params: Self::Params) -> Result; + + /// Drops what is kept only while the record runs, once it has ended, however it ended. + /// + /// It may be called again for a record that already ended, so it must be idempotent. + async fn on_end(&self, _record: &ProvisionerTaskRecord) {} +} + +/// One attempt of a record, as its task sees it. +pub struct Attempt { + pub task_id: Uuid, + pub target: String, + /// Attempts started so far, this one included. + pub attempts: u32, + pub cancel: CancellationToken, + store: DynProvisionerTaskStore, +} + +impl Attempt { + /// Replaces the payload of the running record, so whoever asks about it sees how the attempt goes. + pub async fn report_progress(&self, payload: serde_json::Value) { + if let Err(error) = self.store.update_running(self.task_id, payload).await { + warn!(task.id = %self.task_id, error = format!("{error:#}"), "Failed to record the Task progress"); + } + } +} + +/// [`ProvisionerTask`] with its types erased, so the runner can hold every task together. +#[async_trait] +trait AnyProvisionerTask: Send + Sync { + fn deadline(&self) -> time::Duration; + + fn attempt_timeout(&self) -> std::time::Duration; + + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result; + + async fn on_end(&self, record: &ProvisionerTaskRecord); +} + +#[async_trait] +impl AnyProvisionerTask for T { + fn deadline(&self) -> time::Duration { + ProvisionerTask::deadline(self) + } + + fn attempt_timeout(&self) -> std::time::Duration { + ProvisionerTask::attempt_timeout(self) + } + + async fn run(&self, attempt: &Attempt, params: serde_json::Value) -> Result { + let params = serde_json::from_value(params).map_err(|error| AttemptError::Failed { + reason: "invalid task", + details: format!("invalid parameters: {error}"), + })?; + + let output = ProvisionerTask::run(self, attempt, params).await?; + + serde_json::to_value(output).map_err(|error| AttemptError::Permanent(format!("invalid result: {error}"))) + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) { + ProvisionerTask::on_end(self, record).await; + } +} + +/// Creates provisioner task records and runs them as jobs on the job queue. +#[derive(Clone)] +pub struct ProvisionerTaskRunner { + inner: Arc, +} + +struct Inner { + store: DynProvisionerTaskStore, + job_queue: JobQueueHandle, + /// Kind → the task that runs records of this kind. + tasks: HashMap<&'static str, Arc>, +} + +pub struct ProvisionerTaskRunnerBuilder { + store: DynProvisionerTaskStore, + job_queue: JobQueueHandle, + tasks: HashMap<&'static str, Arc>, +} + +impl ProvisionerTaskRunnerBuilder { + /// Registers `task`, so records of its kind can be created and run. + #[must_use] + pub fn register(mut self, task: T) -> Self { + self.tasks.insert(T::KIND, Arc::new(task)); + self + } + + pub fn build(self) -> ProvisionerTaskRunner { + ProvisionerTaskRunner { + inner: Arc::new(Inner { + store: self.store, + job_queue: self.job_queue, + tasks: self.tasks, + }), + } + } +} + +/// Initial payload of a running record, until its task reports progress. +fn starting_payload() -> serde_json::Value { + json!({ "step": "preparing" }) +} + +impl ProvisionerTaskRunner { + pub fn builder(store: DynProvisionerTaskStore, job_queue: JobQueueHandle) -> ProvisionerTaskRunnerBuilder { + ProvisionerTaskRunnerBuilder { + store, + job_queue, + tasks: HashMap::new(), + } + } + + /// Returns the record with this ID; one past its deadline is failed first. + pub async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result> { + self.inner.store.get(id, now).await + } + + /// Creates a record for the task `T` and queues the job that runs it. + /// + /// `on_created` runs only for a new record, before its job is queued, so what the job needs is in place before it + /// can run; if it fails, the record fails and no job is queued. A known ID or a busy target creates nothing and + /// queues nothing. + pub async fn create( + &self, + id: Uuid, + target: String, + params: &T::Params, + on_created: impl FnOnce(&ProvisionerTaskRecord) -> anyhow::Result<()> + Send, + ) -> anyhow::Result { + let task = self.task(T::KIND).context("unregistered provisioner task")?; + let params = serde_json::to_value(params).context("serialize the task parameters")?; + let now = OffsetDateTime::now_utc(); + let job_token = Uuid::new_v4(); + + let outcome = self + .inner + .store + .create( + NewProvisionerTaskRecord { + id, + kind: T::KIND.to_owned(), + target, + params, + deadline_at: now + task.deadline(), + job_token, + }, + now, + ) + .await?; + + let CreateOutcome::Created(record) = &outcome else { + return Ok(outcome); + }; + + if let Err(error) = on_created(record) { + error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to prepare the Task"); + self.fail(task.as_ref(), id, failure("not queued", &format!("{error:#}"), 0)) + .await; + return Err(error); + } + + if let Err(error) = self.enqueue(id, job_token).await { + error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to queue the Task"); + self.fail(task.as_ref(), id, failure("not queued", &format!("{error:#}"), 0)) + .await; + return Err(error); + } + + info!(task.id = %id, task.kind = T::KIND, "Task queued"); + + Ok(outcome) + } + + /// Hands every unfinished record to a new job queued with `push`: the jobs queued before stop at their next attempt. + /// + /// Call it at startup, before the job runner claims any job, so no earlier job is running an attempt meanwhile. + pub async fn resume(&self, push: F) -> anyhow::Result<()> + where + F: Fn(ProvisionerTaskJob) -> Fut, + Fut: Future>, + { + let now = OffsetDateTime::now_utc(); + let records = self.inner.store.list_unfinished(now).await?; + + for record in records { + let job_token = Uuid::new_v4(); + + if !self.inner.store.replace_job_token(record.id, job_token, now).await? { + continue; + } + + let job = ProvisionerTaskJob { + task_id: record.id, + job_token, + runner: self.clone(), + }; + + match push(job).await { + Ok(()) => info!(task.id = %record.id, task.kind = %record.kind, "Task queued again"), + Err(error) => error!( + task.id = %record.id, + task.kind = %record.kind, + error = format!("{error:#}"), + "Failed to queue the Task again; it fails at its deadline" + ), + } + } + + Ok(()) + } + + #[cfg(test)] + pub(crate) fn store(&self) -> &DynProvisionerTaskStore { + &self.inner.store + } + + /// The task that runs records of `kind`. + fn task(&self, kind: &str) -> Option> { + self.inner.tasks.get(kind).cloned() + } + + async fn enqueue(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { + self.inner + .job_queue + .enqueue(ProvisionerTaskJob { + task_id, + job_token, + runner: self.clone(), + }) + .await + } + + /// Runs one attempt for the job holding `job_token`; returns an error only to ask the job queue for a retry. + async fn run_attempt(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { + let store = &self.inner.store; + let now = OffsetDateTime::now_utc(); + + let record = match store.start_attempt(task_id, job_token, starting_payload(), now).await? { + AttemptStart::Started(record) => record, + AttemptStart::Finished(record) => { + debug!(task.id = %task_id, "Task is over, so this attempt does not run"); + if let Some(task) = self.task(&record.kind) { + task.on_end(&record).await; + } + return Ok(()); + } + AttemptStart::Superseded => { + debug!(task.id = %task_id, "Another job runs this Task, so this one stops"); + return Ok(()); + } + AttemptStart::Unknown => { + debug!(task.id = %task_id, "Task is unknown, so this attempt does not run"); + return Ok(()); + } + }; + + let Some(task) = self.task(&record.kind) else { + error!(task.id = %task_id, task.kind = %record.kind, "Task of an unknown kind"); + let details = format!("unknown task kind {}", record.kind); + if let Err(error) = store + .fail(task_id, failure("invalid task", &details, record.attempts), now) + .await + { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure"); + } + return Ok(()); + }; + + let attempts = record.attempts; + let attempt_timeout = task.attempt_timeout(); + let cancel = CancellationToken::new(); + let attempt = Attempt { + task_id, + target: record.target.clone(), + attempts, + cancel: cancel.clone(), + store: Arc::clone(store), + }; + + info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task attempt started"); + + // The attempt has its own Tokio task, so a panic fails the record instead of leaving it running. + let mut handle = tokio::spawn({ + let task = Arc::clone(&task); + let params = record.params.clone(); + async move { task.run(&attempt, params).await } + }); + + let joined = match tokio::time::timeout(attempt_timeout, &mut handle).await { + Ok(joined) => joined, + Err(_) => { + // Aborting cannot stop blocking steps, so the attempt is cancelled and awaited instead: the next + // attempt never runs alongside it. + cancel.cancel(); + handle.await + } + }; + + let outcome = match joined { + Ok(Err(AttemptError::Cancelled)) => Err(AttemptError::Transient(format!( + "the attempt took longer than {attempt_timeout:?}" + ))), + Ok(outcome) => outcome, + Err(_) => Err(AttemptError::Permanent("the attempt panicked".to_owned())), + }; + + match outcome { + Ok(output) => { + // An error here must not ask for a retry: the work is done. + match store.succeed(task_id, output, OffsetDateTime::now_utc()).await { + Ok(true) => info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"), + Ok(false) => warn!( + task.id = %task_id, + task.kind = %record.kind, + "Task attempt succeeded after the Task had ended, such as past its deadline; the Task keeps its earlier end" + ), + Err(error) => error!( + task.id = %task_id, + error = format!("{error:#}"), + "Failed to record the Task result" + ), + } + + self.end(task.as_ref(), task_id).await; + } + Err(AttemptError::Transient(details)) if attempts < MAX_ATTEMPTS => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task attempt failed; it will be retried"); + anyhow::bail!("Task attempt failed: {details}"); + } + Err(AttemptError::Transient(details)) => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task failed after its last attempt"); + self.fail( + task.as_ref(), + task_id, + failure("attempts exhausted", &details, attempts), + ) + .await; + } + Err(AttemptError::Permanent(details)) => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task failed"); + self.fail(task.as_ref(), task_id, failure("permanent error", &details, attempts)) + .await; + } + Err(AttemptError::Failed { reason, details }) => { + warn!(task.id = %task_id, task.kind = %record.kind, attempts, reason, %details, "Task failed"); + self.fail(task.as_ref(), task_id, failure(reason, &details, attempts)) + .await; + } + Err(AttemptError::Cancelled) => unreachable!("a cancelled attempt is reported as a timeout above"), + } + + Ok(()) + } + + async fn fail(&self, task: &dyn AnyProvisionerTask, task_id: Uuid, payload: serde_json::Value) { + if let Err(error) = self.inner.store.fail(task_id, payload, OffsetDateTime::now_utc()).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure"); + } + + self.end(task, task_id).await; + } + + async fn end(&self, task: &dyn AnyProvisionerTask, task_id: Uuid) { + match self.inner.store.get(task_id, OffsetDateTime::now_utc()).await { + Ok(Some(record)) => task.on_end(&record).await, + Ok(None) => {} + Err(error) => error!(task.id = %task_id, error = format!("{error:#}"), "Failed to read the ended Task"), + } + } +} + +fn failure(reason: &str, details: &str, attempts: u32) -> serde_json::Value { + json!({ "reason": reason, "details": details, "attempts": attempts }) +} + +#[derive(Serialize, serde::Deserialize)] +#[serde(rename_all = "camelCase")] +struct JobDef { + task_id: Uuid, + job_token: Uuid, +} + +/// The job of a provisioner task record: it only wakes the record up, and the record decides what runs. +pub struct ProvisionerTaskJob { + task_id: Uuid, + job_token: Uuid, + runner: ProvisionerTaskRunner, +} + +impl ProvisionerTaskJob { + pub const NAME: &'static str = "provisioner-task"; + + pub fn read_json(json: &str, runner: ProvisionerTaskRunner) -> anyhow::Result { + let def: JobDef = serde_json::from_str(json).context("failed to deserialize ProvisionerTaskJob")?; + + Ok(Self { + task_id: def.task_id, + job_token: def.job_token, + runner, + }) + } +} + +#[async_trait] +impl job_queue::Job for ProvisionerTaskJob { + fn name(&self) -> &str { + Self::NAME + } + + fn write_json(&self) -> anyhow::Result { + serde_json::to_string(&JobDef { + task_id: self.task_id, + job_token: self.job_token, + }) + .context("failed to serialize ProvisionerTaskJob") + } + + async fn run(&mut self) -> anyhow::Result<()> { + self.runner.run_attempt(self.task_id, self.job_token).await + } +} diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs new file mode 100644 index 000000000..a40619393 --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -0,0 +1,301 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use job_queue::Job as _; +use parking_lot::Mutex; +use provisioner_task::ProvisionerTaskState; + +use super::*; +use crate::job_queue::{JobMessage, JobQueueReceiver}; + +const KIND: &str = "test.echo"; + +#[derive(Debug, Serialize, serde::Deserialize)] +struct EchoParams { + value: u32, +} + +/// A task whose attempts return `outcome`, counting its runs and the records it ended. +struct Echo { + outcome: fn(u32) -> Result, + runs: Arc, + ended: Arc>>, +} + +#[async_trait] +impl ProvisionerTask for Echo { + const KIND: &'static str = KIND; + + type Params = EchoParams; + + type Output = serde_json::Value; + + fn deadline(&self) -> time::Duration { + time::Duration::hours(1) + } + + fn attempt_timeout(&self) -> std::time::Duration { + std::time::Duration::from_secs(5) + } + + async fn run(&self, _attempt: &Attempt, params: EchoParams) -> Result { + self.runs.fetch_add(1, Ordering::SeqCst); + (self.outcome)(params.value) + } + + async fn on_end(&self, record: &ProvisionerTaskRecord) { + self.ended.lock().push(record.id); + } +} + +struct Harness { + runner: ProvisionerTaskRunner, + job_queue_rx: JobQueueReceiver, + runs: Arc, + ended: Arc>>, +} + +impl Harness { + async fn new(outcome: fn(u32) -> Result) -> Self { + let conn = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open gateway database") + .connect() + .await + .expect("connect"); + let store = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) + .await + .expect("open task store"), + ); + let (job_queue, job_queue_rx) = JobQueueHandle::new(); + let runs = Arc::new(AtomicUsize::new(0)); + let ended = Arc::new(Mutex::new(Vec::new())); + + let runner = ProvisionerTaskRunner::builder(store, job_queue) + .register(Echo { + outcome, + runs: Arc::clone(&runs), + ended: Arc::clone(&ended), + }) + .build(); + + Self { + runner, + job_queue_rx, + runs, + ended, + } + } + + async fn create(&self, value: u32) -> Uuid { + let id = Uuid::new_v4(); + let outcome = self + .runner + .create::(id, format!("target-{id}"), &EchoParams { value }, |_| Ok(())) + .await + .expect("create"); + assert!(matches!(outcome, CreateOutcome::Created(_)), "{outcome:?}"); + id + } + + fn queued_job(&mut self) -> Box { + let JobMessage { job, .. } = self.job_queue_rx.try_recv().expect("a queued job"); + job + } + + async fn task(&self, id: Uuid) -> ProvisionerTaskRecord { + self.runner + .get(id, OffsetDateTime::now_utc()) + .await + .expect("read task") + .expect("task exists") + } +} + +#[tokio::test] +async fn only_a_new_task_runs_on_created_and_gets_a_job() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = Uuid::new_v4(); + let created = AtomicUsize::new(0); + let create = || { + harness + .runner + .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + created.fetch_add(1, Ordering::SeqCst); + Ok(()) + }) + }; + + assert!(matches!(create().await.expect("create"), CreateOutcome::Created(_))); + assert!(matches!(create().await.expect("again"), CreateOutcome::Existing(_))); + + assert_eq!(created.load(Ordering::SeqCst), 1); + harness.queued_job(); + assert!(harness.job_queue_rx.try_recv().is_err(), "one job only"); +} + +#[tokio::test] +async fn a_task_whose_job_cannot_be_queued_fails() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + harness.job_queue_rx.close(); + let id = Uuid::new_v4(); + + harness + .runner + .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| Ok(())) + .await + .expect_err("the job queue is gone"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "not queued"); + assert_eq!(*harness.ended.lock(), [id], "what was set up for the Task is dropped"); +} + +#[tokio::test] +async fn a_task_that_cannot_be_prepared_fails_without_a_job() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = Uuid::new_v4(); + + harness + .runner + .create::(id, "target".to_owned(), &EchoParams { value: 1 }, |_| { + anyhow::bail!("no room for the secret") + }) + .await + .expect_err("preparing fails"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!(task.payload["reason"], "not queued"); + assert!(harness.job_queue_rx.try_recv().is_err(), "no job is queued"); + assert_eq!(*harness.ended.lock(), [id]); +} + +#[tokio::test] +async fn a_superseded_job_runs_nothing_once_the_task_is_resumed() { + let mut harness = Harness::new(|value| Ok(json!({ "echo": value }))).await; + let id = harness.create(7).await; + let mut stale = harness.queued_job(); + + let resumed = Mutex::new(Vec::new()); + harness + .runner + .resume(|job| { + resumed.lock().push(job); + async { Ok(()) } + }) + .await + .expect("resume"); + let mut resumed = resumed.into_inner(); + assert_eq!(resumed.len(), 1, "the unfinished Task gets a new job"); + + stale.run().await.expect("the stale job stops"); + assert_eq!(harness.runs.load(Ordering::SeqCst), 0); + assert_eq!(harness.task(id).await.attempts, 0, "a stale job counts no attempt"); + assert!(harness.ended.lock().is_empty(), "a stale job ends nothing"); + + resumed[0].run().await.expect("the new job runs"); + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded); + assert_eq!(task.payload, json!({ "echo": 7 })); + assert_eq!(*harness.ended.lock(), [id]); + + let queued_again = AtomicUsize::new(0); + harness + .runner + .resume(|_| { + queued_again.fetch_add(1, Ordering::SeqCst); + async { Ok(()) } + }) + .await + .expect("resume again"); + assert_eq!( + queued_again.load(Ordering::SeqCst), + 0, + "a finished Task gets no new job" + ); +} + +#[tokio::test] +async fn a_panicking_attempt_fails_the_task() { + let mut harness = Harness::new(|_| panic!("the attempt panics")).await; + let id = harness.create(1).await; + + harness.queued_job().run().await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed); + assert_eq!( + task.payload, + json!({ "reason": "permanent error", "details": "the attempt panicked", "attempts": 1 }) + ); + assert_eq!(*harness.ended.lock(), [id]); +} + +#[tokio::test] +async fn a_failure_keeps_the_reason_of_its_kind() { + let mut harness = Harness::new(|_| { + Err(AttemptError::Failed { + reason: "key lost", + details: "gone".to_owned(), + }) + }) + .await; + let id = harness.create(1).await; + + harness.queued_job().run().await.expect("no retry"); + + assert_eq!( + harness.task(id).await.payload, + json!({ "reason": "key lost", "details": "gone", "attempts": 1 }) + ); +} + +#[tokio::test] +async fn invalid_parameters_and_unknown_kinds_fail_the_task() { + let harness = Harness::new(|value| Ok(json!(value))).await; + let store = harness.runner.store(); + let now = OffsetDateTime::now_utc(); + + for (kind, params) in [(KIND, json!({ "value": "not a number" })), ("test.unknown", json!({}))] { + let id = Uuid::new_v4(); + let job_token = Uuid::new_v4(); + store + .create( + NewProvisionerTaskRecord { + id, + kind: kind.to_owned(), + target: id.to_string(), + params, + deadline_at: now + time::Duration::hours(1), + job_token, + }, + now, + ) + .await + .expect("create"); + + harness.runner.run_attempt(id, job_token).await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed, "{kind}"); + assert_eq!(task.payload["reason"], "invalid task", "{kind}"); + } + + assert_eq!(harness.runs.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn the_job_is_read_back_from_its_json() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + let job = harness.queued_job(); + + let json = job.write_json().expect("job JSON"); + let read = ProvisionerTaskJob::read_json(&json, harness.runner.clone()).expect("valid job"); + + assert_eq!(read.task_id, id); + assert_eq!(Some(read.job_token), harness.task(id).await.job_token); + assert_eq!(job.name(), ProvisionerTaskJob::NAME); +} diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 8ece9b9b6..5c5814230 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -30,15 +30,26 @@ const BUFFER_WRITER_SIZE: usize = 64 * 1024; #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -struct JrecFile { +pub(crate) struct JrecFile { file_name: String, start_time: i64, duration: i64, } +impl JrecFile { + pub(crate) fn file_name(&self) -> &str { + &self.file_name + } + + /// Unix seconds. + pub(crate) fn start_time(&self) -> i64 { + self.start_time + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -struct JrecManifest { +pub(crate) struct JrecManifest { session_id: Uuid, start_time: i64, duration: i64, @@ -48,6 +59,20 @@ struct JrecManifest { } impl JrecManifest { + /// Unix seconds. + pub(crate) fn start_time(&self) -> i64 { + self.start_time + } + + /// Seconds. + pub(crate) fn duration(&self) -> i64 { + self.duration + } + + pub(crate) fn files(&self) -> &[JrecFile] { + &self.files + } + fn read_from_file(path: impl AsRef) -> anyhow::Result { let json = std::fs::read(path)?; let manifest = serde_json::from_slice(&json)?; @@ -316,7 +341,7 @@ impl ActiveRecordings { self.0.lock().clone() } - fn insert(&self, id: Uuid) -> usize { + pub(crate) fn insert(&self, id: Uuid) -> usize { let mut guard = self.0.lock(); guard.insert(id); guard.len() @@ -357,6 +382,10 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, + GetFinished { + id: Uuid, + channel: oneshot::Sender>, + }, Disconnect { id: Uuid, }, @@ -377,6 +406,29 @@ enum RecordingManagerMessage { }, } +/// A session that is not recording anymore and has at least one Recording. +#[derive(Debug)] +pub(crate) struct FinishedRecording { + pub(crate) dir: Utf8PathBuf, + pub(crate) manifest: JrecManifest, +} + +impl FinishedRecording { + pub(crate) fn manifest_path(&self) -> Utf8PathBuf { + self.dir.join("recording.json") + } +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum FinishedRecordingError { + #[error("session has no recording")] + NotFound, + #[error("session is still recording")] + Recording, + #[error(transparent)] + Other(anyhow::Error), +} + impl fmt::Debug for RecordingManagerMessage { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { @@ -396,6 +448,9 @@ impl fmt::Debug for RecordingManagerMessage { .field("id", id) .field("kind", kind) .finish_non_exhaustive(), + RecordingManagerMessage::GetFinished { id, channel: _ } => { + f.debug_struct("GetFinished").field("id", id).finish_non_exhaustive() + } RecordingManagerMessage::Disconnect { id } => f.debug_struct("Disconnect").field("id", id).finish(), RecordingManagerMessage::GetState { id, channel: _ } => { f.debug_struct("GetState").field("id", id).finish_non_exhaustive() @@ -459,6 +514,36 @@ impl RecordingMessageSender { rx.await.context("couldn't receive AddArtifact result")? } + pub(crate) async fn get_finished(&self, id: Uuid) -> Result { + let (tx, rx) = oneshot::channel(); + self.channel + .send(RecordingManagerMessage::GetFinished { id, channel: tx }) + .await + .ok() + .context("couldn't send GetFinished message") + .map_err(FinishedRecordingError::Other)?; + let dir = rx + .await + .context("couldn't receive GetFinished result") + .map_err(FinishedRecordingError::Other)??; + + // The manifest is read here, so the recording manager task never waits on the disk for it. + let manifest = match fs::read(dir.join("recording.json")).await { + Ok(manifest) => manifest, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Err(FinishedRecordingError::NotFound), + Err(error) => { + return Err(FinishedRecordingError::Other( + anyhow::Error::new(error).context("read manifest from disk"), + )); + } + }; + let manifest = serde_json::from_slice(&manifest) + .context("parse manifest") + .map_err(FinishedRecordingError::Other)?; + + Ok(FinishedRecording { dir, manifest }) + } + async fn disconnect(&self, id: Uuid) -> anyhow::Result<()> { self.channel .send(RecordingManagerMessage::Disconnect { id }) @@ -819,6 +904,15 @@ impl RecordingManagerTask { Ok(artifact_path) } + fn handle_get_finished(&self, id: Uuid) -> Result { + // The transcript is built from the whole recording, so one still being pushed is not ready. + if self.ongoing_recordings.contains_key(&id) { + return Err(FinishedRecordingError::Recording); + } + + Ok(self.recordings_path.join(id.to_string())) + } + fn handle_remove(&mut self, id: Uuid) { if let Some(ongoing) = self.ongoing_recordings.get(&id) { match ongoing.state { @@ -937,6 +1031,9 @@ async fn recording_manager_task( RecordingManagerMessage::AddArtifact { id, kind, channel } => { let _ = channel.send(manager.handle_add_artifact(id, kind).await); }, + RecordingManagerMessage::GetFinished { id, channel } => { + let _ = channel.send(manager.handle_get_finished(id)); + } RecordingManagerMessage::Disconnect { id } => { if let Err(e) = manager.handle_disconnect(id).await { error!(error = format!("{e:#}"), "handle_disconnect"); @@ -1421,6 +1518,25 @@ mod tests { assert_eq!(harness.sender.get_count().await.expect("count"), 1); } + #[tokio::test] + async fn finished_recordings_are_read_through_the_manager() { + let harness = Harness::start(); + let id = Uuid::new_v4(); + harness.write_manifest(id, MASTER_MANIFEST); + + let finished = harness.sender.get_finished(id).await.expect("finished"); + assert_eq!(finished.dir, harness.recordings_path.join(id.to_string())); + assert_eq!(finished.manifest.files()[0].file_name(), "recording-0.webm"); + + let ongoing = Uuid::new_v4(); + harness.connect(ongoing, WEBM).await; + let result = harness.sender.get_finished(ongoing).await; + assert!(matches!(result, Err(FinishedRecordingError::Recording)), "{result:?}"); + + let missing = harness.sender.get_finished(Uuid::new_v4()).await; + assert!(matches!(missing, Err(FinishedRecordingError::NotFound)), "{missing:?}"); + } + #[test] fn manifest_without_artifacts_round_trips_byte_for_byte() { let dir = tempfile::tempdir().expect("temp dir"); diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index 32abd87a7..eea186dca 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -344,12 +344,26 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { .connect() .await .context("failed to connect the provisioner tasks to the gateway database")?; - let provisioner_tasks = Arc::new( + let task_store: provisioner_task::DynProvisionerTaskStore = Arc::new( provisioner_task_libsql::LibSqlProvisionerTaskStore::open(provisioner_tasks_conn) .await .context("failed to initialize the provisioner task store")?, ); + let provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( + task_store, + job_queue_ctx.job_queue_handle.clone(), + conf_handle.clone(), + recording_manager_handle.clone(), + provisioning.clone(), + ); + + // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. + provisioner_tasks + .resume(|job| job_queue_ctx.push(job)) + .await + .context("failed to queue the unfinished provisioner tasks again")?; + let state = DgwState { conf_handle: conf_handle.clone(), token_cache: Arc::clone(&token_cache), @@ -364,7 +378,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { monitoring_state, traffic_audit_handle: traffic_audit_task.handle(), agent_tunnel_handle, - provisioner_tasks, + provisioner_tasks: provisioner_tasks.clone(), }; for listener in &conf.listeners { @@ -429,7 +443,10 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { job_queue_ctx.job_queue_handle.clone(), )); - tasks.register(devolutions_gateway::job_queue::JobRunnerTask::new(&job_queue_ctx)); + tasks.register(devolutions_gateway::job_queue::JobRunnerTask::new( + &job_queue_ctx, + provisioner_tasks, + )); tasks.register(devolutions_gateway::job_queue::JobQueueTask::new(job_queue_ctx)); tasks.register(traffic_audit_task); From b1d91083200c132e5a9a94ae6c7d243c37cff7e7 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 6 Oct 2026 16:21:34 -0400 Subject: [PATCH 03/21] 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.rs | 249 +++++++- 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, 1034 insertions(+), 4 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.rs b/devolutions-gateway/src/api/jrec.rs index 508d32ff5..ecc3bb5b0 100644 --- a/devolutions-gateway/src/api/jrec.rs +++ b/devolutions-gateway/src/api/jrec.rs @@ -10,9 +10,9 @@ 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::routing::{delete, get}; +use axum::http::header::{CONTENT_DISPOSITION, CONTENT_TYPE, HeaderValue, LOCATION}; +use axum::response::{IntoResponse as _, Response}; +use axum::routing::{delete, get, post}; use axum::{Json, Router}; use bytes::Bytes; use cadeau::xmf; @@ -20,17 +20,24 @@ use camino::{Utf8Path, Utf8PathBuf}; use devolutions_gateway_task::ShutdownSignal; use futures::stream; use hyper::StatusCode; +use secrecy::SecretString; use tokio::io::AsyncReadExt as _; use tracing::Instrument as _; +use url::Url; use uuid::Uuid; use zip::CompressionMethod; use zip::write::SimpleFileOptions; use crate::DgwState; use crate::api::heartbeat::recording_storage_health; +use crate::api::tasks::TaskInfo; use crate::artifacts::JrecArtifacts; -use crate::extract::{JrecToken, RecordingDeleteScope, RecordingsReadScope}; +use crate::extract::{ + JrecToken, ProvisionerTaskReadScope, ProvisionerTaskStartScope, RecordingDeleteScope, RecordingsReadScope, +}; use crate::http::{HttpError, HttpErrorBuilder}; +use crate::provisioner_tasks::ai::{AiProvider, AiSettings}; +use crate::provisioner_tasks::recording_ai_analysis::{self, StartError, StartOutcome, StartRequest}; use crate::recording::{PushOutcome, RecordingMessageSender}; use crate::token::{JrecTokenClaims, RecordingFileType, RecordingOperation}; @@ -59,6 +66,8 @@ pub fn make_router(state: DgwState) -> Router { .route("/play", get(get_player)) .route("/play/{*path}", get(get_player)) .route("/shadow/{id}", get(shadow_recording)) + .route("/{session_id}/ai-analysis", post(start_ai_analysis)) + .route("/{session_id}/ai-analysis/{task_id}", get(get_ai_analysis)) .with_state(state) } @@ -426,6 +435,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, @@ -998,6 +1193,52 @@ mod tests { use super::*; + #[tokio::test] + async fn ai_analysis_of_a_session_still_recording_is_a_conflict() { + use axum::extract::connect_info::MockConnectInfo; + use base64::Engine as _; + use tower::ServiceExt as _; + + let dir = tempfile::tempdir().expect("temp dir"); + let config = serde_json::json!({ + "ProvisionerPublicKeyData": { + "Value": "mMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4vuqLOkl1pWobt6su1XO9VskgCAwevEGs6kkNjJQBwkGnPKYLmNF1E/af1yCocfVn/OnPf9e4x+lXVyZ6LMDJxFxu+axdgOq3Ld392J1iAEbfvwlyRFnEXFOJNyylqg3bY6LvnWHL/XZczVdMD9xYfq2sO9bg3xjRW4s7r9EEYOFjqVT3VFznH9iWJVtcSEKukmS/3uKoO6lGhacvu0HhjXXdgq0R8zvR4XRJ9Fcnf0f9Ypoc+i6L80NVjrRCeVOH+Ld/2fA9bocpfLarcVqG3RjS+qgOtpyCc0jWVFF4zaGQ7LUDFkEIYILkICeMMn2ll29hmZNzsJzZJ9s6NocgQIDAQAB" + }, + "Listeners": [{ "InternalUrl": "http://*:7171", "ExternalUrl": "https://*:7171" }], + "RecordingPath": dir.path().to_str().expect("UTF-8"), + "__debug__": { "disable_token_validation": true }, + }); + let (state, _handles) = DgwState::mock(&config.to_string()).await.expect("mock state"); + let session_id = Uuid::new_v4(); + state.recordings.active_recordings.insert(session_id); + let app = crate::make_http_service(state).layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 3000)))); + + let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD; + let claims = serde_json::json!({ "jti": Uuid::new_v4(), "exp": i64::from(u32::MAX), "scope": "gateway.tasks.recording-ai-analysis.start" }); + let token = format!( + "{}.{}.{}", + engine.encode(r#"{"alg":"RS256","cty":"SCOPE"}"#), + engine.encode(claims.to_string()), + engine.encode(b"signature"), + ); + let body = serde_json::json!({ "taskId": Uuid::new_v4(), "provider": "openai", "model": "gpt-test", "apiKey": "sk-test" }); + let request = axum::http::Request::builder() + .method("POST") + .uri(format!("/jet/jrec/{session_id}/ai-analysis")) + .header(axum::http::header::AUTHORIZATION, format!("Bearer {token}")) + .body(Body::from(body.to_string())) + .expect("request"); + + let response = app.oneshot(request).await.expect("response"); + + assert_eq!(response.status(), StatusCode::CONFLICT); + let body = response.into_body().collect().await.expect("body").to_bytes(); + assert_eq!( + serde_json::from_slice::(&body).expect("JSON"), + serde_json::json!({ "code": "recording_active" }) + ); + } + #[test] fn rejects_unsafe_recording_file_names() { assert!(is_safe_recording_file_name("recording-0.webm")); diff --git a/devolutions-gateway/src/api/mod.rs b/devolutions-gateway/src/api/mod.rs index c7b222ebc..f31a32d15 100644 --- a/devolutions-gateway/src/api/mod.rs +++ b/devolutions-gateway/src/api/mod.rs @@ -13,6 +13,7 @@ pub mod preflight; pub mod rdp; pub mod session; pub mod sessions; +pub mod tasks; pub mod traffic; pub mod tunnel; pub mod update; diff --git a/devolutions-gateway/src/api/tasks.rs b/devolutions-gateway/src/api/tasks.rs new file mode 100644 index 000000000..c82e99076 --- /dev/null +++ b/devolutions-gateway/src/api/tasks.rs @@ -0,0 +1,85 @@ +//! The Task JSON returned by the endpoints of provisioner tasks, such as the AI analysis of a recording. + +use provisioner_task::{ProvisionerTaskRecord, ProvisionerTaskState}; +use time::OffsetDateTime; +use uuid::Uuid; + +/// A Task: the permanent record of one piece of work the provisioner asked Gateway to do +/// +/// The work runs in the background; `state` and `payload` tell how it goes. +/// The shape of `params` and `payload` belongs to the kind. +/// For `recording-ai-analysis`, `payload` is an `AiAnalysisRunningPayload` while `running`, an +/// `AiAnalysisSucceededPayload` once `succeeded`, and an `AiAnalysisFailedPayload` once `failed`. +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct TaskInfo { + /// Task ID, chosen by the provisioner. + id: Uuid, + /// Kind of work, such as `recording-ai-analysis`. + kind: String, + /// What the work is about, such as a session ID. + target: String, + /// What was asked, without secrets. + #[cfg_attr(feature = "openapi", schema(value_type = Object))] + params: serde_json::Value, + state: TaskState, + /// Details of the current state; null while `queued`. + #[cfg_attr(feature = "openapi", schema(value_type = Option))] + payload: serde_json::Value, + /// Number of attempts started so far. + attempts: u32, + /// When the Task was created. + #[serde(with = "time::serde::rfc3339")] + created_at: OffsetDateTime, + /// When its first attempt started. + #[serde(with = "time::serde::rfc3339::option")] + started_at: Option, + /// When it succeeded or failed. + #[serde(with = "time::serde::rfc3339::option")] + finished_at: Option, + /// When an unfinished Task becomes failed, waiting and retries included. + #[serde(with = "time::serde::rfc3339")] + deadline_at: OffsetDateTime, +} + +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum TaskState { + /// Waiting for its first attempt. + Queued, + /// An attempt runs, or the next one waits for a retry. + Running, + Succeeded, + Failed, +} + +impl From for TaskState { + fn from(state: ProvisionerTaskState) -> Self { + match state { + ProvisionerTaskState::Queued => Self::Queued, + ProvisionerTaskState::Running => Self::Running, + ProvisionerTaskState::Succeeded => Self::Succeeded, + ProvisionerTaskState::Failed => Self::Failed, + } + } +} + +impl From for TaskInfo { + fn from(task: ProvisionerTaskRecord) -> Self { + Self { + id: task.id, + kind: task.kind, + target: task.target, + params: task.params, + state: TaskState::from(task.state), + payload: task.payload, + attempts: task.attempts, + created_at: task.created_at, + started_at: task.started_at, + finished_at: task.finished_at, + deadline_at: task.deadline_at, + } + } +} diff --git a/devolutions-gateway/src/extract.rs b/devolutions-gateway/src/extract.rs index b22ee0461..c42d3c14d 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 UpdateScope; diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 5052f6704..db6966f76 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -24,6 +24,8 @@ use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; crate::api::jrec::list_recordings, crate::api::jrec::pull_recording_session, crate::api::jrec::pull_recording_file, + 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, @@ -58,6 +60,16 @@ use crate::config::dto::{DataEncoding, PubKeyFormat, Subscriber}; ConfigPatch, crate::api::jrl::JrlInfo, crate::api::jrec::DeleteManyResult, + 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, @@ -136,6 +148,61 @@ struct SessionInfo { destination_host: Option, } +/// Payload of a `running` AI analysis Task: its progress +#[allow(dead_code)] +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisRunningPayload { + /// Current step: `preparing` (the attempt starts), `reading` (the recordings are read), `describing` (the AI is + /// asked about the recording, part by part), then `writing` (the log is written). + step: String, + /// Work done in the step: bytes read while `reading`, parts described while `describing`. + done: Option, + /// Work the step has in all, in the same unit as `done`. + total: Option, + /// Part of the step done, from 0 to 100. + percent: Option, + /// Short sentence for people following the analysis, such as `asking the AI: 2 of 5 chunks described`. + message: Option, +} + +/// Payload of a `succeeded` AI analysis Task: the log it added to the session +#[allow(dead_code)] +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisSucceededPayload { + /// Name of the new log in the session manifest, such as `ai-analysis-0.slog`. + file_name: String, + /// Model that wrote the log, as reported by the AI provider; the requested model when it reports none. + model: String, + /// Tokens counted by the AI provider; absent when it did not report them for every request. + usage: Option, + /// Version of the instructions sent to the AI provider. + prompt_version: String, +} + +/// Tokens counted by an AI provider +#[allow(dead_code)] +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisUsage { + input_tokens: u64, + output_tokens: u64, +} + +/// Payload of a `failed` AI analysis Task: why it failed +#[allow(dead_code)] +#[derive(Serialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +struct AiAnalysisFailedPayload { + /// Kind of failure, such as `permanent error`, `attempts exhausted`, `key lost` or `timed out`. + reason: String, + /// What went wrong: a text, except for `timed out`, where it is `{ "lastPayload": }`. + details: serde_json::Value, + /// Number of attempts started. + attempts: u32, +} + #[allow(unused)] #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "kebab-case")] diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index 8b17d3746..6ae4bf015 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -14,6 +14,8 @@ use url::Url; use super::AttemptError; use crate::config::Conf; +/// AI provider to ask; `openai-compatible` needs a base URL +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum AiProvider { #[serde(rename = "openai")] diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 25cf3a8f0..48b07b0b4 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -33,6 +33,7 @@ use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; use crate::provisioning::ProvisioningStore; use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; +use crate::token::AccessScope; pub const KIND: &str = "recording-ai-analysis"; @@ -251,6 +252,10 @@ impl RecordingAiAnalysis { impl ProvisionerTask for RecordingAiAnalysis { const KIND: &'static str = KIND; + const START_SCOPE: AccessScope = AccessScope::RecordingAiAnalysisStart; + + const READ_SCOPE: AccessScope = AccessScope::RecordingAiAnalysisRead; + type Params = Params; type Output = Analysis; diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 308ada5e6..4e43724f2 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -30,6 +30,7 @@ use uuid::Uuid; use super::AttemptError; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; +use crate::token::AccessScope; /// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. /// @@ -40,6 +41,12 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// The kind written in the records of this task. const KIND: &'static str; + /// The scope a token needs to start this task, `gateway.tasks..start`. + const START_SCOPE: AccessScope; + + /// The scope a token needs to read the records of this task, `gateway.tasks..read`. + const READ_SCOPE: AccessScope; + /// What the record keeps as asked; never secrets, since records are kept and returned to callers. type Params: Serialize + DeserializeOwned + Send; diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index a40619393..f063a771f 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -6,6 +6,7 @@ use provisioner_task::ProvisionerTaskState; use super::*; use crate::job_queue::{JobMessage, JobQueueReceiver}; +use crate::token::AccessScope; const KIND: &str = "test.echo"; @@ -25,6 +26,10 @@ struct Echo { impl ProvisionerTask for Echo { const KIND: &'static str = KIND; + const START_SCOPE: AccessScope = AccessScope::Wildcard; + + const READ_SCOPE: AccessScope = AccessScope::Wildcard; + type Params = EchoParams; type Output = serde_json::Value; diff --git a/devolutions-gateway/src/token.rs b/devolutions-gateway/src/token.rs index 97d268e44..3ce855cac 100644 --- a/devolutions-gateway/src/token.rs +++ b/devolutions-gateway/src/token.rs @@ -511,6 +511,10 @@ pub enum AccessScope { AgentDelete, #[serde(rename = "gateway.agent.read")] AgentRead, + #[serde(rename = "gateway.tasks.recording-ai-analysis.start")] + RecordingAiAnalysisStart, + #[serde(rename = "gateway.tasks.recording-ai-analysis.read")] + RecordingAiAnalysisRead, } #[derive(Clone, Serialize, Deserialize)] diff --git a/devolutions-gateway/tests/ai_analysis.rs b/devolutions-gateway/tests/ai_analysis.rs new file mode 100644 index 000000000..d69889bf2 --- /dev/null +++ b/devolutions-gateway/tests/ai_analysis.rs @@ -0,0 +1,566 @@ +#![allow(unused_crate_dependencies)] +#![allow(clippy::unwrap_used)] + +use std::io::Write; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; + +use axum::Router; +use axum::body::Body; +use axum::extract::connect_info::MockConnectInfo; +use axum::http::{Request, StatusCode}; +use base64::Engine as _; +use camino::{Utf8Path, Utf8PathBuf}; +use devolutions_gateway::job_queue::{JobMessage, JobQueueReceiver}; +use devolutions_gateway::recording::RecordingManagerTask; +use devolutions_gateway::{DgwState, MockHandles}; +use devolutions_gateway_task::ShutdownHandle; +use http_body_util::BodyExt as _; +use serde_json::{Value, json}; +use time::OffsetDateTime; +use time::format_description::well_known::Rfc3339; +use tower::ServiceExt as _; +use tracing_subscriber::util::SubscriberInitExt as _; +use uuid::Uuid; + +const API_KEY: &str = "sk-ai-analysis-endpoint-test-secret"; + +fn config(recording_path: &Utf8Path) -> String { + json!({ + "ProvisionerPublicKeyData": { + "Value": "mMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA4vuqLOkl1pWobt6su1XO9VskgCAwevEGs6kkNjJQBwkGnPKYLmNF1E/af1yCocfVn/OnPf9e4x+lXVyZ6LMDJxFxu+axdgOq3Ld392J1iAEbfvwlyRFnEXFOJNyylqg3bY6LvnWHL/XZczVdMD9xYfq2sO9bg3xjRW4s7r9EEYOFjqVT3VFznH9iWJVtcSEKukmS/3uKoO6lGhacvu0HhjXXdgq0R8zvR4XRJ9Fcnf0f9Ypoc+i6L80NVjrRCeVOH+Ld/2fA9bocpfLarcVqG3RjS+qgOtpyCc0jWVFF4zaGQ7LUDFkEIYILkICeMMn2ll29hmZNzsJzZJ9s6NocgQIDAQAB" + }, + "Listeners": [{ "InternalUrl": "http://*:7171", "ExternalUrl": "https://*:7171" }], + "Proxy": { "Mode": "Off" }, + "RecordingPath": recording_path, + "__debug__": { "disable_token_validation": true }, + }) + .to_string() +} + +fn scope_token(scope: &str) -> String { + let engine = base64::engine::general_purpose::URL_SAFE_NO_PAD; + let header = engine.encode(r#"{"alg":"RS256","typ":"JWT","cty":"SCOPE"}"#); + let payload = engine.encode( + json!({ + "jti": Uuid::new_v4(), + "exp": OffsetDateTime::now_utc().unix_timestamp() + 3600, + "scope": scope, + }) + .to_string(), + ); + let signature = engine.encode(b"signature"); + format!("{header}.{payload}.{signature}") +} + +/// Collects every log line, so a test can check that no secret reaches the logs. +#[derive(Clone, Default)] +struct LogBuffer(Arc>>); + +impl Write for LogBuffer { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.0.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +impl LogBuffer { + fn text(&self) -> String { + String::from_utf8_lossy(&self.0.lock().unwrap()).into_owned() + } +} + +struct Harness { + app: Router, + job_queue_rx: JobQueueReceiver, + recording_path: Utf8PathBuf, + db_path: Utf8PathBuf, + task_store: provisioner_task::DynProvisionerTaskStore, + _keep: Box, + _dir: tempfile::TempDir, +} + +impl Harness { + async fn new() -> Self { + let dir = tempfile::tempdir().unwrap(); + let root = Utf8PathBuf::from_path_buf(dir.path().to_path_buf()).unwrap(); + let recording_path = root.join("recordings"); + let db_path = root.join("gateway.db"); + + let (mut state, handles) = DgwState::mock(&config(&recording_path)).await.unwrap(); + + // A file, unlike the in-memory default, can be searched for secrets afterwards. + let connection = gateway_db::GatewayDb::open_path(db_path.as_str()) + .await + .unwrap() + .connect() + .await + .unwrap(); + let task_store: provisioner_task::DynProvisionerTaskStore = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(connection) + .await + .unwrap(), + ); + state.provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( + Arc::clone(&task_store), + state.job_queue_handle.clone(), + state.conf_handle.clone(), + state.recordings.clone(), + state.provisioning.clone(), + ); + + let MockHandles { + session_manager_rx, + recording_manager_rx, + subscriber_rx, + job_queue_rx, + traffic_audit_rx, + shutdown_handle, + } = handles; + + let recording_manager = RecordingManagerTask::new( + recording_manager_rx, + recording_path.clone(), + state.sessions.clone(), + state.job_queue_handle.clone(), + ); + let (recording_shutdown_handle, recording_shutdown_signal) = ShutdownHandle::new(); + tokio::spawn(devolutions_gateway_task::Task::run( + recording_manager, + recording_shutdown_signal, + )); + + let app = devolutions_gateway::make_http_service(state) + .layer(MockConnectInfo(SocketAddr::from(([127, 0, 0, 1], 3000)))); + + Self { + app, + job_queue_rx, + recording_path, + db_path, + task_store, + _keep: Box::new(( + session_manager_rx, + subscriber_rx, + traffic_audit_rx, + shutdown_handle, + recording_shutdown_handle, + )), + _dir: dir, + } + } + + fn write_session(&self) -> Uuid { + let session_id = Uuid::new_v4(); + let dir = self.recording_path.join(session_id.to_string()); + std::fs::create_dir_all(&dir).unwrap(); + + let manifest = json!({ + "sessionId": session_id, + "startTime": 1_787_255_035, + "duration": 3, + "files": [{ "fileName": "recording-0.cast", "startTime": 1_787_255_035, "duration": 3 }], + }); + std::fs::write(dir.join("recording.json"), manifest.to_string()).unwrap(); + std::fs::write( + dir.join("recording-0.cast"), + "{\"version\": 2}\n[0.5,\"o\",\"$ ls\\r\\n\"]\n[1.0,\"o\",\"file.txt\\r\\n\"]\n", + ) + .unwrap(); + + session_id + } + + async fn send(&self, request: Request) -> (StatusCode, Value) { + let response = self.app.clone().oneshot(request).await.unwrap(); + let status = response.status(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let body = if body.is_empty() { + Value::Null + } else { + serde_json::from_slice(&body).unwrap() + }; + (status, body) + } + + async fn post_raw(&self, session_id: Uuid, scope: &str, body: String) -> (StatusCode, Value) { + let request = Request::builder() + .method("POST") + .uri(format!("/jet/jrec/{session_id}/ai-analysis")) + .header("content-type", "application/json") + .header("authorization", format!("Bearer {}", scope_token(scope))) + .body(Body::from(body)) + .unwrap(); + self.send(request).await + } + + async fn post(&self, session_id: Uuid, body: &Value) -> (StatusCode, Value) { + self.post_raw( + session_id, + "gateway.tasks.recording-ai-analysis.start", + body.to_string(), + ) + .await + } + + async fn get_analysis(&self, session_id: Uuid, task_id: Uuid, scope: &str) -> (StatusCode, Value) { + let request = Request::builder() + .uri(format!("/jet/jrec/{session_id}/ai-analysis/{task_id}")) + .header("authorization", format!("Bearer {}", scope_token(scope))) + .body(Body::empty()) + .unwrap(); + self.send(request).await + } + + async fn run_queued_job(&mut self) { + let JobMessage { mut job, .. } = self.job_queue_rx.try_recv().expect("a queued job"); + job.run().await.unwrap(); + } + + fn database_bytes(&self) -> Vec { + let mut bytes = Vec::new(); + for suffix in ["", "-wal"] { + if let Ok(content) = std::fs::read(format!("{}{suffix}", self.db_path)) { + bytes.extend(content); + } + } + bytes + } +} + +fn request(task_id: Uuid, model: &str) -> Value { + json!({ "taskId": task_id, "provider": "openai", "model": model, "apiKey": API_KEY }) +} + +/// A mock OpenAI-compatible provider answering every request with one action, or with `status` when it is not 200. +async fn spawn_ai_provider(status: u16) -> String { + let app = Router::new().fallback(axum::routing::post(move || async move { + if status != 200 { + return ( + StatusCode::from_u16(status).unwrap(), + axum::Json(json!({ "error": { "message": "scripted failure" } })), + ); + } + + let answer = json!({ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-test-2026-09-30", + "choices": [{ + "index": 0, + "message": { "role": "assistant", "content": "{\"line\":\"L1\",\"description\":\"Listed files\"}" }, + "finish_reason": "stop" + }], + "usage": { "prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30 } + }); + (StatusCode::OK, axum::Json(answer)) + })); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + + format!("http://{addr}/v1/") +} + +fn assert_rfc3339(value: &Value) { + let text = value.as_str().unwrap_or_else(|| panic!("not a string: {value}")); + OffsetDateTime::parse(text, &Rfc3339).unwrap_or_else(|error| panic!("{text}: {error}")); +} + +#[tokio::test] +async fn start_returns_202_then_200_for_the_same_request() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + + let (status, created) = harness.post(session_id, &request(task_id, "gpt-test")).await; + assert_eq!(status, StatusCode::ACCEPTED, "{created}"); + assert_eq!(created["id"], task_id.to_string()); + assert_eq!(created["kind"], "recording-ai-analysis"); + assert_eq!(created["target"], session_id.to_string()); + assert_eq!( + created["params"], + json!({ "provider": "openai", "model": "gpt-test", "maxOutputTokens": null }) + ); + assert_eq!(created["state"], "queued"); + assert_eq!(created["payload"], Value::Null); + assert_eq!(created["attempts"], 0); + assert_rfc3339(&created["createdAt"]); + assert_rfc3339(&created["deadlineAt"]); + assert_eq!(created["startedAt"], Value::Null); + assert_eq!(created["finishedAt"], Value::Null); + + let (status, again) = harness.post(session_id, &request(task_id, "gpt-test")).await; + assert_eq!(status, StatusCode::OK, "{again}"); + assert_eq!(again, created); +} + +#[tokio::test] +async fn start_conflicts_return_409_with_a_code() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + harness.post(session_id, &request(task_id, "gpt-test")).await; + + let (status, body) = harness.post(session_id, &request(task_id, "gpt-other")).await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(body, json!({ "code": "task_id_conflict" })); + + let (status, body) = harness.post(session_id, &request(Uuid::new_v4(), "gpt-test")).await; + assert_eq!(status, StatusCode::CONFLICT); + assert_eq!(body, json!({ "code": "analysis_in_progress", "taskId": task_id })); +} + +#[tokio::test] +async fn start_without_recording_returns_404() { + let harness = Harness::new().await; + + let (status, _) = harness.post(Uuid::new_v4(), &request(Uuid::new_v4(), "gpt-test")).await; + + assert_eq!(status, StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn invalid_bodies_and_settings_return_400() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + + let mut unknown_field = request(task_id, "gpt-test"); + unknown_field["temperature"] = json!(0.5); + + let mut no_key = request(task_id, "gpt-test"); + no_key.as_object_mut().unwrap().remove("apiKey"); + + let mut empty_key = request(task_id, "gpt-test"); + empty_key["apiKey"] = json!(""); + + let mut no_base_url = request(task_id, "gpt-test"); + no_base_url["provider"] = json!("openai-compatible"); + + let mut bad_base_url = request(task_id, "gpt-test"); + bad_base_url["baseUrl"] = json!("ftp://ai.example/v1/"); + + let mut unknown_provider = request(task_id, "gpt-test"); + unknown_provider["provider"] = json!("skynet"); + + for (case, body) in [ + ("unknown field", unknown_field.to_string()), + ("missing API key", no_key.to_string()), + ("empty API key", empty_key.to_string()), + ("empty model", request(task_id, " ").to_string()), + ("missing base URL", no_base_url.to_string()), + ("unsupported base URL", bad_base_url.to_string()), + ("unknown provider", unknown_provider.to_string()), + ("not JSON", "{".to_owned()), + ] { + let (status, _) = harness + .post_raw(session_id, "gateway.tasks.recording-ai-analysis.start", body) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST, "{case}"); + } + + let (status, _) = harness + .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::NOT_FOUND, "no Task is recorded"); +} + +#[tokio::test] +async fn endpoints_require_their_scope() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let body = request(Uuid::new_v4(), "gpt-test").to_string(); + + for scope in [ + "gateway.tasks.recording-ai-analysis.read", + "gateway.recording.delete", + "gateway.recordings.read", + ] { + let (status, _) = harness.post_raw(session_id, scope, body.clone()).await; + assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); + } + + for scope in ["gateway.tasks.recording-ai-analysis.start", "gateway.sessions.read"] { + let (status, _) = harness.get_analysis(session_id, Uuid::new_v4(), scope).await; + assert_eq!(status, StatusCode::FORBIDDEN, "{scope}"); + } + + let (status, _) = harness.post_raw(session_id, "*", body).await; + assert_eq!(status, StatusCode::ACCEPTED, "the wildcard scope is accepted"); + + let (status, _) = harness.get_analysis(session_id, Uuid::new_v4(), "*").await; + assert_eq!(status, StatusCode::NOT_FOUND, "the wildcard scope is accepted"); +} + +#[tokio::test] +async fn an_analysis_is_found_at_its_location_only() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + let request = Request::builder() + .method("POST") + .uri(format!("/jet/jrec/{session_id}/ai-analysis")) + .header("content-type", "application/json") + .header( + "authorization", + format!("Bearer {}", scope_token("gateway.tasks.recording-ai-analysis.start")), + ) + .body(Body::from(request(task_id, "gpt-test").to_string())) + .unwrap(); + + let response = harness.app.clone().oneshot(request).await.unwrap(); + + assert_eq!(response.status(), StatusCode::ACCEPTED); + let location = response.headers()["location"].to_str().unwrap().to_owned(); + assert_eq!(location, format!("/jet/jrec/{session_id}/ai-analysis/{task_id}")); + + let (status, task) = harness + .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::OK, "{task}"); + assert_eq!(task["id"], task_id.to_string()); + + let (status, _) = harness + .get_analysis(Uuid::new_v4(), task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::NOT_FOUND, "not under another session"); + + // A Task of another kind is not an AI analysis, whatever the scope. + let now = OffsetDateTime::now_utc(); + let other_kind = Uuid::new_v4(); + harness + .task_store + .create( + provisioner_task::NewProvisionerTaskRecord { + id: other_kind, + kind: "test.other-kind".to_owned(), + target: session_id.to_string(), + params: json!({}), + deadline_at: now + time::Duration::hours(1), + job_token: Uuid::new_v4(), + }, + now, + ) + .await + .unwrap(); + let (status, _) = harness.get_analysis(session_id, other_kind, "*").await; + assert_eq!(status, StatusCode::NOT_FOUND); +} +#[tokio::test] +async fn get_analysis_returns_404_for_an_unknown_id() { + let harness = Harness::new().await; + + let (status, _) = harness + .get_analysis( + Uuid::new_v4(), + Uuid::new_v4(), + "gateway.tasks.recording-ai-analysis.read", + ) + .await; + + assert_eq!(status, StatusCode::NOT_FOUND); +} + +#[tokio::test] +async fn analysis_succeeds_without_leaking_the_key_or_the_base_url() { + let logs = LogBuffer::default(); + let _guard = tracing_subscriber::fmt() + .with_writer({ + let logs = logs.clone(); + move || logs.clone() + }) + .with_ansi(false) + .with_max_level(tracing::Level::TRACE) + .set_default(); + + let mut harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_ai_provider(200).await; + + let body = json!({ + "taskId": task_id, + "provider": "openai-compatible", + "model": "gpt-test", + "baseUrl": base_url, + "maxOutputTokens": 1000, + "apiKey": API_KEY, + }); + + let (status, created) = harness.post(session_id, &body).await; + assert_eq!(status, StatusCode::ACCEPTED, "{created}"); + + harness.run_queued_job().await; + + let (status, task) = harness + .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(task["state"], "succeeded", "{task}"); + assert_eq!(task["attempts"], 1); + assert_eq!(task["payload"]["fileName"], "ai-analysis-0.slog"); + assert_eq!(task["payload"]["model"], "gpt-test-2026-09-30"); + assert_eq!( + task["payload"]["usage"], + json!({ "inputTokens": 10, "outputTokens": 20 }) + ); + assert_rfc3339(&task["startedAt"]); + assert_rfc3339(&task["finishedAt"]); + + let (status, again) = harness.post(session_id, &body).await; + assert_eq!(status, StatusCode::OK, "{again}"); + + let responses = [created, task, again].map(|body| body.to_string()).join("\n"); + let logs = logs.text(); + let database = String::from_utf8_lossy(&harness.database_bytes()).into_owned(); + assert!(!database.is_empty(), "gateway.db is read"); + assert!(logs.contains("Task succeeded"), "logs are captured"); + + for secret in [API_KEY, base_url.as_str()] { + assert!(!responses.contains(secret), "{secret} in a response"); + assert!(!logs.contains(secret), "{secret} in the logs"); + assert!(!database.contains(secret), "{secret} in gateway.db"); + } +} + +#[tokio::test] +async fn provider_refusing_the_key_fails_the_task() { + let mut harness = Harness::new().await; + let session_id = harness.write_session(); + let task_id = Uuid::new_v4(); + let base_url = spawn_ai_provider(401).await; + + let body = json!({ + "taskId": task_id, + "provider": "openai-compatible", + "model": "gpt-test", + "baseUrl": base_url, + "apiKey": API_KEY, + }); + let (status, _) = harness.post(session_id, &body).await; + assert_eq!(status, StatusCode::ACCEPTED); + + harness.run_queued_job().await; + + let (status, task) = harness + .get_analysis(session_id, task_id, "gateway.tasks.recording-ai-analysis.read") + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(task["state"], "failed", "{task}"); + assert_eq!(task["payload"]["reason"], "permanent error"); + assert_eq!(task["payload"]["attempts"], 1); + + let (status, _) = harness.post(session_id, &request(Uuid::new_v4(), "gpt-test")).await; + assert_eq!( + status, + StatusCode::ACCEPTED, + "a failed analysis no longer blocks the session" + ); +} diff --git a/devolutions-gateway/tests/dvls_compatibility.rs b/devolutions-gateway/tests/dvls_compatibility.rs index e727c8e01..fa0d080cc 100644 --- a/devolutions-gateway/tests/dvls_compatibility.rs +++ b/devolutions-gateway/tests/dvls_compatibility.rs @@ -301,6 +301,10 @@ mod as_of_v2022_3_0_0 { JrlRead, #[serde(rename = "gateway.config.write")] ConfigWrite, + #[serde(rename = "gateway.tasks.recording-ai-analysis.start")] + RecordingAiAnalysisStart, + #[serde(rename = "gateway.tasks.recording-ai-analysis.read")] + RecordingAiAnalysisRead, } fn dvls_access_scope() -> impl Strategy { @@ -312,6 +316,8 @@ mod as_of_v2022_3_0_0 { Just(DvlsAccessScope::JrlRead), Just(DvlsAccessScope::ConfigWrite), Just(DvlsAccessScope::ConfigWrite), + Just(DvlsAccessScope::RecordingAiAnalysisStart), + Just(DvlsAccessScope::RecordingAiAnalysisRead), ] } From b639b20336a3ac7c83ef377716cab2bdc092d584 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 16:41:35 -0400 Subject: [PATCH 04/21] fix(dgw): store provisioner task jobs before reporting them queued A Task was reported created once its job reached the channel of the job queue task, which gives up after a few failed writes without telling anyone, so a Task could be created without a job. The runner now writes jobs to the job queue database itself, through ProvisionerTaskJobQueue, and only reports a Task created once its job is stored. At startup, a Task whose new job cannot be stored is failed as not queued, instead of waiting for its deadline with no job. A Task is only cleaned up once its end is recorded: when recording a failure fails, the job queue tries the attempt again; when recording a success fails, the Task keeps its key and working files and fails at its deadline, since running the work again would add a second log. Waiting for an earlier AI analysis of the same session now stops when the attempt is cancelled, so the attempt timeout also covers the wait. --- devolutions-gateway/src/job_queue.rs | 21 ++- devolutions-gateway/src/lib.rs | 2 +- .../src/provisioner_tasks/mod.rs | 5 +- .../recording_ai_analysis/mod.rs | 6 +- .../recording_ai_analysis/tests.rs | 2 +- .../src/provisioner_tasks/runner/mod.rs | 161 +++++++++++------- .../src/provisioner_tasks/runner/tests.rs | 52 +++--- devolutions-gateway/src/service.rs | 4 +- devolutions-gateway/tests/ai_analysis.rs | 2 +- 9 files changed, 156 insertions(+), 99 deletions(-) diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 3e4b611fc..60db0a367 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -11,7 +11,7 @@ use job_queue_libsql::libsql; use time::OffsetDateTime; use tokio::sync::{Notify, mpsc}; -use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskRunner}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; /// Runs of a job before the job queue gives up on it. pub const MAX_ATTEMPTS: u32 = 5; @@ -90,12 +90,21 @@ impl JobQueueCtx { }) } - /// Writes a job to the queue right away, without going through [`JobQueueHandle`]. - /// - /// For startup, before the job runner claims anything. - pub async fn push(&self, job: T) -> anyhow::Result<()> { + /// The queue itself, for jobs whose callers must know they are stored. + pub fn durable(&self) -> DurableJobQueue { + DurableJobQueue(Arc::clone(&self.queue)) + } +} + +/// The job queue database: a job pushed here is stored before `push` returns, unlike with [`JobQueueHandle`]. +#[derive(Clone)] +pub struct DurableJobQueue(DynJobQueue); + +#[async_trait] +impl ProvisionerTaskJobQueue for DurableJobQueue { + async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { let job: DynJob = Box::new(job); - self.queue.push_job(&job, None).await + self.0.push_job(&job, None).await } } diff --git a/devolutions-gateway/src/lib.rs b/devolutions-gateway/src/lib.rs index 328d015df..3a0999e6c 100644 --- a/devolutions-gateway/src/lib.rs +++ b/devolutions-gateway/src/lib.rs @@ -101,7 +101,7 @@ impl DgwState { let task_store = Arc::new(provisioner_task_libsql::LibSqlProvisionerTaskStore::open(task_store).await?); let provisioner_tasks = provisioner_tasks::task_runner( task_store, - job_queue_handle.clone(), + Arc::new(job_queue_handle.clone()), conf_handle.clone(), recording_manager_handle.clone(), provisioning.clone(), diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index b8bc845eb..ccb1d4938 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -8,16 +8,15 @@ pub mod runner; use provisioner_task::DynProvisionerTaskStore; use self::recording_ai_analysis::RecordingAiAnalysis; -use self::runner::ProvisionerTaskRunner; +use self::runner::{DynProvisionerTaskJobQueue, ProvisionerTaskRunner}; use crate::config::ConfHandle; -use crate::job_queue::JobQueueHandle; use crate::provisioning::ProvisioningStore; use crate::recording::RecordingMessageSender; /// The provisioner task runner, with every provisioner task Gateway knows. pub fn task_runner( store: DynProvisionerTaskStore, - job_queue: JobQueueHandle, + job_queue: DynProvisionerTaskJobQueue, conf_handle: ConfHandle, recordings: RecordingMessageSender, provisioning: ProvisioningStore, diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 48b07b0b4..dea325249 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -329,7 +329,11 @@ impl RecordingAiAnalysis { .client(&conf, &api_key) .map_err(|error| AttemptError::Permanent(error.to_string()))?; - let _session = lock_session(session_id).await; + // Waiting for an earlier analysis of the session counts in the attempt time, so a timeout cancels it too. + let _session = tokio::select! { + session = lock_session(session_id) => session, + () = attempt.cancel.cancelled() => return Err(AttemptError::Cancelled), + }; let recording = self .recordings diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 3b07f8bbd..fdcee3db4 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -681,7 +681,7 @@ async fn slow_attempt_is_retried() { task.attempt_timeout = Duration::from_millis(500); harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( Arc::clone(harness.state.provisioner_tasks.store()), - harness.state.job_queue_handle.clone(), + Arc::new(harness.state.job_queue_handle.clone()), ) .register(task) .build(); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 4e43724f2..c3b848d39 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -13,7 +13,6 @@ mod tests; use std::collections::HashMap; -use std::future::Future; use std::sync::Arc; use anyhow::Context as _; @@ -30,6 +29,23 @@ use uuid::Uuid; use super::AttemptError; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; + +/// Where the runner queues the jobs of its records. +#[async_trait] +pub trait ProvisionerTaskJobQueue: Send + Sync { + /// Queues `job`, and returns once it is stored, so a record never waits for a job that was lost. + async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()>; +} + +pub type DynProvisionerTaskJobQueue = Arc; + +/// Hands the job to the job queue task, which stores it later: for tests, which take the jobs from its channel. +#[async_trait] +impl ProvisionerTaskJobQueue for JobQueueHandle { + async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { + self.enqueue(job).await + } +} use crate::token::AccessScope; /// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. @@ -135,14 +151,14 @@ pub struct ProvisionerTaskRunner { struct Inner { store: DynProvisionerTaskStore, - job_queue: JobQueueHandle, + job_queue: DynProvisionerTaskJobQueue, /// Kind → the task that runs records of this kind. tasks: HashMap<&'static str, Arc>, } pub struct ProvisionerTaskRunnerBuilder { store: DynProvisionerTaskStore, - job_queue: JobQueueHandle, + job_queue: DynProvisionerTaskJobQueue, tasks: HashMap<&'static str, Arc>, } @@ -171,7 +187,10 @@ fn starting_payload() -> serde_json::Value { } impl ProvisionerTaskRunner { - pub fn builder(store: DynProvisionerTaskStore, job_queue: JobQueueHandle) -> ProvisionerTaskRunnerBuilder { + pub fn builder( + store: DynProvisionerTaskStore, + job_queue: DynProvisionerTaskJobQueue, + ) -> ProvisionerTaskRunnerBuilder { ProvisionerTaskRunnerBuilder { store, job_queue, @@ -221,17 +240,14 @@ impl ProvisionerTaskRunner { return Ok(outcome); }; - if let Err(error) = on_created(record) { - error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to prepare the Task"); - self.fail(task.as_ref(), id, failure("not queued", &format!("{error:#}"), 0)) - .await; - return Err(error); - } + let queued = match on_created(record) { + Ok(()) => self.enqueue(id, job_token).await, + Err(error) => Err(error.context("prepare the Task")), + }; - if let Err(error) = self.enqueue(id, job_token).await { + if let Err(error) = queued { error!(task.id = %id, task.kind = T::KIND, error = format!("{error:#}"), "Failed to queue the Task"); - self.fail(task.as_ref(), id, failure("not queued", &format!("{error:#}"), 0)) - .await; + self.fail_not_queued(task.as_ref(), id, &error, 0).await; return Err(error); } @@ -240,14 +256,11 @@ impl ProvisionerTaskRunner { Ok(outcome) } - /// Hands every unfinished record to a new job queued with `push`: the jobs queued before stop at their next attempt. + /// Hands every unfinished record to a new job: the jobs queued before stop at their next attempt. /// - /// Call it at startup, before the job runner claims any job, so no earlier job is running an attempt meanwhile. - pub async fn resume(&self, push: F) -> anyhow::Result<()> - where - F: Fn(ProvisionerTaskJob) -> Fut, - Fut: Future>, - { + /// Call it at startup, before the job runner claims any job, so no earlier job is running an attempt meanwhile. A + /// record whose new job cannot be queued is failed, since no job is left to run it. + pub async fn resume(&self) -> anyhow::Result<()> { let now = OffsetDateTime::now_utc(); let records = self.inner.store.list_unfinished(now).await?; @@ -258,20 +271,20 @@ impl ProvisionerTaskRunner { continue; } - let job = ProvisionerTaskJob { - task_id: record.id, - job_token, - runner: self.clone(), - }; - - match push(job).await { + match self.enqueue(record.id, job_token).await { Ok(()) => info!(task.id = %record.id, task.kind = %record.kind, "Task queued again"), - Err(error) => error!( - task.id = %record.id, - task.kind = %record.kind, - error = format!("{error:#}"), - "Failed to queue the Task again; it fails at its deadline" - ), + Err(error) => { + error!( + task.id = %record.id, + task.kind = %record.kind, + error = format!("{error:#}"), + "Failed to queue the Task again" + ); + if let Some(task) = self.task(&record.kind) { + self.fail_not_queued(task.as_ref(), record.id, &error, record.attempts) + .await; + } + } } } @@ -291,7 +304,7 @@ impl ProvisionerTaskRunner { async fn enqueue(&self, task_id: Uuid, job_token: Uuid) -> anyhow::Result<()> { self.inner .job_queue - .enqueue(ProvisionerTaskJob { + .push(ProvisionerTaskJob { task_id, job_token, runner: self.clone(), @@ -374,24 +387,28 @@ impl ProvisionerTaskRunner { }; match outcome { - Ok(output) => { - // An error here must not ask for a retry: the work is done. - match store.succeed(task_id, output, OffsetDateTime::now_utc()).await { - Ok(true) => info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"), - Ok(false) => warn!( - task.id = %task_id, - task.kind = %record.kind, - "Task attempt succeeded after the Task had ended, such as past its deadline; the Task keeps its earlier end" - ), - Err(error) => error!( - task.id = %task_id, - error = format!("{error:#}"), - "Failed to record the Task result" - ), + Ok(output) => match store.succeed(task_id, output, OffsetDateTime::now_utc()).await { + Ok(ended) => { + if ended { + info!(task.id = %task_id, task.kind = %record.kind, attempts, "Task succeeded"); + } else { + warn!( + task.id = %task_id, + task.kind = %record.kind, + "Task attempt succeeded after the Task had ended, such as past its deadline; the Task keeps its earlier end" + ); + } + + self.end(task.as_ref(), task_id).await; } - - self.end(task.as_ref(), task_id).await; - } + // No retry: the work is done, and doing it again would repeat it, such as adding a second log. Nothing is + // dropped either, since the record still runs; it fails at its deadline. + Err(error) => error!( + task.id = %task_id, + error = format!("{error:#}"), + "Failed to record the Task result; the Task fails at its deadline" + ), + }, Err(AttemptError::Transient(details)) if attempts < MAX_ATTEMPTS => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task attempt failed; it will be retried"); anyhow::bail!("Task attempt failed: {details}"); @@ -403,17 +420,17 @@ impl ProvisionerTaskRunner { task_id, failure("attempts exhausted", &details, attempts), ) - .await; + .await?; } Err(AttemptError::Permanent(details)) => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, %details, "Task failed"); self.fail(task.as_ref(), task_id, failure("permanent error", &details, attempts)) - .await; + .await?; } Err(AttemptError::Failed { reason, details }) => { warn!(task.id = %task_id, task.kind = %record.kind, attempts, reason, %details, "Task failed"); self.fail(task.as_ref(), task_id, failure(reason, &details, attempts)) - .await; + .await?; } Err(AttemptError::Cancelled) => unreachable!("a cancelled attempt is reported as a timeout above"), } @@ -421,12 +438,40 @@ impl ProvisionerTaskRunner { Ok(()) } - async fn fail(&self, task: &dyn AnyProvisionerTask, task_id: Uuid, payload: serde_json::Value) { - if let Err(error) = self.inner.store.fail(task_id, payload, OffsetDateTime::now_utc()).await { - error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure"); - } + /// Records the failure, then drops what the record kept while running. + /// + /// If the failure cannot be recorded, nothing is dropped, since the record may still run, and the error is returned so + /// the job queue tries again. + async fn fail( + &self, + task: &dyn AnyProvisionerTask, + task_id: Uuid, + payload: serde_json::Value, + ) -> anyhow::Result<()> { + self.inner + .store + .fail(task_id, payload, OffsetDateTime::now_utc()) + .await + .context("record the Task failure")?; self.end(task, task_id).await; + + Ok(()) + } + + /// Fails a record no job will run. + async fn fail_not_queued( + &self, + task: &dyn AnyProvisionerTask, + task_id: Uuid, + error: &anyhow::Error, + attempts: u32, + ) { + let payload = failure("not queued", &format!("{error:#}"), attempts); + + if let Err(error) = self.fail(task, task_id, payload).await { + error!(task.id = %task_id, error = format!("{error:#}"), "Failed to record the Task failure; it fails at its deadline"); + } } async fn end(&self, task: &dyn AnyProvisionerTask, task_id: Uuid) { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index f063a771f..6cbd047ae 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -1,6 +1,5 @@ use std::sync::atomic::{AtomicUsize, Ordering}; -use job_queue::Job as _; use parking_lot::Mutex; use provisioner_task::ProvisionerTaskState; @@ -76,7 +75,7 @@ impl Harness { let runs = Arc::new(AtomicUsize::new(0)); let ended = Arc::new(Mutex::new(Vec::new())); - let runner = ProvisionerTaskRunner::builder(store, job_queue) + let runner = ProvisionerTaskRunner::builder(store, Arc::new(job_queue)) .register(Echo { outcome, runs: Arc::clone(&runs), @@ -183,45 +182,46 @@ async fn a_superseded_job_runs_nothing_once_the_task_is_resumed() { let id = harness.create(7).await; let mut stale = harness.queued_job(); - let resumed = Mutex::new(Vec::new()); - harness - .runner - .resume(|job| { - resumed.lock().push(job); - async { Ok(()) } - }) - .await - .expect("resume"); - let mut resumed = resumed.into_inner(); - assert_eq!(resumed.len(), 1, "the unfinished Task gets a new job"); + harness.runner.resume().await.expect("resume"); + let mut resumed = harness.queued_job(); + assert!( + harness.job_queue_rx.try_recv().is_err(), + "the unfinished Task gets one new job" + ); stale.run().await.expect("the stale job stops"); assert_eq!(harness.runs.load(Ordering::SeqCst), 0); assert_eq!(harness.task(id).await.attempts, 0, "a stale job counts no attempt"); assert!(harness.ended.lock().is_empty(), "a stale job ends nothing"); - resumed[0].run().await.expect("the new job runs"); + resumed.run().await.expect("the new job runs"); let task = harness.task(id).await; assert_eq!(task.state, ProvisionerTaskState::Succeeded); assert_eq!(task.payload, json!({ "echo": 7 })); assert_eq!(*harness.ended.lock(), [id]); - let queued_again = AtomicUsize::new(0); - harness - .runner - .resume(|_| { - queued_again.fetch_add(1, Ordering::SeqCst); - async { Ok(()) } - }) - .await - .expect("resume again"); - assert_eq!( - queued_again.load(Ordering::SeqCst), - 0, + harness.runner.resume().await.expect("resume again"); + assert!( + harness.job_queue_rx.try_recv().is_err(), "a finished Task gets no new job" ); } +#[tokio::test] +async fn a_task_that_cannot_be_queued_again_fails() { + let mut harness = Harness::new(|value| Ok(json!(value))).await; + let id = harness.create(1).await; + harness.queued_job(); + harness.job_queue_rx.close(); + + harness.runner.resume().await.expect("resume"); + + let task = harness.task(id).await; + assert_eq!(task.state, ProvisionerTaskState::Failed, "no job is left to run it"); + assert_eq!(task.payload["reason"], "not queued"); + assert_eq!(*harness.ended.lock(), [id]); +} + #[tokio::test] async fn a_panicking_attempt_fails_the_task() { let mut harness = Harness::new(|_| panic!("the attempt panics")).await; diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index eea186dca..efc0ea707 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -352,7 +352,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { let provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( task_store, - job_queue_ctx.job_queue_handle.clone(), + Arc::new(job_queue_ctx.durable()), conf_handle.clone(), recording_manager_handle.clone(), provisioning.clone(), @@ -360,7 +360,7 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. provisioner_tasks - .resume(|job| job_queue_ctx.push(job)) + .resume() .await .context("failed to queue the unfinished provisioner tasks again")?; diff --git a/devolutions-gateway/tests/ai_analysis.rs b/devolutions-gateway/tests/ai_analysis.rs index d69889bf2..79e28edb6 100644 --- a/devolutions-gateway/tests/ai_analysis.rs +++ b/devolutions-gateway/tests/ai_analysis.rs @@ -107,7 +107,7 @@ impl Harness { ); state.provisioner_tasks = devolutions_gateway::provisioner_tasks::task_runner( Arc::clone(&task_store), - state.job_queue_handle.clone(), + Arc::new(state.job_queue_handle.clone()), state.conf_handle.clone(), state.recordings.clone(), state.provisioning.clone(), From c3eecadaec38cab94fb1f11fbb34122bd29f182b Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:03:51 -0400 Subject: [PATCH 05/21] 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 81bf14fc0..54a512ed6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3996,6 +3996,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 d0f681cec..c3c861d53 100644 --- a/crates/job-queue-libsql/Cargo.toml +++ b/crates/job-queue-libsql/Cargo.toml @@ -19,3 +19,6 @@ tracing = "0.1" typed-builder = "0.21" ulid = { version = "1.2", features = ["uuid"] } uuid = "1.23" + +[dev-dependencies] +tokio = { version = "1.52", features = ["macros", "rt"] } diff --git a/crates/job-queue-libsql/src/lib.rs b/crates/job-queue-libsql/src/lib.rs index 0efbfd23f..04f76b63d 100644 --- a/crates/job-queue-libsql/src/lib.rs +++ b/crates/job-queue-libsql/src/lib.rs @@ -402,6 +402,27 @@ impl JobQueue for LibSqlJobQueue { Ok(Some(scheduled_for)) } + + async fn job_definitions(&self, name: &str) -> anyhow::Result> { + let sql_query = "SELECT json(def) FROM job_queue WHERE name = :name AND failed_attempts < :max_attempts"; + + let params = ((":name", name), (":max_attempts", self.max_attempts)); + + trace!(%sql_query, ?params, "Listing job definitions"); + + let mut rows = self + .conn + .query(sql_query, params) + .await + .context("failed to execute SQL query")?; + + let mut definitions = Vec::new(); + while let Some(row) = rows.next().await.context("failed to read the row")? { + definitions.push(row.get::(0).context("failed to read the job definition")?); + } + + Ok(definitions) + } } // Typically, migrations should not be modified once released, and we should only be appending to this list. @@ -425,3 +446,55 @@ const MIGRATIONS: &[&str] = &[ CREATE INDEX idx_scheduled_for ON job_queue(scheduled_for);", ]; + +#[cfg(test)] +mod tests { + use job_queue::{DynJob, Job, JobQueue as _, RunnerWaker}; + + use super::*; + + struct Named { + name: &'static str, + def: &'static str, + } + + #[async_trait] + impl Job for Named { + fn name(&self) -> &str { + self.name + } + + fn write_json(&self) -> anyhow::Result { + Ok(self.def.to_owned()) + } + + async fn run(&mut self) -> anyhow::Result<()> { + Ok(()) + } + } + + #[tokio::test] + async fn job_definitions_lists_the_stored_jobs_of_one_name() { + let conn = libsql::Builder::new_local(":memory:") + .build() + .await + .expect("build database") + .connect() + .expect("connect"); + let queue = LibSqlJobQueue::builder() + .runner_waker(RunnerWaker::new(|| {})) + .conn(conn) + .build(); + queue.setup().await.expect("setup"); + + for (name, def) in [("a", r#"{"n":1}"#), ("b", r#"{"n":2}"#), ("a", r#"{"n":3}"#)] { + let job: DynJob = Box::new(Named { name, def }); + queue.push_job(&job, None).await.expect("push"); + } + + let mut definitions = queue.job_definitions("a").await.expect("list"); + definitions.sort(); + + assert_eq!(definitions, [r#"{"n":1}"#, r#"{"n":3}"#]); + } +} diff --git a/crates/job-queue/src/lib.rs b/crates/job-queue/src/lib.rs index 93232bd86..5521da528 100644 --- a/crates/job-queue/src/lib.rs +++ b/crates/job-queue/src/lib.rs @@ -62,6 +62,9 @@ pub trait JobQueue: Send + Sync { /// Retrieves the closest future scheduled date async fn next_scheduled_date(&self) -> anyhow::Result>; + + /// Lists the definitions of the stored jobs named `name` that may still run, in no particular order. + async fn job_definitions(&self, name: &str) -> anyhow::Result>; } pub struct JobCtx { diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 60db0a367..42e4f3c65 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -106,6 +106,10 @@ impl ProvisionerTaskJobQueue for DurableJobQueue { let job: DynJob = Box::new(job); self.0.push_job(&job, None).await } + + async fn stored_jobs(&self) -> anyhow::Result> { + self.0.job_definitions(ProvisionerTaskJob::NAME).await + } } impl JobQueueHandle { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index c3b848d39..9e7b08968 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -35,6 +35,9 @@ use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; pub trait ProvisionerTaskJobQueue: Send + Sync { /// Queues `job`, and returns once it is stored, so a record never waits for a job that was lost. async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()>; + + /// Lists the JSON definitions of the provisioner task jobs already stored that may still run. + async fn stored_jobs(&self) -> anyhow::Result>; } pub type DynProvisionerTaskJobQueue = Arc; @@ -45,6 +48,11 @@ impl ProvisionerTaskJobQueue for JobQueueHandle { async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { self.enqueue(job).await } + + /// The channel cannot tell what is stored: every unfinished record looks like it has no job. + async fn stored_jobs(&self) -> anyhow::Result> { + Ok(Vec::new()) + } } use crate::token::AccessScope; @@ -256,15 +264,34 @@ impl ProvisionerTaskRunner { Ok(outcome) } - /// Hands every unfinished record to a new job: the jobs queued before stop at their next attempt. + /// Hands every unfinished record that has no job left to a new job, such as a record created just before a crash. /// - /// Call it at startup, before the job runner claims any job, so no earlier job is running an attempt meanwhile. A - /// record whose new job cannot be queued is failed, since no job is left to run it. + /// A record whose job is still stored keeps it. Call it at startup, before the job runner claims any job, so no job + /// runs an attempt meanwhile. A record whose new job cannot be queued is failed, since no job is left to run it. pub async fn resume(&self) -> anyhow::Result<()> { let now = OffsetDateTime::now_utc(); let records = self.inner.store.list_unfinished(now).await?; + let stored = self + .inner + .job_queue + .stored_jobs() + .await + .context("list the stored provisioner task jobs")? + .iter() + .filter_map(|json| serde_json::from_str::(json).ok()) + .map(|def| (def.task_id, def.job_token)) + .collect::>(); + for record in records { + if record + .job_token + .is_some_and(|job_token| stored.contains(&(record.id, job_token))) + { + debug!(task.id = %record.id, task.kind = %record.kind, "Task still has its job"); + continue; + } + let job_token = Uuid::new_v4(); if !self.inner.store.replace_job_token(record.id, job_token, now).await? { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 6cbd047ae..67e80b2d1 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -222,6 +222,75 @@ async fn a_task_that_cannot_be_queued_again_fails() { assert_eq!(*harness.ended.lock(), [id]); } +/// A job queue that remembers the jobs it stores, as the job queue database does. +#[derive(Default)] +struct StoredJobs(Mutex>); + +#[async_trait] +impl ProvisionerTaskJobQueue for StoredJobs { + async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { + self.0.lock().push(job_queue::Job::write_json(&job)?); + Ok(()) + } + + async fn stored_jobs(&self) -> anyhow::Result> { + Ok(self.0.lock().clone()) + } +} + +#[tokio::test] +async fn resume_queues_only_the_tasks_left_without_a_job() { + let conn = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open gateway database") + .connect() + .await + .expect("connect"); + let store = Arc::new( + provisioner_task_libsql::LibSqlProvisionerTaskStore::open(conn) + .await + .expect("open task store"), + ); + let queue = Arc::new(StoredJobs::default()); + let runner = ProvisionerTaskRunner::builder(store, Arc::clone(&queue) as DynProvisionerTaskJobQueue) + .register(Echo { + outcome: |value| Ok(json!(value)), + runs: Arc::default(), + ended: Arc::default(), + }) + .build(); + let create = |id: Uuid| runner.create::(id, format!("target-{id}"), &EchoParams { value: 1 }, |_| Ok(())); + let with_job = Uuid::new_v4(); + let without_job = Uuid::new_v4(); + create(with_job).await.expect("create"); + create(without_job).await.expect("create"); + let token_with_job = runner + .get(with_job, OffsetDateTime::now_utc()) + .await + .expect("read") + .expect("record") + .job_token; + + // The job of the second record is lost, as when Gateway stops right after creating the record. + queue.0.lock().retain(|json| json.contains(&with_job.to_string())); + + runner.resume().await.expect("resume"); + + let stored = queue.0.lock().clone(); + assert_eq!(stored.len(), 2, "one new job, for the record without one: {stored:?}"); + assert!(stored[1].contains(&without_job.to_string())); + assert_eq!( + runner + .get(with_job, OffsetDateTime::now_utc()) + .await + .expect("read") + .expect("record") + .job_token, + token_with_job, + "a record with its job keeps it" + ); +} + #[tokio::test] async fn a_panicking_attempt_fails_the_task() { let mut harness = Harness::new(|_| panic!("the attempt panics")).await; From c13af1918976244a5ed09980fd2f25b194840b01 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:31:15 -0400 Subject: [PATCH 06/21] docs(dgw): add the provisioner tasks intent --- .../src/provisioner_tasks/INTENT.md | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 devolutions-gateway/src/provisioner_tasks/INTENT.md diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md new file mode 100644 index 000000000..e96647acd --- /dev/null +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -0,0 +1,25 @@ +# Background +Provisioner tasks are tasks issued for Gateway to execute. Tasks's state are pesisted in `gateway.db`, see more at `../crates/provisioner-task-libsql/README.md`. + + +# Invariants + +## Management +`ProvisionerTaskRunner` MUST be the only place touches both the Task records and the job queue. +A provisioner task only defines it's work: what one attempt does, and what to drop once its task ends. +The task record is the source of truth; the `ProvisionerTaskJob` only wakes it up and holds no state. +One task at the same time only has one job to execute task. + + +## Naming +`ProvisionerTask` is the work itself; `ProvisionerTaskRecord` is what is stored about it. + +## Secrets +Secrets a Task needs, such as an API key, never reach the disk: not `gateway.db`, not the job, not the logs. +They are kept encrypted in memory, in the provisioning store, until the Task ends or reaches its deadline. +After a Gateway restart they are gone, and the Tasks that needed them fail; this is accepted. + +## AI analysis +Only a finished recording is analyzed. +A session that mixes terminal and video recordings cannot be analyzed. +Two analyses of the same session never run at the same time. \ No newline at end of file From 0c6388d30b1ab5bd54f2c55ca10bd5bdc1bf73a9 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:40:42 -0400 Subject: [PATCH 07/21] docs(dgw): point the provisioner tasks intent to the task store crate --- devolutions-gateway/src/provisioner_tasks/INTENT.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index e96647acd..b7e8f1245 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -1,5 +1,5 @@ # Background -Provisioner tasks are tasks issued for Gateway to execute. Tasks's state are pesisted in `gateway.db`, see more at `../crates/provisioner-task-libsql/README.md`. +Provisioner tasks are tasks issued for Gateway to execute. Tasks's state are pesisted in `gateway.db`, see more at `../../../crates/provisioner-task-libsql/`. # Invariants From 995e70c3438667943537032dcf80a606fdd44461 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:46:09 -0400 Subject: [PATCH 08/21] docs(dgw): add the jrec intent --- devolutions-gateway/src/api/jrec.intent.md | 10 ++++++++++ 1 file changed, 10 insertions(+) create mode 100644 devolutions-gateway/src/api/jrec.intent.md diff --git a/devolutions-gateway/src/api/jrec.intent.md b/devolutions-gateway/src/api/jrec.intent.md new file mode 100644 index 000000000..e65cc4955 --- /dev/null +++ b/devolutions-gateway/src/api/jrec.intent.md @@ -0,0 +1,10 @@ +# Background + +Jrec is the endpoint to push, streaming, read recordings and its associated files. + + +# AI analysis + +An AI analysis is a task initiated by the provisioner. +It requires the scope token with access `gateway.tasks.recording-ai-analysis.start` to access the url `POST /jet/jrec/{session_id}/ai-analysis`; Naturally in order to read the task's state, it requires the scope token with access `gateway.tasks.recording-ai-analysis.read`. +The analyzed output belongs to the recording reading scope, however, the associated metadata, including but not limited to error messages, token usages etc. belong to the task reading scope. \ No newline at end of file From 42292bd7e88902382e4af6caaaa4bfad05cc5af2 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:57:41 -0400 Subject: [PATCH 09/21] docs(dgw): expand the provisioner tasks intent --- devolutions-gateway/src/provisioner_tasks/INTENT.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/INTENT.md b/devolutions-gateway/src/provisioner_tasks/INTENT.md index b7e8f1245..346de3c4b 100644 --- a/devolutions-gateway/src/provisioner_tasks/INTENT.md +++ b/devolutions-gateway/src/provisioner_tasks/INTENT.md @@ -5,11 +5,16 @@ Provisioner tasks are tasks issued for Gateway to execute. Tasks's state are pes # Invariants ## Management -`ProvisionerTaskRunner` MUST be the only place touches both the Task records and the job queue. -A provisioner task only defines it's work: what one attempt does, and what to drop once its task ends. +`ProvisionerTaskRunner` MUST be the only place touches both the Task records and the job queue. It is also the only place that manages how tasks are executed in relation to the job queue. + +A provisioner task only defines it's work: what one attempt does, and what to drop once its task ends. The coordination belongs to the `ProvisionerTaskRunner`. + The task record is the source of truth; the `ProvisionerTaskJob` only wakes it up and holds no state. + One task at the same time only has one job to execute task. +Jobs doesn't have state, tasks do. + ## Naming `ProvisionerTask` is the work itself; `ProvisionerTaskRecord` is what is stored about it. From 4367f209a3aac8bfe79eeafb36970335c5565f09 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:56 -0400 Subject: [PATCH 10/21] test(dgw): keep the DVLS 2022.3 scope snapshot frozen `as_of_v2022_3_0_0` lists what DVLS 2022.3 sends; later scopes such as `gateway.recording.delete` were never added there. The AI analysis scopes are covered by the `ai_analysis` tests. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/tests/dvls_compatibility.rs | 6 ------ 1 file changed, 6 deletions(-) diff --git a/devolutions-gateway/tests/dvls_compatibility.rs b/devolutions-gateway/tests/dvls_compatibility.rs index fa0d080cc..e727c8e01 100644 --- a/devolutions-gateway/tests/dvls_compatibility.rs +++ b/devolutions-gateway/tests/dvls_compatibility.rs @@ -301,10 +301,6 @@ mod as_of_v2022_3_0_0 { JrlRead, #[serde(rename = "gateway.config.write")] ConfigWrite, - #[serde(rename = "gateway.tasks.recording-ai-analysis.start")] - RecordingAiAnalysisStart, - #[serde(rename = "gateway.tasks.recording-ai-analysis.read")] - RecordingAiAnalysisRead, } fn dvls_access_scope() -> impl Strategy { @@ -316,8 +312,6 @@ mod as_of_v2022_3_0_0 { Just(DvlsAccessScope::JrlRead), Just(DvlsAccessScope::ConfigWrite), Just(DvlsAccessScope::ConfigWrite), - Just(DvlsAccessScope::RecordingAiAnalysisStart), - Just(DvlsAccessScope::RecordingAiAnalysisRead), ] } From 61fc788fa6471fb3b5c9a4fd842254f57f7959a1 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:56 -0400 Subject: [PATCH 11/21] fix(dgw): keep the AI provider URL out of Debug output `AiSettings` and `AiAccess` (and so `StartRequest`) derived `Debug`, which prints the base URL that must never reach the logs. Nothing logs them today; this keeps it that way. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/provisioner_tasks/ai.rs | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/ai.rs b/devolutions-gateway/src/provisioner_tasks/ai.rs index 6ae4bf015..089805ee6 100644 --- a/devolutions-gateway/src/provisioner_tasks/ai.rs +++ b/devolutions-gateway/src/provisioner_tasks/ai.rs @@ -43,7 +43,7 @@ impl From for Provider { } /// How to reach an AI provider, except the API key. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Clone, PartialEq, Eq)] pub struct AiSettings { pub provider: AiProvider, pub model: String, @@ -53,6 +53,17 @@ pub struct AiSettings { pub max_output_tokens: Option, } +impl fmt::Debug for AiSettings { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AiSettings") + .field("provider", &self.provider) + .field("model", &self.model) + .field("base_url", &self.base_url.as_ref().map(|_| "")) + .field("max_output_tokens", &self.max_output_tokens) + .finish() + } +} + /// Why AI settings were refused before a Task was recorded. #[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] pub enum AiSettingsError { @@ -92,12 +103,21 @@ impl fmt::Display for ClientError { } /// What a Task needs to reach its AI provider and must never write to disk. -#[derive(Debug, Clone)] +#[derive(Clone)] pub(crate) struct AiAccess { pub(crate) api_key: SecretString, pub(crate) base_url: Option, } +impl fmt::Debug for AiAccess { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("AiAccess") + .field("api_key", &self.api_key) + .field("base_url", &self.base_url.as_ref().map(|_| "")) + .finish() + } +} + /// [`AiAccess`] as written into its task secret. #[derive(Serialize)] #[serde(rename_all = "camelCase")] From 7b2804bd3adcf3d14a0906605c2b4acce5b0041a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 12/21] 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 db6966f76..eb65f9d50 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -161,6 +161,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 df579485ca12d791e16a3ebb81af150fe0842ad5 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 13/21] fix(dgw): keep Gateway starting when provisioner tasks cannot be resumed A failure to list or hand over the unfinished Tasks stopped the whole service. Unfinished Tasks still fail at their deadline, so the error is logged and startup goes on; one Task that cannot get a new job no longer stops the others. Also follows the provisioner task store changes below this PR: `update_running` takes `now` and `job_token` is no longer optional. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../recording_ai_analysis/tests.rs | 8 ++---- .../src/provisioner_tasks/runner/mod.rs | 27 +++++++++++++------ .../src/provisioner_tasks/runner/tests.rs | 2 +- devolutions-gateway/src/service.rs | 11 +++++--- 4 files changed, 29 insertions(+), 19 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index fdcee3db4..b4157ac7e 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -120,8 +120,7 @@ impl Harness { .await .expect("read task") .expect("task exists") - .job_token - .expect("job token"); + .job_token; let json = json!({ "taskId": task_id, "jobToken": job_token }).to_string(); ProvisionerTaskJob::read_json(&json, self.state.provisioner_tasks.clone()).expect("valid job") @@ -331,10 +330,7 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { assert_eq!(job.name(), ProvisionerTaskJob::NAME); assert_eq!( job.write_json().expect("job JSON"), - format!( - r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, - task.job_token.expect("job token") - ) + format!(r#"{{"taskId":"{task_id}","jobToken":"{}"}}"#, task.job_token) ); assert!(schedule_for.is_none()); } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 9e7b08968..44476192c 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -29,6 +29,7 @@ use uuid::Uuid; use super::AttemptError; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; +use crate::token::AccessScope; /// Where the runner queues the jobs of its records. #[async_trait] @@ -54,7 +55,6 @@ impl ProvisionerTaskJobQueue for JobQueueHandle { Ok(Vec::new()) } } -use crate::token::AccessScope; /// A provisioner task: the work Gateway does for one kind of request, such as `recording-ai-analysis`. /// @@ -107,7 +107,11 @@ pub struct Attempt { impl Attempt { /// Replaces the payload of the running record, so whoever asks about it sees how the attempt goes. pub async fn report_progress(&self, payload: serde_json::Value) { - if let Err(error) = self.store.update_running(self.task_id, payload).await { + if let Err(error) = self + .store + .update_running(self.task_id, payload, OffsetDateTime::now_utc()) + .await + { warn!(task.id = %self.task_id, error = format!("{error:#}"), "Failed to record the Task progress"); } } @@ -284,18 +288,25 @@ impl ProvisionerTaskRunner { .collect::>(); for record in records { - if record - .job_token - .is_some_and(|job_token| stored.contains(&(record.id, job_token))) - { + if stored.contains(&(record.id, record.job_token)) { debug!(task.id = %record.id, task.kind = %record.kind, "Task still has its job"); continue; } let job_token = Uuid::new_v4(); - if !self.inner.store.replace_job_token(record.id, job_token, now).await? { - continue; + match self.inner.store.replace_job_token(record.id, job_token, now).await { + Ok(true) => {} + Ok(false) => continue, + Err(error) => { + error!( + task.id = %record.id, + task.kind = %record.kind, + error = format!("{error:#}"), + "Failed to hand the Task to a new job" + ); + continue; + } } match self.enqueue(record.id, job_token).await { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 67e80b2d1..4c4fc8d33 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -370,6 +370,6 @@ async fn the_job_is_read_back_from_its_json() { let read = ProvisionerTaskJob::read_json(&json, harness.runner.clone()).expect("valid job"); assert_eq!(read.task_id, id); - assert_eq!(Some(read.job_token), harness.task(id).await.job_token); + assert_eq!(read.job_token, harness.task(id).await.job_token); assert_eq!(job.name(), ProvisionerTaskJob::NAME); } diff --git a/devolutions-gateway/src/service.rs b/devolutions-gateway/src/service.rs index efc0ea707..6a33137f7 100644 --- a/devolutions-gateway/src/service.rs +++ b/devolutions-gateway/src/service.rs @@ -359,10 +359,13 @@ async fn spawn_tasks(conf_handle: ConfHandle) -> anyhow::Result { ); // Before the job runner starts, so no job queued earlier runs an attempt meanwhile. - provisioner_tasks - .resume() - .await - .context("failed to queue the unfinished provisioner tasks again")?; + // Unfinished Tasks still fail at their deadline, so this must not keep the other features from starting. + if let Err(error) = provisioner_tasks.resume().await { + error!( + error = format!("{error:#}"), + "Failed to queue the unfinished provisioner tasks again" + ); + } let state = DgwState { conf_handle: conf_handle.clone(), From c4e1032f0e5977269e6d878a250eb9c00b1f9767 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:56:57 -0400 Subject: [PATCH 14/21] 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.rs | 1 + devolutions-gateway/src/extract.rs | 2 ++ 2 files changed, 3 insertions(+) diff --git a/devolutions-gateway/src/api/jrec.rs b/devolutions-gateway/src/api/jrec.rs index ecc3bb5b0..a24dd34e0 100644 --- a/devolutions-gateway/src/api/jrec.rs +++ b/devolutions-gateway/src/api/jrec.rs @@ -621,6 +621,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 c42d3c14d..f6784516e 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 UpdateScope; From 089ecaf2b5c2b27099295bef47a9cc9510a384a6 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 09:48:45 -0400 Subject: [PATCH 15/21] 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.rs | 2 +- devolutions-gateway/tests/ai_analysis.rs | 16 ++++++++++++---- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/devolutions-gateway/src/api/jrec.rs b/devolutions-gateway/src/api/jrec.rs index a24dd34e0..63ae6b5cf 100644 --- a/devolutions-gateway/src/api/jrec.rs +++ b/devolutions-gateway/src/api/jrec.rs @@ -438,7 +438,7 @@ async fn delete_recording(recording_path: &Utf8Path) -> anyhow::Result<()> { /// Settings of an AI analysis #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] #[derive(Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] +#[serde(rename_all = "camelCase")] pub(crate) struct AiAnalysisRequest { /// Task ID, chosen by the caller so it can send the same request again safely. task_id: Uuid, diff --git a/devolutions-gateway/tests/ai_analysis.rs b/devolutions-gateway/tests/ai_analysis.rs index 79e28edb6..d471a42fb 100644 --- a/devolutions-gateway/tests/ai_analysis.rs +++ b/devolutions-gateway/tests/ai_analysis.rs @@ -326,15 +326,24 @@ async fn start_without_recording_returns_404() { assert_eq!(status, StatusCode::NOT_FOUND); } +#[tokio::test] +async fn start_ignores_unknown_fields() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let mut body = request(Uuid::new_v4(), "gpt-test"); + body["temperature"] = json!(0.5); + + let (status, created) = harness.post(session_id, &body).await; + + assert_eq!(status, StatusCode::ACCEPTED, "{created}"); +} + #[tokio::test] async fn invalid_bodies_and_settings_return_400() { let harness = Harness::new().await; let session_id = harness.write_session(); let task_id = Uuid::new_v4(); - let mut unknown_field = request(task_id, "gpt-test"); - unknown_field["temperature"] = json!(0.5); - let mut no_key = request(task_id, "gpt-test"); no_key.as_object_mut().unwrap().remove("apiKey"); @@ -351,7 +360,6 @@ async fn invalid_bodies_and_settings_return_400() { unknown_provider["provider"] = json!("skynet"); for (case, body) in [ - ("unknown field", unknown_field.to_string()), ("missing API key", no_key.to_string()), ("empty API key", empty_key.to_string()), ("empty model", request(task_id, " ").to_string()), From 38325bcd98ec57c95358da9e35408ab5c543a798 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:02:58 -0400 Subject: [PATCH 16/21] fix(dgw): start no AI analysis attempt that cannot finish before the deadline An attempt started near the 2-hour deadline could keep calling the AI for 30 more minutes, add its log to the recording, and still leave the Task failed as timed out. Attempts now start only during the first 2 hours, and the deadline is one attempt timeout after that, so an attempt started in time can finish and its result counts. The runner fails a Task that has no time left for a full attempt as timed out, without running it. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../recording_ai_analysis/mod.rs | 8 ++-- .../recording_ai_analysis/tests.rs | 31 +++++--------- .../src/provisioner_tasks/runner/mod.rs | 30 +++++++++++--- .../src/provisioner_tasks/runner/tests.rs | 41 ++++++++++++++++++- 4 files changed, 78 insertions(+), 32 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index dea325249..f29828ca6 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -37,8 +37,8 @@ use crate::token::AccessScope; pub const KIND: &str = "recording-ai-analysis"; -/// Time an AI analysis gets from its creation to its end, waiting and retries included. -pub const TASK_DEADLINE: time::Duration = time::Duration::hours(2); +/// Time from the creation of an AI analysis during which its attempts may start, waiting and retries included. +pub const ATTEMPT_START_WINDOW: time::Duration = time::Duration::hours(2); /// Time one attempt of an AI analysis gets. pub const ATTEMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60); @@ -260,8 +260,8 @@ impl ProvisionerTask for RecordingAiAnalysis { type Output = Analysis; - fn deadline(&self) -> time::Duration { - TASK_DEADLINE + fn attempt_start_window(&self) -> time::Duration { + ATTEMPT_START_WINDOW } fn attempt_timeout(&self) -> std::time::Duration { diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index b4157ac7e..f0916e583 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -321,7 +321,10 @@ async fn start_records_the_task_keeps_the_key_and_queues_a_job() { json!({ "provider": "openai", "model": "gpt-test", "maxOutputTokens": null }), "neither the key nor the base URL is recorded" ); - assert_eq!(task.deadline_at - task.created_at, TASK_DEADLINE); + assert_eq!( + task.deadline_at - task.created_at, + ATTEMPT_START_WINDOW + ATTEMPT_TIMEOUT + ); let access = harness.access(task_id).expect("key is kept"); assert_eq!(access.base_url.expect("base URL").as_str(), "https://ai.example/v1/"); @@ -694,10 +697,10 @@ async fn slow_attempt_is_retried() { } #[tokio::test] -async fn log_added_after_the_deadline_leaves_the_task_failed() { +async fn task_without_time_for_an_attempt_fails_without_calling_the_ai() { let harness = Harness::new().await; let session_id = harness.write_session(); - let (settings, _requests) = spawn_ai_provider_after(Duration::from_secs(4), |_| 0).await; + let (settings, requests) = spawn_ai_provider(|_| 0).await; let task_id = Uuid::new_v4(); let now = OffsetDateTime::now_utc(); harness @@ -706,32 +709,18 @@ async fn log_added_after_the_deadline_leaves_the_task_failed() { session_id, "openai-compatible", now, - now + time::Duration::seconds(2), + now + ATTEMPT_TIMEOUT - time::Duration::minutes(1), ) .await; harness.insert_key(task_id, settings.base_url, now + time::Duration::hours(1)); - let mut job = harness.job(task_id).await; - let job = tokio::spawn(async move { job.run().await }); - - // Reading the Task past its deadline, as a status request does, ends it while the AI is still answering. - tokio::time::sleep(Duration::from_secs(3)).await; - assert_eq!(harness.task(task_id).await.state, ProvisionerTaskState::Failed); - - job.await.expect("job task").expect("no retry"); + harness.job(task_id).await.run().await.expect("no retry"); let task = harness.task(task_id).await; assert_eq!(task.state, ProvisionerTaskState::Failed); assert_eq!(task.payload["reason"], "timed out"); - assert!( - harness - .recording_path - .join(session_id.to_string()) - .join("ai-analysis-0.slog") - .exists(), - "the log stays in the session" - ); - assert!(!harness.workspace(session_id).exists()); + assert_eq!(task.attempts, 0); + assert!(requests.lock().is_empty()); assert!(!harness.has_key(task_id)); } diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index 44476192c..a5121cc59 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -77,8 +77,10 @@ pub trait ProvisionerTask: Send + Sync + 'static { /// The payload of a `succeeded` record. type Output: Serialize + Send; - /// Time from the creation of a record to its end, waiting and retries included. - fn deadline(&self) -> time::Duration; + /// Time from the creation of a record during which its attempts may start, the first one and retries alike. + /// + /// The record's deadline comes one [`Self::attempt_timeout`] later, so an attempt started in time can finish. + fn attempt_start_window(&self) -> time::Duration; /// Time one attempt gets before it is cancelled and retried. fn attempt_timeout(&self) -> std::time::Duration; @@ -120,7 +122,7 @@ impl Attempt { /// [`ProvisionerTask`] with its types erased, so the runner can hold every task together. #[async_trait] trait AnyProvisionerTask: Send + Sync { - fn deadline(&self) -> time::Duration; + fn attempt_start_window(&self) -> time::Duration; fn attempt_timeout(&self) -> std::time::Duration; @@ -131,8 +133,8 @@ trait AnyProvisionerTask: Send + Sync { #[async_trait] impl AnyProvisionerTask for T { - fn deadline(&self) -> time::Duration { - ProvisionerTask::deadline(self) + fn attempt_start_window(&self) -> time::Duration { + ProvisionerTask::attempt_start_window(self) } fn attempt_timeout(&self) -> std::time::Duration { @@ -241,7 +243,7 @@ impl ProvisionerTaskRunner { kind: T::KIND.to_owned(), target, params, - deadline_at: now + task.deadline(), + deadline_at: now + task.attempt_start_window() + task.attempt_timeout(), job_token, }, now, @@ -355,6 +357,22 @@ impl ProvisionerTaskRunner { let store = &self.inner.store; let now = OffsetDateTime::now_utc(); + // An attempt that cannot finish before the deadline would do work the record never keeps, such as paid AI calls. + if let Some(record) = store.get(task_id, now).await? + && record.job_token == job_token + && !record.state.is_finished() + && let Some(task) = self.task(&record.kind) + && now + task.attempt_timeout() > record.deadline_at + { + warn!(task.id = %task_id, task.kind = %record.kind, attempts = record.attempts, "Task timed out before its next attempt"); + let payload = json!({ + "reason": "timed out", + "details": { "lastPayload": record.payload }, + "attempts": record.attempts, + }); + return self.fail(task.as_ref(), task_id, payload).await; + } + let record = match store.start_attempt(task_id, job_token, starting_payload(), now).await? { AttemptStart::Started(record) => record, AttemptStart::Finished(record) => { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 4c4fc8d33..92f4666ac 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -33,7 +33,7 @@ impl ProvisionerTask for Echo { type Output = serde_json::Value; - fn deadline(&self) -> time::Duration { + fn attempt_start_window(&self) -> time::Duration { time::Duration::hours(1) } @@ -360,6 +360,45 @@ async fn invalid_parameters_and_unknown_kinds_fail_the_task() { assert_eq!(harness.runs.load(Ordering::SeqCst), 0); } +#[tokio::test] +async fn an_attempt_starts_only_if_it_can_finish_before_the_deadline() { + let harness = Harness::new(|value| Ok(json!(value))).await; + let store = harness.runner.store(); + let now = OffsetDateTime::now_utc(); + + // The echo attempt timeout is 5 seconds. + for (left, state) in [(6, ProvisionerTaskState::Succeeded), (4, ProvisionerTaskState::Failed)] { + let id = Uuid::new_v4(); + let job_token = Uuid::new_v4(); + store + .create( + NewProvisionerTaskRecord { + id, + kind: KIND.to_owned(), + target: id.to_string(), + params: json!({ "value": 1 }), + deadline_at: now + time::Duration::seconds(left), + job_token, + }, + now, + ) + .await + .expect("create"); + + harness.runner.run_attempt(id, job_token).await.expect("no retry"); + + let task = harness.task(id).await; + assert_eq!(task.state, state, "{left} seconds left"); + if state == ProvisionerTaskState::Failed { + assert_eq!(task.payload["reason"], "timed out"); + assert_eq!(task.attempts, 0); + assert!(harness.ended.lock().contains(&id)); + } + } + + assert_eq!(harness.runs.load(Ordering::SeqCst), 1); +} + #[tokio::test] async fn the_job_is_read_back_from_its_json() { let mut harness = Harness::new(|value| Ok(json!(value))).await; From 62f707bdf32497a994dd4274d6e6fe5276072087 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:23:39 -0400 Subject: [PATCH 17/21] 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 eb65f9d50..6a8fcd86a 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -160,11 +160,6 @@ struct AiAnalysisRunningPayload { done: Option, /// Work the step has in all, in the same unit as `done`. total: Option, - /// Part of the step done, from 0 to 100. - #[schema(maximum = 100)] - percent: Option, - /// Short sentence for people following the analysis, such as `asking the AI: 2 of 5 chunks described`. - message: Option, } /// Payload of a `succeeded` AI analysis Task: the log it added to the session diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index f29828ca6..39ca2daf3 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -359,7 +359,7 @@ impl RecordingAiAnalysis { let analysis = loop { match events.recv().await { Some(RecordingAnalysisEvent::Progress(progress)) => { - // Reading sends an event every half second: the Task is only updated when the numbers it shows change. + // Reading sends an event every half second, so the Task is updated only when the stage's percent changes. let shown = (progress.stage(), progress.percent()); if recorded != Some(shown) { recorded = Some(shown); @@ -400,8 +400,6 @@ fn progress_payload(progress: &Progress) -> serde_json::Value { "step": progress.stage(), "done": done, "total": total, - "percent": progress.percent(), - "message": progress.to_string(), }) } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index f0916e583..d1845a5ab 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -501,8 +501,6 @@ async fn retry_resumes_at_the_first_chunk_without_a_checkpoint() { "step": "describing", "done": 1, "total": 3, - "percent": 33, - "message": "asking the AI: 1 of 3 chunks described", }) ); assert!(harness.has_key(task_id), "the key is kept for the retry"); From 3a8cbd3d0ac650e97680f91c9b377640523605a7 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 10:40:54 -0400 Subject: [PATCH 18/21] fix(dgw): wait for the remux of the session before the AI analysis reads it A recording is remuxed 60 seconds after it ends, and the remux replaces the video file. An analysis started right after the session ended read the video before it was remuxed. The analysis now waits until the job queue stores no remux job for a file of the session, checking every second, before it reads the recordings. The job queue lists remux jobs that are scheduled or running, including those stored before a Gateway restart. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/job_queue.rs | 4 +- .../src/provisioner_tasks/mod.rs | 11 ++- .../recording_ai_analysis/mod.rs | 42 +++++++++++- .../recording_ai_analysis/tests.rs | 67 ++++++++++++++++++- .../src/provisioner_tasks/runner/mod.rs | 8 +-- .../src/provisioner_tasks/runner/tests.rs | 3 +- devolutions-gateway/src/recording.rs | 10 ++- 7 files changed, 131 insertions(+), 14 deletions(-) diff --git a/devolutions-gateway/src/job_queue.rs b/devolutions-gateway/src/job_queue.rs index 42e4f3c65..9d5d38017 100644 --- a/devolutions-gateway/src/job_queue.rs +++ b/devolutions-gateway/src/job_queue.rs @@ -107,8 +107,8 @@ impl ProvisionerTaskJobQueue for DurableJobQueue { self.0.push_job(&job, None).await } - async fn stored_jobs(&self) -> anyhow::Result> { - self.0.job_definitions(ProvisionerTaskJob::NAME).await + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + self.0.job_definitions(name).await } } diff --git a/devolutions-gateway/src/provisioner_tasks/mod.rs b/devolutions-gateway/src/provisioner_tasks/mod.rs index ccb1d4938..0278157fd 100644 --- a/devolutions-gateway/src/provisioner_tasks/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/mod.rs @@ -5,6 +5,8 @@ pub mod ai; pub mod recording_ai_analysis; pub mod runner; +use std::sync::Arc; + use provisioner_task::DynProvisionerTaskStore; use self::recording_ai_analysis::RecordingAiAnalysis; @@ -21,8 +23,13 @@ pub fn task_runner( recordings: RecordingMessageSender, provisioning: ProvisioningStore, ) -> ProvisionerTaskRunner { - ProvisionerTaskRunner::builder(store, job_queue) - .register(RecordingAiAnalysis::new(conf_handle, recordings, provisioning)) + ProvisionerTaskRunner::builder(store, Arc::clone(&job_queue)) + .register(RecordingAiAnalysis::new( + conf_handle, + recordings, + provisioning, + job_queue, + )) .build() } diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index 39ca2daf3..c3b65141d 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -27,12 +27,12 @@ use uuid::Uuid; use super::AttemptError; use super::ai::{AiAccess, AiProvider, AiSettings, AiSettingsError, ClientError}; -use super::runner::{Attempt, ProvisionerTask}; +use super::runner::{Attempt, DynProvisionerTaskJobQueue, ProvisionerTask}; use crate::DgwState; use crate::artifacts::ArtifactKind; use crate::config::ConfHandle; use crate::provisioning::ProvisioningStore; -use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender}; +use crate::recording::{FinishedRecordingError, JrecManifest, RecordingMessageSender, RemuxJob, remux_pending}; use crate::token::AccessScope; pub const KIND: &str = "recording-ai-analysis"; @@ -43,6 +43,9 @@ pub const ATTEMPT_START_WINDOW: time::Duration = time::Duration::hours(2); /// Time one attempt of an AI analysis gets. pub const ATTEMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60); +/// Time between two checks for a remux of the session. +const REMUX_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1); + const KEY_LOST: &str = "the API key is no longer available, such as after a Gateway restart"; /// Tokens counted by the AI provider, as kept in Task results. @@ -226,15 +229,22 @@ pub struct RecordingAiAnalysis { conf_handle: ConfHandle, recordings: RecordingMessageSender, provisioning: ProvisioningStore, + job_queue: DynProvisionerTaskJobQueue, attempt_timeout: std::time::Duration, } impl RecordingAiAnalysis { - pub fn new(conf_handle: ConfHandle, recordings: RecordingMessageSender, provisioning: ProvisioningStore) -> Self { + pub fn new( + conf_handle: ConfHandle, + recordings: RecordingMessageSender, + provisioning: ProvisioningStore, + job_queue: DynProvisionerTaskJobQueue, + ) -> Self { Self { conf_handle, recordings, provisioning, + job_queue, attempt_timeout: ATTEMPT_TIMEOUT, } } @@ -345,6 +355,8 @@ impl RecordingAiAnalysis { analysis_manifest(&recording.manifest).map_err(|error| AttemptError::Permanent(error.to_string()))?; let manifest_path = recording.manifest_path(); + self.wait_for_remux(attempt, &recording.dir).await?; + let mut request = client .analyze_recording(&manifest, &manifest_path) .cancellation(attempt.cancel.clone()); @@ -392,6 +404,30 @@ impl RecordingAiAnalysis { } } +impl RecordingAiAnalysis { + /// Waits until no remux job of the session is stored, since a remux replaces a video that would be read. + async fn wait_for_remux(&self, attempt: &Attempt, dir: &Utf8Path) -> Result<(), AttemptError> { + loop { + let jobs = self + .job_queue + .stored_jobs(RemuxJob::NAME) + .await + .map_err(|error| AttemptError::Transient(format!("failed to list the remux jobs: {error:#}")))?; + + if !remux_pending(&jobs, dir) { + return Ok(()); + } + + debug!(task.id = %attempt.task_id, %dir, "Waiting for the remux of the session"); + + tokio::select! { + () = tokio::time::sleep(REMUX_POLL_INTERVAL) => {} + () = attempt.cancel.cancelled() => return Err(AttemptError::Cancelled), + } + } + } +} + /// The payload of a running AI analysis Task: its progress, for whoever asks about it. fn progress_payload(progress: &Progress) -> serde_json::Value { let (done, total) = progress.done_and_total().unzip(); diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index d1845a5ab..702989741 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -10,8 +10,9 @@ use parking_lot::Mutex; use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; +use crate::job_queue::JobQueueHandle; use crate::job_queue::{JobMessage, JobQueueReceiver, MAX_ATTEMPTS}; -use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskRunner}; +use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles}; @@ -674,6 +675,7 @@ async fn slow_attempt_is_retried() { harness.state.conf_handle.clone(), harness.state.recordings.clone(), harness.state.provisioning.clone(), + Arc::new(harness.state.job_queue_handle.clone()), ); task.attempt_timeout = Duration::from_millis(500); harness.state.provisioner_tasks = ProvisionerTaskRunner::builder( @@ -789,6 +791,69 @@ async fn finished_task_does_not_run_again() { ); } +/// The job queue channel, plus remux jobs stored as the job queue database would list them. +struct WithRemuxJobs { + queue: JobQueueHandle, + remux_jobs: Arc>>, +} + +#[async_trait] +impl ProvisionerTaskJobQueue for WithRemuxJobs { + async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()> { + self.queue.push(job).await + } + + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + Ok(if name == RemuxJob::NAME { + self.remux_jobs.lock().clone() + } else { + Vec::new() + }) + } +} + +#[tokio::test] +async fn analysis_waits_for_the_remux_of_the_session() { + let mut harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, requests) = spawn_ai_provider(|_| 0).await; + let task_id = harness.start_created(session_id, settings).await; + + let dir = harness.recording_path.join(session_id.to_string()); + let other_dir = harness.recording_path.join(Uuid::new_v4().to_string()); + let remux_jobs = Arc::new(Mutex::new(vec![ + json!({ "input_path": dir.join("recording-1.webm") }).to_string(), + json!({ "input_path": other_dir.join("recording-0.webm") }).to_string(), + ])); + let queue: DynProvisionerTaskJobQueue = Arc::new(WithRemuxJobs { + queue: harness.state.job_queue_handle.clone(), + remux_jobs: Arc::clone(&remux_jobs), + }); + harness.state.provisioner_tasks = + ProvisionerTaskRunner::builder(Arc::clone(harness.state.provisioner_tasks.store()), Arc::clone(&queue)) + .register(RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.state.provisioning.clone(), + queue, + )) + .build(); + + let mut job = harness.job(task_id).await; + let job = tokio::spawn(async move { job.run().await }); + + tokio::time::sleep(REMUX_POLL_INTERVAL * 2).await; + assert!(!job.is_finished(), "the analysis waits"); + assert!(requests.lock().is_empty(), "the AI is not asked yet"); + + remux_jobs.lock().remove(0); + job.await.expect("job task").expect("no retry"); + + let task = harness.task(task_id).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert!(!requests.lock().is_empty()); +} + #[tokio::test] async fn analyses_of_one_session_run_one_at_a_time() { let session_id = Uuid::new_v4(); diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index a5121cc59..f1135f90f 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -37,8 +37,8 @@ pub trait ProvisionerTaskJobQueue: Send + Sync { /// Queues `job`, and returns once it is stored, so a record never waits for a job that was lost. async fn push(&self, job: ProvisionerTaskJob) -> anyhow::Result<()>; - /// Lists the JSON definitions of the provisioner task jobs already stored that may still run. - async fn stored_jobs(&self) -> anyhow::Result>; + /// Lists the JSON definitions of the stored jobs named `name` that may still run. + async fn stored_jobs(&self, name: &str) -> anyhow::Result>; } pub type DynProvisionerTaskJobQueue = Arc; @@ -51,7 +51,7 @@ impl ProvisionerTaskJobQueue for JobQueueHandle { } /// The channel cannot tell what is stored: every unfinished record looks like it has no job. - async fn stored_jobs(&self) -> anyhow::Result> { + async fn stored_jobs(&self, _name: &str) -> anyhow::Result> { Ok(Vec::new()) } } @@ -281,7 +281,7 @@ impl ProvisionerTaskRunner { let stored = self .inner .job_queue - .stored_jobs() + .stored_jobs(ProvisionerTaskJob::NAME) .await .context("list the stored provisioner task jobs")? .iter() diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index 92f4666ac..e962b99a6 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -233,7 +233,8 @@ impl ProvisionerTaskJobQueue for StoredJobs { Ok(()) } - async fn stored_jobs(&self) -> anyhow::Result> { + async fn stored_jobs(&self, name: &str) -> anyhow::Result> { + assert_eq!(name, ProvisionerTaskJob::NAME); Ok(self.0.lock().clone()) } } diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 5c5814230..172ab0880 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -8,7 +8,7 @@ use std::time::Duration; use anyhow::Context as _; use async_trait::async_trait; -use camino::Utf8PathBuf; +use camino::{Utf8Path, Utf8PathBuf}; use devolutions_gateway_task::{ShutdownSignal, Task}; use futures::future::Either; use parking_lot::Mutex; @@ -1127,6 +1127,14 @@ impl RemuxJob { pub const NAME: &'static str = "remux"; } +/// Whether one of the remux job `definitions` is for a recording file in `dir`. +pub(crate) fn remux_pending(definitions: &[String], dir: &Utf8Path) -> bool { + definitions + .iter() + .filter_map(|json| serde_json::from_str::(json).ok()) + .any(|job| job.input_path.parent() == Some(dir)) +} + #[async_trait] impl job_queue::Job for RemuxJob { fn name(&self) -> &str { From a6d4af9cffcbc7bd2e44632be3cbd19677880f58 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:13:21 -0400 Subject: [PATCH 19/21] fix(dgw): tie the AI analysis working files to their Task - The working files of a session were reused by any later Task with the same model, and the late job of a Task that ended without it, such as one failed past its deadline on a status read, could delete the files of a newer Task, which then paid for every chunk again. The analysis now names its Task as the owner of the files: the attempts of one Task resume from them, a new Task starts fresh, and a Task's cleanup leaves another Task's files in place. - An attempt starts only with one minute to spare past its timeout: it ends a little later while it stops, and stored deadlines lose their fraction of a second. This also keeps the runner test clear of rounding. - While describing, progress is written whenever another chunk is done, not when the percent changes, since the payload no longer has a percent. - `deadlineAt` says a Task can time out earlier, and the failure reasons are listed in full. Co-Authored-By: Claude Opus 5.5 (1M context) --- devolutions-gateway/src/api/tasks.rs | 3 ++- devolutions-gateway/src/openapi.rs | 2 +- .../recording_ai_analysis/mod.rs | 13 ++++++--- .../recording_ai_analysis/tests.rs | 27 +++++++++++++++++++ .../src/provisioner_tasks/runner/mod.rs | 6 ++++- .../src/provisioner_tasks/runner/tests.rs | 10 ++++--- 6 files changed, 51 insertions(+), 10 deletions(-) diff --git a/devolutions-gateway/src/api/tasks.rs b/devolutions-gateway/src/api/tasks.rs index c82e99076..d0f2b9504 100644 --- a/devolutions-gateway/src/api/tasks.rs +++ b/devolutions-gateway/src/api/tasks.rs @@ -38,7 +38,8 @@ pub(crate) struct TaskInfo { /// When it succeeded or failed. #[serde(with = "time::serde::rfc3339::option")] finished_at: Option, - /// When an unfinished Task becomes failed, waiting and retries included. + /// Latest time an unfinished Task becomes failed, waiting and retries included; it fails as `timed out` earlier + /// once too little time is left for another attempt. #[serde(with = "time::serde::rfc3339")] deadline_at: OffsetDateTime, } diff --git a/devolutions-gateway/src/openapi.rs b/devolutions-gateway/src/openapi.rs index 6a8fcd86a..55f8e69ca 100644 --- a/devolutions-gateway/src/openapi.rs +++ b/devolutions-gateway/src/openapi.rs @@ -191,7 +191,7 @@ struct AiAnalysisUsage { #[derive(Serialize, utoipa::ToSchema)] #[serde(rename_all = "camelCase")] struct AiAnalysisFailedPayload { - /// Kind of failure, such as `permanent error`, `attempts exhausted`, `key lost` or `timed out`. + /// Kind of failure: `permanent error`, `attempts exhausted`, `key lost`, `timed out`, `invalid task` or `not queued`. reason: String, /// What went wrong: a text, except for `timed out`, where it is `{ "lastPayload": }`. details: serde_json::Value, diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index c3b65141d..f536510ff 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -319,7 +319,7 @@ impl ProvisionerTask for RecordingAiAnalysis { let manifest_path = self.manifest_path(session_id); let _session = lock_session(session_id).await; - if let Err(error) = recording_analysis::discard(&manifest_path).await { + if let Err(error) = recording_analysis::discard(&manifest_path, Some(&record.id.to_string())).await { warn!(task.id = %record.id, session.id = %session_id, %error, "Failed to remove the AI analysis working files"); } } @@ -359,6 +359,7 @@ impl RecordingAiAnalysis { let mut request = client .analyze_recording(&manifest, &manifest_path) + .owner(attempt.task_id.to_string()) .cancellation(attempt.cancel.clone()); if let Some(max_output_tokens) = settings.max_output_tokens { @@ -366,13 +367,17 @@ impl RecordingAiAnalysis { } let mut events = request.start(); - let mut recorded: Option<(&'static str, u8)> = None; + let mut recorded: Option<(&'static str, u64)> = None; let analysis = loop { match events.recv().await { Some(RecordingAnalysisEvent::Progress(progress)) => { - // Reading sends an event every half second, so the Task is updated only when the stage's percent changes. - let shown = (progress.stage(), progress.percent()); + // Reading sends an event every half second, so while reading the Task is updated only when the percent + // changes. + let shown = match progress { + Progress::Reading { .. } => (progress.stage(), u64::from(progress.percent())), + _ => (progress.stage(), progress.done_and_total().map_or(0, |(done, _)| done)), + }; if recorded != Some(shown) { recorded = Some(shown); attempt.report_progress(progress_payload(&progress)).await; diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 702989741..85fa41649 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -892,6 +892,33 @@ async fn working_files_left_by_an_earlier_task_go_with_the_task_that_ends() { assert!(!harness.workspace(session_id).exists()); assert!(!log.exists()); } +#[tokio::test] +async fn an_ended_task_leaves_the_working_files_of_another_task() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 0).await; + let task_id = harness.start_created(session_id, settings).await; + let workspace = harness.workspace(session_id); + std::fs::create_dir_all(&workspace).expect("working files of a newer Task"); + std::fs::write( + workspace.join("analysis.json"), + json!({ "owner": Uuid::new_v4().to_string(), "promptVersion": "p", "model": "m", "chunkLen": 1 }).to_string(), + ) + .expect("settings of a newer Task"); + let record = harness.task(task_id).await; + let task = RecordingAiAnalysis::new( + harness.state.conf_handle.clone(), + harness.state.recordings.clone(), + harness.state.provisioning.clone(), + Arc::new(harness.state.job_queue_handle.clone()), + ); + + task.on_end(&record).await; + + assert!(workspace.exists(), "the newer Task still resumes from them"); + assert!(!harness.has_key(task_id)); +} + #[test] fn unknown_usage_is_left_out_of_the_result() { let analysis = Analysis { diff --git a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs index f1135f90f..2b7e45429 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/mod.rs @@ -31,6 +31,10 @@ use super::AttemptError; use crate::job_queue::{JobQueueHandle, MAX_ATTEMPTS}; use crate::token::AccessScope; +/// Time an attempt may need past its timeout: it ends a little later while it stops, and a stored deadline loses its +/// fraction of a second. +pub(crate) const ATTEMPT_END_GRACE: std::time::Duration = std::time::Duration::from_secs(60); + /// Where the runner queues the jobs of its records. #[async_trait] pub trait ProvisionerTaskJobQueue: Send + Sync { @@ -362,7 +366,7 @@ impl ProvisionerTaskRunner { && record.job_token == job_token && !record.state.is_finished() && let Some(task) = self.task(&record.kind) - && now + task.attempt_timeout() > record.deadline_at + && now + task.attempt_timeout() + ATTEMPT_END_GRACE > record.deadline_at { warn!(task.id = %task_id, task.kind = %record.kind, attempts = record.attempts, "Task timed out before its next attempt"); let payload = json!({ diff --git a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs index e962b99a6..a69558afd 100644 --- a/devolutions-gateway/src/provisioner_tasks/runner/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/runner/tests.rs @@ -367,8 +367,12 @@ async fn an_attempt_starts_only_if_it_can_finish_before_the_deadline() { let store = harness.runner.store(); let now = OffsetDateTime::now_utc(); - // The echo attempt timeout is 5 seconds. - for (left, state) in [(6, ProvisionerTaskState::Succeeded), (4, ProvisionerTaskState::Failed)] { + // The echo attempt timeout is 5 seconds; 2 more or 2 less keep the test clear of whole-second rounding. + let needed = time::Duration::seconds(5) + time::Duration::try_from(ATTEMPT_END_GRACE).expect("grace"); + for (left, state) in [ + (needed + time::Duration::seconds(2), ProvisionerTaskState::Succeeded), + (needed - time::Duration::seconds(2), ProvisionerTaskState::Failed), + ] { let id = Uuid::new_v4(); let job_token = Uuid::new_v4(); store @@ -378,7 +382,7 @@ async fn an_attempt_starts_only_if_it_can_finish_before_the_deadline() { kind: KIND.to_owned(), target: id.to_string(), params: json!({ "value": 1 }), - deadline_at: now + time::Duration::seconds(left), + deadline_at: now + left, job_token, }, now, From 18032301905540c30800daf3c8d54adf0c414260 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:56:31 -0400 Subject: [PATCH 20/21] 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.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.rs b/devolutions-gateway/src/api/jrec.rs index 63ae6b5cf..fc86b89c6 100644 --- a/devolutions-gateway/src/api/jrec.rs +++ b/devolutions-gateway/src/api/jrec.rs @@ -492,7 +492,8 @@ fn ai_analysis_conflict(code: AiAnalysisConflictCode, task_id: Option) -> /// The work runs in the background as a Task; follow it with `GET /jet/jrec/{session_id}/ai-analysis/{task_id}`, /// which the `Location` header of a 202 answer names. /// -/// Sending the same request again returns the same Task and starts nothing new. +/// Sending the same request again returns the same Task and starts nothing new. A new Task for the session starts from +/// scratch: it removes the earlier AI analysis of the session, even a succeeded one, so a session holds at most one. /// The API key and the base URL are kept in memory only, until the Task ends; a Gateway restart fails the unfinished /// Tasks. #[cfg_attr(feature = "openapi", utoipa::path( diff --git a/devolutions-gateway/src/artifacts.rs b/devolutions-gateway/src/artifacts.rs index a2ecd5fa8..d62f5d803 100644 --- a/devolutions-gateway/src/artifacts.rs +++ b/devolutions-gateway/src/artifacts.rs @@ -2,8 +2,8 @@ use serde::{Deserialize, Serialize}; use crate::token::RecordingFileType; -/// Non-recording artifacts, one list per [`ArtifactKind`]. Each list is append-only, like `files`: names -/// are derived from positions. +/// Non-recording artifacts, one list per [`ArtifactKind`]. Each list only grows, like `files`, until it is cleared +/// as a whole: names are derived from positions. #[derive(Debug, Clone, Default, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] pub(crate) struct JrecArtifacts { diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs index f536510ff..acc421195 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/mod.rs @@ -355,6 +355,21 @@ impl RecordingAiAnalysis { analysis_manifest(&recording.manifest).map_err(|error| AttemptError::Permanent(error.to_string()))?; let manifest_path = recording.manifest_path(); + // A new Task is a hard restart: nothing an earlier analysis of the session left stays, not even its result. + // Its later attempts are job retries, which resume from the working files. + if attempt.attempts == 1 { + let cleared = async { + self.recordings + .clear_artifacts(session_id, ArtifactKind::AiAnalysis) + .await?; + recording_analysis::discard(&manifest_path, None).await?; + anyhow::Ok(()) + }; + cleared.await.map_err(|error| { + AttemptError::Permanent(format!("failed to remove the earlier AI analysis: {error:#}")) + })?; + } + self.wait_for_remux(attempt, &recording.dir).await?; let mut request = client diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index 85fa41649..c098fb712 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -892,6 +892,59 @@ async fn working_files_left_by_an_earlier_task_go_with_the_task_that_ends() { assert!(!harness.workspace(session_id).exists()); assert!(!log.exists()); } +#[tokio::test] +async fn a_new_task_replaces_the_result_of_the_earlier_one() { + let harness = Harness::new().await; + let session_id = harness.write_session(); + let (settings, _requests) = spawn_ai_provider(|_| 0).await; + let first = harness.start_created(session_id, settings.clone()).await; + harness.run(first).await.expect("first analysis"); + + let second = harness.start_created(session_id, settings).await; + harness.run(second).await.expect("second analysis"); + + let task = harness.task(second).await; + assert_eq!(task.state, ProvisionerTaskState::Succeeded, "{:?}", task.payload); + assert_eq!(task.payload["fileName"], "ai-analysis-0.slog"); + let dir = harness.recording_path.join(session_id.to_string()); + let manifest: serde_json::Value = + serde_json::from_slice(&std::fs::read(dir.join("recording.json")).expect("manifest")).expect("manifest JSON"); + assert_eq!( + manifest["artifacts"]["ai-analysis"], + json!([{ "fileName": "ai-analysis-0.slog" }]), + "one analysis per recording" + ); + assert!(!dir.join("ai-analysis-1.slog").exists()); +} + +#[tokio::test] +async fn a_new_task_never_resumes_from_an_earlier_task() { + let harness = Harness::new().await; + // About 1 MB of transcript: three chunks. + let session_id = harness.write_long_session(10_000); + let (settings, requests) = spawn_ai_provider(|index| if index == 1 { 503 } else { 0 }).await; + let first = harness.start_created(session_id, settings.clone()).await; + assert!(harness.run(first).await.is_err(), "the outage is retried"); + assert!(harness.workspace(session_id).join("chunk-0000.json").exists()); + harness + .state + .provisioner_tasks + .store() + .fail(first, json!({}), OffsetDateTime::now_utc()) + .await + .expect("end the first Task"); + + let second = harness.start_created(session_id, settings).await; + harness.run(second).await.expect("second analysis"); + + assert_eq!(harness.task(second).await.state, ProvisionerTaskState::Succeeded); + assert_eq!( + requests.lock().len(), + 2 + 3, + "the new Task asks about every chunk again" + ); +} + #[tokio::test] async fn an_ended_task_leaves_the_working_files_of_another_task() { let harness = Harness::new().await; diff --git a/devolutions-gateway/src/recording.rs b/devolutions-gateway/src/recording.rs index 172ab0880..1207f1af6 100644 --- a/devolutions-gateway/src/recording.rs +++ b/devolutions-gateway/src/recording.rs @@ -382,6 +382,11 @@ enum RecordingManagerMessage { kind: ArtifactKind, channel: oneshot::Sender>, }, + ClearArtifacts { + id: Uuid, + kind: ArtifactKind, + channel: oneshot::Sender>, + }, GetFinished { id: Uuid, channel: oneshot::Sender>, @@ -448,6 +453,11 @@ impl fmt::Debug for RecordingManagerMessage { .field("id", id) .field("kind", kind) .finish_non_exhaustive(), + RecordingManagerMessage::ClearArtifacts { id, kind, channel: _ } => f + .debug_struct("ClearArtifacts") + .field("id", id) + .field("kind", kind) + .finish_non_exhaustive(), RecordingManagerMessage::GetFinished { id, channel: _ } => { f.debug_struct("GetFinished").field("id", id).finish_non_exhaustive() } @@ -514,6 +524,19 @@ impl RecordingMessageSender { rx.await.context("couldn't receive AddArtifact result")? } + /// Removes every artifact of `kind` from the session, both its file and its manifest entry. + /// + /// Fails when the session has no recording. + pub(crate) async fn clear_artifacts(&self, id: Uuid, kind: ArtifactKind) -> anyhow::Result<()> { + let (tx, rx) = oneshot::channel(); + self.channel + .send(RecordingManagerMessage::ClearArtifacts { id, kind, channel: tx }) + .await + .ok() + .context("couldn't send ClearArtifacts message")?; + rx.await.context("couldn't receive ClearArtifacts result")? + } + pub(crate) async fn get_finished(&self, id: Uuid) -> Result { let (tx, rx) = oneshot::channel(); self.channel @@ -904,6 +927,39 @@ impl RecordingManagerTask { Ok(artifact_path) } + async fn handle_clear_artifacts(&self, id: Uuid, kind: ArtifactKind) -> anyhow::Result<()> { + let recording_path = self.recordings_path.join(id.to_string()); + let manifest_path = recording_path.join("recording.json"); + + if !manifest_path.exists() { + anyhow::bail!("session {id} has no recording"); + } + + let mut manifest = JrecManifest::read_from_file(&manifest_path).context("read manifest from disk")?; + let artifacts = core::mem::take(manifest.artifacts.of_kind_mut(kind)); + + if artifacts.is_empty() { + return Ok(()); + } + + // The manifest goes first, so it never lists a file that is gone. + manifest + .save_to_file(&manifest_path) + .context("write manifest to disk")?; + + for artifact in artifacts { + let artifact_path = recording_path.join(&artifact.file_name); + match fs::remove_file(&artifact_path).await { + Err(error) if error.kind() != io::ErrorKind::NotFound => { + return Err(anyhow::Error::new(error).context(format!("remove {artifact_path}"))); + } + _ => {} + } + } + + Ok(()) + } + fn handle_get_finished(&self, id: Uuid) -> Result { // The transcript is built from the whole recording, so one still being pushed is not ready. if self.ongoing_recordings.contains_key(&id) { @@ -1031,6 +1087,9 @@ async fn recording_manager_task( RecordingManagerMessage::AddArtifact { id, kind, channel } => { let _ = channel.send(manager.handle_add_artifact(id, kind).await); }, + RecordingManagerMessage::ClearArtifacts { id, kind, channel } => { + let _ = channel.send(manager.handle_clear_artifacts(id, kind).await); + }, RecordingManagerMessage::GetFinished { id, channel } => { let _ = channel.send(manager.handle_get_finished(id)); } From 3e66e697ae91073edd80ac568b7f07fc250694cd Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 14:19:27 -0400 Subject: [PATCH 21/21] style(dgw): format the AI analysis task tests --- .../src/provisioner_tasks/recording_ai_analysis/tests.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs index c098fb712..3c91ad0d5 100644 --- a/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs +++ b/devolutions-gateway/src/provisioner_tasks/recording_ai_analysis/tests.rs @@ -10,8 +10,7 @@ use parking_lot::Mutex; use provisioner_task::{NewProvisionerTaskRecord, ProvisionerTaskState}; use super::*; -use crate::job_queue::JobQueueHandle; -use crate::job_queue::{JobMessage, JobQueueReceiver, MAX_ATTEMPTS}; +use crate::job_queue::{JobMessage, JobQueueHandle, JobQueueReceiver, MAX_ATTEMPTS}; use crate::provisioner_tasks::runner::{ProvisionerTaskJob, ProvisionerTaskJobQueue, ProvisionerTaskRunner}; use crate::recording::RecordingManagerTask; use crate::{DgwState, MockHandles};