task draining / filtering

Cargo.lock

@@ -12066,6 +12066,7 @@ dependencies = [
  "filesys",
  "futures 0.3.29",
  "image",
+ "indexmap 2.9.0",
  "libloading 0.8.6",
  "log",
  "memory_caching",

Cargo.toml

@@ -157,6 +157,7 @@ base64 = "0.22.1"
 chrono = { version = "0.4.40", features = ["serde"] }
 cookie = "0.18.1"
 hex-literal = "1.0.0"
+indexmap = "2.9.0"
 log = "0.4.27"
 once_cell = "1.21.3"
 rand = "0.9.0"

crates/desktop/tauri-artcraft/Cargo.toml

@@ -35,7 +35,7 @@ tokens = { workspace = true }
 
 # External
 anyhow = { workspace = true }
-base64 = "0.22.1"
+base64 = { workspace = true }
 bytes = "1.10"
 chrono = { workspace = true }
 const_format = "0.2.34"
@@ -43,11 +43,11 @@ directories = "6.0.0"
 dirs = "6.0.0"
 futures = "0.3"
 image = "0.25.5"
+indexmap = { workspace = true }
 log = { workspace = true }
 once_cell = { workspace = true }
 rand = { workspace = true }
-# NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
-reqwest = { workspace = true, features = ["multipart", "stream"], default-features = false }
+reqwest = { workspace = true, features = ["multipart", "stream"], default-features = false } # NB(bt,2025-03-10): default-features = false on Windows while I debug openssl / rustls
 serde = { workspace = true }
 serde_derive = { workspace = true }
 serde_json = { workspace = true }

crates/desktop/tauri-artcraft/src/state/sora/mod.rs

@@ -2,4 +2,4 @@ pub mod read_sora_credentials_from_disk;
 pub mod read_sora_credentials_from_disk_with_upgrade;
 pub mod sora_credential_holder;
 pub mod sora_credential_manager;
-mod sora_task_queue;
+pub mod sora_task_queue;

crates/desktop/tauri-artcraft/src/state/sora/sora_task_queue.rs

@@ -1,5 +1,8 @@
-use chrono::{Date, DateTime, Utc};
-use std::collections::HashMap;
+use std::collections::HashSet;
+use chrono::{DateTime, Utc};
+use errors::AnyhowResult;
+use indexmap::IndexMap;
+use openai_sora_client::requests::image_gen::image_gen_status::TaskId;
 use std::sync::{Arc, RwLock};
 
 #[derive(Clone)]
@@ -7,8 +10,222 @@ pub struct TaskDetails {
   pub enqueue_time: DateTime<Utc>
 }
 
+impl TaskDetails {
+  pub fn new() -> Self {
+    Self {
+      enqueue_time: Utc::now()
+    }
+  }
+}
+
+#[derive(Clone)]
+pub struct TaskIdAndDetails {
+  pub task_id: TaskId,
+  pub details: TaskDetails,
+}
+
 #[derive(Clone)]
 pub struct SoraTaskQueue {
-
-  queue: Arc<RwLock<HashMap<String, TaskDetails>>>
+  // Insertion-order map of task IDs to task details
+  queue: Arc<RwLock<IndexMap<TaskId, TaskDetails>>>
 }
+
+impl SoraTaskQueue {
+  pub fn new() -> Self {
+    Self {
+      queue: Arc::new(RwLock::new(IndexMap::new()))
+    }
+  }
+
+  /// Determine if queue is empty
+  /// This is an O(1) operation.
+  pub fn is_empty(&self) -> AnyhowResult<bool> {
+    match self.queue.read() {
+      Err(err) => {
+        Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
+      }
+      Ok(lock) => Ok(lock.is_empty())
+    }
+  }
+
+  /// Get length of set
+  /// This is an O(1) operation.
+  pub fn len(&self) -> AnyhowResult<usize> {
+    match self.queue.read() {
+      Err(err) => {
+        Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
+      }
+      Ok(lock) => Ok(lock.len())
+    }
+  }
+
+  /// Determine if the key is in the queue
+  /// This is an O(1) operation.
+  pub fn contains_key(&self, task_id: &TaskId) -> AnyhowResult<bool> {
+    match self.queue.read() {
+      Err(err) => {
+        Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
+      }
+      Ok(lock) => Ok(lock.contains_key(task_id))
+    }
+  }
+
+  /// Get the first enqueued task
+  /// This is an O(1) operation.
+  pub fn first(&self) -> AnyhowResult<Option<TaskIdAndDetails>> {
+    match self.queue.read() {
+      Err(err) => {
+        Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
+      }
+      Ok(lock) => match lock.first() {
+        None => Ok(None),
+        Some((task_id, task_details)) => Ok(Some(TaskIdAndDetails {
+          task_id: task_id.clone(),
+          details: task_details.clone(),
+        }))
+      }
+    }
+  }
+
+  /// Get the last enqueued task
+  /// This is an O(1) operation.
+  pub fn last(&self) -> AnyhowResult<Option<TaskIdAndDetails>> {
+    match self.queue.read() {
+      Err(err) => {
+        Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
+      }
+      Ok(lock) => match lock.last() {
+        None => Ok(None),
+        Some((task_id, task_details)) => Ok(Some(TaskIdAndDetails {
+          task_id: task_id.clone(),
+          details: task_details.clone(),
+        }))
+      }
+    }
+  }
+
+  /// Insert a new task. If a task already exists, we update its value.
+  /// This is an O(1) operation.
+  pub fn insert(&self, task_id: &TaskId) -> AnyhowResult<Option<TaskDetails>> {
+    match self.queue.write() {
+      Err(err) => {
+        Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err))
+      }
+      Ok(mut lock) => {
+        // NB: Probably fine to compute TaskDetails here for now. We may eventually expose this if we include additional info.
+        let task_details = TaskDetails::new();
+        Ok(lock.insert(task_id.clone(), task_details))
+      }
+    }
+  }
+
+  /// Remove a task by ID.
+  /// This is an O(n) operation.
+  pub fn remove(&self, task_id: &TaskId) -> AnyhowResult<Option<TaskDetails>> {
+    match self.queue.write() {
+      Err(err) => {
+        Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err))
+      }
+      Ok(mut lock) => {
+        Ok(lock.shift_remove(task_id))
+      }
+    }
+  }
+
+  /// Remove a list of tasks by ID.
+  /// This is an O(n) operation.
+  pub fn remove_list(&self, task_ids: &[&TaskId]) -> AnyhowResult<()> {
+    match self.queue.write() {
+      Err(err) => {
+        Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err))
+      }
+      Ok(mut lock) => {
+        // TODO: Cleaner way to do this?
+        let mut set : HashSet<&TaskId, _> = HashSet::with_capacity(task_ids.len());
+        set.extend(task_ids);
+        lock.retain(|k, v| !set.contains(k));
+        Ok(())
+      }
+    }
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use errors::AnyhowResult;
+  use openai_sora_client::requests::image_gen::image_gen_status::TaskId;
+  use crate::state::sora::sora_task_queue::SoraTaskQueue;
+
+  #[test]
+  fn remove_list() -> AnyhowResult<()> {
+    let task_queue = SoraTaskQueue::new();
+
+    assert!(task_queue.is_empty()?);
+    assert_eq!(task_queue.len()?, 0);
+
+    task_queue.insert(&TaskId("a".to_string()))?;
+    task_queue.insert(&TaskId("b".to_string()))?;
+    task_queue.insert(&TaskId("c".to_string()))?;
+    task_queue.insert(&TaskId("d".to_string()))?;
+    task_queue.insert(&TaskId("e".to_string()))?;
+    task_queue.insert(&TaskId("f".to_string()))?;
+
+    assert!(!task_queue.is_empty()?);
+    assert_eq!(task_queue.len()?, 6);
+
+    task_queue.insert(&TaskId("a".to_string()))?;
+    task_queue.insert(&TaskId("b".to_string()))?;
+    task_queue.insert(&TaskId("c".to_string()))?;
+    task_queue.insert(&TaskId("d".to_string()))?;
+    task_queue.insert(&TaskId("e".to_string()))?;
+    task_queue.insert(&TaskId("f".to_string()))?;
+
+    assert!(!task_queue.is_empty()?);
+    assert_eq!(task_queue.len()?, 6);
+
+    task_queue.remove_list(&[
+      &TaskId("a".to_string()),
+      &TaskId("c".to_string()),
+      &TaskId("e".to_string()),
+    ])?;
+
+    assert!(!task_queue.is_empty()?);
+    assert_eq!(task_queue.len()?, 3);
+
+    assert!(task_queue.contains_key(&TaskId("b".to_string()))?);
+    assert!(task_queue.contains_key(&TaskId("d".to_string()))?);
+    assert!(task_queue.contains_key(&TaskId("f".to_string()))?);
+
+    assert!(!task_queue.contains_key(&TaskId("a".to_string()))?);
+    assert!(!task_queue.contains_key(&TaskId("c".to_string()))?);
+    assert!(!task_queue.contains_key(&TaskId("e".to_string()))?);
+
+    task_queue.remove_list(&[
+      &TaskId("foo".to_string()),
+      &TaskId("bar".to_string()),
+      &TaskId("baz".to_string()),
+    ])?;
+
+    assert!(!task_queue.is_empty()?);
+    assert_eq!(task_queue.len()?, 3);
+
+    assert!(task_queue.contains_key(&TaskId("b".to_string()))?);
+    assert!(task_queue.contains_key(&TaskId("d".to_string()))?);
+    assert!(task_queue.contains_key(&TaskId("f".to_string()))?);
+
+    task_queue.remove_list(&[
+      &TaskId("b".to_string()),
+      &TaskId("d".to_string()),
+      &TaskId("f".to_string()),
+    ])?;
+
+    assert!(task_queue.is_empty()?);
+    assert_eq!(task_queue.len()?, 0);
+
+    assert!(!task_queue.contains_key(&TaskId("b".to_string()))?);
+    assert!(!task_queue.contains_key(&TaskId("d".to_string()))?);
+    assert!(!task_queue.contains_key(&TaskId("f".to_string()))?);
+
+    Ok(())
+  }
+}

crates/desktop/tauri-artcraft/src/threads/sora_task_polling_thread.rs

@@ -1,28 +1,44 @@
 use crate::state::app_dir::AppDataRoot;
 use crate::state::sora::sora_credential_manager::SoraCredentialManager;
+use crate::state::sora::sora_task_queue::SoraTaskQueue;
 use errors::AnyhowResult;
 use log::{error, info};
 use openai_sora_client::recipes::list_sora_task_status_with_session_auto_renew::{list_sora_task_status_with_session_auto_renew, StatusRequestArgs};
+use openai_sora_client::requests::image_gen::image_gen_status::{TaskId, TaskStatus};
 use tauri::AppHandle;
 
 pub async fn sora_session_login_thread(
   app_data_root: AppDataRoot,
-  sora_creds_manager: SoraCredentialManager
+  sora_creds_manager: SoraCredentialManager,
+  sora_task_queue: SoraTaskQueue,
 ) -> ! {
   loop {
-    if let Err(err) = polling_loop(&sora_creds_manager).await {
+    let res = polling_loop(&sora_creds_manager, &sora_task_queue).await;
+    if let Err(err) = res {
       error!("An error occurred: {:?}", err);
     }
     tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
   }
 }
 
-async fn polling_loop(sora_creds_manager: &SoraCredentialManager) -> AnyhowResult<()> {
+async fn polling_loop(
+  sora_creds_manager: &SoraCredentialManager,
+  sora_task_queue: &SoraTaskQueue,
+) -> AnyhowResult<()> {
   loop {
+    if sora_task_queue.is_empty()? {
+      // No need to poll if we don't have pending tasks.
+      tokio::time::sleep(std::time::Duration::from_millis(1_000)).await;
+      continue;
+    }
+
+    info!("Task queue has {} pending tasks.", sora_task_queue.len()?);
+
     let creds = sora_creds_manager.get_credentials_required()?;
 
     let (response, maybe_new_creds) = list_sora_task_status_with_session_auto_renew(StatusRequestArgs {
       limit: None,
+      // TODO: How can we use the task id to poll better? Our existing code doesn't seem to illuminate this.
       before: None,
       credentials: creds,
     }).await?;
@@ -32,13 +48,42 @@ async fn polling_loop(sora_creds_manager: &SoraCredentialManager) -> AnyhowResul
       sora_creds_manager.set_credentials(&new_creds)?;
     }
 
+    let mut succeeded_tasks = Vec::new();
+    let mut failed_tasks = Vec::new();
+
     // TODO: The cursoring logic likely needs to improve.
     for task in response.task_responses {
-      info!("Task ID: {}", task.id);
-      info!("Task Status: {:?}", task.status);
-      // Handle the task status as needed
+      let status = TaskStatus::from_str(&task.status);
+
+      match status {
+        TaskStatus::Succeeded => {
+          succeeded_tasks.push(task);
+        }
+        TaskStatus::Failed => {
+          failed_tasks.push(task);
+        }
+        TaskStatus::Queued => {}
+        TaskStatus::Running => {}
+        TaskStatus::Unknown(_) => {}
+      }
     }
 
-    tokio::time::sleep(std::time::Duration::from_millis(10_000)).await;
+    let failed_task_ids: Vec<&TaskId> = failed_tasks
+        .iter()
+        .map(|task| &task.id)
+        .collect();
+
+    sora_task_queue.remove_list(&failed_task_ids)?;
+
+    // TODO: Handle succeeded tasks.
+
+    let succeeded_task_ids : Vec<&TaskId> = succeeded_tasks
+        .iter()
+        .map(|task| &task.id)
+        .collect();
+
+    sora_task_queue.remove_list(&succeeded_task_ids)?;
+
+    tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
   }
 }