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:
2 files changed
+74
-11
No files matched your search
+73
-10
@@ -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(®istry));
|
||||
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 = {
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user