feat: Implement HookEventKind

This commit is contained in:
2026-01-09 19:18:20 +09:00
parent 33f1c218f2
commit 5691b09fc8
15 changed files with 916 additions and 416 deletions
+1 -1
View File
@@ -17,7 +17,7 @@ LLMを用いたワーカーを作成する小型のSDK・ライブラリ。
module構成概念図
```
```plaintext
worker
├── context
├── llm_client
+90
View File
@@ -0,0 +1,90 @@
# 非同期キャンセル設計
Workerの非同期キャンセル機構についての設計ドキュメント。
## 概要
`tokio_util::sync::CancellationToken`を用いて、別タスクからWorkerの実行を安全にキャンセルできる。
```rust
let worker = Arc::new(Mutex::new(Worker::new(client)));
// 実行タスク
let w = worker.clone();
let handle = tokio::spawn(async move {
w.lock().await.run("prompt").await
});
// キャンセル
worker.lock().await.cancel();
```
## キャンセルポイント
キャンセルは以下のタイミングでチェックされる:
1. **ターンループ先頭**`is_cancelled()`で即座にチェック
2. **ストリーム開始前**`client.stream()`呼び出し時
3. **ストリーム受信中**`tokio::select!`で各イベント受信と並行監視
4. **ツール実行中**`join_all()`と並行監視
## キャンセル時の処理フロー
```
キャンセル検知
timeline.abort_current_block() // 進行中ブロックの終端処理
run_on_abort_hooks("Cancelled") // on_abort フック呼び出し
Err(WorkerError::Cancelled) // エラー返却
```
## API
| メソッド | 説明 |
| ---------------------- | --------------------------------------------------------- |
| `cancel()` | キャンセルをトリガー |
| `is_cancelled()` | キャンセル状態を確認 |
| `cancellation_token()` | トークンへの参照を取得(`clone()`してタスク間で共有可能) |
## on_abort フック
`Hook::on_abort(&self, reason: &str)`がキャンセル時に呼ばれる。
クリーンアップ処理やログ記録に使用できる。
```rust
async fn on_abort(&self, reason: &str) -> Result<(), HookError> {
log::info!("Aborted: {}", reason);
Ok(())
}
```
呼び出しタイミング:
- `WorkerError::Cancelled` — reason: `"Cancelled"`
- `ControlFlow::Abort(reason)` — reason: フックが指定した理由
---
## 既知の問題
### 1. キャンセルトークンの再利用不可
`CancellationToken`は一度キャンセルされると永続的にキャンセル状態になる。
同じWorkerインスタンスで再度`run()`を呼ぶと即座に`Cancelled`エラーになる。
**対応案:**
- `run()`開始時に新しいトークンを生成する
- `reset_cancellation()`メソッドを提供する
### 2. Sync バウンドの追加(破壊的変更)
`tokio::select!`使用のため、Handler/Scope型に`Sync`バウンドを追加した。
既存のユーザーコードで`Sync`未実装の型を使用している場合、コンパイルエラーになる。
### 3. エラー時のon_abort呼び出し
現在、`on_abort`はキャンセルとフックAbort時のみ呼ばれる。
ストリームエラー等のその他エラー時には呼ばれないため、一貫性に欠ける可能性がある。
+154 -132
View File
@@ -3,7 +3,8 @@
## 概要
HookはWorker層でのターン制御に介入するためのメカニズムです。
Claude CodeのHooks機能に着想を得ており、メッセージ送信・ツール実行・ターン終了の各ポイントで処理を差し込むことができます。
メッセージ送信・ツール実行・ターン終了等の各ポイントで処理を差し込むことができます。
## コンセプト
@@ -11,76 +12,82 @@ Claude CodeのHooks機能に着想を得ており、メッセージ送信・ツ
- **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
```rust
#[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)
}
pub trait Hook<E: HookEventKind>: Send + Sync {
async fn call(&self, input: &mut E::Input) -> Result<E::Output, HookError>;
}
```
## 制御フロー型
### ControlFlow
### HookEventKind / Result
Hook処理の継続/中断を制御する列挙型
Hookイベントごとに入力/出力型を分離し、意味のない制御フローを排除する
```rust
pub enum ControlFlow {
/// 処理を続行(後続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,
}
```
### TurnResult
### Tool Call Context
ターン終了時の判定結果を表す列挙型
`before_tool_call` / `after_tool_call` は、ツール実行の文脈を含む入力を受け取る
```rust
pub enum TurnResult {
/// ターンを正常終了
Finish,
/// メッセージを追加してターン継続(自己修正など)
ContinueWithMessages(Vec<Message>),
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>,
}
```
@@ -90,28 +97,30 @@ pub enum TurnResult {
Worker::run() ループ
├─▶ on_message_send ──────────────────────────────┐
│ コンテキストの改変、バリデーション、
│ システムプロンプト注入などが可能
├─▶ LLMリクエスト送信 & ストリーム処理
├─▶ ツール呼び出しがある場合:
│ │
│ コンテキストの改変、バリデーション、 │
│ システムプロンプト注入などが可能 │
│ │
├─▶ LLMリクエスト送信 & ストリーム処理 │
│ │
├─▶ ツール呼び出しがある場合: │
│ │ │
│ ├─▶ before_tool_call (各ツールごと・逐次) │
│ │ 実行可否の判定、引数の改変
│ │
│ │ 実行可否の判定、引数の改変 │
│ │ │
│ ├─▶ ツール並列実行 (join_all) │
│ │
│ │ │
│ └─▶ after_tool_call (各結果ごと・逐次) │
│ 結果の確認、加工、ログ出力
│ 結果の確認、加工、ログ出力 │
│ │
├─▶ ツール結果をコンテキストに追加 → ループ先頭へ │
└─▶ ツールなしの場合:
└─▶ on_turn_end ─────────────────────────────┘
│ │
└─▶ ツールなしの場合: │
│ │
└─▶ on_turn_end ─────────────────────────────┘
最終応答のチェック(Lint/Fmt等)
エラーがあればContinueWithMessagesでリトライ
※ 中断時は on_abort が呼ばれる
```
## 各Hookの詳細
@@ -121,10 +130,12 @@ Worker::run() ループ
**呼び出しタイミング**: LLMへリクエスト送信前(ターンループの冒頭)
**用途**:
- コンテキストへのシステムメッセージ注入
- メッセージのバリデーション
- 機密情報のフィルタリング
- リクエスト内容のログ出力
- `OnMessageSendResult::Cancel` による送信キャンセル
**例**: メッセージにタイムスタンプを追加
@@ -132,14 +143,14 @@ Worker::run() ループ
struct TimestampHook;
#[async_trait]
impl WorkerHook for TimestampHook {
async fn on_message_send(
impl Hook<OnMessageSend> for TimestampHook {
async fn call(
&self,
context: &mut Vec<Message>,
) -> Result<ControlFlow, HookError> {
) -> Result<OnMessageSendResult, HookError> {
let timestamp = chrono::Local::now().to_rfc3339();
context.insert(0, Message::user(format!("[{}]", timestamp)));
Ok(ControlFlow::Continue)
Ok(OnMessageSendResult::Continue)
}
}
```
@@ -149,10 +160,15 @@ impl WorkerHook for TimestampHook {
**呼び出しタイミング**: 各ツール実行前(並列実行フェーズの前)
**用途**:
- 危険なツールのブロック
- 引数のサニタイズ
- 確認プロンプトの表示(UIとの連携)
- 実行ログの記録
- `BeforeToolCallResult::Pause` による一時停止
**入力**:
- `ToolCallContext``ToolCall` + `ToolMeta` + `Arc<dyn Tool>`
**例**: 特定ツールをブロック
@@ -162,16 +178,16 @@ struct ToolBlocker {
}
#[async_trait]
impl WorkerHook for ToolBlocker {
async fn before_tool_call(
impl Hook<BeforeToolCall> for ToolBlocker {
async fn 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)
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(ControlFlow::Continue)
Ok(BeforeToolCallResult::Continue)
}
}
}
@@ -182,10 +198,13 @@ impl WorkerHook for ToolBlocker {
**呼び出しタイミング**: 各ツール実行後(並列実行フェーズの後)
**用途**:
- 結果の加工・フォーマット
- 機密情報のマスキング
- 結果のキャッシュ
- 実行結果のログ出力
**入力**:
- `ToolResultContext``ToolResult` + `ToolMeta` + `Arc<dyn Tool>`
**例**: 結果にプレフィックスを追加
@@ -193,15 +212,15 @@ impl WorkerHook for ToolBlocker {
struct ResultFormatter;
#[async_trait]
impl WorkerHook for ResultFormatter {
async fn after_tool_call(
impl Hook<AfterToolCall> for ResultFormatter {
async fn call(
&self,
tool_result: &mut ToolResult,
) -> Result<ControlFlow, HookError> {
if !tool_result.is_error {
tool_result.content = format!("[OK] {}", tool_result.content);
ctx: &mut ToolResultContext,
) -> Result<AfterToolCallResult, HookError> {
if !ctx.result.is_error {
ctx.result.content = format!("[OK] {}", ctx.result.content);
}
Ok(ControlFlow::Continue)
Ok(AfterToolCallResult::Continue)
}
}
```
@@ -211,10 +230,22 @@ impl WorkerHook for ResultFormatter {
**呼び出しタイミング**: ツール呼び出しなしでターンが終了する直前
**用途**:
- 生成されたコードのLint/Fmt
- 出力形式のバリデーション
- 自己修正のためのリトライ指示
- 最終結果のログ出力
- `OnTurnEndResult::Paused` による一時停止
### on_abort
**呼び出しタイミング**: キャンセル/エラー/AbortなどでWorkerが中断された時
**用途**:
- クリーンアップ処理
- 中断理由のログ出力
- 外部システムへの通知
**例**: JSON形式のバリデーション
@@ -222,11 +253,11 @@ impl WorkerHook for ResultFormatter {
struct JsonValidator;
#[async_trait]
impl WorkerHook for JsonValidator {
async fn on_turn_end(
impl Hook<OnTurnEnd> for JsonValidator {
async fn call(
&self,
messages: &[Message],
) -> Result<TurnResult, HookError> {
messages: &mut Vec<Message>,
) -> Result<OnTurnEndResult, HookError> {
// 最後のアシスタントメッセージを取得
let last = messages.iter().rev()
.find(|m| m.role == Role::Assistant);
@@ -236,25 +267,25 @@ impl WorkerHook for JsonValidator {
// JSONとしてパースを試みる
if serde_json::from_str::<serde_json::Value>(text).is_err() {
// 失敗したらリトライ指示
return Ok(TurnResult::ContinueWithMessages(vec![
return Ok(OnTurnEndResult::ContinueWithMessages(vec![
Message::user("Invalid JSON. Please fix and try again.")
]));
}
}
}
Ok(TurnResult::Finish)
Ok(OnTurnEndResult::Finish)
}
}
```
## 複数Hookの実行順序
Hookは**登録順**に実行されます。
Hookは**イベントごとに登録順**に実行されます。
```rust
worker.add_hook(HookA); // 1番目に実行
worker.add_hook(HookB); // 2番目に実行
worker.add_hook(HookC); // 3番目に実行
worker.add_before_tool_call_hook(HookA); // 1番目に実行
worker.add_before_tool_call_hook(HookB); // 2番目に実行
worker.add_before_tool_call_hook(HookC); // 3番目に実行
```
### 制御フローの伝播
@@ -262,6 +293,7 @@ worker.add_hook(HookC); // 3番目に実行
- `Continue`: 後続Hookも実行
- `Skip`: 現在の処理をスキップし、後続Hookは実行しない
- `Abort`: 即座にエラーを返し、処理全体を中断
- `Pause`: Workerを一時停止(再開は`resume`
```
Hook A: Continue → Hook B: Skip → (Hook Cは実行されない)
@@ -271,40 +303,27 @@ 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. デフォルト実装
### 1. イベントごとの実装
全メソッドにデフォルト実装があるため、必要なメソッドだけオーバーライドすれば良い
```rust
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
}
```
必要なイベントのみ `Hook<Event>` を実装する
### 2. 可変参照による改変
`&mut`で引数を受け取るため、直接改変が可能。
```rust
async fn before_tool_call(&self, tool_call: &mut ToolCall) -> ... {
async fn call(&self, ctx: &mut ToolCallContext) -> ... {
// 引数を直接書き換え
tool_call.input["sanitized"] = json!(true);
Ok(ControlFlow::Continue)
ctx.call.input["sanitized"] = json!(true);
Ok(BeforeToolCallResult::Continue)
}
```
@@ -316,7 +335,7 @@ async fn before_tool_call(&self, tool_call: &mut ToolCall) -> ... {
### 4. Send + Sync 要件
`WorkerHook``Send + Sync`を要求するため、スレッドセーフな実装が必要。
`Hook``Send + Sync`を要求するため、スレッドセーフな実装が必要。
状態を持つ場合は`Arc<Mutex<T>>``AtomicUsize`などを使用する。
```rust
@@ -325,24 +344,24 @@ struct CountingHook {
}
#[async_trait]
impl WorkerHook for CountingHook {
async fn before_tool_call(&self, _: &mut ToolCall) -> Result<ControlFlow, HookError> {
impl Hook<BeforeToolCall> for CountingHook {
async fn call(&self, _: &mut ToolCallContext) -> Result<BeforeToolCallResult, HookError> {
self.count.fetch_add(1, Ordering::SeqCst);
Ok(ControlFlow::Continue)
Ok(BeforeToolCallResult::Continue)
}
}
```
## 典型的なユースケース
| ユースケース | 使用Hook | 処理内容 |
|-------------|----------|----------|
| ツール許可制御 | `before_tool_call` | 危険なツールをSkip |
| 実行ログ | `before/after_tool_call` | 呼び出しと結果を記録 |
| 出力バリデーション | `on_turn_end` | 形式チェック、リトライ指示 |
| コンテキスト注入 | `on_message_send` | システムメッセージ追加 |
| 結果のサニタイズ | `after_tool_call` | 機密情報のマスキング |
| レート制限 | `before_tool_call` | 呼び出し頻度の制御 |
| ユースケース | 使用Hook | 処理内容 |
| ------------------ | ------------------------ | -------------------------- |
| ツール許可制御 | `before_tool_call` | 危険なツールをSkip |
| 実行ログ | `before/after_tool_call` | 呼び出しと結果を記録 |
| 出力バリデーション | `on_turn_end` | 形式チェック、リトライ指示 |
| コンテキスト注入 | `on_message_send` | システムメッセージ追加 |
| 結果のサニタイズ | `after_tool_call` | 機密情報のマスキング |
| レート制限 | `before_tool_call` | 呼び出し頻度の制御 |
## TODO
@@ -350,11 +369,14 @@ impl WorkerHook for CountingHook {
現在のHooks実装は基本的なユースケースをカバーしているが、以下の点について将来的に厳密な仕様を定義する必要がある:
- **エラーハンドリングの明確化**: `HookError`発生時のリカバリー戦略、部分的な失敗の扱い
- **エラーハンドリングの明確化**:
`HookError`発生時のリカバリー戦略、部分的な失敗の扱い
- **Hook間の依存関係**: 複数Hookの実行順序が結果に影響する場合のセマンティクス
- **非同期キャンセル**: Hook実行中のキャンセル(タイムアウト等)の振る舞い
- **状態の一貫性**: `on_message_send`で改変されたコンテキストが後続処理で期待通りに反映される保証
- **リトライ制限**: `on_turn_end`での`ContinueWithMessages`による無限ループ防止策
- **状態の一貫性**:
`on_message_send`で改変されたコンテキストが後続処理で期待通りに反映される保証
- **リトライ制限**:
`on_turn_end`での`ContinueWithMessages`による無限ループ防止策
- **Hook優先度**: 登録順以外の優先度指定メカニズムの必要性
- **条件付きHook**: 特定条件でのみ有効化されるHookパターン
- **テスト容易性**: Hookのモック/スタブ作成のためのユーティリティ
+47 -29
View File
@@ -178,41 +178,60 @@ Workerは生成されたラッパー構造体を `Box<dyn Tool>` として保持
```rust
#[async_trait]
pub trait WorkerHook: Send + Sync {
/// メッセージ送信前。
/// リクエストに含まれるメッセージリストを改変できる。
async fn on_message_send(&self, context: &mut Vec<Message>) -> Result<ControlFlow, Error> {
Ok(ControlFlow::Continue)
}
/// ツール実行前。
/// 実行をキャンセルしたり、引数を書き換えることができる。
async fn before_tool_call(&self, tool_call: &mut ToolCall) -> Result<ControlFlow, Error> {
Ok(ControlFlow::Continue)
}
/// ツール実行後。
/// 結果を書き換えたり、隠蔽したりできる。
async fn after_tool_call(&self, tool_result: &mut ToolResult) -> Result<ControlFlow, Error> {
Ok(ControlFlow::Continue)
}
/// ターン終了時。
/// 生成されたメッセージを検査し、必要ならリトライ(ContinueWithMessages)を指示できる。
async fn on_turn_end(&self, messages: &[Message]) -> Result<TurnResult, Error> {
Ok(TurnResult::Finish)
}
pub trait Hook<E: HookEventKind>: Send + Sync {
async fn call(&self, input: &mut E::Input) -> Result<E::Output, Error>;
}
pub enum ControlFlow {
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,
Skip, // Tool実行などをスキップ
Abort(String), // 処理中断
Cancel(String),
}
pub enum TurnResult {
pub enum BeforeToolCallResult {
Continue,
Skip, // Tool実行などをスキップ
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` は、ツール実行の文脈を含む入力を受け取る。
```rust
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>,
}
```
@@ -433,4 +452,3 @@ impl<C: LlmClient> Worker<C> {
3. **選択的購読**: on_*で必要なイベントだけ、またはSubscriberで一括
4. **累積イベントの追加**: Worker層でComplete系イベントを追加提供
5. **後方互換性**: 従来の`run()`も引き続き使用可能