feat: #[tool_registry] and #[tool] macros

This commit is contained in:
2026-01-06 22:42:24 +09:00
parent e82e0a3ed9
commit ddd80908c0
11 changed files with 1262 additions and 125 deletions
+75 -8
View File
@@ -204,26 +204,81 @@ impl EventPlayer {
///
/// 事前に定義されたイベントシーケンスをストリームとして返す。
/// fixtureファイルからロードすることも、直接イベントを渡すこともできる。
///
/// # 複数リクエスト対応
///
/// `with_responses()`を使用して、複数回のリクエストに対して異なるレスポンスを設定できる。
/// リクエスト回数が設定されたレスポンス数を超えた場合は空のストリームを返す。
pub struct MockLlmClient {
events: Vec<Event>,
/// 各リクエストに対するレスポンス(イベントシーケンス)
responses: std::sync::Arc<std::sync::Mutex<Vec<Vec<Event>>>>,
/// 現在のリクエストインデックス
request_index: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}
#[allow(dead_code)]
impl MockLlmClient {
/// イベントリストから直接作成
/// イベントリストから直接作成(単一レスポンス)
///
/// すべてのリクエストに対して同じイベントシーケンスを返す(従来の動作)
pub fn new(events: Vec<Event>) -> Self {
Self { events }
Self {
responses: std::sync::Arc::new(std::sync::Mutex::new(vec![events])),
request_index: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}
}
/// fixtureファイルからロード
/// 複数のレスポンスを設定
///
/// 各リクエストに対して順番にイベントシーケンスを返す。
/// N回目のリクエストにはN番目のレスポンスが使用される。
///
/// # Example
/// ```ignore
/// let client = MockLlmClient::with_responses(vec![
/// // 1回目のリクエスト: ツール呼び出し
/// vec![Event::tool_use_start(0, "call_1", "my_tool"), ...],
/// // 2回目のリクエスト: テキストレスポンス
/// vec![Event::text_block_start(0), ...],
/// ]);
/// ```
pub fn with_responses(responses: Vec<Vec<Event>>) -> Self {
Self {
responses: std::sync::Arc::new(std::sync::Mutex::new(responses)),
request_index: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
}
}
/// fixtureファイルからロード(単一レスポンス)
pub fn from_fixture(path: impl AsRef<Path>) -> std::io::Result<Self> {
let player = EventPlayer::load(path)?;
let events = player.parse_events();
Ok(Self { events })
Ok(Self::new(events))
}
/// 保持しているイベント数を取得
/// 保持しているレスポンス数を取得
pub fn response_count(&self) -> usize {
self.responses.lock().unwrap().len()
}
/// 最初のレスポンスのイベント数を取得(後方互換性)
pub fn event_count(&self) -> usize {
self.events.len()
self.responses
.lock()
.unwrap()
.first()
.map(|v| v.len())
.unwrap_or(0)
}
/// 現在のリクエストインデックスを取得
pub fn current_request_index(&self) -> usize {
self.request_index.load(std::sync::atomic::Ordering::SeqCst)
}
/// リクエストインデックスをリセット
pub fn reset(&self) {
self.request_index.store(0, std::sync::atomic::Ordering::SeqCst);
}
}
@@ -233,8 +288,20 @@ impl LlmClient for MockLlmClient {
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
let events = self.events.clone();
let index = self.request_index.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let events = {
let responses = self.responses.lock().unwrap();
if index < responses.len() {
responses[index].clone()
} else {
// レスポンスが尽きた場合は空のストリーム
Vec::new()
}
};
let stream = futures::stream::iter(events.into_iter().map(Ok));
Ok(Box::pin(stream))
}
}
+254
View File
@@ -0,0 +1,254 @@
//! 並列ツール実行のテスト
//!
//! Workerが複数のツールを並列に実行することを確認する。
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use worker::Worker;
use worker_types::{Event, Message, ResponseStatus, StatusEvent, Tool, ToolError, ToolResult, ToolCall, ControlFlow, HookError, WorkerHook};
mod common;
use common::MockLlmClient;
// =============================================================================
// Parallel Execution Test Tools
// =============================================================================
/// 一定時間待機してから応答するツール
#[derive(Clone)]
struct SlowTool {
name: String,
delay_ms: u64,
call_count: Arc<AtomicUsize>,
}
impl SlowTool {
fn new(name: impl Into<String>, delay_ms: u64) -> Self {
Self {
name: name.into(),
delay_ms,
call_count: Arc::new(AtomicUsize::new(0)),
}
}
fn call_count(&self) -> usize {
self.call_count.load(Ordering::SeqCst)
}
}
#[async_trait]
impl Tool for SlowTool {
fn name(&self) -> &str {
&self.name
}
fn description(&self) -> &str {
"A tool that waits before responding"
}
fn input_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {}
})
}
async fn execute(&self, _input_json: &str) -> Result<String, ToolError> {
self.call_count.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(self.delay_ms)).await;
Ok(format!("Completed after {}ms", self.delay_ms))
}
}
// =============================================================================
// Tests
// =============================================================================
/// 複数のツールが並列に実行されることを確認
///
/// 各ツールが100msかかる場合、逐次実行なら300ms以上かかるが、
/// 並列実行なら100ms程度で完了するはず。
#[tokio::test]
async fn test_parallel_tool_execution() {
// 3つのツール呼び出しを含むイベントシーケンス
let events = vec![
Event::tool_use_start(0, "call_1", "slow_tool_1"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::tool_use_start(1, "call_2", "slow_tool_2"),
Event::tool_input_delta(1, r#"{}"#),
Event::tool_use_stop(1),
Event::tool_use_start(2, "call_3", "slow_tool_3"),
Event::tool_input_delta(2, r#"{}"#),
Event::tool_use_stop(2),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
];
let client = MockLlmClient::new(events);
let mut worker = Worker::new(client);
// 各ツールは100ms待機
let tool1 = SlowTool::new("slow_tool_1", 100);
let tool2 = SlowTool::new("slow_tool_2", 100);
let tool3 = SlowTool::new("slow_tool_3", 100);
let tool1_clone = tool1.clone();
let tool2_clone = tool2.clone();
let tool3_clone = tool3.clone();
worker.register_tool(tool1);
worker.register_tool(tool2);
worker.register_tool(tool3);
let messages = vec![Message::user("Run all tools")];
let start = Instant::now();
let _result = worker.run(messages).await;
let elapsed = start.elapsed();
// 全ツールが呼び出されたことを確認
assert_eq!(tool1_clone.call_count(), 1, "Tool 1 should be called once");
assert_eq!(tool2_clone.call_count(), 1, "Tool 2 should be called once");
assert_eq!(tool3_clone.call_count(), 1, "Tool 3 should be called once");
// 並列実行なら200ms以下で完了するはず(逐次なら300ms以上)
// マージン込みで250msをしきい値とする
assert!(
elapsed < Duration::from_millis(250),
"Parallel execution should complete in ~100ms, but took {:?}",
elapsed
);
println!("Parallel execution completed in {:?}", elapsed);
}
/// Hook: before_tool_call でスキップされたツールは実行されないことを確認
#[tokio::test]
async fn test_before_tool_call_skip() {
let events = vec![
Event::tool_use_start(0, "call_1", "allowed_tool"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::tool_use_start(1, "call_2", "blocked_tool"),
Event::tool_input_delta(1, r#"{}"#),
Event::tool_use_stop(1),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
];
let client = MockLlmClient::new(events);
let mut worker = Worker::new(client);
let allowed_tool = SlowTool::new("allowed_tool", 10);
let blocked_tool = SlowTool::new("blocked_tool", 10);
let allowed_clone = allowed_tool.clone();
let blocked_clone = blocked_tool.clone();
worker.register_tool(allowed_tool);
worker.register_tool(blocked_tool);
// "blocked_tool" をスキップするHook
struct BlockingHook;
#[async_trait]
impl WorkerHook for BlockingHook {
async fn before_tool_call(&self, tool_call: &mut ToolCall) -> Result<ControlFlow, HookError> {
if tool_call.name == "blocked_tool" {
Ok(ControlFlow::Skip)
} else {
Ok(ControlFlow::Continue)
}
}
}
worker.add_hook(BlockingHook);
let messages = vec![Message::user("Test hook")];
let _result = worker.run(messages).await;
// allowed_tool は呼び出されるが、blocked_tool は呼び出されない
assert_eq!(allowed_clone.call_count(), 1, "Allowed tool should be called");
assert_eq!(blocked_clone.call_count(), 0, "Blocked tool should not be called");
}
/// Hook: after_tool_call で結果が改変されることを確認
#[tokio::test]
async fn test_after_tool_call_modification() {
// 複数リクエストに対応するレスポンスを準備
let client = MockLlmClient::with_responses(vec![
// 1回目のリクエスト: ツール呼び出し
vec![
Event::tool_use_start(0, "call_1", "test_tool"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
// 2回目のリクエスト: ツール結果を受けてテキストレスポンス
vec![
Event::text_block_start(0),
Event::text_delta(0, "Done!"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
]);
let mut worker = Worker::new(client);
#[derive(Clone)]
struct SimpleTool;
#[async_trait]
impl Tool for SimpleTool {
fn name(&self) -> &str { "test_tool" }
fn description(&self) -> &str { "Test" }
fn input_schema(&self) -> serde_json::Value { serde_json::json!({}) }
async fn execute(&self, _: &str) -> Result<String, ToolError> {
Ok("Original Result".to_string())
}
}
worker.register_tool(SimpleTool);
// 結果を改変するHook
struct ModifyingHook {
modified_content: Arc<std::sync::Mutex<Option<String>>>,
}
#[async_trait]
impl WorkerHook for ModifyingHook {
async fn after_tool_call(&self, tool_result: &mut ToolResult) -> Result<ControlFlow, HookError> {
tool_result.content = format!("[Modified] {}", tool_result.content);
*self.modified_content.lock().unwrap() = Some(tool_result.content.clone());
Ok(ControlFlow::Continue)
}
}
let modified_content = Arc::new(std::sync::Mutex::new(None));
worker.add_hook(ModifyingHook { modified_content: modified_content.clone() });
let messages = vec![Message::user("Test modification")];
let result = worker.run(messages).await;
assert!(result.is_ok(), "Worker should complete: {:?}", result);
// Hookが呼ばれて内容が改変されたことを確認
let content = modified_content.lock().unwrap().clone();
assert!(content.is_some(), "Hook should have been called");
assert!(
content.unwrap().contains("[Modified]"),
"Result should be modified"
);
}
+212
View File
@@ -0,0 +1,212 @@
//! ツールマクロのテスト
//!
//! `#[tool_registry]` と `#[tool]` マクロの動作を確認する。
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
// マクロ展開に必要なインポート
use schemars;
use serde;
use worker_macros::tool_registry;
use worker_types::Tool;
// =============================================================================
// Test: Basic Tool Generation
// =============================================================================
/// シンプルなコンテキスト構造体
#[derive(Clone)]
struct SimpleContext {
prefix: String,
}
#[tool_registry]
impl SimpleContext {
/// メッセージに挨拶を追加する
///
/// 指定されたメッセージにプレフィックスを付けて返します。
#[tool]
async fn greet(&self, message: String) -> String {
format!("{}: {}", self.prefix, message)
}
/// 二つの数を足す
#[tool]
async fn add(&self, a: i32, b: i32) -> i32 {
a + b
}
/// 引数なしのツール
#[tool]
async fn get_prefix(&self) -> String {
self.prefix.clone()
}
}
#[tokio::test]
async fn test_basic_tool_generation() {
let ctx = SimpleContext {
prefix: "Hello".to_string(),
};
// ファクトリメソッドでツールを取得
let greet_tool = ctx.greet_tool();
// 名前の確認
assert_eq!(greet_tool.name(), "greet");
// 説明の確認(docコメントから取得)
let desc = greet_tool.description();
assert!(desc.contains("メッセージに挨拶を追加する"), "Description should contain doc comment: {}", desc);
// スキーマの確認
let schema = greet_tool.input_schema();
println!("Schema: {}", serde_json::to_string_pretty(&schema).unwrap());
assert!(schema.get("properties").is_some(), "Schema should have properties");
// 実行テスト
let result = greet_tool.execute(r#"{"message": "World"}"#).await;
assert!(result.is_ok(), "Should execute successfully");
let output = result.unwrap();
assert!(output.contains("Hello"), "Output should contain prefix");
assert!(output.contains("World"), "Output should contain message");
}
#[tokio::test]
async fn test_multiple_arguments() {
let ctx = SimpleContext {
prefix: "".to_string(),
};
let add_tool = ctx.add_tool();
assert_eq!(add_tool.name(), "add");
let result = add_tool.execute(r#"{"a": 10, "b": 20}"#).await;
assert!(result.is_ok());
let output = result.unwrap();
assert!(output.contains("30"), "Should contain sum: {}", output);
}
#[tokio::test]
async fn test_no_arguments() {
let ctx = SimpleContext {
prefix: "TestPrefix".to_string(),
};
let get_prefix_tool = ctx.get_prefix_tool();
assert_eq!(get_prefix_tool.name(), "get_prefix");
// 空のJSONオブジェクトで呼び出し
let result = get_prefix_tool.execute(r#"{}"#).await;
assert!(result.is_ok());
let output = result.unwrap();
assert!(output.contains("TestPrefix"), "Should contain prefix: {}", output);
}
#[tokio::test]
async fn test_invalid_arguments() {
let ctx = SimpleContext {
prefix: "".to_string(),
};
let greet_tool = ctx.greet_tool();
// 不正なJSON
let result = greet_tool.execute(r#"{"wrong_field": "value"}"#).await;
assert!(result.is_err(), "Should fail with invalid arguments");
}
// =============================================================================
// Test: Result Return Type
// =============================================================================
#[derive(Clone)]
struct FallibleContext;
#[derive(Debug)]
struct MyError(String);
impl std::fmt::Display for MyError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[tool_registry]
impl FallibleContext {
/// 与えられた値を検証する
#[tool]
async fn validate(&self, value: i32) -> Result<String, MyError> {
if value > 0 {
Ok(format!("Valid: {}", value))
} else {
Err(MyError("Value must be positive".to_string()))
}
}
}
#[tokio::test]
async fn test_result_return_type_success() {
let ctx = FallibleContext;
let validate_tool = ctx.validate_tool();
let result = validate_tool.execute(r#"{"value": 42}"#).await;
assert!(result.is_ok(), "Should succeed for positive value");
let output = result.unwrap();
assert!(output.contains("Valid"), "Should contain Valid: {}", output);
}
#[tokio::test]
async fn test_result_return_type_error() {
let ctx = FallibleContext;
let validate_tool = ctx.validate_tool();
let result = validate_tool.execute(r#"{"value": -1}"#).await;
assert!(result.is_err(), "Should fail for negative value");
let err = result.unwrap_err();
assert!(err.to_string().contains("positive"), "Error should mention positive: {}", err);
}
// =============================================================================
// Test: Synchronous Methods
// =============================================================================
#[derive(Clone)]
struct SyncContext {
counter: Arc<AtomicUsize>,
}
#[tool_registry]
impl SyncContext {
/// カウンターをインクリメントして返す (非async)
#[tool]
fn increment(&self) -> usize {
self.counter.fetch_add(1, Ordering::SeqCst) + 1
}
}
#[tokio::test]
async fn test_sync_method() {
let ctx = SyncContext {
counter: Arc::new(AtomicUsize::new(0)),
};
let increment_tool = ctx.increment_tool();
// 3回実行
let result1 = increment_tool.execute(r#"{}"#).await;
let result2 = increment_tool.execute(r#"{}"#).await;
let result3 = increment_tool.execute(r#"{}"#).await;
assert!(result1.is_ok());
assert!(result2.is_ok());
assert!(result3.is_ok());
// カウンターは3になっているはず
assert_eq!(ctx.counter.load(Ordering::SeqCst), 3);
}
+2 -3
View File
@@ -11,7 +11,7 @@ use std::sync::Arc;
use async_trait::async_trait;
use common::MockLlmClient;
use worker::{Worker, WorkerConfig};
use worker::Worker;
use worker_types::{Tool, ToolError};
/// フィクスチャディレクトリのパス
@@ -163,8 +163,7 @@ async fn test_worker_tool_call() {
let tool_for_check = weather_tool.clone();
worker.register_tool(weather_tool);
// 設定: ツール実行後はターン終了(ループしない)
worker = worker.config(WorkerConfig { max_turns: 1 });
// メッセージを送信
let messages = vec![worker_types::Message::user("What's the weather in Tokyo?")];