bridge_common/types/tasks.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152
// Copyright 2024 StarfleetAI
// SPDX-License-Identifier: Apache-2.0
use std::{
fmt::{self, Display, Formatter},
path::PathBuf,
};
use anyhow::Context;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::types::Result;
#[derive(Serialize, Deserialize, Debug, sqlx::Type, PartialEq, Default, Clone, Copy)]
pub enum Status {
/// Task is in draft and has not been selected for execution yet.
#[default]
Draft,
/// Task is selected for execution.
ToDo,
/// Task is currently being executed.
InProgress,
/// Task is waiting for a user input.
WaitingForUser,
/// Task is completed.
Done,
/// Task execution failed.
Failed,
}
impl Display for Status {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(f, "{self:?}")
}
}
impl From<String> for Status {
fn from(status: String) -> Self {
match status.as_str() {
"ToDo" => Status::ToDo,
"InProgress" => Status::InProgress,
"WaitingForUser" => Status::WaitingForUser,
"Done" => Status::Done,
"Failed" => Status::Failed,
_ => Status::Draft,
}
}
}
#[derive(Serialize, Deserialize, Debug, Default, Clone)]
pub struct Task {
pub id: i32,
pub company_id: i32,
pub user_id: i32,
pub agent_id: i32,
/// Chat from which this task was created.
pub origin_chat_id: Option<i32>,
/// Chat from which this task is being controlled (between the user and the Bridge).
pub control_chat_id: Option<i32>,
/// Chat in which this task is being executed (between the Bridge and the agent).
pub execution_chat_id: Option<i32>,
pub title: String,
pub summary: String,
pub status: Status,
/// Task's parent ids in a form of `1/2/3`. `None` for root tasks.
pub ancestry: Option<String>,
pub ancestry_level: i32,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
impl Task {
/// Returns parent id of the task.
///
/// # Errors
///
/// Returns error if there was a problem while parsing parent id.
pub fn parent_id(&self) -> Result<Option<i32>> {
Ok(match self.ancestry {
Some(ref ancestry) => {
let segment = ancestry
.split('/')
.last()
.context("No segments found in ancestry")?;
Some(
segment.parse::<i32>().with_context(|| {
"Failed to parse parent id from ancestry segment {segment}"
})?,
)
}
None => None,
})
}
/// Returns parent ids of the task.
///
/// # Errors
///
/// Returns error if there was a problem while parsing parent ids.
pub fn parent_ids(&self) -> Result<Option<Vec<i32>>> {
Ok(match self.ancestry {
Some(ref ancestry) => Some(
ancestry
.split('/')
.map(|segment| {
segment.parse::<i32>().with_context(|| {
"Failed to parse parent id from ancestry segment {segment}"
})
})
.collect::<std::result::Result<Vec<i32>, _>>()?,
),
None => None,
})
}
#[must_use]
pub fn children_ancestry(&self) -> String {
match self.ancestry {
Some(ref ancestry) => format!("{}/{}", ancestry, self.id),
None => self.id.to_string(),
}
}
/// Returns workdir for the task.
///
/// # Errors
///
/// Returns error if there was a problem while building workdir path.
pub async fn workdir(&self, root: &PathBuf) -> Result<PathBuf> {
let dir = format!(
"wd-task-{}",
self.workdir_id().context("Failed to get workdir ID")?
);
Ok(root.join(dir))
}
fn workdir_id(&self) -> Result<i32> {
Ok(match self.ancestry {
Some(ref ancestry) => ancestry
.split('/')
.collect::<Vec<_>>()
.first()
.context("No segments found in ancestry")?
.parse::<i32>()
.context("Failed to parse workdir id")?,
None => self.id,
})
}
}