返回 CodeWhale
host.rs
根目录 / crates / tui / src / fleet / host.rs
1 //! Fleet worker host adapters.
2 //!
3 //! Adapters own process boundaries for worker hosts. The manager can lease and
4 //! observe work through this trait without knowing whether the worker is a
5 //! local child process or an SSH-backed remote command.
6
7 #![allow(dead_code)]
8
9 use std::collections::{BTreeMap, BTreeSet};
10 use std::fs::{File, OpenOptions};
11 use std::io::{Read, Seek, SeekFrom};
12 use std::path::{Path, PathBuf};
13 use std::process::{Child, Command, ExitStatus, Stdio};
14 use std::thread;
15 use std::time::{Duration, Instant};
16
17 use codewhale_protocol::fleet::FleetHostSpec;
18 use thiserror::Error;
19
20 #[cfg(unix)]
21 use std::os::unix::process::CommandExt;
22 #[cfg(windows)]
23 use std::os::windows::io::AsRawHandle;
24 #[cfg(unix)]
25 use std::sync::OnceLock;
26 #[cfg(windows)]
27 use windows::Win32::Foundation::{CloseHandle, HANDLE};
28 #[cfg(windows)]
29 use windows::Win32::System::JobObjects::{
30 AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
31 JOBOBJECT_BASIC_ACCOUNTING_INFORMATION, JOBOBJECT_EXTENDED_LIMIT_INFORMATION,
32 JobObjectBasicAccountingInformation, JobObjectExtendedLimitInformation,
33 QueryInformationJobObject, SetInformationJobObject, TerminateJobObject,
34 };
35 #[cfg(windows)]
36 use windows::core::PCWSTR;
37
38 const DEFAULT_LOG_LIMIT_BYTES: usize = 64 * 1024;
39 const DEFAULT_CONNECT_TIMEOUT_SECONDS: u64 = 10;
40 const WORKER_STOP_GRACE: Duration = Duration::from_millis(750);
41
42 pub type FleetHostResult<T> = Result<T, FleetHostError>;
43
44 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
45 pub enum FleetHostErrorKind {
46 Retryable,
47 Terminal,
48 Configuration,
49 }
50
51 #[derive(Debug, Error)]
52 #[error("{kind:?}: {message}")]
53 pub struct FleetHostError {
54 pub kind: FleetHostErrorKind,
55 pub message: String,
56 }
57
58 impl FleetHostError {
59 fn retryable(message: impl Into<String>) -> Self {
60 Self {
61 kind: FleetHostErrorKind::Retryable,
62 message: message.into(),
63 }
64 }
65
66 fn terminal(message: impl Into<String>) -> Self {
67 Self {
68 kind: FleetHostErrorKind::Terminal,
69 message: message.into(),
70 }
71 }
72
73 fn configuration(message: impl Into<String>) -> Self {
74 Self {
75 kind: FleetHostErrorKind::Configuration,
76 message: message.into(),
77 }
78 }
79 }
80
81 #[derive(Debug, Clone, PartialEq, Eq)]
82 pub struct FleetWorkerCommand {
83 pub program: String,
84 pub args: Vec<String>,
85 }
86
87 impl FleetWorkerCommand {
88 pub fn new<S, I, A>(program: S, args: I) -> Self
89 where
90 S: Into<String>,
91 I: IntoIterator<Item = A>,
92 A: Into<String>,
93 {
94 Self {
95 program: program.into(),
96 args: args.into_iter().map(Into::into).collect(),
97 }
98 }
99 }
100
101 #[derive(Debug, Clone)]
102 pub struct FleetWorkerStartRequest {
103 pub worker_id: String,
104 pub command: FleetWorkerCommand,
105 pub cwd: Option<PathBuf>,
106 pub env: BTreeMap<String, String>,
107 pub env_allowlist: BTreeSet<String>,
108 pub log_limit_bytes: usize,
109 }
110
111 impl FleetWorkerStartRequest {
112 pub fn new(worker_id: impl Into<String>, command: FleetWorkerCommand) -> Self {
113 Self {
114 worker_id: worker_id.into(),
115 command,
116 cwd: None,
117 env: BTreeMap::new(),
118 env_allowlist: BTreeSet::new(),
119 log_limit_bytes: DEFAULT_LOG_LIMIT_BYTES,
120 }
121 }
122 }
123
124 #[derive(Debug, Clone, PartialEq, Eq)]
125 pub struct FleetWorkerHandle {
126 pub worker_id: String,
127 pub host_kind: FleetHostKind,
128 pub pid: Option<u32>,
129 pub log_path: PathBuf,
130 }
131
132 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
133 pub enum FleetHostKind {
134 LocalProcess,
135 Ssh,
136 }
137
138 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
139 pub enum FleetHostWorkerState {
140 Running,
141 /// The dispatcher stopped but the owned process session/job is not yet empty.
142 Draining,
143 Exited,
144 Failed,
145 Stopped,
146 Unknown,
147 }
148
149 #[derive(Debug, Clone, PartialEq, Eq)]
150 pub struct FleetHostWorkerStatus {
151 pub worker_id: String,
152 pub state: FleetHostWorkerState,
153 pub pid: Option<u32>,
154 pub exit_code: Option<i32>,
155 pub memory_mb: Option<u64>,
156 pub retryable: bool,
157 }
158
159 pub trait FleetHostAdapter {
160 fn start_worker(
161 &mut self,
162 request: FleetWorkerStartRequest,
163 ) -> FleetHostResult<FleetWorkerHandle>;
164 fn read_status(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus>;
165 fn read_logs(&self, worker_id: &str, max_bytes: usize) -> FleetHostResult<String>;
166 fn interrupt_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus>;
167 fn restart_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetWorkerHandle>;
168 fn stop_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus>;
169 fn cleanup_worker(&mut self, worker_id: &str) -> FleetHostResult<()>;
170 }
171
172 #[derive(Debug)]
173 pub struct LocalProcessFleetHostAdapter {
174 workspace: PathBuf,
175 processes: BTreeMap<String, LocalWorkerProcess>,
176 }
177
178 #[derive(Debug)]
179 struct LocalWorkerProcess {
180 request: FleetWorkerStartRequest,
181 child: Child,
182 #[cfg(unix)]
183 session_id: libc::pid_t,
184 #[cfg(unix)]
185 parent_death_writer: Option<std::io::PipeWriter>,
186 #[cfg(windows)]
187 windows_job: FleetWindowsJob,
188 host_kind: FleetHostKind,
189 log_path: PathBuf,
190 stopped: bool,
191 last_exit: Option<ExitStatus>,
192 last_memory_mb: Option<u64>,
193 }
194
195 impl LocalProcessFleetHostAdapter {
196 pub fn new(workspace: impl AsRef<Path>) -> Self {
197 Self {
198 workspace: workspace.as_ref().to_path_buf(),
199 processes: BTreeMap::new(),
200 }
201 }
202
203 fn start_with_kind(
204 &mut self,
205 request: FleetWorkerStartRequest,
206 host_kind: FleetHostKind,
207 ) -> FleetHostResult<FleetWorkerHandle> {
208 validate_worker_id(&request.worker_id)?;
209 if self.processes.contains_key(&request.worker_id) {
210 let status = self.read_status(&request.worker_id)?;
211 if matches!(status.state, FleetHostWorkerState::Running) {
212 return Err(FleetHostError::terminal(format!(
213 "worker {} is already running",
214 request.worker_id
215 )));
216 }
217 self.processes.remove(&request.worker_id);
218 }
219
220 let env = worker_env(&request.env, &request.env_allowlist)?;
221 let log_path = self.log_path_for(&request.worker_id, host_kind);
222 let log = open_worker_log(&log_path)?;
223 let stderr = log
224 .try_clone()
225 .map_err(|err| FleetHostError::retryable(format!("cloning worker log: {err}")))?;
226
227 let mut command = Command::new(&request.command.program);
228 // Parent-death watch (R7): the worker's stdin is the read end of a pipe
229 // whose write end lives only in this adapter process. If the manager
230 // dies (crash, kill, power loss), the kernel closes the write end, the
231 // worker sees stdin EOF, and `--parent-death-watch` shuts the worker
232 // tree down instead of letting it spend forever. Windows workers are
233 // contained in a Job Object that the OS terminates on parent death.
234 #[cfg(unix)]
235 let (parent_death_writer, stdin) = {
236 let (reader, writer) = std::io::pipe().map_err(|err| {
237 FleetHostError::retryable(format!("creating parent-death pipe: {err}"))
238 })?;
239 (Some(writer), std::process::Stdio::from(reader))
240 };
241 #[cfg(not(unix))]
242 let stdin = Stdio::null();
243 command
244 .args(&request.command.args)
245 .stdin(stdin)
246 .stdout(Stdio::from(log))
247 .stderr(Stdio::from(stderr))
248 .env_clear()
249 .envs(env);
250 if let Some(cwd) = &request.cwd {
251 command.current_dir(cwd);
252 }
253
254 // Fleet owns the complete worker tree, not only the dispatcher PID.
255 // `codewhale` spawns `codewhale-tui`, which can in turn spawn tool
256 // processes; isolating the root prevents a stop from signalling the
257 // operator's own process group.
258 #[cfg(unix)]
259 // SAFETY: `setsid` is async-signal-safe and the closure does not touch
260 // allocator or parent-held state between fork and exec.
261 unsafe {
262 command.pre_exec(|| {
263 if libc::setsid() == -1 {
264 Err(std::io::Error::last_os_error())
265 } else {
266 Ok(())
267 }
268 });
269 }
270
271 let child = command.spawn().map_err(|err| {
272 classify_spawn_error(err, format!("starting worker {}", request.worker_id))
273 })?;
274 #[cfg(windows)]
275 let (child, windows_job) = attach_fleet_windows_job(child).map_err(|err| {
276 FleetHostError::retryable(format!(
277 "containing worker {} in a Windows Job Object: {err}",
278 request.worker_id
279 ))
280 })?;
281 let pid = child.id();
282 let handle = FleetWorkerHandle {
283 worker_id: request.worker_id.clone(),
284 host_kind,
285 pid: Some(pid),
286 log_path: log_path.clone(),
287 };
288 self.processes.insert(
289 request.worker_id.clone(),
290 LocalWorkerProcess {
291 request,
292 child,
293 // SAFETY: `setsid` is async-signal-safe and the closure does not
294 // touch allocator or parent-held state between fork and exec.
295 #[cfg(unix)]
296 session_id: pid as libc::pid_t,
297 #[cfg(unix)]
298 parent_death_writer,
299 #[cfg(windows)]
300 windows_job,
301 host_kind,
302 log_path,
303 stopped: false,
304 last_exit: None,
305 last_memory_mb: None,
306 },
307 );
308 Ok(handle)
309 }
310
311 fn log_path_for(&self, worker_id: &str, host_kind: FleetHostKind) -> PathBuf {
312 let host_dir = match host_kind {
313 FleetHostKind::LocalProcess => "local",
314 FleetHostKind::Ssh => "ssh",
315 };
316 self.workspace
317 .join(".codewhale")
318 .join("fleet-host")
319 .join(host_dir)
320 .join(format!("{}.log", safe_path_segment(worker_id)))
321 }
322 }
323
324 impl FleetHostAdapter for LocalProcessFleetHostAdapter {
325 fn start_worker(
326 &mut self,
327 request: FleetWorkerStartRequest,
328 ) -> FleetHostResult<FleetWorkerHandle> {
329 self.start_with_kind(request, FleetHostKind::LocalProcess)
330 }
331
332 fn read_status(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus> {
333 let process = self
334 .processes
335 .get_mut(worker_id)
336 .ok_or_else(|| FleetHostError::terminal(format!("unknown worker {worker_id}")))?;
337 if let Some(status) = process.last_exit {
338 if local_worker_tree_alive(process)? {
339 return Ok(FleetHostWorkerStatus {
340 worker_id: worker_id.to_string(),
341 state: FleetHostWorkerState::Draining,
342 pid: Some(process.child.id()),
343 exit_code: status.code(),
344 memory_mb: process.last_memory_mb,
345 retryable: true,
346 });
347 }
348 return Ok(status_from_exit(
349 worker_id,
350 Some(process.child.id()),
351 status,
352 process.stopped,
353 process.last_memory_mb,
354 ));
355 }
356 match process.child.try_wait() {
357 Ok(None) => {
358 let pid = process.child.id();
359 let memory_mb = if process.host_kind == FleetHostKind::LocalProcess {
360 sample_process_memory_mb(pid)
361 } else {
362 None
363 };
364 process.last_memory_mb = memory_mb.or(process.last_memory_mb);
365 Ok(FleetHostWorkerStatus {
366 worker_id: worker_id.to_string(),
367 state: FleetHostWorkerState::Running,
368 pid: Some(pid),
369 exit_code: None,
370 // Report the retained value, not the raw sample: a
371 // transient ps failure must not flicker a live worker's
372 // memory to None (the Exited arm already does this).
373 memory_mb: process.last_memory_mb,
374 retryable: false,
375 })
376 }
377 Ok(Some(status)) => {
378 process.last_exit = Some(status);
379 if local_worker_tree_alive(process)? {
380 return Ok(FleetHostWorkerStatus {
381 worker_id: worker_id.to_string(),
382 state: FleetHostWorkerState::Draining,
383 pid: Some(process.child.id()),
384 exit_code: status.code(),
385 memory_mb: process.last_memory_mb,
386 retryable: true,
387 });
388 }
389 Ok(status_from_exit(
390 worker_id,
391 Some(process.child.id()),
392 status,
393 process.stopped,
394 process.last_memory_mb,
395 ))
396 }
397 Err(err) => Err(FleetHostError::retryable(format!(
398 "reading worker {worker_id} status: {err}"
399 ))),
400 }
401 }
402
403 fn read_logs(&self, worker_id: &str, max_bytes: usize) -> FleetHostResult<String> {
404 let process = self
405 .processes
406 .get(worker_id)
407 .ok_or_else(|| FleetHostError::terminal(format!("unknown worker {worker_id}")))?;
408 let max_bytes = max_bytes.min(process.request.log_limit_bytes.max(1));
409 read_bounded_log(&process.log_path, max_bytes)
410 }
411
412 fn interrupt_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus> {
413 {
414 let process = self
415 .processes
416 .get_mut(worker_id)
417 .ok_or_else(|| FleetHostError::terminal(format!("unknown worker {worker_id}")))?;
418 // The direct dispatcher may already be reaped while delegated
419 // session/job descendants remain. Interrupt the containment
420 // boundary unconditionally.
421 interrupt_worker_tree(process)?;
422 }
423 wait_for_exit(self, worker_id, WORKER_STOP_GRACE)
424 }
425
426 fn restart_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetWorkerHandle> {
427 let request = self
428 .processes
429 .get(worker_id)
430 .map(|process| process.request.clone())
431 .ok_or_else(|| FleetHostError::terminal(format!("unknown worker {worker_id}")))?;
432 let _ = self.stop_worker(worker_id);
433 self.processes.remove(worker_id);
434 self.start_worker(request)
435 }
436
437 fn stop_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus> {
438 {
439 let process = self
440 .processes
441 .get_mut(worker_id)
442 .ok_or_else(|| FleetHostError::terminal(format!("unknown worker {worker_id}")))?;
443 process.stopped = true;
444 if process.last_exit.is_none() {
445 match process.child.try_wait() {
446 Ok(Some(status)) => {
447 process.last_exit = Some(status);
448 }
449 Ok(None) => {}
450 Err(err) => {
451 return Err(FleetHostError::retryable(format!(
452 "reading worker {worker_id} status before stop: {err}"
453 )));
454 }
455 }
456 }
457 // Always tear down the containment boundary. A dispatcher can
458 // exit before a delegated TUI/tool child, so direct-child status
459 // is not proof that the complete worker tree is gone.
460 stop_worker_tree(process).map_err(|err| FleetHostError {
461 kind: err.kind,
462 message: format!("stopping worker {worker_id}: {}", err.message),
463 })?;
464 }
465 self.read_status(worker_id)
466 }
467
468 fn cleanup_worker(&mut self, worker_id: &str) -> FleetHostResult<()> {
469 if self.processes.contains_key(worker_id) {
470 // Cleanup is the final containment boundary. Even when the direct
471 // dispatcher already exited, delegated children may still occupy
472 // its Unix session or Windows Job Object.
473 let _ = self.stop_worker(worker_id)?;
474 }
475 self.processes.remove(worker_id);
476 Ok(())
477 }
478 }
479
480 #[derive(Debug, Clone)]
481 pub struct SshFleetHostConfig {
482 pub host: String,
483 pub user: Option<String>,
484 pub port: Option<u16>,
485 pub identity: Option<PathBuf>,
486 pub known_hosts: Option<PathBuf>,
487 pub host_key_fingerprint: Option<String>,
488 pub working_directory: PathBuf,
489 pub env_allowlist: BTreeSet<String>,
490 pub codewhale_binary: String,
491 pub ssh_binary: String,
492 pub connect_timeout_seconds: u64,
493 }
494
495 impl SshFleetHostConfig {
496 pub fn new(host: impl Into<String>, working_directory: impl Into<PathBuf>) -> Self {
497 Self {
498 host: host.into(),
499 user: None,
500 port: None,
501 identity: None,
502 known_hosts: None,
503 host_key_fingerprint: None,
504 working_directory: working_directory.into(),
505 env_allowlist: BTreeSet::new(),
506 codewhale_binary: "codewhale".to_string(),
507 ssh_binary: "ssh".to_string(),
508 connect_timeout_seconds: DEFAULT_CONNECT_TIMEOUT_SECONDS,
509 }
510 }
511
512 pub fn from_host_spec(spec: &FleetHostSpec) -> FleetHostResult<Self> {
513 let FleetHostSpec::Ssh {
514 host,
515 port,
516 user,
517 identity,
518 known_hosts,
519 host_key_fingerprint,
520 working_directory,
521 env_allowlist,
522 codewhale_binary,
523 } = spec
524 else {
525 return Err(FleetHostError::configuration(
526 "expected SSH Fleet host spec",
527 ));
528 };
529 let working_directory = working_directory.clone().ok_or_else(|| {
530 FleetHostError::configuration("SSH Fleet host spec requires working_directory")
531 })?;
532 let codewhale_binary = codewhale_binary.clone().ok_or_else(|| {
533 FleetHostError::configuration("SSH Fleet host spec requires codewhale_binary")
534 })?;
535 let mut config = Self::new(host.clone(), working_directory);
536 config.port = *port;
537 config.user = user.clone();
538 config.identity = identity.clone();
539 config.known_hosts = known_hosts.clone();
540 config.host_key_fingerprint = host_key_fingerprint.clone();
541 config.env_allowlist = env_allowlist.iter().cloned().collect();
542 config.codewhale_binary = codewhale_binary;
543 config.validate()?;
544 Ok(config)
545 }
546
547 fn validate(&self) -> FleetHostResult<()> {
548 if self.host.trim().is_empty() {
549 return Err(FleetHostError::configuration(
550 "SSH Fleet host requires an explicit host",
551 ));
552 }
553 if self.codewhale_binary.trim().is_empty() {
554 return Err(FleetHostError::configuration(
555 "SSH Fleet host requires an explicit codewhale binary path",
556 ));
557 }
558 if self.working_directory.as_os_str().is_empty() {
559 return Err(FleetHostError::configuration(
560 "SSH Fleet host requires an explicit working directory",
561 ));
562 }
563 validate_env_allowlist(&self.env_allowlist)
564 }
565
566 fn target(&self) -> String {
567 self.user
568 .as_ref()
569 .filter(|user| !user.trim().is_empty())
570 .map(|user| format!("{user}@{}", self.host))
571 .unwrap_or_else(|| self.host.clone())
572 }
573 }
574
575 #[derive(Debug)]
576 pub struct SshFleetHostAdapter {
577 config: SshFleetHostConfig,
578 local: LocalProcessFleetHostAdapter,
579 }
580
581 impl SshFleetHostAdapter {
582 pub fn new(workspace: impl AsRef<Path>, config: SshFleetHostConfig) -> FleetHostResult<Self> {
583 config.validate()?;
584 Ok(Self {
585 config,
586 local: LocalProcessFleetHostAdapter::new(workspace),
587 })
588 }
589
590 pub fn build_ssh_command(
591 &self,
592 request: &FleetWorkerStartRequest,
593 ) -> FleetHostResult<FleetWorkerCommand> {
594 self.config.validate()?;
595 let env = filtered_env(&request.env, &self.config.env_allowlist)?;
596 let mut args = vec![
597 "-o".to_string(),
598 "BatchMode=yes".to_string(),
599 "-o".to_string(),
600 format!("ConnectTimeout={}", self.config.connect_timeout_seconds),
601 ];
602 for key in env.keys() {
603 args.push("-o".to_string());
604 args.push(format!("SendEnv={key}"));
605 }
606 if let Some(port) = self.config.port {
607 args.push("-p".to_string());
608 args.push(port.to_string());
609 }
610 if let Some(identity) = &self.config.identity {
611 args.push("-i".to_string());
612 args.push(identity.display().to_string());
613 }
614 args.push(self.config.target());
615 args.push(self.remote_command(request));
616 Ok(FleetWorkerCommand::new(
617 self.config.ssh_binary.clone(),
618 args,
619 ))
620 }
621
622 fn ssh_start_request(
623 &self,
624 request: FleetWorkerStartRequest,
625 ) -> FleetHostResult<FleetWorkerStartRequest> {
626 let command = self.build_ssh_command(&request)?;
627 let mut env = ssh_client_env();
628 env.extend(filtered_env(&request.env, &self.config.env_allowlist)?);
629 let env_allowlist = env.keys().cloned().collect();
630 Ok(FleetWorkerStartRequest {
631 worker_id: request.worker_id,
632 command,
633 cwd: None,
634 env,
635 env_allowlist,
636 log_limit_bytes: request.log_limit_bytes,
637 })
638 }
639
640 fn remote_command(&self, request: &FleetWorkerStartRequest) -> String {
641 let mut parts = vec![
642 "cd".to_string(),
643 shell_quote(&self.config.working_directory.display().to_string()),
644 "&&".to_string(),
645 "exec".to_string(),
646 shell_quote(&self.config.codewhale_binary),
647 ];
648 parts.extend(request.command.args.iter().map(|arg| shell_quote(arg)));
649 parts.join(" ")
650 }
651 }
652
653 impl FleetHostAdapter for SshFleetHostAdapter {
654 fn start_worker(
655 &mut self,
656 request: FleetWorkerStartRequest,
657 ) -> FleetHostResult<FleetWorkerHandle> {
658 let request = self.ssh_start_request(request)?;
659 self.local.start_with_kind(request, FleetHostKind::Ssh)
660 }
661
662 fn read_status(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus> {
663 self.local.read_status(worker_id)
664 }
665
666 fn read_logs(&self, worker_id: &str, max_bytes: usize) -> FleetHostResult<String> {
667 self.local.read_logs(worker_id, max_bytes)
668 }
669
670 fn interrupt_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus> {
671 self.local.interrupt_worker(worker_id)
672 }
673
674 fn restart_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetWorkerHandle> {
675 let request = self
676 .local
677 .processes
678 .get(worker_id)
679 .map(|process| process.request.clone())
680 .ok_or_else(|| FleetHostError::terminal(format!("unknown worker {worker_id}")))?;
681 let _ = self.stop_worker(worker_id);
682 self.local.processes.remove(worker_id);
683 self.local.start_with_kind(request, FleetHostKind::Ssh)
684 }
685
686 fn stop_worker(&mut self, worker_id: &str) -> FleetHostResult<FleetHostWorkerStatus> {
687 self.local.stop_worker(worker_id)
688 }
689
690 fn cleanup_worker(&mut self, worker_id: &str) -> FleetHostResult<()> {
691 self.local.cleanup_worker(worker_id)
692 }
693 }
694
695 fn open_worker_log(path: &Path) -> FleetHostResult<File> {
696 if let Some(parent) = path.parent() {
697 std::fs::create_dir_all(parent).map_err(|err| {
698 FleetHostError::retryable(format!(
699 "creating worker log dir {}: {err}",
700 parent.display()
701 ))
702 })?;
703 }
704 OpenOptions::new()
705 .create(true)
706 .write(true)
707 .truncate(true)
708 .open(path)
709 .map_err(|err| FleetHostError::retryable(format!("opening worker log: {err}")))
710 }
711
712 fn read_bounded_log(path: &Path, max_bytes: usize) -> FleetHostResult<String> {
713 let mut file = File::open(path).map_err(|err| {
714 FleetHostError::retryable(format!("opening worker log {}: {err}", path.display()))
715 })?;
716 let len = file
717 .metadata()
718 .map_err(|err| FleetHostError::retryable(format!("reading worker log metadata: {err}")))?
719 .len();
720 let max_bytes = max_bytes.max(1) as u64;
721 if len > max_bytes {
722 file.seek(SeekFrom::Start(len - max_bytes))
723 .map_err(|err| FleetHostError::retryable(format!("seeking worker log: {err}")))?;
724 }
725 let mut bytes = Vec::new();
726 file.read_to_end(&mut bytes)
727 .map_err(|err| FleetHostError::retryable(format!("reading worker log: {err}")))?;
728 Ok(String::from_utf8_lossy(&bytes).into_owned())
729 }
730
731 fn status_from_exit(
732 worker_id: &str,
733 pid: Option<u32>,
734 status: ExitStatus,
735 stopped: bool,
736 memory_mb: Option<u64>,
737 ) -> FleetHostWorkerStatus {
738 let success = status.success();
739 FleetHostWorkerStatus {
740 worker_id: worker_id.to_string(),
741 state: if stopped {
742 FleetHostWorkerState::Stopped
743 } else if success {
744 FleetHostWorkerState::Exited
745 } else {
746 FleetHostWorkerState::Failed
747 },
748 pid,
749 exit_code: status.code(),
750 memory_mb,
751 retryable: !success && !stopped,
752 }
753 }
754
755 #[cfg(unix)]
756 fn sample_process_memory_mb(pid: u32) -> Option<u64> {
757 // Resolve `ps` via PATH like every other external command in the
758 // codebase: /bin/ps does not exist on NixOS and some minimal containers,
759 // which would silently report permanent None for live workers. Restricted
760 // sandboxes may also deny process-table inspection with EPERM; treat that
761 // as unavailable rather than panicking or inventing a sample.
762 if !process_table_inspection_available() {
763 return None;
764 }
765 let output = match Command::new("ps")
766 .args(["-o", "rss=", "-p", &pid.to_string()])
767 .output()
768 {
769 Ok(output) => output,
770 Err(err) if is_permission_denied(&err) => {
771 mark_process_table_unavailable();
772 return None;
773 }
774 Err(_) => return None,
775 };
776 if !output.status.success() {
777 return None;
778 }
779 let rss_kb = String::from_utf8_lossy(&output.stdout)
780 .split_whitespace()
781 .next()?
782 .parse::<u64>()
783 .ok()?;
784 (rss_kb > 0).then_some(rss_kb.div_ceil(1024))
785 }
786
787 #[cfg(not(unix))]
788 fn sample_process_memory_mb(_pid: u32) -> Option<u64> {
789 None
790 }
791
792 fn classify_spawn_error(err: std::io::Error, context: String) -> FleetHostError {
793 match err.kind() {
794 std::io::ErrorKind::NotFound => FleetHostError::configuration(format!("{context}: {err}")),
795 std::io::ErrorKind::PermissionDenied => {
796 FleetHostError::terminal(format!("{context}: {err}"))
797 }
798 _ => FleetHostError::retryable(format!("{context}: {err}")),
799 }
800 }
801
802 fn wait_for_exit(
803 adapter: &mut LocalProcessFleetHostAdapter,
804 worker_id: &str,
805 timeout: Duration,
806 ) -> FleetHostResult<FleetHostWorkerStatus> {
807 let deadline = Instant::now() + timeout;
808 loop {
809 let status = adapter.read_status(worker_id)?;
810 if !matches!(
811 status.state,
812 FleetHostWorkerState::Running | FleetHostWorkerState::Draining
813 ) {
814 return Ok(status);
815 }
816 if Instant::now() >= deadline {
817 return Ok(status);
818 }
819 thread::sleep(Duration::from_millis(25));
820 }
821 }
822
823 #[cfg(unix)]
824 fn local_worker_tree_alive(process: &LocalWorkerProcess) -> FleetHostResult<bool> {
825 Ok(!unix_session_members(process.session_id, Some(process.session_id))?.is_empty())
826 }
827
828 #[cfg(windows)]
829 fn local_worker_tree_alive(process: &LocalWorkerProcess) -> FleetHostResult<bool> {
830 process.windows_job.has_active_processes().map_err(|err| {
831 FleetHostError::retryable(format!("querying Windows worker job activity: {err}"))
832 })
833 }
834
835 #[cfg(not(any(unix, windows)))]
836 fn local_worker_tree_alive(_process: &LocalWorkerProcess) -> FleetHostResult<bool> {
837 Ok(false)
838 }
839
840 #[cfg(unix)]
841 fn interrupt_worker_tree(process: &mut LocalWorkerProcess) -> FleetHostResult<()> {
842 shutdown_unix_worker_session(process, &[libc::SIGINT, libc::SIGTERM])
843 }
844
845 #[cfg(windows)]
846 fn interrupt_worker_tree(process: &mut LocalWorkerProcess) -> FleetHostResult<()> {
847 process.windows_job.terminate().map_err(|err| {
848 FleetHostError::retryable(format!("interrupting Windows worker tree: {err}"))
849 })
850 }
851
852 #[cfg(not(any(unix, windows)))]
853 fn interrupt_worker_tree(process: &mut LocalWorkerProcess) -> FleetHostResult<()> {
854 process
855 .child
856 .kill()
857 .map_err(|err| FleetHostError::retryable(format!("interrupting worker: {err}")))
858 }
859
860 #[cfg(unix)]
861 fn stop_worker_tree(process: &mut LocalWorkerProcess) -> FleetHostResult<()> {
862 shutdown_unix_worker_session(process, &[libc::SIGTERM])
863 }
864
865 #[cfg(windows)]
866 fn stop_worker_tree(process: &mut LocalWorkerProcess) -> FleetHostResult<()> {
867 process.windows_job.terminate().map_err(|err| {
868 FleetHostError::retryable(format!("terminating Windows worker job: {err}"))
869 })?;
870 if process.last_exit.is_none() {
871 process.last_exit =
872 Some(process.child.wait().map_err(|err| {
873 FleetHostError::retryable(format!("reaping Windows worker: {err}"))
874 })?);
875 }
876 Ok(())
877 }
878
879 #[cfg(not(any(unix, windows)))]
880 fn stop_worker_tree(process: &mut LocalWorkerProcess) -> FleetHostResult<()> {
881 process
882 .child
883 .kill()
884 .map_err(|err| FleetHostError::retryable(format!("killing worker: {err}")))?;
885 process.last_exit = Some(
886 process
887 .child
888 .wait()
889 .map_err(|err| FleetHostError::retryable(format!("reaping worker: {err}")))?,
890 );
891 Ok(())
892 }
893
894 #[cfg(unix)]
895 fn shutdown_unix_worker_session(
896 process: &mut LocalWorkerProcess,
897 graceful_signals: &[libc::c_int],
898 ) -> FleetHostResult<()> {
899 let mut signal_errors = Vec::new();
900 let known_leader = process.session_id;
901 for signal in graceful_signals {
902 signal_errors.extend(signal_unix_session(
903 process.session_id,
904 *signal,
905 Some(known_leader),
906 )?);
907 if wait_for_unix_session_exit(process, WORKER_STOP_GRACE)? {
908 return Ok(());
909 }
910 }
911
912 signal_errors.extend(signal_unix_session(
913 process.session_id,
914 libc::SIGKILL,
915 Some(known_leader),
916 )?);
917 if wait_for_unix_session_exit(process, WORKER_STOP_GRACE)? {
918 return Ok(());
919 }
920
921 // Without a process table we can only reason about the tracked session
922 // leader/dispatcher. Prefer an honest degraded success once that known
923 // pid is gone instead of looping forever on ps EPERM.
924 if !process_table_inspection_available() {
925 if process.last_exit.is_some() && !unix_pid_exists(process.session_id) {
926 return Ok(());
927 }
928 return Err(FleetHostError::retryable(format!(
929 "Fleet session {} still has a live tracked leader after SIGKILL and process-table inspection is unavailable{}",
930 process.session_id,
931 if signal_errors.is_empty() {
932 String::new()
933 } else {
934 format!("; signal errors: {}", signal_errors.join("; "))
935 }
936 )));
937 }
938
939 let alive = unix_session_members(process.session_id, Some(known_leader))?;
940 Err(FleetHostError::retryable(format!(
941 "Fleet session {} still has live processes after SIGKILL: {alive:?}{}",
942 process.session_id,
943 if signal_errors.is_empty() {
944 String::new()
945 } else {
946 format!("; signal errors: {}", signal_errors.join("; "))
947 }
948 )))
949 }
950
951 #[cfg(unix)]
952 fn wait_for_unix_session_exit(
953 process: &mut LocalWorkerProcess,
954 timeout: Duration,
955 ) -> FleetHostResult<bool> {
956 let deadline = Instant::now() + timeout;
957 let known_leader = process.session_id;
958 loop {
959 if process.last_exit.is_none() {
960 process.last_exit = process.child.try_wait().map_err(|err| {
961 FleetHostError::retryable(format!("checking Fleet dispatcher exit: {err}"))
962 })?;
963 }
964 if process.last_exit.is_some() {
965 let members = unix_session_members(process.session_id, Some(known_leader))?;
966 if members.is_empty() {
967 return Ok(true);
968 }
969 // When process-table inspection is denied we can only track the
970 // known session leader. Treat an empty known-pid set as success.
971 if !process_table_inspection_available()
972 && members.iter().all(|pid| !unix_pid_exists(*pid))
973 {
974 return Ok(true);
975 }
976 }
977 if Instant::now() >= deadline {
978 return Ok(false);
979 }
980 thread::sleep(Duration::from_millis(25));
981 }
982 }
983
984 #[cfg(unix)]
985 fn unix_session_members(
986 session_id: libc::pid_t,
987 known_pids: Option<libc::pid_t>,
988 ) -> FleetHostResult<Vec<libc::pid_t>> {
989 match unix_process_ids() {
990 Ok(pids) => {
991 let mut members = Vec::new();
992 for pid in pids {
993 if pid > 0 {
994 // Revalidate against the kernel after parsing the snapshot. A PID
995 // reused by an unrelated process must never receive our signal.
996 // SAFETY: getsid(2) dereferences no pointers.
997 if unsafe { libc::getsid(pid) } == session_id {
998 members.push(pid);
999 }
1000 }
1001 }
1002 Ok(members)
1003 }
1004 Err(err) if process_table_error_is_unavailable(&err) => {
1005 // Restricted sandboxes may deny full process-table walks. Fall back
1006 // to the known session leader so stop/interrupt still reaches the
1007 // tracked dispatcher without inventing a process census.
1008 Ok(known_pids
1009 .into_iter()
1010 .filter(|pid| *pid > 0 && unix_pid_in_session(*pid, session_id))
1011 .collect())
1012 }
1013 Err(err) => Err(err),
1014 }
1015 }
1016
1017 #[cfg(unix)]
1018 fn unix_pid_in_session(pid: libc::pid_t, session_id: libc::pid_t) -> bool {
1019 // SAFETY: getsid(2) dereferences no pointers.
1020 unsafe { libc::getsid(pid) == session_id }
1021 }
1022
1023 #[cfg(unix)]
1024 fn unix_pid_exists(pid: libc::pid_t) -> bool {
1025 if pid <= 0 {
1026 return false;
1027 }
1028 // SAFETY: kill(2) dereferences no pointers; signal 0 sends nothing.
1029 if unsafe { libc::kill(pid, 0) } == 0 {
1030 return true;
1031 }
1032 std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
1033 }
1034
1035 #[cfg(unix)]
1036 fn is_permission_denied(err: &std::io::Error) -> bool {
1037 err.kind() == std::io::ErrorKind::PermissionDenied || err.raw_os_error() == Some(libc::EPERM)
1038 }
1039
1040 #[cfg(unix)]
1041 fn process_table_error_is_unavailable(err: &FleetHostError) -> bool {
1042 err.message.contains("process-table inspection unavailable")
1043 || err.message.contains("Operation not permitted")
1044 || err.message.contains("Permission denied")
1045 || err.message.contains("EPERM")
1046 }
1047
1048 #[cfg(unix)]
1049 fn mark_process_table_unavailable() {
1050 // Record the denial when nothing has been cached yet. OnceLock cannot flip
1051 // a prior true; live call sites still degrade on the immediate EPERM path.
1052 let _ = process_table_probe_cell().get_or_init(|| false);
1053 }
1054
1055 #[cfg(unix)]
1056 fn process_table_probe_cell() -> &'static OnceLock<bool> {
1057 static PROCESS_TABLE_AVAILABLE: OnceLock<bool> = OnceLock::new();
1058 &PROCESS_TABLE_AVAILABLE
1059 }
1060
1061 /// Returns whether full process-table inspection (`ps` / `/proc`) works here.
1062 /// Cached after the first probe so tests and production share one answer.
1063 #[cfg(unix)]
1064 pub(crate) fn process_table_inspection_available() -> bool {
1065 *process_table_probe_cell().get_or_init(|| match unix_process_ids_uncached() {
1066 Ok(_) => true,
1067 Err(err) if process_table_error_is_unavailable(&err) => false,
1068 // Missing `ps` binary is also unavailable inspection, not a transient
1069 // retryable blip for memory sampling / session census.
1070 Err(err)
1071 if err.message.contains("os error 2")
1072 || err.message.contains("No such file")
1073 || err.message.contains("not found") =>
1074 {
1075 false
1076 }
1077 Err(_) => false,
1078 })
1079 }
1080
1081 #[cfg(all(unix, target_os = "linux"))]
1082 fn unix_process_ids() -> FleetHostResult<Vec<libc::pid_t>> {
1083 unix_process_ids_uncached()
1084 }
1085
1086 #[cfg(all(unix, target_os = "linux"))]
1087 fn unix_process_ids_uncached() -> FleetHostResult<Vec<libc::pid_t>> {
1088 let entries = std::fs::read_dir("/proc").map_err(|err| {
1089 if is_permission_denied(&err) {
1090 FleetHostError::retryable(format!(
1091 "listing Fleet session through /proc: process-table inspection unavailable: {err}"
1092 ))
1093 } else {
1094 FleetHostError::retryable(format!("listing Fleet session through /proc: {err}"))
1095 }
1096 })?;
1097 Ok(entries
1098 .filter_map(Result::ok)
1099 .filter_map(|entry| entry.file_name().to_string_lossy().parse().ok())
1100 .collect())
1101 }
1102
1103 #[cfg(all(unix, not(target_os = "linux")))]
1104 fn unix_process_ids() -> FleetHostResult<Vec<libc::pid_t>> {
1105 if let Some(available) = process_table_probe_cell().get()
1106 && !*available
1107 {
1108 return Err(FleetHostError::retryable(
1109 "listing Fleet session with ps: process-table inspection unavailable",
1110 ));
1111 }
1112 match unix_process_ids_uncached() {
1113 Ok(pids) => {
1114 let _ = process_table_probe_cell().get_or_init(|| true);
1115 Ok(pids)
1116 }
1117 Err(err) => {
1118 if process_table_error_is_unavailable(&err) {
1119 let _ = process_table_probe_cell().get_or_init(|| false);
1120 }
1121 Err(err)
1122 }
1123 }
1124 }
1125
1126 #[cfg(all(unix, not(target_os = "linux")))]
1127 fn unix_process_ids_uncached() -> FleetHostResult<Vec<libc::pid_t>> {
1128 let output = Command::new("ps")
1129 .args(["-A", "-o", "pid="])
1130 .output()
1131 .map_err(|err| {
1132 if is_permission_denied(&err) {
1133 FleetHostError::retryable(format!(
1134 "listing Fleet session with ps: process-table inspection unavailable: {err}"
1135 ))
1136 } else {
1137 FleetHostError::retryable(format!("listing Fleet session with ps: {err}"))
1138 }
1139 })?;
1140 if !output.status.success() {
1141 let stderr = String::from_utf8_lossy(&output.stderr);
1142 let denied = stderr.contains("Operation not permitted")
1143 || stderr.contains("Permission denied")
1144 || output.status.code() == Some(1)
1145 && stderr.to_ascii_lowercase().contains("not permitted");
1146 if denied {
1147 return Err(FleetHostError::retryable(format!(
1148 "listing Fleet session with ps: process-table inspection unavailable: {stderr}"
1149 )));
1150 }
1151 return Err(FleetHostError::retryable(format!(
1152 "listing Fleet session with ps exited {:?}",
1153 output.status.code()
1154 )));
1155 }
1156
1157 Ok(String::from_utf8_lossy(&output.stdout)
1158 .lines()
1159 .filter_map(|line| line.trim().parse().ok())
1160 .collect())
1161 }
1162
1163 #[cfg(unix)]
1164 fn signal_unix_session(
1165 session_id: libc::pid_t,
1166 signal: libc::c_int,
1167 known_leader: Option<libc::pid_t>,
1168 ) -> FleetHostResult<Vec<String>> {
1169 // SAFETY: getsid(2) dereferences no pointers.
1170 let own_session = unsafe { libc::getsid(0) };
1171 if session_id <= 0 || session_id == own_session {
1172 return Err(FleetHostError::terminal(format!(
1173 "refusing to signal unsafe Fleet session {session_id}"
1174 )));
1175 }
1176
1177 let mut errors = Vec::new();
1178 // Prefer known leader first so stop/interrupt still works when the full
1179 // process table cannot be enumerated under a restricted sandbox.
1180 let mut candidates = unix_session_members(session_id, known_leader)?;
1181 if candidates.is_empty()
1182 && let Some(leader) = known_leader.filter(|pid| *pid > 0)
1183 {
1184 candidates.push(leader);
1185 }
1186 for pid in candidates {
1187 // Verify identity again immediately before signalling. Session IDs
1188 // remain stable across reparenting and separate process groups.
1189 // SAFETY: getsid(2) dereferences no pointers.
1190 if unsafe { libc::getsid(pid) } != session_id {
1191 // Leader may already be gone; still try kill on known leader when
1192 // getsid fails only with ESRCH-equivalent absence.
1193 if Some(pid) != known_leader || !unix_pid_exists(pid) {
1194 continue;
1195 }
1196 }
1197 // SAFETY: kill(2) dereferences no pointers.
1198 if unsafe { libc::kill(pid, signal) } != 0 {
1199 let err = std::io::Error::last_os_error();
1200 if err.raw_os_error() != Some(libc::ESRCH) {
1201 errors.push(format!("pid {pid}: {err}"));
1202 }
1203 }
1204 }
1205 Ok(errors)
1206 }
1207
1208 #[cfg(unix)]
1209 fn unix_pid_is_running(pid: libc::pid_t) -> bool {
1210 if !unix_pid_exists(pid) {
1211 return false;
1212 }
1213
1214 // `kill(pid, 0)` also succeeds for zombies. The Fleet containment code
1215 // has already finished its job once a descendant is dead; on macOS an
1216 // orphan can remain visible as a zombie briefly while launchd reaps it.
1217 // Ask `ps` for the process state so test assertions do not mistake that
1218 // transient kernel bookkeeping for a live leaked worker. If `ps` itself
1219 // is denied, stay conservative and treat the PID as running.
1220 if !process_table_inspection_available() {
1221 return true;
1222 }
1223 match Command::new("ps")
1224 .args(["-o", "stat=", "-p", &pid.to_string()])
1225 .output()
1226 {
1227 Ok(output) if output.status.success() => String::from_utf8_lossy(&output.stdout)
1228 .split_whitespace()
1229 .next()
1230 .is_some_and(|state| !state.starts_with('Z')),
1231 // A failed status does not prove exit: preserve the positive kernel
1232 // visibility result and let the bounded waiter retry.
1233 Ok(_) => true,
1234 Err(err) if is_permission_denied(&err) => {
1235 let _ = process_table_probe_cell().get_or_init(|| false);
1236 true
1237 }
1238 Err(_) => true,
1239 }
1240 }
1241
1242 #[cfg(all(unix, test))]
1243 fn wait_for_unix_pid_exit(pid: libc::pid_t, timeout: Duration) -> bool {
1244 let deadline = Instant::now() + timeout;
1245 loop {
1246 if !unix_pid_is_running(pid) {
1247 return true;
1248 }
1249 if Instant::now() >= deadline {
1250 return false;
1251 }
1252 thread::sleep(Duration::from_millis(25));
1253 }
1254 }
1255
1256 #[cfg(windows)]
1257 #[derive(Debug)]
1258 struct FleetWindowsJob {
1259 handle: HANDLE,
1260 }
1261
1262 #[cfg(windows)]
1263 // SAFETY: Job handles are process-wide kernel handles. The adapter owns this
1264 // wrapper exclusively and mutates workers through `&mut self`.
1265 unsafe impl Send for FleetWindowsJob {}
1266
1267 #[cfg(windows)]
1268 // SAFETY: The wrapper exposes only kernel job operations; shared access does
1269 // not mutate Rust-owned memory.
1270 unsafe impl Sync for FleetWindowsJob {}
1271
1272 #[cfg(windows)]
1273 impl FleetWindowsJob {
1274 fn attach_to_child(child: &Child) -> std::io::Result<Self> {
1275 // SAFETY: returned handle is owned by the new wrapper.
1276 let handle = unsafe { CreateJobObjectW(None, PCWSTR::null()).map_err(windows_io_error)? };
1277 let job = Self { handle };
1278 let mut limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default();
1279 limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
1280 // SAFETY: `limits` is live with matching size; both handles are live.
1281 unsafe {
1282 SetInformationJobObject(
1283 job.handle,
1284 JobObjectExtendedLimitInformation,
1285 &limits as *const _ as *const core::ffi::c_void,
1286 std::mem::size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
1287 )
1288 .map_err(windows_io_error)?;
1289 AssignProcessToJobObject(job.handle, HANDLE(child.as_raw_handle()))
1290 .map_err(windows_io_error)?;
1291 }
1292 Ok(job)
1293 }
1294
1295 fn terminate(&self) -> std::io::Result<()> {
1296 // SAFETY: `self.handle` is a live owned job handle.
1297 unsafe { TerminateJobObject(self.handle, 1).map_err(windows_io_error) }
1298 }
1299
1300 fn has_active_processes(&self) -> std::io::Result<bool> {
1301 let mut accounting = JOBOBJECT_BASIC_ACCOUNTING_INFORMATION::default();
1302 // SAFETY: `accounting` is live with matching size.
1303 unsafe {
1304 QueryInformationJobObject(
1305 Some(self.handle),
1306 JobObjectBasicAccountingInformation,
1307 &mut accounting as *mut _ as *mut core::ffi::c_void,
1308 std::mem::size_of::<JOBOBJECT_BASIC_ACCOUNTING_INFORMATION>() as u32,
1309 None,
1310 )
1311 .map_err(windows_io_error)?;
1312 }
1313 Ok(accounting.ActiveProcesses > 0)
1314 }
1315 }
1316
1317 #[cfg(windows)]
1318 impl Drop for FleetWindowsJob {
1319 fn drop(&mut self) {
1320 // SAFETY: `self.handle` is owned here; Drop runs once.
1321 unsafe {
1322 let _ = CloseHandle(self.handle);
1323 }
1324 }
1325 }
1326
1327 #[cfg(windows)]
1328 fn attach_fleet_windows_job(mut child: Child) -> std::io::Result<(Child, FleetWindowsJob)> {
1329 match FleetWindowsJob::attach_to_child(&child) {
1330 Ok(job) => Ok((child, job)),
1331 Err(err) => {
1332 let _ = child.kill();
1333 let _ = child.wait();
1334 Err(err)
1335 }
1336 }
1337 }
1338
1339 #[cfg(windows)]
1340 fn windows_io_error(error: windows::core::Error) -> std::io::Error {
1341 std::io::Error::other(error)
1342 }
1343
1344 fn filtered_env(
1345 env: &BTreeMap<String, String>,
1346 allowlist: &BTreeSet<String>,
1347 ) -> FleetHostResult<BTreeMap<String, String>> {
1348 validate_env_allowlist(allowlist)?;
1349 Ok(env
1350 .iter()
1351 .filter(|(key, _)| allowlist.contains(*key))
1352 .map(|(key, value)| (key.clone(), value.clone()))
1353 .collect())
1354 }
1355
1356 fn validate_env_allowlist(allowlist: &BTreeSet<String>) -> FleetHostResult<()> {
1357 for key in allowlist {
1358 if !is_safe_env_key(key) {
1359 return Err(FleetHostError::configuration(format!(
1360 "Fleet host env allowlist key {key} looks secret-bearing; pass secrets through config providers, not worker argv/env"
1361 )));
1362 }
1363 }
1364 Ok(())
1365 }
1366
1367 fn is_safe_env_key(key: &str) -> bool {
1368 let upper = key.to_ascii_uppercase();
1369 ![
1370 "SECRET",
1371 "TOKEN",
1372 "PASSWORD",
1373 "PASSWD",
1374 "API_KEY",
1375 "CREDENTIAL",
1376 "PRIVATE_KEY",
1377 ]
1378 .iter()
1379 .any(|needle| upper.contains(needle))
1380 }
1381
1382 fn ssh_client_env() -> BTreeMap<String, String> {
1383 ["HOME", "PATH", "SSH_AUTH_SOCK"]
1384 .into_iter()
1385 .filter_map(|key| {
1386 std::env::var(key)
1387 .ok()
1388 .map(|value| (key.to_string(), value))
1389 })
1390 .collect()
1391 }
1392
1393 fn process_base_env() -> BTreeMap<String, String> {
1394 let mut env = BTreeMap::new();
1395 for key in [
1396 "HOME",
1397 "PATH",
1398 "SYSTEMROOT",
1399 "SystemRoot",
1400 "COMSPEC",
1401 "ComSpec",
1402 ] {
1403 if let Ok(value) = std::env::var(key) {
1404 env.insert(key.to_string(), value);
1405 }
1406 }
1407 force_worker_telemetry_off(&mut env);
1408 env
1409 }
1410
1411 /// Fleet workers are an implementation detail of the parent session, not
1412 /// sessions of their own — the parent already accounts for the dispatch.
1413 ///
1414 /// The spawn path `env_clear()`s and rebuilds from [`process_base_env`], so an
1415 /// operator's opt-out would otherwise never reach a worker at all. Hard-off
1416 /// here means a worker can never emit, can never inherit an ambient "on", and
1417 /// can never write telemetry state into the operator's home.
1418 fn force_worker_telemetry_off(env: &mut BTreeMap<String, String>) {
1419 env.insert("CODEWHALE_TELEMETRY".to_string(), "false".to_string());
1420 env.insert("DEEPSEEK_TELEMETRY".to_string(), "false".to_string());
1421 }
1422
1423 /// Build the complete environment a worker is spawned with.
1424 ///
1425 /// The caller's allowlisted entries are merged over the process base, then
1426 /// telemetry is forced off again: an allowlist that happens to name
1427 /// `CODEWHALE_TELEMETRY` must not be able to switch a worker back on.
1428 fn worker_env(
1429 request_env: &BTreeMap<String, String>,
1430 allowlist: &BTreeSet<String>,
1431 ) -> FleetHostResult<BTreeMap<String, String>> {
1432 let mut env = process_base_env();
1433 env.extend(filtered_env(request_env, allowlist)?);
1434 force_worker_telemetry_off(&mut env);
1435 Ok(env)
1436 }
1437
1438 fn shell_quote(value: &str) -> String {
1439 if value.is_empty() {
1440 return "''".to_string();
1441 }
1442 format!("'{}'", value.replace('\'', "'\\''"))
1443 }
1444
1445 fn validate_worker_id(worker_id: &str) -> FleetHostResult<()> {
1446 if worker_id.trim().is_empty() {
1447 return Err(FleetHostError::configuration("worker id cannot be empty"));
1448 }
1449 Ok(())
1450 }
1451
1452 fn safe_path_segment(value: &str) -> String {
1453 value
1454 .chars()
1455 .map(|ch| {
1456 if ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.') {
1457 ch
1458 } else {
1459 '_'
1460 }
1461 })
1462 .collect()
1463 }
1464
1465 #[cfg(test)]
1466 mod tests {
1467 use super::*;
1468 use tempfile::TempDir;
1469
1470 #[cfg(unix)]
1471 fn skip_if_process_table_unavailable() -> bool {
1472 if process_table_inspection_available() {
1473 return false;
1474 }
1475 eprintln!("skipping: process-table inspection unavailable (ps/proc denied or missing)");
1476 true
1477 }
1478
1479 #[cfg(unix)]
1480 #[test]
1481 fn sample_process_memory_reports_nonzero_for_self() {
1482 if skip_if_process_table_unavailable() {
1483 return;
1484 }
1485 // The current test process is alive, so its RSS must sample to Some(>0).
1486 let mb = sample_process_memory_mb(std::process::id());
1487 assert!(
1488 matches!(mb, Some(v) if v > 0),
1489 "expected Some(>0) MB for the live self process, got {mb:?}"
1490 );
1491 }
1492
1493 #[cfg(unix)]
1494 #[test]
1495 fn sample_process_memory_is_none_for_dead_pid() {
1496 // Use a PID beyond every mainstream kernel's default pid ceiling
1497 // (Linux pid_max default 4M/32k, macOS ~99998, BSDs 99999): PID 0 is
1498 // kernel_task on macOS and semantically special to `ps -p`, so it is
1499 // not a portable "no such process" probe. When process-table inspection
1500 // is denied the sampler also returns None — same observable result.
1501 assert_eq!(sample_process_memory_mb(999_999_999), None);
1502 }
1503
1504 fn shell_command(script: &str) -> FleetWorkerCommand {
1505 if cfg!(windows) {
1506 FleetWorkerCommand::new("cmd", ["/C", script])
1507 } else {
1508 FleetWorkerCommand::new("sh", ["-c", script])
1509 }
1510 }
1511
1512 #[cfg(unix)]
1513 const DESCENDANT_HELPER_TEST: &str =
1514 "fleet::host::tests::fleet_host_stop_reaps_dispatcher_descendants";
1515
1516 #[cfg(unix)]
1517 fn run_descendant_helper_if_requested() -> bool {
1518 let Ok(mode) = std::env::var("FLEET_DESCENDANT_HELPER") else {
1519 return false;
1520 };
1521 let test_binary = std::env::current_exe().expect("current test binary");
1522 let pid_file = std::env::var("FLEET_DESCENDANT_PID_FILE").expect("helper pid file");
1523 match mode.as_str() {
1524 "dispatcher" | "detached-dispatcher" => {
1525 let mut command = Command::new(&test_binary);
1526 command
1527 .args(["--exact", DESCENDANT_HELPER_TEST, "--nocapture"])
1528 .env("FLEET_DESCENDANT_HELPER", "worker")
1529 .env("FLEET_DESCENDANT_PID_FILE", &pid_file);
1530 if mode == "detached-dispatcher" {
1531 command.spawn().expect("spawn detached dispatcher child");
1532 std::process::exit(0);
1533 }
1534 let status = command.status().expect("spawn dispatcher child");
1535 std::process::exit(status.code().unwrap_or(1));
1536 }
1537 "worker" => {
1538 let mut command = Command::new(&test_binary);
1539 command
1540 .args(["--exact", DESCENDANT_HELPER_TEST, "--nocapture"])
1541 .env("FLEET_DESCENDANT_HELPER", "tool")
1542 .env("FLEET_DESCENDANT_PID_FILE", &pid_file);
1543 // Real shell tools deliberately own a separate process group.
1544 // This makes a root-group-only Fleet stop leak the helper.
1545 command.process_group(0);
1546 let status = command.status().expect("spawn worker tool");
1547 std::process::exit(status.code().unwrap_or(1));
1548 }
1549 "tool" => {
1550 // A shell tool can ignore graceful signals and live in its own
1551 // process group. Fleet's session boundary must still reap it.
1552 unsafe {
1553 libc::signal(libc::SIGINT, libc::SIG_IGN);
1554 libc::signal(libc::SIGTERM, libc::SIG_IGN);
1555 }
1556 std::fs::write(&pid_file, std::process::id().to_string()).expect("write tool pid");
1557 thread::sleep(Duration::from_secs(30));
1558 true
1559 }
1560 other => panic!("unknown descendant helper mode {other}"),
1561 }
1562 }
1563
1564 #[cfg(unix)]
1565 fn start_dispatcher_tree(
1566 adapter: &mut LocalProcessFleetHostAdapter,
1567 tmp: &TempDir,
1568 worker_id: &str,
1569 helper_mode: &str,
1570 ) -> (libc::pid_t, libc::pid_t) {
1571 let pid_file = tmp.path().join(format!("{worker_id}-tool.pid"));
1572 let test_binary = std::env::current_exe().expect("current test binary");
1573 let mut request = FleetWorkerStartRequest::new(
1574 worker_id,
1575 FleetWorkerCommand::new(
1576 test_binary.display().to_string(),
1577 ["--exact", DESCENDANT_HELPER_TEST, "--nocapture"],
1578 ),
1579 );
1580 request.env.insert(
1581 "FLEET_DESCENDANT_HELPER".to_string(),
1582 helper_mode.to_string(),
1583 );
1584 request.env.insert(
1585 "FLEET_DESCENDANT_PID_FILE".to_string(),
1586 pid_file.display().to_string(),
1587 );
1588 request.env_allowlist = BTreeSet::from([
1589 "FLEET_DESCENDANT_HELPER".to_string(),
1590 "FLEET_DESCENDANT_PID_FILE".to_string(),
1591 ]);
1592
1593 let handle = adapter.start_worker(request).expect("start dispatcher");
1594 let root_pid = handle.pid.expect("dispatcher pid") as libc::pid_t;
1595 let tool_pid = wait_for_valid_pid_file(&pid_file, Duration::from_secs(5));
1596 if helper_mode != "detached-dispatcher" {
1597 assert!(unix_pid_is_running(root_pid));
1598 }
1599 assert!(unix_pid_is_running(tool_pid));
1600 (root_pid, tool_pid)
1601 }
1602
1603 #[cfg(unix)]
1604 fn wait_for_host_state(
1605 adapter: &mut LocalProcessFleetHostAdapter,
1606 worker_id: &str,
1607 expected: FleetHostWorkerState,
1608 timeout: Duration,
1609 ) -> FleetHostWorkerStatus {
1610 let deadline = Instant::now() + timeout;
1611 loop {
1612 let status = adapter.read_status(worker_id).expect("worker status");
1613 if status.state == expected || Instant::now() >= deadline {
1614 return status;
1615 }
1616 thread::sleep(Duration::from_millis(25));
1617 }
1618 }
1619
1620 #[cfg(unix)]
1621 fn wait_for_valid_pid_file(pid_file: &Path, timeout: Duration) -> libc::pid_t {
1622 let deadline = Instant::now() + timeout;
1623 let mut last_observation = "file not created".to_string();
1624 loop {
1625 match std::fs::read_to_string(pid_file) {
1626 Ok(contents) => {
1627 let trimmed = contents.trim();
1628 match trimmed.parse::<libc::pid_t>() {
1629 Ok(pid) if pid > 0 => return pid,
1630 _ => last_observation = format!("invalid contents {trimmed:?}"),
1631 }
1632 }
1633 Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
1634 Err(err) => last_observation = format!("read failed: {err}"),
1635 }
1636 if Instant::now() >= deadline {
1637 panic!(
1638 "separate-group tool never published a valid PID to {} ({last_observation})",
1639 pid_file.display()
1640 );
1641 }
1642 thread::sleep(Duration::from_millis(25));
1643 }
1644 }
1645
1646 #[cfg(unix)]
1647 #[test]
1648 fn pid_file_wait_ignores_created_but_incomplete_file() {
1649 let tmp = TempDir::new().unwrap();
1650 let pid_file = tmp.path().join("worker.pid");
1651 std::fs::write(&pid_file, "pid=").unwrap();
1652 let expected_pid = std::process::id() as libc::pid_t;
1653 let writer_path = pid_file.clone();
1654 let writer = thread::spawn(move || {
1655 thread::sleep(Duration::from_millis(50));
1656 std::fs::write(writer_path, expected_pid.to_string()).unwrap();
1657 });
1658
1659 assert_eq!(
1660 wait_for_valid_pid_file(&pid_file, Duration::from_secs(1)),
1661 expected_pid
1662 );
1663 writer.join().unwrap();
1664 }
1665
1666 #[cfg(unix)]
1667 #[test]
1668 fn unix_pid_running_treats_zombie_as_exited() {
1669 if skip_if_process_table_unavailable() {
1670 return;
1671 }
1672 let mut child = Command::new("sh")
1673 .args(["-c", "exit 0"])
1674 .spawn()
1675 .expect("spawn short-lived child");
1676 let pid = child.id() as libc::pid_t;
1677 let deadline = Instant::now() + Duration::from_secs(2);
1678 let saw_zombie = loop {
1679 let state = match Command::new("ps")
1680 .args(["-o", "stat=", "-p", &pid.to_string()])
1681 .output()
1682 {
1683 Ok(output) => output,
1684 Err(err) if is_permission_denied(&err) => {
1685 mark_process_table_unavailable();
1686 child.wait().ok();
1687 eprintln!("skipping: ps denied while inspecting zombie state");
1688 return;
1689 }
1690 Err(err) => panic!("inspect child state: {err}"),
1691 };
1692 let is_zombie = String::from_utf8_lossy(&state.stdout)
1693 .split_whitespace()
1694 .next()
1695 .is_some_and(|state| state.starts_with('Z'));
1696 if is_zombie {
1697 break true;
1698 }
1699 if Instant::now() >= deadline {
1700 break false;
1701 }
1702 thread::sleep(Duration::from_millis(10));
1703 };
1704
1705 let reported_running = unix_pid_is_running(pid);
1706 child.wait().expect("reap zombie child");
1707 assert!(saw_zombie, "child never became a zombie");
1708 assert!(!reported_running, "zombie was reported as running");
1709 }
1710
1711 fn wait_for_log(
1712 adapter: &LocalProcessFleetHostAdapter,
1713 worker_id: &str,
1714 needle: &str,
1715 ) -> String {
1716 let deadline = Instant::now() + Duration::from_secs(3);
1717 loop {
1718 let logs = adapter.read_logs(worker_id, 4096).unwrap();
1719 if logs.contains(needle) || Instant::now() > deadline {
1720 return logs;
1721 }
1722 thread::sleep(Duration::from_millis(25));
1723 }
1724 }
1725
1726 #[test]
1727 fn fleet_host_local_adapter_starts_reads_bounded_logs_and_stops() {
1728 let tmp = TempDir::new().unwrap();
1729 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1730 let script = if cfg!(windows) {
1731 "echo 0123456789abcdef& ping -n 30 127.0.0.1 >NUL"
1732 } else {
1733 "printf 0123456789abcdef; sleep 30"
1734 };
1735 let mut request = FleetWorkerStartRequest::new("local-1", shell_command(script));
1736 let line_ending_bytes = if cfg!(windows) { 2 } else { 0 };
1737 request.log_limit_bytes = 16 + line_ending_bytes;
1738
1739 let handle = adapter.start_worker(request).unwrap();
1740 #[cfg(unix)]
1741 let direct_pid = handle.pid.expect("local worker pid");
1742 assert_eq!(handle.host_kind, FleetHostKind::LocalProcess);
1743 assert!(handle.pid.is_some());
1744 let status = adapter.read_status("local-1").unwrap();
1745 assert_eq!(status.state, FleetHostWorkerState::Running);
1746
1747 let logs = wait_for_log(&adapter, "local-1", "abcdef");
1748 let logs = logs.trim_end_matches(&['\r', '\n'][..]);
1749 assert!(logs.ends_with("0123456789abcdef"), "{logs:?}");
1750 let bounded = adapter.read_logs("local-1", 6 + line_ending_bytes).unwrap();
1751 let bounded = bounded.trim_end_matches(&['\r', '\n'][..]);
1752 assert!(bounded.ends_with("abcdef"), "{bounded:?}");
1753
1754 let status = adapter.stop_worker("local-1").unwrap();
1755 assert_eq!(status.state, FleetHostWorkerState::Stopped);
1756 #[cfg(unix)]
1757 assert!(
1758 wait_for_unix_pid_exit(direct_pid as libc::pid_t, Duration::from_secs(1)),
1759 "stopped direct worker was not reaped"
1760 );
1761 adapter.cleanup_worker("local-1").unwrap();
1762 assert_eq!(
1763 adapter.read_status("local-1").unwrap_err().kind,
1764 FleetHostErrorKind::Terminal
1765 );
1766 }
1767
1768 #[cfg(unix)]
1769 #[test]
1770 fn fleet_host_stop_reaps_dispatcher_descendants() {
1771 if run_descendant_helper_if_requested() {
1772 return;
1773 }
1774 // Full-session reaping of separate process-group tools requires a
1775 // process-table walk; without it production still signals the known
1776 // session leader and these assertions cannot be proven.
1777 if skip_if_process_table_unavailable() {
1778 return;
1779 }
1780
1781 let tmp = TempDir::new().unwrap();
1782 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1783 let (root_pid, tool_pid) =
1784 start_dispatcher_tree(&mut adapter, &tmp, "dispatcher-tree", "dispatcher");
1785
1786 let status = adapter
1787 .stop_worker("dispatcher-tree")
1788 .expect("stop complete worker tree");
1789
1790 assert_eq!(status.state, FleetHostWorkerState::Stopped);
1791 assert!(
1792 wait_for_unix_pid_exit(root_pid, Duration::from_secs(1)),
1793 "dispatcher survived stop"
1794 );
1795 assert!(
1796 wait_for_unix_pid_exit(tool_pid, Duration::from_secs(1)),
1797 "separate-process-group tool survived stop"
1798 );
1799 }
1800
1801 #[cfg(unix)]
1802 #[test]
1803 fn fleet_host_interrupt_reaps_dispatcher_descendants() {
1804 if skip_if_process_table_unavailable() {
1805 return;
1806 }
1807 let tmp = TempDir::new().unwrap();
1808 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1809 let (root_pid, tool_pid) =
1810 start_dispatcher_tree(&mut adapter, &tmp, "interrupt-tree", "dispatcher");
1811
1812 let status = adapter
1813 .interrupt_worker("interrupt-tree")
1814 .expect("interrupt complete worker session");
1815
1816 assert_ne!(status.state, FleetHostWorkerState::Running);
1817 assert!(
1818 wait_for_unix_pid_exit(root_pid, Duration::from_secs(1)),
1819 "dispatcher survived interrupt"
1820 );
1821 assert!(
1822 wait_for_unix_pid_exit(tool_pid, Duration::from_secs(1)),
1823 "separate-process-group tool survived interrupt"
1824 );
1825 }
1826
1827 #[cfg(unix)]
1828 #[test]
1829 fn fleet_host_reports_draining_after_dispatcher_exits_with_live_descendant() {
1830 if skip_if_process_table_unavailable() {
1831 return;
1832 }
1833 let tmp = TempDir::new().unwrap();
1834 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1835 let (root_pid, tool_pid) = start_dispatcher_tree(
1836 &mut adapter,
1837 &tmp,
1838 "draining-dispatcher-tree",
1839 "detached-dispatcher",
1840 );
1841
1842 assert!(
1843 wait_for_unix_pid_exit(root_pid, Duration::from_secs(3)),
1844 "dispatcher did not exit"
1845 );
1846 let status = wait_for_host_state(
1847 &mut adapter,
1848 "draining-dispatcher-tree",
1849 FleetHostWorkerState::Draining,
1850 Duration::from_secs(3),
1851 );
1852 assert_eq!(status.state, FleetHostWorkerState::Draining);
1853 assert!(
1854 unix_pid_is_running(tool_pid),
1855 "descendant exited before draining check"
1856 );
1857
1858 let stopped = adapter
1859 .stop_worker("draining-dispatcher-tree")
1860 .expect("stop draining worker tree");
1861 assert_eq!(stopped.state, FleetHostWorkerState::Stopped);
1862 assert!(
1863 wait_for_unix_pid_exit(tool_pid, Duration::from_secs(1)),
1864 "draining descendant survived bounded stop"
1865 );
1866 }
1867
1868 #[cfg(unix)]
1869 #[test]
1870 fn fleet_host_cleanup_reaps_session_after_dispatcher_exits() {
1871 if skip_if_process_table_unavailable() {
1872 return;
1873 }
1874 let tmp = TempDir::new().unwrap();
1875 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1876 let (root_pid, tool_pid) = start_dispatcher_tree(
1877 &mut adapter,
1878 &tmp,
1879 "exited-dispatcher-tree",
1880 "detached-dispatcher",
1881 );
1882 let deadline = Instant::now() + Duration::from_secs(3);
1883 loop {
1884 let status = adapter.read_status("exited-dispatcher-tree").unwrap();
1885 if status.state != FleetHostWorkerState::Running || Instant::now() >= deadline {
1886 assert_ne!(status.state, FleetHostWorkerState::Running);
1887 break;
1888 }
1889 thread::sleep(Duration::from_millis(25));
1890 }
1891 assert!(
1892 wait_for_unix_pid_exit(root_pid, Duration::from_secs(1)),
1893 "dispatcher should have exited"
1894 );
1895 assert!(
1896 unix_pid_is_running(tool_pid),
1897 "delegated tool exited too early"
1898 );
1899
1900 adapter
1901 .cleanup_worker("exited-dispatcher-tree")
1902 .expect("clean up surviving dispatcher session");
1903
1904 assert!(
1905 wait_for_unix_pid_exit(tool_pid, Duration::from_secs(1)),
1906 "tool survived after its dispatcher exited"
1907 );
1908 assert_eq!(
1909 adapter
1910 .read_status("exited-dispatcher-tree")
1911 .unwrap_err()
1912 .kind,
1913 FleetHostErrorKind::Terminal
1914 );
1915 }
1916
1917 #[test]
1918 fn fleet_host_local_adapter_restarts_worker_with_same_request() {
1919 let tmp = TempDir::new().unwrap();
1920 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1921 let script = if cfg!(windows) {
1922 "echo restart-ready & ping -n 30 127.0.0.1 >NUL"
1923 } else {
1924 "printf restart-ready; sleep 30"
1925 };
1926 let request = FleetWorkerStartRequest::new("local-restart", shell_command(script));
1927 let first = adapter.start_worker(request).unwrap();
1928 let restarted = adapter.restart_worker("local-restart").unwrap();
1929
1930 assert_eq!(restarted.worker_id, first.worker_id);
1931 assert_eq!(restarted.host_kind, FleetHostKind::LocalProcess);
1932 assert_ne!(restarted.pid, first.pid);
1933 let logs = wait_for_log(&adapter, "local-restart", "restart-ready");
1934 assert!(logs.contains("restart-ready"));
1935 adapter.stop_worker("local-restart").unwrap();
1936 }
1937
1938 #[cfg(unix)]
1939 #[test]
1940 fn fleet_host_local_adapter_reports_running_worker_memory_usage() {
1941 if skip_if_process_table_unavailable() {
1942 return;
1943 }
1944 let tmp = TempDir::new().unwrap();
1945 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1946 let request =
1947 FleetWorkerStartRequest::new("local-memory", shell_command("printf ready; sleep 30"));
1948
1949 adapter.start_worker(request).unwrap();
1950 let _ = wait_for_log(&adapter, "local-memory", "ready");
1951
1952 let status = adapter.read_status("local-memory").unwrap();
1953
1954 assert_eq!(status.state, FleetHostWorkerState::Running);
1955 assert!(
1956 status.memory_mb.is_some_and(|memory_mb| memory_mb > 0),
1957 "running local worker status should include RSS memory_mb, got {status:?}"
1958 );
1959
1960 adapter.stop_worker("local-memory").unwrap();
1961 }
1962
1963 #[cfg(unix)]
1964 #[test]
1965 fn fleet_host_stop_signals_known_leader_without_process_table() {
1966 // Even when full session census is unavailable, stop must still reach
1967 // the tracked session leader/dispatcher pid.
1968 let tmp = TempDir::new().unwrap();
1969 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1970 let request = FleetWorkerStartRequest::new(
1971 "known-leader-stop",
1972 shell_command("printf ready; sleep 30"),
1973 );
1974 let handle = adapter.start_worker(request).unwrap();
1975 let pid = handle.pid.expect("pid") as libc::pid_t;
1976 let _ = wait_for_log(&adapter, "known-leader-stop", "ready");
1977
1978 let status = adapter
1979 .stop_worker("known-leader-stop")
1980 .expect("stop known leader without process table");
1981 assert_eq!(status.state, FleetHostWorkerState::Stopped);
1982 assert!(
1983 wait_for_unix_pid_exit(pid, Duration::from_secs(2)),
1984 "known session leader survived stop without process-table census"
1985 );
1986 adapter.cleanup_worker("known-leader-stop").unwrap();
1987 }
1988
1989 #[test]
1990 fn fleet_host_ssh_kind_does_not_report_local_process_memory() {
1991 let tmp = TempDir::new().unwrap();
1992 let mut adapter = LocalProcessFleetHostAdapter::new(tmp.path());
1993 let script = if cfg!(windows) {
1994 "echo ready & ping -n 30 127.0.0.1 >NUL"
1995 } else {
1996 "printf ready; sleep 30"
1997 };
1998 let request = FleetWorkerStartRequest::new("ssh-memory", shell_command(script));
1999
2000 adapter
2001 .start_with_kind(request, FleetHostKind::Ssh)
2002 .unwrap();
2003 let _ = wait_for_log(&adapter, "ssh-memory", "ready");
2004
2005 let status = adapter.read_status("ssh-memory").unwrap();
2006
2007 assert_eq!(status.state, FleetHostWorkerState::Running);
2008 assert_eq!(status.memory_mb, None);
2009
2010 adapter.stop_worker("ssh-memory").unwrap();
2011 }
2012
2013 #[test]
2014 fn fleet_host_rejects_secret_like_env_allowlist_keys() {
2015 let mut env = BTreeMap::new();
2016 env.insert("DEEPSEEK_API_KEY".to_string(), "secret".to_string());
2017 let allowlist = BTreeSet::from(["DEEPSEEK_API_KEY".to_string()]);
2018
2019 let err = filtered_env(&env, &allowlist).unwrap_err();
2020
2021 assert_eq!(err.kind, FleetHostErrorKind::Configuration);
2022 assert!(err.message.contains("looks secret-bearing"));
2023 }
2024
2025 #[test]
2026 fn fleet_host_ssh_command_uses_sendenv_without_argv_secret_values() {
2027 let tmp = TempDir::new().unwrap();
2028 let mut config = SshFleetHostConfig::new("builder.example.test", "/srv/codewhale");
2029 config.user = Some("fleet".to_string());
2030 config.port = Some(2222);
2031 config.identity = Some(PathBuf::from("/tmp/fleet_id"));
2032 config.codewhale_binary = "/usr/local/bin/codewhale".to_string();
2033 config.env_allowlist = BTreeSet::from(["FLEET_PROFILE".to_string()]);
2034 let adapter = SshFleetHostAdapter::new(tmp.path(), config).unwrap();
2035 let mut request = FleetWorkerStartRequest::new(
2036 "ssh-1",
2037 FleetWorkerCommand::new("codewhale", ["fleet-worker", "noop"]),
2038 );
2039 request.env.insert(
2040 "FLEET_PROFILE".to_string(),
2041 "super-secret-profile-value".to_string(),
2042 );
2043
2044 let command = adapter.build_ssh_command(&request).unwrap();
2045 let argv = command.args.join(" ");
2046
2047 assert_eq!(command.program, "ssh");
2048 assert!(argv.contains("BatchMode=yes"));
2049 assert!(argv.contains("SendEnv=FLEET_PROFILE"));
2050 assert!(argv.contains("fleet@builder.example.test"));
2051 assert!(argv.contains("/usr/local/bin/codewhale"));
2052 assert!(argv.contains("fleet-worker"));
2053 assert!(!argv.contains("super-secret-profile-value"));
2054 }
2055
2056 #[test]
2057 fn fleet_host_ssh_config_requires_explicit_safe_fields() {
2058 let tmp = TempDir::new().unwrap();
2059 let mut config = SshFleetHostConfig::new("", "/srv/codewhale");
2060 config.env_allowlist = BTreeSet::from(["SAFE_FLAG".to_string()]);
2061
2062 let err = SshFleetHostAdapter::new(tmp.path(), config).unwrap_err();
2063
2064 assert_eq!(err.kind, FleetHostErrorKind::Configuration);
2065 assert!(err.message.contains("explicit host"));
2066 }
2067
2068 #[test]
2069 fn fleet_host_ssh_config_maps_from_protocol_host_spec() {
2070 let spec = FleetHostSpec::Ssh {
2071 host: "builder.example.test".to_string(),
2072 port: Some(2222),
2073 user: Some("fleet".to_string()),
2074 identity: Some(PathBuf::from("/tmp/fleet_id")),
2075 known_hosts: None,
2076 host_key_fingerprint: None,
2077 working_directory: Some(PathBuf::from("/srv/codewhale")),
2078 env_allowlist: vec!["FLEET_PROFILE".to_string()],
2079 codewhale_binary: Some("/usr/local/bin/codewhale".to_string()),
2080 };
2081
2082 let config = SshFleetHostConfig::from_host_spec(&spec).unwrap();
2083
2084 assert_eq!(config.host, "builder.example.test");
2085 assert_eq!(config.port, Some(2222));
2086 assert_eq!(config.user.as_deref(), Some("fleet"));
2087 assert_eq!(config.working_directory, PathBuf::from("/srv/codewhale"));
2088 assert!(config.env_allowlist.contains("FLEET_PROFILE"));
2089 assert_eq!(config.codewhale_binary, "/usr/local/bin/codewhale");
2090 }
2091
2092 #[test]
2093 fn worker_env_forces_telemetry_off_even_when_the_allowlist_says_otherwise() {
2094 // The spawn path env_clear()s and rebuilds from this map, so anything
2095 // absent here simply does not exist inside the worker.
2096 let base = process_base_env();
2097 assert_eq!(
2098 base.get("CODEWHALE_TELEMETRY").map(String::as_str),
2099 Some("false")
2100 );
2101 assert_eq!(
2102 base.get("DEEPSEEK_TELEMETRY").map(String::as_str),
2103 Some("false")
2104 );
2105
2106 // A caller that allowlists the switch cannot switch it back on.
2107 let request_env = BTreeMap::from([
2108 ("CODEWHALE_TELEMETRY".to_string(), "true".to_string()),
2109 ("DEEPSEEK_TELEMETRY".to_string(), "1".to_string()),
2110 ("FLEET_PROFILE".to_string(), "builder".to_string()),
2111 ]);
2112 let allowlist = BTreeSet::from([
2113 "CODEWHALE_TELEMETRY".to_string(),
2114 "DEEPSEEK_TELEMETRY".to_string(),
2115 "FLEET_PROFILE".to_string(),
2116 ]);
2117
2118 let env = worker_env(&request_env, &allowlist).expect("worker env");
2119 assert_eq!(
2120 env.get("CODEWHALE_TELEMETRY").map(String::as_str),
2121 Some("false")
2122 );
2123 assert_eq!(
2124 env.get("DEEPSEEK_TELEMETRY").map(String::as_str),
2125 Some("false")
2126 );
2127 // Unrelated allowlisted entries still come through.
2128 assert_eq!(
2129 env.get("FLEET_PROFILE").map(String::as_str),
2130 Some("builder")
2131 );
2132 }
2133
2134 /// The env map above is only a claim about a function. This dispatches a
2135 /// real worker through the real spawn path and reads what the worker
2136 /// process actually received, because that is the environment a Codewhale
2137 /// worker would resolve telemetry from.
2138 ///
2139 /// Also asserts the operator's own home stays clean: a worker is an
2140 /// implementation detail of the parent session, and the parent already
2141 /// accounts for the dispatch, so a worker that wrote telemetry state into
2142 /// `$CODEWHALE_HOME` would double-count every fleet run.
2143 #[cfg(unix)]
2144 #[test]
2145 fn fleet_worker_env_carries_telemetry_off() {
2146 let fixture = TempDir::new().expect("fixture root");
2147 let workspace = fixture.path().join("workspace");
2148 let operator_home = fixture.path().join("operator-codewhale-home");
2149 std::fs::create_dir_all(&workspace).expect("workspace");
2150 std::fs::create_dir_all(&operator_home).expect("operator home");
2151 let receipt = fixture.path().join("worker-env.txt");
2152
2153 let mut adapter = LocalProcessFleetHostAdapter::new(&workspace);
2154 let mut request = FleetWorkerStartRequest::new(
2155 "telemetry-env-probe",
2156 FleetWorkerCommand::new(
2157 "/bin/sh",
2158 ["-c".to_string(), format!("env > {}", receipt.display())],
2159 ),
2160 );
2161 // A caller that both sets and allowlists the switch still cannot turn
2162 // a worker on.
2163 request
2164 .env
2165 .insert("CODEWHALE_TELEMETRY".to_string(), "true".to_string());
2166 request.env.insert(
2167 "CODEWHALE_HOME".to_string(),
2168 operator_home.display().to_string(),
2169 );
2170 request
2171 .env_allowlist
2172 .insert("CODEWHALE_TELEMETRY".to_string());
2173 request.env_allowlist.insert("CODEWHALE_HOME".to_string());
2174
2175 adapter.start_worker(request).expect("start worker");
2176 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
2177 let mut dumped = None;
2178 while std::time::Instant::now() < deadline {
2179 if let Ok(contents) = std::fs::read_to_string(&receipt)
2180 && contents.contains("CODEWHALE_TELEMETRY=")
2181 && contents.contains("DEEPSEEK_TELEMETRY=")
2182 {
2183 dumped = Some(contents);
2184 break;
2185 }
2186 std::thread::sleep(std::time::Duration::from_millis(20));
2187 }
2188 let dumped = dumped.unwrap_or_else(|| {
2189 std::fs::read_to_string(&receipt).expect("worker must dump its environment")
2190 });
2191
2192 let value = |key: &str| {
2193 dumped
2194 .lines()
2195 .find_map(|line| line.strip_prefix(&format!("{key}=")))
2196 .map(str::to_string)
2197 };
2198 assert_eq!(
2199 value("CODEWHALE_TELEMETRY").as_deref(),
2200 Some("false"),
2201 "worker environment:\n{dumped}"
2202 );
2203 assert_eq!(
2204 value("DEEPSEEK_TELEMETRY").as_deref(),
2205 Some("false"),
2206 "worker environment:\n{dumped}"
2207 );
2208 assert!(
2209 !operator_home.join("telemetry").exists(),
2210 "a fleet worker must not write telemetry state into the operator's home"
2211 );
2212 }
2213 }
2214
2214 lines RUST