use super::display::{print_tool_call, print_tool_output};
use serde::{Deserialize, Serialize};
use std::sync::Mutex;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum StepStatus {
Todo,
Active,
Done,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Step {
pub description: String,
pub status: StepStatus,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct PlanSnapshot {
#[serde(default)]
pub steps: Vec<Step>,
}
impl PlanSnapshot {
#[must_use]
pub fn is_empty(&self) -> bool {
self.steps.is_empty()
}
#[must_use]
pub fn len(&self) -> usize {
self.steps.len()
}
}
pub(crate) const MAX_STEPS: usize = 30;
pub(crate) const STEP_DESC_CAP: usize = 200;
pub(crate) const PLAN_TOTAL_CAP: usize = 3_000;
pub trait StepLedger: Send + Sync {
fn set_plan(&self, steps: &[String]) -> usize;
fn advance(&self) -> Option<String>;
fn steps(&self) -> Vec<Step>;
fn count(&self) -> u64;
fn done_count(&self) -> u64;
fn clear(&self);
fn snapshot(&self) -> PlanSnapshot;
fn restore(&self, snap: &PlanSnapshot);
}
#[derive(Default)]
pub struct SessionStepLedger {
steps: Mutex<Vec<Step>>,
}
impl StepLedger for SessionStepLedger {
fn set_plan(&self, steps: &[String]) -> usize {
let mut built: Vec<Step> = steps
.iter()
.map(|s| s.trim())
.filter(|s| !s.is_empty())
.take(MAX_STEPS)
.map(|s| Step {
description: s.to_string(),
status: StepStatus::Todo,
})
.collect();
if let Some(first) = built.first_mut() {
first.status = StepStatus::Active;
}
let n = built.len();
*self.steps.lock().unwrap() = built;
n
}
fn advance(&self) -> Option<String> {
let mut steps = self.steps.lock().unwrap();
if let Some(active) = steps.iter_mut().find(|s| s.status == StepStatus::Active) {
active.status = StepStatus::Done;
}
if let Some(next) = steps.iter_mut().find(|s| s.status == StepStatus::Todo) {
next.status = StepStatus::Active;
Some(next.description.clone())
} else {
None
}
}
fn steps(&self) -> Vec<Step> {
self.steps.lock().unwrap().clone()
}
fn count(&self) -> u64 {
self.steps.lock().unwrap().len() as u64
}
fn done_count(&self) -> u64 {
self.steps
.lock()
.unwrap()
.iter()
.filter(|s| s.status == StepStatus::Done)
.count() as u64
}
fn clear(&self) {
self.steps.lock().unwrap().clear();
}
fn snapshot(&self) -> PlanSnapshot {
PlanSnapshot {
steps: self.steps.lock().unwrap().clone(),
}
}
fn restore(&self, snap: &PlanSnapshot) {
*self.steps.lock().unwrap() = snap.steps.clone();
}
}
fn truncate(s: &str, cap: usize) -> String {
if s.chars().count() <= cap {
s.to_string()
} else {
let head: String = s.chars().take(cap).collect();
format!("{head}[…]")
}
}
pub(crate) fn build_plan_block(ledger: &dyn StepLedger, total_cap: usize) -> Option<String> {
let steps = ledger.steps();
if steps.is_empty() {
return None;
}
let mut body = String::from("<plan>\n");
for (i, step) in steps.iter().enumerate() {
let mark = match step.status {
StepStatus::Done => "✓",
StepStatus::Active => "→",
StepStatus::Todo => "☐",
};
let piece = format!(
"{mark} {}. {}\n",
i + 1,
truncate(&step.description, STEP_DESC_CAP)
);
if body.chars().count() + piece.chars().count() + "</plan>".len() > total_cap {
body.push_str("[… plan truncated to fit the budget …]\n");
break;
}
body.push_str(&piece);
}
body.push_str("</plan>");
Some(body)
}
pub fn plan_block(ledger: &dyn StepLedger) -> Option<String> {
build_plan_block(ledger, PLAN_TOTAL_CAP)
}
#[must_use]
pub fn plan_reseat_pointer(ledger: &dyn StepLedger) -> Option<String> {
let steps = ledger.steps();
if steps.len() < 2 {
return None; }
let active = steps.iter().position(|s| s.status == StepStatus::Active)?;
let done = steps
.iter()
.filter(|s| s.status == StepStatus::Done)
.count();
Some(format!(
"[Plan progress: {done}/{} done. ACTIVE \u{2192} step {}: {}. Keep working \
THIS step; earlier steps are already done — do not restart or re-plan them.]",
steps.len(),
active + 1,
truncate(&steps[active].description, STEP_DESC_CAP)
))
}
pub fn update_plan_tool_definition() -> serde_json::Value {
serde_json::json!({
"type": "function",
"function": {
"name": "update_plan",
"description": "Create or update your plan — send the full ordered list each time \
with each step's status. A <plan> checklist is shown at the head of \
every turn; mark the step you are on \"in_progress\" and finished \
steps \"completed\". Prefer this before more investigation for \
multi-step, ambiguous, resumed, or context-compacted work. Replaces \
plan_set/plan_advance.",
"parameters": {
"type": "object",
"properties": {
"plan": {
"type": "array",
"items": {
"type": "object",
"properties": {
"step": {
"type": "string",
"description": "A short imperative phrase for the step."
},
"status": {
"type": "string",
"enum": ["pending", "in_progress", "completed"],
"description": "The step's progress; exactly one step \
should be in_progress."
}
},
"required": ["step", "status"]
},
"description": "The ordered steps, each with its status."
}
},
"required": ["plan"]
}
}
})
}
pub fn plan_get_tool_definition() -> serde_json::Value {
serde_json::json!({
"type": "function",
"function": {
"name": "plan_get",
"description": "Read the current <plan> checklist — the ordered steps and which is \
active. No args. Use it to recover what you were working on after a \
resume; if it reports no active plan, create one with update_plan \
instead of calling plan_get again.",
"parameters": { "type": "object", "properties": {}, "required": [] }
}
})
}
fn parse_step_status(raw: &str) -> StepStatus {
match raw.trim().to_ascii_lowercase().as_str() {
"completed" | "complete" | "done" => StepStatus::Done,
"in_progress" | "in-progress" | "active" | "current" => StepStatus::Active,
_ => StepStatus::Todo, }
}
fn normalize_active(steps: &mut [Step]) {
let active = steps
.iter()
.filter(|s| s.status == StepStatus::Active)
.count();
if active == 1 {
return; }
for s in steps.iter_mut() {
if s.status == StepStatus::Active {
s.status = StepStatus::Todo;
}
}
if let Some(first) = steps.iter_mut().find(|s| s.status != StepStatus::Done) {
first.status = StepStatus::Active;
}
}
pub(crate) fn execute_update_plan(
args: &serde_json::Value,
ledger: &dyn StepLedger,
color: bool,
tool_output_lines: usize,
) -> String {
print_tool_call("update_plan", "", color);
let out = match args["plan"].as_array() {
None => "error: update_plan requires a `plan` array of {\"step\",\"status\"} objects"
.to_string(),
Some(items) => {
let mut steps: Vec<Step> = items
.iter()
.filter_map(|it| {
let desc = it.get("step").and_then(|v| v.as_str())?.trim();
if desc.is_empty() {
return None;
}
let status = it
.get("status")
.and_then(|v| v.as_str())
.map_or(StepStatus::Todo, parse_step_status);
Some(Step {
description: desc.to_string(),
status,
})
})
.take(MAX_STEPS)
.collect();
if steps.is_empty() {
"error: update_plan requires a non-empty `plan` array".to_string()
} else {
normalize_active(&mut steps);
ledger.restore(&PlanSnapshot { steps });
plan_block(ledger).unwrap_or_else(|| "plan updated".to_string())
}
}
};
print_tool_output(&out, tool_output_lines, color);
out
}
pub(crate) fn execute_plan_get(
ledger: &dyn StepLedger,
color: bool,
tool_output_lines: usize,
) -> String {
print_tool_call("plan_get", "", color);
let out = plan_block(ledger).unwrap_or_else(|| {
"no active plan — if this is multi-step, ambiguous, resumed, or context-compacted work, \
call update_plan next with a short 2-6 step ordered plan using statuses \
pending/in_progress/completed; do not call plan_get again until you have created or \
updated a plan"
.to_string()
});
print_tool_output(&out, tool_output_lines, color);
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn set_plan_filters_blanks_caps_and_activates_first() {
let l = SessionStepLedger::default();
let n = l.set_plan(&[
" read the code ".to_string(),
"".to_string(),
" ".to_string(),
"write the fix".to_string(),
]);
assert_eq!(n, 2, "blank steps dropped");
let steps = l.steps();
assert_eq!(steps[0].description, "read the code", "trimmed");
assert_eq!(steps[0].status, StepStatus::Active, "first is active");
assert_eq!(steps[1].status, StepStatus::Todo);
let many: Vec<String> = (0..MAX_STEPS + 10).map(|i| format!("step {i}")).collect();
assert_eq!(l.set_plan(&many), MAX_STEPS, "capped at MAX_STEPS");
l.clear();
assert_eq!(l.count(), 0);
assert!(l.steps().is_empty());
}
#[test]
fn advance_walks_active_to_done_and_reports_complete() {
let l = SessionStepLedger::default();
l.set_plan(&["a".to_string(), "b".to_string()]);
assert_eq!(l.done_count(), 0);
assert_eq!(l.advance().as_deref(), Some("b"));
let steps = l.steps();
assert_eq!(steps[0].status, StepStatus::Done);
assert_eq!(steps[1].status, StepStatus::Active);
assert_eq!(l.done_count(), 1);
assert_eq!(l.advance(), None);
assert_eq!(l.done_count(), 2);
assert_eq!(l.advance(), None);
let empty = SessionStepLedger::default();
assert_eq!(empty.advance(), None);
}
#[test]
fn snapshot_restore_round_trips_full_state_unlike_set_plan() {
let l = SessionStepLedger::default();
l.set_plan(&["a".to_string(), "b".to_string(), "c".to_string()]);
l.advance(); let snap = l.snapshot();
assert_eq!(snap.len(), 3);
assert!(!snap.is_empty());
assert_eq!(snap.steps[0].status, StepStatus::Done);
assert_eq!(snap.steps[1].status, StepStatus::Active);
assert_eq!(snap.steps[2].status, StepStatus::Todo);
let json = serde_json::to_string(&snap).unwrap();
let decoded: PlanSnapshot = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, snap, "snapshot must survive serde verbatim");
assert_eq!(
serde_json::from_str::<PlanSnapshot>("{}").unwrap(),
PlanSnapshot::default(),
"the `{{}}` backfill parses to an empty snapshot"
);
let fresh = SessionStepLedger::default();
fresh.restore(&decoded);
assert_eq!(fresh.snapshot(), snap, "restore reinstates verbatim");
assert_eq!(fresh.done_count(), 1, "the Done step survives restore");
let block = build_plan_block(&fresh, 3000).unwrap();
assert!(block.contains("✓ 1. a"), "{block}");
assert!(block.contains("→ 2. b"), "{block}");
assert!(block.contains("☐ 3. c"), "{block}");
let reset = SessionStepLedger::default();
reset.set_plan(&["a".to_string(), "b".to_string(), "c".to_string()]);
assert_eq!(reset.done_count(), 0, "set_plan resets — proving the gap");
fresh.restore(&PlanSnapshot::default());
assert_eq!(
fresh.count(),
0,
"an empty snapshot leaves the ledger empty"
);
}
#[test]
fn reseat_pointer_compact_and_gated_to_multistep_in_progress() {
let l = SessionStepLedger::default();
assert!(plan_reseat_pointer(&l).is_none(), "empty → none");
l.set_plan(&["only".to_string()]);
assert!(plan_reseat_pointer(&l).is_none(), "single step → none");
l.set_plan(&["a".to_string(), "b".to_string(), "c".to_string()]);
let p = plan_reseat_pointer(&l).expect("multi-step → pointer");
assert!(p.contains("step 1: a"), "names the active step: {p}");
assert!(p.contains("0/3 done"), "shows progress: {p}");
assert_eq!(p.lines().count(), 1, "compact (one line): {p}");
l.advance(); let ptr = plan_reseat_pointer(&l).unwrap();
assert!(
ptr.contains("step 2: b") && ptr.contains("1/3 done"),
"{ptr}"
);
l.advance();
l.advance(); assert!(plan_reseat_pointer(&l).is_none(), "complete plan → none");
}
#[test]
fn build_block_none_when_empty_marks_status_and_caps() {
let l = SessionStepLedger::default();
assert_eq!(build_plan_block(&l, 3000), None, "empty plan → no block");
l.set_plan(&["alpha".to_string(), "beta".to_string(), "gamma".to_string()]);
l.advance(); let block = build_plan_block(&l, 3000).unwrap();
assert!(block.starts_with("<plan>\n") && block.ends_with("</plan>"));
assert!(block.contains("✓ 1. alpha"), "{block}");
assert!(block.contains("→ 2. beta"), "{block}");
assert!(block.contains("☐ 3. gamma"), "{block}");
let long = SessionStepLedger::default();
long.set_plan(&["x".repeat(STEP_DESC_CAP + 50)]);
let lb = build_plan_block(&long, 3000).unwrap();
assert!(
lb.contains("[…]"),
"per-step description truncated: {lb:.80}"
);
let big = SessionStepLedger::default();
let many: Vec<String> = (0..MAX_STEPS)
.map(|i| format!("step number {i} with padding"))
.collect();
big.set_plan(&many);
let capped = build_plan_block(&big, 200).unwrap();
assert!(
capped.chars().count() <= 200 + 60,
"total cap bounds the block"
);
assert!(capped.contains("plan truncated"), "{capped}");
}
#[test]
fn update_plan_sets_statuses_renders_block_and_is_lenient() {
let l = SessionStepLedger::default();
assert!(execute_update_plan(&serde_json::json!({}), &l, false, 20).starts_with("error:"));
assert_eq!(l.count(), 0);
assert!(
execute_update_plan(&serde_json::json!({"plan": []}), &l, false, 20)
.starts_with("error:")
);
assert_eq!(l.count(), 0);
let out = execute_update_plan(
&serde_json::json!({"plan": [
{"step": "scope it", "status": "completed"},
{"step": "build it", "status": "in_progress"},
{"step": "test it", "status": "pending"},
]}),
&l,
false,
20,
);
assert!(
out.starts_with("<plan>\n") && out.ends_with("</plan>"),
"{out}"
);
assert!(out.contains("✓ 1. scope it"), "{out}");
assert!(out.contains("→ 2. build it"), "{out}");
assert!(out.contains("☐ 3. test it"), "{out}");
let steps = l.steps();
assert_eq!(steps[0].status, StepStatus::Done);
assert_eq!(steps[1].status, StepStatus::Active);
assert_eq!(steps[2].status, StepStatus::Todo);
assert_eq!(
steps
.iter()
.filter(|s| s.status == StepStatus::Active)
.count(),
1,
"exactly one in_progress"
);
let none_active = execute_update_plan(
&serde_json::json!({"plan": [
{"step": "a", "status": "completed"},
{"step": "b", "status": "pending"},
{"step": "c", "status": "pending"},
]}),
&l,
false,
20,
);
assert!(
none_active.contains("→ 2. b"),
"first non-completed becomes active: {none_active}"
);
execute_update_plan(
&serde_json::json!({"plan": [
{"step": "x", "status": "in_progress"},
{"step": "y", "status": "in_progress"},
]}),
&l,
false,
20,
);
let s = l.steps();
assert_eq!(
s.iter().filter(|x| x.status == StepStatus::Active).count(),
1,
"multiple in_progress collapse to one"
);
assert_eq!(s[0].status, StepStatus::Active, "the first becomes active");
let dropped = execute_update_plan(
&serde_json::json!({"plan": [
{"step": " ", "status": "pending"},
{"step": "real", "status": "in_progress"},
]}),
&l,
false,
20,
);
assert!(dropped.contains("→ 1. real"), "{dropped}");
assert_eq!(l.count(), 1, "blank step dropped");
}
#[test]
fn tool_definitions_shape() {
assert_eq!(
update_plan_tool_definition()["function"]["name"],
"update_plan"
);
assert_eq!(plan_get_tool_definition()["function"]["name"], "plan_get");
assert!(
update_plan_tool_definition()["function"]["parameters"]["properties"]["plan"]
.is_object()
);
assert_eq!(
plan_get_tool_definition()["function"]["parameters"]["properties"],
serde_json::json!({})
);
}
#[test]
fn plan_get_renders_plan_or_hints_when_empty() {
let l = SessionStepLedger::default();
let empty = execute_plan_get(&l, false, 20);
assert!(empty.starts_with("no active plan"), "{empty}");
assert!(empty.contains("update_plan next"), "{empty}");
assert!(
empty.contains("do not call plan_get again"),
"empty plan_get must steer away from polling: {empty}"
);
assert_eq!(l.count(), 0, "plan_get does not mutate the ledger");
l.set_plan(&["scope it".to_string(), "build it".to_string()]);
let block = execute_plan_get(&l, false, 20);
assert!(
block.starts_with("<plan>\n") && block.ends_with("</plan>"),
"{block}"
);
assert!(block.contains("→ 1. scope it"), "{block}");
assert!(block.contains("☐ 2. build it"), "{block}");
assert_eq!(l.count(), 2, "plan_get is read-only");
}
}