awaken_runtime/loop_runner/
mod.rs

1//! Minimal sequential agent loop driven by state machines.
2//!
3//! Run lifecycle: RunLifecycle (Running → StepCompleted → Done/Waiting)
4//! Tool call lifecycle: ToolCallStates (New → Running → Succeeded/Failed/Suspended)
5
6pub(crate) mod actions;
7mod checkpoint;
8#[cfg(feature = "background")]
9mod compaction;
10mod inference;
11mod logical_inference;
12mod orchestrator;
13#[cfg(feature = "parallel-tools")]
14pub mod parallel_merge;
15mod resume;
16mod setup;
17mod step;
18mod stream_policy;
19
20#[cfg(test)]
21mod tests;
22
23use std::sync::Arc;
24
25use crate::cancellation::CancellationToken;
26use crate::checkpoint_store::RuntimeCheckpointStore;
27use crate::phase::{ExecutionEnv, PhaseRuntime};
28use crate::registry::AgentResolver;
29use crate::state::MutationBatch;
30use async_trait::async_trait;
31use awaken_runtime_contract::StateError;
32use awaken_runtime_contract::contract::event_sink::EventSink;
33use awaken_runtime_contract::contract::identity::RunIdentity;
34use awaken_runtime_contract::contract::inference::InferenceOverride;
35use awaken_runtime_contract::contract::message::{DeliveryBoundary, Message};
36use awaken_runtime_contract::contract::suspension::ToolCallResume;
37use awaken_runtime_contract::contract::tool::{ToolResult, ToolStatus};
38use futures::channel::mpsc;
39use serde_json::Value;
40
41use crate::agent::state::{RunLifecycle, ToolCallStates};
42
43// Re-export submodule items used by external callers
44pub use actions::LoopActionHandlersPlugin;
45pub use checkpoint::CommitWiring;
46pub(crate) use checkpoint::{CommitAppendError, commit_checkpoint_appending};
47pub use resume::prepare_resume;
48
49/// Plugin that registers the core state keys required by the loop runner.
50///
51/// Must be installed on the `StateStore` before running the loop.
52pub struct LoopStatePlugin;
53
54impl crate::plugins::Plugin for LoopStatePlugin {
55    fn descriptor(&self) -> crate::plugins::PluginDescriptor {
56        crate::plugins::PluginDescriptor {
57            name: "__loop_state",
58        }
59    }
60
61    fn register(
62        &self,
63        r: &mut crate::plugins::PluginRegistrar,
64    ) -> Result<(), awaken_runtime_contract::StateError> {
65        use crate::agent::state::{ContextMessageStore, ContextThrottleState};
66        use crate::state::{KeyScope, StateKeyOptions};
67
68        r.register_key::<RunLifecycle>(StateKeyOptions::default())?;
69        r.register_key::<ToolCallStates>(StateKeyOptions {
70            scope: KeyScope::Thread,
71            persistent: true,
72            ..StateKeyOptions::default()
73        })?;
74        r.register_key::<ContextThrottleState>(StateKeyOptions::default())?;
75        r.register_key::<ContextMessageStore>(StateKeyOptions::default())?;
76        r.register_key::<crate::agent::state::PendingWorkKey>(StateKeyOptions::default())?;
77
78        Ok(())
79    }
80}
81
82/// Errors from the agent loop.
83#[derive(Debug, thiserror::Error)]
84pub enum AgentLoopError {
85    #[error("inference failed: {0}")]
86    InferenceFailed(String),
87    /// Structured inference failure that preserves the upstream
88    /// [`InferenceExecutionError`](awaken_runtime_contract::contract::executor::InferenceExecutionError)
89    /// classification. Every production inference fault flows through here
90    /// (see `inference::drive_one_stream`), so downstream dispatch can consult
91    /// [`is_retryable`](awaken_runtime_contract::contract::executor::InferenceExecutionError::is_retryable)
92    /// and [`retry_after`](awaken_runtime_contract::contract::executor::InferenceExecutionError::retry_after)
93    /// to nack-with-backoff vs dead-letter, instead of blindly retrying a
94    /// permanent fault (bad credentials, exhausted quota, context overflow)
95    /// until `max_attempts` is burned.
96    #[error("inference failed: {0}")]
97    Inference(#[from] awaken_runtime_contract::contract::executor::InferenceExecutionError),
98    #[error("storage failed: {0}")]
99    StorageError(String),
100    #[error("phase error: {0}")]
101    PhaseError(#[from] awaken_runtime_contract::StateError),
102    #[error("runtime error: {0}")]
103    RuntimeError(#[from] crate::error::RuntimeError),
104    #[error("invalid activation: {0}")]
105    InvalidActivation(String),
106    #[error("invalid resume: {0}")]
107    InvalidResume(String),
108}
109
110impl From<crate::execution::executor::ToolExecutorError> for AgentLoopError {
111    fn from(e: crate::execution::executor::ToolExecutorError) -> Self {
112        Self::InferenceFailed(e.to_string())
113    }
114}
115
116/// Result of running the agent loop.
117#[derive(Debug)]
118pub struct AgentRunResult {
119    pub run_id: String,
120    pub response: String,
121    pub termination: awaken_runtime_contract::contract::lifecycle::TerminationReason,
122    pub steps: usize,
123}
124
125/// Messages frozen from durable pending state at a loop boundary.
126#[derive(Debug, Clone, Default)]
127pub struct PendingBoundaryFreeze {
128    pub messages: Vec<Message>,
129}
130
131/// Runtime callback for ADR-0042 durable pending consumption.
132///
133/// The loop runner owns `NextStep` / `OnNaturalEnd` injection points but does
134/// not own persistent thread-message storage. Mailbox/server code provides the
135/// implementation when durable pending is enabled.
136#[async_trait]
137pub trait PendingBoundaryHandler: Send + Sync {
138    async fn stage_pending_messages(
139        &self,
140        boundary: DeliveryBoundary,
141        messages: Vec<Message>,
142    ) -> Result<(), AgentLoopError>;
143
144    async fn freeze_pending_boundary(
145        &self,
146        boundary: DeliveryBoundary,
147    ) -> Result<Option<PendingBoundaryFreeze>, AgentLoopError>;
148}
149
150// -- Shared helpers --
151
152pub(crate) use awaken_runtime_contract::now_ms;
153
154fn commit_update<S: crate::state::StateKey>(
155    store: &crate::state::StateStore,
156    update: S::Update,
157) -> Result<(), awaken_runtime_contract::StateError> {
158    let mut patch = MutationBatch::new();
159    patch.update::<S>(update);
160    store.commit(patch)?;
161    clear_pending_scheduled_actions_for_terminal_run::<S>(store)?;
162    Ok(())
163}
164
165fn clear_pending_scheduled_actions_for_terminal_run<S: crate::state::StateKey>(
166    store: &crate::state::StateStore,
167) -> Result<(), awaken_runtime_contract::StateError> {
168    if S::KEY != "__runtime.run_lifecycle" {
169        return Ok(());
170    }
171    let Some(lifecycle) = store.read::<RunLifecycle>() else {
172        return Ok(());
173    };
174    if !lifecycle.status.is_terminal() {
175        return Ok(());
176    }
177    let Some(pending) = store.read::<awaken_runtime_contract::model::PendingScheduledActions>()
178    else {
179        return Ok(());
180    };
181    if pending.is_empty() {
182        return Ok(());
183    }
184
185    let mut cleanup = MutationBatch::new();
186    for action in pending {
187        cleanup.update::<awaken_runtime_contract::model::PendingScheduledActions>(
188            awaken_runtime_contract::model::ScheduledActionQueueUpdate::Remove { id: action.id },
189        );
190    }
191    store.commit(cleanup)?;
192    Ok(())
193}
194
195fn tool_result_to_content(result: &ToolResult) -> String {
196    match &result.message {
197        Some(msg) => msg.clone(),
198        None => serde_json::to_string(&result.data).unwrap_or_default(),
199    }
200}
201
202fn tool_result_to_resume_payload(result: &ToolResult) -> Value {
203    match result.status {
204        ToolStatus::Success => {
205            if result.metadata.is_empty() {
206                result.data.clone()
207            } else {
208                serde_json::json!({
209                    "data": result.data,
210                    "metadata": result.metadata,
211                })
212            }
213        }
214        ToolStatus::Error => {
215            if let Some(message) = result.message.as_ref() {
216                serde_json::json!({ "error": message })
217            } else {
218                result.data.clone()
219            }
220        }
221        ToolStatus::Pending => Value::Null,
222    }
223}
224
225/// All parameters for executing the agent loop.
226pub struct AgentLoopParams<'a> {
227    /// Resolves agent IDs to config + execution environment.
228    pub resolver: &'a dyn AgentResolver,
229    /// Initial agent to resolve at loop start.
230    pub agent_id: &'a str,
231    /// Phase runtime (state store + hook executor).
232    pub runtime: &'a PhaseRuntime,
233    /// Event sink for streaming events to the caller.
234    pub sink: Arc<dyn EventSink>,
235    /// Optional persistent storage for checkpointing.
236    pub checkpoint_store: Option<&'a dyn RuntimeCheckpointStore>,
237    /// Optional commit-coordinator + canonical-draft buffer (ADR-0036).
238    pub commit: checkpoint::CommitWiring<'a>,
239    /// Messages to seed the conversation (history + new user input).
240    pub messages: Vec<Message>,
241    /// Run identity (thread, run, agent IDs).
242    pub run_identity: RunIdentity,
243    /// Cooperative cancellation token.
244    pub cancellation_token: Option<CancellationToken>,
245    /// Live decision channel for suspended tool calls (batched by sender).
246    pub decision_rx: Option<mpsc::UnboundedReceiver<Vec<(String, ToolCallResume)>>>,
247    /// Inference parameter overrides for this run.
248    pub overrides: Option<InferenceOverride>,
249    /// Frontend-defined tool descriptors to merge into the resolved agent.
250    ///
251    /// These are tools defined by the frontend (e.g. CopilotKit `useFrontendTool`)
252    /// whose execution happens client-side. They are made visible to the LLM but
253    /// have no executor — the runtime intercepts them before execution and suspends.
254    pub frontend_tools: Vec<awaken_runtime_contract::contract::tool::ToolDescriptor>,
255    /// Optional inbox receiver for background-task messages.
256    pub inbox: Option<crate::inbox::InboxReceiver>,
257    /// When `true`, the run is a continuation of a previous awaiting_tasks run.
258    /// The orchestrator emits `SetRunning` instead of `Start`.
259    pub is_continuation: bool,
260    /// Initial state to apply to the store before the first step.
261    ///
262    /// Used by child-run helpers to seed a fresh store with parent-derived
263    /// state. Restored with `UnknownKeyPolicy::Error` — every key in the
264    /// seed must be registered in the agent's plugin set.
265    pub initial_state_seed: Option<awaken_runtime_contract::state::PersistedState>,
266}
267
268/// Build an execution environment for the agent loop.
269///
270/// Injects runtime-required default plugins and conditionally adds
271/// context truncation when a policy is provided. All transforms and hooks
272/// flow through the standard plugin registration mechanism.
273///
274/// Prefer `AgentRuntime::run()` for production use.
275pub fn build_agent_env(
276    plugins: &[Arc<dyn crate::plugins::Plugin>],
277    agent: &crate::registry::ResolvedAgent,
278) -> Result<ExecutionEnv, StateError> {
279    let stop_policies = crate::policies::policies_from_specs(agent.stop_conditions());
280    let mut all_plugins = crate::registry::resolve::inject_default_plugins_with_stop_policies(
281        plugins.to_vec(),
282        agent.max_rounds(),
283        stop_policies,
284    );
285
286    if let Some(policy) = agent.context_policy() {
287        let transform_config = agent
288            .spec
289            .config::<crate::context::ContextTransformConfigKey>()
290            .unwrap_or_default();
291        all_plugins.push(Arc::new(
292            crate::context::ContextTransformPlugin::with_config(policy.clone(), transform_config),
293        ));
294    }
295
296    ExecutionEnv::from_plugins(&all_plugins, &std::collections::HashSet::new())
297}
298
299/// Execute the agent loop. Prefer `AgentRuntime::run()` for production use.
300///
301/// Handles both fresh runs and resumed runs (state-driven detection).
302/// Supports dynamic agent handoff via `ActiveAgentIdKey` re-resolve at step boundaries.
303/// Cooperative cancellation via `CancellationToken`.
304pub async fn run_agent_loop(params: AgentLoopParams<'_>) -> Result<AgentRunResult, AgentLoopError> {
305    orchestrator::run_agent_loop_impl(params, None, None).await
306}
307
308pub(crate) async fn run_agent_loop_with_pending_boundary(
309    params: AgentLoopParams<'_>,
310    thread_ctx: Option<crate::ThreadContextSnapshot>,
311    pending_boundary: Option<Arc<dyn PendingBoundaryHandler>>,
312) -> Result<AgentRunResult, AgentLoopError> {
313    orchestrator::run_agent_loop_impl(params, thread_ctx, pending_boundary).await
314}
awaken_runtime/loop_runner/mod.rs

awaken_runtime/loop_runner/
mod.rs
