Skip to main content

vv_agent/tools/
background_agent_task.rs

1use std::collections::BTreeMap;
2use std::panic::{catch_unwind, AssertUnwindSafe};
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::{Arc, Mutex};
5
6use serde_json::{json, Value};
7
8use crate::agent::Agent;
9use crate::result::RunResult;
10use crate::run_config::RunConfig;
11use crate::runner::{NormalizedInput, Runner};
12use crate::runtime::CancellationToken;
13use crate::tools::{Tool, ToolContext, ToolOutput, ToolSpec, ToolSpecKind};
14use crate::types::{AgentStatus, ToolArguments};
15
16static NEXT_BACKGROUND_AGENT_TASK_ID: AtomicU64 = AtomicU64::new(1);
17const BACKGROUND_AGENT_TASK_WORKER_PANIC: &str = "background agent task worker panicked";
18
19#[derive(Clone)]
20pub struct BackgroundAgentTask {
21    agent: Agent,
22    name: String,
23    description: String,
24    parameters_schema: Value,
25    handles: Arc<Mutex<BTreeMap<String, BackgroundAgentTaskHandle>>>,
26}
27
28impl BackgroundAgentTask {
29    pub fn start(
30        &self,
31        runner: &Runner,
32        context: &ToolContext,
33        raw_arguments: Value,
34        run_config: Option<RunConfig>,
35    ) -> Result<BackgroundAgentTaskHandle, String> {
36        let input = self.input_from_arguments(raw_arguments)?;
37        let run_config = inherited_run_config(context, run_config);
38        let task_id = format!(
39            "bg_agent_{:012x}",
40            NEXT_BACKGROUND_AGENT_TASK_ID.fetch_add(1, Ordering::Relaxed)
41        );
42        let state = Arc::new(Mutex::new(BackgroundAgentTaskState {
43            status: AgentStatus::Running,
44            result: None,
45            error: None,
46        }));
47        let state_for_worker = state.clone();
48        let runner = runner.clone();
49        let agent = self.agent.clone();
50        let task_id_for_error = task_id.clone();
51        let _ = std::thread::Builder::new()
52            .name(format!("vv-agent-background-{task_id}"))
53            .spawn(move || {
54                let update = catch_unwind(AssertUnwindSafe(|| {
55                    match runner.run_blocking(
56                        &agent,
57                        NormalizedInput::from(input),
58                        run_config,
59                        None,
60                    ) {
61                        Ok(result) => {
62                            let status = result.status();
63                            let error = result.result().error.clone();
64                            (status, Some(result), error)
65                        }
66                        Err(error) => (AgentStatus::Failed, None, Some(error)),
67                    }
68                }))
69                .unwrap_or_else(|_| {
70                    (
71                        AgentStatus::Failed,
72                        None,
73                        Some(BACKGROUND_AGENT_TASK_WORKER_PANIC.to_string()),
74                    )
75                });
76                if let Ok(mut state) = state_for_worker.lock() {
77                    (state.status, state.result, state.error) = update;
78                }
79            })
80            .map_err(|error| {
81                if let Ok(mut state) = state.lock() {
82                    state.status = AgentStatus::Failed;
83                    state.error = Some(error.to_string());
84                }
85                format!("failed to spawn background agent task {task_id_for_error}: {error}")
86            })?;
87        let handle = BackgroundAgentTaskHandle {
88            task_id,
89            agent_name: self.agent.name().to_string(),
90            state,
91        };
92        self.handles
93            .lock()
94            .map_err(|_| "background task registry lock poisoned".to_string())?
95            .insert(handle.task_id.clone(), handle.clone());
96        Ok(handle)
97    }
98
99    pub fn get_handle(&self, task_id: &str) -> Result<BackgroundAgentTaskHandle, String> {
100        self.handles
101            .lock()
102            .map_err(|_| "background task registry lock poisoned".to_string())?
103            .get(task_id)
104            .cloned()
105            .ok_or_else(|| format!("unknown background agent task: {task_id}"))
106    }
107
108    fn input_from_arguments(&self, raw_arguments: Value) -> Result<String, String> {
109        let object = raw_arguments
110            .as_object()
111            .ok_or_else(|| "background task arguments must be an object".to_string())?;
112        object
113            .get("task_description")
114            .and_then(Value::as_str)
115            .map(str::trim)
116            .filter(|value| !value.is_empty())
117            .map(str::to_string)
118            .ok_or_else(|| "background task requires task_description".to_string())
119    }
120}
121
122impl Tool for BackgroundAgentTask {
123    fn name(&self) -> &str {
124        &self.name
125    }
126
127    fn description(&self) -> &str {
128        &self.description
129    }
130
131    fn parameters_schema(&self) -> &Value {
132        &self.parameters_schema
133    }
134
135    fn as_tool_spec(&self) -> ToolSpec {
136        let name = self.name.clone();
137        let description = self.description.clone();
138        let parameters_schema = self.parameters_schema.clone();
139        let task = self.clone();
140        let mut spec = ToolSpec::new(
141            name.clone(),
142            description.clone(),
143            Arc::new(
144                move |context: &mut ToolContext, arguments: &ToolArguments| {
145                    let raw_arguments = Value::Object(arguments.clone().into_iter().collect());
146                    let task_description = match task.input_from_arguments(raw_arguments.clone()) {
147                        Ok(task_description) => task_description,
148                        Err(error) => {
149                            return ToolOutput::error(error)
150                                .with_code("invalid_background_task_arguments")
151                                .to_result(&context.tool_call_id)
152                        }
153                    };
154                    let Some(model_provider) = context.model_provider.clone() else {
155                        return ToolOutput::error("background agent runtime is not available")
156                            .with_code("background_agent_runtime_unavailable")
157                            .to_result(&context.tool_call_id);
158                    };
159                    let runner = match Runner::builder()
160                        .model_provider_arc(model_provider)
161                        .workspace(context.workspace.clone())
162                        .build()
163                    {
164                        Ok(runner) => runner,
165                        Err(error) => {
166                            return ToolOutput::error(error)
167                                .with_code("background_agent_runtime_unavailable")
168                                .to_result(&context.tool_call_id)
169                        }
170                    };
171                    match task.start(&runner, context, raw_arguments, None) {
172                        Ok(handle) => ToolOutput::json(json!({
173                            "agent_name": task.agent.name(),
174                            "status": "background_task_started",
175                            "task_description": task_description,
176                            "task_id": handle.task_id(),
177                        }))
178                        .to_result(&context.tool_call_id),
179                        Err(error) => ToolOutput::error(error)
180                            .with_code("background_agent_start_failed")
181                            .to_result(&context.tool_call_id),
182                    }
183                },
184            ),
185        );
186        spec.kind = ToolSpecKind::BackgroundAgent;
187        spec.schema = json!({
188            "type": "function",
189            "function": {
190                "name": name,
191                "description": description,
192                "parameters": parameters_schema,
193            }
194        });
195        spec
196    }
197}
198
199pub struct BackgroundAgentTaskBuilder {
200    agent: Agent,
201    name: Option<String>,
202    description: Option<String>,
203}
204
205impl BackgroundAgentTaskBuilder {
206    pub fn new(agent: Agent) -> Self {
207        Self {
208            agent,
209            name: None,
210            description: None,
211        }
212    }
213
214    pub fn name(mut self, name: impl Into<String>) -> Self {
215        self.name = Some(name.into());
216        self
217    }
218
219    pub fn description(mut self, description: impl Into<String>) -> Self {
220        self.description = Some(description.into());
221        self
222    }
223
224    pub fn build(self) -> Result<BackgroundAgentTask, String> {
225        let name = self
226            .name
227            .unwrap_or_else(|| format!("{}_background_task", self.agent.name()));
228        if name.trim().is_empty() {
229            return Err("background task tool name cannot be empty".to_string());
230        }
231        let description = self.description.unwrap_or_else(|| {
232            format!(
233                "Start the {} agent as a background task.",
234                self.agent.name()
235            )
236        });
237        Ok(BackgroundAgentTask {
238            agent: self.agent,
239            name,
240            description,
241            parameters_schema: json!({
242                "type": "object",
243                "properties": {
244                    "task_description": {
245                        "type": "string",
246                        "description": "Task for the background agent."
247                    }
248                },
249                "required": ["task_description"],
250                "additionalProperties": false
251            }),
252            handles: Arc::new(Mutex::new(BTreeMap::new())),
253        })
254    }
255}
256
257pub(crate) fn inherited_run_config(
258    context: &ToolContext,
259    run_config: Option<RunConfig>,
260) -> RunConfig {
261    let explicit_shared_state = run_config
262        .as_ref()
263        .map(|config| config.initial_shared_state.clone())
264        .unwrap_or_default();
265    let explicit_metadata = run_config
266        .as_ref()
267        .map(|config| config.metadata.clone())
268        .unwrap_or_default();
269    let has_projected_parent = context.background_parent_run_config.is_some();
270    let mut config = context
271        .background_parent_run_config
272        .clone()
273        .or(run_config)
274        .unwrap_or_default();
275    if config.workspace.is_none() {
276        config.workspace = Some(context.workspace.clone());
277    }
278    if config.workspace_backend.is_none() {
279        config.workspace_backend = Some(context.workspace_backend.clone());
280    }
281    if config.model_provider.is_none() {
282        config.model_provider = context.model_provider.clone();
283    }
284    if config.execution_backend.is_none() {
285        config.execution_backend = context.execution_backend.clone();
286    }
287    if config.app_state.is_none() {
288        config.app_state = context.app_state.clone();
289    }
290    if config.cancellation_token.is_none() {
291        config.cancellation_token = context
292            .sub_task_turn_snapshot
293            .as_ref()
294            .and_then(|snapshot| snapshot.cancellation_token.as_ref())
295            .cloned();
296    }
297    let scoped_child_cancellation = config
298        .cancellation_token
299        .is_none()
300        .then(CancellationToken::child_of_current)
301        .flatten();
302
303    let mut shared_state = context.shared_state.clone();
304    shared_state.extend(explicit_shared_state);
305
306    if !has_projected_parent {
307        config.metadata = context.metadata.clone();
308        for key in [
309            "agent_name",
310            "session_id",
311            "approved_tool_interruption_ids",
312            "_vv_agent_run_id",
313            "_vv_agent_trace_id",
314            "_vv_agent_agent_name",
315            "_vv_agent_input",
316            "_vv_agent_session_id",
317            "_vv_agent_tool_use_behavior",
318            "_vv_agent_stop_at_tool_names",
319        ] {
320            config.metadata.remove(key);
321        }
322    }
323    config.metadata.extend(explicit_metadata);
324    let mut child = config.for_background_child(shared_state);
325    if child.cancellation_token.is_none() {
326        child.cancellation_token = scoped_child_cancellation;
327    }
328    child
329}
330
331#[derive(Clone)]
332pub struct BackgroundAgentTaskHandle {
333    task_id: String,
334    agent_name: String,
335    state: Arc<Mutex<BackgroundAgentTaskState>>,
336}
337
338impl BackgroundAgentTaskHandle {
339    pub fn task_id(&self) -> &str {
340        &self.task_id
341    }
342
343    pub fn agent_name(&self) -> &str {
344        &self.agent_name
345    }
346
347    pub fn status(&self) -> AgentStatus {
348        self.state
349            .lock()
350            .map(|state| state.status)
351            .unwrap_or(AgentStatus::Failed)
352    }
353
354    pub fn poll(&self) -> Result<BackgroundAgentTaskSnapshot, String> {
355        let state = self
356            .state
357            .lock()
358            .map_err(|_| "background task lock poisoned".to_string())?;
359        Ok(BackgroundAgentTaskSnapshot {
360            task_id: self.task_id.clone(),
361            agent_name: self.agent_name.clone(),
362            status: state.status,
363            final_output: state
364                .result
365                .as_ref()
366                .and_then(|result| result.final_output().map(str::to_string)),
367            error: state.error.clone(),
368        })
369    }
370
371    pub async fn wait(&self) -> Result<BackgroundAgentTaskSnapshot, String> {
372        loop {
373            let snapshot = self.poll()?;
374            if !matches!(snapshot.status, AgentStatus::Running | AgentStatus::Pending) {
375                return Ok(snapshot);
376            }
377            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
378        }
379    }
380
381    pub async fn wait_with_timeout(
382        &self,
383        timeout: std::time::Duration,
384    ) -> Result<BackgroundAgentTaskSnapshot, String> {
385        tokio::time::timeout(timeout, self.wait())
386            .await
387            .map_err(|_| {
388                format!(
389                    "background agent task {} was not ready before timeout",
390                    self.task_id
391                )
392            })?
393    }
394}
395
396struct BackgroundAgentTaskState {
397    status: AgentStatus,
398    result: Option<RunResult>,
399    error: Option<String>,
400}
401
402#[derive(Debug, Clone, PartialEq, Eq)]
403pub struct BackgroundAgentTaskSnapshot {
404    task_id: String,
405    agent_name: String,
406    status: AgentStatus,
407    final_output: Option<String>,
408    error: Option<String>,
409}
410
411impl BackgroundAgentTaskSnapshot {
412    pub fn task_id(&self) -> &str {
413        &self.task_id
414    }
415
416    pub fn agent_name(&self) -> &str {
417        &self.agent_name
418    }
419
420    pub fn status(&self) -> AgentStatus {
421        self.status
422    }
423
424    pub fn final_output(&self) -> Option<&str> {
425        self.final_output.as_deref()
426    }
427
428    pub fn error(&self) -> Option<&str> {
429        self.error.as_deref()
430    }
431}
432
433#[cfg(test)]
434mod tests {
435    use std::sync::Arc;
436
437    use serde_json::json;
438
439    use super::inherited_run_config;
440    use crate::context_providers::{
441        ContextError, ContextFragment, ContextProvider, ContextRequest,
442    };
443    use crate::memory::{
444        MemoryFuture, MemoryProvider, MemorySaveRequest, MemorySaveResult, MemorySearchRequest,
445        MemorySearchResult,
446    };
447    use crate::model::ModelRef;
448    use crate::model_settings::ModelSettings;
449    use crate::run_config::RunConfig;
450    use crate::runtime::{CancellationToken, RuntimeHook, SubTaskManager, SubTaskTurnSnapshot};
451    use crate::sessions::MemorySession;
452    use crate::tools::{ApprovalPolicy, ToolContext, ToolPolicy, ToolRegistry};
453    use crate::types::Message;
454
455    struct NoopContextProvider;
456
457    impl ContextProvider for NoopContextProvider {
458        fn fragments(
459            &self,
460            _request: &ContextRequest<'_>,
461        ) -> Result<Vec<ContextFragment>, ContextError> {
462            Ok(Vec::new())
463        }
464    }
465
466    struct NoopMemoryProvider;
467
468    impl MemoryProvider for NoopMemoryProvider {
469        fn search(&self, _request: MemorySearchRequest) -> MemoryFuture<Vec<MemorySearchResult>> {
470            Box::pin(async { Ok(Vec::new()) })
471        }
472
473        fn save(&self, _request: MemorySaveRequest) -> MemoryFuture<MemorySaveResult> {
474            Box::pin(async { Ok(MemorySaveResult::default()) })
475        }
476    }
477
478    struct NoopHook;
479
480    impl RuntimeHook for NoopHook {}
481
482    #[test]
483    fn approval_resume_snapshot_derives_one_way_child_cancellation() {
484        let snapshot_parent = CancellationToken::default();
485        let fallback_parent = CancellationToken::default();
486        let mut context = ToolContext::new("./workspace");
487        context.sub_task_turn_snapshot = Some(SubTaskTurnSnapshot {
488            cancellation_token: Some(snapshot_parent.clone()),
489            ..SubTaskTurnSnapshot::default()
490        });
491        let child = {
492            let _scope = CancellationToken::enter_scope(Some(&fallback_parent));
493            inherited_run_config(&context, None)
494                .cancellation_token
495                .expect("derived child cancellation")
496        };
497
498        child.cancel();
499        assert!(!snapshot_parent.is_cancelled());
500        assert!(!fallback_parent.is_cancelled());
501
502        let second_child = {
503            let _scope = CancellationToken::enter_scope(Some(&fallback_parent));
504            inherited_run_config(&context, None)
505                .cancellation_token
506                .expect("derived child cancellation")
507        };
508        fallback_parent.cancel();
509        assert!(!second_child.is_cancelled());
510        snapshot_parent.cancel();
511        assert!(second_child.is_cancelled());
512    }
513
514    #[test]
515    fn thread_local_and_explicit_cancellation_are_derived_for_the_child() {
516        let context = ToolContext::new("./workspace");
517        let fallback_parent = CancellationToken::default();
518        let child = {
519            let _scope = CancellationToken::enter_scope(Some(&fallback_parent));
520            inherited_run_config(&context, None)
521                .cancellation_token
522                .expect("derived fallback child cancellation")
523        };
524        fallback_parent.cancel();
525        assert!(child.is_cancelled());
526
527        let unrelated_parent = CancellationToken::default();
528        let explicit = CancellationToken::default();
529        let config = {
530            let _scope = CancellationToken::enter_scope(Some(&unrelated_parent));
531            inherited_run_config(
532                &context,
533                Some(RunConfig {
534                    cancellation_token: Some(explicit.clone()),
535                    ..RunConfig::default()
536                }),
537            )
538        };
539        unrelated_parent.cancel();
540        assert!(!explicit.is_cancelled());
541        let explicit_child = config
542            .cancellation_token
543            .expect("derived explicit child cancellation");
544        explicit_child.cancel();
545        assert!(!explicit.is_cancelled());
546
547        let second_explicit_child = inherited_run_config(
548            &context,
549            Some(RunConfig {
550                cancellation_token: Some(explicit.clone()),
551                ..RunConfig::default()
552            }),
553        )
554        .cancellation_token
555        .expect("second explicit child cancellation");
556        explicit.cancel();
557        assert!(second_explicit_child.is_cancelled());
558        assert!(config.session.is_none());
559        assert!(config.approval_broker.is_none());
560    }
561
562    #[test]
563    fn projected_parent_preserves_capabilities_and_clears_run_instances() {
564        let parent_cancellation = CancellationToken::default();
565        let mut parent = RunConfig {
566            model: Some(ModelRef::named("parent-model")),
567            model_settings: Some(ModelSettings::default()),
568            session: Some(Arc::new(MemorySession::new("parent-session"))),
569            initial_messages: Some(vec![Message::user("parent history")]),
570            max_cycles: Some(7),
571            max_handoffs: Some(4),
572            tool_policy: ToolPolicy {
573                disallowed_tools: vec!["blocked".to_string()],
574                approval: ApprovalPolicy::OnRequest,
575                ..ToolPolicy::default()
576            },
577            cancellation_token: Some(parent_cancellation.clone()),
578            hooks: vec![Arc::new(NoopHook)],
579            context_providers: vec![Arc::new(NoopContextProvider)],
580            max_context_chars: Some(12_345),
581            memory_providers: vec![Arc::new(NoopMemoryProvider)],
582            app_state: Some(Arc::new("parent app state".to_string())),
583            tool_registry_factory: Some(Arc::new(ToolRegistry::default)),
584            log_preview_chars: Some(321),
585            before_cycle_messages: Some(Arc::new(|_, _, _| Vec::new())),
586            interruption_messages: Some(Arc::new(Vec::new)),
587            sub_task_manager: Some(SubTaskManager::default()),
588            stream: Some(Arc::new(|_| {})),
589            ..RunConfig::default()
590        };
591        parent
592            .metadata
593            .insert("parent_metadata".to_string(), json!("retained"));
594        parent
595            .initial_shared_state
596            .insert("stale".to_string(), json!(true));
597
598        let mut context = ToolContext::new("./workspace");
599        context.background_parent_run_config = Some(parent);
600        context
601            .shared_state
602            .insert("live".to_string(), json!("snapshot"));
603        let child = inherited_run_config(&context, None);
604
605        assert!(child.model.is_none());
606        assert!(child.model_settings.is_none());
607        assert!(child.session.is_none());
608        assert!(child.initial_messages.is_none());
609        assert!(child.before_cycle_messages.is_none());
610        assert!(child.interruption_messages.is_none());
611        assert!(child.sub_task_manager.is_none());
612        assert!(child.stream.is_none());
613        assert_eq!(
614            child.initial_shared_state.get("live"),
615            Some(&json!("snapshot"))
616        );
617        assert!(!child.initial_shared_state.contains_key("stale"));
618
619        assert_eq!(child.max_cycles, Some(7));
620        assert_eq!(child.max_handoffs, Some(4));
621        assert_eq!(child.tool_policy.approval, ApprovalPolicy::OnRequest);
622        assert_eq!(child.tool_policy.disallowed_tools, ["blocked"]);
623        assert_eq!(child.hooks.len(), 1);
624        assert_eq!(child.context_providers.len(), 1);
625        assert_eq!(child.max_context_chars, Some(12_345));
626        assert_eq!(child.memory_providers.len(), 1);
627        assert!(child.app_state.is_some());
628        assert!(child.tool_registry_factory.is_some());
629        assert_eq!(child.log_preview_chars, Some(321));
630        assert_eq!(child.metadata["parent_metadata"], json!("retained"));
631
632        let child_cancellation = child
633            .cancellation_token
634            .expect("derived child cancellation");
635        child_cancellation.cancel();
636        assert!(!parent_cancellation.is_cancelled());
637        let second_child = inherited_run_config(&context, None)
638            .cancellation_token
639            .expect("second child cancellation");
640        parent_cancellation.cancel();
641        assert!(second_child.is_cancelled());
642    }
643}