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

vv_agent/runtime/
sub_agent_sessions.rs
