Skip to main content

vv_agent/runtime/backends/
thread.rs

1use std::sync::Arc;
2use std::thread::{self, JoinHandle};
3
4use super::{execute_cycle_loop, execute_cycle_loop_with_state};
5use crate::runtime::CancellationToken;
6use crate::types::{AgentResult, AgentTask, CycleRecord, Message, Metadata};
7
8#[derive(Debug, Clone)]
9pub struct ThreadBackend {
10    max_workers: usize,
11}
12
13impl Default for ThreadBackend {
14    fn default() -> Self {
15        Self::new(4)
16    }
17}
18
19impl ThreadBackend {
20    pub fn new(max_workers: usize) -> Self {
21        Self {
22            max_workers: max_workers.max(1),
23        }
24    }
25
26    pub fn max_workers(&self) -> usize {
27        self.max_workers
28    }
29
30    pub fn execute<F>(
31        &self,
32        _task: &AgentTask,
33        initial_messages: Vec<Message>,
34        shared_state: Metadata,
35        cycle_executor: F,
36        cancellation_token: Option<&CancellationToken>,
37        max_cycles: u32,
38    ) -> AgentResult
39    where
40        F: FnMut(
41            u32,
42            &mut Vec<Message>,
43            &mut Vec<CycleRecord>,
44            &mut Metadata,
45            Option<&CancellationToken>,
46        ) -> Option<AgentResult>,
47    {
48        execute_cycle_loop(
49            initial_messages,
50            shared_state,
51            cycle_executor,
52            cancellation_token,
53            max_cycles,
54        )
55    }
56
57    #[allow(clippy::too_many_arguments)]
58    pub(crate) fn execute_with_state<F>(
59        &self,
60        _task: &AgentTask,
61        initial_messages: Vec<Message>,
62        initial_cycles: Vec<CycleRecord>,
63        shared_state: Metadata,
64        cycle_executor: F,
65        cancellation_token: Option<&CancellationToken>,
66        cycle_index_start: u32,
67        cycle_count: u32,
68    ) -> AgentResult
69    where
70        F: FnMut(
71            u32,
72            &mut Vec<Message>,
73            &mut Vec<CycleRecord>,
74            &mut Metadata,
75            Option<&CancellationToken>,
76        ) -> Option<AgentResult>,
77    {
78        execute_cycle_loop_with_state(
79            initial_messages,
80            initial_cycles,
81            shared_state,
82            cycle_executor,
83            cancellation_token,
84            cycle_index_start,
85            cycle_count,
86        )
87    }
88
89    pub fn submit<R, F>(&self, function: F) -> JoinHandle<R>
90    where
91        R: Send + 'static,
92        F: FnOnce() -> R + Send + 'static,
93    {
94        thread::spawn(function)
95    }
96
97    pub fn parallel_map<T, R, F>(&self, function: F, items: Vec<T>) -> Vec<R>
98    where
99        T: Send + 'static,
100        R: Send + 'static,
101        F: Fn(T) -> R + Send + Sync + 'static,
102    {
103        let function = Arc::new(function);
104        let mut indexed_items = items.into_iter().enumerate();
105        let mut indexed_results = Vec::new();
106
107        loop {
108            let mut handles = Vec::new();
109            for _ in 0..self.max_workers {
110                let Some((index, item)) = indexed_items.next() else {
111                    break;
112                };
113                let function = Arc::clone(&function);
114                handles.push(thread::spawn(move || (index, function(item))));
115            }
116            if handles.is_empty() {
117                break;
118            }
119            for handle in handles {
120                indexed_results.push(handle.join().expect("thread backend worker panicked"));
121            }
122        }
123
124        indexed_results.sort_by_key(|(index, _)| *index);
125        indexed_results
126            .into_iter()
127            .map(|(_, result)| result)
128            .collect()
129    }
130}