Skip to main content

anda_core/
agent.rs

1//! Agent traits and registries.
2//!
3//! This module defines how custom AI agents are described, registered, and
4//! invoked by an Anda runtime. It provides:
5//! - [`Agent`] for strongly typed agent implementations.
6//! - [`DynAgent`] for runtime dispatch through trait objects.
7//! - [`AgentSet`] for name-based registration and lookup.
8//!
9//! Agents may declare tool dependencies and supported resource tags. The
10//! runtime uses those declarations to validate engine configuration and route
11//! resource attachments to the components that can consume them.
12//!
13//! See the `anda_engine` extension modules for concrete agent implementations.
14
15use serde::{Deserialize, Serialize};
16use serde_json::json;
17use std::{any::Any, collections::BTreeMap, future::Future, marker::PhantomData, sync::Arc};
18
19use crate::{
20    BoxError, BoxPinFut, Function,
21    context::AgentContext,
22    model::{AgentOutput, FunctionDefinition, Resource},
23    select_resources, validate_function_name,
24};
25
26/// Default JSON arguments for an agent exposed as a callable function.
27#[derive(Debug, Clone, Deserialize, Serialize)]
28pub struct AgentArgs {
29    /// Self-contained task prompt for the agent.
30    pub prompt: String,
31}
32
33/// Strongly typed interface for an AI agent.
34///
35/// # Type Parameters
36/// - `C`: Runtime context implementing [`AgentContext`].
37pub trait Agent<C>: Send + Sync
38where
39    C: AgentContext + Send + Sync,
40{
41    /// Returns the unique agent name.
42    ///
43    /// Names are registered case-insensitively and stored in lowercase.
44    ///
45    /// # Rules
46    /// - Must not be empty;
47    /// - Must not exceed 64 characters;
48    /// - Must start with a lowercase letter;
49    /// - Can only contain: lowercase letters (a-z), digits (0-9), and underscores (_);
50    /// - Unique within the engine in lowercase.
51    fn name(&self) -> String;
52
53    /// Returns a concise description of the agent's capability.
54    fn description(&self) -> String;
55
56    /// Returns the function definition used for LLM/tool-call integration.
57    ///
58    /// # Returns
59    /// - `FunctionDefinition`: The structured definition of the agent's capabilities.
60    fn definition(&self) -> FunctionDefinition {
61        FunctionDefinition {
62            name: self.name().to_ascii_lowercase(),
63            description: self.description(),
64            parameters: json!({
65                "type": "object",
66                "description": "Run this agent on a focused task. Provide a self-contained prompt with the goal, relevant context, constraints, and expected output.",
67                "properties": {
68                    "prompt": {
69                        "type": "string",
70                        "description": "The task for this agent. Include the objective, relevant context, constraints, preferred workflow or deliverable, and any success criteria needed to complete the work.",
71                        "minLength": 1
72                    },
73                },
74                "required": ["prompt"],
75                "additionalProperties": false
76            }),
77            strict: Some(true),
78        }
79    }
80
81    /// Returns resource tags this agent can consume.
82    ///
83    /// The default implementation returns an empty list, meaning no resources
84    /// are selected for this agent. Return `vec!["*".into()]` to accept all
85    /// attached resources.
86    ///
87    /// # Returns
88    /// Resource tags supported by this agent.
89    fn supported_resource_tags(&self) -> Vec<String> {
90        Vec::new()
91    }
92
93    /// Removes and returns resources matching this agent's supported tags.
94    fn select_resources(&self, resources: &mut Vec<Resource>) -> Vec<Resource> {
95        let supported_tags = self.supported_resource_tags();
96        select_resources(resources, &supported_tags)
97    }
98
99    /// Initializes the agent with the given context.
100    ///
101    /// Runtimes call this once while building the engine.
102    fn init(&self, _ctx: C) -> impl Future<Output = Result<(), BoxError>> + Send {
103        futures::future::ready(Ok(()))
104    }
105
106    /// Returns tool names required by this agent.
107    ///
108    /// Runtimes use this list to validate that required tools are registered.
109    fn tool_dependencies(&self) -> Vec<String> {
110        Vec::new()
111    }
112
113    /// Executes the agent with the given context and inputs.
114    ///
115    /// # Arguments
116    /// - `ctx`: The execution context implementing [`AgentContext`].
117    /// - `prompt`: The input prompt or message for the agent.
118    /// - `resources`: Additional resources selected for this agent. Ignore resources that are not useful.
119    ///
120    /// # Returns
121    /// A future resolving to [`AgentOutput`].
122    fn run(
123        &self,
124        ctx: C,
125        prompt: String,
126        resources: Vec<Resource>,
127    ) -> impl Future<Output = Result<AgentOutput, BoxError>> + Send;
128}
129
130/// Object-safe wrapper around [`Agent`] for runtime dispatch.
131///
132/// Runtime registries store agents through this trait so callers can select and
133/// execute agents by name without knowing their concrete Rust types.
134pub trait DynAgent<C>: Send + Sync
135where
136    C: AgentContext + Send + Sync,
137{
138    fn as_any(&self) -> &(dyn Any + Send + Sync);
139
140    fn into_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;
141
142    fn label(&self) -> &str;
143
144    fn name(&self) -> String;
145
146    fn definition(&self) -> FunctionDefinition;
147
148    fn tool_dependencies(&self) -> Vec<String>;
149
150    fn supported_resource_tags(&self) -> Vec<String>;
151
152    fn init(&self, ctx: C) -> BoxPinFut<Result<(), BoxError>>;
153
154    fn run(
155        &self,
156        ctx: C,
157        prompt: String,
158        resources: Vec<Resource>,
159    ) -> BoxPinFut<Result<AgentOutput, BoxError>>;
160}
161
162impl<C> dyn DynAgent<C>
163where
164    C: AgentContext + Send + Sync + 'static,
165{
166    /// Returns the inner concrete agent type when it matches `T`.
167    pub fn downcast_ref<T>(&self) -> Option<&T>
168    where
169        T: Agent<C> + 'static,
170    {
171        self.as_any().downcast_ref::<T>()
172    }
173
174    /// Returns the inner concrete agent when it matches `T`.
175    pub fn downcast<T>(self: Arc<Self>) -> Result<Arc<T>, Arc<Self>>
176    where
177        T: Agent<C> + 'static,
178    {
179        match self.clone().into_any().downcast::<T>() {
180            Ok(agent) => Ok(agent),
181            Err(_) => Err(self),
182        }
183    }
184}
185
186/// Adapter that exposes a concrete [`Agent`] through [`DynAgent`].
187struct AgentWrapper<T, C>
188where
189    T: Agent<C> + 'static,
190    C: AgentContext + Send + Sync + 'static,
191{
192    inner: Arc<T>,
193    label: String,
194    _phantom: PhantomData<C>,
195}
196
197impl<T, C> DynAgent<C> for AgentWrapper<T, C>
198where
199    T: Agent<C> + 'static,
200    C: AgentContext + Send + Sync + 'static,
201{
202    fn as_any(&self) -> &(dyn Any + Send + Sync) {
203        self.inner.as_ref()
204    }
205
206    fn into_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
207        self.inner.clone()
208    }
209
210    fn label(&self) -> &str {
211        &self.label
212    }
213
214    fn name(&self) -> String {
215        self.inner.name()
216    }
217
218    fn definition(&self) -> FunctionDefinition {
219        self.inner.definition()
220    }
221
222    fn tool_dependencies(&self) -> Vec<String> {
223        self.inner.tool_dependencies()
224    }
225
226    fn supported_resource_tags(&self) -> Vec<String> {
227        self.inner.supported_resource_tags()
228    }
229
230    fn init(&self, ctx: C) -> BoxPinFut<Result<(), BoxError>> {
231        let agent = self.inner.clone();
232        Box::pin(async move { agent.init(ctx).await })
233    }
234
235    fn run(
236        &self,
237        ctx: C,
238        prompt: String,
239        resources: Vec<Resource>,
240    ) -> BoxPinFut<Result<AgentOutput, BoxError>> {
241        let agent = self.inner.clone();
242        Box::pin(async move { agent.run(ctx, prompt, resources).await })
243    }
244}
245
246/// Name-based registry for agents.
247///
248/// # Type Parameters
249/// - `C`: The context type that implements [`AgentContext`].
250#[derive(Default)]
251pub struct AgentSet<C: AgentContext> {
252    pub set: BTreeMap<String, Arc<dyn DynAgent<C>>>,
253}
254
255impl<C> AgentSet<C>
256where
257    C: AgentContext + Send + Sync + 'static,
258{
259    /// Creates a new empty AgentSet.
260    pub fn new() -> Self {
261        Self {
262            set: BTreeMap::new(),
263        }
264    }
265
266    /// Returns whether an agent with the given name exists.
267    pub fn contains(&self, name: &str) -> bool {
268        self.set.contains_key(&name.to_ascii_lowercase())
269    }
270
271    /// Returns whether an agent with the given lowercase name exists.
272    pub fn contains_lowercase(&self, lowercase_name: &str) -> bool {
273        self.set.contains_key(lowercase_name)
274    }
275
276    /// Returns the names of all agents in the set.
277    pub fn names(&self) -> Vec<String> {
278        self.set.keys().cloned().collect()
279    }
280
281    /// Returns the function definition for a specific agent.
282    pub fn definition(&self, name: &str) -> Option<FunctionDefinition> {
283        self.set
284            .get(&name.to_ascii_lowercase())
285            .map(|agent| agent.definition())
286    }
287
288    /// Returns function definitions for all agents or the selected names.
289    ///
290    /// # Arguments
291    /// - `names`: Optional slice of agent names to filter by.
292    ///
293    /// # Returns
294    /// A vector of agent definitions.
295    pub fn definitions(&self, names: Option<&[String]>) -> Vec<FunctionDefinition> {
296        match names {
297            None => self.set.values().map(|agent| agent.definition()).collect(),
298            Some(names) => names
299                .iter()
300                .filter_map(|name| {
301                    self.set
302                        .get(&name.to_ascii_lowercase())
303                        .map(|agent| agent.definition())
304                })
305                .collect(),
306        }
307    }
308
309    /// Returns function metadata for all agents or the selected names.
310    ///
311    /// # Arguments
312    /// - `names`: Optional slice of agent names to filter by.
313    ///
314    /// # Returns
315    /// A vector of agent function metadata.
316    pub fn functions(&self, names: Option<&[String]>) -> Vec<Function> {
317        match names {
318            None => self
319                .set
320                .values()
321                .map(|agent| Function {
322                    definition: agent.definition(),
323                    supported_resource_tags: agent.supported_resource_tags(),
324                })
325                .collect(),
326            Some(names) => names
327                .iter()
328                .filter_map(|name| {
329                    self.set
330                        .get(&name.to_ascii_lowercase())
331                        .map(|agent| Function {
332                            definition: agent.definition(),
333                            supported_resource_tags: agent.supported_resource_tags(),
334                        })
335                })
336                .collect(),
337        }
338    }
339
340    /// Removes and returns resources supported by the named agent.
341    pub fn select_resources(&self, name: &str, resources: &mut Vec<Resource>) -> Vec<Resource> {
342        if resources.is_empty() {
343            return Vec::new();
344        }
345
346        self.set
347            .get(&name.to_ascii_lowercase())
348            .map(|agent| {
349                let supported_tags = agent.supported_resource_tags();
350                select_resources(resources, &supported_tags)
351            })
352            .unwrap_or_default()
353    }
354
355    /// Registers a new agent.
356    ///
357    /// # Arguments
358    /// - `agent`: The agent to register.
359    pub fn add<T>(&mut self, agent: Arc<T>, label: Option<String>) -> Result<(), BoxError>
360    where
361        T: Agent<C> + Send + Sync + 'static,
362    {
363        let name = agent.name().to_ascii_lowercase();
364        validate_function_name(&name)?;
365        if self.set.contains_key(&name) {
366            return Err(format!("agent {} already exists", name).into());
367        }
368
369        let agent_dyn = AgentWrapper {
370            inner: agent,
371            label: label.unwrap_or_else(|| name.clone()),
372            _phantom: PhantomData,
373        };
374        self.set.insert(name, Arc::new(agent_dyn));
375        Ok(())
376    }
377
378    /// Returns an agent by name.
379    pub fn get(&self, name: &str) -> Option<Arc<dyn DynAgent<C>>> {
380        self.set.get(&name.to_ascii_lowercase()).cloned()
381    }
382
383    /// Returns an agent by lowercase name.
384    pub fn get_lowercase(&self, lowercase_name: &str) -> Option<Arc<dyn DynAgent<C>>> {
385        self.set.get(lowercase_name).cloned()
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392    use candid::{CandidType, Principal, utils::ArgumentEncoder};
393    use std::time::Duration;
394
395    use crate::{
396        AgentInput, BaseContext, CacheExpiry, CacheFeatures, CancellationToken, CanisterCaller,
397        CompletionFeatures, CompletionRequest, HttpFeatures, Json, KeysFeatures, ObjectMeta, Path,
398        PutMode, PutResult, RequestMeta, StateFeatures, StoreFeatures, ToolInput, ToolOutput,
399    };
400
401    #[derive(Clone)]
402    struct TestAgentContext {
403        engine_id: Principal,
404        caller: Principal,
405        meta: RequestMeta,
406        cancellation_token: CancellationToken,
407    }
408
409    impl Default for TestAgentContext {
410        fn default() -> Self {
411            Self {
412                engine_id: Principal::management_canister(),
413                caller: Principal::anonymous(),
414                meta: RequestMeta::default(),
415                cancellation_token: CancellationToken::new(),
416            }
417        }
418    }
419
420    impl StateFeatures for TestAgentContext {
421        fn engine_id(&self) -> &Principal {
422            &self.engine_id
423        }
424
425        fn engine_name(&self) -> &str {
426            "test-engine"
427        }
428
429        fn caller(&self) -> &Principal {
430            &self.caller
431        }
432
433        fn meta(&self) -> &RequestMeta {
434            &self.meta
435        }
436
437        fn cancellation_token(&self) -> CancellationToken {
438            self.cancellation_token.clone()
439        }
440
441        fn time_elapsed(&self) -> Duration {
442            Duration::ZERO
443        }
444    }
445
446    impl KeysFeatures for TestAgentContext {
447        async fn a256gcm_key(&self, _derivation_path: Vec<Vec<u8>>) -> Result<[u8; 32], BoxError> {
448            Ok([0; 32])
449        }
450
451        async fn ed25519_sign_message(
452            &self,
453            _derivation_path: Vec<Vec<u8>>,
454            _message: &[u8],
455        ) -> Result<[u8; 64], BoxError> {
456            Ok([0; 64])
457        }
458
459        async fn ed25519_verify(
460            &self,
461            _derivation_path: Vec<Vec<u8>>,
462            _message: &[u8],
463            _signature: &[u8],
464        ) -> Result<(), BoxError> {
465            Ok(())
466        }
467
468        async fn ed25519_public_key(
469            &self,
470            _derivation_path: Vec<Vec<u8>>,
471        ) -> Result<[u8; 32], BoxError> {
472            Ok([0; 32])
473        }
474
475        async fn secp256k1_sign_message_bip340(
476            &self,
477            _derivation_path: Vec<Vec<u8>>,
478            _message: &[u8],
479        ) -> Result<[u8; 64], BoxError> {
480            Ok([0; 64])
481        }
482
483        async fn secp256k1_verify_bip340(
484            &self,
485            _derivation_path: Vec<Vec<u8>>,
486            _message: &[u8],
487            _signature: &[u8],
488        ) -> Result<(), BoxError> {
489            Ok(())
490        }
491
492        async fn secp256k1_sign_message_ecdsa(
493            &self,
494            _derivation_path: Vec<Vec<u8>>,
495            _message: &[u8],
496        ) -> Result<[u8; 64], BoxError> {
497            Ok([0; 64])
498        }
499
500        async fn secp256k1_sign_digest_ecdsa(
501            &self,
502            _derivation_path: Vec<Vec<u8>>,
503            _message_hash: &[u8],
504        ) -> Result<[u8; 64], BoxError> {
505            Ok([0; 64])
506        }
507
508        async fn secp256k1_verify_ecdsa(
509            &self,
510            _derivation_path: Vec<Vec<u8>>,
511            _message_hash: &[u8],
512            _signature: &[u8],
513        ) -> Result<(), BoxError> {
514            Ok(())
515        }
516
517        async fn secp256k1_public_key(
518            &self,
519            _derivation_path: Vec<Vec<u8>>,
520        ) -> Result<[u8; 33], BoxError> {
521            Ok([0; 33])
522        }
523    }
524
525    impl StoreFeatures for TestAgentContext {
526        async fn store_get(&self, _path: &Path) -> Result<(bytes::Bytes, ObjectMeta), BoxError> {
527            Err("not implemented".into())
528        }
529
530        async fn store_list(
531            &self,
532            _prefix: Option<&Path>,
533            _offset: &Path,
534        ) -> Result<Vec<ObjectMeta>, BoxError> {
535            Ok(Vec::new())
536        }
537
538        async fn store_put(
539            &self,
540            _path: &Path,
541            _mode: PutMode,
542            _value: bytes::Bytes,
543        ) -> Result<PutResult, BoxError> {
544            Err("not implemented".into())
545        }
546
547        async fn store_rename_if_not_exists(
548            &self,
549            _from: &Path,
550            _to: &Path,
551        ) -> Result<(), BoxError> {
552            Err("not implemented".into())
553        }
554
555        async fn store_delete(&self, _path: &Path) -> Result<(), BoxError> {
556            Ok(())
557        }
558    }
559
560    impl CacheFeatures for TestAgentContext {
561        fn cache_contains(&self, _key: &str) -> bool {
562            false
563        }
564
565        async fn cache_get<T>(&self, _key: &str) -> Result<T, BoxError>
566        where
567            T: serde::de::DeserializeOwned,
568        {
569            Err("not implemented".into())
570        }
571
572        async fn cache_get_with<T, F>(&self, _key: &str, _init: F) -> Result<T, BoxError>
573        where
574            T: Sized + serde::de::DeserializeOwned + Serialize + Send,
575            F: Future<Output = Result<(T, Option<CacheExpiry>), BoxError>> + Send + 'static,
576        {
577            Err("not implemented".into())
578        }
579
580        async fn cache_set<T>(&self, _key: &str, _val: (T, Option<CacheExpiry>))
581        where
582            T: Sized + Serialize + Send,
583        {
584        }
585
586        async fn cache_set_if_not_exists<T>(
587            &self,
588            _key: &str,
589            _val: (T, Option<CacheExpiry>),
590        ) -> bool
591        where
592            T: Sized + Serialize + Send,
593        {
594            false
595        }
596
597        async fn cache_delete(&self, _key: &str) -> bool {
598            false
599        }
600
601        fn cache_raw_iter(
602            &self,
603        ) -> impl Iterator<Item = (Arc<String>, Arc<(bytes::Bytes, Option<CacheExpiry>)>)> {
604            std::iter::empty()
605        }
606    }
607
608    impl HttpFeatures for TestAgentContext {
609        async fn https_call(
610            &self,
611            _url: &str,
612            _method: http::Method,
613            _headers: Option<http::HeaderMap>,
614            _body: Option<Vec<u8>>,
615        ) -> Result<reqwest::Response, BoxError> {
616            Err("not implemented".into())
617        }
618
619        async fn https_signed_call(
620            &self,
621            _url: &str,
622            _method: http::Method,
623            _message_digest: [u8; 32],
624            _headers: Option<http::HeaderMap>,
625            _body: Option<Vec<u8>>,
626        ) -> Result<reqwest::Response, BoxError> {
627            Err("not implemented".into())
628        }
629
630        async fn https_signed_rpc<T>(
631            &self,
632            _endpoint: &str,
633            _method: &str,
634            _args: impl Serialize + Send,
635        ) -> Result<T, BoxError>
636        where
637            T: serde::de::DeserializeOwned,
638        {
639            Err("not implemented".into())
640        }
641    }
642
643    impl crate::CanisterCaller for TestAgentContext {
644        async fn canister_query<In, Out>(
645            &self,
646            _canister: &Principal,
647            _method: &str,
648            _args: In,
649        ) -> Result<Out, BoxError>
650        where
651            In: ArgumentEncoder + Send,
652            Out: CandidType + for<'a> candid::Deserialize<'a>,
653        {
654            Err("not implemented".into())
655        }
656
657        async fn canister_update<In, Out>(
658            &self,
659            _canister: &Principal,
660            _method: &str,
661            _args: In,
662        ) -> Result<Out, BoxError>
663        where
664            In: ArgumentEncoder + Send,
665            Out: CandidType + for<'a> candid::Deserialize<'a>,
666        {
667            Err("not implemented".into())
668        }
669    }
670
671    impl crate::BaseContext for TestAgentContext {
672        async fn remote_tool_call(
673            &self,
674            _endpoint: &str,
675            _args: ToolInput<Json>,
676        ) -> Result<ToolOutput<Json>, BoxError> {
677            Err("not implemented".into())
678        }
679    }
680
681    impl CompletionFeatures for TestAgentContext {
682        async fn completion(
683            &self,
684            _req: CompletionRequest,
685            _resources: Vec<Resource>,
686        ) -> Result<AgentOutput, BoxError> {
687            Ok(AgentOutput::default())
688        }
689
690        fn model_name(&self) -> String {
691            "test-model".to_string()
692        }
693    }
694
695    impl AgentContext for TestAgentContext {
696        fn tool_definitions(&self, _names: Option<&[String]>) -> Vec<FunctionDefinition> {
697            Vec::new()
698        }
699
700        async fn remote_tool_definitions(
701            &self,
702            _endpoint: Option<&str>,
703            _names: Option<&[String]>,
704        ) -> Result<Vec<FunctionDefinition>, BoxError> {
705            Ok(Vec::new())
706        }
707
708        async fn select_tool_resources(
709            &self,
710            _name: &str,
711            _resources: &mut Vec<Resource>,
712        ) -> Vec<Resource> {
713            Vec::new()
714        }
715
716        fn agent_definitions(&self, _names: Option<&[String]>) -> Vec<FunctionDefinition> {
717            Vec::new()
718        }
719
720        async fn remote_agent_definitions(
721            &self,
722            _endpoint: Option<&str>,
723            _names: Option<&[String]>,
724        ) -> Result<Vec<FunctionDefinition>, BoxError> {
725            Ok(Vec::new())
726        }
727
728        async fn select_agent_resources(
729            &self,
730            _name: &str,
731            _resources: &mut Vec<Resource>,
732        ) -> Vec<Resource> {
733            Vec::new()
734        }
735
736        async fn definitions(&self, _names: Option<&[String]>) -> Vec<FunctionDefinition> {
737            Vec::new()
738        }
739
740        async fn tool_call(
741            &self,
742            _args: ToolInput<Json>,
743        ) -> Result<(ToolOutput<Json>, Option<Principal>), BoxError> {
744            Ok((ToolOutput::new(Json::Null), None))
745        }
746
747        async fn agent_run(
748            self,
749            _args: AgentInput,
750        ) -> Result<(AgentOutput, Option<Principal>), BoxError> {
751            Ok((AgentOutput::default(), None))
752        }
753
754        async fn remote_agent_run(
755            &self,
756            _endpoint: &str,
757            _args: AgentInput,
758        ) -> Result<AgentOutput, BoxError> {
759            Ok(AgentOutput::default())
760        }
761    }
762
763    struct ExampleAgent {
764        id: usize,
765    }
766
767    struct OtherAgent;
768
769    struct TaggedAgent;
770
771    struct InvalidAgent;
772
773    fn resource(id: u64, tags: &[&str]) -> Resource {
774        Resource {
775            _id: id,
776            name: format!("resource-{id}"),
777            tags: tags.iter().map(|tag| tag.to_string()).collect(),
778            ..Default::default()
779        }
780    }
781
782    impl Agent<TestAgentContext> for ExampleAgent {
783        fn name(&self) -> String {
784            "example_agent".to_string()
785        }
786
787        fn description(&self) -> String {
788            "Example agent used for downcast tests".to_string()
789        }
790
791        async fn run(
792            &self,
793            _ctx: TestAgentContext,
794            _prompt: String,
795            _resources: Vec<Resource>,
796        ) -> Result<AgentOutput, BoxError> {
797            Ok(AgentOutput {
798                content: self.id.to_string(),
799                ..AgentOutput::default()
800            })
801        }
802    }
803
804    impl Agent<TestAgentContext> for OtherAgent {
805        fn name(&self) -> String {
806            "other_agent".to_string()
807        }
808
809        fn description(&self) -> String {
810            "Other agent used for downcast tests".to_string()
811        }
812
813        async fn run(
814            &self,
815            _ctx: TestAgentContext,
816            _prompt: String,
817            _resources: Vec<Resource>,
818        ) -> Result<AgentOutput, BoxError> {
819            Ok(AgentOutput {
820                content: "other".to_string(),
821                ..AgentOutput::default()
822            })
823        }
824    }
825
826    impl Agent<TestAgentContext> for TaggedAgent {
827        fn name(&self) -> String {
828            "tagged_agent".to_string()
829        }
830
831        fn description(&self) -> String {
832            "Agent that consumes text and code resources".to_string()
833        }
834
835        fn supported_resource_tags(&self) -> Vec<String> {
836            vec!["text".to_string(), "code".to_string()]
837        }
838
839        fn tool_dependencies(&self) -> Vec<String> {
840            vec!["lookup".to_string(), "summarize".to_string()]
841        }
842
843        async fn run(
844            &self,
845            _ctx: TestAgentContext,
846            prompt: String,
847            resources: Vec<Resource>,
848        ) -> Result<AgentOutput, BoxError> {
849            Ok(AgentOutput {
850                content: format!("{prompt}:{}", resources.len()),
851                ..AgentOutput::default()
852            })
853        }
854    }
855
856    impl Agent<TestAgentContext> for InvalidAgent {
857        fn name(&self) -> String {
858            "bad-agent".to_string()
859        }
860
861        fn description(&self) -> String {
862            "Invalid function name".to_string()
863        }
864
865        async fn run(
866            &self,
867            _ctx: TestAgentContext,
868            _prompt: String,
869            _resources: Vec<Resource>,
870        ) -> Result<AgentOutput, BoxError> {
871            Ok(AgentOutput::default())
872        }
873    }
874
875    #[test]
876    fn dyn_agent_downcast_ref_returns_inner_agent() {
877        let agent = Arc::new(ExampleAgent { id: 7 });
878        let mut agent_set = AgentSet::<TestAgentContext>::new();
879        agent_set
880            .add(agent, Some("test-label".to_string()))
881            .unwrap();
882
883        let dyn_agent = agent_set.get("example_agent").unwrap();
884        let concrete = dyn_agent.downcast_ref::<ExampleAgent>().unwrap();
885
886        assert_eq!(concrete.id, 7);
887        assert!(dyn_agent.downcast_ref::<OtherAgent>().is_none());
888    }
889
890    #[test]
891    fn dyn_agent_downcast_returns_original_arc() {
892        let agent = Arc::new(ExampleAgent { id: 9 });
893        let mut agent_set = AgentSet::<TestAgentContext>::new();
894        agent_set
895            .add(agent.clone(), Some("test-label".to_string()))
896            .unwrap();
897
898        let dyn_agent = agent_set.get("example_agent").unwrap();
899        let concrete = dyn_agent
900            .downcast::<ExampleAgent>()
901            .ok()
902            .expect("expected downcast to ExampleAgent to succeed");
903
904        assert_eq!(concrete.id, 9);
905        assert!(Arc::ptr_eq(&concrete, &agent));
906    }
907
908    #[test]
909    fn dyn_agent_downcast_mismatch_returns_original_arc() {
910        let agent = Arc::new(ExampleAgent { id: 11 });
911        let mut agent_set = AgentSet::<TestAgentContext>::new();
912        agent_set
913            .add(agent, Some("test-label".to_string()))
914            .unwrap();
915
916        let dyn_agent = agent_set.get("example_agent").unwrap();
917        let original = dyn_agent.clone();
918        let err = dyn_agent
919            .downcast::<OtherAgent>()
920            .err()
921            .expect("expected downcast to OtherAgent to fail");
922
923        assert!(Arc::ptr_eq(&err, &original));
924        assert_eq!(err.name(), "example_agent");
925        assert_eq!(err.label(), "test-label");
926    }
927
928    #[test]
929    fn agent_default_methods_and_dyn_wrapper_forward_calls() {
930        futures::executor::block_on(async {
931            let agent = Arc::new(ExampleAgent { id: 42 });
932            let mut resources = vec![resource(1, &["text"])];
933
934            let definition = agent.definition();
935            assert_eq!(definition.name, "example_agent");
936            assert_eq!(definition.description, agent.description());
937            assert_eq!(definition.strict, Some(true));
938            assert_eq!(definition.parameters["type"], "object");
939            assert_eq!(
940                definition.parameters["required"].as_array().unwrap()[0],
941                "prompt"
942            );
943            assert!(agent.supported_resource_tags().is_empty());
944            assert!(agent.select_resources(&mut resources).is_empty());
945            assert_eq!(resources.len(), 1);
946            agent.init(TestAgentContext::default()).await.unwrap();
947            assert!(agent.tool_dependencies().is_empty());
948
949            let mut agent_set = AgentSet::<TestAgentContext>::new();
950            agent_set
951                .add(agent, Some("example label".to_string()))
952                .unwrap();
953            let dyn_agent = agent_set.get("EXAMPLE_AGENT").unwrap();
954
955            assert_eq!(dyn_agent.label(), "example label");
956            assert_eq!(dyn_agent.name(), "example_agent");
957            assert_eq!(dyn_agent.definition().name, "example_agent");
958            assert!(dyn_agent.tool_dependencies().is_empty());
959            assert!(dyn_agent.supported_resource_tags().is_empty());
960            dyn_agent.init(TestAgentContext::default()).await.unwrap();
961
962            let output = dyn_agent
963                .run(
964                    TestAgentContext::default(),
965                    "ignored".to_string(),
966                    Vec::new(),
967                )
968                .await
969                .unwrap();
970            assert_eq!(output.content, "42");
971        });
972    }
973
974    #[test]
975    fn fixture_agents_cover_direct_trait_methods() {
976        futures::executor::block_on(async {
977            let other = OtherAgent;
978            assert_eq!(other.name(), "other_agent");
979            assert_eq!(other.description(), "Other agent used for downcast tests");
980            assert_eq!(
981                other
982                    .run(
983                        TestAgentContext::default(),
984                        "prompt".to_string(),
985                        Vec::new()
986                    )
987                    .await
988                    .unwrap()
989                    .content,
990                "other"
991            );
992
993            let tagged = TaggedAgent;
994            assert_eq!(
995                tagged.tool_dependencies(),
996                vec!["lookup".to_string(), "summarize".to_string()]
997            );
998
999            let invalid = InvalidAgent;
1000            assert_eq!(invalid.name(), "bad-agent");
1001            assert_eq!(invalid.description(), "Invalid function name");
1002            assert!(
1003                invalid
1004                    .run(
1005                        TestAgentContext::default(),
1006                        "prompt".to_string(),
1007                        Vec::new()
1008                    )
1009                    .await
1010                    .unwrap()
1011                    .content
1012                    .is_empty()
1013            );
1014        });
1015    }
1016
1017    #[test]
1018    fn agent_set_registry_filters_resources_and_reports_errors() {
1019        futures::executor::block_on(async {
1020            let mut agent_set = AgentSet::<TestAgentContext>::new();
1021            agent_set
1022                .add(Arc::new(ExampleAgent { id: 1 }), None)
1023                .unwrap();
1024            agent_set
1025                .add(Arc::new(TaggedAgent), Some("tagged label".to_string()))
1026                .unwrap();
1027
1028            assert!(agent_set.contains("EXAMPLE_AGENT"));
1029            assert!(agent_set.contains_lowercase("tagged_agent"));
1030            assert!(!agent_set.contains("missing_agent"));
1031            assert_eq!(
1032                agent_set.names(),
1033                vec!["example_agent".to_string(), "tagged_agent".to_string()]
1034            );
1035
1036            let definition = agent_set.definition("TAGGED_AGENT").unwrap();
1037            assert_eq!(definition.name, "tagged_agent");
1038            assert!(agent_set.definition("missing_agent").is_none());
1039
1040            let selected_names = vec!["TAGGED_AGENT".to_string(), "missing_agent".to_string()];
1041            let selected_definitions = agent_set.definitions(Some(&selected_names));
1042            assert_eq!(selected_definitions.len(), 1);
1043            assert_eq!(selected_definitions[0].name, "tagged_agent");
1044            assert_eq!(agent_set.definitions(None).len(), 2);
1045
1046            let selected_functions = agent_set.functions(Some(&selected_names));
1047            assert_eq!(selected_functions.len(), 1);
1048            assert_eq!(
1049                selected_functions[0].supported_resource_tags,
1050                vec!["text".to_string(), "code".to_string()]
1051            );
1052            assert_eq!(agent_set.functions(None).len(), 2);
1053
1054            let mut empty = Vec::new();
1055            assert!(
1056                agent_set
1057                    .select_resources("tagged_agent", &mut empty)
1058                    .is_empty()
1059            );
1060            let mut resources = vec![
1061                resource(1, &["image"]),
1062                resource(2, &["text"]),
1063                resource(3, &["code", "text"]),
1064                resource(4, &["audio"]),
1065            ];
1066            let selected = agent_set.select_resources("TAGGED_AGENT", &mut resources);
1067            assert_eq!(
1068                selected
1069                    .iter()
1070                    .map(|resource| resource._id)
1071                    .collect::<Vec<_>>(),
1072                vec![2, 3]
1073            );
1074            assert_eq!(
1075                resources
1076                    .iter()
1077                    .map(|resource| resource._id)
1078                    .collect::<Vec<_>>(),
1079                vec![1, 4]
1080            );
1081            assert!(
1082                agent_set
1083                    .select_resources("missing_agent", &mut resources)
1084                    .is_empty()
1085            );
1086
1087            let dyn_agent = agent_set.get_lowercase("tagged_agent").unwrap();
1088            assert_eq!(dyn_agent.label(), "tagged label");
1089            let output = dyn_agent
1090                .run(
1091                    TestAgentContext::default(),
1092                    "prompt".to_string(),
1093                    vec![resource(9, &["text"])],
1094                )
1095                .await
1096                .unwrap();
1097            assert_eq!(output.content, "prompt:1");
1098            assert!(agent_set.get("missing_agent").is_none());
1099            assert!(agent_set.get_lowercase("missing_agent").is_none());
1100
1101            let duplicate = agent_set
1102                .add(Arc::new(ExampleAgent { id: 2 }), None)
1103                .unwrap_err();
1104            assert!(duplicate.to_string().contains("already exists"));
1105
1106            let invalid = agent_set.add(Arc::new(InvalidAgent), None).unwrap_err();
1107            assert!(invalid.to_string().contains("invalid character"));
1108        });
1109    }
1110
1111    #[test]
1112    fn test_agent_context_mock_features_cover_default_paths() {
1113        futures::executor::block_on(async {
1114            let ctx = TestAgentContext::default();
1115            assert_eq!(*ctx.engine_id(), Principal::management_canister());
1116            assert_eq!(ctx.engine_name(), "test-engine");
1117            assert_eq!(*ctx.caller(), Principal::anonymous());
1118            assert!(ctx.meta().user.is_none());
1119            assert!(!ctx.cancellation_token().is_cancelled());
1120            assert_eq!(ctx.time_elapsed(), Duration::ZERO);
1121
1122            assert_eq!(ctx.a256gcm_key(Vec::new()).await.unwrap(), [0; 32]);
1123            assert_eq!(
1124                ctx.ed25519_sign_message(Vec::new(), b"message")
1125                    .await
1126                    .unwrap(),
1127                [0; 64]
1128            );
1129            ctx.ed25519_verify(Vec::new(), b"message", &[0; 64])
1130                .await
1131                .unwrap();
1132            assert_eq!(ctx.ed25519_public_key(Vec::new()).await.unwrap(), [0; 32]);
1133            assert_eq!(
1134                ctx.secp256k1_sign_message_bip340(Vec::new(), b"message")
1135                    .await
1136                    .unwrap(),
1137                [0; 64]
1138            );
1139            ctx.secp256k1_verify_bip340(Vec::new(), b"message", &[0; 64])
1140                .await
1141                .unwrap();
1142            assert_eq!(
1143                ctx.secp256k1_sign_message_ecdsa(Vec::new(), b"message")
1144                    .await
1145                    .unwrap(),
1146                [0; 64]
1147            );
1148            assert_eq!(
1149                ctx.secp256k1_sign_digest_ecdsa(Vec::new(), &[0; 32])
1150                    .await
1151                    .unwrap(),
1152                [0; 64]
1153            );
1154            ctx.secp256k1_verify_ecdsa(Vec::new(), &[0; 32], &[0; 64])
1155                .await
1156                .unwrap();
1157            assert_eq!(ctx.secp256k1_public_key(Vec::new()).await.unwrap(), [0; 33]);
1158
1159            assert!(ctx.store_get(&Path::from("missing")).await.is_err());
1160            assert!(
1161                ctx.store_list(None, &Path::default())
1162                    .await
1163                    .unwrap()
1164                    .is_empty()
1165            );
1166            assert!(
1167                ctx.store_put(&Path::from("file"), PutMode::Overwrite, bytes::Bytes::new())
1168                    .await
1169                    .is_err()
1170            );
1171            assert!(
1172                ctx.store_rename_if_not_exists(&Path::from("a"), &Path::from("b"))
1173                    .await
1174                    .is_err()
1175            );
1176            ctx.store_delete(&Path::from("file")).await.unwrap();
1177
1178            assert!(!ctx.cache_contains("key"));
1179            assert!(ctx.cache_get::<String>("key").await.is_err());
1180            assert!(
1181                ctx.cache_get_with("key", async { Ok(("value".to_string(), None)) })
1182                    .await
1183                    .is_err()
1184            );
1185            ctx.cache_set("key", ("value".to_string(), None)).await;
1186            assert!(
1187                !ctx.cache_set_if_not_exists("key", ("value".to_string(), None))
1188                    .await
1189            );
1190            assert!(!ctx.cache_delete("key").await);
1191            assert_eq!(ctx.cache_raw_iter().count(), 0);
1192
1193            assert!(
1194                ctx.https_call("https://example.test", http::Method::GET, None, None)
1195                    .await
1196                    .is_err()
1197            );
1198            assert!(
1199                ctx.https_signed_call(
1200                    "https://example.test",
1201                    http::Method::POST,
1202                    [0; 32],
1203                    None,
1204                    None,
1205                )
1206                .await
1207                .is_err()
1208            );
1209            let rpc: Result<String, BoxError> = ctx
1210                .https_signed_rpc("https://example.test", "method", &())
1211                .await;
1212            assert!(rpc.is_err());
1213
1214            let query: Result<String, BoxError> = ctx
1215                .canister_query(&Principal::anonymous(), "query", ())
1216                .await;
1217            assert!(query.is_err());
1218            let update: Result<String, BoxError> = ctx
1219                .canister_update(&Principal::anonymous(), "update", ())
1220                .await;
1221            assert!(update.is_err());
1222
1223            assert!(
1224                ctx.remote_tool_call(
1225                    "https://example.test",
1226                    ToolInput::new("tool".to_string(), Json::Null),
1227                )
1228                .await
1229                .is_err()
1230            );
1231            assert_eq!(
1232                ctx.completion(CompletionRequest::default(), Vec::new())
1233                    .await
1234                    .unwrap()
1235                    .content,
1236                ""
1237            );
1238            assert_eq!(ctx.model_name(), "test-model");
1239            assert!(ctx.tool_definitions(None).is_empty());
1240            assert!(
1241                ctx.remote_tool_definitions(None, None)
1242                    .await
1243                    .unwrap()
1244                    .is_empty()
1245            );
1246            assert!(
1247                ctx.select_tool_resources("tool", &mut Vec::new())
1248                    .await
1249                    .is_empty()
1250            );
1251            assert!(ctx.agent_definitions(None).is_empty());
1252            assert!(
1253                ctx.remote_agent_definitions(None, None)
1254                    .await
1255                    .unwrap()
1256                    .is_empty()
1257            );
1258            assert!(
1259                ctx.select_agent_resources("agent", &mut Vec::new())
1260                    .await
1261                    .is_empty()
1262            );
1263            assert!(ctx.definitions(None).await.is_empty());
1264            assert!(
1265                ctx.tool_call(ToolInput::new("tool".to_string(), Json::Null))
1266                    .await
1267                    .unwrap()
1268                    .0
1269                    .output
1270                    .is_null()
1271            );
1272            assert!(
1273                ctx.clone()
1274                    .agent_run(AgentInput {
1275                        name: "agent".to_string(),
1276                        prompt: "prompt".to_string(),
1277                        ..Default::default()
1278                    })
1279                    .await
1280                    .unwrap()
1281                    .0
1282                    .content
1283                    .is_empty()
1284            );
1285            assert!(
1286                ctx.remote_agent_run(
1287                    "https://example.test",
1288                    AgentInput {
1289                        name: "agent".to_string(),
1290                        prompt: "prompt".to_string(),
1291                        ..Default::default()
1292                    },
1293                )
1294                .await
1295                .unwrap()
1296                .content
1297                .is_empty()
1298            );
1299        });
1300    }
1301}