bamboo_subagent/
executor.rs

1//! The `ChildExecutor` seam: how an actor actually runs a task.
2//!
3//! This crate never depends on the agent runtime. The worker process implements
4//! [`ChildExecutor`] backed by the real `agent.execute()`; the transport layer drives it.
5//! [`EchoExecutor`] is a dependency-free stand-in used by the demo worker and tests.
6
7use async_trait::async_trait;
8use tokio::sync::{mpsc, oneshot};
9use tokio_util::sync::CancellationToken;
10
11use crate::proto::{
12    RunSpec, SessionMessageAdmissionConfirmation, SessionMessageDelivery, TerminalStatus,
13};
14
15/// Non-AgentEvent signals an executor sends to its transport.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub enum ExecutorControl {
18    SessionMessageAdmitted(SessionMessageAdmissionConfirmation),
19}
20
21/// Which kind of host callback a [`HostRequest`] is — selects the wire frame.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum HostRequestKind {
24    /// Proxy a gated-tool approval to the host (→ `ChildFrame::ApprovalRequest`).
25    Approval,
26}
27
28/// A single host-callback request: an executor proxies a gated-tool approval
29/// back to the host over the run's WS and awaits the reply.
30pub struct HostRequest {
31    pub kind: HostRequestKind,
32    pub body: serde_json::Value,
33    pub reply: oneshot::Sender<serde_json::Value>,
34}
35
36/// Host-callback bridge handed to an executor (via [`EventSink`]) so a nested
37/// sub-agent can proxy a gated-tool approval to the host — over the same
38/// per-child WebSocket, no broker needed. Absent for tests/[`EchoExecutor`].
39#[derive(Clone)]
40pub struct HostBridge {
41    req_tx: mpsc::UnboundedSender<HostRequest>,
42}
43
44impl HostBridge {
45    /// Create a bridge + the receiver the transport pumps to the wire.
46    pub fn channel() -> (Self, mpsc::UnboundedReceiver<HostRequest>) {
47        let (req_tx, req_rx) = mpsc::unbounded_channel();
48        (HostBridge { req_tx }, req_rx)
49    }
50
51    /// Proxy one gated-tool approval to the host and await the decision JSON
52    /// (`{"approved": bool}`). The worker's permission flow blocks on this so the
53    /// human decides on the parent (Phase 2).
54    pub async fn approval_call(
55        &self,
56        body: serde_json::Value,
57    ) -> Result<serde_json::Value, String> {
58        self.call(HostRequestKind::Approval, body).await
59    }
60
61    async fn call(
62        &self,
63        kind: HostRequestKind,
64        body: serde_json::Value,
65    ) -> Result<serde_json::Value, String> {
66        let (reply, rx) = oneshot::channel();
67        self.req_tx
68            .send(HostRequest { kind, body, reply })
69            .map_err(|_| "host bridge closed".to_string())?;
70        rx.await
71            .map_err(|_| "host bridge dropped reply".to_string())
72    }
73}
74
75/// Sink an executor emits events into; the transport forwards each as a `ChildFrame::Event`.
76#[derive(Clone)]
77pub struct EventSink {
78    tx: mpsc::UnboundedSender<serde_json::Value>,
79    control_tx: Option<mpsc::UnboundedSender<ExecutorControl>>,
80    host: Option<HostBridge>,
81}
82
83impl EventSink {
84    /// Create a sink + the receiver the transport pumps to the wire.
85    pub fn channel() -> (Self, mpsc::UnboundedReceiver<serde_json::Value>) {
86        let (tx, rx) = mpsc::unbounded_channel();
87        (
88            EventSink {
89                tx,
90                control_tx: None,
91                host: None,
92            },
93            rx,
94        )
95    }
96    /// Create a sink with a control channel for protocol-level confirmations.
97    pub fn channel_with_control() -> (
98        Self,
99        mpsc::UnboundedReceiver<serde_json::Value>,
100        mpsc::UnboundedReceiver<ExecutorControl>,
101    ) {
102        let (tx, rx) = mpsc::unbounded_channel();
103        let (control_tx, control_rx) = mpsc::unbounded_channel();
104        (
105            EventSink {
106                tx,
107                control_tx: Some(control_tx),
108                host: None,
109            },
110            rx,
111            control_rx,
112        )
113    }
114    /// Attach a host-callback bridge (the transport wires this for real runs).
115    pub fn with_host_bridge(mut self, bridge: HostBridge) -> Self {
116        self.host = Some(bridge);
117        self
118    }
119    /// The host-callback bridge, if this run was wired with one.
120    pub fn host(&self) -> Option<&HostBridge> {
121        self.host.as_ref()
122    }
123    /// Emit one event (serialized agent event). Dropped silently if the peer is gone.
124    pub fn emit(&self, event: serde_json::Value) {
125        let _ = self.tx.send(event);
126    }
127    /// Confirm a forwarded SessionInbox message only after the executor has
128    /// observed its durable local admitted receipt.
129    pub fn confirm_session_message(&self, confirmation: SessionMessageAdmissionConfirmation) {
130        if let Some(tx) = &self.control_tx {
131            let _ = tx.send(ExecutorControl::SessionMessageAdmitted(confirmation));
132        }
133    }
134}
135
136/// Result of running a task to completion (or suspension).
137#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
138pub struct ChildOutcome {
139    pub status: TerminalStatus,
140    pub result: Option<String>,
141    pub error: Option<String>,
142    /// Full worker transcript, shipped only on suspend so the host can persist
143    /// it onto the child session and rehydrate the worker on resume.
144    pub transcript: Vec<serde_json::Value>,
145}
146
147impl ChildOutcome {
148    pub fn completed(result: impl Into<String>) -> Self {
149        Self {
150            status: TerminalStatus::Completed,
151            result: Some(result.into()),
152            error: None,
153            transcript: Vec::new(),
154        }
155    }
156    pub fn error(msg: impl Into<String>) -> Self {
157        Self {
158            status: TerminalStatus::Error,
159            result: None,
160            error: Some(msg.into()),
161            transcript: Vec::new(),
162        }
163    }
164    pub fn cancelled() -> Self {
165        Self {
166            status: TerminalStatus::Cancelled,
167            result: None,
168            error: None,
169            transcript: Vec::new(),
170        }
171    }
172    /// The worker suspended to wait on its own sub-agents; ship the full
173    /// transcript so the host can resume it later.
174    pub fn suspended(transcript: Vec<serde_json::Value>) -> Self {
175        Self {
176            status: TerminalStatus::Suspended,
177            result: None,
178            error: None,
179            transcript,
180        }
181    }
182}
183
184/// Mid-run steering inbox: `ParentFrame::Message` texts arriving while a run is
185/// active. Executors that support in-band steering admit them at a safe point
186/// (the engine's round boundary); others may simply ignore the inbox.
187pub struct SteerInbox {
188    rx: mpsc::UnboundedReceiver<SteerMessage>,
189}
190
191#[derive(Debug, Clone, PartialEq)]
192pub enum SteerMessage {
193    /// Direct legacy WS text has no durable ingress identity and remains
194    /// explicitly best-effort.
195    Text(String),
196    /// Broker Maildir ingress retains its durable message id so at-least-once
197    /// replay converges on one typed SessionInbox envelope.
198    DurableText {
199        message_id: String,
200        text: String,
201    },
202    SessionMessage(Box<SessionMessageDelivery>),
203}
204
205impl From<String> for SteerMessage {
206    fn from(value: String) -> Self {
207        Self::Text(value)
208    }
209}
210
211impl From<&str> for SteerMessage {
212    fn from(value: &str) -> Self {
213        Self::Text(value.to_string())
214    }
215}
216
217impl SteerInbox {
218    /// Create a sender + inbox pair (the transport holds the sender).
219    pub fn channel() -> (mpsc::UnboundedSender<SteerMessage>, Self) {
220        let (tx, rx) = mpsc::unbounded_channel();
221        (tx, SteerInbox { rx })
222    }
223    /// An already-closed inbox (for tests / executors that don't steer).
224    pub fn disconnected() -> Self {
225        let (_tx, rx) = mpsc::unbounded_channel();
226        SteerInbox { rx }
227    }
228    /// Next steering message, or `None` once the run's sender is gone.
229    pub async fn recv(&mut self) -> Option<String> {
230        match self.rx.recv().await? {
231            SteerMessage::Text(text) => Some(text),
232            SteerMessage::DurableText { text, .. } => Some(text),
233            SteerMessage::SessionMessage(delivery) => serde_json::to_string(&delivery).ok(),
234        }
235    }
236    /// Receive the typed steering value. Runtime-backed workers use this path
237    /// so SessionInbox envelopes never collapse into ad-hoc text.
238    pub async fn recv_message(&mut self) -> Option<SteerMessage> {
239        self.rx.recv().await
240    }
241}
242
243/// What runs inside an actor. Implemented by the worker with the real runtime.
244#[async_trait]
245pub trait ChildExecutor: Send + Sync + 'static {
246    async fn run(
247        &self,
248        spec: RunSpec,
249        events: EventSink,
250        steer: SteerInbox,
251        cancel: CancellationToken,
252    ) -> ChildOutcome;
253}
254
255/// Dependency-free executor: streams one `token` event per word, then completes with an echo.
256/// Used by the demo worker and tests to exercise the full transport without a real LLM.
257///
258/// Test hook: an assignment starting with `__sleep_ms:<n>` sleeps (cancellably)
259/// for `n` milliseconds before echoing the rest — this is what lets cancel /
260/// concurrency e2e tests hold a run open deterministically without an LLM.
261pub struct EchoExecutor;
262
263/// Assignment prefix recognized by [`EchoExecutor`] for a cancellable delay.
264pub const ECHO_SLEEP_PREFIX: &str = "__sleep_ms:";
265
266#[async_trait]
267impl ChildExecutor for EchoExecutor {
268    async fn run(
269        &self,
270        spec: RunSpec,
271        events: EventSink,
272        _steer: SteerInbox,
273        cancel: CancellationToken,
274    ) -> ChildOutcome {
275        // Optional cancellable delay: any token `__sleep_ms:<n>` in the
276        // assignment (scanned, not just the prefix — child creation may wrap
277        // the prompt in a template). The marker token itself is not echoed.
278        let mut sleep_ms: Option<u64> = None;
279        let mut words: Vec<&str> = Vec::new();
280        for word in spec.assignment.split_whitespace() {
281            match word
282                .strip_prefix(ECHO_SLEEP_PREFIX)
283                .and_then(|n| n.parse::<u64>().ok())
284            {
285                Some(ms) if sleep_ms.is_none() => sleep_ms = Some(ms),
286                _ => words.push(word),
287            }
288        }
289        if let Some(ms) = sleep_ms {
290            tokio::select! {
291                _ = tokio::time::sleep(std::time::Duration::from_millis(ms)) => {}
292                _ = cancel.cancelled() => return ChildOutcome::cancelled(),
293            }
294        }
295
296        for word in &words {
297            if cancel.is_cancelled() {
298                return ChildOutcome::cancelled();
299            }
300            events.emit(serde_json::json!({ "type": "token", "content": format!("{word} ") }));
301            // tiny yield so cancellation can interleave; not a real delay
302            tokio::task::yield_now().await;
303        }
304        events.emit(serde_json::json!({ "type": "complete" }));
305        ChildOutcome::completed(format!("echo: {}", words.join(" ")))
306    }
307}
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312
313    #[tokio::test]
314    async fn echo_streams_then_completes() {
315        let (sink, mut rx) = EventSink::channel();
316        let outcome = EchoExecutor
317            .run(
318                RunSpec {
319                    assignment: "alpha beta".into(),
320                    logical_session: None,
321                    project_id: None,
322                    reasoning_effort: None,
323                    permission_policy: None,
324                    messages: Vec::new(),
325                    activation_run_id: None,
326                    initial_session_messages: Vec::new(),
327                    secrets: Default::default(),
328                },
329                sink,
330                SteerInbox::disconnected(),
331                CancellationToken::new(),
332            )
333            .await;
334        assert_eq!(outcome.status, TerminalStatus::Completed);
335        assert_eq!(outcome.result.as_deref(), Some("echo: alpha beta"));
336
337        let mut events = Vec::new();
338        while let Ok(e) = rx.try_recv() {
339            events.push(e);
340        }
341        // two tokens + one complete
342        assert_eq!(events.len(), 3);
343        assert_eq!(events[0]["content"], "alpha ");
344    }
345
346    #[tokio::test]
347    async fn echo_honors_cancel() {
348        let (sink, _rx) = EventSink::channel();
349        let cancel = CancellationToken::new();
350        cancel.cancel();
351        let outcome = EchoExecutor
352            .run(
353                RunSpec {
354                    assignment: "a b c".into(),
355                    logical_session: None,
356                    project_id: None,
357                    reasoning_effort: None,
358                    permission_policy: None,
359                    messages: Vec::new(),
360                    activation_run_id: None,
361                    initial_session_messages: Vec::new(),
362                    secrets: Default::default(),
363                },
364                sink,
365                SteerInbox::disconnected(),
366                cancel,
367            )
368            .await;
369        assert_eq!(outcome.status, TerminalStatus::Cancelled);
370    }
371
372    #[tokio::test]
373    async fn approval_call_sends_approval_kind_and_round_trips_reply() {
374        let (bridge, mut req_rx) = HostBridge::channel();
375        let caller = tokio::spawn(async move {
376            bridge
377                .approval_call(serde_json::json!({"resource": "/tmp/x"}))
378                .await
379        });
380        let req = req_rx.recv().await.expect("a host request");
381        assert_eq!(req.kind, HostRequestKind::Approval);
382        assert_eq!(req.body["resource"], "/tmp/x");
383        let _ = req.reply.send(serde_json::json!({"approved": true}));
384        let reply = caller.await.unwrap().expect("decision");
385        assert_eq!(reply["approved"], true);
386    }
387}
bamboo_subagent/executor.rs

bamboo_subagent/
executor.rs
