Hare/llm_worker_rs

Fork 0

Hare 9547d40538 docs: Add hooks_design.md

2026-01-07 00:26:10 +09:00

11 KiB

Raw Blame History

Hooks 設計

概要

HookはWorker層でのターン制御に介入するためのメカニズムです。 Claude CodeのHooks機能に着想を得ており、メッセージ送信・ツール実行・ターン終了の各ポイントで処理を差し込むことができます。

コンセプト

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

Hook Trait

#[async_trait]
pub trait WorkerHook: Send + Sync {
    /// メッセージ送信前
    /// リクエストに含まれるメッセージリストを改変できる
    async fn on_message_send(
        &self,
        context: &mut Vec<Message>,
    ) -> Result<ControlFlow, HookError> {
        Ok(ControlFlow::Continue)
    }

    /// ツール実行前
    /// 実行をキャンセルしたり、引数を書き換えることができる
    async fn before_tool_call(
        &self,
        tool_call: &mut ToolCall,
    ) -> Result<ControlFlow, HookError> {
        Ok(ControlFlow::Continue)
    }

    /// ツール実行後
    /// 結果を書き換えたり、隠蔽したりできる
    async fn after_tool_call(
        &self,
        tool_result: &mut ToolResult,
    ) -> Result<ControlFlow, HookError> {
        Ok(ControlFlow::Continue)
    }

    /// ターン終了時
    /// 生成されたメッセージを検査し、必要ならリトライを指示できる
    async fn on_turn_end(
        &self,
        messages: &[Message],
    ) -> Result<TurnResult, HookError> {
        Ok(TurnResult::Finish)
    }
}

制御フロー型

ControlFlow

Hook処理の継続/中断を制御する列挙型。

pub enum ControlFlow {
    /// 処理を続行（後続Hookも実行）
    Continue,
    /// 現在の処理をスキップ（ツール実行をスキップ等）
    Skip,
    /// 処理全体を中断（エラーとして扱う）
    Abort(String),
}

TurnResult

ターン終了時の判定結果を表す列挙型。

pub enum TurnResult {
    /// ターンを正常終了
    Finish,
    /// メッセージを追加してターン継続（自己修正など）
    ContinueWithMessages(Vec<Message>),
}

呼び出しタイミング

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

各Hookの詳細

on_message_send

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

用途:

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

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

struct TimestampHook;

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

before_tool_call

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

用途:

危険なツールのブロック
引数のサニタイズ
確認プロンプトの表示（UIとの連携）
実行ログの記録

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

struct ToolBlocker {
    blocked_tools: HashSet<String>,
}

#[async_trait]
impl WorkerHook for ToolBlocker {
    async fn before_tool_call(
        &self,
        tool_call: &mut ToolCall,
    ) -> Result<ControlFlow, HookError> {
        if self.blocked_tools.contains(&tool_call.name) {
            println!("Blocked tool: {}", tool_call.name);
            Ok(ControlFlow::Skip)
        } else {
            Ok(ControlFlow::Continue)
        }
    }
}

after_tool_call

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

用途:

結果の加工・フォーマット
機密情報のマスキング
結果のキャッシュ
実行結果のログ出力

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

struct ResultFormatter;

#[async_trait]
impl WorkerHook for ResultFormatter {
    async fn after_tool_call(
        &self,
        tool_result: &mut ToolResult,
    ) -> Result<ControlFlow, HookError> {
        if !tool_result.is_error {
            tool_result.content = format!("[OK] {}", tool_result.content);
        }
        Ok(ControlFlow::Continue)
    }
}

on_turn_end

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

用途:

生成されたコードのLint/Fmt
出力形式のバリデーション
自己修正のためのリトライ指示
最終結果のログ出力

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

struct JsonValidator;

#[async_trait]
impl WorkerHook for JsonValidator {
    async fn on_turn_end(
        &self,
        messages: &[Message],
    ) -> Result<TurnResult, 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(TurnResult::ContinueWithMessages(vec![
                        Message::user("Invalid JSON. Please fix and try again.")
                    ]));
                }
            }
        }
        Ok(TurnResult::Finish)
    }
}

複数Hookの実行順序

Hookは登録順に実行されます。

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

制御フローの伝播

Continue: 後続Hookも実行
Skip: 現在の処理をスキップし、後続Hookは実行しない
Abort: 即座にエラーを返し、処理全体を中断

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

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

設計上のポイント

1. デフォルト実装

全メソッドにデフォルト実装があるため、必要なメソッドだけオーバーライドすれば良い。

struct SimpleLogger;

#[async_trait]
impl WorkerHook for SimpleLogger {
    // on_message_send だけ実装
    async fn on_message_send(
        &self,
        context: &mut Vec<Message>,
    ) -> Result<ControlFlow, HookError> {
        println!("Sending {} messages", context.len());
        Ok(ControlFlow::Continue)
    }
    // 他のメソッドはデフォルト（Continue/Finish）
}

2. 可変参照による改変

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

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

3. 並列実行との統合

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

4. Send + Sync 要件

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

struct CountingHook {
    count: Arc<AtomicUsize>,
}

#[async_trait]
impl WorkerHook for CountingHook {
    async fn before_tool_call(&self, _: &mut ToolCall) -> Result<ControlFlow, HookError> {
        self.count.fetch_add(1, Ordering::SeqCst);
        Ok(ControlFlow::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のモック/スタブ作成のためのユーティリティ

11 KiB Raw Blame History Unescape Escape