PicoBot/src/tools/traits.rs

use crate::bus::{MediaItem, MediaRef, MessageSource};
use async_trait::async_trait;

/// Session identity supplied by the runtime for tools that own external state.
/// Ordinary stateless tools can ignore it through the default trait method.
#[derive(Debug, Clone)]
pub struct ToolExecutionContext {
    pub session_id: Option<String>,
    pub turn_id: Option<String>,
    pub agent: Option<std::sync::Arc<crate::agent::AgentExecutionContext>>,
    pub cancellation: tokio_util::sync::CancellationToken,
    pub execution_gate: Option<std::sync::Arc<crate::agent::gate::ExecutionGate>>,
}

impl Default for ToolExecutionContext {
    fn default() -> Self {
        Self {
            session_id: None,
            turn_id: None,
            agent: None,
            cancellation: tokio_util::sync::CancellationToken::new(),
            execution_gate: None,
        }
    }
}

impl ToolExecutionContext {
    pub fn for_session(session_id: impl Into<String>) -> Self {
        Self {
            session_id: Some(session_id.into()),
            turn_id: None,
            agent: None,
            cancellation: tokio_util::sync::CancellationToken::new(),
            execution_gate: None,
        }
    }

    pub fn with_turn_id(mut self, turn_id: impl Into<String>) -> Self {
        self.turn_id = Some(turn_id.into());
        self
    }

    pub fn with_agent(
        mut self,
        agent: std::sync::Arc<crate::agent::AgentExecutionContext>,
    ) -> Self {
        self.agent = Some(agent);
        self
    }

    pub fn with_cancellation(mut self, cancellation: tokio_util::sync::CancellationToken) -> Self {
        self.cancellation = cancellation;
        self
    }

    pub fn with_execution_gate(
        mut self,
        gate: std::sync::Arc<crate::agent::gate::ExecutionGate>,
    ) -> Self {
        self.execution_gate = Some(gate);
        self
    }
}

#[derive(Debug, Clone)]
pub struct ToolResult {
    pub success: bool,
    pub output: String,
    pub error: Option<String>,
}

/// The intended consumers of a structured artifact returned by a tool.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolArtifactAudience {
    Model,
    User,
    ModelAndUser,
}

impl ToolArtifactAudience {
    fn includes_model(self) -> bool {
        matches!(self, Self::Model | Self::ModelAndUser)
    }

    fn includes_user(self) -> bool {
        matches!(self, Self::User | Self::ModelAndUser)
    }
}

/// A structured artifact plus its semantic delivery intent. Tools declare
/// intent; the common output processor decides how each consumer receives it.
#[derive(Debug, Clone)]
pub struct ToolArtifact {
    pub media_ref: MediaRef,
    pub audience: ToolArtifactAudience,
}

impl ToolArtifact {
    pub fn model_only(media_ref: MediaRef) -> Self {
        Self {
            media_ref,
            audience: ToolArtifactAudience::Model,
        }
    }

    pub fn model_and_user(media_ref: MediaRef) -> Self {
        Self {
            media_ref,
            audience: ToolArtifactAudience::ModelAndUser,
        }
    }

    pub fn user_only(media_ref: MediaRef) -> Self {
        Self {
            media_ref,
            audience: ToolArtifactAudience::User,
        }
    }
}

/// Unified raw output from every tool. Plain-text tools are converted through
/// `From<ToolResult>`; media-producing tools additionally declare artifacts.
#[derive(Debug, Clone)]
pub struct ToolOutput {
    pub result: ToolResult,
    pub artifacts: Vec<ToolArtifact>,
}

/// Consumer-specific result produced by the common tool-output processor.
#[derive(Debug, Clone)]
pub struct ProcessedToolOutput {
    pub result: ToolResult,
    pub model_media_refs: Vec<MediaRef>,
    pub reply_media_refs: Vec<MediaRef>,
}

pub struct ToolOutputProcessor;

impl ToolOutputProcessor {
    pub fn process(output: ToolOutput) -> ProcessedToolOutput {
        let ToolOutput { result, artifacts } = output;
        let mut model_media_refs = Vec::new();
        let mut reply_media_refs = Vec::new();

        // Failed tools do not publish artifacts that may be incomplete or
        // invalid. Their textual error still follows the ordinary tool path.
        if result.success {
            for artifact in artifacts {
                if artifact.audience.includes_model() {
                    push_unique_media(&mut model_media_refs, &artifact.media_ref);
                }
                if artifact.audience.includes_user() {
                    push_unique_media(&mut reply_media_refs, &artifact.media_ref);
                }
            }
        }

        ProcessedToolOutput {
            result,
            model_media_refs,
            reply_media_refs,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OutboundDelivery {
    Delivered,
    AttachedToCurrentTurn,
}

fn push_unique_media(target: &mut Vec<MediaRef>, media_ref: &MediaRef) {
    if !target.iter().any(|existing| {
        existing.path == media_ref.path && existing.media_type == media_ref.media_type
    }) {
        target.push(media_ref.clone());
    }
}

impl From<ToolResult> for ToolOutput {
    fn from(result: ToolResult) -> Self {
        Self {
            result,
            artifacts: Vec::new(),
        }
    }
}

#[async_trait]
pub trait Tool: Send + Sync + 'static {
    fn name(&self) -> &str;
    fn description(&self) -> &str;
    fn parameters_schema(&self) -> serde_json::Value;
    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult>;

    /// Whether this tool is injected at runtime from the caller context
    /// (delegate targets, signal contract, skill allowlist) and therefore
    /// must never be declared directly in an Agent definition's `tools`
    /// list.  Every ordinary tool returns false: which tools a named Agent
    /// receives is decided solely by its definition file, not by tool-side
    /// delegation flags.
    fn runtime_injected(&self) -> bool {
        false
    }

    /// Execute the tool through the unified output envelope. Most tools return
    /// only text and use this default conversion.
    async fn execute_output(&self, args: serde_json::Value) -> anyhow::Result<ToolOutput> {
        self.execute(args).await.map(Into::into)
    }

    /// Execute with runtime context. Stateful adapters use this to route
    /// external resources without coupling to SessionManager; a configured
    /// single-user adapter may deliberately share state across dialogs.
    async fn execute_with_context(
        &self,
        _context: &ToolExecutionContext,
        args: serde_json::Value,
    ) -> anyhow::Result<ToolOutput> {
        self.execute_output(args).await
    }

    /// Whether this tool is side-effect free and safe to parallelize.
    fn read_only(&self) -> bool {
        false
    }

    /// Whether this tool can run alongside other concurrency-safe tools.
    fn concurrency_safe(&self) -> bool {
        self.read_only() && !self.exclusive()
    }

    /// Whether this tool should run alone even if concurrency is enabled.
    fn exclusive(&self) -> bool {
        false
    }
}

#[async_trait]
pub trait OutboundMessenger: Send + Sync {
    async fn send_message(
        &self,
        channel: &str,
        chat_id: &str,
        dialog_id: Option<&str>,
        content: &str,
        source: MessageSource,
        media: Vec<MediaItem>,
    ) -> Result<OutboundDelivery, String>;
}

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

    fn media(path: &str) -> MediaRef {
        MediaRef {
            path: path.to_string(),
            media_type: "image".to_string(),
        }
    }

    #[test]
    fn processor_routes_and_deduplicates_artifacts_by_audience() {
        let output = ToolOutput {
            result: ToolResult {
                success: true,
                output: "ok".to_string(),
                error: None,
            },
            artifacts: vec![
                ToolArtifact::model_only(media("model.png")),
                ToolArtifact::model_and_user(media("shared.png")),
                ToolArtifact::model_and_user(media("shared.png")),
                ToolArtifact::user_only(media("reply.txt")),
            ],
        };

        let processed = ToolOutputProcessor::process(output);

        assert_eq!(processed.model_media_refs.len(), 2);
        assert_eq!(processed.reply_media_refs.len(), 2);
        assert_eq!(processed.reply_media_refs[0].path, "shared.png");
        assert_eq!(processed.reply_media_refs[1].path, "reply.txt");
    }

    #[test]
    fn processor_discards_artifacts_from_failed_tools() {
        let output = ToolOutput {
            result: ToolResult {
                success: false,
                output: String::new(),
                error: Some("failed".to_string()),
            },
            artifacts: vec![ToolArtifact::model_and_user(media("partial.png"))],
        };

        let processed = ToolOutputProcessor::process(output);

        assert!(processed.model_media_refs.is_empty());
        assert!(processed.reply_media_refs.is_empty());
    }
}