Skip to main content

claude_agent/agent/
executor.rs

1//! Agent core structure and construction.
2
3use std::sync::Arc;
4
5use tokio::sync::RwLock;
6
7use super::config::AgentConfig;
8use crate::Client;
9use crate::budget::{BudgetTracker, TenantBudget};
10use crate::context::PromptOrchestrator;
11use crate::hooks::HookManager;
12use crate::session::ToolState;
13use crate::tools::{ToolRegistry, ToolSearchManager};
14use crate::types::Message;
15
16pub struct Agent {
17    pub(crate) client: Arc<Client>,
18    pub(crate) config: Arc<AgentConfig>,
19    pub(crate) tools: Arc<ToolRegistry>,
20    pub(crate) hooks: Arc<HookManager>,
21    pub(crate) session_id: Arc<str>,
22    pub(crate) state: ToolState,
23    pub(crate) orchestrator: Option<Arc<RwLock<PromptOrchestrator>>>,
24    pub(crate) initial_messages: Option<Vec<Message>>,
25    pub(crate) budget_tracker: Arc<BudgetTracker>,
26    pub(crate) tenant_budget: Option<Arc<TenantBudget>>,
27    pub(crate) mcp_manager: Option<Arc<crate::mcp::McpManager>>,
28    pub(crate) tool_search_manager: Option<Arc<ToolSearchManager>>,
29}
30
31impl Agent {
32    pub fn new(client: Client, mut config: AgentConfig) -> Self {
33        let model_config = &client.config().models;
34        let resolved_primary = model_config.resolve_alias(&config.model.primary);
35        if resolved_primary != config.model.primary {
36            config.model.primary = resolved_primary.to_string();
37        }
38        let resolved_small = model_config.resolve_alias(&config.model.small);
39        if resolved_small != config.model.small {
40            config.model.small = resolved_small.to_string();
41        }
42
43        let tools = ToolRegistry::default_tools(
44            config.security.tool_access.clone(),
45            config.working_dir.clone(),
46            Some(config.security.permission_policy.clone()),
47        );
48        Self::from_parts(
49            Arc::new(client),
50            Arc::new(config),
51            Arc::new(tools),
52            Arc::new(HookManager::new()),
53            None,
54        )
55    }
56
57    pub(crate) fn from_orchestrator(
58        client: Client,
59        config: AgentConfig,
60        tools: Arc<ToolRegistry>,
61        hooks: HookManager,
62        orchestrator: PromptOrchestrator,
63    ) -> Self {
64        Self::from_parts(
65            Arc::new(client),
66            Arc::new(config),
67            tools,
68            Arc::new(hooks),
69            Some(Arc::new(RwLock::new(orchestrator))),
70        )
71    }
72
73    pub(crate) fn from_parts(
74        client: Arc<Client>,
75        config: Arc<AgentConfig>,
76        tools: Arc<ToolRegistry>,
77        hooks: Arc<HookManager>,
78        orchestrator: Option<Arc<RwLock<PromptOrchestrator>>>,
79    ) -> Self {
80        let budget_tracker = match config.budget.max_cost_usd {
81            Some(max) => BudgetTracker::new(max),
82            None => BudgetTracker::unlimited(),
83        };
84
85        let state = tools
86            .tool_state()
87            .cloned()
88            .unwrap_or_else(|| ToolState::new(crate::session::SessionId::new()));
89        let session_id: Arc<str> = state.session_id().to_string().into();
90
91        Self {
92            client,
93            config,
94            tools,
95            hooks,
96            session_id,
97            state,
98            orchestrator,
99            initial_messages: None,
100            budget_tracker: Arc::new(budget_tracker),
101            tenant_budget: None,
102            mcp_manager: None,
103            tool_search_manager: None,
104        }
105    }
106
107    pub(crate) fn tenant_budget(mut self, budget: Arc<TenantBudget>) -> Self {
108        self.tenant_budget = Some(budget);
109        self
110    }
111
112    pub(crate) fn mcp_manager(mut self, manager: Arc<crate::mcp::McpManager>) -> Self {
113        self.mcp_manager = Some(manager);
114        self
115    }
116
117    pub(crate) fn tool_search_manager(mut self, manager: Arc<ToolSearchManager>) -> Self {
118        self.tool_search_manager = Some(manager);
119        self
120    }
121
122    pub(crate) fn initial_messages(mut self, messages: Vec<Message>) -> Self {
123        self.initial_messages = Some(messages);
124        self
125    }
126
127    pub(crate) fn session_id(mut self, id: impl Into<String>) -> Self {
128        self.session_id = id.into().into();
129        self
130    }
131
132    #[must_use]
133    pub fn builder() -> super::AgentBuilder {
134        super::AgentBuilder::new()
135    }
136
137    /// Shortcut for `Agent::builder().model(model)`.
138    pub fn model(model: impl Into<String>) -> super::AgentBuilder {
139        super::AgentBuilder::new().model(model)
140    }
141
142    pub async fn default_agent() -> crate::Result<Self> {
143        Self::builder().build().await
144    }
145
146    #[must_use]
147    pub fn hooks(&self) -> &Arc<HookManager> {
148        &self.hooks
149    }
150
151    #[must_use]
152    pub fn get_session_id(&self) -> &str {
153        &self.session_id
154    }
155
156    #[must_use]
157    pub fn client(&self) -> &Arc<Client> {
158        &self.client
159    }
160
161    pub fn orchestrator(&self) -> Option<&Arc<RwLock<PromptOrchestrator>>> {
162        self.orchestrator.as_ref()
163    }
164
165    #[must_use]
166    pub fn config(&self) -> &AgentConfig {
167        &self.config
168    }
169
170    #[must_use]
171    pub fn tools(&self) -> &Arc<ToolRegistry> {
172        &self.tools
173    }
174
175    #[must_use]
176    pub fn state(&self) -> &ToolState {
177        &self.state
178    }
179}