PicoBot/src/tools/emit_signal.rs

//! Contract-bound `emit_signal` tool for background Agents.
//!
//! The tool only exists inside a run whose definition declares a `signal`
//! contract.  Every limit (total count, rate, burst, severity allowlist,
//! payload size/depth, dedupe cooldown) is enforced here and in the
//! Coordinator; the model only supplies `key`, `severity`, `summary`,
//! `details` and an optional `dedupe_key`.  Signals are durable inbox events
//! delivered to the run's root session lane (queue or steer per contract).
//! The tool returns only after the event is persisted.

use std::collections::VecDeque;
use std::sync::Mutex;

use serde::Deserialize;
use serde_json::Value;

use crate::agent::coordinator::AgentCoordinator;
use crate::agent::definition::SignalContract;
use crate::agent::run::{AgentExecutionContext, EmittedSignal};
use crate::storage::agent_inbox::{AgentEventDelivery, AgentEventType, NewInboxEvent};
use crate::tools::{DelegationPolicy, Tool, ToolExecutionContext, ToolOutput, ToolResult};

#[derive(Debug, Clone)]
pub struct SignalInput {
    pub key: String,
    pub severity: String,
    pub summary: String,
    pub details: Option<Value>,
    pub dedupe_key: Option<String>,
    /// Durable event key computed by the tool (dedupe window included).
    pub event_key: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SignalAcceptedStatus {
    Accepted,
    Deduplicated,
}

#[derive(Debug, Clone)]
pub struct SignalAccepted {
    pub signal_id: String,
    pub status: SignalAcceptedStatus,
    pub delivery: AgentEventDelivery,
}

#[derive(Debug, Default)]
struct SignalRateState {
    total: u32,
    last_at_ms: i64,
    burst_times: VecDeque<i64>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct EmitSignalArgs {
    pub key: String,
    pub severity: String,
    pub summary: String,
    #[serde(default)]
    pub details: Option<Value>,
    #[serde(default)]
    pub dedupe_key: Option<String>,
}

use std::sync::Arc;

pub struct EmitSignalTool {
    coordinator: Arc<AgentCoordinator>,
    contract: Arc<SignalContract>,
    rate: Mutex<SignalRateState>,
}

impl EmitSignalTool {
    pub fn new(coordinator: Arc<AgentCoordinator>, contract: Arc<SignalContract>) -> Self {
        Self {
            coordinator,
            contract,
            rate: Mutex::new(SignalRateState::default()),
        }
    }
}

const MAX_KEY_CHARS: usize = 128;
const MAX_DEDUPE_KEY_CHARS: usize = 128;
const MAX_SUMMARY_CHARS: usize = 1024;

#[async_trait::async_trait]
impl Tool for EmitSignalTool {
    fn name(&self) -> &str {
        "emit_signal"
    }

    fn description(&self) -> &str {
        "向主 Agent 发送一条结构化内部信号，用于重要中间状态或监控告警；普通进度请留在工具调用记录里，不要滥用信号。"
    }

    fn parameters_schema(&self) -> serde_json::Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "key": { "type": "string", "description": "信号键，用于区分不同信号" },
                "severity": { "type": "string", "description": "严重级别" },
                "summary": { "type": "string", "description": "简短摘要" },
                "details": { "type": "object", "description": "结构化详情" },
                "dedupe_key": { "type": "string", "description": "去重键（可选）" }
            },
            "required": ["key", "severity", "summary"]
        })
    }

    async fn execute(&self, _args: serde_json::Value) -> anyhow::Result<ToolResult> {
        Ok(ToolResult {
            success: false,
            output: String::new(),
            error: Some("emit_signal requires a run-bound context".to_string()),
        })
    }

    fn delegation_policy(&self) -> DelegationPolicy {
        DelegationPolicy::RuntimeInjected
    }

    async fn execute_with_context(
        &self,
        context: &ToolExecutionContext,
        args: Value,
    ) -> anyhow::Result<ToolOutput> {
        let args: EmitSignalArgs = match serde_json::from_value(args) {
            Ok(args) => args,
            Err(error) => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some(error.to_string()),
                }
                .into());
            }
        };
        let contract = self.contract.clone();
        let coordinator = self.coordinator.clone();
        let rate = &self.rate;
        let context = context.clone();
        let result: Result<ToolOutput, String> = async move {
            let Some(agent) = context.agent.as_deref() else {
                return Err("emit_signal requires a run-bound Agent context".to_string());
            };
            let now_ms = chrono::Utc::now().timestamp_millis();

            // Structural validation against the contract, independent of the
            // model's cooperation.
            if args.key.trim().is_empty() || args.key.len() > MAX_KEY_CHARS {
                return Err(format!("key must contain 1..={MAX_KEY_CHARS} characters"));
            }
            if let Some(dedupe) = args.dedupe_key.as_deref()
                && (dedupe.trim().is_empty() || dedupe.len() > MAX_DEDUPE_KEY_CHARS)
            {
                return Err(format!(
                    "dedupe_key must contain 1..={MAX_DEDUPE_KEY_CHARS} characters"
                ));
            }
            if !contract
                .severity_allowlist
                .iter()
                .any(|allowed| allowed == &args.severity)
            {
                return Err(format!(
                    "severity '{0}' is not allowed; allowlist: {1}",
                    args.severity,
                    contract.severity_allowlist.join(", ")
                ));
            }
            if args.summary.trim().is_empty() || args.summary.chars().count() > MAX_SUMMARY_CHARS {
                return Err(format!(
                    "summary must contain 1..={MAX_SUMMARY_CHARS} characters"
                ));
            }
            if let Some(details) = args.details.as_ref() {
                let bytes = details.to_string().len();
                if bytes > contract.max_details_bytes {
                    return Err(format!(
                        "details exceed the {} byte contract limit",
                        contract.max_details_bytes
                    ));
                }
                if json_depth(details) > contract.max_payload_depth {
                    return Err(format!(
                        "details exceed the {} level depth limit",
                        contract.max_payload_depth
                    ));
                }
            }

            // Per-run rate limits.  Rate state is in-memory and per-run: the
            // tool instance is created for exactly one run's registry.
            {
                let mut state = rate.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
                if state.total >= contract.max_total {
                    return Err(format!(
                        "signal limit reached: {} signals for this run",
                        contract.max_total
                    ));
                }
                if state.total > 0 && now_ms - state.last_at_ms < contract.min_interval_ms as i64 {
                    return Err(format!(
                        "signal rate limited: wait at least {}ms between signals",
                        contract.min_interval_ms
                    ));
                }
                let window_start = now_ms - contract.burst_window_ms as i64;
                while state
                    .burst_times
                    .front()
                    .is_some_and(|time| *time < window_start)
                {
                    state.burst_times.pop_front();
                }
                if state.burst_times.len() as u32 >= contract.max_burst {
                    return Err(format!(
                        "signal burst limited: at most {} signals per {}ms",
                        contract.max_burst, contract.burst_window_ms
                    ));
                }
            }

            let event_key = match args.dedupe_key.as_deref() {
                Some(key) => {
                    let window = now_ms / contract.dedupe_cooldown_ms as i64;
                    format!("signal:{key}:{window}")
                }
                None => format!("signal:{}", uuid::Uuid::new_v4()),
            };
            let accepted = coordinator
                .emit_signal(
                    agent,
                    SignalInput {
                        key: args.key.clone(),
                        severity: args.severity.clone(),
                        summary: args.summary.clone(),
                        details: args.details.clone(),
                        dedupe_key: args.dedupe_key.clone(),
                        event_key,
                    },
                )
                .await
                .map_err(|error| error.to_string())?;

            if accepted.status == SignalAcceptedStatus::Accepted {
                if let Ok(mut state) = rate.lock() {
                    state.total = state.total.saturating_add(1);
                    state.last_at_ms = now_ms;
                    state.burst_times.push_back(now_ms);
                }
                if let Ok(mut signals) = agent.emitted_signals.lock() {
                    signals.push(EmittedSignal {
                        signal_id: accepted.signal_id.clone(),
                        severity: args.severity.clone(),
                        summary: args.summary.clone(),
                    });
                }
            }

            let output = serde_json::json!({
                "signal_id": accepted.signal_id,
                "status": if accepted.status == SignalAcceptedStatus::Accepted {
                    "accepted"
                } else {
                    "deduplicated"
                },
                "delivery": accepted.delivery.as_str(),
            })
            .to_string();
            Ok(ToolOutput {
                result: ToolResult {
                    success: true,
                    output,
                    error: None,
                },
                artifacts: Vec::new(),
            })
        }
        .await;
        result.map_err(|error| anyhow::anyhow!(error))
    }
}

/// Maximum JSON nesting depth, counting objects/arrays.
fn json_depth(value: &Value) -> usize {
    match value {
        Value::Object(map) => 1 + map.values().map(json_depth).max().unwrap_or(0),
        Value::Array(items) => 1 + items.iter().map(json_depth).max().unwrap_or(0),
        _ => 0,
    }
}

/// Helper for tests and non-coordinator embedders: build the durable inbox
/// event descriptor for a signal without enforcing runtime state.
pub fn build_signal_event(
    context: &AgentExecutionContext,
    input: &SignalInput,
    event_id: String,
    delivery: AgentEventDelivery,
) -> NewInboxEvent {
    NewInboxEvent {
        id: event_id,
        root_session_id: context.root_session_id.clone(),
        scope_kind: "run".to_string(),
        scope_id: context.run_id.clone(),
        run_id: Some(context.run_id.clone()),
        group_id: context.group_id.clone(),
        event_type: AgentEventType::Signal,
        event_key: input.event_key.clone(),
        delivery,
        requires_continuation: true,
        severity: Some(input.severity.clone()),
        payload_json: serde_json::json!({
            "kind": "signal",
            "severity": input.severity,
            "summary": input.summary,
            "details": input.details,
            "dedupe_key": input.dedupe_key,
            "key": input.key,
            "run_id": context.run_id,
            "agent_id": context.current_agent_id,
        })
        .to_string(),
    }
}

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

    #[test]
    fn details_depth_is_measured() {
        assert_eq!(json_depth(&Value::Null), 0);
        assert_eq!(json_depth(&Value::String("x".into())), 0);
        assert_eq!(
            json_depth(&serde_json::json!({"a": {"b": [1, {"c": 2}]}})),
            4
        );
    }

    #[test]
    fn default_contract_limits_are_sane() {
        let contract = SignalContract::default();
        assert_eq!(contract.max_burst, 5);
        assert!(
            contract
                .severity_allowlist
                .contains(&"critical".to_string())
        );
    }
}