Add /graph G2: resumable budget, GraphUpdated status chip, PTY + turn-level coverage

BudgetLimited is now a resumable state: a budget trip demotes in-flight
nodes to Ready (a resource stop, not a verdict — no forever-Running node is
ever persisted) and '/graph resume --budget <tokens>' re-arms the graph with
fresh headroom (new budget = spent-so-far + extra). The tripped node's
partial burn is charged into tokens_spent_nodes at BOTH cascade sites before
the demotion clears current_node, so the top-up arithmetic never runs on an
under-counted ledger. Any input starting with 'resume' resolves to a resume
(case-insensitive; malformed top-ups surface the usage hint) and setup_graph
refuses to replace any non-Complete graph — a typo can no longer silently
destroy a resumable graph. An explicit --budget on a merely-paused graph is
rejected loudly instead of silently discarded; all trip-time messages now
advertise the top-up.

The pager gains a graph status chip: a new GraphUpdated wire variant
(extensions/notification.rs, old pagers degrade via #[serde(other)]) is
emitted from the single persist_graph_state chokepoint — every transition is
both a checkpoint and a badge tick — with a 'cleared' sentinel on /graph
clear and a one-shot re-emit after session restore (the replayed updates log
otherwise shows the pre-shutdown Active state that from_snapshot just
demoted in memory). TUI side: GraphDisplayState, session-notification arm,
and a goal-idiom chip with node progress, clamped current-node title, and
budget-aware spend. Pre-session command availability now advertises /graph
from the flags (it was fail-closed to the in-session path only, so the
welcome-screen slash menu never showed it).

Coverage: GraphUpdated wire round-trip + minimal-payload + unknown-tag
tests; PTY scenarios graph_slash_presession{,_disabled}.yaml (both run
green against the real pager binary); handle_prompt-level e2e for terminal
slash outcomes (/graph status|resume|pause, /goal refusals while the graph
owns the engine); budget top-up e2e driving a BudgetLimited diamond back to
Complete. Not shimmed: pre-G2 persisted snapshots with budget-Failed nodes
(the KIGI_GRAPH flag has never shipped enabled, so none exist).

kigi-shell 4927 and kigi-tui 6610 lib tests green; workspace clippy clean.
This commit is contained in:
2026-07-20 16:35:02 -04:00
parent 4d1e4fdc52
commit 1579558b56
20 changed files with 878 additions and 57 deletions
+8
View File
@@ -110,6 +110,14 @@ edges stay deterministic Rust. The harness appends a terminal
serially on the full goal engine. Concurrency=1 is byte-identical to
the serial G0 path. Ceiling: a worker round exceeding the foreground
subagent await budget (600s) is cancelled and retried via resume.
- G2: `BudgetLimited` is resumable — a budget trip demotes in-flight
nodes to `Ready` (resource stop, not a verdict) and
`/graph resume --budget <tokens>` re-arms with fresh headroom. The
pager shows a graph status chip driven by the `GraphUpdated` wire
variant (`extensions/notification.rs`), emitted from the single
`persist_graph_state` chokepoint (checkpoint ⇔ badge tick); old
pagers degrade via `#[serde(other)] Unknown`. PTY scenarios:
`graph_slash_presession{,_disabled}.yaml`.
## Milestones (PRD §8.3)
@@ -310,8 +310,13 @@ impl MvpAgent {
pub(crate) fn command_availability(
&self,
) -> crate::session::slash_commands::CommandAvailability {
let goal = self.cfg.borrow().resolve_goal().value;
crate::session::slash_commands::CommandAvailability {
goal: self.cfg.borrow().resolve_goal().value,
goal,
// Same convention as /goal: the flag is known at initialize
// time, so advertise pre-session; the in-session path
// re-checks the live toolset.
graph: goal && self.cfg.borrow().resolve_graph().value,
..crate::session::slash_commands::CommandAvailability::default()
}
}
@@ -897,6 +897,36 @@ pub enum SessionUpdate {
#[serde(default, skip_serializing_if = "Option::is_none")]
usage: Option<PromptUsage>,
},
/// Graph mode (`/graph`) progress for the pager's status chip.
/// Wire tag `graph_updated`; `status: "cleared"` tells the pager to
/// drop its graph state (same sentinel convention as `GoalUpdated`).
/// Old pagers degrade to [`Self::Unknown`] silently.
GraphUpdated {
graph_id: String,
objective: String,
/// Goal-status vocabulary (`active`, paused family,
/// `budget_limited`, `complete`) plus `cleared`.
status: String,
/// `idle` | `planning` | `executing`.
phase: String,
plan_version: u32,
total_nodes: u32,
achieved_nodes: u32,
failed_nodes: u32,
running_nodes: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
current_node: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
current_node_title: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
token_budget: Option<i64>,
#[serde(default)]
tokens_spent: i64,
#[serde(default, skip_serializing_if = "Option::is_none")]
last_event: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pause_message: Option<String>,
},
/// Catch-all for unrecognized session update types.
/// Allows forward/backward compatibility when variants are added or removed.
/// All fields from the unrecognized variant are discarded during deserialization.
@@ -2293,3 +2323,88 @@ mod tests {
assert!(serde_json::from_str::<SessionUpdate>(missing_stop_reason).is_err());
}
}
#[cfg(test)]
mod graph_updated_wire_tests {
use super::*;
fn full_graph_updated() -> SessionUpdate {
SessionUpdate::GraphUpdated {
graph_id: "g-1".into(),
objective: "ship it".into(),
status: "active".into(),
phase: "executing".into(),
plan_version: 2,
total_nodes: 5,
achieved_nodes: 2,
failed_nodes: 1,
running_nodes: 1,
current_node: Some("gn-abc".into()),
current_node_title: Some("Node C".into()),
token_budget: Some(10_000),
tokens_spent: 4_200,
last_event: Some("node_achieved".into()),
pause_message: None,
}
}
#[test]
fn graph_updated_round_trips_with_snake_case_tag() {
let update = full_graph_updated();
let json = serde_json::to_value(&update).unwrap();
assert_eq!(json["sessionUpdate"], "graph_updated");
assert_eq!(json["achieved_nodes"], 2);
assert_eq!(json["current_node_title"], "Node C");
// Omitted optionals must not serialize at all.
assert!(json.get("pause_message").is_none());
let back: SessionUpdate = serde_json::from_value(json).unwrap();
assert_eq!(back, update);
}
#[test]
fn graph_updated_minimal_payload_fills_defaults() {
// Only the required fields on the wire: every optional absent,
// `tokens_spent` relies on #[serde(default)].
let json = serde_json::json!({
"sessionUpdate": "graph_updated",
"graph_id": "g-2",
"objective": "o",
"status": "cleared",
"phase": "idle",
"plan_version": 0,
"total_nodes": 0,
"achieved_nodes": 0,
"failed_nodes": 0,
"running_nodes": 0,
});
let update: SessionUpdate = serde_json::from_value(json).unwrap();
match update {
SessionUpdate::GraphUpdated {
status,
tokens_spent,
current_node,
token_budget,
pause_message,
..
} => {
assert_eq!(status, "cleared");
assert_eq!(tokens_spent, 0, "#[serde(default)] must backfill");
assert!(current_node.is_none());
assert!(token_budget.is_none());
assert!(pause_message.is_none());
}
other => panic!("expected GraphUpdated, got {other:?}"),
}
}
/// An OLD pager (this enum before the variant existed) must degrade
/// a graph_updated payload to `Unknown` — pinned by feeding an
/// unknown-tag payload through today's enum, which uses the same
/// #[serde(other)] mechanism.
#[test]
fn unknown_tags_still_degrade_gracefully() {
let json = serde_json::json!({ "sessionUpdate": "graph_updated_v99", "x": 1 });
let update: SessionUpdate = serde_json::from_value(json).unwrap();
assert_eq!(update, SessionUpdate::Unknown);
}
}
@@ -1908,11 +1908,21 @@ impl SessionActor {
// budget, so a node-level budget trip is the graph-level trip.
if self.graph_harness_enabled() && self.graph_tracker.lock().is_active() {
tracing::warn!("graph: node goal budget trip cascades to graph BudgetLimited");
self.graph_tracker.lock().budget_limit();
{
let mut tracker = self.graph_tracker.lock();
// Budget integrity: charge the tripped node's partial burn
// BEFORE budget_limit clears current_node — otherwise the
// top-up arithmetic runs on an under-counted ledger.
if let Some(node_id) = tracker.current_node_id().map(str::to_owned) {
tracker.charge_node_tokens(&node_id, tokens_used);
}
tracker.budget_limit();
}
self.persist_graph_state();
self.send_slash_command_output(&format!(
"Graph token budget reached ({tokens_used} tokens this node) — graph \
stopped. Use /graph clear, then /graph <objective> to start a new one."
stopped. Top up with /graph resume --budget <tokens>, or /graph clear \
to abandon."
))
.await;
return true;
@@ -55,6 +55,98 @@ pub(super) fn node_goal_objective(
)
}
/// Snake-case wire form of a graph/goal status (matches the
/// `GoalUpdated` vocabulary the pager already parses).
fn graph_status_str(status: GoalStatus) -> &'static str {
match status {
GoalStatus::Active => "active",
GoalStatus::UserPaused => "user_paused",
GoalStatus::BackOffPaused => "back_off_paused",
GoalStatus::NoProgressPaused => "no_progress_paused",
GoalStatus::InfraPaused => "infra_paused",
GoalStatus::Blocked => "blocked",
GoalStatus::BudgetLimited => "budget_limited",
GoalStatus::Complete => "complete",
}
}
fn graph_event_as_str(event: &super::super::graph_tracker::GraphEvent) -> &'static str {
use super::super::graph_tracker::GraphEvent;
match event {
GraphEvent::GraphCreated => "graph_created",
GraphEvent::PlanningStarted => "planning_started",
GraphEvent::PlanningCompleted => "planning_completed",
GraphEvent::PlanningFailed => "planning_failed",
GraphEvent::NodeStarted => "node_started",
GraphEvent::NodeAchieved => "node_achieved",
GraphEvent::NodeFailed => "node_failed",
GraphEvent::GraphPaused => "graph_paused",
GraphEvent::GraphResumed => "graph_resumed",
GraphEvent::GraphCompleted => "graph_completed",
GraphEvent::GraphCleared => "graph_cleared",
GraphEvent::BudgetExceeded => "budget_exceeded",
GraphEvent::Unknown => "unknown",
}
}
/// Build the pager-facing `GraphUpdated` badge payload from a snapshot.
pub(crate) fn build_graph_updated(
state: &super::super::graph_tracker::GraphOrchestration,
) -> crate::extensions::notification::SessionUpdate {
use super::super::goal_tracker::GoalPhase;
let count = |s: NodeStatus| state.nodes.iter().filter(|n| n.status == s).count() as u32;
let current = state
.current_node
.as_deref()
.and_then(|id| state.nodes.iter().find(|n| n.id == id));
crate::extensions::notification::SessionUpdate::GraphUpdated {
graph_id: state.graph_id.clone(),
objective: state.objective.clone(),
status: graph_status_str(state.status).to_owned(),
phase: match state.phase {
GoalPhase::Idle => "idle",
GoalPhase::Planning => "planning",
GoalPhase::Executing => "executing",
}
.to_owned(),
plan_version: state.plan_version,
total_nodes: state.nodes.len() as u32,
achieved_nodes: count(NodeStatus::Achieved),
failed_nodes: count(NodeStatus::Failed) + count(NodeStatus::Blocked),
running_nodes: count(NodeStatus::Running) + count(NodeStatus::Verifying),
current_node: current.map(|n| n.id.clone()),
current_node_title: current.map(|n| n.title.clone()),
token_budget: state.token_budget,
tokens_spent: state.tokens_spent_nodes,
last_event: state
.history
.last()
.map(|e| graph_event_as_str(&e.event).to_owned()),
pause_message: state.pause_message.clone(),
}
}
/// `status: "cleared"` sentinel — the pager drops its graph state.
pub(crate) fn build_graph_cleared() -> crate::extensions::notification::SessionUpdate {
crate::extensions::notification::SessionUpdate::GraphUpdated {
graph_id: String::new(),
objective: String::new(),
status: "cleared".to_owned(),
phase: "idle".to_owned(),
plan_version: 0,
total_nodes: 0,
achieved_nodes: 0,
failed_nodes: 0,
running_nodes: 0,
current_node: None,
current_node_title: None,
token_budget: None,
tokens_spent: 0,
last_event: None,
pause_message: None,
}
}
impl SessionActor {
/// Graph feature flag AND the goal harness (nodes execute as goals).
pub(super) fn graph_harness_enabled(&self) -> bool {
@@ -62,13 +154,21 @@ impl SessionActor {
}
/// Send the current graph snapshot (or a tombstone after `clear`) to
/// the persistence actor. Every graph transition is a checkpoint.
/// the persistence actor, AND notify the pager: every graph
/// transition is both a checkpoint and a `GraphUpdated` badge tick
/// (single chokepoint — no transition can persist without also
/// updating the UI, and vice versa).
pub(crate) fn persist_graph_state(&self) {
let snapshot = self.graph_tracker.lock().snapshot().cloned();
let update = match &snapshot {
Some(state) => build_graph_updated(state),
None => build_graph_cleared(),
};
let _ = self
.notifications
.persistence_tx
.send(PersistenceMsg::GraphModeState(snapshot));
self.goal_notify_sender().send_update(update);
}
/// True when the graph occupies the goal engine (any non-terminal
@@ -97,9 +197,12 @@ impl SessionActor {
);
}
let graph_status = self.graph_tracker.lock().status();
if matches!(graph_status, Some(s) if s == GoalStatus::Active || s.is_paused()) {
// Refuse over ANYTHING non-Complete — including BudgetLimited,
// which is resumable and must never be silently replaced.
if matches!(graph_status, Some(s) if s != GoalStatus::Complete) {
return GraphSetupOutcome::Message(
"A graph is already set. Use /graph status, /graph resume, or /graph clear."
"A graph is already set. Use /graph status, /graph resume \
[--budget <tokens>], or /graph clear."
.to_owned(),
);
}
@@ -282,7 +385,8 @@ impl SessionActor {
self.graph_tracker.lock().budget_limit();
self.persist_graph_state();
self.send_slash_command_output(
"Graph token budget exhausted. Use /graph clear, then /graph <objective>.",
"Graph token budget exhausted. Top up with /graph resume --budget <tokens>, \
or /graph clear to abandon.",
)
.await;
return None;
@@ -347,9 +451,21 @@ impl SessionActor {
// Node goals are armed with the remaining graph budget,
// so a node-level trip IS the graph-level trip. The
// enforce-side cascade normally handles this; mirroring
// here is idempotent (budget_limit is Active-only).
// here is idempotent (budget_limit is Active-only), and
// the two charge sites are mutually exclusive with it.
tracing::warn!("graph: node goal budget-limited; graph budget-limited");
self.graph_tracker.lock().budget_limit();
let current_tokens = self.chat_state_handle.get_total_tokens().await as i64;
let node_tokens = self.goal_tokens_used(current_tokens);
{
let mut tracker = self.graph_tracker.lock();
// Budget integrity: the tripped node's partial burn
// was still spent — charge it before the demotion
// clears current_node.
if let Some(node_id) = tracker.current_node_id().map(str::to_owned) {
tracker.charge_node_tokens(&node_id, node_tokens);
}
tracker.budget_limit();
}
self.persist_graph_state();
None
}
@@ -454,7 +570,8 @@ impl SessionActor {
self.graph_tracker.lock().budget_limit();
self.persist_graph_state();
self.send_slash_command_output(
"Graph token budget exhausted. Use /graph clear, then /graph <objective>.",
"Graph token budget exhausted. Top up with /graph resume --budget <tokens>, \
or /graph clear to abandon.",
)
.await;
return None;
@@ -614,7 +731,7 @@ impl SessionActor {
/// included); otherwise (post-restart empty engine) launch the next
/// `Ready` node fresh — the verifier gates completion, so re-running
/// a node is always safe.
pub(super) async fn resume_graph(&self) -> GraphSetupOutcome {
pub(super) async fn resume_graph(&self, extra_budget: Option<i64>) -> GraphSetupOutcome {
use super::goal_support::GoalResumeOutcome;
let status = self.graph_tracker.lock().status();
match status {
@@ -627,10 +744,43 @@ impl SessionActor {
Some(GoalStatus::Complete) => GraphSetupOutcome::Message(
"Graph is already complete. Use /graph <objective> to start a new one.".to_owned(),
),
Some(GoalStatus::BudgetLimited) => GraphSetupOutcome::Message(
"Graph is budget-limited. Use /graph clear, then /graph <objective>.".to_owned(),
),
Some(GoalStatus::BudgetLimited) => {
let Some(extra) = extra_budget else {
return GraphSetupOutcome::Message(
"Graph is budget-limited. Top up with /graph resume --budget \
<tokens>, or /graph clear to abandon."
.to_owned(),
);
};
if !self.graph_tracker.lock().resume_budget_limited(extra) {
return GraphSetupOutcome::Message(
"Budget top-up must be a positive token count.".to_owned(),
);
}
self.persist_graph_state();
tracing::info!(extra, "graph: resumed with budget top-up");
match self.drive_graph().await {
Some(reminder) => GraphSetupOutcome::Inference {
reminder,
user_msg: format!("Graph resumed with {extra} fresh budget tokens."),
},
None => GraphSetupOutcome::Message(
"Graph resumed with fresh budget but did not enter a serial node. \
See /graph status."
.to_owned(),
),
}
}
Some(s) if s.is_paused() => {
if extra_budget.is_some() {
// Never silently discard an explicit flag.
return GraphSetupOutcome::Message(
"--budget only applies to a budget-limited graph; this graph is \
paused. Use /graph resume (no flags) to continue, or /graph \
clear to abandon."
.to_owned(),
);
}
{
let mut tracker = self.graph_tracker.lock();
tracker.resume();
@@ -873,7 +873,7 @@ impl SessionActor {
BuiltinAction::GraphSet { .. } => {
unreachable!("GraphSet is intercepted in handle_prompt")
}
BuiltinAction::GraphResume => {
BuiltinAction::GraphResume { .. } => {
unreachable!("GraphResume is intercepted in handle_prompt")
}
}
@@ -1244,6 +1244,13 @@ pub(crate) async fn spawn_session_actor(
),
)
.await;
// A restored graph was demoted (Active→UserPaused, Running→Ready) IN
// MEMORY after the updates-log replay, whose last GraphUpdated still
// shows the pre-shutdown Active state. Re-emit truth once so a
// reattached pager never renders a stale self-driving chip.
if session.graph_tracker.lock().snapshot().is_some() {
session.persist_graph_state();
}
if let Some(ref display_cwd) = prompt_display_cwd {
session
.agent
@@ -350,16 +350,18 @@ impl SessionActor {
return ok_end_turn(0, None);
}
},
BuiltinAction::GraphResume => match self.resume_graph().await {
super::graph::GraphSetupOutcome::Inference { reminder, user_msg } => {
self.send_slash_command_output(&user_msg).await;
vec![text_block(reminder)]
BuiltinAction::GraphResume { extra_budget } => {
match self.resume_graph(extra_budget).await {
super::graph::GraphSetupOutcome::Inference { reminder, user_msg } => {
self.send_slash_command_output(&user_msg).await;
vec![text_block(reminder)]
}
super::graph::GraphSetupOutcome::Message(msg) => {
self.send_slash_command_output(&msg).await;
return ok_end_turn(0, None);
}
}
super::graph::GraphSetupOutcome::Message(msg) => {
self.send_slash_command_output(&msg).await;
return ok_end_turn(0, None);
}
},
}
_ => return self.execute_builtin_slash_command(action).await,
}
}
@@ -421,7 +421,7 @@ async fn graph_restore_demotes_running_node_and_resume_relaunches_it() {
);
// /graph resume relaunches node b as a fresh goal.
match restored.resume_graph().await {
match restored.resume_graph(None).await {
graph::GraphSetupOutcome::Inference { reminder, .. } => {
assert!(
reminder.contains("Graph node 2/4"),
@@ -474,9 +474,9 @@ async fn node_budget_is_graph_remaining_and_trip_cascades() {
Some(GoalStatus::BudgetLimited),
"node budget trip IS the graph budget trip"
);
// The in-flight node is terminally resolved — never a
// forever-Running node on a budget-dead graph.
assert_eq!(node_statuses(&actor)[0].1, NodeStatus::Failed);
// The in-flight node demotes to Ready — a budget trip is a
// resource stop, not a node verdict; a top-up re-runs it.
assert_eq!(node_statuses(&actor)[0].1, NodeStatus::Ready);
assert_eq!(actor.graph_tracker.lock().current_node_id(), None);
// A terminal graph no longer owns the engine: the user may
@@ -648,7 +648,7 @@ async fn planning_invalid_twice_pauses_and_resume_replans() {
// /graph resume re-plans; a now-valid artifact launches node 1.
let (good_tx, _good_count) = spawn_graph_planner_coordinator(vec![chain_graph_json()]);
actor.tool_context.subagent_event_tx = Some(good_tx);
match actor.resume_graph().await {
match actor.resume_graph(None).await {
graph::GraphSetupOutcome::Inference { reminder, .. } => {
assert!(reminder.contains("Graph node 1/4"), "{reminder}");
}
@@ -1252,7 +1252,7 @@ async fn resume_after_cancelled_batch_demotes_orphaned_running_nodes() {
// /graph resume must demote the orphans and re-dispatch — NOT
// wedge-pause with "no runnable node".
match actor.resume_graph().await {
match actor.resume_graph(None).await {
graph::GraphSetupOutcome::Inference { reminder, .. } => {
assert!(reminder.contains("Node C"), "{reminder}");
}
@@ -1400,3 +1400,222 @@ async fn batch_merges_real_worktrees_and_cleans_them_up() {
.await;
unsafe { std::env::remove_var(ENV_FLAG) };
}
#[tokio::test(flavor = "current_thread")]
#[serial]
async fn budget_top_up_resumes_a_budget_limited_graph_to_completion() {
unsafe { std::env::set_var(ENV_FLAG, "0") };
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let dag = diamond_graph_json();
let (mut actor, _tmp, _prx) = make_graph_actor_detached().await;
let (coord_tx, _captured) = spawn_scripted_coordinator(
_tmp.path().to_path_buf(),
move |req| happy_reply(req, &dag),
false,
);
actor.tool_context.subagent_event_tx = Some(coord_tx);
actor.graph_concurrency = 2;
// Tiny budget: the a+b batch spends 40 (4 spawns × 10) > 30,
// so the inter-batch gate trips before c.
let _ = actor.setup_graph("build the diamond", Some(30)).await;
assert_eq!(
actor.graph_tracker.lock().status(),
Some(GoalStatus::BudgetLimited)
);
// Plain resume must refuse with the top-up hint.
match actor.resume_graph(None).await {
graph::GraphSetupOutcome::Message(msg) => {
assert!(msg.contains("--budget"), "{msg}");
}
graph::GraphSetupOutcome::Inference { .. } => {
panic!("budget-limited graph must not resume without a top-up")
}
}
// Top-up resumes and reaches the serial tail (node c).
match actor.resume_graph(Some(1_000)).await {
graph::GraphSetupOutcome::Inference { reminder, .. } => {
assert!(reminder.contains("Node C"), "{reminder}");
}
graph::GraphSetupOutcome::Message(msg) => {
panic!("top-up must resume the graph, got: {msg}")
}
}
// Finish c + gn-final serially.
for _ in 0..2 {
drive_node_goal_to_complete(&actor).await;
let _ = actor.run_graph_round_end().await;
}
assert_eq!(
actor.graph_tracker.lock().status(),
Some(GoalStatus::Complete)
);
})
.await;
unsafe { std::env::remove_var(ENV_FLAG) };
}
// ── handle_prompt-level coverage (real interception wiring) ────────────
fn agent_text(n: &acp::SessionNotification) -> Option<String> {
match &n.update {
acp::SessionUpdate::AgentMessageChunk(chunk) => match &chunk.content {
acp::ContentBlock::Text(t) => Some(t.text.clone()),
_ => None,
},
_ => None,
}
}
fn capture_gateway(
mut gateway_rx: tokio::sync::mpsc::UnboundedReceiver<kigi_acp_lib::AcpClientMessage>,
) -> StdArc<tokio::sync::Mutex<Vec<acp::SessionNotification>>> {
let sent = StdArc::new(tokio::sync::Mutex::new(Vec::new()));
let sent_for_task = sent.clone();
tokio::task::spawn_local(async move {
while let Some(msg) = gateway_rx.recv().await {
if let kigi_acp_lib::AcpClientMessage::SessionNotification(args) = msg {
sent_for_task.lock().await.push(args.request);
let _ = args.response_tx.send(Ok(()));
}
}
});
sent
}
fn drain_replay(
actor: StdArc<SessionActor>,
mut event_rx: tokio::sync::mpsc::UnboundedReceiver<SessionEvent>,
) {
let settings = actor.buffering_settings.clone();
tokio::task::spawn_local(async move {
let mut replay_buffer = ReplayBuffer::new(settings);
while let Some(event) = event_rx.recv().await {
match event {
SessionEvent::Notification(notification) => {
if let Some((primary, secondary)) = replay_buffer.consume_chunk(notification) {
actor.emit_buffered(primary).await;
if let Some(extra) = secondary {
actor.emit_buffered(extra).await;
}
}
}
SessionEvent::FlushReplay { respond_to } => {
if let Some(notification) = replay_buffer.flush() {
actor.emit_buffered(notification).await;
}
if let Some(tx) = respond_to {
let _ = tx.send(());
}
}
}
}
});
}
/// Actor whose `/graph` slash commands resolve through the REAL
/// `handle_prompt` path: `update_goal` registered (goal gate),
/// `graph_enabled` set, gateway capture + replay drainer wired.
async fn make_graph_turn_actor() -> (
StdArc<SessionActor>,
StdArc<tokio::sync::Mutex<Vec<acp::SessionNotification>>>,
) {
let (gateway_tx, gateway_rx) =
tokio::sync::mpsc::unbounded_channel::<kigi_acp_lib::AcpClientMessage>();
let sent = capture_gateway(gateway_rx);
let (persistence_tx, _persistence_rx) =
tokio::sync::mpsc::unbounded_channel::<PersistenceMsg>();
let (mut actor, event_rx) =
create_test_actor_ex(0, 256_000, 85, gateway_tx, persistence_tx).await;
*actor.agent.borrow_mut() = test_agent_with_goal_tool().await;
actor.goal_enabled = true;
actor.graph_enabled = true;
let actor = StdArc::new(actor);
drain_replay(actor.clone(), event_rx);
(actor, sent)
}
async fn drive_terminal_slash(
actor: &StdArc<SessionActor>,
sent: &StdArc<tokio::sync::Mutex<Vec<acp::SessionNotification>>>,
prompt: &str,
) -> String {
let result = tokio::time::timeout(
std::time::Duration::from_secs(5),
actor.handle_prompt(
&format!("graph-slash-{}", prompt.replace([' ', '/'], "-")),
vec![acp::ContentBlock::Text(acp::TextContent::new(
prompt.to_string(),
))],
PromptMode::Agent,
None,
None,
false,
None,
None,
),
)
.await
.expect("terminal slash outcome must end the turn without inference");
assert!(result.is_ok(), "turn must succeed: {result:?}");
tokio::task::yield_now().await;
let sent = sent.lock().await;
sent.iter()
.filter_map(agent_text)
.collect::<Vec<_>>()
.join("\n")
}
#[tokio::test(flavor = "current_thread")]
#[serial]
async fn graph_slash_terminal_outcomes_through_handle_prompt() {
unsafe { std::env::set_var(ENV_FLAG, "0") };
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let (actor, sent) = make_graph_turn_actor().await;
// /graph status with no graph: terminal message, no inference.
let text = drive_terminal_slash(&actor, &sent, "/graph status").await;
assert!(
text.contains("No graph is set"),
"status message must reach the gateway: {text}"
);
// /graph resume with no graph.
let text = drive_terminal_slash(&actor, &sent, "/graph resume").await;
assert!(text.contains("No graph is set"), "{text}");
// Seed an Active graph that owns the engine; /goal commands
// must be refused through the REAL interception guards.
actor.graph_tracker.lock().create_graph(
"g-1".into(),
"obj".into(),
None,
"2026-01-01T00:00:00Z".into(),
);
let text = drive_terminal_slash(&actor, &sent, "/goal pause").await;
assert!(
text.contains("graph owns the goal engine"),
"/goal pause refusal must surface: {text}"
);
let text = drive_terminal_slash(&actor, &sent, "/goal clear").await;
assert!(text.contains("graph owns the goal engine"), "{text}");
// /graph pause on the planning-phase Active graph pauses it.
let text = drive_terminal_slash(&actor, &sent, "/graph pause").await;
assert!(text.contains("Graph paused"), "{text}");
assert!(
actor
.graph_tracker
.lock()
.status()
.is_some_and(|s| s.is_paused())
);
})
.await;
unsafe { std::env::remove_var(ENV_FLAG) };
}
@@ -562,8 +562,11 @@ impl GraphTracker {
}
/// `Active -> BudgetLimited`; `true` if the transition happened.
/// The in-flight node (if any) is terminally resolved to `Failed` so
/// a budget-dead graph never persists a forever-`Running` node.
/// Every in-flight node (the serial engine node AND any parallel
/// batch node) demotes to `Ready`: a budget trip is a resource
/// stop, not a verdict on the node — so no forever-`Running` node
/// is ever persisted, and a later budget top-up
/// ([`Self::resume_budget_limited`]) re-dispatches it naturally.
pub fn budget_limit(&mut self) -> bool {
self.account_elapsed();
let Some(state) = self.state.as_mut() else {
@@ -573,12 +576,10 @@ impl GraphTracker {
return false;
}
let in_flight = state.current_node.clone();
if let Some(node_id) = &in_flight
&& let Some(node) = state.nodes.iter_mut().find(|n| n.id == *node_id)
&& !node.status.is_terminal()
{
node.status = NodeStatus::Failed;
node.failure = Some("graph token budget exhausted".to_owned());
for node in &mut state.nodes {
if matches!(node.status, NodeStatus::Running | NodeStatus::Verifying) {
node.status = NodeStatus::Ready;
}
}
state.current_node = None;
state.status = GoalStatus::BudgetLimited;
@@ -591,6 +592,31 @@ impl GraphTracker {
true
}
/// `BudgetLimited -> Active` with `extra` fresh headroom: the new
/// budget becomes spent-so-far + extra. `true` if applied.
pub fn resume_budget_limited(&mut self, extra: i64) -> bool {
let Some(state) = self.state.as_mut() else {
return false;
};
if state.status != GoalStatus::BudgetLimited || extra <= 0 {
return false;
}
state.token_budget = Some(state.tokens_spent_nodes.saturating_add(extra));
state.status = GoalStatus::Active;
state.pause_message = None;
push_history(
state,
GraphHistoryEntry::now(
GraphEvent::GraphResumed,
None,
Some(format!("budget topped up by {extra}")),
),
);
self.last_probe = Some(Instant::now());
self.recompute_ready();
true
}
/// Drop all graph state (history records the clear first so a
/// final persisted snapshot, if any, carries it).
pub fn clear(&mut self) {
@@ -860,27 +886,33 @@ mod tests {
}
#[test]
fn budget_limit_terminally_fails_the_running_node() {
fn budget_limit_demotes_in_flight_and_top_up_resumes() {
let mut t = GraphTracker::new(std::env::temp_dir());
t.create_graph("g".into(), "o".into(), Some(10), "t".into());
t.install_nodes(vec![node("a", &[]), node("b", &["a"])]);
t.mark_node_running("a", "goal-1".into());
t.charge_node_tokens("a", 12);
assert!(t.budget_limit());
assert_eq!(t.status(), Some(GoalStatus::BudgetLimited));
let s = t.snapshot().unwrap();
assert_eq!(
s.nodes[0].status,
NodeStatus::Failed,
"a budget-dead graph must not persist a forever-Running node"
);
assert!(
s.nodes[0]
.failure
.as_deref()
.unwrap()
.contains("budget exhausted")
NodeStatus::Ready,
"budget trip is a resource stop, not a node verdict — no \
forever-Running node, runnable again after a top-up"
);
assert_eq!(s.current_node, None);
assert!(!t.resume_budget_limited(0), "top-up must be positive");
assert!(t.resume_budget_limited(100));
assert_eq!(t.status(), Some(GoalStatus::Active));
assert_eq!(
t.snapshot().unwrap().token_budget,
Some(112),
"new budget = spent so far + extra headroom"
);
assert_eq!(t.remaining_budget(), Some(100));
assert_eq!(t.next_ready_node().unwrap().id, "a");
}
#[test]
@@ -263,7 +263,9 @@ pub(super) const BUILTIN_COMMANDS: &[BuiltinCommand] = &[
BuiltinCommand {
name: "graph",
description: "Decompose an objective into a dependency graph of autonomous goals",
argument_hint: Some("<objective> [--budget <tokens>] | status | pause | resume | clear"),
argument_hint: Some(
"<objective> [--budget <tokens>] | status | pause | resume [--budget <tokens>] | clear",
),
aliases: &[],
gate: BuiltinGate::Graph,
resolve: |args| {
@@ -273,9 +275,27 @@ pub(super) const BUILTIN_COMMANDS: &[BuiltinCommand] = &[
// it is an alias for the status tree.
"" | "status" | "show" => BuiltinAction::GraphStatus,
"pause" => BuiltinAction::GraphPause,
"resume" => BuiltinAction::GraphResume,
"resume" => BuiltinAction::GraphResume { extra_budget: None },
"clear" => BuiltinAction::GraphClear,
_ => {
// ANY input starting with `resume` is a resume attempt
// and must NEVER fall through to GraphSet — a typo'd
// top-up would otherwise silently replace a resumable
// BudgetLimited graph. Well-formed `resume --budget
// <tokens>` (case-insensitive keywords) carries the
// top-up; malformed variants resolve to a plain
// resume, whose BudgetLimited arm prints the usage.
let lower = trimmed.to_lowercase();
if let Some(rest) = lower.strip_prefix("resume") {
let extra_budget = rest
.trim()
.strip_prefix("--budget")
.map(str::trim)
.filter(|v| !v.is_empty() && v.bytes().all(|b| b.is_ascii_digit()))
.and_then(|v| v.parse::<i64>().ok())
.filter(|extra| *extra > 0);
return BuiltinAction::GraphResume { extra_budget };
}
let (objective, token_budget) = parse_goal_budget(trimmed);
BuiltinAction::GraphSet {
objective,
@@ -700,7 +720,9 @@ pub(super) enum BuiltinAction {
},
GraphStatus,
GraphPause,
GraphResume,
GraphResume {
extra_budget: Option<i64>,
},
GraphClear,
}
@@ -737,7 +759,7 @@ impl BuiltinAction {
BuiltinAction::GraphSet { .. }
| BuiltinAction::GraphStatus
| BuiltinAction::GraphPause
| BuiltinAction::GraphResume
| BuiltinAction::GraphResume { .. }
| BuiltinAction::GraphClear => "graph",
}
}
@@ -772,10 +794,10 @@ impl BuiltinAction {
| BuiltinAction::GoalResume
| BuiltinAction::GoalClear => false,
BuiltinAction::GraphSet { .. } => true,
BuiltinAction::GraphStatus
| BuiltinAction::GraphPause
| BuiltinAction::GraphResume
| BuiltinAction::GraphClear => false,
BuiltinAction::GraphResume { extra_budget } => extra_budget.is_some(),
BuiltinAction::GraphStatus | BuiltinAction::GraphPause | BuiltinAction::GraphClear => {
false
}
}
}
}
@@ -1714,7 +1736,26 @@ mod tests {
));
assert!(matches!(
resolve_builtin("graph", "resume"),
Some(BuiltinAction::GraphResume)
Some(BuiltinAction::GraphResume { extra_budget: None })
));
assert!(matches!(
resolve_builtin("graph", "resume --budget 800"),
Some(BuiltinAction::GraphResume {
extra_budget: Some(800)
})
));
// Malformed top-ups resolve to a PLAIN resume — never to
// GraphSet, which would silently replace a resumable graph; the
// BudgetLimited resume arm then prints the usage hint.
assert!(matches!(
resolve_builtin("graph", "resume --budget nope"),
Some(BuiltinAction::GraphResume { extra_budget: None })
));
assert!(matches!(
resolve_builtin("graph", "Resume --Budget 800"),
Some(BuiltinAction::GraphResume {
extra_budget: Some(800)
})
));
assert!(matches!(
resolve_builtin("graph", "clear"),
@@ -977,6 +977,37 @@ pub(super) fn handle_session_notification(notif: &acp::ExtNotification, app: &mu
XaiSessionUpdate::InteractionResolved { tool_call_id } => {
agent.dismiss_resolved_interaction(&tool_call_id)
}
XaiSessionUpdate::GraphUpdated {
objective,
status,
total_nodes,
achieved_nodes,
failed_nodes,
running_nodes,
current_node_title,
token_budget,
tokens_spent,
pause_message,
..
} => {
if status == "cleared" {
agent.graph_state.take();
} else {
agent.graph_state = Some(crate::app::agent::GraphDisplayState {
objective,
status: crate::app::agent::GoalDisplayStatus::parse(&status),
total_nodes,
achieved_nodes,
failed_nodes,
running_nodes,
current_node_title,
token_budget,
tokens_spent,
pause_message,
});
}
true
}
_ => {
tracing::trace!(
"Ignoring {}: {:?}",
+19
View File
@@ -385,6 +385,25 @@ impl GoalDisplayPhase {
}
}
}
/// Graph mode display state — the pager-side mirror of the
/// `GraphUpdated` session notification (see the shell's
/// `extensions/notification.rs`). Deliberately lean: the status chip
/// shows counts + the current node; details live in `/graph status`.
#[derive(Debug, Clone)]
pub struct GraphDisplayState {
pub objective: String,
/// Reuses the goal display-status vocabulary (same wire strings).
pub status: GoalDisplayStatus,
pub total_nodes: u32,
pub achieved_nodes: u32,
pub failed_nodes: u32,
pub running_nodes: u32,
pub current_node_title: Option<String>,
pub token_budget: Option<i64>,
pub tokens_spent: i64,
pub pause_message: Option<String>,
}
/// Display state for an active goal, populated from `GoalUpdated`
/// session notifications emitted by the goal orchestrator.
#[derive(Debug, Clone)]
@@ -745,6 +745,9 @@ pub struct AgentView {
/// Current goal orchestration state. Set by `GoalUpdated` session
/// notifications, cleared when a new session starts.
pub goal_state: Option<super::agent::GoalDisplayState>,
/// Current graph orchestration state. Set by `GraphUpdated` session
/// notifications, cleared on the `"cleared"` sentinel / new session.
pub graph_state: Option<super::agent::GraphDisplayState>,
/// The consumed parked-wait marker slot for the current turn, if any.
/// Keyed by prompt id: a new turn naturally invalidates the slot with no
/// explicit clear site. See [`ParkedMarkerSlot`].
@@ -1243,6 +1243,13 @@ impl AgentView {
),
);
}
if let Some(ref graph) = self.graph_state {
let tick = self.tasks.tick_count() as usize;
status.push(
"graph",
crate::views::agent_status::graph_status_line(graph, &theme, tick),
);
}
if let Some(mcp_line) = self.mcp_init_progress.as_ref().and_then(|p| {
crate::views::agent_status::mcp_status_line(p, self.scrollback.animation_tick(), &theme)
}) {
@@ -93,6 +93,7 @@ impl AgentView {
chat_kind: false,
app_chat_mode: false,
goal_state: None,
graph_state: None,
parked_wait_marker_for: None,
end_work_announced: false,
pending_stop_hooks: None,
@@ -282,6 +282,75 @@ pub fn goal_status_line(
])
}
// ---------------------------------------------------------------------------
// Graph status chip
// ---------------------------------------------------------------------------
/// Build the compact `/graph` status chip: node progress, the current
/// node, and spend. Same chip idiom as [`goal_status_line`] — dim
/// brackets, paused chips invert onto `theme.warning`, active chips
/// animate.
pub fn graph_status_line(
graph: &crate::app::agent::GraphDisplayState,
theme: &Theme,
tick: usize,
) -> Line<'static> {
let dim_style = Style::default().fg(theme.gray_dim).bg(theme.bg_base);
let label_style = if graph.status.is_paused() {
Style::default().fg(theme.bg_base).bg(theme.warning)
} else {
Style::default().fg(theme.accent_plan).bg(theme.bg_base)
};
let is_active = matches!(graph.status, GoalDisplayStatus::Active);
let mut progress = format!("{}/{}", graph.achieved_nodes, graph.total_nodes);
if graph.failed_nodes > 0 {
progress.push_str(&format!(" ({} failed)", graph.failed_nodes));
}
// Planner titles are uncapped; clamp so one long title can't push
// the whole status bar off-screen.
let clamped_title = graph.current_node_title.as_deref().map(|t| {
if t.chars().count() > 40 {
let head: String = t.chars().take(39).collect();
format!("{head}")
} else {
t.to_owned()
}
});
let label = match (&graph.status, clamped_title.as_deref()) {
(GoalDisplayStatus::Active, Some(title)) => format!("{progress} · {title}"),
(GoalDisplayStatus::Active, None) if graph.running_nodes > 1 => {
format!("{progress} · {} nodes in flight", graph.running_nodes)
}
(GoalDisplayStatus::Complete, _) => format!("{progress} · complete"),
(GoalDisplayStatus::BudgetLimited, _) => format!("{progress} · budget limit"),
(status, _) if status.is_paused() => format!("{progress} · paused"),
_ => progress,
};
let graph_text = if is_active {
let frames = crate::glyphs::dot_spinner_frames();
let frame = frames[(tick / 4) % frames.len()];
format!("{frame} Graph: {label}")
} else {
format!("Graph: {label}")
};
let tokens_str = format_tokens_compact(graph.tokens_spent.max(0));
let tokens_display = match graph.token_budget {
Some(budget) if budget > 0 => {
format!("{}/{} tokens", tokens_str, format_tokens_compact(budget))
}
_ => format!("{tokens_str} tokens"),
};
Line::from(vec![
Span::styled("[", dim_style),
Span::styled(graph_text, label_style),
Span::styled("]", dim_style),
Span::styled(format!(" {tokens_display}"), dim_style),
])
}
// ---------------------------------------------------------------------------
// MCP connecting indicator
// ---------------------------------------------------------------------------
@@ -0,0 +1,40 @@
name: graph-slash-presession
description: >
With the graph feature flag on (`KIGI_GRAPH=1`, plus `KIGI_GOAL=1` — the
graph gate requires the goal harness), `/graph` must appear in the
slash-command menu on the welcome screen *before* the first user turn
creates a session. Typing `/graph` pre-session must surface the command
and its description in the dropdown.
terminal:
rows: 40
cols: 100
environment:
env:
- key: KIGI_GOAL
value: "1"
- key: KIGI_GRAPH
value: "1"
mock:
response: "unused — this scenario never submits a prompt."
steps:
- action: wait_for_text
text: Quit
timeout_ms: 20000
- action: assert_not_contains
text: panicked
# Sanity: nothing has been submitted yet, so no session exists.
- action: assert_not_contains
text: dependency graph
- action: focus_prompt
- action: type_text
text: "/graph"
- action: wait
millis: 300
# The dropdown row carries the builtin's description, which only renders
# when /graph is actually advertised pre-session.
- action: assert_contains
text: dependency graph
- action: assert_running
- action: screenshot
name: graph-in-slash-menu-presession
note: "/graph advertised in the slash menu on the welcome screen (KIGI_GRAPH=1), before any prompt."
@@ -0,0 +1,45 @@
name: graph-slash-presession-disabled
description: >
Fail-closed counterpart: with `KIGI_GRAPH=0` (goal harness still on),
`/graph` must NOT appear in the pre-session slash menu, while the menu
itself keeps working for other commands.
terminal:
rows: 40
cols: 100
environment:
env:
- key: KIGI_GOAL
value: "1"
- key: KIGI_GRAPH
value: "0"
mock:
response: "unused — this scenario never submits a prompt."
steps:
- action: wait_for_text
text: Quit
timeout_ms: 20000
- action: assert_not_contains
text: panicked
- action: focus_prompt
- action: type_text
text: "/graph"
- action: wait
millis: 300
- action: assert_not_contains
text: dependency graph
- action: assert_running
- action: screenshot
name: graph-absent-from-slash-menu
note: "/graph hidden with KIGI_GRAPH=0 (fail-closed gate)."
# The menu itself still works: clear "/graph" (6 chars) and try /compact.
- action: keys
keys: "<BS><BS><BS><BS><BS><BS>"
- action: type_text
text: "/compact"
- action: wait
millis: 300
- action: assert_contains
text: Compact conversation history
- action: screenshot
name: slash-menu-still-works
note: "Slash menu functional; only /graph is gated off."
@@ -339,6 +339,21 @@ async fn scripted_goal_slash_presession() {
run_scenario("goal_slash_presession.yaml").await;
}
/// `/graph` gating mirror of the `/goal` pre-session scenarios: with
/// `KIGI_GRAPH=1` (+ the goal harness) the command must be advertised on
/// the welcome screen before any session exists.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[ignore = "scripted scenario; run with cargo test -- --ignored"]
async fn scripted_graph_slash_presession() {
run_scenario("graph_slash_presession.yaml").await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[ignore = "scripted scenario; run with cargo test -- --ignored"]
async fn scripted_graph_slash_presession_disabled() {
run_scenario("graph_slash_presession_disabled.yaml").await;
}
/// Counterpart to `scripted_goal_slash_presession`: with the goal flag
/// explicitly off (`KIGI_GOAL=0`; goal mode defaults on), `/goal` must stay
/// hidden pre-session (gate fail-closed) while an
@@ -516,6 +531,8 @@ fn scenarios_parse() {
"path_space_hyperlink.yaml",
"goal_slash_presession.yaml",
"goal_slash_presession_disabled.yaml",
"graph_slash_presession.yaml",
"graph_slash_presession_disabled.yaml",
"folder_trust_prompt.yaml",
"dashboard_model_list_click.yaml",
"paste_chip_double_click.yaml",