zeph_subagent/manager/
secrets.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::collections::HashMap;
5use std::time::Duration;
6
7use zeph_common::secret::Secret;
8use zeph_sanitizer::secret_mask::SecretCategory;
9
10use super::SubAgentManager;
11use crate::error::SubAgentError;
12use crate::grants::{GrantKind, GrantedSecret, SecretRequest};
13
14/// Build the standard hook environment for a sub-agent lifecycle event.
15pub(crate) fn make_hook_env(
16    task_id: &str,
17    agent_name: &str,
18    tool_name: &str,
19) -> HashMap<String, String> {
20    let mut env = HashMap::new();
21    env.insert("ZEPH_AGENT_ID".to_owned(), task_id.to_owned());
22    env.insert("ZEPH_AGENT_NAME".to_owned(), agent_name.to_owned());
23    env.insert("ZEPH_AGENT_TYPE".to_owned(), "subagent".to_owned());
24    env.insert("ZEPH_TOOL_NAME".to_owned(), tool_name.to_owned());
25    env
26}
27
28impl SubAgentManager {
29    /// Approve a secret request for a running sub-agent.
30    ///
31    /// Called after the user approves a vault secret access prompt. The secret
32    /// key must appear in the sub-agent definition's allowed `secrets` list;
33    /// otherwise the request is auto-denied.
34    ///
35    /// # Errors
36    ///
37    /// Returns [`SubAgentError::NotFound`] if the task ID is unknown,
38    /// [`SubAgentError::Invalid`] if the key is not in the definition's allowed list.
39    pub fn approve_secret(
40        &mut self,
41        task_id: &str,
42        secret_key: &str,
43        ttl: Duration,
44    ) -> Result<(), SubAgentError> {
45        let handle = self
46            .agents
47            .get_mut(task_id)
48            .ok_or_else(|| SubAgentError::NotFound(task_id.to_owned()))?;
49
50        handle.grants_lock().sweep_expired();
51
52        if !handle
53            .def
54            .permissions
55            .secrets
56            .iter()
57            .any(|k| k == secret_key)
58        {
59            tracing::warn!(task_id, "secret request denied: key not in allowed list");
60            return Err(SubAgentError::Invalid(format!(
61                "secret is not in the allowed secrets list for '{}'",
62                handle.def.name
63            )));
64        }
65
66        handle.grants_lock().grant_secret(secret_key, ttl);
67        Ok(())
68    }
69
70    /// Deliver a resolved secret value to a waiting sub-agent loop.
71    ///
72    /// Should be called after the user approves the request and the caller has resolved
73    /// `key` to its actual vault value (see [`approve_secret`](Self::approve_secret)).
74    /// Requires an active grant for `key` — delivery is refused if
75    /// [`approve_secret`](Self::approve_secret) was never called or the grant's TTL has
76    /// already elapsed, making
77    /// [`PermissionGrants::is_active`][crate::grants::PermissionGrants::is_active]
78    /// load-bearing rather than unused bookkeeping.
79    ///
80    /// The delivered value is stamped with the grant's expiry (see [`GrantedSecret`]) so the
81    /// sub-agent loop can keep re-validating the TTL locally on every subsequent tool call,
82    /// rather than trusting this one-time gate for the remainder of a long-running turn loop.
83    ///
84    /// # Errors
85    ///
86    /// Returns [`SubAgentError::NotFound`] if the task ID is unknown, or
87    /// [`SubAgentError::Invalid`] if there is no active grant for `key`.
88    pub fn deliver_secret(
89        &mut self,
90        task_id: &str,
91        key: &str,
92        value: Secret,
93    ) -> Result<(), SubAgentError> {
94        let handle = self
95            .agents
96            .get_mut(task_id)
97            .ok_or_else(|| SubAgentError::NotFound(task_id.to_owned()))?;
98
99        let Some(expires_at) = handle
100            .grants_lock()
101            .expires_at(&GrantKind::Secret(key.to_owned()))
102        else {
103            tracing::warn!(
104                task_id,
105                "secret delivery denied: no active grant (missing approval or TTL expired)"
106            );
107            return Err(SubAgentError::Invalid(
108                "no active grant for this secret".to_owned(),
109            ));
110        };
111
112        // Register the delivered value for masking (#6492): closes the forwarding-path leak
113        // (the drain masks against this same registry, which was never populated with the
114        // secret it was supposed to catch) and feeds the agent loop's own tool-result masking
115        // pass, which consults this registry before content reaches the transcript, LLM
116        // context, or debug dump.
117        if let Some(registry) = self.secret_registry.as_ref() {
118            registry.register(key, value.expose(), SecretCategory::from_key_name(key));
119        }
120
121        handle
122            .secret_tx
123            .try_send(Some(GrantedSecret { value, expires_at }))
124            .map_err(|e| SubAgentError::Channel(e.to_string()))
125    }
126
127    /// Deny a pending secret request — sends `None` to unblock the waiting sub-agent loop.
128    ///
129    /// # Errors
130    ///
131    /// Returns [`SubAgentError::NotFound`] if the task ID is unknown,
132    /// [`SubAgentError::Channel`] if the channel is full or closed.
133    pub fn deny_secret(&mut self, task_id: &str) -> Result<(), SubAgentError> {
134        let handle = self
135            .agents
136            .get_mut(task_id)
137            .ok_or_else(|| SubAgentError::NotFound(task_id.to_owned()))?;
138        handle
139            .secret_tx
140            .try_send(None)
141            .map_err(|e| SubAgentError::Channel(e.to_string()))
142    }
143
144    /// Try to receive a pending secret request from any sub-agent (non-blocking).
145    ///
146    /// Polls each active agent's request channel once. Returns `Some((task_id, request))`
147    /// if any agent has a pending request, or `None` if all channels are empty.
148    /// Call this from the main agent loop to surface approval prompts to the user.
149    pub fn try_recv_secret_request(&mut self) -> Option<(String, SecretRequest)> {
150        for handle in self.agents.values_mut() {
151            if let Ok(req) = handle.pending_secret_rx.try_recv() {
152                return Some((handle.task_id.clone(), req));
153            }
154        }
155        None
156    }
157
158    /// Try to receive a pending secret request from one specific sub-agent (non-blocking).
159    ///
160    /// Unlike [`try_recv_secret_request`](Self::try_recv_secret_request), this only polls
161    /// `task_id`'s own request channel, so it never pops and discards an unrelated sibling
162    /// sub-agent's pending request. Use this when the caller already knows which sub-agent
163    /// it wants to act on (e.g. an explicit `/agent approve <id>` command), instead of the
164    /// pop-then-filter pattern of polling [`try_recv_secret_request`](Self::try_recv_secret_request)
165    /// and discarding non-matching results — a discarded result is popped off the channel
166    /// and lost forever, silently starving the sub-agent that actually sent it.
167    ///
168    /// Returns `None` if `task_id` is unknown or has no pending request.
169    pub fn try_recv_secret_request_for(&mut self, task_id: &str) -> Option<SecretRequest> {
170        self.agents
171            .get_mut(task_id)?
172            .pending_secret_rx
173            .try_recv()
174            .ok()
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use std::sync::Arc;
181    use std::time::{Duration, Instant};
182
183    use tokio::sync::{mpsc, watch};
184    use tokio_util::sync::CancellationToken;
185    use zeph_common::secret::Secret;
186    use zeph_sanitizer::secret_mask::SecretMaskRegistry;
187
188    use super::make_hook_env;
189    use crate::def::SubAgentDef;
190    use crate::grants::{GrantedSecret, PermissionGrants};
191    use crate::manager::{SubAgentHandle, SubAgentManager, SubAgentStatus};
192    use crate::state::SubAgentState;
193
194    /// Builds a [`SubAgentHandle`] with a LIVE `secret_tx`/`secret_rx` pair — unlike
195    /// [`SubAgentHandle::for_test`], which immediately drops the receiver (documented as
196    /// valid only for metadata inspection, never for a call that actually sends on the
197    /// channel). `deliver_secret` calls `secret_tx.try_send(..)`, so a dropped receiver
198    /// would make every delivery fail with `Channel("channel closed")` regardless of the
199    /// registration logic under test here.
200    fn handle_with_live_secret_channel(
201        id: &str,
202        def: SubAgentDef,
203    ) -> (SubAgentHandle, mpsc::Receiver<Option<GrantedSecret>>) {
204        let initial_status = SubAgentStatus {
205            state: SubAgentState::Working,
206            last_message: None,
207            turns_used: 0,
208            started_at: Instant::now(),
209        };
210        let (status_tx, status_rx) = watch::channel(initial_status);
211        drop(status_tx);
212        let (pending_secret_rx_tx, pending_secret_rx) = mpsc::channel(1);
213        drop(pending_secret_rx_tx);
214        let (secret_tx, secret_rx) = mpsc::channel(1);
215        let handle = SubAgentHandle {
216            id: id.to_owned(),
217            task_id: id.to_owned(),
218            def,
219            state: SubAgentState::Working,
220            join_handle: None,
221            cancel: CancellationToken::new(),
222            status_rx,
223            grants: std::sync::Arc::new(std::sync::Mutex::new(PermissionGrants::default())),
224            pending_secret_rx,
225            secret_tx,
226            started_at_str: String::new(),
227            transcript_dir: None,
228            mcp_tool_names: Vec::new(),
229        };
230        (handle, secret_rx)
231    }
232
233    #[test]
234    fn make_hook_env_sets_agent_type_subagent() {
235        let env = make_hook_env("task-42", "my-agent", "Shell");
236        assert_eq!(
237            env.get("ZEPH_AGENT_TYPE").map(String::as_str),
238            Some("subagent")
239        );
240        assert_eq!(
241            env.get("ZEPH_AGENT_ID").map(String::as_str),
242            Some("task-42")
243        );
244        assert_eq!(
245            env.get("ZEPH_AGENT_NAME").map(String::as_str),
246            Some("my-agent")
247        );
248        assert_eq!(env.get("ZEPH_TOOL_NAME").map(String::as_str), Some("Shell"));
249    }
250
251    // --- #6492: deliver_secret registers into the shared secret-mask registry ---
252
253    #[test]
254    fn deliver_secret_registers_value_into_secret_mask_registry() {
255        let mut mgr = SubAgentManager::new(4);
256        let registry = Arc::new(SecretMaskRegistry::new());
257        mgr.set_secret_registry(Arc::clone(&registry));
258
259        let (handle, _secret_rx) =
260            handle_with_live_secret_channel("task-1", SubAgentDef::for_test("helper"));
261        handle
262            .grants_lock()
263            .grant_secret("SOME_VAULT_KEY", Duration::from_mins(5));
264        mgr.insert_handle_for_test("task-1".to_owned(), handle);
265
266        mgr.deliver_secret(
267            "task-1",
268            "SOME_VAULT_KEY",
269            Secret::new("the-secret-value-123"),
270        )
271        .expect("delivery must succeed: an active grant exists");
272
273        assert!(
274            registry.would_mask("value is the-secret-value-123"),
275            "a delivered secret must be registered into the shared mask registry — closes \
276             the gap where the forwarding drain masked against a registry that was never \
277             populated with the secret it was supposed to catch"
278        );
279    }
280
281    #[test]
282    fn deliver_secret_without_registry_still_succeeds() {
283        // Regression guard: registering into the mask registry must stay best-effort — a
284        // session with no registry wired (the `None` default) must not lose secret delivery.
285        let mut mgr = SubAgentManager::new(4);
286
287        let (handle, _secret_rx) =
288            handle_with_live_secret_channel("task-1", SubAgentDef::for_test("helper"));
289        handle
290            .grants_lock()
291            .grant_secret("SOME_VAULT_KEY", Duration::from_mins(5));
292        mgr.insert_handle_for_test("task-1".to_owned(), handle);
293
294        let result = mgr.deliver_secret(
295            "task-1",
296            "SOME_VAULT_KEY",
297            Secret::new("the-secret-value-123"),
298        );
299        assert!(
300            result.is_ok(),
301            "delivery must succeed even with no registry wired"
302        );
303    }
304
305    #[test]
306    fn deliver_secret_without_active_grant_is_denied() {
307        // Baseline: delivery must still be refused when there is no active grant, unaffected
308        // by the new registration step (which only runs after the grant check succeeds).
309        let mut mgr = SubAgentManager::new(4);
310        let registry = Arc::new(SecretMaskRegistry::new());
311        mgr.set_secret_registry(Arc::clone(&registry));
312
313        let handle = SubAgentHandle::for_test("task-1", SubAgentDef::for_test("helper"));
314        mgr.insert_handle_for_test("task-1".to_owned(), handle);
315
316        let result = mgr.deliver_secret(
317            "task-1",
318            "SOME_VAULT_KEY",
319            Secret::new("the-secret-value-123"),
320        );
321        assert!(
322            result.is_err(),
323            "delivery without an active grant must be denied"
324        );
325        assert!(
326            !registry.would_mask("value is the-secret-value-123"),
327            "a denied delivery must never register the secret"
328        );
329    }
330}
zeph_subagent/manager/secrets.rs

zeph_subagent/manager/
secrets.rs
