fix(graph): retire the teammate identity before structured extraction

PR #24 follow-up 10. The identity was only retired by the frontier and
chain-exit safety nets, so a peer stayed messageable during
structured::extract after its agent had already stopped. Mark it
finished the moment the agent future resolves, on success and on both
error paths.
This commit is contained in:
Dark-Alex-17 committed 2026-09-11 10:56:27 -06:00
1 parent 874ef7f6ec
commit 0bef0f2c39
2 files changed
+74 -11

No files matched your search

+73 -10
View File
@@ -18,6 +18,7 @@ impl AgentNodeExecutor {
node: &AgentNode,
state_manager: &mut StateManager,
parent_ctx: &mut RequestContext,
retire_peer_on_return: bool,
) -> Result<String> {
let prompt = state_manager
.interpolate(&node.prompt)
@@ -25,19 +26,34 @@ impl AgentNodeExecutor {
let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS));
let raw = timeout(
// run_agent_for_graph takes the identity off the ctx; keep a handle so
// a frontier peer is retired the moment the agent stops, not after
// extraction. Chains re-arm the same identity for later steps and
// leave retirement to the chain runner.
let peer = parent_ctx.peer_registry.clone().zip(
parent_ctx
.peer_assignment
.as_ref()
.map(|(id, _)| id.clone()),
);
let raw_result = timeout(
timeout_dur,
run_agent_for_graph(parent_ctx, &node.agent, &prompt),
)
.await
.with_context(|| {
format!(
"Agent '{}' timed out after {}s",
node.agent,
timeout_dur.as_secs()
)
})?
.with_context(|| format!("Agent '{}' failed", node.agent))?;
.await;
if retire_peer_on_return && let Some((registry, id)) = &peer {
registry.mark_finished(id);
}
let raw = raw_result
.with_context(|| {
format!(
"Agent '{}' timed out after {}s",
node.agent,
timeout_dur.as_secs()
)
})?
.with_context(|| format!("Agent '{}' failed", node.agent))?;
let output_value = match &node.output_schema {
Some(schema) => structured::extract(&raw, schema, parent_ctx)
@@ -70,8 +86,11 @@ fn apply_state_updates(node: &AgentNode, state_manager: &mut StateManager, outpu
mod tests {
use super::super::types::AgentNode;
use super::*;
use crate::config::{AppState, WorkingMode, default_max_agent_depth};
use crate::supervisor::mailbox::{Inbox, PeerRegistry, graph_agent_id};
use serde_json::json;
use std::collections::HashMap;
use std::sync::Arc;
fn manager_with(pairs: &[(&str, Value)]) -> StateManager {
let mut map = HashMap::new();
@@ -92,6 +111,50 @@ mod tests {
}
}
fn ctx_at_max_depth_with_peer() -> (RequestContext, Arc<PeerRegistry>, String) {
let registry = Arc::new(PeerRegistry::new());
let id = graph_agent_id("test_agent");
let inbox = Arc::new(Inbox::new());
registry.insert(id.clone(), "worker[0]".into(), Arc::clone(&inbox));
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
ctx.current_depth = default_max_agent_depth();
ctx.peer_registry = Some(Arc::clone(&registry));
ctx.peer_assignment = Some((id.clone(), inbox));
(ctx, registry, id)
}
async fn execute_past_max_depth(ctx: &mut RequestContext, retire_peer_on_return: bool) {
let mut node = node_with("hi", None);
node.teammates = true;
let mut state = manager_with(&[]);
let err = AgentNodeExecutor::execute(&node, &mut state, ctx, retire_peer_on_return)
.await
.expect_err("agent past max depth should fail before running");
let chain = format!("{err:#}");
assert!(chain.contains("Agent 'test_agent' failed"), "{chain}");
assert!(chain.contains("Max agent depth exceeded"), "{chain}");
}
#[tokio::test]
async fn execute_retires_peer_identity_on_frontier_when_agent_fails() {
let (mut ctx, registry, id) = ctx_at_max_depth_with_peer();
execute_past_max_depth(&mut ctx, true).await;
assert!(registry.is_finished(&id));
}
#[tokio::test]
async fn execute_leaves_peer_identity_live_in_branch_mode_when_agent_fails() {
let (mut ctx, registry, id) = ctx_at_max_depth_with_peer();
execute_past_max_depth(&mut ctx, false).await;
assert!(!registry.is_finished(&id));
}
#[test]
fn state_updates_use_output_placeholder() {
let node = {
+1 -1
View File
@@ -465,7 +465,7 @@ pub(super) async fn step(
) -> Result<StepResult> {
match &node.node_type {
NodeType::Agent(agent_node) => {
AgentNodeExecutor::execute(agent_node, state, ctx).await?;
AgentNodeExecutor::execute(agent_node, state, ctx, !step_ctx.branch_mode).await?;
let targets = static_next_targets(node, current, "agent", step_ctx.branch_mode)?;
Ok(StepResult::Continue(targets))
}