use super::core::{MAX_HISTORY_TASKS, TASK_POOL, Task, TaskPool, TaskStatus};
use super::executor::ExecutableTask;
use crate::workflow::WorkflowCallback;
use crate::{HippoxBatchResult, HippoxBoolResult, HippoxResult, HippoxStringResult, HippoxVoidResult};
use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
pub async fn create_task(task_type: String, input: String) -> HippoxStringResult {
let mut pool = TASK_POOL.write().await;
let task_id = pool.create_task(task_type, input);
HippoxResult::ok(task_id)
}
pub async fn create_task_with_executable(task_type: String, input: String, executable: Arc<dyn ExecutableTask>) -> HippoxStringResult {
let mut pool = TASK_POOL.write().await;
let task_id = pool.create_task_with_executable(task_type, input, executable);
HippoxResult::ok(task_id)
}
pub async fn get_task(task_id: &str) -> HippoxResult<Task> {
let pool = TASK_POOL.read().await;
match pool.get_task(task_id) {
Some(task) => HippoxResult::ok(task),
None => HippoxResult::system_error(format!("Task not found: {}", task_id)),
}
}
pub async fn get_task_status(task_id: &str) -> HippoxResult<TaskStatus> {
let pool = TASK_POOL.read().await;
match pool.get_task(task_id) {
Some(task) => HippoxResult::ok(task.status),
None => HippoxResult::system_error(format!("Task not found: {}", task_id)),
}
}
pub async fn update_task_status(task_id: &str, status: TaskStatus) -> HippoxBoolResult {
let mut pool = TASK_POOL.write().await;
let success = pool.update_task_status(task_id, status);
if success { HippoxResult::ok(true) } else { HippoxResult::system_error(format!("Failed to update task status: {}", task_id)) }
}
pub async fn cancel_task(task_id: &str) -> HippoxBoolResult {
let mut pool = TASK_POOL.write().await;
let success = pool.cancel_task(task_id);
if success { HippoxResult::ok(true) } else { HippoxResult::system_error(format!("Failed to cancel task: {}", task_id)) }
}
pub async fn pause_task(task_id: &str) -> HippoxBoolResult {
let mut pool = TASK_POOL.write().await;
let success = pool.pause_task(task_id);
if success { HippoxResult::ok(true) } else { HippoxResult::system_error(format!("Failed to pause task: {}", task_id)) }
}
pub async fn resume_task(task_id: &str) -> HippoxBoolResult {
let result = {
let mut pool = TASK_POOL.write().await;
pool.resume_task(task_id)
};
if result {
let pool = TASK_POOL.read().await;
if let Some(task) = pool.get_task(task_id) {
if let Some(executable) = task.get_executable() {
if let Some(callback) = executable.get_workflow_callback() {
callback.on_workflow_resumed(task_id, 0, 0).await;
}
}
}
HippoxResult::ok(true)
} else {
HippoxResult::system_error(format!("Failed to resume task: {}", task_id))
}
}
pub async fn retry_task(task_id: &str) -> HippoxBoolResult {
let mut pool = TASK_POOL.write().await;
let success = pool.retry_task(task_id);
if success { HippoxResult::ok(true) } else { HippoxResult::system_error(format!("Failed to retry task: {}", task_id)) }
}
pub async fn get_all_tasks(limit: Option<usize>) -> HippoxBatchResult {
let pool = TASK_POOL.read().await;
let tasks = pool.get_all_tasks(limit);
let task_ids: Vec<String> = tasks.into_iter().map(|t| t.id).collect();
HippoxResult::ok(task_ids)
}
pub async fn get_all_tasks_detailed(limit: Option<usize>) -> HippoxResult<Vec<Task>> {
let pool = TASK_POOL.read().await;
let tasks = pool.get_all_tasks(limit);
HippoxResult::ok(tasks)
}
pub async fn set_max_concurrent(max: usize) -> HippoxVoidResult {
let mut pool = TASK_POOL.write().await;
pool.set_max_concurrent(max);
HippoxResult::ok(())
}
pub async fn running_count() -> HippoxResult<usize> {
let pool = TASK_POOL.read().await;
HippoxResult::ok(pool.running_count())
}
pub async fn pending_count() -> HippoxResult<usize> {
let pool = TASK_POOL.read().await;
HippoxResult::ok(pool.pending_count())
}
pub async fn shutdown_task_pool() -> HippoxVoidResult {
let mut pool = TASK_POOL.write().await;
pool.shutdown();
HippoxResult::ok(())
}
pub async fn wait_task(task_id: &str) -> HippoxStringResult {
loop {
let status = match get_task_status(task_id).await {
HippoxResult { data: Some(s), .. } => s,
HippoxResult { error: Some(e), .. } => {
return HippoxResult::system_error(e);
}
_ => {
return HippoxResult::system_error(format!("Task not found: {}", task_id));
}
};
match status {
TaskStatus::Completed => {
match get_task(task_id).await {
HippoxResult { data: Some(task), .. } => {
let output = task.final_output.unwrap_or_default();
return HippoxResult::ok_with_tokens(output, task.input_token_count, task.output_token_count);
}
HippoxResult { error: Some(e), .. } => {
return HippoxResult::system_error(format!("Task completed but data retrieval failed: {}", e));
}
_ => {
return HippoxResult::system_error(format!("Task completed but data not found: {}", task_id));
}
}
}
TaskStatus::Failed => match get_task(task_id).await {
HippoxResult { data: Some(task), .. } => {
let error = task.error.unwrap_or_else(|| "Unknown error".to_string());
return HippoxResult::system_error(format!("Task failed: {}", error));
}
_ => {
return HippoxResult::system_error(format!("Task failed: {}", task_id));
}
},
TaskStatus::Cancelled => {
return HippoxResult::system_error(format!("Task was cancelled: {}", task_id));
}
TaskStatus::Timeout => {
return HippoxResult::system_error(format!("Task timed out: {}", task_id));
}
TaskStatus::Pending | TaskStatus::Running | TaskStatus::Paused => {
sleep(Duration::from_millis(100)).await;
continue;
}
}
}
}