1use std::collections::HashSet;
2use std::sync::Arc;
3
4use agent_base::{AgentResult, AgentRuntime, Tool};
5
6#[cfg(feature = "multi_agent")]
7use crate::multi_agent::{MultiAgentConfig, MultiAgentRuntime};
8
9#[cfg(feature = "skill")]
10use crate::skill::{LazySkillPrompter, Skill, SkillPrompter};
11
12#[cfg(feature = "multi_agent")]
14pub type MultiAgentToolFactory =
15 Arc<dyn Fn(Arc<MultiAgentRuntime>) -> Vec<Arc<dyn Tool>> + Send + Sync>;
16
17#[cfg(feature = "skill")]
19pub type SkillDetailToolFactory =
20 Arc<dyn Fn(Vec<Arc<dyn Skill>>, String) -> Arc<dyn Tool> + Send + Sync>;
21
22#[cfg(feature = "skill")]
24pub type ListSkillsToolFactory =
25 Arc<dyn Fn(Arc<crate::skill::SkillRegistry>) -> Arc<dyn Tool> + Send + Sync>;
26
27pub struct AgentBuilder {
28 inner: agent_base::AgentBuilder,
29 client: Arc<dyn agent_base::llm_trait::LlmProvider>,
30 system_prompt: Option<String>,
31 tool_names: HashSet<String>,
32 business_tools: Vec<Arc<dyn Tool>>,
34 #[cfg(feature = "multi_agent")]
36 multi_agent_config: Option<MultiAgentConfig>,
37 #[cfg(feature = "multi_agent")]
39 multi_agent_tool_factory: Option<MultiAgentToolFactory>,
40 error_recovery: Option<Arc<dyn agent_base::ToolErrorRecovery>>,
42 language: Option<agent_base::Language>,
44 #[cfg(feature = "skill")]
45 skills: Vec<Arc<dyn Skill>>,
46 #[cfg(feature = "skill")]
47 skill_prompter: Option<Arc<dyn SkillPrompter>>,
48 #[cfg(feature = "skill")]
49 skill_detail_tool_name: String,
50 #[cfg(feature = "skill")]
54 skill_detail_tool_factory: Option<SkillDetailToolFactory>,
55 #[cfg(feature = "skill")]
56 list_skills_tool_factory: Option<ListSkillsToolFactory>,
57 #[cfg(feature = "skill")]
58 disable_skill_prompt_injection: bool,
59 context_compactor: Option<Arc<dyn agent_base::ContextCompaction>>,
61}
62
63impl AgentBuilder {
64 pub fn new(client: Arc<dyn agent_base::llm_trait::LlmProvider>) -> Self {
65 Self {
66 inner: agent_base::AgentBuilder::new(client.clone()),
67 client,
68 system_prompt: None,
69 tool_names: HashSet::new(),
70 business_tools: Vec::new(),
71 #[cfg(feature = "multi_agent")]
72 multi_agent_config: None,
73 #[cfg(feature = "multi_agent")]
74 multi_agent_tool_factory: None,
75 error_recovery: None,
76 language: None,
77 #[cfg(feature = "skill")]
78 skills: Vec::new(),
79 #[cfg(feature = "skill")]
80 skill_prompter: None,
81 #[cfg(feature = "skill")]
82 skill_detail_tool_name: "get_skill_detail".to_string(),
83 #[cfg(feature = "skill")]
84 skill_detail_tool_factory: None,
85 #[cfg(feature = "skill")]
86 list_skills_tool_factory: None,
87 #[cfg(feature = "skill")]
88 disable_skill_prompt_injection: false,
89 context_compactor: None,
90 }
91 }
92
93 #[cfg(feature = "multi_agent")]
98 pub fn with_multi_agent(mut self, config: MultiAgentConfig) -> Self {
99 self.multi_agent_config = Some(config);
100 self
101 }
102
103 #[cfg(feature = "multi_agent")]
108 pub fn without_multi_agent(mut self) -> Self {
109 self.multi_agent_config = None;
110 self.multi_agent_tool_factory = None;
111 self
112 }
113
114 #[cfg(feature = "multi_agent")]
120 pub fn with_multi_agent_tool_factory(mut self, factory: MultiAgentToolFactory) -> Self {
121 self.multi_agent_tool_factory = Some(factory);
122 self
123 }
124
125 #[cfg(feature = "skill")]
130 pub fn with_skill_detail_tool_factory(mut self, factory: SkillDetailToolFactory) -> Self {
131 self.skill_detail_tool_factory = Some(factory);
132 self
133 }
134
135 #[cfg(feature = "skill")]
139 pub fn with_list_skills_tool_factory(mut self, factory: ListSkillsToolFactory) -> Self {
140 self.list_skills_tool_factory = Some(factory);
141 self
142 }
143
144 pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
145 let prompt = prompt.into();
146 self.inner = self.inner.system_prompt(prompt.clone());
147 self.system_prompt = Some(prompt);
148 self
149 }
150
151 pub fn enable_thought(self, enable: bool) -> Self {
152 Self {
153 inner: self.inner.enable_thought(enable),
154 ..self
155 }
156 }
157
158 pub fn reasoning(self, config: agent_base::ReasoningConfig) -> Self {
159 Self {
160 inner: self.inner.reasoning(config),
161 ..self
162 }
163 }
164
165 pub fn enable_thinking(self, enable: bool) -> Self {
166 Self {
167 inner: self.inner.enable_thinking(enable),
168 ..self
169 }
170 }
171
172 pub fn thinking_budget(self, budget: u64) -> Self {
173 Self {
174 inner: self.inner.thinking_budget(budget),
175 ..self
176 }
177 }
178
179 pub fn tool_timeout(self, timeout_ms: u64) -> Self {
180 Self {
181 inner: self.inner.tool_timeout(timeout_ms),
182 ..self
183 }
184 }
185
186 pub fn max_tool_output_chars(self, max_chars: usize) -> Self {
187 Self {
188 inner: self.inner.max_tool_output_chars(max_chars),
189 ..self
190 }
191 }
192
193 pub fn max_sessions(self, max: usize) -> Self {
194 Self {
195 inner: self.inner.max_sessions(max),
196 ..self
197 }
198 }
199
200 pub fn max_turns_per_session(self, max: usize) -> Self {
201 Self {
202 inner: self.inner.max_turns_per_session(max),
203 ..self
204 }
205 }
206
207 pub fn execution_max_turns(self, max: u32) -> Self {
208 Self {
209 inner: self.inner.execution_max_turns(max),
210 ..self
211 }
212 }
213
214 pub fn max_message_tokens(self, max: usize) -> Self {
215 Self {
216 inner: self.inner.max_message_tokens(max),
217 ..self
218 }
219 }
220
221 pub fn register_tool(mut self, tool: impl Tool + 'static) -> Self {
222 let tool_arc: Arc<dyn Tool> = Arc::new(tool);
223 self.tool_names.insert(tool_arc.name().to_string());
224 self.business_tools.push(tool_arc.clone());
225 self.inner = self.inner.register_tool_arc(tool_arc);
226 self
227 }
228
229 pub fn register_tool_arc(mut self, tool: Arc<dyn Tool>) -> Self {
230 self.tool_names.insert(tool.name().to_string());
231 self.business_tools.push(tool.clone());
232 self.inner = self.inner.register_tool_arc(tool);
233 self
234 }
235
236 pub fn approval_handler(self, handler: Arc<dyn agent_base::ApprovalHandler>) -> Self {
237 Self {
238 inner: self.inner.approval_handler(handler),
239 ..self
240 }
241 }
242
243 pub fn tool_policy(self, policy: Arc<dyn agent_base::ToolPolicy>) -> Self {
244 Self {
245 inner: self.inner.tool_policy(policy),
246 ..self
247 }
248 }
249
250 pub fn middleware(self, mw: impl agent_base::Middleware + 'static) -> Self {
251 Self {
252 inner: self.inner.middleware(mw),
253 ..self
254 }
255 }
256
257 pub fn guard(self, guard: impl agent_base::ReactLoopGuard + 'static) -> Self {
258 Self {
259 inner: self.inner.guard(guard),
260 ..self
261 }
262 }
263
264 pub fn context_window(self, max_tokens: usize) -> Self {
265 Self {
266 inner: self.inner.context_window(max_tokens),
267 ..self
268 }
269 }
270
271 pub fn context_window_manager(self, manager: agent_base::ContextWindowManager) -> Self {
272 Self {
273 inner: self.inner.context_window_manager(manager),
274 ..self
275 }
276 }
277
278 pub fn response_format(self, format: agent_base::ResponseFormat) -> Self {
279 Self {
280 inner: self.inner.response_format(format),
281 ..self
282 }
283 }
284
285 pub fn llm_retry(self, retry: agent_base::RetryConfig) -> Self {
286 Self {
287 inner: self.inner.llm_retry(retry),
288 ..self
289 }
290 }
291
292 pub fn session_store(self, store: Arc<dyn agent_base::SessionStore>) -> Self {
293 Self {
294 inner: self.inner.session_store(store),
295 ..self
296 }
297 }
298
299 pub fn error_recovery(mut self, recovery: Arc<dyn agent_base::ToolErrorRecovery>) -> Self {
300 self.error_recovery = Some(recovery.clone());
301 self.inner = self.inner.error_recovery(recovery);
302 self
303 }
304
305 pub fn tool_error_retry_prompt(self, prompt: impl Into<String>) -> Self {
306 Self {
307 inner: self.inner.tool_error_retry_prompt(prompt),
308 ..self
309 }
310 }
311
312 pub fn language(mut self, language: agent_base::Language) -> Self {
313 self.language = Some(language.clone());
314 self.inner = self.inner.language(language);
315 self
316 }
317
318 pub fn context_compactor(mut self, compactor: Arc<dyn agent_base::ContextCompaction>) -> Self {
323 self.context_compactor = Some(compactor);
324 self
325 }
326
327 pub fn event_bus_capacity(self, capacity: usize) -> Self {
328 Self {
329 inner: self.inner.event_bus_capacity(capacity),
330 ..self
331 }
332 }
333
334 pub fn session_id_generator(
335 self,
336 generator: Arc<dyn agent_base::types::SessionIdGenerator>,
337 ) -> Self {
338 Self {
339 inner: self.inner.session_id_generator(generator),
340 ..self
341 }
342 }
343
344 pub fn apply_if<T>(self, value: Option<T>, f: impl FnOnce(Self, T) -> Self) -> Self {
352 match value {
353 Some(v) => f(self, v),
354 None => self,
355 }
356 }
357
358 #[cfg(feature = "skill")]
359 pub fn register_skill(mut self, skill: impl Skill + 'static) -> Self {
360 self.skills.push(Arc::new(skill));
361 self
362 }
363
364 #[cfg(feature = "skill")]
365 pub fn register_skills(mut self, skills: Vec<Arc<dyn Skill>>) -> Self {
366 self.skills.extend(skills);
367 self
368 }
369
370 #[cfg(feature = "skill")]
371 pub fn skill_prompter(mut self, prompter: Arc<dyn SkillPrompter>) -> Self {
372 self.skill_prompter = Some(prompter);
373 self
374 }
375
376 #[cfg(feature = "skill")]
377 pub fn disable_skill_prompt_injection(mut self) -> Self {
378 self.disable_skill_prompt_injection = true;
379 self
380 }
381
382 #[cfg(feature = "skill")]
383 pub fn skill_detail_tool_name(mut self, name: impl Into<String>) -> Self {
384 self.skill_detail_tool_name = name.into();
385 self
386 }
387
388 pub fn build(self) -> AgentResult<AgentRuntime> {
391 #[cfg(feature = "skill")]
392 {
393 self.build_with_skills()
394 }
395 #[cfg(not(feature = "skill"))]
396 {
397 self.build_inner()
398 }
399 }
400
401 #[allow(dead_code, unused_mut)]
402 fn build_inner(mut self) -> AgentResult<AgentRuntime> {
403 if self.inner.get_guard().is_none() {
406 self.inner = self
407 .inner
408 .guard(crate::guard::DefaultGuard::with_llm_client(
409 crate::guard::DefaultGuardConfig::default(),
410 self.client.clone(),
411 ));
412 }
413
414 #[cfg(feature = "multi_agent")]
415 let lang = self.language.clone().unwrap_or_default();
416 #[cfg(feature = "multi_agent")]
417 let business_tools = std::mem::take(&mut self.business_tools);
418 #[cfg(feature = "multi_agent")]
419 let error_recovery = self.error_recovery.clone();
420 #[cfg(feature = "multi_agent")]
421 let tool_names = self.tool_names.clone();
422
423 #[cfg(feature = "multi_agent")]
424 let ma_config = self.multi_agent_config.clone();
425 #[cfg(feature = "multi_agent")]
426 let ma_tool_factory = self.multi_agent_tool_factory.take();
427
428 #[cfg(feature = "multi_agent")]
430 if ma_config.as_ref().map(|c| c.enabled).unwrap_or(false) {
431 let ma_prompt = build_multi_agent_system_prompt();
432 let new_prompt = match self.system_prompt.take() {
433 Some(existing) => format!("{}\n\n---\n\n{}", existing, ma_prompt),
434 None => ma_prompt,
435 };
436 self.inner = self.inner.system_prompt(new_prompt);
437 }
438
439 if let Some(compactor) = self.context_compactor.take() {
441 self.inner = self.inner.context_compactor(compactor);
442 }
443
444 let runtime = self.inner.build()?;
445
446 #[cfg(feature = "multi_agent")]
448 if let Some(config) = ma_config
449 && config.enabled
450 {
451 setup_multi_agent(
452 &runtime,
453 config,
454 lang,
455 business_tools,
456 error_recovery,
457 &tool_names,
458 ma_tool_factory,
459 )?;
460 }
461
462 Ok(runtime)
463 }
464
465 #[cfg(feature = "skill")]
477 fn build_with_skills(mut self) -> AgentResult<AgentRuntime> {
478 if self.inner.get_guard().is_none() {
481 self.inner = self
482 .inner
483 .guard(crate::guard::DefaultGuard::with_llm_client(
484 crate::guard::DefaultGuardConfig::default(),
485 self.client.clone(),
486 ));
487 }
488
489 let mut ab = self.inner;
490 #[cfg(feature = "multi_agent")]
491 let lang = self.language.clone().unwrap_or_default();
492 #[cfg(feature = "multi_agent")]
493 let business_tools = std::mem::take(&mut self.business_tools);
494 #[cfg(feature = "multi_agent")]
495 let error_recovery = self.error_recovery.clone();
496 #[cfg(feature = "multi_agent")]
497 let tool_names = self.tool_names.clone();
498
499 #[cfg(feature = "multi_agent")]
500 let ma_config = self.multi_agent_config.clone();
501 #[cfg(feature = "multi_agent")]
502 let ma_tool_factory = self.multi_agent_tool_factory.take();
503
504 if !self.skills.is_empty() {
506 let prompter: Arc<dyn SkillPrompter> = self
507 .skill_prompter
508 .take()
509 .unwrap_or_else(|| Arc::new(LazySkillPrompter::new()));
510
511 let mut skill_refs: Vec<Arc<dyn Skill>> = Vec::new();
512
513 for skill in self.skills {
514 for tool in skill.tools() {
515 let tool_name = tool.name().to_string();
516 if self.tool_names.contains(&tool_name) {
517 return Err(agent_base::AgentError::internal(format!(
518 "Tool name conflict: `{}` (Skill `{}`)",
519 tool_name,
520 skill.name()
521 )));
522 }
523 self.tool_names.insert(tool_name);
524 ab = ab.register_tool_arc(tool);
525 }
526 skill_refs.push(skill);
527 }
528
529 if !self.disable_skill_prompt_injection {
530 let skill_prompt = prompter.build_prompt(&skill_refs, &self.skill_detail_tool_name);
531 if !skill_prompt.is_empty() {
532 let new_prompt = match self.system_prompt.take() {
533 Some(existing) => format!("{}\n\n---\n\n{}", existing, skill_prompt),
534 None => skill_prompt,
535 };
536 self.system_prompt = Some(new_prompt.clone());
537 ab = ab.system_prompt(new_prompt);
538 }
539 }
540
541 if let Some(factory) = self.skill_detail_tool_factory.take() {
544 let detail_tool = factory(skill_refs.clone(), self.skill_detail_tool_name);
545 ab = ab.register_tool_arc(detail_tool);
546 }
547
548 if let Some(factory) = self.list_skills_tool_factory.take() {
550 let registry = Arc::new(crate::skill::SkillRegistry::new());
551 for skill in &skill_refs {
552 tokio::task::block_in_place(|| {
553 tokio::runtime::Handle::current().block_on(async {
554 registry.register(skill.clone()).await;
555 })
556 });
557 }
558 let list_tool = factory(registry);
559 ab = ab.register_tool_arc(list_tool);
560 }
561 }
562
563 #[cfg(feature = "multi_agent")]
565 if ma_config.as_ref().map(|c| c.enabled).unwrap_or(false) {
566 let ma_prompt = build_multi_agent_system_prompt();
567 let new_prompt = match self.system_prompt.take() {
568 Some(existing) => format!("{}\n\n---\n\n{}", existing, ma_prompt),
569 None => ma_prompt,
570 };
571 ab = ab.system_prompt(new_prompt);
572 }
573
574 if let Some(compactor) = self.context_compactor.take() {
576 ab = ab.context_compactor(compactor);
577 }
578
579 let runtime = ab.build()?;
580
581 #[cfg(feature = "multi_agent")]
583 if let Some(config) = ma_config
584 && config.enabled
585 {
586 setup_multi_agent(
587 &runtime,
588 config,
589 lang,
590 business_tools,
591 error_recovery,
592 &tool_names,
593 ma_tool_factory,
594 )?;
595 }
596
597 Ok(runtime)
598 }
599}
600
601#[cfg(feature = "multi_agent")]
612pub fn setup_multi_agent(
613 runtime: &AgentRuntime,
614 config: MultiAgentConfig,
615 lang: agent_base::Language,
616 business_tools: Vec<Arc<dyn Tool>>,
617 error_recovery: Option<Arc<dyn agent_base::ToolErrorRecovery>>,
618 existing_tool_names: &HashSet<String>,
619 tool_factory: Option<MultiAgentToolFactory>,
620) -> AgentResult<Arc<MultiAgentRuntime>> {
621 let client = runtime.provider();
622 let cancel_token = runtime.cancel_token();
623 let tool_policy = runtime.tool_policy().cloned();
624 let approval_handler = runtime.approval_handler().cloned();
625
626 let ma_runtime = Arc::new(MultiAgentRuntime::new(
627 config.clone(),
628 client,
629 business_tools,
630 cancel_token,
631 error_recovery,
632 lang,
633 tool_policy,
634 approval_handler,
635 ));
636
637 ma_runtime.set_session_manager(Arc::new(runtime.session_manager().clone()));
639
640 let (event_tx, mut event_rx) =
642 tokio::sync::mpsc::unbounded_channel::<agent_base::RuntimeEvent>();
643 ma_runtime.set_event_sender(event_tx);
644 let parent_runtime = runtime.clone();
645 tokio::spawn(async move {
646 while let Some(event) = event_rx.recv().await {
647 parent_runtime.emit_event(event);
648 }
649 });
650
651 if let Some(factory) = tool_factory {
653 let tools = factory(ma_runtime.clone());
654 let registry = runtime.tools_mut();
655 let mut reg = tokio::task::block_in_place(|| registry.blocking_write());
656 for tool in tools {
657 let tool_name = tool.name().to_string();
658 if !existing_tool_names.contains(&tool_name) {
659 reg.register_arc(tool);
660 }
661 }
662 drop(reg);
663 }
664
665 Ok(ma_runtime)
666}
667
668#[cfg(feature = "multi_agent")]
670pub fn build_multi_agent_system_prompt() -> String {
671 r#"## Multi-Agent Capabilities
672
673You have the ability to spawn sub-agents to execute tasks concurrently. Use these tools to delegate work:
674
675- `spawn_agent`: Create a new sub-agent with a specific role. The agent runs independently.
676- `send_message`: Send a message to a sub-agent without triggering execution.
677- `followup_task`: Assign a task to a sub-agent and trigger its execution. Returns immediately.
678- `wait_agent`: Wait for a sub-agent's result. Blocks until the agent completes or timeout.
679- `list_agents`: List all active sub-agents and their status.
680- `close_agent`: Close a sub-agent and release its resources.
681
682### When to Spawn
683
684- Tasks that can run independently and in parallel (e.g., "research X and Y simultaneously")
685- Long-running tasks where you want to check intermediate results
686- Decomposing complex tasks into sub-tasks for focused execution
687
688### When NOT to Spawn
689
690- Simple lookups or single-tool calls (just use the tool directly)
691- Sequential dependencies where the next step requires the previous result
692- Tasks that need your full context or reasoning
693
694### Communication Pattern
695
6961. `spawn_agent` → create the sub-agent
6972. `followup_task` → assign work (can call multiple times)
6983. `wait_agent` → collect results
6994. `close_agent` → clean up when done"#
700 .to_string()
701}
702
703pub fn build_memory_system_prompt() -> String {
712 r#"## Memory
713
714You have a persistent file-based memory at `.phi/memory/`. Use `read_file` and `write_file` to manage it — there are no dedicated memory tools.
715
716### How Memory Works
717
718- `MEMORY.md` is the index — it lists all memories with one-line descriptions. Read it first when you need to recall something.
719- Each memory is a separate `.md` file with YAML frontmatter:
720 ```yaml
721 ---
722 name: <short-kebab-case-slug>
723 description: <one-line summary — used to decide relevance during recall>
724 metadata:
725 node_type: memory
726 type: user | feedback | project | reference
727 ---
728
729 <the fact or instruction>
730 ```
731- The `description` field is the key for recall — write it so you can tell at a glance whether this memory is relevant to the current task.
732- Link related memories with `[[memory-name]]` in the body.
733- `user` type = who the user is (role, expertise, preferences).
734- `feedback` type = guidance the user has given on how you should work.
735- `project` type = ongoing work, goals, or constraints.
736- `reference` type = pointers to external resources (URLs, dashboards, tickets).
737
738### When to Use Memory
739
740- The user explicitly asks you to remember something ("remember this", "save that")
741- You learn something important about the user's preferences or workflow
742- After completing a significant task, save context that would help in future sessions
743- The user gives you feedback on how to work — save it as `feedback` type
744
745### When NOT to Use Memory
746
747- For transient information that won't be useful beyond this session
748- For facts already recorded in the codebase (code structure, git history, config files)
749- For items that only matter to the current conversation
750
751### Pro Tips
752
753- When creating your first memory of a new type, you can read template files for format reference (check `.phi/templates/memory/` if available).
754- Keep the MEMORY.md index concise — it's loaded into context every session.
755- Before writing a new memory, check if an existing file already covers it — update instead of duplicating.
756
757### Workflow
758
759**To recall:** read `MEMORY.md` → find relevant entries by description → read the specific `.md` files you need.
760**To remember:** create a new `.md` file with proper frontmatter → update `MEMORY.md` with a new entry.
761**To update:** edit the existing `.md` file (don't create a duplicate).
762**To forget:** delete the `.md` file → remove its entry from `MEMORY.md`."#
763 .to_string()
764}
765
766#[cfg(test)]
767mod tests {
768 use super::*;
769 use agent_base::Content;
770 use agent_base::llm_trait::response::FinishReason;
771 use agent_base::llm_trait::types::UsageInfo;
772 use agent_base::llm_trait::{
773 Capabilities, ChatRequest, ChatResponse, ChatStream, LlmError, LlmProvider, ProviderInfo,
774 };
775
776 struct StubProvider;
779
780 #[async_trait::async_trait]
781 impl LlmProvider for StubProvider {
782 async fn stream(&self, _request: ChatRequest) -> Result<ChatStream, LlmError> {
783 let chunks = vec![
784 Ok(agent_base::StreamChunk::Text("ok".to_string())),
785 Ok(agent_base::StreamChunk::Stop {
786 finish_reason: Some("stop".to_string()),
787 }),
788 ];
789 Ok(ChatStream::new(Box::pin(futures_util::stream::iter(
790 chunks,
791 ))))
792 }
793
794 async fn chat(&self, _request: ChatRequest) -> Result<ChatResponse, LlmError> {
795 Ok(ChatResponse {
796 content: "ok".to_string(),
797 tool_calls: vec![],
798 usage: UsageInfo::default(),
799 finish_reason: FinishReason::Stop,
800 raw: None,
801 reasoning_content: None,
802 thinking_signature: None,
803 })
804 }
805
806 fn capabilities(&self) -> Capabilities {
807 Capabilities {
808 supports_streaming: true,
809 supports_tools: true,
810 ..Default::default()
811 }
812 }
813
814 fn info(&self) -> ProviderInfo {
815 ProviderInfo {
816 name: "stub".to_string(),
817 model: "stub-model".to_string(),
818 version: None,
819 }
820 }
821 }
822
823 fn make_client() -> Arc<dyn LlmProvider> {
824 Arc::new(StubProvider)
825 }
826
827 #[cfg(feature = "multi_agent")]
830 #[tokio::test(flavor = "multi_thread")]
831 async fn test_setup_multi_agent_without_factory_registers_no_tools() {
832 let client = make_client();
833 let runtime = agent_base::AgentBuilder::new(client.clone())
834 .build()
835 .unwrap();
836 let config = MultiAgentConfig::enabled();
837
838 let result = setup_multi_agent(
839 &runtime,
840 config,
841 agent_base::Language::En,
842 vec![],
843 None,
844 &HashSet::new(),
845 None, );
847 assert!(result.is_ok());
848 let ma_runtime = result.unwrap();
849 let agents = ma_runtime.list_agents();
851 assert!(agents.is_empty());
852 }
853
854 #[cfg(feature = "multi_agent")]
855 #[tokio::test(flavor = "multi_thread")]
856 async fn test_setup_multi_agent_with_factory_registers_tools() {
857 let client = make_client();
858 let runtime = agent_base::AgentBuilder::new(client.clone())
859 .build()
860 .unwrap();
861 let config = MultiAgentConfig::enabled();
862
863 let factory: MultiAgentToolFactory = Arc::new(|_rt| {
864 struct FakeTool;
866 #[async_trait::async_trait]
867 impl Tool for FakeTool {
868 fn name(&self) -> &'static str {
869 "fake_tool"
870 }
871 fn description(&self) -> &'static str {
872 ""
873 }
874 fn schema(&self) -> serde_json::Value {
875 serde_json::json!({})
876 }
877 async fn call(
878 &self,
879 _args: &serde_json::Value,
880 _ctx: &agent_base::ToolContext,
881 ) -> AgentResult<Vec<Content>> {
882 Ok(vec![Content::text("ok")])
883 }
884 }
885 vec![Arc::new(FakeTool)]
886 });
887
888 let result = setup_multi_agent(
889 &runtime,
890 config,
891 agent_base::Language::En,
892 vec![],
893 None,
894 &HashSet::new(),
895 Some(factory),
896 );
897 assert!(result.is_ok());
898
899 let tools: Vec<String> = tokio::task::block_in_place(|| {
901 let tools = runtime.tools_mut();
902 let guard = tools.blocking_read();
903 guard.metadatas().into_iter().map(|m| m.name).collect()
904 });
905 assert!(tools.contains(&"fake_tool".to_string()));
906 }
907
908 #[cfg(feature = "multi_agent")]
909 #[tokio::test(flavor = "multi_thread")]
910 async fn test_setup_multi_agent_skips_duplicate_tool_names() {
911 let client = make_client();
912 let runtime = agent_base::AgentBuilder::new(client.clone())
913 .build()
914 .unwrap();
915
916 struct DupTool;
918 #[async_trait::async_trait]
919 impl Tool for DupTool {
920 fn name(&self) -> &'static str {
921 "dup_tool"
922 }
923 fn description(&self) -> &'static str {
924 ""
925 }
926 fn schema(&self) -> serde_json::Value {
927 serde_json::json!({})
928 }
929 async fn call(
930 &self,
931 _args: &serde_json::Value,
932 _ctx: &agent_base::ToolContext,
933 ) -> AgentResult<Vec<Content>> {
934 Ok(vec![Content::text("ok")])
935 }
936 }
937 {
938 let tools = runtime.tools_mut();
939 let mut reg = tokio::task::block_in_place(|| tools.blocking_write());
940 reg.register(DupTool);
941 }
942
943 let factory: MultiAgentToolFactory = Arc::new(|_rt| {
944 struct FakeTool;
945 #[async_trait::async_trait]
946 impl Tool for FakeTool {
947 fn name(&self) -> &'static str {
948 "dup_tool"
949 }
950 fn description(&self) -> &'static str {
951 ""
952 }
953 fn schema(&self) -> serde_json::Value {
954 serde_json::json!({})
955 }
956 async fn call(
957 &self,
958 _args: &serde_json::Value,
959 _ctx: &agent_base::ToolContext,
960 ) -> AgentResult<Vec<Content>> {
961 Ok(vec![Content::text("ok")])
962 }
963 }
964 vec![Arc::new(FakeTool)]
965 });
966
967 let mut existing = HashSet::new();
968 existing.insert("dup_tool".to_string());
969
970 let result = setup_multi_agent(
971 &runtime,
972 MultiAgentConfig::enabled(),
973 agent_base::Language::En,
974 vec![],
975 None,
976 &existing,
977 Some(factory),
978 );
979 assert!(result.is_ok());
980 let tools = tokio::task::block_in_place(|| {
982 let tools = runtime.tools_mut();
983 let guard = tools.blocking_read();
984 guard
985 .metadatas()
986 .into_iter()
987 .map(|m| m.name)
988 .collect::<Vec<String>>()
989 });
990 let count = tools.iter().filter(|n| n.as_str() == "dup_tool").count();
991 assert_eq!(count, 1);
992 }
993
994 #[cfg(feature = "multi_agent")]
997 #[tokio::test(flavor = "multi_thread")]
998 async fn test_builder_with_multi_agent_without_factory_builds_ok() {
999 let client = make_client();
1000 let runtime = AgentBuilder::new(client)
1001 .with_multi_agent(MultiAgentConfig::enabled())
1002 .build()
1003 .unwrap();
1004 let tools = tokio::task::block_in_place(|| {
1006 let tools = runtime.tools_mut();
1007 let guard = tools.blocking_read();
1008 guard
1009 .metadatas()
1010 .into_iter()
1011 .map(|m| m.name)
1012 .collect::<Vec<String>>()
1013 });
1014 assert!(!tools.contains(&"spawn_agent".to_string()));
1016 }
1017
1018 #[cfg(feature = "multi_agent")]
1019 #[tokio::test(flavor = "multi_thread")]
1020 async fn test_builder_with_factory_registers_tools() {
1021 let client = make_client();
1022 let factory: MultiAgentToolFactory = Arc::new(|_rt| {
1024 struct TestTool;
1025 #[async_trait::async_trait]
1026 impl Tool for TestTool {
1027 fn name(&self) -> &'static str {
1028 "factory_test_tool"
1029 }
1030 fn description(&self) -> &'static str {
1031 ""
1032 }
1033 fn schema(&self) -> serde_json::Value {
1034 serde_json::json!({})
1035 }
1036 async fn call(
1037 &self,
1038 _args: &serde_json::Value,
1039 _ctx: &agent_base::ToolContext,
1040 ) -> AgentResult<Vec<Content>> {
1041 Ok(vec![Content::text("ok")])
1042 }
1043 }
1044 vec![Arc::new(TestTool)]
1045 });
1046
1047 let runtime = AgentBuilder::new(client)
1048 .with_multi_agent(MultiAgentConfig::enabled())
1049 .with_multi_agent_tool_factory(factory)
1050 .build()
1051 .unwrap();
1052
1053 let tools = tokio::task::block_in_place(|| {
1054 let tools = runtime.tools_mut();
1055 let guard = tools.blocking_read();
1056 guard
1057 .metadatas()
1058 .into_iter()
1059 .map(|m| m.name)
1060 .collect::<Vec<String>>()
1061 });
1062 assert!(tools.contains(&"factory_test_tool".to_string()));
1063 }
1064
1065 #[cfg(feature = "multi_agent")]
1066 #[test]
1067 fn test_builder_disabled_multi_agent_skips_factory() {
1068 let client = make_client();
1069 let factory: MultiAgentToolFactory = Arc::new(|_rt| {
1070 panic!("factory should not be called when multi-agent is not configured");
1071 });
1072
1073 let runtime = AgentBuilder::new(client)
1074 .with_multi_agent_tool_factory(factory)
1075 .build()
1077 .unwrap();
1078
1079 let tools = tokio::task::block_in_place(|| {
1080 let tools = runtime.tools_mut();
1081 let guard = tools.blocking_read();
1082 guard
1083 .metadatas()
1084 .into_iter()
1085 .map(|m| m.name)
1086 .collect::<Vec<String>>()
1087 });
1088 assert!(!tools.contains(&"spawn_agent".to_string()));
1089 }
1090
1091 #[cfg(feature = "multi_agent")]
1094 #[test]
1095 fn test_system_prompt_contains_tool_names() {
1096 let prompt = build_multi_agent_system_prompt();
1097 assert!(prompt.contains("spawn_agent"));
1098 assert!(prompt.contains("send_message"));
1099 assert!(prompt.contains("followup_task"));
1100 assert!(prompt.contains("wait_agent"));
1101 assert!(prompt.contains("list_agents"));
1102 assert!(prompt.contains("close_agent"));
1103 }
1104
1105 #[cfg(feature = "multi_agent")]
1106 #[test]
1107 fn test_system_prompt_contains_guidance() {
1108 let prompt = build_multi_agent_system_prompt();
1109 assert!(prompt.contains("When to Spawn"));
1110 assert!(prompt.contains("When NOT to Spawn"));
1111 assert!(prompt.contains("Communication Pattern"));
1112 }
1113
1114 #[cfg(feature = "multi_agent")]
1117 #[tokio::test(flavor = "multi_thread")]
1118 async fn test_without_multi_agent_clears_config_and_factory() {
1119 let client = make_client();
1120
1121 let factory: MultiAgentToolFactory = Arc::new(|_rt| {
1123 panic!("factory should not be called when multi-agent is cleared");
1124 });
1125
1126 let runtime = AgentBuilder::new(client)
1127 .with_multi_agent(MultiAgentConfig::enabled())
1128 .with_multi_agent_tool_factory(factory)
1129 .without_multi_agent() .build()
1131 .unwrap();
1132
1133 let tools = tokio::task::block_in_place(|| {
1134 let tools = runtime.tools_mut();
1135 let guard = tools.blocking_read();
1136 guard
1137 .metadatas()
1138 .into_iter()
1139 .map(|m| m.name)
1140 .collect::<Vec<String>>()
1141 });
1142 assert!(!tools.contains(&"spawn_agent".to_string()));
1143 }
1144
1145 #[test]
1148 fn test_apply_if_some_applies_transformation() {
1149 let client = make_client();
1150 let builder = AgentBuilder::new(client)
1151 .apply_if(Some("custom prompt"), |b, prompt| b.system_prompt(prompt));
1152 assert!(builder.system_prompt.unwrap().contains("custom prompt"));
1154 }
1155
1156 #[test]
1157 fn test_apply_if_none_passes_through() {
1158 let client = make_client();
1159 let builder = AgentBuilder::new(client).apply_if(None as Option<&str>, |_b, _prompt| {
1160 panic!("should not be called when value is None");
1161 });
1162 assert!(builder.system_prompt.is_none());
1163 }
1164
1165 #[test]
1168 fn test_build_memory_system_prompt_non_empty() {
1169 let prompt = build_memory_system_prompt();
1170 assert!(!prompt.is_empty());
1171 assert!(prompt.contains("Memory"));
1172 assert!(prompt.contains("MEMORY.md"));
1173 assert!(prompt.contains("read_file"));
1174 assert!(prompt.contains("write_file"));
1175 }
1176
1177 struct NamedTool(&'static str);
1180
1181 #[async_trait::async_trait]
1182 impl Tool for NamedTool {
1183 fn name(&self) -> &'static str {
1184 self.0
1185 }
1186
1187 fn description(&self) -> &'static str {
1188 ""
1189 }
1190
1191 fn schema(&self) -> serde_json::Value {
1192 serde_json::json!({})
1193 }
1194
1195 async fn call(
1196 &self,
1197 _args: &serde_json::Value,
1198 _ctx: &agent_base::ToolContext,
1199 ) -> AgentResult<Vec<Content>> {
1200 Ok(vec![Content::text("ok")])
1201 }
1202 }
1203
1204 fn runtime_tool_names(runtime: &AgentRuntime) -> Vec<String> {
1205 tokio::task::block_in_place(|| {
1206 let tools = runtime.tools_mut();
1207 let guard = tools.blocking_read();
1208 guard.metadatas().into_iter().map(|m| m.name).collect()
1209 })
1210 }
1211
1212 #[tokio::test(flavor = "multi_thread")]
1213 async fn test_builder_scalar_passthrough_methods() {
1214 let client = make_client();
1215 let runtime = AgentBuilder::new(client)
1216 .enable_thought(true)
1217 .reasoning(agent_base::ReasoningConfig::default())
1218 .enable_thinking(false)
1219 .thinking_budget(1000)
1220 .tool_timeout(5000)
1221 .max_tool_output_chars(4000)
1222 .max_sessions(16)
1223 .max_turns_per_session(20)
1224 .execution_max_turns(10)
1225 .max_message_tokens(8000)
1226 .context_window(64_000)
1227 .context_window_manager(agent_base::ContextWindowManager::new(64_000))
1228 .response_format(agent_base::ResponseFormat::JsonObject)
1229 .llm_retry(agent_base::RetryConfig::default())
1230 .tool_error_retry_prompt("please retry")
1231 .language(agent_base::Language::En)
1232 .event_bus_capacity(256)
1233 .build()
1234 .unwrap();
1235
1236 assert!(runtime.provider().capabilities().supports_streaming);
1237 }
1238
1239 #[tokio::test(flavor = "multi_thread")]
1240 async fn test_register_tool_variants() {
1241 let client = make_client();
1242 let runtime = AgentBuilder::new(client)
1243 .register_tool(NamedTool("tool_by_value"))
1244 .register_tool_arc(Arc::new(NamedTool("tool_by_arc")))
1245 .build()
1246 .unwrap();
1247
1248 let names = runtime_tool_names(&runtime);
1249 assert!(names.contains(&"tool_by_value".to_string()));
1250 assert!(names.contains(&"tool_by_arc".to_string()));
1251 }
1252
1253 #[cfg(feature = "skill")]
1254 mod skill_tests {
1255 use super::*;
1256 use crate::skill::Skill;
1257
1258 struct TestSkill;
1259
1260 impl Skill for TestSkill {
1261 fn name(&self) -> &'static str {
1262 "test_skill"
1263 }
1264
1265 fn brief_description(&self) -> String {
1266 "a test skill".to_string()
1267 }
1268
1269 fn detailed_description(&self) -> String {
1270 "detailed test skill".to_string()
1271 }
1272
1273 fn tools(&self) -> Vec<Arc<dyn Tool>> {
1274 vec![]
1275 }
1276 }
1277
1278 struct ToolSkill;
1279
1280 impl Skill for ToolSkill {
1281 fn name(&self) -> &'static str {
1282 "tool_skill"
1283 }
1284
1285 fn brief_description(&self) -> String {
1286 "skill with a tool".to_string()
1287 }
1288
1289 fn detailed_description(&self) -> String {
1290 "skill that provides a tool".to_string()
1291 }
1292
1293 fn tools(&self) -> Vec<Arc<dyn Tool>> {
1294 vec![Arc::new(NamedTool("skill_provided_tool"))]
1295 }
1296 }
1297
1298 #[tokio::test(flavor = "multi_thread")]
1299 async fn test_register_skill_builds_ok() {
1300 let client = make_client();
1301 let runtime = AgentBuilder::new(client)
1302 .register_skill(TestSkill)
1303 .build()
1304 .unwrap();
1305 assert!(runtime_tool_names(&runtime).is_empty());
1307 }
1308
1309 #[tokio::test(flavor = "multi_thread")]
1310 async fn test_register_skill_with_tool_registers_tool() {
1311 let client = make_client();
1312 let runtime = AgentBuilder::new(client)
1313 .register_skill(ToolSkill)
1314 .build()
1315 .unwrap();
1316 assert!(runtime_tool_names(&runtime).contains(&"skill_provided_tool".to_string()));
1317 }
1318
1319 #[tokio::test(flavor = "multi_thread")]
1320 async fn test_register_skill_tool_name_conflict() {
1321 let client = make_client();
1322 let result = AgentBuilder::new(client)
1323 .register_tool(NamedTool("skill_provided_tool"))
1324 .register_skill(ToolSkill)
1325 .build();
1326 let err = result.err().unwrap();
1327 assert!(format!("{err}").contains("Tool name conflict"));
1328 }
1329
1330 #[tokio::test(flavor = "multi_thread")]
1331 async fn test_skill_detail_tool_factory_registers_tool() {
1332 let client = make_client();
1333 let factory: SkillDetailToolFactory = Arc::new(|_skills, name| {
1334 assert_eq!(name, "get_skill_detail");
1335 Arc::new(NamedTool("detail_tool"))
1336 });
1337 let runtime = AgentBuilder::new(client)
1338 .register_skill(TestSkill)
1339 .with_skill_detail_tool_factory(factory)
1340 .build()
1341 .unwrap();
1342 assert!(runtime_tool_names(&runtime).contains(&"detail_tool".to_string()));
1343 }
1344
1345 #[tokio::test(flavor = "multi_thread")]
1346 async fn test_list_skills_tool_factory_registers_tool() {
1347 let client = make_client();
1348 let factory: ListSkillsToolFactory =
1349 Arc::new(|_registry| Arc::new(NamedTool("list_skills_tool")));
1350 let runtime = AgentBuilder::new(client)
1351 .register_skill(TestSkill)
1352 .with_list_skills_tool_factory(factory)
1353 .build()
1354 .unwrap();
1355 assert!(runtime_tool_names(&runtime).contains(&"list_skills_tool".to_string()));
1356 }
1357
1358 #[tokio::test(flavor = "multi_thread")]
1359 async fn test_disable_skill_prompt_injection_builds() {
1360 let client = make_client();
1361 let runtime = AgentBuilder::new(client)
1362 .register_skill(ToolSkill)
1363 .disable_skill_prompt_injection()
1364 .build()
1365 .unwrap();
1366 assert!(runtime_tool_names(&runtime).contains(&"skill_provided_tool".to_string()));
1368 }
1369 }
1370}