use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use tokio::sync::mpsc;
use zeph_sanitizer::pii::PiiFilter;
use zeph_sanitizer::secret_mask::SecretMaskRegistry;
use zeph_sanitizer::{ContentSanitizer, ContentSource, ContentSourceKind};
use crate::state::SubAgentState;
const FORWARD_CHANNEL_CAPACITY: usize = 128;
const FORWARD_RING_CAPACITY: usize = 200;
const FORWARD_BUFFER_GRACE: Duration = Duration::from_secs(5);
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ForwardSurfaces {
pub tui: bool,
pub bare: bool,
}
impl ForwardSurfaces {
#[must_use]
pub fn any(self) -> bool {
self.tui || self.bare
}
}
#[derive(Debug, Clone)]
pub(crate) struct RawChunk {
task_id: Arc<str>,
def_name: Arc<str>,
seq: u64,
kind: ForwardChunkKind,
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub(crate) enum ForwardChunkKind {
Text(String),
Thinking(String),
Terminal(SubAgentState),
}
#[derive(Debug, Clone)]
pub(crate) struct SanitizedChunk {
pub(crate) task_id: Arc<str>,
pub(crate) def_name: Arc<str>,
pub(crate) seq: u64,
pub(crate) kind: SanitizedChunkKind,
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub(crate) enum SanitizedChunkKind {
Text(String),
Thinking(String),
Terminal(SubAgentState),
}
pub(crate) struct SanitizeLayers {
pub(crate) sanitizer: ContentSanitizer,
pub(crate) secret_registry: Option<Arc<SecretMaskRegistry>>,
pub(crate) pii_filter: Option<PiiFilter>,
}
fn sanitize_text(raw_text: &str, def_name: &str, layers: &SanitizeLayers) -> String {
let source = ContentSource::new(ContentSourceKind::ToolResult).with_identifier(def_name);
let mut body = layers.sanitizer.sanitize(raw_text, source).body;
if let Some(registry) = &layers.secret_registry {
body = registry.mask(&body);
}
if let Some(filter) = &layers.pii_filter {
body = filter.scrub(&body).into_owned();
}
body
}
fn sanitize_chunk(raw: RawChunk, layers: &SanitizeLayers) -> SanitizedChunk {
let kind = match raw.kind {
ForwardChunkKind::Text(text) => {
SanitizedChunkKind::Text(sanitize_text(&text, raw.def_name.as_ref(), layers))
}
ForwardChunkKind::Thinking(text) => {
SanitizedChunkKind::Thinking(sanitize_text(&text, raw.def_name.as_ref(), layers))
}
ForwardChunkKind::Terminal(state) => SanitizedChunkKind::Terminal(state),
};
SanitizedChunk {
task_id: raw.task_id,
def_name: raw.def_name,
seq: raw.seq,
kind,
}
}
pub(crate) struct ForwardSender {
tx: mpsc::Sender<RawChunk>,
task_id: Arc<str>,
def_name: Arc<str>,
seq: AtomicU64,
dropped: AtomicU64,
}
impl ForwardSender {
pub(crate) fn new(tx: mpsc::Sender<RawChunk>, task_id: Arc<str>, def_name: Arc<str>) -> Self {
Self {
tx,
task_id,
def_name,
seq: AtomicU64::new(0),
dropped: AtomicU64::new(0),
}
}
fn try_send(&self, kind: ForwardChunkKind) {
let seq = self.seq.fetch_add(1, Ordering::Relaxed);
let chunk = RawChunk {
task_id: Arc::clone(&self.task_id),
def_name: Arc::clone(&self.def_name),
seq,
kind,
};
if self.tx.try_send(chunk).is_ok() {
tracing::debug!(
task_id = %self.task_id,
seq,
"subagent.forward.emit"
);
} else {
let dropped = self.dropped.fetch_add(1, Ordering::Relaxed) + 1;
tracing::warn!(
task_id = %self.task_id,
seq,
dropped,
"subagent.forward.drop: ingress channel full, chunk dropped"
);
}
}
pub(crate) fn send_text(&self, text: &str) {
if text.is_empty() {
return;
}
self.try_send(ForwardChunkKind::Text(text.to_owned()));
}
pub(crate) fn send_thinking(&self, text: &str) {
if text.is_empty() {
return;
}
self.try_send(ForwardChunkKind::Thinking(text.to_owned()));
}
pub(crate) fn send_terminal(&self, state: SubAgentState) {
tracing::debug!(task_id = %self.task_id, ?state, "subagent.forward.terminal");
self.try_send(ForwardChunkKind::Terminal(state));
}
}
pub(crate) type ForwardBuffer = std::sync::Mutex<HashMap<String, VecDeque<String>>>;
fn display_line(kind: &SanitizedChunkKind) -> Option<String> {
match kind {
SanitizedChunkKind::Text(t) => Some(t.clone()),
SanitizedChunkKind::Thinking(t) => Some(format!("[thinking] {t}")),
SanitizedChunkKind::Terminal(_) => None,
}
}
fn state_str(state: SubAgentState) -> &'static str {
match state {
SubAgentState::Submitted => "submitted",
SubAgentState::Working => "working",
SubAgentState::Completed => "completed",
SubAgentState::Failed => "failed",
SubAgentState::Canceled => "canceled",
}
}
fn emit_bare_line(chunk: &SanitizedChunk) {
#[derive(serde::Serialize)]
struct BareForwardEvent<'a> {
task_id: &'a str,
def_name: &'a str,
seq: u64,
kind: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
content: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
state: Option<&'static str>,
}
let (kind, content, state) = match &chunk.kind {
SanitizedChunkKind::Text(t) => ("text", Some(t.as_str()), None),
SanitizedChunkKind::Thinking(t) => ("thinking", Some(t.as_str()), None),
SanitizedChunkKind::Terminal(s) => ("terminal", None, Some(state_str(*s))),
};
let event = BareForwardEvent {
task_id: &chunk.task_id,
def_name: &chunk.def_name,
seq: chunk.seq,
kind,
content,
state,
};
if let Ok(line) = serde_json::to_string(&event) {
println!("{line}");
}
}
fn dispatch_chunk(chunk: &SanitizedChunk, surfaces: ForwardSurfaces, buffer: &ForwardBuffer) {
if surfaces.tui
&& let Some(line) = display_line(&chunk.kind)
{
let mut guard = buffer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let ring = guard.entry(chunk.task_id.to_string()).or_default();
ring.push_back(line);
while ring.len() > FORWARD_RING_CAPACITY {
ring.pop_front();
}
}
if surfaces.bare {
emit_bare_line(chunk);
}
}
pub(crate) fn new_channel(
task_id: Arc<str>,
def_name: Arc<str>,
) -> (ForwardSender, mpsc::Receiver<RawChunk>) {
let (tx, rx) = mpsc::channel(FORWARD_CHANNEL_CAPACITY);
(ForwardSender::new(tx, task_id, def_name), rx)
}
pub(crate) async fn run_forward_drain(
task_id: Arc<str>,
def_name: Arc<str>,
rx: mpsc::Receiver<RawChunk>,
layers: SanitizeLayers,
surfaces: ForwardSurfaces,
buffer: Arc<ForwardBuffer>,
) {
run_forward_drain_with(
task_id,
def_name,
rx,
layers,
surfaces,
buffer,
dispatch_chunk,
)
.await;
}
async fn run_forward_drain_with(
task_id: Arc<str>,
def_name: Arc<str>,
mut rx: mpsc::Receiver<RawChunk>,
layers: SanitizeLayers,
surfaces: ForwardSurfaces,
buffer: Arc<ForwardBuffer>,
mut dispatch: impl FnMut(&SanitizedChunk, ForwardSurfaces, &ForwardBuffer),
) {
let mut next_seq: u64 = 0;
loop {
if let Some(raw) = rx.recv().await {
next_seq = raw.seq + 1;
let is_terminal = matches!(raw.kind, ForwardChunkKind::Terminal(_));
let chunk = sanitize_chunk(raw, &layers);
dispatch(&chunk, surfaces, &buffer);
if is_terminal {
break;
}
} else {
tracing::warn!(
task_id = %task_id,
"subagent.forward.terminal: ingress channel closed without an explicit \
terminal chunk — synthesizing hard-abort backstop"
);
let synthesized = SanitizedChunk {
task_id: Arc::clone(&task_id),
def_name: Arc::clone(&def_name),
seq: next_seq,
kind: SanitizedChunkKind::Terminal(SubAgentState::Canceled),
};
dispatch(&synthesized, surfaces, &buffer);
break;
}
}
tokio::time::sleep(FORWARD_BUFFER_GRACE).await;
buffer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(task_id.as_ref());
}
pub(crate) fn forwarded_tail(buffer: &ForwardBuffer, task_id: &str, n: usize) -> Vec<String> {
let guard = buffer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
guard.get(task_id).map_or_else(Vec::new, |ring| {
ring.iter().rev().take(n).rev().cloned().collect()
})
}
pub(crate) fn new_buffer() -> Arc<ForwardBuffer> {
Arc::new(std::sync::Mutex::new(HashMap::new()))
}
#[cfg(test)]
mod tests {
use std::sync::atomic::AtomicUsize;
use zeph_config::sanitizer::PiiFilterConfig;
use super::*;
fn layers() -> SanitizeLayers {
SanitizeLayers {
sanitizer: ContentSanitizer::new(&zeph_config::ContentIsolationConfig::default()),
secret_registry: None,
pii_filter: None,
}
}
async fn run_and_count_terminals(
task_id: Arc<str>,
def_name: Arc<str>,
rx: mpsc::Receiver<RawChunk>,
surfaces: ForwardSurfaces,
buffer: Arc<ForwardBuffer>,
) -> usize {
let terminal_dispatches = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&terminal_dispatches);
run_forward_drain_with(
task_id,
def_name,
rx,
layers(),
surfaces,
buffer,
move |chunk, surfaces, buffer| {
if matches!(chunk.kind, SanitizedChunkKind::Terminal(_)) {
counter.fetch_add(1, Ordering::SeqCst);
}
dispatch_chunk(chunk, surfaces, buffer);
},
)
.await;
terminal_dispatches.load(Ordering::SeqCst)
}
#[tokio::test(start_paused = true)]
async fn happy_path_emits_no_spurious_second_terminal() {
let task_id: Arc<str> = Arc::from("task-1");
let def_name: Arc<str> = Arc::from("agent-1");
let (sender, rx) = new_channel(Arc::clone(&task_id), Arc::clone(&def_name));
let buffer = new_buffer();
sender.send_text("hello");
sender.send_terminal(SubAgentState::Completed);
drop(sender);
let terminal_count = run_and_count_terminals(
Arc::clone(&task_id),
def_name,
rx,
ForwardSurfaces {
tui: true,
bare: false,
},
Arc::clone(&buffer),
)
.await;
assert_eq!(
terminal_count, 1,
"exactly one terminal chunk must be dispatched — a second would mean the drain \
looped back to recv() after the explicit terminal (C-new-1 regression)"
);
let tail = forwarded_tail(&buffer, &task_id, 10);
assert!(
tail.is_empty(),
"buffer entry must be evicted after grace window"
);
}
#[tokio::test(start_paused = true)]
async fn hard_abort_without_explicit_terminal_synthesizes_backstop() {
let task_id: Arc<str> = Arc::from("task-2");
let def_name: Arc<str> = Arc::from("agent-2");
let (sender, rx) = new_channel(Arc::clone(&task_id), Arc::clone(&def_name));
let buffer = new_buffer();
sender.send_text("partial output");
drop(sender);
let terminal_count = run_and_count_terminals(
Arc::clone(&task_id),
def_name,
rx,
ForwardSurfaces {
tui: true,
bare: false,
},
buffer,
)
.await;
assert_eq!(
terminal_count, 1,
"exactly one synthesized backstop terminal must be dispatched on hard abort"
);
}
#[tokio::test(start_paused = true)]
async fn zero_consumer_surfaces_still_drains_without_panicking() {
let task_id: Arc<str> = Arc::from("task-3");
let def_name: Arc<str> = Arc::from("agent-3");
let (sender, rx) = new_channel(Arc::clone(&task_id), Arc::clone(&def_name));
let buffer = new_buffer();
sender.send_text("no one is listening");
sender.send_terminal(SubAgentState::Completed);
drop(sender);
run_forward_drain(
task_id,
def_name,
rx,
layers(),
ForwardSurfaces::default(),
buffer,
)
.await;
}
#[tokio::test(start_paused = true)]
async fn secret_registry_masks_forwarded_text_and_thinking() {
use zeph_sanitizer::secret_mask::{SecretCategory, SecretMaskRegistry};
let registry = Arc::new(SecretMaskRegistry::new());
registry.register(
"MY_KEY",
"sk-live-topsecretvalue123",
SecretCategory::ApiKey,
);
let task_id: Arc<str> = Arc::from("task-secret");
let def_name: Arc<str> = Arc::from("agent-secret");
let (sender, rx) = new_channel(Arc::clone(&task_id), Arc::clone(&def_name));
let buffer = new_buffer();
sender.send_text("the key is sk-live-topsecretvalue123, use it wisely");
sender.send_thinking("I will use sk-live-topsecretvalue123 to authenticate");
sender.send_terminal(SubAgentState::Completed);
drop(sender);
let seen = Arc::new(std::sync::Mutex::new(Vec::<SanitizedChunk>::new()));
let collected = Arc::clone(&seen);
let layers = SanitizeLayers {
sanitizer: ContentSanitizer::new(&zeph_config::ContentIsolationConfig::default()),
secret_registry: Some(registry),
pii_filter: None,
};
run_forward_drain_with(
task_id,
def_name,
rx,
layers,
ForwardSurfaces {
tui: true,
bare: false,
},
buffer,
move |chunk, surfaces, buffer| {
collected.lock().unwrap().push(chunk.clone());
dispatch_chunk(chunk, surfaces, buffer);
},
)
.await;
let chunks = seen.lock().unwrap();
for chunk in chunks.iter() {
match &chunk.kind {
SanitizedChunkKind::Text(t) | SanitizedChunkKind::Thinking(t) => {
assert!(
!t.contains("sk-live-topsecretvalue123"),
"forwarded content must not contain the raw secret: {t}"
);
}
SanitizedChunkKind::Terminal(_) => {}
}
}
}
#[tokio::test(start_paused = true)]
async fn pii_filter_scrubs_forwarded_email() {
let task_id: Arc<str> = Arc::from("task-pii");
let def_name: Arc<str> = Arc::from("agent-pii");
let (sender, rx) = new_channel(Arc::clone(&task_id), Arc::clone(&def_name));
let buffer = new_buffer();
sender.send_text("contact me at victim@example.com for details");
sender.send_terminal(SubAgentState::Completed);
drop(sender);
let seen = Arc::new(std::sync::Mutex::new(Vec::<SanitizedChunk>::new()));
let collected = Arc::clone(&seen);
let layers = SanitizeLayers {
sanitizer: ContentSanitizer::new(&zeph_config::ContentIsolationConfig::default()),
secret_registry: None,
pii_filter: Some(PiiFilter::new(PiiFilterConfig::default())),
};
run_forward_drain_with(
task_id,
def_name,
rx,
layers,
ForwardSurfaces {
tui: true,
bare: false,
},
buffer,
move |chunk, surfaces, buffer| {
collected.lock().unwrap().push(chunk.clone());
dispatch_chunk(chunk, surfaces, buffer);
},
)
.await;
let chunks = seen.lock().unwrap();
let text_chunk = chunks
.iter()
.find(|c| matches!(c.kind, SanitizedChunkKind::Text(_)))
.expect("one text chunk must have been dispatched");
let SanitizedChunkKind::Text(ref t) = text_chunk.kind else {
unreachable!()
};
assert!(
!t.contains("victim@example.com"),
"forwarded content must not contain the raw email address: {t}"
);
}
#[tokio::test(start_paused = true)]
async fn buffer_entry_survives_during_grace_window_then_evicted() {
let task_id: Arc<str> = Arc::from("task-grace");
let def_name: Arc<str> = Arc::from("agent-grace");
let (sender, rx) = new_channel(Arc::clone(&task_id), Arc::clone(&def_name));
let buffer = new_buffer();
sender.send_text("visible during the grace window");
sender.send_terminal(SubAgentState::Completed);
drop(sender);
let drain_buffer = Arc::clone(&buffer);
let drain_task_id = Arc::clone(&task_id);
let handle = tokio::spawn(run_forward_drain(
drain_task_id,
def_name,
rx,
layers(),
ForwardSurfaces {
tui: true,
bare: false,
},
drain_buffer,
));
tokio::time::advance(Duration::from_millis(1)).await;
tokio::task::yield_now().await;
let mid_window_tail = forwarded_tail(&buffer, &task_id, 10);
assert_eq!(
mid_window_tail.len(),
1,
"exactly one forwarded line expected"
);
assert!(
mid_window_tail[0].contains("visible during the grace window"),
"the transcript must still be visible during the grace window, got: {:?}",
mid_window_tail[0]
);
tokio::time::advance(FORWARD_BUFFER_GRACE + Duration::from_millis(1)).await;
handle.await.expect("drain task must not panic");
let post_eviction_tail = forwarded_tail(&buffer, &task_id, 10);
assert!(
post_eviction_tail.is_empty(),
"buffer entry must be evicted once the grace window elapses"
);
}
#[test]
fn empty_text_is_not_sent() {
let task_id: Arc<str> = Arc::from("task-4");
let def_name: Arc<str> = Arc::from("agent-4");
let (sender, mut rx) = new_channel(task_id, def_name);
sender.send_text("");
sender.send_thinking("");
drop(sender);
assert!(
rx.try_recv().is_err(),
"empty text/thinking must not be sent onto the ingress channel"
);
}
#[test]
fn channel_full_increments_drop_counter_and_does_not_panic() {
let task_id: Arc<str> = Arc::from("task-5");
let def_name: Arc<str> = Arc::from("agent-5");
let (sender, mut rx) = new_channel(task_id, def_name);
for i in 0..FORWARD_CHANNEL_CAPACITY + 10 {
sender.send_text(&format!("chunk {i}"));
}
let mut received = 0;
while rx.try_recv().is_ok() {
received += 1;
}
assert!(
received > 0,
"at least some chunks must have been delivered"
);
assert!(
received <= FORWARD_CHANNEL_CAPACITY,
"received must never exceed channel capacity"
);
}
#[test]
fn forward_surfaces_any() {
assert!(!ForwardSurfaces::default().any());
assert!(
ForwardSurfaces {
tui: true,
bare: false
}
.any()
);
assert!(
ForwardSurfaces {
tui: false,
bare: true
}
.any()
);
}
}