tiny-agent 0.1.0

use super::AgentRunTime;
use crate::{
    middleware::{MiddlewareCtx, ToolDecision},
    sandbox::Sandbox,
    shared::{ContentBlock, ToolCall, UserInteraction},
    state::AgentState,
    tools::ToolOutcome,
    trajectory::TrajectoryEventKind,
};
use serde_json::{Value, json};
use std::sync::Arc;
use tokio_util::sync::CancellationToken;

/// 单个工具跑完后的产出,交回 `execute_pending_tools` 按序落库。
pub(super) enum SingleOutcome {
    /// 正常完成(含工具自身的 is_error 结果),直接落库。
    Done(ContentBlock),
    /// 工具要求用户输入,需把 agent 停到 `WaitingForUser`。
    NeedsUserInteraction {
        result: ContentBlock,
        interaction: UserInteraction,
    },
    /// 中间件(before/after)硬出错 —— 整批转 Fail(`FailureKind::Middleware`)。
    MiddlewareErr {
        message: String,
        tool_executed: bool,
    },
    /// 工具执行返回 Err —— 该工具失败,整批转 Fail。
    ToolErr(String),
    /// 执行中被取消。
    Cancelled,
}

enum ToolExecution {
    Completed(ContentBlock),
    NeedsUserInteraction {
        result: ContentBlock,
        interaction: UserInteraction,
    },
}

impl ToolExecution {
    fn result(&self) -> &ContentBlock {
        match self {
            Self::Completed(result) | Self::NeedsUserInteraction { result, .. } => result,
        }
    }

    fn result_mut(&mut self) -> &mut ContentBlock {
        match self {
            Self::Completed(result) | Self::NeedsUserInteraction { result, .. } => result,
        }
    }

    fn into_user_interaction(self) -> Option<UserInteraction> {
        match self {
            Self::Completed(_) => None,
            Self::NeedsUserInteraction { interaction, .. } => Some(interaction),
        }
    }
}

impl AgentRunTime {
    /// 跑完单个工具的完整流水线:before_tool → 执行 → after_tool,产出一份 `SingleOutcome`。
    ///
    /// 这里**只产出结果,不碰 transcript / cursor / checkpoint** —— 落库与状态推进由
    /// `execute_pending_tools` 在并发段汇合后按序统一处理,从而既能并发执行又保持有序落库。
    pub(super) async fn run_one(
        &self,
        ctx: &AgentState,
        sandbox: Arc<dyn Sandbox>,
        tool_call: &ToolCall,
        cancellation_token: CancellationToken,
    ) -> SingleOutcome {
        let decision = {
            let mw_ctx = MiddlewareCtx {
                session_id: &ctx.session_id,
                state: ctx,
            };
            match self.middleware.before_tool(&mw_ctx, tool_call).await {
                Ok(decision) => decision,
                Err(e) => {
                    return SingleOutcome::MiddlewareErr {
                        message: e,
                        tool_executed: false,
                    };
                }
            }
        };

        self.emit(
            &ctx.session_id,
            TrajectoryEventKind::ToolStarted {
                call_id: tool_call.call_id.clone(),
                name: tool_call.tool_name.clone(),
                arguments: tool_call.arguments.clone(),
            },
        )
        .await;

        let mut execution = match decision {
            ToolDecision::ShortCircuit { output, is_error } => {
                ToolExecution::Completed(ContentBlock::ToolResult {
                    tool_use_id: tool_call.call_id.clone(),
                    content: output,
                    is_error,
                })
            }
            ToolDecision::Proceed => match self
                .execute_tool(sandbox, tool_call, cancellation_token)
                .await
            {
                Ok(Some(Ok(execution))) => execution,
                Ok(Some(Err(message))) | Err(message) => {
                    return SingleOutcome::ToolErr(message);
                }
                Ok(None) => return SingleOutcome::Cancelled,
            },
        };

        {
            let mw_ctx = MiddlewareCtx {
                session_id: &ctx.session_id,
                state: ctx,
            };
            if let Err(e) = self
                .middleware
                .after_tool(&mw_ctx, tool_call, execution.result_mut())
                .await
            {
                return SingleOutcome::MiddlewareErr {
                    message: e,
                    tool_executed: true,
                };
            }
        }

        let result = execution.result().clone();
        let (output, is_error) = tool_result_parts(&result);
        self.emit(
            &ctx.session_id,
            TrajectoryEventKind::ToolCompleted {
                call_id: tool_call.call_id.clone(),
                name: tool_call.tool_name.clone(),
                output,
                is_error,
            },
        )
        .await;

        match execution.into_user_interaction() {
            Some(interaction) => SingleOutcome::NeedsUserInteraction {
                result,
                interaction,
            },
            None => SingleOutcome::Done(result),
        }
    }

    async fn execute_tool(
        &self,
        sandbox: Arc<dyn Sandbox>,
        tool_call: &ToolCall,
        cancellation_token: CancellationToken,
    ) -> Result<Option<Result<ToolExecution, String>>, String> {
        let tool_future =
            self.tools
                .call(&tool_call.tool_name, tool_call.arguments.clone(), sandbox);
        tokio::pin!(tool_future);

        let output = if self.tool_timeout.is_zero() {
            tokio::select! {
                _ = cancellation_token.cancelled() => return Ok(None),
                output = &mut tool_future => output,
            }
        } else {
            let timeout = tokio::time::sleep(self.tool_timeout);
            tokio::pin!(timeout);
            tokio::select! {
                _ = cancellation_token.cancelled() => return Ok(None),
                _ = &mut timeout => {
                    return Ok(Some(Err(format!(
                        "tool timed out after {:?}",
                        self.tool_timeout
                    ))));
                }
                output = &mut tool_future => output,
            }
        };

        match output {
            Ok(ToolOutcome::Completed(value)) => Ok(Some(Ok(ToolExecution::Completed(
                ContentBlock::ToolResult {
                    tool_use_id: tool_call.call_id.clone(),
                    content: tool_output_blocks(value),
                    is_error: false,
                },
            )))),
            Ok(ToolOutcome::NeedsUserInteraction(interaction)) => {
                Ok(Some(Ok(ToolExecution::NeedsUserInteraction {
                    result: needs_user_interaction_result(&tool_call.call_id, &interaction),
                    interaction,
                })))
            }
            Err(err) => Ok(Some(Err(err))),
        }
    }
}

fn needs_user_interaction_result(tool_use_id: &str, interaction: &UserInteraction) -> ContentBlock {
    ContentBlock::ToolResult {
        tool_use_id: tool_use_id.to_string(),
        content: vec![ContentBlock::Text {
            text: json!({
                "type": "needs_user_interaction",
                "interaction": interaction,
            })
            .to_string(),
        }],
        is_error: false,
    }
}

fn tool_output_blocks(value: Value) -> Vec<ContentBlock> {
    let text = match value {
        Value::String(text) => text,
        other => other.to_string(),
    };
    vec![ContentBlock::Text { text }]
}

/// 从工具结果块里取出 (内容块, 是否出错),用于轨迹事件。
fn tool_result_parts(block: &ContentBlock) -> (Vec<ContentBlock>, bool) {
    match block {
        ContentBlock::ToolResult {
            content, is_error, ..
        } => (content.clone(), *is_error),
        other => (vec![other.clone()], false),
    }
}