1use 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
27pub 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}