theway_core/
agent.rs

1//! The single-agent runtime. The bare `Agent` state machine (this file) is always on —
2//! `prompt()` / `continue()` / `subscribe()` / `abort()`, no harness dependency. The
3//! harness layer (skills, sessions, compaction, permission, …) lives in the submodules
4//! below, behind `#[cfg(feature = "harness")]` (opt-out for embedders that only want the
5//! bare Agent). Orchestration builds on top in `crate::multiagent`.
6//!
7//! Implemented:
8//! - State container + getters/setters (Mutex-protected)
9//! - Listener subscription with unsubscribe fn
10//! - `prompt(...)` / `continue_()` driving the agent loop
11//! - `abort()` via `tokio_util::sync::CancellationToken`
12//! - Steering / follow-up queues (`enqueue_steering` / `enqueue_follow_up`)
13//!
14//! TODO:
15//! - `onPayload` / `onResponse` SimpleStreamOptions surface
16//! - `transformContext` & `getApiKey` hooks (declared, wired up later)
17//! - `prepareNextTurn` model/thinking-level rewrite mid-run
18
19// Harness layer (feature-gated): the bare Agent stays always-on.
20#[cfg(feature = "harness")]
21pub mod assembly;
22#[cfg(feature = "harness")]
23pub mod compaction;
24pub mod context;
25pub mod context_cache;
26#[cfg(feature = "harness")]
27pub mod cost;
28#[cfg(feature = "harness")]
29pub mod messages;
30pub mod model_request;
31#[cfg(feature = "harness")]
32pub mod permission;
33#[cfg(feature = "harness")]
34pub mod runtime_extensions;
35// The loop engine is part of the bare Agent (prompt()/continue_() call it) — always on.
36pub mod run_loop;
37#[cfg(feature = "harness")]
38pub mod session;
39#[cfg(feature = "harness")]
40pub mod skills;
41#[cfg(feature = "harness")]
42pub mod system_prompt;
43#[cfg(feature = "harness")]
44pub mod types;
45use std::sync::Arc;
46
47use parking_lot::Mutex;
48use tokio::sync::{Notify, broadcast};
49use tokio_util::sync::CancellationToken;
50
51use crate::agent::run_loop::{run_agent_loop, run_agent_loop_continue};
52use crate::observability::{
53    ObservationContext, OperationId, RuntimeObserver, noop_runtime_observer,
54};
55use crate::types::*;
56
57use theway_llm_provider::Message;
58
59/// Async listener for lifecycle events. Receives an event and the active cancellation token
60/// for the run. Used for subscribers that need to perform I/O (e.g. session persistence).
61/// For memory-only, sub-microsecond operations prefer [`LoopSyncCallback`].
62pub type LoopListener = Arc<
63    dyn Fn(
64            LoopEvent,
65            CancellationToken,
66        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>
67        + Send
68        + Sync,
69>;
70
71/// Lightweight synchronous callback for lifecycle events. MUST complete in <1µs — no I/O,
72/// no blocking, no allocation beyond simple atomic/counter updates. Each callback is wrapped
73/// in `catch_unwind` during emission so a panic in one does not affect others.
74pub type LoopSyncCallback = Arc<dyn Fn(&LoopEvent) + Send + Sync>;
75
76/// Capacity of the [`LoopEvent`] broadcast channel.
77pub const LOOP_EVENT_BROADCAST_CAPACITY: usize = 256;
78
79/// Options accepted by [`Agent::new`].
80pub struct AgentOptions {
81    pub initial_state: Option<AgentState>,
82    pub convert_to_llm: Option<ConvertToLlm>,
83    pub transform_context: Option<TransformContext>,
84    pub transform_model_request: Option<TransformModelRequest>,
85    pub transform_message: Option<TransformMessage>,
86    pub provider_request_interceptor: Option<theway_llm_provider::ProviderRequestInterceptorHandle>,
87    pub stream_fn: Option<StreamFn>,
88    pub get_api_key: Option<GetApiKey>,
89    pub before_tool_call: Option<BeforeToolCallHook>,
90    pub after_tool_call: Option<AfterToolCallHook>,
91    /// Final tool-result transform after execution-end observation and before
92    /// construction of the model-visible tool-result message.
93    pub transform_tool_result: Option<AfterToolCallHook>,
94    pub on_control_plane_prompt: Option<OnControlPlanePromptHook>,
95    pub should_stop_after_turn: Option<ShouldStopHook>,
96    pub prepare_next_turn: Option<PrepareNextTurnHook>,
97    pub steering_mode: QueueMode,
98    pub follow_up_mode: QueueMode,
99    pub session_id: Option<String>,
100    /// Content-safe runtime observation port supplied by the embedding application.
101    pub observer: Arc<dyn RuntimeObserver>,
102    /// Correlation values inherited by operations created by this agent.
103    pub observation_context: ObservationContext,
104    /// Optional parent operation supplied by an embedding runtime or subagent launcher.
105    pub observation_parent: Option<OperationId>,
106    pub tool_execution: ToolExecutionMode,
107    /// Hard cap on loop iterations (one LLM turn attempt each) for this agent.
108    /// `None` = unbounded (the interactive main agent). Sub-harnesses
109    /// (`subagent` tool, DAG nodes, goal evaluator) set it from their spec's
110    /// `max_iterations`; the cap raises `AgentRunError::Other("max iterations
111    /// (N) exceeded")` before the call that would exceed it.
112    pub max_iterations: Option<u32>,
113}
114
115impl Default for AgentOptions {
116    fn default() -> Self {
117        Self {
118            initial_state: None,
119            convert_to_llm: None,
120            transform_context: None,
121            transform_model_request: None,
122            transform_message: None,
123            provider_request_interceptor: None,
124            stream_fn: None,
125            get_api_key: None,
126            before_tool_call: None,
127            after_tool_call: None,
128            transform_tool_result: None,
129            on_control_plane_prompt: None,
130            should_stop_after_turn: None,
131            prepare_next_turn: None,
132            steering_mode: QueueMode::default(),
133            follow_up_mode: QueueMode::default(),
134            session_id: None,
135            observer: noop_runtime_observer(),
136            observation_context: ObservationContext::default(),
137            observation_parent: None,
138            tool_execution: ToolExecutionMode::default(),
139            max_iterations: None,
140        }
141    }
142}
143
144/// Stateful wrapper around the low-level agent loop.
145pub struct Agent {
146    inner: Arc<AgentInner>,
147}
148
149pub(crate) struct AgentInner {
150    /// Serializes admission and cleanup for one active prompt/continue run.
151    pub run_active: Mutex<bool>,
152    pub state: Mutex<AgentState>,
153    /// Segment 1: synchronous callbacks (memory-only, <1µs). Each wrapped in `catch_unwind`.
154    pub sync_callbacks: Mutex<Vec<LoopSyncCallback>>,
155    /// Segment 2: async await-listeners (persistence, I/O). Emitted sequentially.
156    pub await_listeners: Mutex<Vec<LoopListener>>,
157    /// Segment 3: broadcast channel for external subscribers (UI, gRPC, hooks). Non-blocking send.
158    pub broadcast_tx: broadcast::Sender<LoopEvent>,
159    pub steering: Mutex<PendingMessageQueue>,
160    pub follow_up: Mutex<PendingMessageQueue>,
161    pub options: AgentOptions,
162    pub active_cancel: Mutex<Option<CancellationToken>>,
163    /// Current observation hierarchy. One run is active at a time; parallel tools read the
164    /// same turn parent without mutating it.
165    pub active_run_operation: Mutex<Option<OperationId>>,
166    pub active_turn_operation: Mutex<Option<(OperationId, u32)>>,
167    /// Per-turn cancel token: `interrupt()` cancels the in-flight LLM call only;
168    /// the run survives if a steering message is queued, otherwise it ends.
169    pub turn_cancel: Mutex<Option<CancellationToken>>,
170    pub idle: Notify,
171    /// Hard cap on loop iterations (see [`AgentOptions::max_iterations`]).
172    pub max_iterations: Option<u32>,
173    /// Per-session client-side prefix cache tracker (core-owned, daemon reads
174    /// the results through assistant message `Usage`).
175    pub context_cache: Mutex<crate::agent::context_cache::ContextCacheTracker>,
176}
177
178pub(crate) struct AgentRunPermit {
179    inner: Arc<AgentInner>,
180}
181
182impl AgentRunPermit {
183    pub(crate) fn acquire(inner: Arc<AgentInner>) -> Result<Self, AgentRunError> {
184        let mut active = inner.run_active.lock();
185        if *active {
186            return Err(AgentRunError::AlreadyStreaming);
187        }
188        *active = true;
189        {
190            let mut state = inner.state.lock();
191            state.is_streaming = true;
192            state.error_message = None;
193        }
194        drop(active);
195        Ok(Self { inner })
196    }
197}
198
199impl Drop for AgentRunPermit {
200    fn drop(&mut self) {
201        self.inner.release_run();
202    }
203}
204
205impl AgentInner {
206    pub(crate) fn release_run(&self) {
207        let mut active = self.run_active.lock();
208        if !*active {
209            return;
210        }
211        *self.active_cancel.lock() = None;
212        *self.active_run_operation.lock() = None;
213        *self.active_turn_operation.lock() = None;
214        *self.turn_cancel.lock() = None;
215        self.state.lock().is_streaming = false;
216        *active = false;
217        drop(active);
218        self.idle.notify_waiters();
219    }
220}
221
222pub(crate) struct PendingMessageQueue {
223    mode: QueueMode,
224    items: Vec<AgentMessage>,
225}
226
227impl PendingMessageQueue {
228    fn new(mode: QueueMode) -> Self {
229        Self {
230            mode,
231            items: Vec::new(),
232        }
233    }
234
235    pub fn enqueue(&mut self, m: AgentMessage) {
236        self.items.push(m);
237    }
238
239    pub fn drain(&mut self) -> Vec<AgentMessage> {
240        match self.mode {
241            QueueMode::All => std::mem::take(&mut self.items),
242            QueueMode::OneAtATime => {
243                if self.items.is_empty() {
244                    Vec::new()
245                } else {
246                    vec![self.items.remove(0)]
247                }
248            }
249        }
250    }
251}
252
253impl Agent {
254    pub fn new(mut options: AgentOptions) -> Self {
255        let state = options.initial_state.take().unwrap_or_default();
256        if options.convert_to_llm.is_none() {
257            options.convert_to_llm = Some(default_convert_to_llm());
258        }
259        let max_iterations = options.max_iterations;
260        let (broadcast_tx, _) = broadcast::channel(LOOP_EVENT_BROADCAST_CAPACITY);
261        let inner = AgentInner {
262            run_active: Mutex::new(false),
263            state: Mutex::new(state),
264            sync_callbacks: Mutex::new(Vec::new()),
265            await_listeners: Mutex::new(Vec::new()),
266            broadcast_tx,
267            steering: Mutex::new(PendingMessageQueue::new(options.steering_mode)),
268            follow_up: Mutex::new(PendingMessageQueue::new(options.follow_up_mode)),
269            options,
270            active_cancel: Mutex::new(None),
271            active_run_operation: Mutex::new(None),
272            active_turn_operation: Mutex::new(None),
273            turn_cancel: Mutex::new(None),
274            idle: Notify::new(),
275            max_iterations,
276            context_cache: Mutex::new(crate::agent::context_cache::ContextCacheTracker::new()),
277        };
278        Self {
279            inner: Arc::new(inner),
280        }
281    }
282
283    /// Subscribe an async listener (segment 2 — await path). For persistence/I/O subscribers
284    /// that need the cancellation token. Returns an unsubscribe closure.
285    ///
286    /// For memory-only callbacks (<1µs), use [`Self::subscribe_sync`]. For external
287    /// subscribers that want a broadcast [`tokio::sync::broadcast::Receiver`], use
288    /// [`Self::subscribe_broadcast`].
289    pub fn subscribe(&self, listener: LoopListener) -> impl FnOnce() {
290        let inner = self.inner.clone();
291        inner.await_listeners.lock().push(listener.clone());
292        move || {
293            let mut listeners = inner.await_listeners.lock();
294            if let Some(pos) = listeners.iter().position(|l| Arc::ptr_eq(l, &listener)) {
295                listeners.remove(pos);
296            }
297        }
298    }
299
300    /// Register a synchronous callback (segment 1 — catch_unwind path). The callback MUST
301    /// complete in <1µs — no I/O, no blocking. Returns an unsubscribe closure.
302    pub fn subscribe_sync(&self, callback: LoopSyncCallback) -> impl FnOnce() {
303        let inner = self.inner.clone();
304        inner.sync_callbacks.lock().push(callback.clone());
305        move || {
306            let mut cbs = inner.sync_callbacks.lock();
307            if let Some(pos) = cbs.iter().position(|c| Arc::ptr_eq(c, &callback)) {
308                cbs.remove(pos);
309            }
310        }
311    }
312
313    /// Obtain a new [`tokio::sync::broadcast::Receiver`] for the LoopEvent broadcast
314    /// channel (segment 3). The receiver sees all events emitted after subscription.
315    pub fn subscribe_broadcast(&self) -> broadcast::Receiver<LoopEvent> {
316        self.inner.broadcast_tx.subscribe()
317    }
318
319    /// Inspect the current agent state. The lock guards against concurrent loop mutations.
320    pub fn state(&self) -> parking_lot::MutexGuard<'_, AgentState> {
321        self.inner.state.lock()
322    }
323
324    pub fn is_streaming(&self) -> bool {
325        self.inner.state.lock().is_streaming
326    }
327
328    /// Return the embedder-owned runtime observer used by this agent.
329    pub fn runtime_observer(&self) -> Arc<dyn RuntimeObserver> {
330        Arc::clone(&self.inner.options.observer)
331    }
332
333    /// Return the content-safe correlation context inherited by this agent.
334    pub fn observation_context(&self) -> ObservationContext {
335        self.inner.options.observation_context.clone()
336    }
337
338    /// Return the current agent-run operation, if a prompt is active.
339    pub fn active_run_operation(&self) -> Option<OperationId> {
340        *self.inner.active_run_operation.lock()
341    }
342
343    pub fn enqueue_steering(&self, message: AgentMessage) {
344        self.inner.steering.lock().enqueue(message);
345    }
346
347    pub fn enqueue_follow_up(&self, message: AgentMessage) {
348        self.inner.follow_up.lock().enqueue(message);
349    }
350
351    /// Abort the active run, if any. Subsequent calls are no-ops.
352    pub fn abort(&self) {
353        if let Some(token) = self.inner.active_cancel.lock().as_ref() {
354            token.cancel();
355        }
356    }
357
358    /// Interrupt the current turn: cancels the in-flight LLM call. The run ends
359    /// unless a steering message is queued (then the next turn carries it).
360    pub fn interrupt(&self) {
361        if let Some(token) = self.inner.turn_cancel.lock().as_ref() {
362            token.cancel();
363        }
364    }
365
366    /// Active cancellation token while a run is in flight, otherwise `None`.
367    pub fn active_token(&self) -> Option<CancellationToken> {
368        self.inner.active_cancel.lock().clone()
369    }
370
371    /// Wait until the active run has released admission and all awaited loop
372    /// listeners have completed.
373    pub async fn wait_until_idle(&self) {
374        loop {
375            let notified = self.inner.idle.notified();
376            if !self.is_streaming() {
377                return;
378            }
379            notified.await;
380        }
381    }
382
383    /// Start a new prompt. Appends a user `AgentMessage`, runs the loop, awaits completion.
384    pub async fn prompt(&self, message: AgentMessage) -> Result<(), AgentRunError> {
385        self.prompt_many(vec![message]).await
386    }
387
388    /// Start a new prompt with a batch of messages.
389    pub async fn prompt_many(&self, messages: Vec<AgentMessage>) -> Result<(), AgentRunError> {
390        run_agent_loop(self.inner.clone(), messages).await
391    }
392
393    /// Continue from the current transcript without appending new user messages.
394    pub async fn continue_(&self) -> Result<(), AgentRunError> {
395        run_agent_loop_continue(self.inner.clone()).await
396    }
397}
398
399/// Errors that can short-circuit `prompt` / `continue_`.
400#[derive(Debug, thiserror::Error)]
401pub enum AgentRunError {
402    #[error(
403        "Agent is already processing a prompt. Use enqueue_steering/enqueue_follow_up or wait for completion."
404    )]
405    AlreadyStreaming,
406    /// The current turn was interrupted via [`Agent::interrupt`] and no steering
407    /// message was queued, so the run ended at the turn boundary.
408    #[error("turn interrupted")]
409    TurnInterrupted,
410    #[error("{0}")]
411    Other(String),
412}
413
414impl AgentInner {
415    pub fn convert_to_llm(&self, msgs: &[AgentMessage]) -> Vec<Message> {
416        self.options
417            .convert_to_llm
418            .as_ref()
419            .expect("convert_to_llm is always set in Agent::new")(msgs)
420    }
421}
422
423#[cfg(test)]
424tests_bridge_macro::tests_bridge!("agent");
theway_core/agent.rs

theway_core/
agent.rs
