bamboo-engine 2026.10.13

//! Stop redundant polling after a quiet-command reminder is repeatedly ignored.
//! Only the admitted shell identity and raw successful results participate.
use std::collections::HashMap;

use bamboo_agent_core::tools::{parse_tool_args, ToolCall, ToolResult};

use super::per_call::ToolExecutionOutcome;

const MAX_HANDLES: usize = 64;
const SILENCE_MS: u64 = 90_000;
// The first qualifying readout already carries the reminder; allow three more.
const READOUTS_AFTER_REMINDER: usize = 4;

#[derive(Debug, Clone, Default)]
pub(super) struct QuietPollGuard {
    readouts: HashMap<String, usize>,
    pending: Option<String>,
}

fn poll_args(call: &ToolCall) -> Option<(String, bool)> {
    let args = parse_tool_args(&call.function.arguments).ok()?;
    let object = args.as_object()?;
    if object.keys().any(|key| {
        !matches!(
            key.as_str(),
            "session_id" | "chars" | "eof" | "yield_time_ms" | "max_output_tokens"
        )
    }) {
        return None;
    }
    if args
        .get("yield_time_ms")
        .is_some_and(|v| !v.is_null() && v.as_u64().is_none())
        || args.get("max_output_tokens").is_some_and(|v| {
            !v.is_null() && v.as_u64().and_then(|n| usize::try_from(n).ok()).is_none()
        })
    {
        return None;
    }
    let id = args.get("session_id")?.as_str()?;
    if id.is_empty() || id.len() > 256 {
        return None;
    }
    let chars = match args.get("chars") {
        None => "",
        Some(value) => value.as_str()?,
    };
    let eof = match args.get("eof") {
        None => false,
        Some(value) => value.as_bool()?,
    };
    Some((id.to_owned(), chars.is_empty() && !eof))
}

impl QuietPollGuard {
    pub(super) fn observe(
        &mut self,
        call: &ToolCall,
        outcome: &ToolExecutionOutcome,
        admitted: Option<&str>,
    ) {
        let identity = outcome
            .portable_tool
            .as_ref()
            .map(|tool| tool.execution_name())
            .or(admitted);
        if identity != Some("write_stdin") || !outcome.post_tool_hook_eligible {
            return;
        }
        let Some((id, empty_poll)) = poll_args(call) else {
            return;
        };
        let quiet = empty_poll
            && outcome
                .result
                .as_ref()
                .ok()
                .filter(|r| r.success && r.images.is_empty())
                .and_then(|r| serde_json::from_str::<serde_json::Value>(&r.result).ok())
                .is_some_and(|value| {
                    value["session_id"].as_str() == Some(id.as_str())
                        && value["status"] == "running"
                        && value["output"] == ""
                        && value
                            .get("output_omitted_bytes")
                            .is_none_or(|v| v.as_u64() == Some(0))
                        && value["polling"]["output_idle_ms"]
                            .as_u64()
                            .is_some_and(|ms| ms >= SILENCE_MS)
                        && value["polling"]["consecutive_empty_readouts"]
                            .as_u64()
                            .is_some_and(|n| n >= 3)
                });
        if !quiet {
            self.readouts.remove(&id);
            if self.pending.as_deref() == Some(id.as_str()) {
                self.pending = None;
            }
            return;
        }
        if self.readouts.len() >= MAX_HANDLES && !self.readouts.contains_key(&id) {
            return;
        }
        let count = self.readouts.entry(id.clone()).or_default();
        *count = count.saturating_add(1);
        if *count >= READOUTS_AFTER_REMINDER {
            self.pending = Some(id);
        }
    }

    pub(super) fn pending(&self) -> bool {
        self.pending.is_some()
    }

    pub(super) fn retain_owned_handles(&mut self, handles: &[String]) {
        self.readouts.retain(|id, _| handles.contains(id));
        if self
            .pending
            .as_ref()
            .is_some_and(|id| !handles.contains(id))
        {
            self.pending = None;
        }
    }

    pub(super) fn suppress(&self, call: &ToolCall, admitted: Option<&str>) -> Option<ToolResult> {
        if admitted != Some("write_stdin") {
            return None;
        }
        let (id, empty_poll) = poll_args(call)?;
        if !empty_poll || self.pending.as_deref() != Some(id.as_str()) {
            return None;
        }
        // Pair every planned call, without another wait or capture read. This
        // is a policy observation, not a new process/exit-status observation.
        Some(ToolResult::text(
            true,
            serde_json::json!({
                "session_id":id, "status":"running", "output":"",
                "polling":{"state":"diagnosis_pending", "poll_suppressed":true},
            })
            .to_string(),
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bamboo_agent_core::tools::FunctionCall;

    fn call(name: &str, args: serde_json::Value) -> ToolCall {
        ToolCall {
            id: "probe".into(),
            tool_type: "function".into(),
            function: FunctionCall {
                name: name.into(),
                arguments: args.to_string(),
            },
        }
    }

    fn outcome(tick: u64, output: &str) -> ToolExecutionOutcome {
        ToolExecutionOutcome {
            output_cap: None,
            permission_replay_origin: None,
            portable_tool: None,
            post_tool_hook_eligible: true,
            needs_human: None,
            managed_running_violation: false,
            tool_duration: std::time::Duration::ZERO,
            result: Ok(ToolResult::text(
                true,
                serde_json::json!({
                    "session_id":"handle", "status":"running", "output":output,
                    "polling":{"elapsed_ms":tick, "output_idle_ms":90_000+tick,
                        "consecutive_empty_readouts":3+tick},
                })
                .to_string(),
            )),
        }
    }

    #[test]
    fn timestamps_yield_caps_aliases_and_key_order_cannot_hide_repeated_empty_polls() {
        let mut guard = QuietPollGuard::default();
        for tick in 0..4 {
            let name = if tick % 2 == 0 {
                "write_stdin"
            } else {
                "default::write_stdin"
            };
            let poll = call(
                name,
                serde_json::json!({"session_id":"handle", "chars":"",
                "yield_time_ms":250+tick*1000,"max_output_tokens":tick+1}),
            );
            guard.observe(&poll, &outcome(tick, ""), Some("write_stdin"));
            assert_eq!(guard.pending(), tick == 3);
        }
        let poll = call("write_stdin", serde_json::json!({"session_id":"handle"}));
        let suppressed = guard.suppress(&poll, Some("write_stdin")).unwrap();
        let value: serde_json::Value = serde_json::from_str(&suppressed.result).unwrap();
        assert_eq!(value["status"], "running");
        assert_eq!(value["session_id"], "handle");
        assert_eq!(value["polling"]["poll_suppressed"], true);
        assert!(value.get("exit_code").is_none());
        for invalid in [
            serde_json::json!({"session_id":"handle", "unexpected":true}),
            serde_json::json!({"session_id":"handle", "yield_time_ms":"250"}),
            serde_json::json!({"session_id":"handle", "eof":null}),
        ] {
            assert!(guard
                .suppress(&call("write_stdin", invalid), Some("write_stdin"))
                .is_none());
        }
        assert!(guard.suppress(&poll, Some("vendor::write_stdin")).is_none());
        assert!(guard
            .suppress(
                &call("write_stdin", serde_json::json!({"session_id":"other"})),
                Some("write_stdin")
            )
            .is_none());
        assert!(guard
            .suppress(
                &call("Read", serde_json::json!({"session_id":"handle"})),
                Some("Read")
            )
            .is_none());
        assert!(guard
            .suppress(
                &call(
                    "write_stdin",
                    serde_json::json!({"session_id":"handle","eof":true})
                ),
                Some("write_stdin")
            )
            .is_none());
        // Explicit input remains dispatchable and real progress releases the gate.
        let input = call(
            "write_stdin",
            serde_json::json!({"session_id":"handle","chars":"answer\n"}),
        );
        assert!(guard.suppress(&input, Some("write_stdin")).is_none());
        guard.observe(&input, &outcome(5, ""), Some("write_stdin"));
        assert!(!guard.pending());
        guard.observe(&poll, &outcome(6, "new output"), Some("write_stdin"));
        assert!(!guard.pending());
    }

    #[test]
    fn omitted_output_unknown_metadata_errors_and_foreign_tools_are_not_quiet_progress() {
        let poll = call("write_stdin", serde_json::json!({"session_id":"handle"}));
        for variant in 0..6 {
            let mut guard = QuietPollGuard::default();
            for tick in 0..3 {
                guard.observe(&poll, &outcome(tick, ""), Some("write_stdin"));
            }
            let mut result = outcome(4, "");
            let mut value: serde_json::Value =
                serde_json::from_str(&result.result.as_ref().unwrap().result).unwrap();
            let identity = match variant {
                0 => {
                    value["output_omitted_bytes"] = 1.into();
                    "write_stdin"
                }
                1 => {
                    value.as_object_mut().unwrap().remove("polling");
                    "write_stdin"
                }
                2 => {
                    value["status"] = "completed".into();
                    value["exit_code"] = 7.into();
                    "write_stdin"
                }
                3 => {
                    value["polling"]["output_idle_ms"] = 89_999.into();
                    "write_stdin"
                }
                4 => "vendor::write_stdin",
                _ => "default::write_stdin",
            };
            result.result = Ok(ToolResult::text(true, value.to_string()));
            guard.observe(&poll, &result, Some(identity));
            assert!(!guard.pending());
        }
        let mut guard = QuietPollGuard::default();
        for tick in 0..3 {
            guard.observe(&poll, &outcome(tick, ""), Some("write_stdin"));
        }
        let mut failed = outcome(4, "");
        failed.result = Err("read failed".into());
        guard.observe(&poll, &failed, Some("write_stdin"));
        assert!(!guard.pending());
        guard.observe(&poll, &outcome(5, ""), Some("write_stdin"));
        assert!(!guard.pending());
    }

    #[test]
    fn expired_or_foreign_handles_release_the_bounded_guard_capacity() {
        let mut guard = QuietPollGuard::default();
        let mut handles = Vec::new();
        for index in 0..MAX_HANDLES {
            let id = format!("old-{index}");
            handles.push(id.clone());
            let poll = call("write_stdin", serde_json::json!({"session_id":id}));
            let mut result = outcome(1, "");
            let mut value: serde_json::Value =
                serde_json::from_str(&result.result.as_ref().unwrap().result).unwrap();
            value["session_id"] = id.into();
            result.result = Ok(ToolResult::text(true, value.to_string()));
            for _ in 0..READOUTS_AFTER_REMINDER {
                guard.observe(&poll, &result, Some("write_stdin"));
            }
        }
        guard.retain_owned_handles(&handles);
        assert!(guard.pending());
        let old = call(
            "write_stdin",
            serde_json::json!({"session_id":handles.last().unwrap()}),
        );
        assert!(guard.suppress(&old, Some("write_stdin")).is_some());
        // Cache expiry, terminal consumption, or a different owner snapshot
        // must free policy capacity as well as release any pending suppression.
        guard.retain_owned_handles(&["handle".into()]);
        assert!(!guard.pending());
        assert!(guard.suppress(&old, Some("write_stdin")).is_none());
        let current = call("write_stdin", serde_json::json!({"session_id":"handle"}));
        for tick in 0..READOUTS_AFTER_REMINDER as u64 {
            guard.observe(&current, &outcome(tick, ""), Some("write_stdin"));
        }
        assert!(guard.suppress(&current, Some("write_stdin")).is_some());
        guard.retain_owned_handles(&[]);
        assert!(!guard.pending());
    }
}