microclaw_engine/internal/runtime/
remote.rs

1use std::sync::{Arc, Mutex};
2
3use async_trait::async_trait;
4use microclaw_core::run_protocol::{
5    AgentProfile, RunRequest, RunStatus, RuntimeCapabilities, RuntimeControl, RuntimeError,
6    RuntimeErrorCode, WorkerCommand, WorkerDescriptor, WorkerFrame, WorkerHealth,
7    WorkerHealthStatus, WORKER_PROTOCOL_VERSION,
8};
9
10use crate::internal::runtime::{
11    ControlRequest, ExecutionContext, ExecutionResult, RunExecutor, RunHandle, Runtime, Worker,
12};
13
14/// One ordered, bidirectional connection to a remote Worker endpoint.
15#[async_trait]
16pub trait WorkerConnection: Send {
17    async fn send(&mut self, command: WorkerCommand) -> Result<(), RuntimeError>;
18    async fn receive(&mut self) -> Result<Option<WorkerFrame>, RuntimeError>;
19}
20
21/// Creates independent remote Worker connections for discovery, health, and runs.
22#[async_trait]
23pub trait WorkerTransport: Send + Sync + 'static {
24    async fn connect(&self) -> Result<Box<dyn WorkerConnection>, RuntimeError>;
25}
26
27struct RemoteExecutor {
28    transport: Arc<dyn WorkerTransport>,
29    capabilities: RuntimeCapabilities,
30}
31
32#[async_trait]
33impl RunExecutor for RemoteExecutor {
34    async fn execute(
35        &self,
36        profile: AgentProfile,
37        mut request: RunRequest,
38        mut context: ExecutionContext,
39    ) -> Result<ExecutionResult, RuntimeError> {
40        let run_id = context.run_id().clone();
41        request.run_id = Some(run_id.clone());
42        let mut connection = self.transport.connect().await?;
43        connection
44            .send(WorkerCommand::Submit {
45                protocol_version: WORKER_PROTOCOL_VERSION,
46                profile,
47                request: Box::new(request),
48            })
49            .await?;
50
51        let mut controls_open = true;
52        let mut pending_control: Option<ControlRequest> = None;
53        let mut last_remote_sequence = None;
54        let mut reconnect_attempts = 0_u8;
55        loop {
56            tokio::select! {
57                frame = connection.receive() => {
58                    let frame = match frame {
59                        Ok(Some(frame)) => {
60                            reconnect_attempts = 0;
61                            frame
62                        }
63                        Ok(None) | Err(_) if reconnect_attempts < 3 => {
64                            reconnect_attempts += 1;
65                            connection = self.transport.connect().await?;
66                            connection.send(WorkerCommand::ResumeEvents {
67                                protocol_version: WORKER_PROTOCOL_VERSION,
68                                run_id: run_id.clone(),
69                                after_sequence: last_remote_sequence,
70                            }).await?;
71                            if let Some(control) = pending_control.as_ref() {
72                                connection.send(WorkerCommand::Control {
73                                    protocol_version: WORKER_PROTOCOL_VERSION,
74                                    control: control.control().clone(),
75                                }).await?;
76                            }
77                            continue;
78                        }
79                        Ok(None) => return Err(unavailable("remote Worker disconnected before returning a result")),
80                        Err(error) => return Err(error),
81                    };
82                    frame.validate_protocol()?;
83                    match frame {
84                        WorkerFrame::Accepted { run_id: accepted, .. } => {
85                            ensure_run_id(&run_id, &accepted)?;
86                        }
87                        WorkerFrame::Event { envelope, .. } => {
88                            if envelope.run_id != run_id.as_str() {
89                                return Err(protocol_error("remote Worker returned an event for a different run"));
90                            }
91                            if last_remote_sequence.is_some_and(|sequence| envelope.sequence <= sequence) {
92                                continue;
93                            }
94                            last_remote_sequence = Some(envelope.sequence);
95                            if !matches!(
96                                envelope.event,
97                                microclaw_core::runtime_event::RuntimeEvent::FinalResponse { .. }
98                                    | microclaw_core::runtime_event::RuntimeEvent::Cancelled { .. }
99                            ) {
100                                context.emit(envelope.event)?;
101                            }
102                        }
103                        WorkerFrame::ControlAcknowledged { run_id: acknowledged, .. } => {
104                            ensure_run_id(&run_id, &acknowledged)?;
105                            if let Some(control) = pending_control.take() {
106                                control.accept();
107                            }
108                        }
109                        WorkerFrame::Result { result, .. } => {
110                            ensure_run_id(&run_id, &result.run_id)?;
111                            if let Some(control) = pending_control.take() {
112                                control.reject(unavailable("run completed before control acknowledgement"));
113                            }
114                            return if result.status == RunStatus::Completed {
115                                Ok(ExecutionResult {
116                                    session_id: result.session_id,
117                                    final_text: result.final_text,
118                                    metadata: result.metadata,
119                                })
120                            } else {
121                                Err(result.error.unwrap_or(RuntimeError {
122                                    code: RuntimeErrorCode::Internal,
123                                    message: format!("remote Worker ended with status {:?}", result.status),
124                                    retryable: false,
125                                }))
126                            };
127                        }
128                        WorkerFrame::Error { run_id: frame_run_id, error, .. } => {
129                            if let Some(frame_run_id) = frame_run_id.as_ref() {
130                                ensure_run_id(&run_id, frame_run_id)?;
131                            }
132                            if let Some(control) = pending_control.take() {
133                                control.reject(error);
134                            } else {
135                                return Err(error);
136                            }
137                        }
138                        WorkerFrame::Descriptor { .. } | WorkerFrame::Health { .. } => {}
139                    }
140                }
141                control = context.next_control_request(), if controls_open && pending_control.is_none() => {
142                    match control {
143                        Some(control) => {
144                            let command = control.control().clone();
145                            ensure_control_run_id(&run_id, &command)?;
146                            connection.send(WorkerCommand::Control {
147                                protocol_version: WORKER_PROTOCOL_VERSION,
148                                control: command,
149                            }).await?;
150                            pending_control = Some(control);
151                        }
152                        None => controls_open = false,
153                    }
154                }
155            }
156        }
157    }
158
159    fn capabilities(&self) -> RuntimeCapabilities {
160        self.capabilities.clone()
161    }
162}
163
164/// A Worker backed by a versioned remote transport.
165#[derive(Clone)]
166pub struct RemoteWorker {
167    descriptor: WorkerDescriptor,
168    runtime: Runtime,
169    transport: Arc<dyn WorkerTransport>,
170    health: Arc<Mutex<WorkerHealth>>,
171}
172
173impl RemoteWorker {
174    pub async fn connect(transport: impl WorkerTransport) -> Result<Self, RuntimeError> {
175        let transport: Arc<dyn WorkerTransport> = Arc::new(transport);
176        let mut connection = transport.connect().await?;
177        connection
178            .send(WorkerCommand::Describe {
179                protocol_version: WORKER_PROTOCOL_VERSION,
180            })
181            .await?;
182        let frame = connection
183            .receive()
184            .await?
185            .ok_or_else(|| unavailable("remote Worker disconnected during discovery"))?;
186        frame.validate_protocol()?;
187        let WorkerFrame::Descriptor { descriptor, .. } = frame else {
188            return Err(protocol_error(
189                "remote Worker did not return a descriptor during discovery",
190            ));
191        };
192        let health = WorkerHealth {
193            worker_id: descriptor.id.clone(),
194            status: WorkerHealthStatus::Ready,
195            active_runs: 0,
196            queued_runs: 0,
197            observed_at: String::new(),
198        };
199        let runtime = Runtime::builder()
200            .shared_executor(Arc::new(RemoteExecutor {
201                transport: transport.clone(),
202                capabilities: descriptor.capabilities.clone(),
203            }))
204            .max_concurrent_runs(descriptor.max_concurrent_runs.max(1))
205            .build()
206            .map_err(|error| RuntimeError {
207                code: RuntimeErrorCode::Configuration,
208                message: error.to_string(),
209                retryable: false,
210            })?;
211        Ok(Self {
212            descriptor,
213            runtime,
214            transport,
215            health: Arc::new(Mutex::new(health)),
216        })
217    }
218
219    pub async fn refresh_health(&self) -> Result<WorkerHealth, RuntimeError> {
220        let mut connection = self.transport.connect().await?;
221        connection
222            .send(WorkerCommand::Health {
223                protocol_version: WORKER_PROTOCOL_VERSION,
224            })
225            .await?;
226        let frame = connection
227            .receive()
228            .await?
229            .ok_or_else(|| unavailable("remote Worker disconnected during health check"))?;
230        frame.validate_protocol()?;
231        let WorkerFrame::Health { health, .. } = frame else {
232            return Err(protocol_error(
233                "remote Worker did not return health during health check",
234            ));
235        };
236        if health.worker_id != self.descriptor.id {
237            return Err(protocol_error("remote Worker health identity changed"));
238        }
239        *self
240            .health
241            .lock()
242            .unwrap_or_else(|error| error.into_inner()) = health.clone();
243        Ok(health)
244    }
245}
246
247#[async_trait]
248impl Worker for RemoteWorker {
249    fn descriptor(&self) -> WorkerDescriptor {
250        self.descriptor.clone()
251    }
252
253    fn health(&self) -> WorkerHealth {
254        self.health
255            .lock()
256            .unwrap_or_else(|error| error.into_inner())
257            .clone()
258    }
259
260    async fn submit(
261        &self,
262        profile: AgentProfile,
263        request: RunRequest,
264    ) -> Result<RunHandle, RuntimeError> {
265        if matches!(
266            self.health().status,
267            WorkerHealthStatus::Draining | WorkerHealthStatus::Unavailable
268        ) {
269            return Err(unavailable("remote Worker is not accepting new runs"));
270        }
271        Ok(self.runtime.agent(profile).run(request))
272    }
273}
274
275fn ensure_run_id(
276    expected: &microclaw_core::run_protocol::RunId,
277    actual: &microclaw_core::run_protocol::RunId,
278) -> Result<(), RuntimeError> {
279    if expected == actual {
280        Ok(())
281    } else {
282        Err(protocol_error(
283            "remote Worker returned a frame for a different run",
284        ))
285    }
286}
287
288fn ensure_control_run_id(
289    expected: &microclaw_core::run_protocol::RunId,
290    control: &RuntimeControl,
291) -> Result<(), RuntimeError> {
292    let actual = match control {
293        RuntimeControl::Cancel { run_id }
294        | RuntimeControl::Steer { run_id, .. }
295        | RuntimeControl::ResolveApproval { run_id, .. } => run_id,
296    };
297    ensure_run_id(expected, actual)
298}
299
300fn protocol_error(message: impl Into<String>) -> RuntimeError {
301    RuntimeError {
302        code: RuntimeErrorCode::InvalidRequest,
303        message: message.into(),
304        retryable: false,
305    }
306}
307
308fn unavailable(message: impl Into<String>) -> RuntimeError {
309    RuntimeError {
310        code: RuntimeErrorCode::Unavailable,
311        message: message.into(),
312        retryable: true,
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use std::collections::{BTreeMap, VecDeque};
319
320    use microclaw_core::run_protocol::{RunId, RunResult, SessionId, WorkerId};
321    use microclaw_core::runtime_event::{RuntimeEvent, RuntimeEventEnvelope};
322
323    use super::*;
324
325    #[derive(Clone)]
326    struct ScriptedTransport {
327        connections: Arc<Mutex<VecDeque<VecDeque<WorkerFrame>>>>,
328        sent: Arc<Mutex<Vec<WorkerCommand>>>,
329    }
330
331    struct ScriptedConnection {
332        frames: VecDeque<WorkerFrame>,
333        sent: Arc<Mutex<Vec<WorkerCommand>>>,
334    }
335
336    #[async_trait]
337    impl WorkerTransport for ScriptedTransport {
338        async fn connect(&self) -> Result<Box<dyn WorkerConnection>, RuntimeError> {
339            let frames = self
340                .connections
341                .lock()
342                .unwrap()
343                .pop_front()
344                .ok_or_else(|| unavailable("missing scripted connection"))?;
345            Ok(Box::new(ScriptedConnection {
346                frames,
347                sent: self.sent.clone(),
348            }))
349        }
350    }
351
352    #[async_trait]
353    impl WorkerConnection for ScriptedConnection {
354        async fn send(&mut self, command: WorkerCommand) -> Result<(), RuntimeError> {
355            self.sent.lock().unwrap().push(command);
356            Ok(())
357        }
358
359        async fn receive(&mut self) -> Result<Option<WorkerFrame>, RuntimeError> {
360            Ok(self.frames.pop_front())
361        }
362    }
363
364    fn descriptor() -> WorkerDescriptor {
365        WorkerDescriptor {
366            id: WorkerId::new("remote-1"),
367            name: "Remote one".into(),
368            capabilities: RuntimeCapabilities {
369                streaming: true,
370                remote_workers: true,
371                ..RuntimeCapabilities::default()
372            },
373            max_concurrent_runs: 2,
374            labels: BTreeMap::new(),
375        }
376    }
377
378    #[tokio::test]
379    async fn remote_worker_discovers_streams_and_returns_result() {
380        let run_id = RunId::new("remote-run");
381        let transport = ScriptedTransport {
382            connections: Arc::new(Mutex::new(VecDeque::from([
383                VecDeque::from([WorkerFrame::Descriptor {
384                    protocol_version: WORKER_PROTOCOL_VERSION,
385                    descriptor: descriptor(),
386                }]),
387                VecDeque::from([
388                    WorkerFrame::Accepted {
389                        protocol_version: WORKER_PROTOCOL_VERSION,
390                        run_id: run_id.clone(),
391                    },
392                    WorkerFrame::Event {
393                        protocol_version: WORKER_PROTOCOL_VERSION,
394                        envelope: RuntimeEventEnvelope::new(
395                            run_id.to_string(),
396                            0,
397                            RuntimeEvent::TextDelta {
398                                delta: "remote".into(),
399                            },
400                        ),
401                    },
402                    WorkerFrame::Result {
403                        protocol_version: WORKER_PROTOCOL_VERSION,
404                        result: RunResult {
405                            run_id: run_id.clone(),
406                            session_id: SessionId::new("remote-session"),
407                            status: RunStatus::Completed,
408                            final_text: "done".into(),
409                            error: None,
410                            metadata: BTreeMap::new(),
411                        },
412                    },
413                ]),
414            ]))),
415            sent: Arc::new(Mutex::new(Vec::new())),
416        };
417        let worker = RemoteWorker::connect(transport).await.unwrap();
418        assert_eq!(worker.descriptor().id, WorkerId::new("remote-1"));
419        let mut request = RunRequest::new("work remotely");
420        request.run_id = Some(run_id);
421        let mut run = worker
422            .submit(AgentProfile::default(), request)
423            .await
424            .unwrap();
425        let event = run.next_event().await.unwrap();
426        assert!(matches!(event.event, RuntimeEvent::TextDelta { .. }));
427        assert!(matches!(
428            run.next_event().await.unwrap().event,
429            RuntimeEvent::FinalResponse { .. }
430        ));
431        assert_eq!(run.result().await.unwrap().final_text, "done");
432    }
433
434    #[tokio::test]
435    async fn remote_worker_resumes_without_replaying_delivered_events() {
436        let run_id = RunId::new("resumable-run");
437        let sent = Arc::new(Mutex::new(Vec::new()));
438        let result = WorkerFrame::Result {
439            protocol_version: WORKER_PROTOCOL_VERSION,
440            result: RunResult {
441                run_id: run_id.clone(),
442                session_id: SessionId::new("resumed-session"),
443                status: RunStatus::Completed,
444                final_text: "resumed".into(),
445                error: None,
446                metadata: BTreeMap::new(),
447            },
448        };
449        let event = |sequence, delta: &str| WorkerFrame::Event {
450            protocol_version: WORKER_PROTOCOL_VERSION,
451            envelope: RuntimeEventEnvelope::new(
452                run_id.to_string(),
453                sequence,
454                RuntimeEvent::TextDelta {
455                    delta: delta.into(),
456                },
457            ),
458        };
459        let transport = ScriptedTransport {
460            connections: Arc::new(Mutex::new(VecDeque::from([
461                VecDeque::from([WorkerFrame::Descriptor {
462                    protocol_version: WORKER_PROTOCOL_VERSION,
463                    descriptor: descriptor(),
464                }]),
465                VecDeque::from([
466                    WorkerFrame::Accepted {
467                        protocol_version: WORKER_PROTOCOL_VERSION,
468                        run_id: run_id.clone(),
469                    },
470                    event(0, "first"),
471                ]),
472                VecDeque::from([event(0, "duplicate"), event(1, "second"), result]),
473            ]))),
474            sent: sent.clone(),
475        };
476        let worker = RemoteWorker::connect(transport).await.unwrap();
477        let mut request = RunRequest::new("resume remotely");
478        request.run_id = Some(run_id.clone());
479        let mut run = worker
480            .submit(AgentProfile::default(), request)
481            .await
482            .unwrap();
483
484        let mut deltas = Vec::new();
485        while let Some(envelope) = run.next_event().await {
486            if let RuntimeEvent::TextDelta { delta } = envelope.event {
487                deltas.push(delta);
488            }
489        }
490        assert_eq!(deltas, ["first", "second"]);
491        assert_eq!(run.result().await.unwrap().final_text, "resumed");
492        assert!(sent.lock().unwrap().iter().any(|command| matches!(
493            command,
494            WorkerCommand::ResumeEvents {
495                run_id: resumed,
496                after_sequence: Some(0),
497                ..
498            } if resumed == &run_id
499        )));
500    }
501}
microclaw_engine/internal/runtime/remote.rs

microclaw_engine/internal/runtime/
remote.rs
