Skip to main content

hatchery_node/
restream_g4a_bridge.rs

1//! Node restream g4a bridge (A2).
2//!
3//! Tunnels the local g4a `NodeServer` through a [`crate::restream_session::NodeSession`].
4//! Core speaks the OperatorLink wire (`{"op_id","g4a"}`); C2 only ever sees opaque
5//! Noise ciphertext. The node never learns what sits behind C2.
6
7use crate::restream_session::{NodeSession, UpstreamKeys};
8use crate::uplink_journal::UplinkJournal;
9use gate4agent_node_protocol::{
10    ClientRole, NodeEventEnvelope, NodeFailureCode, NodeId, NodeIncarnationId, NodeRequest,
11    NodeResponse, ServerFrame,
12};
13use gate4agent_node_wire::{LocalNodeClient, NodeClientError};
14use hatchery_core::{
15    C2NodeEvent, C2NodeFailure, C2NodeResponse, ChannelKind, NodeCursor, RoutedNodeEvent,
16    RoutedNodeRequest, RoutedNodeResponse,
17};
18use serde_json::{json, Value};
19use std::net::SocketAddr;
20use std::sync::Arc;
21use std::time::Duration;
22use tokio::sync::watch;
23
24const HANDSHAKE: Duration = Duration::from_secs(10);
25const CONTROL_CHANNEL: u32 = 1;
26
27/// Run until `shutdown` is true: local g4a operator client ↔ restream session.
28pub async fn run(
29    local_endpoint: impl Into<String>,
30    node_id: impl Into<String>,
31    local_token: impl Into<String>,
32    uplink: SocketAddr,
33    keys: UpstreamKeys,
34    journal: Option<Arc<UplinkJournal>>,
35    mut shutdown: watch::Receiver<bool>,
36) -> Result<(), String> {
37    let local_endpoint = local_endpoint.into();
38    let node_id_str = node_id.into();
39    let local_token = local_token.into();
40    let node_id = NodeId::new(node_id_str.clone()).map_err(|e| e.to_string())?;
41
42    let mut client =
43        LocalNodeClient::connect(&local_endpoint, &node_id, ClientRole::Operator, &local_token)
44            .await
45            .map_err(|e| e.to_string())?;
46
47    let mut session = NodeSession::connect(
48        uplink,
49        &node_id_str,
50        CONTROL_CHANNEL,
51        ChannelKind::G4aControl,
52        keys,
53        journal,
54        HANDSHAKE,
55    )
56    .await?;
57
58    let hello = client.hello().clone();
59    let mut incarnation_id = hello.incarnation_id;
60    let connection_id = hello.connection_id;
61    let mut controller_owned = hello
62        .controller
63        .as_ref()
64        .is_some_and(|c| c.connection_id == connection_id);
65
66    tracing::info!(%node_id_str, connection_id, "restream g4a bridge up");
67
68    loop {
69        if *shutdown.borrow() {
70            let _ = release_controller(&mut client, &mut controller_owned).await;
71            return Ok(());
72        }
73        tokio::select! {
74            changed = shutdown.changed() => {
75                if changed.is_err() || *shutdown.borrow() {
76                    let _ = release_controller(&mut client, &mut controller_owned).await;
77                    return Ok(());
78                }
79            }
80            frame = client.recv() => {
81                match frame {
82                    Ok(ServerFrame::Event(envelope)) => {
83                        forward_event(&mut session, &node_id, incarnation_id, envelope).await?;
84                    }
85                    Ok(_) => return Err("local g4a sent an unexpected idle frame".into()),
86                    Err(err) => return Err(err.to_string()),
87                }
88            }
89            plain = session.recv(Duration::from_secs(120)) => {
90                match plain {
91                    Ok(bytes) => {
92                        handle_control(
93                            &mut session,
94                            &mut client,
95                            &node_id,
96                            &mut incarnation_id,
97                            connection_id,
98                            &mut controller_owned,
99                            &bytes,
100                        )
101                        .await?;
102                        while let Some(envelope) = client.take_event() {
103                            forward_event(&mut session, &node_id, incarnation_id, envelope).await?;
104                        }
105                    }
106                    Err(err) if err.contains("timeout") => continue,
107                    Err(err) => return Err(err),
108                }
109            }
110        }
111    }
112}
113
114async fn handle_control(
115    session: &mut NodeSession,
116    client: &mut LocalNodeClient,
117    node_id: &NodeId,
118    incarnation_id: &mut NodeIncarnationId,
119    connection_id: u64,
120    controller_owned: &mut bool,
121    plain: &[u8],
122) -> Result<(), String> {
123    let v: Value = serde_json::from_slice(plain).map_err(|e| e.to_string())?;
124    let Some(op_id) = v.get("op_id").and_then(Value::as_u64) else {
125        return Ok(());
126    };
127    let g4a = v.get("g4a").cloned().unwrap_or(Value::Null);
128    let routed: RoutedNodeRequest = serde_json::from_value(g4a).map_err(|e| e.to_string())?;
129    if routed.route.node_id.as_str() != node_id.as_str() {
130        return send_reply(
131            session,
132            op_id,
133            &failure_response(node_id, *incarnation_id, NodeFailureCode::InvalidRequest, "route node id does not match bridge"),
134        )
135        .await;
136    }
137    if routed.route.expected_incarnation_id != *incarnation_id {
138        return send_reply(
139            session,
140            op_id,
141            &failure_response(
142                node_id,
143                *incarnation_id,
144                NodeFailureCode::InvalidRequest,
145                "node incarnation changed",
146            ),
147        )
148        .await;
149    }
150
151    let request = routed.request;
152    if !is_read_only(&request) && !*controller_owned {
153        match acquire_controller(client, connection_id).await {
154            Ok(true) => *controller_owned = true,
155            Ok(false) => {
156                return send_reply(
157                    session,
158                    op_id,
159                    &failure_response(
160                        node_id,
161                        *incarnation_id,
162                        NodeFailureCode::ControllerBusy,
163                        "node controller lease is unavailable",
164                    ),
165                )
166                .await;
167            }
168            Err(err) => return Err(err.to_string()),
169        }
170    }
171
172    let deadline = request_deadline(&request);
173    let response = match tokio::time::timeout(deadline, client.request(request)).await {
174        Ok(Ok(response)) => Ok(response),
175        Ok(Err(NodeClientError::Node(failure))) => Err(failure),
176        Ok(Err(err)) => return Err(err.to_string()),
177        Err(_) => {
178            return send_reply(
179                session,
180                op_id,
181                &failure_response(
182                    node_id,
183                    *incarnation_id,
184                    NodeFailureCode::HostDirectoryReadTimedOut,
185                    "node request deadline exceeded",
186                ),
187            )
188            .await;
189        }
190    };
191
192    *incarnation_id = client.hello().incarnation_id;
193
194    let mapped = response
195        .map(|r| C2NodeResponse::from_node_response_with_control_detail(&r))
196        .map_err(|f| C2NodeFailure::from(&f));
197    send_reply(
198        session,
199        op_id,
200        &RoutedNodeResponse {
201            node_id: node_id.clone(),
202            incarnation_id: *incarnation_id,
203            response: mapped,
204        },
205    )
206    .await
207}
208
209fn failure_response(
210    node_id: &NodeId,
211    incarnation_id: NodeIncarnationId,
212    code: NodeFailureCode,
213    message: &str,
214) -> RoutedNodeResponse {
215    RoutedNodeResponse {
216        node_id: node_id.clone(),
217        incarnation_id,
218        response: Err(C2NodeFailure {
219            code,
220            message: message.into(),
221        }),
222    }
223}
224
225async fn send_reply(
226    session: &mut NodeSession,
227    op_id: u64,
228    response: &RoutedNodeResponse,
229) -> Result<(), String> {
230    let plain =
231        serde_json::to_vec(&json!({ "op_id": op_id, "g4a": response })).map_err(|e| e.to_string())?;
232    session.send(&plain).await
233}
234
235async fn forward_event(
236    session: &mut NodeSession,
237    node_id: &NodeId,
238    incarnation_id: NodeIncarnationId,
239    envelope: NodeEventEnvelope,
240) -> Result<(), String> {
241    let event = RoutedNodeEvent {
242        node_id: node_id.clone(),
243        cursor: NodeCursor {
244            incarnation_id,
245            sequence: envelope.sequence,
246        },
247        event: C2NodeEvent::from_node_event_with_control_detail(&envelope.event),
248    };
249    let plain = serde_json::to_vec(&json!({ "g4a": event })).map_err(|e| e.to_string())?;
250    session.send(&plain).await
251}
252
253async fn acquire_controller(
254    client: &mut LocalNodeClient,
255    connection_id: u64,
256) -> Result<bool, NodeClientError> {
257    match client
258        .request(NodeRequest::AcquireController {
259            lease_ms: gate4agent_node_protocol::MAX_CONTROLLER_LEASE_MS,
260        })
261        .await?
262    {
263        NodeResponse::Controller { controller } => Ok(controller
264            .as_ref()
265            .is_some_and(|state| state.connection_id == connection_id)),
266        _ => Err(NodeClientError::Protocol(
267            "controller acquisition returned a different response".into(),
268        )),
269    }
270}
271
272async fn release_controller(
273    client: &mut LocalNodeClient,
274    controller_owned: &mut bool,
275) -> Result<(), NodeClientError> {
276    if !*controller_owned {
277        return Ok(());
278    }
279    let _ = client.request(NodeRequest::ReleaseController).await?;
280    *controller_owned = false;
281    Ok(())
282}
283
284fn is_read_only(request: &NodeRequest) -> bool {
285    matches!(
286        request,
287        NodeRequest::Snapshot
288            | NodeRequest::Resync { .. }
289            | NodeRequest::BrowseHostDirectories { .. }
290            | NodeRequest::InspectWorkspace { .. }
291            | NodeRequest::ReadWorkspaceFile { .. }
292            | NodeRequest::ReadGitHistory { .. }
293            | NodeRequest::ReadGitDiff { .. }
294            | NodeRequest::CatalogNativeSessions { .. }
295            | NodeRequest::PageNativeSessions { .. }
296            | NodeRequest::PreviewNativeSession { .. }
297            | NodeRequest::PreviewSessionRecord { .. }
298    )
299}
300
301fn request_deadline(request: &NodeRequest) -> Duration {
302    match request {
303        NodeRequest::CreateStandaloneWorkspace { .. }
304        | NodeRequest::CreateWorktree { .. }
305        | NodeRequest::RemoveWorktree { .. }
306        | NodeRequest::CleanupManagedWorktree { .. } => Duration::from_secs(240),
307        NodeRequest::Spawn { .. }
308        | NodeRequest::Resume { .. }
309        | NodeRequest::Stop { .. }
310        | NodeRequest::SpawnSpec { .. }
311        | NodeRequest::SpawnSpecWithHarnessMcp { .. }
312        | NodeRequest::SpawnManagedWorktree { .. } => Duration::from_secs(60),
313        NodeRequest::InspectWorkspace { .. } => Duration::from_secs(15),
314        _ => Duration::from_secs(30),
315    }
316}