feat: Introduce event timeline and handler system

This commit is contained in:
Keisuke Hirata 2026-01-05 23:10:32 +09:00
parent 14db66e5b0
commit 6d6ae24ffe
10 changed files with 1208 additions and 0 deletions

13
worker-macros/Cargo.toml Normal file
View File

@ -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" }

41
worker-macros/src/lib.rs Normal file
View File

@ -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<WeatherInfo, ToolError> {
/// // ...
/// }
/// ```
#[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
}

8
worker-types/Cargo.toml Normal file
View File

@ -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"

276
worker-types/src/event.rs Normal file
View File

@ -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<u64>,
}
/// 使用量イベント
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct UsageEvent {
/// 入力トークン数
pub input_tokens: Option<u64>,
/// 出力トークン数
pub output_tokens: Option<u64>,
/// 合計トークン数
pub total_tokens: Option<u64>,
/// キャッシュ読み込みトークン数
pub cache_read_input_tokens: Option<u64>,
/// キャッシュ作成トークン数
pub cache_creation_input_tokens: Option<u64>,
}
/// ステータスイベント
#[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<String>,
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<StopReason>,
}
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<String>) -> Self {
Event::BlockDelta(BlockDelta {
index,
delta: DeltaContent::Text(text.into()),
})
}
/// テキストブロック停止イベントを作成
pub fn text_block_stop(index: usize, stop_reason: Option<StopReason>) -> Self {
Event::BlockStop(BlockStop {
index,
block_type: BlockType::Text,
stop_reason,
})
}
/// ツール使用ブロック開始イベントを作成
pub fn tool_use_start(index: usize, id: impl Into<String>, name: impl Into<String>) -> 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<String>) -> 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 })
}
}

141
worker-types/src/handler.rs Normal file
View File

@ -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<K: Kind> {
/// 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<StopReason>,
}
/// 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,
}

12
worker-types/src/lib.rs Normal file
View File

@ -0,0 +1,12 @@
//! worker-types - LLMワーカーで使用される型定義
//!
//! このクレートは以下を提供します:
//! - Event: llm_client層からのフラットなイベント列挙
//! - Kind/Handler: タイムライン層でのイベント処理トレイト
//! - 各種イベント構造体
mod event;
mod handler;
pub use event::*;
pub use handler::*;

10
worker/Cargo.toml Normal file
View File

@ -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" }

View File

@ -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<Event> {
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<UsageKind> 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<String>,
}
impl TextCollector {
fn new() -> Self {
Self { results: Vec::new() }
}
}
impl Handler<TextBlockKind> 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<ToolUseBlockKind> 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),
));
}
}
}
}

10
worker/src/lib.rs Normal file
View File

@ -0,0 +1,10 @@
//! worker - LLMワーカーのメイン実装
//!
//! このクレートは以下を提供します:
//! - Timeline: イベントストリームの状態管理とハンドラーへのディスパッチ
//! - 型消去されたHandler実装
mod timeline;
pub use timeline::*;
pub use worker_types::*;

565
worker/src/timeline.rs Normal file
View File

@ -0,0 +1,565 @@
//! Timeline層の実装
//!
//! イベントストリームを受信し、登録されたHandlerへディスパッチする
use std::marker::PhantomData;
use worker_types::*;
// =============================================================================
// Type-erased Handler
// =============================================================================
/// 型消去されたHandler trait
///
/// 各Handlerは独自のScope型を持つため、Timelineで保持するには型消去が必要
pub trait ErasedHandler<K: Kind>: Send {
/// イベントをディスパッチ
fn dispatch(&mut self, event: &K::Event);
/// スコープを開始Block開始時
fn start_scope(&mut self);
/// スコープを終了Block終了時
fn end_scope(&mut self);
}
/// Handler<K>からErasedHandler<K>へのラッパー
pub struct HandlerWrapper<H, K>
where
H: Handler<K>,
K: Kind,
{
handler: H,
scope: Option<H::Scope>,
// fn() -> K は常にSend+Syncなので、Kの制約に関係なくSendを満たせる
_kind: PhantomData<fn() -> K>,
}
impl<H, K> HandlerWrapper<H, K>
where
H: Handler<K>,
K: Kind,
{
pub fn new(handler: H) -> Self {
Self {
handler,
scope: None,
_kind: PhantomData,
}
}
}
impl<H, K> ErasedHandler<K> for HandlerWrapper<H, K>
where
H: Handler<K> + 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<H>
where
H: Handler<TextBlockKind>,
{
handler: H,
scope: Option<H::Scope>,
}
impl<H> TextBlockHandlerWrapper<H>
where
H: Handler<TextBlockKind>,
{
fn new(handler: H) -> Self {
Self {
handler,
scope: None,
}
}
}
impl<H> ErasedBlockHandler for TextBlockHandlerWrapper<H>
where
H: Handler<TextBlockKind> + 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<H>
where
H: Handler<ThinkingBlockKind>,
{
handler: H,
scope: Option<H::Scope>,
}
impl<H> ThinkingBlockHandlerWrapper<H>
where
H: Handler<ThinkingBlockKind>,
{
fn new(handler: H) -> Self {
Self {
handler,
scope: None,
}
}
}
impl<H> ErasedBlockHandler for ThinkingBlockHandlerWrapper<H>
where
H: Handler<ThinkingBlockKind> + 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<H>
where
H: Handler<ToolUseBlockKind>,
{
handler: H,
scope: Option<H::Scope>,
current_tool: Option<(String, String)>, // (id, name)
}
impl<H> ToolUseBlockHandlerWrapper<H>
where
H: Handler<ToolUseBlockKind>,
{
fn new(handler: H) -> Self {
Self {
handler,
scope: None,
current_tool: None,
}
}
}
impl<H> ErasedBlockHandler for ToolUseBlockHandlerWrapper<H>
where
H: Handler<ToolUseBlockKind> + 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<Box<dyn ErasedHandler<UsageKind>>>,
ping_handlers: Vec<Box<dyn ErasedHandler<PingKind>>>,
status_handlers: Vec<Box<dyn ErasedHandler<StatusKind>>>,
error_handlers: Vec<Box<dyn ErasedHandler<ErrorKind>>>,
// Block系ハンドラーBlockTypeごとにグループ化
text_block_handlers: Vec<Box<dyn ErasedBlockHandler>>,
thinking_block_handlers: Vec<Box<dyn ErasedBlockHandler>>,
tool_use_block_handlers: Vec<Box<dyn ErasedBlockHandler>>,
// 現在アクティブなブロック
current_block: Option<BlockType>,
}
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<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<UsageKind> + 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<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<PingKind> + 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<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<StatusKind> + 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<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<ErrorKind> + 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<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<TextBlockKind> + Send + 'static,
H::Scope: Send,
{
self.text_block_handlers
.push(Box::new(TextBlockHandlerWrapper::new(handler)));
self
}
/// ThinkingBlockKind用のHandlerを登録
pub fn on_thinking_block<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<ThinkingBlockKind> + Send + 'static,
H::Scope: Send,
{
self.thinking_block_handlers
.push(Box::new(ThinkingBlockHandlerWrapper::new(handler)));
self
}
/// ToolUseBlockKind用のHandlerを登録
pub fn on_tool_use_block<H>(&mut self, handler: H) -> &mut Self
where
H: Handler<ToolUseBlockKind> + 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<Box<dyn ErasedBlockHandler>> {
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<BlockType> {
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<Mutex<Vec<UsageEvent>>>,
}
impl Handler<UsageKind> 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));
}
}