microclaw_engine/internal/runtime/
mod.rs

1//! Embeddable, UI-independent MicroClaw run lifecycle.
2//!
3//! Product adapters supply a [`RunExecutor`]. The runtime owns stable Agent and Run handles,
4//! concurrency, control routing, terminal results, and the provider-neutral event stream.
5
6use std::collections::BTreeMap;
7use std::sync::atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering};
8use std::sync::Arc;
9
10use async_trait::async_trait;
11use microclaw_core::run_protocol::{
12    AgentProfile, RunId, RunRequest, RunResult, RunStatus, RuntimeCapabilities, RuntimeControl,
13    RuntimeError, RuntimeErrorCode, SessionId, WorkerDescriptor, WorkerHealth, WorkerHealthStatus,
14    WorkerId,
15};
16use microclaw_core::runtime_event::{RuntimeEvent, RuntimeEventEnvelope};
17use tokio::sync::{mpsc, oneshot, Notify, Semaphore};
18
19mod remote;
20
21pub use remote::{RemoteWorker, RemoteWorkerOptions, WorkerConnection, WorkerTransport};
22
23struct PendingControl {
24    control: RuntimeControl,
25    receipt: Option<oneshot::Sender<Result<(), RuntimeError>>>,
26}
27
28/// A control message received by an executor that can be explicitly accepted or rejected.
29///
30/// Executors should acknowledge only after the underlying engine has accepted the operation.
31/// Dropping a request without responding reports an unavailable control path to the caller.
32pub struct ControlRequest {
33    control: RuntimeControl,
34    receipt: Option<oneshot::Sender<Result<(), RuntimeError>>>,
35}
36
37impl ControlRequest {
38    pub fn control(&self) -> &RuntimeControl {
39        &self.control
40    }
41
42    pub fn into_control(self) -> RuntimeControl {
43        let control = self.control.clone();
44        self.accept();
45        control
46    }
47
48    pub fn accept(mut self) {
49        if let Some(receipt) = self.receipt.take() {
50            let _ = receipt.send(Ok(()));
51        }
52    }
53
54    pub fn reject(mut self, error: RuntimeError) {
55        if let Some(receipt) = self.receipt.take() {
56            let _ = receipt.send(Err(error));
57        }
58    }
59}
60
61impl Drop for ControlRequest {
62    fn drop(&mut self) {
63        if let Some(receipt) = self.receipt.take() {
64            let _ = receipt.send(Err(RuntimeError {
65                code: RuntimeErrorCode::Unavailable,
66                message: "executor did not acknowledge the control request".into(),
67                retryable: false,
68            }));
69        }
70    }
71}
72
73#[derive(Debug, Clone, PartialEq)]
74pub struct ExecutionResult {
75    pub session_id: SessionId,
76    pub final_text: String,
77    pub metadata: BTreeMap<String, serde_json::Value>,
78}
79
80impl ExecutionResult {
81    pub fn new(session_id: impl Into<SessionId>, final_text: impl Into<String>) -> Self {
82        Self {
83            session_id: session_id.into(),
84            final_text: final_text.into(),
85            metadata: BTreeMap::new(),
86        }
87    }
88}
89
90/// Runtime-owned context passed to an execution engine for one run.
91pub struct ExecutionContext {
92    run_id: RunId,
93    event_tx: mpsc::UnboundedSender<RuntimeEventEnvelope>,
94    control_rx: mpsc::UnboundedReceiver<PendingControl>,
95    next_sequence: Arc<AtomicU64>,
96}
97
98impl ExecutionContext {
99    pub fn run_id(&self) -> &RunId {
100        &self.run_id
101    }
102
103    pub fn emit(&self, event: RuntimeEvent) -> Result<(), RuntimeError> {
104        let sequence = self.next_sequence.fetch_add(1, Ordering::Relaxed);
105        self.event_tx
106            .send(RuntimeEventEnvelope::new(
107                self.run_id.to_string(),
108                sequence,
109                event,
110            ))
111            .map_err(|_| RuntimeError {
112                code: RuntimeErrorCode::Unavailable,
113                message: "runtime event consumer disconnected".into(),
114                retryable: false,
115            })
116    }
117
118    pub fn try_next_control(&mut self) -> Option<RuntimeControl> {
119        self.try_next_control_request()
120            .map(ControlRequest::into_control)
121    }
122
123    pub async fn next_control(&mut self) -> Option<RuntimeControl> {
124        self.next_control_request()
125            .await
126            .map(ControlRequest::into_control)
127    }
128
129    pub fn try_next_control_request(&mut self) -> Option<ControlRequest> {
130        self.control_rx
131            .try_recv()
132            .ok()
133            .map(|pending| ControlRequest {
134                control: pending.control,
135                receipt: pending.receipt,
136            })
137    }
138
139    pub async fn next_control_request(&mut self) -> Option<ControlRequest> {
140        self.control_rx.recv().await.map(|pending| ControlRequest {
141            control: pending.control,
142            receipt: pending.receipt,
143        })
144    }
145}
146
147/// Product-independent execution port used by the Runtime.
148///
149/// MicroClaw's shared Agent Engine implements this port. Tests and embedding applications may
150/// provide another implementation without depending on Server or UI code.
151#[async_trait]
152pub trait RunExecutor: Send + Sync + 'static {
153    async fn execute(
154        &self,
155        profile: AgentProfile,
156        request: RunRequest,
157        context: ExecutionContext,
158    ) -> Result<ExecutionResult, RuntimeError>;
159
160    fn capabilities(&self) -> RuntimeCapabilities {
161        RuntimeCapabilities::default()
162    }
163}
164
165#[derive(Debug, Clone, PartialEq, Eq)]
166pub enum RuntimeBuildError {
167    MissingExecutor,
168    InvalidConcurrency,
169}
170
171impl std::fmt::Display for RuntimeBuildError {
172    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        match self {
174            Self::MissingExecutor => f.write_str("a run executor is required"),
175            Self::InvalidConcurrency => f.write_str("max concurrent runs must be greater than 0"),
176        }
177    }
178}
179
180impl std::error::Error for RuntimeBuildError {}
181
182pub struct RuntimeBuilder {
183    executor: Option<Arc<dyn RunExecutor>>,
184    max_concurrent_runs: usize,
185}
186
187impl Default for RuntimeBuilder {
188    fn default() -> Self {
189        Self {
190            executor: None,
191            max_concurrent_runs: 4,
192        }
193    }
194}
195
196impl RuntimeBuilder {
197    pub fn executor(mut self, executor: impl RunExecutor) -> Self {
198        self.executor = Some(Arc::new(executor));
199        self
200    }
201
202    pub fn shared_executor(mut self, executor: Arc<dyn RunExecutor>) -> Self {
203        self.executor = Some(executor);
204        self
205    }
206
207    pub fn max_concurrent_runs(mut self, max: usize) -> Self {
208        self.max_concurrent_runs = max;
209        self
210    }
211
212    pub fn build(self) -> Result<Runtime, RuntimeBuildError> {
213        if self.max_concurrent_runs == 0 {
214            return Err(RuntimeBuildError::InvalidConcurrency);
215        }
216        let executor = self.executor.ok_or(RuntimeBuildError::MissingExecutor)?;
217        Ok(Runtime {
218            inner: Arc::new(RuntimeInner {
219                executor,
220                permits: Arc::new(Semaphore::new(self.max_concurrent_runs)),
221                max_concurrent_runs: self.max_concurrent_runs,
222                active_runs: Arc::new(AtomicUsize::new(0)),
223                queued_runs: AtomicUsize::new(0),
224                idle: Notify::new(),
225            }),
226        })
227    }
228}
229
230struct RuntimeInner {
231    executor: Arc<dyn RunExecutor>,
232    permits: Arc<Semaphore>,
233    max_concurrent_runs: usize,
234    active_runs: Arc<AtomicUsize>,
235    queued_runs: AtomicUsize,
236    idle: Notify,
237}
238
239#[derive(Debug, Clone, Copy, PartialEq, Eq)]
240pub struct RuntimeStats {
241    pub active_runs: usize,
242    pub queued_runs: usize,
243    pub max_concurrent_runs: usize,
244    pub shutting_down: bool,
245}
246
247#[derive(Clone)]
248pub struct Runtime {
249    inner: Arc<RuntimeInner>,
250}
251
252impl Runtime {
253    pub fn builder() -> RuntimeBuilder {
254        RuntimeBuilder::default()
255    }
256
257    pub fn agent(&self, profile: AgentProfile) -> AgentHandle {
258        AgentHandle {
259            runtime: self.clone(),
260            profile,
261        }
262    }
263
264    pub fn capabilities(&self) -> RuntimeCapabilities {
265        self.inner.executor.capabilities()
266    }
267
268    pub fn max_concurrent_runs(&self) -> usize {
269        self.inner.max_concurrent_runs
270    }
271
272    pub fn active_runs(&self) -> usize {
273        self.inner.active_runs.load(Ordering::Acquire)
274    }
275
276    pub fn queued_runs(&self) -> usize {
277        self.inner.queued_runs.load(Ordering::Acquire)
278    }
279
280    pub fn stats(&self) -> RuntimeStats {
281        RuntimeStats {
282            active_runs: self.active_runs(),
283            queued_runs: self.queued_runs(),
284            max_concurrent_runs: self.max_concurrent_runs(),
285            shutting_down: self.is_shutting_down(),
286        }
287    }
288
289    /// Stop accepting new runs and wait until active and queued work has drained.
290    pub async fn shutdown(&self) {
291        self.inner.permits.close();
292        self.wait_for_idle().await;
293    }
294
295    /// Returns true after graceful shutdown has begun.
296    pub fn is_shutting_down(&self) -> bool {
297        self.inner.permits.is_closed()
298    }
299
300    pub async fn wait_for_idle(&self) {
301        loop {
302            let notified = self.inner.idle.notified();
303            if self.active_runs() == 0 && self.queued_runs() == 0 {
304                return;
305            }
306            notified.await;
307        }
308    }
309}
310
311#[derive(Clone)]
312pub struct AgentHandle {
313    runtime: Runtime,
314    profile: AgentProfile,
315}
316
317impl AgentHandle {
318    pub fn profile(&self) -> &AgentProfile {
319        &self.profile
320    }
321
322    pub fn run(&self, request: RunRequest) -> RunHandle {
323        let run_id = request
324            .run_id
325            .clone()
326            .unwrap_or_else(|| RunId::new(uuid::Uuid::new_v4().to_string()));
327        let (event_tx, event_rx) = mpsc::unbounded_channel();
328        let (control_tx, control_rx) = mpsc::unbounded_channel();
329        let (result_tx, result_rx) = oneshot::channel();
330        let next_sequence = Arc::new(AtomicU64::new(0));
331        let executor = self.runtime.inner.executor.clone();
332        let permits = self.runtime.inner.permits.clone();
333        let active_runs = self.runtime.inner.active_runs.clone();
334        let runtime_inner = self.runtime.inner.clone();
335        let profile = self.profile.clone();
336        let task_run_id = run_id.clone();
337        let fallback_session_id = request
338            .session_id
339            .clone()
340            .unwrap_or_else(|| SessionId::new(task_run_id.to_string()));
341
342        runtime_inner.queued_runs.fetch_add(1, Ordering::AcqRel);
343        tokio::spawn(async move {
344            let permit = match permits.acquire_owned().await {
345                Ok(permit) => permit,
346                Err(_) => {
347                    runtime_inner.queued_runs.fetch_sub(1, Ordering::AcqRel);
348                    runtime_inner.idle.notify_waiters();
349                    let _ = result_tx.send(RunResult {
350                        run_id: task_run_id,
351                        session_id: fallback_session_id,
352                        status: RunStatus::Failed,
353                        final_text: String::new(),
354                        error: Some(RuntimeError {
355                            code: RuntimeErrorCode::Unavailable,
356                            message: "runtime is shutting down".into(),
357                            retryable: true,
358                        }),
359                        metadata: BTreeMap::new(),
360                    });
361                    return;
362                }
363            };
364            runtime_inner.queued_runs.fetch_sub(1, Ordering::AcqRel);
365            let _active_run = ActiveRunGuard::new(active_runs, runtime_inner.clone());
366            let context = ExecutionContext {
367                run_id: task_run_id.clone(),
368                event_tx: event_tx.clone(),
369                control_rx,
370                next_sequence: next_sequence.clone(),
371            };
372            let result = executor.execute(profile, request, context).await;
373            drop(permit);
374
375            let terminal = match result {
376                Ok(output) => {
377                    emit_terminal_event(
378                        &event_tx,
379                        &next_sequence,
380                        &task_run_id,
381                        RuntimeEvent::FinalResponse {
382                            text: output.final_text.clone(),
383                        },
384                    );
385                    RunResult {
386                        run_id: task_run_id,
387                        session_id: output.session_id,
388                        status: RunStatus::Completed,
389                        final_text: output.final_text,
390                        error: None,
391                        metadata: output.metadata,
392                    }
393                }
394                Err(error) => {
395                    let status = match error.code {
396                        RuntimeErrorCode::Cancelled => RunStatus::Cancelled,
397                        RuntimeErrorCode::TimedOut => RunStatus::TimedOut,
398                        _ => RunStatus::Failed,
399                    };
400                    if status == RunStatus::Cancelled {
401                        emit_terminal_event(
402                            &event_tx,
403                            &next_sequence,
404                            &task_run_id,
405                            RuntimeEvent::Cancelled {
406                                final_text: error.message.clone(),
407                            },
408                        );
409                    }
410                    RunResult {
411                        run_id: task_run_id,
412                        session_id: fallback_session_id,
413                        status,
414                        final_text: String::new(),
415                        error: Some(error),
416                        metadata: BTreeMap::new(),
417                    }
418                }
419            };
420            let _ = result_tx.send(terminal);
421        });
422
423        RunHandle {
424            run_id,
425            events: event_rx,
426            control_tx,
427            result_rx,
428        }
429    }
430}
431
432struct ActiveRunGuard {
433    active_runs: Arc<AtomicUsize>,
434    runtime: Arc<RuntimeInner>,
435}
436
437impl ActiveRunGuard {
438    fn new(active_runs: Arc<AtomicUsize>, runtime: Arc<RuntimeInner>) -> Self {
439        active_runs.fetch_add(1, Ordering::AcqRel);
440        Self {
441            active_runs,
442            runtime,
443        }
444    }
445}
446
447impl Drop for ActiveRunGuard {
448    fn drop(&mut self) {
449        self.active_runs.fetch_sub(1, Ordering::AcqRel);
450        self.runtime.idle.notify_waiters();
451    }
452}
453
454fn emit_terminal_event(
455    event_tx: &mpsc::UnboundedSender<RuntimeEventEnvelope>,
456    next_sequence: &AtomicU64,
457    run_id: &RunId,
458    event: RuntimeEvent,
459) {
460    let sequence = next_sequence.fetch_add(1, Ordering::Relaxed);
461    let _ = event_tx.send(RuntimeEventEnvelope::new(
462        run_id.to_string(),
463        sequence,
464        event,
465    ));
466}
467
468pub struct RunHandle {
469    run_id: RunId,
470    events: mpsc::UnboundedReceiver<RuntimeEventEnvelope>,
471    control_tx: mpsc::UnboundedSender<PendingControl>,
472    result_rx: oneshot::Receiver<RunResult>,
473}
474
475impl RunHandle {
476    pub fn id(&self) -> &RunId {
477        &self.run_id
478    }
479
480    pub async fn next_event(&mut self) -> Option<RuntimeEventEnvelope> {
481        self.events.recv().await
482    }
483
484    pub fn try_next_event(&mut self) -> Option<RuntimeEventEnvelope> {
485        self.events.try_recv().ok()
486    }
487
488    pub fn cancel(&self) -> Result<(), RuntimeError> {
489        self.controller().cancel()
490    }
491
492    pub fn controller(&self) -> RunController {
493        RunController {
494            run_id: self.run_id.clone(),
495            control_tx: self.control_tx.clone(),
496        }
497    }
498
499    pub async fn cancel_confirmed(&self) -> Result<(), RuntimeError> {
500        self.controller().cancel_confirmed().await
501    }
502
503    pub fn steer(&self, message: impl Into<String>) -> Result<(), RuntimeError> {
504        self.controller().steer(message)
505    }
506
507    pub async fn steer_confirmed(&self, message: impl Into<String>) -> Result<(), RuntimeError> {
508        self.controller().steer_confirmed(message).await
509    }
510
511    pub fn resolve_approval(
512        &self,
513        approval_id: impl Into<String>,
514        decision: impl Into<String>,
515    ) -> Result<(), RuntimeError> {
516        self.controller().resolve_approval(approval_id, decision)
517    }
518
519    pub async fn resolve_approval_confirmed(
520        &self,
521        approval_id: impl Into<String>,
522        decision: impl Into<String>,
523    ) -> Result<(), RuntimeError> {
524        self.controller()
525            .resolve_approval_confirmed(approval_id, decision)
526            .await
527    }
528
529    pub async fn result(self) -> Result<RunResult, RuntimeError> {
530        self.result_rx.await.map_err(|_| RuntimeError {
531            code: RuntimeErrorCode::Internal,
532            message: "run ended without a terminal result".into(),
533            retryable: false,
534        })
535    }
536}
537
538/// Cloneable control plane for a running task.
539///
540/// Keeping this separate from [`RunHandle`] lets one task consume events while another task or
541/// UI callback controls the run without sharing the event receiver.
542#[derive(Clone)]
543pub struct RunController {
544    run_id: RunId,
545    control_tx: mpsc::UnboundedSender<PendingControl>,
546}
547
548impl RunController {
549    pub fn id(&self) -> &RunId {
550        &self.run_id
551    }
552
553    pub fn cancel(&self) -> Result<(), RuntimeError> {
554        self.send_control(RuntimeControl::Cancel {
555            run_id: self.run_id.clone(),
556        })
557    }
558
559    pub fn steer(&self, message: impl Into<String>) -> Result<(), RuntimeError> {
560        self.send_control(RuntimeControl::Steer {
561            run_id: self.run_id.clone(),
562            message: message.into(),
563        })
564    }
565
566    pub fn resolve_approval(
567        &self,
568        approval_id: impl Into<String>,
569        decision: impl Into<String>,
570    ) -> Result<(), RuntimeError> {
571        self.send_control(RuntimeControl::ResolveApproval {
572            run_id: self.run_id.clone(),
573            approval_id: approval_id.into(),
574            decision: decision.into(),
575        })
576    }
577
578    fn send_control(&self, control: RuntimeControl) -> Result<(), RuntimeError> {
579        self.control_tx
580            .send(PendingControl {
581                control,
582                receipt: None,
583            })
584            .map_err(|_| RuntimeError {
585                code: RuntimeErrorCode::Unavailable,
586                message: "run is no longer accepting control messages".into(),
587                retryable: false,
588            })
589    }
590
591    pub async fn cancel_confirmed(&self) -> Result<(), RuntimeError> {
592        self.send_control_confirmed(RuntimeControl::Cancel {
593            run_id: self.run_id.clone(),
594        })
595        .await
596    }
597
598    pub async fn steer_confirmed(&self, message: impl Into<String>) -> Result<(), RuntimeError> {
599        self.send_control_confirmed(RuntimeControl::Steer {
600            run_id: self.run_id.clone(),
601            message: message.into(),
602        })
603        .await
604    }
605
606    pub async fn resolve_approval_confirmed(
607        &self,
608        approval_id: impl Into<String>,
609        decision: impl Into<String>,
610    ) -> Result<(), RuntimeError> {
611        self.send_control_confirmed(RuntimeControl::ResolveApproval {
612            run_id: self.run_id.clone(),
613            approval_id: approval_id.into(),
614            decision: decision.into(),
615        })
616        .await
617    }
618
619    async fn send_control_confirmed(&self, control: RuntimeControl) -> Result<(), RuntimeError> {
620        let (receipt_tx, receipt_rx) = oneshot::channel();
621        self.control_tx
622            .send(PendingControl {
623                control,
624                receipt: Some(receipt_tx),
625            })
626            .map_err(|_| RuntimeError {
627                code: RuntimeErrorCode::Unavailable,
628                message: "run is no longer accepting control messages".into(),
629                retryable: false,
630            })?;
631        receipt_rx.await.map_err(|_| RuntimeError {
632            code: RuntimeErrorCode::Unavailable,
633            message: "run ended before acknowledging the control request".into(),
634            retryable: false,
635        })?
636    }
637}
638
639#[async_trait]
640pub trait Worker: Send + Sync {
641    fn descriptor(&self) -> WorkerDescriptor;
642    fn health(&self) -> WorkerHealth;
643    async fn submit(
644        &self,
645        profile: AgentProfile,
646        request: RunRequest,
647    ) -> Result<RunHandle, RuntimeError>;
648}
649
650const WORKER_READY: u8 = 0;
651const WORKER_DRAINING: u8 = 1;
652const WORKER_UNAVAILABLE: u8 = 2;
653
654#[derive(Clone)]
655pub struct LocalWorker {
656    descriptor: WorkerDescriptor,
657    runtime: Runtime,
658    state: Arc<AtomicU8>,
659}
660
661impl LocalWorker {
662    pub fn new(runtime: Runtime, id: impl Into<WorkerId>, name: impl Into<String>) -> Self {
663        let max_concurrent_runs = runtime.max_concurrent_runs();
664        Self {
665            descriptor: WorkerDescriptor {
666                id: id.into(),
667                name: name.into(),
668                capabilities: runtime.capabilities(),
669                max_concurrent_runs,
670                labels: BTreeMap::new(),
671            },
672            runtime,
673            state: Arc::new(AtomicU8::new(WORKER_READY)),
674        }
675    }
676
677    pub fn with_label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
678        self.descriptor.labels.insert(key.into(), value.into());
679        self
680    }
681
682    pub fn begin_drain(&self) {
683        self.state.store(WORKER_DRAINING, Ordering::Release);
684    }
685
686    pub fn mark_unavailable(&self) {
687        self.state.store(WORKER_UNAVAILABLE, Ordering::Release);
688    }
689
690    pub fn resume(&self) {
691        self.state.store(WORKER_READY, Ordering::Release);
692    }
693
694    pub async fn wait_for_idle(&self) {
695        self.runtime.wait_for_idle().await;
696    }
697}
698
699#[async_trait]
700impl Worker for LocalWorker {
701    fn descriptor(&self) -> WorkerDescriptor {
702        self.descriptor.clone()
703    }
704
705    fn health(&self) -> WorkerHealth {
706        let active_runs = self.runtime.active_runs();
707        let state = self.state.load(Ordering::Acquire);
708        WorkerHealth {
709            worker_id: self.descriptor.id.clone(),
710            status: match state {
711                WORKER_DRAINING => WorkerHealthStatus::Draining,
712                WORKER_UNAVAILABLE => WorkerHealthStatus::Unavailable,
713                _ if active_runs >= self.descriptor.max_concurrent_runs => WorkerHealthStatus::Busy,
714                _ => WorkerHealthStatus::Ready,
715            },
716            active_runs,
717            queued_runs: self.runtime.queued_runs(),
718            observed_at: std::time::SystemTime::now()
719                .duration_since(std::time::UNIX_EPOCH)
720                .unwrap_or_default()
721                .as_secs()
722                .to_string(),
723        }
724    }
725
726    async fn submit(
727        &self,
728        profile: AgentProfile,
729        request: RunRequest,
730    ) -> Result<RunHandle, RuntimeError> {
731        if self.state.load(Ordering::Acquire) != WORKER_READY {
732            return Err(RuntimeError {
733                code: RuntimeErrorCode::Unavailable,
734                message: "Worker is not accepting new runs".into(),
735                retryable: true,
736            });
737        }
738        Ok(self.runtime.agent(profile).run(request))
739    }
740}
741
742#[cfg(test)]
743mod tests {
744    use std::sync::atomic::{AtomicUsize, Ordering};
745    use std::time::Duration;
746
747    use super::*;
748
749    struct EchoExecutor;
750
751    #[async_trait]
752    impl RunExecutor for EchoExecutor {
753        async fn execute(
754            &self,
755            profile: AgentProfile,
756            request: RunRequest,
757            context: ExecutionContext,
758        ) -> Result<ExecutionResult, RuntimeError> {
759            context.emit(RuntimeEvent::TextDelta {
760                delta: format!("{}: ", profile.name),
761            })?;
762            Ok(ExecutionResult::new(
763                request
764                    .session_id
765                    .clone()
766                    .unwrap_or_else(|| SessionId::new("default")),
767                request.prompt,
768            ))
769        }
770
771        fn capabilities(&self) -> RuntimeCapabilities {
772            RuntimeCapabilities {
773                streaming: true,
774                cancellation: true,
775                steering: true,
776                approvals: true,
777                skills: false,
778                subagents: false,
779                remote_workers: false,
780            }
781        }
782    }
783
784    #[tokio::test]
785    async fn runtime_streams_events_and_returns_a_result() {
786        let runtime = Runtime::builder().executor(EchoExecutor).build().unwrap();
787        let agent = runtime.agent(AgentProfile {
788            name: "coder".into(),
789            ..AgentProfile::default()
790        });
791        let mut request = RunRequest::new("hello");
792        request.session_id = Some(SessionId::new("session-1"));
793        let mut run = agent.run(request);
794        let run_id = run.id().clone();
795
796        let delta = run.next_event().await.unwrap();
797        assert_eq!(delta.run_id, run_id.as_str());
798        assert_eq!(delta.sequence, 0);
799        assert!(matches!(delta.event, RuntimeEvent::TextDelta { .. }));
800        let final_event = run.next_event().await.unwrap();
801        assert_eq!(final_event.sequence, 1);
802        assert!(matches!(
803            final_event.event,
804            RuntimeEvent::FinalResponse { ref text } if text == "hello"
805        ));
806
807        let result = run.result().await.unwrap();
808        assert_eq!(result.status, RunStatus::Completed);
809        assert_eq!(result.session_id, SessionId::new("session-1"));
810        assert_eq!(result.final_text, "hello");
811    }
812
813    struct ControlledExecutor;
814
815    #[async_trait]
816    impl RunExecutor for ControlledExecutor {
817        async fn execute(
818            &self,
819            _profile: AgentProfile,
820            _request: RunRequest,
821            mut context: ExecutionContext,
822        ) -> Result<ExecutionResult, RuntimeError> {
823            match context.next_control().await {
824                Some(RuntimeControl::Cancel { .. }) => Err(RuntimeError {
825                    code: RuntimeErrorCode::Cancelled,
826                    message: "cancelled by caller".into(),
827                    retryable: false,
828                }),
829                Some(RuntimeControl::Steer { message, .. }) => {
830                    Ok(ExecutionResult::new("steered", message))
831                }
832                Some(RuntimeControl::ResolveApproval { decision, .. }) => {
833                    Ok(ExecutionResult::new("approved", decision))
834                }
835                None => Err(RuntimeError {
836                    code: RuntimeErrorCode::Unavailable,
837                    message: "control channel closed".into(),
838                    retryable: false,
839                }),
840            }
841        }
842    }
843
844    #[tokio::test]
845    async fn run_handle_routes_cancellation() {
846        let runtime = Runtime::builder()
847            .executor(ControlledExecutor)
848            .build()
849            .unwrap();
850        let run = runtime
851            .agent(AgentProfile::default())
852            .run(RunRequest::new("wait"));
853        run.cancel().unwrap();
854        let result = run.result().await.unwrap();
855        assert_eq!(result.status, RunStatus::Cancelled);
856        assert_eq!(result.error.unwrap().code, RuntimeErrorCode::Cancelled);
857    }
858
859    #[tokio::test]
860    async fn run_handle_routes_steering() {
861        let runtime = Runtime::builder()
862            .executor(ControlledExecutor)
863            .build()
864            .unwrap();
865        let run = runtime
866            .agent(AgentProfile::default())
867            .run(RunRequest::new("wait"));
868        run.steer("new direction").unwrap();
869        let result = run.result().await.unwrap();
870        assert_eq!(result.final_text, "new direction");
871    }
872
873    #[tokio::test]
874    async fn confirmed_control_waits_for_executor_acknowledgement() {
875        let runtime = Runtime::builder()
876            .executor(ControlledExecutor)
877            .build()
878            .unwrap();
879        let run = runtime
880            .agent(AgentProfile::default())
881            .run(RunRequest::new("wait"));
882        let controller = run.controller();
883
884        controller.steer_confirmed("accepted update").await.unwrap();
885        let result = run.result().await.unwrap();
886        assert_eq!(result.final_text, "accepted update");
887    }
888
889    struct RejectingExecutor;
890
891    #[async_trait]
892    impl RunExecutor for RejectingExecutor {
893        async fn execute(
894            &self,
895            _profile: AgentProfile,
896            _request: RunRequest,
897            mut context: ExecutionContext,
898        ) -> Result<ExecutionResult, RuntimeError> {
899            let request = context.next_control_request().await.unwrap();
900            request.reject(RuntimeError {
901                code: RuntimeErrorCode::ApprovalDenied,
902                message: "control rejected by policy".into(),
903                retryable: false,
904            });
905            Ok(ExecutionResult::new("rejected", "unchanged"))
906        }
907    }
908
909    #[tokio::test]
910    async fn confirmed_control_surfaces_executor_rejection() {
911        let runtime = Runtime::builder()
912            .executor(RejectingExecutor)
913            .build()
914            .unwrap();
915        let run = runtime
916            .agent(AgentProfile::default())
917            .run(RunRequest::new("wait"));
918        let controller = run.controller();
919
920        let error = controller.cancel_confirmed().await.unwrap_err();
921        assert_eq!(error.code, RuntimeErrorCode::ApprovalDenied);
922        assert_eq!(run.result().await.unwrap().final_text, "unchanged");
923    }
924
925    struct ConcurrentExecutor {
926        active: Arc<AtomicUsize>,
927        peak: Arc<AtomicUsize>,
928    }
929
930    #[async_trait]
931    impl RunExecutor for ConcurrentExecutor {
932        async fn execute(
933            &self,
934            _profile: AgentProfile,
935            request: RunRequest,
936            _context: ExecutionContext,
937        ) -> Result<ExecutionResult, RuntimeError> {
938            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
939            self.peak.fetch_max(active, Ordering::SeqCst);
940            tokio::time::sleep(Duration::from_millis(20)).await;
941            self.active.fetch_sub(1, Ordering::SeqCst);
942            Ok(ExecutionResult::new("session", request.prompt))
943        }
944    }
945
946    #[tokio::test]
947    async fn runtime_enforces_concurrency_limit() {
948        let active = Arc::new(AtomicUsize::new(0));
949        let peak = Arc::new(AtomicUsize::new(0));
950        let runtime = Runtime::builder()
951            .executor(ConcurrentExecutor {
952                active: active.clone(),
953                peak: peak.clone(),
954            })
955            .max_concurrent_runs(1)
956            .build()
957            .unwrap();
958        let agent = runtime.agent(AgentProfile::default());
959        let first = agent.run(RunRequest::new("one"));
960        let second = agent.run(RunRequest::new("two"));
961        let (a, b) = tokio::join!(first.result(), second.result());
962        assert!(a.is_ok());
963        assert!(b.is_ok());
964        assert_eq!(peak.load(Ordering::SeqCst), 1);
965    }
966
967    #[tokio::test]
968    async fn runtime_shutdown_drains_and_rejects_new_runs() {
969        let runtime = Runtime::builder().executor(EchoExecutor).build().unwrap();
970        runtime.shutdown().await;
971        assert!(runtime.is_shutting_down());
972        assert!(runtime.stats().shutting_down);
973
974        let result = runtime
975            .agent(AgentProfile::default())
976            .run(RunRequest::new("after shutdown"))
977            .result()
978            .await
979            .unwrap();
980        assert_eq!(result.status, RunStatus::Failed);
981        assert_eq!(result.error.unwrap().code, RuntimeErrorCode::Unavailable);
982    }
983
984    #[tokio::test]
985    async fn local_worker_uses_the_same_runtime_contract() {
986        let runtime = Runtime::builder().executor(EchoExecutor).build().unwrap();
987        let worker = LocalWorker::new(runtime, "local", "Local worker");
988        let descriptor = worker.descriptor();
989        assert_eq!(descriptor.id, WorkerId::new("local"));
990        assert_eq!(descriptor.max_concurrent_runs, 4);
991        assert!(descriptor.capabilities.streaming);
992        let run = worker
993            .submit(AgentProfile::default(), RunRequest::new("worker task"))
994            .await
995            .unwrap();
996        assert_eq!(run.result().await.unwrap().final_text, "worker task");
997    }
998
999    #[tokio::test]
1000    async fn local_worker_health_reports_runtime_capacity_and_activity() {
1001        let runtime = Runtime::builder()
1002            .executor(ConcurrentExecutor {
1003                active: Arc::new(AtomicUsize::new(0)),
1004                peak: Arc::new(AtomicUsize::new(0)),
1005            })
1006            .max_concurrent_runs(1)
1007            .build()
1008            .unwrap();
1009        let worker = LocalWorker::new(runtime, "local", "Local worker");
1010        let run = worker
1011            .submit(AgentProfile::default(), RunRequest::new("worker task"))
1012            .await
1013            .unwrap();
1014
1015        tokio::time::timeout(Duration::from_secs(1), async {
1016            loop {
1017                let health = worker.health();
1018                if health.active_runs == 1 {
1019                    assert_eq!(health.status, WorkerHealthStatus::Busy);
1020                    assert!(!health.observed_at.is_empty());
1021                    break;
1022                }
1023                tokio::task::yield_now().await;
1024            }
1025        })
1026        .await
1027        .expect("worker should report its active run");
1028
1029        run.result().await.unwrap();
1030        let health = worker.health();
1031        assert_eq!(health.active_runs, 0);
1032        assert_eq!(health.status, WorkerHealthStatus::Ready);
1033    }
1034
1035    struct RelationshipExecutor;
1036
1037    #[async_trait]
1038    impl RunExecutor for RelationshipExecutor {
1039        async fn execute(
1040            &self,
1041            profile: AgentProfile,
1042            request: RunRequest,
1043            _context: ExecutionContext,
1044        ) -> Result<ExecutionResult, RuntimeError> {
1045            let mut output = ExecutionResult::new("child-session", request.prompt);
1046            output.metadata.insert(
1047                "agent_id".into(),
1048                serde_json::json!(profile.id.map(|id| id.into_inner())),
1049            );
1050            output.metadata.insert(
1051                "parent_run_id".into(),
1052                serde_json::json!(request.parent_run_id.map(|id| id.into_inner())),
1053            );
1054            Ok(output)
1055        }
1056    }
1057
1058    #[tokio::test]
1059    async fn parent_child_identity_survives_local_worker_submission() {
1060        let runtime = Runtime::builder()
1061            .executor(RelationshipExecutor)
1062            .build()
1063            .unwrap();
1064        let worker = LocalWorker::new(runtime, "local", "Local worker");
1065        let profile = AgentProfile {
1066            id: Some(microclaw_core::run_protocol::AgentId::new("reviewer")),
1067            ..AgentProfile::default()
1068        };
1069        let mut request = RunRequest::new("review child task");
1070        request.parent_run_id = Some(RunId::new("parent-1"));
1071
1072        let result = worker
1073            .submit(profile, request)
1074            .await
1075            .unwrap()
1076            .result()
1077            .await
1078            .unwrap();
1079        assert_eq!(result.metadata["agent_id"], "reviewer");
1080        assert_eq!(result.metadata["parent_run_id"], "parent-1");
1081    }
1082
1083    #[tokio::test]
1084    async fn local_worker_drains_without_accepting_new_runs() {
1085        let runtime = Runtime::builder()
1086            .executor(ConcurrentExecutor {
1087                active: Arc::new(AtomicUsize::new(0)),
1088                peak: Arc::new(AtomicUsize::new(0)),
1089            })
1090            .max_concurrent_runs(1)
1091            .build()
1092            .unwrap();
1093        let worker =
1094            LocalWorker::new(runtime, "local", "Local worker").with_label("pool", "foreground");
1095        let run = worker
1096            .submit(AgentProfile::default(), RunRequest::new("existing"))
1097            .await
1098            .unwrap();
1099        worker.begin_drain();
1100
1101        assert_eq!(worker.descriptor().labels["pool"], "foreground");
1102        assert_eq!(worker.health().status, WorkerHealthStatus::Draining);
1103        let rejected = match worker
1104            .submit(AgentProfile::default(), RunRequest::new("new"))
1105            .await
1106        {
1107            Ok(_) => panic!("draining Worker should reject new runs"),
1108            Err(error) => error,
1109        };
1110        assert_eq!(rejected.code, RuntimeErrorCode::Unavailable);
1111
1112        run.result().await.unwrap();
1113        worker.wait_for_idle().await;
1114        assert_eq!(worker.health().active_runs, 0);
1115        worker.mark_unavailable();
1116        assert_eq!(worker.health().status, WorkerHealthStatus::Unavailable);
1117        worker.resume();
1118        assert_eq!(worker.health().status, WorkerHealthStatus::Ready);
1119    }
1120}
microclaw_engine/internal/runtime/mod.rs

microclaw_engine/internal/runtime/
mod.rs
