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