Skip to content

Commit e786d8e

Browse files
ZhiXiao-Linclaude
andauthored
feat(orchestration): pipeline combinator (Phase 3) (#53)
Adds execute_pipeline — the one scheduling shape parallel() can't express: each item flows through a chain of stages independently, with NO barrier between stages, so item A can be in stage 3 while item B is still in stage 1 (wall-clock = slowest single chain, not sum-of-slowest-per-stage). - PipelineStage<I> = Fn(Option<&StepOutcome>, &I) -> Option<AgentStepSpec>: pure spec-builders that can branch on the prior stage's outcome (e.g. "verify the finding the review stage produced"). The executor runs them. - Chains run concurrently bounded by the executor's concurrency_hint; a chain stops early on a None stage or a failed step; a panicking stage isolates to its own chain (-> None) without dropping the others. Reuses run_ordered_parallel_with_limit; built entirely on the AgentExecutor seam. Deferred (Rule 6 — no consumer yet): phase-grouping events (would churn the persisted AgentEvent schema + SDK mapping) and a budget-aware loop_until; both land when the SDK/UI (Phase 5) needs them. Tests: 4 mock-based pipeline tests (chaining + prior-output visibility, stop on failure/None, no-barrier concurrency bound, panic isolation) and a real-LLM \#[ignore] two-stage chain passing against .a3s/config.acl. fmt + clippy --lib --bins clean. Co-authored-by: Claude <claude@anthropic.com>
1 parent 253299a commit e786d8e

4 files changed

Lines changed: 305 additions & 2 deletions

File tree

core/src/lib.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,10 @@ pub use llm::{
135135
ContentBlock, HttpMetricsCallback, HttpMetricsRecord, ImageSource, LlmClient, LlmResponse,
136136
Message, OpenAiClient, TokenUsage,
137137
};
138-
pub use orchestration::{execute_steps_parallel, AgentExecutor, AgentStepSpec, StepOutcome};
138+
pub use orchestration::{
139+
execute_pipeline, execute_steps_parallel, AgentExecutor, AgentStepSpec, PipelineStage,
140+
StepOutcome,
141+
};
139142
pub use prompts::{AgentStyle, DetectionConfidence, PlanningMode, SystemPromptSlots};
140143
pub use run::{
141144
ActiveToolSnapshot, InMemoryRunStore, RunEventRecord, RunHandle, RunRecord, RunSnapshot,
Lines changed: 240 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,240 @@
1+
//! Orchestration combinators built on the [`AgentExecutor`] seam.
2+
//!
3+
//! [`execute_steps_parallel`](super::execute_steps_parallel) (in `executor`)
4+
//! is the barrier (`parallel`) primitive. This module adds `pipeline`: the
5+
//! one genuinely new scheduling shape, where each item flows through a chain
6+
//! of stages independently — no barrier between stages.
7+
8+
use super::executor::{AgentExecutor, AgentStepSpec, StepOutcome};
9+
use crate::agent::AgentEvent;
10+
use crate::ordered_parallel::run_ordered_parallel_with_limit;
11+
use std::sync::Arc;
12+
use tokio::sync::broadcast;
13+
14+
/// A pipeline stage: given the previous stage's outcome (`None` before the
15+
/// first stage) and the original item, produce the next step to run — or
16+
/// `None` to stop this item's chain early.
17+
///
18+
/// Stages are pure spec-builders; the executor runs them. A stage can branch
19+
/// on the prior result (e.g. "verify the finding the review stage produced").
20+
pub type PipelineStage<I> =
21+
Arc<dyn Fn(Option<&StepOutcome>, &I) -> Option<AgentStepSpec> + Send + Sync>;
22+
23+
/// Run each item through `stages` as an independent chain.
24+
///
25+
/// All chains run concurrently, bounded by the executor's
26+
/// [`concurrency_hint`](AgentExecutor::concurrency_hint) — there is **no
27+
/// barrier between stages**, so item A can be in stage 3 while item B is still
28+
/// in stage 1. Wall-clock is the slowest single chain, not the
29+
/// sum-of-slowest-per-stage that a barrier `parallel` per stage would incur.
30+
///
31+
/// A chain stops early when a stage returns `None` or when a step fails
32+
/// (later stages would only build on a failed result). Returns each item's
33+
/// last outcome (`None` if its first stage produced no spec), preserving input
34+
/// order. A stage closure that panics isolates to that one chain (its result
35+
/// becomes `None`) without dropping the others.
36+
pub async fn execute_pipeline<I>(
37+
executor: Arc<dyn AgentExecutor>,
38+
items: Vec<I>,
39+
stages: Vec<PipelineStage<I>>,
40+
event_tx: Option<broadcast::Sender<AgentEvent>>,
41+
) -> Vec<Option<StepOutcome>>
42+
where
43+
I: Send + 'static,
44+
{
45+
let limit = executor.concurrency_hint();
46+
let stages = Arc::new(stages);
47+
48+
let results = run_ordered_parallel_with_limit(items, limit, move |_idx, item| {
49+
let executor = Arc::clone(&executor);
50+
let stages = Arc::clone(&stages);
51+
let event_tx = event_tx.clone();
52+
async move {
53+
let mut prev: Option<StepOutcome> = None;
54+
for stage in stages.iter() {
55+
let Some(spec) = stage(prev.as_ref(), &item) else {
56+
break;
57+
};
58+
let outcome = executor.execute_step(spec, event_tx.clone()).await;
59+
let succeeded = outcome.success;
60+
prev = Some(outcome);
61+
if !succeeded {
62+
break;
63+
}
64+
}
65+
prev
66+
}
67+
})
68+
.await;
69+
70+
// A panicked chain (Err) yields `None`; a normal chain yields its last
71+
// outcome. Order is preserved by `run_ordered_parallel_with_limit`.
72+
results
73+
.into_iter()
74+
.map(|result| result.output.unwrap_or(None))
75+
.collect()
76+
}
77+
78+
#[cfg(test)]
79+
mod tests {
80+
use super::*;
81+
use async_trait::async_trait;
82+
use std::sync::atomic::{AtomicUsize, Ordering};
83+
use std::time::Duration;
84+
85+
/// Echoes the prompt into the output; fails for agent `"fail"`; panics for
86+
/// agent `"boom"`. Records peak concurrency.
87+
struct EchoExecutor {
88+
active: Arc<AtomicUsize>,
89+
max_active: Arc<AtomicUsize>,
90+
}
91+
92+
impl EchoExecutor {
93+
fn new() -> Self {
94+
Self {
95+
active: Arc::new(AtomicUsize::new(0)),
96+
max_active: Arc::new(AtomicUsize::new(0)),
97+
}
98+
}
99+
}
100+
101+
#[async_trait]
102+
impl AgentExecutor for EchoExecutor {
103+
async fn execute_step(
104+
&self,
105+
spec: AgentStepSpec,
106+
_event_tx: Option<broadcast::Sender<AgentEvent>>,
107+
) -> StepOutcome {
108+
let now = self.active.fetch_add(1, Ordering::SeqCst) + 1;
109+
self.max_active.fetch_max(now, Ordering::SeqCst);
110+
tokio::time::sleep(Duration::from_millis(15)).await;
111+
self.active.fetch_sub(1, Ordering::SeqCst);
112+
assert!(spec.agent != "boom", "boom");
113+
StepOutcome {
114+
task_id: spec.task_id.clone(),
115+
session_id: format!("task-run-{}", spec.task_id),
116+
agent: spec.agent.clone(),
117+
output: spec.prompt.clone(),
118+
success: spec.agent != "fail",
119+
structured: None,
120+
}
121+
}
122+
fn concurrency_hint(&self) -> usize {
123+
4
124+
}
125+
}
126+
127+
fn stage<I, F>(f: F) -> PipelineStage<I>
128+
where
129+
F: Fn(Option<&StepOutcome>, &I) -> Option<AgentStepSpec> + Send + Sync + 'static,
130+
{
131+
Arc::new(f)
132+
}
133+
134+
#[tokio::test]
135+
async fn each_item_chains_through_stages_and_later_stages_see_prior_output() {
136+
let exec: Arc<dyn AgentExecutor> = Arc::new(EchoExecutor::new());
137+
// Stage 1: run agent "explore" with the item as the prompt.
138+
// Stage 2: run agent "review" with a prompt derived from stage 1's output.
139+
let stages = vec![
140+
stage(|_prev: Option<&StepOutcome>, item: &&str| {
141+
Some(AgentStepSpec::new("s1", "explore", "d", *item))
142+
}),
143+
stage(|prev: Option<&StepOutcome>, _item: &&str| {
144+
let prior = prev.map(|o| o.output.clone()).unwrap_or_default();
145+
Some(AgentStepSpec::new(
146+
"s2",
147+
"review",
148+
"d",
149+
format!("review of: {prior}"),
150+
))
151+
}),
152+
];
153+
let out = execute_pipeline(exec, vec!["alpha", "beta"], stages, None).await;
154+
155+
assert_eq!(out.len(), 2, "one result per item, order preserved");
156+
// Each item's final outcome is stage 2, whose prompt was derived from
157+
// stage 1's output (the item text).
158+
assert_eq!(out[0].as_ref().unwrap().output, "review of: alpha");
159+
assert_eq!(out[1].as_ref().unwrap().output, "review of: beta");
160+
assert!(out.iter().all(|o| o.as_ref().unwrap().success));
161+
}
162+
163+
#[tokio::test]
164+
async fn chain_stops_on_failure_and_on_none_stage() {
165+
let exec: Arc<dyn AgentExecutor> = Arc::new(EchoExecutor::new());
166+
// First item: stage 1 fails (agent "fail") → stage 2 must not run.
167+
// Second item: stage 1 ok, stage 2 returns None → chain stops at stage 1.
168+
let stages = vec![
169+
stage(|_p: Option<&StepOutcome>, item: &&str| {
170+
let agent = if *item == "x" { "fail" } else { "explore" };
171+
Some(AgentStepSpec::new("s1", agent, "d", *item))
172+
}),
173+
stage(|_p: Option<&StepOutcome>, item: &&str| {
174+
if *item == "y" {
175+
None // stop the second item's chain at stage 1
176+
} else {
177+
Some(AgentStepSpec::new("s2", "review", "d", "second"))
178+
}
179+
}),
180+
];
181+
let out = execute_pipeline(exec, vec!["x", "y"], stages, None).await;
182+
183+
let first = out[0].as_ref().unwrap();
184+
assert!(!first.success, "failed stage 1 surfaces");
185+
assert_eq!(
186+
first.output, "x",
187+
"stage 2 did not run after stage 1 failed"
188+
);
189+
190+
let second = out[1].as_ref().unwrap();
191+
assert!(second.success);
192+
assert_eq!(
193+
second.output, "y",
194+
"stage 2 returned None → chain stopped at stage 1"
195+
);
196+
}
197+
198+
#[tokio::test]
199+
async fn no_barrier_between_stages_bounded_by_hint() {
200+
let echo = EchoExecutor::new();
201+
let max_active = Arc::clone(&echo.max_active);
202+
let exec: Arc<dyn AgentExecutor> = Arc::new(echo);
203+
let stages = vec![
204+
stage(|_p: Option<&StepOutcome>, item: &usize| {
205+
Some(AgentStepSpec::new(
206+
format!("s1-{item}"),
207+
"explore",
208+
"d",
209+
"p",
210+
))
211+
}),
212+
stage(|_p: Option<&StepOutcome>, item: &usize| {
213+
Some(AgentStepSpec::new(format!("s2-{item}"), "review", "d", "p"))
214+
}),
215+
];
216+
let items: Vec<usize> = (0..8).collect();
217+
let out = execute_pipeline(exec, items, stages, None).await;
218+
assert_eq!(out.len(), 8);
219+
assert!(out.iter().all(|o| o.is_some()));
220+
// concurrency_hint is 4: chains run concurrently but never exceed it.
221+
assert!(
222+
max_active.load(Ordering::SeqCst) <= 4,
223+
"concurrency never exceeds the executor's hint"
224+
);
225+
}
226+
227+
#[tokio::test]
228+
async fn panicking_stage_isolates_to_its_chain() {
229+
let exec: Arc<dyn AgentExecutor> = Arc::new(EchoExecutor::new());
230+
let stages = vec![stage(|_p: Option<&StepOutcome>, item: &&str| {
231+
// The middle item routes to the panicking agent.
232+
Some(AgentStepSpec::new("s1", *item, "d", "p"))
233+
})];
234+
let out = execute_pipeline(exec, vec!["explore", "boom", "review"], stages, None).await;
235+
assert_eq!(out.len(), 3);
236+
assert!(out[0].as_ref().unwrap().success);
237+
assert!(out[1].is_none(), "panicked chain becomes None, not a drop");
238+
assert!(out[2].as_ref().unwrap().success, "later chains unaffected");
239+
}
240+
}

core/src/orchestration/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
//!
2222
//! The in-box implementation is [`crate::tools::TaskExecutor`].
2323
24+
mod combinators;
2425
mod executor;
2526

27+
pub use combinators::{execute_pipeline, PipelineStage};
2628
pub use executor::{execute_steps_parallel, AgentExecutor, AgentStepSpec, StepOutcome};

core/tests/test_orchestration_real_llm.rs

Lines changed: 59 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@ use std::sync::Arc;
1212

1313
use a3s_code_core::config::CodeConfig;
1414
use a3s_code_core::llm::create_client_with_config;
15-
use a3s_code_core::orchestration::{AgentExecutor, AgentStepSpec};
15+
use a3s_code_core::orchestration::{
16+
execute_pipeline, AgentExecutor, AgentStepSpec, PipelineStage, StepOutcome,
17+
};
1618
use a3s_code_core::subagent::AgentRegistry;
1719
use a3s_code_core::tools::TaskExecutor;
1820

@@ -91,3 +93,59 @@ async fn real_execute_step_with_schema_returns_validated_object() {
9193
"object has a boolean `is_systems_language`: {object}"
9294
);
9395
}
96+
97+
/// Phase 3: a two-stage pipeline chains live agents — stage 2's prompt is
98+
/// derived from stage 1's output, with no barrier between them.
99+
#[tokio::test(flavor = "multi_thread")]
100+
#[ignore = "requires real provider credentials and network access"]
101+
async fn real_pipeline_chains_two_agent_stages() {
102+
let (executor, _workspace) = local_executor();
103+
let exec: Arc<dyn AgentExecutor> = Arc::new(executor);
104+
105+
let stages: Vec<PipelineStage<&str>> = vec![
106+
Arc::new(|_prev: Option<&StepOutcome>, topic: &&str| {
107+
Some(
108+
AgentStepSpec::new(
109+
"real-p1",
110+
"general",
111+
"summarize",
112+
format!("In one sentence, what is {topic}?"),
113+
)
114+
.with_max_steps(2),
115+
)
116+
}),
117+
Arc::new(|prev: Option<&StepOutcome>, _topic: &&str| {
118+
let summary = prev.map(|o| o.output.clone()).unwrap_or_default();
119+
Some(
120+
AgentStepSpec::new(
121+
"real-p2",
122+
"general",
123+
"classify",
124+
format!(
125+
"Reply with exactly one word, YES or NO: does this describe a \
126+
programming language?\n\nText: {summary}"
127+
),
128+
)
129+
.with_max_steps(2),
130+
)
131+
}),
132+
];
133+
134+
let out = execute_pipeline(exec, vec!["the Rust programming language"], stages, None).await;
135+
136+
assert_eq!(out.len(), 1);
137+
let final_outcome = out[0].as_ref().expect("the chain produced a final outcome");
138+
assert!(
139+
final_outcome.success,
140+
"pipeline chain succeeded: {}",
141+
final_outcome.output
142+
);
143+
assert_eq!(
144+
final_outcome.task_id, "real-p2",
145+
"the returned outcome is the last stage's"
146+
);
147+
assert!(
148+
!final_outcome.output.trim().is_empty(),
149+
"stage 2 produced output derived from stage 1"
150+
);
151+
}

0 commit comments

Comments
 (0)