From 210a9a99bacb37408cad45ac5fd4348c9df6e9be Mon Sep 17 00:00:00 2001 From: "Andrei G." Date: Tue, 14 Jul 2026 22:29:30 +0200 Subject: [PATCH] fix(core): parallelize PreToolUse hooks, backfill extraction, and dedupe MAGE signal mapping Three independent architecture-review findings in the agent turn loop: PreToolUse hooks fired sequentially per tier (mirrors the already-fixed PostToolUse twin, #6128); graph_backfill extracted entities strictly sequentially despite the store's per-entity upsert already making concurrent extraction safe; begin_turn re-derived the MAGE signal type from a raw code instead of matching the already-computed RiskSignal enum. Closes #6259 Closes #6261 Closes #6272 --- CHANGELOG.md | 17 + .../zeph-core/src/agent/agent_access_impl.rs | 326 +++++++++++++++--- crates/zeph-core/src/agent/mod.rs | 32 +- .../agent/tests/mage_signal_mapping_tests.rs | 136 ++++++++ crates/zeph-core/src/agent/tests/mod.rs | 2 + .../src/agent/tool_execution/tests/mod.rs | 1 + .../tests/pre_tool_use_concurrency_tests.rs | 173 ++++++++++ .../src/agent/tool_execution/tier_loop.rs | 158 +++++---- 8 files changed, 719 insertions(+), 126 deletions(-) create mode 100644 crates/zeph-core/src/agent/tests/mage_signal_mapping_tests.rs create mode 100644 crates/zeph-core/src/agent/tool_execution/tests/pre_tool_use_concurrency_tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 99ee2dced..cfa3b6290 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -210,6 +210,23 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Fixed +- **core**: `build_tier_call_futures` fired each tier's `PreToolUse` hooks sequentially, + adding `N × hook_latency` of purely serial blocking on the agent turn loop before the + tier's already-parallelized tool execution even began — the same defect class already + fixed for the `PostToolUse` side (#6128) but never mirrored to `PreToolUse`. Hooks now + fire concurrently, bounded by the tier semaphore, with the per-call invariant preserved + (each call's own hook still fires before that call's own gate check) (#6259). +- **core**: `AgentAccess::graph_backfill` extracted entities/edges from each unprocessed + message strictly sequentially — one LLM call plus SQLite/Qdrant write at a time — despite + the store's `UNIQUE(canonical_name, entity_type)` upsert already making concurrent + extraction across messages safe. Now uses `futures::stream::iter(...).buffer_unordered(4)`, + matching the existing `semantic_scan_plugin_add` pattern, cutting backfill wall time + roughly 4x with no correctness change (#6261). +- **core**: `Agent::begin_turn` re-derived the MAGE `(AuditSignalType, Severity)` pair from + the raw trajectory-signal `u8` code via an independent hand-rolled match, duplicating the + code-to-meaning table already authoritative in `RiskSignal::from_code` — the two tables + were not compiler-coupled and could silently drift. Now matches on the already-computed + `RiskSignal` enum value instead; zero behavior change (#6272). - **Security (`ShadowSentinel`)**: `check_tool_call` awaited its two pre-tool-dispatch DB reads (`get_trajectory`, `get_tool_history`) with no timeout, so a stalled DB connection (e.g. a slow/unresponsive Postgres backend) could block dispatch of every `Shell`/`FileWrite`/ diff --git a/crates/zeph-core/src/agent/agent_access_impl.rs b/crates/zeph-core/src/agent/agent_access_impl.rs index ceb3aa49f..2f81cc248 100644 --- a/crates/zeph-core/src/agent/agent_access_impl.rs +++ b/crates/zeph-core/src/agent/agent_access_impl.rs @@ -599,62 +599,91 @@ impl AgentAccess for Agent { let ids: Vec = messages.iter().map(|(id, _)| *id).collect(); - for (_id, content) in &messages { - if content.trim().is_empty() { - continue; - } - let extraction_cfg = GraphExtractionConfig { - max_entities: graph_cfg.max_entities_per_message, - max_edges: graph_cfg.max_edges_per_message, - extraction_timeout_secs: graph_cfg.extraction_timeout_secs, - community_refresh_interval: 0, - expired_edge_retention_days: graph_cfg.expired_edge_retention_days, - max_entities_cap: graph_cfg.max_entities, - community_summary_max_prompt_bytes: graph_cfg - .community_summary_max_prompt_bytes, - community_summary_concurrency: graph_cfg.community_summary_concurrency, - lpa_edge_chunk_size: graph_cfg.lpa_edge_chunk_size, - note_linking: zeph_memory::NoteLinkingConfig::default(), - link_weight_decay_lambda: graph_cfg.link_weight_decay_lambda, - link_weight_decay_interval_secs: graph_cfg - .link_weight_decay_interval_secs, - belief_revision_enabled: graph_cfg.belief_revision.enabled, - belief_revision_similarity_threshold: graph_cfg - .belief_revision - .similarity_threshold, - conversation_id: None, - apex_mem_enabled: graph_cfg.apex_mem.enabled, - llm_timeout_secs: graph_cfg.llm_timeout_secs, - embed_timeout_secs, - turn_index: None, - write_gate_min_relevance: graph_cfg - .write_gate - .enabled - .then_some(graph_cfg.write_gate.min_edge_relevance), - benna_fast_rate: graph_cfg.spreading_activation.benna_fast_rate, - benna_slow_rate: graph_cfg.spreading_activation.benna_slow_rate, - provenance: None, - system_prompt: None, - recall_include_imported: graph_cfg.recall_include_imported, - }; - let pool = store.pool().clone(); - match extract_and_store( - content.clone(), - vec![], - provider.clone(), - pool, - extraction_cfg, - None, - None, - ) - .await - { - Ok(result) => { - total_entities += result.stats.entities_upserted; - total_edges += result.stats.edges_inserted; - } - Err(e) => { - tracing::warn!("backfill extraction error: {e:#}"); + // extraction_cfg is loop-invariant (derived only from graph_cfg / + // embed_timeout_secs, never from message content), so it is built once per + // batch and cloned per message below. + let extraction_cfg = GraphExtractionConfig { + max_entities: graph_cfg.max_entities_per_message, + max_edges: graph_cfg.max_edges_per_message, + extraction_timeout_secs: graph_cfg.extraction_timeout_secs, + community_refresh_interval: 0, + expired_edge_retention_days: graph_cfg.expired_edge_retention_days, + max_entities_cap: graph_cfg.max_entities, + community_summary_max_prompt_bytes: graph_cfg + .community_summary_max_prompt_bytes, + community_summary_concurrency: graph_cfg.community_summary_concurrency, + lpa_edge_chunk_size: graph_cfg.lpa_edge_chunk_size, + note_linking: zeph_memory::NoteLinkingConfig::default(), + link_weight_decay_lambda: graph_cfg.link_weight_decay_lambda, + link_weight_decay_interval_secs: graph_cfg.link_weight_decay_interval_secs, + belief_revision_enabled: graph_cfg.belief_revision.enabled, + belief_revision_similarity_threshold: graph_cfg + .belief_revision + .similarity_threshold, + conversation_id: None, + apex_mem_enabled: graph_cfg.apex_mem.enabled, + llm_timeout_secs: graph_cfg.llm_timeout_secs, + embed_timeout_secs, + turn_index: None, + write_gate_min_relevance: graph_cfg + .write_gate + .enabled + .then_some(graph_cfg.write_gate.min_edge_relevance), + benna_fast_rate: graph_cfg.spreading_activation.benna_fast_rate, + benna_slow_rate: graph_cfg.spreading_activation.benna_slow_rate, + provenance: None, + system_prompt: None, + recall_include_imported: graph_cfg.recall_include_imported, + }; + + // Extract concurrently, bounded to 4 in-flight — matches + // semantic_scan_plugin_add's existing batched-LLM-call bound. Safe because + // `extract_and_store` builds a fresh `EntityResolver` per call (its + // `lock_name` guard does not span calls), so the actual concurrency-safety + // mechanism is the DB-level `UNIQUE(canonical_name, entity_type)` constraint + // and `ON CONFLICT ... DO UPDATE ... RETURNING id` upsert in + // `GraphStore::upsert_entity` (plus `add_alias`'s `INSERT OR IGNORE`), which + // makes concurrent entity creation for the same name idempotent regardless of + // in-process locking. + { + use futures::stream::StreamExt as _; + + let extraction_futs: Vec<_> = messages + .iter() + .filter_map(|(_id, content)| { + if content.trim().is_empty() { + return None; + } + let content = content.clone(); + let provider = provider.clone(); + let pool = store.pool().clone(); + let extraction_cfg = extraction_cfg.clone(); + Some(extract_and_store( + content, + vec![], + provider, + pool, + extraction_cfg, + None, + None, + )) + }) + .collect(); + + let results: Vec<_> = futures::stream::iter(extraction_futs) + .buffer_unordered(4) + .collect() + .await; + + for result in results { + match result { + Ok(result) => { + total_entities += result.stats.entities_upserted; + total_edges += result.stats.edges_inserted; + } + Err(e) => { + tracing::warn!("backfill extraction error: {e:#}"); + } } } } @@ -2486,6 +2515,191 @@ mod tests { ); } + // #6261: graph_backfill extracts each batch's unprocessed messages concurrently via + // `futures::stream::iter(...).buffer_unordered(4)` instead of a sequential per-message + // loop. buffer_unordered completes futures in an order that need not match input order, so + // this test asserts on aggregate totals (immune to completion order) and on per-entity / + // per-message presence, proving the concurrent rewrite neither drops nor double-counts + // results relative to the pre-#6261 sequential behavior. + #[tokio::test] + async fn graph_backfill_concurrent_extraction_aggregates_stats_without_dropping_results() { + let n = 6; + let cfg = crate::config::GraphConfig { + enabled: true, + ..Default::default() + }; + let mut memory = memory_without_qdrant().await; + let store = install_graph_store(&mut memory); + let cid = memory.sqlite().create_conversation().await.unwrap(); + + for i in 0..n { + sqlx::query(zeph_db::sql!( + "INSERT INTO messages (conversation_id, role, content) VALUES (?1, 'user', ?2)" + )) + .bind(cid.0) + .bind(format!("message body {i}")) + .execute(memory.sqlite().pool()) + .await + .unwrap(); + } + + // One canned extraction response per message, each yielding exactly one distinct + // entity. MockProvider serves responses in call order (not message order), which + // mirrors buffer_unordered's out-of-order completion. + let responses: Vec = (0..n) + .map(|i| { + format!( + r#"{{"entities":[{{"name":"Entity{i}","type":"concept","summary":""}}],"edges":[]}}"# + ) + }) + .collect(); + + let mut agent = Agent::new( + mock_provider(responses), + MockChannel::new(vec![]), + create_test_registry(), + None, + 5, + MockToolExecutor::no_tools(), + ) + .with_memory(std::sync::Arc::new(memory), cid, 50, 5, 100) + .with_graph_config(cfg); + + let mut progress = vec![]; + let result = agent + .graph_backfill(None, &mut |msg| progress.push(msg)) + .await + .unwrap(); + + assert!( + result.contains(&format!("{n} entities")), + "expected all {n} entities aggregated in the result, got: {result}" + ); + assert!( + result.contains(&format!("from {n} messages")), + "expected all {n} messages counted as processed, got: {result}" + ); + + // No drops/double-counts at the store level: every entity must be present exactly once. + for i in 0..n { + let name = format!("entity{i}"); + let found = store + .find_entity(&name, zeph_memory::EntityType::Concept) + .await + .unwrap(); + assert!(found.is_some(), "entity{i} must have been upserted"); + } + + // Every message in the batch must be marked processed — none left behind by a + // buffer_unordered future that was dropped or never polled to completion. + let remaining = store.unprocessed_message_count().await.unwrap(); + assert_eq!(remaining, 0, "all messages must be marked graph_processed"); + } + + // #6261 follow-up (impl-critic finding): the aggregation test above uses an in-memory + // SQLite database, which `zeph-db`'s pool forces to a single connection + // (`connect_sqlite`'s `effective_max = if path == ":memory:" { 1 }`, see + // `crates/zeph-db/src/pool.rs`) — so it never actually exercises concurrent writers racing + // for the SQLite write lock. This test uses a real file-backed database instead (default + // pool_size = 5, WAL journal mode + 5s busy_timeout — see `DbConfig::connect_sqlite`) with + // more unprocessed messages than the `buffer_unordered(4)` bound, so multiple pooled + // connections genuinely contend for writes concurrently. It confirms `extract_and_store`'s + // upserts — relying on WAL mode + busy_timeout + `EntityResolver`'s per-entity-name locking, + // the same assumption `semantic_scan_plugin_add`'s existing `buffer_unordered(4)` usage + // relies on — complete without a "database is locked" error under real multi-connection + // write contention. + #[tokio::test] + async fn graph_backfill_concurrent_extraction_survives_real_sqlite_write_contention() { + let n = 8; + let tmp = tempfile::NamedTempFile::new().expect("tempfile"); + let path = tmp.path().to_str().expect("valid utf-8 path").to_owned(); + + let cfg = crate::config::GraphConfig { + enabled: true, + ..Default::default() + }; + let mut memory = SemanticMemory::new( + &path, + "http://127.0.0.1:1", + None, + zeph_llm::any::AnyProvider::Mock(zeph_llm::mock::MockProvider::default()), + "test-model", + ) + .await + .unwrap(); + let store = install_graph_store(&mut memory); + let cid = memory.sqlite().create_conversation().await.unwrap(); + + for i in 0..n { + sqlx::query(zeph_db::sql!( + "INSERT INTO messages (conversation_id, role, content) VALUES (?1, 'user', ?2)" + )) + .bind(cid.0) + .bind(format!("contention message body {i}")) + .execute(memory.sqlite().pool()) + .await + .unwrap(); + } + + // A small per-call delay forces the (up to 4) concurrently in-flight extraction futures + // to genuinely overlap their subsequent SQLite writes, rather than happening to resolve + // one at a time fast enough to never actually race. + let responses: Vec = (0..n) + .map(|i| { + format!( + r#"{{"entities":[{{"name":"ContentionEntity{i}","type":"concept","summary":""}}],"edges":[]}}"# + ) + }) + .collect(); + let mut provider = zeph_llm::mock::MockProvider::with_responses(responses); + provider.delay_ms = 15; + let provider = zeph_llm::any::AnyProvider::Mock(provider); + + let mut agent = Agent::new( + provider, + MockChannel::new(vec![]), + create_test_registry(), + None, + 5, + MockToolExecutor::no_tools(), + ) + .with_memory(std::sync::Arc::new(memory), cid, 50, 5, 100) + .with_graph_config(cfg); + + let mut progress = vec![]; + let result = agent + .graph_backfill(None, &mut |msg| progress.push(msg)) + .await + .unwrap(); + + assert!( + result.contains(&format!("{n} entities")), + "expected all {n} entities aggregated despite concurrent SQLite writers, got: {result}" + ); + + // The decisive assertion: if a concurrent writer had hit "database is locked" + // (SQLITE_BUSY surfacing as an error instead of the busy_timeout retry succeeding), + // extract_and_store logs a warning and skips that message's upsert (the `Err(e) => + // tracing::warn!(...)` arm in graph_backfill) rather than failing the whole batch — so a + // missing entity here is the observable symptom of exactly the failure mode flagged. + for i in 0..n { + let name = format!("contentionentity{i}"); + let found = store + .find_entity(&name, zeph_memory::EntityType::Concept) + .await + .unwrap(); + assert!( + found.is_some(), + "entity {i} must have been upserted; a missing entity indicates a dropped/failed \ + concurrent write (e.g. a 'database is locked' error) under real multi-connection \ + contention" + ); + } + + let remaining = store.unprocessed_message_count().await.unwrap(); + assert_eq!(remaining, 0, "all messages must be marked graph_processed"); + } + // R-4139: graph_entities with enabled graph but no store (Qdrant unreachable) must // report unavailable, not panic or hang. #[tokio::test] diff --git a/crates/zeph-core/src/agent/mod.rs b/crates/zeph-core/src/agent/mod.rs index 2f97c3f5e..adc9862df 100644 --- a/crates/zeph-core/src/agent/mod.rs +++ b/crates/zeph-core/src/agent/mod.rs @@ -942,6 +942,7 @@ impl Agent { // Spec 050 §2: drain pending risk signals from executor layers before advancing. // Also advance MAGE accumulator (spec 004-16 FR-009) and ingest mapped signals. { + use crate::agent::trajectory::{RiskSignal, VigilRiskLevel}; use zeph_memory::shadow::{AuditSignalType as MageSignal, Severity as MageSev}; let pending: Vec = { let mut q = self.services.security.trajectory_signal_queue.lock(); @@ -949,17 +950,26 @@ impl Agent { }; self.services.security.mage_accumulator.advance_turn(); for code in pending { - self.services - .security - .trajectory - .record(crate::agent::trajectory::RiskSignal::from_code(code)); - // Map signal codes to MAGE AuditSignalType + Severity (spec 004-16 FR-002, FR-007). - // Code 1=PolicyDeny, 6=VigilMedium, 7=VigilHigh, 2=ExfiltrationRedaction. - let mage_signal: Option<(MageSignal, MageSev)> = match code { - 1 => Some((MageSignal::PolicyViolation, MageSev::Medium)), - 2 => Some((MageSignal::ToolChainAnomaly, MageSev::Medium)), - 6 => Some((MageSignal::PromptInjectionPattern, MageSev::Medium)), - 7 => Some((MageSignal::PromptInjectionPattern, MageSev::High)), + let signal = RiskSignal::from_code(code); + self.services.security.trajectory.record(signal); + // Map RiskSignal to MAGE AuditSignalType + Severity (spec 004-16 FR-002, FR-007). + // Matching on the already-decoded `RiskSignal` (rather than the raw `code`) + // keeps this in sync with `RiskSignal::from_code`, the single source of truth + // for the code-to-meaning table. Only the four spec-004-16 signal classes have a + // MAGE equivalent; the remaining RiskSignal variants (OutOfScope, PiiRedaction, + // ToolFailure, HighCallRate, UnusualReadVolume, ToolPairTransition, and + // VigilFlagged(Low)) are trajectory-only and intentionally not surfaced to MAGE. + let mage_signal: Option<(MageSignal, MageSev)> = match signal { + RiskSignal::PolicyDeny => Some((MageSignal::PolicyViolation, MageSev::Medium)), + RiskSignal::ExfiltrationRedaction => { + Some((MageSignal::ToolChainAnomaly, MageSev::Medium)) + } + RiskSignal::VigilFlagged(VigilRiskLevel::Medium) => { + Some((MageSignal::PromptInjectionPattern, MageSev::Medium)) + } + RiskSignal::VigilFlagged(VigilRiskLevel::High) => { + Some((MageSignal::PromptInjectionPattern, MageSev::High)) + } _ => None, }; if let Some((sig, sev)) = mage_signal { diff --git a/crates/zeph-core/src/agent/tests/mage_signal_mapping_tests.rs b/crates/zeph-core/src/agent/tests/mage_signal_mapping_tests.rs new file mode 100644 index 000000000..3b75a06c1 --- /dev/null +++ b/crates/zeph-core/src/agent/tests/mage_signal_mapping_tests.rs @@ -0,0 +1,136 @@ +// SPDX-FileCopyrightText: 2026 Andrei G +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Tests for #6272: `Agent::begin_turn` maps drained `RiskSignal`s to MAGE +//! `(AuditSignalType, Severity)` pairs by matching on the already-decoded `RiskSignal` enum +//! rather than re-deriving the mapping from the raw `u8` signal code. These tests pin the +//! resulting mapping table (spec 004-16 FR-002/FR-007) so a future refactor of either +//! `RiskSignal::from_code` or the MAGE match arm cannot silently desync the two. + +use zeph_config::TrajectoryRiskAccumulatorConfig; +use zeph_memory::shadow::{AuditSignalType, Severity}; + +use crate::agent::agent_tests::{ + MockChannel, MockToolExecutor, create_test_registry, mock_provider, +}; +use crate::agent::turn::TurnInput; +use crate::agent::{Agent, trajectory::RiskSignal}; + +fn make_agent_with_mage() -> Agent { + let agent = Agent::new( + mock_provider(vec![]), + MockChannel::new(vec![]), + create_test_registry(), + None, + 5, + MockToolExecutor::no_tools(), + ); + agent.with_mage_accumulator_config(TrajectoryRiskAccumulatorConfig { + enabled: true, + ..Default::default() + }) +} + +/// Push a raw signal code into the trajectory queue the same way `RiskSignalSink` callbacks +/// do, then drive one turn so `begin_turn` drains and maps it. +fn drain_one_code(agent: &mut Agent, code: u8) { + agent + .services + .security + .trajectory_signal_queue + .lock() + .push(code); + let _turn = agent.begin_turn(TurnInput::new("hi".to_owned(), vec![])); +} + +/// Codes 1, 2, 6, 7 (`PolicyDeny`, `ExfiltrationRedaction`, `VigilFlagged(Medium)`, +/// `VigilFlagged(High)`) are the only `RiskSignal` variants with a MAGE equivalent +/// (spec 004-16 FR-002). Each must ingest into `mage_accumulator` with the exact +/// `AuditSignalType`/`Severity` pair documented at the match site in `begin_turn`. +#[test] +fn begin_turn_maps_known_risk_codes_to_mage_signals() { + let cases: [(u8, AuditSignalType, Severity); 4] = [ + (1, AuditSignalType::PolicyViolation, Severity::Medium), + (2, AuditSignalType::ToolChainAnomaly, Severity::Medium), + (6, AuditSignalType::PromptInjectionPattern, Severity::Medium), + (7, AuditSignalType::PromptInjectionPattern, Severity::High), + ]; + + for (code, expected_type, expected_severity) in cases { + let mut agent = make_agent_with_mage(); + drain_one_code(&mut agent, code); + + assert!( + agent.services.security.mage_accumulator.current_risk() > 0.0, + "code {code} must ingest a non-zero-weight MAGE signal" + ); + let top = agent.services.security.mage_accumulator.top_signals(1); + assert_eq!( + top.len(), + 1, + "code {code} must record exactly one MAGE signal event" + ); + assert_eq!( + top[0].signal_type, expected_type, + "code {code} mapped to the wrong AuditSignalType" + ); + assert_eq!( + top[0].severity, expected_severity, + "code {code} mapped to the wrong Severity" + ); + } +} + +/// The remaining `RiskSignal` variants — `OutOfScope` (3), `PiiRedaction` (4), +/// `ToolFailure` (5), and the `VigilFlagged(Low)` fallback (any unmapped code, e.g. 99) — +/// are trajectory-only per the doc comment above the match in `begin_turn` and must NOT +/// surface to MAGE. +#[test] +fn begin_turn_no_mage_signal_for_trajectory_only_codes() { + for code in [3u8, 4, 5, 99] { + let mut agent = make_agent_with_mage(); + drain_one_code(&mut agent, code); + + // trajectory_risk only ever accumulates non-negative contributions, so `<= 0.0` is + // equivalent to `== 0.0` here without tripping clippy::float_cmp on exact equality. + assert!( + agent.services.security.mage_accumulator.current_risk() <= 0.0, + "code {code} must not ingest any MAGE signal (trajectory-only)" + ); + assert!( + agent + .services + .security + .mage_accumulator + .top_signals(1) + .is_empty(), + "code {code} must leave MAGE signal history empty" + ); + } +} + +/// Sanity guard: `RiskSignal::from_code` itself must still decode these codes to the +/// variants this test file assumes — if this fails, the MAGE-mapping tests above are +/// exercising the wrong `RiskSignal`, not the mapping logic. +#[test] +fn risk_signal_from_code_matches_assumed_variants() { + use crate::agent::trajectory::VigilRiskLevel; + + assert_eq!(RiskSignal::from_code(1), RiskSignal::PolicyDeny); + assert_eq!(RiskSignal::from_code(2), RiskSignal::ExfiltrationRedaction); + assert_eq!( + RiskSignal::from_code(6), + RiskSignal::VigilFlagged(VigilRiskLevel::Medium) + ); + assert_eq!( + RiskSignal::from_code(7), + RiskSignal::VigilFlagged(VigilRiskLevel::High) + ); + assert_eq!(RiskSignal::from_code(3), RiskSignal::OutOfScope); + assert_eq!(RiskSignal::from_code(4), RiskSignal::PiiRedaction); + assert_eq!(RiskSignal::from_code(5), RiskSignal::ToolFailure); + assert_eq!( + RiskSignal::from_code(99), + RiskSignal::VigilFlagged(VigilRiskLevel::Low) + ); +} diff --git a/crates/zeph-core/src/agent/tests/mod.rs b/crates/zeph-core/src/agent/tests/mod.rs index 946776cd5..bca394f0e 100644 --- a/crates/zeph-core/src/agent/tests/mod.rs +++ b/crates/zeph-core/src/agent/tests/mod.rs @@ -20,6 +20,8 @@ mod flush_orphaned_tests; #[cfg(test)] mod inline_tool_loop_tests; #[cfg(test)] +mod mage_signal_mapping_tests; +#[cfg(test)] mod pre_execution_audit_tests; #[cfg(test)] mod provider_override_masking_tests; diff --git a/crates/zeph-core/src/agent/tool_execution/tests/mod.rs b/crates/zeph-core/src/agent/tool_execution/tests/mod.rs index 88e9c5228..89454aa53 100644 --- a/crates/zeph-core/src/agent/tool_execution/tests/mod.rs +++ b/crates/zeph-core/src/agent/tool_execution/tests/mod.rs @@ -8,6 +8,7 @@ mod hook_block_cap_tests; mod mage_escalation_tests; mod native_tests; mod parallel_and_handle_tests; +mod pre_tool_use_concurrency_tests; mod pure_helpers_tests; mod retry_and_skill_env_tests; mod sanitize_and_native_tests; diff --git a/crates/zeph-core/src/agent/tool_execution/tests/pre_tool_use_concurrency_tests.rs b/crates/zeph-core/src/agent/tool_execution/tests/pre_tool_use_concurrency_tests.rs new file mode 100644 index 000000000..56cd3a420 --- /dev/null +++ b/crates/zeph-core/src/agent/tool_execution/tests/pre_tool_use_concurrency_tests.rs @@ -0,0 +1,173 @@ +// SPDX-FileCopyrightText: 2026 Andrei G +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Tests for #6259: `build_tier_call_futures` must fire `PreToolUse` hooks for every tool +//! call in a tier concurrently (Phase 1), not serially, before running the sequential +//! per-index gate-check loop (Phase 2). Mirrors `apply_tier_results_tests.rs`'s coverage of +//! the already-fixed `PostToolUse`/`RuntimeLayer::after_tool` twin (#6128). + +use std::time::{Duration, Instant}; + +use zeph_config::{HookAction, HookDef, HookMatcher}; +use zeph_llm::provider::{Message, MessagePart, Role, ToolUseRequest}; + +use crate::agent::agent_tests::{ + MockChannel, MockToolExecutor, create_test_registry, mock_provider, +}; + +fn make_tool_use_request(id: &str, name: &str) -> ToolUseRequest { + ToolUseRequest { + id: id.into(), + name: name.into(), + input: serde_json::json!({}), + } +} + +fn sleep_hook(secs: f64) -> HookDef { + HookDef { + action: HookAction::Command { + command: format!("sleep {secs}"), + }, + timeout_secs: 5, + fail_closed: false, + r#if: None, + } +} + +/// N tool calls land in a single tier, each matching a `PreToolUse` hook that sleeps. +/// Serial hook dispatch (the pre-#6259 behavior) would take N * delay; concurrent dispatch +/// (Phase 1, bounded by the tier semaphore) should stay close to a single delay regardless +/// of N. +#[tokio::test] +async fn pre_tool_use_hooks_fire_concurrently_across_tier_indices() { + let n = 4; + let delay_secs = 0.06; + let delay = Duration::from_millis(60); + + let provider = mock_provider(vec![]); + let channel = MockChannel::new(vec![]); + let registry = create_test_registry(); + let executor = MockToolExecutor::new((0..n).map(|_| Ok(None)).collect()); + let mut agent = crate::agent::Agent::new(provider, channel, registry, None, 5, executor); + agent.runtime.config.timeouts.max_parallel_tools = n; + agent.services.session.hooks_config.pre_tool_use = vec![HookMatcher { + matcher: "noop".to_owned(), + hooks: vec![sleep_hook(delay_secs)], + }]; + agent + .msg + .messages + .push(Message::from_legacy(Role::System, "system")); + + let tool_calls: Vec = (0..n) + .map(|i| make_tool_use_request(&format!("id-{i}"), "noop")) + .collect(); + + let start = Instant::now(); + agent + .handle_native_tool_calls(None, &tool_calls) + .await + .unwrap(); + let elapsed = start.elapsed(); + + assert!( + elapsed < delay * u32::try_from(n).unwrap(), + "PreToolUse hooks appear to have fired serially: took {elapsed:?} for {n} x {delay:?}" + ); + + // Every call must still have proceeded to execution (hook is fail_open and succeeds). + let tool_result_count = agent + .msg + .messages + .iter() + .flat_map(|m| m.parts.iter()) + .filter(|p| matches!(p, MessagePart::ToolResult { .. })) + .count(); + assert_eq!( + tool_result_count, n, + "every tool call must get a persisted ToolResult after its PreToolUse hook fires" + ); +} + +/// Order invariant: a `fail_closed` `PreToolUse` hook block on one tier index must not affect +/// sibling indices in the same tier. Regression guard for the Phase 1 / Phase 2 split — Phase +/// 1 collects all blocked indices into a `HashMap` up front, and Phase 2 must +/// only consult the entry for its own idx. +#[tokio::test] +async fn pre_tool_use_hook_block_on_one_index_does_not_affect_siblings() { + let provider = mock_provider(vec![]); + let channel = MockChannel::new(vec![]); + let registry = create_test_registry(); + // Only one call ("read") will actually reach the executor; "shell" is blocked before + // dispatch by its fail_closed PreToolUse hook. + let executor = MockToolExecutor::new(vec![Ok(None)]); + let mut agent = crate::agent::Agent::new(provider, channel, registry, None, 5, executor); + agent.runtime.config.timeouts.max_parallel_tools = 2; + agent.services.session.hooks_config.pre_tool_use = vec![HookMatcher { + matcher: "shell".to_owned(), + hooks: vec![HookDef { + action: HookAction::Command { + command: "exit 1".to_owned(), + }, + timeout_secs: 5, + fail_closed: true, + r#if: None, + }], + }]; + agent + .msg + .messages + .push(Message::from_legacy(Role::System, "system")); + + let tool_calls = vec![ + make_tool_use_request("id-shell", "shell"), + make_tool_use_request("id-read", "read"), + ]; + + agent + .handle_native_tool_calls(None, &tool_calls) + .await + .unwrap(); + + assert_eq!( + agent.tool_orchestrator.hook_block_count, 1, + "exactly one call (shell) must be blocked by its own fail_closed hook" + ); + + let tool_results: Vec<(&str, &str)> = agent + .msg + .messages + .iter() + .flat_map(|m| m.parts.iter()) + .filter_map(|p| { + if let MessagePart::ToolResult { + tool_use_id, + content, + .. + } = p + { + Some((tool_use_id.as_str(), content.as_str())) + } else { + None + } + }) + .collect(); + + let shell_result = tool_results + .iter() + .find(|(id, _)| *id == "id-shell") + .expect("shell ToolResult must be present"); + assert!( + shell_result.1.contains("[blocked]"), + "shell call must be blocked by its own hook: {shell_result:?}" + ); + + let read_result = tool_results + .iter() + .find(|(id, _)| *id == "id-read") + .expect("read ToolResult must be present"); + assert!( + !read_result.1.contains("[blocked]"), + "read call has no matching hook and must NOT be blocked by shell's hook: {read_result:?}" + ); +} diff --git a/crates/zeph-core/src/agent/tool_execution/tier_loop.rs b/crates/zeph-core/src/agent/tool_execution/tier_loop.rs index 75e855e2e..8624c93d7 100644 --- a/crates/zeph-core/src/agent/tool_execution/tier_loop.rs +++ b/crates/zeph-core/src/agent/tool_execution/tier_loop.rs @@ -1980,6 +1980,96 @@ impl Agent { .rate_limiter .check_batch(&tier_tool_names); + // Phase 1: fire PreToolUse hooks for every call in the tier concurrently, bounded by the + // same tier semaphore used for tool execution below — mirrors apply_tier_results Phase 2 + // (#6128), which already parallelized the PostToolUse/RuntimeLayer side of this same + // per-index hook-firing pattern. There is no ordering constraint between different + // tool-call indices' hook firing; the only required invariant — a call's own gate checks + // must observe that same call's own hook result — holds because this whole phase + // completes before Phase 2's per-index gate checks below begin. + // + // Focus/compress tools are synthetic internal calls that never reach hooks or gates (see + // the matching `continue` in Phase 2 below), so they are excluded here too. + let pre_hooks = self.services.session.hooks_config.pre_tool_use.clone(); + let mut pre_hook_blocked: std::collections::HashMap = + std::collections::HashMap::new(); + if !pre_hooks.is_empty() { + let conv_id_str = self + .services + .memory + .persistence + .conversation_id + .map(|id| id.0.to_string()); + let dispatch = self.mcp_dispatch(); + let mcp: Option<&dyn zeph_subagent::McpDispatch> = dispatch + .as_ref() + .map(|d| d as &dyn zeph_subagent::McpDispatch); + + let futs = tier_indices.iter().filter_map(|&idx| { + let tc = &tool_calls[idx]; + if tc.name == "compress_context" + || tc.name == "request_compaction" + || (self.services.focus.config.enabled + && (tc.name == "start_focus" || tc.name == "complete_focus")) + { + return None; + } + let matched: Vec<&zeph_config::HookDef> = + zeph_subagent::matching_hooks(&pre_hooks, tc.name.as_str()); + if matched.is_empty() { + return None; + } + let has_fail_closed = matched.iter().any(|h| h.fail_closed); + let owned: Vec = matched.into_iter().cloned().collect(); + let env = make_tool_hook_env(tc.name.as_str(), &tc.input, conv_id_str.as_deref()); + let sem = std::sync::Arc::clone(semaphore); + let tool_name = tc.name.clone(); + Some(async move { + let Ok(_permit) = sem.acquire().await else { + tracing::warn!( + tool = %tool_name, + "semaphore closed during pre-tool hook firing, skipping \ + PreToolUse hook for this call" + ); + return (idx, None); + }; + let result = zeph_subagent::hooks::fire_hooks(&owned, &env, mcp, None) + .instrument(tracing::info_span!( + "core.hooks.pre_tool_use", + tool = %tool_name + )) + .await; + (idx, Some(result.map_err(|e| (e, has_fail_closed)))) + }) + }); + + for (idx, outcome) in futures::future::join_all(futs).await { + let Some(Err((e, has_fail_closed))) = outcome else { + continue; + }; + let tool_name = tool_calls[idx].name.as_str(); + if has_fail_closed { + self.tool_orchestrator.hook_block_count += 1; + tracing::warn!( + error = %e, + tool = %tool_name, + hook_block_count = self.tool_orchestrator.hook_block_count, + hook_block_cap = self.tool_orchestrator.hook_block_cap, + "PreToolUse hook blocked tool (fail_closed)" + ); + pre_hook_blocked.insert( + idx, + format!("[blocked] PreToolUse hook blocked tool `{tool_name}`: {e}"), + ); + } else { + tracing::warn!(error = %e, tool = %tool_name, "PreToolUse hook failed"); + } + } + } + + // Phase 2: per-index gate checks, cache lookups, and execution-future construction. + // Stays sequential — it needs `&mut self` throughout, and each idx's control flow + // depends on that same idx's own PreToolUse hook outcome computed in Phase 1 above. let mut tier_futs: Vec<(usize, ToolExecFut)> = Vec::with_capacity(tier_indices.len()); for (tier_local_idx, &idx) in tier_indices.iter().enumerate() { let tc = &tool_calls[idx]; @@ -1994,65 +2084,15 @@ impl Agent { continue; } - // Fire PreToolUse hooks before any gate check so the hook always observes the - // LLM's tool request, even when a gate (utility, quota, dep, repeat) intercepts it. - // Focus/compress tools are excluded by the early `continue` above — they are - // synthetic internal tools that must never surface to the hook system. - let pre_hooks = self.services.session.hooks_config.pre_tool_use.clone(); - if !pre_hooks.is_empty() { - let matched: Vec<&zeph_config::HookDef> = - zeph_subagent::matching_hooks(&pre_hooks, tc.name.as_str()); - if !matched.is_empty() { - let conv_id_str = self - .services - .memory - .persistence - .conversation_id - .map(|id| id.0.to_string()); - let env = - make_tool_hook_env(tc.name.as_str(), &tc.input, conv_id_str.as_deref()); - let has_fail_closed = matched.iter().any(|h| h.fail_closed); - let owned: Vec = matched.into_iter().cloned().collect(); - let dispatch = self.mcp_dispatch(); - let mcp: Option<&dyn zeph_subagent::McpDispatch> = dispatch - .as_ref() - .map(|d| d as &dyn zeph_subagent::McpDispatch); - if let Err(e) = zeph_subagent::hooks::fire_hooks(&owned, &env, mcp, None) - .instrument(tracing::info_span!( - "core.hooks.pre_tool_use", - tool = %tc.name - )) - .await - { - if has_fail_closed { - self.tool_orchestrator.hook_block_count += 1; - tracing::warn!( - error = %e, - tool = %tc.name, - hook_block_count = self.tool_orchestrator.hook_block_count, - hook_block_cap = self.tool_orchestrator.hook_block_cap, - "PreToolUse hook blocked tool (fail_closed)" - ); - let msg = format!( - "[blocked] PreToolUse hook blocked tool `{}`: {e}", - tc.name - ); - tier_futs.push(( - idx, - Box::pin(std::future::ready(Ok(Some(skipped_output( - tc.name.clone(), - msg, - ))))), - )); - continue; - } - tracing::warn!( - error = %e, - tool = %tc.name, - "PreToolUse hook failed" - ); - } - } + if let Some(msg) = pre_hook_blocked.remove(&idx) { + tier_futs.push(( + idx, + Box::pin(std::future::ready(Ok(Some(skipped_output( + tc.name.clone(), + msg, + ))))), + )); + continue; } // Check static gates: dep failure, quota, pre-exec block, utility gate, repeat.