Adding a New Tool

Relevant source files

The following files were used as context for generating this wiki page:

• tools/shared/base_tool.py
• tools/simple/base.py
• tools/workflow/base.py
• server.py
• tools/docgen.py
• README.md
• systemprompts/docgen_prompt.py

Tools are the fundamental building blocks of functionality within the Zen MCP Server. They encapsulate specific AI-driven tasks, ranging from simple, single-turn interactions like chat to complex, multi-step processes like codereview or docgen. The server is designed with an extensible architecture that allows developers to easily add new tools to expand its capabilities.

This guide provides a comprehensive overview of the tool architecture and a step-by-step process for creating and integrating a new tool into the Zen MCP Server. All tools inherit from a common base class, ensuring a consistent interface for registration and execution.

Tool Architecture

The core of the tool system is the BaseTool abstract class, which defines the contract that all tools must adhere to. The system provides two primary specializations of this base class: SimpleTool for stateless, single-response tasks, and WorkflowTool for stateful, multi-step tasks.

Tool Inheritance Hierarchy

The following diagram illustrates the class inheritance structure for tools.

This diagram shows that SimpleTool and WorkflowTool both inherit from BaseTool, and a concrete tool like DocgenTool implements the WorkflowTool interface.
Sources: tools/shared/base_tool.py, tools/simple/base.py, tools/workflow/base.py, tools/docgen.py

The BaseTool Abstract Class

All tools must inherit from BaseTool and implement its abstract methods. This ensures that the server can discover, register, and execute any tool in a standardized way.

Abstract Method	Description
get_name()	Returns the unique, lowercase name of the tool (e.g., "chat", "docgen").
get_description()	Provides a user-facing description of the tool's purpose and capabilities.
get_input_schema()	Defines the JSON schema for the tool's input parameters, which is exposed to the client.
get_model_category()	Specifies the type of AI model best suited for the tool's task (e.g., CHAT, EXTENDED_REASONING).
get_system_prompt()	Returns the system prompt string that instructs the AI model on its role and behavior for this tool.
execute()	The main entry point for running the tool's logic.

Sources: tools/shared/base_tool.py:27-128

Tool Types

There are two distinct types of tools you can create, each suited for different kinds of tasks.

SimpleTool

SimpleTool is designed for stateless, single-turn interactions. It takes a request, prepares a prompt for an AI model, sends it, and processes the model's response into a final output. It is ideal for tasks that don't require maintaining state across multiple steps.

Key Characteristics:

• Inherits from tools.simple.base.SimpleTool.
• Implements prepare_prompt() to construct the query for the AI model.
• Implements process_response() to format the AI model's output for the user.
• The entire operation is handled within a single execute() call.

Sources: tools/simple/base.py

WorkflowTool

WorkflowTool is designed for complex, multi-step, and stateful processes. It guides an AI agent through a sequence of actions, maintaining context and tracking progress along the way. This is suitable for tasks like code reviews, debugging, or documentation generation, which involve multiple phases of analysis and interaction.

Key Characteristics:

• Inherits from tools.workflow.base.WorkflowTool.
• Defines a Pydantic model for its request structure (e.g., DocgenRequest for the docgen tool).
• Implements execute_workflow() to manage the step-by-step process.
• Uses a continuation_id to maintain state across multiple client calls.
• Often involves a step_number and total_steps to track progress.

An example is the docgen tool, which follows a structured workflow:

1. Discovery: Identify all files needing documentation.
2. Documentation: Process one file at a time, adding docstrings and comments.
3. Verification: Ensure all files are documented before completion.

Sources: tools/workflow/base.py, tools/docgen.py:10-100

Tool Registration and Loading

Tools are automatically discovered and registered by the server at startup. The process involves scanning the tools/ directory, instantiating tool classes, and filtering them based on configuration.

The following diagram shows the tool loading and registration flow.

A critical part of this process is the DISABLED_TOOLS environment variable. This comma-separated string allows operators to disable specific tools to save resources or restrict functionality without modifying the code.

Configuration	Description
DISABLED_TOOLS	A comma-separated list of tool names to disable. For example: analyze,refactor,testgen. If a tool is in this list, it will not be available to clients.

Sources: README.md

Creating a New Tool: A Step-by-Step Guide

Follow these steps to create and integrate a new tool.

1. Choose a Tool Type

First, decide whether your tool is a simple, single-shot action (SimpleTool) or a complex, multi-step process (WorkflowTool).

2. Create the Tool Class

Create a new Python file in the tools/ directory (e.g., tools/my_new_tool.py). In this file, define a class that inherits from your chosen base class.

# tools/my_new_tool.py
from tools.simple.base import SimpleTool
from tools.models import ToolModelCategory

class MyNewTool(SimpleTool):
    def get_name(self) -> str:
        return "my_new_tool"

    def get_description(self) -> str:
        return "This is a description for my new tool."

    def get_model_category(self) -> ToolModelCategory:
        return ToolModelCategory.CHAT

    # ... implement other required methods

Sources: tools/simple/base.py, tools/shared/base_tool.py

3. Define the Input Schema

The get_input_schema() method is crucial. It defines the parameters your tool accepts from the client. The server uses this to generate a JSON schema. For WorkflowTool, you typically define a Pydantic model to represent the request structure.

For example, the docgen tool defines a DocgenRequest class with fields like step, step_number, and total_files_to_document.

# From tools/docgen.py
class DocgenRequest(WorkflowRequest):
    """Request model for documentation generation steps"""
    step: str = Field(..., description="...")
    step_number: int = Field(..., description="...")
    total_steps: int = Field(..., description="...")
    # ... other fields

Sources: tools/docgen.py:53-76, tools/shared/base_tool.py:41-45

4. Implement the Core Logic

Place your tool's main functionality inside the execute() method.

• For a SimpleTool, this involves calling self.prepare_prompt(), invoking the model, and then self.process_response().
• For a WorkflowTool, this involves calling self.execute_workflow(), which orchestrates the multi-step logic.

5. Create the System Prompt

Each tool requires a system prompt that guides the AI model's behavior. Create a corresponding Python file in the systemprompts/ directory (e.g., systemprompts/my_new_tool_prompt.py). The get_system_prompt() method in your tool class should return the content of this prompt.

For instance, DocgenTool uses DOCGEN_PROMPT from systemprompts/docgen_prompt.py.
Sources: tools/docgen.py:108-110, systemprompts/docgen_prompt.py

6. Enable the Tool

Your tool will be automatically discovered at startup. To enable it, ensure its name is not in the DISABLED_TOOLS environment variable in your .env file. If DISABLED_TOOLS is not set or is empty, all tools will be enabled by default.
Sources: README.md, server.py:101-118

Conclusion

The Zen MCP Server's tool architecture is designed for modularity and extensibility. By adhering to the BaseTool interface and choosing between the SimpleTool and WorkflowTool patterns, developers can seamlessly integrate new, powerful AI-driven capabilities. The automatic discovery and configuration-based enabling/disabling of tools provide a flexible system for managing the server's feature set.