1use std::collections::HashSet;
4use std::path::PathBuf;
5use std::sync::Arc;
6
7use talos_core::provider::ToolDefinition;
8use talos_core::tool::{AgentTool, ToolPresentationPolicy, ToolProtocol, ToolRegistry};
9use talos_permission::PermissionEngine;
10use talos_plugin::HookRegistry;
11use talos_sandbox::SandboxProvider;
12use talos_skill::SkillIndex;
13use tokio_util::sync::CancellationToken;
14
15use crate::prompt::{ActivatedSkillContext, ContextFile, SystemPromptBuilder, ToolDescription};
16use crate::{
17 Agent, MemoryProviderCallback, RequestBudgetSpec, SandboxFallbackHandler,
18 SandboxFallbackPolicy, TodoSectionProviderCallback, prompt,
19};
20
21impl Agent {
22 #[deprecated(
36 note = "Agent::new() has NO permission engine and NO sandbox; use Agent::with_security(). See docs/decisions/007-process-hardening-unsafe.md and ARCH review."
37 )]
38 #[must_use]
39 pub fn new(
40 provider: Arc<dyn talos_core::provider::LanguageModel>,
41 tools: ToolRegistry,
42 ) -> Self {
43 Self {
44 provider,
45 tools,
46 permission_engine: None,
47 sandbox: None,
48 sandbox_fallback_policy: SandboxFallbackPolicy::Deny,
49 sandbox_fallback_handler: None,
50 workspace_root: PathBuf::from("."),
51 prompt_builder: SystemPromptBuilder::new().with_workspace_info("Workspace root: ."),
52 hook_registry: Arc::new(HookRegistry::new()),
53 workspace_context: None,
54 tool_definitions: Vec::new(),
55 presented_tool_names: HashSet::new(),
56 enforce_tool_presentation_policy: false,
57 tool_presentation_policy: ToolPresentationPolicy::full(),
58 cached_stable_prefix: std::sync::Mutex::new(None),
59 memory_provider: None,
60 todo_section_provider: None,
61 provider_key: None,
62 model_id: None,
63 replay_reasoning: true,
64 bash_compression_enabled: false,
65 tool_output_threshold: 4000,
66 image_input_supported: false,
67 request_budget_spec: RequestBudgetSpec::default(),
68 }
69 }
70
71 #[must_use]
84 pub fn with_security(
85 provider: Arc<dyn talos_core::provider::LanguageModel>,
86 tools: ToolRegistry,
87 permission_engine: Option<Arc<PermissionEngine>>,
88 sandbox: Option<Box<dyn SandboxProvider>>,
89 workspace_root: PathBuf,
90 ) -> Self {
91 Self::with_security_and_hooks(
92 provider,
93 tools,
94 permission_engine,
95 sandbox,
96 workspace_root,
97 Arc::new(HookRegistry::new()),
98 )
99 }
100
101 #[must_use]
103 pub fn with_security_and_sandbox_fallback(
104 provider: Arc<dyn talos_core::provider::LanguageModel>,
105 tools: ToolRegistry,
106 permission_engine: Option<Arc<PermissionEngine>>,
107 sandbox: Option<Box<dyn SandboxProvider>>,
108 workspace_root: PathBuf,
109 sandbox_fallback_policy: SandboxFallbackPolicy,
110 sandbox_fallback_handler: Option<Arc<dyn SandboxFallbackHandler>>,
111 ) -> Self {
112 Self::with_security_and_hooks_and_sandbox_fallback(
113 provider,
114 tools,
115 permission_engine,
116 sandbox,
117 workspace_root,
118 Arc::new(HookRegistry::new()),
119 sandbox_fallback_policy,
120 sandbox_fallback_handler,
121 )
122 }
123
124 #[must_use]
126 pub fn with_security_and_hooks(
127 provider: Arc<dyn talos_core::provider::LanguageModel>,
128 tools: ToolRegistry,
129 permission_engine: Option<Arc<PermissionEngine>>,
130 sandbox: Option<Box<dyn SandboxProvider>>,
131 workspace_root: PathBuf,
132 hook_registry: Arc<HookRegistry>,
133 ) -> Self {
134 Self::with_security_and_hooks_and_sandbox_fallback(
135 provider,
136 tools,
137 permission_engine,
138 sandbox,
139 workspace_root,
140 hook_registry,
141 SandboxFallbackPolicy::Deny,
142 None,
143 )
144 }
145
146 #[must_use]
148 #[allow(clippy::too_many_arguments)]
149 pub fn with_security_and_hooks_and_sandbox_fallback(
150 provider: Arc<dyn talos_core::provider::LanguageModel>,
151 tools: ToolRegistry,
152 permission_engine: Option<Arc<PermissionEngine>>,
153 sandbox: Option<Box<dyn SandboxProvider>>,
154 workspace_root: PathBuf,
155 hook_registry: Arc<HookRegistry>,
156 sandbox_fallback_policy: SandboxFallbackPolicy,
157 sandbox_fallback_handler: Option<Arc<dyn SandboxFallbackHandler>>,
158 ) -> Self {
159 let tool_presentation_policy = ToolPresentationPolicy::runtime_default();
160 let (descriptions, tool_definitions, presented_tool_names) =
161 describe_presented_tools(&tools, &tool_presentation_policy);
162
163 let descriptions: Vec<_> = descriptions
167 .into_iter()
168 .filter(|d| d.name != "read_image")
169 .collect();
170 let tool_definitions: Vec<_> = tool_definitions
171 .into_iter()
172 .filter(|td| td.name != "read_image")
173 .collect();
174 let presented_tool_names: HashSet<_> = presented_tool_names
175 .into_iter()
176 .filter(|n| n != "read_image")
177 .collect();
178
179 let prompt_builder = SystemPromptBuilder::new()
180 .with_workspace_info(format!("Workspace root: {}", workspace_root.display()))
181 .with_tools(descriptions.clone());
182
183 Self {
184 provider,
185 tools,
186 permission_engine,
187 sandbox: sandbox.map(Arc::from),
188 sandbox_fallback_policy,
189 sandbox_fallback_handler,
190 workspace_root,
191 prompt_builder,
192 hook_registry,
193 workspace_context: None,
194 tool_definitions,
195 presented_tool_names,
196 enforce_tool_presentation_policy: true,
197 tool_presentation_policy,
198 cached_stable_prefix: std::sync::Mutex::new(None),
199 memory_provider: None,
200 todo_section_provider: None,
201 provider_key: None,
202 model_id: None,
203 replay_reasoning: true,
204 bash_compression_enabled: false,
205 tool_output_threshold: 4000,
206 image_input_supported: false,
207 request_budget_spec: RequestBudgetSpec::default(),
208 }
209 }
210
211 #[must_use]
213 pub fn with_reasoning_identity(
214 mut self,
215 provider_key: Option<String>,
216 model_id: Option<String>,
217 replay: bool,
218 ) -> Self {
219 self.provider_key = provider_key;
220 self.model_id = model_id;
221 self.replay_reasoning = replay;
222 self
223 }
224
225 pub fn set_request_budget_spec(&mut self, spec: RequestBudgetSpec) {
227 self.request_budget_spec = spec;
228 }
229
230 #[must_use]
231 pub fn request_budget_spec(&self) -> RequestBudgetSpec {
232 self.request_budget_spec
233 }
234
235 pub fn set_memory_provider(&mut self, provider: Arc<MemoryProviderCallback>) {
240 self.memory_provider = Some(provider);
241 }
242
243 pub fn set_todo_section_provider(&mut self, provider: Arc<TodoSectionProviderCallback>) {
248 self.todo_section_provider = Some(provider);
249 }
250
251 #[must_use]
259 pub fn with_bash_compression(mut self, enabled: bool) -> Self {
260 self.bash_compression_enabled = enabled;
261 self
262 }
263
264 #[must_use]
271 pub fn with_image_input_supported(mut self, supported: bool) -> Self {
272 self.image_input_supported = supported;
273 self
274 }
275
276 pub fn set_image_input_supported(&mut self, supported: bool) {
280 self.image_input_supported = supported;
281 let (descs, defs, names) =
282 describe_presented_tools(&self.tools, &self.tool_presentation_policy);
283 let descs: Vec<_> = descs
284 .into_iter()
285 .filter(|d| supported || d.name != "read_image")
286 .collect();
287 self.tool_definitions = defs
288 .into_iter()
289 .filter(|td| supported || td.name != "read_image")
290 .collect();
291 self.presented_tool_names = names
292 .into_iter()
293 .filter(|n| supported || n != "read_image")
294 .collect();
295 self.enforce_tool_presentation_policy = true;
296 self.update_prompt_builder(true, |builder| builder.with_tools(descs));
297 }
298
299 pub fn set_tools(&mut self, tools: Vec<ToolDescription>) {
304 self.tool_definitions = tools
305 .iter()
306 .map(|tool| ToolDefinition {
307 name: tool.name.clone(),
308 description: tool.description.clone(),
309 parameters: tool.parameters.clone(),
310 })
311 .collect();
312 self.presented_tool_names = tools.iter().map(|tool| tool.name.clone()).collect();
313 self.enforce_tool_presentation_policy = true;
314 self.update_prompt_builder(true, |builder| builder.with_tools(tools));
315 }
316
317 pub fn set_tool_presentation_policy(&mut self, policy: ToolPresentationPolicy) {
323 self.tool_presentation_policy = policy;
324 let (descriptions, tool_definitions, presented_tool_names) =
325 describe_presented_tools(&self.tools, &self.tool_presentation_policy);
326 self.tool_definitions = tool_definitions;
327 self.presented_tool_names = presented_tool_names;
328 self.enforce_tool_presentation_policy = true;
329 self.update_prompt_builder(true, |builder| builder.with_tools(descriptions));
330 }
331
332 pub fn set_tool_protocol(&mut self, protocol: ToolProtocol) {
334 self.update_prompt_builder(true, |builder| match protocol {
335 ToolProtocol::TalosStrict => builder.with_strict_tool_format(),
336 ToolProtocol::Compat => builder.with_tool_format(prompt::TOOL_CALLING_FORMAT),
337 ToolProtocol::Native => builder.with_tool_format(""),
338 });
339 }
340
341 pub fn set_skill_index(&mut self, skills: Vec<SkillIndex>) {
345 self.update_prompt_builder(true, |builder| builder.with_skill_index(skills));
346 }
347
348 pub fn set_activated_skill_context(&mut self, context: Option<ActivatedSkillContext>) {
353 self.update_prompt_builder(true, |builder| builder.with_activated_skill(context));
354 }
355
356 pub fn set_context_files(&mut self, files: Vec<ContextFile>) {
360 self.update_prompt_builder(false, |builder| builder.with_context_files(files));
361 }
362
363 pub fn set_user_preferences(&mut self, prefs: String) {
365 self.update_prompt_builder(false, |builder| builder.with_user_preferences(prefs));
366 }
367
368 pub fn set_custom_prompt(&mut self, prompt: String) {
370 self.update_prompt_builder(true, |builder| builder.with_custom_prompt(prompt));
371 }
372
373 pub fn set_append_prompt(&mut self, prompt: String) {
375 self.update_prompt_builder(false, |builder| builder.with_append_prompt(prompt));
376 }
377
378 pub fn clear_append_prompt(&mut self) {
380 self.prompt_builder.clear_append_prompt();
381 }
382
383 pub fn set_append_prompt_opt(&mut self, prompt: Option<String>) {
387 self.prompt_builder.set_append_prompt_opt(prompt);
388 }
389
390 #[must_use]
396 pub fn build_system_prompt(&self) -> String {
397 self.prompt_builder.build()
398 }
399
400 #[must_use]
406 pub fn cancellation_token(&self) -> CancellationToken {
407 CancellationToken::new()
408 }
409
410 fn update_prompt_builder(
411 &mut self,
412 invalidate_stable_prefix: bool,
413 update: impl FnOnce(SystemPromptBuilder) -> SystemPromptBuilder,
414 ) {
415 self.prompt_builder = update(std::mem::take(&mut self.prompt_builder));
416 if invalidate_stable_prefix {
417 self.invalidate_stable_prefix_cache();
418 }
419 }
420
421 pub(crate) fn invalidate_stable_prefix_cache(&self) {
422 *self
423 .cached_stable_prefix
424 .lock()
425 .expect("cache lock poisoned") = None;
426 }
427}
428
429pub(crate) fn describe_presented_tools(
430 tools: &ToolRegistry,
431 policy: &ToolPresentationPolicy,
432) -> (Vec<ToolDescription>, Vec<ToolDefinition>, HashSet<String>) {
433 let mut selected: Vec<&dyn AgentTool> = tools
434 .list()
435 .into_iter()
436 .filter(|tool| policy.allows_tool(*tool))
437 .collect();
438 selected.sort_by(|a, b| a.name().cmp(b.name()));
439
440 let descriptions: Vec<ToolDescription> = selected
441 .iter()
442 .map(|tool| {
443 let backends = policy.backend_set_for(tool.name());
444 ToolDescription {
445 name: tool.name().to_string(),
446 description: tool.description_for_backends(&backends),
447 parameters: tool.parameters_for_backends(&backends),
448 family: tool.family(),
449 }
450 })
451 .collect();
452
453 let tool_definitions = descriptions
454 .iter()
455 .map(|tool| ToolDefinition {
456 name: tool.name.clone(),
457 description: tool.description.clone(),
458 parameters: tool.parameters.clone(),
459 })
460 .collect();
461 let presented_tool_names = descriptions.iter().map(|tool| tool.name.clone()).collect();
462
463 (descriptions, tool_definitions, presented_tool_names)
464}