quorum_rs/agents/

user_tools.rs

use crate::agents::{DeliberationPhase, PendingToolCall, ToolCallStatus, UserToolHandlerTrait};
use crate::nats_utils::ensure_kv_bucket;
use anyhow::Result;
use async_trait::async_trait;
use futures_util::StreamExt;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tracing::{info, warn};
use uuid::Uuid;

/// Escape XML special characters to prevent breaking XML-like wrapper elements.
fn escape_xml(s: &str) -> String {
    s.replace('&', "&")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
}

/// Escape XML attribute values — additionally escapes quotes to prevent attribute injection.
fn escape_xml_attr(s: &str) -> String {
    s.replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
        .replace('"', "&quot;")
        .replace('\'', "&apos;")
}

/// Compute finalization reserve from phase budget and config parameters.
/// Returns the smaller of ratio-based and fixed-cap reserves.
/// Inputs are sanitized: NaN, infinite, and negative values are clamped to 0.0,
/// and `reserve_ratio` is capped at 1.0 to avoid nonsensical reserves.
fn compute_finalization_reserve(
    phase_budget: Duration,
    reserve_secs: f64,
    reserve_ratio: f64,
) -> Duration {
    let safe_secs = if reserve_secs.is_finite() && reserve_secs > 0.0 {
        reserve_secs
    } else {
        0.0
    };
    let safe_ratio = if reserve_ratio.is_finite() && reserve_ratio > 0.0 {
        reserve_ratio.min(1.0)
    } else {
        0.0
    };
    let ratio_based = Duration::from_secs_f64(phase_budget.as_secs_f64() * safe_ratio);
    let fixed = Duration::from_secs_f64(safe_secs);
    ratio_based.min(fixed)
}

/// Name of the KV bucket carrying one session's pending user tool calls.
///
/// The single definition on purpose. The agent writes these records and the
/// orchestrator polls them, and each used to build the name itself — the reader
/// with `nsed` hardcoded, the writer from its configured prefix. They agree only
/// while the prefix is the default: change it and the agent creates a second
/// bucket, writes a question nobody is watching, waits out the whole deadline,
/// and the orchestrator's cleanup deletes the other name and leaks this one.
/// Sharing the function makes that divergence unrepresentable rather than
/// merely detectable.
pub fn toolcalls_bucket_name(subject_prefix: &str, session_id: &str) -> String {
    format!(
        "{}_toolcalls_{}",
        crate::nats_utils::sanitize_subject_component(subject_prefix),
        crate::nats_utils::sanitize_subject_component(session_id)
    )
}

/// Encapsulates the NATS state needed to handle user tool calls within the
/// react loop. Created by the NatsNsedWorker and passed down to the agent.
#[derive(Clone, Debug)]
pub struct UserToolHandler {
    nats_client: async_nats::Client,
    js_context: async_nats::jetstream::Context,
    session_id: String,
    agent_id: String,
    /// NATS subject/bucket prefix (e.g. "nsed"). Avoids hardcoding.
    subject_prefix: String,
    /// Phase start time — used to compute remaining budget dynamically.
    phase_start: Instant,
    /// Phase budget as originally allocated.
    phase_budget: Duration,
    /// Configurable limits
    max_pending_per_agent: usize,
    finalization_reserve_secs: f64,
    finalization_reserve_ratio: f64,
}

enum WaitResult {
    Responded(String),
    Timeout,
    Error(String),
}

impl UserToolHandler {
    pub fn new(
        nats_client: async_nats::Client,
        js_context: async_nats::jetstream::Context,
        session_id: String,
        agent_id: String,
        phase_budget_remaining_secs: f64,
    ) -> Self {
        // Sanitize: NaN and infinite values would panic in Duration::from_secs_f64
        let safe_budget = if phase_budget_remaining_secs.is_finite() {
            phase_budget_remaining_secs.max(0.0)
        } else {
            0.0
        };
        Self {
            nats_client,
            js_context,
            session_id,
            agent_id,
            subject_prefix: "nsed".to_string(),
            phase_start: Instant::now(),
            phase_budget: Duration::from_secs_f64(safe_budget),
            max_pending_per_agent: 3,
            finalization_reserve_secs: 30.0,
            finalization_reserve_ratio: 0.15,
        }
    }

    /// Set the NATS subject/bucket prefix (default: "nsed").
    pub fn with_subject_prefix(mut self, prefix: String) -> Self {
        self.subject_prefix = prefix;
        self
    }

    /// Configure rate limit for pending calls per agent.
    pub fn with_max_pending_per_agent(mut self, max: usize) -> Self {
        self.max_pending_per_agent = max;
        self
    }

    /// Configure finalization reserve parameters.
    pub fn with_finalization_reserve(mut self, secs: f64, ratio: f64) -> Self {
        self.finalization_reserve_secs = secs;
        self.finalization_reserve_ratio = ratio;
        self
    }

    /// Returns the remaining time in this phase, accounting for time already
    /// spent in the react loop since the handler was created.
    fn remaining_budget(&self) -> Duration {
        self.phase_budget.saturating_sub(self.phase_start.elapsed())
    }

    /// Compute the finalization reserve: the time reserved for the agent to
    /// produce a proposal after a tool call, even without a user response.
    fn finalization_reserve(&self) -> Duration {
        compute_finalization_reserve(
            self.phase_budget,
            self.finalization_reserve_secs,
            self.finalization_reserve_ratio,
        )
    }

    fn now_epoch_millis() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64
    }

    fn bucket_name(&self) -> String {
        toolcalls_bucket_name(&self.subject_prefix, &self.session_id)
    }

    /// Handle a user tool call: publish to KV, wait for response, return result.
    pub async fn handle_call(
        &self,
        tool_name: &str,
        arguments_json: &str,
        round: u32,
        phase: DeliberationPhase,
    ) -> String {
        // 1. Refuse before touching anything if there is no time to wait for an
        //    answer. This used to be asked after the bucket, the record and the
        //    SSE event had all gone out, so a question that was already dead left
        //    a trail: a poller could only ever catch it Expired, a live client
        //    rendered and withdrew it in one breath, and an empty bucket outlived
        //    both. The budget needs no NATS to read, so nothing has to happen
        //    before this is known.
        //
        //    Reachable whenever a call lands inside the finalization reserve —
        //    `force_finalize` withdraws these tools at three average iterations of
        //    remaining budget, a shorter window than the reserve when iterations
        //    are quick.
        let remaining = self.remaining_budget();
        let reserve = self.finalization_reserve();
        if remaining <= reserve {
            info!(
                agent = %self.agent_id,
                tool = %tool_name,
                remaining_secs = remaining.as_secs_f64(),
                reserve_secs = reserve.as_secs_f64(),
                "Not asking: no budget beyond the finalization reserve."
            );
            return "[No response — phase budget exhausted. Proceed immediately.]".to_string();
        }

        // 2. Parse arguments — propagate parse error instead of silently defaulting
        let arguments: serde_json::Value = match serde_json::from_str(arguments_json) {
            Ok(v) => v,
            Err(e) => {
                return format!("Error: Invalid JSON arguments: {}", e);
            }
        };

        // 3. Get or create the toolcalls KV bucket
        let bucket_name = self.bucket_name();
        let toolcall_store = match self.get_or_create_bucket(&bucket_name).await {
            Ok(store) => store,
            Err(e) => {
                warn!("Failed to access toolcall bucket: {}", e);
                return format!("Error: Failed to register tool call: {}", e);
            }
        };

        // 3. Rate limiting — count pending calls for this agent
        match self.count_pending_for_agent(&toolcall_store).await {
            Ok(count) if count >= self.max_pending_per_agent => {
                return format!(
                    "Error: Maximum pending tool calls ({}) reached for this agent. \
                     Wait for existing calls to be answered before making new ones.",
                    self.max_pending_per_agent
                );
            }
            Err(e) => {
                warn!("Failed to check pending count: {}", e);
                // Continue anyway — rate limiting is best-effort
            }
            _ => {}
        }

        // 4. Create PendingToolCall
        let call_id = Uuid::new_v4().to_string();
        let pending_call = PendingToolCall {
            call_id: call_id.clone(),
            job_id: self.session_id.clone(),
            agent_id: self.agent_id.clone(),
            tool_name: tool_name.to_string(),
            arguments: arguments.clone(),
            round,
            phase,
            status: ToolCallStatus::Pending,
            created_at: Self::now_epoch_millis(),
            responded_at: None,
            result: None,
        };

        // 5. Store in KV
        let key = format!("call_{}", call_id);
        let data = match serde_json::to_vec(&pending_call) {
            Ok(d) => d,
            Err(e) => return format!("Error: Failed to serialize tool call: {}", e),
        };
        if let Err(e) = toolcall_store.put(&key, data.into()).await {
            return format!("Error: Failed to store pending tool call: {}", e);
        }

        // 6. Publish SSE event
        self.publish_sse_event(
            "tool_call_pending",
            &serde_json::json!({
                "call_id": &call_id,
                "agent_id": &self.agent_id,
                "tool_name": tool_name,
                "arguments": &arguments,
                "round": round,
                "phase": &phase,
            }),
        )
        .await;

        info!(
            agent = %self.agent_id,
            tool = %tool_name,
            call_id = %call_id,
            "User tool call published. Waiting for response..."
        );

        // Deadline from the budget measured before publishing. Re-reading it here
        // would only shave off the microseconds the publish took.
        let deadline = remaining.saturating_sub(reserve);

        // 8. Watch for response with timeout
        let result = self
            .wait_for_response(&toolcall_store, &key, deadline)
            .await;

        match result {
            WaitResult::Responded(response_text) => {
                self.publish_sse_event(
                    "tool_call_responded",
                    &serde_json::json!({
                        "call_id": &call_id,
                        "agent_id": &self.agent_id,
                        "tool_name": tool_name,
                    }),
                )
                .await;
                info!(
                    agent = %self.agent_id,
                    call_id = %call_id,
                    "User tool call responded."
                );
                // Wrap in tags to delimit untrusted content; escape to prevent XML breakout.
                // Attribute values use escape_xml_attr (includes quote escaping).
                let escaped = escape_xml(&response_text);
                let safe_tool = escape_xml_attr(tool_name);
                let safe_call = escape_xml_attr(&call_id);
                format!(
                    "<user_tool_result tool=\"{}\" call_id=\"{}\">{}</user_tool_result>",
                    safe_tool, safe_call, escaped
                )
            }
            WaitResult::Timeout => {
                self.expire_call(&toolcall_store, &key, &call_id, tool_name)
                    .await;
                let remaining_after = self.remaining_budget();
                format!(
                    "[No response yet — you have {:.0}s remaining to finalize your proposal \
                     with your best judgment. The user may respond later and the result will \
                     be available next round.]",
                    remaining_after.as_secs_f64()
                )
            }
            WaitResult::Error(e) => {
                warn!(call_id = %call_id, error = %e, "Error waiting for tool call response");
                format!("Error waiting for user response: {}", e)
            }
        }
    }

    async fn get_or_create_bucket(
        &self,
        bucket_name: &str,
    ) -> Result<async_nats::jetstream::kv::Store> {
        ensure_kv_bucket(
            &self.js_context,
            async_nats::jetstream::kv::Config {
                bucket: bucket_name.to_string(),
                history: 5,
                max_age: Duration::from_secs(86400 * 3),
                storage: async_nats::jetstream::stream::StorageType::File,
                ..Default::default()
            },
        )
        .await
    }

    async fn count_pending_for_agent(
        &self,
        store: &async_nats::jetstream::kv::Store,
    ) -> Result<usize> {
        let scan_start = Instant::now();
        let mut count = 0;
        let mut total_keys = 0u32;
        let mut keys = store.keys().await?;
        while let Some(key_result) = keys.next().await {
            let Ok(key) = key_result else { continue };
            if !key.starts_with("call_") {
                continue;
            }
            total_keys += 1;
            let Ok(Some(entry)) = store.get(&key).await else {
                continue;
            };
            let Ok(call) = serde_json::from_slice::<PendingToolCall>(&entry) else {
                continue;
            };
            if call.agent_id == self.agent_id && call.status == ToolCallStatus::Pending {
                count += 1;
            }
        }
        let scan_ms = scan_start.elapsed().as_millis();
        if total_keys > 50 || scan_ms > 100 {
            warn!(
                total_keys = total_keys,
                pending = count,
                agent = %self.agent_id,
                scan_ms = scan_ms,
                "Tool call bucket scan is growing — consider secondary counter if this persists"
            );
        }
        Ok(count)
    }

    async fn wait_for_response(
        &self,
        store: &async_nats::jetstream::kv::Store,
        key: &str,
        timeout_duration: Duration,
    ) -> WaitResult {
        // Use watch_with_history so the most recent entry is replayed before live updates.
        // This prevents the race where a response is written between our initial put()
        // and the watcher subscription being established.
        let mut watcher = match store.watch_with_history(key).await {
            Ok(w) => w,
            Err(e) => return WaitResult::Error(format!("Failed to create KV watcher: {}", e)),
        };

        tokio::select! {
            result = async {
                while let Some(entry) = watcher.next().await {
                    let Ok(entry) = entry else { continue };
                    let Ok(call) = serde_json::from_slice::<PendingToolCall>(&entry.value) else {
                        continue;
                    };
                    if call.status == ToolCallStatus::Responded {
                        return WaitResult::Responded(call.result.unwrap_or_default());
                    }
                }
                WaitResult::Error("KV watcher stream ended unexpectedly".to_string())
            } => result,
            _ = tokio::time::sleep(timeout_duration) => {
                WaitResult::Timeout
            }
        }
    }

    async fn expire_call(
        &self,
        store: &async_nats::jetstream::kv::Store,
        key: &str,
        call_id: &str,
        tool_name: &str,
    ) {
        // Use entry() to get the revision for CAS, and only expire if still Pending
        if let Ok(Some(entry)) = store.entry(key).await
            && let Ok(mut call) = serde_json::from_slice::<PendingToolCall>(&entry.value)
        {
            if call.status != ToolCallStatus::Pending {
                // Already responded or expired concurrently — don't clobber
                return;
            }
            call.status = ToolCallStatus::Expired;
            match serde_json::to_vec(&call) {
                Ok(data) => {
                    // CAS update: only succeeds if entry hasn't been modified since we read it
                    match store.update(key, data.into(), entry.revision).await {
                        Ok(_) => {
                            // CAS succeeded — publish SSE only after confirmed expiration
                            self.publish_sse_event(
                                "tool_call_expired",
                                &serde_json::json!({
                                    "call_id": call_id,
                                    "agent_id": &self.agent_id,
                                    "tool_name": tool_name,
                                    "timeout_secs": self.phase_budget.as_secs_f64(),
                                }),
                            )
                            .await;
                        }
                        Err(e) => {
                            warn!(
                                call_id = %call_id,
                                error = %e,
                                "CAS update failed for expire_call (concurrent modification?)"
                            );
                        }
                    }
                }
                Err(e) => {
                    warn!(
                        call_id = %call_id,
                        error = %e,
                        "Failed to serialize expired tool call"
                    );
                }
            }
        }
    }

    async fn publish_sse_event<T: serde::Serialize>(&self, suffix: &str, payload: &T) {
        let data = match serde_json::to_vec(payload) {
            Ok(d) => d,
            Err(e) => {
                warn!("Failed to serialize SSE event: {}", e);
                return;
            }
        };
        let safe_session = crate::nats_utils::sanitize_subject_component(&self.session_id);
        let safe_prefix = crate::nats_utils::sanitize_subject_component(&self.subject_prefix);
        let subject = format!("{}.{}.result.event.{}", safe_prefix, safe_session, suffix);
        if let Err(e) = self.nats_client.publish(subject.clone(), data.into()).await {
            warn!("Failed to publish SSE event to {}: {}", subject, e);
        }
    }
}

/// Implement the SDK trait so `UserToolHandler` can be stored as `Arc<dyn UserToolHandlerTrait>`.
#[async_trait]
impl UserToolHandlerTrait for UserToolHandler {
    async fn handle_call(
        &self,
        tool_name: &str,
        arguments_json: &str,
        round: u32,
        phase: DeliberationPhase,
    ) -> String {
        self.handle_call(tool_name, arguments_json, round, phase)
            .await
    }
}

/// Factory that creates [`UserToolHandler`] instances for each task execution.
///
/// Reference implementation of the [`UserToolHandlerFactory`](crate::workers::UserToolHandlerFactory)
/// trait — the worker uses this to materialise a handler per task without
/// depending on NATS internals at the trait level.
#[derive(Debug)]
pub struct NatsUserToolHandlerFactory;

impl crate::workers::UserToolHandlerFactory for NatsUserToolHandlerFactory {
    fn create(
        &self,
        nats: async_nats::Client,
        js: async_nats::jetstream::Context,
        session_id: String,
        agent_id: String,
        budget_remaining_secs: f64,
        subject_prefix: String,
    ) -> std::sync::Arc<dyn UserToolHandlerTrait> {
        std::sync::Arc::new(
            UserToolHandler::new(nats, js, session_id, agent_id, budget_remaining_secs)
                .with_subject_prefix(subject_prefix),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agents::{PendingToolCall, ToolCallStatus, UserToolDefinition};

    #[test]
    fn test_user_tool_definition_serde_roundtrip() {
        let def = UserToolDefinition {
            name: "dm_user".to_string(),
            description: "Send a DM to the user".to_string(),
            parameters: Some(serde_json::json!({
                "type": "object",
                "properties": {
                    "message": { "type": "string" }
                },
                "required": ["message"]
            })),
            strict: Some(true),
        };

        let json = serde_json::to_string(&def).unwrap();
        let parsed: UserToolDefinition = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.name, "dm_user");
        assert_eq!(parsed.strict, Some(true));
        assert!(parsed.parameters.is_some());
    }

    #[test]
    fn test_user_tool_definition_minimal() {
        let json = r#"{"name": "ping", "description": "Ping the user"}"#;
        let parsed: UserToolDefinition = serde_json::from_str(json).unwrap();
        assert_eq!(parsed.name, "ping");
        assert!(parsed.parameters.is_none());
        assert!(parsed.strict.is_none());
    }

    #[test]
    fn test_pending_tool_call_serde_roundtrip() {
        let call = PendingToolCall {
            call_id: "abc-123".to_string(),
            job_id: "job-1".to_string(),
            agent_id: "agent-1".to_string(),
            tool_name: "user_dm_user".to_string(),
            arguments: serde_json::json!({"message": "hello"}),
            round: 1,
            phase: DeliberationPhase::Proposing,
            status: ToolCallStatus::Pending,
            created_at: 1234567890,
            responded_at: None,
            result: None,
        };

        let json = serde_json::to_string(&call).unwrap();
        let parsed: PendingToolCall = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.call_id, "abc-123");
        assert_eq!(parsed.status, ToolCallStatus::Pending);
        assert!(parsed.responded_at.is_none());
        assert!(parsed.result.is_none());
    }

    #[test]
    fn test_pending_tool_call_responded() {
        let call = PendingToolCall {
            call_id: "abc-123".to_string(),
            job_id: "job-1".to_string(),
            agent_id: "agent-1".to_string(),
            tool_name: "user_dm_user".to_string(),
            arguments: serde_json::json!({}),
            round: 2,
            phase: DeliberationPhase::Evaluating,
            status: ToolCallStatus::Responded,
            created_at: 1234567890,
            responded_at: Some(1234567900),
            result: Some("The answer is 42".to_string()),
        };

        let json = serde_json::to_string(&call).unwrap();
        let parsed: PendingToolCall = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.status, ToolCallStatus::Responded);
        assert_eq!(parsed.result, Some("The answer is 42".to_string()));
    }

    #[test]
    fn test_tool_call_status_all_variants() {
        for (status, expected) in [
            (ToolCallStatus::Pending, "\"Pending\""),
            (ToolCallStatus::Responded, "\"Responded\""),
            (ToolCallStatus::Expired, "\"Expired\""),
        ] {
            let json = serde_json::to_string(&status).unwrap();
            assert_eq!(json, expected);
            let parsed: ToolCallStatus = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed, status);
        }
    }

    /// A JetStream context over a live server, or `None` so the test skips.
    async fn js() -> Option<(async_nats::Client, async_nats::jetstream::Context)> {
        let url = std::env::var("NATS_URL").unwrap_or_else(|_| "nats://localhost:4222".to_string());
        let client = async_nats::connect(&url).await.ok()?;
        let js = async_nats::jetstream::new(client.clone());
        Some((client, js))
    }

    #[tokio::test]
    async fn a_call_with_no_time_to_wait_is_refused_before_anything_is_published() {
        let Some((client, js)) = js().await else {
            eprintln!("Skipping: NATS unavailable");
            return;
        };
        let session = format!("test-refuse-{}", Uuid::new_v4());
        // reserve = min(budget × 1.0, 1000s) = the whole budget, so the call lands
        // inside the reserve on its first attempt — the same position a call
        // reaches late in a real phase, without waiting for one.
        let handler = UserToolHandler::new(
            client,
            js.clone(),
            session.clone(),
            "ALPHA".to_string(),
            60.0,
        )
        .with_finalization_reserve(1000.0, 1.0);

        let answer = handler
            .handle_call(
                "user_dm_user",
                r#"{"message":"hi"}"#,
                1,
                DeliberationPhase::Proposing,
            )
            .await;
        assert!(
            answer.contains("phase budget exhausted"),
            "the model is told to proceed: {answer}"
        );

        // Nothing was published: no bucket, so no record and no pending row for a
        // client to render and immediately withdraw.
        let bucket = format!("nsed_toolcalls_{session}");
        assert!(
            js.get_key_value(&bucket).await.is_err(),
            "a question that cannot be waited for must not create {bucket}"
        );
    }

    #[test]
    fn the_bucket_name_follows_the_configured_prefix() {
        assert_eq!(
            toolcalls_bucket_name("nsed", "room-f2205792"),
            "nsed_toolcalls_room-f2205792"
        );
        // A non-default prefix must reach the name, or the reader and writer end
        // up on different buckets.
        assert_eq!(
            toolcalls_bucket_name("staging", "room-1"),
            "staging_toolcalls_room-1"
        );
        // Both components are sanitized, so a session id carrying subject
        // metacharacters cannot escape into the bucket name.
        assert_eq!(
            toolcalls_bucket_name("nsed", "room.1 *>"),
            toolcalls_bucket_name("nsed", "room.1 *>"),
        );
        let odd = toolcalls_bucket_name("ns.ed", "a>b");
        assert!(!odd.contains('.'), "{odd}");
        assert!(!odd.contains('>'), "{odd}");
    }

    #[test]
    fn test_finalization_reserve_computation() {
        // 200 * 0.15 = 30, min(30, 30) = 30
        assert_eq!(
            compute_finalization_reserve(Duration::from_secs(200), 30.0, 0.15),
            Duration::from_secs(30)
        );
    }

    #[test]
    fn test_finalization_reserve_small_budget() {
        // 60 * 0.15 = 9, min(9, 30) = 9
        assert_eq!(
            compute_finalization_reserve(Duration::from_secs(60), 30.0, 0.15),
            Duration::from_secs(9)
        );
    }

    // ---- Regression tests for PR review fixes ----

    /// escape_xml must escape &, <, > to prevent XML breakout in user_tool_result
    #[test]
    fn test_escape_xml_basic_entities() {
        assert_eq!(escape_xml("hello"), "hello");
        assert_eq!(escape_xml("<script>"), "&lt;script&gt;");
        assert_eq!(escape_xml("a & b"), "a &amp; b");
        assert_eq!(escape_xml(""), "");
    }

    /// Verifies that the escaped form doesn't break the user_tool_result wrapper
    #[test]
    fn test_escape_xml_preserves_wrapper_integrity() {
        let malicious = "</user_tool_result><injected>evil</injected>";
        let escaped = escape_xml(malicious);
        let wrapped = format!(
            "<user_tool_result tool=\"test\" call_id=\"c1\">{}</user_tool_result>",
            escaped
        );
        // The wrapper tags should remain intact
        assert!(wrapped.starts_with("<user_tool_result tool=\"test\" call_id=\"c1\">"));
        assert!(wrapped.ends_with("</user_tool_result>"));
        // The injected tags should be escaped, not present as raw XML
        assert!(!wrapped.contains("<injected>"));
        assert!(wrapped.contains("&lt;injected&gt;"));
    }

    /// escape_xml handles all three special chars in a single string
    #[test]
    fn test_escape_xml_combined() {
        let input = "x < 5 & y > 3";
        let expected = "x &lt; 5 &amp; y &gt; 3";
        assert_eq!(escape_xml(input), expected);
    }

    /// escape_xml handles ampersand-first correctly (no double-escaping)
    #[test]
    fn test_escape_xml_ampersand_first() {
        // Since & is replaced first, &lt; in input becomes &amp;lt;
        let input = "&lt;";
        let expected = "&amp;lt;";
        assert_eq!(escape_xml(input), expected);
    }

    // ---- Tests for escape_xml_attr ----

    /// escape_xml_attr escapes quotes in addition to &, <, >
    #[test]
    fn test_escape_xml_attr_quotes() {
        assert_eq!(escape_xml_attr(r#"he said "hi""#), "he said &quot;hi&quot;");
        assert_eq!(escape_xml_attr("it's"), "it&apos;s");
    }

    /// escape_xml_attr prevents attribute injection via tool_name
    #[test]
    fn test_escape_xml_attr_prevents_attribute_injection() {
        // An attacker tries to close the tool attr and inject a new one
        let malicious_tool = r#"evil" onclick="alert(1)"#;
        let safe = escape_xml_attr(malicious_tool);
        // The escaped string must not contain unescaped quotes
        assert!(!safe.contains('"'));
        assert!(safe.contains("&quot;"));
    }

    // ---- Tests for compute_finalization_reserve sanitization ----

    /// NaN inputs should not panic, should return Duration::ZERO
    #[test]
    fn test_finalization_reserve_nan_inputs() {
        let result = compute_finalization_reserve(Duration::from_secs(100), f64::NAN, 0.15);
        assert_eq!(result, Duration::ZERO);

        let result = compute_finalization_reserve(Duration::from_secs(100), 30.0, f64::NAN);
        assert_eq!(result, Duration::ZERO);
    }

    /// Negative inputs should not panic, should be treated as zero
    #[test]
    fn test_finalization_reserve_negative_inputs() {
        let result = compute_finalization_reserve(Duration::from_secs(100), -10.0, 0.15);
        assert_eq!(result, Duration::ZERO);

        let result = compute_finalization_reserve(Duration::from_secs(100), 30.0, -0.5);
        assert_eq!(result, Duration::ZERO);
    }

    /// Infinite inputs should not panic
    #[test]
    fn test_finalization_reserve_infinite_inputs() {
        let result = compute_finalization_reserve(Duration::from_secs(100), f64::INFINITY, 0.15);
        // ratio_based = 100*0.15 = 15, fixed = 0 (inf sanitized to 0), min(15,0) = 0
        assert_eq!(result, Duration::ZERO);

        let result =
            compute_finalization_reserve(Duration::from_secs(100), 30.0, f64::NEG_INFINITY);
        assert_eq!(result, Duration::ZERO);
    }

    /// reserve_ratio capped at 1.0
    #[test]
    fn test_finalization_reserve_ratio_capped() {
        // ratio 2.0 should be capped to 1.0 → 100*1.0 = 100, min(100, 200) = 100
        let result = compute_finalization_reserve(Duration::from_secs(100), 200.0, 2.0);
        assert_eq!(result, Duration::from_secs(100));
    }

    // ---- Constructor phase_budget sanitization tests ----

    /// The same sanitization pattern used in UserToolHandler::new
    /// (extracted for direct testing since the constructor requires NATS)
    #[test]
    fn test_phase_budget_sanitization_nan() {
        let val: f64 = f64::NAN;
        let safe = if val.is_finite() { val.max(0.0) } else { 0.0 };
        assert_eq!(safe, 0.0);
        // Must not panic
        let _ = Duration::from_secs_f64(safe);
    }

    #[test]
    fn test_phase_budget_sanitization_infinity() {
        let val: f64 = f64::INFINITY;
        let safe = if val.is_finite() { val.max(0.0) } else { 0.0 };
        assert_eq!(safe, 0.0);
        let _ = Duration::from_secs_f64(safe);
    }

    #[test]
    fn test_phase_budget_sanitization_neg_infinity() {
        let val: f64 = f64::NEG_INFINITY;
        let safe = if val.is_finite() { val.max(0.0) } else { 0.0 };
        assert_eq!(safe, 0.0);
        let _ = Duration::from_secs_f64(safe);
    }

    #[test]
    fn test_phase_budget_sanitization_negative() {
        let val: f64 = -100.0;
        let safe = if val.is_finite() { val.max(0.0) } else { 0.0 };
        assert_eq!(safe, 0.0);
        let _ = Duration::from_secs_f64(safe);
    }

    #[test]
    fn test_phase_budget_sanitization_valid() {
        let val: f64 = 42.5;
        let safe = if val.is_finite() { val.max(0.0) } else { 0.0 };
        assert_eq!(safe, 42.5);
        assert_eq!(Duration::from_secs_f64(safe), Duration::from_millis(42500));
    }

    #[test]
    fn test_publish_sse_event_sanitizes_session_id() {
        // Verify sanitize_subject_component strips NATS-invalid chars
        let raw = "my.session>with*wildcards";
        let safe = crate::nats_utils::sanitize_subject_component(raw);
        assert!(!safe.contains('.'), "Dots should be removed");
        assert!(!safe.contains('>'), "Greater-than should be removed");
        assert!(!safe.contains('*'), "Wildcards should be removed");
        assert!(!safe.is_empty(), "Sanitized result should not be empty");
    }

    // =========================================================================
    // escape_xml — additional edge cases
    // =========================================================================

    #[test]
    fn test_escape_xml_empty_string() {
        assert_eq!(escape_xml(""), "");
    }

    #[test]
    fn test_escape_xml_no_special_chars() {
        let input = "Hello, world! 123 test";
        assert_eq!(escape_xml(input), input);
    }

    #[test]
    fn test_escape_xml_all_special_chars() {
        let input = "&<>";
        assert_eq!(escape_xml(input), "&amp;&lt;&gt;");
    }

    #[test]
    fn test_escape_xml_preserves_quotes() {
        // escape_xml does NOT escape quotes (only escape_xml_attr does)
        let input = "He said \"hello\" and it's fine";
        assert_eq!(escape_xml(input), "He said \"hello\" and it's fine");
    }

    #[test]
    fn test_escape_xml_already_escaped() {
        // Already-escaped content should be double-escaped
        let input = "&amp; &lt; &gt;";
        let result = escape_xml(input);
        assert_eq!(result, "&amp;amp; &amp;lt; &amp;gt;");
    }

    #[test]
    fn test_escape_xml_multiline() {
        let input = "line1 <b>bold</b>\nline2 & more\nline3 > end";
        let expected = "line1 &lt;b&gt;bold&lt;/b&gt;\nline2 &amp; more\nline3 &gt; end";
        assert_eq!(escape_xml(input), expected);
    }

    // =========================================================================
    // escape_xml_attr — additional edge cases
    // =========================================================================

    #[test]
    fn test_escape_xml_attr_empty() {
        assert_eq!(escape_xml_attr(""), "");
    }

    #[test]
    fn test_escape_xml_attr_all_five_special_chars() {
        let input = "&<>\"'";
        assert_eq!(escape_xml_attr(input), "&amp;&lt;&gt;&quot;&apos;");
    }

    #[test]
    fn test_escape_xml_attr_no_special_chars() {
        let input = "simple text 123";
        assert_eq!(escape_xml_attr(input), input);
    }

    #[test]
    fn test_escape_xml_attr_mixed_quotes_and_entities() {
        let input = "tool_name=\"bad\" & 'evil' <injected>";
        let expected = "tool_name=&quot;bad&quot; &amp; &apos;evil&apos; &lt;injected&gt;";
        assert_eq!(escape_xml_attr(input), expected);
    }

    // =========================================================================
    // compute_finalization_reserve — additional edge cases
    // =========================================================================

    #[test]
    fn test_finalization_reserve_zero_budget() {
        let result = compute_finalization_reserve(Duration::ZERO, 30.0, 0.15);
        assert_eq!(result, Duration::ZERO);
    }

    #[test]
    fn test_finalization_reserve_both_params_zero() {
        let result = compute_finalization_reserve(Duration::from_secs(100), 0.0, 0.0);
        assert_eq!(result, Duration::ZERO);
    }

    #[test]
    fn test_finalization_reserve_ratio_exactly_one() {
        // ratio=1.0 means reserve = entire budget, min(budget, fixed_secs)
        let result = compute_finalization_reserve(Duration::from_secs(100), 200.0, 1.0);
        // ratio_based = 100 * 1.0 = 100s, fixed = 200s, min(100, 200) = 100s
        assert_eq!(result, Duration::from_secs(100));
    }

    #[test]
    fn test_finalization_reserve_fixed_smaller_than_ratio() {
        // fixed = 10s, ratio_based = 100 * 0.5 = 50s → min(50, 10) = 10s
        let result = compute_finalization_reserve(Duration::from_secs(100), 10.0, 0.5);
        assert_eq!(result, Duration::from_secs(10));
    }

    #[test]
    fn test_finalization_reserve_very_large_budget() {
        // Large budget with small ratio
        let result = compute_finalization_reserve(Duration::from_secs(10000), 60.0, 0.01);
        // ratio_based = 10000 * 0.01 = 100s, fixed = 60s → min(100, 60) = 60s
        assert_eq!(result, Duration::from_secs(60));
    }

    #[test]
    fn test_finalization_reserve_very_small_budget() {
        // Sub-second budget should still produce a safe reserve
        let result = compute_finalization_reserve(Duration::from_millis(100), 30.0, 0.15);
        // ratio_based = 0.1s * 0.15 = 0.015s, fixed = 30, min(0.015, 30) = 0.015s
        assert_eq!(result, Duration::from_millis(15));
    }

    #[test]
    fn test_finalization_reserve_both_nan() {
        let result = compute_finalization_reserve(Duration::from_secs(100), f64::NAN, f64::NAN);
        assert_eq!(result, Duration::ZERO);
    }

    #[test]
    fn test_finalization_reserve_both_infinite() {
        let result =
            compute_finalization_reserve(Duration::from_secs(100), f64::INFINITY, f64::INFINITY);
        // Both sanitized to 0 → Duration::ZERO
        assert_eq!(result, Duration::ZERO);
    }

    #[test]
    fn test_finalization_reserve_neg_infinity_secs() {
        let result =
            compute_finalization_reserve(Duration::from_secs(100), f64::NEG_INFINITY, 0.15);
        // safe_secs = 0 (neg_inf is not finite), ratio_based = 15, fixed = 0, min(15, 0) = 0
        assert_eq!(result, Duration::ZERO);
    }

    #[test]
    fn test_finalization_reserve_fractional_duration() {
        // Test with sub-second precision
        let result = compute_finalization_reserve(Duration::from_millis(500), 1.0, 0.1);
        // ratio_based = 0.5 * 0.1 = 0.05s = 50ms, fixed = 1.0s = 1000ms, min(50, 1000) = 50ms
        assert_eq!(result, Duration::from_millis(50));
    }
}