Skip to main content

zeph_subagent/manager/
collect.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use 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    /// Collect the result from a completed sub-agent, removing it from the active set.
20    ///
21    /// Writes a final `TranscriptMeta` sidecar with the terminal state and turn count.
22    ///
23    /// # Errors
24    ///
25    /// Returns [`SubAgentError::NotFound`] if the task ID is unknown,
26    /// [`SubAgentError::Spawn`] if the task panicked.
27    #[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    /// Resolve the effective transcript directory from config or default.
102    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    /// Look up the definition name for a resumable transcript without spawning.
113    ///
114    /// Used by callers that need to resolve skills before calling `resume()`.
115    /// Offloads the blocking FS reads to a `spawn_blocking` thread.
116    ///
117    /// # Errors
118    ///
119    /// Returns the same errors as [`crate::transcript::TranscriptReader::find_by_prefix`] and
120    /// [`crate::transcript::TranscriptReader::load_meta`].
121    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    /// Return a snapshot of all active sub-agent statuses.
139    #[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    /// Return the definition for a specific agent by `task_id`.
154    #[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    /// Return the transcript directory for a specific agent by `task_id`.
160    #[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    /// Create a transcript writer if transcripts are enabled.
168    ///
169    /// All three blocking FS operations (sweep, file open, meta write) are offloaded via
170    /// [`tokio::task::block_in_place`] on multi-thread runtimes so the Tokio executor
171    /// thread is not stalled. Falls back to a direct call on `current_thread` runtimes
172    /// (e.g. single-threaded unit tests) where `block_in_place` would panic.
173    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
220/// Run a blocking closure without stalling the Tokio executor.
221///
222/// On a multi-thread runtime, delegates to [`tokio::task::block_in_place`] so other
223/// tasks can continue running while the blocking work executes. On a `current_thread`
224/// runtime (unit tests, single-threaded entry points) calls the closure directly,
225/// since there is no thread pool to offload to and `block_in_place` would panic.
226fn 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}