Hare/llm_worker_rs

Fork 0

Hare 5691b09fc8 feat: Implement HookEventKind

2026-01-09 19:18:20 +09:00

12 KiB

Raw Blame History

Hooks 設計

概要

HookはWorker層でのターン制御に介入するためのメカニズムです。

メッセージ送信・ツール実行・ターン終了等の各ポイントで処理を差し込むことができます。

コンセプト

制御の介入: ターンの進行、メッセージの内容、ツールの実行に対して介入
Contextへのアクセス: メッセージ履歴を読み書き可能
非破壊的チェーン: 複数のHookを登録順に実行、後続Hookへの影響を制御

Hook一覧

Hook	タイミング	主な用途	戻り値
`on_message_send`	LLM送信前	コンテキスト改変/検証	`OnMessageSendResult`
`before_tool_call`	ツール実行前	実行許可/引数改変	`BeforeToolCallResult`
`after_tool_call`	ツール実行後	結果加工/マスキング	`AfterToolCallResult`
`on_turn_end`	ツールなしでターン終了直前	検証/リトライ指示	`OnTurnEndResult`
`on_abort`	中断時	クリーンアップ/通知	`()`

Hook Trait

#[async_trait]
pub trait Hook<E: HookEventKind>: Send + Sync {
    async fn call(&self, input: &mut E::Input) -> Result<E::Output, HookError>;
}

制御フロー型

HookEventKind / Result

Hookイベントごとに入力/出力型を分離し、意味のない制御フローを排除する。

pub trait HookEventKind {
    type Input;
    type Output;
}

pub struct OnMessageSend;
pub struct BeforeToolCall;
pub struct AfterToolCall;
pub struct OnTurnEnd;
pub struct OnAbort;

pub enum OnMessageSendResult {
    Continue,
    Cancel(String),
}

pub enum BeforeToolCallResult {
    Continue,
    Skip,
    Abort(String),
    Pause,
}

pub enum AfterToolCallResult {
    Continue,
    Abort(String),
}

pub enum OnTurnEndResult {
    Finish,
    ContinueWithMessages(Vec<Message>),
    Paused,
}

Tool Call Context

before_tool_call / after_tool_call は、ツール実行の文脈を含む入力を受け取る。

pub struct ToolCallContext {
    pub call: ToolCall,
    pub meta: ToolMeta,      // 不変メタデータ
    pub tool: Arc<dyn Tool>, // 状態アクセス用
}

pub struct ToolResultContext {
    pub result: ToolResult,
    pub meta: ToolMeta,
    pub tool: Arc<dyn Tool>,
}

呼び出しタイミング

Worker::run() ループ
│
├─▶ on_message_send ──────────────────────────────┐
│   コンテキストの改変、バリデーション、          │
│   システムプロンプト注入などが可能              │
│                                                 │
├─▶ LLMリクエスト送信 & ストリーム処理            │
│                                                 │
├─▶ ツール呼び出しがある場合:                     │
│   │                                             │
│   ├─▶ before_tool_call (各ツールごと・逐次)     │
│   │   実行可否の判定、引数の改変                │
│   │                                             │
│   ├─▶ ツール並列実行 (join_all)                 │
│   │                                             │
│   └─▶ after_tool_call (各結果ごと・逐次)        │
│       結果の確認、加工、ログ出力                │
│                                                 │
├─▶ ツール結果をコンテキストに追加 → ループ先頭へ │
│                                                 │
└─▶ ツールなしの場合:                             │
    │                                             │
    └─▶ on_turn_end  ─────────────────────────────┘
        最終応答のチェック（Lint/Fmt等）
        エラーがあればContinueWithMessagesでリトライ

※ 中断時は on_abort が呼ばれる

各Hookの詳細

on_message_send

呼び出しタイミング: LLMへリクエスト送信前（ターンループの冒頭）

用途:

コンテキストへのシステムメッセージ注入
メッセージのバリデーション
機密情報のフィルタリング
リクエスト内容のログ出力
OnMessageSendResult::Cancel による送信キャンセル

例: メッセージにタイムスタンプを追加

struct TimestampHook;

#[async_trait]
impl Hook<OnMessageSend> for TimestampHook {
    async fn call(
        &self,
        context: &mut Vec<Message>,
    ) -> Result<OnMessageSendResult, HookError> {
        let timestamp = chrono::Local::now().to_rfc3339();
        context.insert(0, Message::user(format!("[{}]", timestamp)));
        Ok(OnMessageSendResult::Continue)
    }
}

before_tool_call

呼び出しタイミング: 各ツール実行前（並列実行フェーズの前）

用途:

危険なツールのブロック
引数のサニタイズ
確認プロンプトの表示（UIとの連携）
実行ログの記録
BeforeToolCallResult::Pause による一時停止

入力:

ToolCallContext（ToolCall + ToolMeta + Arc<dyn Tool>）

例: 特定ツールをブロック

struct ToolBlocker {
    blocked_tools: HashSet<String>,
}

#[async_trait]
impl Hook<BeforeToolCall> for ToolBlocker {
    async fn call(
        &self,
        ctx: &mut ToolCallContext,
    ) -> Result<BeforeToolCallResult, HookError> {
        if self.blocked_tools.contains(&ctx.call.name) {
            println!("Blocked tool: {}", ctx.call.name);
            Ok(BeforeToolCallResult::Skip)
        } else {
            Ok(BeforeToolCallResult::Continue)
        }
    }
}

after_tool_call

呼び出しタイミング: 各ツール実行後（並列実行フェーズの後）

用途:

結果の加工・フォーマット
機密情報のマスキング
結果のキャッシュ
実行結果のログ出力入力:
ToolResultContext（ToolResult + ToolMeta + Arc<dyn Tool>）

例: 結果にプレフィックスを追加

struct ResultFormatter;

#[async_trait]
impl Hook<AfterToolCall> for ResultFormatter {
    async fn call(
        &self,
        ctx: &mut ToolResultContext,
    ) -> Result<AfterToolCallResult, HookError> {
        if !ctx.result.is_error {
            ctx.result.content = format!("[OK] {}", ctx.result.content);
        }
        Ok(AfterToolCallResult::Continue)
    }
}

on_turn_end

呼び出しタイミング: ツール呼び出しなしでターンが終了する直前

用途:

生成されたコードのLint/Fmt
出力形式のバリデーション
自己修正のためのリトライ指示
最終結果のログ出力
OnTurnEndResult::Paused による一時停止

on_abort

呼び出しタイミング: キャンセル/エラー/AbortなどでWorkerが中断された時

用途:

クリーンアップ処理
中断理由のログ出力
外部システムへの通知

例: JSON形式のバリデーション

struct JsonValidator;

#[async_trait]
impl Hook<OnTurnEnd> for JsonValidator {
    async fn call(
        &self,
        messages: &mut Vec<Message>,
    ) -> Result<OnTurnEndResult, HookError> {
        // 最後のアシスタントメッセージを取得
        let last = messages.iter().rev()
            .find(|m| m.role == Role::Assistant);
        
        if let Some(msg) = last {
            if let MessageContent::Text(text) = &msg.content {
                // JSONとしてパースを試みる
                if serde_json::from_str::<serde_json::Value>(text).is_err() {
                    // 失敗したらリトライ指示
                    return Ok(OnTurnEndResult::ContinueWithMessages(vec![
                        Message::user("Invalid JSON. Please fix and try again.")
                    ]));
                }
            }
        }
        Ok(OnTurnEndResult::Finish)
    }
}

複数Hookの実行順序

Hookはイベントごとに登録順に実行されます。

worker.add_before_tool_call_hook(HookA);  // 1番目に実行
worker.add_before_tool_call_hook(HookB);  // 2番目に実行
worker.add_before_tool_call_hook(HookC);  // 3番目に実行

制御フローの伝播

Continue: 後続Hookも実行
Skip: 現在の処理をスキップし、後続Hookは実行しない
Abort: 即座にエラーを返し、処理全体を中断
Pause: Workerを一時停止（再開はresume）

Hook A: Continue → Hook B: Skip → (Hook Cは実行されない)
                                   ↓
                            処理をスキップ

Hook A: Continue → Hook B: Abort("reason")
                           ↓
                     WorkerError::Aborted

Hook A: Continue → Hook B: Pause
                           ↓
                     WorkerResult::Paused

設計上のポイント

1. イベントごとの実装

必要なイベントのみ Hook<Event> を実装する。

2. 可変参照による改変

&mutで引数を受け取るため、直接改変が可能。

async fn call(&self, ctx: &mut ToolCallContext) -> ... {
    // 引数を直接書き換え
    ctx.call.input["sanitized"] = json!(true);
    Ok(BeforeToolCallResult::Continue)
}

3. 並列実行との統合

before_tool_call: 並列実行前に逐次実行（許可判定のため）
ツール実行: join_allで並列実行
after_tool_call: 並列実行後に逐次実行（結果加工のため）

4. Send + Sync 要件

HookはSend + Syncを要求するため、スレッドセーフな実装が必要。状態を持つ場合はArc<Mutex<T>>やAtomicUsizeなどを使用する。

struct CountingHook {
    count: Arc<AtomicUsize>,
}

#[async_trait]
impl Hook<BeforeToolCall> for CountingHook {
    async fn call(&self, _: &mut ToolCallContext) -> Result<BeforeToolCallResult, HookError> {
        self.count.fetch_add(1, Ordering::SeqCst);
        Ok(BeforeToolCallResult::Continue)
    }
}

典型的なユースケース

ユースケース	使用Hook	処理内容
ツール許可制御	`before_tool_call`	危険なツールをSkip
実行ログ	`before/after_tool_call`	呼び出しと結果を記録
出力バリデーション	`on_turn_end`	形式チェック、リトライ指示
コンテキスト注入	`on_message_send`	システムメッセージ追加
結果のサニタイズ	`after_tool_call`	機密情報のマスキング
レート制限	`before_tool_call`	呼び出し頻度の制御

TODO

Hooks仕様の厳密な再定義

現在のHooks実装は基本的なユースケースをカバーしているが、以下の点について将来的に厳密な仕様を定義する必要がある：

エラーハンドリングの明確化: HookError発生時のリカバリー戦略、部分的な失敗の扱い
Hook間の依存関係: 複数Hookの実行順序が結果に影響する場合のセマンティクス
非同期キャンセル: Hook実行中のキャンセル（タイムアウト等）の振る舞い
状態の一貫性: on_message_sendで改変されたコンテキストが後続処理で期待通りに反映される保証
リトライ制限: on_turn_endでのContinueWithMessagesによる無限ループ防止策
Hook優先度: 登録順以外の優先度指定メカニズムの必要性
条件付きHook: 特定条件でのみ有効化されるHookパターン
テスト容易性: Hookのモック/スタブ作成のためのユーティリティ

12 KiB Raw Blame History Unescape Escape