Skip to main content

fhd/
lib.rs

1// Unsafe is allowed only at the FFI boundaries listed in CONTRIBUTING.md
2// (gethostname, kill(2), platform resource probes), each with a SAFETY
3// contract. Connection handling, sync, execution, and session logic must
4// stay pure safe Rust.
5#![deny(unsafe_code)]
6
7pub mod metrics;
8pub mod metrics_server;
9
10// Module layout: `lib.rs` owns the server lifecycle (accept loop, TLS
11// dispatch, per-connection run pipeline); everything else lives in a focused
12// module. See each module's docs for its responsibility.
13mod active;
14mod exec;
15mod session;
16mod stream;
17
18pub use exec::parse_custom_shell;
19pub use session::{get_hostname, is_exposed_bind, validate_start_config, ServerContext};
20
21use active::{next_run_id, ActiveBuildGuard};
22use session::{deny_control_request, DEFAULT_MAX_CONNECTIONS, UNLIMITED_CONNECTIONS};
23use stream::{execute_and_stream, execute_raw_command_and_stream};
24
25use protocol::{
26    decode_json, read_frame, read_frame_limited, write_frame, write_json_frame, FrameError,
27    HelloAckPayload, HelloPayload, ManifestPayload, MsgType, NeedPayload, ResultPayload,
28    RunPayload, CURRENT_PROTOCOL_VERSION,
29};
30use std::collections::HashMap;
31use std::path::PathBuf;
32use std::sync::Arc;
33use tokio::io::{AsyncRead, AsyncWrite};
34use tokio::net::TcpListener;
35use tokio::sync::Mutex;
36use tracing::{error, info, warn};
37
38/// Accept loop and per-connection dispatch for the agent daemon.
39///
40/// The parameter list is long but each item is a distinct daemon policy
41/// (limits, tags, storage, TLS, queueing); grouping them behind a config
42/// struct would only move the same fields one call frame away. See
43/// `ServerContext` for the subset the connection handler needs.
44#[allow(clippy::too_many_arguments)]
45pub async fn run_server(
46    listener: TcpListener,
47    expected_token: Option<String>,
48    workdir: PathBuf,
49    custom_shell: Option<String>,
50    max_concurrent_runs: Option<usize>,
51    tags: Vec<String>,
52    min_disk_bytes: Option<u64>,
53    cas_dir: Option<PathBuf>,
54    no_cas: bool,
55    tls_acceptor: Option<protocol::TlsAcceptor>,
56    max_connections: Option<usize>,
57    max_queued_runs: Option<usize>,
58    lock_manager: Option<workspace::WorkspaceLockManager>,
59    // When set, serve Prometheus metrics on this port, alongside the agent
60    // protocol on the main listener.
61    metrics_port: Option<u16>,
62) -> Result<(), Box<dyn std::error::Error>> {
63    let max_runs = max_concurrent_runs.unwrap_or_else(|| {
64        std::thread::available_parallelism()
65            .map(|n| n.get())
66            .unwrap_or(4)
67    });
68    let semaphore = Arc::new(tokio::sync::Semaphore::new(max_runs));
69    let lock_manager = lock_manager.unwrap_or_default();
70    let queue_depth = Arc::new(std::sync::atomic::AtomicUsize::new(0));
71    let connection_limiter = Arc::new(tokio::sync::Semaphore::new(match max_connections {
72        Some(0) => UNLIMITED_CONNECTIONS,
73        Some(n) => n,
74        None => DEFAULT_MAX_CONNECTIONS,
75    }));
76
77    let cas_store = if !no_cas {
78        let base = cas_dir.unwrap_or_else(|| workdir.clone());
79        Some(workspace::CasStore::new(&base))
80    } else {
81        None
82    };
83
84    let ctx = Arc::new(ServerContext {
85        expected_token,
86        workdir_root: workdir,
87        custom_shell,
88        semaphore,
89        lock_manager,
90        tags,
91        queue_depth,
92        max_runs,
93        min_disk_bytes: min_disk_bytes.unwrap_or(2_500_000_000), // Default: 2.5 GB
94        cas_store,
95        start_time: std::time::Instant::now(),
96        active_builds: Arc::new(std::sync::Mutex::new(HashMap::new())),
97        connection_limiter,
98        max_queued_runs: max_queued_runs.unwrap_or(16),
99    });
100
101    if let Some(port) = metrics_port {
102        match tokio::net::TcpListener::bind(("0.0.0.0", port)).await {
103            Ok(metrics_listener) => {
104                let metrics_ctx = Arc::clone(&ctx);
105                tokio::spawn(async move {
106                    metrics_server::serve_metrics(metrics_listener, metrics_ctx).await;
107                });
108            }
109            Err(e) => {
110                // A metrics port that cannot bind must not take the agent down
111                // with it: the agent is what people depend on.
112                error!("failed to bind metrics port {port} ({e}); metrics disabled");
113            }
114        }
115    }
116
117    loop {
118        match listener.accept().await {
119            Ok((stream, addr)) => {
120                // Bound concurrent connections: excess sockets are closed
121                // immediately so a flood cannot spawn unbounded tasks.
122                let permit = match ctx.connection_limiter.clone().try_acquire_owned() {
123                    Ok(permit) => permit,
124                    Err(_) => {
125                        warn!(
126                            "Connection limit reached; dropping connection from {}",
127                            addr
128                        );
129                        continue;
130                    }
131                };
132                info!("Accepted connection from {}", addr);
133                let ctx_clone = Arc::clone(&ctx);
134                let tls_acceptor_clone = tls_acceptor.clone();
135                tokio::spawn(async move {
136                    let _permit = permit; // released when the connection task ends
137                    let res = match tls_acceptor_clone {
138                        Some(acceptor) => match acceptor.accept(stream).await {
139                            Ok(tls_stream) => {
140                                handle_connection(
141                                    protocol::MaybeTlsStream::Server(tls_stream),
142                                    ctx_clone,
143                                    addr.to_string(),
144                                )
145                                .await
146                            }
147                            Err(e) => {
148                                error!("TLS handshake failed with {}: {}", addr, e);
149                                Err(e.into())
150                            }
151                        },
152                        None => {
153                            handle_connection(
154                                protocol::MaybeTlsStream::Plain(stream),
155                                ctx_clone,
156                                addr.to_string(),
157                            )
158                            .await
159                        }
160                    };
161                    if let Err(e) = res {
162                        error!("Connection from {} error: {}", addr, e);
163                    }
164                    info!("Connection from {} closed", addr);
165                });
166            }
167            Err(e) => {
168                warn!("Accept failed: {}", e);
169            }
170        }
171    }
172}
173
174pub async fn handle_connection<S: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
175    mut stream: S,
176    ctx: Arc<ServerContext>,
177    client_addr: String,
178) -> Result<(), Box<dyn std::error::Error>> {
179    // 1. First frame: can be STATUS probe, HISTORY query, or HELLO handshake.
180    //    Pre-authentication, so a strict size cap applies (see MAX_PRE_AUTH_PAYLOAD).
181    let (msg_type, payload) =
182        read_frame_limited(&mut stream, protocol::MAX_PRE_AUTH_PAYLOAD).await?;
183
184    if msg_type == MsgType::Status {
185        let status_req: protocol::StatusRequestPayload = decode_json(&payload)?;
186        if !ctx.authorize(Some(status_req.token.as_str())) {
187            return deny_control_request(&mut stream, "STATUS").await;
188        }
189        let active_runs = ctx
190            .max_runs
191            .saturating_sub(ctx.semaphore.available_permits());
192        let depth = ctx.queue_depth.load(std::sync::atomic::Ordering::Relaxed);
193        let (disk_free_bytes, disk_total_bytes) = match workspace::get_disk_space(&ctx.workdir_root)
194        {
195            Ok(space) => (Some(space.available_bytes), Some(space.total_bytes)),
196            Err(_) => (None, None),
197        };
198        let (memory_used_bytes, memory_total_bytes) = metrics::get_memory_info();
199        let cpu_count = Some(metrics::get_cpu_count());
200        let load_averages = metrics::get_load_averages();
201        let uptime_secs = Some(ctx.start_time.elapsed().as_secs());
202        let workspaces_count = metrics::get_workspaces_count(&ctx.workdir_root);
203
204        let active_builds: Vec<protocol::ActiveBuildInfo> = {
205            let map = ctx
206                .active_builds
207                .lock()
208                .unwrap_or_else(std::sync::PoisonError::into_inner);
209            map.iter()
210                .map(|(id, (project, argv, start_instant, client_addr))| {
211                    protocol::ActiveBuildInfo {
212                        id: id.clone(),
213                        project: project.clone(),
214                        argv: argv.clone(),
215                        elapsed_ms: start_instant.elapsed().as_millis() as u64,
216                        client_addr: client_addr.clone(),
217                    }
218                })
219                .collect()
220        };
221        // Guard is gone here — never held across an await.
222
223        let resp = protocol::StatusResponsePayload {
224            active_runs,
225            max_runs: ctx.max_runs,
226            queue_depth: depth,
227            hostname: get_hostname(),
228            tags: ctx.tags.clone(),
229            disk_free_bytes,
230            disk_total_bytes,
231            cpu_count,
232            load_averages,
233            memory_used_bytes,
234            memory_total_bytes,
235            uptime_secs,
236            active_builds: Some(active_builds),
237            workspaces_count,
238        };
239        write_json_frame(&mut stream, MsgType::StatusResp, &resp).await?;
240        return Ok(());
241    }
242
243    if msg_type == MsgType::History {
244        let history_req: protocol::HistoryRequestPayload = decode_json(&payload)?;
245        if !ctx.authorize(Some(history_req.token.as_str())) {
246            return deny_control_request(&mut stream, "HISTORY").await;
247        }
248        let workspace_dir =
249            workspace::resolve_workspace_dir(&ctx.workdir_root, &history_req.project);
250        let runs = workspace::history::get_recent_runs(&workspace_dir, history_req.limit)
251            .unwrap_or_default();
252        let resp = protocol::HistoryResponsePayload {
253            project: history_req.project,
254            runs,
255        };
256        write_json_frame(&mut stream, MsgType::HistoryResp, &resp).await?;
257        return Ok(());
258    }
259
260    if msg_type == MsgType::Clean {
261        let clean_req: protocol::CleanRequestPayload = decode_json(&payload)?;
262        if !ctx.authorize(Some(clean_req.token.as_str())) {
263            return deny_control_request(&mut stream, "CLEAN").await;
264        }
265
266        let mut bytes_freed = 0u64;
267        let msg;
268
269        if clean_req.all_branches {
270            let base_name = workspace::parse_base_project_name(&clean_req.project)
271                .unwrap_or(&clean_req.project);
272            let workspaces = workspace::scan_workspaces(&ctx.workdir_root);
273            let mut count = 0;
274            let mut skipped_busy = 0;
275            for ws in workspaces {
276                if !ws.is_canonical && ws.name.starts_with(base_name) {
277                    // Never delete a workspace whose project is mid-run.
278                    if ctx.lock_manager.is_locked(&ws.name).await {
279                        skipped_busy += 1;
280                        continue;
281                    }
282                    bytes_freed += ws.size_bytes;
283                    if let Err(e) = std::fs::remove_dir_all(&ws.path) {
284                        warn!("CLEAN: failed to remove {}: {}", ws.path.display(), e);
285                    } else {
286                        count += 1;
287                    }
288                }
289            }
290            msg = format!(
291                "Purged {} branch workspaces for project '{}'{}",
292                count,
293                base_name,
294                if skipped_busy > 0 {
295                    format!(" ({} skipped: active run)", skipped_busy)
296                } else {
297                    String::new()
298                }
299            );
300        } else {
301            let ws_dir = workspace::resolve_workspace_dir(&ctx.workdir_root, &clean_req.project);
302            if ws_dir.is_dir() {
303                if ctx.lock_manager.is_locked(&clean_req.project).await {
304                    msg = format!(
305                        "Workspace '{}' is busy (active run); try again after it finishes",
306                        clean_req.project
307                    );
308                    let resp = protocol::CleanResponsePayload {
309                        ok: false,
310                        message: msg,
311                        bytes_freed: 0,
312                    };
313                    write_json_frame(&mut stream, MsgType::CleanResp, &resp).await?;
314                    return Ok(());
315                }
316                if clean_req.caches_only {
317                    bytes_freed = workspace::trim_workspace_caches(&ws_dir);
318                    msg = format!(
319                        "Trimmed volatile caches for workspace '{}'",
320                        clean_req.project
321                    );
322                } else {
323                    bytes_freed = workspace::calculate_dir_size(&ws_dir);
324                    if let Err(e) = std::fs::remove_dir_all(&ws_dir) {
325                        warn!("CLEAN: failed to remove {}: {}", ws_dir.display(), e);
326                        let resp = protocol::CleanResponsePayload {
327                            ok: false,
328                            message: format!(
329                                "Failed to remove workspace '{}': {}",
330                                clean_req.project, e
331                            ),
332                            bytes_freed: 0,
333                        };
334                        write_json_frame(&mut stream, MsgType::CleanResp, &resp).await?;
335                        return Ok(());
336                    }
337                    msg = format!("Removed workspace '{}'", clean_req.project);
338                }
339            } else {
340                msg = format!("Workspace '{}' does not exist remotely", clean_req.project);
341            }
342        }
343
344        let resp = protocol::CleanResponsePayload {
345            ok: true,
346            message: msg,
347            bytes_freed,
348        };
349        write_json_frame(&mut stream, MsgType::CleanResp, &resp).await?;
350        return Ok(());
351    }
352
353    if msg_type != MsgType::Hello {
354        let ack = HelloAckPayload {
355            ok: false,
356            error: Some("Expected HELLO, STATUS, HISTORY, or CLEAN frame".into()),
357            compression: None,
358            remote_workdir: None,
359        };
360        write_json_frame(&mut stream, MsgType::HelloAck, &ack).await?;
361        return Err("Protocol error: expected HELLO, STATUS, HISTORY, or CLEAN".into());
362    }
363
364    let hello: HelloPayload = decode_json(&payload)?;
365    if hello.protocol_version != CURRENT_PROTOCOL_VERSION {
366        let ack = HelloAckPayload {
367            ok: false,
368            error: Some(format!(
369                "Protocol version mismatch: expected {}, got {}",
370                CURRENT_PROTOCOL_VERSION, hello.protocol_version
371            )),
372            compression: None,
373            remote_workdir: None,
374        };
375        write_json_frame(&mut stream, MsgType::HelloAck, &ack).await?;
376        return Err("Protocol version mismatch".into());
377    }
378
379    if !ctx.authorize(Some(hello.token.as_str())) {
380        let ack = HelloAckPayload {
381            ok: false,
382            error: Some("Unauthorized: invalid auth token".into()),
383            compression: None,
384            remote_workdir: None,
385        };
386        write_json_frame(&mut stream, MsgType::HelloAck, &ack).await?;
387        return Err("Unauthorized".into());
388    }
389
390    // Pre-flight disk space guard: verify host volume has sufficient free space
391    if ctx.min_disk_bytes > 0 {
392        if let Ok(space) = workspace::get_disk_space(&ctx.workdir_root) {
393            if space.available_bytes < ctx.min_disk_bytes {
394                let needed = ctx.min_disk_bytes.saturating_sub(space.available_bytes);
395                info!(
396                    "Available disk space ({} bytes) is below minimum threshold ({} bytes). Running emergency GC...",
397                    space.available_bytes, ctx.min_disk_bytes
398                );
399                // Snapshot locked projects so emergency GC never deletes
400                // workspaces with active runs.
401                let locked: std::collections::HashSet<String> = ctx
402                    .lock_manager
403                    .locked_projects()
404                    .await
405                    .into_iter()
406                    .collect();
407                let root = ctx.workdir_root.clone();
408                let gc_report = tokio::task::spawn_blocking(move || {
409                    workspace::run_emergency_disk_gc(&root, needed, &|name: &str| {
410                        locked.contains(name)
411                    })
412                })
413                .await
414                .unwrap_or_default();
415                info!(
416                    "Emergency GC pruned {} workspaces, trimmed {} bytes.",
417                    gc_report.workspaces_deleted, gc_report.caches_trimmed_bytes
418                );
419
420                if let Ok(new_space) = workspace::get_disk_space(&ctx.workdir_root) {
421                    if new_space.available_bytes < ctx.min_disk_bytes {
422                        let free_gb = new_space.available_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
423                        let req_gb = ctx.min_disk_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
424                        let err_msg = format!(
425                            "Remote agent disk low ({:.2} GB free, required >= {:.2} GB). Run 'fh clean' or free host disk.",
426                            free_gb, req_gb
427                        );
428                        warn!("{}", err_msg);
429                        let ack = HelloAckPayload {
430                            ok: false,
431                            error: Some(err_msg.clone()),
432                            compression: None,
433                            remote_workdir: None,
434                        };
435                        write_json_frame(&mut stream, MsgType::HelloAck, &ack).await?;
436                        return Err(err_msg.into());
437                    }
438                }
439            }
440        }
441    }
442
443    // Negotiate compression algorithm from client's offered list
444    let negotiated_compression = if let Some(client_algos) = &hello.compressions {
445        if client_algos.iter().any(|a| a.eq_ignore_ascii_case("zstd")) {
446            "zstd".to_string()
447        } else if client_algos
448            .iter()
449            .any(|a| a.eq_ignore_ascii_case("gzip") || a.eq_ignore_ascii_case("gz"))
450        {
451            "gzip".to_string()
452        } else if client_algos.iter().any(|a| a.eq_ignore_ascii_case("none")) {
453            "none".to_string()
454        } else {
455            "gzip".to_string()
456        }
457    } else {
458        "gzip".to_string()
459    };
460
461    // Resolve persistent workspace directory (forks from seed via APFS CoW if branch)
462    let workspace_dir = match workspace::ensure_workspace_dir(&ctx.workdir_root, &hello.project) {
463        Ok(dir) => dir,
464        Err(e) => {
465            let ack = HelloAckPayload {
466                ok: false,
467                error: Some(format!("Failed to prepare workspace: {}", e)),
468                compression: None,
469                remote_workdir: None,
470            };
471            write_json_frame(&mut stream, MsgType::HelloAck, &ack).await?;
472            return Err(e.into());
473        }
474    };
475    info!("Using persistent workspace: {}", workspace_dir.display());
476
477    // Acknowledge handshake
478    let ack = HelloAckPayload {
479        ok: true,
480        error: None,
481        compression: Some(negotiated_compression.clone()),
482        remote_workdir: Some(workspace_dir.display().to_string()),
483    };
484    write_json_frame(&mut stream, MsgType::HelloAck, &ack).await?;
485    info!(
486        "Handshake successful for project '{}' (negotiated compression: '{}')",
487        hello.project, negotiated_compression
488    );
489
490    // Split the stream up front: writes go through the shared writer from
491    // here on, and a watchdog can hold the read half while we are queued to
492    // detect client disconnects (a dead client must free its queue slot).
493    let (read_half, write_half) = tokio::io::split(stream);
494    let shared_writer = Arc::new(Mutex::new(write_half));
495    let mut read_half = read_half;
496    let mut queued_frames: Vec<(MsgType, Vec<u8>)> = Vec::new();
497
498    // 2. Receive next frame: PUT_TEMPLATE or MANIFEST
499    let (msg_type, payload) = read_frame(&mut read_half).await?;
500    let manifest: ManifestPayload = if msg_type == MsgType::PutTemplate {
501        let put_req: protocol::PutTemplatePayload = decode_json(&payload)?;
502        info!(
503            "Received PUT_TEMPLATE for '{}' (scope: {})",
504            put_req.name, put_req.scope
505        );
506        match templates::save_template(
507            Some(&workspace_dir),
508            &put_req.name,
509            &put_req.yaml,
510            &put_req.scope,
511        ) {
512            Ok(saved_path) => {
513                info!("Saved template to {}", saved_path.display());
514                let ack = HelloAckPayload {
515                    ok: true,
516                    error: None,
517                    compression: None,
518                    remote_workdir: None,
519                };
520                {
521                    let mut w = shared_writer.lock().await;
522                    write_json_frame(&mut *w, MsgType::HelloAck, &ack).await?;
523                }
524                return Ok(());
525            }
526            Err(e) => {
527                warn!("Failed to save template: {}", e);
528                let ack = HelloAckPayload {
529                    ok: false,
530                    error: Some(e.to_string()),
531                    compression: None,
532                    remote_workdir: None,
533                };
534                {
535                    let mut w = shared_writer.lock().await;
536                    write_json_frame(&mut *w, MsgType::HelloAck, &ack).await?;
537                }
538                return Err(format!("Failed to save template: {}", e).into());
539            }
540        }
541    } else if msg_type == MsgType::Manifest {
542        decode_json(&payload)?
543    } else {
544        return Err(format!(
545            "Expected MANIFEST or PUT_TEMPLATE frame, got {:?}",
546            msg_type
547        )
548        .into());
549    };
550
551    /// Wait for `acquire` while watching the client connection for disconnects.
552    ///
553    /// While queued, the daemon normally does not read the socket — a client that
554    /// dies while waiting would hold its queue slot forever. The watchdog owns
555    /// the read half during the wait, buffers any frames that (against today's
556    /// protocol) arrive early, and detects EOF. On acquisition the read half and
557    /// any buffered frames are handed back; `Ok(None)` means the client
558    /// disconnected and the caller must abort.
559    #[allow(clippy::type_complexity)]
560    async fn wait_with_disconnect_watch<S, T, Fut>(
561        read_half: tokio::io::ReadHalf<S>,
562        queued: protocol::QueuedPayload,
563        writer: &Arc<Mutex<tokio::io::WriteHalf<S>>>,
564        acquire: Fut,
565    ) -> Result<Option<(tokio::io::ReadHalf<S>, T, Vec<(MsgType, Vec<u8>)>)>, String>
566    where
567        S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
568        Fut: std::future::Future<Output = T>,
569    {
570        let (acquired_tx, acquired_rx) = tokio::sync::oneshot::channel::<()>();
571        let watchdog: tokio::task::JoinHandle<
572            Result<
573                (tokio::io::ReadHalf<S>, Vec<(MsgType, Vec<u8>)>, bool),
574                std::convert::Infallible,
575            >,
576        > = tokio::spawn(async move {
577            let mut read_half = read_half;
578            let mut buffered: Vec<(MsgType, Vec<u8>)> = Vec::new();
579            tokio::pin!(acquired_rx);
580            loop {
581                tokio::select! {
582                    _ = &mut acquired_rx => {
583                        break Ok((read_half, buffered, false));
584                    }
585                    res = read_frame(&mut read_half) => match res {
586                        Ok(frame) => {
587                            // Early frames are a protocol violation today; buffer
588                            // a bounded amount and keep consuming so EOF
589                            // detection keeps working.
590                            if buffered.len() < 16 {
591                                buffered.push(frame);
592                            }
593                            continue;
594                        }
595                        Err(_) => {
596                            break Ok((read_half, buffered, true));
597                        }
598                    }
599                }
600            }
601        });
602
603        {
604            let mut w = writer.lock().await;
605            write_json_frame(&mut *w, MsgType::Queued, &queued)
606                .await
607                .map_err(|e| e.to_string())?;
608        }
609
610        let acquired = acquire.await;
611        let _ = acquired_tx.send(());
612        let (read_half, buffered, disconnected) = match watchdog.await {
613            Ok(Ok(t)) => t,
614            Ok(Err(infallible)) => match infallible {},
615            Err(join_err) => return Err(format!("queue watchdog panicked: {}", join_err)),
616        };
617
618        if disconnected {
619            return Ok(None);
620        }
621        Ok(Some((read_half, acquired, buffered)))
622    }
623
624    // Acquire per-project lock to ensure serialized execution on the same project workspace
625    let project_mutex = ctx.lock_manager.get_lock(&hello.project).await;
626    let _project_guard = match project_mutex.clone().try_lock_owned() {
627        Ok(guard) => guard,
628        Err(_) => {
629            info!(
630                "Project '{}' is busy. Sending QUEUED frame...",
631                hello.project
632            );
633            // Admission control: a full queue rejects immediately instead of
634            // letting disconnected/queued connections grow memory forever.
635            if ctx.queue_depth.load(std::sync::atomic::Ordering::SeqCst) >= ctx.max_queued_runs {
636                let ack = HelloAckPayload {
637                    ok: false,
638                    error: Some(format!(
639                        "Agent queue is full ({} queued runs). Try again later.",
640                        ctx.max_queued_runs
641                    )),
642                    compression: None,
643                    remote_workdir: None,
644                };
645                let mut w = shared_writer.lock().await;
646                write_json_frame(&mut *w, MsgType::HelloAck, &ack).await?;
647                return Err("queue full".into());
648            }
649            ctx.queue_depth
650                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
651            let pos = ctx.queue_depth.load(std::sync::atomic::Ordering::SeqCst);
652            let queued = protocol::QueuedPayload {
653                position: pos,
654                reason: "project_busy".to_string(),
655            };
656            let outcome = wait_with_disconnect_watch(
657                read_half,
658                queued,
659                &shared_writer,
660                project_mutex.lock_owned(),
661            )
662            .await?;
663            let (rh, guard, frames) = match outcome {
664                Some(t) => t,
665                None => {
666                    ctx.queue_depth
667                        .fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
668                    info!("Client disconnected while queued (project busy); freeing slot");
669                    return Ok(());
670                }
671            };
672            read_half = rh;
673            queued_frames.extend(frames);
674            ctx.queue_depth
675                .fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
676            guard
677        }
678    };
679
680    // Acquire global concurrency permit
681    let _permit = match ctx.semaphore.clone().try_acquire_owned() {
682        Ok(permit) => permit,
683        Err(_) => {
684            info!("Agent concurrency limit reached. Sending QUEUED frame...");
685            if ctx.queue_depth.load(std::sync::atomic::Ordering::SeqCst) >= ctx.max_queued_runs {
686                let ack = HelloAckPayload {
687                    ok: false,
688                    error: Some(format!(
689                        "Agent queue is full ({} queued runs). Try again later.",
690                        ctx.max_queued_runs
691                    )),
692                    compression: None,
693                    remote_workdir: None,
694                };
695                let mut w = shared_writer.lock().await;
696                write_json_frame(&mut *w, MsgType::HelloAck, &ack).await?;
697                return Err("queue full".into());
698            }
699            ctx.queue_depth
700                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
701            let pos = ctx.queue_depth.load(std::sync::atomic::Ordering::SeqCst);
702            let queued = protocol::QueuedPayload {
703                position: pos,
704                reason: "concurrency_limit".to_string(),
705            };
706            let outcome = wait_with_disconnect_watch(
707                read_half,
708                queued,
709                &shared_writer,
710                ctx.semaphore.clone().acquire_owned(),
711            )
712            .await?;
713            let (rh, permit_res, frames) = match outcome {
714                Some(t) => t,
715                None => {
716                    ctx.queue_depth
717                        .fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
718                    info!("Client disconnected while queued (concurrency limit); freeing slot");
719                    return Ok(());
720                }
721            };
722            read_half = rh;
723            queued_frames.extend(frames);
724            ctx.queue_depth
725                .fetch_sub(1, std::sync::atomic::Ordering::SeqCst);
726            permit_res.map_err(|e| e.to_string())?
727        }
728    };
729
730    info!(
731        "Received client manifest with {} files. Diffing against workspace cache...",
732        manifest.files.len()
733    );
734
735    let extra_ignores = templates::resolve_template_extra_ignores(&workspace_dir, None);
736    // Full-workspace scan + hashing is blocking I/O — keep it off the async
737    // runtime (workdir/manifest/ignores moved in owned form).
738    let diff_dir = workspace_dir.clone();
739    let diff_manifest = manifest.clone();
740    let diff_ignores = extra_ignores.clone();
741    let mut diff = tokio::task::spawn_blocking(move || {
742        workspace::diff_manifests(&diff_dir, &diff_manifest, &diff_ignores)
743    })
744    .await
745    .map_err(|e| -> Box<dyn std::error::Error> { Box::new(std::io::Error::other(e)) })??;
746
747    if let Some(cas) = &ctx.cas_store {
748        let manifest_map: std::collections::HashMap<&str, &str> = manifest
749            .files
750            .iter()
751            .map(|f| (f.path.as_str(), f.hash.as_str()))
752            .collect();
753
754        let mut remaining_want = Vec::new();
755        let mut hydrated_count = 0;
756
757        for rel_path in diff.want {
758            if let Some(hash) = manifest_map.get(rel_path.as_str()) {
759                if let Ok(rel_buf) = protocol::from_wire_path(&rel_path) {
760                    let target_path = workspace_dir.join(rel_buf);
761                    if let Ok(true) = cas.materialize_to(hash, &target_path) {
762                        hydrated_count += 1;
763                        continue;
764                    }
765                }
766            }
767            remaining_want.push(rel_path);
768        }
769
770        if hydrated_count > 0 {
771            info!(
772                "Hydrated {} file(s) from global CAS without network transfer",
773                hydrated_count
774            );
775        }
776        diff.want = remaining_want;
777    }
778
779    info!(
780        "Diff computed: {} files needed, {} extraneous files flagged for deletion",
781        diff.want.len(),
782        diff.delete_extraneous.len()
783    );
784
785    // 3. Send NEED frame
786    let need = NeedPayload {
787        want: diff.want.clone(),
788        delete_extraneous: diff.delete_extraneous.clone(),
789    };
790    {
791        let mut w = shared_writer.lock().await;
792        write_json_frame(&mut *w, MsgType::Need, &need).await?;
793    }
794
795    // 4. Receive FILES frame (delta archive). If a queued watchdog buffered
796    // early frames, consume them first (defensive; the client sends nothing
797    // while queued in today's protocol).
798    let (msg_type, payload) = loop {
799        if !queued_frames.is_empty() {
800            let frame = queued_frames.remove(0);
801            if frame.0 == MsgType::Files {
802                break frame;
803            }
804            warn!(
805                "Discarding unexpected buffered frame {:?} from queue window",
806                frame.0
807            );
808            continue;
809        }
810        match read_frame(&mut read_half).await {
811            Ok(frame) => break frame,
812            Err(e) => {
813                // `fh sync --dry-run` reads NEED to learn what would move and
814                // then hangs up without sending FILES. Nothing was transferred
815                // and the workspace is untouched, which is exactly what the
816                // client asked for — not a protocol error.
817                if matches!(e, FrameError::Io(_) | FrameError::UnexpectedEof) {
818                    info!("Client hung up after NEED (dry run); no files were transferred");
819                    return Ok(());
820                }
821                return Err(e.into());
822            }
823        }
824    };
825    if msg_type != MsgType::Files {
826        return Err(format!("Expected FILES frame, got {:?}", msg_type).into());
827    }
828
829    let bytes_synced = payload.len() as u64;
830    if !payload.is_empty() {
831        info!("Unpacking {} delta bytes into workspace", payload.len());
832        let unpack_dir = workspace_dir.clone();
833        let unpack_payload = payload;
834        tokio::task::spawn_blocking(move || fileset::unpack_tar(&unpack_dir, &unpack_payload))
835            .await
836            .map_err(|e| -> Box<dyn std::error::Error> { Box::new(std::io::Error::other(e)) })??;
837
838        if let Some(cas) = &ctx.cas_store {
839            for entry in &manifest.files {
840                if let Ok(rel_buf) = protocol::from_wire_path(&entry.path) {
841                    let local_file = workspace_dir.join(rel_buf);
842                    if local_file.is_file() {
843                        let _ = cas.put_file(&entry.hash, &local_file);
844                    }
845                }
846            }
847        }
848    } else {
849        info!("Zero delta bytes uploaded (workspace up to date)");
850    }
851
852    // Apply deletions of extraneous files
853    if !diff.delete_extraneous.is_empty() {
854        let deleted = workspace::apply_deletions(&workspace_dir, &diff.delete_extraneous)?;
855        info!("Deleted {} extraneous files from workspace", deleted);
856    }
857
858    // 5. Receive RUN frame
859    let (msg_type, payload) = loop {
860        if !queued_frames.is_empty() {
861            let frame = queued_frames.remove(0);
862            if frame.0 == MsgType::Run {
863                break frame;
864            }
865            warn!(
866                "Discarding unexpected buffered frame {:?} before RUN",
867                frame.0
868            );
869            continue;
870        }
871        match read_frame(&mut read_half).await {
872            Ok(frame) => break frame,
873            Err(e) => {
874                // A client that sent its files and then hung up has finished a
875                // sync-only session (`fh sync`): the delta is already applied
876                // and there is no command to run. That is a successful outcome,
877                // not a protocol error — the workspace really is up to date.
878                if matches!(e, FrameError::Io(_) | FrameError::UnexpectedEof) {
879                    info!("Client completed a sync-only session; workspace is up to date");
880                    return Ok(());
881                }
882                return Err(e.into());
883            }
884        }
885    };
886    if msg_type != MsgType::Run {
887        return Err(format!("Expected RUN frame, got {:?}", msg_type).into());
888    }
889    let run: RunPayload = decode_json(&payload)?;
890    info!("Executing command: {:?}", run.argv);
891
892    let run_start = std::time::Instant::now();
893    let run_id = next_run_id();
894
895    ctx.active_builds
896        .lock()
897        .unwrap_or_else(std::sync::PoisonError::into_inner)
898        .insert(
899            run_id.clone(),
900            (
901                hello.project.clone(),
902                run.argv.clone(),
903                std::time::Instant::now(),
904                client_addr.clone(),
905            ),
906        );
907    let _active_build_guard = ActiveBuildGuard {
908        active_builds: ctx.active_builds.clone(),
909        run_id: run_id.clone(),
910    };
911
912    // 6. Pre-build dependency caching hook
913    // (the stream was already split after the handshake; read/write halves
914    // and the shared writer are in scope)
915
916    let matched_templates = templates::match_templates(&workspace_dir, run.template.as_deref());
917    let hook_template = matched_templates
918        .into_iter()
919        .find(|t| t.hints.install_command.is_some());
920
921    if let Some(template) = hook_template {
922        let install_cmd = template.hints.install_command.as_ref().unwrap();
923        let current_lock_hash =
924            workspace::state::compute_lockfiles_hash(&workspace_dir, &template.hints.lockfiles);
925        let prev_state = workspace::state::read_state(&workspace_dir);
926
927        let need_install = if let Some(ref current_hash) = current_lock_hash {
928            run.no_cache
929                || match &prev_state {
930                    Some(s) => s.last_success_lockfile_hash != *current_hash,
931                    None => true,
932                }
933        } else if template.hints.lockfiles.is_empty() {
934            run.no_cache || prev_state.is_none()
935        } else {
936            // Lockfiles were declared in the template, but none exist in the workspace
937            false
938        };
939
940        if need_install {
941            let start_banner = format!(
942                "=== [farhand] Running dependency hook: {} ===\n",
943                install_cmd
944            );
945            {
946                let mut writer = shared_writer.lock().await;
947                let log = protocol::LogPayload {
948                    stream: "stdout".into(),
949                    data: start_banner,
950                };
951                let _ = write_json_frame(&mut *writer, MsgType::Log, &log).await;
952            }
953
954            let hook_exit = execute_raw_command_and_stream(
955                shared_writer.clone(),
956                &mut read_half,
957                &workspace_dir,
958                install_cmd,
959                ctx.custom_shell.as_deref(),
960                run.env.as_ref(),
961                run.toolchain.as_ref(),
962            )
963            .await?;
964
965            if hook_exit != 0 {
966                info!(
967                    "Dependency install hook failed with exit code {}",
968                    hook_exit
969                );
970                let result = ResultPayload {
971                    exit_code: hook_exit,
972                    error: Some(format!(
973                        "dependency hook '{}' failed with exit code {}",
974                        install_cmd, hook_exit
975                    )),
976                };
977                let record = protocol::RunRecord {
978                    id: run_id,
979                    timestamp_rfc3339: workspace::format_rfc3339(std::time::SystemTime::now()),
980                    project: hello.project.clone(),
981                    argv: run.argv.clone(),
982                    exit_code: hook_exit,
983                    duration_ms: run_start.elapsed().as_millis() as u64,
984                    bytes_synced,
985                    artifact_size: 0,
986                    client_addr: client_addr.clone(),
987                    error: result.error.clone(),
988                };
989                if let Err(e) = tokio::task::spawn_blocking({
990                    let dir = workspace_dir.clone();
991                    move || workspace::history::save_run(&dir, &record)
992                })
993                .await
994                .map_err(|e| -> Box<dyn std::error::Error> { Box::new(std::io::Error::other(e)) })?
995                {
996                    warn!("Failed to save run record: {}", e);
997                }
998
999                let mut writer = shared_writer.lock().await;
1000                write_json_frame(&mut *writer, MsgType::Result, &result).await?;
1001                drop(read_half);
1002                return Ok(());
1003            }
1004
1005            // Install succeeded: record state
1006            let new_state = workspace::WorkspaceState {
1007                version: 1,
1008                last_success_lockfile_hash: current_lock_hash.unwrap_or_default(),
1009                last_installed_at: std::time::SystemTime::now(),
1010                template: template.name.clone(),
1011            };
1012            if let Err(e) = workspace::state::write_state(&workspace_dir, &new_state) {
1013                warn!("Failed to write workspace state: {}", e);
1014            }
1015
1016            let end_banner =
1017                "=== [farhand] Dependencies up to date. Proceeding to user command ===\n"
1018                    .to_string();
1019            {
1020                let mut writer = shared_writer.lock().await;
1021                let log = protocol::LogPayload {
1022                    stream: "stdout".into(),
1023                    data: end_banner,
1024                };
1025                let _ = write_json_frame(&mut *writer, MsgType::Log, &log).await;
1026            }
1027        } else {
1028            info!(
1029                "Lockfiles unchanged or not present ({:?}). Skipping dependency install hook '{}'.",
1030                current_lock_hash, install_cmd
1031            );
1032        }
1033    }
1034
1035    // 7. Execute user command and stream output
1036    let exit_code = execute_and_stream(
1037        shared_writer.clone(),
1038        &mut read_half,
1039        &workspace_dir,
1040        &run.argv,
1041        ctx.custom_shell.as_deref(),
1042        run.env.as_ref(),
1043        run.toolchain.as_ref(),
1044        run.tty,
1045        run.cols,
1046        run.rows,
1047        run.raw_stdio.unwrap_or(false),
1048    )
1049    .await?;
1050
1051    info!("Command exited with status code {}", exit_code);
1052
1053    // 8. If command succeeded, resolve artifacts before reporting result & saving history
1054    let mut artifact_size = 0u64;
1055    let mut artifact_payload = None;
1056    if exit_code == 0 {
1057        let artifact_paths = workspace::resolve_artifact_paths(
1058            &workspace_dir,
1059            run.outputs.as_deref(),
1060            run.template.as_deref(),
1061        );
1062        if !artifact_paths.is_empty() {
1063            info!(
1064                "Packing {} artifact paths using compression '{}'",
1065                artifact_paths.len(),
1066                negotiated_compression
1067            );
1068            let algo = fileset::CompressionAlgo::from_str_opt(Some(&negotiated_compression));
1069            let tar_bytes = fileset::pack_tar_with_algo(&workspace_dir, &artifact_paths, algo)?;
1070            artifact_size = tar_bytes.len() as u64;
1071            artifact_payload = Some(tar_bytes);
1072        }
1073    }
1074
1075    let record = protocol::RunRecord {
1076        id: run_id,
1077        timestamp_rfc3339: workspace::format_rfc3339(std::time::SystemTime::now()),
1078        project: hello.project.clone(),
1079        argv: run.argv.clone(),
1080        exit_code,
1081        duration_ms: run_start.elapsed().as_millis() as u64,
1082        bytes_synced,
1083        artifact_size,
1084        client_addr,
1085        error: if exit_code != 0 {
1086            Some(format!("command exited with status code {}", exit_code))
1087        } else {
1088            None
1089        },
1090    };
1091    if let Err(e) = tokio::task::spawn_blocking({
1092        let dir = workspace_dir.clone();
1093        move || workspace::history::save_run(&dir, &record)
1094    })
1095    .await
1096    .map_err(|e| -> Box<dyn std::error::Error> { Box::new(std::io::Error::other(e)) })?
1097    {
1098        warn!("Failed to save run record: {}", e);
1099    }
1100
1101    // 9. Send RESULT and ARTIFACTS frames
1102    let result = ResultPayload {
1103        exit_code,
1104        error: None,
1105    };
1106    {
1107        let mut writer = shared_writer.lock().await;
1108        write_json_frame(&mut *writer, MsgType::Result, &result).await?;
1109        if let Some(tar_gz) = artifact_payload {
1110            write_frame(&mut *writer, MsgType::Artifacts, &tar_gz).await?;
1111            info!("Sent ARTIFACTS frame ({} bytes)", tar_gz.len());
1112        }
1113    }
1114
1115    drop(read_half);
1116    Ok(())
1117}
1118
1119#[cfg(test)]
1120mod tests {
1121    use super::*;
1122
1123    fn ctx_with(token: Option<&str>) -> ServerContext {
1124        ServerContext {
1125            expected_token: token.map(str::to_string),
1126            workdir_root: std::env::temp_dir(),
1127            custom_shell: None,
1128            semaphore: Arc::new(tokio::sync::Semaphore::new(1)),
1129            lock_manager: workspace::WorkspaceLockManager::new(),
1130            tags: vec![],
1131            queue_depth: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
1132            max_runs: 1,
1133            min_disk_bytes: 0,
1134            cas_store: None,
1135            start_time: std::time::Instant::now(),
1136            active_builds: Arc::new(std::sync::Mutex::new(HashMap::new())),
1137            connection_limiter: Arc::new(tokio::sync::Semaphore::new(UNLIMITED_CONNECTIONS)),
1138            max_queued_runs: 16,
1139        }
1140    }
1141
1142    #[test]
1143    fn authorize_accepts_matching_token() {
1144        let ctx = ctx_with(Some("s3cret"));
1145        assert!(ctx.authorize(Some("s3cret")));
1146    }
1147
1148    #[test]
1149    fn authorize_rejects_wrong_and_missing_tokens() {
1150        let ctx = ctx_with(Some("s3cret"));
1151        assert!(!ctx.authorize(Some("wrong")));
1152        assert!(!ctx.authorize(Some("")));
1153        assert!(!ctx.authorize(None));
1154        // Every single-byte mutation must be rejected.
1155        for i in 0.."s3cret".len() {
1156            let mut tampered = "s3cret".to_string();
1157            let replacement = if i % 2 == 0 { "x" } else { "y" };
1158            tampered.replace_range(i..i + 1, replacement);
1159            assert!(
1160                !ctx.authorize(Some(&tampered)),
1161                "mutation at byte {i} leaked"
1162            );
1163        }
1164    }
1165
1166    #[test]
1167    fn authorize_unauthenticated_mode_allows_all_callers() {
1168        let ctx = ctx_with(None);
1169        assert!(ctx.authorize(None));
1170        assert!(ctx.authorize(Some("anything")));
1171    }
1172
1173    #[test]
1174    fn authorize_treats_empty_configured_token_as_legacy_open() {
1175        // Direct run_server callers only; the CLI rejects empty tokens.
1176        let ctx = ctx_with(Some(""));
1177        assert!(ctx.authorize(Some("any")));
1178        assert!(ctx.authorize(None));
1179    }
1180
1181    #[test]
1182    fn validate_start_config_requires_token_or_opt_in() {
1183        assert!(validate_start_config(Some("tok"), false).is_ok());
1184        assert!(validate_start_config(Some("tok"), true).is_ok());
1185
1186        let no_token = validate_start_config(None, false).unwrap_err();
1187        assert!(no_token.contains("--token"));
1188        assert!(no_token.contains("--allow-unauthenticated"));
1189
1190        assert!(validate_start_config(None, true).is_ok());
1191
1192        let empty = validate_start_config(Some("   "), false).unwrap_err();
1193        assert!(empty.contains("empty string"));
1194    }
1195
1196    #[test]
1197    fn is_exposed_bind_classifies_addresses() {
1198        assert!(is_exposed_bind("0.0.0.0:9876"));
1199        assert!(is_exposed_bind("[::]:9876"));
1200        assert!(is_exposed_bind("10.1.2.3:9876"));
1201        assert!(is_exposed_bind("192.168.1.10:9876"));
1202        assert!(!is_exposed_bind("127.0.0.1:9876"));
1203        assert!(!is_exposed_bind("[::1]:9876"));
1204        // Unparseable input is treated conservatively as exposed.
1205        assert!(is_exposed_bind("not-an-address"));
1206    }
1207
1208    #[test]
1209    fn parse_custom_shell_rejects_empty_invocations() {
1210        // Regression: `--shell " "` panicked the connection task on parts[0].
1211        assert_eq!(parse_custom_shell("  "), None);
1212        assert_eq!(parse_custom_shell(""), None);
1213        assert_eq!(
1214            parse_custom_shell("/bin/sh -c"),
1215            Some(vec!["/bin/sh".to_string(), "-c".to_string()])
1216        );
1217    }
1218
1219    #[test]
1220    fn next_run_id_is_collision_free_under_stress() {
1221        // Regression: `millis ^ pid` collided for two runs in the same
1222        // millisecond. Nanos + counter must never repeat within a process.
1223        let mut ids = std::collections::HashSet::with_capacity(10_000);
1224        for _ in 0..10_000 {
1225            let id = next_run_id();
1226            assert!(ids.insert(id), "run_id collision generated");
1227        }
1228        // Distinct lengths/format sanity (16 hex nanos + 4 hex counter).
1229        assert!(ids.iter().all(|id| id.len() == 20));
1230    }
1231
1232    #[tokio::test]
1233    async fn active_build_guard_drop_removes_entry_synchronously() {
1234        let ctx = ctx_with(Some("tok"));
1235        let map = ctx.active_builds.clone();
1236        map.lock().unwrap().insert(
1237            "run-1".to_string(),
1238            (
1239                "proj".to_string(),
1240                Vec::new(),
1241                std::time::Instant::now(),
1242                "127.0.0.1".to_string(),
1243            ),
1244        );
1245
1246        {
1247            let _guard = ActiveBuildGuard {
1248                active_builds: map.clone(),
1249                run_id: "run-1".to_string(),
1250            };
1251            assert!(map.lock().unwrap().contains_key("run-1"));
1252        }
1253        // Synchronous removal: no detached task, no window where STATUS sees
1254        // a finished run as active.
1255        assert!(!map.lock().unwrap().contains_key("run-1"));
1256    }
1257}