Integrate host escalation with production sessions

crates/lazaret-host-coordinator/tests/escalation.rs

@@ -272,6 +272,31 @@ fn drive_terminal(
     Err("terminal result deadline exhausted".into())
 }
 
+fn acknowledge_one_chunk(
+    mailboxes: &Mailboxes,
+    coordinator: &mut EscalationCoordinator,
+    request: &mut ClientRequest<'_>,
+    now: Instant,
+) -> TestResult<Vec<u8>> {
+    for _ in 0..500 {
+        coordinator.progress(request.request_id(), now)?;
+        match request.poll()? {
+            ClientEvent::Output {
+                sequence, bytes, ..
+            } => {
+                let ack = mailboxes.consume_ack(request.request_id(), sequence)?;
+                coordinator.acknowledge(request.request_id(), &ack, now)?;
+                return Ok(bytes);
+            }
+            ClientEvent::Terminal(_) => {
+                return Err("request completed before partial output".into());
+            }
+            ClientEvent::Pending => std::thread::sleep(Duration::from_millis(1)),
+        }
+    }
+    Err("partial output deadline exhausted".into())
+}
+
 fn completed_auto_and_developer_requests_preserve_attribution() -> TestResult {
     for (seed, auto) in [(10, true), (30, false)] {
         let fixture = make_fixture(seed)?;
@@ -336,6 +361,19 @@ fn completed_auto_and_developer_requests_preserve_attribution() -> TestResult {
                 ApprovalSource::Developer
             }
         );
+        assert!(
+            RegularFileMailbox::open(&mailboxes.outbox_path)?
+                .names()?
+                .contains(&lazaret_core::protocol::result_file_name(
+                    request.request_id()
+                ))
+        );
+        coordinator.session_lost(fixture.session.session_binding().session_id)?;
+        assert!(
+            RegularFileMailbox::open(&mailboxes.outbox_path)?
+                .names()?
+                .is_empty()
+        );
         assert!(recover_process(ownership.last.unwrap().identity()).is_err());
     }
     Ok(())
@@ -558,6 +596,73 @@ fn output_window_never_exceeds_eight_unacknowledged_chunks() -> TestResult {
     Ok(())
 }
 
+fn interrupted_execution_commits_acknowledged_output() -> TestResult {
+    let fixture = make_fixture(115)?;
+    let mailboxes = Mailboxes::create("cancel-after-output")?;
+    let client = mailboxes.client(&fixture)?;
+    let argv = vec![
+        mailboxes.command("/usr/bin/head", "head")?,
+        b"-c".to_vec(),
+        b"1000000".to_vec(),
+        b"/dev/zero".to_vec(),
+    ];
+    let mut request = client.request(argv.clone(), "cancel after output".into())?;
+    let signed = mailboxes.consume_request(request.request_id())?;
+    let now = Instant::now();
+    let mut coordinator =
+        mailboxes.coordinator(&fixture, Duration::from_secs(5), Duration::from_secs(2))?;
+    coordinator.receive_auto_approved(
+        fixture.session.clone(),
+        &signed,
+        executable(&argv)?,
+        argv,
+        &mut RecordingOwnership::default(),
+        now,
+    )?;
+    assert!(!acknowledge_one_chunk(&mailboxes, &mut coordinator, &mut request, now)?.is_empty());
+    request.cancel()?;
+    let cancellation = RegularFileMailbox::open(&mailboxes.inbox_path)?.consume(
+        &lazaret_core::protocol::cancellation_file_name(request.request_id()),
+        MAX_SIGNED_RECORD_BYTES,
+    )?;
+    coordinator.cancel(request.request_id(), &cancellation)?;
+    assert!(
+        matches!(request.poll()?, ClientEvent::Terminal(result) if result.outcome == EscalationOutcome::Cancelled)
+    );
+
+    let fixture = make_fixture(116)?;
+    let mailboxes = Mailboxes::create("timeout-after-output")?;
+    let client = mailboxes.client(&fixture)?;
+    let argv = vec![
+        mailboxes.command("/usr/bin/head", "head")?,
+        b"-c".to_vec(),
+        b"1000000".to_vec(),
+        b"/dev/zero".to_vec(),
+    ];
+    let mut request = client.request(argv.clone(), "timeout after output".into())?;
+    let signed = mailboxes.consume_request(request.request_id())?;
+    let start = Instant::now();
+    let mut coordinator =
+        mailboxes.coordinator(&fixture, Duration::from_millis(10), Duration::from_secs(2))?;
+    coordinator.receive_auto_approved(
+        fixture.session.clone(),
+        &signed,
+        executable(&argv)?,
+        argv,
+        &mut RecordingOwnership::default(),
+        start,
+    )?;
+    assert!(!acknowledge_one_chunk(&mailboxes, &mut coordinator, &mut request, start)?.is_empty());
+    assert_eq!(
+        coordinator.expire(start + Duration::from_millis(20))?,
+        vec![request.request_id()]
+    );
+    assert!(
+        matches!(request.poll()?, ClientEvent::Terminal(result) if result.outcome == EscalationOutcome::TimedOut)
+    );
+    Ok(())
+}
+
 fn replay_and_session_tuple_mismatch_refuse() -> TestResult {
     let fixture = make_fixture(120)?;
     let mailboxes = Mailboxes::create("replay")?;
@@ -767,6 +872,7 @@ fn run() -> TestResult {
     denial_failure_and_request_timeout_publish_terminal_outcomes()?;
     cancellation_session_loss_and_ack_timeout_reap_exact_children()?;
     output_window_never_exceeds_eight_unacknowledged_chunks()?;
+    interrupted_execution_commits_acknowledged_output()?;
     replay_and_session_tuple_mismatch_refuse()?;
     signed_cancellation_decision_and_result_tuple_mismatches_refuse()?;
     exact_descriptor_launcher_survives_path_replacement()?;