diff --git a/docs/modules/graph/goals.md b/docs/modules/graph/goals.md index 53afdec5..a7e8e354 100644 --- a/docs/modules/graph/goals.md +++ b/docs/modules/graph/goals.md @@ -61,8 +61,32 @@ it spent into `State`, and the caller's `progress` / `run_turn` closure reports it. `made_progress == false` is the graph analogue of OpenHuman's "the turn produced no tool calls". +## Budget enforcement + +`store::account_usage` is the raw write; `graph::goals::budget` is the policy +around it — the two halves of enforcement a host would otherwise reimplement. + +- `account_turn(store, thread_id, input, output, secs, user_initiated)` charges + a finished turn against the thread's **active** goal (a paused, complete, or + budget-limited goal accrues nothing from incidental conversation) and returns + the goal as it stands afterwards, including a flip to `BudgetLimited`. A + `user_initiated` turn also clears the one-shot `continuation_suppressed` flag; + a continuation must not clear its own, or it would loop. +- `GoalBudgetGuard::for_goal(goal)` arms only for an active goal that has a + budget. `check(store, in_flight_tokens)` adds the turn's spend so far to the + accounted total and returns `BudgetVerdict::Stop` once it reaches the ceiling + — checked mid-turn, this bounds a run to a small overshoot instead of + discovering the overrun after the fact. `verdict_for` is the same decision + with no store read. The guard captures the `goal_id` it was armed for, so a + goal replaced mid-turn quietly disarms it rather than enforcing a ceiling that + no longer describes the work. + +Neither aborts anything: `account_turn` reports, the guard returns a verdict, +and wiring a stop into a turn stays the host's call. + ## Testing -Unit tests in `src/graph/goals/test.rs` (types, store, tools, and the gate loop -on `InMemoryStore`); an end-to-end self-driving loop in -`tests/e2e_graph_goals.rs`. +Unit tests in `src/graph/goals/test.rs` (types, store, tools, the gate loop, and +budget enforcement on `InMemoryStore`); an end-to-end self-driving loop in +`tests/e2e_graph_goals.rs`; feature coverage for budget accounting and the +mid-turn guard in `tests/feature_graph_goal_budget.rs`. diff --git a/docs/modules/graph/todos.md b/docs/modules/graph/todos.md index 92b3bc13..bc617101 100644 --- a/docs/modules/graph/todos.md +++ b/docs/modules/graph/todos.md @@ -55,7 +55,33 @@ bound to `ToolExecutionContext::thread_id` (never a tool argument). Domain error (unknown id, invariant violation) are surfaced to the model as tool errors rather than failing the run. +## Runs and dispatch + +Two layers sit on top of the board for hosts that run cards autonomously. + +`graph::todos::runs` is the claim/heartbeat/reclaim log (`task_run_store`). A +worker claims a card, opens a `TaskRun`, ticks a heartbeat while it works, and +closes the run with a `RunOutcome`. `reclaim_stale` sweeps runs whose heartbeat +or claim aged out under `RunLimits`, closes them `Reclaimed`, and returns their +card to `Todo` — or parks it at `Blocked` once the card has exceeded +`max_reclaim_count`, so a card that keeps killing workers stops cycling. The +staleness policy itself is the pure, clock-injected `staleness_reason`. See +[`src/graph/todos/runs/README.md`](../../../src/graph/todos/runs/README.md). + +`graph::todos::dispatch` is the scheduling policy: `pick_next_card` (urgency, +then board order, optionally agent-assigned only), `requires_plan_approval` +(the card's own mode outranks the global gate), `PollCadence` (idle backoff), +`build_task_prompt` / `build_progress_instruction`, and `ActiveRunRegistry` +(in-flight runs with race-free removal, so a terminal write-back happens once). +See [`src/graph/todos/dispatch/README.md`](../../../src/graph/todos/dispatch/README.md). + +Executing a card is out of scope for the crate — that needs a host's agent and +tool belt. `tests/e2e_graph_task_dispatch.rs` is the reference assembly. + ## Testing -Unit tests in `src/graph/todos/test.rs` (types, store invariants, tool); an -end-to-end model-driven tool run in `tests/e2e_graph_todos.rs`. +Unit tests in `src/graph/todos/test.rs` (types, store invariants, tool), +`src/graph/todos/runs/test.rs`, and `src/graph/todos/dispatch/test.rs`; an +end-to-end model-driven tool run in `tests/e2e_graph_todos.rs`; feature coverage +for the run lifecycle in `tests/feature_graph_task_runs.rs`; and the full +dispatch loop in `tests/e2e_graph_task_dispatch.rs`. diff --git a/src/graph/goals/README.md b/src/graph/goals/README.md index a3ef11ae..f24ee280 100644 --- a/src/graph/goals/README.md +++ b/src/graph/goals/README.md @@ -123,4 +123,5 @@ let exec = graph.run_with_thread("thread-1", St::default()).await?; | `store.rs` | `Store`-backed CRUD, per-thread RMW lock, budget + CAS guards. | | `tool.rs` | `GoalTool` / `GoalToolKind` harness tools. | | `continuation.rs` | `goal_gate_node`, `run_continuation_tick`, `note_user_turn`. | -| `test.rs` | Unit tests (types, store, tools, continuation loop). | +| `budget.rs` | `account_turn`, `GoalBudgetGuard`, `BudgetVerdict` — charging a finished turn and stopping an overrunning one. | +| `test.rs` | Unit tests (types, store, tools, continuation loop, budget). | diff --git a/src/graph/goals/budget.rs b/src/graph/goals/budget.rs new file mode 100644 index 00000000..dbe7c1e8 --- /dev/null +++ b/src/graph/goals/budget.rs @@ -0,0 +1,178 @@ +//! Charging a turn against a goal, and stopping one that would overrun it. +//! +//! [`store::account_usage`](super::store::account_usage) is the raw write. +//! This module is the policy around it — the two halves of budget enforcement a +//! host otherwise reimplements: +//! +//! - [`account_turn`] — *after* a turn: fold its usage into the thread's active +//! goal, and clear the one-shot continuation suppression when the turn was +//! user-initiated (a person re-engaging means a later idle period may +//! auto-continue again). +//! - [`GoalBudgetGuard`] — *during* a turn: given the tokens spent so far, +//! decide whether to stop now. Checked mid-turn, this bounds an autonomous +//! run to a small overshoot past its ceiling instead of discovering the +//! overrun only once the turn is over. +//! +//! Neither aborts anything itself. `account_turn` reports the goal as it stands +//! afterwards and the guard returns a [`BudgetVerdict`]; wiring a stop into a +//! turn is the host's call, because only the host knows whether a graceful +//! wrap-up or a hard cut is wanted. + +use std::sync::Arc; + +use super::store; +use super::types::{ThreadGoal, ThreadGoalStatus}; +use crate::error::Result; +use crate::harness::store::Store; + +/// Total tokens a turn spent — the quantity charged against a goal's budget. +pub fn turn_tokens(input: u64, output: u64) -> u64 { + input.saturating_add(output) +} + +/// Whether an in-flight turn should be stopped to stay inside its goal's budget. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum BudgetVerdict { + /// Keep going. + Continue, + /// The projected spend meets or exceeds the budget. + Stop { + /// Human-readable reason, suitable for a transcript or a log line. + reason: String, + }, +} + +impl BudgetVerdict { + /// Whether this verdict calls for stopping. + pub fn is_stop(&self) -> bool { + matches!(self, Self::Stop { .. }) + } +} + +/// Fold a finished turn's usage into the thread's goal. +/// +/// Only an **active** goal is charged: a paused, complete, or budget-limited +/// goal does not accrue usage from incidental conversation. Returns the goal as +/// it stands afterwards (including a flip to +/// [`BudgetLimited`](ThreadGoalStatus::BudgetLimited)), or `None` when the +/// thread has no goal or its goal is not active. +/// +/// `user_initiated` distinguishes a person's turn from an autonomous +/// continuation. A user turn clears the one-shot `continuation_suppressed` +/// flag; a continuation must not clear its own suppression, or it would loop. +pub async fn account_turn( + store: &Arc, + thread_id: &str, + input_tokens: u64, + output_tokens: u64, + elapsed_secs: u64, + user_initiated: bool, +) -> Result> { + let Some(goal) = store::get(store, thread_id).await? else { + return Ok(None); + }; + if !goal.status.is_active() { + return Ok(None); + } + + let mut current = goal; + if current.continuation_suppressed + && user_initiated + && let Some(updated) = + store::set_continuation_suppressed_if(store, thread_id, ¤t.goal_id, false).await? + { + current = updated; + } + + let delta = turn_tokens(input_tokens, output_tokens); + if delta == 0 && elapsed_secs == 0 { + return Ok(Some(current)); + } + store::account_usage(store, thread_id, ¤t.goal_id, delta, elapsed_secs).await +} + +/// Mid-turn budget check, armed for one specific version of a goal. +/// +/// Built from a goal that is active and has a budget; a goal with neither is +/// nothing to enforce. The captured `goal_id` is what makes the guard safe to +/// hold across a long turn: if the objective is replaced while the turn runs, +/// the new goal has a new id, the guard stops matching, and it quietly stands +/// down instead of enforcing a budget that no longer applies. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GoalBudgetGuard { + thread_id: String, + goal_id: String, + budget: u64, +} + +impl GoalBudgetGuard { + /// A guard for `goal`, or `None` when it is inactive or has no budget. + pub fn for_goal(goal: &ThreadGoal) -> Option { + if !goal.status.is_active() { + return None; + } + Some(Self { + thread_id: goal.thread_id.clone(), + goal_id: goal.goal_id.clone(), + budget: goal.token_budget?, + }) + } + + /// The thread this guard watches. + pub fn thread_id(&self) -> &str { + &self.thread_id + } + + /// The goal version this guard was armed for. + pub fn goal_id(&self) -> &str { + &self.goal_id + } + + /// The ceiling being enforced. + pub fn budget(&self) -> u64 { + self.budget + } + + /// Verdict for a turn that has spent `in_flight_tokens` so far. + /// + /// Reads the goal's already-accounted usage, adds the in-flight spend, and + /// stops once the total reaches the budget. Returns + /// [`Continue`](BudgetVerdict::Continue) whenever there is nothing left to + /// enforce: the goal is gone, was replaced, or is no longer active. + pub async fn check( + &self, + store: &Arc, + in_flight_tokens: u64, + ) -> Result { + let Some(goal) = store::get(store, &self.thread_id).await? else { + return Ok(BudgetVerdict::Continue); + }; + if goal.goal_id != self.goal_id || !goal.status.is_active() { + return Ok(BudgetVerdict::Continue); + } + Ok(self.verdict_for(goal.tokens_used, in_flight_tokens)) + } + + /// The pure half of [`check`](Self::check): the verdict for a given + /// accounted and in-flight spend, with no store read. + pub fn verdict_for(&self, accounted_tokens: u64, in_flight_tokens: u64) -> BudgetVerdict { + let projected = accounted_tokens.saturating_add(in_flight_tokens); + if projected >= self.budget { + BudgetVerdict::Stop { + reason: format!( + "thread goal budget reached: {projected} tokens >= {} budget — stopping to \ + summarise progress", + self.budget + ), + } + } else { + BudgetVerdict::Continue + } + } +} + +/// Whether `status` still accrues usage. Kept next to the accounting policy so +/// callers do not re-derive the rule. +pub fn accrues_usage(status: ThreadGoalStatus) -> bool { + status.is_active() +} diff --git a/src/graph/goals/mod.rs b/src/graph/goals/mod.rs index c438f418..7472bf08 100644 --- a/src/graph/goals/mod.rs +++ b/src/graph/goals/mod.rs @@ -6,18 +6,22 @@ //! is exhausted, or a host pauses it. This module owns the data model //! ([`types`]), harness-[`Store`](crate::harness::store::Store)-backed //! persistence ([`store`]), the model-facing controls exposed as harness tools -//! ([`tool`]), and the graph-native continuation surface ([`continuation`]). +//! ([`tool`]), the graph-native continuation surface ([`continuation`]), and +//! budget enforcement ([`budget`]) — charging a finished turn against the goal +//! and stopping an in-flight one that would overrun its ceiling. //! //! It is the graph analogue of OpenHuman's `thread_goals`, minus the //! app-specific coupling (event bus, RPC envelopes, heartbeat scheduler): the //! primitive is provider-neutral and drives off the graph runtime. +pub mod budget; mod continuation; mod prompt; pub mod store; mod tool; mod types; +pub use budget::{BudgetVerdict, GoalBudgetGuard, account_turn, accrues_usage, turn_tokens}; pub use continuation::{goal_gate_node, note_user_turn, run_continuation_tick}; pub use prompt::active_goal_context_block; pub use tool::{GoalTool, GoalToolKind, goal_tools, register_goal_tools}; diff --git a/src/graph/goals/test.rs b/src/graph/goals/test.rs index 839ec8b9..773e14e0 100644 --- a/src/graph/goals/test.rs +++ b/src/graph/goals/test.rs @@ -632,3 +632,196 @@ mod continuation_tests { assert!(note_user_turn(&s, "missing").await.unwrap().is_none()); } } + +mod budget_tests { + use std::sync::Arc; + + use super::super::budget::{ + BudgetVerdict, GoalBudgetGuard, account_turn, accrues_usage, turn_tokens, + }; + use super::super::store; + use super::{ThreadGoalStatus, goal}; + use crate::harness::store::{InMemoryStore, Store}; + + fn kv() -> Arc { + Arc::new(InMemoryStore::default()) + } + + #[test] + fn a_turns_charge_is_its_input_plus_output() { + assert_eq!(turn_tokens(120, 80), 200); + // Saturating, so an absurd provider report cannot wrap to a tiny charge. + assert_eq!(turn_tokens(u64::MAX, 1), u64::MAX); + } + + #[test] + fn only_an_active_goal_accrues_usage() { + assert!(accrues_usage(ThreadGoalStatus::Active)); + assert!(!accrues_usage(ThreadGoalStatus::Paused)); + assert!(!accrues_usage(ThreadGoalStatus::BudgetLimited)); + assert!(!accrues_usage(ThreadGoalStatus::Complete)); + } + + #[tokio::test] + async fn a_turn_is_charged_against_the_active_goal() { + let kv = kv(); + store::set(&kv, "t1", "ship it", Some(1_000)).await.unwrap(); + + let updated = account_turn(&kv, "t1", 100, 50, 12, true) + .await + .unwrap() + .expect("goal charged"); + assert_eq!(updated.tokens_used, 150); + assert_eq!(updated.time_used_seconds, 12); + assert_eq!(updated.status, ThreadGoalStatus::Active); + } + + #[tokio::test] + async fn crossing_the_budget_flips_the_goal_to_budget_limited() { + let kv = kv(); + store::set(&kv, "t1", "ship it", Some(100)).await.unwrap(); + + let updated = account_turn(&kv, "t1", 60, 60, 1, true) + .await + .unwrap() + .unwrap(); + assert_eq!(updated.status, ThreadGoalStatus::BudgetLimited); + assert!(updated.over_budget()); + + // A limited goal stops accruing: incidental chat afterwards is free. + assert!( + account_turn(&kv, "t1", 500, 500, 5, true) + .await + .unwrap() + .is_none() + ); + let after = store::get(&kv, "t1").await.unwrap().unwrap(); + assert_eq!(after.tokens_used, 120); + } + + #[tokio::test] + async fn a_thread_with_no_goal_or_an_idle_turn_changes_nothing() { + let kv = kv(); + assert!( + account_turn(&kv, "missing", 10, 10, 1, true) + .await + .unwrap() + .is_none() + ); + + store::set(&kv, "t1", "ship it", None).await.unwrap(); + let unchanged = account_turn(&kv, "t1", 0, 0, 0, true) + .await + .unwrap() + .unwrap(); + assert_eq!(unchanged.tokens_used, 0); + } + + #[tokio::test] + async fn a_user_turn_clears_suppression_but_a_continuation_does_not() { + let kv = kv(); + let goal = store::set(&kv, "t1", "ship it", None).await.unwrap(); + store::set_continuation_suppressed_if(&kv, "t1", &goal.goal_id, true) + .await + .unwrap(); + + // The continuation's own accounting must not clear its one-shot flag, + // or the loop would never stop. + account_turn(&kv, "t1", 10, 10, 1, false).await.unwrap(); + assert!( + store::get(&kv, "t1") + .await + .unwrap() + .unwrap() + .continuation_suppressed + ); + + // A person re-engaging re-arms the next idle continuation. + account_turn(&kv, "t1", 10, 10, 1, true).await.unwrap(); + assert!( + !store::get(&kv, "t1") + .await + .unwrap() + .unwrap() + .continuation_suppressed + ); + } + + #[test] + fn a_guard_is_only_armed_for_an_active_goal_with_a_budget() { + assert!(GoalBudgetGuard::for_goal(&goal(ThreadGoalStatus::Active, Some(100), 0)).is_some()); + assert!(GoalBudgetGuard::for_goal(&goal(ThreadGoalStatus::Active, None, 0)).is_none()); + assert!(GoalBudgetGuard::for_goal(&goal(ThreadGoalStatus::Paused, Some(100), 0)).is_none()); + assert!( + GoalBudgetGuard::for_goal(&goal(ThreadGoalStatus::Complete, Some(100), 0)).is_none() + ); + } + + #[test] + fn the_verdict_counts_in_flight_spend_against_the_ceiling() { + let guard = + GoalBudgetGuard::for_goal(&goal(ThreadGoalStatus::Active, Some(1_000), 0)).unwrap(); + assert_eq!(guard.budget(), 1_000); + assert_eq!(guard.goal_id(), "goal-0"); + assert_eq!(guard.thread_id(), "t"); + + assert_eq!(guard.verdict_for(400, 300), BudgetVerdict::Continue); + // Reaching the budget stops the turn — it does not have to exceed it. + let verdict = guard.verdict_for(400, 600); + assert!(verdict.is_stop()); + match verdict { + BudgetVerdict::Stop { reason } => { + assert!(reason.contains("1000 tokens >= 1000 budget"), "{reason}") + } + BudgetVerdict::Continue => unreachable!(), + } + } + + #[tokio::test] + async fn a_guard_stops_a_turn_that_would_overrun() { + let kv = kv(); + let stored = store::set(&kv, "t1", "ship it", Some(500)).await.unwrap(); + let guard = GoalBudgetGuard::for_goal(&stored).unwrap(); + account_turn(&kv, "t1", 200, 100, 1, true).await.unwrap(); + + assert_eq!( + guard.check(&kv, 100).await.unwrap(), + BudgetVerdict::Continue + ); + assert!(guard.check(&kv, 200).await.unwrap().is_stop()); + } + + #[tokio::test] + async fn a_guard_stands_down_when_its_goal_is_replaced_or_gone() { + let kv = kv(); + let stored = store::set(&kv, "t1", "ship it", Some(10)).await.unwrap(); + let guard = GoalBudgetGuard::for_goal(&stored).unwrap(); + // Armed and biting on the goal it was built for. + assert!(guard.check(&kv, 50).await.unwrap().is_stop()); + + // A fresh objective mints a new goal id: the old guard no longer applies. + store::set(&kv, "t1", "a different objective", Some(10)) + .await + .unwrap(); + assert_eq!(guard.check(&kv, 50).await.unwrap(), BudgetVerdict::Continue); + + // And a cleared goal leaves nothing to enforce. + store::clear(&kv, "t1").await.unwrap(); + assert_eq!(guard.check(&kv, 50).await.unwrap(), BudgetVerdict::Continue); + } + + #[tokio::test] + async fn a_guard_stands_down_once_its_goal_is_no_longer_active() { + let kv = kv(); + let stored = store::set(&kv, "t1", "ship it", Some(10)).await.unwrap(); + let guard = GoalBudgetGuard::for_goal(&stored).unwrap(); + store::pause(&kv, "t1").await.unwrap(); + + // A paused goal is not burning a live budget, so a user-present turn is + // never hard-stopped by it. + assert_eq!( + guard.check(&kv, 1_000).await.unwrap(), + BudgetVerdict::Continue + ); + } +} diff --git a/src/graph/mod.rs b/src/graph/mod.rs index 17bfc3cc..66e86c53 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -15,7 +15,8 @@ //! streaming/events ([`stream`]), run-status snapshots ([`status`]), graph //! export/visualization ([`export`]), subgraph embedding ([`subgraph`]), and //! per-thread productivity primitives — a durable goal ([`goals`]) and a kanban -//! task board ([`todos`]) — exposed as harness tools. +//! task board ([`todos`], with its claim/heartbeat run log and dispatch policy) +//! — exposed as harness tools. //! //! Each concern lives in its own submodule with `types.rs` (definitions), //! `mod.rs` (implementations), and `test.rs` (unit tests). @@ -64,9 +65,10 @@ pub use export::{ blueprint_to_mermaid, blueprint_to_topology, from_json, to_json, to_mermaid, }; pub use goals::{ - GoalProgress, GoalTool, GoalToolKind, ThreadGoal, ThreadGoalStatus, TurnOutcome, - active_goal_context_block, goal_gate_node, goal_tools, note_user_turn, register_goal_tools, - run_continuation_tick, + BudgetVerdict, GoalBudgetGuard, GoalProgress, GoalTool, GoalToolKind, ThreadGoal, + ThreadGoalStatus, TurnOutcome, account_turn, accrues_usage, active_goal_context_block, + goal_gate_node, goal_tools, note_user_turn, register_goal_tools, run_continuation_tick, + turn_tokens, }; pub use observability::{ GraphEventJournal, GraphHealthSummary, GraphLangfuseExporter, GraphLatencyMetrics, @@ -102,6 +104,13 @@ pub use testkit::{ assert_graph, failing_node, fanout_node, interrupting_node, noop_node, run_recorded, scripted_route_node, scripted_update_node, subagent_fake_node, subgraph_test_node, }; +pub use todos::dispatch::{ + ActiveRun, ActiveRunRegistry, PollCadence, TaskPromptTools, build_progress_instruction, + build_task_prompt, card_urgency, has_card_in_progress, pick_next_card, requires_plan_approval, +}; +pub use todos::runs::{ + ReclaimDetail, ReclaimResult, RunLimits, RunOutcome, TaskRun, staleness_reason, +}; pub use todos::{ CardPatch, TaskApprovalMode, TaskBoard, TaskBoardCard, TaskCardStatus, TodoTool, TodosSnapshot, normalise_board, parse_status, register_todo_tools, render_markdown, todo_tools, diff --git a/src/graph/todos/README.md b/src/graph/todos/README.md index 291b19b7..719e49a8 100644 --- a/src/graph/todos/README.md +++ b/src/graph/todos/README.md @@ -93,4 +93,6 @@ let tool = TodoTool::new(store.clone()); | `types.rs` | Card/board model, `parse_status`, `render_markdown`, `normalise_board`, `CardPatch`, `TodosSnapshot`. | | `store.rs` | `Store`-backed CRUD, per-thread RMW lock, single-in-progress invariant, CAS `claim_card`. | | `tool.rs` | The `todo` multiplexer tool. | +| `runs/` | Claim / heartbeat / reclaim log over the board — see [`runs/README.md`](runs/README.md). | +| `dispatch/` | Selection, approval gate, poll cadence, prompts, in-flight run registry — see [`dispatch/README.md`](dispatch/README.md). | | `test.rs` | Unit tests (types, store, tool). | diff --git a/src/graph/todos/dispatch/README.md b/src/graph/todos/dispatch/README.md new file mode 100644 index 00000000..de420b4f --- /dev/null +++ b/src/graph/todos/dispatch/README.md @@ -0,0 +1,82 @@ +# graph::todos::dispatch + +Dispatch policy for a task board: **what runs next, what it is told to do, and +how a run in flight is tracked and cancelled.** `store` owns the board and +`runs` owns claims; this is the layer between them a scheduler is built from. + +Ported from OpenHuman's `agent::task_dispatcher`, keeping the parts that are +policy and leaving behind the parts that are a host (its agent registry, config +loading, personality profiles, and event bus). + +## Selection (`select.rs`) + +Pure functions over a board snapshot — no store, so they are trivially testable +and can be applied to cards a host already holds. + +- `pick_next_card(cards, agent_assigned_only)` — the highest-urgency + dispatchable card (`Todo` or approved `Ready`). Urgency comes from + `source_metadata.urgency` (`card_urgency`, default `0.0`); ties break toward + the lower board `order`, so equal-priority work runs in planned order. + `agent_assigned_only` restricts the pick to cards with an `assigned_agent`, + which is how a host keeps an autonomous sweep off a person's own todos. +- `has_card_in_progress(cards)` — the board already has a card being worked, so + there is nothing to claim this tick. +- `requires_plan_approval(global_required, approval_mode)` — the card's own + mode is authoritative; `global_required` is only the fallback. `Required` + parks the card **even when the global gate is off**, or a plan stamped for + interactive review would execute before anyone saw it. +- `PollCadence { base, max_backoff, grace_ticks }` — diminishing-returns + polling. `next_delay(idle_ticks)` holds `base` through the grace window, then + doubles per idle tick, saturating at `max_backoff`. Monotonic and + overflow-free for any `u32` streak. Defaults: 60s base, 15min ceiling, 2 ticks + of grace. + +## Prompts (`prompt.rs`) + +- `build_task_prompt(card, tools)` — objective (falling back to the title), then + numbered plan steps and acceptance criteria. A card carrying + `source_metadata` with a `provider` also gets its provenance, a pointer at the + host's memory-recall tool, and an instruction to record the outcome back on + the upstream item. An id-only card gets no provenance block at all, since a + bare `#123` tells the model nothing. +- `build_progress_instruction(card_id, thread_id, tools)` — the addendum that + asks the run to append notes/evidence as it works and, crucially, to **block + rather than guess** when it needs a decision it cannot make. A run that blocks + leaves the card paused for a human instead of force-completed. +- `TaskPromptTools { memory_recall, update_task }` names the tools those prompts + point at, so a host that registers them under other names (or has no memory + tool) still gets coherent text. + +## Registry (`registry.rs`) + +`ActiveRunRegistry` maps a session thread id to an `ActiveRun` +(`run_id`, `card_id`, an `AbortHandle`, a heartbeat cancel sender, and the +host's own `Context`). + +Its real job is **deciding who cleans up**. A run finishing naturally and a +cancel arriving at the same moment both want to write the card's terminal +state; both must go through `take` / `take_if`, and only one gets `Some`. +`take_if(thread_id, Some(run_id))` matches and removes under a single lock, so +a cancel for a superseded request cannot tear down the run that replaced it. + +## Composing a dispatcher + +```text +reclaim_stale → has_card_in_progress? → pick_next_card + → requires_plan_approval? → park at AwaitingApproval + → claim_card → create_run → build_task_prompt → spawn + register + → complete_run → card write-back → take from the registry +``` + +Executing the card is deliberately out of scope: that needs an agent, a model, +and a host's tool belt. `tests/e2e_graph_task_dispatch.rs` wires the whole loop +against a `MockModel` and is the reference assembly. + +## Files + +| File | Role | +| --- | --- | +| `select.rs` | Card selection, approval gate, polling cadence. | +| `prompt.rs` | Task prompt and progress-instruction rendering. | +| `registry.rs` | In-flight run tracking with race-free removal. | +| `test.rs` | Unit tests for all three. | diff --git a/src/graph/todos/dispatch/mod.rs b/src/graph/todos/dispatch/mod.rs new file mode 100644 index 00000000..f41e8aed --- /dev/null +++ b/src/graph/todos/dispatch/mod.rs @@ -0,0 +1,34 @@ +//! Dispatch policy for a task board: **what runs next, what it is told to do, +//! and how a run in flight is tracked and cancelled**. +//! +//! [`store`](crate::graph::todos::store) owns the board and +//! [`runs`](crate::graph::todos::runs) owns claims; this module is the layer +//! between them that a scheduler is built from: +//! +//! - [`select`] — pure policy. Pick the highest-urgency dispatchable card, +//! decide whether it needs plan approval first, and pace a polling sweep so +//! an idle board is not swept at full rate forever. +//! - [`prompt`] — render a card into the prompt its run works from, plus the +//! addendum that keeps the card current while the run works. +//! - [`registry`] — track in-flight runs so a cancel can reach a detached task, +//! with race-free removal so the card's terminal write-back happens once. +//! +//! Actually *executing* a card is deliberately not here: that needs an agent, +//! a model, and a host's own tool belt. The intended shape is a loop that asks +//! [`select::pick_next_card`] for work, claims it with +//! [`store::claim_card`](crate::graph::todos::store::claim_card), opens a run +//! with [`runs::create_run`](crate::graph::todos::runs::create_run), spawns the +//! work with [`prompt::build_task_prompt`], and registers the handle. + +pub mod prompt; +pub mod registry; +pub mod select; + +pub use prompt::{TaskPromptTools, build_progress_instruction, build_task_prompt}; +pub use registry::{ActiveRun, ActiveRunRegistry}; +pub use select::{ + PollCadence, card_urgency, has_card_in_progress, pick_next_card, requires_plan_approval, +}; + +#[cfg(test)] +mod test; diff --git a/src/graph/todos/dispatch/prompt.rs b/src/graph/todos/dispatch/prompt.rs new file mode 100644 index 00000000..0af363d3 --- /dev/null +++ b/src/graph/todos/dispatch/prompt.rs @@ -0,0 +1,164 @@ +//! Turning a task card into the prompt an autonomous run works from. +//! +//! Two pieces, both pure: +//! +//! - [`build_task_prompt`] — the goal prompt: the card's objective, plan, and +//! acceptance criteria, plus provenance for a card ingested from an external +//! source. +//! - [`build_progress_instruction`] — the addendum that tells the run how to +//! keep its own card current while it works, and how to *stop* by blocking +//! rather than guessing. +//! +//! The instruction text names two tools by convention — a memory-recall tool +//! and the board's own card-update tool. A host that registers them under +//! different names should pass its own names via [`TaskPromptTools`]. + +use crate::graph::todos::types::TaskBoardCard; + +/// Tool names the generated prompts point the model at. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TaskPromptTools { + /// Tool that pulls related context out of memory. Omit to drop the + /// "recall related context" sentence entirely. + pub memory_recall: Option, + /// Tool that edits a card on a named board by id. + pub update_task: String, +} + +impl Default for TaskPromptTools { + fn default() -> Self { + Self { + memory_recall: Some("memory_recall".to_string()), + update_task: "update_task".to_string(), + } + } +} + +/// Render `card` into the goal prompt handed to an autonomous run. +/// +/// Leads with the card's `objective` (falling back to its title), then the +/// `plan` steps and `acceptance_criteria` that define done. When the card +/// carries `source_metadata` naming a provider, the prompt also points the run +/// at the originating item — so it can pull the item's prior discussion out of +/// memory before it starts, and record the outcome back on the source when it +/// finishes. +pub fn build_task_prompt(card: &TaskBoardCard, tools: &TaskPromptTools) -> String { + let mut lines: Vec = Vec::new(); + + let objective = card + .objective + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| card.title.trim()); + lines.push(format!( + "You are autonomously executing one task to completion. Objective:\n{objective}" + )); + + if !card.plan.is_empty() { + lines.push("\nPlan:".to_string()); + for (index, step) in card.plan.iter().enumerate() { + lines.push(format!("{}. {}", index + 1, step.trim())); + } + } + + if !card.acceptance_criteria.is_empty() { + lines.push("\nAcceptance criteria (the task is done only when all hold):".to_string()); + for criterion in &card.acceptance_criteria { + lines.push(format!("- {}", criterion.trim())); + } + } + + if let Some(meta) = &card.source_metadata { + let provider = meta.get("provider").and_then(|v| v.as_str()); + let external_id = meta.get("external_id").and_then(|v| v.as_str()); + let url = meta.get("url").and_then(|v| v.as_str()); + let origin = source_origin( + provider, + meta.get("repo").and_then(|v| v.as_str()), + external_id, + ); + + // Gated on a known provider so the origin string is always meaningful: + // an id-only card would otherwise render a bare "#123". + if provider.is_some() { + if let Some(recall) = &tools.memory_recall { + lines.push(format!( + "\nThis task originates from {origin}. Its activity has been ingested into \ + memory — use your {recall} tool to pull related context (prior discussion, \ + linked items) before and while you work." + )); + } else { + lines.push(format!("\nThis task originates from {origin}.")); + } + } + if let Some(url) = url { + lines.push(format!("Source link: {url}")); + } + // When the upstream item is addressable, close the loop on it: the run + // reports back through whatever integration tools it already holds, + // under their existing write scope. + if provider.is_some() && external_id.is_some() { + lines.push(format!( + "\nWhen the task is complete, record the outcome on the upstream source \ + ({origin}): use your integration tools to add a comment summarising the \ + resolution and, if the work fully addresses it, close/resolve the item. If you \ + lack the permission or connection to do so, say so in your final summary instead \ + of guessing." + )); + } + } + + lines.push( + "\nWork the task to completion. Do not pick up unrelated work. When finished, your final \ + message should summarise what you did and the evidence (commits, PRs, results)." + .to_string(), + ); + + lines.join("\n") +} + +/// A one-line `provider repo#id` origin, with the blank parts left out. +fn source_origin(provider: Option<&str>, repo: Option<&str>, external_id: Option<&str>) -> String { + let mut origin = String::new(); + if let Some(provider) = provider { + origin.push_str(provider); + } + if let Some(repo) = repo { + origin.push(' '); + origin.push_str(repo); + } + if let Some(external_id) = external_id { + origin.push('#'); + origin.push_str(external_id); + } + origin.trim().to_string() +} + +/// The addendum appended to a run's prompt so it keeps its own card current. +/// +/// The card is already `InProgress` — the dispatcher claimed it before spawning +/// the run — and is addressed by exact id and board, because a card-update tool +/// that defaults to some other board would silently miss it. Two behaviours are +/// asked for: append progress as it happens, and **block instead of guessing** +/// when the run needs a decision it cannot make. A run that blocks leaves the +/// card paused for a human rather than force-completed. +pub fn build_progress_instruction( + card_id: &str, + thread_id: &str, + tools: &TaskPromptTools, +) -> String { + let update = &tools.update_task; + format!( + "\n\nThis task is tracked as card `{card_id}` on the `{thread_id}` board. As you work, \ + call the `{update}` tool (id `{card_id}`, threadId `{thread_id}`) to keep the card \ + current — append `notes`/`evidence` as you make progress.\n\nIf you need a decision or \ + information from the user, or you genuinely cannot proceed (missing access, ambiguous \ + requirement, an action that needs the user's confirmation), call `{update}` with \ + `status: blocked` and a `blocker` that states exactly what you need from the user. The \ + task will stay paused in that blocked state until the user responds — do NOT guess, \ + fabricate, or take a risky irreversible action just to avoid blocking. If instead you \ + finish the work, end with a summary of what you did and the evidence; completion is \ + recorded automatically." + ) +} diff --git a/src/graph/todos/dispatch/registry.rs b/src/graph/todos/dispatch/registry.rs new file mode 100644 index 00000000..85af1895 --- /dev/null +++ b/src/graph/todos/dispatch/registry.rs @@ -0,0 +1,148 @@ +//! Registry of in-flight autonomous runs, keyed by the session they stream into. +//! +//! Autonomous card runs are detached tasks, not turns a chat channel knows +//! about, so a "stop" arriving through the normal path has nothing to cancel. +//! Registering each run's [`AbortHandle`](tokio::task::AbortHandle) here gives +//! that path a handle to pull. +//! +//! The registry's real job is **deciding who cleans up**. A run that finishes +//! naturally and a cancel that arrives at the same moment both try to write the +//! card's terminal state. Both must go through [`ActiveRunRegistry::take`] (or +//! [`take_if`](ActiveRunRegistry::take_if)), and only one of them gets `Some` — +//! so the write-back happens exactly once. + +use std::collections::HashMap; +use std::sync::Mutex; + +use tokio::sync::watch; +use tokio::task::AbortHandle; + +/// A live run's cancellation handles plus whatever context its host needs to +/// finish the card off. +/// +/// `Context` is the host's own payload — typically the board coordinates and +/// anything else its write-back needs. The crate never inspects it. +#[derive(Debug)] +pub struct ActiveRun { + /// The run this entry belongs to. + pub run_id: String, + /// The card it is executing. + pub card_id: String, + /// Aborts the detached run task. + pub abort: AbortHandle, + /// Stops the run's background heartbeat. + pub heartbeat_cancel: watch::Sender, + /// Host-supplied context, returned with the entry when it is taken. + pub context: Context, +} + +impl ActiveRun { + /// Abort the run task and stop its heartbeat. + /// + /// Does **not** write the card back — that is the caller's job, because + /// only the caller knows what terminal state the card should land in. + pub fn cancel(&self) { + self.abort.abort(); + let _ = self.heartbeat_cancel.send(true); + } +} + +/// A `session_thread_id → ActiveRun` map with race-free removal. +#[derive(Debug, Default)] +pub struct ActiveRunRegistry { + runs: Mutex>>, +} + +impl ActiveRunRegistry { + /// An empty registry. + pub fn new() -> Self { + Self { + runs: Mutex::new(HashMap::new()), + } + } + + /// Record `run` as the live run on `thread_id`, returning any entry it + /// displaced (which the caller should cancel). + pub fn register( + &self, + thread_id: impl Into, + run: ActiveRun, + ) -> Option> { + self.lock().insert(thread_id.into(), run) + } + + /// Remove and return the run on `thread_id`, if any. + /// + /// Whoever gets `Some` owns the terminal card write-back. + pub fn take(&self, thread_id: &str) -> Option> { + self.lock().remove(thread_id) + } + + /// Remove and return the run on `thread_id`, but only when it is the run + /// `run_id` names. `None` for `run_id` removes whatever is there. + /// + /// The match and the removal happen under one lock acquisition. A + /// peek-then-take would leave a window in which the matched run finishes + /// and a *newer* run takes its place before removal — cancelling the new + /// run instead of the intended one. + pub fn take_if(&self, thread_id: &str, run_id: Option<&str>) -> Option> { + let mut runs = self.lock(); + if let Some(run_id) = run_id { + match runs.get(thread_id) { + None => { + tracing::debug!( + thread_id = %thread_id, + request_run_id = %run_id, + "[graph:todos:dispatch] scoped cancel ignored: no active run on thread" + ); + return None; + } + Some(active) if active.run_id != run_id => { + tracing::debug!( + thread_id = %thread_id, + request_run_id = %run_id, + active_run_id = %active.run_id, + "[graph:todos:dispatch] scoped cancel ignored: run id mismatch" + ); + return None; + } + _ => {} + } + } + runs.remove(thread_id) + } + + /// Whether a run is registered for `thread_id`. + pub fn contains(&self, thread_id: &str) -> bool { + self.lock().contains_key(thread_id) + } + + /// Number of live runs. + pub fn len(&self) -> usize { + self.lock().len() + } + + /// Whether no runs are live. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// The session thread ids with a live run, in unspecified order. + pub fn thread_ids(&self) -> Vec { + self.lock().keys().cloned().collect() + } + + /// Remove every entry, returning them so the caller can cancel each. + pub fn drain(&self) -> Vec<(String, ActiveRun)> { + self.lock().drain().collect() + } + + fn lock(&self) -> std::sync::MutexGuard<'_, HashMap>> { + // A panic inside the map would leave it poisoned; the map holds no + // invariant that a panic could break, so recovering is safe and keeps + // one bad run from wedging cancellation for every other thread. + self.runs + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} diff --git a/src/graph/todos/dispatch/select.rs b/src/graph/todos/dispatch/select.rs new file mode 100644 index 00000000..47eeaeda --- /dev/null +++ b/src/graph/todos/dispatch/select.rs @@ -0,0 +1,135 @@ +//! Which card runs next, whether it needs approval first, and how often to +//! look — the scheduling policy of a task board, as pure functions. +//! +//! None of this touches a [`Store`](crate::harness::store::Store): a caller +//! reads a board snapshot, asks these functions what to do, and performs the +//! claim itself. That keeps the policy trivially testable and lets a host apply +//! it to boards it holds in memory. + +use std::time::Duration; + +use crate::graph::todos::types::{TaskApprovalMode, TaskBoardCard, TaskCardStatus}; + +/// A card's urgency, read from `source_metadata.urgency`. Cards without one +/// sort as `0.0`. +pub fn card_urgency(card: &TaskBoardCard) -> f64 { + card.source_metadata + .as_ref() + .and_then(|meta| meta.get("urgency")) + .and_then(serde_json::Value::as_f64) + .unwrap_or(0.0) +} + +/// Whether the board already has a card being worked. +/// +/// The board store caps a board at one `InProgress` card, so this answers +/// "is there anything for a dispatcher to claim right now?". +pub fn has_card_in_progress(cards: &[TaskBoardCard]) -> bool { + cards + .iter() + .any(|card| card.status == TaskCardStatus::InProgress) +} + +/// The highest-urgency dispatchable card, or `None` when the board has none. +/// +/// Dispatchable means `Todo` (not yet triaged) or `Ready` (approved). Ties on +/// urgency break toward the lower board `order`, so equal-priority work runs in +/// the order it was planned. +/// +/// `agent_assigned_only` restricts the pick to cards with an `assigned_agent`. +/// A host uses it for boards that mix human-authored and agent-authored cards, +/// so an autonomous sweep never picks up a card a person wrote for themselves. +pub fn pick_next_card(cards: &[TaskBoardCard], agent_assigned_only: bool) -> Option { + cards + .iter() + .filter(|card| matches!(card.status, TaskCardStatus::Todo | TaskCardStatus::Ready)) + .filter(|card| !agent_assigned_only || is_agent_assigned(card)) + .max_by(|a, b| { + card_urgency(a) + .partial_cmp(&card_urgency(b)) + .unwrap_or(std::cmp::Ordering::Equal) + // Reversed, so the *lower* order wins an urgency tie. + .then(b.order.cmp(&a.order)) + }) + .cloned() +} + +fn is_agent_assigned(card: &TaskBoardCard) -> bool { + card.assigned_agent + .as_deref() + .is_some_and(|agent| !agent.trim().is_empty()) +} + +/// Whether a card must be parked at +/// [`AwaitingApproval`](TaskCardStatus::AwaitingApproval) before it runs. +/// +/// The card's own [`TaskApprovalMode`] is authoritative when set; `global_required` +/// is only the fallback for cards that express no preference: +/// +/// - [`Required`](TaskApprovalMode::Required) → always park, *even when the +/// global default is off*. A card stamped by an interactive plan review must +/// still be reviewed, or the plan would execute before anyone saw it. +/// - [`NotRequired`](TaskApprovalMode::NotRequired) → never park; the card has +/// already cleared review. +/// - unset → `global_required`. +pub fn requires_plan_approval( + global_required: bool, + approval_mode: Option<&TaskApprovalMode>, +) -> bool { + match approval_mode { + Some(TaskApprovalMode::Required) => true, + Some(TaskApprovalMode::NotRequired) => false, + None => global_required, + } +} + +/// Diminishing-returns polling cadence for a board sweep. +/// +/// A dispatcher that sweeps a board on a timer should not keep sweeping an idle +/// board at full rate forever. [`PollCadence::next_delay`] holds the base +/// interval while there is work (and for a short grace period after it dries +/// up), then doubles per idle tick up to a ceiling — an effective self-suspend +/// that still rechecks often enough to pick up newly-arrived work. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PollCadence { + /// Interval used while the board has work. + pub base: Duration, + /// Ceiling the backoff saturates at. + pub max_backoff: Duration, + /// Consecutive idle ticks tolerated at `base` before backing off, so a + /// briefly-empty board does not immediately slow down. + pub grace_ticks: u32, +} + +impl Default for PollCadence { + fn default() -> Self { + Self { + base: Duration::from_secs(60), + max_backoff: Duration::from_secs(15 * 60), + grace_ticks: 2, + } + } +} + +impl PollCadence { + /// How long to wait before the next sweep, given the number of consecutive + /// idle ticks so far (`0` right after a dispatch). + /// + /// Monotonic non-decreasing in `idle_ticks`, never above `max_backoff`, and + /// overflow-free for any `u32` streak. + pub fn next_delay(&self, idle_ticks: u32) -> Duration { + let over = idle_ticks.saturating_sub(self.grace_ticks); + if over == 0 { + return self.base; + } + // Doubling per idle tick past the grace window. The shift is clamped so + // a long streak saturates instead of wrapping to a tiny delay. + let factor = 1u64.checked_shl(over.min(20)).unwrap_or(u64::MAX); + let secs = self + .base + .as_secs() + .saturating_mul(factor) + .min(self.max_backoff.as_secs()); + Duration::from_secs(secs) + } +} diff --git a/src/graph/todos/dispatch/test.rs b/src/graph/todos/dispatch/test.rs new file mode 100644 index 00000000..e036f8a7 --- /dev/null +++ b/src/graph/todos/dispatch/test.rs @@ -0,0 +1,440 @@ +//! Unit tests for the dispatch policy: selection, approval, cadence, prompts, +//! and the in-flight run registry. + +use std::time::Duration; + +use serde_json::json; + +use super::prompt::{TaskPromptTools, build_progress_instruction, build_task_prompt}; +use super::registry::{ActiveRun, ActiveRunRegistry}; +use super::select::{ + PollCadence, card_urgency, has_card_in_progress, pick_next_card, requires_plan_approval, +}; +use crate::graph::todos::types::{TaskApprovalMode, TaskBoardCard, TaskCardStatus}; + +fn card(id: &str, status: TaskCardStatus, order: u32) -> TaskBoardCard { + TaskBoardCard { + id: id.to_string(), + status, + order, + ..TaskBoardCard::new(id) + } +} + +fn with_urgency(mut card: TaskBoardCard, urgency: f64) -> TaskBoardCard { + card.source_metadata = Some(json!({ "urgency": urgency })); + card +} + +fn assigned(mut card: TaskBoardCard, agent: &str) -> TaskBoardCard { + card.assigned_agent = Some(agent.to_string()); + card +} + +// ── Selection ─────────────────────────────────────────────────────────────── + +#[test] +fn only_todo_and_ready_cards_are_dispatchable() { + let cards = vec![ + card("done", TaskCardStatus::Done, 0), + card("blocked", TaskCardStatus::Blocked, 1), + card("awaiting", TaskCardStatus::AwaitingApproval, 2), + card("rejected", TaskCardStatus::Rejected, 3), + card("running", TaskCardStatus::InProgress, 4), + ]; + assert!(pick_next_card(&cards, false).is_none()); + + let mut cards = cards; + cards.push(card("ready", TaskCardStatus::Ready, 5)); + assert_eq!(pick_next_card(&cards, false).unwrap().id, "ready"); +} + +#[test] +fn the_most_urgent_card_wins() { + let cards = vec![ + with_urgency(card("low", TaskCardStatus::Todo, 0), 0.1), + with_urgency(card("high", TaskCardStatus::Todo, 1), 0.9), + card("none", TaskCardStatus::Todo, 2), + ]; + assert_eq!(pick_next_card(&cards, false).unwrap().id, "high"); +} + +#[test] +fn equal_urgency_runs_in_board_order() { + let cards = vec![ + with_urgency(card("second", TaskCardStatus::Todo, 5), 0.5), + with_urgency(card("first", TaskCardStatus::Todo, 1), 0.5), + ]; + assert_eq!(pick_next_card(&cards, false).unwrap().id, "first"); + + // Unscored cards tie at 0.0 and follow the same rule. + let cards = vec![ + card("later", TaskCardStatus::Todo, 9), + card("earlier", TaskCardStatus::Todo, 2), + ]; + assert_eq!(pick_next_card(&cards, false).unwrap().id, "earlier"); +} + +#[test] +fn agent_assigned_only_skips_human_authored_cards() { + let cards = vec![ + with_urgency(card("mine", TaskCardStatus::Todo, 0), 0.9), + with_urgency( + assigned(card("agents", TaskCardStatus::Todo, 1), "researcher"), + 0.1, + ), + ]; + + // Unfiltered, urgency wins; filtered, the unassigned card is invisible even + // though it is the more urgent one. + assert_eq!(pick_next_card(&cards, false).unwrap().id, "mine"); + assert_eq!(pick_next_card(&cards, true).unwrap().id, "agents"); +} + +#[test] +fn a_blank_assignee_does_not_count_as_assigned() { + let cards = vec![assigned(card("blank", TaskCardStatus::Todo, 0), " ")]; + assert!(pick_next_card(&cards, true).is_none()); + assert_eq!(pick_next_card(&cards, false).unwrap().id, "blank"); +} + +#[test] +fn an_empty_board_has_nothing_to_dispatch() { + assert!(pick_next_card(&[], false).is_none()); + assert!(!has_card_in_progress(&[])); +} + +#[test] +fn in_progress_detection_gates_a_sweep() { + let idle = vec![card("a", TaskCardStatus::Todo, 0)]; + let busy = vec![card("a", TaskCardStatus::InProgress, 0)]; + assert!(!has_card_in_progress(&idle)); + assert!(has_card_in_progress(&busy)); +} + +#[test] +fn urgency_defaults_to_zero_for_odd_metadata() { + assert_eq!(card_urgency(&card("plain", TaskCardStatus::Todo, 0)), 0.0); + + let mut wrong_type = card("odd", TaskCardStatus::Todo, 0); + wrong_type.source_metadata = Some(json!({ "urgency": "very" })); + assert_eq!(card_urgency(&wrong_type), 0.0); + + let mut absent = card("absent", TaskCardStatus::Todo, 0); + absent.source_metadata = Some(json!({ "provider": "github" })); + assert_eq!(card_urgency(&absent), 0.0); +} + +// ── Approval ──────────────────────────────────────────────────────────────── + +#[test] +fn a_cards_own_approval_mode_outranks_the_global_default() { + // Required holds even when the global switch is off — otherwise a plan + // stamped for review would execute before anyone saw it. + assert!(requires_plan_approval( + false, + Some(&TaskApprovalMode::Required) + )); + assert!(requires_plan_approval( + true, + Some(&TaskApprovalMode::Required) + )); + + // NotRequired means review already happened. + assert!(!requires_plan_approval( + true, + Some(&TaskApprovalMode::NotRequired) + )); + assert!(!requires_plan_approval( + false, + Some(&TaskApprovalMode::NotRequired) + )); +} + +#[test] +fn a_card_with_no_preference_follows_the_global_default() { + assert!(requires_plan_approval(true, None)); + assert!(!requires_plan_approval(false, None)); +} + +// ── Cadence ───────────────────────────────────────────────────────────────── + +#[test] +fn cadence_holds_the_base_interval_through_the_grace_window() { + let cadence = PollCadence::default(); + assert_eq!(cadence.next_delay(0), cadence.base); + assert_eq!(cadence.next_delay(cadence.grace_ticks), cadence.base); +} + +#[test] +fn cadence_doubles_past_the_grace_window() { + let cadence = PollCadence::default(); + let base = cadence.base.as_secs(); + assert_eq!( + cadence.next_delay(cadence.grace_ticks + 1), + Duration::from_secs(base * 2) + ); + assert_eq!( + cadence.next_delay(cadence.grace_ticks + 2), + Duration::from_secs(base * 4) + ); + assert_eq!( + cadence.next_delay(cadence.grace_ticks + 3), + Duration::from_secs(base * 8) + ); +} + +#[test] +fn cadence_is_monotonic_and_never_exceeds_its_ceiling() { + let cadence = PollCadence::default(); + assert_eq!(cadence.next_delay(50), cadence.max_backoff); + // A long idle streak saturates rather than overflowing back to a tiny delay. + assert_eq!(cadence.next_delay(u32::MAX), cadence.max_backoff); + + let mut previous = cadence.next_delay(0); + for idle in 1..40u32 { + let delay = cadence.next_delay(idle); + assert!(delay >= previous, "backoff must not shrink as idle grows"); + assert!( + delay <= cadence.max_backoff, + "backoff must not exceed the ceiling" + ); + previous = delay; + } +} + +#[test] +fn a_zero_grace_cadence_backs_off_from_the_first_idle_tick() { + let cadence = PollCadence { + base: Duration::from_secs(10), + max_backoff: Duration::from_secs(40), + grace_ticks: 0, + }; + assert_eq!(cadence.next_delay(0), Duration::from_secs(10)); + assert_eq!(cadence.next_delay(1), Duration::from_secs(20)); + assert_eq!(cadence.next_delay(2), Duration::from_secs(40)); + assert_eq!(cadence.next_delay(3), Duration::from_secs(40)); +} + +// ── Prompts ───────────────────────────────────────────────────────────────── + +#[test] +fn the_prompt_leads_with_the_objective_and_falls_back_to_the_title() { + let mut card = TaskBoardCard::new("Fix the flaky test"); + let tools = TaskPromptTools::default(); + + let prompt = build_task_prompt(&card, &tools); + assert!(prompt.contains("Fix the flaky test"), "{prompt}"); + + card.objective = Some("Make CI green on main".to_string()); + let prompt = build_task_prompt(&card, &tools); + assert!(prompt.contains("Make CI green on main"), "{prompt}"); + + // A whitespace-only objective is not an objective. + card.objective = Some(" ".to_string()); + assert!(build_task_prompt(&card, &tools).contains("Fix the flaky test")); +} + +#[test] +fn the_prompt_numbers_plan_steps_and_lists_acceptance_criteria() { + let mut card = TaskBoardCard::new("Ship it"); + card.plan = vec!["Reproduce".to_string(), "Fix".to_string()]; + card.acceptance_criteria = vec!["CI is green".to_string()]; + + let prompt = build_task_prompt(&card, &TaskPromptTools::default()); + assert!(prompt.contains("1. Reproduce"), "{prompt}"); + assert!(prompt.contains("2. Fix"), "{prompt}"); + assert!(prompt.contains("- CI is green"), "{prompt}"); + assert!(prompt.contains("Acceptance criteria"), "{prompt}"); +} + +#[test] +fn a_sourced_card_gets_provenance_and_a_write_back_instruction() { + let mut card = TaskBoardCard::new("Triage issue"); + card.source_metadata = Some(json!({ + "provider": "github", + "repo": "tinyhumansai/tinyagents", + "external_id": "412", + "url": "https://example.invalid/412", + })); + + let prompt = build_task_prompt(&card, &TaskPromptTools::default()); + assert!( + prompt.contains("github tinyhumansai/tinyagents#412"), + "{prompt}" + ); + assert!(prompt.contains("memory_recall"), "{prompt}"); + assert!(prompt.contains("https://example.invalid/412"), "{prompt}"); + assert!( + prompt.contains("record the outcome on the upstream source"), + "{prompt}" + ); +} + +#[test] +fn an_id_only_card_gets_no_provenance_line() { + // Without a provider the origin would render as a bare "#7", which tells + // the model nothing — so the whole block is skipped. + let mut card = TaskBoardCard::new("Mystery task"); + card.source_metadata = Some(json!({ "external_id": "7" })); + + let prompt = build_task_prompt(&card, &TaskPromptTools::default()); + assert!(!prompt.contains("#7"), "{prompt}"); + assert!(!prompt.contains("originates from"), "{prompt}"); +} + +#[test] +fn a_host_without_a_memory_tool_still_gets_provenance() { + let mut card = TaskBoardCard::new("Triage issue"); + card.source_metadata = Some(json!({ "provider": "linear", "external_id": "ENG-1" })); + let tools = TaskPromptTools { + memory_recall: None, + ..TaskPromptTools::default() + }; + + let prompt = build_task_prompt(&card, &tools); + assert!(prompt.contains("originates from linear#ENG-1"), "{prompt}"); + assert!(!prompt.contains("memory_recall"), "{prompt}"); +} + +#[test] +fn the_progress_instruction_names_the_card_board_and_tool() { + let tools = TaskPromptTools { + update_task: "board_update".to_string(), + ..TaskPromptTools::default() + }; + let instruction = build_progress_instruction("task-9", "user-tasks", &tools); + + assert!(instruction.contains("task-9"), "{instruction}"); + assert!(instruction.contains("user-tasks"), "{instruction}"); + assert!(instruction.contains("board_update"), "{instruction}"); + // Blocking is the sanctioned way out, and it must be spelled out. + assert!(instruction.contains("status: blocked"), "{instruction}"); + assert!(instruction.contains("do NOT guess"), "{instruction}"); +} + +// ── Registry ──────────────────────────────────────────────────────────────── + +fn spawn_pending() -> tokio::task::JoinHandle<()> { + tokio::spawn(async { std::future::pending::<()>().await }) +} + +async fn active_run( + run_id: &str, + card_id: &str, +) -> (ActiveRun<&'static str>, tokio::task::JoinHandle<()>) { + let handle = spawn_pending(); + let (heartbeat_cancel, _rx) = tokio::sync::watch::channel(false); + ( + ActiveRun { + run_id: run_id.to_string(), + card_id: card_id.to_string(), + abort: handle.abort_handle(), + heartbeat_cancel, + context: "board", + }, + handle, + ) +} + +#[tokio::test] +async fn only_one_taker_owns_a_runs_cleanup() { + let registry = ActiveRunRegistry::new(); + let (run, _handle) = active_run("run-1", "task-1").await; + assert!(registry.register("thread-1", run).is_none()); + assert!(registry.contains("thread-1")); + assert_eq!(registry.len(), 1); + + // The natural completion and a concurrent cancel both call `take`; exactly + // one gets the entry, so the terminal write-back happens once. + let first = registry.take("thread-1"); + let second = registry.take("thread-1"); + assert!(first.is_some()); + assert!(second.is_none()); + assert!(registry.is_empty()); +} + +#[tokio::test] +async fn a_scoped_cancel_ignores_a_superseded_run() { + let registry = ActiveRunRegistry::new(); + let (run, _handle) = active_run("run-new", "task-1").await; + registry.register("thread-1", run); + + // A cancel for a run that has already been replaced must not tear down the + // run that took its place. + assert!(registry.take_if("thread-1", Some("run-old")).is_none()); + assert!(registry.contains("thread-1")); + + let taken = registry + .take_if("thread-1", Some("run-new")) + .expect("matching run"); + assert_eq!(taken.run_id, "run-new"); +} + +#[tokio::test] +async fn an_unscoped_cancel_takes_whatever_is_running() { + let registry = ActiveRunRegistry::new(); + let (run, _handle) = active_run("run-1", "task-1").await; + registry.register("thread-1", run); + + assert!(registry.take_if("thread-1", None).is_some()); + // And on an idle thread it is simply a no-op. + assert!(registry.take_if("thread-1", None).is_none()); + assert!(registry.take_if("unknown-thread", Some("run-1")).is_none()); +} + +#[tokio::test] +async fn registering_over_a_live_run_hands_back_the_displaced_one() { + let registry = ActiveRunRegistry::new(); + let (first, _h1) = active_run("run-1", "task-1").await; + let (second, _h2) = active_run("run-2", "task-2").await; + registry.register("thread-1", first); + + let displaced = registry + .register("thread-1", second) + .expect("displaced run"); + assert_eq!(displaced.run_id, "run-1"); + assert_eq!(registry.len(), 1); + assert_eq!(registry.thread_ids(), vec!["thread-1".to_string()]); +} + +#[tokio::test] +async fn cancelling_aborts_the_task_and_stops_the_heartbeat() { + let registry = ActiveRunRegistry::new(); + let handle = spawn_pending(); + let (heartbeat_cancel, mut heartbeat_rx) = tokio::sync::watch::channel(false); + registry.register( + "thread-1", + ActiveRun { + run_id: "run-1".to_string(), + card_id: "task-1".to_string(), + abort: handle.abort_handle(), + heartbeat_cancel, + context: (), + }, + ); + + let run = registry.take("thread-1").expect("live run"); + run.cancel(); + + assert!(handle.await.unwrap_err().is_cancelled()); + assert!(heartbeat_rx.changed().await.is_ok()); + assert!(*heartbeat_rx.borrow()); +} + +#[tokio::test] +async fn draining_returns_every_live_run() { + let registry = ActiveRunRegistry::new(); + let (first, _h1) = active_run("run-1", "task-1").await; + let (second, _h2) = active_run("run-2", "task-2").await; + registry.register("thread-1", first); + registry.register("thread-2", second); + + let mut drained: Vec = registry.drain().into_iter().map(|(id, _)| id).collect(); + drained.sort(); + assert_eq!( + drained, + vec!["thread-1".to_string(), "thread-2".to_string()] + ); + assert!(registry.is_empty()); +} diff --git a/src/graph/todos/mod.rs b/src/graph/todos/mod.rs index 04f7c0de..17b06fed 100644 --- a/src/graph/todos/mod.rs +++ b/src/graph/todos/mod.rs @@ -8,10 +8,18 @@ //! single-`InProgress` invariant ([`store`]), and the model-facing multiplexer //! tool ([`tool`]). //! +//! Two layers sit on top of the board for hosts that run cards autonomously: +//! [`runs`] records who claimed a card, heartbeats while they work, and hands +//! the card back when a worker goes silent; [`dispatch`] is the scheduling +//! policy — which card is next, whether it needs approval first, what prompt +//! its run gets, and how an in-flight run is cancelled. +//! //! Ported from OpenHuman's task board / `todos` modules, minus the app-specific //! coupling (progress events, RPC envelopes, in-memory scratch fallback): a //! board is always `(Store, thread_id)`. +pub mod dispatch; +pub mod runs; pub mod store; mod tool; mod types; diff --git a/src/graph/todos/runs/README.md b/src/graph/todos/runs/README.md new file mode 100644 index 00000000..ad2b6124 --- /dev/null +++ b/src/graph/todos/runs/README.md @@ -0,0 +1,94 @@ +# graph::todos::runs + +The **claim / heartbeat / reclaim** layer over a task board. A +[`TaskBoardCard`](../types.rs) says what needs doing; a `TaskRun` says who is +doing it right now, since when, and how it ended. + +Ported from OpenHuman's `threads::todos::runs`, minus the app coupling (domain +event bus, workspace-file layout). Runs are addressed like boards — +`(Store, thread_id)` — and share the board's timestamp format (unix-epoch +milliseconds as a string), so a board and its run log live side by side in one +store. + +## Why it exists + +A worker that dies mid-card leaves the card `InProgress` forever, and the +board's single-in-progress rule then wedges the whole thread. The run log makes +that recoverable: a sweep notices the silence, closes the run, and hands the +card back. + +The bound matters as much as the recovery. A card that keeps killing its workers +would otherwise cycle through them forever, so after +`RunLimits::max_reclaim_count` reclaims it parks at `Blocked` with a blocker +naming the last reason. + +## Data model (`types.rs`) + +- `TaskRun { run_id, card_id, claimed_by, claim_token, started_at, + last_heartbeat_at, completed_at, outcome, error, evidence }` (serde + `camelCase`); `is_active()` is "no `completed_at`". +- `RunOutcome { Success, Failed, Reclaimed }`. +- `RunLimits { heartbeat_stale_secs, claim_ttl_secs, max_reclaim_count }`, + defaulting to 300s / 3600s / 3. +- `ReclaimResult { reclaimed_count, blocked_count, details }` and + `ReclaimDetail { run_id, card_id, reason, new_card_status }` — the crate emits + no events, so a host derives its own from `details`. +- `staleness_reason(run, now_ms, limits)` — the policy as a pure, + clock-injected function: TTL is checked before heartbeat (a run that is both + reports the more fundamental reason), and an unparsable stamp reads as + healthy so a corrupt record never yanks a card from a live worker. + +## Store (`store.rs`) + +One `Vec` per thread under the `graph.todos.runs` namespace, keyed by +the hex-encoded thread id. Every mutation is `load → mutate → put` under a +per-thread async mutex (single-process, like the board store); the run lock map +is separate from the board's, so a run write never blocks a card write. + +| Function | Role | +| --- | --- | +| `create_run` | Open a claim. Rejects a duplicate caller-supplied `run_id`. | +| `update_heartbeat` | Tick liveness. Errors on an unknown or finished run. | +| `complete_run` | Close with an outcome, error, and evidence. | +| `list_runs` / `get_run` | Read the log, optionally filtered to one card. | +| `find_stale_runs` | Active runs judged stale, each with its reason. | +| `count_reclaims_for_card` | How often a card has been handed back. | +| `reclaim_stale` | The sweep: close stale runs, return or park their cards. | +| `spawn_heartbeat_task` | Background ticker; stops on cancel or completion. | + +Claiming does **not** touch the card. The caller moves it through +`todo_store::claim_card`, which is what enforces the single-`InProgress` rule. + +## Example + +```rust,ignore +use tinyagents::{RunLimits, RunOutcome, TaskCardStatus, task_run_store, todo_store}; + +todo_store::claim_card(&store, thread, &card_id, + &[TaskCardStatus::Todo], TaskCardStatus::InProgress).await?; +let run = task_run_store::create_run(&store, thread, None, &card_id, "worker-a").await?; + +task_run_store::spawn_heartbeat_task( + store.clone(), thread.into(), run.run_id.clone(), cancel_rx, + task_run_store::DEFAULT_HEARTBEAT_TICK, +); + +// … work … +task_run_store::complete_run( + &store, thread, &run.run_id, RunOutcome::Success, None, vec!["pr #12".into()], +).await?; + +// Elsewhere, on a timer: hand back anything whose worker went quiet. +task_run_store::reclaim_stale(&store, thread, &RunLimits::default()).await?; +``` + +## Files + +| File | Role | +| --- | --- | +| `types.rs` | `TaskRun`, `RunOutcome`, `RunLimits`, reclaim reports, `staleness_reason`. | +| `store.rs` | `Store`-backed lifecycle, the reclaim sweep, the heartbeat task. | +| `test.rs` | Unit tests (lifecycle, staleness policy, reclaim, serialization). | + +Feature coverage lives in `tests/feature_graph_task_runs.rs`; the dispatch loop +that drives it end to end is in `tests/e2e_graph_task_dispatch.rs`. diff --git a/src/graph/todos/runs/mod.rs b/src/graph/todos/runs/mod.rs new file mode 100644 index 00000000..dfe05cb0 --- /dev/null +++ b/src/graph/todos/runs/mod.rs @@ -0,0 +1,47 @@ +//! **Task runs**: the claim, heartbeat, and reclaim layer over a task board. +//! +//! A [`TaskBoardCard`](super::TaskBoardCard) records what needs doing; +//! [`TaskRun`] records an attempt at doing it. A worker claims a card, writes a +//! run, ticks a heartbeat while it works, and closes the run with an outcome. +//! If it dies without closing, the run goes stale and +//! [`store::reclaim_stale`] hands the card back to the queue — bounded by +//! [`RunLimits::max_reclaim_count`] so a card that keeps killing its workers +//! parks at `Blocked` rather than cycling forever. +//! +//! Runs are addressed exactly like boards — `(Store, thread_id)` — and share +//! the board's timestamp format (unix-epoch milliseconds as a string). The +//! staleness policy itself is a pure, clock-injected function +//! ([`staleness_reason`]) so it can be tested without waiting. +//! +//! ```no_run +//! use std::sync::Arc; +//! use tinyagents::graph::todos::runs::{RunLimits, RunOutcome, store as runs}; +//! use tinyagents::harness::store::{InMemoryStore, Store}; +//! +//! # async fn demo() -> tinyagents::error::Result<()> { +//! let store: Arc = Arc::new(InMemoryStore::new()); +//! let run = runs::create_run(&store, "thread-1", None, "task-1", "worker-a").await?; +//! runs::update_heartbeat(&store, "thread-1", &run.run_id).await?; +//! runs::complete_run(&store, "thread-1", &run.run_id, RunOutcome::Success, None, vec![]).await?; +//! // A sweep finds nothing to reclaim: the run reached a terminal state. +//! let swept = runs::reclaim_stale(&store, "thread-1", &RunLimits::default()).await?; +//! assert_eq!(swept.reclaimed_count, 0); +//! # Ok(()) +//! # } +//! ``` + +pub mod store; +mod types; + +pub use store::{ + DEFAULT_HEARTBEAT_TICK, RUNS_NAMESPACE, complete_run, count_reclaims_for_card, create_run, + find_stale_runs, get_run, import_if_absent, list_runs, reclaim_stale, spawn_heartbeat_task, + update_heartbeat, +}; +pub use types::{ + DEFAULT_CLAIM_TTL_SECS, DEFAULT_HEARTBEAT_STALE_SECS, DEFAULT_MAX_RECLAIM_COUNT, ReclaimDetail, + ReclaimResult, RunLimits, RunOutcome, TaskRun, staleness_reason, +}; + +#[cfg(test)] +mod test; diff --git a/src/graph/todos/runs/store.rs b/src/graph/todos/runs/store.rs new file mode 100644 index 00000000..507ca7fa --- /dev/null +++ b/src/graph/todos/runs/store.rs @@ -0,0 +1,421 @@ +//! Persistence and lifecycle for [`TaskRun`]s, on the harness +//! [`Store`](crate::harness::store::Store). +//! +//! Each thread's runs are a single serialized `Vec` under the +//! [`RUNS_NAMESPACE`] namespace, keyed by the hex-encoded thread id — the same +//! addressing the board itself uses in +//! [`graph::todos::store`](crate::graph::todos::store), so a board and its run +//! log live side by side in one store. Every mutation runs +//! `load → mutate → put` under a **per-thread async mutex**, so the +//! read-modify-write is atomic within the process (the same single-process +//! caveat as the board store). +//! +//! The run log is what makes a wedged worker recoverable: [`reclaim_stale`] +//! sweeps runs whose heartbeat or claim has aged out, closes them as +//! [`RunOutcome::Reclaimed`], and moves their card back to `Todo` — or parks it +//! at `Blocked` once the card has burned through +//! [`RunLimits::max_reclaim_count`] attempts, so a permanently poisonous card +//! stops cycling through workers. + +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use tokio::sync::Mutex; + +use super::types::{ + ReclaimDetail, ReclaimResult, RunLimits, RunOutcome, TaskRun, new_claim_token, new_run_id, + staleness_reason, +}; +use crate::error::{Result, TinyAgentsError}; +use crate::graph::thread_locks::ThreadLockMap; +use crate::graph::todos::store as board; +use crate::graph::todos::types::{CardPatch, TaskCardStatus, now_millis, now_stamp}; +use crate::harness::store::Store; + +/// The [`Store`] namespace holding one `Vec` per thread. +pub const RUNS_NAMESPACE: &str = "graph.todos.runs"; + +/// Default cadence of the background heartbeat spawned by +/// [`spawn_heartbeat_task`]. +pub const DEFAULT_HEARTBEAT_TICK: Duration = Duration::from_secs(30); + +/// Serialises `load → mutate → put` per thread. Kept separate from the board's +/// lock map so a run write never blocks a card write (and vice versa) — the +/// reclaim path deliberately takes them one after the other, never nested. +fn runs_lock(thread_id: &str) -> Arc> { + static LOCKS: OnceLock = OnceLock::new(); + LOCKS + .get_or_init(|| ThreadLockMap::new("task run lock map")) + .lock_for(thread_id) +} + +/// Hex-encodes the thread id into a [`Store`]-safe key. +fn key(thread_id: &str) -> String { + thread_id + .as_bytes() + .iter() + .map(|b| format!("{b:02x}")) + .collect() +} + +fn validate_thread_id(thread_id: &str) -> Result { + let trimmed = thread_id.trim(); + if trimmed.is_empty() { + return Err(TinyAgentsError::Validation( + "task run thread_id must not be empty or whitespace".to_string(), + )); + } + Ok(trimmed.to_string()) +} + +async fn load(store: &Arc, thread_id: &str) -> Result> { + match store.get(RUNS_NAMESPACE, &key(thread_id)).await? { + Some(value) => Ok(serde_json::from_value(value)?), + None => Ok(Vec::new()), + } +} + +async fn save(store: &Arc, thread_id: &str, runs: &[TaskRun]) -> Result<()> { + store + .put(RUNS_NAMESPACE, &key(thread_id), serde_json::to_value(runs)?) + .await +} + +/// Claim `card_id` for `claimed_by` and record the claim as a fresh active run. +/// +/// `run_id` lets a caller supply its own id (to correlate the run with an +/// external session); `None` mints one. Claiming does **not** touch the card — +/// the caller moves it to `InProgress` through the board store, which is what +/// enforces the single-`InProgress` invariant. +pub async fn create_run( + store: &Arc, + thread_id: &str, + run_id: Option<&str>, + card_id: &str, + claimed_by: &str, +) -> Result { + let thread_id = validate_thread_id(thread_id)?; + let lock = runs_lock(&thread_id); + let _guard = lock.lock().await; + + let now = now_stamp(); + let run = TaskRun { + run_id: run_id.map(str::to_string).unwrap_or_else(new_run_id), + card_id: card_id.to_string(), + claimed_by: claimed_by.to_string(), + claim_token: new_claim_token(), + started_at: now.clone(), + last_heartbeat_at: now, + completed_at: None, + outcome: None, + error: None, + evidence: Vec::new(), + }; + + let mut runs = load(store, &thread_id).await?; + if runs.iter().any(|existing| existing.run_id == run.run_id) { + return Err(TinyAgentsError::Validation(format!( + "task run '{}' already exists on thread '{thread_id}'", + run.run_id + ))); + } + runs.push(run.clone()); + save(store, &thread_id, &runs).await?; + + tracing::info!( + thread_id = %thread_id, + run_id = %run.run_id, + card_id = %card_id, + claimed_by = %claimed_by, + "[graph:todos:runs] claim created" + ); + Ok(run) +} + +/// Import a run log only when the thread has none. +/// +/// The existence check and the write share the normal per-thread lock, and an +/// existing log is left untouched even when it uses a newer or undecodable +/// schema — which makes this suitable for a one-time migration off a legacy +/// store. Merging is deliberately not offered: two histories for one thread +/// would double-count the reclaims [`reclaim_stale`]'s budget is derived from. +/// +/// Returns whether the runs were written. +pub async fn import_if_absent( + store: &Arc, + thread_id: &str, + runs: Vec, +) -> Result { + let thread_id = validate_thread_id(thread_id)?; + let lock = runs_lock(&thread_id); + let _guard = lock.lock().await; + if store.get(RUNS_NAMESPACE, &key(&thread_id)).await?.is_some() { + return Ok(false); + } + save(store, &thread_id, &runs).await?; + Ok(true) +} + +/// Refresh the liveness tick of an **active** run. +/// +/// Errors when the run is unknown or already terminal — the heartbeat loop +/// treats that as its stop signal rather than resurrecting a finished run. +pub async fn update_heartbeat(store: &Arc, thread_id: &str, run_id: &str) -> Result<()> { + let thread_id = validate_thread_id(thread_id)?; + let lock = runs_lock(&thread_id); + let _guard = lock.lock().await; + + let mut runs = load(store, &thread_id).await?; + let run = runs + .iter_mut() + .find(|run| run.run_id == run_id && run.is_active()) + .ok_or_else(|| { + TinyAgentsError::Validation(format!( + "active task run '{run_id}' not found on thread '{thread_id}'" + )) + })?; + run.last_heartbeat_at = now_stamp(); + save(store, &thread_id, &runs).await +} + +/// Close an **active** run with a terminal outcome, returning the closed record. +pub async fn complete_run( + store: &Arc, + thread_id: &str, + run_id: &str, + outcome: RunOutcome, + error: Option, + evidence: Vec, +) -> Result { + let thread_id = validate_thread_id(thread_id)?; + let lock = runs_lock(&thread_id); + let _guard = lock.lock().await; + + let mut runs = load(store, &thread_id).await?; + let run = runs + .iter_mut() + .find(|run| run.run_id == run_id && run.is_active()) + .ok_or_else(|| { + TinyAgentsError::Validation(format!( + "active task run '{run_id}' not found on thread '{thread_id}'" + )) + })?; + run.completed_at = Some(now_stamp()); + run.outcome = Some(outcome); + run.error = error; + run.evidence = evidence; + let completed = run.clone(); + save(store, &thread_id, &runs).await?; + + tracing::info!( + thread_id = %thread_id, + run_id = %run_id, + outcome = ?completed.outcome, + "[graph:todos:runs] run completed" + ); + Ok(completed) +} + +/// Every run recorded for the thread, oldest first; filtered to one card when +/// `card_id` is given. +pub async fn list_runs( + store: &Arc, + thread_id: &str, + card_id: Option<&str>, +) -> Result> { + let thread_id = validate_thread_id(thread_id)?; + let lock = runs_lock(&thread_id); + let _guard = lock.lock().await; + + let runs = load(store, &thread_id).await?; + Ok(match card_id { + Some(card_id) => runs.into_iter().filter(|r| r.card_id == card_id).collect(), + None => runs, + }) +} + +/// One run by id, or `None` when the thread has never recorded it. +pub async fn get_run( + store: &Arc, + thread_id: &str, + run_id: &str, +) -> Result> { + let thread_id = validate_thread_id(thread_id)?; + let lock = runs_lock(&thread_id); + let _guard = lock.lock().await; + + Ok(load(store, &thread_id) + .await? + .into_iter() + .find(|run| run.run_id == run_id)) +} + +/// Active runs judged stale under `limits`, each paired with the reason. +pub async fn find_stale_runs( + store: &Arc, + thread_id: &str, + limits: &RunLimits, +) -> Result> { + let thread_id = validate_thread_id(thread_id)?; + let lock = runs_lock(&thread_id); + let _guard = lock.lock().await; + + let now = now_millis(); + Ok(load(store, &thread_id) + .await? + .into_iter() + .filter(TaskRun::is_active) + .filter_map(|run| staleness_reason(&run, now, limits).map(|reason| (run, reason))) + .collect()) +} + +/// How many times `card_id` has already been reclaimed on this thread. +pub async fn count_reclaims_for_card( + store: &Arc, + thread_id: &str, + card_id: &str, +) -> Result { + Ok(list_runs(store, thread_id, Some(card_id)) + .await? + .iter() + .filter(|run| run.outcome == Some(RunOutcome::Reclaimed)) + .count() as u32) +} + +/// Sweep the thread's stale runs: close each as [`RunOutcome::Reclaimed`], then +/// move its card back to `Todo` so a later dispatch can pick it up — or park it +/// at `Blocked` with a diagnostic blocker once the card has exceeded +/// [`RunLimits::max_reclaim_count`] reclaims. +/// +/// Best-effort per run: a card write that fails is logged and skipped, leaving +/// the rest of the sweep to proceed. The returned [`ReclaimResult`] is the +/// crate's only report — it publishes no events of its own, so a host that +/// needs them derives them from `details`. +pub async fn reclaim_stale( + store: &Arc, + thread_id: &str, + limits: &RunLimits, +) -> Result { + let thread_id = validate_thread_id(thread_id)?; + let stale = find_stale_runs(store, &thread_id, limits).await?; + if stale.is_empty() { + return Ok(ReclaimResult::default()); + } + + let mut result = ReclaimResult::default(); + for (run, reason) in &stale { + if let Err(error) = complete_run( + store, + &thread_id, + &run.run_id, + RunOutcome::Reclaimed, + Some(reason.clone()), + Vec::new(), + ) + .await + { + tracing::warn!( + thread_id = %thread_id, + run_id = %run.run_id, + %error, + "[graph:todos:runs] could not close stale run" + ); + continue; + } + + // Counted *after* closing this run, so the current reclaim is included: + // the card parks at `Blocked` on the reclaim that reaches the limit. + let reclaims = count_reclaims_for_card(store, &thread_id, &run.card_id).await?; + let park = reclaims >= limits.max_reclaim_count; + let status = if park { + TaskCardStatus::Blocked + } else { + TaskCardStatus::Todo + }; + let patch = CardPatch { + status: Some(status), + blocker: park.then(|| { + format!( + "Reclaimed {reclaims} time(s), exceeding limit of {}. Last reclaim reason: {reason}", + limits.max_reclaim_count + ) + }), + ..Default::default() + }; + + match board::edit(store, &thread_id, &run.card_id, patch).await { + Ok(_) => { + if park { + result.blocked_count += 1; + } else { + result.reclaimed_count += 1; + } + result.details.push(ReclaimDetail { + run_id: run.run_id.clone(), + card_id: run.card_id.clone(), + reason: reason.clone(), + new_card_status: status.as_str().to_string(), + }); + tracing::info!( + thread_id = %thread_id, + run_id = %run.run_id, + card_id = %run.card_id, + new_status = status.as_str(), + reclaims, + %reason, + "[graph:todos:runs] card reclaimed" + ); + } + Err(error) => tracing::warn!( + thread_id = %thread_id, + run_id = %run.run_id, + card_id = %run.card_id, + %error, + "[graph:todos:runs] could not move card after reclaim" + ), + } + } + Ok(result) +} + +/// Spawn the background heartbeat for an in-flight run. +/// +/// Ticks [`update_heartbeat`] every `tick` until either the run stops being +/// active (the tick errors, which is the normal end after +/// [`complete_run`]) or `cancel` fires. The first immediate tick is skipped so +/// the freshly-created run is not written twice. +pub fn spawn_heartbeat_task( + store: Arc, + thread_id: String, + run_id: String, + mut cancel: tokio::sync::watch::Receiver, + tick: Duration, +) { + tokio::spawn(async move { + let mut ticker = tokio::time::interval(tick); + ticker.tick().await; // `interval` fires immediately; skip that one. + loop { + tokio::select! { + _ = ticker.tick() => { + if let Err(error) = update_heartbeat(&store, &thread_id, &run_id).await { + tracing::debug!( + thread_id = %thread_id, + run_id = %run_id, + %error, + "[graph:todos:runs] heartbeat stopped (run is no longer active)" + ); + break; + } + } + _ = cancel.changed() => { + tracing::debug!( + thread_id = %thread_id, + run_id = %run_id, + "[graph:todos:runs] heartbeat cancelled" + ); + break; + } + } + } + }); +} diff --git a/src/graph/todos/runs/test.rs b/src/graph/todos/runs/test.rs new file mode 100644 index 00000000..c97458cf --- /dev/null +++ b/src/graph/todos/runs/test.rs @@ -0,0 +1,441 @@ +//! Unit tests for the task-run claim/heartbeat/reclaim layer. + +use std::sync::Arc; + +use super::store::{ + complete_run, count_reclaims_for_card, create_run, find_stale_runs, get_run, list_runs, + reclaim_stale, update_heartbeat, +}; +use super::types::{RunLimits, RunOutcome, TaskRun, staleness_reason}; +use crate::graph::todos::store as board; +use crate::graph::todos::types::TaskCardStatus; +use crate::harness::store::{InMemoryStore, Store}; + +fn store() -> Arc { + Arc::new(InMemoryStore::new()) +} + +/// A run with hand-set timestamps, so staleness is testable without waiting. +fn run_at(run_id: &str, card_id: &str, started_ms: u64, heartbeat_ms: u64) -> TaskRun { + TaskRun { + run_id: run_id.to_string(), + card_id: card_id.to_string(), + claimed_by: "worker".to_string(), + claim_token: "token".to_string(), + started_at: started_ms.to_string(), + last_heartbeat_at: heartbeat_ms.to_string(), + completed_at: None, + outcome: None, + error: None, + evidence: Vec::new(), + } +} + +async fn seed_card(store: &Arc, thread_id: &str, title: &str) -> String { + let snapshot = board::add(store, thread_id, title, Default::default()) + .await + .expect("add card"); + snapshot.cards.last().expect("card added").id.clone() +} + +#[tokio::test] +async fn create_records_an_active_run_listable_by_card() { + let store = store(); + let run = create_run(&store, "thread-1", None, "task-1", "worker-a") + .await + .unwrap(); + + assert_eq!(run.card_id, "task-1"); + assert_eq!(run.claimed_by, "worker-a"); + assert!(run.is_active()); + assert!(!run.claim_token.is_empty()); + assert!(!run.run_id.is_empty()); + + assert_eq!(list_runs(&store, "thread-1", None).await.unwrap().len(), 1); + assert_eq!( + list_runs(&store, "thread-1", Some("task-1")) + .await + .unwrap() + .len(), + 1 + ); + assert!( + list_runs(&store, "thread-1", Some("task-other")) + .await + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn runs_are_scoped_to_their_thread() { + let store = store(); + create_run(&store, "thread-a", Some("run-a"), "task-1", "worker") + .await + .unwrap(); + create_run(&store, "thread-b", Some("run-b"), "task-1", "worker") + .await + .unwrap(); + + let a = list_runs(&store, "thread-a", None).await.unwrap(); + assert_eq!(a.len(), 1); + assert_eq!(a[0].run_id, "run-a"); + assert!( + get_run(&store, "thread-a", "run-b") + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn a_supplied_run_id_cannot_be_reused_on_one_thread() { + let store = store(); + create_run(&store, "thread-1", Some("run-1"), "task-1", "worker") + .await + .unwrap(); + let error = create_run(&store, "thread-1", Some("run-1"), "task-2", "worker") + .await + .expect_err("duplicate run id must be rejected"); + assert!(error.to_string().contains("already exists")); +} + +#[tokio::test] +async fn create_rejects_a_blank_thread_id() { + let store = store(); + assert!( + create_run(&store, " ", None, "task-1", "worker") + .await + .is_err() + ); +} + +#[tokio::test] +async fn heartbeat_advances_the_liveness_stamp() { + let store = store(); + let run = create_run(&store, "thread-1", None, "task-1", "worker") + .await + .unwrap(); + + update_heartbeat(&store, "thread-1", &run.run_id) + .await + .unwrap(); + + let after = get_run(&store, "thread-1", &run.run_id) + .await + .unwrap() + .unwrap(); + assert!(after.last_heartbeat_at >= run.last_heartbeat_at); + assert!(after.is_active()); +} + +#[tokio::test] +async fn heartbeat_and_completion_reject_a_finished_run() { + let store = store(); + let run = create_run(&store, "thread-1", None, "task-1", "worker") + .await + .unwrap(); + complete_run( + &store, + "thread-1", + &run.run_id, + RunOutcome::Success, + None, + vec!["pr #12".to_string()], + ) + .await + .unwrap(); + + // A finished run is not resurrected, and it is not completed twice. + assert!( + update_heartbeat(&store, "thread-1", &run.run_id) + .await + .is_err() + ); + assert!( + complete_run( + &store, + "thread-1", + &run.run_id, + RunOutcome::Failed, + None, + vec![] + ) + .await + .is_err() + ); +} + +#[tokio::test] +async fn completion_records_outcome_error_and_evidence() { + let store = store(); + let run = create_run(&store, "thread-1", None, "task-1", "worker") + .await + .unwrap(); + let done = complete_run( + &store, + "thread-1", + &run.run_id, + RunOutcome::Failed, + Some("provider timed out".to_string()), + vec!["log line".to_string()], + ) + .await + .unwrap(); + + assert!(!done.is_active()); + assert_eq!(done.outcome, Some(RunOutcome::Failed)); + assert_eq!(done.error.as_deref(), Some("provider timed out")); + assert_eq!(done.evidence, vec!["log line".to_string()]); + assert!(done.completed_at.is_some()); +} + +#[test] +fn staleness_reports_ttl_before_heartbeat() { + let limits = RunLimits { + heartbeat_stale_secs: 10, + claim_ttl_secs: 60, + max_reclaim_count: 3, + }; + let now = 1_000_000u64; + + // Healthy: young claim, fresh heartbeat. + let healthy = run_at("r", "c", now - 5_000, now - 1_000); + assert!(staleness_reason(&healthy, now, &limits).is_none()); + + // Silent worker, claim still inside its TTL. + let silent = run_at("r", "c", now - 30_000, now - 30_000); + let reason = staleness_reason(&silent, now, &limits).expect("stale heartbeat"); + assert!(reason.contains("heartbeat stale"), "{reason}"); + + // Both aged out: the TTL is the reason reported. + let expired = run_at("r", "c", now - 120_000, now - 120_000); + let reason = staleness_reason(&expired, now, &limits).expect("expired claim"); + assert!(reason.contains("claim TTL expired"), "{reason}"); +} + +#[test] +fn staleness_treats_an_unparsable_stamp_as_healthy() { + // A corrupt record must never cause a live worker's card to be yanked away. + let mut run = run_at("r", "c", 0, 0); + run.started_at = "not-a-timestamp".to_string(); + assert!(staleness_reason(&run, 9_999_999, &RunLimits::default()).is_none()); +} + +#[tokio::test] +async fn find_stale_runs_ignores_completed_runs() { + let store = store(); + let run = create_run(&store, "thread-1", None, "task-1", "worker") + .await + .unwrap(); + complete_run( + &store, + "thread-1", + &run.run_id, + RunOutcome::Success, + None, + vec![], + ) + .await + .unwrap(); + + // Zero limits would make any *active* run stale; this one is finished. + let limits = RunLimits { + heartbeat_stale_secs: 0, + claim_ttl_secs: 0, + max_reclaim_count: 3, + }; + assert!( + find_stale_runs(&store, "thread-1", &limits) + .await + .unwrap() + .is_empty() + ); +} + +/// Force `run_id` to look ancient by rewriting its stamps in place. +async fn age_run(store: &Arc, thread_id: &str, run_id: &str) { + let mut runs = list_runs(store, thread_id, None).await.unwrap(); + for run in runs.iter_mut().filter(|r| r.run_id == run_id) { + run.started_at = "0".to_string(); + run.last_heartbeat_at = "0".to_string(); + } + let key: String = thread_id + .as_bytes() + .iter() + .map(|b| format!("{b:02x}")) + .collect(); + store + .put( + super::store::RUNS_NAMESPACE, + &key, + serde_json::to_value(&runs).unwrap(), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn reclaim_returns_a_wedged_card_to_the_queue() { + let store = store(); + let thread_id = "thread-reclaim"; + let card_id = seed_card(&store, thread_id, "ship the thing").await; + board::claim_card( + &store, + thread_id, + &card_id, + &[TaskCardStatus::Todo], + TaskCardStatus::InProgress, + ) + .await + .unwrap(); + let run = create_run(&store, thread_id, None, &card_id, "worker") + .await + .unwrap(); + age_run(&store, thread_id, &run.run_id).await; + + let result = reclaim_stale(&store, thread_id, &RunLimits::default()) + .await + .unwrap(); + + assert_eq!(result.reclaimed_count, 1); + assert_eq!(result.blocked_count, 0); + assert_eq!(result.details.len(), 1); + assert_eq!(result.details[0].card_id, card_id); + assert_eq!(result.details[0].new_card_status, "todo"); + + let closed = get_run(&store, thread_id, &run.run_id) + .await + .unwrap() + .unwrap(); + assert_eq!(closed.outcome, Some(RunOutcome::Reclaimed)); + assert!(closed.error.is_some(), "the reason is recorded on the run"); + + let snapshot = board::list(&store, thread_id).await.unwrap(); + assert_eq!(snapshot.cards[0].status, TaskCardStatus::Todo); +} + +#[tokio::test] +async fn a_card_that_keeps_wedging_workers_parks_as_blocked() { + let store = store(); + let thread_id = "thread-poison"; + let card_id = seed_card(&store, thread_id, "poison card").await; + let limits = RunLimits { + max_reclaim_count: 2, + ..RunLimits::default() + }; + + // Two reclaims are tolerated; the second reaches the limit and parks it. + for expected_status in ["todo", "blocked"] { + board::claim_card( + &store, + thread_id, + &card_id, + &[TaskCardStatus::Todo, TaskCardStatus::Blocked], + TaskCardStatus::InProgress, + ) + .await + .unwrap(); + let run = create_run(&store, thread_id, None, &card_id, "worker") + .await + .unwrap(); + age_run(&store, thread_id, &run.run_id).await; + let result = reclaim_stale(&store, thread_id, &limits).await.unwrap(); + assert_eq!(result.details[0].new_card_status, expected_status); + } + + assert_eq!( + count_reclaims_for_card(&store, thread_id, &card_id) + .await + .unwrap(), + 2 + ); + let snapshot = board::list(&store, thread_id).await.unwrap(); + assert_eq!(snapshot.cards[0].status, TaskCardStatus::Blocked); + let blocker = snapshot.cards[0].blocker.as_deref().unwrap_or_default(); + assert!(blocker.contains("exceeding limit of 2"), "{blocker}"); +} + +#[tokio::test] +async fn reclaim_leaves_a_healthy_run_alone() { + let store = store(); + let thread_id = "thread-healthy"; + let card_id = seed_card(&store, thread_id, "in flight").await; + board::claim_card( + &store, + thread_id, + &card_id, + &[TaskCardStatus::Todo], + TaskCardStatus::InProgress, + ) + .await + .unwrap(); + create_run(&store, thread_id, None, &card_id, "worker") + .await + .unwrap(); + + let result = reclaim_stale(&store, thread_id, &RunLimits::default()) + .await + .unwrap(); + + assert_eq!(result, Default::default()); + let snapshot = board::list(&store, thread_id).await.unwrap(); + assert_eq!(snapshot.cards[0].status, TaskCardStatus::InProgress); +} + +#[tokio::test] +async fn reclaim_on_a_thread_with_no_runs_is_a_no_op() { + let store = store(); + let result = reclaim_stale(&store, "quiet-thread", &RunLimits::default()) + .await + .unwrap(); + assert_eq!(result.reclaimed_count, 0); + assert!(result.details.is_empty()); +} + +#[test] +fn a_run_round_trips_through_json() { + let run = run_at("run-1", "task-1", 10, 20); + let json = serde_json::to_value(&run).unwrap(); + // camelCase on the wire, so a host's UI and RPC layer read it directly. + assert_eq!(json["runId"], "run-1"); + assert_eq!(json["cardId"], "task-1"); + assert!( + json.get("completedAt").is_none(), + "absent fields stay absent" + ); + assert_eq!(serde_json::from_value::(json).unwrap(), run); +} + +#[test] +fn default_limits_are_the_documented_policy() { + let limits = RunLimits::default(); + assert_eq!( + limits.heartbeat_stale_secs, + super::DEFAULT_HEARTBEAT_STALE_SECS + ); + assert_eq!(limits.claim_ttl_secs, super::DEFAULT_CLAIM_TTL_SECS); + assert_eq!(limits.max_reclaim_count, super::DEFAULT_MAX_RECLAIM_COUNT); + assert!(limits.claim_ttl_secs > limits.heartbeat_stale_secs); +} + +#[tokio::test] +async fn import_if_absent_never_replaces_an_existing_log() { + let store = store(); + let imported = vec![run_at("legacy-run", "task-1", 10, 20)]; + assert!( + super::store::import_if_absent(&store, "thread-1", imported.clone()) + .await + .unwrap() + ); + assert_eq!(list_runs(&store, "thread-1", None).await.unwrap(), imported); + + // A second import is refused, so a re-run of a migration cannot duplicate + // the reclaim history the sweep's budget is counted from. + assert!( + !super::store::import_if_absent(&store, "thread-1", vec![run_at("other", "task-2", 1, 2)]) + .await + .unwrap() + ); + assert_eq!(list_runs(&store, "thread-1", None).await.unwrap(), imported); +} diff --git a/src/graph/todos/runs/types.rs b/src/graph/todos/runs/types.rs new file mode 100644 index 00000000..d714a708 --- /dev/null +++ b/src/graph/todos/runs/types.rs @@ -0,0 +1,169 @@ +//! Domain types for **task runs**: the claim/heartbeat/outcome record a worker +//! writes while it executes one [`TaskBoardCard`](super::super::TaskBoardCard). +//! +//! A board card says *what* to do; a [`TaskRun`] says *who is doing it right +//! now, since when, and whether they are still alive*. Ported from OpenHuman's +//! `threads::todos::runs`, minus the app coupling (domain-event bus, +//! workspace-file layout): a run is always `(Store, thread_id)`-addressed like +//! the board it belongs to. + +use serde::{Deserialize, Serialize}; + +use crate::harness::ids::{next_seq, process_nonce}; + +/// Default staleness threshold for a run's heartbeat, in seconds. +/// +/// A healthy worker ticks [`update_heartbeat`](super::store::update_heartbeat) +/// well inside this window; one that has not is presumed wedged. +pub const DEFAULT_HEARTBEAT_STALE_SECS: u64 = 300; + +/// Default ceiling on a single claim's total age, in seconds. A run older than +/// this is reclaimed even if it is still heartbeating. +pub const DEFAULT_CLAIM_TTL_SECS: u64 = 3600; + +/// Default number of reclaims a card tolerates before it parks as +/// [`Blocked`](super::super::TaskCardStatus::Blocked) instead of returning to +/// the queue. +pub const DEFAULT_MAX_RECLAIM_COUNT: u32 = 3; + +/// How a run ended. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RunOutcome { + /// The worker finished the card's work. + Success, + /// The worker ran and failed; `error` carries the reason. + Failed, + /// The run went stale and was reclaimed by a sweep; the worker never + /// reported an outcome of its own. + Reclaimed, +} + +/// One claim on a card: who took it, when, and how it ended. +/// +/// Timestamps are unix-epoch **milliseconds** rendered as strings, matching the +/// board's [`updated_at`](super::super::TaskBoardCard::updated_at). +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TaskRun { + /// Stable run id, unique within the thread. + pub run_id: String, + /// The card this run is executing. + pub card_id: String, + /// Opaque worker identity (an agent id, a host label, …). + pub claimed_by: String, + /// Freshly minted per claim, so a reclaimed worker's late write-back can be + /// told apart from the current claim's. + pub claim_token: String, + /// When the claim was taken. + pub started_at: String, + /// Last liveness tick. + pub last_heartbeat_at: String, + /// When the run reached a terminal state; `None` while it is active. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub completed_at: Option, + /// Terminal outcome; `None` while the run is active. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub outcome: Option, + /// Failure or reclaim reason. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error: Option, + /// Evidence the run gathered toward the card's acceptance criteria. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub evidence: Vec, +} + +impl TaskRun { + /// Whether the run has not yet reached a terminal state. + pub fn is_active(&self) -> bool { + self.completed_at.is_none() + } +} + +/// Staleness and reclaim policy applied by a sweep. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RunLimits { + /// A run whose last heartbeat is older than this is stale. + pub heartbeat_stale_secs: u64, + /// A run older than this is stale regardless of its heartbeat. + pub claim_ttl_secs: u64, + /// Reclaims a card tolerates before it parks as `Blocked`. + pub max_reclaim_count: u32, +} + +impl Default for RunLimits { + fn default() -> Self { + Self { + heartbeat_stale_secs: DEFAULT_HEARTBEAT_STALE_SECS, + claim_ttl_secs: DEFAULT_CLAIM_TTL_SECS, + max_reclaim_count: DEFAULT_MAX_RECLAIM_COUNT, + } + } +} + +/// What one [`reclaim_stale`](super::store::reclaim_stale) sweep did. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ReclaimResult { + /// Cards returned to `Todo` and re-dispatchable. + pub reclaimed_count: usize, + /// Cards parked at `Blocked` for exceeding + /// [`RunLimits::max_reclaim_count`]. + pub blocked_count: usize, + /// One entry per reclaimed run, in sweep order. Hosts that publish domain + /// events read them from here — the crate emits none of its own. + pub details: Vec, +} + +/// One reclaimed run within a [`ReclaimResult`]. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ReclaimDetail { + /// The run that was reclaimed. + pub run_id: String, + /// The card it had claimed. + pub card_id: String, + /// Why it was judged stale. + pub reason: String, + /// The status the card was moved to (`todo` or `blocked`). + pub new_card_status: String, +} + +/// Why `run` is stale at `now_ms` under `limits`, or `None` when it is healthy. +/// +/// Pure and clock-injected so the policy is testable without sleeping. The TTL +/// check is evaluated first, so a run that is both too old *and* silent reports +/// the more fundamental reason. An unparsable timestamp is treated as healthy: +/// a corrupt record must not cause a live worker's card to be yanked away. +pub fn staleness_reason(run: &TaskRun, now_ms: u64, limits: &RunLimits) -> Option { + let started = run.started_at.parse::().ok()?; + let last_heartbeat = run.last_heartbeat_at.parse::().ok()?; + + let age_secs = now_ms.saturating_sub(started) / 1000; + let heartbeat_age_secs = now_ms.saturating_sub(last_heartbeat) / 1000; + + if age_secs > limits.claim_ttl_secs { + return Some(format!( + "claim TTL expired (age {age_secs}s > limit {}s)", + limits.claim_ttl_secs + )); + } + if heartbeat_age_secs > limits.heartbeat_stale_secs { + return Some(format!( + "heartbeat stale (last heartbeat {heartbeat_age_secs}s ago > limit {}s)", + limits.heartbeat_stale_secs + )); + } + None +} + +/// Mints a fresh, process-unique run id (`run--`). +pub(crate) fn new_run_id() -> String { + format!("run-{:x}-{}", process_nonce(), next_seq()) +} + +/// Mints a fresh claim token, unique across processes and restarts. +pub(crate) fn new_claim_token() -> String { + format!("claim-{:x}-{}", process_nonce(), next_seq()) +} diff --git a/src/graph/todos/types.rs b/src/graph/todos/types.rs index a02d695f..494ddae3 100644 --- a/src/graph/todos/types.rs +++ b/src/graph/todos/types.rs @@ -372,11 +372,16 @@ pub(crate) fn new_card_id() -> String { format!("task-{}", next_seq()) } -/// Current unix time in milliseconds, as a string. Dependency-free (no `chrono`). -pub(crate) fn now_stamp() -> String { +/// Current unix time in milliseconds. Dependency-free (no `chrono`). +pub(crate) fn now_millis() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_millis() as u64) .unwrap_or(0) - .to_string() +} + +/// Current unix time in milliseconds, as a string — the timestamp format every +/// board and run record uses. +pub(crate) fn now_stamp() -> String { + now_millis().to_string() } diff --git a/src/lib.rs b/src/lib.rs index 7d4dbcca..98cf7296 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -225,18 +225,25 @@ pub use graph::{ // the tools and continuation helpers are re-exported flat for discoverability. pub use graph::goals::store as goal_store; pub use graph::{ - GoalProgress, GoalTool, GoalToolKind, ThreadGoal, ThreadGoalStatus, TurnOutcome, - active_goal_context_block, goal_gate_node, goal_tools, note_user_turn, register_goal_tools, - run_continuation_tick, + BudgetVerdict, GoalBudgetGuard, GoalProgress, GoalTool, GoalToolKind, ThreadGoal, + ThreadGoalStatus, TurnOutcome, account_turn, accrues_usage, active_goal_context_block, + goal_gate_node, goal_tools, note_user_turn, register_goal_tools, run_continuation_tick, + turn_tokens, }; // --- Graph: per-thread task board (kanban todos) --- -// `todo_store` is the programmatic CRUD surface (add/edit/claim_card/...); the -// tool and data model are re-exported flat for discoverability. +// `todo_store` is the programmatic CRUD surface (add/edit/claim_card/...); +// `task_run_store` is the claim/heartbeat/reclaim log a dispatcher writes +// alongside it. The tool, data model, and dispatch policy are re-exported flat +// for discoverability. +pub use graph::todos::runs::store as task_run_store; pub use graph::todos::store as todo_store; pub use graph::{ - CardPatch, TaskApprovalMode, TaskBoard, TaskBoardCard, TaskCardStatus, TodoTool, TodosSnapshot, - normalise_board, parse_status, register_todo_tools, render_markdown, todo_tools, + ActiveRun, ActiveRunRegistry, CardPatch, PollCadence, ReclaimDetail, ReclaimResult, RunLimits, + RunOutcome, TaskApprovalMode, TaskBoard, TaskBoardCard, TaskCardStatus, TaskPromptTools, + TaskRun, TodoTool, TodosSnapshot, build_progress_instruction, build_task_prompt, card_urgency, + has_card_in_progress, normalise_board, parse_status, pick_next_card, register_todo_tools, + render_markdown, requires_plan_approval, staleness_reason, todo_tools, }; // --- Graph: parallel map/reduce helper --- diff --git a/tests/e2e_graph_task_dispatch.rs b/tests/e2e_graph_task_dispatch.rs new file mode 100644 index 00000000..78afd11a --- /dev/null +++ b/tests/e2e_graph_task_dispatch.rs @@ -0,0 +1,471 @@ +//! End-to-end coverage for autonomous task dispatch, assembled from the public +//! crate surface exactly as a host would assemble it: +//! +//! ```text +//! pick_next_card → requires_plan_approval → claim_card → create_run +//! → build_task_prompt → agent turn → complete_run → card write-back +//! ``` +//! +//! The agent side is a `MockModel` driving the real `todo` tool, so the card a +//! run reports on is the same card the dispatcher claimed. What these tests +//! pin down is the *loop*: the right card is picked, an unapproved plan never +//! runs, a claimed card is invisible to the next tick, a cancelled run does not +//! strand its card `in_progress`, and an abandoned worker's card comes back. + +use std::sync::Arc; + +use serde_json::json; + +use tinyagents::harness::context::RunConfig; +use tinyagents::harness::message::{AssistantMessage, ContentBlock, Message}; +use tinyagents::harness::model::ModelResponse; +use tinyagents::harness::providers::MockModel; +use tinyagents::harness::runtime::AgentHarness; +use tinyagents::harness::store::{InMemoryStore, Store}; +use tinyagents::harness::tool::ToolCall; +use tinyagents::harness::usage::Usage; +use tinyagents::{ + ActiveRun, ActiveRunRegistry, CardPatch, PollCadence, RunLimits, RunOutcome, TaskApprovalMode, + TaskBoardCard, TaskCardStatus, TaskPromptTools, TodoTool, build_progress_instruction, + build_task_prompt, has_card_in_progress, pick_next_card, requires_plan_approval, + task_run_store, todo_store, +}; + +const THREAD: &str = "user-tasks"; + +fn tool_call_response(id: &str, arguments: serde_json::Value) -> ModelResponse { + ModelResponse { + message: AssistantMessage { + id: Some(format!("msg-{id}")), + content: Vec::new(), + tool_calls: vec![ToolCall::new(id, "todo", arguments)], + usage: Some(Usage::new(7, 3)), + }, + usage: Some(Usage::new(7, 3)), + finish_reason: Some("tool_calls".to_string()), + raw: None, + resolved_model: None, + continue_turn: None, + served_from_cache: false, + } +} + +fn text_response(text: &str) -> ModelResponse { + ModelResponse { + message: AssistantMessage { + id: None, + content: vec![ContentBlock::Text(text.to_string())], + tool_calls: Vec::new(), + usage: Some(Usage::new(4, 2)), + }, + usage: Some(Usage::new(4, 2)), + finish_reason: Some("stop".to_string()), + raw: None, + resolved_model: None, + continue_turn: None, + served_from_cache: false, + } +} + +/// Add a card, returning its id. +async fn add_card(store: &Arc, title: &str, patch: CardPatch) -> String { + let snapshot = todo_store::add(store, THREAD, title, patch) + .await + .expect("add card"); + snapshot + .cards + .iter() + .find(|card| card.title == title) + .expect("card present") + .id + .clone() +} + +fn agent_card(agent: &str, urgency: f64) -> CardPatch { + CardPatch { + assigned_agent: Some(agent.to_string()), + source_metadata: Some(json!({ "urgency": urgency })), + ..CardPatch::default() + } +} + +/// What one dispatcher tick decided to do. +#[derive(Debug, PartialEq)] +enum Tick { + /// Nothing to claim: the board is empty, busy, or holds only work this + /// dispatcher may not run. + Idle, + /// The card was parked for a human to approve its plan. + Parked(String), + /// The card was claimed and a run opened for it. + Dispatched { card_id: String, run_id: String }, +} + +/// One sweep of the board, wired from the crate's dispatch policy: reclaim +/// what has gone stale, refuse to double-book a busy board, pick the most +/// urgent agent-assigned card, and either park it for approval or claim it. +async fn tick(store: &Arc, approval_required: bool) -> Tick { + task_run_store::reclaim_stale(store, THREAD, &RunLimits::default()) + .await + .expect("sweep"); + + let board = todo_store::list(store, THREAD).await.expect("list board"); + if has_card_in_progress(&board.cards) { + return Tick::Idle; + } + let Some(card) = pick_next_card(&board.cards, true) else { + return Tick::Idle; + }; + + if card.status == TaskCardStatus::Todo + && requires_plan_approval(approval_required, card.approval_mode.as_ref()) + { + todo_store::update_status(store, THREAD, &card.id, TaskCardStatus::AwaitingApproval) + .await + .expect("park for approval"); + return Tick::Parked(card.id); + } + + todo_store::claim_card( + store, + THREAD, + &card.id, + &[TaskCardStatus::Todo, TaskCardStatus::Ready], + TaskCardStatus::InProgress, + ) + .await + .expect("claim card"); + let run = task_run_store::create_run(store, THREAD, None, &card.id, "dispatcher") + .await + .expect("open run"); + Tick::Dispatched { + card_id: card.id, + run_id: run.run_id, + } +} + +/// Run one card through a mock agent that marks it done through the `todo` +/// tool, then close the run. Returns the prompt the agent was given. +async fn run_card(store: &Arc, card: &TaskBoardCard, run_id: &str) -> String { + let tools = TaskPromptTools::default(); + let prompt = format!( + "{}{}", + build_task_prompt(card, &tools), + build_progress_instruction(&card.id, THREAD, &tools) + ); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness + .register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response( + "call-1", + json!({ + "op": "edit", + "id": card.id, + "evidence": ["ran the migration"], + }), + ), + tool_call_response( + "call-2", + json!({ "op": "update_status", "id": card.id, "status": "done" }), + ), + text_response("migration applied"), + ])), + ) + .set_default_model("mock") + .register_tool(Arc::new(TodoTool::new(store.clone()))); + + let outcome = harness + .invoke( + &(), + (), + RunConfig::new(run_id).with_thread(THREAD), + vec![Message::user(prompt.clone())], + ) + .await + .expect("agent run succeeds"); + assert_eq!(outcome.tool_calls, 2); + + task_run_store::complete_run( + store, + THREAD, + run_id, + RunOutcome::Success, + None, + vec!["migration applied".to_string()], + ) + .await + .expect("close run"); + + prompt +} + +#[tokio::test] +async fn the_dispatcher_runs_the_most_urgent_agent_card_and_leaves_human_work_alone() { + let store: Arc = Arc::new(InMemoryStore::default()); + + // A human's own todo (unassigned, most urgent), and two agent-assigned + // cards. Only the agent cards are the dispatcher's to run. + let mine = add_card( + &store, + "call the dentist", + CardPatch { + source_metadata: Some(json!({ "urgency": 0.99 })), + ..CardPatch::default() + }, + ) + .await; + let low = add_card(&store, "tidy the changelog", agent_card("scribe", 0.1)).await; + let high = add_card(&store, "apply the migration", agent_card("dba", 0.8)).await; + + // Tick 1: the urgent agent card is claimed; the human's card is not touched. + let Tick::Dispatched { card_id, run_id } = tick(&store, false).await else { + panic!("expected a dispatch"); + }; + assert_eq!(card_id, high); + + // While it runs the board is busy, so the next tick claims nothing — + // the single-`in_progress` rule holds across ticks, not just within one. + assert_eq!(tick(&store, false).await, Tick::Idle); + + let board = todo_store::list(&store, THREAD).await.expect("board"); + let card = board + .cards + .iter() + .find(|card| card.id == card_id) + .expect("claimed card") + .clone(); + let prompt = run_card(&store, &card, &run_id).await; + assert!(prompt.contains("apply the migration"), "{prompt}"); + assert!( + prompt.contains(&card_id), + "the run is told which card it owns" + ); + + // Tick 2: with the first card done, the remaining agent card goes next. + let Tick::Dispatched { card_id, .. } = tick(&store, false).await else { + panic!("expected the second agent card"); + }; + assert_eq!(card_id, low); + + let board = todo_store::list(&store, THREAD).await.expect("board"); + let statuses: Vec<_> = board + .cards + .iter() + .map(|card| (card.id.clone(), card.status)) + .collect(); + assert!(statuses.contains(&(mine.clone(), TaskCardStatus::Todo))); + assert!(statuses.contains(&(high, TaskCardStatus::Done))); + assert!(statuses.contains(&(low, TaskCardStatus::InProgress))); + + // The evidence the run reported is on the card it was working. + let done = board + .cards + .iter() + .find(|c| c.status == TaskCardStatus::Done) + .unwrap(); + assert_eq!(done.evidence, vec!["ran the migration".to_string()]); +} + +#[tokio::test] +async fn a_plan_awaiting_approval_never_runs_until_it_is_approved() { + let store: Arc = Arc::new(InMemoryStore::default()); + let card_id = add_card(&store, "delete the old bucket", agent_card("ops", 0.5)).await; + + // With approval on, the tick parks the card instead of claiming it, and + // keeps parking nothing afterwards: an awaiting card is not dispatchable. + assert_eq!(tick(&store, true).await, Tick::Parked(card_id.clone())); + assert_eq!(tick(&store, true).await, Tick::Idle); + assert!( + task_run_store::list_runs(&store, THREAD, None) + .await + .expect("runs") + .is_empty(), + "no run is opened for an unapproved plan" + ); + + // A human approves it; now the same tick claims it. + todo_store::decide_plan(&store, THREAD, &card_id, true) + .await + .expect("approve plan"); + let Tick::Dispatched { + card_id: claimed, .. + } = tick(&store, true).await + else { + panic!("an approved plan runs"); + }; + assert_eq!(claimed, card_id); +} + +#[tokio::test] +async fn a_card_stamped_required_is_parked_even_with_the_global_gate_off() { + let store: Arc = Arc::new(InMemoryStore::default()); + let card_id = add_card( + &store, + "email the customer", + CardPatch { + approval_mode: Some(Some(TaskApprovalMode::Required)), + ..agent_card("support", 0.5) + }, + ) + .await; + + // The card's own stamp outranks the global default — an interactive plan + // review must hold regardless of how the host is configured. + assert_eq!(tick(&store, false).await, Tick::Parked(card_id)); +} + +#[tokio::test] +async fn a_cancelled_run_leaves_its_card_blocked_rather_than_stranded() { + let store: Arc = Arc::new(InMemoryStore::default()); + let card_id = add_card(&store, "long crawl", agent_card("crawler", 0.5)).await; + let registry: ActiveRunRegistry = ActiveRunRegistry::new(); + + let Tick::Dispatched { run_id, .. } = tick(&store, false).await else { + panic!("expected a dispatch"); + }; + + // The run is a detached task; the registry is how a cancel reaches it. + let work = tokio::spawn(async { std::future::pending::<()>().await }); + let (heartbeat_cancel, mut heartbeat_rx) = tokio::sync::watch::channel(false); + registry.register( + THREAD, + ActiveRun { + run_id: run_id.clone(), + card_id: card_id.clone(), + abort: work.abort_handle(), + heartbeat_cancel, + context: THREAD.to_string(), + }, + ); + + // A cancel for some *other* run must not tear this one down. + assert!( + registry + .take_if(THREAD, Some("run-from-a-previous-request")) + .is_none() + ); + + let active = registry + .take_if(THREAD, Some(&run_id)) + .expect("the live run"); + active.cancel(); + assert!(work.await.unwrap_err().is_cancelled()); + assert!(heartbeat_rx.changed().await.is_ok()); + + // The aborted task never reaches its own write-back, so the canceller owns + // it: close the run and park the card, rather than leaving it in progress. + task_run_store::complete_run( + &store, + THREAD, + &run_id, + RunOutcome::Failed, + Some("cancelled by user".to_string()), + vec![], + ) + .await + .expect("close run"); + todo_store::edit( + &store, + THREAD, + &card_id, + CardPatch { + status: Some(TaskCardStatus::Blocked), + blocker: Some("cancelled by user".to_string()), + ..CardPatch::default() + }, + ) + .await + .expect("park card"); + + let board = todo_store::list(&store, THREAD).await.expect("board"); + assert_eq!(board.cards[0].status, TaskCardStatus::Blocked); + assert_eq!(board.cards[0].blocker.as_deref(), Some("cancelled by user")); + assert!(registry.is_empty()); + + // And the board is free again: a later tick is not blocked by a ghost. + assert_eq!( + tick(&store, false).await, + Tick::Idle, + "a blocked card is not re-run" + ); +} + +#[tokio::test] +async fn an_abandoned_run_is_reclaimed_by_the_next_tick() { + let store: Arc = Arc::new(InMemoryStore::default()); + let card_id = add_card(&store, "flaky job", agent_card("runner", 0.5)).await; + + let Tick::Dispatched { run_id, .. } = tick(&store, false).await else { + panic!("expected a dispatch"); + }; + + // The worker vanishes: age its stamps so the next sweep judges it dead. + let mut runs = task_run_store::list_runs(&store, THREAD, None) + .await + .expect("runs"); + for run in runs.iter_mut() { + run.started_at = "0".to_string(); + run.last_heartbeat_at = "0".to_string(); + } + let key: String = THREAD + .as_bytes() + .iter() + .map(|b| format!("{b:02x}")) + .collect(); + store + .put( + task_run_store::RUNS_NAMESPACE, + &key, + serde_json::to_value(&runs).expect("serialize"), + ) + .await + .expect("write runs"); + + // The next tick reclaims the dead run and re-dispatches the card in the + // same sweep — recovery needs no operator intervention. + let Tick::Dispatched { + card_id: reclaimed, + run_id: fresh, + } = tick(&store, false).await + else { + panic!("expected a re-dispatch"); + }; + assert_eq!(reclaimed, card_id); + assert_ne!(fresh, run_id, "a new claim, not the dead one"); + + let history = task_run_store::list_runs(&store, THREAD, Some(&card_id)) + .await + .expect("history"); + assert_eq!(history.len(), 2); + assert_eq!(history[0].outcome, Some(RunOutcome::Reclaimed)); + assert!(history[1].is_active()); +} + +#[tokio::test] +async fn an_idle_board_backs_the_sweep_off_and_fresh_work_resets_it() { + let store: Arc = Arc::new(InMemoryStore::default()); + let cadence = PollCadence::default(); + + // Nothing to do: the tick is idle and the interval starts stretching once + // the grace window is used up. + let mut idle_ticks = 0u32; + for _ in 0..6 { + if tick(&store, false).await == Tick::Idle { + idle_ticks += 1; + } + } + assert_eq!(idle_ticks, 6); + assert!( + cadence.next_delay(idle_ticks) > cadence.base, + "a persistently empty board is not swept at full rate forever" + ); + + // Work arrives, the tick dispatches, and the cadence snaps back. + add_card(&store, "new work", agent_card("worker", 0.5)).await; + assert!(matches!(tick(&store, false).await, Tick::Dispatched { .. })); + idle_ticks = 0; + assert_eq!(cadence.next_delay(idle_ticks), cadence.base); +} diff --git a/tests/feature_graph_goal_budget.rs b/tests/feature_graph_goal_budget.rs new file mode 100644 index 00000000..082f0ce5 --- /dev/null +++ b/tests/feature_graph_goal_budget.rs @@ -0,0 +1,203 @@ +//! Feature coverage for goal budget enforcement (`graph::goals::budget`) through +//! the public crate surface. +//! +//! Two halves of the same policy: [`account_turn`] charges a finished turn +//! against the thread's goal, and a [`GoalBudgetGuard`] decides mid-turn +//! whether the turn should stop before it overruns. Together they are what +//! keeps an autonomous run from spending without a ceiling — and what keeps a +//! *user-present* conversation from being hard-stopped once there is no live +//! budget left to protect. + +use std::sync::Arc; + +use tinyagents::harness::store::{InMemoryStore, Store}; +use tinyagents::{ + BudgetVerdict, GoalBudgetGuard, ThreadGoalStatus, account_turn, goal_store, turn_tokens, +}; + +fn store() -> Arc { + Arc::new(InMemoryStore::default()) +} + +#[tokio::test] +async fn an_autonomous_run_spends_down_its_budget_and_stops() { + let store = store(); + let goal = goal_store::set(&store, "thread-1", "reindex everything", Some(300)) + .await + .expect("set goal"); + let guard = GoalBudgetGuard::for_goal(&goal).expect("budgeted goal is enforceable"); + + // Two turns inside the ceiling: charged, still active, still runnable. + for _ in 0..2 { + let updated = account_turn(&store, "thread-1", 60, 40, 5, false) + .await + .expect("account") + .expect("goal charged"); + assert_eq!(updated.status, ThreadGoalStatus::Active); + assert_eq!( + guard.check(&store, 0).await.unwrap(), + BudgetVerdict::Continue + ); + } + let goal = goal_store::get(&store, "thread-1").await.unwrap().unwrap(); + assert_eq!(goal.tokens_used, 200); + assert_eq!(goal.budget_remaining(), Some(100)); + + // Mid-turn, the guard sees the projected spend cross the ceiling and calls + // a stop before the turn burns through it. + assert_eq!( + guard.check(&store, 50).await.unwrap(), + BudgetVerdict::Continue + ); + let verdict = guard.check(&store, 100).await.unwrap(); + assert!(verdict.is_stop(), "{verdict:?}"); + + // The turn wraps up and is accounted; the goal becomes budget-limited. + let final_goal = account_turn(&store, "thread-1", 60, 60, 3, false) + .await + .expect("account") + .expect("goal charged"); + assert_eq!(final_goal.status, ThreadGoalStatus::BudgetLimited); + assert!(final_goal.over_budget()); + assert_eq!(final_goal.budget_remaining(), Some(0)); +} + +#[tokio::test] +async fn a_limited_goal_stops_charging_and_stops_hard_stopping() { + let store = store(); + let goal = goal_store::set(&store, "thread-1", "ship it", Some(100)) + .await + .expect("set goal"); + let guard = GoalBudgetGuard::for_goal(&goal).expect("guard"); + account_turn(&store, "thread-1", 80, 40, 1, false) + .await + .expect("account"); + + let limited = goal_store::get(&store, "thread-1").await.unwrap().unwrap(); + assert_eq!(limited.status, ThreadGoalStatus::BudgetLimited); + let used = limited.tokens_used; + + // The user keeps talking. Nothing further is charged against the exhausted + // goal, and the guard stands down — the injected goal context is what + // steers the model now, not a hard stop on a user-present turn. + assert!( + account_turn(&store, "thread-1", 500, 500, 60, true) + .await + .expect("account") + .is_none() + ); + assert_eq!( + goal_store::get(&store, "thread-1") + .await + .unwrap() + .unwrap() + .tokens_used, + used + ); + assert_eq!( + guard.check(&store, 10_000).await.unwrap(), + BudgetVerdict::Continue + ); +} + +#[tokio::test] +async fn a_goal_with_no_budget_is_charged_but_never_stopped() { + let store = store(); + let goal = goal_store::set(&store, "thread-1", "keep an eye on things", None) + .await + .expect("set goal"); + + // Nothing to enforce, so there is no guard to arm. + assert!(GoalBudgetGuard::for_goal(&goal).is_none()); + + // Usage is still tracked — a host may want the number even with no cap. + let updated = account_turn(&store, "thread-1", 5_000, 5_000, 600, true) + .await + .expect("account") + .expect("goal charged"); + assert_eq!(updated.tokens_used, turn_tokens(5_000, 5_000)); + assert_eq!(updated.time_used_seconds, 600); + assert_eq!(updated.status, ThreadGoalStatus::Active); + assert_eq!(updated.budget_remaining(), None); +} + +#[tokio::test] +async fn a_replaced_objective_starts_a_fresh_budget_and_disarms_the_old_guard() { + let store = store(); + let first = goal_store::set(&store, "thread-1", "objective one", Some(100)) + .await + .expect("set goal"); + let stale_guard = GoalBudgetGuard::for_goal(&first).expect("guard"); + account_turn(&store, "thread-1", 50, 40, 1, true) + .await + .expect("account"); + + // A new objective mints a new goal id and resets the counters. + let second = goal_store::set(&store, "thread-1", "objective two", Some(100)) + .await + .expect("replace goal"); + assert_ne!(second.goal_id, first.goal_id); + assert_eq!(second.tokens_used, 0); + + // The guard armed for the previous objective quietly stands down instead of + // enforcing a ceiling that no longer describes the work. + assert_eq!( + stale_guard.check(&store, 10_000).await.unwrap(), + BudgetVerdict::Continue + ); + + // A guard for the current goal enforces the fresh budget. + let guard = GoalBudgetGuard::for_goal(&second).expect("guard"); + assert!(guard.check(&store, 100).await.unwrap().is_stop()); +} + +#[tokio::test] +async fn a_user_turn_re_arms_continuation_but_the_continuation_itself_does_not() { + let store = store(); + let goal = goal_store::set(&store, "thread-1", "watch the queue", None) + .await + .expect("set goal"); + + // An idle-period continuation fired and suppressed itself one-shot. + goal_store::set_continuation_suppressed_if(&store, "thread-1", &goal.goal_id, true) + .await + .expect("suppress"); + + // Accounting for the continuation's own turn must not clear that flag, or + // the loop would drive itself forever. + account_turn(&store, "thread-1", 100, 50, 10, false) + .await + .expect("account"); + assert!( + goal_store::get(&store, "thread-1") + .await + .unwrap() + .unwrap() + .continuation_suppressed + ); + + // A person replying re-arms it: the next idle period may continue again. + account_turn(&store, "thread-1", 100, 50, 10, true) + .await + .expect("account"); + let goal = goal_store::get(&store, "thread-1").await.unwrap().unwrap(); + assert!(!goal.continuation_suppressed); + assert_eq!(goal.tokens_used, 300, "both turns were charged"); +} + +#[tokio::test] +async fn a_thread_with_no_goal_is_untouched_and_unguarded() { + let store = store(); + assert!( + account_turn(&store, "no-goal-here", 100, 100, 10, true) + .await + .expect("account") + .is_none() + ); + assert!( + goal_store::get(&store, "no-goal-here") + .await + .unwrap() + .is_none() + ); +} diff --git a/tests/feature_graph_task_runs.rs b/tests/feature_graph_task_runs.rs new file mode 100644 index 00000000..aca52a42 --- /dev/null +++ b/tests/feature_graph_task_runs.rs @@ -0,0 +1,361 @@ +//! Feature coverage for the task-run layer (`graph::todos::runs`) through the +//! public crate surface. +//! +//! The unit tests in `src/graph/todos/runs/test.rs` cover each operation in +//! isolation. These exercise the *feature* a host actually depends on: a card +//! is claimed, worked, and either finished or — when its worker dies without +//! saying so — handed back to the queue by a sweep, with a poisonous card +//! eventually parked instead of cycling forever. + +use std::sync::Arc; +use std::time::Duration; + +use tinyagents::harness::store::{InMemoryStore, Store}; +use tinyagents::{ + CardPatch, RunLimits, RunOutcome, TaskCardStatus, TaskRun, staleness_reason, task_run_store, + todo_store, +}; + +fn store() -> Arc { + Arc::new(InMemoryStore::new()) +} + +/// Add one card and return its id. +async fn seed_card(store: &Arc, thread_id: &str, title: &str) -> String { + todo_store::add(store, thread_id, title, CardPatch::default()) + .await + .expect("add card") + .cards + .last() + .expect("card present") + .id + .clone() +} + +/// Rewrite a run's timestamps so it looks abandoned, without waiting for real +/// time to pass. +async fn abandon(store: &Arc, thread_id: &str, run_id: &str) { + let mut runs = task_run_store::list_runs(store, thread_id, None) + .await + .expect("list runs"); + for run in runs.iter_mut().filter(|run| run.run_id == run_id) { + run.started_at = "0".to_string(); + run.last_heartbeat_at = "0".to_string(); + } + let key: String = thread_id + .as_bytes() + .iter() + .map(|b| format!("{b:02x}")) + .collect(); + store + .put( + task_run_store::RUNS_NAMESPACE, + &key, + serde_json::to_value(&runs).expect("serialize runs"), + ) + .await + .expect("write runs"); +} + +#[tokio::test] +async fn a_worker_claims_works_and_finishes_a_card() { + let store = store(); + let thread = "feature-happy-path"; + let card_id = seed_card(&store, thread, "write the release notes").await; + + // Claim: the card moves to InProgress and a run opens alongside it. + let card = todo_store::claim_card( + &store, + thread, + &card_id, + &[TaskCardStatus::Todo], + TaskCardStatus::InProgress, + ) + .await + .expect("claim card"); + assert_eq!(card.status, TaskCardStatus::InProgress); + + let run = task_run_store::create_run(&store, thread, None, &card_id, "writer-agent") + .await + .expect("open run"); + assert!(run.is_active()); + + // Work: heartbeats keep the claim alive, so no sweep touches it. + for _ in 0..3 { + task_run_store::update_heartbeat(&store, thread, &run.run_id) + .await + .expect("heartbeat"); + } + let swept = task_run_store::reclaim_stale(&store, thread, &RunLimits::default()) + .await + .expect("sweep"); + assert_eq!(swept.reclaimed_count, 0); + assert_eq!(swept.blocked_count, 0); + + // Finish: the run closes with its evidence and the card is marked done. + let closed = task_run_store::complete_run( + &store, + thread, + &run.run_id, + RunOutcome::Success, + None, + vec!["notes.md".to_string()], + ) + .await + .expect("close run"); + assert_eq!(closed.outcome, Some(RunOutcome::Success)); + assert_eq!(closed.evidence, vec!["notes.md".to_string()]); + + todo_store::update_status(&store, thread, &card_id, TaskCardStatus::Done) + .await + .expect("mark done"); + let board = todo_store::list(&store, thread).await.expect("list board"); + assert_eq!(board.cards[0].status, TaskCardStatus::Done); + + // The run log is the audit trail: one attempt, and it succeeded. + let history = task_run_store::list_runs(&store, thread, Some(&card_id)) + .await + .expect("history"); + assert_eq!(history.len(), 1); + assert!(!history[0].is_active()); +} + +#[tokio::test] +async fn a_dead_worker_hands_its_card_back_to_the_queue() { + let store = store(); + let thread = "feature-dead-worker"; + let card_id = seed_card(&store, thread, "reindex the archive").await; + todo_store::claim_card( + &store, + thread, + &card_id, + &[TaskCardStatus::Todo], + TaskCardStatus::InProgress, + ) + .await + .expect("claim card"); + let run = task_run_store::create_run(&store, thread, None, &card_id, "worker-a") + .await + .expect("open run"); + + // The worker dies mid-card: no completion, no more heartbeats. + abandon(&store, thread, &run.run_id).await; + + let swept = task_run_store::reclaim_stale(&store, thread, &RunLimits::default()) + .await + .expect("sweep"); + assert_eq!(swept.reclaimed_count, 1); + assert_eq!(swept.details[0].new_card_status, "todo"); + + // The card is dispatchable again, and a second worker can take it cleanly. + let board = todo_store::list(&store, thread).await.expect("list board"); + assert_eq!(board.cards[0].status, TaskCardStatus::Todo); + + todo_store::claim_card( + &store, + thread, + &card_id, + &[TaskCardStatus::Todo], + TaskCardStatus::InProgress, + ) + .await + .expect("second claim"); + let second = task_run_store::create_run(&store, thread, None, &card_id, "worker-b") + .await + .expect("second run"); + task_run_store::complete_run( + &store, + thread, + &second.run_id, + RunOutcome::Success, + None, + vec![], + ) + .await + .expect("second run finishes"); + + let history = task_run_store::list_runs(&store, thread, Some(&card_id)) + .await + .expect("history"); + assert_eq!(history.len(), 2, "both attempts are recorded"); + assert_eq!(history[0].outcome, Some(RunOutcome::Reclaimed)); + assert_eq!(history[1].outcome, Some(RunOutcome::Success)); +} + +#[tokio::test] +async fn a_card_that_keeps_killing_workers_stops_cycling() { + let store = store(); + let thread = "feature-poison"; + let card_id = seed_card(&store, thread, "run the cursed migration").await; + let limits = RunLimits { + max_reclaim_count: 2, + ..RunLimits::default() + }; + + let mut statuses = Vec::new(); + for _ in 0..3 { + todo_store::claim_card( + &store, + thread, + &card_id, + &[TaskCardStatus::Todo, TaskCardStatus::Blocked], + TaskCardStatus::InProgress, + ) + .await + .expect("claim card"); + let run = task_run_store::create_run(&store, thread, None, &card_id, "worker") + .await + .expect("open run"); + abandon(&store, thread, &run.run_id).await; + let swept = task_run_store::reclaim_stale(&store, thread, &limits) + .await + .expect("sweep"); + statuses.push(swept.details[0].new_card_status.clone()); + } + + // Two attempts get the card back; the one that reaches the limit parks it, + // and it stays parked from then on. + assert_eq!(statuses, vec!["todo", "blocked", "blocked"]); + let board = todo_store::list(&store, thread).await.expect("list board"); + assert_eq!(board.cards[0].status, TaskCardStatus::Blocked); + assert!( + board.cards[0] + .blocker + .as_deref() + .unwrap_or_default() + .contains("exceeding limit of 2") + ); +} + +#[tokio::test] +async fn one_thread_sweep_does_not_disturb_another() { + let store = store(); + let quiet_card = seed_card(&store, "thread-quiet", "leave me alone").await; + todo_store::claim_card( + &store, + "thread-quiet", + &quiet_card, + &[TaskCardStatus::Todo], + TaskCardStatus::InProgress, + ) + .await + .expect("claim quiet card"); + task_run_store::create_run(&store, "thread-quiet", None, &quiet_card, "worker") + .await + .expect("quiet run"); + + let noisy_card = seed_card(&store, "thread-noisy", "abandoned work").await; + todo_store::claim_card( + &store, + "thread-noisy", + &noisy_card, + &[TaskCardStatus::Todo], + TaskCardStatus::InProgress, + ) + .await + .expect("claim noisy card"); + let noisy_run = task_run_store::create_run(&store, "thread-noisy", None, &noisy_card, "worker") + .await + .expect("noisy run"); + abandon(&store, "thread-noisy", &noisy_run.run_id).await; + + let swept = task_run_store::reclaim_stale(&store, "thread-noisy", &RunLimits::default()) + .await + .expect("sweep"); + assert_eq!(swept.reclaimed_count, 1); + + let quiet = todo_store::list(&store, "thread-quiet") + .await + .expect("quiet board"); + assert_eq!( + quiet.cards[0].status, + TaskCardStatus::InProgress, + "a healthy run on another thread is untouched" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn a_heartbeat_task_keeps_a_long_run_alive_and_stops_on_cancel() { + let store = store(); + let thread = "feature-heartbeat"; + let run = task_run_store::create_run(&store, thread, None, "task-1", "worker") + .await + .expect("open run"); + let (cancel_tx, cancel_rx) = tokio::sync::watch::channel(false); + + // A fast tick keeps the test quick; the cadence is a parameter precisely so + // a caller (or a test) is not stuck with the 30s production default. + let tick = Duration::from_millis(20); + task_run_store::spawn_heartbeat_task( + store.clone(), + thread.to_string(), + run.run_id.clone(), + cancel_rx, + tick, + ); + + // Staleness is judged against the wall clock, so the run is aged by + // rewriting its stamps; a tick proves itself by writing a fresh one back. + let limits = RunLimits { + heartbeat_stale_secs: 60, + claim_ttl_secs: u64::MAX, + max_reclaim_count: 3, + }; + abandon(&store, thread, &run.run_id).await; + tokio::time::sleep(tick * 10).await; + assert!( + task_run_store::find_stale_runs(&store, thread, &limits) + .await + .expect("stale check") + .is_empty(), + "a heartbeating run is never stale" + ); + + // Once cancelled the ticks stop, so the next silence is not papered over. + cancel_tx.send(true).expect("cancel heartbeat"); + tokio::time::sleep(tick * 2).await; + abandon(&store, thread, &run.run_id).await; + tokio::time::sleep(tick * 10).await; + assert_eq!( + task_run_store::find_stale_runs(&store, thread, &limits) + .await + .expect("stale check") + .len(), + 1, + "a cancelled heartbeat stops keeping the claim alive" + ); +} + +#[test] +fn the_staleness_policy_is_pure_and_clock_injected() { + // A host can evaluate the policy against its own clock — no store, no wait. + let run = TaskRun { + run_id: "run-1".to_string(), + card_id: "task-1".to_string(), + claimed_by: "worker".to_string(), + claim_token: "token".to_string(), + started_at: "0".to_string(), + last_heartbeat_at: "0".to_string(), + completed_at: None, + outcome: None, + error: None, + evidence: Vec::new(), + }; + let limits = RunLimits { + heartbeat_stale_secs: 10, + claim_ttl_secs: 100, + max_reclaim_count: 3, + }; + + assert!(staleness_reason(&run, 5_000, &limits).is_none()); + assert!( + staleness_reason(&run, 20_000, &limits) + .expect("silent worker") + .contains("heartbeat stale") + ); + assert!( + staleness_reason(&run, 200_000, &limits) + .expect("expired claim") + .contains("claim TTL expired") + ); +}