microvms_core/session/exec.rs
1// SPDX-License-Identifier: Apache-2.0
2//! The handle a caller holds for one running exec.
3//!
4//! The whole reason this module exists is [`ExecHandle::stream`]. Everything else here
5//! is a thin wrapper over one route.
6//!
7//! # The cursor is the contract
8//!
9//! A reconnect passes `?offset=<bytes consumed>`, and the daemon replays from there.
10//! Without it a dropped connection means either losing everything after the drop or
11//! replaying from zero, and neither is distinguishable from correct behaviour with
12//! nothing to check against — E2B's cursorless `connect(pid)` and their issue #1352 are
13//! exactly that.
14//!
15//! Two rules the cursor arithmetic has to obey, and both are properties of *what the
16//! caller has seen* rather than of what arrived:
17//!
18//! * Advance only past bytes actually handed to the caller, so a reconnect never
19//! re-delivers and never skips.
20//! * Advance past a `gap` too. The daemon has already moved its own cursor past evicted
21//! bytes; if ours did not follow, a reconnect would ask for them again and be told
22//! about the same gap forever.
23//!
24//! # The reconnect condition is a typed event, not a byte count
25//!
26//! A stream that ended *without* an `exit` event was cut. A stream that ended *with* one
27//! is over. Those two are the identical byte sequence on a raw stream, which is the
28//! whole argument for SSE framing here and why this is not a byte stream. The
29//! reconnect decision reads that distinction and nothing else.
30
31use std::sync::Arc;
32use std::time::Duration;
33
34use base64::Engine as _;
35use futures_util::Stream;
36
37use super::sse::{ExecEvent, SseParser, decode};
38use super::{HttpRequest, Transport};
39use crate::error::{Error, ErrorKind, WireKind};
40
41/// Backoff between reconnect attempts. Short, because the offset makes a reconnect
42/// cheap — the daemon replays from the cursor, so nothing already delivered is refetched.
43const RECONNECT_BACKOFF: [Duration; 5] = [
44 Duration::from_millis(250),
45 Duration::from_millis(500),
46 Duration::from_secs(1),
47 Duration::from_secs(2),
48 Duration::from_secs(4),
49];
50
51/// How long an attached stream may be silent before it is treated as dead.
52///
53/// Four times the daemon's fifteen-second SSE keepalive, so three missed keepalives are
54/// tolerated before a reconnect. Tighter than that turns a slow proxy into a reconnect
55/// loop.
56pub const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
57
58/// How often to re-poll while waiting for an exec to finish.
59const POLL_INTERVAL: Duration = Duration::from_secs(1);
60
61/// An exec's phase and, once it has one, its outcome.
62///
63/// A thin wrapper over the daemon's [`protocol::exec::PollResponse`] rather than a
64/// re-modelling of it, so the two cannot disagree. What it adds is [`ExecResult::done`],
65/// which is the question every caller actually asks.
66#[derive(Debug)]
67pub struct ExecResult {
68 pub exec_id: String,
69 pub phase: protocol::exec::Phase,
70 /// `None` while running. Present once the child has exited.
71 pub outcome: Option<protocol::exec::Outcome>,
72}
73
74impl From<protocol::exec::PollResponse> for ExecResult {
75 fn from(response: protocol::exec::PollResponse) -> Self {
76 Self {
77 exec_id: response.exec_id,
78 phase: response.phase,
79 outcome: response.result,
80 }
81 }
82}
83
84impl ExecResult {
85 /// Whether the exec has finished, whichever way.
86 pub fn done(&self) -> bool {
87 matches!(
88 self.phase,
89 protocol::exec::Phase::Exited | protocol::exec::Phase::Acked
90 )
91 }
92
93 /// The exit code, when the child exited with one rather than dying to a signal.
94 pub fn exit_code(&self) -> Option<i32> {
95 self.outcome.as_ref().and_then(|outcome| outcome.exit_code)
96 }
97
98 /// Whether the command succeeded. `false` for a signal death and for a still-running
99 /// exec, since neither is a success.
100 pub fn succeeded(&self) -> bool {
101 self.outcome
102 .as_ref()
103 .is_some_and(|outcome| outcome.exit_code == Some(0) && !outcome.timed_out)
104 }
105
106 pub fn stdout(&self) -> &str {
107 self.outcome
108 .as_ref()
109 .map_or("", |outcome| outcome.stdout.as_str())
110 }
111
112 pub fn stderr(&self) -> &str {
113 self.outcome
114 .as_ref()
115 .map_or("", |outcome| outcome.stderr.as_str())
116 }
117}
118
119/// What a stdin write accomplished.
120pub type StdinAck = protocol::exec::StdinResponse;
121
122/// Why a stream stopped, and where a resume would pick up.
123///
124/// Returned by [`ExecHandle::for_each_event`] and
125/// [`ExecHandle::for_each_event_async`] so the three endings are distinguishable
126/// *without* a caller having to reason about which event it last saw. That distinction is
127/// the whole point of this module — a body that ended without an `exit` event was cut, and
128/// the byte sequence alone cannot say so — and asking every caller to re-derive it from a
129/// tally is how one of them gets it wrong and reports a truncated stream as a clean pass.
130#[derive(Clone, Copy, Debug, Eq, PartialEq)]
131pub struct StreamEnd {
132 pub reason: EndReason,
133 /// The offset a resume would pass as [`StreamOptions::offset`].
134 ///
135 /// **Core's cursor, not a count of what the callback saw.** The two rules in this
136 /// module's docs (advance only past delivered bytes, and advance past a `gap`) are
137 /// obeyed by the state machine below; handing the number back means a caller resuming
138 /// does not maintain a second cursor that disagrees with this one exactly when a
139 /// reconnect happened.
140 pub cursor: u64,
141}
142
143/// The three ways a stream ends without an error.
144#[derive(Clone, Copy, Debug, Eq, PartialEq)]
145pub enum EndReason {
146 /// The terminal `exit` event arrived and was delivered. The command is over.
147 Exited,
148 /// The callback answered [`ControlFlow::Break`]. Nothing is wrong; the caller stopped
149 /// reading, and [`StreamEnd::cursor`] is where it stopped.
150 Stopped,
151 /// The body ended with **no** `exit` event and `reconnect` was off, so the connection
152 /// was cut and this client was told not to re-attach. The command's outcome is
153 /// unknown — not zero — and a caller reporting success here would pass a CI step on
154 /// evidence it never received.
155 Cut,
156}
157
158/// How a stream should behave. See [`ExecHandle::stream_with`].
159#[derive(Clone, Debug)]
160pub struct StreamOptions {
161 /// The byte to start at. Non-zero resumes a stream a previous process was reading.
162 pub offset: u64,
163 /// Whether to reconnect after a cut. `false` ends the stream at the cut instead,
164 /// which is what a caller doing its own reconnection wants.
165 pub reconnect: bool,
166 /// How many reconnects before giving up. A bound rather than forever, because a
167 /// stream that drops every time is a condition a caller needs reported.
168 pub max_reconnects: u32,
169 /// Turns a `gap` into an error instead of an event. What a caller that must have
170 /// complete output wants.
171 pub error_on_gap: bool,
172 /// How long the body may be silent before the connection is treated as dead.
173 pub idle_timeout: Duration,
174}
175
176impl Default for StreamOptions {
177 fn default() -> Self {
178 Self {
179 offset: 0,
180 reconnect: true,
181 max_reconnects: 20,
182 error_on_gap: false,
183 idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT,
184 }
185 }
186}
187
188/// A fresh exec id: `x-` plus 16 hex characters, the Python client's shape.
189///
190/// Here rather than in each binding, because both bindings carried a byte-identical
191/// copy of this function — and the id's shape is part of the session story ("the exec
192/// id is caller-minted"), so the crate that tells that story is the crate that mints.
193///
194/// The id needs to be distinct rather than unguessable — it is an idempotency key,
195/// not a credential, and the daemon rejects an unknown one — so the nanosecond clock
196/// mixed with a counter is enough. (`uuid` v7 would also be enough and is on the
197/// dependency-sweep list; the shape below is the Python client's, kept for wire
198/// familiarity until that sweep decides.)
199pub fn mint_exec_id() -> String {
200 use std::sync::atomic::{AtomicU64, Ordering};
201 static SEQUENCE: AtomicU64 = AtomicU64::new(0);
202 let nanos = std::time::SystemTime::now()
203 .duration_since(std::time::UNIX_EPOCH)
204 .map(|since| since.as_nanos() as u64)
205 .unwrap_or_default();
206 let sequence = SEQUENCE.fetch_add(1, Ordering::Relaxed);
207 // The counter folds into the clock's high bits so two ids minted in the same
208 // nanosecond still differ.
209 format!("x-{:016x}", nanos ^ (sequence << 40))
210}
211
212/// One exec, addressed by its caller-minted id.
213///
214/// The id is the idempotency key, so a handle survives a process restart: rebuild it
215/// with the same id and every method still addresses the same server-side exec.
216pub struct ExecHandle {
217 transport: Arc<Transport>,
218 exec_id: String,
219}
220
221impl ExecHandle {
222 pub(crate) fn new(transport: Arc<Transport>, exec_id: String) -> Self {
223 Self { transport, exec_id }
224 }
225
226 pub fn exec_id(&self) -> &str {
227 &self.exec_id
228 }
229
230 /// Reads current status and output. Read-only server-side; safe to spin on.
231 pub async fn poll(&self) -> Result<ExecResult, Error> {
232 let response: protocol::exec::PollResponse = self
233 .transport
234 .send_json(HttpRequest::new(
235 "GET",
236 format!("/v1/exec/{}", self.exec_id),
237 ))
238 .await?;
239 Ok(response.into())
240 }
241
242 /// Polls until the exec is done, or fails with a client-side timeout.
243 ///
244 /// A timeout here has not touched the exec — polling is read-only and the output
245 /// lives until it is acked — so a caller that gives up can come back and poll again,
246 /// and the message says so.
247 ///
248 /// A retryable error mid-wait is swallowed rather than raised: a VM under load drops
249 /// a connection occasionally, and the whole point of a read-only poll is that
250 /// repeating it costs nothing. A fatal one ends the wait.
251 pub async fn wait(&self, timeout: Duration) -> Result<ExecResult, Error> {
252 let deadline = tokio::time::Instant::now() + timeout;
253 let mut last_phase: Option<protocol::exec::Phase> = None;
254 loop {
255 match self.poll().await {
256 Ok(result) => {
257 if result.done() {
258 return Ok(result);
259 }
260 last_phase = Some(result.phase);
261 }
262 Err(err) if err.retryable() => {}
263 Err(err) => return Err(err),
264 }
265 let now = tokio::time::Instant::now();
266 if now >= deadline {
267 let phase = match last_phase {
268 Some(protocol::exec::Phase::Running) => "running",
269 Some(protocol::exec::Phase::Exited) => "exited",
270 Some(protocol::exec::Phase::Acked) => "acked",
271 None => "unknown",
272 };
273 return Err(Error::wire(
274 WireKind::ExecTimeout,
275 format!(
276 "exec {} was still {phase} after {}s; the record and its output \
277 are untouched and can be re-polled",
278 self.exec_id,
279 timeout.as_secs()
280 ),
281 ));
282 }
283 tokio::time::sleep(POLL_INTERVAL.min(deadline - now)).await;
284 }
285 }
286
287 /// Yields output as it arrives, reconnecting at the last good offset.
288 pub fn stream(&self) -> impl Stream<Item = Result<ExecEvent, Error>> + '_ {
289 self.stream_with(StreamOptions::default())
290 }
291
292 /// [`Self::stream`] with the knobs. See [`StreamOptions`].
293 ///
294 /// Written as a generator over an explicit state machine rather than as a hand-rolled
295 /// `Stream` impl: the reconnect logic is a loop with an `await` in the middle of it,
296 /// and expressing that as a `poll_next` would mean storing the in-flight attach as a
297 /// pinned field — which is where a self-referential-future bug lives.
298 pub fn stream_with(
299 &self,
300 options: StreamOptions,
301 ) -> impl Stream<Item = Result<ExecEvent, Error>> + '_ {
302 futures_util::stream::unfold(
303 StreamState::Reconnect {
304 cursor: options.offset,
305 attempts: 0,
306 },
307 move |state| {
308 let options = options.clone();
309 async move { self.advance(state, &options).await }
310 },
311 )
312 }
313
314 /// Drives a stream to its end, handing every event to `on_event`.
315 ///
316 /// The same state machine [`Self::stream_with`] runs, driven by an ordinary loop
317 /// instead of wrapped in a [`Stream`]. Everything about the wire behaviour is
318 /// identical — same reconnects, same cursor, same terminal-event rule — and both share
319 /// [`Self::advance`], so there is one implementation of the property and two ways to
320 /// consume it.
321 ///
322 /// # Why a callback exists beside the `Stream`
323 ///
324 /// Because `Stream` is not in `std`. A consumer of [`Self::stream_with`] has to name
325 /// the crate that defines the trait in order to call `poll_next`, so a client that
326 /// wants to print a line per event acquires a dependency to advance a loop —
327 /// `microvms-cli` carried `futures-util` for exactly that and nothing else, and
328 /// `tests/thinness.rs` (CLI-2) asserts that crate's dependency set exactly. This
329 /// method is that dependency's replacement: `ControlFlow` and `async fn` are both std,
330 /// so a caller needs nothing but this crate.
331 ///
332 /// A **synchronous** callback is the right shape for a caller whose per-event work is
333 /// synchronous — the CLI writes a line and returns. A caller whose per-event work is an
334 /// `await` (both bindings: a bounded-channel `send`) wants
335 /// [`Self::for_each_event_async`], which is the same loop with an awaited callback.
336 /// Using this one there would mean `blocking_send` on a runtime worker, which is why
337 /// the overload exists.
338 ///
339 /// # What the return value says that a tally cannot
340 ///
341 /// [`StreamEnd::reason`] distinguishes an `exit` event from a cut body from a caller
342 /// that stopped reading. A caller counting events cannot: a cut stream and a finished
343 /// one differ only in the terminal event, which is the argument for SSE framing here in
344 /// the first place. [`StreamEnd::cursor`] is the state machine's own offset, so a
345 /// caller resuming does not maintain a second one.
346 ///
347 /// An error ends the drive with `Err` and the events already handed over **stay handed
348 /// over**. That asymmetry is deliberate: the bytes a callback already wrote are real
349 /// output the caller received, and there is nothing to unwind them with.
350 pub async fn for_each_event<F>(
351 &self,
352 options: StreamOptions,
353 mut on_event: F,
354 ) -> Result<StreamEnd, Error>
355 where
356 F: FnMut(ExecEvent) -> std::ops::ControlFlow<()>,
357 {
358 // Delegated rather than a second copy of the loop. A synchronous callback *is* an
359 // async one whose future is already complete, and `std::future::ready` says exactly
360 // that — so the ordering, the cursor read, the `Break` arm, and the three endings
361 // have one implementation instead of two that agree until one is edited.
362 self.for_each_event_async(options, |event| std::future::ready(on_event(event)))
363 .await
364 }
365
366 /// [`Self::for_each_event`] for a callback that **awaits**.
367 ///
368 /// The same loop and the same [`Self::advance`] state machine; the only difference is
369 /// that the per-event callback answers a future, which is `.await`ed before the next
370 /// step is taken. So the events are still delivered strictly in order and still one at a
371 /// time — a callback that awaits between events holds the drive rather than racing it,
372 /// which is what makes this usable as backpressure.
373 ///
374 /// # Why the overload exists, and what it closed
375 ///
376 /// Both bindings consume a stream by pushing into a **capacity-1** `mpsc` channel that a
377 /// foreign-language iterator drains (`microvms-py/src/exec.rs`, `microvms-js/src/exec.rs`).
378 /// The channel bound is deliberate — the daemon's SSE body is the backpressure signal —
379 /// and it means the channel is full whenever the host-language consumer is even slightly
380 /// behind, i.e. normally. With [`Self::for_each_event`] the only available send is
381 /// `blocking_send`, which would block the very runtime worker the driver is running on;
382 /// so both bindings kept the `Stream` path, and `futures-util` with it, for a reason
383 /// about backpressure rather than about the trait. `send(..).await` inside an
384 /// `AsyncFnMut` yields the worker instead, and both bindings dropped the dependency on
385 /// 2026-08-09.
386 ///
387 /// # `FnMut(ExecEvent) -> Fut` rather than `AsyncFnMut`
388 ///
389 /// Measured, not stylistic, and the measurement is the reason the signature looks older
390 /// than the edition. `AsyncFnMut` is the obvious spelling and it compiles here — but a
391 /// caller cannot `tokio::spawn` a drive that uses it. Proving the returned future `Send`
392 /// requires naming `F::CallRefFuture<'a>` under a `for<'a>` bound, which is the unstable
393 /// `async_fn_traits` feature; without it the spawn fails with *"`Send` would have to be
394 /// implemented for the type `&ExecHandle`… but `Send` is actually implemented for
395 /// `&'0 ExecHandle`, for some specific lifetime"*. Both bindings spawn this drive onto a
396 /// runtime, so that is not a corner they can avoid, and the unstable bound additionally
397 /// forces the callback's captures to `'static`.
398 ///
399 /// So `Fut` is a plain type parameter. The cost is that the callback's future cannot
400 /// *borrow* from the closure's captures — it has no lifetime to name them with — so a
401 /// caller sending on a channel clones the sender per event rather than holding `&sender`.
402 /// For a `tokio::sync::mpsc::Sender` that is one atomic increment, which is the right
403 /// price for a signature that can cross a spawn.
404 ///
405 /// The return value and the error asymmetry are [`Self::for_each_event`]'s, unchanged.
406 pub async fn for_each_event_async<F, Fut>(
407 &self,
408 options: StreamOptions,
409 mut on_event: F,
410 ) -> Result<StreamEnd, Error>
411 where
412 F: FnMut(ExecEvent) -> Fut,
413 Fut: std::future::Future<Output = std::ops::ControlFlow<()>>,
414 {
415 let mut state = StreamState::Reconnect {
416 cursor: options.offset,
417 attempts: 0,
418 };
419 // Seeded from the caller's own starting offset, so a drive that ends before any
420 // event arrives reports where it began rather than zero.
421 let mut cursor = options.offset;
422 loop {
423 let Some((item, next)) = self.advance(state, &options).await else {
424 // `advance` yields nothing only when the body ended without an `exit`
425 // event and reconnecting was refused — the `exit` case returns below,
426 // before the machine is stepped again. So this is the cut.
427 return Ok(StreamEnd {
428 reason: EndReason::Cut,
429 cursor,
430 });
431 };
432 // Read off the machine rather than recomputed from the event: the two rules in
433 // this module's docs live in `advance`, and a second `cursor.max(end)` here
434 // would be a second implementation that agrees until a `gap` arrives.
435 if let Some(advanced) = next.cursor() {
436 cursor = advanced;
437 }
438 state = next;
439
440 let event = item?;
441 let terminal = matches!(event, ExecEvent::Exit(_));
442 // Awaited before the loop advances, which is the whole point: the next attach
443 // read does not start until the callback's own work is done, so a slow consumer
444 // slows the stream rather than buffering behind it.
445 if on_event(event).await.is_break() {
446 return Ok(StreamEnd {
447 reason: EndReason::Stopped,
448 cursor,
449 });
450 }
451 if terminal {
452 return Ok(StreamEnd {
453 reason: EndReason::Exited,
454 cursor,
455 });
456 }
457 }
458 }
459
460 /// One step of the stream state machine.
461 ///
462 /// Returns `None` to end the stream, or the next item plus the state after it.
463 async fn advance(
464 &self,
465 state: StreamState,
466 options: &StreamOptions,
467 ) -> Option<(Result<ExecEvent, Error>, StreamState)> {
468 let mut state = state;
469 loop {
470 match state {
471 StreamState::Done => return None,
472 StreamState::Reconnect { cursor, attempts } => {
473 if attempts > 0 {
474 if !options.reconnect {
475 return None;
476 }
477 if attempts > options.max_reconnects {
478 return Some((
479 Err(Error::new(
480 ErrorKind::Retryable,
481 format!(
482 "the stream of {} dropped {attempts} times without \
483 an exit event; last good offset {cursor}",
484 self.exec_id
485 ),
486 )),
487 StreamState::Done,
488 ));
489 }
490 let backoff = RECONNECT_BACKOFF
491 [((attempts - 1) as usize).min(RECONNECT_BACKOFF.len() - 1)];
492 tokio::time::sleep(backoff).await;
493 }
494 match self.attach(cursor, options.idle_timeout).await {
495 Ok(attach) => {
496 state = StreamState::Attached {
497 attach: Box::new(attach),
498 cursor,
499 attempts,
500 };
501 }
502 // A cut connection or a failed mint. Neither says anything
503 // about the exec, which is still running server-side, so both
504 // are reconnected through rather than surfaced.
505 Err(err) if err.retryable() && options.reconnect => {
506 state = StreamState::Reconnect {
507 cursor,
508 attempts: attempts + 1,
509 };
510 }
511 // Anything fatal — a 404 on a collected entry above all, where
512 // reconnecting can never succeed — ends the stream with the
513 // error rather than looping on it.
514 Err(err) => return Some((Err(err), StreamState::Done)),
515 }
516 }
517 StreamState::Attached {
518 mut attach,
519 mut cursor,
520 attempts,
521 } => match attach.next_event().await {
522 Ok(Some(ExecEvent::Output {
523 stream,
524 offset,
525 data,
526 })) => {
527 // Advance only past bytes actually handed over.
528 let end = offset + data.len() as u64;
529 cursor = cursor.max(end);
530 return Some((
531 Ok(ExecEvent::Output {
532 stream,
533 offset,
534 data,
535 }),
536 StreamState::Attached {
537 attach,
538 cursor,
539 attempts,
540 },
541 ));
542 }
543 Ok(Some(ExecEvent::Gap { from, to })) => {
544 // Follow the daemon past the evicted range, or a reconnect asks
545 // for it again and is told about the same gap forever.
546 cursor = cursor.max(to);
547 let item = if options.error_on_gap {
548 Err(Error::wire(
549 WireKind::OutputGap,
550 format!("output bytes [{from}, {to}) are unrecoverable"),
551 ))
552 } else {
553 Ok(ExecEvent::Gap { from, to })
554 };
555 let next = if options.error_on_gap {
556 StreamState::Done
557 } else {
558 StreamState::Attached {
559 attach,
560 cursor,
561 attempts,
562 }
563 };
564 return Some((item, next));
565 }
566 // The terminal event. The stream is over, and this is the only thing
567 // that says so.
568 Ok(Some(ExecEvent::Exit(exit))) => {
569 return Some((Ok(ExecEvent::Exit(exit)), StreamState::Done));
570 }
571 // The body ended with no exit event, so the connection was cut.
572 Ok(None) => {
573 if !options.reconnect {
574 return None;
575 }
576 state = StreamState::Reconnect {
577 cursor,
578 attempts: attempts + 1,
579 };
580 }
581 Err(err) if err.retryable() && options.reconnect => {
582 state = StreamState::Reconnect {
583 cursor,
584 attempts: attempts + 1,
585 };
586 }
587 Err(err) => return Some((Err(err), StreamState::Done)),
588 },
589 }
590 }
591 }
592
593 /// Opens one attach at `offset`.
594 async fn attach(&self, offset: u64, idle_timeout: Duration) -> Result<Attach, Error> {
595 let path = format!("/v1/exec/{}/stream?offset={offset}", self.exec_id);
596 let mut request = HttpRequest::new("GET", path.clone());
597 request
598 .headers
599 .push(("accept".into(), "text/event-stream".into()));
600 // Headers are built here rather than in `Transport::request`, because the
601 // streaming path does not go through it — and the mint has to happen anyway,
602 // which is what makes a mid-stream reconnect re-mint an expired token.
603 request.headers.extend(self.transport.headers(None).await?);
604 request.timeout = None;
605
606 let (head, chunks) = self
607 .transport
608 .backend
609 .open_stream(request, idle_timeout)
610 .await?;
611 // The status is checked before any body byte, so a 404 on an unknown exec id
612 // surfaces as NotFound rather than as an empty stream.
613 head.error_for_status("GET", &path)?;
614 Ok(Attach {
615 chunks,
616 parser: SseParser::new(),
617 frames: std::collections::VecDeque::new(),
618 })
619 }
620
621 /// Writes to the child's stdin.
622 ///
623 /// Requires the exec to have been started with `stdin: true`, or the daemon answers
624 /// 409. `eof` in the same call is the common case for feeding a prompt: two round
625 /// trips would leave a window where the child has the bytes but not the EOF that
626 /// tells it the input is complete.
627 pub async fn write_stdin(&self, data: &[u8], eof: bool) -> Result<StdinAck, Error> {
628 let body = protocol::exec::StdinRequest {
629 data_b64: (!data.is_empty())
630 .then(|| base64::engine::general_purpose::STANDARD.encode(data)),
631 signal: eof.then(|| "eof".to_string()),
632 };
633 let mut request = HttpRequest::new("POST", format!("/v1/exec/{}/stdin", self.exec_id));
634 request
635 .headers
636 .push(("content-type".into(), "application/json".into()));
637 request.body = serde_json::to_vec(&body).map_err(|err| {
638 Error::invalid_arg(format!("the stdin body will not serialize: {err}"))
639 })?;
640 self.transport.send_json(request).await
641 }
642
643 /// Sends EOF.
644 ///
645 /// Nothing else closes stdin: the daemon's copy of the pipe outlives `Child::wait()`,
646 /// so a child blocked reading stdin hangs until its timeout unless someone calls
647 /// this.
648 pub async fn close_stdin(&self) -> Result<StdinAck, Error> {
649 self.write_stdin(&[], true).await
650 }
651
652 /// Releases the buffered output and starts the TTL clock.
653 ///
654 /// 409 means either the exec has not exited — output is still being written — or an
655 /// earlier ack already took it. Both are real states, not the same state, and the
656 /// daemon's detail string distinguishes them.
657 pub async fn ack(&self) -> Result<ExecResult, Error> {
658 let response: protocol::exec::PollResponse = self
659 .transport
660 .send_json(HttpRequest::new(
661 "POST",
662 format!("/v1/exec/{}/ack", self.exec_id),
663 ))
664 .await?;
665 Ok(response.into())
666 }
667
668 /// Signals the whole process group, not just the direct child.
669 ///
670 /// Returns whether anything was signalled. `false` means no pgid was ever captured,
671 /// i.e. the child had already been reaped — which is the outcome a kill wanted, and
672 /// is why it is a 200 rather than an error.
673 pub async fn kill(&self) -> Result<bool, Error> {
674 let response: protocol::exec::KillResponse = self
675 .transport
676 .send_json(HttpRequest::new(
677 "POST",
678 format!("/v1/exec/{}/kill", self.exec_id),
679 ))
680 .await?;
681 Ok(response.killed)
682 }
683
684 /// Waits, then acks, returning the result that carries the output.
685 ///
686 /// Which result is returned matters: the ack response carries the released output,
687 /// and a poll issued after the ack reports `acked` with no output at all. Returning
688 /// the wrong one is a silent empty-output bug, so the sequencing lives here rather
689 /// than at every call site.
690 pub async fn wait_and_ack(&self, timeout: Duration) -> Result<ExecResult, Error> {
691 let done = self.wait(timeout).await?;
692 if done.phase == protocol::exec::Phase::Acked {
693 return Ok(done);
694 }
695 self.ack().await
696 }
697}
698
699impl std::fmt::Debug for ExecHandle {
700 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
701 f.debug_struct("ExecHandle")
702 .field("exec_id", &self.exec_id)
703 .finish()
704 }
705}
706
707/// One live attach: a chunk source plus the parser holding its partial frame.
708struct Attach {
709 chunks: Box<dyn super::ChunkSource>,
710 parser: SseParser,
711 frames: std::collections::VecDeque<super::Frame>,
712}
713
714impl Attach {
715 /// The next event this attach can produce, or `None` when its body ends.
716 ///
717 /// A frame that decodes to nothing is skipped rather than returned, so an unknown
718 /// event name does not surface as a spurious end-of-stream.
719 async fn next_event(&mut self) -> Result<Option<ExecEvent>, Error> {
720 loop {
721 while let Some(frame) = self.frames.pop_front() {
722 if let Some(event) = decode(&frame)? {
723 return Ok(Some(event));
724 }
725 }
726 let Some(chunk) = self.chunks.next_chunk().await? else {
727 return Ok(None);
728 };
729 // A parse failure here is the undelimited-data ceiling, which is
730 // `ErrorKind::Protocol` and therefore not retryable — so `advance` ends the
731 // stream with it rather than reconnecting into the same non-SSE body. That is
732 // the point: a proxy answering an error page would otherwise be retried
733 // `max_reconnects` times, filling the buffer again each time.
734 self.frames.extend(self.parser.feed(&chunk)?);
735 }
736 }
737}
738
739/// Where the stream is. Boxed attach because the enum otherwise carries a trait object
740/// inline and every state pays its size.
741enum StreamState {
742 /// Needs an attach. `attempts` is zero for the first one, which is why the backoff
743 /// and the max-reconnect check are both skipped there.
744 Reconnect {
745 cursor: u64,
746 attempts: u32,
747 },
748 Attached {
749 attach: Box<Attach>,
750 cursor: u64,
751 attempts: u32,
752 },
753 Done,
754}
755
756impl StreamState {
757 /// The cursor this state carries, or `None` for [`StreamState::Done`].
758 ///
759 /// `Done` has none on purpose: it is reached from three different places and inventing
760 /// one there would mean picking a number, where the caller
761 /// ([`ExecHandle::for_each_event`]) already holds the last real one. `None` says "this
762 /// state advanced nothing" rather than "the cursor is zero".
763 fn cursor(&self) -> Option<u64> {
764 match self {
765 StreamState::Reconnect { cursor, .. } | StreamState::Attached { cursor, .. } => {
766 Some(*cursor)
767 }
768 StreamState::Done => None,
769 }
770 }
771}
772
773#[cfg(test)]
774mod tests {
775 use super::*;
776 use crate::session::testing::{Recorder, Reply, session_with};
777 use futures_util::StreamExt as _;
778
779 /// One SSE `output` frame on the wire.
780 fn output(offset: u64, bytes: &[u8]) -> Vec<u8> {
781 let encoded = base64::engine::general_purpose::STANDARD.encode(bytes);
782 format!(
783 "event: output\ndata: {{\"offset\":{offset},\"stream\":\"stdout\",\
784 \"output\":\"{encoded}\"}}\n\n"
785 )
786 .into_bytes()
787 }
788
789 fn exit(total: u64) -> Vec<u8> {
790 format!(
791 "event: exit\ndata: {{\"exit_code\":0,\"signal\":null,\"truncated\":false,\
792 \"writers_may_be_alive\":false,\"offset\":{total}}}\n\n"
793 )
794 .into_bytes()
795 }
796
797 fn gap(from: u64, to: u64) -> Vec<u8> {
798 format!("event: gap\ndata: {{\"from\":{from},\"to\":{to}}}\n\n").into_bytes()
799 }
800
801 /// Collects a stream's bytes and its events.
802 async fn collect(
803 handle: &ExecHandle,
804 options: StreamOptions,
805 ) -> (Vec<u8>, Vec<Result<ExecEvent, Error>>) {
806 let mut bytes = Vec::new();
807 let mut events = Vec::new();
808 let mut stream = std::pin::pin!(handle.stream_with(options));
809 while let Some(item) = stream.next().await {
810 if let Ok(ExecEvent::Output { data, .. }) = &item {
811 bytes.extend_from_slice(data);
812 }
813 events.push(item);
814 }
815 (bytes, events)
816 }
817
818 /// The offset a stream request asked for.
819 fn requested_offset(request: &HttpRequest) -> u64 {
820 request
821 .path
822 .split("offset=")
823 .nth(1)
824 .expect("a stream request always carries an offset")
825 .parse()
826 .expect("the offset is a number")
827 }
828
829 /// The reconnect property, asserted on the reassembled byte sequence.
830 ///
831 /// The verdict is the bytes, not the absence of an error: a client that reconnected
832 /// at zero would deliver every byte too, and only the seam shows the difference. So
833 /// this checks the join *and* the offset the second attach asked for.
834 #[tokio::test(start_paused = true)]
835 async fn a_cut_mid_stream_reconnects_at_the_cursor_losing_and_duplicating_nothing() {
836 let recorder = Recorder::with([
837 // First attach: two output frames, then the body ends with no exit event.
838 Reply::Chunks(200, vec![output(0, b"AAAA\n"), output(5, b"BBBB\n")]),
839 // Second attach: the daemon replays from the requested offset.
840 Reply::Chunks(200, vec![output(10, b"CCCC\n"), exit(15)]),
841 ]);
842 let (session, _, _) = session_with(Arc::clone(&recorder));
843 let handle = session.exec("resume");
844
845 let (bytes, events) = collect(&handle, StreamOptions::default()).await;
846
847 assert_eq!(
848 String::from_utf8_lossy(&bytes),
849 "AAAA\nBBBB\nCCCC\n",
850 "the two attaches did not reconstruct the output, so the resume \
851 duplicated or lost bytes at the seam"
852 );
853 assert!(
854 events.iter().all(Result::is_ok),
855 "a cut surfaced as an error rather than a reconnect: {events:#?}"
856 );
857 assert!(
858 matches!(events.last(), Some(Ok(ExecEvent::Exit(_)))),
859 "the stream did not end on the terminal event: {events:#?}"
860 );
861
862 let seen = recorder.requests();
863 assert_eq!(
864 seen.len(),
865 2,
866 "the cut did not produce exactly one reconnect"
867 );
868 assert_eq!(requested_offset(&seen[0]), 0);
869 assert_eq!(
870 requested_offset(&seen[1]),
871 10,
872 "the reconnect asked for the wrong byte, so the seam is wrong"
873 );
874 }
875
876 /// A stream that ends *with* an exit event does not reconnect.
877 ///
878 /// The mirror of the test above, and the reason the transport is framed: the two
879 /// cases are the same byte count and differ only in the terminal event.
880 #[tokio::test(start_paused = true)]
881 async fn a_stream_that_ended_with_an_exit_event_does_not_reconnect() {
882 let recorder = Recorder::with([Reply::Chunks(200, vec![output(0, b"hi\n"), exit(3)])]);
883 let (session, _, _) = session_with(Arc::clone(&recorder));
884 let handle = session.exec("finished");
885
886 let (bytes, _) = collect(&handle, StreamOptions::default()).await;
887 assert_eq!(bytes, b"hi\n");
888 assert_eq!(
889 recorder.requests().len(),
890 1,
891 "a finished stream was reattached, which would replay a completed exec \
892 forever"
893 );
894 }
895
896 /// A gap advances the cursor past the evicted range.
897 ///
898 /// Without this the reconnect asks for bytes the daemon has already dropped and is
899 /// told about the same gap forever — a livelock that looks like a slow stream.
900 #[tokio::test(start_paused = true)]
901 async fn a_gap_advances_the_cursor_so_a_reconnect_does_not_ask_for_evicted_bytes() {
902 let recorder = Recorder::with([
903 Reply::Chunks(200, vec![output(0, b"AA"), gap(2, 900)]),
904 Reply::Chunks(200, vec![output(900, b"ZZ"), exit(902)]),
905 ]);
906 let (session, _, _) = session_with(Arc::clone(&recorder));
907 let handle = session.exec("lagged");
908
909 let (bytes, events) = collect(&handle, StreamOptions::default()).await;
910 assert_eq!(bytes, b"AAZZ");
911 assert!(
912 events
913 .iter()
914 .any(|event| matches!(event, Ok(ExecEvent::Gap { from: 2, to: 900 }))),
915 "the gap was swallowed, so a truncated log reads as a complete one: \
916 {events:#?}"
917 );
918 assert_eq!(
919 requested_offset(&recorder.requests()[1]),
920 900,
921 "the reconnect asked for bytes the daemon had already evicted"
922 );
923 }
924
925 /// `error_on_gap` turns a gap into a typed error and ends the stream.
926 #[tokio::test(start_paused = true)]
927 async fn error_on_gap_surfaces_the_lost_range_as_an_error() {
928 let recorder = Recorder::with([Reply::Chunks(200, vec![output(0, b"AA"), gap(2, 900)])]);
929 let (session, _, _) = session_with(recorder);
930 let handle = session.exec("lagged");
931
932 let (_, events) = collect(
933 &handle,
934 StreamOptions {
935 error_on_gap: true,
936 ..StreamOptions::default()
937 },
938 )
939 .await;
940 let err = events
941 .last()
942 .expect("an event")
943 .as_ref()
944 .expect_err("the gap is an error");
945 assert_eq!(err.wire_kind(), Some(WireKind::OutputGap));
946 assert!(err.to_string().contains("[2, 900)"), "{err}");
947 }
948
949 /// A transport failure opening the attach is reconnected through, and the mint that
950 /// the reconnect performs is what carries a fresh token mid-stream (TRAP-9).
951 #[tokio::test(start_paused = true)]
952 async fn a_failed_attach_is_retried_and_the_reconnect_re_mints_an_expired_token() {
953 let recorder = Recorder::with([
954 Reply::Cut("connection reset"),
955 Reply::Chunks(200, vec![output(0, b"ok\n"), exit(3)]),
956 ]);
957 let (session, auth, clock) = session_with(Arc::clone(&recorder));
958 let handle = session.exec("retried");
959
960 // Age the cached token past the refresh window before the stream starts, so the
961 // attach has to mint. The reconnect then crosses the window again.
962 clock.advance(super::super::DEFAULT_REFRESH_AFTER + Duration::from_secs(1));
963
964 let (bytes, events) = collect(&handle, StreamOptions::default()).await;
965 assert_eq!(bytes, b"ok\n");
966 assert!(events.iter().all(Result::is_ok), "{events:#?}");
967 assert_eq!(
968 auth.mint_count(),
969 1,
970 "the attach did not mint inside the request path"
971 );
972 assert_eq!(recorder.requests().len(), 2);
973 }
974
975 /// A 404 is fatal even with reconnect on: the entry was collected, so reattaching
976 /// can never succeed.
977 #[tokio::test(start_paused = true)]
978 async fn a_collected_exec_ends_the_stream_rather_than_reconnecting_forever() {
979 let recorder = Recorder::with([Reply::Body(404, b"{\"error\":\"unknown_exec\"}".to_vec())]);
980 let (session, _, _) = session_with(Arc::clone(&recorder));
981 let handle = session.exec("gone");
982
983 let (_, events) = collect(&handle, StreamOptions::default()).await;
984 assert_eq!(events.len(), 1);
985 let err = events[0].as_ref().expect_err("404 is an error");
986 assert_eq!(err.wire_kind(), Some(WireKind::NotFound));
987 assert_eq!(
988 recorder.requests().len(),
989 1,
990 "a fatal status was reattached, which can never succeed"
991 );
992 }
993
994 /// The reconnect budget is bounded, and running out is reported with the last good
995 /// offset rather than as a silent end.
996 #[tokio::test(start_paused = true)]
997 async fn a_stream_that_never_completes_gives_up_naming_its_last_good_offset() {
998 let mut replies = vec![Reply::Chunks(200, vec![output(0, b"AA")])];
999 for _ in 0..3 {
1000 replies.push(Reply::Chunks(200, vec![]));
1001 }
1002 let recorder = Recorder::with(replies);
1003 let (session, _, _) = session_with(Arc::clone(&recorder));
1004 let handle = session.exec("flaky");
1005
1006 let (bytes, events) = collect(
1007 &handle,
1008 StreamOptions {
1009 max_reconnects: 2,
1010 ..StreamOptions::default()
1011 },
1012 )
1013 .await;
1014 assert_eq!(bytes, b"AA");
1015 let err = events
1016 .last()
1017 .expect("an event")
1018 .as_ref()
1019 .expect_err("running out of reconnects is an error");
1020 assert!(err.retryable(), "the exec is still alive server-side");
1021 assert!(
1022 err.to_string().contains("last good offset 2"),
1023 "the give-up must name where to resume: {err}"
1024 );
1025 assert_eq!(recorder.requests().len(), 3, "the budget was not respected");
1026 }
1027
1028 /// `reconnect: false` ends the stream at the cut, for a caller doing its own.
1029 #[tokio::test(start_paused = true)]
1030 async fn reconnect_off_ends_the_stream_at_the_cut() {
1031 let recorder = Recorder::with([Reply::Chunks(200, vec![output(0, b"AA")])]);
1032 let (session, _, _) = session_with(Arc::clone(&recorder));
1033 let handle = session.exec("once");
1034
1035 let (bytes, events) = collect(
1036 &handle,
1037 StreamOptions {
1038 reconnect: false,
1039 ..StreamOptions::default()
1040 },
1041 )
1042 .await;
1043 assert_eq!(bytes, b"AA");
1044 assert!(events.iter().all(Result::is_ok));
1045 assert_eq!(recorder.requests().len(), 1);
1046 }
1047
1048 /// A non-zero starting offset is what a second process resuming a stream passes.
1049 #[tokio::test(start_paused = true)]
1050 async fn a_stream_can_start_at_an_offset_a_previous_process_left_off_at() {
1051 let recorder = Recorder::with([Reply::Chunks(200, vec![output(64, b"tail"), exit(68)])]);
1052 let (session, _, _) = session_with(Arc::clone(&recorder));
1053 let handle = session.exec("resumed");
1054
1055 let (bytes, _) = collect(
1056 &handle,
1057 StreamOptions {
1058 offset: 64,
1059 ..StreamOptions::default()
1060 },
1061 )
1062 .await;
1063 assert_eq!(bytes, b"tail");
1064 assert_eq!(requested_offset(&recorder.last()), 64);
1065 }
1066
1067 /// `wait_and_ack` returns the ack's result, which is the one carrying the output.
1068 ///
1069 /// Returning the poll's instead is a silent empty-output bug: a poll after the ack
1070 /// reports `acked` with nothing in it.
1071 #[tokio::test(start_paused = true)]
1072 async fn wait_and_ack_returns_the_ack_result_that_carries_the_output() {
1073 let recorder = Recorder::with([
1074 Reply::ok(serde_json::json!({"exec_id":"e1","phase":"running"})),
1075 Reply::ok(serde_json::json!({"exec_id":"e1","phase":"exited"})),
1076 Reply::ok(serde_json::json!({
1077 "exec_id":"e1","phase":"exited","exit_code":0,"signal":null,
1078 "stdout":"the output","stderr":"","truncated":false,
1079 "writers_may_be_alive":false
1080 })),
1081 ]);
1082 let (session, _, _) = session_with(Arc::clone(&recorder));
1083 let handle = session.exec("e1");
1084
1085 let result = handle
1086 .wait_and_ack(Duration::from_secs(30))
1087 .await
1088 .expect("waits then acks");
1089 assert_eq!(result.stdout(), "the output");
1090 assert!(result.succeeded());
1091
1092 let seen = recorder.requests();
1093 assert_eq!(seen.len(), 3);
1094 assert_eq!(seen[2].path, "/v1/exec/e1/ack");
1095 assert_eq!(seen[2].method, "POST");
1096 }
1097
1098 /// An already-acked exec is not acked twice: the second ack is a 409, and the phase
1099 /// is what tells the difference before the request goes out.
1100 #[tokio::test(start_paused = true)]
1101 async fn wait_and_ack_does_not_ack_an_already_acked_exec() {
1102 let recorder = Recorder::with([Reply::ok(serde_json::json!({
1103 "exec_id":"e1","phase":"acked","exit_code":0,"signal":null,
1104 "stdout":"","stderr":"","truncated":false,"writers_may_be_alive":false
1105 }))]);
1106 let (session, _, _) = session_with(Arc::clone(&recorder));
1107
1108 session
1109 .exec("e1")
1110 .wait_and_ack(Duration::from_secs(30))
1111 .await
1112 .expect("an acked exec is done");
1113 assert_eq!(
1114 recorder.requests().len(),
1115 1,
1116 "a second ack went out, which the daemon answers 409"
1117 );
1118 }
1119
1120 /// A retryable failure mid-wait is swallowed, because polling is read-only and
1121 /// repeating it costs nothing.
1122 #[tokio::test(start_paused = true)]
1123 async fn a_dropped_connection_mid_wait_is_retried_rather_than_raised() {
1124 let recorder = Recorder::with([
1125 Reply::Cut("connection reset"),
1126 Reply::ok(serde_json::json!({
1127 "exec_id":"e1","phase":"exited","exit_code":3,"signal":null,
1128 "stdout":"","stderr":"boom","truncated":false,"writers_may_be_alive":false
1129 })),
1130 ]);
1131 let (session, _, _) = session_with(recorder);
1132
1133 let result = session
1134 .exec("e1")
1135 .wait(Duration::from_secs(30))
1136 .await
1137 .expect("the retry lands");
1138 assert_eq!(result.exit_code(), Some(3));
1139 assert!(!result.succeeded());
1140 assert_eq!(result.stderr(), "boom");
1141 }
1142
1143 /// A wait that times out says the exec is untouched, because it is.
1144 #[tokio::test(start_paused = true)]
1145 async fn a_wait_timeout_says_the_record_can_be_re_polled() {
1146 let replies: Vec<Reply> = (0..10)
1147 .map(|_| Reply::ok(serde_json::json!({"exec_id":"e1","phase":"running"})))
1148 .collect();
1149 let recorder = Recorder::with(replies);
1150 let (session, _, _) = session_with(recorder);
1151
1152 let err = session
1153 .exec("e1")
1154 .wait(Duration::from_secs(3))
1155 .await
1156 .expect_err("the exec never finishes");
1157 assert_eq!(err.wire_kind(), Some(WireKind::ExecTimeout));
1158 assert_eq!(err.kind(), ErrorKind::Timeout);
1159 assert!(err.to_string().contains("still running"), "{err}");
1160 assert!(err.to_string().contains("re-polled"), "{err}");
1161 }
1162
1163 /// A stdin write base64s its data and names the eof signal, and an empty write with
1164 /// eof is a valid close.
1165 #[tokio::test]
1166 async fn a_stdin_write_carries_base64_and_the_eof_signal() {
1167 let recorder = Recorder::with([
1168 Reply::ok(serde_json::json!({"exec_id":"e1","written":3,"eof":true})),
1169 Reply::ok(serde_json::json!({"exec_id":"e1","written":0,"eof":true})),
1170 ]);
1171 let (session, _, _) = session_with(Arc::clone(&recorder));
1172 let handle = session.exec("e1");
1173
1174 let ack = handle.write_stdin(b"go\n", true).await.expect("writes");
1175 assert_eq!(ack.written, 3);
1176 assert!(ack.eof);
1177 let body: serde_json::Value =
1178 serde_json::from_slice(&recorder.last().body).expect("json body");
1179 assert_eq!(body["data_b64"], "Z28K");
1180 assert_eq!(body["signal"], "eof");
1181
1182 handle.close_stdin().await.expect("closes");
1183 let body: serde_json::Value =
1184 serde_json::from_slice(&recorder.last().body).expect("json body");
1185 assert!(
1186 body.get("data_b64").is_none_or(serde_json::Value::is_null),
1187 "a bare close must not send an empty data field: {body}"
1188 );
1189 assert_eq!(body["signal"], "eof");
1190 }
1191
1192 /// `kill` reports whether anything was signalled, and `false` with a 200 is a real
1193 /// answer rather than a failure.
1194 #[tokio::test]
1195 async fn kill_reports_false_for_a_group_that_had_already_exited() {
1196 let recorder = Recorder::with([
1197 Reply::ok(serde_json::json!({"exec_id":"e1","killed":true})),
1198 Reply::ok(serde_json::json!({"exec_id":"e1","killed":false})),
1199 ]);
1200 let (session, _, _) = session_with(Arc::clone(&recorder));
1201
1202 assert!(session.kill("e1").await.expect("kills"));
1203 assert!(!session.kill("e1").await.expect("already gone"));
1204 assert_eq!(recorder.last().path, "/v1/exec/e1/kill");
1205 }
1206
1207 /// A running poll has no outcome, and that absence is not an error — it is the shape
1208 /// of every poll before the last one.
1209 #[tokio::test]
1210 async fn a_running_poll_carries_no_outcome_and_is_not_an_error() {
1211 let recorder = Recorder::with([Reply::ok(
1212 serde_json::json!({"exec_id":"e1","phase":"running"}),
1213 )]);
1214 let (session, _, _) = session_with(recorder);
1215
1216 let result = session.exec("e1").poll().await.expect("polls");
1217 assert!(!result.done());
1218 assert!(result.outcome.is_none());
1219 assert_eq!(result.stdout(), "");
1220 assert_eq!(result.exit_code(), None);
1221 assert!(!result.succeeded());
1222 }
1223
1224 // ── the callback driver (`for_each_event`) ───────────────────────────────
1225 //
1226 // Four tests, mirroring the four `stream_with` properties that matter to a caller: the
1227 // events arrive in order, `Break` stops the drive, a reconnect still joins at the
1228 // cursor, and a cut without an exit event is *reported as a cut*. They are separate
1229 // tests rather than a re-run of the loop above on purpose — the driver is what the CLI
1230 // and (next) the bindings actually call, and "the `Stream` version works" is not a
1231 // statement about it.
1232
1233 /// Drives a stream with `for_each_event` and returns what the callback saw.
1234 async fn drive(
1235 handle: &ExecHandle,
1236 options: StreamOptions,
1237 ) -> (Vec<ExecEvent>, Result<StreamEnd, Error>) {
1238 let mut seen = Vec::new();
1239 let end = handle
1240 .for_each_event(options, |event| {
1241 seen.push(event);
1242 std::ops::ControlFlow::Continue(())
1243 })
1244 .await;
1245 (seen, end)
1246 }
1247
1248 /// **Events reach the callback in wire order, and the end names the terminal event.**
1249 ///
1250 /// Order is the property a callback API can silently break where a `Stream` cannot —
1251 /// buffering to hand over in a batch would reverse or coalesce, and the reader of a
1252 /// child's stdout would see the second chunk first. Asserted on the *offsets* rather
1253 /// than only on the reassembled bytes, because two chunks concatenated the wrong way
1254 /// round still yield the right total length.
1255 ///
1256 /// **Falsification** — push events onto a `Vec` inside `for_each_event` and call the
1257 /// callback after the loop, and the reason still reads `Exited` while the offsets come
1258 /// out of order only if the wire order was wrong; swap the two `Reply` chunks and this
1259 /// is red on `[0, 5]`. Verified: reversing the callback's argument order (delivering
1260 /// the second frame first) fails the offsets assertion.
1261 #[tokio::test(start_paused = true)]
1262 async fn the_driver_hands_every_event_to_the_callback_in_order() {
1263 let recorder = Recorder::with([Reply::Chunks(
1264 200,
1265 vec![output(0, b"AAAAA"), output(5, b"BBBBB"), exit(10)],
1266 )]);
1267 let (session, _, _) = session_with(Arc::clone(&recorder));
1268 let handle = session.exec("ordered");
1269
1270 let (seen, end) = drive(&handle, StreamOptions::default()).await;
1271 let end = end.expect("the stream completes");
1272
1273 let offsets: Vec<u64> = seen
1274 .iter()
1275 .filter_map(|event| match event {
1276 ExecEvent::Output { offset, .. } => Some(*offset),
1277 _ => None,
1278 })
1279 .collect();
1280 assert_eq!(offsets, [0, 5], "the callback saw the frames out of order");
1281 assert!(
1282 matches!(seen.last(), Some(ExecEvent::Exit(_))),
1283 "the terminal event has to be delivered, not merely observed: {seen:#?}"
1284 );
1285 assert_eq!(
1286 end,
1287 StreamEnd {
1288 reason: EndReason::Exited,
1289 cursor: 10,
1290 },
1291 "the end has to name the exit event and the total"
1292 );
1293 assert_eq!(recorder.requests().len(), 1);
1294 }
1295
1296 /// **`ControlFlow::Break` stops the drive, and says so rather than reporting a cut.**
1297 ///
1298 /// Two assertions and the second is the load-bearing one. Stopping has to be
1299 /// distinguishable from a truncated stream: both end without an exit event, and a
1300 /// caller that read `Cut` after its own `Break` would report "the command's outcome is
1301 /// unknown" about a command it chose not to watch. The cursor is where the callback
1302 /// stopped, so passing it back as `--from-offset` resumes exactly there.
1303 ///
1304 /// **Falsification** — return `EndReason::Cut` for the break arm and the reason
1305 /// assertion is red; drop the `is_break()` check and the drive runs to the exit event,
1306 /// so `seen.len()` reads 3 instead of 1. Verified: removing the early return makes the
1307 /// event-count assertion fail with 3.
1308 #[tokio::test(start_paused = true)]
1309 async fn control_flow_break_stops_the_stream_and_is_not_reported_as_a_cut() {
1310 let recorder = Recorder::with([Reply::Chunks(
1311 200,
1312 vec![output(0, b"first\n"), output(6, b"second\n"), exit(13)],
1313 )]);
1314 let (session, _, _) = session_with(Arc::clone(&recorder));
1315 let handle = session.exec("halted");
1316
1317 let mut seen = Vec::new();
1318 let end = handle
1319 .for_each_event(StreamOptions::default(), |event| {
1320 seen.push(event);
1321 // Stop on the first event, which is what a consumer with a `head -1` does.
1322 std::ops::ControlFlow::Break(())
1323 })
1324 .await
1325 .expect("stopping is not a failure");
1326
1327 assert_eq!(
1328 seen.len(),
1329 1,
1330 "the drive kept going after Break, so a caller that stopped reading is still \
1331 being read to: {seen:#?}"
1332 );
1333 assert_eq!(
1334 end,
1335 StreamEnd {
1336 reason: EndReason::Stopped,
1337 cursor: 6,
1338 },
1339 "a caller's own stop must not read as a cut stream, and the cursor is where it \
1340 stopped"
1341 );
1342 }
1343
1344 /// **A cut mid-drive still reconnects at the cursor, losing and duplicating nothing.**
1345 ///
1346 /// The same property as the `Stream` version's first test, asserted through the
1347 /// driver, because that is the path the CLI takes: if the driver dropped the cursor the
1348 /// reconnect would ask for byte zero and the caller would see `AAAA` twice, and the
1349 /// total byte count would look right. So the verdict is the reassembled bytes *and*
1350 /// the offset the second attach asked for.
1351 ///
1352 /// **Falsification** — seed the driver's `cursor` from zero instead of
1353 /// `options.offset`, or drop the `next.cursor()` read, and the reported end cursor is
1354 /// wrong while the bytes still join (the machine's own cursor is what drives the
1355 /// reconnect). Verified: replacing `next.cursor()` with `None` fails the `StreamEnd`
1356 /// assertion at cursor 0.
1357 #[tokio::test(start_paused = true)]
1358 async fn the_driver_reconnects_at_the_cursor_across_a_cut() {
1359 let recorder = Recorder::with([
1360 Reply::Chunks(200, vec![output(0, b"AAAA\n"), output(5, b"BBBB\n")]),
1361 Reply::Chunks(200, vec![output(10, b"CCCC\n"), exit(15)]),
1362 ]);
1363 let (session, _, _) = session_with(Arc::clone(&recorder));
1364 let handle = session.exec("resumed-driver");
1365
1366 let (seen, end) = drive(&handle, StreamOptions::default()).await;
1367 let end = end.expect("the reconnect lands");
1368
1369 let mut bytes = Vec::new();
1370 for event in &seen {
1371 if let ExecEvent::Output { data, .. } = event {
1372 bytes.extend_from_slice(data);
1373 }
1374 }
1375 assert_eq!(
1376 String::from_utf8_lossy(&bytes),
1377 "AAAA\nBBBB\nCCCC\n",
1378 "the two attaches did not reconstruct the output through the driver"
1379 );
1380 assert_eq!(end.reason, EndReason::Exited);
1381 assert_eq!(end.cursor, 15);
1382
1383 let seen_requests = recorder.requests();
1384 assert_eq!(seen_requests.len(), 2, "the cut produced no reconnect");
1385 assert_eq!(
1386 requested_offset(&seen_requests[1]),
1387 10,
1388 "the reconnect asked for the wrong byte, so the seam is wrong"
1389 );
1390 }
1391
1392 /// **A body that ended with no exit event and no reconnect is reported as `Cut`.**
1393 ///
1394 /// The one ending a byte count cannot detect, and the reason this returns a typed
1395 /// reason at all: `reconnect: false` is what a caller doing its own reconnection
1396 /// passes, and it must be able to tell "the command finished" from "the connection
1397 /// dropped and you told me not to retry". `cursor` is where its own reconnect starts.
1398 ///
1399 /// **Falsification** — report `EndReason::Exited` for the `advance` returned `None`
1400 /// arm and this is red on the reason; a caller acting on that would report exit code
1401 /// 0 for a command whose outcome nobody knows. Verified.
1402 #[tokio::test(start_paused = true)]
1403 async fn a_cut_with_reconnect_off_ends_the_drive_naming_the_cut() {
1404 let recorder = Recorder::with([Reply::Chunks(200, vec![output(0, b"partial")])]);
1405 let (session, _, _) = session_with(Arc::clone(&recorder));
1406 let handle = session.exec("truncated");
1407
1408 let (seen, end) = drive(
1409 &handle,
1410 StreamOptions {
1411 reconnect: false,
1412 ..StreamOptions::default()
1413 },
1414 )
1415 .await;
1416 let end = end.expect("a cut is not an error when reconnect is off");
1417
1418 assert_eq!(seen.len(), 1);
1419 assert!(
1420 !seen.iter().any(|event| matches!(event, ExecEvent::Exit(_))),
1421 "there was no terminal event to deliver"
1422 );
1423 assert_eq!(
1424 end,
1425 StreamEnd {
1426 reason: EndReason::Cut,
1427 cursor: 7,
1428 },
1429 "a stream with no exit event is a cut, and reporting anything else would let a \
1430 caller pass a CI step on output it never received"
1431 );
1432 assert_eq!(recorder.requests().len(), 1);
1433 }
1434
1435 /// A fatal error ends the drive with `Err`, and the events already delivered stay
1436 /// delivered.
1437 ///
1438 /// The asymmetry `for_each_event`'s docs state, asserted: a 404 mid-stream is fatal
1439 /// (the entry was collected, so reattaching can never succeed) and there is nothing to
1440 /// unwind the bytes a callback already wrote with.
1441 #[tokio::test(start_paused = true)]
1442 async fn a_fatal_status_ends_the_drive_with_an_error() {
1443 let recorder = Recorder::with([Reply::Body(404, b"{\"error\":\"unknown_exec\"}".to_vec())]);
1444 let (session, _, _) = session_with(Arc::clone(&recorder));
1445 let handle = session.exec("collected");
1446
1447 let (seen, end) = drive(&handle, StreamOptions::default()).await;
1448 assert!(seen.is_empty(), "nothing was delivered: {seen:#?}");
1449 let error = end.expect_err("a 404 on the attach is fatal");
1450 assert_eq!(error.wire_kind(), Some(WireKind::NotFound));
1451 assert_eq!(
1452 recorder.requests().len(),
1453 1,
1454 "a fatal status was reattached, which can never succeed"
1455 );
1456 }
1457
1458 // ── the async callback driver (`for_each_event_async`) ────────────────────
1459 //
1460 // Four tests. Three mirror `for_each_event`'s — order, `Break` stops, a reconnect joins at
1461 // the cursor — because "the sync driver works" is not a statement about this one even though
1462 // the sync driver is now written in terms of it: the delegation could be inverted tomorrow,
1463 // and the bindings call *this* method. The fourth is the property only an awaited callback
1464 // has, and it is the reason the overload exists at all: a callback that awaits between
1465 // events loses none of them.
1466
1467 /// Drives a stream with the async driver and returns what the callback saw.
1468 async fn drive_async(
1469 handle: &ExecHandle,
1470 options: StreamOptions,
1471 ) -> (Vec<ExecEvent>, Result<StreamEnd, Error>) {
1472 let mut seen = Vec::new();
1473 let end = handle
1474 .for_each_event_async(options, |event| {
1475 seen.push(event);
1476 std::future::ready(std::ops::ControlFlow::Continue(()))
1477 })
1478 .await;
1479 (seen, end)
1480 }
1481
1482 /// **Events reach an async callback in wire order, and the end names the terminal event.**
1483 ///
1484 /// Asserted on the *offsets* rather than only on the reassembled bytes, because two chunks
1485 /// concatenated the wrong way round still yield the right total length. Order is the
1486 /// property an awaited callback could break where the sync one cannot: a driver that
1487 /// launched each callback future without awaiting it, or awaited them out of order, would
1488 /// deliver a child's second stdout chunk before its first.
1489 ///
1490 /// **Falsification** — remove the `.await` on `on_event(event)` and the loop does not
1491 /// compile; replace it with a driver that collects the futures and joins them at the end
1492 /// and the offsets come out in completion order rather than wire order. Verified by
1493 /// reversing the `Reply::Chunks` frame order, which fails this assertion on `[5, 0]`.
1494 #[tokio::test(start_paused = true)]
1495 async fn the_async_driver_hands_every_event_to_the_callback_in_order() {
1496 let recorder = Recorder::with([Reply::Chunks(
1497 200,
1498 vec![output(0, b"AAAAA"), output(5, b"BBBBB"), exit(10)],
1499 )]);
1500 let (session, _, _) = session_with(Arc::clone(&recorder));
1501 let handle = session.exec("ordered-async");
1502
1503 let (seen, end) = drive_async(&handle, StreamOptions::default()).await;
1504 let end = end.expect("the stream completes");
1505
1506 let offsets: Vec<u64> = seen
1507 .iter()
1508 .filter_map(|event| match event {
1509 ExecEvent::Output { offset, .. } => Some(*offset),
1510 _ => None,
1511 })
1512 .collect();
1513 assert_eq!(offsets, [0, 5], "the callback saw the frames out of order");
1514 assert!(
1515 matches!(seen.last(), Some(ExecEvent::Exit(_))),
1516 "the terminal event has to be delivered, not merely observed: {seen:#?}"
1517 );
1518 assert_eq!(
1519 end,
1520 StreamEnd {
1521 reason: EndReason::Exited,
1522 cursor: 10,
1523 },
1524 "the end has to name the exit event and the total"
1525 );
1526 assert_eq!(recorder.requests().len(), 1);
1527 }
1528
1529 /// **`ControlFlow::Break` from an async callback stops the drive and is not a cut.**
1530 ///
1531 /// The load-bearing half is the reason: stopping has to be distinguishable from a truncated
1532 /// stream, because both end without an exit event. A binding whose host-language iterator
1533 /// was dropped answers `Break` here — a failed channel send — and it must not read as "the
1534 /// command's outcome is unknown".
1535 ///
1536 /// **Falsification** — return `EndReason::Cut` from the break arm and the reason assertion
1537 /// is red; drop the `is_break()` check and the drive runs to the exit event, so `seen.len()`
1538 /// reads 3 instead of 1. Verified: deleting the early return fails the event-count assertion
1539 /// with 3.
1540 #[tokio::test(start_paused = true)]
1541 async fn an_async_callback_breaking_stops_the_stream_and_is_not_reported_as_a_cut() {
1542 let recorder = Recorder::with([Reply::Chunks(
1543 200,
1544 vec![output(0, b"first\n"), output(6, b"second\n"), exit(13)],
1545 )]);
1546 let (session, _, _) = session_with(Arc::clone(&recorder));
1547 let handle = session.exec("halted-async");
1548
1549 let mut seen = Vec::new();
1550 let end = handle
1551 .for_each_event_async(StreamOptions::default(), |event| {
1552 seen.push(event);
1553 // Stop on the first event, which is what a dropped binding iterator does.
1554 std::future::ready(std::ops::ControlFlow::Break(()))
1555 })
1556 .await
1557 .expect("stopping is not a failure");
1558
1559 assert_eq!(
1560 seen.len(),
1561 1,
1562 "the drive kept going after Break, so a consumer that stopped reading is still \
1563 being read to: {seen:#?}"
1564 );
1565 assert_eq!(
1566 end,
1567 StreamEnd {
1568 reason: EndReason::Stopped,
1569 cursor: 6,
1570 },
1571 "a consumer's own stop must not read as a cut stream, and the cursor is where it \
1572 stopped"
1573 );
1574 }
1575
1576 /// **A cut mid-drive still reconnects at the cursor through the async driver.**
1577 ///
1578 /// The same seam property, asserted through the path the bindings take. If the driver lost
1579 /// the cursor the reconnect would ask for byte zero and the consumer would see `AAAA` twice
1580 /// while the total byte count still looked plausible — so the verdict is the reassembled
1581 /// bytes *and* the offset the second attach asked for.
1582 ///
1583 /// **Falsification** — seed `cursor` from zero instead of `options.offset`, or drop the
1584 /// `next.cursor()` read, and the reported end cursor is wrong while the bytes still join.
1585 /// Verified: replacing `next.cursor()` with `None` fails the end-cursor assertion at 0.
1586 #[tokio::test(start_paused = true)]
1587 async fn the_async_driver_reconnects_at_the_cursor_across_a_cut() {
1588 let recorder = Recorder::with([
1589 Reply::Chunks(200, vec![output(0, b"AAAA\n"), output(5, b"BBBB\n")]),
1590 Reply::Chunks(200, vec![output(10, b"CCCC\n"), exit(15)]),
1591 ]);
1592 let (session, _, _) = session_with(Arc::clone(&recorder));
1593 let handle = session.exec("resumed-async");
1594
1595 let (seen, end) = drive_async(&handle, StreamOptions::default()).await;
1596 let end = end.expect("the reconnect lands");
1597
1598 let mut bytes = Vec::new();
1599 for event in &seen {
1600 if let ExecEvent::Output { data, .. } = event {
1601 bytes.extend_from_slice(data);
1602 }
1603 }
1604 assert_eq!(
1605 String::from_utf8_lossy(&bytes),
1606 "AAAA\nBBBB\nCCCC\n",
1607 "the two attaches did not reconstruct the output through the async driver"
1608 );
1609 assert_eq!(end.reason, EndReason::Exited);
1610 assert_eq!(end.cursor, 15);
1611
1612 let seen_requests = recorder.requests();
1613 assert_eq!(seen_requests.len(), 2, "the cut produced no reconnect");
1614 assert_eq!(
1615 requested_offset(&seen_requests[1]),
1616 10,
1617 "the reconnect asked for the wrong byte, so the seam is wrong"
1618 );
1619 }
1620
1621 /// **A callback that awaits between events loses none of them.**
1622 ///
1623 /// The one property only this overload has, set up as the exact shape both bindings use: a
1624 /// **capacity-1** `mpsc` channel the callback `send`s into and a separate consumer drains,
1625 /// with the consumer deliberately slower than the producer so the channel is full for every
1626 /// event after the first. That is the configuration the sync driver could not serve — its
1627 /// only available send is `blocking_send`, which would park the runtime worker the driver
1628 /// itself is running on — and it is why `microvms-py` and `microvms-js` kept the `Stream`
1629 /// path until this method existed.
1630 ///
1631 /// Five events go in and five come out, in wire order, terminal event included. The
1632 /// verdict is the sequence and not a count: a driver that dropped an event under
1633 /// backpressure would still report `Exited` with the right cursor, because the cursor is the
1634 /// state machine's and not the callback's.
1635 ///
1636 /// **Falsification** — replace the awaited `send` with `try_send` and ignore the failure
1637 /// (the shape of a driver that does not respect backpressure) and this is red with three
1638 /// of five events delivered. Verified.
1639 #[tokio::test(start_paused = true)]
1640 async fn an_async_callback_that_awaits_between_events_loses_none_of_them() {
1641 let recorder = Recorder::with([Reply::Chunks(
1642 200,
1643 vec![
1644 output(0, b"AA"),
1645 output(2, b"BB"),
1646 output(4, b"CC"),
1647 output(6, b"DD"),
1648 exit(8),
1649 ],
1650 )]);
1651 let (session, _, _) = session_with(Arc::clone(&recorder));
1652 let handle = session.exec("backpressured");
1653
1654 // Capacity 1, the bindings' own bound: the SSE body is the backpressure signal, and
1655 // buffering here would defeat the cursor the core reconnects at.
1656 let (sender, mut receiver) = tokio::sync::mpsc::channel::<String>(1);
1657 let consumer = tokio::spawn(async move {
1658 let mut drained = Vec::new();
1659 while let Some(described) = receiver.recv().await {
1660 // Slower than the producer, so every send after the first waits.
1661 tokio::time::sleep(Duration::from_millis(50)).await;
1662 drained.push(described);
1663 }
1664 drained
1665 });
1666
1667 let end = handle
1668 .for_each_event_async(StreamOptions::default(), |event| {
1669 let described = match &event {
1670 ExecEvent::Output { offset, data, .. } => {
1671 format!("out@{offset}:{}", String::from_utf8_lossy(data))
1672 }
1673 ExecEvent::Gap { from, to } => format!("gap@{from}..{to}"),
1674 ExecEvent::Exit(exit) => format!("exit:{:?}", exit.exit_code),
1675 };
1676 // Cloned per event rather than borrowed, which is what the plain-`Fut`
1677 // signature requires and what both bindings do. One atomic increment.
1678 let sender = sender.clone();
1679 async move {
1680 match sender.send(described).await {
1681 Ok(()) => std::ops::ControlFlow::Continue(()),
1682 // A closed receiver is the dropped-iterator case, not this path.
1683 Err(_) => std::ops::ControlFlow::Break(()),
1684 }
1685 }
1686 })
1687 .await
1688 .expect("the stream completes even though the consumer lags");
1689 // Closing the sender is what ends the consumer's `recv` loop.
1690 drop(sender);
1691
1692 let drained = consumer.await.expect("the consumer task does not panic");
1693 assert_eq!(
1694 drained,
1695 [
1696 "out@0:AA",
1697 "out@2:BB",
1698 "out@4:CC",
1699 "out@6:DD",
1700 "exit:Some(0)",
1701 ],
1702 "a lagging capacity-1 consumer lost or reordered events, which is the failure the \
1703 awaited callback exists to prevent"
1704 );
1705 assert_eq!(
1706 end,
1707 StreamEnd {
1708 reason: EndReason::Exited,
1709 cursor: 8,
1710 },
1711 "the ending must still name the exit event and the total"
1712 );
1713 }
1714}