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