fix: correct task scheduling and simplify agent internals

Fix queue cancellation, concurrency limits and worker status updates.

Use direct Explorer imports, split the Diagnoser loop into helpers and remove unused tool registration and prompt-builder code.

Correct model caching, endpoint selection and retry handling.

artemis/tools/base.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Universal Tool Protocol and Registry for ARTEMIS.
+"""Universal Tool Protocol for ARTEMIS.
 
 Enables tools to be authored once using standard Pydantic schemas and async handlers,
 then automatically exported to LangChain Tools (Pro Graph), Google GenAI FunctionDeclarations
@@ -197,91 +197,3 @@ class ArtemisTool:
                 required=required,
             ),
         )
-
-
-class ToolRegistry:
-    """Global registry managing universal tools."""
-
-    _tools: dict[str, ArtemisTool] = {}
-
-    @classmethod
-    def register(cls, tool: ArtemisTool) -> ArtemisTool:
-        """Registers a tool in the global registry."""
-        cls._tools[tool.name] = tool
-        return tool
-
-    @classmethod
-    def get(cls, name: str) -> ArtemisTool | None:
-        """Retrieves a tool by name from the registry."""
-        normalized = name.split(":")[-1] if ":" in name else name
-        return cls._tools.get(normalized)
-
-    @classmethod
-    def list_tools(
-        cls, category: ToolCategory | None = None, available_only: bool = True
-    ) -> list[ArtemisTool]:
-        """Lists all registered tools, optionally filtered by category and availability."""
-        tools = [t for t in cls._tools.values() if not available_only or t.is_available()]
-        if category:
-            return [t for t in tools if t.category == category]
-        return list(tools)
-
-    @classmethod
-    def get_langchain_tools(
-        cls, ctx: Any, names: list[str] | None = None, available_only: bool = True
-    ) -> list[BaseTool]:
-        """Returns LangChain tool wrappers for the requested or all registered tools."""
-        tools = (
-            [cls._tools[n] for n in names if n in cls._tools]
-            if names
-            else list(cls._tools.values())
-        )
-        if available_only:
-            tools = [t for t in tools if t.is_available(ctx)]
-        return [t.to_langchain_tool(ctx) for t in tools]
-
-    @classmethod
-    def get_genai_declarations(
-        cls, names: list[str] | None = None, available_only: bool = True
-    ) -> list[genai_types.FunctionDeclaration]:
-        """Returns Google GenAI function declarations for registered tools."""
-        tools = (
-            [cls._tools[n] for n in names if n in cls._tools]
-            if names
-            else list(cls._tools.values())
-        )
-        if available_only:
-            tools = [t for t in tools if t.is_available()]
-        return [t.to_genai_declaration() for t in tools]
-
-    @classmethod
-    async def execute(
-        cls, name: str, args: dict[str, Any], driver: BaseDeviceDriver, ctx: Any
-    ) -> Any:
-        """Executes a registered tool by name with provided arguments, driver, and context."""
-        tool = cls.get(name)
-        if not tool:
-            raise ValueError(f"Tool '{name}' not found in ToolRegistry.")
-        return await tool.execute(driver=driver, ctx=ctx, **args)
-
-
-def artemis_tool(
-    name: str,
-    description: str,
-    args_schema: type[BaseModel],
-    category: ToolCategory = "action",
-):
-    """Decorator to define and register a universal ARTEMIS tool."""
-
-    def decorator(fn: Callable[..., Any]) -> ArtemisTool:
-        tool_obj = ArtemisTool(
-            name=name,
-            description=description,
-            args_schema=args_schema,
-            handler=fn,
-            category=category,
-        )
-        ToolRegistry.register(tool_obj)
-        return tool_obj
-
-    return decorator

artemis/tools/command_tool.py

@@ -54,7 +54,7 @@ from artemis.controllers.platform_specific_commands_controller import (
 )
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
-from artemis.tools.base import ArtemisTool, ToolRegistry
+from artemis.tools.base import ArtemisTool
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.utils.logger import get_logger
 
@@ -832,7 +832,6 @@ class RunAdbCommandTool(ArtemisTool):
 # Universal tool instance & aliases
 run_adb_command = RunAdbCommandTool()
 RunAdbCommand = RunAdbCommandTool
-ToolRegistry.register(run_adb_command)
 
 
 def get_run_adb_command_tool(ctx: ArtemisContext) -> BaseTool:
@@ -978,7 +977,6 @@ class ManageTaskTool(ArtemisTool):
 # Universal tool instance & aliases
 manage_task = ManageTaskTool()
 ManageTask = ManageTaskTool
-ToolRegistry.register(manage_task)
 
 
 def get_manage_task_tool(ctx: ArtemisContext) -> BaseTool:
@@ -1111,7 +1109,6 @@ class RunShortAdbCommandTool(ArtemisTool):
 # Universal tool instance & aliases
 run_short_adb_command = RunShortAdbCommandTool()
 RunShortAdbCommand = RunShortAdbCommandTool
-ToolRegistry.register(run_short_adb_command)
 
 
 def get_run_short_adb_command_tool(ctx: ArtemisContext) -> BaseTool:
@@ -1198,7 +1195,6 @@ class AnalyzeTaskOutputTool(ArtemisTool):
 # Universal tool instance & aliases
 analyze_task_output = AnalyzeTaskOutputTool()
 AnalyzeTaskOutput = AnalyzeTaskOutputTool
-ToolRegistry.register(analyze_task_output)
 
 
 def get_analyze_task_output_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/committee_tool.py

@@ -36,7 +36,7 @@ from artemis.data_engine.trace import (
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
 from artemis.services.llm import get_llm, invoke_llm_with_timeout_message
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.log_tool import get_analyze_logs_tool
 from artemis.tools.scratchpad import get_list_notes_tool, get_read_note_tool
 from artemis.tools.tool_wrapper import (
@@ -108,7 +108,6 @@ class AskCommitteeTool(ArtemisTool):
 ask_committee = AskCommitteeTool()
 AskCommittee = AskCommitteeTool
 AskCommitteeToolAlias = AskCommitteeTool
-ToolRegistry.register(ask_committee)
 
 
 def get_ask_committee_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/diagnoser_submit_answer_tool.py

@@ -22,7 +22,7 @@ from pydantic import BaseModel, Field
 from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
-from artemis.tools.base import ArtemisTool, ToolRegistry
+from artemis.tools.base import ArtemisTool
 from artemis.utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -80,7 +80,6 @@ class SubmitAnswerTool(ArtemisTool):
 submit_answer = SubmitAnswerTool()
 SubmitAnswer = SubmitAnswerTool
 DiagnoserSubmitAnswerTool = SubmitAnswerTool
-ToolRegistry.register(submit_answer)
 
 
 def get_submit_answer_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/diagnostic_tool.py

@@ -24,7 +24,7 @@ from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolRegistry
+from artemis.tools.base import ArtemisTool
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.logger import get_logger
@@ -100,7 +100,6 @@ class AskDiagnoserTool(ArtemisTool):
 ask_diagnoser = AskDiagnoserTool()
 AskDiagnoser = AskDiagnoserTool
 DiagnosticTool = AskDiagnoserTool
-ToolRegistry.register(ask_diagnoser)
 
 
 def get_ask_diagnostic_agent_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/explorer_tool.py

@@ -56,7 +56,7 @@ from artemis.config import ExplorerVersion, resolve_explorer_version, settings
 from artemis.context import ArtemisContext
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.element_hit_test import find_element_at_point
@@ -527,7 +527,6 @@ class AskExplorerTool(ArtemisTool):
 
 
 ask_explorer = AskExplorerTool()
-ToolRegistry.register(ask_explorer)
 
 
 def get_ask_explorer_tool(

artemis/tools/history_recall.py

@@ -45,7 +45,7 @@ from pydantic import BaseModel, Field
 
 from artemis.core.tool_declaration import ToolDeclaration
 from artemis.data_engine.trace import trace_langchain_tool
-from artemis.tools.base import ArtemisTool, ToolRegistry
+from artemis.tools.base import ArtemisTool
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.utils.logger import get_logger
 
@@ -547,7 +547,6 @@ def _recall_config(warn: bool = True) -> Any:
 
 
 recall_history = RecallHistoryTool()
-ToolRegistry.register(recall_history)
 
 _DECLARATION_TYPES = {
     "query": {"type": "string"},

artemis/tools/image_processor_tool.py

@@ -23,7 +23,7 @@ from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace, trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolRegistry
+from artemis.tools.base import ArtemisTool
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.logger import get_logger
@@ -90,7 +90,6 @@ ask_image_processor = AskImageProcessorTool()
 AskImageProcessor = AskImageProcessorTool
 AskVisionCoderTool = AskImageProcessorTool
 ImageProcessorTool = AskImageProcessorTool
-ToolRegistry.register(ask_image_processor)
 
 
 def get_ask_image_processor_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/log_tool.py

@@ -24,7 +24,7 @@ from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolRegistry
+from artemis.tools.base import ArtemisTool
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.logger import get_logger
@@ -93,7 +93,6 @@ class AnalyzeLogsTool(ArtemisTool):
 analyze_logs = AnalyzeLogsTool()
 AnalyzeLogs = AnalyzeLogsTool
 LogTool = AnalyzeLogsTool
-ToolRegistry.register(analyze_logs)
 
 
 def get_analyze_logs_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/memory/notes.py

@@ -1,66 +0,0 @@
-# Copyright 2026 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Memory and Scratchpad Note Management tools."""
-
-from typing import Any
-
-from pydantic import BaseModel, Field
-
-from artemis.tools.base import artemis_tool
-
-_IN_MEMORY_NOTES: dict[str, str] = {}
-
-
-class ReadNoteArgs(BaseModel):
-    """Arguments schema for reading a note."""
-
-    title: str = Field(..., description="Note title to read.")
-
-
-class SaveNoteArgs(BaseModel):
-    """Arguments schema for saving a note."""
-
-    title: str = Field(..., description="Note title to create/save.")
-    content: str = Field(..., description="Content text of the note.")
-
-
-@artemis_tool(
-    name="read_note",
-    description="[MEMORY] Reads content of an existing scratchpad note.",
-    args_schema=ReadNoteArgs,
-    category="memory",
-)
-async def read_note(
-    title: str,
-    ctx: Any = None,  # pylint: disable=unused-argument
-) -> str:
-    """Reads content of an existing scratchpad note."""
-    return _IN_MEMORY_NOTES.get(title, f"Note '{title}' not found.")
-
-
-@artemis_tool(
-    name="save_note",
-    description="[MEMORY] Saves content into the scratchpad notes system.",
-    args_schema=SaveNoteArgs,
-    category="memory",
-)
-async def save_note(
-    title: str,
-    content: str,
-    ctx: Any = None,  # pylint: disable=unused-argument
-) -> str:
-    """Saves content into the scratchpad notes system."""
-    _IN_MEMORY_NOTES[title] = content
-    return f"Note '{title}' saved successfully."

artemis/tools/mobile/launch_app.py

@@ -26,7 +26,7 @@ from artemis.controllers.platform_specific_commands_controller import (
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.app_launch_utils import launch_app_with_retries
@@ -189,7 +189,6 @@ class LaunchAppTool(ArtemisTool):
 # Universal tool instance & aliases
 launch_app = LaunchAppTool()
 LaunchApp = LaunchAppTool
-ToolRegistry.register(launch_app)
 
 
 def get_launch_app_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/mobile/ocr.py

@@ -25,7 +25,7 @@ from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.utils.logger import get_logger
 from artemis.utils.ocr_api import is_ocr_configured, perform_ocr
@@ -187,7 +187,6 @@ ocr_recognition = OcrRecognitionTool()
 OcrRecognition = OcrRecognitionTool
 OCRRecognition = OcrRecognitionTool
 OcrTool = OcrRecognitionTool
-ToolRegistry.register(ocr_recognition)
 
 
 def get_ocr_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/mobile/read_hierarchy.py

@@ -22,7 +22,7 @@ from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.logger import get_logger
@@ -91,7 +91,6 @@ class GetUiHierarchyTool(ArtemisTool):
 get_ui_hierarchy = GetUiHierarchyTool()
 GetUiHierarchy = GetUiHierarchyTool
 GetUIHierarchy = GetUiHierarchyTool
-ToolRegistry.register(get_ui_hierarchy)
 
 
 def get_ui_hierarchy_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/mobile/read_logs.py

@@ -21,7 +21,7 @@ from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.mobile.log_utils import fetch_and_filter_logs
 from artemis.utils.logger import get_logger
 
@@ -108,7 +108,6 @@ class ReadLogsTool(ArtemisTool):
 read_logs = ReadLogsTool()
 ReadLogs = ReadLogsTool
 ReadLogsToolAlias = ReadLogsTool
-ToolRegistry.register(read_logs)
 
 
 def get_read_logs_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/mobile/search_logs.py

@@ -22,7 +22,7 @@ from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.mobile.log_utils import fetch_and_filter_logs
 from artemis.utils.logger import get_logger
 
@@ -229,7 +229,6 @@ class SearchLogsTool(ArtemisTool):
 search_logs = SearchLogsTool()
 SearchLogs = SearchLogsTool
 SearchLogsToolAlias = SearchLogsTool
-ToolRegistry.register(search_logs)
 
 
 def get_search_logs_tool(ctx: ArtemisContext) -> BaseTool:

artemis/tools/object_detection_tool.py

@@ -33,7 +33,7 @@ from artemis.data_engine.trace import (
 )
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.logger import get_logger
@@ -220,8 +220,6 @@ operator_object_detection = OperatorObjectDetectionTool()
 OperatorObjectDetection = OperatorObjectDetectionTool
 OperatorObjectDetectorTool = OperatorObjectDetectionTool
 
-ToolRegistry.register(object_detection)
-
 
 def get_object_detector_tool(ctx: ArtemisContext) -> BaseTool:
     """Exports object_detection as a LangChain BaseTool."""

artemis/tools/perception/object_detect.py

@@ -1,43 +0,0 @@
-# Copyright 2026 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Visual Object and Icon Detection perception tool."""
-
-from pydantic import BaseModel, Field
-from artemis.tools.base import artemis_tool
-from artemis.drivers.base import BaseDeviceDriver
-
-
-class ObjectDetectArgs(BaseModel):
-    """Arguments schema for object detection."""
-
-    icon_description: str = Field(
-        ..., description="Description of the icon, button, or graphic element."
-    )
-
-
-@artemis_tool(
-    name="detect_object",
-    description=(
-        "[PERCEPTION] Detects visual icon or UI element on screen using multimodal localization."
-    ),
-    args_schema=ObjectDetectArgs,
-    category="perception",
-)
-async def detect_object(
-    icon_description: str,
-    driver: BaseDeviceDriver | None = None,  # pylint: disable=unused-argument
-) -> str:
-    """Detects visual icon or UI element on screen using multimodal localization."""
-    return f"Object Detect: Found '{icon_description}' at normalized [500, 500]"

artemis/tools/perception/ocr.py

@@ -1,46 +0,0 @@
-# Copyright 2026 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Optical Character Recognition (OCR) perception tool."""
-
-from pydantic import BaseModel, Field
-from artemis.tools.base import artemis_tool
-from artemis.drivers.base import BaseDeviceDriver
-
-
-class OCRRecognitionArgs(BaseModel):
-    """Arguments schema for OCR recognition."""
-
-    query_text: str | None = Field(
-        default=None,
-        description="Optional target text snippet to search on screen.",
-    )
-
-
-@artemis_tool(
-    name="ocr_recognize",
-    description=(
-        "[PERCEPTION] Performs OCR on the active screen to extract visible text and coordinates."
-    ),
-    args_schema=OCRRecognitionArgs,
-    category="perception",
-)
-async def ocr_recognize(
-    query_text: str | None = None,
-    driver: BaseDeviceDriver | None = None,  # pylint: disable=unused-argument
-) -> str:
-    """Performs OCR on the active screen to extract visible text and coordinates."""
-    if query_text:
-        return f"OCR: Text '{query_text}' located at center [500, 500]"
-    return "OCR: Extracted visible text items from screen."

artemis/tools/perception/xml_search.py

@@ -1,45 +0,0 @@
-# Copyright 2026 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""XML Layout Search perception tool."""
-
-from pydantic import BaseModel, Field
-from artemis.tools.base import artemis_tool
-from artemis.drivers.base import BaseDeviceDriver
-
-
-class XMLSearchArgs(BaseModel):
-    """Arguments schema for XML layout search."""
-
-    query: str = Field(
-        ...,
-        description="Text, resource ID, or class pattern to search in UI hierarchy.",
-    )
-
-
-@artemis_tool(
-    name="xml_search",
-    description=(
-        "[PERCEPTION] Searches Android UI layout hierarchy for elements"
-        " matching text or resource ID."
-    ),
-    args_schema=XMLSearchArgs,
-    category="perception",
-)
-async def xml_search(
-    query: str,
-    driver: BaseDeviceDriver | None = None,  # pylint: disable=unused-argument
-) -> str:
-    """Searches Android UI layout hierarchy for elements matching text or resource ID."""
-    return f"XML Search: Found element matching '{query}'"

artemis/tools/scratchpad.py

@@ -22,7 +22,7 @@ from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolRegistry
+from artemis.tools.base import ArtemisTool
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.logger import get_logger
@@ -166,7 +166,6 @@ class SaveNoteTool(ArtemisTool):
 # Universal tool instance & aliases
 save_note = SaveNoteTool()
 SaveNote = SaveNoteTool
-ToolRegistry.register(save_note)
 
 
 def get_save_note_tool(ctx: ArtemisContext) -> BaseTool:
@@ -246,7 +245,6 @@ class AppendNoteTool(ArtemisTool):
 # Universal tool instance & aliases
 append_note = AppendNoteTool()
 AppendNote = AppendNoteTool
-ToolRegistry.register(append_note)
 
 
 def get_append_note_tool(ctx: ArtemisContext) -> BaseTool:
@@ -331,7 +329,6 @@ class UpdateNoteTool(ArtemisTool):
 # Universal tool instance & aliases
 update_note = UpdateNoteTool()
 UpdateNote = UpdateNoteTool
-ToolRegistry.register(update_note)
 
 
 def get_update_note_tool(ctx: ArtemisContext) -> BaseTool:
@@ -416,7 +413,6 @@ class ReadNoteTool(ArtemisTool):
 # Universal tool instance & aliases
 read_note = ReadNoteTool()
 ReadNote = ReadNoteTool
-ToolRegistry.register(read_note)
 
 
 def get_read_note_tool(ctx: ArtemisContext) -> BaseTool:
@@ -478,7 +474,6 @@ class ListNotesTool(ArtemisTool):
 # Universal tool instance & aliases
 list_notes = ListNotesTool()
 ListNotes = ListNotesTool
-ToolRegistry.register(list_notes)
 
 
 def get_list_notes_tool(ctx: ArtemisContext) -> BaseTool:
@@ -545,7 +540,6 @@ class ReadNotePureTool(ArtemisTool):
 # Universal pure tool instance & aliases
 read_note_pure = ReadNotePureTool()
 ReadNotePure = ReadNotePureTool
-ToolRegistry.register(read_note_pure)
 
 
 def get_read_note_tool_pure(ctx: ArtemisContext) -> BaseTool:
@@ -591,7 +585,6 @@ class ListNotesPureTool(ArtemisTool):
 # Universal pure tool instance & aliases
 list_notes_pure = ListNotesPureTool()
 ListNotesPure = ListNotesPureTool
-ToolRegistry.register(list_notes_pure)
 
 
 def get_list_notes_tool_pure(ctx: ArtemisContext) -> BaseTool:
@@ -656,7 +649,6 @@ class SaveNotePureTool(ArtemisTool):
 # Universal pure tool instance & aliases
 save_note_pure = SaveNotePureTool()
 SaveNotePure = SaveNotePureTool
-ToolRegistry.register(save_note_pure)
 
 
 def get_save_note_tool_pure(ctx: ArtemisContext) -> BaseTool:
@@ -725,7 +717,6 @@ class UpdateNotePureTool(ArtemisTool):
 # Universal pure tool instance & aliases
 update_note_pure = UpdateNotePureTool()
 UpdateNotePure = UpdateNotePureTool
-ToolRegistry.register(update_note_pure)
 
 
 def get_update_note_tool_pure(ctx: ArtemisContext) -> BaseTool:
@@ -790,7 +781,6 @@ class AppendNotePureTool(ArtemisTool):
 # Universal pure tool instance & aliases
 append_note_pure = AppendNotePureTool()
 AppendNotePure = AppendNotePureTool
-ToolRegistry.register(append_note_pure)
 
 
 def get_append_note_tool_pure(ctx: ArtemisContext) -> BaseTool:

artemis/tools/tool_wrapper.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 from collections.abc import Callable
+import functools
 import inspect
 import time
 from typing import Annotated, Any, get_args, get_origin
-from unittest.mock import MagicMock, Mock
 import uuid
 
 from langchain_core.messages import ToolMessage
@@ -64,27 +64,19 @@ async def invoke_tool_with_injection(
     """
     final_args = dict(args)
 
-    if isinstance(tool, (Mock, MagicMock)):
-        has_real_coro = hasattr(tool, "coroutine") and not isinstance(
-            tool.coroutine, (Mock, MagicMock)
-        )
-        has_real_func = hasattr(tool, "func") and not isinstance(tool.func, (Mock, MagicMock))
-        if not has_real_coro and not has_real_func:
-            if (
-                state is not None
-                and hasattr(tool, "args")
-                and isinstance(tool.args, dict)
-                and "state" in tool.args
-            ):
-                final_args["state"] = state
-            return await tool.ainvoke(final_args)
+    def _real_callable(candidate: Any) -> Callable | None:
+        """Return a routine or partial whose signature supports direct injection."""
+        target = candidate.func if isinstance(candidate, functools.partial) else candidate
+        if inspect.isroutine(target):
+            return candidate
+        return None
 
     # For StructuredTool, the original function might be in tool.func or tool.coroutine
     func = None
-    if hasattr(tool, "coroutine") and tool.coroutine:
-        func = tool.coroutine
-    elif hasattr(tool, "func") and tool.func:
-        func = tool.func
+    if getattr(tool, "coroutine", None):
+        func = _real_callable(tool.coroutine)
+    if func is None and getattr(tool, "func", None):
+        func = _real_callable(tool.func)
 
     if func:
         try:
@@ -163,7 +155,11 @@ async def invoke_tool_with_injection(
                 )
             CURRENT_TRACE_ID.reset(token)
     else:
-        # Fallback to ainvoke if no direct func available
+        # Fallback to ainvoke if no direct func available; tools that declare
+        # a "state" argument still receive the graph state.
+        tool_args = getattr(tool, "args", None)
+        if state is not None and isinstance(tool_args, dict) and "state" in tool_args:
+            final_args["state"] = state
         return await tool.ainvoke(final_args)
 
 

artemis/tools/video_tool.py

@@ -31,7 +31,7 @@ from artemis.core.tool_declaration import ToolDeclaration
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
 from artemis.graph.state import State
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.logger import get_logger
@@ -180,7 +180,6 @@ class VideoAnalyzerTool(ArtemisTool):
 
 # Universal tool instance & aliases
 video_analyzer = VideoAnalyzerTool()
-ToolRegistry.register(video_analyzer)
 
 
 class VideoAnalyzerPureTool(ArtemisTool):
@@ -231,7 +230,6 @@ class VideoAnalyzerPureTool(ArtemisTool):
 # Universal pure tool instance & aliases
 video_analyzer_pure = VideoAnalyzerPureTool()
 VideoAnalyzerPure = VideoAnalyzerPureTool
-ToolRegistry.register(video_analyzer_pure)
 
 
 def get_video_analyzer_tool(

artemis/tools/wait_tool.py

@@ -21,7 +21,7 @@ from pydantic import BaseModel, Field
 from artemis.context import ArtemisContext
 from artemis.data_engine.trace import trace_langchain_tool
 from artemis.drivers.base import BaseDeviceDriver
-from artemis.tools.base import ArtemisTool, ToolCategory, ToolRegistry
+from artemis.tools.base import ArtemisTool, ToolCategory
 from artemis.tools.tool_wrapper import ToolWrapper
 from artemis.tools.types import CyFunctionDetector
 from artemis.utils.logger import get_logger
@@ -82,7 +82,6 @@ class WaitTool(ArtemisTool):
 # Universal tool instance & aliases
 wait = WaitTool()
 Wait = WaitTool
-ToolRegistry.register(wait)
 
 
 def get_wait_tool(ctx: ArtemisContext) -> BaseTool: