//! Durable task-budget accounting for the session runner.
//!
//! The runner gives a task *slices*. Each finished slice is recorded as an event,
//! and the next admission is decided from those records - not from an in-memory
//! counter - so a restart cannot hand the task a fresh quota.
use anyhow::Result;
use klbr_core::runtime_turn::budget::{
BudgetSpend, PauseReason, SliceEnd, TaskAdmission, TaskBudget,
};
use klbr_runtime::{
protocol::SessionId,
store::{now_ms, Store},
};
use serde_json::json;
/// Event kinds the budget accounting owns.
pub const SLICE_EVENT: &str = "runtime.slice";
pub const PAUSE_EVENT: &str = "runtime.task_paused";
/// One slice's outcome, as committed.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SliceRecord {
pub task: String,
pub attempt_id: String,
pub slice: usize,
pub allowance: usize,
pub steps_used: usize,
pub end: SliceEnd,
}
/// Rebuild a task's spend from its committed slice records.
///
/// Only slices that used at least one step count towards the total, so a slice
/// that ended before doing anything cannot be used to rearm a budget.
pub async fn task_spend(store: &Store, session: &SessionId, task: &str) -> Result<BudgetSpend> {
let mut spend = BudgetSpend::default();
let mut after = 0;
loop {
let page = store.history(session, after, 1024).await?;
if page.is_empty() {
break;
}
if let Some(last) = page.last() {
after = last.id;
}
for event in &page {
if event.kind != SLICE_EVENT {
continue;
}
if event.payload.get("work_item_id").and_then(|v| v.as_str()) != Some(task) {
continue;
}
let used = event
.payload
.get("steps_used")
.and_then(|v| v.as_u64())
.unwrap_or(0) as usize;
spend = spend.after_slice(used);
}
if page.len() < 1024 {
break;
}
}
Ok(spend)
}
/// Commit one slice's outcome before the host decides what happens next.
pub async fn record_slice(store: &Store, session: &SessionId, record: &SliceRecord) -> Result<i64> {
store
.append(
session,
SLICE_EVENT,
json!({
"work_item_id": record.task,
"attempt_id": record.attempt_id,
"slice": record.slice,
"allowance": record.allowance,
"steps_used": record.steps_used,
"end": record.end.as_str(),
}),
)
.await
}
/// Commit the reason a task stopped, with the usage that justifies it.
pub async fn record_pause(
store: &Store,
session: &SessionId,
task: &str,
reason: PauseReason,
spend: BudgetSpend,
) -> Result<i64> {
let (used, limit) = match reason {
PauseReason::BudgetExhausted { used, limit } => (used, limit),
PauseReason::Stopped => (spend.steps, spend.steps),
};
store
.append(
session,
PAUSE_EVENT,
json!({
"work_item_id": task,
"reason": reason.as_str(),
"used": used,
"limit": limit,
"slices": spend.slices,
"at_ms": now_ms().unwrap_or_default(),
}),
)
.await
}
/// True when the task's most recent budget record is already a pause, so a
/// repeated wake does not pile up duplicate pause evidence.
pub async fn pause_recorded(store: &Store, session: &SessionId, task: &str) -> Result<bool> {
let mut last: Option<String> = None;
let mut after = 0;
loop {
let page = store.history(session, after, 1024).await?;
if page.is_empty() {
break;
}
if let Some(last_event) = page.last() {
after = last_event.id;
}
for event in &page {
if event.kind != SLICE_EVENT && event.kind != PAUSE_EVENT {
continue;
}
if event.payload.get("work_item_id").and_then(|v| v.as_str()) != Some(task) {
continue;
}
last = Some(event.kind.clone());
}
if page.len() < 1024 {
break;
}
}
Ok(last.as_deref() == Some(PAUSE_EVENT))
}
/// The task key a session's budget is accounted against.
#[must_use]
pub fn task_key(work_item_id: Option<&str>, session: &SessionId) -> String {
match work_item_id {
Some(id) => id.to_string(),
None => format!("session:{session}"),
}
}
/// Decide the next slice for a task from its durable records.
pub async fn admit_next(
store: &Store,
session: &SessionId,
task: &str,
budget: TaskBudget,
) -> Result<(BudgetSpend, TaskAdmission)> {
let spend = task_spend(store, session, task).await?;
Ok((spend, spend.admission(budget)))
}