//! Task budget: a task runs in *slices*, and every slice is bounded.
//!
//! Slice exhaustion is its own outcome. It is not success (the task is not
//! finished) and not failure (nothing went wrong) - it is a resumable attempt
//! that the host continues on its own. The total budget is cumulative across
//! slices, so automatic reslicing can never evade it.
use serde::{Deserialize, Serialize};
/// What a task may spend in total, and how much of that one slice may use.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TaskBudget {
/// Steps one slice may use before it must yield to a continuation.
pub slice_steps: usize,
/// Steps the whole task may use, across every slice it is given.
pub total_steps: usize,
}
impl TaskBudget {
#[must_use]
pub fn new(slice_steps: usize, total_steps: usize) -> Self {
Self {
slice_steps,
total_steps,
}
}
}
impl Default for TaskBudget {
fn default() -> Self {
Self {
slice_steps: 10,
total_steps: 60,
}
}
}
/// Durable spend for one task, rebuilt from its committed slice records.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct BudgetSpend {
/// Slices the task has been given so far.
pub slices: usize,
/// Steps those slices actually used.
pub steps: usize,
}
/// Why a slice ended.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SliceEnd {
/// The model finished its response inside the allowance.
Finished,
/// The allowance ran out first. The task is incomplete, not failed.
Exhausted,
}
impl SliceEnd {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::Finished => "finished",
Self::Exhausted => "exhausted",
}
}
/// True when the host must continue the task without new user input.
#[must_use]
pub fn requires_continuation(self) -> bool {
matches!(self, Self::Exhausted)
}
}
/// Why a task stopped making progress.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PauseReason {
/// The cumulative budget is spent. A new slice would reset nothing.
BudgetExhausted { used: usize, limit: usize },
/// The task was stopped by the operator.
Stopped,
}
impl PauseReason {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::BudgetExhausted { .. } => "budget_exhausted",
Self::Stopped => "stopped",
}
}
}
/// The host's decision about the next slice.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TaskAdmission {
Admit { slice: usize, allowance: usize },
Paused { reason: PauseReason },
}
impl BudgetSpend {
#[must_use]
pub fn steps_remaining(self, budget: TaskBudget) -> usize {
budget.total_steps.saturating_sub(self.steps)
}
/// Decide whether this task gets another slice, and how big it may be.
///
/// The allowance is what is left of the *total*, so a slice may never be
/// larger than the task's remaining budget. This is the property that stops
/// reslicing from evading the quota.
#[must_use]
pub fn admission(self, budget: TaskBudget) -> TaskAdmission {
let remaining = self.steps_remaining(budget);
if remaining == 0 {
return TaskAdmission::Paused {
reason: PauseReason::BudgetExhausted {
used: self.steps,
limit: budget.total_steps,
},
};
}
TaskAdmission::Admit {
slice: self.slices + 1,
allowance: budget.slice_steps.min(remaining),
}
}
/// Fold one completed slice into the task's spend.
#[must_use]
pub fn after_slice(self, steps_used: usize) -> Self {
Self {
slices: self.slices + 1,
steps: self.steps.saturating_add(steps_used),
}
}
}