Skip to main content

vv_agent/tools/
background_agent_task.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2use std::sync::{Arc, Mutex};
3
4use serde_json::{json, Value};
5
6use crate::agent::Agent;
7use crate::result::RunResult;
8use crate::runner::{NormalizedInput, Runner};
9use crate::tools::{Tool, ToolContext, ToolOutput, ToolSpec, ToolSpecKind};
10use crate::types::{AgentStatus, ToolArguments};
11
12static NEXT_BACKGROUND_AGENT_TASK_ID: AtomicU64 = AtomicU64::new(1);
13
14#[derive(Clone)]
15pub struct BackgroundAgentTask {
16    agent: Agent,
17    name: String,
18    description: String,
19    parameters_schema: Value,
20}
21
22impl BackgroundAgentTask {
23    pub fn start(
24        &self,
25        runner: &Runner,
26        _context: &mut ToolContext,
27        raw_arguments: Value,
28    ) -> Result<BackgroundAgentTaskHandle, String> {
29        let input = self.input_from_arguments(raw_arguments)?;
30        let task_id = format!(
31            "bg_agent_{:012x}",
32            NEXT_BACKGROUND_AGENT_TASK_ID.fetch_add(1, Ordering::Relaxed)
33        );
34        let state = Arc::new(Mutex::new(BackgroundAgentTaskState {
35            status: AgentStatus::Running,
36            result: None,
37            error: None,
38        }));
39        let state_for_worker = state.clone();
40        let runner = runner.clone();
41        let agent = self.agent.clone();
42        let task_id_for_error = task_id.clone();
43        let _ = std::thread::Builder::new()
44            .name(format!("vv-agent-background-{task_id}"))
45            .spawn(move || {
46                let result = runner.run_blocking(
47                    &agent,
48                    NormalizedInput::from(input),
49                    crate::run_config::RunConfig::default(),
50                    None,
51                );
52                if let Ok(mut state) = state_for_worker.lock() {
53                    match result {
54                        Ok(result) => {
55                            state.status = result.status();
56                            state.result = Some(result);
57                        }
58                        Err(error) => {
59                            state.status = AgentStatus::Failed;
60                            state.error = Some(error);
61                        }
62                    }
63                }
64            })
65            .map_err(|error| {
66                if let Ok(mut state) = state.lock() {
67                    state.status = AgentStatus::Failed;
68                    state.error = Some(error.to_string());
69                }
70                format!("failed to spawn background agent task {task_id_for_error}: {error}")
71            })?;
72        Ok(BackgroundAgentTaskHandle {
73            task_id,
74            agent_name: self.agent.name().to_string(),
75            state,
76        })
77    }
78
79    fn input_from_arguments(&self, raw_arguments: Value) -> Result<String, String> {
80        let object = raw_arguments
81            .as_object()
82            .ok_or_else(|| "background task arguments must be an object".to_string())?;
83        object
84            .get("task_description")
85            .or_else(|| object.get("task"))
86            .or_else(|| object.get("input"))
87            .and_then(Value::as_str)
88            .map(str::trim)
89            .filter(|value| !value.is_empty())
90            .map(str::to_string)
91            .ok_or_else(|| "background task requires task_description".to_string())
92    }
93}
94
95impl Tool for BackgroundAgentTask {
96    fn name(&self) -> &str {
97        &self.name
98    }
99
100    fn description(&self) -> &str {
101        &self.description
102    }
103
104    fn parameters_schema(&self) -> &Value {
105        &self.parameters_schema
106    }
107
108    fn as_tool_spec(&self) -> ToolSpec {
109        let name = self.name.clone();
110        let description = self.description.clone();
111        let parameters_schema = self.parameters_schema.clone();
112        let task = self.clone();
113        let mut spec = ToolSpec::new(
114            name.clone(),
115            description.clone(),
116            Arc::new(
117                move |_context: &mut ToolContext, arguments: &ToolArguments| {
118                    let raw_arguments = Value::Object(arguments.clone().into_iter().collect());
119                    match task.input_from_arguments(raw_arguments) {
120                        Ok(task_description) => ToolOutput::json(json!({
121                            "agent_name": task.agent.name(),
122                            "status": "background_task_requested",
123                            "task_description": task_description,
124                        }))
125                        .to_result(""),
126                        Err(error) => ToolOutput::error(error)
127                            .with_code("invalid_background_task_arguments")
128                            .to_result(""),
129                    }
130                },
131            ),
132        );
133        spec.kind = ToolSpecKind::BackgroundAgent;
134        spec.schema = json!({
135            "type": "function",
136            "function": {
137                "name": name,
138                "description": description,
139                "parameters": parameters_schema,
140            }
141        });
142        spec
143    }
144}
145
146pub struct BackgroundAgentTaskBuilder {
147    agent: Agent,
148    name: Option<String>,
149    description: Option<String>,
150}
151
152impl BackgroundAgentTaskBuilder {
153    pub fn new(agent: Agent) -> Self {
154        Self {
155            agent,
156            name: None,
157            description: None,
158        }
159    }
160
161    pub fn name(mut self, name: impl Into<String>) -> Self {
162        self.name = Some(name.into());
163        self
164    }
165
166    pub fn description(mut self, description: impl Into<String>) -> Self {
167        self.description = Some(description.into());
168        self
169    }
170
171    pub fn build(self) -> Result<BackgroundAgentTask, String> {
172        let name = self
173            .name
174            .unwrap_or_else(|| format!("{}_background_task", self.agent.name()));
175        if name.trim().is_empty() {
176            return Err("background task tool name cannot be empty".to_string());
177        }
178        let description = self.description.unwrap_or_else(|| {
179            format!(
180                "Start the {} agent as a background task.",
181                self.agent.name()
182            )
183        });
184        Ok(BackgroundAgentTask {
185            agent: self.agent,
186            name,
187            description,
188            parameters_schema: json!({
189                "type": "object",
190                "properties": {
191                    "task_description": {
192                        "type": "string",
193                        "description": "Task for the background agent."
194                    }
195                },
196                "required": ["task_description"]
197            }),
198        })
199    }
200}
201
202#[derive(Clone)]
203pub struct BackgroundAgentTaskHandle {
204    task_id: String,
205    agent_name: String,
206    state: Arc<Mutex<BackgroundAgentTaskState>>,
207}
208
209impl BackgroundAgentTaskHandle {
210    pub fn task_id(&self) -> &str {
211        &self.task_id
212    }
213
214    pub fn agent_name(&self) -> &str {
215        &self.agent_name
216    }
217
218    pub fn status(&self) -> AgentStatus {
219        self.state
220            .lock()
221            .map(|state| state.status)
222            .unwrap_or(AgentStatus::Failed)
223    }
224
225    pub fn poll(&self) -> Result<BackgroundAgentTaskSnapshot, String> {
226        let state = self
227            .state
228            .lock()
229            .map_err(|_| "background task lock poisoned".to_string())?;
230        Ok(BackgroundAgentTaskSnapshot {
231            task_id: self.task_id.clone(),
232            agent_name: self.agent_name.clone(),
233            status: state.status,
234            final_output: state
235                .result
236                .as_ref()
237                .and_then(|result| result.final_output().map(str::to_string)),
238            error: state.error.clone(),
239        })
240    }
241
242    pub async fn wait(&self) -> Result<BackgroundAgentTaskSnapshot, String> {
243        loop {
244            let snapshot = self.poll()?;
245            if !matches!(snapshot.status, AgentStatus::Running | AgentStatus::Pending) {
246                return Ok(snapshot);
247            }
248            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
249        }
250    }
251}
252
253struct BackgroundAgentTaskState {
254    status: AgentStatus,
255    result: Option<RunResult>,
256    error: Option<String>,
257}
258
259#[derive(Debug, Clone, PartialEq, Eq)]
260pub struct BackgroundAgentTaskSnapshot {
261    task_id: String,
262    agent_name: String,
263    status: AgentStatus,
264    final_output: Option<String>,
265    error: Option<String>,
266}
267
268impl BackgroundAgentTaskSnapshot {
269    pub fn task_id(&self) -> &str {
270        &self.task_id
271    }
272
273    pub fn agent_name(&self) -> &str {
274        &self.agent_name
275    }
276
277    pub fn status(&self) -> AgentStatus {
278        self.status
279    }
280
281    pub fn final_output(&self) -> Option<&str> {
282        self.final_output.as_deref()
283    }
284
285    pub fn error(&self) -> Option<&str> {
286        self.error.as_deref()
287    }
288}