phi-agent 0.18.1

//! Custom Policy — demonstrate ToolPolicy, Middleware, and RuntimeEvent hooks.
//!
//! This example shows the complete hook/event system without needing an API key:
//! 1. ToolPolicy  — evaluate_approval / before_call / after_call
//! 2. Middleware   — on_user_message / on_post_llm
//! 3. RuntimeEvent — collecting and printing the full event stream
//!
//! Uses a Mock LLM client so it runs offline.
//!
//! Run with:
//! ```bash
//! cargo run --example custom-policy
//! ```

use std::collections::VecDeque;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU32, Ordering};
use std::task::{Context, Poll};

use agent_base::llm_trait::response::{FinishReason, ToolCall};
use agent_base::llm_trait::{Capabilities, ChatRequest, ChatResponse, ChatStream, LlmError, LlmProvider, ProviderInfo};
use agent_base::{
    AgentResult, ApprovalRequest, Content, Middleware, PostLlmCtx, RuntimeEvent, StreamChunk, Tool, ToolContext,
    ToolDecision, ToolPolicy,
};
use async_trait::async_trait;
use futures_core::Stream;
use phi_agent::{ApprovalMode, AutoApprovalHandler, base_agent_builder};
use serde_json::{Value, json};

// ── Mock LLM that returns a tool call then text ──

/// A simple stream that yields items from a queue.
struct QueueStream {
    items: VecDeque<Result<StreamChunk, LlmError>>,
}

impl Stream for QueueStream {
    type Item = Result<StreamChunk, LlmError>;
    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Poll::Ready(self.items.pop_front())
    }
}

struct MockLlmClient {
    call_count: AtomicU32,
}

impl MockLlmClient {
    fn new() -> Self {
        Self { call_count: AtomicU32::new(0) }
    }
}

#[async_trait]
impl LlmProvider for MockLlmClient {
    async fn stream(&self, _request: ChatRequest) -> Result<ChatStream, LlmError> {
        let count = self.call_count.fetch_add(1, Ordering::SeqCst);
        let mut items = VecDeque::new();
        if count == 0 {
            // First call: emit a tool call chunk
            items.push_back(Ok(StreamChunk::ToolCall(json!({
                "delta": {
                    "tool_calls": [{
                        "index": 0,
                        "id": "call_1",
                        "function": {
                            "name": "echo",
                            "arguments": "{\"msg\": \"hello from hook demo\"}"
                        }
                    }]
                }
            }))));
        } else {
            // Second call: emit text
            items.push_back(Ok(StreamChunk::Text("All hooks and events demonstrated!".to_string())));
        }
        items.push_back(Ok(StreamChunk::Stop { finish_reason: Some("stop".to_string()) }));
        Ok(ChatStream::new(Box::pin(QueueStream { items })))
    }

    async fn chat(&self, _request: ChatRequest) -> Result<ChatResponse, LlmError> {
        let count = self.call_count.fetch_add(1, Ordering::SeqCst);
        if count == 0 {
            Ok(ChatResponse {
                content: String::new(),
                reasoning_content: None,
                tool_calls: vec![ToolCall {
                    id: "call_1".to_string(),
                    name: "echo".to_string(),
                    arguments: "{\"msg\": \"hello from hook demo\"}".to_string(),
                }],
                usage: agent_base::UsageInfo::default(),
                finish_reason: FinishReason::ToolCalls,
                raw: None,
                thinking_signature: None,
            })
        } else {
            Ok(ChatResponse {
                content: "All hooks and events demonstrated!".to_string(),
                reasoning_content: None,
                tool_calls: vec![],
                usage: agent_base::UsageInfo::default(),
                finish_reason: FinishReason::Stop,
                raw: None,
                thinking_signature: None,
            })
        }
    }

    fn capabilities(&self) -> Capabilities {
        Capabilities { supports_streaming: true, supports_tools: true, ..Default::default() }
    }

    fn info(&self) -> ProviderInfo {
        ProviderInfo { name: "mock".into(), model: "mock".into(), version: None }
    }
}

// ── A simple echo tool ──

struct EchoTool;

#[async_trait]
impl Tool for EchoTool {
    fn name(&self) -> &'static str {
        "echo"
    }

    fn description(&self) -> &'static str {
        "Echo a message back"
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "msg": { "type": "string", "description": "The message to echo" }
            },
            "required": ["msg"]
        })
    }

    async fn call(&self, args: &Value, _ctx: &ToolContext) -> AgentResult<Vec<Content>> {
        let msg = args["msg"].as_str().unwrap_or("ok");
        println!("    → EchoTool.call() executing with msg=\"{msg}\"");
        Ok(vec![Content::text(format!("EchoTool executed: {msg}"))])
    }
}

// ── Custom ToolPolicy ──

struct DemoPolicy {
    before_count: AtomicU32,
    after_count: AtomicU32,
}

impl DemoPolicy {
    fn new() -> Self {
        Self { before_count: AtomicU32::new(0), after_count: AtomicU32::new(0) }
    }
}

#[async_trait]
impl ToolPolicy for DemoPolicy {
    /// 1. evaluate_approval — decide if tool needs user approval (async)
    async fn evaluate_approval(&self, tool_name: &str, args: &Value) -> Option<ApprovalRequest> {
        let msg = args.get("msg").and_then(Value::as_str).unwrap_or("");
        println!("  [ToolPolicy] evaluate_approval('{tool_name}', msg=\"{msg}\") → no approval needed");
        None
    }

    /// 2. before_call — check before execution, return Block to abort
    fn before_call(&self, tool_name: &str, _args: &Value, _ctx: &ToolContext) -> AgentResult<ToolDecision> {
        self.before_count.fetch_add(1, Ordering::SeqCst);
        println!("  [ToolPolicy] before_call('{tool_name}') — about to execute");
        Ok(ToolDecision::Proceed)
    }

    /// 3. after_call — callback after successful execution
    fn after_call(&self, tool_name: &str, _args: &Value, result: &[Content], _ctx: &ToolContext) -> AgentResult<()> {
        self.after_count.fetch_add(1, Ordering::SeqCst);
        println!("  [ToolPolicy] after_call('{tool_name}') — result: {}", agent_base::tool::content_text(result));
        Ok(())
    }
}

// ── Custom Middleware ──

struct DemoMiddleware;

#[async_trait]
impl Middleware for DemoMiddleware {
    async fn on_user_message(&self, ctx: &mut agent_base::UserMessageCtx) -> AgentResult<()> {
        println!("  [Middleware] on_user_message — input: \"{}\"", ctx.user_input);
        Ok(())
    }

    async fn on_post_llm(&self, ctx: &mut PostLlmCtx) -> AgentResult<()> {
        if ctx.is_tool_call {
            println!("  [Middleware] on_post_llm — is_tool_call=true, {} tool(s) requested", ctx.tool_calls.len(),);
        } else {
            println!("  [Middleware] on_post_llm — is_tool_call=false, text=\"{}\"", ctx.full_text,);
        }
        Ok(())
    }
}

// ── Main ──

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    println!("=== Custom Policy Demo ===\n");
    println!("This example demonstrates ToolPolicy, Middleware, and RuntimeEvent hooks.\n");

    let llm_client = Arc::new(MockLlmClient::new());
    let policy = Arc::new(DemoPolicy::new());

    let runtime = base_agent_builder(llm_client)
        .system_prompt("You are a helpful assistant.".to_string())
        .register_tool(EchoTool)
        .tool_policy(policy.clone())
        .middleware(DemoMiddleware)
        .approval_handler(Arc::new(AutoApprovalHandler::new(ApprovalMode::Auto)))
        .build()?;

    let session_id = runtime.create_session().await;

    println!("--- Running turn: 'echo hello from hook demo' ---\n");

    let events = Arc::new(Mutex::new(Vec::new()));
    let events_clone = events.clone();
    runtime
        .run_turn(session_id, "echo hello from hook demo", move |event| {
            match &event {
                RuntimeEvent::ToolCallStarted { tool_name, args_json, .. } => {
                    println!("  [Event] ToolCallStarted   — tool=\"{tool_name}\", args={args_json}");
                },
                RuntimeEvent::ToolCallFinished { tool_name, summary, .. } => {
                    println!("  [Event] ToolCallFinished  — tool=\"{tool_name}\", summary=\"{summary}\"");
                },
                RuntimeEvent::TextDelta { text, .. } => {
                    println!("  [Event] TextDelta         — \"{text}\"");
                },
                RuntimeEvent::AwaitingApproval { request, .. } => {
                    println!(
                        "  [Event] AwaitingApproval  — risk={:?}, title=\"{}\"",
                        request.risk_level, request.title,
                    );
                },
                RuntimeEvent::RunFinished { .. } => {
                    println!("  [Event] RunFinished");
                },
                RuntimeEvent::UserEvent { event: agent_base::UserEvent::Progress { text }, .. } => {
                    println!("  [Event] UserEvent(Progress) — \"{text}\"");
                },
                _ => {},
            }
            events_clone.lock().unwrap().push(event);
            Ok(())
        })
        .await?;

    println!("\n--- Summary ---");
    println!("Total events captured: {}", events.lock().unwrap().len());
    println!("ToolPolicy before_call count: {}", policy.before_count.load(Ordering::SeqCst),);
    println!("ToolPolicy after_call count:  {}", policy.after_count.load(Ordering::SeqCst),);
    println!("\nHook execution order:");
    println!("  on_user_message → on_post_llm → evaluate_approval → before_call →");
    println!("  ToolCallStarted → tool.call() → ToolCallFinished → after_call →");
    println!("  on_post_llm → TextDelta → RunFinished");
    println!("\n=== Done ===");

    Ok(())
}