1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
#[cfg(feature = "skills")]
use crate::skill_shim::load_skill_index;
use crate::skill_shim::{SelectionPolicy, SkillIndex};
use adk_core::{
AfterAgentCallback, Agent, BeforeAgentCallback, CallbackContext, Event, EventStream,
InvocationContext, Result, SharedState,
};
use async_stream::stream;
use async_trait::async_trait;
use std::sync::Arc;
use super::branch_context::{BranchContext, derive_sub_branch};
use super::shared_state_context::SharedStateContext;
/// Parallel agent executes sub-agents concurrently
pub struct ParallelAgent {
name: String,
description: String,
sub_agents: Vec<Arc<dyn Agent>>,
skills_index: Option<Arc<SkillIndex>>,
skill_policy: SelectionPolicy,
max_skill_chars: usize,
before_callbacks: Arc<Vec<BeforeAgentCallback>>,
after_callbacks: Arc<Vec<AfterAgentCallback>>,
shared_state_enabled: bool,
}
impl ParallelAgent {
/// Create a new parallel agent with the given name and sub-agents.
pub fn new(name: impl Into<String>, sub_agents: Vec<Arc<dyn Agent>>) -> Self {
Self {
name: name.into(),
description: String::new(),
sub_agents,
skills_index: None,
skill_policy: SelectionPolicy::default(),
max_skill_chars: 2000,
before_callbacks: Arc::new(Vec::new()),
after_callbacks: Arc::new(Vec::new()),
shared_state_enabled: false,
}
}
/// Set the agent description.
pub fn with_description(mut self, desc: impl Into<String>) -> Self {
self.description = desc.into();
self
}
/// Add a before-agent callback.
pub fn before_callback(mut self, callback: BeforeAgentCallback) -> Self {
if let Some(callbacks) = Arc::get_mut(&mut self.before_callbacks) {
callbacks.push(callback);
}
self
}
/// Add an after-agent callback.
pub fn after_callback(mut self, callback: AfterAgentCallback) -> Self {
if let Some(callbacks) = Arc::get_mut(&mut self.after_callbacks) {
callbacks.push(callback);
}
self
}
/// Set a preloaded skills index for this agent.
#[cfg(feature = "skills")]
pub fn with_skills(mut self, index: SkillIndex) -> Self {
self.skills_index = Some(Arc::new(index));
self
}
/// Auto-load skills from `.skills/` in the current working directory.
#[cfg(feature = "skills")]
pub fn with_auto_skills(self) -> Result<Self> {
self.with_skills_from_root(".")
}
/// Auto-load skills from `.skills/` under a custom root directory.
#[cfg(feature = "skills")]
pub fn with_skills_from_root(mut self, root: impl AsRef<std::path::Path>) -> Result<Self> {
let index = load_skill_index(root).map_err(|e| adk_core::AdkError::agent(e.to_string()))?;
self.skills_index = Some(Arc::new(index));
Ok(self)
}
/// Customize skill selection behavior.
#[cfg(feature = "skills")]
pub fn with_skill_policy(mut self, policy: SelectionPolicy) -> Self {
self.skill_policy = policy;
self
}
/// Limit injected skill content length.
#[cfg(feature = "skills")]
pub fn with_skill_budget(mut self, max_chars: usize) -> Self {
self.max_skill_chars = max_chars;
self
}
/// Enables shared state coordination for sub-agents.
///
/// When enabled, a fresh `SharedState` instance is created for each
/// `run()` invocation and injected into each sub-agent's context.
/// Sub-agents can then use `ctx.shared_state()` to access the store.
pub fn with_shared_state(mut self) -> Self {
self.shared_state_enabled = true;
self
}
}
#[async_trait]
impl Agent for ParallelAgent {
fn name(&self) -> &str {
&self.name
}
fn description(&self) -> &str {
&self.description
}
fn sub_agents(&self) -> &[Arc<dyn Agent>] {
&self.sub_agents
}
fn supports_agent_transfer(&self) -> bool {
// Deterministic workflow agent: on cross-turn resumption the runner
// must restart from this root so every sub-agent runs again, rather
// than resuming a single sub-agent that responded last.
false
}
async fn run(&self, ctx: Arc<dyn InvocationContext>) -> Result<EventStream> {
let sub_agents = self.sub_agents.clone();
let run_ctx = super::skill_context::with_skill_injected_context(
ctx,
self.skills_index.as_ref(),
&self.skill_policy,
self.max_skill_chars,
);
let before_callbacks = self.before_callbacks.clone();
let after_callbacks = self.after_callbacks.clone();
let agent_name = self.name.clone();
let invocation_id = run_ctx.invocation_id().to_string();
let shared_state_enabled = self.shared_state_enabled;
let s = stream! {
use futures::stream::{StreamExt, select_all};
for callback in before_callbacks.as_ref() {
match callback(run_ctx.clone() as Arc<dyn CallbackContext>).await {
Ok(Some(content)) => {
let mut early_event = Event::new(&invocation_id);
early_event.author = agent_name.clone();
early_event.llm_response.content = Some(content);
yield Ok(early_event);
for after_callback in after_callbacks.as_ref() {
match after_callback(run_ctx.clone() as Arc<dyn CallbackContext>).await {
Ok(Some(after_content)) => {
let mut after_event = Event::new(&invocation_id);
after_event.author = agent_name.clone();
after_event.llm_response.content = Some(after_content);
yield Ok(after_event);
return;
}
Ok(None) => continue,
Err(e) => {
yield Err(e);
return;
}
}
}
return;
}
Ok(None) => continue,
Err(e) => {
yield Err(e);
return;
}
}
}
// Create shared state if enabled (fresh per run)
let shared = if shared_state_enabled {
Some(Arc::new(SharedState::new()))
} else {
None
};
// Each sub-agent gets its own stream that resolves `run()` and drains
// the resulting events. Merging these with `select_all` polls every
// sub-agent concurrently, which is what makes this agent parallel:
// `Agent::run` only *builds* an `EventStream`, so awaiting the run
// futures together is not enough — the streams themselves have to be
// polled together. Draining one stream to completion before touching
// the next made nominally parallel branches run one at a time.
//
// Polling from a single task also gives the backpressure the ADK
// Python and Go implementations arrange explicitly (a resume signal
// and an ack channel respectively): a sub-agent cannot run ahead
// while an already-produced event is still being consumed upstream,
// so the runner's per-event persistence stays in step with execution.
//
// Dropping the merged stream drops every sub-agent stream with it, so
// a consumer that stops early tears down in-flight sub-agents instead
// of leaving them running.
let mut merged = {
// Item is (sub-agent index, event result). The index lets a failure
// be attributed to the branch that produced it.
type BranchStream =
std::pin::Pin<Box<dyn futures::Stream<Item = (usize, Result<Event>)> + Send>>;
let mut per_agent: Vec<BranchStream> = Vec::with_capacity(sub_agents.len());
for (index, agent) in sub_agents.into_iter().enumerate() {
let base: Arc<dyn InvocationContext> = if let Some(ref shared) = shared {
Arc::new(SharedStateContext::new(run_ctx.clone(), shared.clone()))
} else {
run_ctx.clone()
};
// Each sub-agent runs on its own branch, so a history read
// scoped by branch excludes what its siblings produced while
// still seeing the conversation that led to the fan-out. The
// shape mirrors ADK Python (`{parent}.{agent}.{sub_agent}`)
// and ADK Go.
let branch = derive_sub_branch(
base.branch(),
&format!("{agent_name}.{}", agent.name()),
);
let ctx: Arc<dyn InvocationContext> =
Arc::new(BranchContext::new(base, branch.clone()));
per_agent.push(Box::pin(stream! {
match agent.run(ctx).await {
Ok(mut events) => {
while let Some(event_result) = events.next().await {
let failed = event_result.is_err();
// Record which branch produced the event so a
// later branch-scoped history read can exclude
// it from siblings. A nested workflow may have
// already stamped a deeper branch; leave it.
let event_result = event_result.map(|mut event| {
if event.branch.is_empty() {
event.branch = branch.clone();
}
event
});
yield (index, event_result);
if failed {
// Abandon this branch, leave the others running.
break;
}
}
}
Err(e) => yield (index, Err(e)),
}
}));
}
select_all(per_agent)
};
// Errors are collected with their sub-agent index so the reported
// error stays deterministic. With branches running concurrently,
// "whichever failed first" would be a race; the lowest index matches
// the declared sub-agent order this agent was constructed with.
let mut failures: Vec<(usize, adk_core::AdkError)> = Vec::new();
while let Some((index, event_result)) = merged.next().await {
match event_result {
Ok(event) => yield Ok(event),
Err(e) => failures.push((index, e)),
}
}
// After all agents complete, propagate the first error if any
if let Some((_, e)) = failures.into_iter().min_by_key(|(index, _)| *index) {
yield Err(e);
return;
}
for callback in after_callbacks.as_ref() {
match callback(run_ctx.clone() as Arc<dyn CallbackContext>).await {
Ok(Some(content)) => {
let mut after_event = Event::new(&invocation_id);
after_event.author = agent_name.clone();
after_event.llm_response.content = Some(content);
yield Ok(after_event);
break;
}
Ok(None) => continue,
Err(e) => {
yield Err(e);
return;
}
}
}
};
Ok(Box::pin(s))
}
}