Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

Workerの自動キャッシュロック

Browse Source
This commit is contained in:
Hare
2026-04-11 18:47:33 +09:00
parent f241dafac8
commit 9b78c51d0a
23 changed files with 375 additions and 352 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/llm-worker/examples/worker_cancel_demo.rs
+21 -33
View File
@@ -4,9 +4,7 @@
use llm_worker::llm_client::providers::anthropic::AnthropicClient;
use llm_worker::{Worker, WorkerResult};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -25,48 +23,38 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
std::env::var("ANTHROPIC_API_KEY").expect("ANTHROPIC_API_KEY environment variable not set");
let client = AnthropicClient::new(&api_key, "claude-sonnet-4-20250514");
let worker = Arc::new(Mutex::new(Worker::new(client)));
let worker = Worker::new(client);
println!("🚀 Starting Worker...");
println!("💡 Will cancel after 2 seconds\n");
// Get cancel sender first (without holding lock)
let cancel_tx = {
let w = worker.lock().await;
w.cancel_sender()
};
// Get cancel sender before run (Mutable state)
let cancel_tx = worker.cancel_sender();
// Task 1: Run Worker
let worker_clone = worker.clone();
let task = tokio::spawn(async move {
let mut w = worker_clone.lock().await;
println!("📡 Sending request to LLM...");
match w.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await {
Ok(WorkerResult::Finished) => {
println!("✅ Task completed normally");
}
Ok(WorkerResult::Paused) => {
println!("⏸️ Task paused");
}
Ok(WorkerResult::LimitReached) => {
println!("🔒 Turn limit reached");
}
Err(e) => {
println!("❌ Task error: {}", e);
}
}
});
// Task 2: Cancel after 2 seconds
// Task: Cancel after 2 seconds
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(2)).await;
println!("\n🛑 Cancelling worker...");
let _ = cancel_tx.send(()).await;
});
// Wait for task completion
task.await?;
println!("📡 Sending request to LLM...");
// Mutable::run consumes self → (Locked, WorkerResult)
match worker.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await {
Ok((_locked, WorkerResult::Finished)) => {
println!("✅ Task completed normally");
}
Ok((_locked, WorkerResult::Paused)) => {
println!("⏸️ Task paused");
}
Ok((_locked, WorkerResult::LimitReached)) => {
println!("🔒 Turn limit reached");
}
Err(e) => {
println!("❌ Task error: {}", e);
}
}
println!("\n✨ Demo complete!");
crates/llm-worker/examples/worker_cli.rs
+24 -7
View File
@@ -438,10 +438,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Register tools (unless --no-tools)
if !args.no_tools {
let app = AppContext;
worker
.register_tool(app.get_current_time_definition())
.unwrap();
worker.register_tool(app.calculate_definition()).unwrap();
worker.register_tool(app.get_current_time_definition());
worker.register_tool(app.calculate_definition());
}
// Register streaming display handlers
@@ -465,7 +463,27 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
return Ok(());
}
// Interactive loop
// Interactive loop — first input transitions Mutable → Locked
print!("\n👤 You: ");
io::stdout().flush()?;
let mut first_input = String::new();
io::stdin().read_line(&mut first_input)?;
let first_input = first_input.trim();
if first_input == "quit" || first_input == "exit" || first_input.is_empty() {
println!("\n👋 Goodbye!");
return Ok(());
}
let (mut locked, _) = match worker.run(first_input).await {
Ok(pair) => pair,
Err(e) => {
eprintln!("\n❌ Error: {}", e);
return Ok(());
}
};
loop {
print!("\n👤 You: ");
io::stdout().flush()?;
@@ -483,8 +501,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
break;
}
// Run Worker (Worker manages history)
match worker.run(input).await {
match locked.run(input).await {
Ok(_) => {}
Err(e) => {
eprintln!("\n❌ Error: {}", e);
crates/llm-worker/src/lib.rs
+6 -4
View File
@@ -27,12 +27,14 @@
//!
//! # Cache Protection
//!
//! To maximize KV cache hit rate, transition to the locked state
//! with [`Worker::lock()`] before execution.
//! `run()` automatically locks the cache. To edit state between turns,
//! call `unlock_cache()` first; the next `run()` re-locks automatically.
//!
//! ```ignore
//! let mut locked = worker.lock();
//! locked.run("user input").await?;
//! worker.run("user input").await?;
//! worker.unlock_cache();
//! worker.set_system_prompt("new prompt");
//! worker.run("next input").await?;
//! ```
mod handler;
crates/llm-worker/src/state.rs
+5 -5
View File
@@ -1,7 +1,7 @@
//! Worker State
//!
//! State marker types for cache protection using the Type-state pattern.
//! Worker has state transitions from `Mutable` → `CacheLocked`.
//! Worker has state transitions from `Mutable` → `Locked`.
/// Marker trait representing Worker state
///
@@ -19,7 +19,7 @@ mod private {
/// - Editing message history (add, delete, clear)
/// - Registering tools and hooks
///
/// Can transition to [`CacheLocked`] state via `Worker::lock()`.
/// Can transition to [`Locked`] state via `Worker::lock()`.
///
/// # Examples
///
@@ -54,7 +54,7 @@ impl WorkerState for Mutable {}
/// Can return to [`Mutable`] state via `Worker::unlock()`,
/// but note that cache protection will be released.
#[derive(Debug, Clone, Copy, Default)]
pub struct CacheLocked;
pub struct Locked;
impl private::Sealed for CacheLocked {}
impl WorkerState for CacheLocked {}
impl private::Sealed for Locked {}
impl WorkerState for Locked {}
crates/llm-worker/src/tool_server.rs
+80 -27
View File
@@ -26,6 +26,7 @@ pub enum ToolServerError {
#[derive(Clone, Default)]
pub struct ToolServer {
tools: Arc<Mutex<ToolMap>>,
pending: Arc<Mutex<Vec<WorkerToolDefinition>>>,
}
impl ToolServer {
@@ -38,6 +39,7 @@ impl ToolServer {
pub fn handle(&self) -> ToolServerHandle {
ToolServerHandle {
tools: Arc::clone(&self.tools),
pending: Arc::clone(&self.pending),
}
}
}
@@ -46,32 +48,57 @@ impl ToolServer {
#[derive(Clone, Default)]
pub struct ToolServerHandle {
tools: Arc<Mutex<ToolMap>>,
pending: Arc<Mutex<Vec<WorkerToolDefinition>>>,
}
impl ToolServerHandle {
/// Register one tool.
pub(crate) fn register_tool(
&self,
factory: WorkerToolDefinition,
) -> Result<(), ToolServerError> {
let (meta, instance) = factory();
let mut guard = self.tools.lock().unwrap_or_else(|e| e.into_inner());
if guard.contains_key(&meta.name) {
return Err(ToolServerError::DuplicateName(meta.name));
}
guard.insert(meta.name.clone(), (meta, instance));
Ok(())
/// Queue a tool factory for deferred initialization.
///
/// The factory is **not** called here; it is stored and executed
/// when [`flush_pending`](Self::flush_pending) is called (typically
/// at the start of `Worker::run()`).
pub(crate) fn register_tool(&self, factory: WorkerToolDefinition) {
self.pending
.lock()
.unwrap_or_else(|e| e.into_inner())
.push(factory);
}
/// Register many tools.
/// Queue many tool factories for deferred initialization.
pub(crate) fn register_tools(
&self,
factories: impl IntoIterator<Item = WorkerToolDefinition>,
) -> Result<(), ToolServerError> {
for factory in factories {
self.register_tool(factory)?;
) {
let mut guard = self.pending.lock().unwrap_or_else(|e| e.into_inner());
guard.extend(factories);
}
/// Execute all pending factories and register the resulting tools.
///
/// # Panics
///
/// Panics if any factory produces a tool whose name collides with
/// an already-registered tool. Duplicate names are a programming
/// error and should be caught during development.
pub(crate) fn flush_pending(&self) {
let pending: Vec<_> = {
let mut guard = self.pending.lock().unwrap_or_else(|e| e.into_inner());
std::mem::take(&mut *guard)
};
if pending.is_empty() {
return;
}
// Execute all factories first, then validate and insert atomically.
let materialized: Vec<_> = pending.into_iter().map(|f| f()).collect();
let mut tools = self.tools.lock().unwrap_or_else(|e| e.into_inner());
for (meta, instance) in materialized {
assert!(
!tools.contains_key(&meta.name),
"duplicate tool name: '{}'",
meta.name,
);
tools.insert(meta.name.clone(), (meta, instance));
}
Ok(())
}
/// Get a tool by name for hook contexts.
@@ -143,19 +170,37 @@ mod tests {
}
#[test]
fn register_duplicate_name_fails() {
fn flush_pending_registers_tools() {
let handle = ToolServer::new().handle();
handle.register_tool(def("alpha")).expect("first register");
let err = handle
.register_tool(def("alpha"))
.expect_err("duplicate should fail");
assert_eq!(err, ToolServerError::DuplicateName("alpha".to_string()));
handle.register_tool(def("alpha"));
handle.register_tool(def("beta"));
// Before flush, no tools are available
assert!(handle.get_tool("alpha").is_none());
handle.flush_pending();
// After flush, tools are available
assert!(handle.get_tool("alpha").is_some());
assert!(handle.get_tool("beta").is_some());
}
#[test]
#[should_panic(expected = "duplicate tool name: 'alpha'")]
fn flush_pending_duplicate_name_panics() {
let handle = ToolServer::new().handle();
handle.register_tool(def("alpha"));
handle.flush_pending();
handle.register_tool(def("alpha"));
handle.flush_pending(); // panics
}
#[tokio::test]
async fn call_tool_success_and_not_found() {
let handle = ToolServer::new().handle();
handle.register_tool(def("echo")).expect("register");
handle.register_tool(def("echo"));
handle.flush_pending();
let out = handle.call_tool("echo", r#"{"x":1}"#).await.expect("call");
assert_eq!(out, r#"{"x":1}"#);
@@ -170,9 +215,10 @@ mod tests {
#[test]
fn tool_definitions_are_sorted() {
let handle = ToolServer::new().handle();
handle.register_tool(def("zeta")).expect("register zeta");
handle.register_tool(def("alpha")).expect("register alpha");
handle.register_tool(def("beta")).expect("register beta");
handle.register_tool(def("zeta"));
handle.register_tool(def("alpha"));
handle.register_tool(def("beta"));
handle.flush_pending();
let names: Vec<_> = handle
.tool_definitions_sorted()
@@ -181,4 +227,11 @@ mod tests {
.collect();
assert_eq!(names, vec!["alpha", "beta", "zeta"]);
}
#[test]
fn flush_pending_is_noop_when_empty() {
let handle = ToolServer::new().handle();
handle.flush_pending();
handle.flush_pending();
}
}
crates/llm-worker/src/worker.rs
+109 -111
View File
@@ -13,7 +13,7 @@ use crate::{
DefaultInterceptor, Interceptor, PostToolAction, PreRequestAction, PreToolAction,
PromptAction, ToolCallInfo, ToolResultInfo, TurnEndAction,
},
state::{CacheLocked, Mutable, WorkerState},
state::{Locked, Mutable, WorkerState},
callback::{
ClosureMetaHandler, ClosureTextBlockHandler, ClosureToolUseBlockHandler, TextBlockScope,
ToolUseBlockScope,
@@ -22,12 +22,9 @@ use crate::{
timeline::{TextBlockCollector, Timeline, ToolCallCollector},
timeline::event::{ErrorEvent, StatusEvent, UsageEvent},
tool::{ToolCall, ToolDefinition as WorkerToolDefinition, ToolError, ToolOutputProcessor, ToolResult},
tool_server::{ToolServer, ToolServerError, ToolServerHandle},
tool_server::{ToolServer, ToolServerHandle},
};
// =============================================================================
// Worker Error
// =============================================================================
/// Worker errors
#[derive(Debug, thiserror::Error)]
@@ -57,9 +54,6 @@ pub enum ToolRegistryError {
DuplicateName(String),
}
// =============================================================================
// Worker Config
// =============================================================================
/// Worker configuration
#[derive(Debug, Clone, Default)]
@@ -68,9 +62,6 @@ pub struct WorkerConfig {
_private: (),
}
// =============================================================================
// Worker Result Types
// =============================================================================
/// Worker execution result (status)
#[derive(Debug)]
@@ -89,9 +80,6 @@ enum ToolExecutionResult {
Paused,
}
// =============================================================================
// Worker
// =============================================================================
/// Central component for managing LLM interactions
///
@@ -100,32 +88,28 @@ enum ToolExecutionResult {
///
/// # State Transitions (Type-state)
///
/// - [`Mutable`]: Initial state. System prompt and history can be freely edited.
/// - [`CacheLocked`]: Cache-protected state. Transition via `lock()`. Prefix context is immutable.
/// - [`Mutable`]: Initial state. System prompt, history, and tools can be freely edited.
/// - [`Locked`]: Cache-protected state. Prefix context is immutable; only `run()` / `resume()` are available.
///
/// # Examples
/// Calling `run()` on a `Mutable` Worker consumes it and returns a
/// `Locked` Worker together with the result. This ensures the
/// cache prefix is fixed for optimal KV cache hit rate.
///
/// ```ignore
/// use llm_worker::{Worker, Item};
///
/// // Create a Worker and register tools
/// let mut worker = Worker::new(client)
/// .system_prompt("You are a helpful assistant.");
/// worker.register_tool(my_tool);
///
/// // Run the interaction
/// let history = worker.run("Hello!").await?;
/// ```
/// // Mutable::run() consumes self → Locked
/// let (mut worker, _result) = worker.run("Hello").await?;
///
/// # When Cache Protection is Needed
/// // Locked::run() borrows &mut self
/// worker.run("Follow-up").await?;
///
/// ```ignore
/// let mut worker = Worker::new(client)
/// .system_prompt("...");
///
/// // After setting history, lock to protect cache
/// let mut locked = worker.lock();
/// locked.run("user input").await?;
/// // To edit between turns, unlock back to Mutable
/// let mut worker = worker.unlock();
/// worker.history_mut().truncate(5);
/// let (mut worker, _result) = worker.run("Continue").await?;
/// ```
pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
/// LLM client
@@ -144,7 +128,7 @@ pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
system_prompt: Option<String>,
/// Item history (owned by Worker)
history: Vec<Item>,
/// History length at lock time (only meaningful in CacheLocked state)
/// History length at lock time (only meaningful in Locked state)
locked_prefix_len: usize,
/// Turn count
turn_count: usize,
@@ -167,40 +151,12 @@ pub struct Worker<C: LlmClient, S: WorkerState = Mutable> {
_state: PhantomData<S>,
}
// =============================================================================
// Common Implementation (available in all states)
// =============================================================================
impl<C: LlmClient, S: WorkerState> Worker<C, S> {
fn reset_interruption_state(&mut self) {
self.last_run_interrupted = false;
}
/// Execute a turn
///
/// Adds a new user message to history and sends a request to the LLM.
/// Automatically loops if there are tool calls.
pub async fn run(
&mut self,
user_input: impl Into<String>,
) -> Result<WorkerResult, WorkerError> {
self.reset_interruption_state();
// Interceptor: on_prompt_submit
let mut user_item = Item::user_message(user_input);
match self.interceptor.on_prompt_submit(&mut user_item).await {
PromptAction::Cancel(reason) => {
self.last_run_interrupted = true;
return self
.finalize_interruption(Err(WorkerError::Aborted(reason)))
.await;
}
PromptAction::Continue => {}
}
self.history.push(user_item);
let result = self.run_turn_loop().await;
self.finalize_interruption(result).await
}
fn drain_cancel_queue(&mut self) {
use tokio::sync::mpsc::error::TryRecvError;
loop {
@@ -892,19 +848,8 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
}
}
/// Resume execution (from Paused state)
///
/// Resumes turn processing from current state without adding a new user message to history.
pub async fn resume(&mut self) -> Result<WorkerResult, WorkerError> {
self.reset_interruption_state();
let result = self.run_turn_loop().await;
self.finalize_interruption(result).await
}
}
// =============================================================================
// Mutable State-Specific Implementation
// =============================================================================
impl<C: LlmClient> Worker<C, Mutable> {
/// Create a new Worker (in Mutable state)
@@ -941,43 +886,21 @@ impl<C: LlmClient> Worker<C, Mutable> {
}
}
/// Register a tool
/// Register a tool factory for deferred initialization.
///
/// Registered tools are automatically executed when called by the LLM.
/// Registering a tool with the same name will result in an error.
///
/// Available only in Mutable state.
pub fn register_tool(
&mut self,
factory: WorkerToolDefinition,
) -> Result<(), ToolRegistryError> {
match self.tool_server.register_tool(factory) {
Ok(()) => Ok(()),
Err(ToolServerError::DuplicateName(name)) => {
Err(ToolRegistryError::DuplicateName(name))
}
Err(ToolServerError::ToolNotFound(_) | ToolServerError::ToolExecution(_)) => {
unreachable!("register_tool should only fail with DuplicateName")
}
}
/// The factory is queued and executed at the next `run()` or `resume()` call.
/// Duplicate name detection occurs at that point and surfaces as
/// [`WorkerError::ToolRegistry`].
pub fn register_tool(&mut self, factory: WorkerToolDefinition) {
self.tool_server.register_tool(factory);
}
/// Register multiple tools
///
/// Available only in Mutable state.
/// Register multiple tool factories for deferred initialization.
pub fn register_tools(
&mut self,
factories: impl IntoIterator<Item = WorkerToolDefinition>,
) -> Result<(), ToolRegistryError> {
match self.tool_server.register_tools(factories) {
Ok(()) => Ok(()),
Err(ToolServerError::DuplicateName(name)) => {
Err(ToolRegistryError::DuplicateName(name))
}
Err(ToolServerError::ToolNotFound(_) | ToolServerError::ToolExecution(_)) => {
unreachable!("register_tools should only fail with DuplicateName")
}
}
) {
self.tool_server.register_tools(factories);
}
/// Set system prompt (builder pattern)
@@ -1082,40 +1005,47 @@ impl<C: LlmClient> Worker<C, Mutable> {
/// Get a mutable reference to history
///
/// Available only in Mutable state. History can be freely edited.
/// Available only in Mutable state.
pub fn history_mut(&mut self) -> &mut Vec<Item> {
&mut self.history
}
/// Set history
pub fn set_history(&mut self, items: Vec<Item>) {
self.history = items;
}
/// Add an item to history (builder pattern)
pub fn with_item(mut self, item: Item) -> Self {
self.history.push(item);
self
}
/// Add an item to history
pub fn push_item(&mut self, item: Item) {
self.history.push(item);
}
/// Add multiple items to history (builder pattern)
pub fn with_items(mut self, items: impl IntoIterator<Item = Item>) -> Self {
self.history.extend(items);
self
}
/// Add multiple items to history
pub fn extend_history(&mut self, items: impl IntoIterator<Item = Item>) {
self.history.extend(items);
}
/// Clear history
pub fn clear_history(&mut self) {
self.history.clear();
}
@@ -1148,11 +1078,48 @@ impl<C: LlmClient> Worker<C, Mutable> {
self
}
/// Lock and transition to CacheLocked state
/// Execute a turn, consuming self and transitioning to Locked.
///
/// This operation fixes the current system prompt and history as a "committed prefix".
/// After this, only appending to history is allowed, ensuring cache hits.
pub fn lock(self) -> Worker<C, CacheLocked> {
/// This is the primary entry point for first use. Equivalent to
/// `self.lock()` followed by `locked.run(user_input)`.
///
/// Subsequent runs can use [`Worker<C, Locked>::run()`] directly.
/// To edit state between turns, call [`unlock()`](Worker::unlock) first.
pub async fn run(
self,
user_input: impl Into<String>,
) -> Result<(Worker<C, Locked>, WorkerResult), WorkerError> {
let mut locked = self.lock();
let result = locked.run(user_input).await?;
Ok((locked, result))
}
/// Resume from Paused, consuming self and transitioning to Locked.
///
/// Used after `unlock()` → edit → resume.
pub async fn resume(
self,
) -> Result<(Worker<C, Locked>, WorkerResult), WorkerError> {
let mut locked = self.lock();
let result = locked.resume().await?;
Ok((locked, result))
}
/// Lock and transition to Locked state
///
/// Flushes pending tool factories, then fixes the current system prompt
/// and history as a "committed prefix". After this, only `run()` / `resume()`
/// may append to history, ensuring cache hits.
///
/// Most callers should use [`run()`](Self::run) instead, which calls
/// this internally. Use `lock()` directly only when you need the
/// `Locked` worker back on error (e.g. in a persistence layer).
///
/// # Panics
///
/// Panics if a pending tool factory produces a duplicate name.
pub fn lock(self) -> Worker<C, Locked> {
self.tool_server.flush_pending();
let locked_prefix_len = self.history.len();
Worker {
client: self.client,
@@ -1178,11 +1145,42 @@ impl<C: LlmClient> Worker<C, Mutable> {
}
}
// =============================================================================
// CacheLocked State-Specific Implementation
// =============================================================================
impl<C: LlmClient> Worker<C, CacheLocked> {
impl<C: LlmClient> Worker<C, Locked> {
/// Execute a turn
///
/// Adds a new user message to history and sends a request to the LLM.
/// Automatically loops if there are tool calls.
pub async fn run(
&mut self,
user_input: impl Into<String>,
) -> Result<WorkerResult, WorkerError> {
self.reset_interruption_state();
// Interceptor: on_prompt_submit
let mut user_item = Item::user_message(user_input);
match self.interceptor.on_prompt_submit(&mut user_item).await {
PromptAction::Cancel(reason) => {
self.last_run_interrupted = true;
return self
.finalize_interruption(Err(WorkerError::Aborted(reason)))
.await;
}
PromptAction::Continue => {}
}
self.history.push(user_item);
let result = self.run_turn_loop().await;
self.finalize_interruption(result).await
}
/// Resume execution (from Paused state)
///
/// Resumes turn processing from current state without adding a new user message.
pub async fn resume(&mut self) -> Result<WorkerResult, WorkerError> {
self.reset_interruption_state();
let result = self.run_turn_loop().await;
self.finalize_interruption(result).await
}
/// Get the prefix length at lock time
pub fn locked_prefix_len(&self) -> usize {
self.locked_prefix_len
crates/llm-worker/tests/callback_test.rs
+5 -1
View File
@@ -46,8 +46,9 @@ async fn test_callback_text_block_events() {
});
});
// Mutable::run consumes self, returns (Locked, WorkerResult)
let result = worker.run("Greet me").await;
assert!(result.is_ok(), "Worker should complete: {:?}", result);
assert!(result.is_ok(), "Worker should complete");
let deltas = text_deltas.lock().unwrap();
assert_eq!(deltas.len(), 2);
@@ -91,6 +92,7 @@ async fn test_callback_tool_call_complete() {
});
});
// Mutable::run consumes self, returns (Locked, WorkerResult)
let _ = worker.run("Weather please").await;
let starts = tool_starts.lock().unwrap();
@@ -133,6 +135,7 @@ async fn test_callback_turn_events() {
ends.lock().unwrap().push(turn);
});
// Mutable::run consumes self, returns (Locked, WorkerResult)
let result = worker.run("Do something").await;
assert!(result.is_ok());
@@ -169,6 +172,7 @@ async fn test_callback_usage_events() {
usages.lock().unwrap().push(event.clone());
});
// Mutable::run consumes self, returns (Locked, WorkerResult)
let _ = worker.run("Hello").await;
let usages = usage_events.lock().unwrap();
crates/llm-worker/tests/compile_fail.rs
+1 -1
View File
@@ -1,6 +1,6 @@
#[test]
fn compile_fail_state_constraints() {
let t = trybuild::TestCases::new();
t.compile_fail("tests/ui/cache_locked_register_tool.rs");
t.compile_fail("tests/ui/locked_register_tool.rs");
t.compile_fail("tests/ui/tool_server_handle_register_tool.rs");
}
crates/llm-worker/tests/parallel_execution_test.rs
+10 -7
View File
@@ -102,11 +102,12 @@ async fn test_parallel_tool_execution() {
let tool2_clone = tool2.clone();
let tool3_clone = tool3.clone();
worker.register_tool(tool1.definition()).unwrap();
worker.register_tool(tool2.definition()).unwrap();
worker.register_tool(tool3.definition()).unwrap();
worker.register_tool(tool1.definition());
worker.register_tool(tool2.definition());
worker.register_tool(tool3.definition());
let start = Instant::now();
// Mutable::run consumes self, returns (Locked, WorkerResult)
let _result = worker.run("Run all tools").await;
let elapsed = start.elapsed();
@@ -150,8 +151,8 @@ async fn test_before_tool_call_skip() {
let allowed_clone = allowed_tool.clone();
let blocked_clone = blocked_tool.clone();
worker.register_tool(allowed_tool.definition()).unwrap();
worker.register_tool(blocked_tool.definition()).unwrap();
worker.register_tool(allowed_tool.definition());
worker.register_tool(blocked_tool.definition());
// Policy to skip "blocked_tool"
struct BlockingPolicy;
@@ -169,6 +170,7 @@ async fn test_before_tool_call_skip() {
worker.set_interceptor(BlockingPolicy);
// Mutable::run consumes self, returns (Locked, WorkerResult)
let _result = worker.run("Test hook").await;
// allowed_tool is called, but blocked_tool is not
@@ -230,7 +232,7 @@ async fn test_post_tool_call_modification() {
})
}
worker.register_tool(simple_tool_definition()).unwrap();
worker.register_tool(simple_tool_definition());
// Policy to modify results
struct ModifyingPolicy {
@@ -251,9 +253,10 @@ async fn test_post_tool_call_modification() {
modified_content: modified_content.clone(),
});
// Mutable::run consumes self, returns (Locked, WorkerResult)
let result = worker.run("Test modification").await;
assert!(result.is_ok(), "Worker should complete: {:?}", result);
assert!(result.is_ok(), "Worker should complete");
// Verify hook was called and content was modified
let content = modified_content.lock().unwrap().clone();
crates/llm-worker/tests/ui/cache_locked_register_tool.rs → crates/llm-worker/tests/ui/locked_register_tool.rs
View File
crates/llm-worker/tests/ui/cache_locked_register_tool.stderr → crates/llm-worker/tests/ui/locked_register_tool.stderr
+3 -3
View File
@@ -1,8 +1,8 @@
error[E0599]: no method named `register_tool` found for struct `Worker<OllamaClient, CacheLocked>` in the current scope
--> tests/ui/cache_locked_register_tool.rs:10:20
error[E0599]: no method named `register_tool` found for struct `Worker<OllamaClient, Locked>` in the current scope
--> tests/ui/locked_register_tool.rs:10:20
|
10 | let _ = locked.register_tool(def);
| ^^^^^^^^^^^^^ method not found in `Worker<OllamaClient, CacheLocked>`
| ^^^^^^^^^^^^^ method not found in `Worker<OllamaClient, Locked>`
|
= note: the method was found for
- `Worker<C>`
crates/llm-worker/tests/ui/tool_server_handle_register_tool.stderr
+4 -7
View File
@@ -1,13 +1,10 @@
error[E0624]: method `register_tool` is private
--> tests/ui/tool_server_handle_register_tool.rs:10:20
|
10 | let _ = handle.register_tool(def);
| ^^^^^^^^^^^^^ private method
10 | let _ = handle.register_tool(def);
| ^^^^^^^^^^^^^ private method
|
::: src/tool_server.rs
|
| / pub(crate) fn register_tool(
| | &self,
| | factory: WorkerToolDefinition,
| | ) -> Result<(), ToolServerError> {
| |____________________________________- private method defined here
| pub(crate) fn register_tool(&self, factory: WorkerToolDefinition) {
| ----------------------------------------------------------------- private method defined here
crates/llm-worker/tests/worker_fixtures.rs
+6 -5
View File
@@ -129,9 +129,9 @@ async fn test_worker_simple_text_response() {
}
let client = MockLlmClient::from_fixture(&fixture_path).unwrap();
let mut worker = Worker::new(client);
let worker = Worker::new(client);
// Send a simple message
// Send a simple message (Mutable::run consumes self, returns tuple)
let result = worker.run("Hello").await;
assert!(result.is_ok(), "Worker should complete successfully");
@@ -156,9 +156,9 @@ async fn test_worker_tool_call() {
// Register tool
let weather_tool = MockWeatherTool::new();
let tool_for_check = weather_tool.clone();
worker.register_tool(weather_tool.definition()).unwrap();
worker.register_tool(weather_tool.definition());
// Send message
// Send message (Mutable::run consumes self, returns tuple)
let _result = worker.run("What's the weather in Tokyo?").await;
// Verify tool was called
@@ -190,8 +190,9 @@ async fn test_worker_with_programmatic_events() {
];
let client = MockLlmClient::new(events);
let mut worker = Worker::new(client);
let worker = Worker::new(client);
// Mutable::run consumes self, returns tuple
let result = worker.run("Greet me").await;
assert!(result.is_ok(), "Worker should complete successfully");
crates/llm-worker/tests/worker_state_test.rs
+25 -24
View File
@@ -1,6 +1,6 @@
//! Worker state management tests
//!
//! Tests for state transitions using the Type-state pattern (Mutable/CacheLocked)
//! Tests for state transitions using the Type-state pattern (Mutable/Locked)
//! and state preservation between turns.
mod common;
@@ -11,7 +11,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use common::MockLlmClient;
use llm_worker::Item;
use llm_worker::Worker;
use llm_worker::{Worker, WorkerError};
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta};
@@ -147,15 +147,15 @@ fn test_mutable_can_register_tool() {
let mut worker = Worker::new(client);
let tool = CountingTool::new("count_tool");
let result = worker.register_tool(tool.definition());
assert!(result.is_ok(), "Mutable should allow tool registration");
// register_tool is infallible (factory deferred to run-time flush)
worker.register_tool(tool.definition());
}
// =============================================================================
// State Transition Tests
// =============================================================================
/// Verify that lock() transitions from Mutable -> CacheLocked state
/// Verify that lock() transitions from Mutable -> Locked state
#[test]
fn test_lock_transition() {
let client = MockLlmClient::new(vec![]);
@@ -168,13 +168,13 @@ fn test_lock_transition() {
// Lock
let locked_worker = worker.lock();
// History and system prompt are still accessible in CacheLocked state
// History and system prompt are still accessible in Locked state
assert_eq!(locked_worker.get_system_prompt(), Some("System"));
assert_eq!(locked_worker.history().len(), 2);
assert_eq!(locked_worker.locked_prefix_len(), 2);
}
/// Verify that unlock() transitions from CacheLocked -> Mutable state
/// Verify that unlock() transitions from Locked -> Mutable state
#[test]
fn test_unlock_transition() {
let client = MockLlmClient::new(vec![]);
@@ -198,7 +198,7 @@ fn test_unlock_transition() {
/// Verify that history is correctly updated after running a turn in Mutable state
#[tokio::test]
async fn test_mutable_run_updates_history() {
async fn test_mutable_run_updates_history() -> Result<(), WorkerError> {
let events = vec![
Event::text_block_start(0),
Event::text_delta(0, "Hello, I'm an assistant!"),
@@ -209,11 +209,10 @@ async fn test_mutable_run_updates_history() {
];
let client = MockLlmClient::new(events);
let mut worker = Worker::new(client);
let worker = Worker::new(client);
// Execute
let result = worker.run("Hi there").await;
assert!(result.is_ok());
// Execute (Mutable::run consumes self, returns (Locked, WorkerResult))
let (worker, _result) = worker.run("Hi there").await?;
// History is updated
let history = worker.history();
@@ -224,9 +223,11 @@ async fn test_mutable_run_updates_history() {
// Assistant message
assert_eq!(history[1].as_text(), Some("Hello, I'm an assistant!"));
Ok(())
}
/// Verify that history accumulates correctly over multiple turns in CacheLocked state
/// Verify that history accumulates correctly over multiple turns in Locked state
#[tokio::test]
async fn test_locked_multi_turn_history_accumulation() {
// Prepare responses for 2 requests
@@ -327,7 +328,7 @@ async fn test_locked_prefix_len_tracking() {
/// Verify that turn count is correctly incremented
#[tokio::test]
async fn test_turn_count_increment() {
async fn test_turn_count_increment() -> Result<(), WorkerError> {
let client = MockLlmClient::with_responses(vec![
vec![
Event::text_block_start(0),
@@ -347,15 +348,19 @@ async fn test_turn_count_increment() {
],
]);
let mut worker = Worker::new(client);
let worker = Worker::new(client);
assert_eq!(worker.turn_count(), 0);
worker.run("First").await.unwrap();
// First run consumes Mutable, returns Locked
let (mut worker, _) = worker.run("First").await?;
assert_eq!(worker.turn_count(), 1);
worker.run("Second").await.unwrap();
// Subsequent runs on Locked take &mut self
worker.run("Second").await?;
assert_eq!(worker.turn_count(), 2);
Ok(())
}
/// Verify that history can be edited after unlock and re-locked
@@ -430,9 +435,7 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
let mut worker = Worker::new(client);
let tool_a = CountingTool::new("tool_a");
worker
.register_tool(tool_a.definition())
.expect("register tool_a should succeed");
worker.register_tool(tool_a.definition());
let mut locked = worker.lock();
locked.run("first").await.expect("first run");
@@ -440,9 +443,7 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
let mut unlocked = locked.unlock();
let tool_b = CountingTool::new("tool_b");
unlocked
.register_tool(tool_b.definition())
.expect("register tool_b after unlock should succeed");
unlocked.register_tool(tool_b.definition());
let mut relocked = unlocked.lock();
relocked.run("second").await.expect("second run");
@@ -455,7 +456,7 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
// System Prompt Preservation Tests
// =============================================================================
/// Verify that system prompt is preserved in CacheLocked state
/// Verify that system prompt is preserved in Locked state
#[test]
fn test_system_prompt_preserved_in_locked_state() {
let client = MockLlmClient::new(vec![]);