vv_agent/runtime/sub_task_manager/
sessions.rs1use std::collections::BTreeMap;
2use std::panic::{catch_unwind, AssertUnwindSafe};
3use std::sync::Arc;
4use std::thread;
5
6use crate::runtime::sub_agent_sessions::{
7 register_sub_agent_session, unregister_sub_agent_session, SubAgentSession,
8 SubAgentSessionListener,
9};
10use crate::types::{AgentStatus, SubTaskOutcome};
11use crate::workspace::WorkspaceBackend;
12
13use super::helpers::{now_iso, panic_payload_to_string};
14use super::manager::SubTaskManager;
15use super::record::ManagedSubTask;
16use super::types::SubTaskSessionAttachment;
17
18impl SubTaskManager {
19 pub fn attach_session(
20 &self,
21 task_id: impl Into<String>,
22 session_id: impl Into<String>,
23 agent_name: impl Into<String>,
24 task_title: impl Into<String>,
25 workspace_backend: Arc<dyn WorkspaceBackend>,
26 session: Arc<dyn SubAgentSession>,
27 ) {
28 self.attach_session_with_resolved(SubTaskSessionAttachment {
29 task_id: task_id.into(),
30 session_id: session_id.into(),
31 agent_name: agent_name.into(),
32 task_title: task_title.into(),
33 workspace_backend,
34 session,
35 resolved: BTreeMap::new(),
36 });
37 }
38
39 pub fn attach_session_with_resolved(&self, attachment: SubTaskSessionAttachment) {
40 let SubTaskSessionAttachment {
41 task_id,
42 session_id,
43 agent_name,
44 task_title,
45 workspace_backend,
46 session,
47 resolved,
48 } = attachment;
49 let should_attach_listener = {
50 let mut tasks = self.tasks.lock().expect("sub-task manager poisoned");
51 match tasks.get_mut(&task_id) {
52 Some(record) => {
53 record.session_id = session_id;
54 record.agent_name = agent_name;
55 if !task_title.is_empty() {
56 record.task_title = task_title;
57 }
58 record.workspace_backend = Some(workspace_backend);
59 record.session = Some(session.clone());
60 if !resolved.is_empty() {
61 record.resolved = resolved;
62 }
63 record.updated_at = now_iso();
64 let should_attach_listener = !record.manager_listener_attached;
65 record.manager_listener_attached = true;
66 should_attach_listener
67 }
68 None => {
69 tasks.insert(
70 task_id.clone(),
71 ManagedSubTask {
72 task_id: task_id.clone(),
73 session_id,
74 agent_name,
75 task_title,
76 workspace_backend: Some(workspace_backend),
77 session: Some(session.clone()),
78 outcome: None,
79 resolved,
80 current_cycle_index: None,
81 recent_activity: None,
82 latest_cycle: None,
83 latest_tool_call: None,
84 handle: None,
85 updated_at: now_iso(),
86 manager_listener_attached: true,
87 },
88 );
89 true
90 }
91 }
92 };
93
94 if should_attach_listener {
95 let manager = self.clone();
96 let task_id = task_id.clone();
97 let listener: SubAgentSessionListener = Arc::new(move |event, payload| {
98 manager.handle_session_event(&task_id, event, payload);
99 });
100 let _ = session.subscribe(listener);
101 }
102 }
103
104 pub fn continue_task(&self, task_id: &str, prompt: &str) -> Result<(), String> {
105 let prompt = prompt.trim();
106 if prompt.is_empty() {
107 return Err("Follow-up prompt cannot be empty.".to_string());
108 }
109
110 let task_id = task_id.trim();
111 let (session_id, agent_name, session) = {
112 let mut tasks = self.tasks.lock().expect("sub-task manager poisoned");
113 let Some(record) = tasks.get_mut(task_id) else {
114 return Err(format!("Sub-task {task_id} not found."));
115 };
116 if record.is_running() {
117 return Err(format!("Sub-task {task_id} is already running."));
118 }
119 if record
120 .outcome
121 .as_ref()
122 .is_some_and(|outcome| outcome.status == AgentStatus::MaxCycles)
123 {
124 return Err(format!(
125 "Sub-task {task_id} reached max cycles and cannot continue."
126 ));
127 }
128 if record.session_id.trim().is_empty() {
129 return Err(format!("Sub-task {task_id} session is not available."));
130 }
131 let Some(session) = record.session.clone() else {
132 return Err(format!("Sub-task {task_id} session is not attached."));
133 };
134
135 record.task_title = prompt.to_string();
136 record.outcome = None;
137 record.recent_activity = Some(prompt.to_string());
138 record.updated_at = now_iso();
139 (
140 record.session_id.clone(),
141 record.agent_name.clone(),
142 session,
143 )
144 };
145
146 session.sanitize_for_resume();
147
148 let tasks = self.tasks.clone();
149 let task_id_for_thread = task_id.to_string();
150 let prompt_for_thread = prompt.to_string();
151 let session_id_for_thread = session_id.clone();
152 let agent_name_for_thread = agent_name.clone();
153 let handle = thread::spawn(move || {
154 register_sub_agent_session(session_id_for_thread.clone(), session.clone());
155 let outcome = match catch_unwind(AssertUnwindSafe(|| {
156 session.continue_run(&prompt_for_thread)
157 })) {
158 Ok(Ok(outcome)) => outcome,
159 Ok(Err(error)) => SubTaskOutcome {
160 task_id: task_id_for_thread.clone(),
161 agent_name: agent_name_for_thread.clone(),
162 status: AgentStatus::Failed,
163 session_id: Some(session_id_for_thread.clone()),
164 final_answer: None,
165 wait_reason: None,
166 error: Some(error),
167 cycles: 0,
168 todo_list: Vec::new(),
169 resolved: BTreeMap::new(),
170 },
171 Err(payload) => SubTaskOutcome {
172 task_id: task_id_for_thread.clone(),
173 agent_name: agent_name_for_thread.clone(),
174 status: AgentStatus::Failed,
175 session_id: Some(session_id_for_thread.clone()),
176 final_answer: None,
177 wait_reason: None,
178 error: Some(panic_payload_to_string(payload.as_ref())),
179 cycles: 0,
180 todo_list: Vec::new(),
181 resolved: BTreeMap::new(),
182 },
183 };
184 unregister_sub_agent_session(&session_id_for_thread);
185 let mut tasks = tasks.lock().expect("sub-task manager poisoned");
186 if let Some(record) = tasks.get_mut(&task_id_for_thread) {
187 let mut outcome = outcome;
188 if outcome.resolved.is_empty() && !record.resolved.is_empty() {
189 outcome.resolved = record.resolved.clone();
190 }
191 record.session_id = outcome
192 .session_id
193 .clone()
194 .unwrap_or_else(|| record.session_id.clone());
195 record.agent_name = outcome.agent_name.clone();
196 record.update_from_outcome(&outcome);
197 record.outcome = Some(outcome);
198 record.updated_at = now_iso();
199 record.handle = None;
200 }
201 });
202
203 let mut tasks = self.tasks.lock().expect("sub-task manager poisoned");
204 if let Some(record) = tasks.get_mut(task_id) {
205 record.handle = Some(handle);
206 record.updated_at = now_iso();
207 }
208 Ok(())
209 }
210}