Skip to main content

microsandbox_agentd/
agent.rs

1//! Main agent loop: serial I/O, session management, heartbeat.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4use std::env;
5use std::fs::{File, OpenOptions};
6use std::os::fd::AsRawFd;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::task::{Context, Poll};
10use std::time::Instant;
11
12use bytes::BytesMut;
13use chrono::Utc;
14use tokio::io::unix::AsyncFd;
15use tokio::sync::{OwnedSemaphorePermit, Semaphore, watch};
16use tokio::time::{self, Duration};
17
18use microsandbox_protocol::AGENT_TRANSPORT_DUAL_PORT_CMDLINE;
19use microsandbox_protocol::bootstrap::GuestBootstrap;
20use microsandbox_protocol::bulk::{
21    BULK_HEADER_SIZE, BULK_PROTOCOL_VERSION, BulkCancel, BulkCancelReason, BulkCredit, BulkFinish,
22    BulkFlow, BulkKind, BulkRecord, DEFAULT_BULK_WINDOW, DEFAULT_FILESYSTEM_BULK_RECORD_PAYLOAD,
23    MAX_BULK_RECORD_PAYLOAD, MIN_BULK_RECORD_PAYLOAD,
24};
25use microsandbox_protocol::codec::{self, DecodedFrame, MAX_FRAME_SIZE};
26use microsandbox_protocol::core::{
27    ClockSync, CoreError, CoreErrorKind, InitAck, InitFailureReason, InitResolved, Ping, Pong,
28    Ready, RelayClientDisconnected, ResolvedUser, Touch, Touched,
29    WORKLOAD_TRANSPORT_BARRIER_VERSION, WORKLOAD_TRANSPORT_BULK_BYTES,
30    WORKLOAD_TRANSPORT_BULK_FRAMES, WORKLOAD_TRANSPORT_CONTROL_BYTES,
31    WORKLOAD_TRANSPORT_CONTROL_FRAMES, WorkloadFailure, WorkloadFailureDisposition, WorkloadFreeze,
32    WorkloadFrozen, WorkloadThaw, WorkloadThawed, WorkloadTransportCredit,
33    WorkloadTransportPosition,
34};
35use microsandbox_protocol::exec::{
36    ExecExited, ExecFailed, ExecFailureKind, ExecRequest, ExecResize, ExecSignal, ExecStarted,
37    ExecStderr, ExecStdin, ExecStdinError, ExecStdout,
38};
39use microsandbox_protocol::fs::{FS_CHUNK_SIZE, FsData, FsRequest, FsResponse};
40use microsandbox_protocol::heartbeat::{ActivityCounters, Heartbeat};
41use microsandbox_protocol::message::{FRAME_HEADER_SIZE, Message, MessageType};
42use microsandbox_protocol::tcp::{TcpClose, TcpConnect, TcpData, TcpEof, TcpFailed};
43use microsandbox_protocol::transport::{
44    BULK_BINDING_SIZE, BulkTransportReady, CLIENT_INCARNATION_SIZE, ClientIncarnation,
45    IncarnatedBulkFrame, RelayClientConnected, RelayClientDisconnectedAck, RelayLeaseReady,
46    decode_bulk_ack, encode_bulk_hello, encode_relay_client_disconnected_ack,
47    relay_client_id_range, relay_client_slot, try_decode_incarnated_bulk_from_bytes,
48    try_decode_relay_client_connected_from_bytes,
49    try_decode_relay_client_disconnected_ack_from_bytes,
50    validate_relay_client_range as canonical_relay_client_range,
51};
52
53use crate::config::{AgentdConfig, scripts_path};
54use crate::error::{AgentdError, AgentdResult};
55use crate::fs::{FsReadSession, FsState, FsStreamSession, FsWriteSession};
56use crate::process::ProcessManager;
57use crate::serial::{AGENT_BULK_PORT_NAME, AGENT_PORT_NAME, InputCharge, InputLane, InputWindow};
58use crate::session::{
59    BulkOutputCommand, ExecSession, RawActivity, RawSessionCompletion, RawSessionOutput,
60    SessionOutput, SessionOutputEnvelope, SessionOutputSender, resolve_default_user,
61};
62use crate::tcp::TcpSession;
63use crate::workload::{WorkloadLatch, WorkloadLatchError};
64use crate::{clock, fs, handoff, heartbeat, serial};
65
66//--------------------------------------------------------------------------------------------------
67// Constants
68//--------------------------------------------------------------------------------------------------
69
70/// Heartbeat interval in seconds.
71///
72/// Keep this short so small idle timeouts (for example `--idle-timeout 1`)
73/// can be enforced without multi-second scheduling drift.
74const HEARTBEAT_INTERVAL_SECS: u64 = 1;
75
76/// Read buffer size for the serial port.
77/// Control-only serial reads stay cache-friendly because their frames are small.
78const CONTROL_SERIAL_READ_BUF_SIZE: usize = 64 * 1024;
79
80/// Combined transport can receive any negotiated raw record without avoidable read cycles.
81const COMBINED_SERIAL_READ_BUF_SIZE: usize = DEFAULT_FILESYSTEM_BULK_RECORD_PAYLOAD as usize + 32;
82
83/// Dedicated bulk reads can consume one negotiated record in a single syscall when available.
84const BULK_SERIAL_READ_BUF_SIZE: usize =
85    DEFAULT_FILESYSTEM_BULK_RECORD_PAYLOAD as usize + CLIENT_INCARNATION_SIZE + 32;
86
87/// Maximum allowed input buffer size (frame size limit + 4 bytes for length prefix).
88const MAX_INPUT_BUF_SIZE: usize = MAX_FRAME_SIZE as usize + 4;
89
90/// Dedicated records additionally carry the transport-level client incarnation.
91const MAX_BULK_INPUT_BUF_SIZE: usize = MAX_INPUT_BUF_SIZE + CLIENT_INCARNATION_SIZE;
92
93/// Bound actual console work between runtime scheduling points, independently of wire records.
94/// Partial records count too: a readable bulk port must not hide fresh primary/PTY readiness.
95const AGENT_READ_QUANTUM_BYTES: usize = 256 * 1024;
96const AGENT_READ_QUANTUM_CALLS: usize = 64;
97
98/// Maximum time to wait for the host to acknowledge the init context.
99const INIT_ACK_TIMEOUT_SECS: u64 = 60;
100
101/// Startup window shared by bulk-port discovery and binding.
102const BULK_BINDING_TIMEOUT_SECS: u64 = 60;
103
104/// Per-correlation quantum used by the dedicated guest-to-host scheduler.
105const BULK_SCHEDULER_QUANTUM: usize = 256 * 1024;
106
107/// Maximum bytes one correlation may emit in one scheduler round.
108const BULK_SCHEDULER_MAX_BURST: usize = MAX_BULK_RECORD_PAYLOAD as usize;
109
110/// Maximum bytes queued for one correlation outside the console backend.
111const BULK_SCHEDULER_FLOW_CAPACITY: usize = 8 * 1024 * 1024;
112
113/// Maximum active bulk correlations owned by one relay client.
114const BULK_SCHEDULER_MAX_FLOWS_PER_CLIENT: usize = 64;
115
116/// Whole records processed before the bulk reader yields on a one-vCPU guest.
117const BULK_READER_MAX_RECORDS_PER_TURN: usize = 64;
118
119/// Payload bytes processed before the bulk reader yields on a one-vCPU guest.
120const BULK_READER_MAX_BYTES_PER_TURN: usize = 1024 * 1024;
121
122/// Enough data slots for the default byte window at the smallest negotiated record size.
123const FS_BULK_INPUT_ITEM_CAPACITY: usize =
124    DEFAULT_BULK_WINDOW as usize / MIN_BULK_RECORD_PAYLOAD as usize;
125
126/// Aggregate host-to-guest payload retained across every relay client and bulk destination.
127const BULK_INPUT_BYTE_CAPACITY: usize = 32 * 1024 * 1024;
128
129/// Filesystem activity bytes coalesced before publishing an otherwise empty worker event.
130const FS_ACTIVITY_BATCH_BYTES: usize = 4 * 1024 * 1024;
131
132/// Maximum contiguous payload coalesced into one guest filesystem effect.
133const FS_BULK_WRITE_COALESCE_BYTES: usize = FS_CHUNK_SIZE;
134
135/// Byte and time bounds for coalescing bulk-only heartbeat publications.
136const BULK_ACTIVITY_PUBLISH_BYTES: usize = 4 * 1024 * 1024;
137const BULK_ACTIVITY_PUBLISH_INTERVAL: Duration = Duration::from_millis(100);
138
139/// Best-effort window for publishing typed cancellation before agentd terminates a failed lane.
140const BULK_FAILURE_FLUSH_TIMEOUT: Duration = Duration::from_secs(2);
141
142//--------------------------------------------------------------------------------------------------
143// Types
144//--------------------------------------------------------------------------------------------------
145
146struct AgentState {
147    input_window: InputWindow,
148    pending_freeze: Option<Message>,
149    aborted_transport_attempt: Option<String>,
150    output_parked: bool,
151    resume_output_after_flush: bool,
152    frozen_host_input: Option<WorkloadTransportPosition>,
153    frozen_guest_bulk_bytes: u64,
154    // Retain stdin/PTY and process registrations until inherited output readers finish.
155    detached_sessions: HashMap<(u64, u32), ExecSession>,
156    stdin_poll_offset: usize,
157    restored_attempt: Option<String>,
158    client_incarnations: HashMap<u32, ClientIncarnation>,
159    bulk_input_budget: Arc<Semaphore>,
160    sessions: HashMap<u32, ExecSession>,
161    write_sessions: HashMap<u32, FsWriteSession>,
162    bulk_write_workers: HashMap<u32, FsBulkWriteWorker>,
163    read_sessions: HashMap<u32, FsReadSession>,
164    tcp_sessions: HashMap<u32, TcpSession>,
165    bulk_received_offsets: HashMap<u32, u64>,
166    pending_bulk_finishes: HashMap<u32, BulkFinish>,
167    fs: FsState,
168}
169
170/// Bounded command path for one filesystem bulk write, isolated from the control loop.
171struct FsBulkWriteWorker {
172    records: tokio::sync::mpsc::Sender<AdmittedBulkRecord>,
173    finish: tokio::sync::mpsc::Sender<BulkFinish>,
174    task: tokio::task::JoinHandle<()>,
175}
176
177/// One host-to-guest record whose global input capacity remains charged until the destination
178/// filesystem or TCP socket has consumed its payload.
179pub(crate) struct AdmittedBulkRecord {
180    record: BulkRecord,
181    _permit: OwnedSemaphorePermit,
182    _transport_charge: Option<InputCharge>,
183}
184
185/// Both budgets follow the payload all the way into asynchronous TCP/filesystem consumption.
186pub(crate) struct BulkInputPermit {
187    _payload: OwnedSemaphorePermit,
188    _transport: Option<InputCharge>,
189}
190
191/// Dedicated-lane input waiting for aggregate client capacity while the control actor stays live.
192struct PendingBulkInput {
193    frame: IncarnatedBulkFrame,
194    budget: Arc<Semaphore>,
195    charge: InputCharge,
196}
197
198struct ActivityTracker {
199    activity_seq: u64,
200    counters: ActivityCounters,
201}
202
203/// Buffered console input retained across the bootstrap and init handshakes.
204///
205/// A single device read may contain multiple frames, so boot phases must share
206/// this buffer instead of dropping bytes after decoding their expected frame.
207#[derive(Default)]
208pub struct BootConsoleState {
209    input: Vec<u8>,
210    init_failure_ack: bool,
211}
212
213#[derive(Clone)]
214struct HeartbeatSnapshot {
215    activity_seq: u64,
216    active_exec_sessions: u32,
217    active_fs_streams: u32,
218    active_tcp_streams: u32,
219    counters: ActivityCounters,
220}
221
222/// Successfully bound second console port and its per-boot identity.
223pub struct BoundBulkPort {
224    file: File,
225    connection_id: [u8; 16],
226}
227
228/// Dedicated-lane input owned by the main agent actor.
229///
230/// Keeping framing here removes the reader-task/channel round trip while bounded turns still give
231/// control work a scheduling point between filesystem or TCP bursts.
232struct BulkInputState {
233    port: AsyncFd<File>,
234    read_buf: Vec<u8>,
235    input: BytesMut,
236}
237
238/// Shared primary/bulk work retained across select turns, including incomplete wire records.
239#[derive(Default)]
240struct AgentReadBudget {
241    bytes: usize,
242    calls: usize,
243}
244
245/// One correlation's pending records in the guest-to-host DRR scheduler.
246struct BulkWriteFlow {
247    queue: VecDeque<SessionOutputEnvelope>,
248    queued_bytes: usize,
249    deficit: usize,
250}
251
252/// Deferred acknowledgement after the scheduler has discarded already-queued producer output.
253enum BulkOutputCleanup {
254    Park {
255        completion: tokio::sync::oneshot::Sender<u64>,
256        wire_bytes: u64,
257    },
258    Flow(tokio::sync::oneshot::Sender<()>),
259    Incarnation {
260        incarnation: ClientIncarnation,
261        completion: tokio::sync::oneshot::Sender<()>,
262    },
263}
264
265#[derive(Default)]
266struct BulkOutputPosition {
267    parked: bool,
268    wire_bytes: u64,
269}
270
271//--------------------------------------------------------------------------------------------------
272// Methods
273//--------------------------------------------------------------------------------------------------
274
275impl Default for AgentState {
276    fn default() -> Self {
277        Self {
278            input_window: InputWindow::new(initial_input_credit()),
279            pending_freeze: None,
280            aborted_transport_attempt: None,
281            output_parked: false,
282            resume_output_after_flush: false,
283            frozen_host_input: None,
284            frozen_guest_bulk_bytes: 0,
285            detached_sessions: HashMap::new(),
286            stdin_poll_offset: 0,
287            restored_attempt: None,
288            client_incarnations: HashMap::new(),
289            bulk_input_budget: Arc::new(Semaphore::new(BULK_INPUT_BYTE_CAPACITY)),
290            sessions: HashMap::new(),
291            write_sessions: HashMap::new(),
292            bulk_write_workers: HashMap::new(),
293            read_sessions: HashMap::new(),
294            tcp_sessions: HashMap::new(),
295            bulk_received_offsets: HashMap::new(),
296            pending_bulk_finishes: HashMap::new(),
297            fs: FsState::default(),
298        }
299    }
300}
301
302impl AdmittedBulkRecord {
303    pub(crate) fn record(&self) -> &BulkRecord {
304        &self.record
305    }
306
307    pub(crate) fn into_parts(self) -> (BulkRecord, BulkInputPermit) {
308        (
309            self.record,
310            BulkInputPermit {
311                _payload: self._permit,
312                _transport: self._transport_charge,
313            },
314        )
315    }
316
317    #[cfg(test)]
318    pub(crate) fn for_test(record: BulkRecord) -> Self {
319        let payload_len = record.payload.len();
320        let budget = Arc::new(Semaphore::new(payload_len));
321        let permit = budget
322            .try_acquire_many_owned(payload_len as u32)
323            .expect("test record fits its exact input budget");
324        Self {
325            record,
326            _permit: permit,
327            _transport_charge: None,
328        }
329    }
330}
331
332impl ActivityTracker {
333    fn new() -> Self {
334        Self {
335            activity_seq: 0,
336            counters: ActivityCounters::default(),
337        }
338    }
339
340    fn record_host_message(&mut self) {
341        self.touch();
342        self.counters.host_messages = self.counters.host_messages.saturating_add(1);
343    }
344
345    fn record_guest_message(&mut self) {
346        self.touch();
347        self.counters.guest_messages = self.counters.guest_messages.saturating_add(1);
348    }
349
350    fn add_exec_output_bytes(&mut self, len: usize) {
351        self.counters.exec_output_bytes =
352            self.counters.exec_output_bytes.saturating_add(len as u64);
353    }
354
355    fn add_fs_bytes(&mut self, len: usize) {
356        self.counters.fs_bytes = self.counters.fs_bytes.saturating_add(len as u64);
357    }
358
359    fn add_tcp_bytes(&mut self, len: usize) {
360        self.counters.tcp_bytes = self.counters.tcp_bytes.saturating_add(len as u64);
361    }
362
363    fn touch(&mut self) {
364        self.activity_seq = self.activity_seq.saturating_add(1);
365    }
366}
367
368impl BulkInputState {
369    fn new(file: File) -> AgentdResult<Self> {
370        Ok(Self {
371            port: AsyncFd::new(file)?,
372            read_buf: vec![0u8; BULK_SERIAL_READ_BUF_SIZE],
373            input: BytesMut::new(),
374        })
375    }
376
377    /// Read at most once, then return one bounded batch already available to the actor.
378    async fn read_turn(
379        &mut self,
380        read_budget: &mut AgentReadBudget,
381    ) -> AgentdResult<Vec<IncarnatedBulkFrame>> {
382        let buffered = self.drain_turn()?;
383        if !buffered.is_empty() {
384            return Ok(buffered);
385        }
386
387        let mut guard = self.port.readable().await?;
388        match guard
389            .try_io(|inner| read_budget.read_fd(inner.get_ref().as_raw_fd(), &mut self.read_buf))
390        {
391            Ok(Ok(0)) => {
392                return Err(AgentdError::ExecSession(
393                    "dedicated bulk port closed".into(),
394                ));
395            }
396            Ok(Ok(read)) => {
397                self.input.extend_from_slice(&self.read_buf[..read]);
398                if self.input.len() > MAX_BULK_INPUT_BUF_SIZE {
399                    return Err(AgentdError::ExecSession(
400                        "dedicated bulk input exceeded maximum frame buffer".into(),
401                    ));
402                }
403            }
404            Ok(Err(error)) if error.kind() == std::io::ErrorKind::Interrupted => {}
405            Ok(Err(error)) => return Err(error.into()),
406            Err(_would_block) => {}
407        }
408        self.drain_turn()
409    }
410
411    fn drain_turn(&mut self) -> AgentdResult<Vec<IncarnatedBulkFrame>> {
412        let mut frames = Vec::new();
413        let mut fs_bytes = 0usize;
414        let mut tcp_bytes = 0usize;
415        while frames.len() < BULK_READER_MAX_RECORDS_PER_TURN
416            && fs_bytes < MAX_BULK_RECORD_PAYLOAD as usize
417            && tcp_bytes < BULK_READER_MAX_BYTES_PER_TURN
418        {
419            let Some(frame) = try_decode_incarnated_bulk_from_bytes(&mut self.input)
420                .map_err(|error| AgentdError::ExecSession(format!("decode agent-bulk: {error}")))?
421            else {
422                break;
423            };
424            match frame.record.kind {
425                BulkKind::Filesystem => {
426                    fs_bytes = fs_bytes.saturating_add(frame.record.payload.len());
427                }
428                BulkKind::Tcp => {
429                    tcp_bytes = tcp_bytes.saturating_add(frame.record.payload.len());
430                }
431            }
432            frames.push(frame);
433        }
434        Ok(frames)
435    }
436}
437
438impl AgentReadBudget {
439    fn read_fd(&mut self, fd: i32, buf: &mut [u8]) -> std::io::Result<usize> {
440        let result = read_from_fd(fd, buf);
441        self.record_read(result.as_ref().copied().unwrap_or(0));
442        result
443    }
444
445    fn record_read(&mut self, bytes: usize) {
446        self.calls = self.calls.saturating_add(1);
447        self.bytes = self.bytes.saturating_add(bytes);
448    }
449
450    fn exhausted(&self) -> bool {
451        self.bytes >= AGENT_READ_QUANTUM_BYTES || self.calls >= AGENT_READ_QUANTUM_CALLS
452    }
453
454    /// Call outside the competing select futures, after decoded records have an owning queue.
455    /// AsyncFd::readable can stay immediately ready without spending Tokio's cooperative budget;
456    /// returning to select alone does not let the driver discover a writable PTY or new control IO.
457    async fn yield_if_exhausted(&mut self) {
458        if self.exhausted() {
459            tokio::task::yield_now().await;
460            *self = Self::default();
461        }
462    }
463}
464
465//--------------------------------------------------------------------------------------------------
466// Functions
467//--------------------------------------------------------------------------------------------------
468
469/// Runs the main agent loop.
470///
471/// Reuses the already-open virtio serial port, sends `core.ready` with boot timing data,
472/// then enters the main select loop handling serial I/O, process output, and heartbeat.
473///
474/// - `boot_time_ns`: `CLOCK_BOOTTIME` at `main()` start (kernel boot duration).
475/// - `init_time_ns`: nanoseconds spent in `init::init()`.
476pub async fn run(
477    boot_time_ns: u64,
478    init_time_ns: u64,
479    config: &AgentdConfig,
480    port_file: File,
481    boot_console: BootConsoleState,
482    bulk_port: Option<BoundBulkPort>,
483) -> AgentdResult<()> {
484    let process_manager = ProcessManager::get()?;
485    let mut process_manager_failure = process_manager.subscribe_failure()?;
486
487    // Set non-blocking for async I/O. Early boot handshakes use the same fd
488    // in blocking mode before it is moved into the async loop.
489    let port_fd = port_file.as_raw_fd();
490    set_nonblocking(port_fd)?;
491
492    // A single AsyncFd tracks both readable and writable readiness.
493    let async_port = AsyncFd::new(port_file)?;
494
495    // A combined port carries full raw records; a dual-port control stream does not. Size the
496    // allocation from the selected topology so combined throughput does not pay four reads per
497    // default record while dual mode retains the small control-only working set.
498    let mut read_buf = vec![0u8; primary_serial_read_buf_size(bulk_port.is_some())];
499    // The final bootstrap read may also contain the first ordinary frame. Preserve that tail while
500    // moving into the cursor-based buffer used by the bounded main-loop decoder.
501    let mut serial_in_buf = BytesMut::from(boot_console.input.as_slice());
502    let mut serial_out_buf = Vec::new();
503
504    let mut state = AgentState::default();
505    let mut workload = if handoff::is_pid_1() {
506        WorkloadLatch::initialize()
507    } else {
508        WorkloadLatch::unavailable(
509            "PID 1 handoff workloads are not wholly owned by agentd's cgroup",
510        )
511    };
512    if let Some(reason) = workload.unavailable_reason() {
513        eprintln!("checkpoint workload freezer unavailable: {reason}");
514    }
515
516    // Channel for session output events.
517    let (mut session_tx, mut session_rx, bulk_session_rx, bulk_command_rx) =
518        SessionOutputSender::split_channel();
519
520    // Heartbeat/activity state.
521    let mut activity = ActivityTracker::new();
522    let (heartbeat_tx, heartbeat_rx) = watch::channel(heartbeat_snapshot(&state, &activity));
523    // The liveness pulse runs on a dedicated OS thread, NOT a Tokio task. On the
524    // single-threaded agent runtime a flood of exec output can monopolize the
525    // executor and starve a heartbeat *task*, freezing the pulse even though the
526    // agent is alive — which makes the host wrongly declare it unresponsive and
527    // kill the sandbox. A plain OS thread is scheduled by the guest kernel
528    // independently of the async runtime, so the pulse keeps ticking under load.
529    let heartbeat_shutdown = Arc::new(AtomicBool::new(false));
530    let heartbeat_control = Arc::new(heartbeat::HeartbeatControl::default());
531    let heartbeat_thread = spawn_heartbeat_thread(
532        heartbeat_rx,
533        Arc::clone(&heartbeat_shutdown),
534        Arc::clone(&heartbeat_control),
535    );
536
537    // Send core.ready with boot timing data.
538    let ready_time_ns = clock::boottime_ns();
539    let ready_msg = Message::with_payload(
540        MessageType::Ready,
541        0,
542        &Ready {
543            boot_time_ns,
544            init_time_ns,
545            ready_time_ns,
546            agent_version: env!("CARGO_PKG_VERSION").to_string(),
547            bulk_transport: bulk_port
548                .as_ref()
549                .map(|port| BulkTransportReady::dual_port_v1(port.connection_id)),
550            relay_lease: Some(RelayLeaseReady::range_lease_v1()),
551            // Shared arenas are negotiated only on the local SDK-to-runtime hop. Agentd speaks to
552            // the runtime over the guest consoles, so the runtime injects this capability later.
553            local_transport: None,
554            workload_transport_barrier_version: Some(WORKLOAD_TRANSPORT_BARRIER_VERSION),
555        },
556    )
557    .map_err(|e| AgentdError::ExecSession(format!("encode ready: {e}")))?;
558    codec::encode_to_buf(&ready_msg, &mut serial_out_buf)
559        .map_err(|e| AgentdError::ExecSession(format!("encode ready frame: {e}")))?;
560    flush_write_buf(&async_port, &mut serial_out_buf).await?;
561
562    let (mut combined_bulk_rx, mut bulk_input, mut bulk_failure_rx, mut bulk_activity_rx) =
563        match bulk_port {
564            Some(port) => {
565                let input = BulkInputState::new(port.file.try_clone()?)?;
566                let (failure_rx, activity_rx) =
567                    spawn_bulk_writer_task(port.file, bulk_session_rx, bulk_command_rx);
568                (None, Some(input), Some(failure_rx), Some(activity_rx))
569            }
570            None => {
571                session_tx.disable_bulk_scheduler();
572                (Some(bulk_session_rx), None, None, None)
573            }
574        };
575    let dual_port_active = bulk_input.is_some();
576    let mut pending_bulk_inputs = VecDeque::<PendingBulkInput>::new();
577    let mut bulk_input_bytes_since_snapshot = 0usize;
578    let mut last_bulk_input_snapshot = Instant::now();
579    let input_refunds = state.input_window.clone();
580    let mut credit_deadline = None;
581    let mut last_input_credit = state.input_window.credit()?;
582    let mut read_budget = AgentReadBudget::default();
583
584    // Main loop.
585    'agent: loop {
586        // All consumed bytes are now in persistent input buffers or destination-owned records.
587        // Never yield after decoding inside read_turn: select cancellation could drop its frames.
588        read_budget.yield_if_exhausted().await;
589        if state.resume_output_after_flush {
590            // The private Thawed reply crosses the primary lane before either ordinary writer
591            // becomes eligible again. A partial old record is never abandoned by this release.
592            flush_write_buf(&async_port, &mut serial_out_buf).await?;
593            session_tx
594                .resume_bulk_output()
595                .await
596                .map_err(|error| AgentdError::ExecSession(error.into()))?;
597            state.resume_output_after_flush = false;
598            state.output_parked = false;
599        }
600        if state.pending_freeze.as_ref().is_some_and(|message| {
601            message.payload::<WorkloadFreeze>().is_ok_and(|request| {
602                let position = state.input_window.position();
603                position.bulk_bytes >= request.host_input.bulk_bytes
604                    && position.bulk_frames >= request.host_input.bulk_frames
605            })
606        }) {
607            let message = state.pending_freeze.take().expect("ready pending freeze");
608            handle_message_with_charge(
609                message,
610                &mut state,
611                &mut activity,
612                &mut session_tx,
613                &mut serial_out_buf,
614                config,
615                &mut workload,
616                &heartbeat_control,
617                None,
618            )
619            .await?;
620            flush_write_buf(&async_port, &mut serial_out_buf).await?;
621        }
622        if !state.output_parked {
623            let credit = state.input_window.credit()?;
624            if input_credit_update_due(&last_input_credit, &credit) {
625                encode_input_credit(credit, &mut serial_out_buf)?;
626                flush_write_buf(&async_port, &mut serial_out_buf).await?;
627                last_input_credit = credit;
628                credit_deadline = None;
629            }
630        }
631        // A control-lane disconnect cancels the pending acquire future at the select boundary.
632        // Remove its stale record before rebuilding that future against the global budget.
633        pending_bulk_inputs.retain(|pending| {
634            client_incarnation_for_id(&state, pending.frame.record.id)
635                == Some(pending.frame.incarnation)
636        });
637        let pending_bulk_admission = pending_bulk_inputs.front().map(|pending| {
638            (
639                Arc::clone(&pending.budget),
640                pending.frame.record.payload.len(),
641            )
642        });
643        let has_pending_bulk_admission = pending_bulk_admission.is_some();
644        tokio::select! {
645            _ = input_refunds.refunded(), if !state.output_parked && credit_deadline.is_none() => {
646                if state.input_window.credit()? != last_input_credit {
647                    credit_deadline = Some(time::Instant::now() + Duration::from_millis(5));
648                }
649            }
650
651            _ = wait_input_credit_deadline(credit_deadline), if !state.output_parked => {
652                credit_deadline = None;
653                let credit = state.input_window.credit()?;
654                if credit != last_input_credit {
655                    encode_input_credit(credit, &mut serial_out_buf)?;
656                    flush_write_buf(&async_port, &mut serial_out_buf).await?;
657                    last_input_credit = credit;
658                }
659            }
660
661            (id, inherited, result) = std::future::poll_fn(|cx| poll_pending_stdin(&mut state, cx)),
662                if !state.output_parked && !workload.is_frozen() => {
663                if !inherited && let Err(error) = result {
664                    encode_stdin_error(id, &AgentdError::Io(error), &mut serial_out_buf)?;
665                    flush_write_buf(&async_port, &mut serial_out_buf).await?;
666                }
667            }
668            failure = process_manager_failure.changed() => {
669                let error = match failure {
670                    Ok(()) => process_manager_failure
671                        .borrow()
672                        .clone()
673                        .unwrap_or_else(|| "process manager stopped without an error".to_string()),
674                    Err(error) => format!("process manager failure channel closed: {error}"),
675                };
676                return Err(AgentdError::ExecSession(error));
677            }
678
679            Some(error) = recv_optional(&mut bulk_failure_rx) => {
680                if state.output_parked {
681                    return Err(AgentdError::ExecSession(format!(
682                        "dedicated bulk transport failed while frozen: {error}"
683                    )));
684                }
685                cancel_all_bulk_correlations(
686                    &mut state,
687                    &session_tx,
688                    &mut serial_out_buf,
689                    "dedicated bulk transport failed",
690                )?;
691                let _ = time::timeout(
692                    BULK_FAILURE_FLUSH_TIMEOUT,
693                    flush_write_buf(&async_port, &mut serial_out_buf),
694                ).await;
695                return Err(AgentdError::ExecSession(format!(
696                    "dedicated bulk transport failed: {error}"
697                )));
698            }
699
700            Some(output_activity) = recv_optional(&mut bulk_activity_rx) => {
701                apply_raw_activity(output_activity, &mut activity);
702                publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
703            }
704
705            permit = acquire_bulk_input_permit(pending_bulk_admission), if has_pending_bulk_admission => {
706                let pending = pending_bulk_inputs
707                    .pop_front()
708                    .expect("guarded pending bulk input exists");
709                let permit = match permit {
710                    Ok(permit) => permit,
711                    Err(error) => {
712                        // The process-wide budget is never normally closed. Preserve the
713                        // ownership check so a future epoch-scoped budget can cancel stale waits
714                        // without allowing old bytes into a recycled correlation range.
715                        if validate_bulk_client_incarnation(
716                            &state,
717                            pending.frame.record.id,
718                            pending.frame.incarnation,
719                        )? {
720                            return Err(error);
721                        }
722                        continue;
723                    }
724                };
725                if !validate_bulk_client_incarnation(
726                    &state,
727                    pending.frame.record.id,
728                    pending.frame.incarnation,
729                )? {
730                    continue;
731                }
732                let payload_len = pending.frame.record.payload.len();
733                let bulk_session_tx = session_tx.with_incarnation(Some(pending.frame.incarnation));
734                handle_bulk_record(
735                    AdmittedBulkRecord {
736                        record: pending.frame.record,
737                        _permit: permit,
738                        _transport_charge: Some(pending.charge),
739                    },
740                    &mut state,
741                    &mut activity,
742                    &mut serial_out_buf,
743                    &bulk_session_tx,
744                ).await?;
745                bulk_input_bytes_since_snapshot =
746                    bulk_input_bytes_since_snapshot.saturating_add(payload_len);
747                if bulk_input_bytes_since_snapshot >= BULK_ACTIVITY_PUBLISH_BYTES
748                    || last_bulk_input_snapshot.elapsed() >= BULK_ACTIVITY_PUBLISH_INTERVAL
749                {
750                    publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
751                    bulk_input_bytes_since_snapshot = 0;
752                    last_bulk_input_snapshot = Instant::now();
753                }
754                if !serial_out_buf.is_empty() {
755                    flush_write_buf(&async_port, &mut serial_out_buf).await?;
756                }
757            }
758
759            Some(envelope) = recv_optional(&mut combined_bulk_rx), if !state.output_parked => {
760                if discard_inherited_output(&mut state, &envelope, session_tx.generation()) {
761                    continue;
762                }
763                if envelope.incarnation.is_some()
764                    && client_incarnation_for_id(&state, envelope.id) != envelope.incarnation
765                {
766                    continue;
767                }
768                let SessionOutput::Bulk(output) = envelope.output else {
769                    return Err(AgentdError::ExecSession(
770                        "non-bulk event entered combined bulk queue".into(),
771                    ));
772                };
773                apply_raw_activity(output.activity, &mut activity);
774                if !serial_out_buf.is_empty() {
775                    flush_write_buf(&async_port, &mut serial_out_buf).await?;
776                }
777                write_bulk_record_async_fd(&async_port, &output.record).await?;
778                publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
779            }
780
781            turn = async {
782                bulk_input
783                    .as_mut()
784                    .expect("guarded dedicated bulk input")
785                    .read_turn(&mut read_budget)
786                    .await
787            }, if bulk_input.is_some() && pending_bulk_inputs.is_empty() => {
788                let frames = match turn {
789                    Ok(frames) => frames,
790                    Err(error) => {
791                        cancel_all_bulk_correlations(
792                            &mut state,
793                            &session_tx,
794                            &mut serial_out_buf,
795                            "dedicated bulk transport failed",
796                        )?;
797                        let _ = time::timeout(
798                            BULK_FAILURE_FLUSH_TIMEOUT,
799                            flush_write_buf(&async_port, &mut serial_out_buf),
800                        ).await;
801                        return Err(AgentdError::ExecSession(format!(
802                            "dedicated bulk transport failed: {error}"
803                        )));
804                    }
805                };
806                let mut capacity_deferred = false;
807                if state.output_parked && !frames.is_empty() {
808                    return Err(AgentdError::ExecSession("bulk input crossed the frozen transport cut".into()));
809                }
810                for frame in frames {
811                    let charge = state.input_window.admit(
812                        InputLane::Bulk, bulk_wire_bytes(&frame.record, true),
813                    )?;
814                    if !validate_bulk_client_incarnation(
815                        &state,
816                        frame.record.id,
817                        frame.incarnation,
818                    )? {
819                        // The range has been disconnected or recycled since these bytes entered
820                        // the other physical lane. They must not affect the new owner's ID.
821                        continue;
822                    }
823                    let budget = Arc::clone(&state.bulk_input_budget);
824                    if capacity_deferred {
825                        pending_bulk_inputs.push_back(PendingBulkInput { frame, budget, charge });
826                        continue;
827                    }
828                    let payload_len = frame.record.payload.len();
829                    let charged = u32::try_from(payload_len).map_err(|_| {
830                        AgentdError::ExecSession("bulk input payload budget overflow".into())
831                    })?;
832                    let permit = match Arc::clone(&budget).try_acquire_many_owned(charged) {
833                        Ok(permit) => permit,
834                        Err(_) if !budget.is_closed() => {
835                            capacity_deferred = true;
836                            pending_bulk_inputs.push_back(PendingBulkInput { frame, budget, charge });
837                            continue;
838                        }
839                        Err(error) => {
840                            return Err(AgentdError::ExecSession(format!(
841                                "bulk input budget closed unexpectedly: {error}"
842                            )));
843                        }
844                    };
845                    let bulk_session_tx = session_tx.with_incarnation(Some(frame.incarnation));
846                    handle_bulk_record(
847                        AdmittedBulkRecord {
848                            record: frame.record,
849                            _permit: permit,
850                            _transport_charge: Some(charge),
851                        },
852                        &mut state,
853                        &mut activity,
854                        &mut serial_out_buf,
855                        &bulk_session_tx,
856                    ).await?;
857                    bulk_input_bytes_since_snapshot =
858                        bulk_input_bytes_since_snapshot.saturating_add(payload_len);
859                }
860                if bulk_input_bytes_since_snapshot >= BULK_ACTIVITY_PUBLISH_BYTES
861                    || last_bulk_input_snapshot.elapsed() >= BULK_ACTIVITY_PUBLISH_INTERVAL
862                {
863                    publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
864                    bulk_input_bytes_since_snapshot = 0;
865                    last_bulk_input_snapshot = Instant::now();
866                }
867                if !serial_out_buf.is_empty() {
868                    flush_write_buf(&async_port, &mut serial_out_buf).await?;
869                }
870            }
871
872            // Read from serial port.
873            result = async_port.readable() => {
874                let Ok(mut guard) = result else {
875                    break;
876                };
877                let mut combined_bulk_records = 0usize;
878                let mut combined_bulk_payload_bytes = 0usize;
879                let mut combined_turn_exhausted = false;
880
881                loop {
882                    match guard.try_io(|inner| {
883                        read_budget.read_fd(inner.get_ref().as_raw_fd(), &mut read_buf)
884                    }) {
885                        Ok(Ok(0)) => {
886                            // EOF on serial — host disconnected.
887                            if !handoff::is_pid_1() {
888                                guard.clear_ready();
889                                drop(guard);
890                                time::sleep(Duration::from_millis(100)).await;
891                                break;
892                            }
893                            break 'agent;
894                        }
895                        Ok(Ok(n)) => {
896                            serial_in_buf.extend_from_slice(&read_buf[..n]);
897
898                            // Guard against unbounded buffer growth.
899                            if serial_in_buf.len() > MAX_INPUT_BUF_SIZE {
900                                return Err(AgentdError::ExecSession(
901                                    "serial input buffer exceeded maximum size".into(),
902                                ));
903                            }
904
905                            // Try to parse complete frames. Recoverable
906                            // message-level failures are reported on the same
907                            // correlation ID with `core.error`; unrecoverable
908                            // frame-level failures still close the agent loop.
909                            loop {
910                                let input_before = serial_in_buf.len();
911                                if let Some(connected) =
912                                        try_decode_relay_client_connected_from_bytes(&mut serial_in_buf)
913                                            .map_err(|e| AgentdError::ExecSession(format!(
914                                                "decode relay client lease: {e}"
915                                            )))?
916                                {
917                                    if state.output_parked {
918                                        return Err(AgentdError::ExecSession("relay lease crossed the frozen transport cut".into()));
919                                    }
920                                    let _charge = state.input_window.admit(
921                                        InputLane::Control, input_before - serial_in_buf.len(),
922                                    )?;
923                                    establish_relay_client(&mut state, connected)?;
924                                    continue;
925                                }
926                                let Some(frame) = codec::try_decode_frame_from_bytes(&mut serial_in_buf)
927                                    .map_err(|e| AgentdError::ExecSession(format!("decode frame: {e}")))?
928                                else {
929                                    break;
930                                };
931                                let wire_bytes = input_before - serial_in_buf.len();
932                                let DecodedFrame::Control(msg) = frame else {
933                                    let DecodedFrame::Bulk(record) = frame else {
934                                        unreachable!();
935                                    };
936                                    if dual_port_active {
937                                        return Err(AgentdError::ExecSession(
938                                            "raw bulk record arrived on the bound control port"
939                                                .into(),
940                                        ));
941                                    }
942                                    if state.output_parked {
943                                        return Err(AgentdError::ExecSession("raw input crossed the frozen transport cut".into()));
944                                    }
945                                    let bulk_session_tx = session_tx.with_incarnation(
946                                        client_incarnation_for_id(&state, record.id),
947                                    );
948                                    let payload_len = record.payload.len();
949                                    let charge = state.input_window.admit(InputLane::Bulk, wire_bytes)?;
950                                    let budget = Arc::clone(&state.bulk_input_budget);
951                                    let permit = acquire_bulk_input_permit(Some((
952                                        budget,
953                                        payload_len,
954                                    ))).await?;
955                                    handle_bulk_record(
956                                        AdmittedBulkRecord {
957                                            record,
958                                            _permit: permit,
959                                            _transport_charge: Some(charge),
960                                        },
961                                        &mut state,
962                                        &mut activity,
963                                        &mut serial_out_buf,
964                                        &bulk_session_tx,
965                                    ).await?;
966                                    combined_bulk_records = combined_bulk_records.saturating_add(1);
967                                    combined_bulk_payload_bytes =
968                                        combined_bulk_payload_bytes.saturating_add(payload_len);
969                                    bulk_input_bytes_since_snapshot =
970                                        bulk_input_bytes_since_snapshot.saturating_add(payload_len);
971                                    if bulk_input_bytes_since_snapshot >= BULK_ACTIVITY_PUBLISH_BYTES
972                                        || last_bulk_input_snapshot.elapsed()
973                                            >= BULK_ACTIVITY_PUBLISH_INTERVAL
974                                    {
975                                        publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
976                                        bulk_input_bytes_since_snapshot = 0;
977                                        last_bulk_input_snapshot = Instant::now();
978                                    }
979                                    if bulk_reader_turn_exhausted(
980                                        combined_bulk_records,
981                                        combined_bulk_payload_bytes,
982                                    ) {
983                                        // A task-level yield is insufficient here: this branch is
984                                        // still inside the serial drain loop, so agent session
985                                        // output (including bulk credit) cannot be selected. Finish
986                                        // decoding the bytes already read, then return to the outer
987                                        // select before reading another batch from the port.
988                                        combined_turn_exhausted = true;
989                                    }
990                                    continue;
991                                };
992                                let charge = if private_lifecycle_message(&msg) {
993                                    None
994                                } else {
995                                    if state.output_parked {
996                                        return Err(AgentdError::ExecSession("ordinary input crossed the frozen transport cut".into()));
997                                    }
998                                    let lane = if msg.t.uses_workload_data_credit() {
999                                        InputLane::Bulk
1000                                    } else {
1001                                        InputLane::Control
1002                                    };
1003                                    Some(state.input_window.admit(lane, wire_bytes)?)
1004                                };
1005                                if msg.flags != msg.t.flags() {
1006                                    let out_before = serial_out_buf.len();
1007                                    encode_core_error_if_supported(
1008                                        &msg,
1009                                        msg.id,
1010                                        CoreErrorKind::InvalidFlags,
1011                                        format!(
1012                                            "invalid flags for {}: got {}, expected {}",
1013                                            msg.t.as_str(),
1014                                            msg.flags,
1015                                            msg.t.flags()
1016                                        ),
1017                                        Some(msg.t.as_str().to_string()),
1018                                        &mut serial_out_buf,
1019                                    )?;
1020                                    record_encoded_guest_messages(
1021                                        &serial_out_buf,
1022                                        out_before,
1023                                        &mut activity,
1024                                    );
1025                                    publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
1026                                    continue;
1027                                }
1028
1029                                if message_refreshes_idle_timer(&msg.t) {
1030                                    activity.record_host_message();
1031                                    publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
1032                                }
1033
1034                                let out_before = serial_out_buf.len();
1035                                handle_message_with_charge(
1036                                    msg,
1037                                    &mut state,
1038                                    &mut activity,
1039                                    &mut session_tx,
1040                                    &mut serial_out_buf,
1041                                    config,
1042                                    &mut workload,
1043                                    &heartbeat_control,
1044                                    charge,
1045                                ).await?;
1046                                record_encoded_guest_messages(
1047                                    &serial_out_buf,
1048                                    out_before,
1049                                    &mut activity,
1050                                );
1051                                publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
1052                            }
1053
1054                            // Flush any outgoing messages.
1055                            if !serial_out_buf.is_empty() {
1056                                flush_write_buf(&async_port, &mut serial_out_buf).await?;
1057                            }
1058                            // Thawed has now crossed the wire. A host can immediately resume
1059                            // ordinary input, so release our parked writers at the outer-loop
1060                            // boundary before draining another readable batch.
1061                            if combined_turn_exhausted
1062                                || state.resume_output_after_flush
1063                                || read_budget.exhausted()
1064                            {
1065                                break;
1066                            }
1067                        }
1068                        Ok(Err(e)) if e.kind() == std::io::ErrorKind::Interrupted => {
1069                            if read_budget.exhausted() {
1070                                break;
1071                            }
1072                            continue;
1073                        }
1074                        Ok(Err(_)) if !handoff::is_pid_1() => {
1075                            guard.clear_ready();
1076                            drop(guard);
1077                            time::sleep(Duration::from_millis(100)).await;
1078                            break;
1079                        }
1080                        Ok(Err(e)) => return Err(e.into()),
1081                        Err(_would_block) => break,
1082                    }
1083                }
1084            }
1085
1086            // Receive output events from session reader tasks.
1087            Some(envelope) = session_rx.recv(), if !state.output_parked => {
1088                if discard_inherited_output(&mut state, &envelope, session_tx.generation()) {
1089                    continue;
1090                }
1091                if envelope.incarnation.is_some()
1092                    && client_incarnation_for_id(&state, envelope.id) != envelope.incarnation
1093                {
1094                    // Background output carries the owner captured at session creation. Dropping
1095                    // stale output here protects control frames as well as dedicated raw records.
1096                    continue;
1097                }
1098                let id = envelope.id;
1099                match envelope.output {
1100                    SessionOutput::Stdout(data) => {
1101                        let len = data.len();
1102                        let msg = Message::with_payload(MessageType::ExecStdout, id, &ExecStdout { data })
1103                            .map_err(|e| AgentdError::ExecSession(format!("encode stdout: {e}")))?;
1104                        codec::encode_to_buf(&msg, &mut serial_out_buf)
1105                            .map_err(|e| AgentdError::ExecSession(format!("encode stdout frame: {e}")))?;
1106                        activity.record_guest_message();
1107                        activity.add_exec_output_bytes(len);
1108                    }
1109                    SessionOutput::Stderr(data) => {
1110                        let len = data.len();
1111                        let msg = Message::with_payload(MessageType::ExecStderr, id, &ExecStderr { data })
1112                            .map_err(|e| AgentdError::ExecSession(format!("encode stderr: {e}")))?;
1113                        codec::encode_to_buf(&msg, &mut serial_out_buf)
1114                            .map_err(|e| AgentdError::ExecSession(format!("encode stderr frame: {e}")))?;
1115                        activity.record_guest_message();
1116                        activity.add_exec_output_bytes(len);
1117                    }
1118                    SessionOutput::Exited(code) => {
1119                        let msg = Message::with_payload(MessageType::ExecExited, id, &ExecExited { code })
1120                            .map_err(|e| AgentdError::ExecSession(format!("encode exited: {e}")))?;
1121                        codec::encode_to_buf(&msg, &mut serial_out_buf)
1122                            .map_err(|e| AgentdError::ExecSession(format!("encode exited frame: {e}")))?;
1123                        state.sessions.remove(&id);
1124                        activity.record_guest_message();
1125                    }
1126                    SessionOutput::Raw(output) => {
1127                        apply_raw_activity(output.activity, &mut activity);
1128                        let completion = output.completion;
1129                        complete_raw_session(
1130                            id,
1131                            completion,
1132                            &mut state.read_sessions,
1133                            &mut state.tcp_sessions,
1134                        );
1135                        if matches!(completion, Some(RawSessionCompletion::FsWrite))
1136                            && let Some(worker) = state.bulk_write_workers.remove(&id)
1137                        {
1138                            worker.task.abort();
1139                        }
1140                        if completion.is_some() {
1141                            clear_bulk_receive_state(&mut state, id);
1142                        }
1143                        // The producer already owns an encoded frame. Write from that allocation
1144                        // directly so multi-megabyte FS/TCP frames are not copied into a second
1145                        // serial staging buffer.
1146                        if !serial_out_buf.is_empty() {
1147                            flush_write_buf(&async_port, &mut serial_out_buf).await?;
1148                        }
1149                        if !output.frame.is_empty() {
1150                            write_all_async_fd(&async_port, &output.frame).await?;
1151                        }
1152                    }
1153                    SessionOutput::Bulk(output) => {
1154                        apply_raw_activity(output.activity, &mut activity);
1155                        if !serial_out_buf.is_empty() {
1156                            flush_write_buf(&async_port, &mut serial_out_buf).await?;
1157                        }
1158                        write_bulk_record_async_fd(&async_port, &output.record).await?;
1159                    }
1160                }
1161                publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
1162
1163                if !serial_out_buf.is_empty() {
1164                    flush_write_buf(&async_port, &mut serial_out_buf).await?;
1165                }
1166            }
1167        }
1168    }
1169
1170    heartbeat_shutdown.store(true, Ordering::Relaxed);
1171    let _ = heartbeat_thread.join();
1172
1173    Ok(())
1174}
1175
1176/// Start the dedicated-lane output actor; input stays fused into the main control actor.
1177fn spawn_bulk_writer_task(
1178    writer_file: File,
1179    output_rx: tokio::sync::mpsc::Receiver<SessionOutputEnvelope>,
1180    command_rx: tokio::sync::mpsc::Receiver<BulkOutputCommand>,
1181) -> (
1182    tokio::sync::mpsc::Receiver<String>,
1183    tokio::sync::mpsc::Receiver<RawActivity>,
1184) {
1185    let (failure_tx, failure_rx) = tokio::sync::mpsc::channel(1);
1186    let (activity_tx, activity_rx) = tokio::sync::mpsc::channel(128);
1187
1188    tokio::spawn(async move {
1189        if let Err(error) = bulk_writer_task(writer_file, output_rx, command_rx, activity_tx).await
1190        {
1191            let _ = failure_tx.send(error.to_string()).await;
1192        }
1193    });
1194
1195    (failure_rx, activity_rx)
1196}
1197
1198/// Select the primary port's read working set from whether bulk has its own physical lane.
1199fn primary_serial_read_buf_size(has_dedicated_bulk_port: bool) -> usize {
1200    if has_dedicated_bulk_port {
1201        CONTROL_SERIAL_READ_BUF_SIZE
1202    } else {
1203        COMBINED_SERIAL_READ_BUF_SIZE
1204    }
1205}
1206
1207/// Bound one raw-input scheduling turn so a single-thread guest can run the destination worker.
1208fn bulk_reader_turn_exhausted(records: usize, payload_bytes: usize) -> bool {
1209    records >= BULK_READER_MAX_RECORDS_PER_TURN || payload_bytes >= BULK_READER_MAX_BYTES_PER_TURN
1210}
1211
1212/// Receive from an optional channel without making combined mode spin.
1213async fn recv_optional<T>(receiver: &mut Option<tokio::sync::mpsc::Receiver<T>>) -> Option<T> {
1214    let Some(active) = receiver.as_mut() else {
1215        return std::future::pending().await;
1216    };
1217    match active.recv().await {
1218        Some(value) => Some(value),
1219        None => {
1220            // Disable a closed optional branch. Returning `None` immediately on every select-loop
1221            // iteration would turn a failed worker into a one-vCPU busy loop while its companion
1222            // failure channel is still trying to report the real error.
1223            *receiver = None;
1224            std::future::pending().await
1225        }
1226    }
1227}
1228
1229/// Schedule guest-to-host records fairly by correlation before writing the bulk port.
1230async fn bulk_writer_task(
1231    file: File,
1232    mut output_rx: tokio::sync::mpsc::Receiver<SessionOutputEnvelope>,
1233    mut command_rx: tokio::sync::mpsc::Receiver<BulkOutputCommand>,
1234    activity_tx: tokio::sync::mpsc::Sender<RawActivity>,
1235) -> AgentdResult<()> {
1236    let async_port = AsyncFd::new(file)?;
1237    let mut generation = 0;
1238    let mut flows = HashMap::<(ClientIncarnation, u32), BulkWriteFlow>::new();
1239    let mut active = VecDeque::<(ClientIncarnation, u32)>::new();
1240    let mut retired = HashMap::<ClientIncarnation, Vec<u64>>::new();
1241    let mut retiring_incarnations = HashSet::<ClientIncarnation>::new();
1242    let mut pending_activity = RawActivity::default();
1243    let mut transport = BulkOutputPosition::default();
1244
1245    loop {
1246        let mut cleanups = Vec::new();
1247        tokio::select! {
1248            biased;
1249            Some(command) = command_rx.recv() => {
1250                cleanups.push(apply_bulk_output_command(
1251                    command,
1252                    &mut generation,
1253                    &mut flows,
1254                    &mut active,
1255                    &mut retired,
1256                    &mut retiring_incarnations,
1257                    &mut transport,
1258                )?);
1259            }
1260            Some(envelope) = output_rx.recv(), if !transport.parked => {
1261                enqueue_bulk_output(
1262                    envelope,
1263                    generation,
1264                    &mut flows,
1265                    &mut active,
1266                    &retired,
1267                    &retiring_incarnations,
1268                )?;
1269            }
1270            else => break,
1271        }
1272        while let Ok(command) = command_rx.try_recv() {
1273            cleanups.push(apply_bulk_output_command(
1274                command,
1275                &mut generation,
1276                &mut flows,
1277                &mut active,
1278                &mut retired,
1279                &mut retiring_incarnations,
1280                &mut transport,
1281            )?);
1282        }
1283        while let Ok(envelope) = output_rx.try_recv() {
1284            enqueue_bulk_output(
1285                envelope,
1286                generation,
1287                &mut flows,
1288                &mut active,
1289                &retired,
1290                &retiring_incarnations,
1291            )?;
1292        }
1293        complete_bulk_output_cleanups(cleanups, &mut retired, &mut retiring_incarnations);
1294
1295        while !transport.parked && !active.is_empty() {
1296            let round_len = active.len();
1297            let quantum = if round_len == 1 {
1298                BULK_SCHEDULER_MAX_BURST
1299            } else {
1300                BULK_SCHEDULER_QUANTUM
1301            };
1302            for _ in 0..round_len {
1303                if transport.parked || active.is_empty() {
1304                    break;
1305                }
1306                let key = active.pop_front().expect("active flow exists");
1307                if let Some(flow) = flows.get_mut(&key) {
1308                    flow.deficit = flow
1309                        .deficit
1310                        .saturating_add(quantum)
1311                        .min(BULK_SCHEDULER_MAX_BURST);
1312                }
1313
1314                let mut burst = 0usize;
1315                loop {
1316                    let next_len = flows
1317                        .get(&key)
1318                        .and_then(|flow| flow.queue.front())
1319                        .map(bulk_envelope_payload_len)
1320                        .transpose()?
1321                        .unwrap_or(0);
1322                    let can_send = flows.get(&key).is_some_and(|flow| {
1323                        next_len != 0
1324                            && next_len <= flow.deficit
1325                            && burst.saturating_add(next_len) <= BULK_SCHEDULER_MAX_BURST
1326                    });
1327                    if !can_send {
1328                        break;
1329                    }
1330
1331                    let envelope = {
1332                        let flow = flows.get_mut(&key).expect("scheduled flow exists");
1333                        let envelope = flow.queue.pop_front().expect("scheduled record exists");
1334                        flow.queued_bytes = flow.queued_bytes.saturating_sub(next_len);
1335                        flow.deficit = flow.deficit.saturating_sub(next_len);
1336                        envelope
1337                    };
1338                    let incarnation = envelope.incarnation.ok_or_else(|| {
1339                        AgentdError::ExecSession(
1340                            "dedicated bulk output is missing client incarnation".into(),
1341                        )
1342                    })?;
1343                    let SessionOutput::Bulk(output) = envelope.output else {
1344                        unreachable!("bulk scheduler accepted a non-bulk event");
1345                    };
1346                    let output_activity = output.activity;
1347                    write_incarnated_bulk_record_async_fd(&async_port, incarnation, &output.record)
1348                        .await?;
1349                    transport.wire_bytes = transport
1350                        .wire_bytes
1351                        .checked_add(bulk_wire_bytes(&output.record, true) as u64)
1352                        .ok_or_else(|| {
1353                            AgentdError::ExecSession("bulk output counter exhausted".into())
1354                        })?;
1355                    pending_activity.guest_messages = pending_activity
1356                        .guest_messages
1357                        .saturating_add(output_activity.guest_messages);
1358                    pending_activity.fs_bytes = pending_activity
1359                        .fs_bytes
1360                        .saturating_add(output_activity.fs_bytes);
1361                    pending_activity.tcp_bytes = pending_activity
1362                        .tcp_bytes
1363                        .saturating_add(output_activity.tcp_bytes);
1364                    if pending_activity
1365                        .fs_bytes
1366                        .saturating_add(pending_activity.tcp_bytes)
1367                        >= BULK_ACTIVITY_PUBLISH_BYTES
1368                    {
1369                        publish_bulk_activity(&activity_tx, &mut pending_activity)?;
1370                    }
1371                    burst = burst.saturating_add(next_len);
1372                    // Lifecycle commands cut only between complete wire records. In particular,
1373                    // never wait for the rest of this flow's burst before acknowledging Park.
1374                    let mut cleanups = Vec::new();
1375                    while let Ok(command) = command_rx.try_recv() {
1376                        cleanups.push(apply_bulk_output_command(
1377                            command,
1378                            &mut generation,
1379                            &mut flows,
1380                            &mut active,
1381                            &mut retired,
1382                            &mut retiring_incarnations,
1383                            &mut transport,
1384                        )?);
1385                    }
1386                    while let Ok(envelope) = output_rx.try_recv() {
1387                        enqueue_bulk_output(
1388                            envelope,
1389                            generation,
1390                            &mut flows,
1391                            &mut active,
1392                            &retired,
1393                            &retiring_incarnations,
1394                        )?;
1395                    }
1396                    complete_bulk_output_cleanups(
1397                        cleanups,
1398                        &mut retired,
1399                        &mut retiring_incarnations,
1400                    );
1401                    if transport.parked {
1402                        break;
1403                    }
1404                }
1405
1406                if flows.get(&key).is_some_and(|flow| flow.queue.is_empty()) {
1407                    flows.remove(&key);
1408                } else if flows.contains_key(&key) && !active.contains(&key) {
1409                    active.push_back(key);
1410                }
1411            }
1412
1413            let mut cleanups = Vec::new();
1414            while let Ok(command) = command_rx.try_recv() {
1415                cleanups.push(apply_bulk_output_command(
1416                    command,
1417                    &mut generation,
1418                    &mut flows,
1419                    &mut active,
1420                    &mut retired,
1421                    &mut retiring_incarnations,
1422                    &mut transport,
1423                )?);
1424            }
1425            while let Ok(envelope) = output_rx.try_recv() {
1426                enqueue_bulk_output(
1427                    envelope,
1428                    generation,
1429                    &mut flows,
1430                    &mut active,
1431                    &retired,
1432                    &retiring_incarnations,
1433                )?;
1434            }
1435            complete_bulk_output_cleanups(cleanups, &mut retired, &mut retiring_incarnations);
1436            if active.is_empty() && pending_activity.guest_messages != 0 {
1437                publish_bulk_activity(&activity_tx, &mut pending_activity)?;
1438            }
1439            tokio::task::yield_now().await;
1440        }
1441    }
1442
1443    Ok(())
1444}
1445
1446/// Activity must never prevent a restore command from reaching the scheduler.
1447fn publish_bulk_activity(
1448    sender: &tokio::sync::mpsc::Sender<RawActivity>,
1449    pending: &mut RawActivity,
1450) -> AgentdResult<()> {
1451    match sender.try_send(std::mem::take(pending)) {
1452        Ok(()) => Ok(()),
1453        Err(tokio::sync::mpsc::error::TrySendError::Full(activity)) => {
1454            *pending = activity;
1455            Ok(())
1456        }
1457        Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => Err(AgentdError::ExecSession(
1458            "dedicated bulk activity consumer stopped".into(),
1459        )),
1460    }
1461}
1462
1463fn apply_bulk_output_command(
1464    command: BulkOutputCommand,
1465    generation: &mut u64,
1466    flows: &mut HashMap<(ClientIncarnation, u32), BulkWriteFlow>,
1467    active: &mut VecDeque<(ClientIncarnation, u32)>,
1468    retired: &mut HashMap<ClientIncarnation, Vec<u64>>,
1469    retiring_incarnations: &mut HashSet<ClientIncarnation>,
1470    transport: &mut BulkOutputPosition,
1471) -> AgentdResult<BulkOutputCleanup> {
1472    match command {
1473        BulkOutputCommand::Park { completion } => {
1474            transport.parked = true;
1475            Ok(BulkOutputCleanup::Park {
1476                completion,
1477                wire_bytes: transport.wire_bytes,
1478            })
1479        }
1480        BulkOutputCommand::Resume { completion } => {
1481            transport.parked = false;
1482            Ok(BulkOutputCleanup::Flow(completion))
1483        }
1484        BulkOutputCommand::Restore {
1485            generation: next,
1486            completion,
1487        } => {
1488            *generation = next;
1489            flows.clear();
1490            active.clear();
1491            retired.clear();
1492            retiring_incarnations.clear();
1493            Ok(BulkOutputCleanup::Flow(completion))
1494        }
1495        BulkOutputCommand::DropFlow {
1496            incarnation,
1497            id,
1498            completion,
1499        } => {
1500            let key = (incarnation, id);
1501            flows.remove(&key);
1502            active.retain(|active_key| *active_key != key);
1503            retire_bulk_output(retired, incarnation, id)?;
1504            Ok(BulkOutputCleanup::Flow(completion))
1505        }
1506        BulkOutputCommand::DropIncarnation {
1507            incarnation,
1508            completion,
1509        } => {
1510            flows.retain(|(owner, _), _| *owner != incarnation);
1511            active.retain(|(owner, _)| *owner != incarnation);
1512            retiring_incarnations.insert(incarnation);
1513            Ok(BulkOutputCleanup::Incarnation {
1514                incarnation,
1515                completion,
1516            })
1517        }
1518    }
1519}
1520
1521/// Complete lifecycle cuts only after the data receiver has been drained through its tombstones.
1522fn complete_bulk_output_cleanups(
1523    cleanups: Vec<BulkOutputCleanup>,
1524    retired: &mut HashMap<ClientIncarnation, Vec<u64>>,
1525    retiring_incarnations: &mut HashSet<ClientIncarnation>,
1526) {
1527    for cleanup in cleanups {
1528        match cleanup {
1529            BulkOutputCleanup::Park {
1530                completion,
1531                wire_bytes,
1532            } => {
1533                let _ = completion.send(wire_bytes);
1534            }
1535            BulkOutputCleanup::Flow(completion) => {
1536                let _ = completion.send(());
1537            }
1538            BulkOutputCleanup::Incarnation {
1539                incarnation,
1540                completion,
1541            } => {
1542                retired.remove(&incarnation);
1543                retiring_incarnations.remove(&incarnation);
1544                let _ = completion.send(());
1545            }
1546        }
1547    }
1548}
1549
1550fn enqueue_bulk_output(
1551    envelope: SessionOutputEnvelope,
1552    generation: u64,
1553    flows: &mut HashMap<(ClientIncarnation, u32), BulkWriteFlow>,
1554    active: &mut VecDeque<(ClientIncarnation, u32)>,
1555    retired: &HashMap<ClientIncarnation, Vec<u64>>,
1556    retiring_incarnations: &HashSet<ClientIncarnation>,
1557) -> AgentdResult<()> {
1558    if envelope.generation != generation {
1559        return Ok(());
1560    }
1561    let id = envelope.id;
1562    let incarnation = envelope.incarnation.ok_or_else(|| {
1563        AgentdError::ExecSession("dedicated bulk output is missing client incarnation".into())
1564    })?;
1565    let payload_len = bulk_envelope_payload_len(&envelope)?;
1566    if id == 0 {
1567        return Err(AgentdError::ExecSession(
1568            "bulk record cannot use correlation ID zero".into(),
1569        ));
1570    }
1571    if payload_len == 0 || payload_len > MAX_BULK_RECORD_PAYLOAD as usize {
1572        return Err(AgentdError::ExecSession(format!(
1573            "bulk output payload {payload_len} exceeds schedulable maximum {MAX_BULK_RECORD_PAYLOAD}"
1574        )));
1575    }
1576
1577    let key = (incarnation, id);
1578    if retiring_incarnations.contains(&incarnation)
1579        || bulk_output_is_retired(retired, incarnation, id)
1580    {
1581        // A producer can race cancellation after the lifecycle command overtakes the data queue.
1582        // Retaining the owner-scoped tombstone makes that late record release its permit here.
1583        return Ok(());
1584    }
1585    if !flows.contains_key(&key) {
1586        let client_flows = flows
1587            .keys()
1588            .filter(|(owner, _)| *owner == incarnation)
1589            .count();
1590        if client_flows >= BULK_SCHEDULER_MAX_FLOWS_PER_CLIENT {
1591            return Err(AgentdError::ExecSession(format!(
1592                "relay client exceeded active bulk-flow limit for correlation {id}"
1593            )));
1594        }
1595        flows.insert(
1596            key,
1597            BulkWriteFlow {
1598                queue: VecDeque::new(),
1599                queued_bytes: 0,
1600                deficit: 0,
1601            },
1602        );
1603        active.push_back(key);
1604    }
1605
1606    let flow = flows.get_mut(&key).expect("newly inserted flow exists");
1607    let queued_bytes = flow
1608        .queued_bytes
1609        .checked_add(payload_len)
1610        .ok_or_else(|| AgentdError::ExecSession("bulk flow byte budget overflow".into()))?;
1611    if queued_bytes > BULK_SCHEDULER_FLOW_CAPACITY {
1612        return Err(AgentdError::ExecSession(format!(
1613            "bulk flow {id} exceeded queued byte budget"
1614        )));
1615    }
1616    flow.queued_bytes = queued_bytes;
1617    flow.queue.push_back(envelope);
1618    Ok(())
1619}
1620
1621fn bulk_envelope_payload_len(envelope: &SessionOutputEnvelope) -> AgentdResult<usize> {
1622    match &envelope.output {
1623        SessionOutput::Bulk(output) => Ok(output.record.payload.len()),
1624        _ => Err(AgentdError::ExecSession(
1625            "non-bulk event entered dedicated bulk scheduler".into(),
1626        )),
1627    }
1628}
1629
1630/// Discover and bind the dedicated bulk port when the internal boot hint requests it.
1631pub fn open_and_bind_bulk_port() -> AgentdResult<Option<BoundBulkPort>> {
1632    let cmdline = std::fs::read_to_string("/proc/cmdline")?;
1633    if !cmdline_requests_dual_port(&cmdline) {
1634        return Ok(None);
1635    }
1636
1637    let deadline = Instant::now() + Duration::from_secs(BULK_BINDING_TIMEOUT_SECS);
1638    let port_path = loop {
1639        match serial::find_serial_port(AGENT_BULK_PORT_NAME) {
1640            Ok(path) => break path,
1641            Err(error) if Instant::now() < deadline => {
1642                let _ = error;
1643                std::thread::sleep(Duration::from_millis(10));
1644            }
1645            Err(error) => return Err(error),
1646        }
1647    };
1648    let file = OpenOptions::new().read(true).write(true).open(port_path)?;
1649    let fd = file.as_raw_fd();
1650    set_nonblocking(fd)?;
1651
1652    let connection_id = random_connection_id()?;
1653    write_all_to_fd(fd, &encode_bulk_hello(connection_id), Some(deadline))?;
1654
1655    let mut ack = [0u8; BULK_BINDING_SIZE];
1656    read_exact_from_fd(fd, &mut ack, deadline, "bulk binding acknowledgement")?;
1657    decode_bulk_ack(&ack, connection_id)
1658        .map_err(|error| AgentdError::ExecSession(format!("bind agent-bulk: {error}")))?;
1659
1660    Ok(Some(BoundBulkPort {
1661        file,
1662        connection_id,
1663    }))
1664}
1665
1666/// Opens the agent virtio-serial port once for early boot handshakes and the agent loop.
1667pub fn open_serial_port() -> AgentdResult<File> {
1668    // Discover serial port.
1669    let port_path = serial::find_serial_port(AGENT_PORT_NAME)?;
1670
1671    // Open the port once with read+write. Virtio-console multiport devices
1672    // only allow a single open; a second open returns EBUSY.
1673    Ok(OpenOptions::new().read(true).write(true).open(&port_path)?)
1674}
1675
1676/// Receive and validate the first host-to-guest bootstrap frame.
1677pub fn receive_bootstrap(port_file: &File) -> AgentdResult<(GuestBootstrap, BootConsoleState)> {
1678    let fd = port_file.as_raw_fd();
1679    set_nonblocking(fd)?;
1680    let deadline = init_ack_deadline();
1681    let mut state = BootConsoleState::default();
1682    let msg = read_boot_message(fd, &mut state, Some(deadline), "guest bootstrap")?;
1683    let bootstrap = decode_bootstrap_message(msg)?;
1684    state.init_failure_ack = bootstrap.init_failure_ack;
1685    Ok((bootstrap, state))
1686}
1687
1688fn decode_bootstrap_message(msg: Message) -> AgentdResult<GuestBootstrap> {
1689    if msg.id != 0 || msg.flags != 0 {
1690        return Err(AgentdError::Config(format!(
1691            "guest bootstrap requires id=0 and flags=0, got id={} flags={}",
1692            msg.id, msg.flags
1693        )));
1694    }
1695    if msg.t != MessageType::Bootstrap {
1696        return Err(AgentdError::Config(format!(
1697            "expected core.bootstrap as first console frame, got {}",
1698            msg.t.as_str()
1699        )));
1700    }
1701    let min_version = MessageType::Bootstrap.min_protocol_version();
1702    if msg.v < min_version {
1703        return Err(AgentdError::Config(format!(
1704            "guest bootstrap requires protocol generation {min_version} or newer, got {}",
1705            msg.v
1706        )));
1707    }
1708
1709    msg.payload::<GuestBootstrap>()
1710        .map_err(|e| AgentdError::Config(format!("decode guest bootstrap payload: {e}")))
1711}
1712
1713/// Reports init-time guest context to the host and waits for an acknowledgement.
1714pub fn report_init_context(
1715    port_file: &File,
1716    boot_console: &mut BootConsoleState,
1717    default_user: Option<&str>,
1718) -> AgentdResult<()> {
1719    let (uid, gid) = resolve_default_user(default_user)?;
1720    let deadline = init_ack_deadline();
1721    let fd = port_file.as_raw_fd();
1722    set_nonblocking(fd)?;
1723
1724    let msg = Message::with_payload(
1725        MessageType::InitResolved,
1726        0,
1727        &InitResolved {
1728            default_user: ResolvedUser { uid, gid },
1729        },
1730    )
1731    .map_err(|e| AgentdError::ExecSession(format!("encode init context: {e}")))?;
1732
1733    let mut out = Vec::new();
1734    codec::encode_to_buf(&msg, &mut out)
1735        .map_err(|e| AgentdError::ExecSession(format!("encode init context frame: {e}")))?;
1736    write_all_to_fd(fd, &out, Some(deadline))?;
1737    wait_for_init_ack(fd, boot_console, deadline)
1738}
1739
1740/// Reports a fatal startup error before exiting the guest.
1741///
1742/// The host acknowledges only after saving the diagnostic for CLI and SDK
1743/// callers. Supporting hosts control shutdown; older hosts discard the
1744/// report, so only their wait is bounded.
1745pub fn report_init_failure(
1746    port_file: &File,
1747    boot_console: &mut BootConsoleState,
1748    message: &str,
1749    error: &AgentdError,
1750) -> AgentdResult<()> {
1751    let deadline = if boot_console.init_failure_ack {
1752        None
1753    } else {
1754        Some(Instant::now() + std::time::Duration::from_secs(2))
1755    };
1756    let fd = port_file.as_raw_fd();
1757    set_nonblocking(fd)?;
1758    let msg = Message::with_payload(
1759        MessageType::CoreError,
1760        0,
1761        &CoreError {
1762            kind: CoreErrorKind::InitializationFailed,
1763            message: message.to_owned(),
1764            offending_type: None,
1765            init_failure: match error {
1766                AgentdError::UserNotFound(_) => Some(InitFailureReason::UserNotFound),
1767                AgentdError::GroupNotFound(_) => Some(InitFailureReason::GroupNotFound),
1768                _ => None,
1769            },
1770            workload_failure: None,
1771        },
1772    )?;
1773    let mut out = Vec::new();
1774    codec::encode_to_buf(&msg, &mut out)?;
1775    write_all_to_fd(fd, &out, deadline)?;
1776
1777    loop {
1778        let ack = read_boot_message(fd, boot_console, deadline, "init failure ack")?;
1779        if ack.t != MessageType::InitAck || ack.id != 0 || ack.flags != 0 {
1780            return Err(AgentdError::Init(
1781                "expected init failure acknowledgement".into(),
1782            ));
1783        }
1784        if ack.payload::<InitAck>()?.failure {
1785            return Ok(());
1786        }
1787        // A late success acknowledgement must not release a failing guest.
1788    }
1789}
1790
1791//--------------------------------------------------------------------------------------------------
1792// Functions: Helpers
1793//--------------------------------------------------------------------------------------------------
1794
1795fn ensure_fs_bulk_write_worker(
1796    id: u32,
1797    state: &mut AgentState,
1798    session_tx: &SessionOutputSender,
1799) -> Result<(), String> {
1800    if state.bulk_write_workers.contains_key(&id) {
1801        return Ok(());
1802    }
1803
1804    let session = state
1805        .write_sessions
1806        .remove(&id)
1807        .ok_or_else(|| format!("unknown filesystem write session: {id}"))?;
1808    if !session.is_bulk() {
1809        state.write_sessions.insert(id, session);
1810        return Err("raw bulk record sent to a generation-6 filesystem write".into());
1811    }
1812
1813    let (records, record_rx) = tokio::sync::mpsc::channel(FS_BULK_INPUT_ITEM_CAPACITY);
1814    // Finish is a lifecycle cut, not another data record. Its dedicated slot cannot be consumed by
1815    // a legal full 8 MiB data window.
1816    let (finish, finish_rx) = tokio::sync::mpsc::channel(1);
1817    let output = session_tx.clone();
1818    let task = tokio::spawn(run_fs_bulk_write_worker(
1819        id, session, record_rx, finish_rx, output,
1820    ));
1821    state.bulk_write_workers.insert(
1822        id,
1823        FsBulkWriteWorker {
1824            records,
1825            finish,
1826            task,
1827        },
1828    );
1829    Ok(())
1830}
1831
1832fn enqueue_fs_bulk_record(
1833    id: u32,
1834    record: AdmittedBulkRecord,
1835    state: &mut AgentState,
1836    session_tx: &SessionOutputSender,
1837) -> Result<(), String> {
1838    ensure_fs_bulk_write_worker(id, state, session_tx)?;
1839    let result = state
1840        .bulk_write_workers
1841        .get(&id)
1842        .expect("filesystem bulk worker was just established")
1843        .records
1844        .try_send(record)
1845        .map_err(|error| format!("filesystem bulk input queue is unavailable: {error}"));
1846    if result.is_err()
1847        && let Some(worker) = state.bulk_write_workers.remove(&id)
1848    {
1849        worker.task.abort();
1850    }
1851    result
1852}
1853
1854fn finish_fs_bulk_write(
1855    id: u32,
1856    finish: BulkFinish,
1857    state: &mut AgentState,
1858    session_tx: &SessionOutputSender,
1859) -> Result<(), String> {
1860    ensure_fs_bulk_write_worker(id, state, session_tx)?;
1861    let result = state
1862        .bulk_write_workers
1863        .get(&id)
1864        .expect("filesystem bulk worker was just established")
1865        .finish
1866        .try_send(finish)
1867        .map_err(|error| format!("filesystem bulk finish path is unavailable: {error}"));
1868    if result.is_err()
1869        && let Some(worker) = state.bulk_write_workers.remove(&id)
1870    {
1871        worker.task.abort();
1872    }
1873    result
1874}
1875
1876async fn run_fs_bulk_write_worker(
1877    id: u32,
1878    mut session: FsWriteSession,
1879    mut record_rx: tokio::sync::mpsc::Receiver<AdmittedBulkRecord>,
1880    mut finish_rx: tokio::sync::mpsc::Receiver<BulkFinish>,
1881    output_tx: SessionOutputSender,
1882) {
1883    let mut pending_finish = None;
1884    let mut finish_open = true;
1885    let mut received_offset = 0u64;
1886    let mut pending_activity_bytes = 0usize;
1887    let mut pending_record = None;
1888
1889    loop {
1890        let record = if pending_record.is_some() {
1891            pending_record.take()
1892        } else {
1893            tokio::select! {
1894                biased;
1895                finish = finish_rx.recv(), if pending_finish.is_none() && finish_open => {
1896                    match finish {
1897                        Some(finish) => pending_finish = Some(finish),
1898                        None => finish_open = false,
1899                    }
1900                    None
1901                }
1902                record = record_rx.recv() => record,
1903            }
1904        };
1905
1906        let Some(record) = record else {
1907            if pending_finish
1908                .as_ref()
1909                .is_some_and(|finish| finish.final_offset <= received_offset)
1910            {
1911                let finish = pending_finish
1912                    .take()
1913                    .expect("checked pending finish exists");
1914                let mut frame = Vec::new();
1915                if let Err(error) =
1916                    fs::handle_fs_bulk_finish(id, finish, &mut session, &mut frame).await
1917                    && encode_bulk_fs_failure(id, error, &mut frame).is_err()
1918                {
1919                    return;
1920                }
1921                let activity = RawActivity {
1922                    guest_messages: 1,
1923                    fs_bytes: std::mem::take(&mut pending_activity_bytes),
1924                    ..RawActivity::default()
1925                };
1926                let output = crate::session::RawSessionOutput::new(
1927                    frame,
1928                    activity,
1929                    Some(RawSessionCompletion::FsWrite),
1930                );
1931                let _ = output_tx.send(id, SessionOutput::Raw(output)).await;
1932                return;
1933            }
1934            if record_rx.is_closed() {
1935                return;
1936            }
1937            continue;
1938        };
1939
1940        // Drain only records already admitted to this flow. Waiting for a larger batch would add
1941        // latency and could deadlock at a credit boundary; the next input scheduling turn will
1942        // naturally form the next batch.
1943        let mut records = Vec::with_capacity(
1944            FS_BULK_WRITE_COALESCE_BYTES.div_ceil(DEFAULT_FILESYSTEM_BULK_RECORD_PAYLOAD as usize),
1945        );
1946        let mut retained_permits = Vec::with_capacity(records.capacity());
1947        let mut payload_len = record.record().payload.len();
1948        let (record, permit) = record.into_parts();
1949        records.push(record);
1950        retained_permits.push(permit);
1951        while payload_len < FS_BULK_WRITE_COALESCE_BYTES {
1952            match record_rx.try_recv() {
1953                Ok(record) => {
1954                    let next_len = payload_len.saturating_add(record.record().payload.len());
1955                    if next_len > FS_BULK_WRITE_COALESCE_BYTES {
1956                        pending_record = Some(record);
1957                        break;
1958                    }
1959                    let (record, permit) = record.into_parts();
1960                    payload_len = next_len;
1961                    records.push(record);
1962                    retained_permits.push(permit);
1963                }
1964                Err(tokio::sync::mpsc::error::TryRecvError::Empty) => break,
1965                Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => break,
1966            }
1967        }
1968
1969        let mut frame = Vec::new();
1970        let last_record = records
1971            .last()
1972            .expect("filesystem bulk batch contains its initial record");
1973        let record_end = last_record
1974            .offset
1975            .saturating_add(last_record.payload.len() as u64);
1976        let mut completed =
1977            match fs::handle_fs_bulk_records(id, &records, &mut session, &mut frame).await {
1978                Ok(completed) => {
1979                    if !completed {
1980                        received_offset = record_end;
1981                        pending_activity_bytes = pending_activity_bytes.saturating_add(payload_len);
1982                    }
1983                    completed
1984                }
1985                Err(error) => {
1986                    if encode_bulk_fs_failure(id, error, &mut frame).is_err() {
1987                        return;
1988                    }
1989                    true
1990                }
1991            };
1992        // Filesystem effects have consumed every payload in this batch. Release aggregate input
1993        // capacity before potentially waiting to publish credit or heartbeat output.
1994        drop(retained_permits);
1995
1996        if !completed
1997            && pending_finish
1998                .as_ref()
1999                .is_some_and(|finish| finish.final_offset <= received_offset)
2000        {
2001            let finish = pending_finish
2002                .take()
2003                .expect("checked pending finish exists");
2004            completed = match fs::handle_fs_bulk_finish(id, finish, &mut session, &mut frame).await
2005            {
2006                Ok(completed) => completed,
2007                Err(error) => {
2008                    if encode_bulk_fs_failure(id, error, &mut frame).is_err() {
2009                        return;
2010                    }
2011                    true
2012                }
2013            };
2014        }
2015
2016        // Credits and terminal results already need an outward event. Pure byte accounting is
2017        // coalesced so a successful 256 KiB record does not make a second trip through the agent
2018        // main loop solely to update heartbeat counters.
2019        if frame.is_empty() && pending_activity_bytes < FS_ACTIVITY_BATCH_BYTES && !completed {
2020            continue;
2021        }
2022        let activity = RawActivity {
2023            guest_messages: usize::from(!frame.is_empty()),
2024            fs_bytes: std::mem::take(&mut pending_activity_bytes),
2025            ..RawActivity::default()
2026        };
2027        let completion = completed.then_some(RawSessionCompletion::FsWrite);
2028        let output = crate::session::RawSessionOutput::new(frame, activity, completion);
2029        if !output_tx.send(id, SessionOutput::Raw(output)).await || completed {
2030            return;
2031        }
2032    }
2033}
2034
2035/// Handles a single incoming message from the host.
2036async fn handle_bulk_record(
2037    record: AdmittedBulkRecord,
2038    state: &mut AgentState,
2039    activity: &mut ActivityTracker,
2040    out_buf: &mut Vec<u8>,
2041    session_tx: &SessionOutputSender,
2042) -> AgentdResult<()> {
2043    activity.record_host_message();
2044    let id = record.record().id;
2045    let record_end = record
2046        .record()
2047        .offset
2048        .checked_add(record.record().payload.len() as u64)
2049        .ok_or_else(|| AgentdError::ExecSession("bulk record offset overflow".into()))?;
2050    let record_payload_len = record.record().payload.len();
2051    let kind = record.record().kind;
2052    let flow = record.record().flow;
2053    let mut accepted = false;
2054    match kind {
2055        BulkKind::Filesystem => {
2056            if flow != BulkFlow::HostToGuest {
2057                encode_bulk_fs_failure(
2058                    id,
2059                    "host sent a filesystem record in the guest-to-host flow".into(),
2060                    out_buf,
2061                )?;
2062                cancel_bulk_correlation(id, BulkKind::Filesystem, state, session_tx);
2063                return Ok(());
2064            }
2065
2066            if let Err(error) = enqueue_fs_bulk_record(id, record, state, session_tx) {
2067                encode_bulk_fs_failure(id, error, out_buf)?;
2068                cancel_bulk_correlation(id, BulkKind::Filesystem, state, session_tx);
2069            } else {
2070                accepted = true;
2071            }
2072        }
2073        BulkKind::Tcp => {
2074            if flow != BulkFlow::HostToGuest {
2075                encode_bulk_tcp_failure(
2076                    id,
2077                    "host sent a TCP record in the guest-to-host flow".into(),
2078                    out_buf,
2079                )?;
2080                cancel_bulk_correlation(id, BulkKind::Tcp, state, session_tx);
2081                return Ok(());
2082            }
2083            let result = match state.tcp_sessions.get(&id) {
2084                Some(session) => session.write_bulk(record).await,
2085                None => Err(format!("unknown TCP session: {id}")),
2086            };
2087            if let Err(error) = result {
2088                encode_bulk_tcp_failure(id, error, out_buf)?;
2089                cancel_bulk_correlation(id, BulkKind::Tcp, state, session_tx);
2090            } else {
2091                accepted = true;
2092                activity.add_tcp_bytes(record_payload_len);
2093            }
2094        }
2095    }
2096    if accepted {
2097        state.bulk_received_offsets.insert(id, record_end);
2098        if state
2099            .pending_bulk_finishes
2100            .get(&id)
2101            .is_some_and(|finish| finish.final_offset <= record_end)
2102        {
2103            let finish = state
2104                .pending_bulk_finishes
2105                .remove(&id)
2106                .expect("checked pending finish exists");
2107            dispatch_bulk_finish(id, finish, state, out_buf, session_tx).await?;
2108        }
2109    }
2110    Ok(())
2111}
2112
2113/// Apply an exact host-to-guest end marker after all raw records through its offset arrived.
2114async fn dispatch_bulk_finish(
2115    id: u32,
2116    finish: BulkFinish,
2117    state: &mut AgentState,
2118    out_buf: &mut Vec<u8>,
2119    session_tx: &SessionOutputSender,
2120) -> AgentdResult<()> {
2121    match finish.kind {
2122        BulkKind::Filesystem => {
2123            if let Err(error) = finish_fs_bulk_write(id, finish, state, session_tx) {
2124                encode_bulk_fs_failure(id, error, out_buf)?;
2125                cancel_bulk_correlation(id, BulkKind::Filesystem, state, session_tx);
2126            }
2127        }
2128        BulkKind::Tcp => {
2129            let result = match state.tcp_sessions.get(&id) {
2130                Some(session) => session.finish_bulk(finish).await,
2131                None => Err(format!("unknown TCP session: {id}")),
2132            };
2133            if let Err(error) = result {
2134                encode_bulk_tcp_failure(id, error, out_buf)?;
2135                if let Some(session) = state.tcp_sessions.remove(&id) {
2136                    session.close();
2137                }
2138            }
2139        }
2140    }
2141    clear_bulk_receive_state(state, id);
2142    Ok(())
2143}
2144
2145/// Return whether the finish belongs to a live host-to-guest raw receiver.
2146fn has_bulk_receive_session(state: &AgentState, id: u32, kind: BulkKind) -> bool {
2147    match kind {
2148        BulkKind::Filesystem => {
2149            state
2150                .write_sessions
2151                .get(&id)
2152                .is_some_and(FsWriteSession::is_bulk)
2153                || state.bulk_write_workers.contains_key(&id)
2154        }
2155        BulkKind::Tcp => state.tcp_sessions.get(&id).is_some_and(TcpSession::is_bulk),
2156    }
2157}
2158
2159/// Remove offset/finish state after a correlation reaches any terminal path.
2160fn clear_bulk_receive_state(state: &mut AgentState, id: u32) {
2161    state.bulk_received_offsets.remove(&id);
2162    state.pending_bulk_finishes.remove(&id);
2163}
2164
2165/// Cancel every operation resource owned by one bulk correlation.
2166fn cancel_bulk_correlation(
2167    id: u32,
2168    kind: BulkKind,
2169    state: &mut AgentState,
2170    session_tx: &SessionOutputSender,
2171) {
2172    let owner_tx = session_tx.with_incarnation(client_incarnation_for_id(state, id));
2173    // Stop every producer before publishing the scheduler cut. Any send that completed before
2174    // teardown is already bounded in the scheduler input queue; an in-progress send is cancelled
2175    // with its producer task and cannot appear after the cleanup acknowledgement.
2176    match kind {
2177        BulkKind::Filesystem => {
2178            state.write_sessions.remove(&id);
2179            if let Some(worker) = state.bulk_write_workers.remove(&id) {
2180                worker.task.abort();
2181            }
2182            if let Some(session) = state.read_sessions.remove(&id) {
2183                session.abort();
2184            }
2185        }
2186        BulkKind::Tcp => {
2187            if let Some(session) = state.tcp_sessions.remove(&id) {
2188                session.close();
2189            }
2190        }
2191    }
2192    clear_bulk_receive_state(state, id);
2193    if let Err(error) = owner_tx.drop_bulk_flow(id) {
2194        eprintln!("agentd: failed to purge cancelled bulk output id={id}: {error}");
2195    }
2196}
2197
2198fn cancel_all_bulk_correlations(
2199    state: &mut AgentState,
2200    session_tx: &SessionOutputSender,
2201    out_buf: &mut Vec<u8>,
2202    message: &str,
2203) -> AgentdResult<()> {
2204    let mut correlations = HashMap::<u32, BulkKind>::new();
2205    correlations.extend(
2206        state
2207            .read_sessions
2208            .iter()
2209            .filter(|(_, session)| session.is_bulk())
2210            .map(|(id, _)| (*id, BulkKind::Filesystem)),
2211    );
2212    correlations.extend(
2213        state
2214            .write_sessions
2215            .iter()
2216            .filter(|(_, session)| session.is_bulk())
2217            .map(|(id, _)| (*id, BulkKind::Filesystem)),
2218    );
2219    correlations.extend(
2220        state
2221            .bulk_write_workers
2222            .keys()
2223            .map(|id| (*id, BulkKind::Filesystem)),
2224    );
2225    correlations.extend(
2226        state
2227            .tcp_sessions
2228            .iter()
2229            .filter(|(_, session)| session.is_bulk())
2230            .map(|(id, _)| (*id, BulkKind::Tcp)),
2231    );
2232
2233    for (id, kind) in correlations {
2234        encode_bulk_cancel(
2235            id,
2236            kind,
2237            BulkCancelReason::TransportFailure,
2238            message.to_string(),
2239            out_buf,
2240        )?;
2241        cancel_bulk_correlation(id, kind, state, session_tx);
2242        encode_bulk_terminal_failure(id, kind, message.to_string(), out_buf)?;
2243    }
2244    Ok(())
2245}
2246
2247/// Release reorder metadata when the relay recycles an SDK client's ID range.
2248fn clear_bulk_receive_range(state: &mut AgentState, id_start: u32, id_end_exclusive: u32) {
2249    state
2250        .bulk_received_offsets
2251        .retain(|id, _| *id < id_start || *id >= id_end_exclusive);
2252    state
2253        .pending_bulk_finishes
2254        .retain(|id, _| *id < id_start || *id >= id_end_exclusive);
2255}
2256
2257fn client_incarnation_for_id(state: &AgentState, id: u32) -> Option<ClientIncarnation> {
2258    relay_client_slot(id).and_then(|slot| state.client_incarnations.get(&slot).copied())
2259}
2260
2261fn bulk_output_retired_bit(id: u32) -> Option<(usize, u64)> {
2262    let slot = relay_client_slot(id)?;
2263    let (id_start, _) = relay_client_id_range(slot)?;
2264    let local = usize::try_from(id.checked_sub(id_start)?).ok()?;
2265    Some((
2266        local / u64::BITS as usize,
2267        1u64 << (local % u64::BITS as usize),
2268    ))
2269}
2270
2271fn bulk_output_is_retired(
2272    retired: &HashMap<ClientIncarnation, Vec<u64>>,
2273    incarnation: ClientIncarnation,
2274    id: u32,
2275) -> bool {
2276    let Some((word, mask)) = bulk_output_retired_bit(id) else {
2277        return false;
2278    };
2279    retired
2280        .get(&incarnation)
2281        .and_then(|bitmap| bitmap.get(word))
2282        .is_some_and(|bits| bits & mask != 0)
2283}
2284
2285fn retire_bulk_output(
2286    retired: &mut HashMap<ClientIncarnation, Vec<u64>>,
2287    incarnation: ClientIncarnation,
2288    id: u32,
2289) -> AgentdResult<()> {
2290    let (word, mask) = bulk_output_retired_bit(id).ok_or_else(|| {
2291        AgentdError::ExecSession(format!("cannot retire unassigned bulk correlation {id}"))
2292    })?;
2293    let bitmap = retired.entry(incarnation).or_default();
2294    if bitmap.len() <= word {
2295        bitmap.resize(word + 1, 0);
2296    }
2297    bitmap[word] |= mask;
2298    Ok(())
2299}
2300
2301/// Accept current ownership, silently reject stale ownership, and fail cross-client spoofing.
2302fn validate_bulk_client_incarnation(
2303    state: &AgentState,
2304    id: u32,
2305    claimed: ClientIncarnation,
2306) -> AgentdResult<bool> {
2307    if client_incarnation_for_id(state, id) == Some(claimed) {
2308        return Ok(true);
2309    }
2310    if state
2311        .client_incarnations
2312        .values()
2313        .any(|current| *current == claimed)
2314    {
2315        return Err(AgentdError::ExecSession(format!(
2316            "dedicated bulk record correlation {id} lies outside its client incarnation range"
2317        )));
2318    }
2319    Ok(false)
2320}
2321
2322async fn acquire_bulk_input_permit(
2323    admission: Option<(Arc<Semaphore>, usize)>,
2324) -> AgentdResult<OwnedSemaphorePermit> {
2325    let (budget, payload_len) = admission.expect("bulk admission future is guarded by queue state");
2326    let charged = u32::try_from(payload_len)
2327        .map_err(|_| AgentdError::ExecSession("bulk input payload budget overflow".into()))?;
2328    budget
2329        .acquire_many_owned(charged)
2330        .await
2331        .map_err(|error| AgentdError::ExecSession(format!("bulk input budget closed: {error}")))
2332}
2333
2334/// Validate that an internal lifecycle message names one complete relay-owned range.
2335fn validate_relay_client_range(id_start: u32, id_end_exclusive: u32) -> AgentdResult<u32> {
2336    canonical_relay_client_range(id_start, id_end_exclusive).ok_or_else(|| {
2337        AgentdError::ExecSession(format!(
2338            "invalid relay client range [{id_start}, {id_end_exclusive})"
2339        ))
2340    })
2341}
2342
2343fn establish_relay_client(
2344    state: &mut AgentState,
2345    connected: RelayClientConnected,
2346) -> AgentdResult<()> {
2347    let slot = validate_relay_client_range(connected.id_start, connected.id_end_exclusive)?;
2348    if connected.incarnation == [0; CLIENT_INCARNATION_SIZE] {
2349        return Err(AgentdError::ExecSession(
2350            "relay client incarnation cannot be zero".into(),
2351        ));
2352    }
2353    if state
2354        .client_incarnations
2355        .get(&slot)
2356        .is_some_and(|current| *current != connected.incarnation)
2357    {
2358        return Err(AgentdError::ExecSession(format!(
2359            "relay client slot {slot} was replaced before acknowledged disconnect"
2360        )));
2361    }
2362    state
2363        .client_incarnations
2364        .insert(slot, connected.incarnation);
2365    Ok(())
2366}
2367
2368/// Remove one current range owner; return false for a stale disconnect that was ignored.
2369fn disconnect_relay_client(
2370    state: &mut AgentState,
2371    disconnected: RelayClientDisconnected,
2372    session_tx: &SessionOutputSender,
2373) -> AgentdResult<(bool, Option<tokio::sync::oneshot::Receiver<()>>)> {
2374    let slot = validate_relay_client_range(disconnected.id_start, disconnected.id_end_exclusive)?;
2375    if let Some(incarnation) = disconnected.incarnation {
2376        if state.client_incarnations.get(&slot) != Some(&incarnation) {
2377            return Ok((false, None));
2378        }
2379        state.client_incarnations.remove(&slot);
2380    } else if state.client_incarnations.contains_key(&slot) {
2381        return Err(AgentdError::ExecSession(
2382            "bound dual-port range cleanup omitted its incarnation".into(),
2383        ));
2384    }
2385    cleanup_relay_client_range(state, disconnected.id_start, disconnected.id_end_exclusive);
2386    let scheduler_cleanup = match disconnected.incarnation {
2387        Some(incarnation) => session_tx
2388            .drop_bulk_incarnation(incarnation)
2389            .map_err(|error| AgentdError::ExecSession(error.into()))?,
2390        None => None,
2391    };
2392    Ok((true, scheduler_cleanup))
2393}
2394
2395/// Tear down resources and reorder state belonging to one relay-owned correlation range.
2396fn cleanup_relay_client_range(state: &mut AgentState, id_start: u32, id_end_exclusive: u32) {
2397    state.sessions.retain(|id, session| {
2398        let keep = *id < id_start || *id >= id_end_exclusive;
2399        if !keep {
2400            // The SDK owner disappeared, so no later stdin/signal request can terminate this tree.
2401            // Kill the process group now and drop its PTY/stdin handles with the session entry.
2402            let _ = session.send_signal(9);
2403        }
2404        keep
2405    });
2406    state.fs.close_owner_range(id_start, id_end_exclusive);
2407    abort_read_sessions_in_owner_range(&mut state.read_sessions, id_start, id_end_exclusive);
2408    state.write_sessions.retain(|_, session| {
2409        let owner_id = session.owner_id();
2410        owner_id < id_start || owner_id >= id_end_exclusive
2411    });
2412    state.bulk_write_workers.retain(|id, worker| {
2413        let keep = *id < id_start || *id >= id_end_exclusive;
2414        if !keep {
2415            worker.task.abort();
2416        }
2417        keep
2418    });
2419    close_tcp_sessions_in_owner_range(&mut state.tcp_sessions, id_start, id_end_exclusive);
2420    clear_bulk_receive_range(state, id_start, id_end_exclusive);
2421}
2422
2423/// Detach transport ownership while retaining the complete guest process lifetime.
2424async fn restore_client_state(
2425    state: &mut AgentState,
2426    sender: &mut SessionOutputSender,
2427) -> AgentdResult<()> {
2428    let generation = sender.generation();
2429    state
2430        .detached_sessions
2431        .extend(
2432            std::mem::take(&mut state.sessions)
2433                .into_iter()
2434                .map(|(id, mut session)| {
2435                    session.detach_stdin();
2436                    ((generation, id), session)
2437                }),
2438        );
2439    state.client_incarnations.clear();
2440    for (_, session) in state.read_sessions.drain() {
2441        session.abort();
2442    }
2443    for (_, worker) in state.bulk_write_workers.drain() {
2444        worker.task.abort();
2445        let _ = worker.task.await;
2446    }
2447    state.write_sessions.clear();
2448    for (_, session) in state.tcp_sessions.drain() {
2449        session.close();
2450    }
2451    state.fs.clear();
2452    state.bulk_received_offsets.clear();
2453    state.pending_bulk_finishes.clear();
2454    // The host drains both physical ports during this barrier. A partially written old record
2455    // must finish, not be interrupted midway and desynchronize the bulk byte stream.
2456    time::timeout(Duration::from_secs(5), sender.restore_generation())
2457        .await
2458        .map_err(|_| {
2459            AgentdError::ExecSession(
2460                "restore bulk cleanup timed out; workloads remain frozen".into(),
2461            )
2462        })?
2463        .map_err(|error| AgentdError::ExecSession(error.into()))
2464}
2465
2466/// Retired producers still drain stdout/stderr so inherited children never hit a closed pipe.
2467fn discard_inherited_output(
2468    state: &mut AgentState,
2469    envelope: &SessionOutputEnvelope,
2470    generation: u64,
2471) -> bool {
2472    if envelope.generation == generation {
2473        return false;
2474    }
2475    if matches!(envelope.output, SessionOutput::Exited(_)) {
2476        state
2477            .detached_sessions
2478            .remove(&(envelope.generation, envelope.id));
2479    }
2480    true
2481}
2482
2483// Keep the loop-owned latch and output generation explicit at this dispatch boundary; merging
2484// them into AgentState would obscure the ownership needed by restore and background producers.
2485#[allow(clippy::too_many_arguments)]
2486async fn handle_message_with_charge(
2487    msg: Message,
2488    state: &mut AgentState,
2489    activity: &mut ActivityTracker,
2490    root_session_tx: &mut SessionOutputSender,
2491    out_buf: &mut Vec<u8>,
2492    config: &AgentdConfig,
2493    workload: &mut WorkloadLatch,
2494    heartbeat_control: &heartbeat::HeartbeatControl,
2495    mut input_charge: Option<InputCharge>,
2496) -> AgentdResult<()> {
2497    // Background producers retain the range owner that opened them. The main loop can then drop
2498    // queued output after a disconnect instead of relabelling it with a recycled correlation ID.
2499    let session_tx = root_session_tx.with_incarnation(client_incarnation_for_id(state, msg.id));
2500    match msg.t {
2501        MessageType::RootDiskPrepare | MessageType::RootDiskGrow => {
2502            let Some(request) = decode_payload_or_core_error::<
2503                microsandbox_protocol::core::RootDiskGrow,
2504            >(&msg, out_buf)?
2505            else {
2506                return Ok(());
2507            };
2508            let apply = msg.t == MessageType::RootDiskGrow;
2509            let result = tokio::task::spawn_blocking(move || {
2510                crate::root_disk::resize(request.size_bytes, apply)
2511            })
2512            .await
2513            .map_err(|e| AgentdError::ExecSession(format!("root grow worker: {e}")))?;
2514            let reply = match result {
2515                Ok(state) => Message::with_payload(MessageType::RootDiskState, msg.id, &state),
2516                Err(message) => Message::with_payload(
2517                    MessageType::CoreError,
2518                    msg.id,
2519                    &CoreError {
2520                        kind: CoreErrorKind::CapabilityUnavailable,
2521                        message,
2522                        offending_type: Some(msg.t.as_str().into()),
2523                        init_failure: None,
2524                        workload_failure: None,
2525                    },
2526                ),
2527            }
2528            .map_err(|e| AgentdError::ExecSession(format!("encode root capacity: {e}")))?;
2529            codec::encode_to_buf(&reply, out_buf).map_err(|e| {
2530                AgentdError::ExecSession(format!("encode root capacity frame: {e}"))
2531            })?;
2532        }
2533
2534        MessageType::Ping => {
2535            let Some(_) = decode_payload_or_core_error::<Ping>(&msg, out_buf)? else {
2536                return Ok(());
2537            };
2538            let reply = Message::with_payload(MessageType::Pong, msg.id, &Pong {})
2539                .map_err(|e| AgentdError::ExecSession(format!("encode pong: {e}")))?;
2540            codec::encode_to_buf(&reply, out_buf)
2541                .map_err(|e| AgentdError::ExecSession(format!("encode pong frame: {e}")))?;
2542        }
2543
2544        MessageType::Touch => {
2545            let Some(_) = decode_payload_or_core_error::<Touch>(&msg, out_buf)? else {
2546                return Ok(());
2547            };
2548            activity.record_host_message();
2549            let reply = Message::with_payload(
2550                MessageType::Touched,
2551                msg.id,
2552                &Touched {
2553                    activity_seq: activity.activity_seq,
2554                },
2555            )
2556            .map_err(|e| AgentdError::ExecSession(format!("encode touched: {e}")))?;
2557            codec::encode_to_buf(&reply, out_buf)
2558                .map_err(|e| AgentdError::ExecSession(format!("encode touched frame: {e}")))?;
2559        }
2560
2561        MessageType::WorkloadFreeze => {
2562            let Some(request) = decode_payload_or_core_error::<WorkloadFreeze>(&msg, out_buf)?
2563            else {
2564                return Ok(());
2565            };
2566            let position = state.input_window.position();
2567            if state.pending_freeze.as_ref().is_some_and(|pending| {
2568                pending
2569                    .payload::<WorkloadFreeze>()
2570                    .is_ok_and(|old| old != request)
2571            }) {
2572                encode_workload_error(
2573                    &msg,
2574                    &request.attempt_id,
2575                    WorkloadLatchError::Conflict("another transport cut is pending".into()),
2576                    out_buf,
2577                )?;
2578                return Ok(());
2579            }
2580            if workload.is_frozen() {
2581                if let Err(error) = workload.require_frozen_attempt(&request.attempt_id) {
2582                    encode_workload_error(&msg, &request.attempt_id, error, out_buf)?;
2583                    return Ok(());
2584                }
2585                if state
2586                    .frozen_host_input
2587                    .is_some_and(|cut| cut != request.host_input)
2588                {
2589                    encode_workload_error(
2590                        &msg,
2591                        &request.attempt_id,
2592                        WorkloadLatchError::Conflict("frozen transport cut cannot change".into()),
2593                        out_buf,
2594                    )?;
2595                    return Ok(());
2596                }
2597            }
2598            if request.host_input.control_bytes != position.control_bytes
2599                || request.host_input.control_frames != position.control_frames
2600                || request.host_input.bulk_bytes < position.bulk_bytes
2601                || request.host_input.bulk_frames < position.bulk_frames
2602            {
2603                encode_workload_error(
2604                    &msg,
2605                    &request.attempt_id,
2606                    WorkloadLatchError::Conflict(
2607                        "transport input cut differs from received complete frames".into(),
2608                    ),
2609                    out_buf,
2610                )?;
2611                return Ok(());
2612            }
2613            state.aborted_transport_attempt = None;
2614            if request.host_input != position {
2615                // Retain the private request while the independently ordered bulk prefix
2616                // arrives. Decoding owns its bytes; application consumption is not required.
2617                state.pending_freeze = Some(msg);
2618                return Ok(());
2619            }
2620            state.pending_freeze = None;
2621            let was_frozen = workload.is_frozen();
2622            match workload.freeze(&request.attempt_id) {
2623                Ok(()) => {
2624                    if !was_frozen {
2625                        // A new capture can reuse a human-selected checkpoint name.
2626                        state.restored_attempt = None;
2627                    }
2628                    if let Err(pause_error) = heartbeat_control.pause() {
2629                        let rollback = workload.thaw(&request.attempt_id).err();
2630                        let message = match rollback {
2631                            Some(error) => format!(
2632                                "heartbeat checkpoint gate failed: {pause_error}; workload rollback failed: {error}"
2633                            ),
2634                            None => format!("heartbeat checkpoint gate failed: {pause_error}"),
2635                        };
2636                        encode_workload_error(
2637                            &msg,
2638                            &request.attempt_id,
2639                            WorkloadLatchError::Io(std::io::Error::other(message)),
2640                            out_buf,
2641                        )?;
2642                    } else {
2643                        state.output_parked = true;
2644                        state.frozen_host_input = Some(position);
2645                        state.frozen_guest_bulk_bytes = root_session_tx
2646                            .park_bulk_output()
2647                            .await
2648                            .map_err(|error| AgentdError::ExecSession(error.into()))?;
2649                        // Agent upload workers are outside the workload cgroup. A capture
2650                        // cannot certify clean external pages while one may still write.
2651                        // Resident pause remains available; only mount capture needs this proof.
2652                        let external_mount_tags = request.external_mount_tags.clone();
2653                        let external_mounts_synced = if external_mount_tags.is_empty() {
2654                            true
2655                        } else if state.bulk_write_workers.is_empty() {
2656                            let permit = crate::mount_checkpoint::try_start_sync();
2657                            if let Some(permit) = permit {
2658                                matches!(
2659                                    time::timeout(
2660                                        std::time::Duration::from_secs(20),
2661                                        tokio::task::spawn_blocking(move || {
2662                                            // A timed-out join does not cancel syncfs. Keep the
2663                                            // reservation in the worker until the actual flush ends.
2664                                            let _permit = permit;
2665                                            crate::mount_checkpoint::sync_external_mounts(
2666                                                external_mount_tags,
2667                                            )
2668                                        }),
2669                                    )
2670                                    .await,
2671                                    Ok(Ok(Ok(())))
2672                                )
2673                            } else {
2674                                false
2675                            }
2676                        } else {
2677                            false
2678                        };
2679                        let reply = Message::with_payload(
2680                            MessageType::WorkloadFrozen,
2681                            msg.id,
2682                            &WorkloadFrozen {
2683                                attempt_id: request.attempt_id,
2684                                guest_bulk_bytes_target: state.frozen_guest_bulk_bytes,
2685                                input_credit: state.input_window.credit()?,
2686                                external_mounts_synced,
2687                            },
2688                        )
2689                        .map_err(|error| {
2690                            AgentdError::ExecSession(format!(
2691                                "encode workload-frozen response: {error}"
2692                            ))
2693                        })?;
2694                        codec::encode_to_buf(&reply, out_buf).map_err(|error| {
2695                            AgentdError::ExecSession(format!(
2696                                "encode workload-frozen frame: {error}"
2697                            ))
2698                        })?;
2699                    }
2700                }
2701                Err(error) => encode_workload_error(&msg, &request.attempt_id, error, out_buf)?,
2702            }
2703        }
2704
2705        MessageType::WorkloadThaw => {
2706            let Some(request) = decode_payload_or_core_error::<WorkloadThaw>(&msg, out_buf)? else {
2707                return Ok(());
2708            };
2709            if state.pending_freeze.as_ref().is_some_and(|pending| {
2710                pending
2711                    .payload::<WorkloadFreeze>()
2712                    .is_ok_and(|freeze| freeze.attempt_id != request.attempt_id)
2713            }) {
2714                encode_workload_error(
2715                    &msg,
2716                    &request.attempt_id,
2717                    WorkloadLatchError::Conflict("another transport cut is pending".into()),
2718                    out_buf,
2719                )?;
2720                return Ok(());
2721            }
2722            if state.pending_freeze.as_ref().is_some_and(|pending| {
2723                pending
2724                    .payload::<WorkloadFreeze>()
2725                    .is_ok_and(|freeze| freeze.attempt_id == request.attempt_id)
2726            }) && !workload.is_frozen()
2727                || (!workload.is_frozen()
2728                    && state.aborted_transport_attempt.as_deref() == Some(&request.attempt_id))
2729            {
2730                if request.mode != microsandbox_protocol::core::WorkloadThawMode::Continue {
2731                    encode_workload_error(
2732                        &msg,
2733                        &request.attempt_id,
2734                        WorkloadLatchError::Conflict(
2735                            "restore requires a completed transport cut".into(),
2736                        ),
2737                        out_buf,
2738                    )?;
2739                    return Ok(());
2740                }
2741                if let Some(pending) = state.pending_freeze.take() {
2742                    // The abandoned Freeze RPC still owns a private host reply slot. Complete
2743                    // it explicitly before acknowledging the independent recovery Thaw RPC.
2744                    encode_workload_error(
2745                        &pending,
2746                        &request.attempt_id,
2747                        WorkloadLatchError::Conflict(
2748                            "transport cut aborted by source continuation".into(),
2749                        ),
2750                        out_buf,
2751                    )?;
2752                }
2753                state.aborted_transport_attempt = Some(request.attempt_id.clone());
2754                encode_input_credit(state.input_window.credit()?, out_buf)?;
2755                let reply = Message::with_payload(
2756                    MessageType::WorkloadThawed,
2757                    msg.id,
2758                    &WorkloadThawed {
2759                        attempt_id: request.attempt_id,
2760                    },
2761                )?;
2762                codec::encode_to_buf(&reply, out_buf)?;
2763                return Ok(());
2764            }
2765            if request.mode == microsandbox_protocol::core::WorkloadThawMode::Restore
2766                && state.restored_attempt.as_deref() != Some(&request.attempt_id)
2767            {
2768                if let Err(error) = workload.require_frozen_attempt(&request.attempt_id) {
2769                    encode_workload_error(&msg, &request.attempt_id, error, out_buf)?;
2770                    return Ok(());
2771                }
2772                // Never call ordinary disconnect here: it kills the very processes we captured.
2773                restore_client_state(state, root_session_tx).await?;
2774                state.restored_attempt = Some(request.attempt_id.clone());
2775            }
2776            match workload.thaw(&request.attempt_id) {
2777                Ok(()) => {
2778                    heartbeat_control.resume();
2779                    // Cumulative grants include refunds from detached transfer cleanup; accepted
2780                    // stdin remains charged until its inherited process consumes it after thaw.
2781                    encode_input_credit(state.input_window.credit()?, out_buf)?;
2782                    state.resume_output_after_flush = true;
2783                    state.frozen_host_input = None;
2784                    let reply = Message::with_payload(
2785                        MessageType::WorkloadThawed,
2786                        msg.id,
2787                        &WorkloadThawed {
2788                            attempt_id: request.attempt_id,
2789                        },
2790                    )
2791                    .map_err(|error| {
2792                        AgentdError::ExecSession(format!(
2793                            "encode workload-thawed response: {error}"
2794                        ))
2795                    })?;
2796                    codec::encode_to_buf(&reply, out_buf).map_err(|error| {
2797                        AgentdError::ExecSession(format!("encode workload-thawed frame: {error}"))
2798                    })?;
2799                }
2800                Err(error) => encode_workload_error(&msg, &request.attempt_id, error, out_buf)?,
2801            }
2802        }
2803
2804        MessageType::ExecRequest => {
2805            let Some(mut req) = decode_payload_or_core_error::<ExecRequest>(&msg, out_buf)? else {
2806                return Ok(());
2807            };
2808            if req.cwd.is_none() {
2809                req.cwd = config.default_cwd().map(str::to_string);
2810            }
2811            prepend_scripts_to_path(&mut req);
2812            if workload.is_frozen() {
2813                encode_exec_failed(
2814                    msg.id,
2815                    ExecFailed {
2816                        kind: ExecFailureKind::Other,
2817                        errno: None,
2818                        errno_name: None,
2819                        message: "sandbox workload is frozen for checkpoint activation".into(),
2820                        stage: Some("workload_latch".into()),
2821                    },
2822                    out_buf,
2823                )?;
2824                return Ok(());
2825            }
2826            let workload_placement = match workload.placement() {
2827                Ok(placement) => placement,
2828                Err(error) => {
2829                    encode_exec_failed(
2830                        msg.id,
2831                        ExecFailed {
2832                            kind: ExecFailureKind::Other,
2833                            errno: None,
2834                            errno_name: None,
2835                            message: error.to_string(),
2836                            stage: Some("workload_cgroup".into()),
2837                        },
2838                        out_buf,
2839                    )?;
2840                    return Ok(());
2841                }
2842            };
2843            match ExecSession::spawn(
2844                msg.id,
2845                &req,
2846                session_tx.clone(),
2847                config.user.as_deref(),
2848                config.security_profile,
2849                workload_placement,
2850            ) {
2851                Ok(session) => {
2852                    let reply = Message::with_payload(
2853                        MessageType::ExecStarted,
2854                        msg.id,
2855                        &ExecStarted { pid: session.pid() },
2856                    )
2857                    .map_err(|e| AgentdError::ExecSession(format!("encode started: {e}")))?;
2858                    codec::encode_to_buf(&reply, out_buf).map_err(|e| {
2859                        AgentdError::ExecSession(format!("encode started frame: {e}"))
2860                    })?;
2861                    state.sessions.insert(msg.id, session);
2862                }
2863                Err(e) => {
2864                    // Send a typed `ExecFailed` so the host can render a
2865                    // useful message + hint. `ExecSpawnFailed` already
2866                    // carries the structured payload; other error
2867                    // variants (free-form `ExecSession(_)` etc.) get
2868                    // wrapped as `Other` with the message preserved.
2869                    let payload = match &e {
2870                        AgentdError::ExecSpawnFailed(p) => p.clone(),
2871                        other => ExecFailed {
2872                            kind: ExecFailureKind::Other,
2873                            errno: None,
2874                            errno_name: None,
2875                            message: other.to_string(),
2876                            stage: None,
2877                        },
2878                    };
2879                    let reply = Message::with_payload(MessageType::ExecFailed, msg.id, &payload)
2880                        .map_err(|e| AgentdError::ExecSession(format!("encode failed: {e}")))?;
2881                    codec::encode_to_buf(&reply, out_buf).map_err(|e| {
2882                        AgentdError::ExecSession(format!("encode failed frame: {e}"))
2883                    })?;
2884                    eprintln!("failed to spawn exec session {}: {e}", msg.id);
2885                }
2886            }
2887        }
2888
2889        MessageType::ExecStdin => {
2890            let Some(stdin) = decode_payload_or_core_error::<ExecStdin>(&msg, out_buf)? else {
2891                return Ok(());
2892            };
2893            if let Some(session) = state.sessions.get_mut(&msg.id)
2894                && let Err(error) = session.enqueue_stdin(stdin.data, input_charge.take())
2895            {
2896                encode_stdin_error(msg.id, &AgentdError::Io(error), out_buf)?;
2897            }
2898        }
2899
2900        MessageType::ExecResize => {
2901            let Some(resize) = decode_payload_or_core_error::<ExecResize>(&msg, out_buf)? else {
2902                return Ok(());
2903            };
2904            if let Some(session) = state.sessions.get(&msg.id) {
2905                let _ = session.resize(resize.rows, resize.cols);
2906            }
2907        }
2908
2909        MessageType::ExecSignal => {
2910            let Some(signal) = decode_payload_or_core_error::<ExecSignal>(&msg, out_buf)? else {
2911                return Ok(());
2912            };
2913            if let Some(session) = state.sessions.get(&msg.id) {
2914                let _ = session.send_signal(signal.signal);
2915            }
2916        }
2917
2918        MessageType::FsRequest => {
2919            let Some(req) = decode_payload_or_core_error::<FsRequest>(&msg, out_buf)? else {
2920                return Ok(());
2921            };
2922            match fs::handle_fs_request(msg.id, msg.v, req, &mut state.fs, out_buf, &session_tx)
2923                .await
2924            {
2925                Ok(Some(FsStreamSession::Read(rs))) => {
2926                    state.read_sessions.insert(msg.id, rs);
2927                }
2928                Ok(Some(FsStreamSession::Write(ws))) => {
2929                    state.write_sessions.insert(msg.id, ws);
2930                }
2931                Ok(None) => {}
2932                Err(e) => {
2933                    eprintln!("fs request error for {}: {e}", msg.id);
2934                }
2935            }
2936        }
2937
2938        MessageType::FsData => {
2939            let Some(data) = decode_payload_or_core_error::<FsData>(&msg, out_buf)? else {
2940                return Ok(());
2941            };
2942            let len = data.data.len();
2943            if let Some(session) = state.write_sessions.get_mut(&msg.id) {
2944                match fs::handle_fs_data(msg.id, data, session, out_buf).await {
2945                    Ok(true) => {
2946                        // Session complete — remove it.
2947                        state.write_sessions.remove(&msg.id);
2948                        clear_bulk_receive_state(state, msg.id);
2949                    }
2950                    Ok(false) => {
2951                        activity.add_fs_bytes(len);
2952                    }
2953                    Err(e) => {
2954                        eprintln!("fs data error for {}: {e}", msg.id);
2955                        state.write_sessions.remove(&msg.id);
2956                        clear_bulk_receive_state(state, msg.id);
2957                    }
2958                }
2959            } else {
2960                // No write session for this ID — send error response.
2961                let resp = microsandbox_protocol::fs::FsResponse {
2962                    ok: false,
2963                    error: Some(format!("unknown write session: {}", msg.id)),
2964                    data: None,
2965                };
2966                let reply = Message::with_payload(MessageType::FsResponse, msg.id, &resp)
2967                    .map_err(|e| AgentdError::ExecSession(format!("encode fs error: {e}")))?;
2968                codec::encode_to_buf(&reply, out_buf)
2969                    .map_err(|e| AgentdError::ExecSession(format!("encode fs error frame: {e}")))?;
2970            }
2971        }
2972
2973        MessageType::BulkCredit => {
2974            let Some(credit) = decode_payload_or_core_error::<BulkCredit>(&msg, out_buf)? else {
2975                return Ok(());
2976            };
2977            match credit.kind {
2978                BulkKind::Filesystem => {
2979                    let result = state
2980                        .read_sessions
2981                        .get(&msg.id)
2982                        .ok_or_else(|| format!("unknown filesystem read session: {}", msg.id))
2983                        .and_then(|session| session.apply_credit(credit));
2984                    if let Err(error) = result {
2985                        encode_bulk_fs_failure(msg.id, error, out_buf)?;
2986                        cancel_bulk_correlation(msg.id, BulkKind::Filesystem, state, &session_tx);
2987                    }
2988                }
2989                BulkKind::Tcp => {
2990                    let result = match state.tcp_sessions.get(&msg.id) {
2991                        Some(session) => session.apply_credit(credit).await,
2992                        // A terminal may retire the producer before the other physical lane
2993                        // finishes delivering its output. Late credit has no recipient and must
2994                        // not cancel that tail or manufacture a second terminal response.
2995                        None => Ok(()),
2996                    };
2997                    if let Err(error) = result {
2998                        encode_bulk_tcp_failure(msg.id, error, out_buf)?;
2999                        cancel_bulk_correlation(msg.id, BulkKind::Tcp, state, &session_tx);
3000                    }
3001                }
3002            }
3003        }
3004
3005        MessageType::BulkFinish => {
3006            let Some(finish) = decode_payload_or_core_error::<BulkFinish>(&msg, out_buf)? else {
3007                return Ok(());
3008            };
3009            let received_offset = state
3010                .bulk_received_offsets
3011                .get(&msg.id)
3012                .copied()
3013                .unwrap_or(0);
3014            if has_bulk_receive_session(state, msg.id, finish.kind)
3015                && finish.final_offset > received_offset
3016            {
3017                if state.pending_bulk_finishes.insert(msg.id, finish).is_some() {
3018                    encode_bulk_cancel(
3019                        msg.id,
3020                        finish.kind,
3021                        BulkCancelReason::ProtocolState,
3022                        "duplicate deferred bulk finish".into(),
3023                        out_buf,
3024                    )?;
3025                    cancel_bulk_correlation(msg.id, finish.kind, state, &session_tx);
3026                }
3027            } else {
3028                dispatch_bulk_finish(msg.id, finish, state, out_buf, &session_tx).await?;
3029            }
3030        }
3031
3032        MessageType::BulkCancel => {
3033            let Some(cancel) = decode_payload_or_core_error::<BulkCancel>(&msg, out_buf)? else {
3034                return Ok(());
3035            };
3036            cancel_bulk_correlation(msg.id, cancel.kind, state, &session_tx);
3037            // Cancellation owns the correlation through its ordinary terminal response. This
3038            // gives SDK dispatch a precise point at which it may retire the route while late raw
3039            // records admitted under pre-cancel credit are still being discarded.
3040            encode_bulk_terminal_failure(msg.id, cancel.kind, cancel.message, out_buf)?;
3041        }
3042
3043        MessageType::BulkAccepted => {
3044            encode_core_error_if_supported(
3045                &msg,
3046                msg.id,
3047                CoreErrorKind::UnsupportedMessageType,
3048                "host cannot accept a guest-initiated bulk offer".into(),
3049                Some(msg.t.as_str().to_string()),
3050                out_buf,
3051            )?;
3052        }
3053
3054        MessageType::TcpConnect => {
3055            let Some(req) = decode_payload_or_core_error::<TcpConnect>(&msg, out_buf)? else {
3056                return Ok(());
3057            };
3058            if req.bulk.is_some() && msg.v < BULK_PROTOCOL_VERSION {
3059                encode_tcp_failed(
3060                    msg.id,
3061                    format!("raw bulk offer requires protocol generation {BULK_PROTOCOL_VERSION}"),
3062                    out_buf,
3063                )?;
3064                return Ok(());
3065            }
3066            // The connect runs inside the session task; the agent loop never
3067            // blocks on it. Success or failure arrives later as a tcp frame.
3068            let session = TcpSession::open(msg.id, req, &session_tx);
3069            state.tcp_sessions.insert(msg.id, session);
3070        }
3071
3072        MessageType::TcpData => {
3073            let Some(data) = decode_payload_or_core_error::<TcpData>(&msg, out_buf)? else {
3074                return Ok(());
3075            };
3076            let len = data.data.len();
3077            if let Some(session) = state.tcp_sessions.get(&msg.id) {
3078                if let Err(e) = session
3079                    .write_data_charged(data.data, input_charge.take())
3080                    .await
3081                {
3082                    state.tcp_sessions.remove(&msg.id);
3083                    clear_bulk_receive_state(state, msg.id);
3084                    encode_tcp_failed(msg.id, e, out_buf)?;
3085                } else {
3086                    activity.add_tcp_bytes(len);
3087                }
3088            } else {
3089                encode_tcp_failed(msg.id, format!("unknown TCP session: {}", msg.id), out_buf)?;
3090            }
3091        }
3092
3093        MessageType::TcpEof => {
3094            let Some(_) = decode_payload_or_core_error::<TcpEof>(&msg, out_buf)? else {
3095                return Ok(());
3096            };
3097            if let Some(session) = state.tcp_sessions.get(&msg.id)
3098                && let Err(e) = session.close_write_charged(input_charge.take()).await
3099            {
3100                state.tcp_sessions.remove(&msg.id);
3101                clear_bulk_receive_state(state, msg.id);
3102                encode_tcp_failed(msg.id, e, out_buf)?;
3103            }
3104        }
3105
3106        MessageType::TcpClose => {
3107            let Some(_) = decode_payload_or_core_error::<TcpClose>(&msg, out_buf)? else {
3108                return Ok(());
3109            };
3110            if let Some(session) = state.tcp_sessions.remove(&msg.id) {
3111                session.close();
3112            }
3113            clear_bulk_receive_state(state, msg.id);
3114        }
3115
3116        MessageType::RelayClientDisconnected => {
3117            let Some(disconnected) =
3118                decode_payload_or_core_error::<RelayClientDisconnected>(&msg, out_buf)?
3119            else {
3120                return Ok(());
3121            };
3122            // A late cleanup for an old owner must never affect a recycled range.
3123            let disconnect_ack =
3124                disconnected
3125                    .incarnation
3126                    .map(|incarnation| RelayClientDisconnectedAck {
3127                        id_start: disconnected.id_start,
3128                        id_end_exclusive: disconnected.id_end_exclusive,
3129                        incarnation,
3130                    });
3131            let (removed, scheduler_cleanup) =
3132                disconnect_relay_client(state, disconnected, &session_tx)?;
3133            if removed && let Some(disconnect_ack) = disconnect_ack {
3134                if let Some(cleanup) = scheduler_cleanup {
3135                    // Keep the single-vCPU control actor available while the bulk actor reaches
3136                    // its cleanup cut. The acknowledgement is emitted through the ordinary
3137                    // control queue only after matching queued records release their permits.
3138                    let output_tx = session_tx.clone();
3139                    tokio::spawn(async move {
3140                        if cleanup.await.is_ok() {
3141                            let frame =
3142                                encode_relay_client_disconnected_ack(disconnect_ack).to_vec();
3143                            let _ = output_tx
3144                                .send(
3145                                    0,
3146                                    SessionOutput::Raw(RawSessionOutput::new(
3147                                        frame,
3148                                        RawActivity::default(),
3149                                        None,
3150                                    )),
3151                                )
3152                                .await;
3153                        }
3154                    });
3155                } else {
3156                    // Combined leased mode has no dedicated scheduler, so resource cleanup itself
3157                    // is the acknowledgement cut.
3158                    out_buf
3159                        .extend_from_slice(&encode_relay_client_disconnected_ack(disconnect_ack));
3160                }
3161            }
3162        }
3163
3164        MessageType::ClockSync => {
3165            let Some(sync) = decode_payload_or_core_error::<ClockSync>(&msg, out_buf)? else {
3166                return Ok(());
3167            };
3168            if let Err(e) = clock::sync_realtime_unix_nanos(sync.unix_time_nanos) {
3169                eprintln!("clock: failed to sync realtime clock: {e}");
3170            }
3171        }
3172
3173        MessageType::Shutdown => {
3174            // Graceful shutdown — signal all sessions, then ask the guest
3175            // kernel to power off so block-root filesystems can shut down
3176            // cleanly instead of leaving ext4 journal recovery pending.
3177            for session in state
3178                .sessions
3179                .drain()
3180                .map(|(_, session)| session)
3181                .chain(state.detached_sessions.drain().map(|(_, session)| session))
3182            {
3183                let _ = session.send_signal(15); // SIGTERM
3184            }
3185            state.write_sessions.clear();
3186            for (_, worker) in state.bulk_write_workers.drain() {
3187                worker.task.abort();
3188            }
3189            for (_, session) in state.tcp_sessions.drain() {
3190                session.close();
3191            }
3192            state.fs.clear();
3193
3194            // A request is not completion. In handoff mode the foreign init owns the
3195            // shutdown schedule; keep serving until the kernel actually powers off.
3196            // Exiting agentd here can itself tear down the VM and bypass that schedule.
3197            if let Err(error) = request_guest_poweroff() {
3198                eprintln!("agentd: graceful poweroff request failed: {error}");
3199            }
3200        }
3201
3202        _ => {
3203            // Ignore unknown or unexpected message types.
3204        }
3205    }
3206
3207    Ok(())
3208}
3209
3210/// Prepends `/.msb/scripts` to PATH in the exec request's environment.
3211///
3212/// If the request already has a PATH entry, prepends to it. Otherwise
3213/// inherits from agentd's environment and prepends.
3214/// Returns whether a host message should refresh the sandbox idle timer.
3215///
3216/// Maintenance traffic such as clock synchronization and reachability checks
3217/// must not count as user activity, otherwise periodic host tasks would keep an
3218/// idle sandbox alive. `core.touch` is excluded here too because it refreshes
3219/// idleness explicitly in its handler, after its payload has been validated.
3220fn message_refreshes_idle_timer(t: &MessageType) -> bool {
3221    !matches!(
3222        t,
3223        MessageType::ClockSync
3224            | MessageType::Ping
3225            | MessageType::Touch
3226            | MessageType::WorkloadFreeze
3227            | MessageType::WorkloadThaw
3228    )
3229}
3230
3231/// Returns whether an agent reply should refresh the sandbox idle timer.
3232///
3233/// Most guest output still represents useful sandbox activity. Maintenance
3234/// replies to `core.ping` and `core.touch` are excluded so `ping` is a pure
3235/// health check and `touch` advances activity exactly once. `core.error` is
3236/// also excluded because valid work already records activity on the incoming
3237/// request, while malformed maintenance traffic should not become a keepalive.
3238fn guest_message_refreshes_idle_timer(t: &MessageType) -> bool {
3239    !matches!(
3240        t,
3241        MessageType::Pong
3242            | MessageType::Touched
3243            | MessageType::WorkloadFrozen
3244            | MessageType::WorkloadThawed
3245            | MessageType::WorkloadTransportCredit
3246            | MessageType::CoreError
3247    )
3248}
3249
3250/// Spawns the heartbeat pulse on a dedicated OS thread.
3251///
3252/// This thread is intentionally outside the Tokio runtime: it reads the latest
3253/// [`HeartbeatSnapshot`] (a lock-free `watch` borrow) and writes the heartbeat
3254/// file with blocking `std::fs` once per [`HEARTBEAT_INTERVAL_SECS`]. Because it
3255/// is an ordinary kernel-scheduled thread, a CPU-bound or I/O-saturated async
3256/// runtime cannot delay the pulse — which is exactly the starvation that made
3257/// the host kill busy-but-healthy sandboxes. The sleep is chunked so the thread
3258/// observes the shutdown flag promptly when the agent loop exits.
3259fn spawn_heartbeat_thread(
3260    snapshot_rx: watch::Receiver<HeartbeatSnapshot>,
3261    shutdown: Arc<AtomicBool>,
3262    control: Arc<heartbeat::HeartbeatControl>,
3263) -> std::thread::JoinHandle<()> {
3264    std::thread::Builder::new()
3265        .name("agentd-heartbeat".to_string())
3266        .spawn(move || {
3267            let mut heartbeat_seq = 0u64;
3268            let mut last_activity_seq = snapshot_rx.borrow().activity_seq;
3269            let mut last_activity = Utc::now();
3270
3271            let interval = Duration::from_secs(HEARTBEAT_INTERVAL_SECS);
3272            let step = Duration::from_millis(100);
3273
3274            while !shutdown.load(Ordering::Relaxed) {
3275                let mut slept = Duration::ZERO;
3276                while slept < interval {
3277                    if shutdown.load(Ordering::Relaxed) {
3278                        return;
3279                    }
3280                    std::thread::sleep(step);
3281                    slept += step;
3282                }
3283
3284                if !heartbeat::heartbeat_dir_exists() {
3285                    continue;
3286                }
3287
3288                heartbeat_seq = heartbeat_seq.saturating_add(1);
3289                let snapshot = snapshot_rx.borrow().clone();
3290                let timestamp = Utc::now();
3291                if snapshot.activity_seq != last_activity_seq {
3292                    last_activity_seq = snapshot.activity_seq;
3293                    last_activity = timestamp;
3294                }
3295                let heartbeat = Heartbeat {
3296                    heartbeat_seq,
3297                    activity_seq: snapshot.activity_seq,
3298                    timestamp,
3299                    last_activity,
3300                    active_exec_sessions: snapshot.active_exec_sessions,
3301                    active_fs_streams: snapshot.active_fs_streams,
3302                    active_tcp_streams: snapshot.active_tcp_streams,
3303                    activity_counters: snapshot.counters,
3304                };
3305                let Some(_write_guard) = control.begin_write() else {
3306                    continue;
3307                };
3308                let _ = heartbeat::write_heartbeat(&heartbeat);
3309            }
3310        })
3311        .expect("failed to spawn agentd heartbeat thread")
3312}
3313
3314fn heartbeat_snapshot(state: &AgentState, activity: &ActivityTracker) -> HeartbeatSnapshot {
3315    HeartbeatSnapshot {
3316        activity_seq: activity.activity_seq,
3317        active_exec_sessions: (state.sessions.len() + state.detached_sessions.len()) as u32,
3318        active_fs_streams: state
3319            .read_sessions
3320            .len()
3321            .saturating_add(state.write_sessions.len())
3322            .saturating_add(state.bulk_write_workers.len()) as u32,
3323        active_tcp_streams: state.tcp_sessions.len() as u32,
3324        counters: activity.counters,
3325    }
3326}
3327
3328fn initial_input_credit() -> WorkloadTransportCredit {
3329    WorkloadTransportCredit {
3330        control_bytes: WORKLOAD_TRANSPORT_CONTROL_BYTES,
3331        control_frames: WORKLOAD_TRANSPORT_CONTROL_FRAMES,
3332        bulk_bytes: WORKLOAD_TRANSPORT_BULK_BYTES,
3333        bulk_frames: WORKLOAD_TRANSPORT_BULK_FRAMES,
3334    }
3335}
3336
3337fn bulk_wire_bytes(record: &BulkRecord, dedicated: bool) -> usize {
3338    // The counter follows complete wire records, not destination payload consumption. Combined
3339    // raw frames use the same bulk allowance without the dedicated lane's incarnation prefix.
3340    4 + FRAME_HEADER_SIZE
3341        + BULK_HEADER_SIZE
3342        + record.payload.len()
3343        + if dedicated {
3344            CLIENT_INCARNATION_SIZE
3345        } else {
3346            0
3347        }
3348}
3349
3350fn private_lifecycle_message(message: &Message) -> bool {
3351    message.id == u32::MAX
3352        && matches!(
3353            message.t,
3354            MessageType::WorkloadFreeze | MessageType::WorkloadThaw
3355        )
3356        && message.flags == message.t.flags()
3357}
3358
3359fn input_credit_update_due(
3360    previous: &WorkloadTransportCredit,
3361    current: &WorkloadTransportCredit,
3362) -> bool {
3363    current.control_bytes.saturating_sub(previous.control_bytes) >= 64 * 1024
3364        || current.bulk_bytes.saturating_sub(previous.bulk_bytes) >= 64 * 1024
3365        || current
3366            .control_frames
3367            .saturating_sub(previous.control_frames)
3368            >= 16
3369        || current.bulk_frames.saturating_sub(previous.bulk_frames) >= 16
3370}
3371
3372async fn wait_input_credit_deadline(deadline: Option<time::Instant>) {
3373    match deadline {
3374        Some(deadline) => time::sleep_until(deadline).await,
3375        None => std::future::pending().await,
3376    }
3377}
3378
3379fn encode_input_credit(credit: WorkloadTransportCredit, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3380    // The host intercepts the reserved runtime correlation before SDK output admission, so a
3381    // stalled client cannot stop refunds needed by unrelated stdin and bulk producers.
3382    let message = Message::with_payload(MessageType::WorkloadTransportCredit, u32::MAX, &credit)?;
3383    codec::encode_to_buf(&message, out_buf)?;
3384    Ok(())
3385}
3386
3387fn poll_pending_stdin(
3388    state: &mut AgentState,
3389    cx: &mut Context<'_>,
3390) -> Poll<(u32, bool, std::io::Result<()>)> {
3391    // Rotate the first descriptor considered so a continuously writable active session cannot
3392    // starve inherited input. The cursor needs no per-payload task or additional ownership queue.
3393    let total = state.sessions.len() + state.detached_sessions.len();
3394    let start = state.stdin_poll_offset.checked_rem(total).unwrap_or(0);
3395    for pass in 0..2 {
3396        let sessions = state
3397            .sessions
3398            .iter_mut()
3399            .map(|(id, session)| (*id, false, session))
3400            .chain(
3401                state
3402                    .detached_sessions
3403                    .iter_mut()
3404                    .map(|((_, id), session)| (*id, true, session)),
3405            );
3406        for (index, (id, inherited, session)) in sessions.enumerate() {
3407            if (pass == 0 && index < start) || (pass == 1 && index >= start) {
3408                continue;
3409            }
3410            if session.has_pending_stdin()
3411                && let Poll::Ready(result) = session.poll_pending_stdin(cx)
3412            {
3413                state.stdin_poll_offset = index + 1;
3414                return Poll::Ready((id, inherited, result));
3415            }
3416        }
3417    }
3418    Poll::Pending
3419}
3420
3421fn encode_stdin_error(id: u32, error: &AgentdError, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3422    let message =
3423        Message::with_payload(MessageType::ExecStdinError, id, &stdin_error_payload(error))?;
3424    codec::encode_to_buf(&message, out_buf)?;
3425    Ok(())
3426}
3427
3428fn publish_heartbeat_snapshot(
3429    heartbeat_tx: &watch::Sender<HeartbeatSnapshot>,
3430    state: &AgentState,
3431    activity: &ActivityTracker,
3432) {
3433    let _ = heartbeat_tx.send(heartbeat_snapshot(state, activity));
3434}
3435
3436fn record_encoded_guest_messages(out_buf: &[u8], start: usize, activity: &mut ActivityTracker) {
3437    let mut offset = start;
3438    while offset + 4 <= out_buf.len() {
3439        let frame_len = u32::from_be_bytes([
3440            out_buf[offset],
3441            out_buf[offset + 1],
3442            out_buf[offset + 2],
3443            out_buf[offset + 3],
3444        ]) as usize;
3445        let total = 4usize.saturating_add(frame_len);
3446        if offset.saturating_add(total) > out_buf.len() {
3447            break;
3448        }
3449
3450        if encoded_guest_message_refreshes_idle_timer(out_buf, offset, frame_len) {
3451            activity.record_guest_message();
3452        }
3453        offset += total;
3454    }
3455}
3456
3457fn encoded_guest_message_refreshes_idle_timer(
3458    out_buf: &[u8],
3459    offset: usize,
3460    frame_len: usize,
3461) -> bool {
3462    let frame_end = offset.saturating_add(4).saturating_add(frame_len);
3463    if frame_end <= out_buf.len() {
3464        let mut frame = BytesMut::from(&out_buf[offset..frame_end]);
3465        if try_decode_relay_client_disconnected_ack_from_bytes(&mut frame)
3466            .is_ok_and(|ack| ack.is_some())
3467        {
3468            return false;
3469        }
3470    }
3471    if frame_len < microsandbox_protocol::message::FRAME_HEADER_SIZE {
3472        return true;
3473    }
3474
3475    let id_start = offset + 4;
3476    let flags_index = id_start + 4;
3477    let body_start = flags_index + 1;
3478    let body_end = offset + 4 + frame_len;
3479    if body_end > out_buf.len() || body_start > body_end {
3480        return true;
3481    }
3482
3483    let id = u32::from_be_bytes([
3484        out_buf[id_start],
3485        out_buf[id_start + 1],
3486        out_buf[id_start + 2],
3487        out_buf[id_start + 3],
3488    ]);
3489    let frame = codec::RawFrame {
3490        id,
3491        flags: out_buf[flags_index],
3492        body: out_buf[body_start..body_end].to_vec(),
3493    };
3494
3495    codec::raw_frame_to_message(frame)
3496        .map(|msg| guest_message_refreshes_idle_timer(&msg.t))
3497        .unwrap_or(true)
3498}
3499
3500fn apply_raw_activity(raw: RawActivity, activity: &mut ActivityTracker) {
3501    if raw.guest_messages != 0 {
3502        activity.activity_seq = activity
3503            .activity_seq
3504            .saturating_add(raw.guest_messages as u64);
3505        activity.counters.guest_messages = activity
3506            .counters
3507            .guest_messages
3508            .saturating_add(raw.guest_messages as u64);
3509    }
3510    if raw.fs_bytes > 0 {
3511        activity.add_fs_bytes(raw.fs_bytes);
3512    }
3513    if raw.tcp_bytes > 0 {
3514        activity.add_tcp_bytes(raw.tcp_bytes);
3515    }
3516}
3517
3518fn complete_raw_session(
3519    id: u32,
3520    completion: Option<RawSessionCompletion>,
3521    read_sessions: &mut HashMap<u32, FsReadSession>,
3522    tcp_sessions: &mut HashMap<u32, TcpSession>,
3523) {
3524    match completion {
3525        Some(RawSessionCompletion::FsRead) => {
3526            read_sessions.remove(&id);
3527        }
3528        Some(RawSessionCompletion::FsWrite) => {}
3529        Some(RawSessionCompletion::Tcp) => {
3530            tcp_sessions.remove(&id);
3531        }
3532        None => {}
3533    }
3534}
3535
3536fn abort_read_sessions_in_owner_range(
3537    read_sessions: &mut HashMap<u32, FsReadSession>,
3538    id_start: u32,
3539    id_end_exclusive: u32,
3540) {
3541    let mut retained = HashMap::new();
3542    for (id, session) in read_sessions.drain() {
3543        let owner_id = session.owner_id();
3544        if owner_id >= id_start && owner_id < id_end_exclusive {
3545            session.abort();
3546        } else {
3547            retained.insert(id, session);
3548        }
3549    }
3550    *read_sessions = retained;
3551}
3552
3553fn close_tcp_sessions_in_owner_range(
3554    tcp_sessions: &mut HashMap<u32, TcpSession>,
3555    id_start: u32,
3556    id_end_exclusive: u32,
3557) {
3558    let mut retained = HashMap::new();
3559    for (id, session) in tcp_sessions.drain() {
3560        let owner_id = session.owner_id();
3561        if owner_id >= id_start && owner_id < id_end_exclusive {
3562            session.close();
3563        } else {
3564            retained.insert(id, session);
3565        }
3566    }
3567    *tcp_sessions = retained;
3568}
3569
3570fn encode_tcp_failed(id: u32, error: String, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3571    let reply = Message::with_payload(MessageType::TcpFailed, id, &TcpFailed { error })
3572        .map_err(|e| AgentdError::ExecSession(format!("encode tcp failed: {e}")))?;
3573    codec::encode_to_buf(&reply, out_buf)
3574        .map_err(|e| AgentdError::ExecSession(format!("encode tcp failed frame: {e}")))?;
3575    Ok(())
3576}
3577
3578fn encode_bulk_fs_failure(id: u32, error: String, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3579    encode_bulk_cancel(
3580        id,
3581        BulkKind::Filesystem,
3582        BulkCancelReason::ProtocolState,
3583        error.clone(),
3584        out_buf,
3585    )?;
3586    let response = Message::with_payload(
3587        MessageType::FsResponse,
3588        id,
3589        &FsResponse {
3590            ok: false,
3591            error: Some(error),
3592            data: None,
3593        },
3594    )
3595    .map_err(|error| AgentdError::ExecSession(format!("encode fs failure: {error}")))?;
3596    codec::encode_to_buf(&response, out_buf)
3597        .map_err(|error| AgentdError::ExecSession(format!("encode fs failure frame: {error}")))
3598}
3599
3600fn encode_bulk_tcp_failure(id: u32, error: String, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3601    encode_bulk_cancel(
3602        id,
3603        BulkKind::Tcp,
3604        BulkCancelReason::ProtocolState,
3605        error.clone(),
3606        out_buf,
3607    )?;
3608    encode_tcp_failed(id, error, out_buf)
3609}
3610
3611/// Emit the operation family's existing terminal failure without recursively sending a cancel.
3612fn encode_bulk_terminal_failure(
3613    id: u32,
3614    kind: BulkKind,
3615    error: String,
3616    out_buf: &mut Vec<u8>,
3617) -> AgentdResult<()> {
3618    match kind {
3619        BulkKind::Filesystem => {
3620            let response = Message::with_payload(
3621                MessageType::FsResponse,
3622                id,
3623                &FsResponse {
3624                    ok: false,
3625                    error: Some(error),
3626                    data: None,
3627                },
3628            )
3629            .map_err(|error| {
3630                AgentdError::ExecSession(format!("encode filesystem terminal failure: {error}"))
3631            })?;
3632            codec::encode_to_buf(&response, out_buf).map_err(|error| {
3633                AgentdError::ExecSession(format!(
3634                    "encode filesystem terminal failure frame: {error}"
3635                ))
3636            })
3637        }
3638        BulkKind::Tcp => encode_tcp_failed(id, error, out_buf),
3639    }
3640}
3641
3642fn encode_bulk_cancel(
3643    id: u32,
3644    kind: BulkKind,
3645    reason: BulkCancelReason,
3646    message: String,
3647    out_buf: &mut Vec<u8>,
3648) -> AgentdResult<()> {
3649    let cancel = Message::with_payload(
3650        MessageType::BulkCancel,
3651        id,
3652        &BulkCancel {
3653            kind,
3654            reason,
3655            message,
3656        },
3657    )
3658    .map_err(|error| AgentdError::ExecSession(format!("encode bulk cancel: {error}")))?;
3659    codec::encode_to_buf(&cancel, out_buf)
3660        .map_err(|error| AgentdError::ExecSession(format!("encode bulk cancel frame: {error}")))
3661}
3662
3663fn encode_core_error_if_supported(
3664    source: &Message,
3665    id: u32,
3666    kind: CoreErrorKind,
3667    message: String,
3668    offending_type: Option<String>,
3669    out_buf: &mut Vec<u8>,
3670) -> AgentdResult<()> {
3671    if !MessageType::CoreError.is_available_at(source.v) {
3672        return Err(AgentdError::ExecSession(format!(
3673            "cannot send core.error to protocol generation {}",
3674            source.v
3675        )));
3676    }
3677
3678    encode_core_error(id, kind, message, offending_type, out_buf)
3679}
3680
3681fn encode_core_error(
3682    id: u32,
3683    kind: CoreErrorKind,
3684    message: String,
3685    offending_type: Option<String>,
3686    out_buf: &mut Vec<u8>,
3687) -> AgentdResult<()> {
3688    let reply = Message::with_payload(
3689        MessageType::CoreError,
3690        id,
3691        &CoreError {
3692            kind,
3693            message,
3694            offending_type,
3695            init_failure: None,
3696            workload_failure: None,
3697        },
3698    )
3699    .map_err(|e| AgentdError::ExecSession(format!("encode core error: {e}")))?;
3700    codec::encode_to_buf(&reply, out_buf)
3701        .map_err(|e| AgentdError::ExecSession(format!("encode core error frame: {e}")))?;
3702    Ok(())
3703}
3704
3705fn encode_workload_error(
3706    source: &Message,
3707    attempt_id: &str,
3708    error: WorkloadLatchError,
3709    out_buf: &mut Vec<u8>,
3710) -> AgentdResult<()> {
3711    let kind = match &error {
3712        WorkloadLatchError::Unavailable(_) | WorkloadLatchError::Io(_) => {
3713            CoreErrorKind::CapabilityUnavailable
3714        }
3715        WorkloadLatchError::InvalidAttempt(_) => CoreErrorKind::InvalidPayload,
3716        WorkloadLatchError::Conflict(_) => CoreErrorKind::InvalidSession,
3717    };
3718    // Keep the existing error category readable by older hosts. Only the additive,
3719    // attempt-scoped detail proves that a basic-pause fallback is safe.
3720    let disposition = match &error {
3721        WorkloadLatchError::Unavailable(_) => WorkloadFailureDisposition::Unavailable,
3722        _ => WorkloadFailureDisposition::RecoveryRequired,
3723    };
3724    if !MessageType::CoreError.is_available_at(source.v) {
3725        return Err(AgentdError::ExecSession(
3726            "peer cannot receive workload errors".into(),
3727        ));
3728    }
3729    let reply = Message::with_payload(
3730        MessageType::CoreError,
3731        source.id,
3732        &CoreError {
3733            kind,
3734            message: error.to_string(),
3735            offending_type: Some(source.t.as_str().to_string()),
3736            init_failure: None,
3737            workload_failure: Some(WorkloadFailure {
3738                attempt_id: attempt_id.to_string(),
3739                disposition,
3740            }),
3741        },
3742    )
3743    .map_err(|error| AgentdError::ExecSession(format!("encode workload error: {error}")))?;
3744    codec::encode_to_buf(&reply, out_buf)
3745        .map_err(|error| AgentdError::ExecSession(format!("encode workload error frame: {error}")))
3746}
3747
3748fn encode_exec_failed(id: u32, payload: ExecFailed, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3749    let reply = Message::with_payload(MessageType::ExecFailed, id, &payload)
3750        .map_err(|error| AgentdError::ExecSession(format!("encode exec failure: {error}")))?;
3751    codec::encode_to_buf(&reply, out_buf)
3752        .map_err(|error| AgentdError::ExecSession(format!("encode exec failure frame: {error}")))?;
3753    Ok(())
3754}
3755
3756fn decode_payload_or_core_error<T>(msg: &Message, out_buf: &mut Vec<u8>) -> AgentdResult<Option<T>>
3757where
3758    T: serde::de::DeserializeOwned,
3759{
3760    match msg.payload::<T>() {
3761        Ok(payload) => Ok(Some(payload)),
3762        Err(error) => {
3763            encode_core_error_if_supported(
3764                msg,
3765                msg.id,
3766                CoreErrorKind::InvalidPayload,
3767                format!("decode payload for {}: {error}", msg.t.as_str()),
3768                Some(msg.t.as_str().to_string()),
3769                out_buf,
3770            )?;
3771            Ok(None)
3772        }
3773    }
3774}
3775
3776/// Build an `ExecStdinError` payload from a failed `write_stdin` result.
3777fn stdin_error_payload(err: &AgentdError) -> ExecStdinError {
3778    let io_err = match err {
3779        AgentdError::Io(e) => Some(e),
3780        _ => None,
3781    };
3782    let errno = io_err.and_then(|e| e.raw_os_error());
3783    ExecStdinError {
3784        errno,
3785        errno_name: errno.and_then(errno_name),
3786        message: err.to_string(),
3787    }
3788}
3789
3790/// Map common errno values to their standard names. Returns `None` for
3791/// codes we don't recognize; callers fall back to the numeric `errno`.
3792fn errno_name(code: i32) -> Option<String> {
3793    let name = match code {
3794        libc::EPIPE => "EPIPE",
3795        libc::EBADF => "EBADF",
3796        libc::EINVAL => "EINVAL",
3797        libc::EIO => "EIO",
3798        libc::ENOSPC => "ENOSPC",
3799        libc::EFBIG => "EFBIG",
3800        _ => return None,
3801    };
3802    Some(name.to_string())
3803}
3804
3805fn prepend_scripts_to_path(req: &mut microsandbox_protocol::exec::ExecRequest) {
3806    // Check if the request already specifies PATH.
3807    if let Some(entry) = req.env.iter_mut().find(|e| e.starts_with("PATH=")) {
3808        let existing = &entry["PATH=".len()..];
3809        *entry = format!("PATH={}", scripts_path(Some(existing)));
3810    } else {
3811        let inherited = env::var("PATH").ok();
3812        req.env
3813            .push(format!("PATH={}", scripts_path(inherited.as_deref())));
3814    }
3815}
3816
3817/// Sets a file descriptor to non-blocking mode.
3818fn set_nonblocking(fd: i32) -> AgentdResult<()> {
3819    let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
3820    if flags < 0 {
3821        return Err(std::io::Error::last_os_error().into());
3822    }
3823    let ret = unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) };
3824    if ret < 0 {
3825        return Err(std::io::Error::last_os_error().into());
3826    }
3827    Ok(())
3828}
3829
3830fn cmdline_requests_dual_port(cmdline: &str) -> bool {
3831    cmdline
3832        .split_ascii_whitespace()
3833        .any(|argument| argument == AGENT_TRANSPORT_DUAL_PORT_CMDLINE)
3834}
3835
3836fn random_connection_id() -> AgentdResult<[u8; 16]> {
3837    let mut id = [0u8; 16];
3838    let mut filled = 0;
3839    while filled < id.len() {
3840        let result =
3841            unsafe { libc::getrandom(id[filled..].as_mut_ptr().cast(), id.len() - filled, 0) };
3842        if result > 0 {
3843            filled += result as usize;
3844            continue;
3845        }
3846        let error = std::io::Error::last_os_error();
3847        if error.kind() == std::io::ErrorKind::Interrupted {
3848            continue;
3849        }
3850        return Err(error.into());
3851    }
3852    Ok(id)
3853}
3854
3855fn read_exact_from_fd(
3856    fd: i32,
3857    mut buf: &mut [u8],
3858    deadline: Instant,
3859    label: &str,
3860) -> AgentdResult<()> {
3861    while !buf.is_empty() {
3862        if !poll_fd_until(fd, libc::POLLIN, Some(deadline))? {
3863            return Err(AgentdError::ExecSession(format!(
3864                "timed out waiting for {label}"
3865            )));
3866        }
3867        match read_from_fd(fd, buf) {
3868            Ok(0) => {
3869                return Err(AgentdError::ExecSession(format!(
3870                    "serial port closed while waiting for {label}"
3871                )));
3872            }
3873            Ok(read) => {
3874                let (_, remainder) = std::mem::take(&mut buf).split_at_mut(read);
3875                buf = remainder;
3876            }
3877            Err(error)
3878                if matches!(
3879                    error.kind(),
3880                    std::io::ErrorKind::Interrupted | std::io::ErrorKind::WouldBlock
3881                ) => {}
3882            Err(error) => return Err(error.into()),
3883        }
3884    }
3885    Ok(())
3886}
3887
3888fn init_ack_deadline() -> Instant {
3889    Instant::now() + std::time::Duration::from_secs(INIT_ACK_TIMEOUT_SECS)
3890}
3891
3892fn init_ack_timeout() -> AgentdError {
3893    AgentdError::ExecSession("timed out waiting for init ack".into())
3894}
3895
3896fn wait_for_init_ack(
3897    fd: i32,
3898    boot_console: &mut BootConsoleState,
3899    deadline: Instant,
3900) -> AgentdResult<()> {
3901    let msg = read_boot_message(fd, boot_console, Some(deadline), "init ack")?;
3902    if msg.t == MessageType::InitAck {
3903        let _: InitAck = msg
3904            .payload()
3905            .map_err(|e| AgentdError::ExecSession(format!("decode init ack payload: {e}")))?;
3906        return Ok(());
3907    }
3908
3909    Err(AgentdError::ExecSession(format!(
3910        "expected core.init.ack, got {}",
3911        msg.t.as_str()
3912    )))
3913}
3914
3915fn read_boot_message(
3916    fd: i32,
3917    state: &mut BootConsoleState,
3918    deadline: Option<Instant>,
3919    context: &str,
3920) -> AgentdResult<Message> {
3921    let mut read_buf = [0u8; 4096];
3922    loop {
3923        if let Some(msg) = codec::try_decode_from_buf(&mut state.input)
3924            .map_err(|e| AgentdError::ExecSession(format!("decode {context}: {e}")))?
3925        {
3926            return Ok(msg);
3927        }
3928        if state.input.len() > MAX_INPUT_BUF_SIZE {
3929            return Err(AgentdError::ExecSession(format!(
3930                "serial input buffer exceeded maximum size while waiting for {context}"
3931            )));
3932        }
3933        if !poll_fd_until(fd, libc::POLLIN, deadline)? {
3934            return Err(if context == "init ack" {
3935                init_ack_timeout()
3936            } else {
3937                AgentdError::ExecSession(format!("timed out waiting for {context}"))
3938            });
3939        }
3940        let n = match read_from_fd(fd, &mut read_buf) {
3941            Ok(n) => n,
3942            Err(error)
3943                if matches!(
3944                    error.kind(),
3945                    std::io::ErrorKind::Interrupted | std::io::ErrorKind::WouldBlock
3946                ) =>
3947            {
3948                continue;
3949            }
3950            Err(error) => return Err(error.into()),
3951        };
3952        if n == 0 {
3953            return Err(AgentdError::ExecSession(format!(
3954                "serial port closed while waiting for {context}"
3955            )));
3956        }
3957        state.input.extend_from_slice(&read_buf[..n]);
3958    }
3959}
3960
3961fn poll_fd_until(fd: i32, events: i16, deadline: Option<Instant>) -> AgentdResult<bool> {
3962    loop {
3963        let timeout_ms = match deadline {
3964            Some(deadline) => {
3965                let remaining = deadline.saturating_duration_since(Instant::now());
3966                if remaining.is_zero() {
3967                    return Ok(false);
3968                }
3969                remaining.as_millis().clamp(1, i32::MAX as u128) as i32
3970            }
3971            None => -1,
3972        };
3973        let mut pfd = libc::pollfd {
3974            fd,
3975            events,
3976            revents: 0,
3977        };
3978        let ret = unsafe { libc::poll(&mut pfd, 1, timeout_ms) };
3979        if ret > 0 {
3980            return Ok(true);
3981        }
3982        if ret == 0 {
3983            return Ok(false);
3984        }
3985        let err = std::io::Error::last_os_error();
3986        if err.raw_os_error() == Some(libc::EINTR) {
3987            continue;
3988        }
3989        return Err(err.into());
3990    }
3991}
3992
3993/// Reads from a raw fd (non-blocking).
3994fn read_from_fd(fd: i32, buf: &mut [u8]) -> std::io::Result<usize> {
3995    let n = unsafe { libc::read(fd, buf.as_mut_ptr() as *mut libc::c_void, buf.len()) };
3996    if n < 0 {
3997        Err(std::io::Error::last_os_error())
3998    } else {
3999        Ok(n as usize)
4000    }
4001}
4002
4003fn write_all_to_fd(fd: i32, mut buf: &[u8], deadline: Option<Instant>) -> AgentdResult<()> {
4004    while !buf.is_empty() {
4005        match write_to_fd(fd, buf) {
4006            Ok(0) => return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into()),
4007            Ok(n) => buf = &buf[n..],
4008            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
4009            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
4010                if !poll_fd_until(fd, libc::POLLOUT, deadline)? {
4011                    return Err(init_ack_timeout());
4012                }
4013            }
4014            Err(e) => return Err(e.into()),
4015        }
4016    }
4017
4018    Ok(())
4019}
4020
4021/// Flushes the write buffer to the async fd.
4022async fn flush_write_buf(fd: &AsyncFd<std::fs::File>, buf: &mut Vec<u8>) -> AgentdResult<()> {
4023    write_all_async_fd(fd, buf).await?;
4024    buf.clear();
4025    Ok(())
4026}
4027
4028/// Write an immutable region to the nonblocking serial descriptor with cursor advancement.
4029async fn write_all_async_fd(fd: &AsyncFd<std::fs::File>, buf: &[u8]) -> AgentdResult<()> {
4030    let mut written = 0;
4031    while written < buf.len() {
4032        let mut guard = fd.writable().await?;
4033        match guard.try_io(|inner| write_to_fd(inner.get_ref().as_raw_fd(), &buf[written..])) {
4034            Ok(Ok(n)) => {
4035                if n == 0 {
4036                    return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into());
4037                }
4038                written += n;
4039            }
4040            Ok(Err(e)) if e.kind() == std::io::ErrorKind::Interrupted => continue,
4041            Ok(Err(e)) => return Err(e.into()),
4042            Err(_would_block) => continue,
4043        }
4044    }
4045    Ok(())
4046}
4047
4048/// Writes a raw bulk header and payload with cursor-safe `writev` calls.
4049async fn write_bulk_record_async_fd(
4050    fd: &AsyncFd<std::fs::File>,
4051    record: &BulkRecord,
4052) -> AgentdResult<()> {
4053    let header = codec::encode_bulk_header(record)
4054        .map_err(|error| AgentdError::ExecSession(format!("encode bulk header: {error}")))?;
4055    write_bulk_parts_async_fd(fd, &header, &record.payload).await
4056}
4057
4058/// Prefix a dedicated-lane record without copying its opaque payload.
4059async fn write_incarnated_bulk_record_async_fd(
4060    fd: &AsyncFd<std::fs::File>,
4061    incarnation: ClientIncarnation,
4062    record: &BulkRecord,
4063) -> AgentdResult<()> {
4064    let public_header = codec::encode_bulk_header(record)
4065        .map_err(|error| AgentdError::ExecSession(format!("encode bulk header: {error}")))?;
4066    let mut header = [0u8; CLIENT_INCARNATION_SIZE + 4 + FRAME_HEADER_SIZE + BULK_HEADER_SIZE];
4067    header[..CLIENT_INCARNATION_SIZE].copy_from_slice(&incarnation);
4068    header[CLIENT_INCARNATION_SIZE..].copy_from_slice(&public_header);
4069    write_bulk_parts_async_fd(fd, &header, &record.payload).await
4070}
4071
4072async fn write_bulk_parts_async_fd(
4073    fd: &AsyncFd<std::fs::File>,
4074    header: &[u8],
4075    payload: &[u8],
4076) -> AgentdResult<()> {
4077    let mut header_offset = 0;
4078    let mut payload_offset = 0;
4079
4080    while header_offset < header.len() || payload_offset < payload.len() {
4081        let mut guard = fd.writable().await?;
4082        let result = guard.try_io(|inner| {
4083            write_vectored_to_fd(
4084                inner.get_ref().as_raw_fd(),
4085                &header[header_offset..],
4086                &payload[payload_offset..],
4087            )
4088        });
4089        let written = match result {
4090            Ok(Ok(0)) => {
4091                return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into());
4092            }
4093            Ok(Ok(written)) => written,
4094            Ok(Err(error)) if error.kind() == std::io::ErrorKind::Interrupted => continue,
4095            Ok(Err(error)) => return Err(error.into()),
4096            Err(_would_block) => continue,
4097        };
4098
4099        let header_remaining = header.len() - header_offset;
4100        if written < header_remaining {
4101            header_offset += written;
4102        } else {
4103            header_offset = header.len();
4104            payload_offset += written - header_remaining;
4105        }
4106    }
4107
4108    Ok(())
4109}
4110
4111/// Writes to a raw fd (non-blocking).
4112fn write_to_fd(fd: i32, buf: &[u8]) -> std::io::Result<usize> {
4113    let n = unsafe { libc::write(fd, buf.as_ptr() as *const libc::c_void, buf.len()) };
4114    if n < 0 {
4115        Err(std::io::Error::last_os_error())
4116    } else {
4117        Ok(n as usize)
4118    }
4119}
4120
4121fn write_vectored_to_fd(fd: i32, header: &[u8], payload: &[u8]) -> std::io::Result<usize> {
4122    if header.is_empty() {
4123        return write_to_fd(fd, payload);
4124    }
4125
4126    let vectors = [
4127        libc::iovec {
4128            iov_base: header.as_ptr().cast_mut().cast(),
4129            iov_len: header.len(),
4130        },
4131        libc::iovec {
4132            iov_base: payload.as_ptr().cast_mut().cast(),
4133            iov_len: payload.len(),
4134        },
4135    ];
4136    let vector_count = if payload.is_empty() { 1 } else { 2 };
4137    let written = unsafe { libc::writev(fd, vectors.as_ptr(), vector_count) };
4138    if written < 0 {
4139        Err(std::io::Error::last_os_error())
4140    } else {
4141        Ok(written as usize)
4142    }
4143}
4144
4145fn request_guest_poweroff() -> AgentdResult<()> {
4146    if crate::handoff::is_pid_1() {
4147        // PID 1 mode (no handoff): tear down filesystems so block-backed
4148        // mounts reach a clean terminal state, then power the kernel off.
4149        crate::teardown::teardown_filesystems(true);
4150        let ret = unsafe { libc::reboot(libc::RB_POWER_OFF) };
4151        if ret != 0 {
4152            return Err(std::io::Error::last_os_error().into());
4153        }
4154        return Ok(());
4155    }
4156
4157    unsafe {
4158        libc::sync();
4159    }
4160
4161    // Foreign init owns both shutdown duration and filesystem teardown. No timer or
4162    // fallback signal may convert this graceful request into forced termination.
4163    crate::handoff::signal_init_shutdown()
4164}
4165
4166//--------------------------------------------------------------------------------------------------
4167// Tests
4168//--------------------------------------------------------------------------------------------------
4169
4170#[cfg(test)]
4171mod tests {
4172    use super::*;
4173    use bytes::Bytes;
4174    use microsandbox_protocol::message::PROTOCOL_VERSION;
4175
4176    #[allow(clippy::too_many_arguments)]
4177    async fn handle_message(
4178        msg: Message,
4179        state: &mut AgentState,
4180        activity: &mut ActivityTracker,
4181        sender: &mut SessionOutputSender,
4182        out_buf: &mut Vec<u8>,
4183        config: &AgentdConfig,
4184        workload: &mut WorkloadLatch,
4185        heartbeat: &heartbeat::HeartbeatControl,
4186    ) -> AgentdResult<()> {
4187        handle_message_with_charge(
4188            msg, state, activity, sender, out_buf, config, workload, heartbeat, None,
4189        )
4190        .await
4191    }
4192
4193    fn decode_reply_skipping_credit(bytes: &mut BytesMut) -> Message {
4194        loop {
4195            let Some(DecodedFrame::Control(reply)) =
4196                codec::try_decode_frame_from_bytes(bytes).unwrap()
4197            else {
4198                panic!("missing control reply");
4199            };
4200            if reply.t != MessageType::WorkloadTransportCredit {
4201                return reply;
4202            }
4203            assert_eq!(reply.id, u32::MAX);
4204        }
4205    }
4206
4207    #[test]
4208    fn admitted_transport_window_fits_each_filesystem_input_queue() {
4209        // Each accepted raw record occupies one queue slot until consumed. Even if one flow
4210        // receives the whole allowance from both lanes, admission cannot stall the serial actor.
4211        assert!(
4212            WORKLOAD_TRANSPORT_CONTROL_FRAMES + WORKLOAD_TRANSPORT_BULK_FRAMES
4213                <= FS_BULK_INPUT_ITEM_CAPACITY as u64
4214        );
4215        assert!(WORKLOAD_TRANSPORT_CONTROL_BYTES <= BULK_INPUT_BYTE_CAPACITY as u64);
4216        assert!(WORKLOAD_TRANSPORT_BULK_BYTES <= BULK_INPUT_BYTE_CAPACITY as u64);
4217    }
4218
4219    #[test]
4220    fn private_admission_and_credit_batching_use_the_complete_wire_contract() {
4221        let mut request = Message::with_payload(
4222            MessageType::WorkloadFreeze,
4223            u32::MAX,
4224            &WorkloadFreeze {
4225                external_mount_tags: Vec::new(),
4226                attempt_id: "cut".into(),
4227                host_input: Default::default(),
4228            },
4229        )
4230        .unwrap();
4231        assert!(private_lifecycle_message(&request));
4232        request.id = 0;
4233        assert!(!private_lifecycle_message(&request));
4234        request.id = u32::MAX;
4235        request.t = MessageType::ClockSync;
4236        assert!(!private_lifecycle_message(&request));
4237        let record = BulkRecord {
4238            id: 1,
4239            kind: BulkKind::Tcp,
4240            flow: BulkFlow::HostToGuest,
4241            offset: 0,
4242            payload: Bytes::from_static(b"payload"),
4243        };
4244        let mut encoded = codec::encode_bulk_header(&record).unwrap().to_vec();
4245        encoded.extend_from_slice(&record.payload);
4246        assert_eq!(bulk_wire_bytes(&record, false), encoded.len());
4247        assert_eq!(
4248            bulk_wire_bytes(&record, true),
4249            encoded.len() + CLIENT_INCARNATION_SIZE
4250        );
4251        let previous = initial_input_credit();
4252        let mut current = previous;
4253        current.control_frames += 15;
4254        assert!(!input_credit_update_due(&previous, &current));
4255        current.control_frames += 1;
4256        assert!(input_credit_update_due(&previous, &current));
4257        let mut credit = Vec::new();
4258        encode_input_credit(current, &mut credit).unwrap();
4259        let mut bytes = BytesMut::from(credit.as_slice());
4260        let Some(DecodedFrame::Control(reply)) =
4261            codec::try_decode_frame_from_bytes(&mut bytes).unwrap()
4262        else {
4263            panic!()
4264        };
4265        assert_eq!(reply.id, u32::MAX);
4266        assert_eq!(reply.payload::<WorkloadTransportCredit>().unwrap(), current);
4267        assert!(!guest_message_refreshes_idle_timer(&reply.t));
4268    }
4269
4270    #[tokio::test]
4271    async fn transport_cut_waits_for_decoded_prefix_and_abort_is_owned_and_repeatable() {
4272        use microsandbox_protocol::core::WorkloadThawMode::{Continue, Restore};
4273        let mut state = AgentState::default();
4274        let (mut sender, _output) = SessionOutputSender::channel();
4275        let mut activity = ActivityTracker::new();
4276        let config = AgentdConfig {
4277            user: None,
4278            security_profile: Default::default(),
4279            default_cwd: None,
4280            default_env: Vec::new(),
4281        };
4282        let heartbeat = heartbeat::HeartbeatControl::default();
4283        let mut workload = crate::workload::tests::fake_latch();
4284        let target = WorkloadTransportPosition {
4285            bulk_bytes: 80,
4286            bulk_frames: 1,
4287            ..Default::default()
4288        };
4289        let freeze = || {
4290            Message::with_payload(
4291                MessageType::WorkloadFreeze,
4292                u32::MAX,
4293                &WorkloadFreeze {
4294                    external_mount_tags: Vec::new(),
4295                    attempt_id: "prefix".into(),
4296                    host_input: target,
4297                },
4298            )
4299            .unwrap()
4300        };
4301        for _ in 0..2 {
4302            let mut out = Vec::new();
4303            handle_message(
4304                freeze(),
4305                &mut state,
4306                &mut activity,
4307                &mut sender,
4308                &mut out,
4309                &config,
4310                &mut workload,
4311                &heartbeat,
4312            )
4313            .await
4314            .unwrap();
4315            assert!(out.is_empty());
4316            assert!(!workload.is_frozen());
4317            assert!(state.pending_freeze.is_some());
4318        }
4319        for (attempt, mode, expected) in [
4320            ("other", Continue, MessageType::CoreError),
4321            ("prefix", Restore, MessageType::CoreError),
4322            ("prefix", Continue, MessageType::WorkloadThawed),
4323            ("prefix", Continue, MessageType::WorkloadThawed),
4324        ] {
4325            let thaw = Message::with_payload(
4326                MessageType::WorkloadThaw,
4327                u32::MAX,
4328                &WorkloadThaw {
4329                    attempt_id: attempt.into(),
4330                    mode,
4331                },
4332            )
4333            .unwrap();
4334            let aborts_pending =
4335                expected == MessageType::WorkloadThawed && state.pending_freeze.is_some();
4336            let mut out = Vec::new();
4337            handle_message(
4338                thaw,
4339                &mut state,
4340                &mut activity,
4341                &mut sender,
4342                &mut out,
4343                &config,
4344                &mut workload,
4345                &heartbeat,
4346            )
4347            .await
4348            .unwrap();
4349            let mut bytes = BytesMut::from(out.as_slice());
4350            if aborts_pending {
4351                let cancelled = decode_reply_skipping_credit(&mut bytes);
4352                let error = cancelled.payload::<CoreError>().unwrap();
4353                assert_eq!(cancelled.id, u32::MAX);
4354                assert_eq!(
4355                    error.offending_type.as_deref(),
4356                    Some(MessageType::WorkloadFreeze.as_str())
4357                );
4358                assert_eq!(error.workload_failure.unwrap().attempt_id, "prefix");
4359            }
4360            assert_eq!(decode_reply_skipping_credit(&mut bytes).t, expected);
4361            assert_eq!(
4362                state.pending_freeze.is_some(),
4363                expected == MessageType::CoreError
4364            );
4365        }
4366        let charge = state.input_window.admit(InputLane::Bulk, 80).unwrap();
4367        for _ in 0..2 {
4368            let mut out = Vec::new();
4369            handle_message(
4370                freeze(),
4371                &mut state,
4372                &mut activity,
4373                &mut sender,
4374                &mut out,
4375                &config,
4376                &mut workload,
4377                &heartbeat,
4378            )
4379            .await
4380            .unwrap();
4381            let reply = decode_reply_skipping_credit(&mut BytesMut::from(out.as_slice()));
4382            let frozen = reply.payload::<WorkloadFrozen>().unwrap();
4383            assert_eq!(reply.t, MessageType::WorkloadFrozen);
4384            assert!(workload.is_frozen());
4385            assert!(state.output_parked);
4386            assert_eq!(
4387                frozen.input_credit,
4388                initial_input_credit(),
4389                "decode is not consumption"
4390            );
4391            assert_eq!(state.input_window.position(), target);
4392        }
4393        drop(charge);
4394    }
4395
4396    #[tokio::test]
4397    async fn mixed_primary_data_cut_waits_for_the_complete_dedicated_prefix() {
4398        let mut state = AgentState::default();
4399        let (mut sender, _output) = SessionOutputSender::channel();
4400        let mut activity = ActivityTracker::new();
4401        let config = AgentdConfig {
4402            user: None,
4403            security_profile: Default::default(),
4404            default_cwd: None,
4405            default_env: Vec::new(),
4406        };
4407        let heartbeat = heartbeat::HeartbeatControl::default();
4408        let mut workload = crate::workload::tests::fake_latch();
4409
4410        // Primary stdin and its empty EOF share the logical data ledger with the dedicated
4411        // port. Retain every charge to model a consumer that has not accepted any input yet.
4412        let mut charges = Vec::new();
4413        for data in [vec![0x31; 1024], Vec::new()] {
4414            let message =
4415                Message::with_payload(MessageType::ExecStdin, 1, &ExecStdin { data }).unwrap();
4416            assert!(message.t.uses_workload_data_credit());
4417            let mut wire = Vec::new();
4418            codec::encode_to_buf(&message, &mut wire).unwrap();
4419            charges.push(
4420                state
4421                    .input_window
4422                    .admit(InputLane::Bulk, wire.len())
4423                    .unwrap(),
4424            );
4425        }
4426        let primary_position = state.input_window.position();
4427        assert_eq!(primary_position.control_bytes, 0);
4428        assert_eq!(primary_position.control_frames, 0);
4429        assert_eq!(primary_position.bulk_frames, 2);
4430
4431        let incarnation = [0x43; CLIENT_INCARNATION_SIZE];
4432        let record = BulkRecord {
4433            id: 2,
4434            kind: BulkKind::Filesystem,
4435            flow: BulkFlow::HostToGuest,
4436            offset: 0,
4437            payload: Bytes::from(vec![0x52; 512]),
4438        };
4439        let mut dedicated_wire = incarnation.to_vec();
4440        codec::encode_bulk_to_buf(&record, &mut dedicated_wire).unwrap();
4441        let target = WorkloadTransportPosition {
4442            bulk_bytes: primary_position.bulk_bytes + dedicated_wire.len() as u64,
4443            bulk_frames: primary_position.bulk_frames + 1,
4444            ..primary_position
4445        };
4446        let freeze = Message::with_payload(
4447            MessageType::WorkloadFreeze,
4448            u32::MAX,
4449            &WorkloadFreeze {
4450                external_mount_tags: Vec::new(),
4451                attempt_id: "mixed-prefix".into(),
4452                host_input: target,
4453            },
4454        )
4455        .unwrap();
4456
4457        let mut dedicated_input = BytesMut::new();
4458        for prefix in [
4459            &dedicated_wire[..0],
4460            &dedicated_wire[..dedicated_wire.len() - 1],
4461        ] {
4462            dedicated_input.extend_from_slice(prefix);
4463            assert!(
4464                try_decode_incarnated_bulk_from_bytes(&mut dedicated_input)
4465                    .unwrap()
4466                    .is_none()
4467            );
4468            let mut out = Vec::new();
4469            handle_message(
4470                freeze.clone(),
4471                &mut state,
4472                &mut activity,
4473                &mut sender,
4474                &mut out,
4475                &config,
4476                &mut workload,
4477                &heartbeat,
4478            )
4479            .await
4480            .unwrap();
4481            assert!(
4482                out.is_empty(),
4483                "primary data cannot cover missing dedicated bytes"
4484            );
4485            assert!(!workload.is_frozen());
4486            assert!(!state.output_parked);
4487            assert!(state.pending_freeze.is_some());
4488            assert_eq!(state.input_window.position(), primary_position);
4489        }
4490
4491        dedicated_input.extend_from_slice(&dedicated_wire[dedicated_wire.len() - 1..]);
4492        let decoded = try_decode_incarnated_bulk_from_bytes(&mut dedicated_input)
4493            .unwrap()
4494            .unwrap();
4495        assert_eq!(decoded.incarnation, incarnation);
4496        assert_eq!(decoded.record, record);
4497        assert!(dedicated_input.is_empty());
4498        charges.push(
4499            state
4500                .input_window
4501                .admit(InputLane::Bulk, bulk_wire_bytes(&decoded.record, true))
4502                .unwrap(),
4503        );
4504        assert_eq!(state.input_window.position(), target);
4505
4506        // Resume the retained request as the outer actor does once both cumulative counters
4507        // reach the cut. Decoding suffices; none of the primary or dedicated data is consumed.
4508        let pending = state.pending_freeze.take().unwrap();
4509        let mut out = Vec::new();
4510        handle_message(
4511            pending,
4512            &mut state,
4513            &mut activity,
4514            &mut sender,
4515            &mut out,
4516            &config,
4517            &mut workload,
4518            &heartbeat,
4519        )
4520        .await
4521        .unwrap();
4522        let mut bytes = BytesMut::from(out.as_slice());
4523        let reply = decode_reply_skipping_credit(&mut bytes);
4524        let frozen = reply.payload::<WorkloadFrozen>().unwrap();
4525        assert_eq!(reply.t, MessageType::WorkloadFrozen);
4526        assert_eq!(reply.id, u32::MAX);
4527        assert_eq!(frozen.attempt_id, "mixed-prefix");
4528        assert_eq!(frozen.input_credit, initial_input_credit());
4529        assert_eq!(state.frozen_host_input, Some(target));
4530        assert!(workload.is_frozen());
4531        assert!(state.output_parked);
4532        assert!(state.pending_freeze.is_none());
4533        assert!(bytes.is_empty());
4534        drop(charges);
4535    }
4536
4537    #[tokio::test]
4538    async fn late_tcp_credit_preserves_raw_tail_before_and_after_terminal_retirement() {
4539        use std::collections::hash_map::DefaultHasher;
4540        use std::hash::Hasher;
4541
4542        use microsandbox_protocol::bulk::BulkOffer;
4543        use tokio::io::{AsyncReadExt, AsyncWriteExt};
4544        use tokio::net::TcpListener;
4545
4546        tokio::time::timeout(Duration::from_secs(10), async {
4547            let mut state = AgentState::default();
4548            let incarnation = [0x57; CLIENT_INCARNATION_SIZE];
4549            establish_relay_client(
4550                &mut state,
4551                RelayClientConnected {
4552                    id_start: 1,
4553                    id_end_exclusive: microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP,
4554                    incarnation,
4555                },
4556            )
4557            .unwrap();
4558            // Leave the dedicated scheduler's input queued: primary completion is allowed to
4559            // overtake these bytes, but no late credit may send DropFlow to discard their tail.
4560            let (mut sender, mut control, mut bulk, mut scheduler_commands) =
4561                SessionOutputSender::split_channel();
4562            let producer = sender.with_incarnation(Some(incarnation));
4563            let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
4564            state.tcp_sessions.insert(
4565                1,
4566                TcpSession::open(
4567                    1,
4568                    TcpConnect {
4569                        host: "127.0.0.1".into(),
4570                        port: listener.local_addr().unwrap().port(),
4571                        bulk: Some(BulkOffer::tcp()),
4572                    },
4573                    &producer,
4574                ),
4575            );
4576            let (mut peer, _) = listener.accept().await.unwrap();
4577            for expected in [MessageType::TcpConnected, MessageType::BulkAccepted] {
4578                let envelope = control.recv().await.unwrap();
4579                let SessionOutput::Raw(mut output) = envelope.output else {
4580                    panic!()
4581                };
4582                assert_eq!(
4583                    codec::try_decode_from_buf(&mut output.frame)
4584                        .unwrap()
4585                        .unwrap()
4586                        .t,
4587                    expected
4588                );
4589            }
4590            state
4591                .tcp_sessions
4592                .get(&1)
4593                .unwrap()
4594                .finish_bulk(BulkFinish {
4595                    kind: BulkKind::Tcp,
4596                    flow: BulkFlow::HostToGuest,
4597                    final_offset: 0,
4598                })
4599                .await
4600                .unwrap();
4601            assert_eq!(peer.read(&mut [0]).await.unwrap(), 0);
4602            let payload = Bytes::from(
4603                (0..6 * 1024 * 1024)
4604                    .map(|index| (index % 251) as u8)
4605                    .collect::<Vec<_>>(),
4606            );
4607            let peer_payload = payload.clone();
4608            let peer_task = tokio::spawn(async move {
4609                peer.write_all(&peer_payload).await.unwrap();
4610                peer.shutdown().await.unwrap();
4611            });
4612            let mut received = Vec::new();
4613            let mut consumed = 0;
4614            while consumed < 4 * 1024 * 1024 {
4615                let envelope = bulk.recv().await.unwrap();
4616                let SessionOutput::Bulk(output) = &envelope.output else {
4617                    panic!()
4618                };
4619                consumed += output.record.payload.len();
4620                received.push(envelope);
4621            }
4622            peer_task.await.unwrap();
4623            while !state.tcp_sessions.get(&1).unwrap().is_finished() {
4624                tokio::task::yield_now().await;
4625            }
4626            assert!(consumed < payload.len());
4627            assert!(
4628                !bulk.is_empty(),
4629                "the dedicated output tail must still be queued"
4630            );
4631            let credit = BulkCredit {
4632                kind: BulkKind::Tcp,
4633                flow: BulkFlow::GuestToHost,
4634                consumed_offset: consumed as u64,
4635                credit_limit: consumed as u64 + DEFAULT_BULK_WINDOW,
4636            };
4637            let mut activity = ActivityTracker::new();
4638            let config = AgentdConfig {
4639                user: None,
4640                security_profile: Default::default(),
4641                default_cwd: None,
4642                default_env: Vec::new(),
4643            };
4644            let mut workload = crate::workload::tests::fake_latch();
4645            let heartbeat = heartbeat::HeartbeatControl::default();
4646            for retired in [false, true] {
4647                if retired {
4648                    for expected in [MessageType::BulkFinish, MessageType::TcpClosed] {
4649                        let envelope = control.try_recv().unwrap();
4650                        let SessionOutput::Raw(mut output) = envelope.output else {
4651                            panic!()
4652                        };
4653                        let message = codec::try_decode_from_buf(&mut output.frame)
4654                            .unwrap()
4655                            .unwrap();
4656                        assert_eq!(message.t, expected);
4657                        if expected == MessageType::BulkFinish {
4658                            assert_eq!(
4659                                message.payload::<BulkFinish>().unwrap().final_offset,
4660                                payload.len() as u64
4661                            );
4662                        } else {
4663                            assert_ne!(
4664                                message.flags & microsandbox_protocol::message::FLAG_TERMINAL,
4665                                0
4666                            );
4667                            assert!(matches!(output.completion, Some(RawSessionCompletion::Tcp)));
4668                            complete_raw_session(
4669                                1,
4670                                output.completion,
4671                                &mut state.read_sessions,
4672                                &mut state.tcp_sessions,
4673                            );
4674                            clear_bulk_receive_state(&mut state, 1);
4675                        }
4676                    }
4677                }
4678                assert_eq!(state.tcp_sessions.contains_key(&1), !retired);
4679                let mut out = Vec::new();
4680                handle_message(
4681                    Message::with_payload(MessageType::BulkCredit, 1, &credit).unwrap(),
4682                    &mut state,
4683                    &mut activity,
4684                    &mut sender,
4685                    &mut out,
4686                    &config,
4687                    &mut workload,
4688                    &heartbeat,
4689                )
4690                .await
4691                .unwrap();
4692                assert!(
4693                    out.is_empty(),
4694                    "late credit must not emit cancellation or another terminal"
4695                );
4696                assert!(
4697                    matches!(
4698                        scheduler_commands.try_recv(),
4699                        Err(tokio::sync::mpsc::error::TryRecvError::Empty)
4700                    ),
4701                    "late credit must not purge raw output"
4702                );
4703                assert!(!bulk.is_empty());
4704            }
4705            assert!(matches!(
4706                control.try_recv(),
4707                Err(tokio::sync::mpsc::error::TryRecvError::Empty)
4708            ));
4709            while let Ok(envelope) = bulk.try_recv() {
4710                received.push(envelope);
4711            }
4712            let mut offset = 0;
4713            let mut actual_hash = DefaultHasher::new();
4714            let mut expected_hash = DefaultHasher::new();
4715            expected_hash.write(&payload);
4716            for envelope in received {
4717                assert_eq!(envelope.id, 1);
4718                assert_eq!(envelope.incarnation, Some(incarnation));
4719                let SessionOutput::Bulk(output) = envelope.output else {
4720                    panic!()
4721                };
4722                let record = output.record;
4723                assert_eq!(record.offset, offset as u64);
4724                let end = offset + record.payload.len();
4725                assert_eq!(record.payload.as_ref(), &payload[offset..end]);
4726                actual_hash.write(&record.payload);
4727                offset = end;
4728            }
4729            assert_eq!(offset, payload.len());
4730            assert_eq!(actual_hash.finish(), expected_hash.finish());
4731        })
4732        .await
4733        .expect("late-credit TCP tail did not complete");
4734    }
4735
4736    #[tokio::test]
4737    async fn dedicated_bulk_park_finishes_one_record_and_retains_the_next() {
4738        use std::os::fd::OwnedFd;
4739        use tokio::io::AsyncReadExt;
4740
4741        let (writer, reader) = std::os::unix::net::UnixStream::pair().unwrap();
4742        writer.set_nonblocking(true).unwrap();
4743        reader.set_nonblocking(true).unwrap();
4744        let send_buffer: libc::c_int = 4096;
4745        assert_eq!(
4746            unsafe {
4747                libc::setsockopt(
4748                    writer.as_raw_fd(),
4749                    libc::SOL_SOCKET,
4750                    libc::SO_SNDBUF,
4751                    (&send_buffer as *const libc::c_int).cast(),
4752                    std::mem::size_of_val(&send_buffer) as libc::socklen_t,
4753                )
4754            },
4755            0
4756        );
4757        let file = File::from(OwnedFd::from(writer));
4758        let mut reader = tokio::net::UnixStream::from_std(reader).unwrap();
4759        let (sender, _control, bulk, commands) = SessionOutputSender::split_channel();
4760        let (activity, _activities) = tokio::sync::mpsc::channel(8);
4761        let writer = tokio::spawn(bulk_writer_task(file, bulk, commands, activity));
4762        let sender = sender.with_incarnation(Some([0x62; CLIENT_INCARNATION_SIZE]));
4763        let first = BulkRecord {
4764            id: 1,
4765            kind: BulkKind::Filesystem,
4766            flow: BulkFlow::GuestToHost,
4767            offset: 0,
4768            payload: Bytes::from(vec![0x6a; 1024 * 1024]),
4769        };
4770        let first_len = bulk_wire_bytes(&first, true);
4771        let second = BulkRecord {
4772            offset: first.payload.len() as u64,
4773            payload: Bytes::from_static(b"next"),
4774            ..first.clone()
4775        };
4776        let second_len = bulk_wire_bytes(&second, true);
4777        for record in [first, second] {
4778            assert!(
4779                sender
4780                    .send(
4781                        1,
4782                        SessionOutput::Bulk(crate::session::BulkSessionOutput::new(
4783                            record,
4784                            RawActivity::default()
4785                        ))
4786                    )
4787                    .await
4788            );
4789        }
4790        let mut bytes = vec![0; first_len];
4791        time::timeout(Duration::from_secs(5), reader.read_exact(&mut bytes[..37]))
4792            .await
4793            .unwrap()
4794            .unwrap();
4795        let park_sender = sender.clone();
4796        let mut park = tokio::spawn(async move { park_sender.park_bulk_output().await });
4797        assert!(
4798            time::timeout(Duration::from_millis(20), &mut park)
4799                .await
4800                .is_err(),
4801            "park cut a partial record"
4802        );
4803        time::timeout(Duration::from_secs(5), reader.read_exact(&mut bytes[37..]))
4804            .await
4805            .unwrap()
4806            .unwrap();
4807        assert_eq!(
4808            time::timeout(Duration::from_secs(5), park)
4809                .await
4810                .unwrap()
4811                .unwrap()
4812                .unwrap(),
4813            first_len as u64
4814        );
4815        let mut decoded = BytesMut::from(bytes.as_slice());
4816        let frame = try_decode_incarnated_bulk_from_bytes(&mut decoded)
4817            .unwrap()
4818            .unwrap();
4819        assert_eq!(frame.record.payload, Bytes::from(vec![0x6a; 1024 * 1024]));
4820        assert!(decoded.is_empty());
4821        assert!(
4822            time::timeout(Duration::from_millis(20), reader.read(&mut [0; 1]))
4823                .await
4824                .is_err()
4825        );
4826        sender.resume_bulk_output().await.unwrap();
4827        let mut bytes = vec![0; second_len];
4828        time::timeout(Duration::from_secs(5), reader.read_exact(&mut bytes))
4829            .await
4830            .unwrap()
4831            .unwrap();
4832        assert_eq!(
4833            sender.park_bulk_output().await.unwrap(),
4834            (first_len + second_len) as u64
4835        );
4836        writer.abort();
4837        let _ = writer.await;
4838    }
4839
4840    #[test]
4841    fn init_failure_waits_for_failure_ack_after_late_success_ack() {
4842        use std::io::{Read, Write};
4843        use std::os::fd::OwnedFd;
4844        use std::os::unix::net::UnixStream;
4845
4846        for ack_failure in [false, true] {
4847            let (guest, mut host) = UnixStream::pair().unwrap();
4848            host.set_read_timeout(Some(std::time::Duration::from_secs(5)))
4849                .unwrap();
4850            let guest = File::from(OwnedFd::from(guest));
4851            let worker = std::thread::spawn(move || {
4852                report_init_failure(
4853                    &guest,
4854                    &mut BootConsoleState::default(),
4855                    "missing guest user",
4856                    &AgentdError::UserNotFound("iggy".into()),
4857                )
4858            });
4859            let mut header = [0u8; 4];
4860            host.read_exact(&mut header).unwrap();
4861            let mut frame = vec![0; u32::from_be_bytes(header) as usize];
4862            host.read_exact(&mut frame).unwrap();
4863            let mut input = header.to_vec();
4864            input.extend_from_slice(&frame);
4865            let report = codec::try_decode_from_buf(&mut input).unwrap().unwrap();
4866            assert_eq!(report.t, MessageType::CoreError);
4867            assert_eq!(
4868                report.payload::<CoreError>().unwrap().init_failure,
4869                Some(InitFailureReason::UserNotFound)
4870            );
4871            assert_eq!(
4872                report.payload::<CoreError>().unwrap().message,
4873                "missing guest user"
4874            );
4875
4876            let mut acks = Vec::new();
4877            for failure in [false].into_iter().chain(ack_failure.then_some(true)) {
4878                codec::encode_to_buf(
4879                    &Message::with_payload(MessageType::InitAck, 0, &InitAck { failure }).unwrap(),
4880                    &mut acks,
4881                )
4882                .unwrap();
4883            }
4884            host.write_all(&acks).unwrap();
4885            host.shutdown(std::net::Shutdown::Write).unwrap();
4886            assert_eq!(worker.join().unwrap().is_ok(), ack_failure);
4887        }
4888    }
4889
4890    #[test]
4891    fn init_failure_waits_for_slow_supporting_host_but_bounds_legacy_wait() {
4892        use std::io::Write;
4893        use std::os::fd::OwnedFd;
4894        use std::os::unix::net::UnixStream;
4895        use std::sync::mpsc;
4896
4897        for supported in [false, true] {
4898            let (guest, mut host) = UnixStream::pair().unwrap();
4899            let guest = File::from(OwnedFd::from(guest));
4900            let mut bootstrap = Vec::new();
4901            codec::encode_to_buf(
4902                &Message::with_payload(
4903                    MessageType::Bootstrap,
4904                    0,
4905                    &GuestBootstrap {
4906                        init_failure_ack: supported,
4907                        ..Default::default()
4908                    },
4909                )
4910                .unwrap(),
4911                &mut bootstrap,
4912            )
4913            .unwrap();
4914            host.write_all(&bootstrap).unwrap();
4915            let (done_tx, done_rx) = mpsc::channel();
4916            let worker = std::thread::spawn(move || {
4917                let (_, mut state) = receive_bootstrap(&guest).unwrap();
4918                let result = report_init_failure(
4919                    &guest,
4920                    &mut state,
4921                    "missing user",
4922                    &AgentdError::UserNotFound("iggy".into()),
4923                );
4924                done_tx.send(result).unwrap();
4925            });
4926            let report = read_boot_message(
4927                host.as_raw_fd(),
4928                &mut BootConsoleState::default(),
4929                Some(Instant::now() + std::time::Duration::from_secs(5)),
4930                "startup error",
4931            )
4932            .unwrap();
4933            assert_eq!(report.t, MessageType::CoreError);
4934
4935            let result = done_rx.recv_timeout(std::time::Duration::from_secs(3));
4936            if supported {
4937                assert!(
4938                    matches!(result, Err(mpsc::RecvTimeoutError::Timeout)),
4939                    "supporting host must be allowed to save beyond the legacy timeout"
4940                );
4941                let mut ack = Vec::new();
4942                codec::encode_to_buf(
4943                    &Message::with_payload(MessageType::InitAck, 0, &InitAck { failure: true })
4944                        .unwrap(),
4945                    &mut ack,
4946                )
4947                .unwrap();
4948                host.write_all(&ack).unwrap();
4949                done_rx
4950                    .recv_timeout(std::time::Duration::from_secs(5))
4951                    .unwrap()
4952                    .unwrap();
4953            } else {
4954                assert!(result.expect("legacy host wait must be bounded").is_err());
4955            }
4956            worker.join().unwrap();
4957        }
4958    }
4959
4960    #[test]
4961    fn coalesced_bootstrap_and_init_ack_retain_the_second_frame() {
4962        let bootstrap = GuestBootstrap::default();
4963        let bootstrap_message =
4964            Message::with_payload(MessageType::Bootstrap, 0, &bootstrap).unwrap();
4965        let ack_message =
4966            Message::with_payload(MessageType::InitAck, 0, &InitAck::default()).unwrap();
4967        let mut state = BootConsoleState::default();
4968        codec::encode_to_buf(&bootstrap_message, &mut state.input).unwrap();
4969        codec::encode_to_buf(&ack_message, &mut state.input).unwrap();
4970
4971        let decoded = read_boot_message(
4972            -1,
4973            &mut state,
4974            Some(Instant::now() + std::time::Duration::from_secs(1)),
4975            "guest bootstrap",
4976        )
4977        .unwrap();
4978        assert_eq!(decode_bootstrap_message(decoded).unwrap(), bootstrap);
4979        assert!(
4980            !state.input.is_empty(),
4981            "init ack frame should remain buffered"
4982        );
4983
4984        wait_for_init_ack(
4985            -1,
4986            &mut state,
4987            Instant::now() + std::time::Duration::from_secs(1),
4988        )
4989        .unwrap();
4990        assert!(state.input.is_empty());
4991    }
4992
4993    #[test]
4994    fn bootstrap_rejects_non_control_correlation_fields() {
4995        let mut message =
4996            Message::with_payload(MessageType::Bootstrap, 1, &GuestBootstrap::default()).unwrap();
4997        message.flags = 1;
4998
4999        let error = decode_bootstrap_message(message).unwrap_err();
5000        assert!(error.to_string().contains("requires id=0 and flags=0"));
5001    }
5002
5003    #[test]
5004    fn bootstrap_rejects_wrong_first_message_type() {
5005        let message = Message::with_payload(MessageType::Ping, 0, &Ping {}).unwrap();
5006
5007        let error = decode_bootstrap_message(message).unwrap_err();
5008        assert!(error.to_string().contains("expected core.bootstrap"));
5009    }
5010
5011    #[test]
5012    fn bootstrap_rejects_older_protocol_generation() {
5013        let mut message =
5014            Message::with_payload(MessageType::Bootstrap, 0, &GuestBootstrap::default()).unwrap();
5015        let min_version = MessageType::Bootstrap.min_protocol_version();
5016        assert!(min_version > 0);
5017        message.v = min_version - 1;
5018
5019        let error = decode_bootstrap_message(message).unwrap_err();
5020        assert!(error.to_string().contains("or newer"));
5021    }
5022
5023    #[test]
5024    fn bootstrap_accepts_minimum_supported_protocol_generation() {
5025        let mut message =
5026            Message::with_payload(MessageType::Bootstrap, 0, &GuestBootstrap::default()).unwrap();
5027        message.v = MessageType::Bootstrap.min_protocol_version();
5028
5029        assert_eq!(
5030            decode_bootstrap_message(message).unwrap(),
5031            GuestBootstrap::default()
5032        );
5033    }
5034
5035    #[test]
5036    fn bootstrap_accepts_newer_additive_protocol_generation() {
5037        let mut message =
5038            Message::with_payload(MessageType::Bootstrap, 0, &GuestBootstrap::default()).unwrap();
5039        message.v = PROTOCOL_VERSION + 1;
5040
5041        assert_eq!(
5042            decode_bootstrap_message(message).unwrap(),
5043            GuestBootstrap::default()
5044        );
5045    }
5046
5047    #[test]
5048    fn bootstrap_rejects_malformed_payload() {
5049        let message = Message::new(MessageType::Bootstrap, 0, vec![0xff]);
5050
5051        let error = decode_bootstrap_message(message).unwrap_err();
5052        assert!(error.to_string().contains("decode guest bootstrap payload"));
5053    }
5054    #[test]
5055    fn dual_port_cmdline_hint_requires_an_exact_argument() {
5056        assert!(cmdline_requests_dual_port(
5057            "console=hvc0 microsandbox.agent_transport=dual-port-v1 quiet"
5058        ));
5059        assert!(!cmdline_requests_dual_port(
5060            "microsandbox.agent_transport=dual-port-v10"
5061        ));
5062        assert!(!cmdline_requests_dual_port(
5063            "prefix=microsandbox.agent_transport=dual-port-v1"
5064        ));
5065    }
5066
5067    #[test]
5068    fn primary_read_buffer_tracks_whether_it_carries_bulk() {
5069        assert_eq!(
5070            primary_serial_read_buf_size(false),
5071            COMBINED_SERIAL_READ_BUF_SIZE
5072        );
5073        assert_eq!(
5074            primary_serial_read_buf_size(true),
5075            CONTROL_SERIAL_READ_BUF_SIZE
5076        );
5077    }
5078
5079    #[test]
5080    fn bulk_reader_turn_is_bounded_by_records_or_payload_bytes() {
5081        assert!(!bulk_reader_turn_exhausted(1, 1));
5082        assert!(bulk_reader_turn_exhausted(
5083            BULK_READER_MAX_RECORDS_PER_TURN,
5084            1
5085        ));
5086        assert!(bulk_reader_turn_exhausted(
5087            1,
5088            BULK_READER_MAX_BYTES_PER_TURN
5089        ));
5090    }
5091
5092    #[test]
5093    fn agent_read_budget_counts_bytes_and_calls_independently() {
5094        let mut bytes = AgentReadBudget::default();
5095        bytes.record_read(AGENT_READ_QUANTUM_BYTES - 1);
5096        assert!(!bytes.exhausted());
5097        bytes.record_read(1);
5098        assert!(bytes.exhausted());
5099
5100        let mut calls = AgentReadBudget::default();
5101        for _ in 1..AGENT_READ_QUANTUM_CALLS {
5102            assert!(calls.read_fd(-1, &mut [0]).is_err());
5103            assert!(!calls.exhausted());
5104        }
5105        assert!(calls.read_fd(-1, &mut [0]).is_err());
5106        assert!(calls.exhausted(), "failed reads also bound a retry loop");
5107        assert_eq!(calls.bytes, 0);
5108
5109        let mut large = AgentReadBudget::default();
5110        large.record_read(BULK_SERIAL_READ_BUF_SIZE);
5111        assert!(
5112            large.exhausted(),
5113            "a large read is not a smaller wire record"
5114        );
5115        assert_eq!(large.bytes, BULK_SERIAL_READ_BUF_SIZE);
5116    }
5117
5118    #[tokio::test(flavor = "current_thread")]
5119    async fn agent_read_budget_services_driver_during_partial_bulk_records() {
5120        // The control demonstrates the failure without relying on wall-clock delays: cached
5121        // readable bulk input never lets the runtime observe the newly readable/writable peers.
5122        assert!(!observe_driver_during_partial_bulk(false).await);
5123        assert!(observe_driver_during_partial_bulk(true).await);
5124    }
5125
5126    async fn observe_driver_during_partial_bulk(yield_at_boundary: bool) -> bool {
5127        use std::io::Write;
5128        use std::os::fd::OwnedFd;
5129        use std::os::unix::net::UnixStream;
5130        use std::sync::atomic::AtomicUsize;
5131
5132        let incarnation = [0x42; CLIENT_INCARNATION_SIZE];
5133        let first = BulkRecord {
5134            id: 1,
5135            kind: BulkKind::Filesystem,
5136            flow: BulkFlow::HostToGuest,
5137            offset: 0,
5138            payload: Bytes::from(vec![0xa5; 512]),
5139        };
5140        let second = BulkRecord {
5141            offset: first.payload.len() as u64,
5142            payload: Bytes::from(vec![0x5a; 512]),
5143            ..first.clone()
5144        };
5145        let mut wire = Vec::new();
5146        for record in [&first, &second] {
5147            wire.extend_from_slice(&incarnation);
5148            codec::encode_bulk_to_buf(record, &mut wire).unwrap();
5149        }
5150        let (mut source, input) = UnixStream::pair().unwrap();
5151        input.set_nonblocking(true).unwrap();
5152        source.write_all(&wire).unwrap();
5153        let mut input = BulkInputState::new(File::from(OwnedFd::from(input))).unwrap();
5154        // Model fragmented console reads while keeping the complete wire records unchanged.
5155        input.read_buf.truncate(1);
5156
5157        let (mut control_source, control) = UnixStream::pair().unwrap();
5158        control.set_nonblocking(true).unwrap();
5159        let control = AsyncFd::new(control).unwrap();
5160        let (pipe_reader, pipe_writer) = nix::unistd::pipe2(nix::fcntl::OFlag::O_NONBLOCK).unwrap();
5161        let pipe_writer = AsyncFd::new(pipe_writer).unwrap();
5162        // Prime the reactor, then clear the pipe's cached writable event with a real EAGAIN.
5163        let mut writable = pipe_writer.writable().await.unwrap();
5164        loop {
5165            match writable.try_io(|fd| write_to_fd(fd.get_ref().as_raw_fd(), &[0; 4096])) {
5166                Ok(Ok(count)) => assert!(count > 0),
5167                Ok(Err(error)) => panic!("fill stdin pipe: {error}"),
5168                Err(_) => break,
5169            }
5170        }
5171        drop(writable);
5172        let mut drained = [0; 4096];
5173        loop {
5174            match read_from_fd(pipe_reader.as_raw_fd(), &mut drained) {
5175                Ok(count) => assert!(count > 0),
5176                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => break,
5177                Err(error) => panic!("drain stdin pipe: {error}"),
5178            }
5179        }
5180        control_source.write_all(b"c").unwrap();
5181        let serviced = Arc::new(AtomicUsize::new(0));
5182        let control_serviced = Arc::clone(&serviced);
5183        let control_task = tokio::spawn(async move {
5184            loop {
5185                let mut ready = control.readable().await.unwrap();
5186                let mut byte = [0];
5187                match ready.try_io(|fd| read_from_fd(fd.get_ref().as_raw_fd(), &mut byte)) {
5188                    Ok(Ok(1)) => {
5189                        assert_eq!(byte, *b"c");
5190                        control_serviced.fetch_or(1, Ordering::Relaxed);
5191                        break;
5192                    }
5193                    Ok(result) => panic!("read control marker: {result:?}"),
5194                    Err(_) => continue,
5195                }
5196            }
5197        });
5198        let stdin_serviced = Arc::clone(&serviced);
5199        let stdin_task = tokio::spawn(async move {
5200            std::future::poll_fn(|cx| {
5201                loop {
5202                    let mut ready = std::task::ready!(pipe_writer.poll_write_ready(cx)).unwrap();
5203                    match ready.try_io(|fd| write_to_fd(fd.get_ref().as_raw_fd(), b"i")) {
5204                        Ok(result) => return Poll::Ready(result),
5205                        Err(_) => continue,
5206                    }
5207                }
5208            })
5209            .await
5210            .unwrap();
5211            stdin_serviced.fetch_or(2, Ordering::Relaxed);
5212        });
5213
5214        let mut budget = AgentReadBudget::default();
5215        let mut received = Vec::new();
5216        let mut serviced_before_first_record = false;
5217        while received.len() < 2 {
5218            let frames = tokio::select! {
5219                frames = input.read_turn(&mut budget) => frames.unwrap(),
5220                _ = std::future::pending::<()>() => unreachable!(),
5221            };
5222            for frame in frames {
5223                assert_eq!(frame.incarnation, incarnation);
5224                received.push(frame.record);
5225            }
5226            // This is the production cancellation-safe boundary: no decoded frames are local to
5227            // a competing select future, and even partial-record reads have spent the budget.
5228            if yield_at_boundary {
5229                budget.yield_if_exhausted().await;
5230            }
5231            if received.is_empty() && serviced.load(Ordering::Relaxed) == 3 {
5232                serviced_before_first_record = true;
5233            }
5234        }
5235        assert_eq!(
5236            received,
5237            [first, second],
5238            "yield preserves record bytes and FIFO"
5239        );
5240        assert!(input.input.is_empty());
5241        control_task.await.unwrap();
5242        stdin_task.await.unwrap();
5243        let mut marker = [0];
5244        assert_eq!(
5245            read_from_fd(pipe_reader.as_raw_fd(), &mut marker).unwrap(),
5246            1
5247        );
5248        assert_eq!(marker, *b"i");
5249        serviced_before_first_record
5250    }
5251
5252    #[tokio::test(flavor = "current_thread")]
5253    async fn cancelled_bulk_read_keeps_partial_record_and_read_budget() {
5254        use std::io::Write;
5255        use std::os::fd::OwnedFd;
5256        use std::os::unix::net::UnixStream;
5257
5258        let incarnation = [0x29; CLIENT_INCARNATION_SIZE];
5259        let record = BulkRecord {
5260            id: 1,
5261            kind: BulkKind::Filesystem,
5262            flow: BulkFlow::HostToGuest,
5263            offset: 0,
5264            payload: Bytes::from(vec![0x51; 512]),
5265        };
5266        let mut wire = incarnation.to_vec();
5267        codec::encode_bulk_to_buf(&record, &mut wire).unwrap();
5268        let (mut source, input) = UnixStream::pair().unwrap();
5269        input.set_nonblocking(true).unwrap();
5270        let mut input = BulkInputState::new(File::from(OwnedFd::from(input))).unwrap();
5271        let mut budget = AgentReadBudget::default();
5272        source.write_all(&wire[..40]).unwrap();
5273        assert!(input.read_turn(&mut budget).await.unwrap().is_empty());
5274        // Consume the stale readable hint so the following select truly cancels a pending read.
5275        assert!(input.read_turn(&mut budget).await.unwrap().is_empty());
5276        let calls = budget.calls;
5277        tokio::select! {
5278            biased;
5279            result = input.read_turn(&mut budget) => panic!("unexpected read: {result:?}"),
5280            _ = std::future::ready(()) => {},
5281        }
5282        assert_eq!(input.input.as_ref(), &wire[..40]);
5283        assert_eq!(budget.bytes, 40);
5284        assert_eq!(budget.calls, calls);
5285        source.write_all(&wire[40..]).unwrap();
5286        let frames = input.read_turn(&mut budget).await.unwrap();
5287        assert_eq!(frames.len(), 1);
5288        assert_eq!(frames[0].record, record);
5289        assert_eq!(frames[0].incarnation, incarnation);
5290        assert!(input.input.is_empty());
5291        assert_eq!(budget.bytes, wire.len());
5292    }
5293
5294    #[test]
5295    fn disconnect_cleanup_keeps_other_clients_bulk_offsets() {
5296        let mut state = AgentState::default();
5297        for id in [1, 10, 20] {
5298            state.bulk_received_offsets.insert(id, id as u64);
5299            state.pending_bulk_finishes.insert(
5300                id,
5301                BulkFinish {
5302                    kind: BulkKind::Filesystem,
5303                    flow: BulkFlow::HostToGuest,
5304                    final_offset: id as u64,
5305                },
5306            );
5307        }
5308
5309        clear_bulk_receive_range(&mut state, 10, 20);
5310
5311        assert_eq!(state.bulk_received_offsets.len(), 2);
5312        assert!(state.bulk_received_offsets.contains_key(&1));
5313        assert!(state.bulk_received_offsets.contains_key(&20));
5314        assert_eq!(state.pending_bulk_finishes.len(), 2);
5315        assert!(state.pending_bulk_finishes.contains_key(&1));
5316        assert!(state.pending_bulk_finishes.contains_key(&20));
5317    }
5318
5319    #[test]
5320    fn recycled_range_rejects_old_incarnation_and_stale_disconnect() {
5321        let id_start = 1;
5322        let id_end_exclusive = microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP;
5323        let old = [0x11; CLIENT_INCARNATION_SIZE];
5324        let new = [0x22; CLIENT_INCARNATION_SIZE];
5325        let mut state = AgentState::default();
5326
5327        establish_relay_client(
5328            &mut state,
5329            RelayClientConnected {
5330                id_start,
5331                id_end_exclusive,
5332                incarnation: old,
5333            },
5334        )
5335        .unwrap();
5336        state.bulk_received_offsets.insert(id_start, 4096);
5337        let (session_tx, _session_rx) = SessionOutputSender::channel();
5338        let (removed, cleanup) = disconnect_relay_client(
5339            &mut state,
5340            RelayClientDisconnected {
5341                id_start,
5342                id_end_exclusive,
5343                incarnation: Some(old),
5344            },
5345            &session_tx,
5346        )
5347        .unwrap();
5348        assert!(removed);
5349        assert!(cleanup.is_none());
5350        establish_relay_client(
5351            &mut state,
5352            RelayClientConnected {
5353                id_start,
5354                id_end_exclusive,
5355                incarnation: new,
5356            },
5357        )
5358        .unwrap();
5359
5360        assert_eq!(client_incarnation_for_id(&state, id_start), Some(new));
5361        assert!(validate_bulk_client_incarnation(&state, id_start, new).unwrap());
5362        assert!(!validate_bulk_client_incarnation(&state, id_start, old).unwrap());
5363
5364        let other_start = microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP + 1;
5365        let other_end = microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP * 2;
5366        establish_relay_client(
5367            &mut state,
5368            RelayClientConnected {
5369                id_start: other_start,
5370                id_end_exclusive: other_end,
5371                incarnation: [0x33; CLIENT_INCARNATION_SIZE],
5372            },
5373        )
5374        .unwrap();
5375        assert!(validate_bulk_client_incarnation(&state, other_start, new).is_err());
5376        assert!(!state.bulk_received_offsets.contains_key(&id_start));
5377        assert!(
5378            !disconnect_relay_client(
5379                &mut state,
5380                RelayClientDisconnected {
5381                    id_start,
5382                    id_end_exclusive,
5383                    incarnation: Some(old),
5384                },
5385                &session_tx,
5386            )
5387            .unwrap()
5388            .0
5389        );
5390        assert_eq!(client_incarnation_for_id(&state, id_start), Some(new));
5391    }
5392
5393    #[test]
5394    fn aggregate_bulk_input_budget_is_bounded_and_recoverable() {
5395        let state = AgentState::default();
5396        let budget = state.bulk_input_budget;
5397        let held = Arc::clone(&budget)
5398            .try_acquire_many_owned(BULK_INPUT_BYTE_CAPACITY as u32)
5399            .unwrap();
5400
5401        assert!(Arc::clone(&budget).try_acquire_owned().is_err());
5402        drop(held);
5403        assert_eq!(budget.available_permits(), BULK_INPUT_BYTE_CAPACITY);
5404    }
5405
5406    #[test]
5407    fn disconnect_ack_does_not_refresh_idle_activity() {
5408        let ack = encode_relay_client_disconnected_ack(RelayClientDisconnectedAck {
5409            id_start: 1,
5410            id_end_exclusive: microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP,
5411            incarnation: [0x4a; CLIENT_INCARNATION_SIZE],
5412        });
5413        assert!(!encoded_guest_message_refreshes_idle_timer(
5414            &ack,
5415            0,
5416            ack.len() - 4,
5417        ));
5418    }
5419
5420    #[tokio::test]
5421    async fn closed_optional_receiver_disables_itself_without_spinning() {
5422        let (tx, rx) = tokio::sync::mpsc::channel::<u8>(1);
5423        drop(tx);
5424        let mut receiver = Some(rx);
5425
5426        assert!(
5427            tokio::time::timeout(Duration::from_millis(10), recv_optional(&mut receiver))
5428                .await
5429                .is_err()
5430        );
5431        assert!(receiver.is_none());
5432    }
5433
5434    #[tokio::test]
5435    async fn bulk_output_tombstone_discards_late_producer_record() {
5436        let incarnation = [0x5a; CLIENT_INCARNATION_SIZE];
5437        let (session_tx, _control_rx, mut bulk_rx, _command_rx) =
5438            SessionOutputSender::split_channel();
5439        let owner_tx = session_tx.with_incarnation(Some(incarnation));
5440        let record = || BulkRecord {
5441            id: 17,
5442            kind: BulkKind::Filesystem,
5443            flow: BulkFlow::GuestToHost,
5444            offset: 0,
5445            payload: bytes::Bytes::from_static(b"late"),
5446        };
5447        assert!(
5448            owner_tx
5449                .send(
5450                    17,
5451                    SessionOutput::Bulk(crate::session::BulkSessionOutput::new(
5452                        record(),
5453                        RawActivity::default(),
5454                    )),
5455                )
5456                .await
5457        );
5458
5459        let mut generation = 0;
5460        let mut flows = HashMap::new();
5461        let mut active = VecDeque::new();
5462        let mut retired = HashMap::new();
5463        let mut retiring_incarnations = HashSet::new();
5464        enqueue_bulk_output(
5465            bulk_rx.recv().await.unwrap(),
5466            0,
5467            &mut flows,
5468            &mut active,
5469            &retired,
5470            &retiring_incarnations,
5471        )
5472        .unwrap();
5473        let (completion, mut completed) = tokio::sync::oneshot::channel();
5474        let cleanup = apply_bulk_output_command(
5475            BulkOutputCommand::DropFlow {
5476                incarnation,
5477                id: 17,
5478                completion,
5479            },
5480            &mut generation,
5481            &mut flows,
5482            &mut active,
5483            &mut retired,
5484            &mut retiring_incarnations,
5485            &mut BulkOutputPosition::default(),
5486        )
5487        .unwrap();
5488        assert!(matches!(
5489            completed.try_recv(),
5490            Err(tokio::sync::oneshot::error::TryRecvError::Empty)
5491        ));
5492        complete_bulk_output_cleanups(vec![cleanup], &mut retired, &mut retiring_incarnations);
5493        assert_eq!(completed.try_recv(), Ok(()));
5494
5495        assert!(
5496            owner_tx
5497                .send(
5498                    17,
5499                    SessionOutput::Bulk(crate::session::BulkSessionOutput::new(
5500                        record(),
5501                        RawActivity::default(),
5502                    )),
5503                )
5504                .await
5505        );
5506        enqueue_bulk_output(
5507            bulk_rx.recv().await.unwrap(),
5508            0,
5509            &mut flows,
5510            &mut active,
5511            &retired,
5512            &retiring_incarnations,
5513        )
5514        .unwrap();
5515
5516        assert!(flows.is_empty());
5517        assert!(active.is_empty());
5518        assert!(bulk_output_is_retired(&retired, incarnation, 17));
5519    }
5520
5521    #[test]
5522    fn record_encoded_guest_messages_counts_only_appended_frames() {
5523        let mut out_buf = Vec::new();
5524        let existing =
5525            Message::with_payload(MessageType::ExecStarted, 1, &ExecStarted { pid: 123 }).unwrap();
5526        codec::encode_to_buf(&existing, &mut out_buf).unwrap();
5527        let start = out_buf.len();
5528
5529        let appended =
5530            Message::with_payload(MessageType::ExecStarted, 2, &ExecStarted { pid: 456 }).unwrap();
5531        codec::encode_to_buf(&appended, &mut out_buf).unwrap();
5532
5533        let mut activity = ActivityTracker::new();
5534        record_encoded_guest_messages(&out_buf, start, &mut activity);
5535
5536        assert_eq!(activity.activity_seq, 1);
5537        assert_eq!(activity.counters.guest_messages, 1);
5538    }
5539
5540    #[test]
5541    fn apply_raw_activity_updates_guest_and_byte_counters() {
5542        let mut activity = ActivityTracker::new();
5543
5544        apply_raw_activity(RawActivity::fs_bytes(42), &mut activity);
5545        apply_raw_activity(RawActivity::tcp_bytes(7), &mut activity);
5546
5547        assert_eq!(activity.activity_seq, 2);
5548        assert_eq!(activity.counters.guest_messages, 2);
5549        assert_eq!(activity.counters.fs_bytes, 42);
5550        assert_eq!(activity.counters.tcp_bytes, 7);
5551    }
5552
5553    #[test]
5554    fn apply_raw_activity_preserves_coalesced_message_count() {
5555        let mut activity = ActivityTracker::new();
5556
5557        apply_raw_activity(
5558            RawActivity {
5559                guest_messages: 16,
5560                fs_bytes: 4 * 1024 * 1024,
5561                tcp_bytes: 0,
5562            },
5563            &mut activity,
5564        );
5565
5566        assert_eq!(activity.activity_seq, 16);
5567        assert_eq!(activity.counters.guest_messages, 16);
5568        assert_eq!(activity.counters.fs_bytes, 4 * 1024 * 1024);
5569    }
5570
5571    #[tokio::test]
5572    async fn filesystem_finish_remains_admissible_after_a_full_data_window() {
5573        let (records, _record_rx) = tokio::sync::mpsc::channel(FS_BULK_INPUT_ITEM_CAPACITY);
5574        let (finish, mut finish_rx) = tokio::sync::mpsc::channel(1);
5575        let worker = FsBulkWriteWorker {
5576            records,
5577            finish,
5578            task: tokio::spawn(std::future::pending::<()>()),
5579        };
5580        let payload = Bytes::from(vec![0u8; MIN_BULK_RECORD_PAYLOAD as usize]);
5581
5582        for index in 0..FS_BULK_INPUT_ITEM_CAPACITY {
5583            worker
5584                .records
5585                .try_send(AdmittedBulkRecord::for_test(BulkRecord {
5586                    id: 17,
5587                    kind: BulkKind::Filesystem,
5588                    flow: BulkFlow::HostToGuest,
5589                    offset: (index * MIN_BULK_RECORD_PAYLOAD as usize) as u64,
5590                    payload: payload.clone(),
5591                }))
5592                .unwrap();
5593        }
5594        assert_eq!(worker.records.capacity(), 0);
5595
5596        let expected = BulkFinish {
5597            kind: BulkKind::Filesystem,
5598            flow: BulkFlow::HostToGuest,
5599            final_offset: microsandbox_protocol::bulk::DEFAULT_BULK_WINDOW,
5600        };
5601        worker.finish.try_send(expected).unwrap();
5602        assert_eq!(finish_rx.recv().await, Some(expected));
5603        worker.task.abort();
5604    }
5605
5606    #[test]
5607    fn maintenance_messages_do_not_implicitly_refresh_idle_timer() {
5608        assert!(!message_refreshes_idle_timer(&MessageType::ClockSync));
5609        assert!(!message_refreshes_idle_timer(&MessageType::Ping));
5610        assert!(!message_refreshes_idle_timer(&MessageType::Touch));
5611        assert!(message_refreshes_idle_timer(&MessageType::ExecRequest));
5612    }
5613
5614    #[test]
5615    fn maintenance_replies_do_not_refresh_idle_timer() {
5616        assert!(!guest_message_refreshes_idle_timer(&MessageType::Pong));
5617        assert!(!guest_message_refreshes_idle_timer(&MessageType::Touched));
5618        assert!(!guest_message_refreshes_idle_timer(&MessageType::CoreError));
5619        assert!(guest_message_refreshes_idle_timer(&MessageType::ExecStdout));
5620    }
5621
5622    #[test]
5623    fn record_encoded_guest_messages_ignores_pong_and_touched() {
5624        let mut out_buf = Vec::new();
5625        let pong = Message::with_payload(MessageType::Pong, 1, &Pong {}).unwrap();
5626        codec::encode_to_buf(&pong, &mut out_buf).unwrap();
5627
5628        let touched =
5629            Message::with_payload(MessageType::Touched, 2, &Touched { activity_seq: 42 }).unwrap();
5630        codec::encode_to_buf(&touched, &mut out_buf).unwrap();
5631
5632        let mut activity = ActivityTracker::new();
5633        record_encoded_guest_messages(&out_buf, 0, &mut activity);
5634
5635        assert_eq!(activity.activity_seq, 0);
5636        assert_eq!(activity.counters.guest_messages, 0);
5637    }
5638
5639    #[test]
5640    fn bulk_cancellation_uses_each_operations_existing_terminal_type() {
5641        for (kind, expected) in [
5642            (BulkKind::Filesystem, MessageType::FsResponse),
5643            (BulkKind::Tcp, MessageType::TcpFailed),
5644        ] {
5645            let mut encoded = Vec::new();
5646            encode_bulk_terminal_failure(17, kind, "cancelled".into(), &mut encoded).unwrap();
5647            let mut bytes = BytesMut::from(encoded.as_slice());
5648            let frame = codec::try_decode_frame_from_bytes(&mut bytes)
5649                .unwrap()
5650                .expect("terminal frame");
5651            let DecodedFrame::Control(message) = frame else {
5652                panic!("cancellation terminal must be control");
5653            };
5654            assert_eq!(message.id, 17);
5655            assert_eq!(message.t, expected);
5656            assert_eq!(message.flags, microsandbox_protocol::message::FLAG_TERMINAL);
5657            assert!(bytes.is_empty());
5658        }
5659    }
5660    #[tokio::test]
5661    async fn restore_thaw_checks_ownership_is_idempotent_and_leaves_source_clients_alone() {
5662        use microsandbox_protocol::core::WorkloadThawMode;
5663
5664        let mut state = AgentState::default();
5665        let (mut sender, _output) = SessionOutputSender::channel();
5666        let mut activity = ActivityTracker::new();
5667        let heartbeat = heartbeat::HeartbeatControl::default();
5668        let config = AgentdConfig {
5669            user: None,
5670            security_profile: Default::default(),
5671            default_cwd: None,
5672            default_env: Vec::new(),
5673        };
5674        let mut workload = crate::workload::tests::fake_latch();
5675        let owner = [0x51; CLIENT_INCARNATION_SIZE];
5676        let connect = |state: &mut AgentState| {
5677            establish_relay_client(
5678                state,
5679                RelayClientConnected {
5680                    id_start: 1,
5681                    id_end_exclusive: microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP,
5682                    incarnation: owner,
5683                },
5684            )
5685            .unwrap();
5686        };
5687        connect(&mut state);
5688        workload.freeze("capture").unwrap();
5689
5690        // Wrong attempts cannot detach a client's state. Source-side thaw must also preserve it.
5691        for (attempt, mode, expected) in [
5692            ("wrong", WorkloadThawMode::Restore, MessageType::CoreError),
5693            (
5694                "capture",
5695                WorkloadThawMode::Continue,
5696                MessageType::WorkloadThawed,
5697            ),
5698        ] {
5699            let request = Message::with_payload(
5700                MessageType::WorkloadThaw,
5701                0,
5702                &WorkloadThaw {
5703                    attempt_id: attempt.into(),
5704                    mode,
5705                },
5706            )
5707            .unwrap();
5708            let mut encoded = Vec::new();
5709            handle_message(
5710                request,
5711                &mut state,
5712                &mut activity,
5713                &mut sender,
5714                &mut encoded,
5715                &config,
5716                &mut workload,
5717                &heartbeat,
5718            )
5719            .await
5720            .unwrap();
5721            let mut bytes = BytesMut::from(encoded.as_slice());
5722            let reply = decode_reply_skipping_credit(&mut bytes);
5723            assert_eq!(reply.t, expected);
5724            assert_eq!(sender.generation(), 0);
5725            assert_eq!(client_incarnation_for_id(&state, 1), Some(owner));
5726        }
5727
5728        for generation in 1..=2 {
5729            // Reuse the same human-readable name for a subsequent capture, then retry its thaw.
5730            let request = Message::with_payload(
5731                MessageType::WorkloadFreeze,
5732                0,
5733                &WorkloadFreeze {
5734                    external_mount_tags: Vec::new(),
5735                    attempt_id: "capture".into(),
5736                    host_input: WorkloadTransportPosition::default(),
5737                },
5738            )
5739            .unwrap();
5740            handle_message(
5741                request,
5742                &mut state,
5743                &mut activity,
5744                &mut sender,
5745                &mut Vec::new(),
5746                &config,
5747                &mut workload,
5748                &heartbeat,
5749            )
5750            .await
5751            .unwrap();
5752            for retry in [false, true] {
5753                let request = Message::with_payload(
5754                    MessageType::WorkloadThaw,
5755                    0,
5756                    &WorkloadThaw {
5757                        attempt_id: "capture".into(),
5758                        mode: WorkloadThawMode::Restore,
5759                    },
5760                )
5761                .unwrap();
5762                let mut encoded = Vec::new();
5763                handle_message(
5764                    request,
5765                    &mut state,
5766                    &mut activity,
5767                    &mut sender,
5768                    &mut encoded,
5769                    &config,
5770                    &mut workload,
5771                    &heartbeat,
5772                )
5773                .await
5774                .unwrap();
5775                let mut bytes = BytesMut::from(encoded.as_slice());
5776                let reply = decode_reply_skipping_credit(&mut bytes);
5777                assert_eq!(reply.t, MessageType::WorkloadThawed);
5778                assert!(!workload.is_frozen());
5779                assert_eq!(sender.generation(), generation);
5780                assert_eq!(client_incarnation_for_id(&state, 1), retry.then_some(owner));
5781                if !retry {
5782                    // A retry must not detach clients which connected after the first success.
5783                    connect(&mut state);
5784                }
5785            }
5786        }
5787    }
5788
5789    #[tokio::test]
5790    async fn saturated_stdin_allows_thaw_and_preserves_input_and_eof() {
5791        use microsandbox_protocol::core::WorkloadThawMode::{Continue, Restore};
5792
5793        // SIGSTOP provides a deterministic blocked pipe/PTY without a privileged cgroup mount.
5794        // Thaw must complete before external consumer progress, then accepted data drains exactly
5795        // once. Restore closes inherited pipe input after data even when no EOF was accepted.
5796        struct StoppedProcessGuard(i32);
5797        impl Drop for StoppedProcessGuard {
5798            fn drop(&mut self) {
5799                unsafe {
5800                    libc::kill(-self.0, libc::SIGCONT);
5801                    libc::kill(-self.0, libc::SIGKILL);
5802                }
5803            }
5804        }
5805
5806        for (tty, mode, accepted_eof) in [
5807            (false, Continue, false),
5808            (false, Continue, true),
5809            (false, Restore, false),
5810            (false, Restore, true),
5811            (true, Continue, true),
5812            (true, Restore, true),
5813        ] {
5814            let mut state = AgentState::default();
5815            let (mut sender, mut output) = SessionOutputSender::channel();
5816            let mut activity = ActivityTracker::new();
5817            let heartbeat = heartbeat::HeartbeatControl::default();
5818            let config = AgentdConfig {
5819                user: None,
5820                security_profile: Default::default(),
5821                default_cwd: None,
5822                default_env: Vec::new(),
5823            };
5824            let mut workload = crate::workload::tests::fake_latch();
5825            workload.freeze("stdin-cut").unwrap();
5826            let request = ExecRequest {
5827                cmd: "/bin/sh".into(),
5828                args: vec![
5829                    "-c".into(),
5830                    if tty {
5831                        "stty raw -echo; kill -STOP $$; exec cat"
5832                    } else {
5833                        "kill -STOP $$; exec cat"
5834                    }
5835                    .into(),
5836                ],
5837                env: vec![],
5838                cwd: None,
5839                user: None,
5840                tty,
5841                rows: 24,
5842                cols: 80,
5843                rlimits: vec![],
5844            };
5845            let session = ExecSession::spawn(
5846                1,
5847                &request,
5848                sender.clone(),
5849                None,
5850                crate::config::SecurityProfile::Default,
5851                None,
5852            )
5853            .unwrap();
5854            let pid = session.pid() as i32;
5855            state.sessions.insert(1, session);
5856            let _cleanup = StoppedProcessGuard(pid);
5857            time::timeout(Duration::from_secs(5), async {
5858                loop {
5859                    let status = std::fs::read_to_string(format!("/proc/{pid}/status")).unwrap();
5860                    if status.lines().any(|line| line.starts_with("State:\tT")) {
5861                        break;
5862                    }
5863                    time::sleep(Duration::from_millis(1)).await;
5864                }
5865            })
5866            .await
5867            .expect("workload should stop before stdin is admitted");
5868
5869            let mut encoded = Vec::new();
5870            let ledger = state.input_window.clone();
5871            let initial = ledger.credit().unwrap();
5872            let data = vec![0x61; 1024 * 1024];
5873            let charge = ledger.admit(InputLane::Bulk, data.len() + 32).unwrap();
5874            let mut stdin =
5875                Message::with_payload(MessageType::ExecStdin, 1, &ExecStdin { data }).unwrap();
5876            // A generation-8 SDK still travels over the bundled private transport contract.
5877            // Dispatch must not fall back to blocking writes based on the client message version.
5878            stdin.v = 8;
5879            time::timeout(
5880                Duration::from_millis(100),
5881                handle_message_with_charge(
5882                    stdin,
5883                    &mut state,
5884                    &mut activity,
5885                    &mut sender,
5886                    &mut encoded,
5887                    &config,
5888                    &mut workload,
5889                    &heartbeat,
5890                    Some(charge),
5891                ),
5892            )
5893            .await
5894            .expect("stdin dispatch waited on its blocked consumer")
5895            .unwrap();
5896            assert!(encoded.is_empty());
5897            assert!(state.sessions[&1].has_pending_stdin());
5898            if accepted_eof {
5899                let charge = ledger.admit(InputLane::Bulk, 32).unwrap();
5900                let mut eof = Message::with_payload(
5901                    MessageType::ExecStdin,
5902                    1,
5903                    &ExecStdin { data: Vec::new() },
5904                )
5905                .unwrap();
5906                eof.v = 8;
5907                time::timeout(
5908                    Duration::from_millis(100),
5909                    handle_message_with_charge(
5910                        eof,
5911                        &mut state,
5912                        &mut activity,
5913                        &mut sender,
5914                        &mut encoded,
5915                        &config,
5916                        &mut workload,
5917                        &heartbeat,
5918                        Some(charge),
5919                    ),
5920                )
5921                .await
5922                .expect("EOF dispatch waited on preceding blocked data")
5923                .unwrap();
5924            }
5925            assert_eq!(
5926                ledger.credit().unwrap(),
5927                initial,
5928                "blocked input refunded before consumption"
5929            );
5930            let position = ledger.position();
5931            let thaw = Message::with_payload(
5932                MessageType::WorkloadThaw,
5933                u32::MAX,
5934                &WorkloadThaw {
5935                    attempt_id: "stdin-cut".into(),
5936                    mode,
5937                },
5938            )
5939            .unwrap();
5940            time::timeout(
5941                Duration::from_millis(100),
5942                handle_message(
5943                    thaw,
5944                    &mut state,
5945                    &mut activity,
5946                    &mut sender,
5947                    &mut encoded,
5948                    &config,
5949                    &mut workload,
5950                    &heartbeat,
5951                ),
5952            )
5953            .await
5954            .expect("thaw waited on saturated stdin")
5955            .unwrap();
5956            assert!(!workload.is_frozen());
5957            let mut bytes = BytesMut::from(encoded.as_slice());
5958            let reply = decode_reply_skipping_credit(&mut bytes);
5959            assert_eq!(reply.t, MessageType::WorkloadThawed);
5960            assert_eq!(
5961                ledger.position(),
5962                position,
5963                "restore reset cumulative input position"
5964            );
5965            assert_eq!(
5966                ledger.credit().unwrap(),
5967                initial,
5968                "restore discarded accepted input"
5969            );
5970            assert_eq!(unsafe { libc::kill(-pid, libc::SIGCONT) }, 0);
5971            let mut received = Vec::new();
5972            let mut exited = false;
5973            time::timeout(Duration::from_secs(5), async {
5974                while received.len() < 1024 * 1024
5975                    || state.sessions.values().chain(state.detached_sessions.values()).any(ExecSession::has_pending_stdin)
5976                    || (!tty && (accepted_eof || mode == Restore) && !exited) {
5977                    tokio::select! {
5978                        (_, _, result) = std::future::poll_fn(|cx| poll_pending_stdin(&mut state, cx)) => result.unwrap(),
5979                        envelope = output.recv() => match envelope.unwrap().output {
5980                            SessionOutput::Stdout(data) => received.extend(data),
5981                            SessionOutput::Exited(code) => { assert_eq!(code, 0); exited = true; },
5982                            _ => {},
5983                        },
5984                    }
5985                }
5986            })
5987            .await
5988            .expect("accepted stdin or ordered EOF did not drain");
5989            assert_eq!(received, vec![0x61; 1024 * 1024]);
5990            assert_eq!(
5991                ledger.credit().unwrap().bulk_bytes,
5992                initial.bulk_bytes + position.bulk_bytes
5993            );
5994            if !tty && mode == Continue && !accepted_eof {
5995                // Source Continue must not synthesize EOF. The same owner can still send input.
5996                assert!(!exited);
5997                state
5998                    .sessions
5999                    .get_mut(&1)
6000                    .unwrap()
6001                    .enqueue_stdin(b"tail".to_vec(), None)
6002                    .unwrap();
6003                let envelope = time::timeout(Duration::from_secs(5), output.recv())
6004                    .await
6005                    .unwrap()
6006                    .unwrap();
6007                assert!(
6008                    matches!(envelope.output, SessionOutput::Stdout(ref data) if data == b"tail")
6009                );
6010            }
6011        }
6012    }
6013
6014    #[tokio::test]
6015    async fn restore_detaches_piped_and_pty_workloads_without_killing_or_reusing_output() {
6016        for tty in [false, true] {
6017            let mut state = AgentState::default();
6018            let (mut sender, mut output) = SessionOutputSender::channel();
6019            let id = 1;
6020            let end = microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP;
6021            let old_owner = [0x31; CLIENT_INCARNATION_SIZE];
6022            let new_owner = [0x32; CLIENT_INCARNATION_SIZE];
6023            establish_relay_client(
6024                &mut state,
6025                RelayClientConnected {
6026                    id_start: id,
6027                    id_end_exclusive: end,
6028                    incarnation: old_owner,
6029                },
6030            )
6031            .unwrap();
6032            let request = |script: &str| ExecRequest {
6033                cmd: "/bin/sh".into(),
6034                args: vec!["-c".into(), script.into()],
6035                env: vec![],
6036                cwd: None,
6037                user: None,
6038                tty,
6039                rows: 24,
6040                cols: 80,
6041                rlimits: vec![],
6042            };
6043            let session = ExecSession::spawn(
6044                id,
6045                &request("echo before; sleep 0.3; echo inherited; exit 23"),
6046                sender.with_incarnation(Some(old_owner)),
6047                None,
6048                crate::config::SecurityProfile::Default,
6049                None,
6050            )
6051            .unwrap();
6052            let pid = session.pid();
6053            state.sessions.insert(id, session);
6054            let before = time::timeout(Duration::from_secs(5), output.recv())
6055                .await
6056                .unwrap()
6057                .unwrap();
6058            assert!(matches!(before.output, SessionOutput::Stdout(_)));
6059            let old_generation = sender.generation();
6060
6061            restore_client_state(&mut state, &mut sender).await.unwrap();
6062            assert_eq!(
6063                unsafe { libc::kill(pid as i32, 0) },
6064                0,
6065                "restore killed the captured process"
6066            );
6067            assert_eq!(state.detached_sessions.len(), 1);
6068            assert!(state.sessions.is_empty());
6069            establish_relay_client(
6070                &mut state,
6071                RelayClientConnected {
6072                    id_start: id,
6073                    id_end_exclusive: end,
6074                    incarnation: new_owner,
6075                },
6076            )
6077            .unwrap();
6078            let fresh = ExecSession::spawn(
6079                id,
6080                &request("echo fresh; sleep 0.8; exit 24"),
6081                sender.with_incarnation(Some(new_owner)),
6082                None,
6083                crate::config::SecurityProfile::Default,
6084                None,
6085            )
6086            .unwrap();
6087            state.sessions.insert(id, fresh);
6088            let mut saw_old_exit = false;
6089            let mut saw_new_exit = false;
6090            let mut saw_inherited = false;
6091            time::timeout(Duration::from_secs(5), async {
6092                while !saw_old_exit || !saw_new_exit {
6093                    let envelope = output.recv().await.unwrap();
6094                    if envelope.generation == old_generation {
6095                        if let SessionOutput::Stdout(bytes) = &envelope.output {
6096                            saw_inherited |= String::from_utf8_lossy(bytes).contains("inherited");
6097                        }
6098                        if let SessionOutput::Exited(code) = &envelope.output {
6099                            assert_eq!(*code, 23);
6100                            saw_old_exit = true;
6101                        }
6102                        assert!(discard_inherited_output(
6103                            &mut state,
6104                            &envelope,
6105                            sender.generation()
6106                        ));
6107                        assert!(
6108                            state.sessions.contains_key(&id),
6109                            "old exit removed new session"
6110                        );
6111                    } else {
6112                        assert!(!discard_inherited_output(
6113                            &mut state,
6114                            &envelope,
6115                            sender.generation()
6116                        ));
6117                        if let SessionOutput::Exited(code) = envelope.output {
6118                            assert_eq!(code, 24);
6119                            saw_new_exit = true;
6120                        }
6121                    }
6122                }
6123            })
6124            .await
6125            .unwrap();
6126            assert!(saw_inherited, "inherited output reader stopped draining");
6127            assert!(
6128                state.detached_sessions.is_empty(),
6129                "exited inherited handles leaked"
6130            );
6131        }
6132    }
6133
6134    #[tokio::test]
6135    async fn restore_closes_inherited_filesystem_handles_and_tcp_transfers() {
6136        use microsandbox_protocol::fs::{FsOp, FsOpenOptions, FsResponse, FsResponseData};
6137        use tokio::io::AsyncReadExt;
6138
6139        async fn fs_request(
6140            state: &mut AgentState,
6141            sender: &SessionOutputSender,
6142            op: FsOp,
6143        ) -> FsResponse {
6144            let mut encoded = Vec::new();
6145            crate::fs::handle_fs_request(
6146                1,
6147                PROTOCOL_VERSION,
6148                FsRequest { op, bulk: None },
6149                &mut state.fs,
6150                &mut encoded,
6151                sender,
6152            )
6153            .await
6154            .unwrap();
6155            let mut bytes = BytesMut::from(encoded.as_slice());
6156            let Some(DecodedFrame::Control(reply)) =
6157                codec::try_decode_frame_from_bytes(&mut bytes).unwrap()
6158            else {
6159                panic!("missing filesystem response");
6160            };
6161            reply.payload::<FsResponse>().unwrap()
6162        }
6163
6164        let mut state = AgentState::default();
6165        let (mut sender, _output) = SessionOutputSender::channel();
6166        let open = || FsOp::OpenFile {
6167            path: "/dev/null".into(),
6168            options: FsOpenOptions {
6169                read: true,
6170                ..Default::default()
6171            },
6172        };
6173        let Some(FsResponseData::Handle(old_handle)) =
6174            fs_request(&mut state, &sender, open()).await.data
6175        else {
6176            panic!("file did not open");
6177        };
6178        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6179        state.tcp_sessions.insert(
6180            2,
6181            TcpSession::open(
6182                2,
6183                TcpConnect {
6184                    host: "127.0.0.1".into(),
6185                    port: listener.local_addr().unwrap().port(),
6186                    bulk: None,
6187                },
6188                &sender,
6189            ),
6190        );
6191        let (mut peer, _) = time::timeout(Duration::from_secs(5), listener.accept())
6192            .await
6193            .unwrap()
6194            .unwrap();
6195
6196        restore_client_state(&mut state, &mut sender).await.unwrap();
6197        assert!(state.tcp_sessions.is_empty());
6198        assert_eq!(
6199            time::timeout(Duration::from_secs(5), peer.read(&mut [0u8; 1]))
6200                .await
6201                .unwrap()
6202                .unwrap(),
6203            0
6204        );
6205        assert!(
6206            !fs_request(&mut state, &sender, FsOp::FStat { handle: old_handle })
6207                .await
6208                .ok
6209        );
6210        let Some(FsResponseData::Handle(new_handle)) =
6211            fs_request(&mut state, &sender, open()).await.data
6212        else {
6213            panic!("fresh file did not open");
6214        };
6215        // Closing handles must not reset their allocator and alias a stale handle to a new file.
6216        assert_ne!(new_handle, old_handle);
6217        assert!(
6218            fs_request(&mut state, &sender, FsOp::FStat { handle: new_handle })
6219                .await
6220                .ok
6221        );
6222    }
6223
6224    #[tokio::test]
6225    async fn restore_retires_unleased_output_and_can_repeat_with_the_same_ids() {
6226        let mut state = AgentState::default();
6227        let (mut sender, mut output) = SessionOutputSender::channel();
6228        for _ in 0..3 {
6229            let old = sender.clone();
6230            old.send(1, SessionOutput::Exited(7)).await;
6231            restore_client_state(&mut state, &mut sender).await.unwrap();
6232            let stale = output.recv().await.unwrap();
6233            assert!(discard_inherited_output(
6234                &mut state,
6235                &stale,
6236                sender.generation()
6237            ));
6238            // A late producer is as stale as an event that was already queued at the cut.
6239            old.send(1, SessionOutput::Exited(7)).await;
6240            assert!(discard_inherited_output(
6241                &mut state,
6242                &output.recv().await.unwrap(),
6243                sender.generation()
6244            ));
6245            sender.send(1, SessionOutput::Exited(8)).await;
6246            assert!(!discard_inherited_output(
6247                &mut state,
6248                &output.recv().await.unwrap(),
6249                sender.generation()
6250            ));
6251        }
6252    }
6253
6254    #[tokio::test]
6255    async fn restore_bulk_cut_drops_queued_and_late_records_and_releases_capacity() {
6256        let (mut sender, _control, mut bulk, mut commands) = SessionOutputSender::split_channel();
6257        let old = sender.with_incarnation(Some([0x41; CLIENT_INCARNATION_SIZE]));
6258        let send = |sender: SessionOutputSender| async move {
6259            sender
6260                .send(
6261                    1,
6262                    SessionOutput::Bulk(crate::session::BulkSessionOutput::new(
6263                        BulkRecord {
6264                            id: 1,
6265                            kind: BulkKind::Filesystem,
6266                            flow: BulkFlow::GuestToHost,
6267                            offset: 0,
6268                            payload: bytes::Bytes::from_static(b"old"),
6269                        },
6270                        RawActivity::default(),
6271                    )),
6272                )
6273                .await
6274        };
6275        assert!(send(old.clone()).await);
6276        let reset = tokio::spawn(async move {
6277            sender.restore_generation().await.unwrap();
6278            sender
6279        });
6280        let command = commands.recv().await.unwrap();
6281        let mut generation = 0;
6282        let mut flows = HashMap::new();
6283        let mut active = VecDeque::new();
6284        let mut retired = HashMap::new();
6285        let mut retiring = HashSet::new();
6286        let cleanup = apply_bulk_output_command(
6287            command,
6288            &mut generation,
6289            &mut flows,
6290            &mut active,
6291            &mut retired,
6292            &mut retiring,
6293            &mut BulkOutputPosition::default(),
6294        )
6295        .unwrap();
6296        enqueue_bulk_output(
6297            bulk.recv().await.unwrap(),
6298            generation,
6299            &mut flows,
6300            &mut active,
6301            &retired,
6302            &retiring,
6303        )
6304        .unwrap();
6305        assert!(flows.is_empty());
6306        complete_bulk_output_cleanups(vec![cleanup], &mut retired, &mut retiring);
6307        let fresh = reset.await.unwrap();
6308        assert!(send(old).await);
6309        enqueue_bulk_output(
6310            bulk.recv().await.unwrap(),
6311            generation,
6312            &mut flows,
6313            &mut active,
6314            &retired,
6315            &retiring,
6316        )
6317        .unwrap();
6318        assert!(flows.is_empty());
6319        assert!(send(fresh.with_incarnation(Some([0x42; CLIENT_INCARNATION_SIZE]))).await);
6320        enqueue_bulk_output(
6321            bulk.recv().await.unwrap(),
6322            generation,
6323            &mut flows,
6324            &mut active,
6325            &retired,
6326            &retiring,
6327        )
6328        .unwrap();
6329        assert_eq!(flows.len(), 1);
6330    }
6331
6332    #[test]
6333    fn restore_activity_backpressure_does_not_block_scheduler_control() {
6334        let (sender, mut receiver) = tokio::sync::mpsc::channel(1);
6335        sender.try_send(RawActivity::default()).unwrap();
6336        let mut pending = RawActivity {
6337            guest_messages: 3,
6338            fs_bytes: 17,
6339            ..Default::default()
6340        };
6341        publish_bulk_activity(&sender, &mut pending).unwrap();
6342        assert_eq!(pending.fs_bytes, 17);
6343        receiver.try_recv().unwrap();
6344        publish_bulk_activity(&sender, &mut pending).unwrap();
6345        assert_eq!(receiver.try_recv().unwrap().fs_bytes, 17);
6346        assert_eq!(pending.fs_bytes, 0);
6347    }
6348}