From d63b4ea4709447c6a4ca47aabb5f6db32b40e2f8 Mon Sep 17 00:00:00 2001
From: Hare <kei.hiracchi.0928@gmail.com>
Date: Thu, 19 Feb 2026 18:29:02 +0900
Subject: [PATCH] restrict tool registration to mutable state

---
 llm-worker/src/tool_server.rs |  7 +++-
 llm-worker/src/worker.rs      | 75 ++++++++++++++++-------------------
 2 files changed, 40 insertions(+), 42 deletions(-)

diff --git a/llm-worker/src/tool_server.rs b/llm-worker/src/tool_server.rs
index 8ebad01..746a063 100644
--- a/llm-worker/src/tool_server.rs
+++ b/llm-worker/src/tool_server.rs
@@ -50,7 +50,10 @@ pub struct ToolServerHandle {
 
 impl ToolServerHandle {
     /// Register one tool.
-    pub fn register_tool(&self, factory: WorkerToolDefinition) -> Result<(), ToolServerError> {
+    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) {
@@ -61,7 +64,7 @@ impl ToolServerHandle {
     }
 
     /// Register many tools.
-    pub fn register_tools(
+    pub(crate) fn register_tools(
         &self,
         factories: impl IntoIterator<Item = WorkerToolDefinition>,
     ) -> Result<(), ToolServerError> {
diff --git a/llm-worker/src/worker.rs b/llm-worker/src/worker.rs
index 7b42087..afc5d02 100644
--- a/llm-worker/src/worker.rs
+++ b/llm-worker/src/worker.rs
@@ -300,46 +300,6 @@ impl<C: LlmClient, S: WorkerState> Worker<C, S> {
             .push(Box::new(SubscriberTurnNotifier { subscriber }));
     }
 
-    /// Register a tool
-    ///
-    /// Registered tools are automatically executed when called by the LLM.
-    /// Registering a tool with the same name will result in an error.
-    ///
-    /// # Examples
-    ///
-    /// ```ignore
-    /// use llm_worker::tool::{ToolMeta, ToolDefinition, Tool};
-    /// use std::sync::Arc;
-    ///
-    /// let def: ToolDefinition = Arc::new(|| {
-    ///     (ToolMeta::new("search").description("..."), Arc::new(MyTool) as Arc<dyn Tool>)
-    /// });
-    /// worker.register_tool(def)?;
-    /// ```
-    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")
-            }
-        }
-    }
-
-    /// Register multiple tools
-    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")
-            }
-        }
-    }
-
     /// Get a shared tool server handle.
     pub fn tool_server_handle(&self) -> ToolServerHandle {
         self.tool_server.clone()
@@ -1059,6 +1019,41 @@ impl<C: LlmClient> Worker<C, Mutable> {
         }
     }
 
+    /// Register a tool
+    ///
+    /// 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")
+            }
+        }
+    }
+
+    /// Register multiple tools
+    ///
+    /// Available only in Mutable state.
+    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")
+            }
+        }
+    }
+
     /// Set system prompt (builder pattern)
     pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
         self.system_prompt = Some(prompt.into());