1use 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#[derive(Debug, Clone, Deserialize, Serialize)]
28pub struct AgentArgs {
29 pub prompt: String,
31}
32
33pub trait Agent<C>: Send + Sync
38where
39 C: AgentContext + Send + Sync,
40{
41 fn name(&self) -> String;
52
53 fn description(&self) -> String;
55
56 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 fn supported_resource_tags(&self) -> Vec<String> {
90 Vec::new()
91 }
92
93 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 fn init(&self, _ctx: C) -> impl Future<Output = Result<(), BoxError>> + Send {
103 futures::future::ready(Ok(()))
104 }
105
106 fn tool_dependencies(&self) -> Vec<String> {
110 Vec::new()
111 }
112
113 fn run(
123 &self,
124 ctx: C,
125 prompt: String,
126 resources: Vec<Resource>,
127 ) -> impl Future<Output = Result<AgentOutput, BoxError>> + Send;
128}
129
130pub 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 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 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
186struct 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#[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 pub fn new() -> Self {
261 Self {
262 set: BTreeMap::new(),
263 }
264 }
265
266 pub fn contains(&self, name: &str) -> bool {
268 self.set.contains_key(&name.to_ascii_lowercase())
269 }
270
271 pub fn contains_lowercase(&self, lowercase_name: &str) -> bool {
273 self.set.contains_key(lowercase_name)
274 }
275
276 pub fn names(&self) -> Vec<String> {
278 self.set.keys().cloned().collect()
279 }
280
281 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 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 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 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 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 pub fn get(&self, name: &str) -> Option<Arc<dyn DynAgent<C>>> {
380 self.set.get(&name.to_ascii_lowercase()).cloned()
381 }
382
383 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}