feat: Introduce WorkerSubscriber system with adapters

2026-01-07 00:30:58 +09:00
parent 9547d40538
commit 89b12d277a
6 changed files with 693 additions and 6 deletions
@@ -0,0 +1,237 @@
+//! WorkerSubscriberのテスト
+//!
+//! WorkerSubscriberを使ってイベントを購読するテスト
+
+mod common;
+
+use std::sync::{Arc, Mutex};
+
+use common::MockLlmClient;
+use worker::{Worker, WorkerSubscriber};
+use worker_types::{
+    ErrorEvent, Event, Message, ResponseStatus, StatusEvent, TextBlockEvent, ToolCall,
+    ToolUseBlockEvent, UsageEvent,
+};
+
+// =============================================================================
+// Test Subscriber
+// =============================================================================
+
+/// テスト用のシンプルなSubscriber実装
+struct TestSubscriber {
+    // 記録用のバッファ
+    text_deltas: Arc<Mutex<Vec<String>>>,
+    text_completes: Arc<Mutex<Vec<String>>>,
+    tool_call_completes: Arc<Mutex<Vec<ToolCall>>>,
+    usage_events: Arc<Mutex<Vec<UsageEvent>>>,
+    status_events: Arc<Mutex<Vec<StatusEvent>>>,
+    turn_starts: Arc<Mutex<Vec<usize>>>,
+    turn_ends: Arc<Mutex<Vec<usize>>>,
+}
+
+impl TestSubscriber {
+    fn new() -> Self {
+        Self {
+            text_deltas: Arc::new(Mutex::new(Vec::new())),
+            text_completes: Arc::new(Mutex::new(Vec::new())),
+            tool_call_completes: Arc::new(Mutex::new(Vec::new())),
+            usage_events: Arc::new(Mutex::new(Vec::new())),
+            status_events: Arc::new(Mutex::new(Vec::new())),
+            turn_starts: Arc::new(Mutex::new(Vec::new())),
+            turn_ends: Arc::new(Mutex::new(Vec::new())),
+        }
+    }
+}
+
+impl WorkerSubscriber for TestSubscriber {
+    type TextBlockScope = String;
+    type ToolUseBlockScope = ();
+
+    fn on_text_block(&mut self, buffer: &mut String, event: &TextBlockEvent) {
+        if let TextBlockEvent::Delta(text) = event {
+            buffer.push_str(text);
+            self.text_deltas.lock().unwrap().push(text.clone());
+        }
+    }
+
+    fn on_text_complete(&mut self, text: &str) {
+        self.text_completes.lock().unwrap().push(text.to_string());
+    }
+
+    fn on_tool_use_block(&mut self, _scope: &mut (), _event: &ToolUseBlockEvent) {
+        // 必要に応じて処理
+    }
+
+    fn on_tool_call_complete(&mut self, call: &ToolCall) {
+        self.tool_call_completes.lock().unwrap().push(call.clone());
+    }
+
+    fn on_usage(&mut self, event: &UsageEvent) {
+        self.usage_events.lock().unwrap().push(event.clone());
+    }
+
+    fn on_status(&mut self, event: &StatusEvent) {
+        self.status_events.lock().unwrap().push(event.clone());
+    }
+
+    fn on_error(&mut self, _event: &ErrorEvent) {
+        // 必要に応じて処理
+    }
+
+    fn on_turn_start(&mut self, turn: usize) {
+        self.turn_starts.lock().unwrap().push(turn);
+    }
+
+    fn on_turn_end(&mut self, turn: usize) {
+        self.turn_ends.lock().unwrap().push(turn);
+    }
+}
+
+// =============================================================================
+// Tests
+// =============================================================================
+
+/// WorkerSubscriberがテキストブロックイベントを正しく受け取ることを確認
+#[tokio::test]
+async fn test_subscriber_text_block_events() {
+    // テキストレスポンスを含むイベントシーケンス
+    let events = vec![
+        Event::text_block_start(0),
+        Event::text_delta(0, "Hello, "),
+        Event::text_delta(0, "World!"),
+        Event::text_block_stop(0, None),
+        Event::Status(StatusEvent {
+            status: ResponseStatus::Completed,
+        }),
+    ];
+
+    let client = MockLlmClient::new(events);
+    let mut worker = Worker::new(client);
+
+    // Subscriberを登録
+    let subscriber = TestSubscriber::new();
+    let text_deltas = subscriber.text_deltas.clone();
+    let text_completes = subscriber.text_completes.clone();
+    worker.subscribe(subscriber);
+
+    // 実行
+    let messages = vec![Message::user("Greet me")];
+    let result = worker.run(messages).await;
+
+    assert!(result.is_ok(), "Worker should complete: {:?}", result);
+
+    // デルタが収集されていることを確認
+    let deltas = text_deltas.lock().unwrap();
+    assert_eq!(deltas.len(), 2);
+    assert_eq!(deltas[0], "Hello, ");
+    assert_eq!(deltas[1], "World!");
+
+    // 完了テキストが収集されていることを確認
+    let completes = text_completes.lock().unwrap();
+    assert_eq!(completes.len(), 1);
+    assert_eq!(completes[0], "Hello, World!");
+}
+
+/// WorkerSubscriberがツール呼び出し完了イベントを正しく受け取ることを確認
+#[tokio::test]
+async fn test_subscriber_tool_call_complete() {
+    // ツール呼び出しを含むイベントシーケンス
+    let events = vec![
+        Event::tool_use_start(0, "call_123", "get_weather"),
+        Event::tool_input_delta(0, r#"{"city":"#),
+        Event::tool_input_delta(0, r#""Tokyo"}"#),
+        Event::tool_use_stop(0),
+        Event::Status(StatusEvent {
+            status: ResponseStatus::Completed,
+        }),
+    ];
+
+    let client = MockLlmClient::new(events);
+    let mut worker = Worker::new(client);
+
+    // Subscriberを登録
+    let subscriber = TestSubscriber::new();
+    let tool_call_completes = subscriber.tool_call_completes.clone();
+    worker.subscribe(subscriber);
+
+    // 実行
+    let messages = vec![Message::user("Weather please")];
+    let _ = worker.run(messages).await;
+
+    // ツール呼び出し完了が収集されていることを確認
+    let completes = tool_call_completes.lock().unwrap();
+    assert_eq!(completes.len(), 1);
+    assert_eq!(completes[0].name, "get_weather");
+    assert_eq!(completes[0].id, "call_123");
+    assert_eq!(completes[0].input["city"], "Tokyo");
+}
+
+/// WorkerSubscriberがターンイベントを正しく受け取ることを確認
+#[tokio::test]
+async fn test_subscriber_turn_events() {
+    let events = vec![
+        Event::text_block_start(0),
+        Event::text_delta(0, "Done!"),
+        Event::text_block_stop(0, None),
+        Event::Status(StatusEvent {
+            status: ResponseStatus::Completed,
+        }),
+    ];
+
+    let client = MockLlmClient::new(events);
+    let mut worker = Worker::new(client);
+
+    // Subscriberを登録
+    let subscriber = TestSubscriber::new();
+    let turn_starts = subscriber.turn_starts.clone();
+    let turn_ends = subscriber.turn_ends.clone();
+    worker.subscribe(subscriber);
+
+    // 実行
+    let messages = vec![Message::user("Do something")];
+    let result = worker.run(messages).await;
+
+    assert!(result.is_ok());
+
+    // ターンイベントが収集されていることを確認
+    let starts = turn_starts.lock().unwrap();
+    let ends = turn_ends.lock().unwrap();
+
+    assert_eq!(starts.len(), 1);
+    assert_eq!(starts[0], 0); // 最初のターン
+
+    assert_eq!(ends.len(), 1);
+    assert_eq!(ends[0], 0);
+}
+
+/// WorkerSubscriberがUsageイベントを正しく受け取ることを確認
+#[tokio::test]
+async fn test_subscriber_usage_events() {
+    let events = vec![
+        Event::text_block_start(0),
+        Event::text_delta(0, "Hello"),
+        Event::text_block_stop(0, None),
+        Event::usage(100, 50),
+        Event::Status(StatusEvent {
+            status: ResponseStatus::Completed,
+        }),
+    ];
+
+    let client = MockLlmClient::new(events);
+    let mut worker = Worker::new(client);
+
+    // Subscriberを登録
+    let subscriber = TestSubscriber::new();
+    let usage_events = subscriber.usage_events.clone();
+    worker.subscribe(subscriber);
+
+    // 実行
+    let messages = vec![Message::user("Hello")];
+    let _ = worker.run(messages).await;
+
+    // Usageイベントが収集されていることを確認
+    let usages = usage_events.lock().unwrap();
+    assert_eq!(usages.len(), 1);
+    assert_eq!(usages[0].input_tokens, Some(100));
+    assert_eq!(usages[0].output_tokens, Some(50));
+}