zeph_subagent/manager/
collect.rs1use std::path::PathBuf;
5
6use zeph_config::SubAgentConfig;
7
8use super::SubAgentManager;
9use super::SubAgentStatus;
10use crate::def::SubAgentDef;
11use crate::error::SubAgentError;
12use crate::fleet::FleetSessionStatus;
13use crate::hooks::fire_hooks;
14use crate::manager::secrets::make_hook_env;
15use crate::state::SubAgentState;
16use crate::transcript::{TranscriptMeta, TranscriptWriter, sweep_old_transcripts};
17
18impl SubAgentManager {
19 #[tracing::instrument(name = "subagent.manager.collect", skip_all, fields(task_id = task_id))]
28 pub async fn collect(&mut self, task_id: &str) -> Result<String, SubAgentError> {
29 let mut handle = self
30 .agents
31 .remove(task_id)
32 .ok_or_else(|| SubAgentError::NotFound(task_id.to_owned()))?;
33
34 if !self.stop_hooks.is_empty() {
35 let stop_hooks = self.stop_hooks.clone();
36 let stop_env = make_hook_env(task_id, &handle.def.name, "");
37 self.spawn_hook_task(async move {
38 if let Err(e) = fire_hooks(&stop_hooks, &stop_env, None, None).await {
39 tracing::warn!(error = %e, "SubagentStop hook failed");
40 }
41 });
42 }
43
44 handle.grants.revoke_all();
45
46 let result = if let Some(jh) = handle.join_handle.take() {
47 jh.join()
48 .await
49 .map_err(|e| SubAgentError::Spawn(e.to_string()))?
50 } else {
51 Ok(String::new())
52 };
53
54 let final_state = {
55 let status = handle.status_rx.borrow();
56 if result.is_err() {
57 SubAgentState::Failed
58 } else if status.state == SubAgentState::Canceled {
59 SubAgentState::Canceled
60 } else {
61 SubAgentState::Completed
62 }
63 };
64
65 if let Some(ref registry) = self.fleet_registry {
66 let registry = std::sync::Arc::clone(registry);
67 let tid = task_id.to_owned();
68 let fleet_status = match final_state {
69 SubAgentState::Failed => FleetSessionStatus::Failed,
70 SubAgentState::Canceled => FleetSessionStatus::Cancelled,
71 _ => FleetSessionStatus::Completed,
72 };
73 self.spawn_hook_task(async move {
74 if let Err(e) = registry.mark_terminal(&tid, fleet_status).await {
75 tracing::warn!(error = %e, task_id = %tid, "fleet: mark_terminal failed");
76 }
77 });
78 }
79
80 if let Some(ref dir) = handle.transcript_dir.clone() {
81 let turns_used = handle.status_rx.borrow().turns_used;
82 let meta = TranscriptMeta {
83 agent_id: task_id.to_owned(),
84 agent_name: handle.def.name.clone(),
85 def_name: handle.def.name.clone(),
86 status: final_state,
87 started_at: handle.started_at_str.clone(),
88 finished_at: Some(crate::transcript::utc_now()),
89 resumed_from: None,
90 turns_used,
91 mcp_tool_names: handle.mcp_tool_names.clone(),
92 };
93 if let Err(e) = TranscriptWriter::write_meta_async(dir, task_id, &meta).await {
94 tracing::warn!(error = %e, task_id, "failed to write final transcript meta");
95 }
96 }
97
98 result
99 }
100
101 pub(crate) fn effective_transcript_dir(&self, config: &SubAgentConfig) -> PathBuf {
103 if let Some(ref dir) = self.transcript_dir {
104 dir.clone()
105 } else if let Some(ref dir) = config.transcript_dir {
106 dir.clone()
107 } else {
108 PathBuf::from(".zeph/subagents")
109 }
110 }
111
112 pub async fn def_name_for_resume(
122 &self,
123 id_prefix: &str,
124 config: &SubAgentConfig,
125 ) -> Result<String, SubAgentError> {
126 let dir = self.effective_transcript_dir(config);
127 let id_prefix = id_prefix.to_owned();
128 tokio::task::spawn_blocking(move || {
129 let original_id =
130 crate::transcript::TranscriptReader::find_by_prefix(&dir, &id_prefix)?;
131 let meta = crate::transcript::TranscriptReader::load_meta(&dir, &original_id)?;
132 Ok(meta.def_name)
133 })
134 .await
135 .map_err(|e| SubAgentError::Spawn(format!("spawn_blocking panicked: {e}")))?
136 }
137
138 #[must_use]
140 pub fn statuses(&self) -> Vec<(String, SubAgentStatus)> {
141 self.agents
142 .values()
143 .map(|h| {
144 let mut status = h.status_rx.borrow().clone();
145 if h.state == SubAgentState::Canceled {
146 status.state = SubAgentState::Canceled;
147 }
148 (h.task_id.clone(), status)
149 })
150 .collect()
151 }
152
153 #[must_use]
155 pub fn agents_def(&self, task_id: &str) -> Option<&SubAgentDef> {
156 self.agents.get(task_id).map(|h| &h.def)
157 }
158
159 #[must_use]
161 pub fn agent_transcript_dir(&self, task_id: &str) -> Option<&std::path::Path> {
162 self.agents
163 .get(task_id)
164 .and_then(|h| h.transcript_dir.as_deref())
165 }
166
167 pub(crate) fn create_transcript_writer(
174 &mut self,
175 config: &SubAgentConfig,
176 task_id: &str,
177 agent_name: &str,
178 resumed_from: Option<&str>,
179 ) -> Option<TranscriptWriter> {
180 if !config.transcript_enabled {
181 return None;
182 }
183 let dir = self.effective_transcript_dir(config);
184 let max_files = self.transcript_max_files;
185 let path = dir.join(format!("{task_id}.jsonl"));
186 let meta = TranscriptMeta {
187 agent_id: task_id.to_owned(),
188 agent_name: agent_name.to_owned(),
189 def_name: agent_name.to_owned(),
190 status: SubAgentState::Submitted,
191 started_at: crate::transcript::utc_now(),
192 finished_at: None,
193 resumed_from: resumed_from.map(str::to_owned),
194 turns_used: 0,
195 mcp_tool_names: Vec::new(),
196 };
197 let task_id = task_id.to_owned();
198 run_blocking(move || {
199 if max_files > 0
200 && let Err(e) = sweep_old_transcripts(&dir, max_files)
201 {
202 tracing::warn!(error = %e, "transcript sweep failed");
203 }
204 match TranscriptWriter::new(&path) {
205 Ok(w) => {
206 if let Err(e) = TranscriptWriter::write_meta(&dir, &task_id, &meta) {
207 tracing::warn!(error = %e, "failed to write initial transcript meta");
208 }
209 Some(w)
210 }
211 Err(e) => {
212 tracing::warn!(error = %e, "failed to create transcript writer");
213 None
214 }
215 }
216 })
217 }
218}
219
220fn run_blocking<T>(f: impl FnOnce() -> T) -> T {
227 if tokio::runtime::Handle::try_current()
228 .is_ok_and(|h| h.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread)
229 {
230 tokio::task::block_in_place(f)
231 } else {
232 f()
233 }
234}