microclaw_engine/internal/
worker.rs

1//! Authenticated WebSocket transport and host for MicroClaw Workers.
2
3use std::collections::{HashMap, VecDeque};
4use std::sync::{Arc, Mutex};
5
6use crate::internal::runtime::{RunController, Worker, WorkerConnection, WorkerTransport};
7use async_trait::async_trait;
8use futures_util::{SinkExt, StreamExt};
9use http::header::AUTHORIZATION;
10use microclaw_core::run_protocol::{
11    RuntimeError, RuntimeErrorCode, WorkerCommand, WorkerFrame, WORKER_PROTOCOL_VERSION,
12};
13use tokio::net::{TcpListener, TcpStream};
14use tokio_tungstenite::tungstenite::client::IntoClientRequest;
15use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
16use tokio_tungstenite::tungstenite::http::StatusCode;
17use tokio_tungstenite::tungstenite::{Error as WebSocketError, Message};
18use tokio_tungstenite::{accept_hdr_async, connect_async, MaybeTlsStream, WebSocketStream};
19
20/// Creates authenticated WebSocket connections to a MicroClaw Worker host.
21#[derive(Debug, Clone)]
22pub struct WebSocketWorkerTransport {
23    url: String,
24    bearer_token: Option<String>,
25}
26
27impl WebSocketWorkerTransport {
28    pub fn new(url: impl Into<String>) -> Self {
29        Self {
30            url: url.into(),
31            bearer_token: None,
32        }
33    }
34
35    pub fn bearer_token(mut self, token: impl Into<String>) -> Self {
36        self.bearer_token = Some(token.into());
37        self
38    }
39}
40
41struct WebSocketConnection {
42    socket: WebSocketStream<MaybeTlsStream<TcpStream>>,
43}
44
45#[async_trait]
46impl WorkerTransport for WebSocketWorkerTransport {
47    async fn connect(&self) -> Result<Box<dyn WorkerConnection>, RuntimeError> {
48        let mut request = self
49            .url
50            .as_str()
51            .into_client_request()
52            .map_err(configuration_error)?;
53        if let Some(token) = &self.bearer_token {
54            let value = format!("Bearer {token}")
55                .parse()
56                .map_err(|_| configuration_error("invalid Worker bearer token"))?;
57            request.headers_mut().insert(AUTHORIZATION, value);
58        }
59        let (socket, _) = connect_async(request).await.map_err(connection_error)?;
60        Ok(Box::new(WebSocketConnection { socket }))
61    }
62}
63
64#[async_trait]
65impl WorkerConnection for WebSocketConnection {
66    async fn send(&mut self, command: WorkerCommand) -> Result<(), RuntimeError> {
67        let payload = serde_json::to_string(&command).map_err(protocol_error)?;
68        self.socket
69            .send(Message::Text(payload))
70            .await
71            .map_err(connection_error)
72    }
73
74    async fn receive(&mut self) -> Result<Option<WorkerFrame>, RuntimeError> {
75        while let Some(message) = self.socket.next().await {
76            match message.map_err(connection_error)? {
77                Message::Text(payload) => {
78                    return serde_json::from_str(&payload)
79                        .map(Some)
80                        .map_err(protocol_error);
81                }
82                Message::Binary(payload) => {
83                    return serde_json::from_slice(&payload)
84                        .map(Some)
85                        .map_err(protocol_error);
86                }
87                Message::Ping(payload) => self
88                    .socket
89                    .send(Message::Pong(payload))
90                    .await
91                    .map_err(connection_error)?,
92                Message::Close(_) => return Ok(None),
93                Message::Pong(_) | Message::Frame(_) => {}
94            }
95        }
96        Ok(None)
97    }
98}
99
100/// Authenticated WebSocket host for a Local or custom Worker implementation.
101pub struct WorkerHost {
102    worker: Arc<dyn Worker>,
103    bearer_token: Option<String>,
104    runs: Arc<tokio::sync::Mutex<HostedRunRegistry>>,
105    max_retained_runs: usize,
106    max_events_per_run: usize,
107}
108
109#[derive(Default)]
110struct HostedRunRegistry {
111    runs: HashMap<String, Arc<HostedRun>>,
112    completed: VecDeque<String>,
113}
114
115struct HostedRun {
116    request_fingerprint: String,
117    controller: RunController,
118    state: Mutex<HostedRunState>,
119    changed: tokio::sync::Notify,
120}
121
122#[derive(Default)]
123struct HostedRunState {
124    events: VecDeque<microclaw_core::runtime_event::RuntimeEventEnvelope>,
125    discarded_through: Option<u64>,
126    result: Option<microclaw_core::run_protocol::RunResult>,
127}
128
129impl WorkerHost {
130    pub fn new(worker: impl Worker + 'static) -> Self {
131        Self {
132            worker: Arc::new(worker),
133            bearer_token: None,
134            runs: Arc::new(tokio::sync::Mutex::new(HostedRunRegistry::default())),
135            max_retained_runs: 256,
136            max_events_per_run: 4096,
137        }
138    }
139
140    pub fn shared(worker: Arc<dyn Worker>) -> Self {
141        Self {
142            worker,
143            bearer_token: None,
144            runs: Arc::new(tokio::sync::Mutex::new(HostedRunRegistry::default())),
145            max_retained_runs: 256,
146            max_events_per_run: 4096,
147        }
148    }
149
150    pub fn bearer_token(mut self, token: impl Into<String>) -> Self {
151        self.bearer_token = Some(token.into());
152        self
153    }
154
155    pub fn replay_limits(mut self, retained_runs: usize, events_per_run: usize) -> Self {
156        self.max_retained_runs = retained_runs.max(1);
157        self.max_events_per_run = events_per_run.max(1);
158        self
159    }
160
161    /// Serve until the listener is closed or the task is cancelled.
162    pub async fn serve(self, listener: TcpListener) -> Result<(), RuntimeError> {
163        let worker = self.worker;
164        let runs = self.runs;
165        let max_retained_runs = self.max_retained_runs;
166        let max_events_per_run = self.max_events_per_run;
167        let bearer_token = self.bearer_token.map(|token| format!("Bearer {token}"));
168        loop {
169            let (stream, _) = listener.accept().await.map_err(connection_error)?;
170            let worker = worker.clone();
171            let runs = runs.clone();
172            let bearer_token = bearer_token.clone();
173            tokio::spawn(async move {
174                let _ = serve_connection(
175                    stream,
176                    worker,
177                    runs,
178                    bearer_token,
179                    max_retained_runs,
180                    max_events_per_run,
181                )
182                .await;
183            });
184        }
185    }
186}
187
188#[allow(clippy::result_large_err)]
189async fn serve_connection(
190    stream: TcpStream,
191    worker: Arc<dyn Worker>,
192    runs: Arc<tokio::sync::Mutex<HostedRunRegistry>>,
193    bearer_token: Option<String>,
194    max_retained_runs: usize,
195    max_events_per_run: usize,
196) -> Result<(), RuntimeError> {
197    let socket = accept_hdr_async(stream, move |request: &Request, mut response: Response| {
198        let authorized = bearer_token.as_ref().is_none_or(|expected| {
199            request
200                .headers()
201                .get(AUTHORIZATION)
202                .and_then(|value| value.to_str().ok())
203                .is_some_and(|actual| constant_time_eq(actual.as_bytes(), expected.as_bytes()))
204        });
205        if authorized {
206            Ok(response)
207        } else {
208            *response.status_mut() = StatusCode::UNAUTHORIZED;
209            Err(response.map(|_| Some("unauthorized".into())))
210        }
211    })
212    .await
213    .map_err(connection_error)?;
214    let (mut sink, mut stream) = socket.split();
215
216    while let Some(message) = stream.next().await {
217        let command = decode_command(message.map_err(connection_error)?)?;
218        command.validate_protocol()?;
219        match command {
220            WorkerCommand::Describe { .. } => {
221                send_frame(
222                    &mut sink,
223                    WorkerFrame::Descriptor {
224                        protocol_version: WORKER_PROTOCOL_VERSION,
225                        descriptor: worker.descriptor(),
226                    },
227                )
228                .await?;
229            }
230            WorkerCommand::Health { .. } => {
231                send_frame(
232                    &mut sink,
233                    WorkerFrame::Health {
234                        protocol_version: WORKER_PROTOCOL_VERSION,
235                        health: worker.health(),
236                    },
237                )
238                .await?;
239            }
240            WorkerCommand::Submit {
241                profile, request, ..
242            } => {
243                let run_id = request
244                    .run_id
245                    .clone()
246                    .ok_or_else(|| protocol_error("remote submission requires run_id"))?;
247                let request_fingerprint =
248                    serde_json::to_string(&(profile.clone(), &request)).map_err(protocol_error)?;
249                let mut registry = runs.lock().await;
250                let hosted = if let Some(existing) = registry.runs.get(run_id.as_str()).cloned() {
251                    if existing.request_fingerprint != request_fingerprint {
252                        send_frame(
253                            &mut sink,
254                            WorkerFrame::Error {
255                                protocol_version: WORKER_PROTOCOL_VERSION,
256                                run_id: Some(run_id),
257                                error: protocol_error(
258                                    "run_id was already submitted with a different request",
259                                ),
260                            },
261                        )
262                        .await?;
263                        continue;
264                    }
265                    existing
266                } else {
267                    let run = worker.submit(profile, *request).await?;
268                    let hosted = Arc::new(HostedRun {
269                        request_fingerprint,
270                        controller: run.controller(),
271                        state: Mutex::new(HostedRunState::default()),
272                        changed: tokio::sync::Notify::new(),
273                    });
274                    registry.runs.insert(run_id.to_string(), hosted.clone());
275                    tokio::spawn(pump_run(
276                        run,
277                        hosted.clone(),
278                        runs.clone(),
279                        max_retained_runs,
280                        max_events_per_run,
281                    ));
282                    hosted
283                };
284                drop(registry);
285                send_frame(
286                    &mut sink,
287                    WorkerFrame::Accepted {
288                        protocol_version: WORKER_PROTOCOL_VERSION,
289                        run_id: run_id.clone(),
290                    },
291                )
292                .await?;
293                stream_hosted_run(&mut sink, &mut stream, hosted, run_id, None).await?;
294                return Ok(());
295            }
296            WorkerCommand::Control { .. } => {
297                return Err(protocol_error("control command has no active run"));
298            }
299            WorkerCommand::ResumeEvents {
300                run_id,
301                after_sequence,
302                ..
303            } => {
304                let hosted = runs.lock().await.runs.get(run_id.as_str()).cloned();
305                let Some(hosted) = hosted else {
306                    send_frame(
307                        &mut sink,
308                        WorkerFrame::Error {
309                            protocol_version: WORKER_PROTOCOL_VERSION,
310                            run_id: Some(run_id),
311                            error: RuntimeError {
312                                code: RuntimeErrorCode::Unavailable,
313                                message: "run is not known by this Worker host".into(),
314                                retryable: true,
315                            },
316                        },
317                    )
318                    .await?;
319                    continue;
320                };
321                stream_hosted_run(&mut sink, &mut stream, hosted, run_id, after_sequence).await?;
322                return Ok(());
323            }
324        }
325    }
326    Ok(())
327}
328
329async fn pump_run(
330    mut run: crate::internal::runtime::RunHandle,
331    hosted: Arc<HostedRun>,
332    registry: Arc<tokio::sync::Mutex<HostedRunRegistry>>,
333    max_retained_runs: usize,
334    max_events_per_run: usize,
335) {
336    while let Some(envelope) = run.next_event().await {
337        let mut state = hosted
338            .state
339            .lock()
340            .unwrap_or_else(|error| error.into_inner());
341        state.events.push_back(envelope);
342        while state.events.len() > max_events_per_run {
343            state.discarded_through = state.events.pop_front().map(|event| event.sequence);
344        }
345        drop(state);
346        hosted.changed.notify_waiters();
347    }
348    let result =
349        run.result()
350            .await
351            .unwrap_or_else(|error| microclaw_core::run_protocol::RunResult {
352                run_id: hosted.controller.id().clone(),
353                session_id: microclaw_core::run_protocol::SessionId::new(""),
354                status: microclaw_core::run_protocol::RunStatus::Failed,
355                final_text: String::new(),
356                error: Some(error),
357                metadata: Default::default(),
358            });
359    hosted
360        .state
361        .lock()
362        .unwrap_or_else(|error| error.into_inner())
363        .result = Some(result);
364    hosted.changed.notify_waiters();
365    let run_id = hosted.controller.id().to_string();
366    let mut registry = registry.lock().await;
367    registry.completed.push_back(run_id);
368    while registry.completed.len() > max_retained_runs {
369        if let Some(expired) = registry.completed.pop_front() {
370            registry.runs.remove(&expired);
371        }
372    }
373}
374
375async fn stream_hosted_run<S, R>(
376    sink: &mut S,
377    stream: &mut R,
378    hosted: Arc<HostedRun>,
379    run_id: microclaw_core::run_protocol::RunId,
380    after_sequence: Option<u64>,
381) -> Result<(), RuntimeError>
382where
383    S: futures_util::Sink<Message, Error = WebSocketError> + Unpin,
384    R: futures_util::Stream<Item = Result<Message, WebSocketError>> + Unpin,
385{
386    let mut cursor = after_sequence;
387    loop {
388        let (events, result, history_lost) = {
389            let state = hosted
390                .state
391                .lock()
392                .unwrap_or_else(|error| error.into_inner());
393            (
394                state
395                    .events
396                    .iter()
397                    .filter(|event| cursor.is_none_or(|value| event.sequence > value))
398                    .cloned()
399                    .collect::<Vec<_>>(),
400                state.result.clone(),
401                state
402                    .discarded_through
403                    .is_some_and(|discarded| cursor.is_none_or(|sequence| sequence < discarded)),
404            )
405        };
406        if history_lost {
407            send_frame(
408                sink,
409                WorkerFrame::Error {
410                    protocol_version: WORKER_PROTOCOL_VERSION,
411                    run_id: Some(run_id),
412                    error: RuntimeError {
413                        code: RuntimeErrorCode::Unavailable,
414                        message: "requested Worker event history has expired".into(),
415                        retryable: true,
416                    },
417                },
418            )
419            .await?;
420            return Ok(());
421        }
422        for envelope in events {
423            cursor = Some(envelope.sequence);
424            send_frame(
425                sink,
426                WorkerFrame::Event {
427                    protocol_version: WORKER_PROTOCOL_VERSION,
428                    envelope,
429                },
430            )
431            .await?;
432        }
433        if let Some(result) = result {
434            send_frame(
435                sink,
436                WorkerFrame::Result {
437                    protocol_version: WORKER_PROTOCOL_VERSION,
438                    result,
439                },
440            )
441            .await?;
442            return Ok(());
443        }
444        tokio::select! {
445            _ = hosted.changed.notified() => {}
446            message = stream.next() => {
447                let Some(message) = message else { return Err(connection_error("Worker client disconnected")); };
448                let command = decode_command(message.map_err(connection_error)?)?;
449                command.validate_protocol()?;
450                let WorkerCommand::Control { control, .. } = command else {
451                    send_frame(sink, WorkerFrame::Error { protocol_version: WORKER_PROTOCOL_VERSION, run_id: Some(run_id.clone()), error: protocol_error("only control commands are accepted during a run") }).await?;
452                    continue;
453                };
454                let outcome = match control {
455                    microclaw_core::run_protocol::RuntimeControl::Cancel { run_id: target } if target == run_id => hosted.controller.cancel_confirmed().await,
456                    microclaw_core::run_protocol::RuntimeControl::Steer { run_id: target, message } if target == run_id => hosted.controller.steer_confirmed(message).await,
457                    microclaw_core::run_protocol::RuntimeControl::ResolveApproval { run_id: target, approval_id, decision } if target == run_id => hosted.controller.resolve_approval_confirmed(approval_id, decision).await,
458                    _ => Err(protocol_error("control command targets a different run")),
459                };
460                match outcome {
461                    Ok(()) => send_frame(sink, WorkerFrame::ControlAcknowledged { protocol_version: WORKER_PROTOCOL_VERSION, run_id: run_id.clone() }).await?,
462                    Err(error) => send_frame(sink, WorkerFrame::Error { protocol_version: WORKER_PROTOCOL_VERSION, run_id: Some(run_id.clone()), error }).await?,
463                }
464            }
465        }
466    }
467}
468
469fn decode_command(message: Message) -> Result<WorkerCommand, RuntimeError> {
470    match message {
471        Message::Text(payload) => serde_json::from_str(&payload).map_err(protocol_error),
472        Message::Binary(payload) => serde_json::from_slice(&payload).map_err(protocol_error),
473        _ => Err(protocol_error("expected a JSON Worker command")),
474    }
475}
476
477async fn send_frame<S>(sink: &mut S, frame: WorkerFrame) -> Result<(), RuntimeError>
478where
479    S: futures_util::Sink<Message, Error = WebSocketError> + Unpin,
480{
481    let payload = serde_json::to_string(&frame).map_err(protocol_error)?;
482    sink.send(Message::Text(payload))
483        .await
484        .map_err(connection_error)
485}
486
487fn constant_time_eq(left: &[u8], right: &[u8]) -> bool {
488    if left.len() != right.len() {
489        return false;
490    }
491    left.iter()
492        .zip(right)
493        .fold(0_u8, |diff, (left, right)| diff | (left ^ right))
494        == 0
495}
496
497fn configuration_error(error: impl std::fmt::Display) -> RuntimeError {
498    RuntimeError {
499        code: RuntimeErrorCode::Configuration,
500        message: error.to_string(),
501        retryable: false,
502    }
503}
504
505fn protocol_error(error: impl std::fmt::Display) -> RuntimeError {
506    RuntimeError {
507        code: RuntimeErrorCode::InvalidRequest,
508        message: error.to_string(),
509        retryable: false,
510    }
511}
512
513fn connection_error(error: impl std::fmt::Display) -> RuntimeError {
514    RuntimeError {
515        code: RuntimeErrorCode::Unavailable,
516        message: error.to_string(),
517        retryable: true,
518    }
519}
520
521#[cfg(test)]
522mod tests {
523    use crate::internal::runtime::{
524        ExecutionContext, ExecutionResult, LocalWorker, RemoteWorker, RunExecutor, Runtime,
525    };
526    use microclaw_core::run_protocol::{
527        AgentProfile, RunRequest, RuntimeCapabilities, SessionId, WorkerHealthStatus,
528    };
529    use microclaw_core::runtime_event::RuntimeEvent;
530    use std::sync::atomic::{AtomicUsize, Ordering};
531
532    use super::*;
533
534    struct Echo;
535
536    #[async_trait]
537    impl RunExecutor for Echo {
538        async fn execute(
539            &self,
540            _profile: AgentProfile,
541            request: RunRequest,
542            context: ExecutionContext,
543        ) -> Result<ExecutionResult, RuntimeError> {
544            context.emit(RuntimeEvent::TextDelta {
545                delta: "from worker".into(),
546            })?;
547            Ok(ExecutionResult::new(
548                request
549                    .session_id
550                    .unwrap_or_else(|| SessionId::new("remote-session")),
551                request.prompt,
552            ))
553        }
554
555        fn capabilities(&self) -> RuntimeCapabilities {
556            RuntimeCapabilities {
557                streaming: true,
558                cancellation: true,
559                remote_workers: true,
560                ..RuntimeCapabilities::default()
561            }
562        }
563    }
564
565    struct CountingEcho(Arc<AtomicUsize>);
566
567    #[async_trait]
568    impl RunExecutor for CountingEcho {
569        async fn execute(
570            &self,
571            _profile: AgentProfile,
572            request: RunRequest,
573            context: ExecutionContext,
574        ) -> Result<ExecutionResult, RuntimeError> {
575            self.0.fetch_add(1, Ordering::SeqCst);
576            tokio::time::sleep(std::time::Duration::from_millis(30)).await;
577            context.emit(RuntimeEvent::TextDelta {
578                delta: "persisted".into(),
579            })?;
580            Ok(ExecutionResult::new(
581                SessionId::new("resumed-session"),
582                request.prompt,
583            ))
584        }
585    }
586
587    async fn start_host(token: &str) -> (String, tokio::task::JoinHandle<()>) {
588        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
589        let address = listener.local_addr().unwrap();
590        let runtime = Runtime::builder().executor(Echo).build().unwrap();
591        let worker = LocalWorker::new(runtime, "worker-1", "Test worker");
592        let host = WorkerHost::new(worker).bearer_token(token);
593        let task = tokio::spawn(async move {
594            let _ = host.serve(listener).await;
595        });
596        (format!("ws://{address}"), task)
597    }
598
599    #[tokio::test]
600    async fn authenticated_websocket_worker_runs_and_reports_health() {
601        let (url, host) = start_host("secret").await;
602        let transport = WebSocketWorkerTransport::new(url).bearer_token("secret");
603        let worker = RemoteWorker::connect(transport).await.unwrap();
604        assert_eq!(worker.descriptor().name, "Test worker");
605
606        let health = worker.refresh_health().await.unwrap();
607        assert_eq!(health.status, WorkerHealthStatus::Ready);
608        let mut run = worker
609            .submit(AgentProfile::default(), RunRequest::new("hello remote"))
610            .await
611            .unwrap();
612        assert!(matches!(
613            run.next_event().await.unwrap().event,
614            RuntimeEvent::TextDelta { .. }
615        ));
616        assert!(matches!(
617            run.next_event().await.unwrap().event,
618            RuntimeEvent::FinalResponse { .. }
619        ));
620        assert_eq!(run.result().await.unwrap().final_text, "hello remote");
621        host.abort();
622    }
623
624    #[tokio::test]
625    async fn websocket_worker_rejects_invalid_credentials() {
626        let (url, host) = start_host("secret").await;
627        let transport = WebSocketWorkerTransport::new(url).bearer_token("wrong");
628        let error = match RemoteWorker::connect(transport).await {
629            Ok(_) => panic!("invalid Worker credentials should fail"),
630            Err(error) => error,
631        };
632        assert_eq!(error.code, RuntimeErrorCode::Unavailable);
633        host.abort();
634    }
635
636    #[tokio::test]
637    async fn websocket_worker_keeps_run_alive_and_resumes_after_disconnect() {
638        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
639        let address = listener.local_addr().unwrap();
640        let executions = Arc::new(AtomicUsize::new(0));
641        let runtime = Runtime::builder()
642            .executor(CountingEcho(executions.clone()))
643            .build()
644            .unwrap();
645        let host = WorkerHost::new(LocalWorker::new(runtime, "worker-1", "Test worker"));
646        let host_task = tokio::spawn(async move {
647            let _ = host.serve(listener).await;
648        });
649        let transport = WebSocketWorkerTransport::new(format!("ws://{address}"));
650        let run_id = microclaw_core::run_protocol::RunId::new("durable-run");
651        let mut first = transport.connect().await.unwrap();
652        let mut request = RunRequest::new("survive disconnect");
653        request.run_id = Some(run_id.clone());
654        first
655            .send(WorkerCommand::Submit {
656                protocol_version: WORKER_PROTOCOL_VERSION,
657                profile: AgentProfile::default(),
658                request: Box::new(request),
659            })
660            .await
661            .unwrap();
662        assert!(matches!(
663            first.receive().await.unwrap(),
664            Some(WorkerFrame::Accepted { .. })
665        ));
666        drop(first);
667
668        tokio::time::sleep(std::time::Duration::from_millis(60)).await;
669        let mut resumed = transport.connect().await.unwrap();
670        resumed
671            .send(WorkerCommand::ResumeEvents {
672                protocol_version: WORKER_PROTOCOL_VERSION,
673                run_id: run_id.clone(),
674                after_sequence: None,
675            })
676            .await
677            .unwrap();
678        let result = loop {
679            match resumed.receive().await.unwrap() {
680                Some(WorkerFrame::Event { .. }) => {}
681                Some(WorkerFrame::Result { result, .. }) => break result,
682                frame => panic!("unexpected resumed frame: {frame:?}"),
683            }
684        };
685        assert_eq!(result.run_id, run_id);
686        assert_eq!(executions.load(Ordering::SeqCst), 1);
687        host_task.abort();
688    }
689}
microclaw_engine/internal/worker.rs

microclaw_engine/internal/
worker.rs
