diff --git a/crates/zeph-core/src/agent/mod.rs b/crates/zeph-core/src/agent/mod.rs index 618596f5..895b3ad7 100644 --- a/crates/zeph-core/src/agent/mod.rs +++ b/crates/zeph-core/src/agent/mod.rs @@ -500,6 +500,17 @@ impl Agent Agent match item { + Some(r) => r, + None => break, + }, + () = super::shutdown_signal(&mut self.shutdown) => { + tracing::info!("streaming interrupted by shutdown"); + break; + } + }; let chunk: String = chunk_result?; response.push_str(&chunk); let display = self.maybe_redact(&chunk); @@ -382,6 +392,11 @@ impl Agent = clients.drain().collect(); for (id, client) in drained { tracing::info!(server_id = id, "shutting down MCP client"); - client.shutdown().await; + if tokio::time::timeout(Duration::from_secs(5), client.shutdown()) + .await + .is_err() + { + tracing::warn!(server_id = id, "MCP client shutdown timed out"); + } } } } diff --git a/src/main.rs b/src/main.rs index 22cfcf9a..97bb1a54 100644 --- a/src/main.rs +++ b/src/main.rs @@ -565,9 +565,11 @@ async fn main() -> anyhow::Result<()> { tokio::select! { result = tui_task => { + agent.shutdown().await; return result?; } result = agent_future => { + agent.shutdown().await; return result; } } @@ -576,7 +578,9 @@ async fn main() -> anyhow::Result<()> { warmup_provider(&warmup_provider_clone).await; tokio::spawn(forward_status_to_stderr(status_rx)); // Box::pin avoids clippy::large_futures on non-TUI path - Box::pin(agent.run()).await + let result = Box::pin(agent.run()).await; + agent.shutdown().await; + result } async fn forward_status_to_stderr(mut rx: tokio::sync::mpsc::UnboundedReceiver) {