Skip to main content

wabot_testing/
agent.rs

1//! Drive a real agent in a test. Port of
2//! `wabot-ts/src/testing/agentHarness.ts`.
3
4use std::sync::Arc;
5
6use wabot_core::injection::Container;
7use wabot_feature_agent::{register_agents, Agent, AgentBuilder, AgentFactory, AgentSession};
8use wabot_feature_chat_bot::ChatAdapter;
9
10use crate::mock_adapter::MockChatAdapter;
11
12/// Runs a real agent — real [`AgentFactory`], real session loop, real
13/// tool gating and answer validation — against a scriptable adapter.
14///
15/// It is a **thin wrapper over the production path**: [`for_agent`]
16/// hands back the same [`AgentBuilder`] production uses, so there is no
17/// parallel session logic that could drift from what ships. That is the
18/// property worth protecting; everything else here is convenience.
19///
20/// ```ignore
21/// let harness = AgentHarness::new(triage_agent);
22/// harness.adapter().call_tool(ANSWER_TOOL_NAME, json!({ "urgency": "high" }));
23///
24/// let triage: Triage = harness
25///     .for_agent()
26///     .for_mindset()               // the delegation path, gating and all
27///     .allow_tools(["read_order"])
28///     .session()
29///     .await
30///     .ask("How urgent is this?")
31///     .await?;
32/// ```
33///
34/// [`for_agent`]: AgentHarness::for_agent
35pub struct AgentHarness {
36    agent: Arc<dyn Agent>,
37    adapter: Arc<MockChatAdapter>,
38    factory: Arc<AgentFactory>,
39    container: Container,
40}
41
42impl AgentHarness {
43    pub fn new(agent: Arc<dyn Agent>) -> Self {
44        Self::builder(agent).build()
45    }
46
47    pub fn builder(agent: Arc<dyn Agent>) -> AgentHarnessBuilder {
48        AgentHarnessBuilder {
49            agent,
50            container: None,
51            adapter: None,
52        }
53    }
54
55    pub fn adapter(&self) -> &Arc<MockChatAdapter> {
56        &self.adapter
57    }
58
59    /// The container the agent's tools resolve from — also where the
60    /// agent-tools provider is registered, so a mindset built from it
61    /// can delegate.
62    pub fn container(&self) -> &Container {
63        &self.container
64    }
65
66    pub fn factory(&self) -> &Arc<AgentFactory> {
67        &self.factory
68    }
69
70    /// The production builder for this agent: chain `for_mindset()`,
71    /// `allow_tools()`, `deny_tools()`, `with_budget()`,
72    /// `with_context()`, then `session()`.
73    pub fn for_agent(&self) -> AgentBuilder {
74        self.factory.for_agent(self.agent.clone())
75    }
76
77    /// Shortcut for a session with default gating and budget.
78    pub async fn session(&self) -> AgentSession {
79        self.for_agent().session().await
80    }
81}
82
83pub struct AgentHarnessBuilder {
84    agent: Arc<dyn Agent>,
85    container: Option<Container>,
86    adapter: Option<Arc<MockChatAdapter>>,
87}
88
89impl AgentHarnessBuilder {
90    /// The container the agent's tools resolve from. Register their
91    /// dependencies there before building.
92    pub fn container(mut self, container: Container) -> Self {
93        self.container = Some(container);
94        self
95    }
96
97    pub fn adapter(mut self, adapter: Arc<MockChatAdapter>) -> Self {
98        self.adapter = Some(adapter);
99        self
100    }
101
102    pub fn build(self) -> AgentHarness {
103        let container = self.container.unwrap_or_default();
104        let adapter = self.adapter.unwrap_or_else(MockChatAdapter::arc);
105        // The real installer, so a harness-built container can also
106        // back a delegating mindset — the same call an app makes.
107        let factory = register_agents(&container, adapter.clone() as Arc<dyn ChatAdapter>);
108
109        AgentHarness {
110            agent: self.agent,
111            adapter,
112            factory,
113            container,
114        }
115    }
116}