From 6d6ae24ffe3b7c40ee0cc7c2639c9ab610d71094 Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 5 Jan 2026 23:10:32 +0900 Subject: [PATCH] feat: Introduce event timeline and handler system --- worker-macros/Cargo.toml | 13 + worker-macros/src/lib.rs | 41 +++ worker-types/Cargo.toml | 8 + worker-types/src/event.rs | 276 +++++++++++++++ worker-types/src/handler.rs | 141 ++++++++ worker-types/src/lib.rs | 12 + worker/Cargo.toml | 10 + worker/examples/timeline_basic.rs | 132 +++++++ worker/src/lib.rs | 10 + worker/src/timeline.rs | 565 ++++++++++++++++++++++++++++++ 10 files changed, 1208 insertions(+) create mode 100644 worker-macros/Cargo.toml create mode 100644 worker-macros/src/lib.rs create mode 100644 worker-types/Cargo.toml create mode 100644 worker-types/src/event.rs create mode 100644 worker-types/src/handler.rs create mode 100644 worker-types/src/lib.rs create mode 100644 worker/Cargo.toml create mode 100644 worker/examples/timeline_basic.rs create mode 100644 worker/src/lib.rs create mode 100644 worker/src/timeline.rs diff --git a/worker-macros/Cargo.toml b/worker-macros/Cargo.toml new file mode 100644 index 0000000..16e04da --- /dev/null +++ b/worker-macros/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "worker-macros" +version = "0.1.0" +edition = "2024" + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = "1" +quote = "1" +syn = { version = "2", features = ["full"] } +worker-types = { path = "../worker-types" } diff --git a/worker-macros/src/lib.rs b/worker-macros/src/lib.rs new file mode 100644 index 0000000..7ae263a --- /dev/null +++ b/worker-macros/src/lib.rs @@ -0,0 +1,41 @@ +//! worker-macros - LLMワーカー用のProcedural Macros +//! +//! このクレートはTools/Hooksを定義するためのマクロを提供する予定です。 +//! +//! TODO: Tool定義マクロの実装 +//! TODO: Hook定義マクロの実装 + +use proc_macro::TokenStream; + +/// ツール定義マクロ（未実装） +/// +/// # Example +/// ```ignore +/// #[tool( +/// name = "get_weather", +/// description = "Get weather information for a city" +/// )] +/// fn get_weather(city: String) -> Result { +/// // ... +/// } +/// ``` +#[proc_macro_attribute] +pub fn tool(_attr: TokenStream, item: TokenStream) -> TokenStream { + // TODO: 実装 + item +} + +/// フック定義マクロ（未実装） +/// +/// # Example +/// ```ignore +/// #[hook(on = "before_tool_call")] +/// fn log_tool_call(tool_name: &str) { +/// println!("Calling tool: {}", tool_name); +/// } +/// ``` +#[proc_macro_attribute] +pub fn hook(_attr: TokenStream, item: TokenStream) -> TokenStream { + // TODO: 実装 + item +} diff --git a/worker-types/Cargo.toml b/worker-types/Cargo.toml new file mode 100644 index 0000000..9ba1934 --- /dev/null +++ b/worker-types/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "worker-types" +version = "0.1.0" +edition = "2024" + +[dependencies] +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" diff --git a/worker-types/src/event.rs b/worker-types/src/event.rs new file mode 100644 index 0000000..066713e --- /dev/null +++ b/worker-types/src/event.rs @@ -0,0 +1,276 @@ +//! イベント型定義 +//! +//! llm_client層が出力するフラットなイベント列挙と関連型 + +use serde::{Deserialize, Serialize}; + +// ============================================================================= +// Core Event Types (from llm_client layer) +// ============================================================================= + +/// llm_client層が出力するフラットなイベント列挙 +/// +/// Timeline層がこのイベントストリームを受け取り、ブロック構造化を行う +#[derive(Debug, Clone, PartialEq)] +pub enum Event { + // Meta events (not tied to a block) + Ping(PingEvent), + Usage(UsageEvent), + Status(StatusEvent), + Error(ErrorEvent), + + // Block lifecycle events + BlockStart(BlockStart), + BlockDelta(BlockDelta), + BlockStop(BlockStop), + BlockAbort(BlockAbort), +} + +// ============================================================================= +// Meta Events +// ============================================================================= + +/// Pingイベント（ハートビート） +#[derive(Debug, Clone, PartialEq, Default)] +pub struct PingEvent { + pub timestamp: Option, +} + +/// 使用量イベント +#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)] +pub struct UsageEvent { + /// 入力トークン数 + pub input_tokens: Option, + /// 出力トークン数 + pub output_tokens: Option, + /// 合計トークン数 + pub total_tokens: Option, + /// キャッシュ読み込みトークン数 + pub cache_read_input_tokens: Option, + /// キャッシュ作成トークン数 + pub cache_creation_input_tokens: Option, +} + +/// ステータスイベント +#[derive(Debug, Clone, PartialEq)] +pub struct StatusEvent { + pub status: ResponseStatus, +} + +/// レスポンスステータス +#[derive(Debug, Clone, PartialEq)] +pub enum ResponseStatus { + /// ストリーム開始 + Started, + /// 正常完了 + Completed, + /// キャンセルされた + Cancelled, + /// エラー発生 + Failed, +} + +/// エラーイベント +#[derive(Debug, Clone, PartialEq)] +pub struct ErrorEvent { + pub code: Option, + pub message: String, +} + +// ============================================================================= +// Block Types +// ============================================================================= + +/// ブロックの種別 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum BlockType { + /// テキスト生成 + Text, + /// 思考 (Claude Extended Thinking等) + Thinking, + /// ツール呼び出し + ToolUse, + /// ツール結果 + ToolResult, +} + +/// ブロック開始イベント +#[derive(Debug, Clone, PartialEq)] +pub struct BlockStart { + /// ブロックのインデックス + pub index: usize, + /// ブロックの種別 + pub block_type: BlockType, + /// ブロック固有のメタデータ + pub metadata: BlockMetadata, +} + +impl BlockStart { + pub fn block_type(&self) -> BlockType { + self.block_type + } +} + +/// ブロックのメタデータ +#[derive(Debug, Clone, PartialEq)] +pub enum BlockMetadata { + Text, + Thinking, + ToolUse { id: String, name: String }, + ToolResult { tool_use_id: String }, +} + +/// ブロックデルタイベント +#[derive(Debug, Clone, PartialEq)] +pub struct BlockDelta { + /// ブロックのインデックス + pub index: usize, + /// デルタの内容 + pub delta: DeltaContent, +} + +/// デルタの内容 +#[derive(Debug, Clone, PartialEq)] +pub enum DeltaContent { + /// テキストデルタ + Text(String), + /// 思考デルタ + Thinking(String), + /// ツール引数のJSON部分文字列 + InputJson(String), +} + +impl DeltaContent { + /// デルタのブロック種別を取得 + pub fn block_type(&self) -> BlockType { + match self { + DeltaContent::Text(_) => BlockType::Text, + DeltaContent::Thinking(_) => BlockType::Thinking, + DeltaContent::InputJson(_) => BlockType::ToolUse, + } + } +} + +/// ブロック停止イベント +#[derive(Debug, Clone, PartialEq)] +pub struct BlockStop { + /// ブロックのインデックス + pub index: usize, + /// ブロックの種別 + pub block_type: BlockType, + /// 停止理由 + pub stop_reason: Option, +} + +impl BlockStop { + pub fn block_type(&self) -> BlockType { + self.block_type + } +} + +/// ブロック中断イベント +#[derive(Debug, Clone, PartialEq)] +pub struct BlockAbort { + /// ブロックのインデックス + pub index: usize, + /// ブロックの種別 + pub block_type: BlockType, + /// 中断理由 + pub reason: String, +} + +impl BlockAbort { + pub fn block_type(&self) -> BlockType { + self.block_type + } +} + +/// 停止理由 +#[derive(Debug, Clone, PartialEq)] +pub enum StopReason { + /// 自然終了 + EndTurn, + /// 最大トークン数到達 + MaxTokens, + /// ストップシーケンス到達 + StopSequence, + /// ツール使用 + ToolUse, +} + +// ============================================================================= +// Builder / Factory helpers +// ============================================================================= + +impl Event { + /// テキストブロック開始イベントを作成 + pub fn text_block_start(index: usize) -> Self { + Event::BlockStart(BlockStart { + index, + block_type: BlockType::Text, + metadata: BlockMetadata::Text, + }) + } + + /// テキストデルタイベントを作成 + pub fn text_delta(index: usize, text: impl Into) -> Self { + Event::BlockDelta(BlockDelta { + index, + delta: DeltaContent::Text(text.into()), + }) + } + + /// テキストブロック停止イベントを作成 + pub fn text_block_stop(index: usize, stop_reason: Option) -> Self { + Event::BlockStop(BlockStop { + index, + block_type: BlockType::Text, + stop_reason, + }) + } + + /// ツール使用ブロック開始イベントを作成 + pub fn tool_use_start(index: usize, id: impl Into, name: impl Into) -> Self { + Event::BlockStart(BlockStart { + index, + block_type: BlockType::ToolUse, + metadata: BlockMetadata::ToolUse { + id: id.into(), + name: name.into(), + }, + }) + } + + /// ツール引数デルタイベントを作成 + pub fn tool_input_delta(index: usize, json: impl Into) -> Self { + Event::BlockDelta(BlockDelta { + index, + delta: DeltaContent::InputJson(json.into()), + }) + } + + /// ツール使用ブロック停止イベントを作成 + pub fn tool_use_stop(index: usize) -> Self { + Event::BlockStop(BlockStop { + index, + block_type: BlockType::ToolUse, + stop_reason: Some(StopReason::ToolUse), + }) + } + + /// 使用量イベントを作成 + pub fn usage(input_tokens: u64, output_tokens: u64) -> Self { + Event::Usage(UsageEvent { + input_tokens: Some(input_tokens), + output_tokens: Some(output_tokens), + total_tokens: Some(input_tokens + output_tokens), + cache_read_input_tokens: None, + cache_creation_input_tokens: None, + }) + } + + /// Pingイベントを作成 + pub fn ping() -> Self { + Event::Ping(PingEvent { timestamp: None }) + } +} diff --git a/worker-types/src/handler.rs b/worker-types/src/handler.rs new file mode 100644 index 0000000..d47dcc4 --- /dev/null +++ b/worker-types/src/handler.rs @@ -0,0 +1,141 @@ +//! Handler/Kind関連の型定義 +//! +//! Timeline層でのイベント処理に使用するトレイトとKind定義 + +use crate::event::*; + +// ============================================================================= +// Kind Trait +// ============================================================================= + +/// Kindはイベント型のみを定義する +/// +/// スコープはHandler側で定義するため、同じKindに対して +/// 異なるスコープを持つHandlerを登録できる +pub trait Kind { + /// このKindに対応するイベント型 + type Event; +} + +// ============================================================================= +// Handler Trait +// ============================================================================= + +/// Kindに対する処理を定義し、自身のスコープ型も決定する +pub trait Handler { + /// Handler固有のスコープ型 + type Scope: Default; + + /// イベントを処理する + fn on_event(&mut self, scope: &mut Self::Scope, event: &K::Event); +} + +// ============================================================================= +// Meta Kind Definitions +// ============================================================================= + +/// Usage Kind - 使用量イベント用 +pub struct UsageKind; +impl Kind for UsageKind { + type Event = UsageEvent; +} + +/// Ping Kind - Pingイベント用 +pub struct PingKind; +impl Kind for PingKind { + type Event = PingEvent; +} + +/// Status Kind - ステータスイベント用 +pub struct StatusKind; +impl Kind for StatusKind { + type Event = StatusEvent; +} + +/// Error Kind - エラーイベント用 +pub struct ErrorKind; +impl Kind for ErrorKind { + type Event = ErrorEvent; +} + +// ============================================================================= +// Block Kind Definitions +// ============================================================================= + +/// TextBlock Kind - テキストブロック用 +pub struct TextBlockKind; +impl Kind for TextBlockKind { + type Event = TextBlockEvent; +} + +/// テキストブロックのイベント +#[derive(Debug, Clone, PartialEq)] +pub enum TextBlockEvent { + Start(TextBlockStart), + Delta(String), + Stop(TextBlockStop), +} + +#[derive(Debug, Clone, PartialEq)] +pub struct TextBlockStart { + pub index: usize, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct TextBlockStop { + pub index: usize, + pub stop_reason: Option, +} + +/// ThinkingBlock Kind - 思考ブロック用 +pub struct ThinkingBlockKind; +impl Kind for ThinkingBlockKind { + type Event = ThinkingBlockEvent; +} + +/// 思考ブロックのイベント +#[derive(Debug, Clone, PartialEq)] +pub enum ThinkingBlockEvent { + Start(ThinkingBlockStart), + Delta(String), + Stop(ThinkingBlockStop), +} + +#[derive(Debug, Clone, PartialEq)] +pub struct ThinkingBlockStart { + pub index: usize, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct ThinkingBlockStop { + pub index: usize, +} + +/// ToolUseBlock Kind - ツール使用ブロック用 +pub struct ToolUseBlockKind; +impl Kind for ToolUseBlockKind { + type Event = ToolUseBlockEvent; +} + +/// ツール使用ブロックのイベント +#[derive(Debug, Clone, PartialEq)] +pub enum ToolUseBlockEvent { + Start(ToolUseBlockStart), + /// ツール引数のJSON部分文字列 + InputJsonDelta(String), + Stop(ToolUseBlockStop), +} + +#[derive(Debug, Clone, PartialEq)] +pub struct ToolUseBlockStart { + pub index: usize, + pub id: String, + pub name: String, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct ToolUseBlockStop { + pub index: usize, + pub id: String, + pub name: String, +} diff --git a/worker-types/src/lib.rs b/worker-types/src/lib.rs new file mode 100644 index 0000000..96566c4 --- /dev/null +++ b/worker-types/src/lib.rs @@ -0,0 +1,12 @@ +//! worker-types - LLMワーカーで使用される型定義 +//! +//! このクレートは以下を提供します: +//! - Event: llm_client層からのフラットなイベント列挙 +//! - Kind/Handler: タイムライン層でのイベント処理トレイト +//! - 各種イベント構造体 + +mod event; +mod handler; + +pub use event::*; +pub use handler::*; diff --git a/worker/Cargo.toml b/worker/Cargo.toml new file mode 100644 index 0000000..5dcdbe0 --- /dev/null +++ b/worker/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "worker" +version = "0.1.0" +edition = "2024" + +[dependencies] +serde_json = "1.0" +thiserror = "1.0" +worker-macros = { path = "../worker-macros" } +worker-types = { path = "../worker-types" } diff --git a/worker/examples/timeline_basic.rs b/worker/examples/timeline_basic.rs new file mode 100644 index 0000000..bb3cff2 --- /dev/null +++ b/worker/examples/timeline_basic.rs @@ -0,0 +1,132 @@ +//! Timeline使用例 +//! +//! 設計ドキュメントに基づいたTimelineの使用パターンを示すサンプル + +use worker::{ + Event, Handler, TextBlockEvent, TextBlockKind, Timeline, + ToolUseBlockEvent, ToolUseBlockKind, UsageEvent, UsageKind, +}; + +fn main() { + // シミュレートされたイベントストリーム + let events = simulate_llm_response(); + + // Timelineを作成し、ハンドラーを登録 + let mut timeline = Timeline::new(); + + // Usage収集ハンドラー + timeline.on_usage(UsageAccumulator::new()); + + // テキスト収集ハンドラー + timeline.on_text_block(TextCollector::new()); + + // ツール呼び出し収集ハンドラー + timeline.on_tool_use_block(ToolCallCollector::new()); + + // イベントをディスパッチ + for event in &events { + timeline.dispatch(event); + } + + println!("Timeline example completed!"); + println!("Events processed: {}", events.len()); +} + +/// LLMレスポンスをシミュレート +fn simulate_llm_response() -> Vec { + vec![ + // テキストブロック + Event::text_block_start(0), + Event::text_delta(0, "Hello, "), + Event::text_delta(0, "I can help you with that."), + Event::text_block_stop(0, None), + // 使用量 + Event::usage(100, 50), + // ツール呼び出し + Event::tool_use_start(1, "call_abc123", "get_weather"), + Event::tool_input_delta(1, r#"{"city":"#), + Event::tool_input_delta(1, r#""Tokyo"}"#), + Event::tool_use_stop(1), + // 最終的な使用量 + Event::usage(100, 75), + ] +} + +// ============================================================================= +// Example Handlers (defined in example, not in library) +// ============================================================================= + +/// 使用量を累積するハンドラー +struct UsageAccumulator { + total_tokens: u64, +} + +impl UsageAccumulator { + fn new() -> Self { + Self { total_tokens: 0 } + } +} + +impl Handler for UsageAccumulator { + type Scope = (); + fn on_event(&mut self, _scope: &mut (), usage: &UsageEvent) { + self.total_tokens += usage.total_tokens.unwrap_or(0); + } +} + +/// テキストを収集するハンドラー +struct TextCollector { + results: Vec, +} + +impl TextCollector { + fn new() -> Self { + Self { results: Vec::new() } + } +} + +impl Handler for TextCollector { + type Scope = String; + fn on_event(&mut self, buffer: &mut String, event: &TextBlockEvent) { + match event { + TextBlockEvent::Start(_) => {} + TextBlockEvent::Delta(s) => buffer.push_str(s), + TextBlockEvent::Stop(_) => { + self.results.push(std::mem::take(buffer)); + } + } + } +} + +/// ツール呼び出しを収集するハンドラー +struct ToolCallCollector { + calls: Vec<(String, String)>, // (name, args) +} + +impl ToolCallCollector { + fn new() -> Self { + Self { calls: Vec::new() } + } +} + +#[derive(Default)] +struct ToolCallScope { + name: String, + args: String, +} + +impl Handler for ToolCallCollector { + type Scope = ToolCallScope; + fn on_event(&mut self, scope: &mut ToolCallScope, event: &ToolUseBlockEvent) { + match event { + ToolUseBlockEvent::Start(s) => scope.name = s.name.clone(), + ToolUseBlockEvent::InputJsonDelta(json) => scope.args.push_str(json), + ToolUseBlockEvent::Stop(_) => { + self.calls.push(( + std::mem::take(&mut scope.name), + std::mem::take(&mut scope.args), + )); + } + } + } +} diff --git a/worker/src/lib.rs b/worker/src/lib.rs new file mode 100644 index 0000000..3b13109 --- /dev/null +++ b/worker/src/lib.rs @@ -0,0 +1,10 @@ +//! worker - LLMワーカーのメイン実装 +//! +//! このクレートは以下を提供します: +//! - Timeline: イベントストリームの状態管理とハンドラーへのディスパッチ +//! - 型消去されたHandler実装 + +mod timeline; + +pub use timeline::*; +pub use worker_types::*; diff --git a/worker/src/timeline.rs b/worker/src/timeline.rs new file mode 100644 index 0000000..f6e7612 --- /dev/null +++ b/worker/src/timeline.rs @@ -0,0 +1,565 @@ +//! Timeline層の実装 +//! +//! イベントストリームを受信し、登録されたHandlerへディスパッチする + +use std::marker::PhantomData; + +use worker_types::*; + +// ============================================================================= +// Type-erased Handler +// ============================================================================= + +/// 型消去されたHandler trait +/// +/// 各Handlerは独自のScope型を持つため、Timelineで保持するには型消去が必要 +pub trait ErasedHandler: Send { + /// イベントをディスパッチ + fn dispatch(&mut self, event: &K::Event); + /// スコープを開始（Block開始時） + fn start_scope(&mut self); + /// スコープを終了（Block終了時） + fn end_scope(&mut self); +} + +/// HandlerからErasedHandlerへのラッパー +pub struct HandlerWrapper +where + H: Handler, + K: Kind, +{ + handler: H, + scope: Option, + // fn() -> K は常にSend+Syncなので、Kの制約に関係なくSendを満たせる + _kind: PhantomData K>, +} + +impl HandlerWrapper +where + H: Handler, + K: Kind, +{ + pub fn new(handler: H) -> Self { + Self { + handler, + scope: None, + _kind: PhantomData, + } + } +} + +impl ErasedHandler for HandlerWrapper +where + H: Handler + Send, + K: Kind, + H::Scope: Send, +{ + fn dispatch(&mut self, event: &K::Event) { + if let Some(scope) = &mut self.scope { + self.handler.on_event(scope, event); + } + } + + fn start_scope(&mut self) { + self.scope = Some(H::Scope::default()); + } + + fn end_scope(&mut self) { + self.scope = None; + } +} + +// ============================================================================= +// Block Handler Registry +// ============================================================================= + +/// ブロックハンドラーの型消去trait +trait ErasedBlockHandler: Send { + fn dispatch_start(&mut self, start: &BlockStart); + fn dispatch_delta(&mut self, delta: &BlockDelta); + fn dispatch_stop(&mut self, stop: &BlockStop); + fn dispatch_abort(&mut self, abort: &BlockAbort); + fn start_scope(&mut self); + fn end_scope(&mut self); +} + +/// TextBlockKind用のラッパー +struct TextBlockHandlerWrapper +where + H: Handler, +{ + handler: H, + scope: Option, +} + +impl TextBlockHandlerWrapper +where + H: Handler, +{ + fn new(handler: H) -> Self { + Self { + handler, + scope: None, + } + } +} + +impl ErasedBlockHandler for TextBlockHandlerWrapper +where + H: Handler + Send, + H::Scope: Send, +{ + fn dispatch_start(&mut self, start: &BlockStart) { + if let Some(scope) = &mut self.scope { + self.handler.on_event( + scope, + &TextBlockEvent::Start(TextBlockStart { index: start.index }), + ); + } + } + + fn dispatch_delta(&mut self, delta: &BlockDelta) { + if let Some(scope) = &mut self.scope { + if let DeltaContent::Text(text) = &delta.delta { + self.handler.on_event(scope, &TextBlockEvent::Delta(text.clone())); + } + } + } + + fn dispatch_stop(&mut self, stop: &BlockStop) { + if let Some(scope) = &mut self.scope { + self.handler.on_event( + scope, + &TextBlockEvent::Stop(TextBlockStop { + index: stop.index, + stop_reason: stop.stop_reason.clone(), + }), + ); + } + } + + fn dispatch_abort(&mut self, _abort: &BlockAbort) { + // TextBlockはabortを特別扱いしない（スコープ終了のみ） + } + + fn start_scope(&mut self) { + self.scope = Some(H::Scope::default()); + } + + fn end_scope(&mut self) { + self.scope = None; + } +} + +/// ThinkingBlockKind用のラッパー +struct ThinkingBlockHandlerWrapper +where + H: Handler, +{ + handler: H, + scope: Option, +} + +impl ThinkingBlockHandlerWrapper +where + H: Handler, +{ + fn new(handler: H) -> Self { + Self { + handler, + scope: None, + } + } +} + +impl ErasedBlockHandler for ThinkingBlockHandlerWrapper +where + H: Handler + Send, + H::Scope: Send, +{ + fn dispatch_start(&mut self, start: &BlockStart) { + if let Some(scope) = &mut self.scope { + self.handler.on_event( + scope, + &ThinkingBlockEvent::Start(ThinkingBlockStart { index: start.index }), + ); + } + } + + fn dispatch_delta(&mut self, delta: &BlockDelta) { + if let Some(scope) = &mut self.scope { + if let DeltaContent::Thinking(text) = &delta.delta { + self.handler.on_event(scope, &ThinkingBlockEvent::Delta(text.clone())); + } + } + } + + fn dispatch_stop(&mut self, stop: &BlockStop) { + if let Some(scope) = &mut self.scope { + self.handler.on_event( + scope, + &ThinkingBlockEvent::Stop(ThinkingBlockStop { index: stop.index }), + ); + } + } + + fn dispatch_abort(&mut self, _abort: &BlockAbort) {} + + fn start_scope(&mut self) { + self.scope = Some(H::Scope::default()); + } + + fn end_scope(&mut self) { + self.scope = None; + } +} + +/// ToolUseBlockKind用のラッパー +struct ToolUseBlockHandlerWrapper +where + H: Handler, +{ + handler: H, + scope: Option, + current_tool: Option<(String, String)>, // (id, name) +} + +impl ToolUseBlockHandlerWrapper +where + H: Handler, +{ + fn new(handler: H) -> Self { + Self { + handler, + scope: None, + current_tool: None, + } + } +} + +impl ErasedBlockHandler for ToolUseBlockHandlerWrapper +where + H: Handler + Send, + H::Scope: Send, +{ + fn dispatch_start(&mut self, start: &BlockStart) { + if let Some(scope) = &mut self.scope { + if let BlockMetadata::ToolUse { id, name } = &start.metadata { + self.current_tool = Some((id.clone(), name.clone())); + self.handler.on_event( + scope, + &ToolUseBlockEvent::Start(ToolUseBlockStart { + index: start.index, + id: id.clone(), + name: name.clone(), + }), + ); + } + } + } + + fn dispatch_delta(&mut self, delta: &BlockDelta) { + if let Some(scope) = &mut self.scope { + if let DeltaContent::InputJson(json) = &delta.delta { + self.handler + .on_event(scope, &ToolUseBlockEvent::InputJsonDelta(json.clone())); + } + } + } + + fn dispatch_stop(&mut self, stop: &BlockStop) { + if let Some(scope) = &mut self.scope { + if let Some((id, name)) = self.current_tool.take() { + self.handler.on_event( + scope, + &ToolUseBlockEvent::Stop(ToolUseBlockStop { + index: stop.index, + id, + name, + }), + ); + } + } + } + + fn dispatch_abort(&mut self, _abort: &BlockAbort) { + self.current_tool = None; + } + + fn start_scope(&mut self) { + self.scope = Some(H::Scope::default()); + } + + fn end_scope(&mut self) { + self.scope = None; + self.current_tool = None; + } +} + +// ============================================================================= +// Timeline +// ============================================================================= + +/// Timeline - イベントストリームの状態管理とディスパッチ +/// +/// # 責務 +/// 1. Eventストリームを受信 +/// 2. Block系イベントをBlockKindごとのライフサイクルイベントに変換 +/// 3. 各Handlerごとのスコープの生成・管理 +/// 4. 登録されたHandlerへの登録順ディスパッチ +pub struct Timeline { + // Meta系ハンドラー + usage_handlers: Vec>>, + ping_handlers: Vec>>, + status_handlers: Vec>>, + error_handlers: Vec>>, + + // Block系ハンドラー（BlockTypeごとにグループ化） + text_block_handlers: Vec>, + thinking_block_handlers: Vec>, + tool_use_block_handlers: Vec>, + + // 現在アクティブなブロック + current_block: Option, +} + +impl Default for Timeline { + fn default() -> Self { + Self::new() + } +} + +impl Timeline { + pub fn new() -> Self { + Self { + usage_handlers: Vec::new(), + ping_handlers: Vec::new(), + status_handlers: Vec::new(), + error_handlers: Vec::new(), + text_block_handlers: Vec::new(), + thinking_block_handlers: Vec::new(), + tool_use_block_handlers: Vec::new(), + current_block: None, + } + } + + // ========================================================================= + // Handler Registration + // ========================================================================= + + /// UsageKind用のHandlerを登録 + pub fn on_usage(&mut self, handler: H) -> &mut Self + where + H: Handler + Send + 'static, + H::Scope: Send, + { + // Meta系はデフォルトでスコープを開始しておく + let mut wrapper = HandlerWrapper::new(handler); + wrapper.start_scope(); + self.usage_handlers.push(Box::new(wrapper)); + self + } + + /// PingKind用のHandlerを登録 + pub fn on_ping(&mut self, handler: H) -> &mut Self + where + H: Handler + Send + 'static, + H::Scope: Send, + { + let mut wrapper = HandlerWrapper::new(handler); + wrapper.start_scope(); + self.ping_handlers.push(Box::new(wrapper)); + self + } + + /// StatusKind用のHandlerを登録 + pub fn on_status(&mut self, handler: H) -> &mut Self + where + H: Handler + Send + 'static, + H::Scope: Send, + { + let mut wrapper = HandlerWrapper::new(handler); + wrapper.start_scope(); + self.status_handlers.push(Box::new(wrapper)); + self + } + + /// ErrorKind用のHandlerを登録 + pub fn on_error(&mut self, handler: H) -> &mut Self + where + H: Handler + Send + 'static, + H::Scope: Send, + { + let mut wrapper = HandlerWrapper::new(handler); + wrapper.start_scope(); + self.error_handlers.push(Box::new(wrapper)); + self + } + + /// TextBlockKind用のHandlerを登録 + pub fn on_text_block(&mut self, handler: H) -> &mut Self + where + H: Handler + Send + 'static, + H::Scope: Send, + { + self.text_block_handlers + .push(Box::new(TextBlockHandlerWrapper::new(handler))); + self + } + + /// ThinkingBlockKind用のHandlerを登録 + pub fn on_thinking_block(&mut self, handler: H) -> &mut Self + where + H: Handler + Send + 'static, + H::Scope: Send, + { + self.thinking_block_handlers + .push(Box::new(ThinkingBlockHandlerWrapper::new(handler))); + self + } + + /// ToolUseBlockKind用のHandlerを登録 + pub fn on_tool_use_block(&mut self, handler: H) -> &mut Self + where + H: Handler + Send + 'static, + H::Scope: Send, + { + self.tool_use_block_handlers + .push(Box::new(ToolUseBlockHandlerWrapper::new(handler))); + self + } + + // ========================================================================= + // Event Dispatch + // ========================================================================= + + /// メインのディスパッチエントリポイント + pub fn dispatch(&mut self, event: &Event) { + match event { + // Meta系: 即時ディスパッチ（登録順） + Event::Usage(u) => self.dispatch_usage(u), + Event::Ping(p) => self.dispatch_ping(p), + Event::Status(s) => self.dispatch_status(s), + Event::Error(e) => self.dispatch_error(e), + + // Block系: スコープ管理しながらディスパッチ + Event::BlockStart(s) => self.handle_block_start(s), + Event::BlockDelta(d) => self.handle_block_delta(d), + Event::BlockStop(s) => self.handle_block_stop(s), + Event::BlockAbort(a) => self.handle_block_abort(a), + } + } + + fn dispatch_usage(&mut self, event: &UsageEvent) { + for handler in &mut self.usage_handlers { + handler.dispatch(event); + } + } + + fn dispatch_ping(&mut self, event: &PingEvent) { + for handler in &mut self.ping_handlers { + handler.dispatch(event); + } + } + + fn dispatch_status(&mut self, event: &StatusEvent) { + for handler in &mut self.status_handlers { + handler.dispatch(event); + } + } + + fn dispatch_error(&mut self, event: &ErrorEvent) { + for handler in &mut self.error_handlers { + handler.dispatch(event); + } + } + + fn handle_block_start(&mut self, start: &BlockStart) { + self.current_block = Some(start.block_type); + + let handlers = self.get_block_handlers_mut(start.block_type); + for handler in handlers { + handler.start_scope(); + handler.dispatch_start(start); + } + } + + fn handle_block_delta(&mut self, delta: &BlockDelta) { + let block_type = delta.delta.block_type(); + let handlers = self.get_block_handlers_mut(block_type); + for handler in handlers { + handler.dispatch_delta(delta); + } + } + + fn handle_block_stop(&mut self, stop: &BlockStop) { + let handlers = self.get_block_handlers_mut(stop.block_type); + for handler in handlers { + handler.dispatch_stop(stop); + handler.end_scope(); + } + self.current_block = None; + } + + fn handle_block_abort(&mut self, abort: &BlockAbort) { + let handlers = self.get_block_handlers_mut(abort.block_type); + for handler in handlers { + handler.dispatch_abort(abort); + handler.end_scope(); + } + self.current_block = None; + } + + fn get_block_handlers_mut(&mut self, block_type: BlockType) -> &mut Vec> { + match block_type { + BlockType::Text => &mut self.text_block_handlers, + BlockType::Thinking => &mut self.thinking_block_handlers, + BlockType::ToolUse => &mut self.tool_use_block_handlers, + BlockType::ToolResult => &mut self.text_block_handlers, // ToolResultはTextとして扱う + } + } + + /// 現在アクティブなブロックタイプを取得 + pub fn current_block(&self) -> Option { + self.current_block + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Arc, Mutex}; + + #[test] + fn test_timeline_creation() { + let timeline = Timeline::new(); + assert!(timeline.current_block().is_none()); + } + + #[test] + fn test_meta_event_dispatch() { + // シンプルなテスト用構造体 + struct TestUsageHandler { + calls: Arc>>, + } + + impl Handler for TestUsageHandler { + type Scope = (); + fn on_event(&mut self, _scope: &mut (), event: &UsageEvent) { + self.calls.lock().unwrap().push(event.clone()); + } + } + + let calls = Arc::new(Mutex::new(Vec::new())); + let handler = TestUsageHandler { calls: calls.clone() }; + + let mut timeline = Timeline::new(); + timeline.on_usage(handler); + + timeline.dispatch(&Event::usage(100, 50)); + + let recorded = calls.lock().unwrap(); + assert_eq!(recorded.len(), 1); + assert_eq!(recorded[0].input_tokens, Some(100)); + } +}