awaken_runtime/backend/
mod.rs

1//! Runtime execution backends and canonical request/result types.
2
3mod checkpoint;
4mod local;
5
6mod capabilities;
7pub use capabilities::{
8    BackendCancellationCapability, BackendContinuationCapability, BackendOutputCapability,
9    BackendTranscriptCapability, BackendWaitCapability,
10};
11use std::sync::Arc;
12
13use async_trait::async_trait;
14use awaken_runtime_contract::contract::event::AgentEvent;
15use awaken_runtime_contract::contract::event_sink::EventSink;
16use awaken_runtime_contract::contract::identity::RunIdentity;
17use awaken_runtime_contract::contract::lifecycle::{RunStatus, TerminationReason};
18use awaken_runtime_contract::contract::message::{Message, gen_message_id};
19use awaken_runtime_contract::contract::suspension::ToolCallResume;
20use awaken_runtime_contract::contract::tool::ToolDescriptor;
21use awaken_runtime_contract::now_ms;
22use awaken_runtime_contract::registry_spec::RemoteEndpoint;
23use awaken_runtime_contract::state::PersistedState;
24use futures::channel::mpsc;
25use serde::{Deserialize, Serialize};
26use serde_json::{Value, json};
27
28use crate::cancellation::CancellationToken;
29use crate::checkpoint_store::RuntimeCheckpointStore;
30use crate::inbox::{InboxReceiver, InboxSender};
31use crate::loop_runner::{AgentLoopError, AgentRunResult, PendingBoundaryHandler};
32use crate::phase::PhaseRuntime;
33use crate::{
34    registry::{AgentResolver, ResolvedBackendAgent},
35    resolution::ExecutionPlan,
36};
37use checkpoint::persist_remote_root_checkpoint;
38
39pub use local::LocalBackend;
40
41const BACKEND_OUTPUT_STATE_KEY: &str = "__runtime_backend_output";
42
43/// Optional parent lineage for a backend run.
44#[derive(Debug, Clone, Default)]
45pub struct BackendParentContext {
46    pub parent_run_id: Option<String>,
47    pub parent_thread_id: Option<String>,
48    pub parent_tool_call_id: Option<String>,
49}
50/// Cooperative runtime controls exposed to a backend implementation.
51#[derive(Default)]
52pub struct BackendControl {
53    pub cancellation_token: Option<CancellationToken>,
54    pub decision_rx: Option<mpsc::UnboundedReceiver<Vec<(String, ToolCallResume)>>>,
55    pub pending_boundary: Option<Arc<dyn PendingBoundaryHandler>>,
56}
57
58/// Root execution request shared by local and remote root execution.
59pub struct BackendRootRunRequest<'a> {
60    pub agent_id: &'a str,
61    pub messages: Vec<Message>,
62    pub new_messages: Vec<Message>,
63    pub sink: Arc<dyn EventSink>,
64    pub resolver: &'a dyn AgentResolver,
65    pub run_identity: RunIdentity,
66    pub checkpoint_store: Option<&'a dyn RuntimeCheckpointStore>,
67    pub commit: crate::loop_runner::CommitWiring<'a>,
68    pub control: BackendControl,
69    pub decisions: Vec<(String, ToolCallResume)>,
70    pub overrides: Option<awaken_runtime_contract::contract::inference::InferenceOverride>,
71    pub frontend_tools: Vec<ToolDescriptor>,
72    pub local: Option<BackendLocalRootContext<'a>>,
73    pub inbox: Option<InboxReceiver>,
74    pub is_continuation: bool,
75}
76
77/// Local-only dependencies carried by the root request context.
78#[derive(Clone, Copy)]
79pub struct BackendLocalRootContext<'a> {
80    pub phase_runtime: &'a PhaseRuntime,
81}
82
83/// Delegate execution persistence policy.
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum BackendDelegatePersistence {
86    Ephemeral,
87}
88
89/// Delegate execution continuation policy.
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum BackendDelegateContinuation {
92    Disabled,
93}
94
95/// Explicit policy for delegated agent tool calls.
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub struct BackendDelegatePolicy {
98    pub persistence: BackendDelegatePersistence,
99    pub continuation: BackendDelegateContinuation,
100}
101
102impl Default for BackendDelegatePolicy {
103    fn default() -> Self {
104        Self {
105            persistence: BackendDelegatePersistence::Ephemeral,
106            continuation: BackendDelegateContinuation::Disabled,
107        }
108    }
109}
110
111/// Delegate execution request. Delegates are explicitly child invocations.
112pub struct BackendDelegateRunRequest<'a> {
113    pub agent_id: &'a str,
114    pub messages: Vec<Message>,
115    pub new_messages: Vec<Message>,
116    pub sink: Arc<dyn EventSink>,
117    pub resolver: &'a dyn AgentResolver,
118    pub parent: BackendParentContext,
119    pub control: BackendControl,
120    pub policy: BackendDelegatePolicy,
121    /// Initial state to seed the child run with before the first step.
122    pub state_seed: Option<PersistedState>,
123}
124
125/// Best-effort abort request for an in-flight backend execution.
126pub struct BackendAbortRequest<'a> {
127    pub agent_id: &'a str,
128    pub run_identity: &'a RunIdentity,
129    pub parent: Option<&'a BackendParentContext>,
130    pub persisted_state: Option<&'a PersistedState>,
131    pub is_continuation: bool,
132}
133
134/// Structured output preserved by a backend result.
135#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
136pub struct BackendRunOutput {
137    #[serde(default, skip_serializing_if = "Option::is_none")]
138    pub text: Option<String>,
139    #[serde(default, skip_serializing_if = "Vec::is_empty")]
140    pub artifacts: Vec<BackendOutputArtifact>,
141    #[serde(default, skip_serializing_if = "Option::is_none")]
142    pub raw: Option<Value>,
143}
144
145impl BackendRunOutput {
146    #[must_use]
147    pub fn from_text(text: Option<String>) -> Self {
148        Self {
149            text,
150            artifacts: Vec::new(),
151            raw: None,
152        }
153    }
154
155    #[must_use]
156    pub fn text_or<'a>(&'a self, fallback: &'a Option<String>) -> Option<String> {
157        self.text.clone().or_else(|| fallback.clone())
158    }
159}
160
161/// Backend artifact in a transport-neutral shape.
162#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
163pub struct BackendOutputArtifact {
164    #[serde(default, skip_serializing_if = "Option::is_none")]
165    pub id: Option<String>,
166    #[serde(default, skip_serializing_if = "Option::is_none")]
167    pub name: Option<String>,
168    #[serde(default, skip_serializing_if = "Option::is_none")]
169    pub media_type: Option<String>,
170    pub content: Value,
171}
172
173/// Result of executing an agent through a runtime backend.
174#[derive(Debug, Clone)]
175pub struct BackendRunResult {
176    pub agent_id: String,
177    pub status: BackendRunStatus,
178    pub termination: TerminationReason,
179    pub status_reason: Option<String>,
180    pub response: Option<String>,
181    pub output: BackendRunOutput,
182    pub steps: usize,
183    pub run_id: Option<String>,
184    pub inbox: Option<InboxSender>,
185    /// Run-scoped persisted state, written to the run record.
186    pub state: Option<PersistedState>,
187    /// Thread-scoped persisted state (ADR-0038 C4), written to the per-thread
188    /// `thread_state` store in the same commit. Only set by backends that own a
189    /// live state registry (e.g. the in-process local backend) and can classify
190    /// keys by scope; opaque remote/A2A backends leave this `None` and carry all
191    /// keys on `state`.
192    pub thread_state: Option<PersistedState>,
193}
194
195/// Terminal status of a backend run.
196#[derive(Debug, Clone)]
197pub enum BackendRunStatus {
198    Completed,
199    WaitingInput(Option<String>),
200    WaitingAuth(Option<String>),
201    Suspended(Option<String>),
202    Failed(String),
203    Cancelled,
204    Timeout,
205}
206
207impl std::fmt::Display for BackendRunStatus {
208    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209        match self {
210            Self::Completed => write!(f, "completed"),
211            Self::WaitingInput(Some(msg)) => write!(f, "waiting_input: {msg}"),
212            Self::WaitingInput(None) => write!(f, "waiting_input"),
213            Self::WaitingAuth(Some(msg)) => write!(f, "waiting_auth: {msg}"),
214            Self::WaitingAuth(None) => write!(f, "waiting_auth"),
215            Self::Suspended(Some(msg)) => write!(f, "suspended: {msg}"),
216            Self::Suspended(None) => write!(f, "suspended"),
217            Self::Failed(msg) => write!(f, "failed: {msg}"),
218            Self::Cancelled => write!(f, "cancelled"),
219            Self::Timeout => write!(f, "timeout"),
220        }
221    }
222}
223
224impl BackendRunStatus {
225    #[must_use]
226    pub fn durable_run_status(&self, termination: &TerminationReason) -> RunStatus {
227        match self {
228            Self::WaitingInput(_) | Self::WaitingAuth(_) | Self::Suspended(_) => RunStatus::Waiting,
229            Self::Completed => termination.to_run_status().0,
230            Self::Failed(_) | Self::Cancelled | Self::Timeout => RunStatus::Done,
231        }
232    }
233
234    #[must_use]
235    pub fn durable_status_reason(&self, termination: &TerminationReason) -> Option<String> {
236        match self {
237            Self::WaitingInput(_) => Some("input_required".to_string()),
238            Self::WaitingAuth(_) => Some("auth_required".to_string()),
239            Self::Suspended(_) => Some("suspended".to_string()),
240            Self::Timeout => Some("timeout".to_string()),
241            Self::Failed(_) => Some("error".to_string()),
242            Self::Cancelled => Some("cancelled".to_string()),
243            Self::Completed => termination.to_run_status().1,
244        }
245    }
246
247    #[must_use]
248    pub fn result_status_label(&self, termination: &TerminationReason) -> &'static str {
249        match self {
250            Self::Completed => run_status_label(termination.to_run_status().0),
251            Self::WaitingInput(_) => "waiting_input",
252            Self::WaitingAuth(_) => "waiting_auth",
253            Self::Suspended(_) => "suspended",
254            Self::Failed(_) => "failed",
255            Self::Cancelled => "cancelled",
256            Self::Timeout => "timeout",
257        }
258    }
259}
260
261/// Backend for executing an agent, either locally or through a remote transport.
262#[async_trait]
263pub trait ExecutionBackend: Send + Sync {
264    fn capabilities(&self) -> crate::resolution::BackendProfile {
265        crate::resolution::BackendProfile::remote_stateless_text()
266    }
267
268    async fn abort(&self, _request: BackendAbortRequest<'_>) -> Result<(), ExecutionBackendError> {
269        Ok(())
270    }
271
272    async fn execute_root(
273        &self,
274        _request: BackendRootRunRequest<'_>,
275    ) -> Result<BackendRunResult, ExecutionBackendError> {
276        Err(ExecutionBackendError::ExecutionFailed(
277            "backend does not support root execution".into(),
278        ))
279    }
280
281    async fn execute_delegate(
282        &self,
283        _request: BackendDelegateRunRequest<'_>,
284    ) -> Result<BackendRunResult, ExecutionBackendError> {
285        Err(ExecutionBackendError::ExecutionFailed(
286            "backend does not support delegated execution".into(),
287        ))
288    }
289}
290
291/// JSON schema metadata for one version of a backend-specific agent config.
292#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
293pub struct ExecutionBackendConfigSchema {
294    pub version: u32,
295    pub schema: Value,
296    #[serde(default, skip_serializing_if = "Option::is_none")]
297    pub display_name: Option<String>,
298    #[serde(default, skip_serializing_if = "Option::is_none")]
299    pub description: Option<String>,
300    #[serde(default)]
301    pub default_config: Value,
302    #[serde(default, skip_serializing_if = "Option::is_none")]
303    pub ui_schema: Option<Value>,
304}
305
306impl ExecutionBackendConfigSchema {
307    #[must_use]
308    pub fn generic_remote() -> Self {
309        Self {
310            version: 1,
311            schema: json!({
312                "type": "object",
313                "title": "Backend config",
314                "additionalProperties": true,
315            }),
316            display_name: None,
317            description: None,
318            default_config: json!({}),
319            ui_schema: None,
320        }
321    }
322}
323
324/// Schema for the built-in in-process Awaken backend.
325#[must_use]
326pub fn awaken_backend_config_schema() -> ExecutionBackendConfigSchema {
327    ExecutionBackendConfigSchema {
328        version: 1,
329        schema: json!({
330            "type": "object",
331            "title": "Awaken backend config",
332            "description": "In-process Awaken agent execution.",
333            "additionalProperties": false,
334            "required": ["model_id", "system_prompt", "max_rounds"],
335            "properties": {
336                "model_id": {
337                    "type": "string",
338                    "title": "Model",
339                    "minLength": 1
340                },
341                "system_prompt": {
342                    "type": "string",
343                    "title": "System prompt"
344                },
345                "max_rounds": {
346                    "type": "integer",
347                    "title": "Max rounds",
348                    "minimum": 1,
349                    "default": 10
350                }
351            }
352        }),
353        display_name: Some("Awaken".to_string()),
354        description: Some("Run the agent inside this Awaken runtime.".to_string()),
355        default_config: json!({
356            "model_id": "",
357            "system_prompt": "",
358            "max_rounds": 10,
359        }),
360        ui_schema: Some(json!({
361            "system_prompt": {
362                "ui:widget": "textarea"
363            }
364        })),
365    }
366}
367
368/// Factory for backend implementations backed by canonical `RemoteEndpoint` config.
369pub trait ExecutionBackendFactory: Send + Sync {
370    fn backend(&self) -> &str;
371
372    fn config_schema(&self) -> ExecutionBackendConfigSchema {
373        ExecutionBackendConfigSchema::generic_remote()
374    }
375
376    fn validate(&self, endpoint: &RemoteEndpoint) -> Result<(), ExecutionBackendFactoryError> {
377        self.build(endpoint).map(|_| ())
378    }
379
380    fn build(
381        &self,
382        endpoint: &RemoteEndpoint,
383    ) -> Result<Arc<dyn ExecutionBackend>, ExecutionBackendFactoryError>;
384}
385
386#[derive(Debug, thiserror::Error)]
387pub enum ExecutionBackendFactoryError {
388    #[error("{0}")]
389    InvalidConfig(String),
390}
391
392#[derive(Debug, thiserror::Error)]
393pub enum ExecutionBackendError {
394    #[error("agent not found: {0}")]
395    AgentNotFound(String),
396    #[error("execution failed: {0}")]
397    ExecutionFailed(String),
398    #[error("remote error: {0}")]
399    RemoteError(String),
400    #[error(transparent)]
401    Loop(#[from] AgentLoopError),
402}
403
404pub fn execution_capabilities(
405    execution: &ExecutionPlan,
406) -> Result<crate::resolution::BackendProfile, ExecutionBackendError> {
407    match execution {
408        ExecutionPlan::Local(_) => Ok(LocalBackend::new().capabilities()),
409        ExecutionPlan::Remote(agent) => Ok(agent.backend()?.capabilities()),
410    }
411}
412
413fn root_request_requirements(
414    request: &BackendRootRunRequest<'_>,
415) -> crate::resolution::BackendRequirements {
416    use crate::resolution::{
417        BackendRequirements, ContinuationCapability, DecisionCapability, FrontendToolCapability,
418        OverrideCapability,
419    };
420    let has_seeded = !request.decisions.is_empty();
421    let has_live = request.control.decision_rx.is_some();
422    let decisions = match (has_seeded, has_live) {
423        (false, false) => None,
424        (false, true) => Some(DecisionCapability::LiveOnly),
425        (true, false) => Some(DecisionCapability::DurableResume),
426        (true, true) => Some(DecisionCapability::LiveAndDurable),
427    };
428    BackendRequirements {
429        cancellation: None,
430        continuation: request
431            .is_continuation
432            .then_some(ContinuationCapability::InProcessState),
433        decisions,
434        overrides: request
435            .overrides
436            .as_ref()
437            .map(|_| OverrideCapability::InferenceParams),
438        frontend_tools: (!request.frontend_tools.is_empty())
439            .then_some(FrontendToolCapability::DescriptorsOnly),
440        persistence: None,
441        waits: None,
442        transcript: None,
443        output: None,
444    }
445}
446
447fn delegate_request_requirements(
448    request: &BackendDelegateRunRequest<'_>,
449) -> crate::resolution::BackendRequirements {
450    use crate::resolution::{BackendRequirements, ContinuationCapability};
451    BackendRequirements {
452        cancellation: None,
453        continuation: (request.policy.continuation != BackendDelegateContinuation::Disabled)
454            .then_some(ContinuationCapability::InProcessState),
455        decisions: None,
456        overrides: None,
457        frontend_tools: None,
458        persistence: None,
459        waits: None,
460        transcript: None,
461        output: None,
462    }
463}
464
465fn check_profile(
466    profile: &crate::resolution::BackendProfile,
467    req: &crate::resolution::BackendRequirements,
468    agent_id: &str,
469) -> Result<(), ExecutionBackendError> {
470    use crate::resolution::CapabilityDecision;
471    if let CapabilityDecision::Unsupported(mismatches) = profile.check(req) {
472        let listed = mismatches
473            .iter()
474            .map(|m| m.capability)
475            .collect::<Vec<_>>()
476            .join(", ");
477        return Err(ExecutionBackendError::ExecutionFailed(format!(
478            "agent '{agent_id}' backend does not support: {listed}"
479        )));
480    }
481    Ok(())
482}
483
484pub fn validate_root_execution_request(
485    execution: &ExecutionPlan,
486    request: &BackendRootRunRequest<'_>,
487) -> Result<(), ExecutionBackendError> {
488    let profile = execution_capabilities(execution)?;
489    check_profile(
490        &profile,
491        &root_request_requirements(request),
492        request.agent_id,
493    )
494}
495
496pub fn validate_delegate_execution_request(
497    execution: &ExecutionPlan,
498    request: &BackendDelegateRunRequest<'_>,
499) -> Result<(), ExecutionBackendError> {
500    if request.policy.persistence != BackendDelegatePersistence::Ephemeral {
501        return Err(ExecutionBackendError::ExecutionFailed(format!(
502            "agent '{}' backend does not support: delegate_persistence",
503            request.agent_id
504        )));
505    }
506    if request.state_seed.is_some() && !matches!(execution, ExecutionPlan::Local(_)) {
507        return Err(ExecutionBackendError::ExecutionFailed(format!(
508            "agent '{}' backend does not support: delegate_state_seed",
509            request.agent_id
510        )));
511    }
512    let profile = execution_capabilities(execution)?;
513    check_profile(
514        &profile,
515        &delegate_request_requirements(request),
516        request.agent_id,
517    )
518}
519
520pub async fn execute_resolved_delegate_execution(
521    execution: &ExecutionPlan,
522    request: BackendDelegateRunRequest<'_>,
523) -> Result<BackendRunResult, ExecutionBackendError> {
524    validate_delegate_execution_request(execution, &request)?;
525    match execution {
526        ExecutionPlan::Local(agent) => {
527            LocalBackend::execute_resolved(agent.as_ref(), request).await
528        }
529        ExecutionPlan::Remote(agent) => agent.backend()?.execute_delegate(request).await,
530    }
531}
532
533/// Execute a remote root run including canonical runtime lifecycle events and persistence.
534pub async fn execute_remote_root_lifecycle(
535    agent: &ResolvedBackendAgent,
536    request: BackendRootRunRequest<'_>,
537    run_created_at: u64,
538    runtime_cancellation_token: CancellationToken,
539    previous_state: Option<PersistedState>,
540) -> Result<AgentRunResult, AgentLoopError> {
541    let backend = agent.backend().map_err(|error| {
542        AgentLoopError::RuntimeError(crate::RuntimeError::ResolveFailed {
543            message: error.to_string(),
544        })
545    })?;
546    let run_identity = request.run_identity.clone();
547    let sink = request.sink.clone();
548    let checkpoint_store = request.checkpoint_store;
549    let commit = request.commit;
550    let mut messages = request.messages.clone();
551    let input_message_count = messages.len();
552    let request_is_continuation = request.is_continuation;
553
554    sink.emit(AgentEvent::RunStart {
555        thread_id: run_identity.thread_id.clone(),
556        run_id: run_identity.run_id.clone(),
557        parent_run_id: run_identity.parent_run_id.clone(),
558        identity: Some(run_identity.clone()),
559    })
560    .await;
561    sink.emit(AgentEvent::StepStart {
562        message_id: gen_message_id(),
563    })
564    .await;
565
566    let execution_started_at = now_ms();
567    let backend_execution = backend.execute_root(request);
568    tokio::pin!(backend_execution);
569    let delegate_result = tokio::select! {
570        result = &mut backend_execution => {
571            match result {
572                Ok(result) => result,
573                Err(error) => {
574                    let error_message = remote_backend_error_message(error);
575                    let termination = TerminationReason::Error(error_message.clone());
576                    let latest_state = load_checkpoint_state_for_active_remote_run(
577                        checkpoint_store,
578                        &run_identity.run_id,
579                        previous_state.clone(),
580                    )
581                    .await?;
582                    return finish_remote_root_run(
583                        checkpoint_store,
584                        &run_identity.thread_id,
585                        &run_identity.run_id,
586                        &run_identity.agent_id,
587                        run_identity.parent_run_id.clone(),
588                        &run_identity,
589                        run_created_at,
590                        messages,
591                        input_message_count,
592                        BackendRunStatus::Failed(error_message),
593                        termination,
594                        None,
595                        0,
596                        String::new(),
597                        BackendRunOutput::default(),
598                        latest_state,
599                        None,
600                        commit,
601                        &sink,
602                    )
603                    .await;
604                }
605            }
606        }
607        _ = runtime_cancellation_token.cancelled() => {
608            let latest_state = load_checkpoint_state_for_active_remote_run(
609                checkpoint_store,
610                &run_identity.run_id,
611                previous_state.clone(),
612            )
613            .await?;
614            if backend.capabilities().cancellation.supports_remote_abort()
615                && let Err(error) = backend
616                    .abort(BackendAbortRequest {
617                        agent_id: &run_identity.agent_id,
618                        run_identity: &run_identity,
619                        parent: None,
620                        persisted_state: latest_state.as_ref(),
621                        is_continuation: request_is_continuation,
622                    })
623                    .await
624            {
625                tracing::warn!(
626                    agent_id = %run_identity.agent_id,
627                    run_id = %run_identity.run_id,
628                    error = %error,
629                    "non-local backend abort hook failed after cancellation"
630                );
631            }
632            return finish_remote_root_run(
633                checkpoint_store,
634                &run_identity.thread_id,
635                &run_identity.run_id,
636                &run_identity.agent_id,
637                run_identity.parent_run_id.clone(),
638                &run_identity,
639                run_created_at,
640                messages,
641                input_message_count,
642                BackendRunStatus::Cancelled,
643                TerminationReason::Cancelled,
644                None,
645                0,
646                String::new(),
647                BackendRunOutput::default(),
648                latest_state,
649                None,
650                commit,
651                &sink,
652            )
653            .await;
654        }
655    };
656
657    let termination = delegate_result.termination.clone();
658    let status_reason = delegate_result.status_reason.clone();
659    let mut output = delegate_result.output.clone();
660    let response = output
661        .text_or(&delegate_result.response)
662        .unwrap_or_default();
663    if output.text.is_none() && !response.is_empty() {
664        output.text = Some(response.clone());
665    }
666    let status = delegate_result.status;
667    let steps = delegate_result.steps;
668    let thread_state = delegate_result.thread_state;
669    let state = delegate_result.state.or(previous_state);
670    if !response.is_empty() {
671        sink.emit(AgentEvent::TextDelta {
672            delta: response.clone(),
673        })
674        .await;
675        messages.push(Message::assistant(response.clone()));
676    }
677
678    if matches!(
679        termination,
680        TerminationReason::NaturalEnd | TerminationReason::BehaviorRequested
681    ) {
682        sink.emit(AgentEvent::InferenceComplete {
683            model: agent.spec.model_id.clone(),
684            usage: None,
685            duration_ms: now_ms().saturating_sub(execution_started_at),
686        })
687        .await;
688    }
689
690    finish_remote_root_run(
691        checkpoint_store,
692        &run_identity.thread_id,
693        &run_identity.run_id,
694        &run_identity.agent_id,
695        run_identity.parent_run_id.clone(),
696        &run_identity,
697        run_created_at,
698        messages,
699        input_message_count,
700        status,
701        termination,
702        status_reason,
703        steps,
704        response,
705        output,
706        state,
707        thread_state,
708        commit,
709        &sink,
710    )
711    .await
712}
713
714async fn load_checkpoint_state(
715    storage: Option<&dyn RuntimeCheckpointStore>,
716    run_id: &str,
717    fallback: Option<PersistedState>,
718) -> Result<Option<PersistedState>, AgentLoopError> {
719    let Some(storage) = storage else {
720        return Ok(fallback);
721    };
722    match storage.load_run(run_id).await {
723        Ok(Some(run)) => Ok(run.state),
724        Ok(None) => Err(AgentLoopError::StorageError(format!(
725            "checkpoint state for run '{run_id}' was not found"
726        ))),
727        Err(error) => Err(AgentLoopError::StorageError(format!(
728            "failed to load latest checkpoint state for run '{run_id}': {error}"
729        ))),
730    }
731}
732
733async fn load_checkpoint_state_for_active_remote_run(
734    storage: Option<&dyn RuntimeCheckpointStore>,
735    run_id: &str,
736    fallback: Option<PersistedState>,
737) -> Result<Option<PersistedState>, AgentLoopError> {
738    if fallback.is_none() {
739        return Ok(None);
740    }
741    load_checkpoint_state(storage, run_id, fallback).await
742}
743
744#[allow(clippy::too_many_arguments)]
745async fn finish_remote_root_run(
746    storage: Option<&dyn RuntimeCheckpointStore>,
747    thread_id: &str,
748    run_id: &str,
749    agent_id: &str,
750    parent_run_id: Option<String>,
751    run_identity: &RunIdentity,
752    run_created_at: u64,
753    messages: Vec<Message>,
754    input_message_count: usize,
755    backend_status: BackendRunStatus,
756    termination: TerminationReason,
757    status_reason_override: Option<String>,
758    steps: usize,
759    response: String,
760    output: BackendRunOutput,
761    state: Option<PersistedState>,
762    thread_state: Option<PersistedState>,
763    commit: crate::loop_runner::CommitWiring<'_>,
764    sink: &Arc<dyn EventSink>,
765) -> Result<AgentRunResult, AgentLoopError> {
766    let status = backend_status.durable_run_status(&termination);
767    let status_reason =
768        status_reason_override.or_else(|| backend_status.durable_status_reason(&termination));
769    let state = state_with_backend_output(state, &output);
770    let mut result_json = json!({
771        "response": response,
772        "status": backend_status.result_status_label(&termination),
773    });
774    if output != BackendRunOutput::default() {
775        result_json["output"] = serde_json::to_value(&output).unwrap_or(Value::Null);
776    }
777    if let Some(reason) = &status_reason {
778        result_json["status_reason"] = Value::String(reason.clone());
779    }
780
781    let terminal_termination = status.is_terminal().then_some(termination.clone());
782    let terminal_final_output = status.is_terminal().then(|| response.clone());
783    let terminal_error_payload = status
784        .is_terminal()
785        .then(|| match &termination {
786            TerminationReason::Error(message) => Some(json!({ "message": message })),
787            _ => None,
788        })
789        .flatten();
790
791    persist_remote_root_checkpoint(
792        storage,
793        thread_id,
794        run_id,
795        agent_id,
796        parent_run_id,
797        run_created_at,
798        &messages,
799        input_message_count,
800        status,
801        terminal_termination,
802        status_reason.clone(),
803        terminal_final_output.filter(|output| !output.is_empty()),
804        terminal_error_payload,
805        run_identity,
806        steps,
807        state,
808        thread_state,
809        commit,
810    )
811    .await?;
812
813    sink.emit(AgentEvent::StepEnd).await;
814    sink.emit(AgentEvent::RunFinish {
815        thread_id: thread_id.to_string(),
816        run_id: run_id.to_string(),
817        identity: Some(run_identity.clone()),
818        result: Some(result_json),
819        termination: termination.clone(),
820    })
821    .await;
822
823    Ok(AgentRunResult {
824        run_id: run_id.to_string(),
825        response,
826        termination,
827        steps,
828    })
829}
830
831fn state_with_backend_output(
832    state: Option<PersistedState>,
833    output: &BackendRunOutput,
834) -> Option<PersistedState> {
835    if output == &BackendRunOutput::default() {
836        return state;
837    }
838
839    let mut state = state.unwrap_or(PersistedState {
840        revision: 0,
841        extensions: std::collections::HashMap::new(),
842    });
843    if let Ok(value) = serde_json::to_value(output) {
844        state
845            .extensions
846            .insert(BACKEND_OUTPUT_STATE_KEY.to_string(), value);
847    }
848    Some(state)
849}
850
851fn remote_backend_error_message(error: ExecutionBackendError) -> String {
852    match error {
853        ExecutionBackendError::AgentNotFound(message)
854        | ExecutionBackendError::ExecutionFailed(message)
855        | ExecutionBackendError::RemoteError(message) => message,
856        ExecutionBackendError::Loop(error) => error.to_string(),
857    }
858}
859
860fn run_status_label(status: RunStatus) -> &'static str {
861    match status {
862        RunStatus::Created => "created",
863        RunStatus::Running => "running",
864        RunStatus::Waiting => "waiting",
865        RunStatus::Done => "done",
866    }
867}
868
869#[cfg(test)]
870mod tests {
871    use super::*;
872    use awaken_runtime_contract::contract::storage::{RunRecord, StorageError};
873    use awaken_runtime_contract::thread::Thread;
874
875    struct FailingCheckpointStore;
876
877    #[async_trait::async_trait]
878    impl RuntimeCheckpointStore for FailingCheckpointStore {
879        async fn load_thread(&self, _thread_id: &str) -> Result<Option<Thread>, StorageError> {
880            Ok(None)
881        }
882
883        async fn load_messages(
884            &self,
885            _thread_id: &str,
886        ) -> Result<Option<Vec<Message>>, StorageError> {
887            Ok(None)
888        }
889
890        async fn load_committed_messages(
891            &self,
892            _thread_id: &str,
893        ) -> Result<Option<Vec<Message>>, StorageError> {
894            Ok(None)
895        }
896
897        async fn load_run(&self, _run_id: &str) -> Result<Option<RunRecord>, StorageError> {
898            Err(StorageError::Io("injected read failure".into()))
899        }
900
901        async fn latest_run(&self, _thread_id: &str) -> Result<Option<RunRecord>, StorageError> {
902            Ok(None)
903        }
904    }
905
906    #[test]
907    fn backend_status_timeout_is_first_class_at_runtime_boundary() {
908        let status = BackendRunStatus::Timeout;
909
910        assert_eq!(
911            status.durable_run_status(&TerminationReason::Error("polling timeout exceeded".into())),
912            RunStatus::Done
913        );
914        assert_eq!(
915            status
916                .durable_status_reason(&TerminationReason::Error("polling timeout exceeded".into()))
917                .as_deref(),
918            Some("timeout")
919        );
920        assert_eq!(
921            status
922                .result_status_label(&TerminationReason::Error("polling timeout exceeded".into())),
923            "timeout"
924        );
925    }
926
927    #[test]
928    fn backend_status_waiting_is_first_class_at_runtime_boundary() {
929        let status = BackendRunStatus::WaitingInput(Some("need details".into()));
930
931        assert_eq!(
932            status.durable_run_status(&TerminationReason::Error("should not win".into())),
933            RunStatus::Waiting
934        );
935        assert_eq!(
936            status
937                .durable_status_reason(&TerminationReason::Error("should not win".into()))
938                .as_deref(),
939            Some("input_required")
940        );
941        assert_eq!(
942            status.result_status_label(&TerminationReason::Error("should not win".into())),
943            "waiting_input"
944        );
945    }
946
947    #[tokio::test]
948    async fn load_checkpoint_state_propagates_durable_read_errors() {
949        let error = load_checkpoint_state(Some(&FailingCheckpointStore), "run-1", None)
950            .await
951            .unwrap_err();
952
953        assert!(
954            matches!(error, AgentLoopError::StorageError(ref message)
955                if message.contains("failed to load latest checkpoint state")
956                    && message.contains("injected read failure")),
957            "unexpected error: {error}"
958        );
959    }
960
961    struct MissingCheckpointStore;
962
963    #[async_trait]
964    impl RuntimeCheckpointStore for MissingCheckpointStore {
965        async fn load_thread(&self, _thread_id: &str) -> Result<Option<Thread>, StorageError> {
966            Ok(None)
967        }
968
969        async fn load_messages(
970            &self,
971            _thread_id: &str,
972        ) -> Result<Option<Vec<Message>>, StorageError> {
973            Ok(None)
974        }
975
976        async fn load_committed_messages(
977            &self,
978            _thread_id: &str,
979        ) -> Result<Option<Vec<Message>>, StorageError> {
980            Ok(None)
981        }
982
983        async fn load_run(&self, _run_id: &str) -> Result<Option<RunRecord>, StorageError> {
984            Ok(None)
985        }
986
987        async fn latest_run(&self, _thread_id: &str) -> Result<Option<RunRecord>, StorageError> {
988            Ok(None)
989        }
990    }
991
992    #[tokio::test]
993    async fn load_checkpoint_state_rejects_missing_durable_run_even_with_fallback() {
994        let fallback = PersistedState {
995            revision: 9,
996            extensions: std::collections::HashMap::new(),
997        };
998        let error = load_checkpoint_state(Some(&MissingCheckpointStore), "run-1", Some(fallback))
999            .await
1000            .unwrap_err();
1001
1002        assert!(
1003            matches!(error, AgentLoopError::StorageError(ref message)
1004                if message.contains("checkpoint state for run 'run-1' was not found")),
1005            "unexpected error: {error}"
1006        );
1007    }
1008}
awaken_runtime/backend/mod.rs

awaken_runtime/backend/
mod.rs
