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, 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}
245
246#[derive(Clone)]
247pub struct Runtime {
248    inner: Arc<RuntimeInner>,
249}
250
251impl Runtime {
252    pub fn builder() -> RuntimeBuilder {
253        RuntimeBuilder::default()
254    }
255
256    pub fn agent(&self, profile: AgentProfile) -> AgentHandle {
257        AgentHandle {
258            runtime: self.clone(),
259            profile,
260        }
261    }
262
263    pub fn capabilities(&self) -> RuntimeCapabilities {
264        self.inner.executor.capabilities()
265    }
266
267    pub fn max_concurrent_runs(&self) -> usize {
268        self.inner.max_concurrent_runs
269    }
270
271    pub fn active_runs(&self) -> usize {
272        self.inner.active_runs.load(Ordering::Acquire)
273    }
274
275    pub fn queued_runs(&self) -> usize {
276        self.inner.queued_runs.load(Ordering::Acquire)
277    }
278
279    pub fn stats(&self) -> RuntimeStats {
280        RuntimeStats {
281            active_runs: self.active_runs(),
282            queued_runs: self.queued_runs(),
283            max_concurrent_runs: self.max_concurrent_runs(),
284        }
285    }
286
287    pub async fn wait_for_idle(&self) {
288        loop {
289            let notified = self.inner.idle.notified();
290            if self.active_runs() == 0 && self.queued_runs() == 0 {
291                return;
292            }
293            notified.await;
294        }
295    }
296}
297
298#[derive(Clone)]
299pub struct AgentHandle {
300    runtime: Runtime,
301    profile: AgentProfile,
302}
303
304impl AgentHandle {
305    pub fn profile(&self) -> &AgentProfile {
306        &self.profile
307    }
308
309    pub fn run(&self, request: RunRequest) -> RunHandle {
310        let run_id = request
311            .run_id
312            .clone()
313            .unwrap_or_else(|| RunId::new(uuid::Uuid::new_v4().to_string()));
314        let (event_tx, event_rx) = mpsc::unbounded_channel();
315        let (control_tx, control_rx) = mpsc::unbounded_channel();
316        let (result_tx, result_rx) = oneshot::channel();
317        let next_sequence = Arc::new(AtomicU64::new(0));
318        let executor = self.runtime.inner.executor.clone();
319        let permits = self.runtime.inner.permits.clone();
320        let active_runs = self.runtime.inner.active_runs.clone();
321        let runtime_inner = self.runtime.inner.clone();
322        let profile = self.profile.clone();
323        let task_run_id = run_id.clone();
324        let fallback_session_id = request
325            .session_id
326            .clone()
327            .unwrap_or_else(|| SessionId::new(task_run_id.to_string()));
328
329        runtime_inner.queued_runs.fetch_add(1, Ordering::AcqRel);
330        tokio::spawn(async move {
331            let permit = match permits.acquire_owned().await {
332                Ok(permit) => permit,
333                Err(_) => {
334                    runtime_inner.queued_runs.fetch_sub(1, Ordering::AcqRel);
335                    runtime_inner.idle.notify_waiters();
336                    let _ = result_tx.send(RunResult {
337                        run_id: task_run_id,
338                        session_id: fallback_session_id,
339                        status: RunStatus::Failed,
340                        final_text: String::new(),
341                        error: Some(RuntimeError {
342                            code: RuntimeErrorCode::Unavailable,
343                            message: "runtime is shutting down".into(),
344                            retryable: true,
345                        }),
346                        metadata: BTreeMap::new(),
347                    });
348                    return;
349                }
350            };
351            runtime_inner.queued_runs.fetch_sub(1, Ordering::AcqRel);
352            let _active_run = ActiveRunGuard::new(active_runs, runtime_inner.clone());
353            let context = ExecutionContext {
354                run_id: task_run_id.clone(),
355                event_tx: event_tx.clone(),
356                control_rx,
357                next_sequence: next_sequence.clone(),
358            };
359            let result = executor.execute(profile, request, context).await;
360            drop(permit);
361
362            let terminal = match result {
363                Ok(output) => {
364                    emit_terminal_event(
365                        &event_tx,
366                        &next_sequence,
367                        &task_run_id,
368                        RuntimeEvent::FinalResponse {
369                            text: output.final_text.clone(),
370                        },
371                    );
372                    RunResult {
373                        run_id: task_run_id,
374                        session_id: output.session_id,
375                        status: RunStatus::Completed,
376                        final_text: output.final_text,
377                        error: None,
378                        metadata: output.metadata,
379                    }
380                }
381                Err(error) => {
382                    let status = match error.code {
383                        RuntimeErrorCode::Cancelled => RunStatus::Cancelled,
384                        RuntimeErrorCode::TimedOut => RunStatus::TimedOut,
385                        _ => RunStatus::Failed,
386                    };
387                    if status == RunStatus::Cancelled {
388                        emit_terminal_event(
389                            &event_tx,
390                            &next_sequence,
391                            &task_run_id,
392                            RuntimeEvent::Cancelled {
393                                final_text: error.message.clone(),
394                            },
395                        );
396                    }
397                    RunResult {
398                        run_id: task_run_id,
399                        session_id: fallback_session_id,
400                        status,
401                        final_text: String::new(),
402                        error: Some(error),
403                        metadata: BTreeMap::new(),
404                    }
405                }
406            };
407            let _ = result_tx.send(terminal);
408        });
409
410        RunHandle {
411            run_id,
412            events: event_rx,
413            control_tx,
414            result_rx,
415        }
416    }
417}
418
419struct ActiveRunGuard {
420    active_runs: Arc<AtomicUsize>,
421    runtime: Arc<RuntimeInner>,
422}
423
424impl ActiveRunGuard {
425    fn new(active_runs: Arc<AtomicUsize>, runtime: Arc<RuntimeInner>) -> Self {
426        active_runs.fetch_add(1, Ordering::AcqRel);
427        Self {
428            active_runs,
429            runtime,
430        }
431    }
432}
433
434impl Drop for ActiveRunGuard {
435    fn drop(&mut self) {
436        self.active_runs.fetch_sub(1, Ordering::AcqRel);
437        self.runtime.idle.notify_waiters();
438    }
439}
440
441fn emit_terminal_event(
442    event_tx: &mpsc::UnboundedSender<RuntimeEventEnvelope>,
443    next_sequence: &AtomicU64,
444    run_id: &RunId,
445    event: RuntimeEvent,
446) {
447    let sequence = next_sequence.fetch_add(1, Ordering::Relaxed);
448    let _ = event_tx.send(RuntimeEventEnvelope::new(
449        run_id.to_string(),
450        sequence,
451        event,
452    ));
453}
454
455pub struct RunHandle {
456    run_id: RunId,
457    events: mpsc::UnboundedReceiver<RuntimeEventEnvelope>,
458    control_tx: mpsc::UnboundedSender<PendingControl>,
459    result_rx: oneshot::Receiver<RunResult>,
460}
461
462impl RunHandle {
463    pub fn id(&self) -> &RunId {
464        &self.run_id
465    }
466
467    pub async fn next_event(&mut self) -> Option<RuntimeEventEnvelope> {
468        self.events.recv().await
469    }
470
471    pub fn try_next_event(&mut self) -> Option<RuntimeEventEnvelope> {
472        self.events.try_recv().ok()
473    }
474
475    pub fn cancel(&self) -> Result<(), RuntimeError> {
476        self.controller().cancel()
477    }
478
479    pub fn controller(&self) -> RunController {
480        RunController {
481            run_id: self.run_id.clone(),
482            control_tx: self.control_tx.clone(),
483        }
484    }
485
486    pub async fn cancel_confirmed(&self) -> Result<(), RuntimeError> {
487        self.controller().cancel_confirmed().await
488    }
489
490    pub fn steer(&self, message: impl Into<String>) -> Result<(), RuntimeError> {
491        self.controller().steer(message)
492    }
493
494    pub async fn steer_confirmed(&self, message: impl Into<String>) -> Result<(), RuntimeError> {
495        self.controller().steer_confirmed(message).await
496    }
497
498    pub fn resolve_approval(
499        &self,
500        approval_id: impl Into<String>,
501        decision: impl Into<String>,
502    ) -> Result<(), RuntimeError> {
503        self.controller().resolve_approval(approval_id, decision)
504    }
505
506    pub async fn resolve_approval_confirmed(
507        &self,
508        approval_id: impl Into<String>,
509        decision: impl Into<String>,
510    ) -> Result<(), RuntimeError> {
511        self.controller()
512            .resolve_approval_confirmed(approval_id, decision)
513            .await
514    }
515
516    pub async fn result(self) -> Result<RunResult, RuntimeError> {
517        self.result_rx.await.map_err(|_| RuntimeError {
518            code: RuntimeErrorCode::Internal,
519            message: "run ended without a terminal result".into(),
520            retryable: false,
521        })
522    }
523}
524
525/// Cloneable control plane for a running task.
526///
527/// Keeping this separate from [`RunHandle`] lets one task consume events while another task or
528/// UI callback controls the run without sharing the event receiver.
529#[derive(Clone)]
530pub struct RunController {
531    run_id: RunId,
532    control_tx: mpsc::UnboundedSender<PendingControl>,
533}
534
535impl RunController {
536    pub fn id(&self) -> &RunId {
537        &self.run_id
538    }
539
540    pub fn cancel(&self) -> Result<(), RuntimeError> {
541        self.send_control(RuntimeControl::Cancel {
542            run_id: self.run_id.clone(),
543        })
544    }
545
546    pub fn steer(&self, message: impl Into<String>) -> Result<(), RuntimeError> {
547        self.send_control(RuntimeControl::Steer {
548            run_id: self.run_id.clone(),
549            message: message.into(),
550        })
551    }
552
553    pub fn resolve_approval(
554        &self,
555        approval_id: impl Into<String>,
556        decision: impl Into<String>,
557    ) -> Result<(), RuntimeError> {
558        self.send_control(RuntimeControl::ResolveApproval {
559            run_id: self.run_id.clone(),
560            approval_id: approval_id.into(),
561            decision: decision.into(),
562        })
563    }
564
565    fn send_control(&self, control: RuntimeControl) -> Result<(), RuntimeError> {
566        self.control_tx
567            .send(PendingControl {
568                control,
569                receipt: None,
570            })
571            .map_err(|_| RuntimeError {
572                code: RuntimeErrorCode::Unavailable,
573                message: "run is no longer accepting control messages".into(),
574                retryable: false,
575            })
576    }
577
578    pub async fn cancel_confirmed(&self) -> Result<(), RuntimeError> {
579        self.send_control_confirmed(RuntimeControl::Cancel {
580            run_id: self.run_id.clone(),
581        })
582        .await
583    }
584
585    pub async fn steer_confirmed(&self, message: impl Into<String>) -> Result<(), RuntimeError> {
586        self.send_control_confirmed(RuntimeControl::Steer {
587            run_id: self.run_id.clone(),
588            message: message.into(),
589        })
590        .await
591    }
592
593    pub async fn resolve_approval_confirmed(
594        &self,
595        approval_id: impl Into<String>,
596        decision: impl Into<String>,
597    ) -> Result<(), RuntimeError> {
598        self.send_control_confirmed(RuntimeControl::ResolveApproval {
599            run_id: self.run_id.clone(),
600            approval_id: approval_id.into(),
601            decision: decision.into(),
602        })
603        .await
604    }
605
606    async fn send_control_confirmed(&self, control: RuntimeControl) -> Result<(), RuntimeError> {
607        let (receipt_tx, receipt_rx) = oneshot::channel();
608        self.control_tx
609            .send(PendingControl {
610                control,
611                receipt: Some(receipt_tx),
612            })
613            .map_err(|_| RuntimeError {
614                code: RuntimeErrorCode::Unavailable,
615                message: "run is no longer accepting control messages".into(),
616                retryable: false,
617            })?;
618        receipt_rx.await.map_err(|_| RuntimeError {
619            code: RuntimeErrorCode::Unavailable,
620            message: "run ended before acknowledging the control request".into(),
621            retryable: false,
622        })?
623    }
624}
625
626#[async_trait]
627pub trait Worker: Send + Sync {
628    fn descriptor(&self) -> WorkerDescriptor;
629    fn health(&self) -> WorkerHealth;
630    async fn submit(
631        &self,
632        profile: AgentProfile,
633        request: RunRequest,
634    ) -> Result<RunHandle, RuntimeError>;
635}
636
637const WORKER_READY: u8 = 0;
638const WORKER_DRAINING: u8 = 1;
639const WORKER_UNAVAILABLE: u8 = 2;
640
641#[derive(Clone)]
642pub struct LocalWorker {
643    descriptor: WorkerDescriptor,
644    runtime: Runtime,
645    state: Arc<AtomicU8>,
646}
647
648impl LocalWorker {
649    pub fn new(runtime: Runtime, id: impl Into<WorkerId>, name: impl Into<String>) -> Self {
650        let max_concurrent_runs = runtime.max_concurrent_runs();
651        Self {
652            descriptor: WorkerDescriptor {
653                id: id.into(),
654                name: name.into(),
655                capabilities: runtime.capabilities(),
656                max_concurrent_runs,
657                labels: BTreeMap::new(),
658            },
659            runtime,
660            state: Arc::new(AtomicU8::new(WORKER_READY)),
661        }
662    }
663
664    pub fn with_label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
665        self.descriptor.labels.insert(key.into(), value.into());
666        self
667    }
668
669    pub fn begin_drain(&self) {
670        self.state.store(WORKER_DRAINING, Ordering::Release);
671    }
672
673    pub fn mark_unavailable(&self) {
674        self.state.store(WORKER_UNAVAILABLE, Ordering::Release);
675    }
676
677    pub fn resume(&self) {
678        self.state.store(WORKER_READY, Ordering::Release);
679    }
680
681    pub async fn wait_for_idle(&self) {
682        self.runtime.wait_for_idle().await;
683    }
684}
685
686#[async_trait]
687impl Worker for LocalWorker {
688    fn descriptor(&self) -> WorkerDescriptor {
689        self.descriptor.clone()
690    }
691
692    fn health(&self) -> WorkerHealth {
693        let active_runs = self.runtime.active_runs();
694        let state = self.state.load(Ordering::Acquire);
695        WorkerHealth {
696            worker_id: self.descriptor.id.clone(),
697            status: match state {
698                WORKER_DRAINING => WorkerHealthStatus::Draining,
699                WORKER_UNAVAILABLE => WorkerHealthStatus::Unavailable,
700                _ if active_runs >= self.descriptor.max_concurrent_runs => WorkerHealthStatus::Busy,
701                _ => WorkerHealthStatus::Ready,
702            },
703            active_runs,
704            queued_runs: self.runtime.queued_runs(),
705            observed_at: std::time::SystemTime::now()
706                .duration_since(std::time::UNIX_EPOCH)
707                .unwrap_or_default()
708                .as_secs()
709                .to_string(),
710        }
711    }
712
713    async fn submit(
714        &self,
715        profile: AgentProfile,
716        request: RunRequest,
717    ) -> Result<RunHandle, RuntimeError> {
718        if self.state.load(Ordering::Acquire) != WORKER_READY {
719            return Err(RuntimeError {
720                code: RuntimeErrorCode::Unavailable,
721                message: "Worker is not accepting new runs".into(),
722                retryable: true,
723            });
724        }
725        Ok(self.runtime.agent(profile).run(request))
726    }
727}
728
729#[cfg(test)]
730mod tests {
731    use std::sync::atomic::{AtomicUsize, Ordering};
732    use std::time::Duration;
733
734    use super::*;
735
736    struct EchoExecutor;
737
738    #[async_trait]
739    impl RunExecutor for EchoExecutor {
740        async fn execute(
741            &self,
742            profile: AgentProfile,
743            request: RunRequest,
744            context: ExecutionContext,
745        ) -> Result<ExecutionResult, RuntimeError> {
746            context.emit(RuntimeEvent::TextDelta {
747                delta: format!("{}: ", profile.name),
748            })?;
749            Ok(ExecutionResult::new(
750                request
751                    .session_id
752                    .clone()
753                    .unwrap_or_else(|| SessionId::new("default")),
754                request.prompt,
755            ))
756        }
757
758        fn capabilities(&self) -> RuntimeCapabilities {
759            RuntimeCapabilities {
760                streaming: true,
761                cancellation: true,
762                steering: true,
763                approvals: true,
764                skills: false,
765                subagents: false,
766                remote_workers: false,
767            }
768        }
769    }
770
771    #[tokio::test]
772    async fn runtime_streams_events_and_returns_a_result() {
773        let runtime = Runtime::builder().executor(EchoExecutor).build().unwrap();
774        let agent = runtime.agent(AgentProfile {
775            name: "coder".into(),
776            ..AgentProfile::default()
777        });
778        let mut request = RunRequest::new("hello");
779        request.session_id = Some(SessionId::new("session-1"));
780        let mut run = agent.run(request);
781        let run_id = run.id().clone();
782
783        let delta = run.next_event().await.unwrap();
784        assert_eq!(delta.run_id, run_id.as_str());
785        assert_eq!(delta.sequence, 0);
786        assert!(matches!(delta.event, RuntimeEvent::TextDelta { .. }));
787        let final_event = run.next_event().await.unwrap();
788        assert_eq!(final_event.sequence, 1);
789        assert!(matches!(
790            final_event.event,
791            RuntimeEvent::FinalResponse { ref text } if text == "hello"
792        ));
793
794        let result = run.result().await.unwrap();
795        assert_eq!(result.status, RunStatus::Completed);
796        assert_eq!(result.session_id, SessionId::new("session-1"));
797        assert_eq!(result.final_text, "hello");
798    }
799
800    struct ControlledExecutor;
801
802    #[async_trait]
803    impl RunExecutor for ControlledExecutor {
804        async fn execute(
805            &self,
806            _profile: AgentProfile,
807            _request: RunRequest,
808            mut context: ExecutionContext,
809        ) -> Result<ExecutionResult, RuntimeError> {
810            match context.next_control().await {
811                Some(RuntimeControl::Cancel { .. }) => Err(RuntimeError {
812                    code: RuntimeErrorCode::Cancelled,
813                    message: "cancelled by caller".into(),
814                    retryable: false,
815                }),
816                Some(RuntimeControl::Steer { message, .. }) => {
817                    Ok(ExecutionResult::new("steered", message))
818                }
819                Some(RuntimeControl::ResolveApproval { decision, .. }) => {
820                    Ok(ExecutionResult::new("approved", decision))
821                }
822                None => Err(RuntimeError {
823                    code: RuntimeErrorCode::Unavailable,
824                    message: "control channel closed".into(),
825                    retryable: false,
826                }),
827            }
828        }
829    }
830
831    #[tokio::test]
832    async fn run_handle_routes_cancellation() {
833        let runtime = Runtime::builder()
834            .executor(ControlledExecutor)
835            .build()
836            .unwrap();
837        let run = runtime
838            .agent(AgentProfile::default())
839            .run(RunRequest::new("wait"));
840        run.cancel().unwrap();
841        let result = run.result().await.unwrap();
842        assert_eq!(result.status, RunStatus::Cancelled);
843        assert_eq!(result.error.unwrap().code, RuntimeErrorCode::Cancelled);
844    }
845
846    #[tokio::test]
847    async fn run_handle_routes_steering() {
848        let runtime = Runtime::builder()
849            .executor(ControlledExecutor)
850            .build()
851            .unwrap();
852        let run = runtime
853            .agent(AgentProfile::default())
854            .run(RunRequest::new("wait"));
855        run.steer("new direction").unwrap();
856        let result = run.result().await.unwrap();
857        assert_eq!(result.final_text, "new direction");
858    }
859
860    #[tokio::test]
861    async fn confirmed_control_waits_for_executor_acknowledgement() {
862        let runtime = Runtime::builder()
863            .executor(ControlledExecutor)
864            .build()
865            .unwrap();
866        let run = runtime
867            .agent(AgentProfile::default())
868            .run(RunRequest::new("wait"));
869        let controller = run.controller();
870
871        controller.steer_confirmed("accepted update").await.unwrap();
872        let result = run.result().await.unwrap();
873        assert_eq!(result.final_text, "accepted update");
874    }
875
876    struct RejectingExecutor;
877
878    #[async_trait]
879    impl RunExecutor for RejectingExecutor {
880        async fn execute(
881            &self,
882            _profile: AgentProfile,
883            _request: RunRequest,
884            mut context: ExecutionContext,
885        ) -> Result<ExecutionResult, RuntimeError> {
886            let request = context.next_control_request().await.unwrap();
887            request.reject(RuntimeError {
888                code: RuntimeErrorCode::ApprovalDenied,
889                message: "control rejected by policy".into(),
890                retryable: false,
891            });
892            Ok(ExecutionResult::new("rejected", "unchanged"))
893        }
894    }
895
896    #[tokio::test]
897    async fn confirmed_control_surfaces_executor_rejection() {
898        let runtime = Runtime::builder()
899            .executor(RejectingExecutor)
900            .build()
901            .unwrap();
902        let run = runtime
903            .agent(AgentProfile::default())
904            .run(RunRequest::new("wait"));
905        let controller = run.controller();
906
907        let error = controller.cancel_confirmed().await.unwrap_err();
908        assert_eq!(error.code, RuntimeErrorCode::ApprovalDenied);
909        assert_eq!(run.result().await.unwrap().final_text, "unchanged");
910    }
911
912    struct ConcurrentExecutor {
913        active: Arc<AtomicUsize>,
914        peak: Arc<AtomicUsize>,
915    }
916
917    #[async_trait]
918    impl RunExecutor for ConcurrentExecutor {
919        async fn execute(
920            &self,
921            _profile: AgentProfile,
922            request: RunRequest,
923            _context: ExecutionContext,
924        ) -> Result<ExecutionResult, RuntimeError> {
925            let active = self.active.fetch_add(1, Ordering::SeqCst) + 1;
926            self.peak.fetch_max(active, Ordering::SeqCst);
927            tokio::time::sleep(Duration::from_millis(20)).await;
928            self.active.fetch_sub(1, Ordering::SeqCst);
929            Ok(ExecutionResult::new("session", request.prompt))
930        }
931    }
932
933    #[tokio::test]
934    async fn runtime_enforces_concurrency_limit() {
935        let active = Arc::new(AtomicUsize::new(0));
936        let peak = Arc::new(AtomicUsize::new(0));
937        let runtime = Runtime::builder()
938            .executor(ConcurrentExecutor {
939                active: active.clone(),
940                peak: peak.clone(),
941            })
942            .max_concurrent_runs(1)
943            .build()
944            .unwrap();
945        let agent = runtime.agent(AgentProfile::default());
946        let first = agent.run(RunRequest::new("one"));
947        let second = agent.run(RunRequest::new("two"));
948        let (a, b) = tokio::join!(first.result(), second.result());
949        assert!(a.is_ok());
950        assert!(b.is_ok());
951        assert_eq!(peak.load(Ordering::SeqCst), 1);
952    }
953
954    #[tokio::test]
955    async fn local_worker_uses_the_same_runtime_contract() {
956        let runtime = Runtime::builder().executor(EchoExecutor).build().unwrap();
957        let worker = LocalWorker::new(runtime, "local", "Local worker");
958        let descriptor = worker.descriptor();
959        assert_eq!(descriptor.id, WorkerId::new("local"));
960        assert_eq!(descriptor.max_concurrent_runs, 4);
961        assert!(descriptor.capabilities.streaming);
962        let run = worker
963            .submit(AgentProfile::default(), RunRequest::new("worker task"))
964            .await
965            .unwrap();
966        assert_eq!(run.result().await.unwrap().final_text, "worker task");
967    }
968
969    #[tokio::test]
970    async fn local_worker_health_reports_runtime_capacity_and_activity() {
971        let runtime = Runtime::builder()
972            .executor(ConcurrentExecutor {
973                active: Arc::new(AtomicUsize::new(0)),
974                peak: Arc::new(AtomicUsize::new(0)),
975            })
976            .max_concurrent_runs(1)
977            .build()
978            .unwrap();
979        let worker = LocalWorker::new(runtime, "local", "Local worker");
980        let run = worker
981            .submit(AgentProfile::default(), RunRequest::new("worker task"))
982            .await
983            .unwrap();
984
985        tokio::time::timeout(Duration::from_secs(1), async {
986            loop {
987                let health = worker.health();
988                if health.active_runs == 1 {
989                    assert_eq!(health.status, WorkerHealthStatus::Busy);
990                    assert!(!health.observed_at.is_empty());
991                    break;
992                }
993                tokio::task::yield_now().await;
994            }
995        })
996        .await
997        .expect("worker should report its active run");
998
999        run.result().await.unwrap();
1000        let health = worker.health();
1001        assert_eq!(health.active_runs, 0);
1002        assert_eq!(health.status, WorkerHealthStatus::Ready);
1003    }
1004
1005    struct RelationshipExecutor;
1006
1007    #[async_trait]
1008    impl RunExecutor for RelationshipExecutor {
1009        async fn execute(
1010            &self,
1011            profile: AgentProfile,
1012            request: RunRequest,
1013            _context: ExecutionContext,
1014        ) -> Result<ExecutionResult, RuntimeError> {
1015            let mut output = ExecutionResult::new("child-session", request.prompt);
1016            output.metadata.insert(
1017                "agent_id".into(),
1018                serde_json::json!(profile.id.map(|id| id.into_inner())),
1019            );
1020            output.metadata.insert(
1021                "parent_run_id".into(),
1022                serde_json::json!(request.parent_run_id.map(|id| id.into_inner())),
1023            );
1024            Ok(output)
1025        }
1026    }
1027
1028    #[tokio::test]
1029    async fn parent_child_identity_survives_local_worker_submission() {
1030        let runtime = Runtime::builder()
1031            .executor(RelationshipExecutor)
1032            .build()
1033            .unwrap();
1034        let worker = LocalWorker::new(runtime, "local", "Local worker");
1035        let profile = AgentProfile {
1036            id: Some(microclaw_core::run_protocol::AgentId::new("reviewer")),
1037            ..AgentProfile::default()
1038        };
1039        let mut request = RunRequest::new("review child task");
1040        request.parent_run_id = Some(RunId::new("parent-1"));
1041
1042        let result = worker
1043            .submit(profile, request)
1044            .await
1045            .unwrap()
1046            .result()
1047            .await
1048            .unwrap();
1049        assert_eq!(result.metadata["agent_id"], "reviewer");
1050        assert_eq!(result.metadata["parent_run_id"], "parent-1");
1051    }
1052
1053    #[tokio::test]
1054    async fn local_worker_drains_without_accepting_new_runs() {
1055        let runtime = Runtime::builder()
1056            .executor(ConcurrentExecutor {
1057                active: Arc::new(AtomicUsize::new(0)),
1058                peak: Arc::new(AtomicUsize::new(0)),
1059            })
1060            .max_concurrent_runs(1)
1061            .build()
1062            .unwrap();
1063        let worker =
1064            LocalWorker::new(runtime, "local", "Local worker").with_label("pool", "foreground");
1065        let run = worker
1066            .submit(AgentProfile::default(), RunRequest::new("existing"))
1067            .await
1068            .unwrap();
1069        worker.begin_drain();
1070
1071        assert_eq!(worker.descriptor().labels["pool"], "foreground");
1072        assert_eq!(worker.health().status, WorkerHealthStatus::Draining);
1073        let rejected = match worker
1074            .submit(AgentProfile::default(), RunRequest::new("new"))
1075            .await
1076        {
1077            Ok(_) => panic!("draining Worker should reject new runs"),
1078            Err(error) => error,
1079        };
1080        assert_eq!(rejected.code, RuntimeErrorCode::Unavailable);
1081
1082        run.result().await.unwrap();
1083        worker.wait_for_idle().await;
1084        assert_eq!(worker.health().active_runs, 0);
1085        worker.mark_unavailable();
1086        assert_eq!(worker.health().status, WorkerHealthStatus::Unavailable);
1087        worker.resume();
1088        assert_eq!(worker.health().status, WorkerHealthStatus::Ready);
1089    }
1090}
microclaw_engine/internal/runtime/mod.rs

microclaw_engine/internal/runtime/
mod.rs
