microclaw_engine/internal/runtime/
remote.rs

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

microclaw_engine/internal/runtime/
remote.rs
