Skip to main content

mj_controller/
worker_client.rs

1//! Controller-side client for a session relay's JSON-lines proxy.
2
3use std::collections::{BTreeMap, BTreeSet, VecDeque};
4use std::path::Path;
5use std::process::Stdio;
6use std::sync::{Arc, PoisonError};
7use std::time::{Duration, Instant};
8
9use anyhow::{Context, Result, anyhow, bail};
10use base64::Engine as _;
11use base64::engine::general_purpose::STANDARD as BASE64;
12use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt, BufReader};
13use tokio::process::{Child, ChildStdin, ChildStdout, Command};
14use tokio::sync::{mpsc, watch};
15
16use crate::targets::{
17    CommandSpec, SSH_RETRY_ATTEMPTS, SshAdmission, SshPermit, is_transport_rejection,
18};
19use mj_core::config::harness_authentication_marker;
20use mj_core::credentials::{
21    CredentialSnapshot, CredentialSyncAction, CredentialSyncHandle, CredentialSyncOutcome,
22    CredentialSyncResult, CredentialSyncTarget, SYNC_INTERVAL, SyncAction, SyncTrigger, enqueue,
23    profiles_with_targets, read_credential_file, reconcile, validate_credential_payload,
24    write_credential_file,
25};
26use mj_core::elicitation::ElicitationResponse;
27use mj_core::relay::{
28    MAX_FRAME_BYTES, RELAY_EVENT_GENESIS_DIGEST, RELAY_MIN_PROTOCOL_VERSION,
29    RELAY_PROTOCOL_VERSION, RelayCommand, RelayCursor, RelayErrorCode, RelayEvent,
30    RelayOperationalState, RelayProtocolError, RelayRequest, RelayRequestEnvelope,
31    RelayResponseBody, RelayResponseEnvelope, RelayResponsePayload, RelayVersionRange,
32    ReviewerRequest, validate_relay_event,
33};
34
35pub use mj_client::session::{RelayAttachment, StartedReviewer};
36use mj_core::worker_launch::ReviewerLaunchConfig;
37
38const RELAY_RPC_TIMEOUT: Duration = Duration::from_secs(15);
39const RELAY_SLOW_OPERATION_WARNING: Duration = Duration::from_secs(5);
40/// Starting a target-side proxy may page the full worker executable in and
41/// traverse a container runtime before the relay sees `hello`. That is worker
42/// startup latency, not an ordinary in-connection RPC.
43const RELAY_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(300);
44/// An attachment can decompress a transport-sized page from cold journal
45/// segments. It remains bounded by the relay frame budget, but cold or loaded
46/// storage needs a filesystem deadline rather than an in-memory RPC deadline.
47const RELAY_HISTORY_TIMEOUT: Duration = Duration::from_secs(900);
48/// Advancing an acknowledgement can durably prune a large relay journal. The
49/// worker performs that maintenance before replying, so it needs a deadline
50/// sized for filesystem work rather than ordinary relay bookkeeping.
51const RELAY_ACKNOWLEDGE_TIMEOUT: Duration = Duration::from_secs(300);
52/// Capturing a review delta runs Git over every workspace repository, which is
53/// filesystem work on a possibly large tree rather than relay bookkeeping.
54const REVIEW_CAPTURE_TIMEOUT: Duration = Duration::from_secs(300);
55/// Bifrost's semantic diff analysis has its own 600-second budget inside the
56/// worker; this leaves room for it to report a timeout as an error rather than
57/// having the call time out underneath it.
58const REVIEW_ANALYSIS_TIMEOUT: Duration = Duration::from_secs(660);
59const RELAY_PROXY_DETACH_GRACE: Duration = Duration::from_millis(500);
60const RELAY_PROXY_REAP_POLL: Duration = Duration::from_millis(10);
61
62/// How many trailing stderr lines a failed connect reports back to its caller.
63const RELAY_PROXY_STDERR_TAIL: usize = 10;
64
65/// The proxy's last [`RELAY_PROXY_STDERR_TAIL`] non-empty stderr lines, shared
66/// with whoever has to report them.
67///
68/// The drain publishes each line here as it reads it, rather than returning
69/// the whole tail when it finishes. A failed connect has to bound how long it
70/// waits for the drain, because a proxy that leaves a grandchild holding
71/// stderr never reaches EOF. Reading the tail from here means that bound costs
72/// only the lines not yet read, instead of discarding every line already
73/// collected.
74type ProxyStderrTail = Arc<std::sync::Mutex<VecDeque<String>>>;
75
76/// Forward a relay proxy's stderr to the log, one line at a time, until the
77/// child closes it, keeping the tail in `tail`. Reporting rather than dropping
78/// keeps connect failures diagnosable now that the controller no longer shares
79/// its terminal, and lets a failed connect put the proxy's own complaint in
80/// the error the caller sees rather than only in the log.
81async fn drain_proxy_stderr(
82    errors: tokio::process::ChildStderr,
83    purpose: String,
84    session_id: String,
85    tail: ProxyStderrTail,
86) {
87    let mut lines = BufReader::new(errors).lines();
88    loop {
89        match lines.next_line().await {
90            Ok(Some(line)) if line.trim().is_empty() => continue,
91            Ok(Some(line)) => {
92                tracing::warn!(%session_id, %purpose, %line, "relay proxy stderr");
93                let mut tail = tail.lock().unwrap_or_else(PoisonError::into_inner);
94                if tail.len() == RELAY_PROXY_STDERR_TAIL {
95                    tail.pop_front();
96                }
97                tail.push_back(line);
98            }
99            Ok(None) => return,
100            Err(error) => {
101                tracing::warn!(%session_id, %purpose, %error, "read relay proxy stderr");
102                return;
103            }
104        }
105    }
106}
107
108mod errors;
109pub use errors::*;
110mod connect;
111mod exchange;
112mod relay;
113mod reviewer;
114mod transport;
115use transport::*;
116mod credential_sync;
117pub use credential_sync::*;
118
119/// Controller-side connection to the durable ACP relay protocol.
120///
121/// This type does not construct transcript state or request unbounded history.
122/// Callers persist bounded attachment pages, then acknowledge only a frontier
123/// that is already durable locally.
124pub struct RelayClient {
125    child: Option<Child>,
126    input: Option<ChildStdin>,
127    output: BufReader<ChildStdout>,
128    request_timeout: Duration,
129    /// Why this connection can no longer be used, once a call gave up on a
130    /// reply that is still in flight. See [`RelayClient::exchange`].
131    abandoned: Option<String>,
132    next_request: u64,
133    connection_nonce: u64,
134    protocol_version: u32,
135    session_id: String,
136    relay_version: String,
137    /// Content address of the executable the worker is running, as reported in
138    /// hello. `None` from a worker built before the field existed.
139    worker_build: Option<String>,
140    latest_ordinal: u64,
141    latest_digest: String,
142}
143
144#[cfg(test)]
145mod tests;