use crate::restream_session::{NodeSession, SendPrio, UpstreamKeys};
use crate::uplink_journal::UplinkJournal;
use gate4agent_node_protocol::{
ClientRole, NodeEventEnvelope, NodeFailureCode, NodeId, NodeIncarnationId, NodeRequest,
NodeResponse, ServerFrame,
};
use gate4agent_node_wire::{LocalNodeClient, NodeClientError};
use hatchery_core::{
BrowserClientOp, BrowserOpReply, BrowserStreamEvent, C2NodeEvent, C2NodeFailure,
C2NodeResponse, ChannelKind, NodeCursor, RoutedNodeEvent, RoutedNodeRequest,
RoutedNodeResponse,
};
use serde_json::{json, Value};
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, watch};
const HANDSHAKE: Duration = Duration::from_secs(10);
const CONTROL_CHANNEL: u32 = 1;
pub async fn run(
local_endpoint: impl Into<String>,
node_id: impl Into<String>,
local_token: impl Into<String>,
uplink: SocketAddr,
keys: UpstreamKeys,
journal: Option<Arc<UplinkJournal>>,
browser_profile_root: Option<PathBuf>,
calls: Option<crate::uplink_calls::NodeCalls>,
mut shutdown: watch::Receiver<bool>,
) -> Result<(), String> {
let local_endpoint = local_endpoint.into();
let node_id_str = node_id.into();
let local_token = local_token.into();
let node_id = NodeId::new(node_id_str.clone()).map_err(|e| e.to_string())?;
let browser_profile_root = browser_profile_root.unwrap_or_else(|| {
std::env::temp_dir().join("hatchery-browser-profiles")
});
let mut client =
LocalNodeClient::connect(&local_endpoint, &node_id, ClientRole::Operator, &local_token)
.await
.map_err(|e| e.to_string())?;
let mut session = NodeSession::connect(
uplink,
&node_id_str,
CONTROL_CHANNEL,
ChannelKind::G4aControl,
keys,
journal,
HANDSHAKE,
)
.await?;
let hello = client.hello().clone();
let incarnation_id = hello.incarnation_id;
let connection_id = hello.connection_id;
tracing::info!(%node_id_str, connection_id, "restream g4a bridge up");
if let Some(c) = &calls {
c.set_uplink(Some(session.uplink())).await;
}
let result = run_loop(&mut client, &mut session, &node_id, incarnation_id, &local_endpoint, &local_token, browser_profile_root, &mut shutdown).await;
if let Some(c) = &calls {
c.set_uplink(None).await;
}
result
}
#[allow(clippy::too_many_arguments)]
async fn run_loop(
client: &mut LocalNodeClient,
session: &mut NodeSession,
node_id: &NodeId,
incarnation_id: NodeIncarnationId,
local_endpoint: &str,
local_token: &str,
browser_profile_root: PathBuf,
shutdown: &mut watch::Receiver<bool>,
) -> Result<(), String> {
let worker_client =
LocalNodeClient::connect(&local_endpoint, &node_id, ClientRole::Operator, &local_token)
.await
.map_err(|e| e.to_string())?;
let (req_tx, req_rx) = mpsc::unbounded_channel::<Vec<u8>>();
let (out_tx, mut out_rx) = mpsc::unbounded_channel::<Vec<u8>>();
let (browser_ev_tx, mut browser_ev_rx) = mpsc::unbounded_channel::<BrowserStreamEvent>();
let (worker_stop_tx, worker_stop) = watch::channel(false);
let mut worker = tokio::spawn(request_worker(
worker_client,
node_id.clone(),
req_rx,
out_tx,
browser_profile_root.clone(),
browser_ev_tx,
worker_stop,
));
let mut last_pty_input: Option<std::time::Instant> = None;
let uplink_watch = session.uplink();
let result = loop {
if *shutdown.borrow() {
break Ok(());
}
tokio::select! {
changed = shutdown.changed() => {
if changed.is_err() || *shutdown.borrow() { break Ok(()); }
}
done = &mut worker => {
break match done { Ok(r) => r.and(Err("request worker ended".into())), Err(e) => Err(e.to_string()) };
}
frame = client.recv() => {
match frame {
Ok(ServerFrame::Event(envelope)) => {
let is_pty = matches!(envelope.event, gate4agent_node_protocol::NodeEvent::TerminalFrame { .. });
if is_pty {
if let Some(at) = last_pty_input.take() {
tracing::debug!(target: "pty_lat", bridge_us = at.elapsed().as_micros() as u64, "bridge: pty input received to echo frame forwarded");
}
}
if envelope.event.requires_harness_mcp_proxy_capability() { continue; }
let plain = event_bytes(&node_id, incarnation_id, envelope)?;
let prio = if is_pty { SendPrio::Pty } else { SendPrio::Control };
if let Err(e) = session.send_prio(&plain, prio).await { break Err(e); }
}
Ok(_) => break Err("local g4a sent an unexpected frame on the event connection".into()),
Err(err) => break Err(err.to_string()),
}
}
out = out_rx.recv() => {
let Some(out) = out else { break Err("request worker closed".into()) };
if let Err(e) = session.send_prio(&out, SendPrio::Control).await { break Err(e); }
}
plain = session.recv(Duration::from_secs(120)) => {
match plain {
Ok(bytes) => {
if last_pty_input.is_none() && is_pty_input(&bytes) { last_pty_input = Some(std::time::Instant::now()); }
if req_tx.send(bytes).is_err() { break Err("request worker closed".into()); }
}
Err(err) if err.contains("timeout") => continue,
Err(err) => break Err(err),
}
}
_ = uplink_watch.closed() => break Err("uplink write stalled; reconnecting".into()),
ev = browser_ev_rx.recv() => {
let Some(ev) = ev else { continue };
let plain = serde_json::to_vec(&json!({ "browser": ev })).map_err(|e| e.to_string())?;
if let Err(e) = session.send_prio(&plain, SendPrio::Browser).await { break Err(e); }
}
}
};
let _ = worker_stop_tx.send(true);
let _ = tokio::time::timeout(Duration::from_secs(3), worker).await;
result
}
fn is_pty_input(plain: &[u8]) -> bool {
let s = String::from_utf8_lossy(&plain[..plain.len().min(512)]);
s.contains("TerminalBytes") || s.contains("terminal_bytes")
}
type Sink = mpsc::UnboundedSender<Vec<u8>>;
async fn request_worker(
mut client: LocalNodeClient,
node_id: NodeId,
mut rx: mpsc::UnboundedReceiver<Vec<u8>>,
out: Sink,
browser_profile_root: PathBuf,
browser_ev_tx: mpsc::UnboundedSender<BrowserStreamEvent>,
mut stop: watch::Receiver<bool>,
) -> Result<(), String> {
#[cfg(feature = "d2b")]
let browsers = crate::browser_capture::shared_map();
let hello = client.hello().clone();
let mut incarnation_id = hello.incarnation_id;
let connection_id = hello.connection_id;
let mut controller_owned = hello.controller.as_ref().is_some_and(|c| c.connection_id == connection_id);
let mut inflight: std::collections::HashMap<u64, (u64, NodeRequest)> = std::collections::HashMap::new();
let mut retry: Vec<(u64, NodeRequest)> = Vec::new();
let fwd = |out: &Sink, incarnation_id: NodeIncarnationId, envelope: NodeEventEnvelope| -> Result<(), String> {
if envelope.event.requires_harness_mcp_proxy_capability() {
out.send(event_bytes(&node_id, incarnation_id, envelope)?).map_err(|_| "bridge closed".to_string())?;
}
Ok(())
};
loop {
tokio::select! {
_ = stop.changed() => {
let _ = release_controller(&mut client, &mut controller_owned).await;
return Ok(());
}
frame = client.recv() => {
match frame.map_err(|e| e.to_string())? {
ServerFrame::Event(envelope) => fwd(&out, incarnation_id, envelope)?,
ServerFrame::Reply(reply) => {
on_pipelined_reply(&out, &node_id, &mut incarnation_id, &mut controller_owned, &mut inflight, &mut retry, reply)?;
}
_ => return Err("local g4a sent an unexpected idle frame".into()),
}
}
plain = rx.recv() => {
let Some(bytes) = plain else {
let _ = release_controller(&mut client, &mut controller_owned).await;
return Ok(());
};
if controller_owned && retry.is_empty() {
if let Some((op_id, request)) = pipelinable(&bytes, &node_id, incarnation_id) {
let request_id = client.send(request.clone()).await.map_err(|e| e.to_string())?;
inflight.insert(request_id, (op_id, request));
continue;
}
}
while !inflight.is_empty() {
match client.recv().await.map_err(|e| e.to_string())? {
ServerFrame::Event(envelope) => fwd(&out, incarnation_id, envelope)?,
ServerFrame::Reply(reply) => on_pipelined_reply(&out, &node_id, &mut incarnation_id, &mut controller_owned, &mut inflight, &mut retry, reply)?,
_ => return Err("local g4a sent an unexpected idle frame".into()),
}
}
let mut queue: Vec<Vec<u8>> = Vec::new();
for (op_id, request) in std::mem::take(&mut retry) {
let routed = RoutedNodeRequest {
route: hatchery_core::NodeRoute {
node_id: hatchery_core::NodeId::new(node_id.as_str().to_owned()).map_err(|e| e.to_string())?,
expected_incarnation_id: NodeIncarnationId::from_bytes([0; 16]),
},
request,
};
queue.push(serde_json::to_vec(&json!({ "op_id": op_id, "g4a": routed })).map_err(|e| e.to_string())?);
}
queue.push(bytes);
for plain in queue {
let started = std::time::Instant::now();
let kind: String = serde_json::from_slice::<Value>(&plain).ok().and_then(|v| v.pointer("/g4a/request").and_then(|r| r.as_object().and_then(|o| o.keys().next().cloned()).or_else(|| r.as_str().map(str::to_owned)))).unwrap_or_else(|| "browser".into());
handle_control(&out, &mut client, &node_id, &mut incarnation_id, connection_id, &mut controller_owned, &plain, &browser_profile_root, browser_ev_tx.clone(), #[cfg(feature = "d2b")] &browsers).await?;
while let Some(envelope) = client.take_event() {
fwd(&out, incarnation_id, envelope)?;
}
if started.elapsed() > Duration::from_millis(500) {
tracing::warn!(request = %kind, elapsed_ms = started.elapsed().as_millis() as u64, "bridge: slow node request held the request queue");
}
}
}
}
}
}
fn pipelinable(plain: &[u8], node_id: &NodeId, incarnation_id: NodeIncarnationId) -> Option<(u64, NodeRequest)> {
let v: Value = serde_json::from_slice(plain).ok()?;
let op_id = v.get("op_id")?.as_u64()?;
let routed: RoutedNodeRequest = serde_json::from_value(v.get("g4a")?.clone()).ok()?;
if routed.route.node_id.as_str() != node_id.as_str() { return None; }
let expected = routed.route.expected_incarnation_id;
if expected != NodeIncarnationId::from_bytes([0; 16]) && expected != incarnation_id { return None; }
match routed.request {
NodeRequest::TerminalBytes { .. } | NodeRequest::TerminalControl { .. } | NodeRequest::Interrupt { .. } => Some((op_id, routed.request)),
_ => None,
}
}
fn on_pipelined_reply(
out: &Sink,
node_id: &NodeId,
incarnation_id: &mut NodeIncarnationId,
controller_owned: &mut bool,
inflight: &mut std::collections::HashMap<u64, (u64, NodeRequest)>,
retry: &mut Vec<(u64, NodeRequest)>,
reply: gate4agent_node_protocol::ResponseEnvelope,
) -> Result<(), String> {
let Some((op_id, request)) = inflight.remove(&reply.request_id) else {
return Err(format!("unexpected local reply id {}", reply.request_id));
};
if let Err(f) = &reply.result {
if f.code == NodeFailureCode::ControllerRequired {
*controller_owned = false;
retry.push((op_id, request));
return Ok(());
}
}
if !retry.is_empty() {
retry.push((op_id, request));
return Ok(());
}
let mapped = reply
.result
.map(|r| C2NodeResponse::from_node_response_with_control_detail(&r))
.map_err(|f| C2NodeFailure::from(&f));
let plain = serde_json::to_vec(&json!({ "op_id": op_id, "g4a": RoutedNodeResponse { node_id: node_id.clone(), incarnation_id: *incarnation_id, response: mapped } })).map_err(|e| e.to_string())?;
out.send(plain).map_err(|_| "bridge closed".to_string())
}
async fn handle_control(
session: &Sink,
client: &mut LocalNodeClient,
node_id: &NodeId,
incarnation_id: &mut NodeIncarnationId,
connection_id: u64,
controller_owned: &mut bool,
plain: &[u8],
browser_profile_root: &std::path::Path,
browser_ev_tx: mpsc::UnboundedSender<BrowserStreamEvent>,
#[cfg(feature = "d2b")] browsers: &std::sync::Arc<crate::browser_capture::BrowserSessionMap>,
) -> Result<(), String> {
let v: Value = serde_json::from_slice(plain).map_err(|e| e.to_string())?;
let Some(op_id) = v.get("op_id").and_then(Value::as_u64) else {
return Ok(());
};
if let Some(browser) = v.get("browser") {
let op: BrowserClientOp = serde_json::from_value(browser.clone()).map_err(|e| e.to_string())?;
let reply = handle_browser_op(
&node_id_str_from(node_id),
op,
browser_profile_root,
browser_ev_tx,
#[cfg(feature = "d2b")]
browsers,
)
.await;
let plain = serde_json::to_vec(&json!({ "op_id": op_id, "browser": reply })).map_err(|e| e.to_string())?;
return session.send(plain).map_err(|_| "bridge closed".to_string());
}
let g4a = v.get("g4a").cloned().unwrap_or(Value::Null);
let routed: RoutedNodeRequest = serde_json::from_value(g4a).map_err(|e| e.to_string())?;
if routed.route.node_id.as_str() != node_id.as_str() {
return send_reply(
session,
op_id,
&failure_response(node_id, *incarnation_id, NodeFailureCode::InvalidRequest, "route node id does not match bridge"),
)
.await;
}
let expected = routed.route.expected_incarnation_id;
if expected != NodeIncarnationId::from_bytes([0; 16]) && expected != *incarnation_id {
return send_reply(
session,
op_id,
&failure_response(
node_id,
*incarnation_id,
NodeFailureCode::InvalidRequest,
"node incarnation changed",
),
)
.await;
}
let request = routed.request;
let controller_gated = !is_read_only(&request);
if controller_gated && !*controller_owned {
match acquire_controller(client, connection_id).await {
Ok(true) => *controller_owned = true,
Ok(false) => {
return send_reply(
session,
op_id,
&failure_response(
node_id,
*incarnation_id,
NodeFailureCode::ControllerBusy,
"node controller lease is unavailable",
),
)
.await;
}
Err(err) => return Err(err.to_string()),
}
}
let deadline = request_deadline(&request);
let mut remaining_retries = 1u8;
let response = loop {
match tokio::time::timeout(deadline, client.request(request.clone())).await {
Ok(Ok(response)) => break Ok(response),
Ok(Err(NodeClientError::Node(failure))) => {
if controller_gated
&& failure.code == NodeFailureCode::ControllerRequired
&& remaining_retries > 0
{
remaining_retries -= 1;
*controller_owned = false;
match acquire_controller(client, connection_id).await {
Ok(true) => {
*controller_owned = true;
continue;
}
Ok(false) => break Err(failure),
Err(err) => return Err(err.to_string()),
}
}
break Err(failure);
}
Ok(Err(err)) => return Err(err.to_string()),
Err(_) => {
return send_reply(
session,
op_id,
&failure_response(
node_id,
*incarnation_id,
NodeFailureCode::HostDirectoryReadTimedOut,
"node request deadline exceeded",
),
)
.await;
}
}
};
*incarnation_id = client.hello().incarnation_id;
let mut mapped = response
.map(|r| C2NodeResponse::from_node_response_with_control_detail(&r))
.map_err(|f| C2NodeFailure::from(&f));
if let Ok(C2NodeResponse::Snapshot { snapshot, .. }) = &mut mapped {
snapshot.machine = Some(crate::collect_machine_environment());
}
send_reply(
session,
op_id,
&RoutedNodeResponse {
node_id: node_id.clone(),
incarnation_id: *incarnation_id,
response: mapped,
},
)
.await
}
fn failure_response(
node_id: &NodeId,
incarnation_id: NodeIncarnationId,
code: NodeFailureCode,
message: &str,
) -> RoutedNodeResponse {
RoutedNodeResponse {
node_id: node_id.clone(),
incarnation_id,
response: Err(C2NodeFailure {
code,
message: message.into(),
}),
}
}
async fn send_reply(
session: &Sink,
op_id: u64,
response: &RoutedNodeResponse,
) -> Result<(), String> {
let plain =
serde_json::to_vec(&json!({ "op_id": op_id, "g4a": response })).map_err(|e| e.to_string())?;
session.send(plain).map_err(|_| "bridge closed".to_string())
}
fn event_bytes(
node_id: &NodeId,
incarnation_id: NodeIncarnationId,
envelope: NodeEventEnvelope,
) -> Result<Vec<u8>, String> {
let event = RoutedNodeEvent {
node_id: node_id.clone(),
cursor: NodeCursor {
incarnation_id,
sequence: envelope.sequence,
},
event: C2NodeEvent::from_node_event_with_control_detail(&envelope.event),
};
serde_json::to_vec(&json!({ "g4a": event })).map_err(|e| e.to_string())
}
async fn acquire_controller(
client: &mut LocalNodeClient,
connection_id: u64,
) -> Result<bool, NodeClientError> {
match client
.request(NodeRequest::AcquireController {
lease_ms: gate4agent_node_protocol::MAX_CONTROLLER_LEASE_MS,
})
.await?
{
NodeResponse::Controller { controller } => Ok(controller
.as_ref()
.is_some_and(|state| state.connection_id == connection_id)),
_ => Err(NodeClientError::Protocol(
"controller acquisition returned a different response".into(),
)),
}
}
async fn release_controller(
client: &mut LocalNodeClient,
controller_owned: &mut bool,
) -> Result<(), NodeClientError> {
if !*controller_owned {
return Ok(());
}
let _ = client.request(NodeRequest::ReleaseController).await?;
*controller_owned = false;
Ok(())
}
fn is_read_only(request: &NodeRequest) -> bool {
matches!(
request,
NodeRequest::Snapshot
| NodeRequest::Resync { .. }
| NodeRequest::BrowseHostDirectories { .. }
| NodeRequest::InspectWorkspace { .. }
| NodeRequest::ReadWorkspaceFile { .. }
| NodeRequest::ReadGitHistory { .. }
| NodeRequest::ReadGitDiff { .. }
| NodeRequest::CatalogNativeSessions { .. }
| NodeRequest::PageNativeSessions { .. }
| NodeRequest::PreviewNativeSession { .. }
| NodeRequest::PreviewSessionRecord { .. }
)
}
fn request_deadline(request: &NodeRequest) -> Duration {
match request {
NodeRequest::CreateStandaloneWorkspace { .. }
| NodeRequest::CreateWorktree { .. }
| NodeRequest::RemoveWorktree { .. }
| NodeRequest::CleanupManagedWorktree { .. } => Duration::from_secs(240),
NodeRequest::Spawn { .. }
| NodeRequest::Resume { .. }
| NodeRequest::Stop { .. }
| NodeRequest::SpawnSpec { .. }
| NodeRequest::SpawnSpecWithHarnessMcp { .. }
| NodeRequest::SpawnManagedWorktree { .. } => Duration::from_secs(60),
NodeRequest::InspectWorkspace { .. } => Duration::from_secs(15),
_ => Duration::from_secs(30),
}
}
fn node_id_str_from(node_id: &NodeId) -> String {
node_id.as_str().to_owned()
}
async fn handle_browser_op(
node_id: &str,
op: BrowserClientOp,
browser_profile_root: &std::path::Path,
_browser_ev_tx: mpsc::UnboundedSender<BrowserStreamEvent>,
#[cfg(feature = "d2b")] browsers: &std::sync::Arc<crate::browser_capture::BrowserSessionMap>,
) -> BrowserOpReply {
#[cfg(feature = "local-server")]
{
if let Some(reply) = handle_local_server_op(node_id, &op).await {
return reply;
}
}
#[cfg(not(feature = "d2b"))]
{
let _ = (node_id, op, browser_profile_root, _browser_ev_tx);
return BrowserOpReply {
ok: false,
session: String::new(),
message: "dig2browser feature (d2b) is not enabled in this node build".into(),
body: None,
};
}
#[cfg(feature = "d2b")]
{
match op {
BrowserClientOp::Start {
session,
url,
profile_dir,
} => {
let dir = match profile_dir {
Some(p) if !p.is_empty() => PathBuf::from(p),
_ => browser_profile_root.join(&session),
};
match browsers
.start(node_id.to_owned(), session.clone(), url, dir, _browser_ev_tx)
.await
{
Ok(()) => BrowserOpReply {
ok: true,
session,
message: String::new(),
body: None,
},
Err(message) => BrowserOpReply {
ok: false,
session,
message,
body: None,
},
}
}
BrowserClientOp::Input { session, bytes } => match browsers.input(&session, &bytes).await
{
Ok(()) => BrowserOpReply {
ok: true,
session,
message: String::new(),
body: None,
},
Err(message) => BrowserOpReply {
ok: false,
session,
message,
body: None,
},
},
BrowserClientOp::Stop { session } => {
#[cfg(feature = "local-server")]
{
crate::local_server::LocalServerPool::global()
.stop_session(&session)
.await;
}
match browsers.stop(&session).await {
Ok(()) => BrowserOpReply {
ok: true,
session,
message: String::new(),
body: None,
},
Err(message) => BrowserOpReply {
ok: false,
session,
message,
body: None,
},
}
}
BrowserClientOp::LocalServerStart { session, .. }
| BrowserClientOp::LocalServerStop { session, .. }
| BrowserClientOp::LocalServerList { session } => BrowserOpReply {
ok: false,
session,
message: "local_server op was not handled".into(),
body: None,
},
}
}
}
#[cfg(feature = "local-server")]
async fn handle_local_server_op(node_id: &str, op: &BrowserClientOp) -> Option<BrowserOpReply> {
let _ = node_id;
let (session, door_op, payload) = match op {
BrowserClientOp::LocalServerList { session } => (
session.clone(),
"local_server.list",
serde_json::json!({ "session": session }),
),
BrowserClientOp::LocalServerStop { session, id } => (
session.clone(),
"local_server.stop",
serde_json::json!({ "session": session, "id": id }),
),
BrowserClientOp::LocalServerStart {
session,
kind,
root,
cwd,
bin,
open,
spa,
} => {
let mut payload = serde_json::json!({
"session": session,
"kind": kind,
"open": open,
});
if let Some(r) = root {
payload["root"] = serde_json::json!(r);
}
if let Some(c) = cwd {
payload["cwd"] = serde_json::json!(c);
}
if let Some(b) = bin {
payload["bin"] = serde_json::json!(b);
}
if let Some(s) = spa {
payload["spa"] = serde_json::json!(s);
}
(session.clone(), "local_server.start", payload)
}
_ => return None,
};
match crate::local_server::handle_op(door_op, &payload).await {
Ok(body) => Some(BrowserOpReply {
ok: true,
session,
message: String::new(),
body: Some(body),
}),
Err(message) => Some(BrowserOpReply {
ok: false,
session,
message,
body: None,
}),
}
}