1#![allow(
8 deprecated,
9 reason = "deprecated AgentRegistry shims retained for one-release migration; internal use until loader cutover"
10)]
11
12use crate::AgentConfig;
13use crate::{Agent, AgentResponse, ExecutionMetadata};
14use ares_llm::coordinator::ConversationMessage;
15use ares_llm::observability::{LlmCallRecord, ObservabilitySink, ToolCallRecord};
16#[cfg(feature = "postgres")]
17use ares_llm::compact::{CompactConfig, CompactionState, TurnEntry};
18use ares_llm::compact::Compactor;
19use ares_llm::{LLMClient, LLMResponse};
20use ares_tools::Tools;
21use ares_types::types::{AgentContext, AgentType, AppError, Result, ToolDefinition};
22use async_trait::async_trait;
23use cordis::{Context, CordisError, EventsService};
24use std::future::Future;
25use std::sync::Arc;
26
27use cordis::Service;
33
34pub struct PostgresService;
40impl Service for PostgresService {
41 fn check(&self) -> bool {
42 cfg!(feature = "postgres")
43 }
44}
45
46struct ProviderLlm {
47 provider_name: String,
48 llm: Box<dyn LLMClient>,
49}
50
51struct LlmAttemptResponse {
52 response: LLMResponse,
53 provider_name: String,
54 model_name: String,
55}
56
57const SESSION_COMPACTOR_CAP: usize = 256;
64
65#[derive(Default)]
73pub struct SessionCompactors {
74 inner: parking_lot::Mutex<std::collections::HashMap<(String, String), Arc<Compactor>>>,
75}
76
77fn global_session_compactors() -> &'static SessionCompactors {
84 static COMPACTORS: std::sync::LazyLock<SessionCompactors> =
85 std::sync::LazyLock::new(SessionCompactors::new);
86 &COMPACTORS
87}
88
89fn log_compact_event(stage: &'static str, event: &ares_llm::compact::CompactEvent) {
91 use ares_llm::compact::CompactEvent;
92 match event {
93 CompactEvent::Scored { seq, score } => {
94 tracing::debug!(stage, seq, score, "compaction scored turn");
95 }
96 CompactEvent::Audited {
97 critical_kept,
98 memory_chars,
99 dropped_seqs,
100 } => {
101 tracing::debug!(
102 stage,
103 critical_kept,
104 memory_chars,
105 dropped = dropped_seqs.len(),
106 "compaction audit finished"
107 );
108 }
109 CompactEvent::Skipped { reason } => {
110 tracing::debug!(stage, reason, "compaction step skipped");
111 }
112 }
113}
114
115impl SessionCompactors {
116 pub fn new() -> Self {
118 Self::default()
119 }
120
121 #[cfg(feature = "postgres")]
125 async fn get_or_create(
126 &self,
127 tenant_id: &str,
128 session_id: &str,
129 llm: &ares_llm::Llm,
130 ctx: &Arc<Context>,
131 ) -> Option<Arc<Compactor>> {
132 if let Some(existing) = self.lock().get(&(tenant_id.to_string(), session_id.to_string())) {
133 return Some(Arc::clone(existing));
134 }
135
136 let client = llm.get_client(ctx, ares_llm::CapabilityRequirements::default()).await.ok()?;
137 let mut state = CompactionState::default();
138 if let Some(db) = ctx.get::<ares_store::TenantDb>() {
139 let store = ares_store::postgres::PostgresClient {
140 pool: db.pool().clone(),
141 };
142 match store.get_conversation_snapshot(session_id).await {
143 Ok(Some(row)) => {
144 let entries: Vec<TurnEntry> =
145 serde_json::from_value(row.entries).unwrap_or_default();
146 let critical: Vec<String> =
147 serde_json::from_value(row.critical).unwrap_or_default();
148 state = CompactionState::from_parts(
149 entries,
150 critical,
151 row.memory,
152 row.last_audit_seq.max(0) as u64,
153 );
154 }
155 Ok(None) => {}
156 Err(e) => {
157 tracing::debug!(session_id, error = %e, "snapshot hydrate failed; starting cold");
158 }
159 }
160 }
161
162 let compactor = Arc::new(Compactor::hydrate(
163 CompactConfig::default(),
164 client,
165 state,
166 ));
167 let mut map = self.lock();
168 if map.len() >= SESSION_COMPACTOR_CAP {
169 map.clear();
171 }
172 map.insert(
173 (tenant_id.to_string(), session_id.to_string()),
174 Arc::clone(&compactor),
175 );
176 Some(compactor)
177 }
178
179 #[cfg(not(feature = "postgres"))]
181 async fn get_or_create(
182 &self,
183 tenant_id: &str,
184 session_id: &str,
185 llm: &ares_llm::Llm,
186 ctx: &Arc<Context>,
187 ) -> Option<Arc<Compactor>> {
188 if let Some(existing) = self.lock().get(&(tenant_id.to_string(), session_id.to_string())) {
189 return Some(Arc::clone(existing));
190 }
191 let client = llm.get_client(ctx, ares_llm::CapabilityRequirements::default()).await.ok()?;
192 let compactor = Arc::new(Compactor::with_client(client));
193 let mut map = self.lock();
194 if map.len() >= SESSION_COMPACTOR_CAP {
195 map.clear();
196 }
197 map.insert(
198 (tenant_id.to_string(), session_id.to_string()),
199 Arc::clone(&compactor),
200 );
201 Some(compactor)
202 }
203
204 fn lock(
205 &self,
206 ) -> parking_lot::MutexGuard<
207 '_,
208 std::collections::HashMap<(String, String), Arc<Compactor>>,
209 > {
210 self.inner.lock()
211 }
212}
213
214fn tenant_key_for_compaction(ctx: &Context, fallback_user: &str) -> String {
220 for tid in [
221 std::any::TypeId::of::<crate::Execute>(),
222 std::any::TypeId::of::<ares_tools::Tools>(),
223 ] {
224 if let Some(label) = ctx.isolate_label(tid) {
225 let trimmed = label
226 .strip_prefix("tenant:")
227 .or_else(|| label.strip_prefix("user:"))
228 .unwrap_or(&label);
229 if !trimmed.is_empty() {
230 return trimmed.to_string();
231 }
232 }
233 }
234 if let Some(tc) = ctx.get::<ares_types::models::TenantContext>() {
235 if !tc.tenant_id.is_empty() {
236 return tc.tenant_id.clone();
237 }
238 }
239 fallback_user.to_string()
240}
241
242#[cfg(feature = "postgres")]
247async fn persist_session_snapshot(compactor: &Compactor, pool: sqlx::PgPool, session_key: &str) {
248 let state = compactor.export();
249 let entries = match serde_json::to_value(state.entries()) {
250 Ok(v) => v,
251 Err(e) => {
252 tracing::debug!(session_key, error = %e, "snapshot entries serialize failed");
253 return;
254 }
255 };
256 let critical = match serde_json::to_value(state.critical()) {
257 Ok(v) => v,
258 Err(e) => {
259 tracing::debug!(session_key, error = %e, "snapshot critical serialize failed");
260 return;
261 }
262 };
263 let store = ares_store::postgres::PostgresClient { pool };
264 if let Err(e) = store
265 .upsert_conversation_snapshot(
266 session_key,
267 entries,
268 critical,
269 state.memory(),
270 state.last_audit_seq() as i64,
271 )
272 .await
273 {
274 tracing::debug!(session_key, error = %e, "conversation snapshot persist failed");
275 }
276}
277
278pub struct ConfigurableAgent {
280 name: String,
282 agent_type: AgentType,
284 llm: Box<dyn LLMClient>,
286 provider_name: String,
288 system_prompt: String,
290 tools: Option<Arc<Tools>>,
292 cordis_ctx: Option<Arc<Context>>,
294 allowed_tools: Option<Vec<String>>,
297 max_tool_iterations: usize,
299 parallel_tools: bool,
301 observability: Option<Arc<dyn ObservabilitySink>>,
303 fallback_llms: Vec<ProviderLlm>,
305 run_id: Option<String>,
307 compaction: bool,
309}
310
311fn is_prebuilt_connector_tool(name: &str) -> bool {
312 matches!(
313 name,
314 "google_calendar_list_events"
315 | "google_calendar_create_event"
316 | "google_calendar_delete_event"
317 | "google_calendar_get_free_busy"
318 | "gmail_send_email"
319 | "gmail_list_messages"
320 | "gmail_get_message"
321 | "hubspot_get_contact"
322 | "hubspot_create_contact"
323 | "hubspot_list_deals"
324 | "hubspot_create_deal"
325 | "linkedin_create_share"
326 | "linkedin_get_company_updates"
327 | "salesforce_soql_query"
328 | "salesforce_get_record"
329 | "salesforce_create_record"
330 | "slack_send_message"
331 | "slack_list_channels"
332 | "slack_upload_file"
333 )
334}
335
336fn history_messages_from_payload(payload: &serde_json::Value) -> Vec<(String, String)> {
337 payload
338 .get("messages")
339 .and_then(|v| v.as_array())
340 .map(|arr| {
341 arr.iter()
342 .filter_map(|m| {
343 Some((
344 m.get("role")?.as_str()?.to_string(),
345 m.get("content")?.as_str()?.to_string(),
346 ))
347 })
348 .collect()
349 })
350 .unwrap_or_default()
351}
352
353fn conversation_messages_from_payload(payload: &serde_json::Value) -> Vec<ConversationMessage> {
354 payload
355 .get("messages")
356 .cloned()
357 .and_then(|v| serde_json::from_value(v).ok())
358 .unwrap_or_default()
359}
360
361fn tools_from_payload(payload: &serde_json::Value) -> Vec<ToolDefinition> {
362 let Some(v) = payload.get("tools") else {
363 return Vec::new();
364 };
365 if let Ok(defs) = serde_json::from_value::<Vec<ToolDefinition>>(v.clone()) {
366 return defs;
367 }
368 v.as_array()
369 .map(|arr| {
370 arr.iter()
371 .filter_map(|t| {
372 let name = t.as_str().or_else(|| t.get("name")?.as_str())?;
373 Some(ToolDefinition {
374 name: name.to_string(),
375 description: String::new(),
376 parameters: serde_json::json!({}),
377 })
378 })
379 .collect()
380 })
381 .unwrap_or_default()
382}
383
384fn attempt_to_generate_json(attempt: &LlmAttemptResponse) -> serde_json::Value {
385 serde_json::json!({
386 "content": attempt.response.content,
387 "usage": attempt.response.usage,
388 "model_name": attempt.model_name,
389 "provider_name": attempt.provider_name,
390 "tool_calls": attempt.response.tool_calls,
391 "finish_reason": attempt.response.finish_reason,
392 })
393}
394
395fn generate_denied(out: &serde_json::Value) -> bool {
396 match out.get("deny") {
397 Some(serde_json::Value::Bool(true)) => true,
398 Some(serde_json::Value::String(s)) if !s.is_empty() => true,
399 _ => false,
400 }
401}
402
403async fn run_events_waterfall<F, Fut>(
407 events: &EventsService,
408 event: &str,
409 payload: serde_json::Value,
410 core: F,
411) -> std::result::Result<serde_json::Value, CordisError>
412where
413 F: FnOnce(serde_json::Value) -> Fut,
414 Fut: Future<Output = std::result::Result<serde_json::Value, CordisError>>,
415{
416 let (req_tx, req_rx) = tokio::sync::oneshot::channel();
417 let (res_tx, res_rx) = tokio::sync::oneshot::channel();
418 let wf = events.waterfall_around(event.to_string(), payload, move |p| async move {
419 let _ = req_tx.send(p);
420 match res_rx.await {
421 Ok(r) => r,
422 Err(_) => Err(CordisError::Fiber("llm generate core dropped".into())),
423 }
424 });
425 tokio::pin!(wf);
426 tokio::select! {
427 wf_res = &mut wf => wf_res,
428 req = req_rx => {
429 match req {
430 Ok(p) => {
431 let out = core(p).await;
432 let _ = res_tx.send(out);
433 wf.await
434 }
435 Err(_) => wf.await,
436 }
437 }
438 }
439}
440
441impl ConfigurableAgent {
442 #[deprecated(note = "use new_with_tool_service with ctx.get::<dyn ToolService>()")]
448 pub fn new(
449 name: &str,
450 config: &AgentConfig,
451 llm: Box<dyn LLMClient>,
452 tools: Option<Arc<Tools>>,
453 ) -> Self {
454 Self::new_with_provider(name, config, llm, tools, config.model.clone())
455 }
456
457 fn resolve_common_fields(
462 name: &str,
463 config: &AgentConfig,
464 ) -> (AgentType, String, Option<Vec<String>>) {
465 let agent_type = Self::name_to_type(name);
466 let system_prompt = config
467 .system_prompt
468 .clone()
469 .unwrap_or_else(|| Self::default_system_prompt(name));
470 let allowed_tools = config.allowed_tools.clone().or_else(|| {
472 if config.tools.is_empty() {
473 None
474 } else {
475 Some(config.tools.clone())
476 }
477 });
478 (agent_type, system_prompt, allowed_tools)
479 }
480
481 #[deprecated(note = "use new_with_provider_and_tool_service with ctx.get::<dyn ToolService>()")]
485 pub fn new_with_provider(
486 name: &str,
487 config: &AgentConfig,
488 llm: Box<dyn LLMClient>,
489 tools: Option<Arc<Tools>>,
490 provider_name: String,
491 ) -> Self {
492 let (agent_type, system_prompt, allowed_tools) = Self::resolve_common_fields(name, config);
493
494 Self {
495 name: name.to_string(),
496 agent_type,
497 llm,
498 provider_name,
499 system_prompt,
500 tools,
501 cordis_ctx: None,
502 allowed_tools,
503 max_tool_iterations: config.max_tool_iterations,
504 parallel_tools: config.parallel_tools,
505 observability: None,
506 fallback_llms: Vec::new(),
507 run_id: None,
508 compaction: config.compaction_enabled.unwrap_or(false),
509 }
510 }
511
512 #[deprecated(note = "use with_tool_service_params with ctx.get::<dyn ToolService>()")]
516 #[allow(clippy::too_many_arguments)]
517 pub fn with_params(
518 name: &str,
519 agent_type: AgentType,
520 llm: Box<dyn LLMClient>,
521 system_prompt: String,
522 tools: Option<Arc<Tools>>,
523 allowed_tools: Option<Vec<String>>,
524 max_tool_iterations: usize,
525 parallel_tools: bool,
526 ) -> Self {
527 Self {
528 name: name.to_string(),
529 agent_type,
530 llm,
531 provider_name: "unknown".to_string(),
532 system_prompt,
533 tools,
534 cordis_ctx: None,
535 allowed_tools,
536 max_tool_iterations,
537 parallel_tools,
538 observability: None,
539 fallback_llms: Vec::new(),
540 run_id: None,
541 compaction: false,
542 }
543 }
544
545 pub fn new_with_tool_service(
553 name: &str,
554 config: &AgentConfig,
555 llm: Box<dyn LLMClient>,
556 tool_service: Option<Arc<Tools>>,
557 ) -> Self {
558 Self::new_with_provider_and_tool_service(
559 name,
560 config,
561 llm,
562 tool_service,
563 config.model.clone(),
564 )
565 }
566
567 pub fn new_with_provider_and_tool_service(
569 name: &str,
570 config: &AgentConfig,
571 llm: Box<dyn LLMClient>,
572 tool_service: Option<Arc<Tools>>,
573 provider_name: String,
574 ) -> Self {
575 let (agent_type, system_prompt, allowed_tools) = Self::resolve_common_fields(name, config);
576 Self {
577 name: name.to_string(),
578 agent_type,
579 llm,
580 provider_name,
581 system_prompt,
582 tools: tool_service,
583 cordis_ctx: None,
584 allowed_tools,
585 max_tool_iterations: config.max_tool_iterations,
586 parallel_tools: config.parallel_tools,
587 observability: None,
588 fallback_llms: Vec::new(),
589 run_id: None,
590 compaction: config.compaction_enabled.unwrap_or(false),
591 }
592 }
593
594 pub fn set_compaction(&mut self, enabled: bool) {
596 self.compaction = enabled;
597 }
598
599 #[allow(clippy::too_many_arguments)]
601 pub fn with_tool_service_params(
602 name: &str,
603 agent_type: AgentType,
604 llm: Box<dyn LLMClient>,
605 system_prompt: String,
606 tool_service: Option<Arc<Tools>>,
607 allowed_tools: Option<Vec<String>>,
608 max_tool_iterations: usize,
609 parallel_tools: bool,
610 ) -> Self {
611 Self {
612 name: name.to_string(),
613 agent_type,
614 llm,
615 provider_name: "unknown".to_string(),
616 system_prompt,
617 tools: tool_service,
618 cordis_ctx: None,
619 allowed_tools,
620 max_tool_iterations,
621 parallel_tools,
622 observability: None,
623 fallback_llms: Vec::new(),
624 run_id: None,
625 compaction: false,
626 }
627 }
628
629 pub fn new_from_context(
633 ctx: &Arc<Context>,
634 name: &str,
635 config: &AgentConfig,
636 llm: Box<dyn LLMClient>,
637 ) -> Self {
638 let mut agent = Self::new_with_tool_service(name, config, llm, None);
639 if let Some(tools) = ctx.get::<Tools>() {
640 agent.set_tools(tools);
641 }
642 agent.bind_request_ctx(ctx.clone());
643 agent
644 }
645
646 pub fn new_from_context_with_provider(
648 ctx: &Arc<Context>,
649 name: &str,
650 config: &AgentConfig,
651 llm: Box<dyn LLMClient>,
652 provider_name: String,
653 ) -> Self {
654 let mut agent =
655 Self::new_with_provider_and_tool_service(name, config, llm, None, provider_name);
656 if let Some(tools) = ctx.get::<Tools>() {
657 agent.set_tools(tools);
658 }
659 agent.bind_request_ctx(ctx.clone());
660 agent
661 }
662
663 fn name_to_type(name: &str) -> AgentType {
665 AgentType::from_string(name)
666 }
667
668 fn default_system_prompt(name: &str) -> String {
670 match name.to_lowercase().as_str() {
671 "router" => r#"You are a routing agent that classifies user queries.
672Available agents: product, invoice, sales, finance, hr, orchestrator.
673Respond with ONLY the agent name (one word, lowercase)."#
674 .to_string(),
675
676 "orchestrator" => r#"You are an orchestrator agent for complex queries.
677Break down requests, delegate to specialists, and synthesize results."#
678 .to_string(),
679
680 "product" => r#"You are a Product Agent for product-related queries.
681Handle catalog, specifications, inventory, and pricing questions."#
682 .to_string(),
683
684 "invoice" => r#"You are an Invoice Agent for billing queries.
685Handle invoices, payments, and billing history."#
686 .to_string(),
687
688 "sales" => r#"You are a Sales Agent for sales analytics.
689Handle performance metrics, revenue, and customer data."#
690 .to_string(),
691
692 "finance" => r#"You are a Finance Agent for financial analysis.
693Handle statements, budgets, and expense management."#
694 .to_string(),
695
696 "hr" => r#"You are an HR Agent for human resources.
697Handle employee info, policies, and benefits."#
698 .to_string(),
699
700 _ => format!("You are a {} agent.", name),
701 }
702 }
703
704 pub fn name(&self) -> &str {
706 &self.name
707 }
708
709 pub fn max_tool_iterations(&self) -> usize {
711 self.max_tool_iterations
712 }
713
714 pub fn parallel_tools(&self) -> bool {
716 self.parallel_tools
717 }
718
719 pub fn has_tools(&self) -> bool {
721 self.tools.is_some()
722 }
723
724 pub fn tools(&self) -> Option<&Arc<Tools>> {
726 self.tools.as_ref()
727 }
728
729 pub fn set_tools(&mut self, tools: Arc<Tools>) {
731 self.tools = Some(tools);
732 }
733
734 pub fn bind_request_ctx(&mut self, ctx: Arc<Context>) {
737 if let Some(tools) = ctx.get::<Tools>() {
738 self.tools = Some(tools);
739 }
740 self.cordis_ctx = Some(ctx);
741 }
742
743 pub fn allowed_tools(&self) -> Option<&[String]> {
746 self.allowed_tools.as_deref()
747 }
748
749 pub fn set_allowed_tools(&mut self, allowed_tools: Option<Vec<String>>) {
752 self.allowed_tools = allowed_tools;
753 }
754
755 pub fn set_observability(&mut self, obs: Arc<dyn ObservabilitySink>) {
757 self.observability = Some(obs);
758 }
759
760 pub fn set_fallback_llms(&mut self, fallbacks: Vec<Box<dyn LLMClient>>) {
762 self.fallback_llms = fallbacks
763 .into_iter()
764 .map(|llm| ProviderLlm {
765 provider_name: "unknown".to_string(),
766 llm,
767 })
768 .collect();
769 }
770
771 pub fn set_fallback_llms_with_providers(
773 &mut self,
774 fallbacks: Vec<(String, Box<dyn LLMClient>)>,
775 ) {
776 self.fallback_llms = fallbacks
777 .into_iter()
778 .map(|(provider_name, llm)| ProviderLlm { provider_name, llm })
779 .collect();
780 }
781
782 pub fn set_run_id(&mut self, run_id: String) {
784 self.run_id = Some(run_id);
785 }
786
787 async fn session_compactor(&self, context: &AgentContext) -> Option<Arc<Compactor>> {
792 let ctx = self.cordis_ctx.as_ref()?;
793 let llm = ctx.get::<ares_llm::Llm>()?;
794 let registry = global_session_compactors();
795 let tenant = tenant_key_for_compaction(ctx.as_ref(), &context.user_id);
796 registry
797 .get_or_create(&tenant, &context.session_id, llm.as_ref(), ctx)
798 .await
799 }
800
801 async fn preflight_budget_check(&self, tenant_id: &str) -> Result<()> {
802 let Some(ctx) = self.cordis_ctx.as_ref() else {
803 return Ok(());
804 };
805 #[cfg(feature = "postgres")]
806 if let Some(db) = ctx.get::<ares_store::TenantDb>() {
807 let store = ares_store::token_budgets::TokenBudgetStore::new(db.pool());
808 let status = store.check_budget(tenant_id).await?;
809 if status.would_exceed {
810 return Err(AppError::RateLimited(format!(
811 "Tenant {} token budget exceeded ({} / {})",
812 tenant_id, status.tokens_used, status.token_limit
813 )));
814 }
815 }
816 Ok(())
817 }
818
819 async fn record_and_check_budget(
820 &self,
821 tenant_id: &str,
822 prompt_tokens: i64,
823 completion_tokens: i64,
824 ) -> Result<()> {
825 let Some(ctx) = self.cordis_ctx.as_ref() else {
826 return Ok(());
827 };
828 #[cfg(feature = "postgres")]
829 if let Some(db) = ctx.get::<ares_store::TenantDb>() {
830 let store = ares_store::token_budgets::TokenBudgetStore::new(db.pool());
831 store
832 .record_usage(
833 tenant_id,
834 self.run_id.as_deref(),
835 &self.name,
836 self.llm.model_name(),
837 prompt_tokens,
838 completion_tokens,
839 )
840 .await?;
841 let status = store.check_budget(tenant_id).await?;
842 if status.percentage >= status.alert_threshold {
843 tracing::warn!(
844 tenant_id,
845 usage_pct = status.percentage,
846 threshold = status.alert_threshold,
847 "Token budget alert threshold crossed"
848 );
849 }
850 if status.would_exceed {
851 tracing::warn!(
852 tenant_id,
853 remaining = status.remaining,
854 "Tenant token budget would be exceeded"
855 );
856 }
857 }
858 Ok(())
859 }
860
861 async fn try_generate_with_history(
863 &self,
864 messages: &[(String, String)],
865 ) -> Result<LlmAttemptResponse> {
866 let ctx = self.cordis_ctx.clone().unwrap_or_else(Context::new_root);
867 let Some(events) = ctx.get::<EventsService>() else {
868 return self.generate_with_history_direct(messages).await;
869 };
870 let orig: Vec<(String, String)> = messages.to_vec();
871 let payload = serde_json::to_value(cordis::LlmGeneratePayload {
872 messages: orig
873 .iter()
874 .map(|(role, content)| cordis::LlmMessage {
875 role: role.clone(),
876 content: serde_json::Value::String(content.clone()),
877 })
878 .collect(),
879 })
880 .unwrap_or(serde_json::Value::Null);
881 let out = run_events_waterfall(
882 &events,
883 cordis::events_catalog::ev::LLM_GENERATE,
884 payload,
885 |payload| async move {
886 let parsed = history_messages_from_payload(&payload);
887 let msgs = if parsed.is_empty() { orig } else { parsed };
888 match self.generate_with_history_direct(&msgs).await {
889 Ok(attempt) => Ok(attempt_to_generate_json(&attempt)),
890 Err(e) => Err(CordisError::Fiber(e.to_string())),
891 }
892 },
893 )
894 .await
895 .map_err(|e| AppError::Internal(e.to_string()))?;
896 self.generate_attempt_from_payload(out, cordis::events_catalog::ev::LLM_GENERATE)
897 }
898
899 async fn generate_with_history_direct(
900 &self,
901 messages: &[(String, String)],
902 ) -> Result<LlmAttemptResponse> {
903 match self.llm.generate_with_history(messages).await {
904 Ok(response) => Ok(LlmAttemptResponse {
905 response,
906 provider_name: self.provider_name.clone(),
907 model_name: self.llm.model_name().to_string(),
908 }),
909 Err(e) => {
910 let primary_error = e.to_string();
911 let mut fallback_errors = Vec::new();
912 for (i, fallback) in self.fallback_llms.iter().enumerate() {
913 tracing::warn!(
914 agent = %self.name,
915 fallback_idx = %i,
916 "Primary LLM failed, trying fallback"
917 );
918 match fallback.llm.generate_with_history(messages).await {
919 Ok(response) => {
920 tracing::info!(
921 agent = %self.name,
922 fallback_idx = %i,
923 provider = %fallback.provider_name,
924 "Fallback LLM succeeded"
925 );
926 return Ok(LlmAttemptResponse {
927 response,
928 provider_name: fallback.provider_name.clone(),
929 model_name: fallback.llm.model_name().to_string(),
930 });
931 }
932 Err(fallback_error) => {
933 fallback_errors.push(format!("fallback[{i}]: {fallback_error}"));
934 }
935 }
936 }
937 if fallback_errors.is_empty() {
938 Err(e)
939 } else {
940 Err(AppError::LLM(format!(
941 "All LLM providers failed for agent '{}'; primary: {}; {}",
942 self.name,
943 primary_error,
944 fallback_errors.join("; ")
945 )))
946 }
947 }
948 }
949 }
950
951 async fn try_generate_with_tools_and_history(
953 &self,
954 messages: &[ConversationMessage],
955 tools: &[ToolDefinition],
956 ) -> Result<LlmAttemptResponse> {
957 let ctx = self.cordis_ctx.clone().unwrap_or_else(Context::new_root);
958 let Some(events) = ctx.get::<EventsService>() else {
959 return self
960 .generate_with_tools_and_history_direct(messages, tools)
961 .await;
962 };
963 let orig_messages = messages.to_vec();
964 let orig_tools = tools.to_vec();
965 let payload = serde_json::to_value(cordis::LlmGenerateToolsPayload {
966 messages: orig_messages
967 .iter()
968 .map(|m| serde_json::to_value(m).unwrap_or(serde_json::Value::Null))
969 .collect(),
970 tools: orig_tools
971 .iter()
972 .map(|t| serde_json::to_value(t).unwrap_or(serde_json::Value::Null))
973 .collect(),
974 })
975 .unwrap_or(serde_json::Value::Null);
976 let out = run_events_waterfall(
977 &events,
978 cordis::events_catalog::ev::LLM_GENERATE_TOOLS,
979 payload,
980 |payload| async move {
981 let parsed_msgs = conversation_messages_from_payload(&payload);
982 let msgs = if parsed_msgs.is_empty() {
983 orig_messages
984 } else {
985 parsed_msgs
986 };
987 let parsed_tools = tools_from_payload(&payload);
988 let tool_defs = if parsed_tools.is_empty() && payload.get("tools").is_none() {
989 orig_tools
990 } else {
991 parsed_tools
992 };
993 match self
994 .generate_with_tools_and_history_direct(&msgs, &tool_defs)
995 .await
996 {
997 Ok(attempt) => Ok(attempt_to_generate_json(&attempt)),
998 Err(e) => Err(CordisError::Fiber(e.to_string())),
999 }
1000 },
1001 )
1002 .await
1003 .map_err(|e| AppError::Internal(e.to_string()))?;
1004 self.generate_attempt_from_payload(out, cordis::events_catalog::ev::LLM_GENERATE_TOOLS)
1005 }
1006
1007 async fn generate_with_tools_and_history_direct(
1008 &self,
1009 messages: &[ConversationMessage],
1010 tools: &[ToolDefinition],
1011 ) -> Result<LlmAttemptResponse> {
1012 match self
1013 .llm
1014 .generate_with_tools_and_history(messages, tools)
1015 .await
1016 {
1017 Ok(response) => Ok(LlmAttemptResponse {
1018 response,
1019 provider_name: self.provider_name.clone(),
1020 model_name: self.llm.model_name().to_string(),
1021 }),
1022 Err(e) => {
1023 let primary_error = e.to_string();
1024 let mut fallback_errors = Vec::new();
1025 for (i, fallback) in self.fallback_llms.iter().enumerate() {
1026 tracing::warn!(
1027 agent = %self.name,
1028 fallback_idx = %i,
1029 "Primary LLM (tools) failed, trying fallback"
1030 );
1031 match fallback
1032 .llm
1033 .generate_with_tools_and_history(messages, tools)
1034 .await
1035 {
1036 Ok(response) => {
1037 tracing::info!(
1038 agent = %self.name,
1039 fallback_idx = %i,
1040 provider = %fallback.provider_name,
1041 "Fallback LLM (tools) succeeded"
1042 );
1043 return Ok(LlmAttemptResponse {
1044 response,
1045 provider_name: fallback.provider_name.clone(),
1046 model_name: fallback.llm.model_name().to_string(),
1047 });
1048 }
1049 Err(fallback_error) => {
1050 fallback_errors.push(format!("fallback[{i}]: {fallback_error}"));
1051 }
1052 }
1053 }
1054 if fallback_errors.is_empty() {
1055 Err(e)
1056 } else {
1057 Err(AppError::LLM(format!(
1058 "All LLM providers failed for agent '{}'; primary: {}; {}",
1059 self.name,
1060 primary_error,
1061 fallback_errors.join("; ")
1062 )))
1063 }
1064 }
1065 }
1066 }
1067
1068 fn generate_attempt_from_payload(
1069 &self,
1070 out: serde_json::Value,
1071 event: &str,
1072 ) -> Result<LlmAttemptResponse> {
1073 if generate_denied(&out) {
1074 let reason = out
1075 .get("reason")
1076 .and_then(|v| v.as_str())
1077 .unwrap_or("denied");
1078 return Err(AppError::InvalidInput(format!("{event} {reason}")));
1079 }
1080 let content = out
1081 .get("content")
1082 .and_then(|v| v.as_str())
1083 .unwrap_or("")
1084 .to_string();
1085 let usage = out
1086 .get("usage")
1087 .cloned()
1088 .and_then(|v| serde_json::from_value(v).ok());
1089 let model_name = out
1090 .get("model_name")
1091 .and_then(|v| v.as_str())
1092 .map(str::to_string)
1093 .unwrap_or_else(|| self.llm.model_name().to_string());
1094 let provider_name = out
1095 .get("provider_name")
1096 .and_then(|v| v.as_str())
1097 .map(str::to_string)
1098 .unwrap_or_else(|| self.provider_name.clone());
1099 let tool_calls = out
1100 .get("tool_calls")
1101 .cloned()
1102 .and_then(|v| serde_json::from_value(v).ok())
1103 .unwrap_or_default();
1104 let finish_reason = out
1105 .get("finish_reason")
1106 .and_then(|v| v.as_str())
1107 .unwrap_or("stop")
1108 .to_string();
1109 Ok(LlmAttemptResponse {
1110 response: LLMResponse {
1111 content,
1112 tool_calls,
1113 finish_reason,
1114 usage,
1115 },
1116 provider_name,
1117 model_name,
1118 })
1119 }
1120
1121 pub fn get_filtered_tool_definitions(&self) -> Vec<ToolDefinition> {
1126 let (Some(tools), Some(allowed)) = (&self.tools, &self.allowed_tools) else {
1127 return Vec::new();
1128 };
1129 if allowed.is_empty() {
1130 return Vec::new();
1131 }
1132 let ctx = self.cordis_ctx.clone().unwrap_or_else(Context::new_root);
1133 tools
1134 .list(&ctx)
1135 .into_iter()
1136 .filter(|def| allowed.iter().any(|a| a == &def.name))
1137 .collect()
1138 }
1139
1140 pub fn can_use_tool(&self, tool_name: &str) -> bool {
1143 let whitelisted = match &self.allowed_tools {
1144 Some(allowed) => allowed.iter().any(|allowed| allowed == tool_name),
1145 None => false,
1146 };
1147 if !whitelisted {
1148 return false;
1149 }
1150 self.tools.is_some()
1152 }
1153
1154 fn tenant_scoped_builtin_args(
1155 &self,
1156 name: &str,
1157 mut args: serde_json::Value,
1158 ) -> Result<serde_json::Value> {
1159 if !is_prebuilt_connector_tool(name) {
1160 return Ok(args);
1161 }
1162 let Some(tenant_id) = self.connector_tenant_id() else {
1163 return Err(AppError::InvalidInput(
1164 "tenant_id is required for connector tool execution".to_string(),
1165 ));
1166 };
1167 let serde_json::Value::Object(map) = &mut args else {
1168 return Err(AppError::InvalidInput(
1169 "connector tool arguments must be an object".to_string(),
1170 ));
1171 };
1172 if let Some(provided) = map.get("tenant_id").and_then(|value| value.as_str()) {
1173 if provided != tenant_id {
1174 return Err(AppError::Auth(
1175 "connector tenant_id does not match executing tenant".to_string(),
1176 ));
1177 }
1178 }
1179 map.insert(
1180 "tenant_id".to_string(),
1181 serde_json::Value::String(tenant_id),
1182 );
1183 Ok(args)
1184 }
1185
1186 fn connector_tenant_id(&self) -> Option<String> {
1187 let ctx = self.cordis_ctx.as_ref()?;
1188 let id = crate::user_id_from_ctx(ctx, "");
1189 if id.is_empty() {
1190 None
1191 } else {
1192 Some(id)
1193 }
1194 }
1195
1196 async fn dispatch_tool(
1198 &self,
1199 name: &str,
1200 args: serde_json::Value,
1201 ) -> Result<serde_json::Value> {
1202 let Some(tools) = &self.tools else {
1203 return Err(AppError::NotFound(format!("Tool not found: {name}")));
1204 };
1205 let ctx = self.cordis_ctx.clone().unwrap_or_else(Context::new_root);
1206 let args = self.tenant_scoped_builtin_args(name, args)?;
1207 tools.execute(&ctx, name, args).await
1208 }
1209
1210 fn observed_tool_type(&self, name: &str, is_builtin: bool) -> String {
1211 if is_builtin {
1212 return "builtin".to_string();
1213 }
1214 let _ = name;
1215 "runtime".to_string()
1216 }
1217
1218 fn effective_system_prompt(&self) -> String {
1219 if let Some(ctx) = &self.cordis_ctx {
1220 if let Some(ext) = ctx.get::<crate::external_context::ExternalContext>() {
1221 if !ext.0.is_empty() {
1222 tracing::debug!(
1223 agent = %self.name,
1224 ctx_len = ext.0.len(),
1225 "External context injected into system prompt"
1226 );
1227 return format!(
1228 "{}
1229
1230{}
1231
1232When referencing facts above, cite [E1], [E2] etc.",
1233 ext.0, self.system_prompt
1234 );
1235 }
1236 }
1237 }
1238 self.system_prompt.clone()
1239 }
1240
1241 async fn execute_with_tools(
1243 &self,
1244 input: &str,
1245 context: &AgentContext,
1246 ) -> Result<AgentResponse> {
1247 use ares_llm::client::TokenUsage;
1248
1249 let tools = self.get_filtered_tool_definitions();
1250 tracing::debug!(
1251 agent = %self.name,
1252 allowed_tools = ?self.allowed_tools,
1253 tool_count = tools.len(),
1254 "execute_with_tools: tool definitions loaded"
1255 );
1256
1257 let mut messages: Vec<ConversationMessage> = Vec::new();
1258
1259 let effective_prompt = self.effective_system_prompt();
1262 messages.push(ConversationMessage::system(&effective_prompt));
1263
1264 let session_compactor = if self.compaction {
1268 self.session_compactor(context).await
1269 } else {
1270 None
1271 };
1272 #[cfg(feature = "postgres")]
1274 let snapshot_pool: Option<sqlx::PgPool> = if session_compactor.is_some() {
1275 self.cordis_ctx
1276 .as_ref()
1277 .and_then(|ctx| ctx.get::<ares_store::TenantDb>())
1278 .map(|db| db.pool().clone())
1279 } else {
1280 None
1281 };
1282 #[cfg(feature = "postgres")]
1283 let session_key = context.session_id.clone();
1284 #[cfg(feature = "postgres")]
1285 let tenant_key: Option<String> = self.cordis_ctx.as_ref().and_then(|ctx| {
1286 let key = tenant_key_for_compaction(ctx, &context.user_id);
1287 if key.is_empty() { None } else { Some(key) }
1288 });
1289 #[cfg(not(feature = "postgres"))]
1290 let _ = &context;
1291 if let Some(compactor) = &session_compactor {
1292 for (role, content) in compactor.build_context("", 5).into_iter().skip(1) {
1293 if role == "assistant" {
1294 messages.push(ConversationMessage::assistant(&content, vec![]));
1295 } else if role == "user" {
1296 messages.push(ConversationMessage::user(&content));
1297 } else {
1298 messages.push(ConversationMessage::system(&content));
1299 }
1300 }
1301 } else {
1302 for msg in context.conversation_history.iter().rev().take(5).rev() {
1304 let cm = match msg.role {
1305 ares_types::types::MessageRole::User => ConversationMessage::user(&msg.content),
1306 ares_types::types::MessageRole::Assistant => {
1307 ConversationMessage::assistant(&msg.content, vec![])
1308 }
1309 _ => ConversationMessage::system(&msg.content),
1310 };
1311 messages.push(cm);
1312 }
1313 }
1314
1315 messages.push(ConversationMessage::user(input));
1316
1317 let mut total_usage = TokenUsage::default();
1318 let mut last_provider_name = self.provider_name.clone();
1319 let mut last_model_name = self.llm.model_name().to_string();
1320
1321 for iteration in 0..self.max_tool_iterations {
1322 self.preflight_budget_check(&context.user_id).await?;
1323
1324 let llm_start = std::time::Instant::now();
1325 let attempt = self
1326 .try_generate_with_tools_and_history(&messages, &tools)
1327 .await?;
1328 let llm_latency = llm_start.elapsed().as_millis() as i64;
1329 last_provider_name = attempt.provider_name;
1330 last_model_name = attempt.model_name;
1331 let response = attempt.response;
1332
1333 {
1334 let prompt_tok = response
1335 .usage
1336 .as_ref()
1337 .map(|u| u.prompt_tokens as i64)
1338 .unwrap_or(0);
1339 let completion_tok = response
1340 .usage
1341 .as_ref()
1342 .map(|u| u.completion_tokens as i64)
1343 .unwrap_or(0);
1344 self.record_and_check_budget(&context.user_id, prompt_tok, completion_tok)
1345 .await?;
1346 }
1347
1348 if let Some(obs) = &self.observability {
1350 let prompt_tok = response
1351 .usage
1352 .as_ref()
1353 .map(|u| u.prompt_tokens as i64)
1354 .unwrap_or(0);
1355 let completion_tok = response
1356 .usage
1357 .as_ref()
1358 .map(|u| u.completion_tokens as i64)
1359 .unwrap_or(0);
1360 let record = LlmCallRecord {
1361 step_index: iteration as i32,
1362 provider: last_provider_name.clone(),
1363 model: last_model_name.clone(),
1364 prompt_tokens: prompt_tok,
1365 completion_tokens: completion_tok,
1366 latency_ms: llm_latency,
1367 status: "success".to_string(),
1368 cached_tokens: response.usage.as_ref().and_then(|u| u.cached_tokens),
1369 total_time_ms: Some(llm_latency),
1370 };
1371 let _ = obs.log_llm_call(record).await;
1372 }
1373
1374 if let Some(usage) = &response.usage {
1375 total_usage = TokenUsage::new(
1376 total_usage.prompt_tokens + usage.prompt_tokens,
1377 total_usage.completion_tokens + usage.completion_tokens,
1378 );
1379 }
1380
1381 if response.tool_calls.is_empty() {
1382 if let Some(compactor) = &session_compactor {
1385 let compactor = Arc::clone(compactor);
1386 let input = input.to_string();
1387 let content = response.content.clone();
1388 tokio::spawn(async move {
1389 let event = compactor.record_turn(input, content).await;
1390 log_compact_event("record_turn", &event);
1391 for event in compactor.audit_if_due().await {
1392 log_compact_event("audit", &event);
1393 }
1394 #[cfg(feature = "postgres")]
1395 if let Some(pool) = snapshot_pool.clone() {
1396 persist_session_snapshot(&compactor, pool, &session_key).await;
1397 }
1398 });
1399 }
1400 return Ok(AgentResponse {
1401 content: response.content,
1402 usage: Some(total_usage),
1403 metadata: Some(ExecutionMetadata {
1404 model_name: last_model_name,
1405 provider_name: last_provider_name,
1406 }),
1407 });
1408 }
1409
1410 messages.push(ConversationMessage::assistant(
1412 &response.content,
1413 response.tool_calls.clone(),
1414 ));
1415
1416 for tc in &response.tool_calls {
1418 if !self.can_use_tool(&tc.name) {
1420 tracing::warn!(
1421 agent = %self.name,
1422 tool = %tc.name,
1423 allowed_tools = ?self.allowed_tools,
1424 "Tool not in allowed_tools list — denying execution"
1425 );
1426 return Err(AppError::Auth(format!(
1427 "Tool '{}' is not allowed for this agent",
1428 tc.name
1429 )));
1430 }
1431
1432 let tool_start = std::time::Instant::now();
1433 let is_builtin = {
1434 let ctx = self.cordis_ctx.clone().unwrap_or_else(Context::new_root);
1435 self.tools
1436 .as_ref()
1437 .and_then(|t| t.resolve(&ctx, &tc.name))
1438 .is_some()
1439 };
1440 let tool_type = self.observed_tool_type(&tc.name, is_builtin);
1441 let result = self.dispatch_tool(&tc.name, tc.arguments.clone()).await;
1442 let tool_latency = tool_start.elapsed().as_millis() as i64;
1443 let result_value = match result {
1444 Ok(v) => v,
1445 Err(e) => serde_json::json!({"error": e.to_string()}),
1446 };
1447
1448 if let Some(obs) = &self.observability {
1450 let status = if result_value.get("error").is_some() {
1451 "error".to_string()
1452 } else {
1453 "success".to_string()
1454 };
1455 let tool_record = ToolCallRecord {
1456 step_index: iteration as i32,
1457 tool_name: tc.name.clone(),
1458 tool_type,
1459 arguments: tc.arguments.clone(),
1460 result: Some(result_value.clone()),
1461 latency_ms: tool_latency,
1462 status,
1463 };
1464 let _ = obs.log_tool_call(tool_record).await;
1465 }
1466
1467 messages.push(ConversationMessage::tool_result(&tc.id, &result_value));
1468 }
1469 }
1470
1471 tracing::warn!(
1475 agent = %self.name,
1476 "Max tool iterations ({}) reached — making final synthesis call",
1477 self.max_tool_iterations
1478 );
1479 self.preflight_budget_check(&context.user_id).await?;
1480
1481 let synth_start = std::time::Instant::now();
1482 let final_response = self
1483 .try_generate_with_tools_and_history(&messages, &[])
1484 .await;
1485 let synth_latency = synth_start.elapsed().as_millis() as i64;
1486 if let Ok(attempt) = &final_response {
1487 last_provider_name = attempt.provider_name.clone();
1488 last_model_name = attempt.model_name.clone();
1489 }
1490
1491 if let Ok(attempt) = &final_response {
1492 let prompt_tok = attempt
1493 .response
1494 .usage
1495 .as_ref()
1496 .map(|u| u.prompt_tokens as i64)
1497 .unwrap_or(0);
1498 let completion_tok = attempt
1499 .response
1500 .usage
1501 .as_ref()
1502 .map(|u| u.completion_tokens as i64)
1503 .unwrap_or(0);
1504 let _ = self
1505 .record_and_check_budget(&context.user_id, prompt_tok, completion_tok)
1506 .await;
1507 }
1508
1509 if let Some(obs) = &self.observability {
1511 let (prompt_tok, completion_tok, status) = match &final_response {
1512 Ok(attempt) => (
1513 attempt
1514 .response
1515 .usage
1516 .as_ref()
1517 .map(|u| u.prompt_tokens as i64)
1518 .unwrap_or(0),
1519 attempt
1520 .response
1521 .usage
1522 .as_ref()
1523 .map(|u| u.completion_tokens as i64)
1524 .unwrap_or(0),
1525 "success".to_string(),
1526 ),
1527 Err(_) => (0, 0, "error".to_string()),
1528 };
1529 let record = LlmCallRecord {
1530 step_index: self.max_tool_iterations as i32,
1531 provider: last_provider_name.clone(),
1532 model: last_model_name.clone(),
1533 prompt_tokens: prompt_tok,
1534 completion_tokens: completion_tok,
1535 latency_ms: synth_latency,
1536 cached_tokens: final_response
1537 .as_ref()
1538 .ok()
1539 .and_then(|attempt| attempt.response.usage.as_ref())
1540 .and_then(|u| u.cached_tokens),
1541 total_time_ms: Some(synth_latency),
1542 status,
1543 };
1544 let _ = obs.log_llm_call(record).await;
1545 }
1546
1547 if let (Ok(attempt), Some(compactor)) = (&final_response, &session_compactor) {
1550 if !attempt.response.content.is_empty() {
1551 let compactor = Arc::clone(compactor);
1552 let input = input.to_string();
1553 let content = attempt.response.content.clone();
1554 tokio::spawn(async move {
1555 let event = compactor.record_turn(input, content).await;
1556 log_compact_event("record_turn", &event);
1557 for event in compactor.audit_if_due().await {
1558 log_compact_event("audit", &event);
1559 }
1560 #[cfg(feature = "postgres")]
1561 if let Some(pool) = snapshot_pool.clone() {
1562 persist_session_snapshot(&compactor, pool, &session_key).await;
1563 }
1564 });
1565 }
1566 }
1567
1568 let content = match final_response {
1569 Ok(attempt) if !attempt.response.content.is_empty() => attempt.response.content,
1570 Ok(_) => {
1571 messages
1573 .iter()
1574 .rev()
1575 .find(|m| {
1576 m.role == ares_llm::coordinator::MessageRole::Assistant
1577 && !m.content.is_empty()
1578 })
1579 .map(|m| m.content.clone())
1580 .unwrap_or_else(|| {
1581 "Agent completed tool calls but could not generate a final response."
1582 .to_string()
1583 })
1584 }
1585 Err(e) => {
1586 tracing::error!(error = %e, "Final synthesis call failed");
1587 messages
1589 .iter()
1590 .rev()
1591 .find(|m| {
1592 m.role == ares_llm::coordinator::MessageRole::Assistant
1593 && !m.content.is_empty()
1594 })
1595 .map(|m| m.content.clone())
1596 .unwrap_or_else(|| format!("Agent completed but synthesis failed: {}", e))
1597 }
1598 };
1599
1600 Ok(AgentResponse {
1601 content,
1602 usage: Some(total_usage),
1603 metadata: Some(ExecutionMetadata {
1604 model_name: last_model_name,
1605 provider_name: last_provider_name,
1606 }),
1607 })
1608 }
1609}
1610
1611#[async_trait]
1612impl Agent for ConfigurableAgent {
1613 async fn execute(&self, input: &str, context: &AgentContext) -> Result<AgentResponse> {
1614 if self.has_tools() {
1615 tracing::debug!(agent = %self.name, "execute: using tool-calling path");
1616 return self.execute_with_tools(input, context).await;
1617 }
1618 tracing::debug!(agent = %self.name, "execute: no tools, using simple path");
1619
1620 let effective_prompt = self.effective_system_prompt();
1623 let mut messages = vec![("system".to_string(), effective_prompt)];
1624
1625 if let Some(memory) = &context.user_memory {
1627 let memory_context = format!(
1628 "User preferences: {}",
1629 memory
1630 .preferences
1631 .iter()
1632 .map(|p| format!("{}: {}", p.key, p.value))
1633 .collect::<Vec<_>>()
1634 .join(", ")
1635 );
1636 messages.push(("system".to_string(), memory_context));
1637 }
1638
1639 let session_compactor = if self.compaction {
1643 self.session_compactor(context).await
1644 } else {
1645 None
1646 };
1647 #[cfg(feature = "postgres")]
1649 let snapshot_pool: Option<sqlx::PgPool> = if session_compactor.is_some() {
1650 self.cordis_ctx
1651 .as_ref()
1652 .and_then(|ctx| ctx.get::<ares_store::TenantDb>())
1653 .map(|db| db.pool().clone())
1654 } else {
1655 None
1656 };
1657 #[cfg(feature = "postgres")]
1658 let session_key = context.session_id.clone();
1659 #[cfg(feature = "postgres")]
1660 let tenant_key: Option<String> = self.cordis_ctx.as_ref().and_then(|ctx| {
1661 let key = tenant_key_for_compaction(ctx, &context.user_id);
1662 if key.is_empty() { None } else { Some(key) }
1663 });
1664 #[cfg(not(feature = "postgres"))]
1665 let _ = &context;
1666 if let Some(compactor) = &session_compactor {
1667 for (role, content) in compactor.build_context("", 5).into_iter().skip(1) {
1668 messages.push((role, content));
1669 }
1670 } else {
1671 for msg in context.conversation_history.iter().rev().take(5).rev() {
1673 let role = match msg.role {
1674 ares_types::types::MessageRole::User => "user",
1675 ares_types::types::MessageRole::Assistant => "assistant",
1676 _ => "system",
1677 };
1678 messages.push((role.to_string(), msg.content.clone()));
1679 }
1680 }
1681
1682 messages.push(("user".to_string(), input.to_string()));
1683
1684 self.preflight_budget_check(&context.user_id).await?;
1685
1686 let llm_start = std::time::Instant::now();
1687 let attempt = self.try_generate_with_history(&messages).await?;
1688 let llm_latency = llm_start.elapsed().as_millis() as i64;
1689 let provider_name = attempt.provider_name;
1690 let model_name = attempt.model_name;
1691 let llm_response = attempt.response;
1692
1693 {
1694 let prompt_tok = llm_response
1695 .usage
1696 .as_ref()
1697 .map(|u| u.prompt_tokens as i64)
1698 .unwrap_or(0);
1699 let completion_tok = llm_response
1700 .usage
1701 .as_ref()
1702 .map(|u| u.completion_tokens as i64)
1703 .unwrap_or(0);
1704 self.record_and_check_budget(&context.user_id, prompt_tok, completion_tok)
1705 .await?;
1706 }
1707
1708 if let Some(obs) = &self.observability {
1710 let prompt_tok = llm_response
1711 .usage
1712 .as_ref()
1713 .map(|u| u.prompt_tokens as i64)
1714 .unwrap_or(0);
1715 let completion_tok = llm_response
1716 .usage
1717 .as_ref()
1718 .map(|u| u.completion_tokens as i64)
1719 .unwrap_or(0);
1720 let record = LlmCallRecord {
1721 step_index: 0,
1722 provider: provider_name.clone(),
1723 model: model_name.clone(),
1724 prompt_tokens: prompt_tok,
1725 completion_tokens: completion_tok,
1726 latency_ms: llm_latency,
1727 status: "success".to_string(),
1728 cached_tokens: llm_response.usage.as_ref().and_then(|u| u.cached_tokens),
1729 total_time_ms: Some(llm_latency),
1730 };
1731 let _ = obs.log_llm_call(record).await;
1732 }
1733
1734 if let Some(compactor) = &session_compactor {
1737 let compactor = Arc::clone(compactor);
1738 let input = input.to_string();
1739 let content = llm_response.content.clone();
1740 tokio::spawn(async move {
1741 let event = compactor.record_turn(input, content).await;
1742 log_compact_event("record_turn", &event);
1743 for event in compactor.audit_if_due().await {
1744 log_compact_event("audit", &event);
1745 }
1746 #[cfg(feature = "postgres")]
1747 if let Some(pool) = snapshot_pool.clone() {
1748 persist_session_snapshot(&compactor, pool, &session_key).await;
1749 }
1750 });
1751 }
1752
1753 Ok(AgentResponse {
1754 content: llm_response.content,
1755 usage: llm_response.usage,
1756 metadata: Some(ExecutionMetadata {
1757 model_name,
1758 provider_name,
1759 }),
1760 })
1761 }
1762
1763 fn system_prompt(&self) -> String {
1764 self.system_prompt.clone()
1765 }
1766
1767 fn agent_type(&self) -> AgentType {
1768 self.agent_type.clone()
1769 }
1770}
1771
1772#[cfg(feature = "postgres")]
1776pub fn agent_config_from_user_agent(user_agent: &ares_store::postgres::UserAgent) -> AgentConfig {
1777 AgentConfig {
1778 model: user_agent.model.clone(),
1779 system_prompt: user_agent.system_prompt.clone(),
1780 tools: user_agent.tools_vec(),
1781 max_tool_iterations: user_agent.max_tool_iterations as usize,
1782 parallel_tools: user_agent.parallel_tools,
1783 compaction_enabled: None,
1784 allowed_tools: None,
1785 extra: std::collections::HashMap::new(),
1786 }
1787}
1788
1789#[cfg(test)]
1790mod tests {
1791 use super::*;
1792 use crate::AgentConfig;
1793 use ares_llm::client::TokenUsage;
1794 use ares_llm::LLMResponse;
1795 use ares_tools::Tool;
1796 use ares_types::types::{Message, MessageRole, Preference, ToolCall, UserMemory};
1797 use chrono::Utc;
1798 use std::collections::{HashMap, VecDeque};
1799 use std::sync::atomic::{AtomicBool, Ordering};
1800 use std::sync::{Arc, Mutex};
1801
1802 struct MockLLM {
1810 content: String,
1811 tool_responses: Arc<Mutex<VecDeque<LLMResponse>>>,
1812 generated: Arc<AtomicBool>,
1813 echo_last: bool,
1814 }
1815
1816 impl MockLLM {
1817 fn new() -> Self {
1818 Self::with_content("mock")
1819 }
1820
1821 fn with_content(content: &str) -> Self {
1822 Self {
1823 content: content.to_string(),
1824 tool_responses: Arc::new(Mutex::new(VecDeque::new())),
1825 generated: Arc::new(AtomicBool::new(false)),
1826 echo_last: false,
1827 }
1828 }
1829
1830 fn with_tool_responses(responses: Vec<LLMResponse>) -> Self {
1833 Self {
1834 content: "mock".to_string(),
1835 tool_responses: Arc::new(Mutex::new(responses.into())),
1836 generated: Arc::new(AtomicBool::new(false)),
1837 echo_last: false,
1838 }
1839 }
1840
1841 fn echo_last() -> Self {
1842 let mut llm = Self::new();
1843 llm.echo_last = true;
1844 llm
1845 }
1846
1847 fn with_generated_flag() -> (Self, Arc<AtomicBool>) {
1848 let generated = Arc::new(AtomicBool::new(false));
1849 (
1850 Self {
1851 content: "should-not-appear".to_string(),
1852 tool_responses: Arc::new(Mutex::new(VecDeque::new())),
1853 generated: Arc::clone(&generated),
1854 echo_last: false,
1855 },
1856 generated,
1857 )
1858 }
1859 }
1860
1861 #[async_trait]
1862 impl LLMClient for MockLLM {
1863 async fn generate(&self, _: &str) -> Result<String> {
1864 Ok(self.content.clone())
1865 }
1866 async fn generate_with_system(&self, _: &str, _: &str) -> Result<String> {
1867 Ok(self.content.clone())
1868 }
1869 async fn generate_with_history(
1870 &self,
1871 messages: &[(String, String)],
1872 ) -> Result<LLMResponse> {
1873 self.generated.store(true, Ordering::SeqCst);
1874 let content = if self.echo_last {
1875 messages
1876 .last()
1877 .map(|(_, c)| c.clone())
1878 .unwrap_or_else(|| self.content.clone())
1879 } else {
1880 self.content.clone()
1881 };
1882 Ok(LLMResponse {
1883 content,
1884 tool_calls: vec![],
1885 finish_reason: "stop".to_string(),
1886 usage: None,
1887 })
1888 }
1889 async fn generate_with_tools(&self, _: &str, _: &[ToolDefinition]) -> Result<LLMResponse> {
1890 Ok(LLMResponse {
1891 content: self.content.clone(),
1892 tool_calls: vec![],
1893 finish_reason: "stop".to_string(),
1894 usage: None,
1895 })
1896 }
1897 async fn generate_with_tools_and_history(
1898 &self,
1899 _: &[ares_llm::coordinator::ConversationMessage],
1900 _: &[ToolDefinition],
1901 ) -> Result<LLMResponse> {
1902 let mut q = self.tool_responses.lock().unwrap();
1903 Ok(q.pop_front().unwrap_or(LLMResponse {
1904 content: self.content.clone(),
1905 tool_calls: vec![],
1906 finish_reason: "stop".to_string(),
1907 usage: None,
1908 }))
1909 }
1910 async fn stream(
1911 &self,
1912 _: &str,
1913 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
1914 Ok(Box::new(futures::stream::empty()))
1915 }
1916 async fn stream_with_system(
1917 &self,
1918 _: &str,
1919 _: &str,
1920 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
1921 Ok(Box::new(futures::stream::empty()))
1922 }
1923 async fn stream_with_history(
1924 &self,
1925 _: &[(String, String)],
1926 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
1927 Ok(Box::new(futures::stream::empty()))
1928 }
1929 fn model_name(&self) -> &str {
1930 "mock"
1931 }
1932 }
1933
1934 struct MockTool {
1937 name: String,
1938 description: String,
1939 }
1940
1941 struct EchoArgsTool {
1942 name: String,
1943 }
1944
1945 impl MockTool {
1946 fn new(name: &str) -> Self {
1947 Self {
1948 name: name.to_string(),
1949 description: format!("Mock tool: {}", name),
1950 }
1951 }
1952 }
1953
1954 #[async_trait]
1955 impl ares_tools::Tool for EchoArgsTool {
1956 fn name(&self) -> &str {
1957 &self.name
1958 }
1959 fn description(&self) -> &str {
1960 "Echoes arguments"
1961 }
1962 fn parameters_schema(&self) -> serde_json::Value {
1963 serde_json::json!({"type":"object"})
1964 }
1965 async fn execute(&self, args: serde_json::Value) -> Result<serde_json::Value> {
1966 Ok(args)
1967 }
1968 }
1969
1970 #[async_trait]
1971 impl ares_tools::Tool for MockTool {
1972 fn name(&self) -> &str {
1973 &self.name
1974 }
1975 fn description(&self) -> &str {
1976 &self.description
1977 }
1978 fn parameters_schema(&self) -> serde_json::Value {
1979 serde_json::json!({})
1980 }
1981 async fn execute(&self, _args: serde_json::Value) -> Result<serde_json::Value> {
1982 Ok(serde_json::json!({"result": "ok"}))
1983 }
1984 }
1985
1986 fn make_config(tools: Vec<&str>, system_prompt: Option<&str>) -> AgentConfig {
1989 AgentConfig {
1990 model: "default".to_string(),
1991 system_prompt: system_prompt.map(String::from),
1992 tools: tools.into_iter().map(String::from).collect(),
1993 allowed_tools: None,
1994 max_tool_iterations: 5,
1995 parallel_tools: false,
1996 compaction_enabled: None,
1997 extra: HashMap::new(),
1998 }
1999 }
2000
2001 fn make_context() -> AgentContext {
2002 AgentContext {
2003 user_id: "test-user".to_string(),
2004 session_id: "test-session".to_string(),
2005 conversation_history: vec![],
2006 user_memory: None,
2007 }
2008 }
2009
2010 fn make_context_with_history(history: Vec<(MessageRole, &str)>) -> AgentContext {
2011 AgentContext {
2012 user_id: "test-user".to_string(),
2013 session_id: "test-session".to_string(),
2014 conversation_history: history
2015 .into_iter()
2016 .map(|(role, content)| Message {
2017 role,
2018 content: content.to_string(),
2019 timestamp: Utc::now(),
2020 })
2021 .collect(),
2022 user_memory: None,
2023 }
2024 }
2025
2026 fn make_tool_response(content: &str, calls: Vec<ToolCall>) -> LLMResponse {
2027 let finish_reason = if calls.is_empty() {
2028 "stop"
2029 } else {
2030 "tool_calls"
2031 };
2032 LLMResponse {
2033 content: content.to_string(),
2034 tool_calls: calls,
2035 finish_reason: finish_reason.to_string(),
2036 usage: Some(TokenUsage::new(10, 5)),
2037 }
2038 }
2039
2040 fn make_tools_with_tool(name: &str) -> Arc<Tools> {
2041 Arc::new(Tools::from_static([
2042 Arc::new(MockTool::new(name)) as Arc<dyn Tool>
2043 ]))
2044 }
2045
2046 fn make_tools_with_echo_tool(name: &str) -> Arc<Tools> {
2047 Arc::new(Tools::from_static([Arc::new(EchoArgsTool {
2048 name: name.to_string(),
2049 }) as Arc<dyn Tool>]))
2050 }
2051
2052 #[test]
2057 fn test_default_system_prompt_router() {
2058 let p = ConfigurableAgent::default_system_prompt("router");
2059 assert!(
2060 p.contains("routing"),
2061 "router prompt should mention 'routing'"
2062 );
2063 }
2064
2065 #[test]
2066 fn test_default_system_prompt_orchestrator() {
2067 let p = ConfigurableAgent::default_system_prompt("orchestrator");
2068 assert!(p.contains("orchestrator"));
2069 }
2070
2071 #[test]
2072 fn test_default_system_prompt_product() {
2073 let p = ConfigurableAgent::default_system_prompt("product");
2074 assert!(p.contains("Product"));
2075 }
2076
2077 #[test]
2078 fn test_default_system_prompt_invoice() {
2079 let p = ConfigurableAgent::default_system_prompt("invoice");
2080 assert!(p.contains("Invoice"));
2081 }
2082
2083 #[test]
2084 fn test_default_system_prompt_sales() {
2085 let p = ConfigurableAgent::default_system_prompt("sales");
2086 assert!(p.contains("Sales"));
2087 }
2088
2089 #[test]
2090 fn test_default_system_prompt_finance() {
2091 let p = ConfigurableAgent::default_system_prompt("finance");
2092 assert!(p.contains("Finance"));
2093 }
2094
2095 #[test]
2096 fn test_default_system_prompt_hr() {
2097 let p = ConfigurableAgent::default_system_prompt("hr");
2098 assert!(p.contains("HR"));
2099 }
2100
2101 #[test]
2102 fn test_default_system_prompt_unknown_name() {
2103 let p = ConfigurableAgent::default_system_prompt("unknown_name");
2104 assert_eq!(p, "You are a unknown_name agent.");
2105 }
2106
2107 #[test]
2108 fn test_default_system_prompt_case_insensitive() {
2109 let p = ConfigurableAgent::default_system_prompt("ROUTER");
2110 assert!(p.contains("routing"), "ROUTER should match router branch");
2111 }
2112
2113 #[test]
2118 fn test_name_to_type_router() {
2119 assert!(matches!(
2120 ConfigurableAgent::name_to_type("router"),
2121 AgentType::Router
2122 ));
2123 }
2124
2125 #[test]
2126 fn test_name_to_type_invoice() {
2127 assert!(matches!(
2128 ConfigurableAgent::name_to_type("invoice"),
2129 AgentType::Invoice
2130 ));
2131 }
2132
2133 #[test]
2134 fn test_name_to_type_sales() {
2135 assert!(matches!(
2136 ConfigurableAgent::name_to_type("sales"),
2137 AgentType::Sales
2138 ));
2139 }
2140
2141 #[test]
2142 fn test_name_to_type_finance() {
2143 assert!(matches!(
2144 ConfigurableAgent::name_to_type("finance"),
2145 AgentType::Finance
2146 ));
2147 }
2148
2149 #[test]
2150 fn test_name_to_type_hr() {
2151 assert!(matches!(
2152 ConfigurableAgent::name_to_type("hr"),
2153 AgentType::HR
2154 ));
2155 }
2156
2157 #[test]
2158 fn test_name_to_type_orchestrator() {
2159 assert!(matches!(
2160 ConfigurableAgent::name_to_type("orchestrator"),
2161 AgentType::Orchestrator
2162 ));
2163 }
2164
2165 #[test]
2166 fn test_name_to_type_product_upper() {
2167 assert!(matches!(
2168 ConfigurableAgent::name_to_type("PRODUCT"),
2169 AgentType::Product
2170 ));
2171 }
2172
2173 #[test]
2174 fn test_name_to_type_unknown() {
2175 assert!(matches!(
2176 ConfigurableAgent::name_to_type("unknown"),
2177 AgentType::Custom(_)
2178 ));
2179 }
2180
2181 #[test]
2182 fn test_name_to_type_custom_preserves_name() {
2183 if let AgentType::Custom(name) = ConfigurableAgent::name_to_type("my-custom-agent") {
2184 assert_eq!(name, "my-custom-agent");
2185 } else {
2186 panic!("Expected Custom variant");
2187 }
2188 }
2189
2190 #[test]
2191 fn test_name_to_type_empty_string() {
2192 assert!(matches!(
2193 ConfigurableAgent::name_to_type(""),
2194 AgentType::Custom(ref s) if s.is_empty()
2195 ));
2196 }
2197
2198 #[test]
2203 fn test_name_accessor() {
2204 let config = make_config(vec![], None);
2205 let agent = ConfigurableAgent::new("product", &config, Box::new(MockLLM::new()), None);
2206 assert_eq!(agent.name(), "product");
2207 }
2208
2209 #[test]
2210 fn test_max_tool_iterations_accessor() {
2211 let mut config = make_config(vec![], None);
2212 config.max_tool_iterations = 42;
2213 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), None);
2214 assert_eq!(agent.max_tool_iterations(), 42);
2215 }
2216
2217 #[test]
2218 fn test_parallel_tools_accessor() {
2219 let mut config = make_config(vec![], None);
2220 config.parallel_tools = true;
2221 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), None);
2222 assert!(agent.parallel_tools());
2223 }
2224
2225 #[test]
2226 fn test_tools_returns_some_when_provided() {
2227 let config = make_config(vec![], None);
2228 let tools = Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()));
2229 let agent = ConfigurableAgent::new(
2230 "router",
2231 &config,
2232 Box::new(MockLLM::new()),
2233 Some(tools.clone()),
2234 );
2235 assert!(agent.tools().is_some());
2236 assert!(Arc::ptr_eq(agent.tools().unwrap(), &tools));
2237 }
2238
2239 #[test]
2240 fn test_tools_returns_none_when_absent() {
2241 let config = make_config(vec![], None);
2242 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), None);
2243 assert!(agent.tools().is_none());
2244 }
2245
2246 #[test]
2247 fn test_allowed_tools_from_config() {
2248 let config = make_config(vec!["calculator", "web_search"], None);
2249 let agent = ConfigurableAgent::new("orchestrator", &config, Box::new(MockLLM::new()), None);
2250 let allowed = agent.allowed_tools().expect("should have allowed tools");
2251 assert_eq!(allowed.len(), 2);
2252 assert!(allowed.contains(&"calculator".to_string()));
2253 assert!(allowed.contains(&"web_search".to_string()));
2254 }
2255
2256 #[test]
2261 fn test_can_use_tool_in_allowed_but_no_registry() {
2262 let config = make_config(vec!["calculator"], None);
2263 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), None);
2264 assert!(!agent.can_use_tool("calculator"), "no registry → false");
2265 }
2266
2267 #[test]
2268 fn test_can_use_tool_not_in_allowed_list() {
2269 let config = make_config(vec!["calculator"], None);
2270 let reg = Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()));
2271 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), Some(reg));
2272 assert!(!agent.can_use_tool("web_search"), "not in allowed → false");
2273 }
2274
2275 #[test]
2276 fn test_can_use_tool_both_empty() {
2277 let config = make_config(vec![], None);
2278 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), None);
2279 assert!(!agent.can_use_tool("anything"));
2280 }
2281
2282 #[test]
2287 fn test_get_filtered_tool_definitions_no_registry() {
2288 let config = make_config(vec!["calculator"], None);
2289 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), None);
2290 assert!(agent.get_filtered_tool_definitions().is_empty());
2291 }
2292
2293 #[test]
2294 fn test_get_filtered_tool_definitions_with_registry() {
2295 let reg = make_tools_with_tool("calculator");
2296 let config = make_config(vec!["calculator"], None);
2297 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), Some(reg));
2298 let defs = agent.get_filtered_tool_definitions();
2299 assert_eq!(defs.len(), 1);
2300 assert_eq!(defs[0].name, "calculator");
2301 }
2302
2303 #[test]
2308 fn test_with_params_sets_all_fields() {
2309 let agent = ConfigurableAgent::with_params(
2310 "my-agent",
2311 AgentType::Finance,
2312 Box::new(MockLLM::new()),
2313 "Custom prompt".to_string(),
2314 Some(Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()))),
2315 Some(vec!["tool_a".to_string()]),
2316 10,
2317 true,
2318 );
2319 assert_eq!(agent.name(), "my-agent");
2320 assert!(matches!(agent.agent_type(), AgentType::Finance));
2321 assert_eq!(agent.system_prompt(), "Custom prompt");
2322 assert!(agent.tools().is_some());
2323 let allowed = agent.allowed_tools().expect("should have allowed tools");
2324 assert_eq!(allowed.len(), 1);
2325 assert_eq!(agent.max_tool_iterations(), 10);
2326 assert!(agent.parallel_tools());
2327 }
2328
2329 #[test]
2330 fn test_with_params_provider_name_is_unknown() {
2331 let agent = ConfigurableAgent::with_params(
2332 "x",
2333 AgentType::Router,
2334 Box::new(MockLLM::new()),
2335 "p".to_string(),
2336 None,
2337 None,
2338 1,
2339 false,
2340 );
2341 assert_eq!(agent.provider_name, "unknown");
2342 }
2343
2344 #[test]
2349 fn test_new_with_provider_uses_explicit_name() {
2350 let config = make_config(vec![], None);
2351 let agent = ConfigurableAgent::new_with_provider(
2352 "sales",
2353 &config,
2354 Box::new(MockLLM::new()),
2355 None,
2356 "my-provider".to_string(),
2357 );
2358 assert_eq!(agent.name(), "sales");
2359 assert!(matches!(agent.agent_type(), AgentType::Sales));
2360 assert_eq!(agent.provider_name, "my-provider");
2361 }
2362
2363 #[test]
2364 fn test_new_with_provider_falls_back_to_default_prompt() {
2365 let config = make_config(vec![], None); let agent = ConfigurableAgent::new_with_provider(
2367 "invoice",
2368 &config,
2369 Box::new(MockLLM::new()),
2370 None,
2371 "p".to_string(),
2372 );
2373 assert!(
2374 agent.system_prompt().contains("Invoice"),
2375 "should use default_system_prompt for 'invoice'"
2376 );
2377 }
2378
2379 #[test]
2380 fn test_new_with_provider_uses_config_prompt_when_some() {
2381 let config = make_config(vec![], Some("Custom system prompt"));
2382 let agent = ConfigurableAgent::new_with_provider(
2383 "invoice",
2384 &config,
2385 Box::new(MockLLM::new()),
2386 None,
2387 "p".to_string(),
2388 );
2389 assert_eq!(agent.system_prompt(), "Custom system prompt");
2390 }
2391
2392 #[test]
2397 fn test_agent_trait_system_prompt() {
2398 let config = make_config(vec![], Some("Hello from config"));
2399 let agent = ConfigurableAgent::new("router", &config, Box::new(MockLLM::new()), None);
2400 assert_eq!(Agent::system_prompt(&agent), "Hello from config");
2401 }
2402
2403 #[test]
2404 fn test_agent_trait_agent_type() {
2405 let config = make_config(vec![], None);
2406 let agent = ConfigurableAgent::new("finance", &config, Box::new(MockLLM::new()), None);
2407 assert!(matches!(Agent::agent_type(&agent), AgentType::Finance));
2408 }
2409
2410 #[tokio::test]
2415 async fn test_execute_simple_calls_generate_with_history() {
2416 let config = make_config(vec![], Some("You are helpful"));
2417 let agent = ConfigurableAgent::new(
2418 "router",
2419 &config,
2420 Box::new(MockLLM::with_content("hello world")),
2421 None,
2422 );
2423 let ctx = make_context();
2424 let resp = Agent::execute(&agent, "hi", &ctx).await.unwrap();
2425 assert_eq!(resp.content, "hello world");
2426 }
2427
2428 #[tokio::test]
2429 async fn test_execute_simple_returns_metadata() {
2430 let config = make_config(vec![], None);
2431 let agent = ConfigurableAgent::new(
2432 "router",
2433 &config,
2434 Box::new(MockLLM::with_content("ok")),
2435 None,
2436 );
2437 let ctx = make_context();
2438 let resp = Agent::execute(&agent, "test", &ctx).await.unwrap();
2439 let meta = resp.metadata.unwrap();
2440 assert_eq!(meta.model_name, "mock");
2441 assert_eq!(meta.provider_name, "default"); }
2443
2444 #[tokio::test]
2445 async fn test_execute_simple_empty_conversation_history() {
2446 let config = make_config(vec![], Some("system"));
2447 let agent = ConfigurableAgent::new(
2448 "router",
2449 &config,
2450 Box::new(MockLLM::with_content("reply")),
2451 None,
2452 );
2453 let ctx = AgentContext {
2454 user_id: "u".to_string(),
2455 session_id: "s".to_string(),
2456 conversation_history: vec![],
2457 user_memory: None,
2458 };
2459 let resp = Agent::execute(&agent, "q", &ctx).await.unwrap();
2460 assert_eq!(resp.content, "reply");
2461 }
2462
2463 #[tokio::test]
2464 async fn test_execute_simple_with_conversation_history() {
2465 let config = make_config(vec![], Some("system"));
2466 let agent = ConfigurableAgent::new(
2467 "router",
2468 &config,
2469 Box::new(MockLLM::with_content("contextual reply")),
2470 None,
2471 );
2472 let ctx = make_context_with_history(vec![
2473 (MessageRole::User, "first question"),
2474 (MessageRole::Assistant, "first answer"),
2475 (MessageRole::User, "follow up"),
2476 ]);
2477 let resp = Agent::execute(&agent, "final", &ctx).await.unwrap();
2478 assert_eq!(resp.content, "contextual reply");
2479 }
2480
2481 #[tokio::test]
2482 async fn test_execute_simple_with_user_memory() {
2483 let config = make_config(vec![], Some("system"));
2484 let agent = ConfigurableAgent::new(
2485 "router",
2486 &config,
2487 Box::new(MockLLM::with_content("memory reply")),
2488 None,
2489 );
2490 let ctx = AgentContext {
2491 user_id: "u".to_string(),
2492 session_id: "s".to_string(),
2493 conversation_history: vec![],
2494 user_memory: Some(UserMemory {
2495 user_id: "u".to_string(),
2496 preferences: vec![Preference {
2497 category: "communication".to_string(),
2498 key: "style".to_string(),
2499 value: "concise".to_string(),
2500 confidence: 0.9,
2501 }],
2502 facts: vec![],
2503 }),
2504 };
2505 let resp = Agent::execute(&agent, "q", &ctx).await.unwrap();
2506 assert_eq!(resp.content, "memory reply");
2507 }
2508
2509 #[tokio::test]
2514 async fn test_execute_tool_path_no_tool_calls_returns_final() {
2515 let reg = make_tools_with_tool("calculator");
2516 let mut config = make_config(vec!["calculator"], Some("system"));
2517 config.max_tool_iterations = 3;
2518 let agent = ConfigurableAgent::new(
2520 "orchestrator",
2521 &config,
2522 Box::new(MockLLM::with_content("final answer")),
2523 Some(reg),
2524 );
2525 let ctx = make_context();
2526 let resp = Agent::execute(&agent, "compute 2+2", &ctx).await.unwrap();
2527 assert_eq!(resp.content, "final answer");
2528 let meta = resp.metadata.unwrap();
2529 assert_eq!(meta.model_name, "mock");
2530 }
2531
2532 #[tokio::test]
2533 async fn test_execute_tool_path_tool_calls_then_final() {
2534 let reg = make_tools_with_tool("calculator");
2535 let mut config = make_config(vec!["calculator"], Some("system"));
2536 config.max_tool_iterations = 3;
2537
2538 let tool_call = ToolCall {
2540 id: "tc_1".to_string(),
2541 name: "calculator".to_string(),
2542 arguments: serde_json::json!({"expression": "2+2"}),
2543 };
2544 let responses = vec![
2545 make_tool_response("", vec![tool_call]),
2546 make_tool_response("The answer is 4", vec![]),
2547 ];
2548
2549 let agent = ConfigurableAgent::new(
2550 "orchestrator",
2551 &config,
2552 Box::new(MockLLM::with_tool_responses(responses)),
2553 Some(reg),
2554 );
2555 let ctx = make_context();
2556 let resp = Agent::execute(&agent, "2+2?", &ctx).await.unwrap();
2557 assert_eq!(resp.content, "The answer is 4");
2558 }
2559
2560 #[tokio::test]
2561 async fn test_execute_tool_path_max_iterations_reaches_synthesis() {
2562 let reg = make_tools_with_tool("calculator");
2563 let mut config = make_config(vec!["calculator"], Some("system"));
2564 config.max_tool_iterations = 2; let tc = ToolCall {
2567 id: "tc_1".to_string(),
2568 name: "calculator".to_string(),
2569 arguments: serde_json::json!({}),
2570 };
2571 let responses = vec![
2573 make_tool_response("", vec![tc.clone()]),
2574 make_tool_response("", vec![tc]),
2575 ];
2576
2577 let agent = ConfigurableAgent::new(
2578 "orchestrator",
2579 &config,
2580 Box::new(MockLLM::with_tool_responses(responses)),
2581 Some(reg),
2582 );
2583 let ctx = make_context();
2584 let resp = Agent::execute(&agent, "compute", &ctx).await.unwrap();
2585 assert!(!resp.content.is_empty());
2588 }
2589
2590 #[tokio::test]
2591 async fn test_execute_tool_path_tool_execution_error() {
2592 struct FailingTool;
2594 #[async_trait]
2595 impl ares_tools::Tool for FailingTool {
2596 fn name(&self) -> &str {
2597 "fail_tool"
2598 }
2599 fn description(&self) -> &str {
2600 "always fails"
2601 }
2602 fn parameters_schema(&self) -> serde_json::Value {
2603 serde_json::json!({})
2604 }
2605 async fn execute(&self, _args: serde_json::Value) -> Result<serde_json::Value> {
2606 Err(ares_types::AppError::Internal("tool crashed".to_string()))
2607 }
2608 }
2609
2610 let reg = Arc::new(Tools::from_static([Arc::new(FailingTool) as Arc<dyn Tool>]));
2611
2612 let mut config = make_config(vec!["fail_tool"], Some("system"));
2613 config.max_tool_iterations = 3;
2614
2615 let tc = ToolCall {
2616 id: "tc_err".to_string(),
2617 name: "fail_tool".to_string(),
2618 arguments: serde_json::json!({}),
2619 };
2620 let responses = vec![
2621 make_tool_response("", vec![tc]),
2622 make_tool_response("Error handled", vec![]),
2623 ];
2624
2625 let agent = ConfigurableAgent::new(
2626 "orchestrator",
2627 &config,
2628 Box::new(MockLLM::with_tool_responses(responses)),
2629 Some(reg),
2630 );
2631 let ctx = make_context();
2632 let resp = Agent::execute(&agent, "do it", &ctx).await.unwrap();
2633 assert_eq!(resp.content, "Error handled");
2636 }
2637
2638 #[test]
2643 fn test_can_use_tool_empty_list_denies_all() {
2644 let reg = Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()));
2645 let agent = ConfigurableAgent::with_params(
2646 "router",
2647 AgentType::Router,
2648 Box::new(MockLLM::new()),
2649 "system".to_string(),
2650 Some(reg),
2651 Some(vec![]),
2652 1,
2653 false,
2654 );
2655 assert!(
2656 !agent.can_use_tool("anything"),
2657 "empty allowed_tools → deny all"
2658 );
2659 }
2660
2661 #[test]
2662 fn prebuilt_connector_tool_detection_covers_registered_connectors() {
2663 assert!(is_prebuilt_connector_tool("slack_send_message"));
2664 assert!(is_prebuilt_connector_tool("google_calendar_list_events"));
2665 assert!(is_prebuilt_connector_tool("salesforce_create_record"));
2666 assert!(!is_prebuilt_connector_tool("calculator"));
2667 }
2668
2669 #[tokio::test]
2670 async fn dispatch_prebuilt_connector_injects_runtime_tenant() {
2671 let reg = make_tools_with_echo_tool("slack_send_message");
2672 let mut agent = ConfigurableAgent::with_params(
2673 "orchestrator",
2674 AgentType::Orchestrator,
2675 Box::new(MockLLM::new()),
2676 "system".to_string(),
2677 Some(reg),
2678 Some(vec!["slack_send_message".to_string()]),
2679 3,
2680 false,
2681 );
2682 let ctx = crate::tenant_scope(&cordis::Context::new_root(), "tenant-a");
2683 agent.bind_request_ctx(ctx);
2684
2685 let result = agent
2686 .dispatch_tool("slack_send_message", serde_json::json!({"channel":"ops"}))
2687 .await
2688 .expect("connector dispatch");
2689
2690 assert_eq!(result["tenant_id"], "tenant-a");
2691 assert_eq!(result["channel"], "ops");
2692 }
2693
2694 #[tokio::test]
2695 async fn dispatch_prebuilt_connector_rejects_cross_tenant_arg() {
2696 let reg = make_tools_with_echo_tool("slack_send_message");
2697 let mut agent = ConfigurableAgent::with_params(
2698 "orchestrator",
2699 AgentType::Orchestrator,
2700 Box::new(MockLLM::new()),
2701 "system".to_string(),
2702 Some(reg),
2703 Some(vec!["slack_send_message".to_string()]),
2704 3,
2705 false,
2706 );
2707 let ctx = crate::tenant_scope(&cordis::Context::new_root(), "tenant-a");
2708 agent.bind_request_ctx(ctx);
2709
2710 let err = agent
2711 .dispatch_tool(
2712 "slack_send_message",
2713 serde_json::json!({"tenant_id":"tenant-b","channel":"ops"}),
2714 )
2715 .await
2716 .unwrap_err();
2717
2718 assert!(err.to_string().contains("does not match executing tenant"));
2719 }
2720
2721 #[tokio::test]
2722 async fn test_execute_tool_allowed_tool_succeeds() {
2723 let reg = make_tools_with_tool("http");
2724 let mut config = make_config(vec!["http"], Some("system"));
2725 config.max_tool_iterations = 3;
2726
2727 let tc = ToolCall {
2728 id: "tc_1".to_string(),
2729 name: "http".to_string(),
2730 arguments: serde_json::json!({}),
2731 };
2732 let responses = vec![
2733 make_tool_response("", vec![tc]),
2734 make_tool_response("HTTP result", vec![]),
2735 ];
2736
2737 let agent = ConfigurableAgent::new(
2738 "orchestrator",
2739 &config,
2740 Box::new(MockLLM::with_tool_responses(responses)),
2741 Some(reg),
2742 );
2743 let ctx = make_context();
2744 let resp = Agent::execute(&agent, "fetch", &ctx).await.unwrap();
2745 assert_eq!(resp.content, "HTTP result");
2746 }
2747
2748 #[tokio::test]
2749 async fn test_execute_tool_disallowed_tool_returns_error() {
2750 let reg = Arc::new(Tools::from_static([
2751 Arc::new(MockTool::new("http")) as Arc<dyn Tool>,
2752 Arc::new(MockTool::new("sql")) as Arc<dyn Tool>,
2753 ]));
2754
2755 let agent = ConfigurableAgent::with_params(
2756 "orchestrator",
2757 AgentType::Orchestrator,
2758 Box::new(MockLLM::with_tool_responses(vec![make_tool_response(
2759 "",
2760 vec![ToolCall {
2761 id: "tc_1".to_string(),
2762 name: "sql".to_string(),
2763 arguments: serde_json::json!({}),
2764 }],
2765 )])),
2766 "system".to_string(),
2767 Some(reg),
2768 Some(vec!["http".to_string()]),
2769 3,
2770 false,
2771 );
2772
2773 let ctx = make_context();
2774 let result = Agent::execute(&agent, "query", &ctx).await;
2775 assert!(result.is_err(), "disallowed tool should return error");
2776 let err = result.err().unwrap().to_string();
2777 assert!(
2778 err.contains("sql"),
2779 "error should mention denied tool: {}",
2780 err
2781 );
2782 assert!(
2783 err.contains("not allowed"),
2784 "error should say tool is not allowed: {}",
2785 err
2786 );
2787 }
2788
2789 #[tokio::test]
2790 async fn test_execute_tool_empty_allowed_tools_denies_all() {
2791 let reg = make_tools_with_tool("http");
2792 let agent = ConfigurableAgent::with_params(
2793 "orchestrator",
2794 AgentType::Orchestrator,
2795 Box::new(MockLLM::with_tool_responses(vec![make_tool_response(
2796 "",
2797 vec![ToolCall {
2798 id: "tc_1".to_string(),
2799 name: "http".to_string(),
2800 arguments: serde_json::json!({}),
2801 }],
2802 )])),
2803 "system".to_string(),
2804 Some(reg),
2805 Some(vec![]),
2806 3,
2807 false,
2808 );
2809
2810 let ctx = make_context();
2811 let result = Agent::execute(&agent, "fetch", &ctx).await;
2812 assert!(result.is_err(), "empty allowed_tools should deny all");
2813 let err = result.err().unwrap().to_string();
2814 assert!(
2815 err.contains("http"),
2816 "error should mention denied tool: {}",
2817 err
2818 );
2819 }
2820
2821 #[tokio::test]
2822 async fn test_execute_tool_none_allowed_tools_denies_all() {
2823 let reg = Arc::new(Tools::from_static([
2824 Arc::new(MockTool::new("http")) as Arc<dyn Tool>,
2825 Arc::new(MockTool::new("sql")) as Arc<dyn Tool>,
2826 ]));
2827
2828 let agent = ConfigurableAgent::with_params(
2829 "orchestrator",
2830 AgentType::Orchestrator,
2831 Box::new(MockLLM::with_tool_responses(vec![
2832 make_tool_response(
2833 "",
2834 vec![ToolCall {
2835 id: "tc_1".to_string(),
2836 name: "sql".to_string(),
2837 arguments: serde_json::json!({}),
2838 }],
2839 ),
2840 make_tool_response("SQL result", vec![]),
2841 ])),
2842 "system".to_string(),
2843 Some(reg),
2844 None,
2845 3,
2846 false,
2847 );
2848
2849 let ctx = make_context();
2850 let result = Agent::execute(&agent, "query", &ctx).await;
2851 assert!(result.is_err(), "missing allowed_tools should deny all");
2852 let err = result.err().unwrap().to_string();
2853 assert!(
2854 err.contains("sql"),
2855 "error should mention denied tool: {err}"
2856 );
2857 }
2858
2859 #[test]
2860 fn test_observed_tool_type_marks_builtins() {
2861 let agent = ConfigurableAgent::with_params(
2862 "router",
2863 AgentType::Router,
2864 Box::new(MockLLM::new()),
2865 "system".to_string(),
2866 None,
2867 Some(vec!["http".to_string()]),
2868 1,
2869 false,
2870 );
2871
2872 assert_eq!(agent.observed_tool_type("http", true), "builtin");
2873 }
2874
2875 #[test]
2876 fn test_observed_tool_type_falls_back_for_runtime_tools() {
2877 let agent = ConfigurableAgent::with_params(
2878 "router",
2879 AgentType::Router,
2880 Box::new(MockLLM::new()),
2881 "system".to_string(),
2882 None,
2883 Some(vec!["tenant_http".to_string()]),
2884 1,
2885 false,
2886 );
2887
2888 assert_eq!(agent.observed_tool_type("tenant_http", false), "runtime");
2889 }
2890
2891 #[test]
2892 fn test_set_allowed_tools_intersection() {
2893 let reg = Arc::new(Tools::from_static(Vec::<Arc<dyn Tool>>::new()));
2894 let mut agent = ConfigurableAgent::with_params(
2895 "router",
2896 AgentType::Router,
2897 Box::new(MockLLM::new()),
2898 "system".to_string(),
2899 Some(reg),
2900 Some(vec!["http".to_string(), "sql".to_string()]),
2901 1,
2902 false,
2903 );
2904 assert!(agent.can_use_tool("http"));
2905 assert!(agent.can_use_tool("sql"));
2906 agent.set_allowed_tools(Some(vec!["http".to_string()]));
2907 assert!(agent.can_use_tool("http"));
2908 assert!(!agent.can_use_tool("sql"));
2909 agent.set_allowed_tools(Some(vec![]));
2910 assert!(!agent.can_use_tool("http"));
2911 agent.set_allowed_tools(None);
2912 assert!(!agent.can_use_tool("sql"));
2913 }
2914
2915 struct FailingMockLLM {
2918 error_msg: String,
2919 }
2920
2921 #[async_trait]
2922 impl LLMClient for FailingMockLLM {
2923 async fn generate(&self, _: &str) -> Result<String> {
2924 Err(ares_types::types::AppError::LLM(self.error_msg.clone()))
2925 }
2926 async fn generate_with_system(&self, _: &str, _: &str) -> Result<String> {
2927 Err(ares_types::types::AppError::LLM(self.error_msg.clone()))
2928 }
2929 async fn generate_with_history(&self, _: &[(String, String)]) -> Result<LLMResponse> {
2930 Err(ares_types::types::AppError::LLM(self.error_msg.clone()))
2931 }
2932 async fn generate_with_tools(&self, _: &str, _: &[ToolDefinition]) -> Result<LLMResponse> {
2933 Err(ares_types::types::AppError::LLM(self.error_msg.clone()))
2934 }
2935 async fn generate_with_tools_and_history(
2936 &self,
2937 _: &[ares_llm::coordinator::ConversationMessage],
2938 _: &[ToolDefinition],
2939 ) -> Result<LLMResponse> {
2940 Err(ares_types::types::AppError::LLM(self.error_msg.clone()))
2941 }
2942 async fn stream(
2943 &self,
2944 _: &str,
2945 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
2946 Ok(Box::new(futures::stream::empty()))
2947 }
2948 async fn stream_with_system(
2949 &self,
2950 _: &str,
2951 _: &str,
2952 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
2953 Ok(Box::new(futures::stream::empty()))
2954 }
2955 async fn stream_with_history(
2956 &self,
2957 _: &[(String, String)],
2958 ) -> Result<Box<dyn futures::Stream<Item = Result<String>> + Send + Unpin>> {
2959 Ok(Box::new(futures::stream::empty()))
2960 }
2961 fn model_name(&self) -> &str {
2962 "failing-mock"
2963 }
2964 }
2965
2966 #[tokio::test]
2967 async fn test_fallback_used_when_primary_fails() {
2968 let mut agent = ConfigurableAgent::with_params(
2969 "test",
2970 AgentType::Product,
2971 Box::new(FailingMockLLM {
2972 error_msg: "primary failed".to_string(),
2973 }),
2974 "system".to_string(),
2975 None,
2976 None,
2977 1,
2978 false,
2979 );
2980
2981 let fallback = MockLLM::with_content("fallback-success");
2982 agent.set_fallback_llms_with_providers(vec![(
2983 "fallback-provider".to_string(),
2984 Box::new(fallback),
2985 )]);
2986
2987 let ctx = make_context();
2988 let resp = Agent::execute(&agent, "hello", &ctx).await.unwrap();
2989 assert_eq!(resp.content, "fallback-success");
2990 let metadata = resp.metadata.expect("metadata");
2991 assert_eq!(metadata.provider_name, "fallback-provider");
2992 assert_eq!(metadata.model_name, "mock");
2993 }
2994
2995 #[tokio::test]
2996 async fn test_primary_succeeds_without_fallback() {
2997 let mut agent = ConfigurableAgent::with_params(
2998 "test",
2999 AgentType::Product,
3000 Box::new(MockLLM::with_content("primary-success")),
3001 "system".to_string(),
3002 None,
3003 None,
3004 1,
3005 false,
3006 );
3007
3008 let fallback = FailingMockLLM {
3009 error_msg: "fallback should not run".to_string(),
3010 };
3011 agent.set_fallback_llms(vec![Box::new(fallback)]);
3012
3013 let ctx = make_context();
3014 let resp = Agent::execute(&agent, "hello", &ctx).await.unwrap();
3015 assert_eq!(resp.content, "primary-success");
3016 }
3017
3018 #[tokio::test]
3019 async fn test_all_fallbacks_fail_reports_every_error() {
3020 let mut agent = ConfigurableAgent::with_params(
3021 "test",
3022 AgentType::Product,
3023 Box::new(FailingMockLLM {
3024 error_msg: "primary failed".to_string(),
3025 }),
3026 "system".to_string(),
3027 None,
3028 None,
3029 1,
3030 false,
3031 );
3032
3033 agent.set_fallback_llms(vec![
3034 Box::new(FailingMockLLM {
3035 error_msg: "fallback-0 failed".to_string(),
3036 }),
3037 Box::new(FailingMockLLM {
3038 error_msg: "fallback-1 failed".to_string(),
3039 }),
3040 ]);
3041
3042 let ctx = make_context();
3043 let err = match Agent::execute(&agent, "hello", &ctx).await {
3044 Ok(_) => panic!("expected all LLMs to fail"),
3045 Err(err) => err.to_string(),
3046 };
3047 assert!(err.contains("primary failed"), "got: {err}");
3048 assert!(err.contains("fallback[0]"), "got: {err}");
3049 assert!(err.contains("fallback-0 failed"), "got: {err}");
3050 assert!(err.contains("fallback[1]"), "got: {err}");
3051 assert!(err.contains("fallback-1 failed"), "got: {err}");
3052 }
3053
3054 #[test]
3055 fn user_id_from_ctx_reads_intercept() {
3056 use ares_types::models::{TenantContext, TenantTier};
3057
3058 let root: Arc<cordis::Context> = cordis::Context::new_root();
3059 assert_eq!(crate::user_id_from_ctx(&root, ""), "");
3060
3061 let ctx = root.with_intercept(TenantContext::new("acme".into(), TenantTier::Pro));
3062 assert_eq!(crate::user_id_from_ctx(&ctx, "anon"), "acme");
3063 }
3064
3065 #[test]
3066 fn user_id_from_ctx_isolate_wins_over_intercept() {
3067 use ares_types::models::{TenantContext, TenantTier};
3068
3069 let root: Arc<cordis::Context> = cordis::Context::new_root();
3070 let intercepted =
3071 root.with_intercept(TenantContext::new("from-intercept".into(), TenantTier::Pro));
3072 let isolated = crate::tenant_scope(&intercepted, "from-isolate");
3073 assert_eq!(crate::user_id_from_ctx(&isolated, "anon"), "from-isolate");
3074 }
3075
3076 #[tokio::test]
3077 async fn configurable_generate_waterfall_rewrites_last_message() {
3078 let ctx = Context::new_root();
3079 let events = ctx.provide(EventsService::new());
3080 events.on_waterfall(
3081 cordis::events_catalog::ev::LLM_GENERATE.to_string(),
3082 |mut payload, next| async move {
3083 if let Some(arr) = payload.get_mut("messages").and_then(|v| v.as_array_mut()) {
3084 if let Some(last) = arr.last_mut() {
3085 last["content"] = serde_json::json!("rewritten-hello");
3086 }
3087 }
3088 next(payload).await
3089 },
3090 );
3091
3092 let mut agent = ConfigurableAgent::new(
3093 "router",
3094 &make_config(vec![], Some("system")),
3095 Box::new(MockLLM::echo_last()),
3096 None,
3097 );
3098 agent.bind_request_ctx(ctx);
3099
3100 let resp = Agent::execute(&agent, "original", &make_context())
3101 .await
3102 .expect("execute");
3103 assert_eq!(resp.content, "rewritten-hello");
3104 }
3105
3106 #[tokio::test]
3107 async fn configurable_generate_short_circuit_skips_llm() {
3108 let ctx = Context::new_root();
3109 let events = ctx.provide(EventsService::new());
3110 events.on_waterfall(
3111 cordis::events_catalog::ev::LLM_GENERATE.to_string(),
3112 |_payload, _next| async move { Ok(serde_json::json!({ "content": "cached" })) },
3113 );
3114
3115 let (llm, generated) = MockLLM::with_generated_flag();
3116 let mut agent = ConfigurableAgent::new(
3117 "router",
3118 &make_config(vec![], Some("system")),
3119 Box::new(llm),
3120 None,
3121 );
3122 agent.bind_request_ctx(ctx);
3123
3124 let resp = Agent::execute(&agent, "would-call-llm", &make_context())
3125 .await
3126 .expect("execute");
3127 assert_eq!(resp.content, "cached");
3128 assert!(
3129 !generated.load(Ordering::SeqCst),
3130 "dummy generate must stay false when handler skips next"
3131 );
3132 }
3133
3134 #[tokio::test]
3135 async fn configurable_generate_tools_short_circuit_skips_llm() {
3136 let ctx = Context::new_root();
3139 let events = ctx.provide(EventsService::new());
3140 events.on_waterfall(
3141 cordis::events_catalog::ev::LLM_GENERATE_TOOLS.to_string(),
3142 |_payload, _next| async move {
3143 Ok(serde_json::json!({
3144 "content": "tools-cached",
3145 "provider": "cache",
3146 "model": "cached-model",
3147 "tool_calls": [],
3148 }))
3149 },
3150 );
3151
3152 let (llm, generated) = MockLLM::with_generated_flag();
3153 ctx.provide(ares_tools::Tools::from_static(Vec::<
3156 Arc<dyn ares_tools::Tool>,
3157 >::new()));
3158 let mut agent = ConfigurableAgent::new(
3159 "router",
3160 &make_config(vec![], Some("system")),
3161 Box::new(llm),
3162 None,
3163 );
3164 agent.bind_request_ctx(ctx);
3165 let resp = Agent::execute(&agent, "would-call-llm", &make_context())
3166 .await
3167 .expect("execute");
3168 assert_eq!(resp.content, "tools-cached");
3169 assert!(
3170 !generated.load(Ordering::SeqCst),
3171 "mock client must never be called when generate_tools handler skips next"
3172 );
3173 }
3174}