modular_agent_core/
agent.rs

1use std::any::Any;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use serde_json::Value;
6
7use crate::config::AgentConfigs;
8use crate::context::AgentContext;
9use crate::error::AgentError;
10use crate::modular_agent::ModularAgent;
11use crate::runtime::runtime;
12use crate::spec::AgentSpec;
13use crate::value::AgentValue;
14
15/// The lifecycle status of an agent.
16#[derive(Debug, Default, Clone, PartialEq)]
17pub enum AgentStatus {
18    #[default]
19    Init,
20    Start,
21    Stop,
22}
23
24/// Internal messages sent to agents.
25pub(crate) enum AgentMessage {
26    /// Input value received on a port.
27    Input {
28        ctx: AgentContext,
29        port: String,
30        value: AgentValue,
31    },
32
33    /// Configuration value update.
34    Config { key: String, value: AgentValue },
35
36    /// Full configuration update.
37    Configs { configs: AgentConfigs },
38
39    /// Stop the agent.
40    Stop,
41}
42
43/// The core trait for all agents.
44///
45/// All agents implement this trait. Defines lifecycle management,
46/// configuration access, and message processing.
47#[async_trait]
48pub trait Agent: Send + Sync + 'static {
49    /// Constructs a new agent instance.
50    fn new(ma: ModularAgent, id: String, spec: AgentSpec) -> Result<Self, AgentError>
51    where
52        Self: Sized;
53
54    /// Returns the `ModularAgent`.
55    fn ma(&self) -> &ModularAgent;
56
57    /// Returns the unique agent ID.
58    fn id(&self) -> &str;
59
60    /// Returns the current lifecycle status.
61    fn status(&self) -> &AgentStatus;
62
63    /// Returns the agent specification.
64    fn spec(&self) -> &AgentSpec;
65
66    /// Updates the agent specification.
67    fn update_spec(&mut self, spec_update: &Value) -> Result<(), AgentError>;
68
69    /// Returns the agent definition name.
70    fn def_name(&self) -> &str;
71
72    /// Returns the agent's configuration.
73    ///
74    /// # Errors
75    ///
76    /// Returns `NoConfig` if no configuration is available.
77    fn configs(&self) -> Result<&AgentConfigs, AgentError>;
78
79    /// Sets a configuration value.
80    fn set_config(&mut self, key: String, value: AgentValue) -> Result<(), AgentError>;
81
82    /// Sets the entire configuration.
83    fn set_configs(&mut self, configs: AgentConfigs) -> Result<(), AgentError>;
84
85    /// Gets global configuration for this agent.
86    fn get_global_configs(&self) -> Option<AgentConfigs> {
87        self.ma().get_global_configs(self.def_name())
88    }
89
90    /// Returns the preset ID this agent belongs to.
91    fn preset_id(&self) -> &str;
92
93    /// Sets the preset ID.
94    fn set_preset_id(&mut self, preset_id: String);
95
96    /// Starts the agent.
97    ///
98    /// Called when the workflow starts. Use for initialization and initial output.
99    async fn start(&mut self) -> Result<(), AgentError>;
100
101    /// Stops the agent.
102    async fn stop(&mut self) -> Result<(), AgentError>;
103
104    /// Processes an input message.
105    ///
106    /// Called when the agent receives a value on an input port.
107    async fn process(
108        &mut self,
109        ctx: AgentContext,
110        port: String,
111        value: AgentValue,
112    ) -> Result<(), AgentError>;
113
114    /// Returns the tokio runtime.
115    fn runtime(&self) -> &tokio::runtime::Runtime {
116        runtime()
117    }
118
119    fn as_any(&self) -> &dyn Any;
120
121    fn as_any_mut(&mut self) -> &mut dyn Any;
122}
123
124impl dyn Agent {
125    pub fn as_agent<T: Agent>(&self) -> Option<&T> {
126        self.as_any().downcast_ref::<T>()
127    }
128
129    pub fn as_agent_mut<T: Agent>(&mut self) -> Option<&mut T> {
130        self.as_any_mut().downcast_mut::<T>()
131    }
132}
133
134/// Core data structure for an agent.
135///
136/// Used by agents implementing `AsAgent` to store common state.
137/// The `#[modular_agent]` macro generates a struct with this as a field.
138pub struct AgentData {
139    /// The ModularAgent instance.
140    pub ma: ModularAgent,
141
142    /// The unique identifier for this agent.
143    pub id: String,
144
145    /// The specification of the agent (definition, config, etc.).
146    pub spec: AgentSpec,
147
148    /// The preset identifier for the agent.
149    /// Empty string when the agent does not belong to any preset.
150    pub preset_id: String,
151
152    /// The current lifecycle status of the agent.
153    pub status: AgentStatus,
154}
155
156impl AgentData {
157    /// Creates a new `AgentData` instance.
158    ///
159    /// Removes any `_`-prefixed config keys that were preserved by
160    /// `AgentDefinition::reconcile_spec()` for lazy migration.
161    /// Agents can read these keys from the `spec` parameter in `AsAgent::new()`
162    /// before calling this method.
163    pub fn new(ma: ModularAgent, id: String, mut spec: AgentSpec) -> Self {
164        if let Some(ref mut configs) = spec.configs {
165            configs.retain(|key, _| !key.starts_with('_'));
166        }
167        Self {
168            ma,
169            id,
170            spec,
171            preset_id: String::new(),
172            status: AgentStatus::Init,
173        }
174    }
175}
176
177/// Trait for types that contain `AgentData`.
178///
179/// Required by `AsAgent`. Usually implemented automatically via `#[modular_agent]` macro.
180pub trait HasAgentData {
181    fn data(&self) -> &AgentData;
182
183    fn mut_data(&mut self) -> &mut AgentData;
184}
185
186/// Simplified trait for implementing custom agents.
187///
188/// Implement this trait instead of `Agent` directly.
189/// The `Agent` trait is automatically implemented for all types that implement `AsAgent`.
190///
191/// # Cancellation safety
192///
193/// The agent loop races [`process()`](Self::process) against the agent's
194/// cancellation token, which fires when the agent (or its whole preset) is
195/// stopped. On cancellation the in-flight `process()` future is **dropped at
196/// whatever await point it has reached** — implementations must not rely on
197/// running to completion. In particular, outputs emitted before the drop
198/// stay emitted, and internal bookkeeping updated across await points (e.g.
199/// entries in a pending map) may be left behind; keep such state consistent
200/// at every await point or clean it up in [`stop()`](Self::stop).
201///
202/// Flow-level aborts ([`ModularAgent::abort_context`](crate::ModularAgent::abort_context))
203/// are cooperative: the context's token fires, but messages carrying it are
204/// still delivered so that wind-down outputs (e.g. an aborted final message
205/// replacing a dangling partial in history) can traverse the graph.
206/// Implementations that initiate external work (network requests, DB writes,
207/// message posts) must therefore check
208/// [`ctx.is_cancelled()`](crate::AgentContext::is_cancelled) before starting
209/// it, and may select on
210/// [`ctx.cancel_token()`](crate::AgentContext::cancel_token) at long awaits
211/// to wind down gracefully.
212#[async_trait]
213pub trait AsAgent: HasAgentData + Send + Sync + 'static {
214    /// Constructs a new agent instance.
215    fn new(ma: ModularAgent, id: String, spec: AgentSpec) -> Result<Self, AgentError>
216    where
217        Self: Sized;
218
219    /// Called when configuration values change.
220    ///
221    /// Override to react to configuration changes at runtime.
222    fn configs_changed(&mut self) -> Result<(), AgentError> {
223        Ok(())
224    }
225
226    /// Called when the agent starts.
227    ///
228    /// Override for initialization logic or to emit initial values.
229    async fn start(&mut self) -> Result<(), AgentError> {
230        Ok(())
231    }
232
233    /// Called when the agent stops.
234    ///
235    /// Override for cleanup logic.
236    async fn stop(&mut self) -> Result<(), AgentError> {
237        Ok(())
238    }
239
240    /// Processes an input message.
241    ///
242    /// Override to implement the agent's main logic.
243    ///
244    /// This method may be cancelled by being dropped at any await point (see
245    /// the [trait-level docs](AsAgent#cancellation-safety)). Long-running
246    /// implementations can additionally observe
247    /// [`AgentContext::cancel_token`] to abort gracefully when the flow is
248    /// cancelled via [`ModularAgent::abort_context`], recording the
249    /// interruption as [`AgentError::Cancelled`].
250    async fn process(
251        &mut self,
252        _ctx: AgentContext,
253        _port: String,
254        _value: AgentValue,
255    ) -> Result<(), AgentError> {
256        Ok(())
257    }
258}
259
260#[async_trait]
261impl<T: AsAgent> Agent for T {
262    fn new(ma: ModularAgent, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
263        let mut agent = T::new(ma, id, spec)?;
264        agent.mut_data().status = AgentStatus::Init;
265        Ok(agent)
266    }
267
268    fn ma(&self) -> &ModularAgent {
269        &self.data().ma
270    }
271
272    fn id(&self) -> &str {
273        &self.data().id
274    }
275
276    fn spec(&self) -> &AgentSpec {
277        &self.data().spec
278    }
279
280    fn update_spec(&mut self, value: &Value) -> Result<(), AgentError> {
281        self.mut_data().spec.update(value)
282    }
283
284    fn status(&self) -> &AgentStatus {
285        &self.data().status
286    }
287
288    fn def_name(&self) -> &str {
289        self.data().spec.def_name.as_str()
290    }
291
292    fn configs(&self) -> Result<&AgentConfigs, AgentError> {
293        self.data()
294            .spec
295            .configs
296            .as_ref()
297            .ok_or(AgentError::NoConfig)
298    }
299
300    fn set_config(&mut self, key: String, value: AgentValue) -> Result<(), AgentError> {
301        if let Some(configs) = &mut self.mut_data().spec.configs {
302            configs.set(key, value);
303            self.configs_changed()?;
304        }
305        Ok(())
306    }
307
308    fn set_configs(&mut self, configs: AgentConfigs) -> Result<(), AgentError> {
309        self.mut_data().spec.configs = Some(configs);
310        self.configs_changed()
311    }
312
313    fn preset_id(&self) -> &str {
314        &self.data().preset_id
315    }
316
317    fn set_preset_id(&mut self, preset_id: String) {
318        self.mut_data().preset_id = preset_id;
319    }
320
321    async fn start(&mut self) -> Result<(), AgentError> {
322        self.mut_data().status = AgentStatus::Start;
323
324        if let Err(e) = <T as AsAgent>::start(self).await {
325            self.ma()
326                .emit_agent_error(self.id().to_string(), e.to_string());
327            return Err(e);
328        }
329
330        Ok(())
331    }
332
333    async fn stop(&mut self) -> Result<(), AgentError> {
334        self.mut_data().status = AgentStatus::Stop;
335        <T as AsAgent>::stop(self).await?;
336        self.mut_data().status = AgentStatus::Init;
337        Ok(())
338    }
339
340    async fn process(
341        &mut self,
342        ctx: AgentContext,
343        port: String,
344        value: AgentValue,
345    ) -> Result<(), AgentError> {
346        if let Err(e) = <T as AsAgent>::process(self, ctx.clone(), port, value).await {
347            self.ma()
348                .emit_agent_error(self.id().to_string(), e.to_string());
349            self.ma()
350                .send_agent_out(
351                    self.id().to_string(),
352                    ctx,
353                    "err".to_string(),
354                    AgentValue::Error(Arc::new(e.clone())),
355                )
356                .await
357                .unwrap_or_else(|e| {
358                    log::error!("Failed to send error message for {}: {}", self.id(), e);
359                });
360            return Err(e);
361        }
362        Ok(())
363    }
364
365    fn get_global_configs(&self) -> Option<AgentConfigs> {
366        self.ma().get_global_configs(self.def_name())
367    }
368
369    fn as_any(&self) -> &dyn Any {
370        self
371    }
372
373    fn as_any_mut(&mut self) -> &mut dyn Any {
374        self
375    }
376}
377
378/// Creates a boxed agent instance from a concrete type.
379#[doc(hidden)]
380pub fn new_agent_boxed<T: Agent>(
381    ma: ModularAgent,
382    id: String,
383    spec: AgentSpec,
384) -> Result<Box<dyn Agent>, AgentError> {
385    Ok(Box::new(T::new(ma, id, spec)?))
386}
387
388/// Creates an agent based on its definition.
389///
390/// Looks up the agent definition by name and calls the appropriate constructor.
391pub(crate) fn agent_new(
392    ma: ModularAgent,
393    agent_id: String,
394    mut spec: AgentSpec,
395) -> Result<Box<dyn Agent>, AgentError> {
396    let def;
397    {
398        let def_name = &spec.def_name;
399        let defs = ma.defs.lock().unwrap();
400        def = defs
401            .get(def_name)
402            .ok_or_else(|| AgentError::UnknownDefName(def_name.to_string()))?
403            .clone();
404    }
405
406    def.reconcile_spec(&mut spec);
407
408    if let Some(new_boxed) = def.new_boxed {
409        return new_boxed(ma, agent_id, spec);
410    }
411
412    Err(AgentError::UnknownDefKind(def.kind.to_string()))
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418    use crate::config::AgentConfigs;
419    use crate::value::AgentValue;
420
421    #[test]
422    fn test_agent_data_new_strips_prefixed_keys() {
423        let ma = ModularAgent::init().unwrap();
424        let mut configs = AgentConfigs::new();
425        configs.set("name".into(), AgentValue::string("hello"));
426        configs.set("count".into(), AgentValue::integer(10));
427        configs.set("_old_key".into(), AgentValue::string("stale"));
428        configs.set("_removed".into(), AgentValue::integer(42));
429
430        let spec = AgentSpec {
431            configs: Some(configs),
432            ..Default::default()
433        };
434
435        let data = AgentData::new(ma.clone(), "test_id".into(), spec);
436
437        let c = data.spec.configs.as_ref().unwrap();
438        assert_eq!(c.get_string_or_default("name"), "hello");
439        assert_eq!(c.get_integer_or_default("count"), 10);
440        assert!(c.get("_old_key").is_err());
441        assert!(c.get("_removed").is_err());
442
443        ma.quit();
444    }
445}
modular_agent_core/agent.rs

modular_agent_core/
agent.rs
