vv_agent/runtime/
sub_agent_sessions.rs

1use std::collections::BTreeMap;
2use std::sync::{Arc, Mutex, OnceLock};
3
4use serde_json::Value;
5
6use crate::runtime::sub_task_manager::SubTaskTurnSnapshot;
7use crate::types::SubTaskOutcome;
8
9pub type SubAgentSessionListener =
10    Arc<dyn Fn(&str, &BTreeMap<String, Value>) + Send + Sync + 'static>;
11pub type SubAgentSessionUnsubscribe = Box<dyn FnOnce() + Send + 'static>;
12
13pub trait SubAgentSession: Send + Sync {
14    fn steer(&self, prompt: &str) -> Result<(), String>;
15
16    fn cancel(&self) -> bool {
17        false
18    }
19
20    fn sanitize_for_resume(&self) -> usize {
21        0
22    }
23
24    fn continue_run(&self, _prompt: &str) -> Result<SubTaskOutcome, String> {
25        Err("Sub-agent session continuation is not supported.".to_string())
26    }
27
28    fn continue_run_with_snapshot(
29        &self,
30        prompt: &str,
31        _snapshot: SubTaskTurnSnapshot,
32    ) -> Result<SubTaskOutcome, String> {
33        self.continue_run(prompt)
34    }
35
36    fn subscribe(&self, _listener: SubAgentSessionListener) -> Option<SubAgentSessionUnsubscribe> {
37        None
38    }
39}
40
41#[derive(Default)]
42pub struct SubAgentSessionRegistry {
43    sessions: Mutex<BTreeMap<String, Arc<dyn SubAgentSession>>>,
44}
45
46impl SubAgentSessionRegistry {
47    pub fn register(&self, session_id: impl Into<String>, session: Arc<dyn SubAgentSession>) {
48        let session_id = session_id.into();
49        let session_id = session_id.trim();
50        if session_id.is_empty() {
51            return;
52        }
53        self.sessions
54            .lock()
55            .unwrap_or_else(|poisoned| poisoned.into_inner())
56            .insert(session_id.to_string(), session);
57    }
58
59    pub fn unregister(&self, session_id: &str) {
60        let session_id = session_id.trim();
61        if session_id.is_empty() {
62            return;
63        }
64        self.sessions
65            .lock()
66            .unwrap_or_else(|poisoned| poisoned.into_inner())
67            .remove(session_id);
68    }
69
70    pub fn unregister_if_matches(
71        &self,
72        session_id: &str,
73        session: Option<&Arc<dyn SubAgentSession>>,
74    ) -> bool {
75        let session_id = session_id.trim();
76        if session_id.is_empty() {
77            return false;
78        }
79        let mut sessions = self
80            .sessions
81            .lock()
82            .unwrap_or_else(|poisoned| poisoned.into_inner());
83        let Some(registered) = sessions.get(session_id) else {
84            return false;
85        };
86        if let Some(expected) = session {
87            if !Arc::ptr_eq(registered, expected) {
88                return false;
89            }
90        }
91        sessions.remove(session_id).is_some()
92    }
93
94    pub fn get(&self, session_id: &str) -> Option<Arc<dyn SubAgentSession>> {
95        let session_id = session_id.trim();
96        if session_id.is_empty() {
97            return None;
98        }
99        self.sessions
100            .lock()
101            .unwrap_or_else(|poisoned| poisoned.into_inner())
102            .get(session_id)
103            .cloned()
104    }
105
106    pub fn clear(&self) {
107        self.sessions
108            .lock()
109            .unwrap_or_else(|poisoned| poisoned.into_inner())
110            .clear();
111    }
112}
113
114static ACTIVE_SUB_AGENT_SESSIONS: OnceLock<SubAgentSessionRegistry> = OnceLock::new();
115
116pub fn sub_agent_session_registry() -> &'static SubAgentSessionRegistry {
117    ACTIVE_SUB_AGENT_SESSIONS.get_or_init(SubAgentSessionRegistry::default)
118}
119
120pub fn register_sub_agent_session(
121    session_id: impl Into<String>,
122    session: Arc<dyn SubAgentSession>,
123) {
124    sub_agent_session_registry().register(session_id, session);
125}
126
127pub fn unregister_sub_agent_session(session_id: &str) {
128    sub_agent_session_registry().unregister(session_id);
129}
130
131pub fn _register_sub_agent_session(
132    session_id: impl Into<String>,
133    session: Arc<dyn SubAgentSession>,
134) {
135    register_sub_agent_session(session_id, session);
136}
137
138pub fn _unregister_sub_agent_session(
139    session_id: &str,
140    session: Option<Arc<dyn SubAgentSession>>,
141) -> bool {
142    sub_agent_session_registry().unregister_if_matches(session_id, session.as_ref())
143}
144
145pub fn get_sub_agent_session(session_id: &str) -> Option<Arc<dyn SubAgentSession>> {
146    sub_agent_session_registry().get(session_id)
147}
148
149pub fn steer_sub_agent_session(session_id: &str, prompt: &str) -> bool {
150    let session_id = session_id.trim();
151    let prompt = prompt.trim();
152    if session_id.is_empty() || prompt.is_empty() {
153        return false;
154    }
155    let Some(session) = get_sub_agent_session(session_id) else {
156        return false;
157    };
158    session.steer(prompt).is_ok()
159}
160
161pub fn continue_sub_agent_session(
162    session_id: &str,
163    prompt: &str,
164) -> Result<SubTaskOutcome, String> {
165    let session_id = session_id.trim();
166    let prompt = prompt.trim();
167    if session_id.is_empty() {
168        return Err("Sub-agent session id cannot be empty.".to_string());
169    }
170    if prompt.is_empty() {
171        return Err("Follow-up prompt cannot be empty.".to_string());
172    }
173    let Some(session) = get_sub_agent_session(session_id) else {
174        return Err(format!("Sub-agent session {session_id} is not registered."));
175    };
176    session.continue_run(prompt)
177}
178
179pub fn subscribe_sub_agent_session(
180    session_id: &str,
181    listener: SubAgentSessionListener,
182) -> Option<SubAgentSessionUnsubscribe> {
183    get_sub_agent_session(session_id)?.subscribe(listener)
184}
vv_agent/runtime/sub_agent_sessions.rs

vv_agent/runtime/
sub_agent_sessions.rs
