llm-worker 0.2.0

//! Hook関連の型定義
//!
//! Worker層でのターン制御・介入に使用される型

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;

// =============================================================================
// Hook Event Kinds
// =============================================================================

pub trait HookEventKind: Send + Sync + 'static {
    type Input;
    type Output;
}

pub struct OnPromptSubmit;
pub struct PreLlmRequest;
pub struct PreToolCall;
pub struct PostToolCall;
pub struct OnTurnEnd;
pub struct OnAbort;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OnPromptSubmitResult {
    Continue,
    Cancel(String),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PreLlmRequestResult {
    Continue,
    Cancel(String),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PreToolCallResult {
    Continue,
    Skip,
    Abort(String),
    Pause,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PostToolCallResult {
    Continue,
    Abort(String),
}

#[derive(Debug, Clone)]
pub enum OnTurnEndResult {
    Finish,
    ContinueWithMessages(Vec<crate::Message>),
    Paused,
}

use std::sync::Arc;

use crate::tool::{Tool, ToolMeta};

/// PreToolCall の入力コンテキスト
pub struct ToolCallContext {
    /// ツール呼び出し情報（改変可能）
    pub call: ToolCall,
    /// ツールメタ情報（不変）
    pub meta: ToolMeta,
    /// ツールインスタンス（状態アクセス用）
    pub tool: Arc<dyn Tool>,
}

/// PostToolCall の入力コンテキスト
pub struct PostToolCallContext {
    /// ツール呼び出し情報
    pub call: ToolCall,
    /// ツール実行結果（改変可能）
    pub result: ToolResult,
    /// ツールメタ情報（不変）
    pub meta: ToolMeta,
    /// ツールインスタンス（状態アクセス用）
    pub tool: Arc<dyn Tool>,
}

impl HookEventKind for OnPromptSubmit {
    type Input = crate::Message;
    type Output = OnPromptSubmitResult;
}

impl HookEventKind for PreLlmRequest {
    type Input = Vec<crate::Message>;
    type Output = PreLlmRequestResult;
}

impl HookEventKind for PreToolCall {
    type Input = ToolCallContext;
    type Output = PreToolCallResult;
}

impl HookEventKind for PostToolCall {
    type Input = PostToolCallContext;
    type Output = PostToolCallResult;
}

impl HookEventKind for OnTurnEnd {
    type Input = Vec<crate::Message>;
    type Output = OnTurnEndResult;
}

impl HookEventKind for OnAbort {
    type Input = String;
    type Output = ();
}

// =============================================================================
// Tool Call / Result Types
// =============================================================================

/// ツール呼び出し情報
///
/// LLMからのToolUseブロックを表現し、Hook処理で改変可能
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    /// ツール呼び出しID（レスポンスとの紐付けに使用）
    pub id: String,
    /// ツール名
    pub name: String,
    /// 入力引数（JSON）
    pub input: Value,
}

/// ツール実行結果
///
/// ツール実行後の結果を表現し、Hook処理で改変可能
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
    /// 対応するツール呼び出しID
    pub tool_use_id: String,
    /// 結果コンテンツ
    pub content: String,
    /// エラーかどうか
    #[serde(default)]
    pub is_error: bool,
}

impl ToolResult {
    /// 成功結果を作成
    pub fn success(tool_use_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            tool_use_id: tool_use_id.into(),
            content: content.into(),
            is_error: false,
        }
    }

    /// エラー結果を作成
    pub fn error(tool_use_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            tool_use_id: tool_use_id.into(),
            content: content.into(),
            is_error: true,
        }
    }
}

// =============================================================================
// Hook Error
// =============================================================================

/// Hookエラー
#[derive(Debug, Error)]
pub enum HookError {
    /// 処理が中断された
    #[error("Aborted: {0}")]
    Aborted(String),
    /// 内部エラー
    #[error("Hook error: {0}")]
    Internal(String),
}

// =============================================================================
// Hook Trait
// =============================================================================

/// Hookイベントの処理を行うトレイト
///
/// 各イベント種別は戻り値型が異なるため、`HookEventKind`を介して型を制約する。
#[async_trait]
pub trait Hook<E: HookEventKind>: Send + Sync {
    async fn call(&self, input: &mut E::Input) -> Result<E::Output, HookError>;
}

// =============================================================================
// Hook Registry
// =============================================================================

/// 全 Hook を保持するレジストリ
///
/// Worker 内部で使用され、各種 Hook を一括管理する。
pub struct HookRegistry {
    /// on_prompt_submit Hook
    pub(crate) on_prompt_submit: Vec<Box<dyn Hook<OnPromptSubmit>>>,
    /// pre_llm_request Hook
    pub(crate) pre_llm_request: Vec<Box<dyn Hook<PreLlmRequest>>>,
    /// pre_tool_call Hook
    pub(crate) pre_tool_call: Vec<Box<dyn Hook<PreToolCall>>>,
    /// post_tool_call Hook
    pub(crate) post_tool_call: Vec<Box<dyn Hook<PostToolCall>>>,
    /// on_turn_end Hook
    pub(crate) on_turn_end: Vec<Box<dyn Hook<OnTurnEnd>>>,
    /// on_abort Hook
    pub(crate) on_abort: Vec<Box<dyn Hook<OnAbort>>>,
}

impl Default for HookRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl HookRegistry {
    /// 空の HookRegistry を作成
    pub fn new() -> Self {
        Self {
            on_prompt_submit: Vec::new(),
            pre_llm_request: Vec::new(),
            pre_tool_call: Vec::new(),
            post_tool_call: Vec::new(),
            on_turn_end: Vec::new(),
            on_abort: Vec::new(),
        }
    }
}