use async_trait::async_trait;
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use tokio::sync::{mpsc, oneshot, Semaphore};
use tokio_util::sync::CancellationToken;
use crate::proto::{
RunSpec, SessionMessageAdmissionConfirmation, SessionMessageDelivery, TerminalStatus,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExecutorControl {
SessionMessageAdmitted(SessionMessageAdmissionConfirmation),
InitialInputReleaseRequest(crate::proto::InitialInputReleaseRequest),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HostRequestKind {
Approval,
OwnedTree,
SubAgent,
}
pub struct HostRequest {
pub kind: HostRequestKind,
pub body: serde_json::Value,
pub reply: oneshot::Sender<serde_json::Value>,
}
#[derive(Clone)]
pub struct HostBridge {
req_tx: mpsc::UnboundedSender<HostRequest>,
tree_permits: Arc<Semaphore>,
tree_calls: Arc<AtomicUsize>,
subagent_permits: Arc<Semaphore>,
subagent_calls: Arc<AtomicUsize>,
}
impl HostBridge {
pub fn channel() -> (Self, mpsc::UnboundedReceiver<HostRequest>) {
let (req_tx, req_rx) = mpsc::unbounded_channel();
(
HostBridge {
req_tx,
tree_permits: Arc::new(Semaphore::new(1)),
tree_calls: Arc::new(AtomicUsize::new(0)),
subagent_permits: Arc::new(Semaphore::new(1)),
subagent_calls: Arc::new(AtomicUsize::new(0)),
},
req_rx,
)
}
pub async fn approval_call(
&self,
body: serde_json::Value,
) -> Result<serde_json::Value, String> {
self.call(HostRequestKind::Approval, body).await
}
pub async fn owned_tree_call(&self, cursor: Option<&str>) -> Result<serde_json::Value, String> {
const MAX_CURSOR_BYTES: usize = 128;
const MAX_RESULT_BYTES: usize = 8192;
const MAX_CALLS_PER_RUN: usize = 16;
const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
if cursor.is_some_and(|cursor| cursor.is_empty() || cursor.len() > MAX_CURSOR_BYTES) {
return Err("invalid owned tree cursor".into());
}
let _permit = tokio::time::timeout(TIMEOUT, self.tree_permits.acquire())
.await
.map_err(|_| "owned tree request timed out")?
.map_err(|_| "owned tree inspection unavailable")?;
if self.tree_calls.fetch_add(1, Ordering::AcqRel) >= MAX_CALLS_PER_RUN {
return Err("owned tree request limit reached".into());
}
let reply = tokio::time::timeout(
TIMEOUT,
self.call(
HostRequestKind::OwnedTree,
serde_json::json!({"cursor":cursor}),
),
)
.await
.map_err(|_| "owned tree request timed out")??;
let page = reply
.get("page")
.filter(|page| !page.is_null())
.ok_or_else(|| "owned tree inspection rejected".to_string())?;
if serde_json::to_vec(page).map_or(true, |bytes| bytes.len() > MAX_RESULT_BYTES) {
return Err("owned tree observation exceeds its limit".into());
}
Ok(page.clone())
}
pub async fn subagent_call(
&self,
args: serde_json::Value,
tool_call_id: &str,
) -> Result<serde_json::Value, String> {
const MAX_ARGS_BYTES: usize = 64 * 1024;
const MAX_REPLY_BYTES: usize = 16 * 1024;
const MAX_CALLS_PER_RUN: usize = 64;
const TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
if tool_call_id.is_empty()
|| tool_call_id.len() > 128
|| serde_json::to_vec(&args).map_or(true, |bytes| bytes.len() > MAX_ARGS_BYTES)
{
return Err("invalid canonical SubAgent request".into());
}
let _permit = tokio::time::timeout(TIMEOUT, self.subagent_permits.acquire())
.await
.map_err(|_| "canonical SubAgent request timed out")?
.map_err(|_| "canonical SubAgent unavailable")?;
if self.subagent_calls.fetch_add(1, Ordering::AcqRel) >= MAX_CALLS_PER_RUN {
return Err("canonical SubAgent request limit reached".into());
}
let reply = tokio::time::timeout(
TIMEOUT,
self.call(
HostRequestKind::SubAgent,
serde_json::json!({"args":args,"tool_call_id":tool_call_id}),
),
)
.await
.map_err(|_| "canonical SubAgent request timed out")??;
if serde_json::to_vec(&reply).map_or(true, |bytes| bytes.len() > MAX_REPLY_BYTES) {
return Err("canonical SubAgent result exceeds its limit".into());
}
if let Some(error) = reply.get("error").and_then(serde_json::Value::as_str) {
return Err(error.to_owned());
}
reply
.get("result")
.cloned()
.ok_or_else(|| "canonical SubAgent operation rejected".into())
}
async fn call(
&self,
kind: HostRequestKind,
body: serde_json::Value,
) -> Result<serde_json::Value, String> {
let (reply, rx) = oneshot::channel();
self.req_tx
.send(HostRequest { kind, body, reply })
.map_err(|_| "host bridge closed".to_string())?;
rx.await
.map_err(|_| "host bridge dropped reply".to_string())
}
}
#[derive(Clone)]
pub struct EventSink {
tx: mpsc::Sender<serde_json::Value>,
control_tx: Option<mpsc::Sender<ExecutorControl>>,
host: Option<HostBridge>,
}
impl EventSink {
pub const EVENT_CAPACITY: usize = 256;
pub const CONTROL_CAPACITY: usize = 32;
pub fn channel() -> (Self, mpsc::Receiver<serde_json::Value>) {
let (tx, rx) = mpsc::channel(Self::EVENT_CAPACITY);
(
EventSink {
tx,
control_tx: None,
host: None,
},
rx,
)
}
pub fn channel_with_control() -> (
Self,
mpsc::Receiver<serde_json::Value>,
mpsc::Receiver<ExecutorControl>,
) {
let (tx, rx) = mpsc::channel(Self::EVENT_CAPACITY);
let (control_tx, control_rx) = mpsc::channel(Self::CONTROL_CAPACITY);
(
EventSink {
tx,
control_tx: Some(control_tx),
host: None,
},
rx,
control_rx,
)
}
pub fn with_host_bridge(mut self, bridge: HostBridge) -> Self {
self.host = Some(bridge);
self
}
pub fn host(&self) -> Option<&HostBridge> {
self.host.as_ref()
}
pub async fn emit(&self, event: serde_json::Value) {
let _ = self.tx.send(event).await;
}
pub async fn request_initial_release(
&self,
request: crate::proto::InitialInputReleaseRequest,
) -> Result<(), String> {
self.control_tx
.as_ref()
.ok_or("initial release channel missing")?
.send(ExecutorControl::InitialInputReleaseRequest(request))
.await
.map_err(|_| "initial release channel closed".into())
}
pub async fn confirm_session_message(&self, confirmation: SessionMessageAdmissionConfirmation) {
if let Some(tx) = &self.control_tx {
let _ = tx
.send(ExecutorControl::SessionMessageAdmitted(confirmation))
.await;
}
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ChildOutcome {
pub status: TerminalStatus,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub final_event_watermark: Option<crate::proto::ActorEventWatermark>,
pub result: Option<String>,
pub error: Option<String>,
#[serde(default)]
pub transcript: Vec<serde_json::Value>,
}
impl ChildOutcome {
pub fn completed(result: impl Into<String>) -> Self {
Self {
status: TerminalStatus::Completed,
final_event_watermark: None,
result: Some(result.into()),
error: None,
transcript: Vec::new(),
}
}
pub fn error(msg: impl Into<String>) -> Self {
Self {
status: TerminalStatus::Error,
final_event_watermark: None,
result: None,
error: Some(msg.into()),
transcript: Vec::new(),
}
}
pub fn cancelled() -> Self {
Self {
status: TerminalStatus::Cancelled,
final_event_watermark: None,
result: None,
error: None,
transcript: Vec::new(),
}
}
pub fn suspended(transcript: Vec<serde_json::Value>) -> Self {
Self {
status: TerminalStatus::Suspended,
final_event_watermark: None,
result: None,
error: None,
transcript,
}
}
}
pub struct SteerInbox {
rx: mpsc::UnboundedReceiver<SteerMessage>,
accepted_initial_releases: Vec<crate::proto::InitialInputRelease>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum SteerMessage {
Text(String),
DurableText {
message_id: String,
text: String,
},
SessionMessage(Box<SessionMessageDelivery>),
InitialInputRelease(crate::proto::InitialInputRelease),
}
impl From<String> for SteerMessage {
fn from(value: String) -> Self {
Self::Text(value)
}
}
impl From<&str> for SteerMessage {
fn from(value: &str) -> Self {
Self::Text(value.to_string())
}
}
impl SteerInbox {
pub fn channel() -> (mpsc::UnboundedSender<SteerMessage>, Self) {
let (tx, rx) = mpsc::unbounded_channel();
(
tx,
SteerInbox {
rx,
accepted_initial_releases: Vec::new(),
},
)
}
pub fn disconnected() -> Self {
let (_tx, rx) = mpsc::unbounded_channel();
SteerInbox {
rx,
accepted_initial_releases: Vec::new(),
}
}
pub async fn recv(&mut self) -> Option<String> {
match self.recv_message().await? {
SteerMessage::Text(text) => Some(text),
SteerMessage::DurableText { text, .. } => Some(text),
SteerMessage::SessionMessage(delivery) => serde_json::to_string(&delivery).ok(),
SteerMessage::InitialInputRelease(_) => None, }
}
pub async fn wait_initial_release(
&mut self,
request: &crate::proto::InitialInputReleaseRequest,
cancel: &CancellationToken,
events: &EventSink,
) -> Result<chrono::DateTime<chrono::Utc>, String> {
self.wait_initial_release_bounded(
request,
cancel,
events,
std::time::Duration::from_secs(60),
)
.await
}
async fn wait_initial_release_bounded(
&mut self,
request: &crate::proto::InitialInputReleaseRequest,
cancel: &CancellationToken,
events: &EventSink,
wait: std::time::Duration,
) -> Result<chrono::DateTime<chrono::Utc>, String> {
let timeout = tokio::time::sleep(wait);
tokio::pin!(timeout);
let mut retry = tokio::time::interval(std::time::Duration::from_secs(5));
retry.tick().await;
loop {
tokio::select! {
biased;
_ = cancel.cancelled() => return Err("initial release cancelled".into()),
_ = &mut timeout => return Err("initial release timed out".into()),
message = self.recv_message() => return match message {
Some(SteerMessage::InitialInputRelease(release)) if release.permits(request, chrono::Utc::now()) => {
if self.accepted_initial_releases.iter().any(|accepted| accepted.request == release.request) {
return Err("initial release missing, expired or mismatched".into());
}
if self.accepted_initial_releases.len() >= bamboo_domain::SessionInboxLimits::default().max_claim_batch {
return Err("initial release completed batch limit exceeded".into());
}
let deadline = release.expires_at;
self.accepted_initial_releases.push(release);
Ok(deadline)
},
_ => Err("initial release missing, expired or mismatched".into()),
},
_ = retry.tick() => events.request_initial_release(request.clone()).await?,
}
}
}
pub async fn recv_message(&mut self) -> Option<SteerMessage> {
loop {
let message = self.rx.recv().await?;
if let SteerMessage::InitialInputRelease(release) = &message {
if chrono::Utc::now() < release.expires_at
&& self.accepted_initial_releases.contains(release)
{
continue;
}
}
return Some(message);
}
}
}
#[async_trait]
pub trait ChildExecutor: Send + Sync + 'static {
fn requires_contiguous_events(&self) -> bool {
false
}
fn supports_environment_lease_v1(&self) -> bool {
false
}
fn max_parallel_executions(&self) -> usize {
1
}
async fn run(
&self,
spec: RunSpec,
events: EventSink,
steer: SteerInbox,
cancel: CancellationToken,
) -> ChildOutcome;
}
pub struct EchoExecutor;
pub const ECHO_SLEEP_PREFIX: &str = "__sleep_ms:";
#[async_trait]
impl ChildExecutor for EchoExecutor {
fn max_parallel_executions(&self) -> usize {
256
}
async fn run(
&self,
spec: RunSpec,
events: EventSink,
_steer: SteerInbox,
cancel: CancellationToken,
) -> ChildOutcome {
let mut sleep_ms: Option<u64> = None;
let mut words: Vec<&str> = Vec::new();
for word in spec.assignment.split_whitespace() {
match word
.strip_prefix(ECHO_SLEEP_PREFIX)
.and_then(|n| n.parse::<u64>().ok())
{
Some(ms) if sleep_ms.is_none() => sleep_ms = Some(ms),
_ => words.push(word),
}
}
if let Some(ms) = sleep_ms {
tokio::select! {
_ = tokio::time::sleep(std::time::Duration::from_millis(ms)) => {}
_ = cancel.cancelled() => return ChildOutcome::cancelled(),
}
}
for word in &words {
if cancel.is_cancelled() {
return ChildOutcome::cancelled();
}
events
.emit(serde_json::json!({ "type": "token", "content": format!("{word} ") }))
.await;
tokio::task::yield_now().await;
}
events.emit(serde_json::json!({ "type": "complete" })).await;
ChildOutcome::completed(format!("echo: {}", words.join(" ")))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn duplicate_release_request(slot: u64) -> crate::proto::InitialInputReleaseRequest {
crate::proto::InitialInputReleaseRequest {
version: 1,
nonce: uuid::Uuid::new_v4().to_string(),
child_id: "duplicate-child".into(),
parent_id: "parent".into(),
root_id: "parent".into(),
created_at: chrono::Utc::now(),
spawn_depth: 1,
project_id: None,
envelope_id: format!("input-{slot}"),
generation: slot + 1,
activation_run_id: "same-physical-run".into(),
execution_epoch: 1,
}
}
#[tokio::test(start_paused = true)]
async fn initial_release_duplicate_retry_between_batch_items_is_harmless() {
let first = duplicate_release_request(0);
let second = duplicate_release_request(1);
let first_release = crate::proto::InitialInputRelease {
request: first.clone(),
expires_at: chrono::Utc::now() + chrono::Duration::seconds(30),
};
let second_release = crate::proto::InitialInputRelease {
request: second.clone(),
expires_at: first_release.expires_at - chrono::Duration::seconds(1),
};
let (events, _events, mut controls) = EventSink::channel_with_control();
let (sender, mut inbox) = SteerInbox::channel();
let cancel = CancellationToken::new();
let mut waiting = Box::pin(inbox.wait_initial_release(&first, &cancel, &events));
assert!(futures_util::poll!(waiting.as_mut()).is_pending());
tokio::time::advance(std::time::Duration::from_secs(5)).await;
assert!(futures_util::poll!(waiting.as_mut()).is_pending());
assert_eq!(
controls.try_recv().unwrap(),
ExecutorControl::InitialInputReleaseRequest(first.clone())
);
sender
.send(SteerMessage::InitialInputRelease(first_release.clone()))
.unwrap();
assert_eq!(waiting.await.unwrap(), first_release.expires_at);
sender
.send(SteerMessage::InitialInputRelease(first_release.clone()))
.unwrap();
sender
.send(SteerMessage::InitialInputRelease(second_release.clone()))
.unwrap();
assert_eq!(
inbox
.wait_initial_release(&second, &cancel, &events)
.await
.unwrap(),
second_release.expires_at
);
sender
.send(SteerMessage::InitialInputRelease(first_release))
.unwrap();
sender
.send(SteerMessage::InitialInputRelease(second_release))
.unwrap();
sender
.send(SteerMessage::Text("next ordinary steer".into()))
.unwrap();
assert_eq!(
inbox.recv_message().await,
Some(SteerMessage::Text("next ordinary steer".into()))
);
}
#[tokio::test]
async fn initial_release_changed_prior_reply_cannot_be_filtered_or_release_next_input() {
for changed in [
"nonce",
"epoch",
"birth",
"envelope",
"generation",
"deadline",
"expired",
] {
let first = duplicate_release_request(0);
let second = duplicate_release_request(1);
let original = crate::proto::InitialInputRelease {
request: first.clone(),
expires_at: chrono::Utc::now() + chrono::Duration::seconds(30),
};
let (events, _events, _controls) = EventSink::channel_with_control();
let (sender, mut inbox) = SteerInbox::channel();
let cancel = CancellationToken::new();
sender
.send(SteerMessage::InitialInputRelease(original.clone()))
.unwrap();
inbox
.wait_initial_release(&first, &cancel, &events)
.await
.unwrap();
let mut altered = original;
match changed {
"nonce" => altered.request.nonce = uuid::Uuid::new_v4().to_string(),
"epoch" => altered.request.execution_epoch += 1,
"birth" => altered.request.created_at += chrono::Duration::seconds(1),
"envelope" => altered.request.envelope_id = "other-input".into(),
"generation" => altered.request.generation += 1,
"deadline" => altered.expires_at += chrono::Duration::seconds(1),
"expired" => altered.expires_at = chrono::Utc::now() - chrono::Duration::seconds(1),
_ => unreachable!(),
}
sender
.send(SteerMessage::InitialInputRelease(altered))
.unwrap();
let error = inbox
.wait_initial_release(&second, &cancel, &events)
.await
.unwrap_err();
assert!(error.contains("mismatched"), "{changed}: {error}");
}
}
#[tokio::test(start_paused = true)]
async fn initial_release_wait_observes_timeout_disconnect_and_cancel() {
let request = crate::proto::InitialInputReleaseRequest {
version: 1,
nonce: uuid::Uuid::new_v4().to_string(),
child_id: "wait-child".into(),
parent_id: "parent".into(),
root_id: "parent".into(),
created_at: chrono::Utc::now(),
spawn_depth: 1,
project_id: None,
envelope_id: "input".into(),
generation: 1,
activation_run_id: "run".into(),
execution_epoch: 1,
};
let (events, _rx, _controls) = EventSink::channel_with_control();
for case in ["timeout", "disconnect", "cancel"] {
let (sender, mut inbox) = SteerInbox::channel();
let cancel = CancellationToken::new();
if case == "disconnect" {
drop(sender);
} else if case == "cancel" {
cancel.cancel();
}
let error = inbox
.wait_initial_release_bounded(
&request,
&cancel,
&events,
std::time::Duration::from_millis(2),
)
.await
.unwrap_err();
assert!(
error.contains(match case {
"timeout" => "timed out",
"cancel" => "cancelled",
_ => "missing",
}),
"initial release case={case}, actual error={error:?}"
);
}
}
#[tokio::test]
async fn full_event_queue_backpressures_without_losing_durable_order_or_blocking_control() {
let (sink, mut events, mut controls) = EventSink::channel_with_control();
for index in 0..EventSink::EVENT_CAPACITY {
sink.emit(serde_json::json!({"type":"tool_start","index":index}))
.await;
}
assert_eq!(events.len(), EventSink::EVENT_CAPACITY);
let terminal = sink.emit(serde_json::json!({"type":"complete"}));
tokio::pin!(terminal);
assert!(futures_util::poll!(&mut terminal).is_pending());
let confirmation = SessionMessageAdmissionConfirmation {
target_session_id: "bounded-child".into(),
envelope_id: "message".into(),
canonical_claim_generation: 1,
activation_run_id: "activation".into(),
};
sink.confirm_session_message(confirmation.clone()).await;
assert_eq!(
controls.recv().await.unwrap(),
ExecutorControl::SessionMessageAdmitted(confirmation)
);
assert_eq!(events.recv().await.unwrap()["index"], 0);
terminal.await;
for index in 1..EventSink::EVENT_CAPACITY {
assert_eq!(events.recv().await.unwrap()["index"], index);
}
assert_eq!(events.recv().await.unwrap()["type"], "complete");
}
#[test]
fn echo_explicitly_opts_into_high_parallelism() {
assert!(EchoExecutor.max_parallel_executions() >= 200);
}
#[tokio::test]
async fn echo_streams_then_completes() {
let (sink, mut rx) = EventSink::channel();
let outcome = EchoExecutor
.run(
RunSpec {
assignment: "alpha beta".into(),
logical_session: None,
project_id: None,
reasoning_effort: None,
permission_policy: None,
messages: Vec::new(),
activation_run_id: None,
execution_epoch: 0,
initial_session_messages: Vec::new(),
secrets: Default::default(),
},
sink,
SteerInbox::disconnected(),
CancellationToken::new(),
)
.await;
assert_eq!(outcome.status, TerminalStatus::Completed);
assert_eq!(outcome.result.as_deref(), Some("echo: alpha beta"));
let mut events = Vec::new();
while let Ok(e) = rx.try_recv() {
events.push(e);
}
assert_eq!(events.len(), 3);
assert_eq!(events[0]["content"], "alpha ");
}
#[tokio::test]
async fn echo_honors_cancel() {
let (sink, _rx) = EventSink::channel();
let cancel = CancellationToken::new();
cancel.cancel();
let outcome = EchoExecutor
.run(
RunSpec {
assignment: "a b c".into(),
logical_session: None,
project_id: None,
reasoning_effort: None,
permission_policy: None,
messages: Vec::new(),
activation_run_id: None,
execution_epoch: 0,
initial_session_messages: Vec::new(),
secrets: Default::default(),
},
sink,
SteerInbox::disconnected(),
cancel,
)
.await;
assert_eq!(outcome.status, TerminalStatus::Cancelled);
}
#[tokio::test]
async fn approval_call_sends_approval_kind_and_round_trips_reply() {
let (bridge, mut req_rx) = HostBridge::channel();
let caller = tokio::spawn(async move {
bridge
.approval_call(serde_json::json!({"resource": "/tmp/x"}))
.await
});
let req = req_rx.recv().await.expect("a host request");
assert_eq!(req.kind, HostRequestKind::Approval);
assert_eq!(req.body["resource"], "/tmp/x");
let _ = req.reply.send(serde_json::json!({"approved": true}));
let reply = caller.await.unwrap().expect("decision");
assert_eq!(reply["approved"], true);
}
}