DEV Community

Cover image for Code Agent Anatomy (20): Extending from Scratch — Adding a New Tool to the Agent
WonderLab
WonderLab

Posted on

Code Agent Anatomy (20): Extending from Scratch — Adding a New Tool to the Agent

Starting from a Leftover Question in Part 4

The first four Parts have been "anatomy": reading code, understanding design, studying principles. This article starts turning toward "hands-on" — using MyCodeAgent as a starting point to extend your own things.

First, let's answer a question: what exactly is an agent "tool"?

From the model's perspective, a tool is a function definition in Function Calling: a name, a description, and a set of parameters. The model selects this tool, the framework executes it, and the result is stuffed back into the conversation history as an observation.

From the framework's perspective, a tool is a Python class implementing a specific interface: it has parameter definitions, a run() method, and run() returns a result object in a standard format.

Understanding these two perspectives makes adding a new tool a well-defined process.


The Conclusion First

Adding a new tool to MyCodeAgent requires four steps:

Step What to Do Files Involved
1. Inherit the Tool base class Define parameters, implement run() tools/builtin/your_tool.py
2. Register to Registry Let the framework "know" this tool exists runtime/host.py or app/bootstrap.py
3. Write a Prompt Tell the model when to use it and how prompts/tools_prompts/your_tool_prompt.py
4. Write tests Verify protocol compliance, verify logic correctness tests/test_your_tool.py

I. The Tool Base Class: Everything Starts Here

Open tools/base.py and you'll see the foundational structure of the entire tool system.

# tools/base.py
class Tool(ABC):
    def __init__(self, name, description, project_root=None, working_dir=None):
        self.name = name
        self.description = description
        self._project_root = Path(project_root).resolve() if project_root else None
        self._working_dir = ...

    @abstractmethod
    def run(self, parameters: Dict[str, Any]) -> ToolResult:
        pass

    @abstractmethod
    def get_parameters(self) -> List[ToolParameter]:
        pass
Enter fullscreen mode Exit fullscreen mode

Two abstract methods must be implemented:

  • get_parameters(): tells the framework what parameters this tool accepts
  • run(): the actual logic of the tool, returns a ToolResult

ToolResult is also in base.py. It's an immutable dataclass encapsulating a standard response envelope:

@dataclass(frozen=True)
class ToolResult:
    status: ToolStatus      # success / partial / error
    text: str               # text summary for the model
    data: Dict[str, Any]    # core payload
    error_code: ...         # only has value on error
    stats: Dict[str, Any]   # timing and other stats
    context: Dict[str, Any] # cwd, params_input, etc.
Enter fullscreen mode Exit fullscreen mode

Note frozen=True: ToolResult cannot be modified after creation, preventing accidental changes in the tool execution pipeline.


II. Hands-on: Writing a WordCount Tool

Let's walk through the complete path with a concrete example. The tool we'll write: count the number of lines, words, and characters in a file.

Step One: Implement the Tool Class

# tools/builtin/word_count.py
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from ..base import Tool, ToolParameter, ToolResult, ErrorCode


class WordCountTool(Tool):
    """Count lines, words, and characters in a file."""

    def __init__(
        self,
        name: str = "WordCount",
        project_root: Optional[Path] = None,
        working_dir: Optional[Path] = None,
    ):
        if project_root is None:
            raise ValueError("project_root must be provided by the framework")
        super().__init__(
            name=name,
            description="Count lines, words, and characters in a file.",
            project_root=project_root,
            working_dir=working_dir or project_root,
        )

    def get_parameters(self) -> List[ToolParameter]:
        return [
            ToolParameter(
                name="path",
                type="string",
                description="Path to the file (relative to project root).",
                required=True,
            ),
        ]

    def run(self, parameters: Dict[str, Any]) -> ToolResult:
        start_time = time.monotonic()
        params_input = dict(parameters)
        path_str = parameters.get("path")

        # Parameter validation
        if not path_str:
            return self.error_result(
                error_code=ErrorCode.INVALID_PARAM,
                message="Parameter 'path' is required.",
                params_input=params_input,
            )

        # Sandbox: ensure path is within project_root
        target = (self._project_root / path_str).resolve()
        try:
            target.relative_to(self._project_root)
        except ValueError:
            return self.error_result(
                error_code=ErrorCode.ACCESS_DENIED,
                message=f"Path '{path_str}' is outside project root.",
                params_input=params_input,
            )

        if not target.exists():
            return self.error_result(
                error_code=ErrorCode.NOT_FOUND,
                message=f"File '{path_str}' does not exist.",
                params_input=params_input,
            )
        if target.is_dir():
            return self.error_result(
                error_code=ErrorCode.IS_DIRECTORY,
                message=f"Path '{path_str}' is a directory, not a file.",
                params_input=params_input,
            )

        # Core logic
        content = target.read_text(encoding="utf-8", errors="replace")
        line_count = len(content.splitlines())
        word_count = len(content.split())
        char_count = len(content)
        elapsed_ms = int((time.monotonic() - start_time) * 1000)
        rel_path = str(target.relative_to(self._project_root))

        return self.success_result(
            data={
                "lines": line_count,
                "words": word_count,
                "characters": char_count,
            },
            text=(
                f"'{rel_path}': {line_count} lines, "
                f"{word_count} words, {char_count} characters."
            ),
            params_input=params_input,
            time_ms=elapsed_ms,
            path_resolved=rel_path,
        )
Enter fullscreen mode Exit fullscreen mode

A few notable details:

Sandbox check: target.relative_to(self._project_root) throws ValueError if the path escapes the project root directory. This line is a necessary protection for every tool that involves the filesystem.

Parameter validation first: Validate parameters before doing any IO. This way, when the model passes bad parameters, it immediately gets a clear error message instead of a mysterious exception from the IO layer.

success_result() helper methods: The base class already provides success_result(), partial_result(), and error_result() helpers — no need to manually construct ToolResult. They automatically assemble fixed fields like stats.time_ms and context.cwd.


III. Registration: Making the Framework "See" This Tool

The tool class is written, but the framework still doesn't know it exists. Add it to the tool registration area in runtime/host.py:

# runtime/host.py — add these two lines in the built-in tool registration area
from tools.builtin.word_count import WordCountTool

# Inside _build_tool_registry() or __init__:
registry.register_tool(WordCountTool(
    project_root=self._project_root,
    working_dir=self._working_dir,
))
Enter fullscreen mode Exit fullscreen mode

After registration, the tool will appear in the list returned by registry.get_openai_tools(), and the model will see this tool's schema on the next request.


IV. Prompt: Telling the Model When and How to Use It

The tool can be executed, but the model doesn't necessarily know when to use it. Write a Prompt file:

# prompts/tools_prompts/word_count_prompt.py
word_count_prompt = """Count lines, words, and characters in a file.

Use this tool when you need to:
- Know the size of a file before deciding whether to read it in full
- Get a quick overview of a file's content volume

Parameters:
- path (required): Relative path to the file

Returns:
- lines: Number of lines
- words: Number of words  
- characters: Number of characters

Example:
  WordCount(path="src/main.py")
  → "src/main.py: 312 lines, 1847 words, 14203 characters."
"""
Enter fullscreen mode Exit fullscreen mode

Then reference it in the tool class's description parameter:

from prompts.tools_prompts.word_count_prompt import word_count_prompt

super().__init__(
    name=name,
    description=word_count_prompt,  # this description goes into the Function Calling schema
    ...
)
Enter fullscreen mode Exit fullscreen mode

This description is the sole basis on which the model decides "should I use this tool." Write it clearly, and the model will select it at the right moment; write it vaguely, and the model either won't use it or will use it in the wrong context.


V. Tests: Verifying Protocol Compliance

A new tool needs at least two categories of tests:

# tests/test_word_count_tool.py
from pathlib import Path
import pytest
from tools.builtin.word_count import WordCountTool
from tools.base import ToolStatus, ErrorCode


@pytest.fixture
def tool(tmp_path):
    return WordCountTool(project_root=tmp_path)


def test_success(tool, tmp_path):
    (tmp_path / "hello.txt").write_text("hello world\nfoo bar baz\n")
    result = tool.run({"path": "hello.txt"})
    assert result.status == ToolStatus.SUCCESS
    assert result.data["lines"] == 2
    assert result.data["words"] == 5
    assert "stats" in result.__dataclass_fields__
    assert result.stats["time_ms"] >= 0


def test_not_found(tool):
    result = tool.run({"path": "nonexistent.txt"})
    assert result.status == ToolStatus.ERROR
    assert result.error_code == ErrorCode.NOT_FOUND


def test_sandbox_escape(tool):
    result = tool.run({"path": "../../../etc/passwd"})
    assert result.status == ToolStatus.ERROR
    assert result.error_code == ErrorCode.ACCESS_DENIED


def test_missing_param(tool):
    result = tool.run({})
    assert result.status == ToolStatus.ERROR
    assert result.error_code == ErrorCode.INVALID_PARAM
Enter fullscreen mode Exit fullscreen mode

The sandbox escape test (../../../etc/passwd) is a required test. A tool that can be used by the model to read files outside the project is a security vulnerability.


Design Highlights

1. Framework injects, tools don't guess paths

project_root is passed in by the framework at registration time; the tool itself doesn't decide "where to start." This ensures all path operations are within a controllable scope, and also lets tools use tmp_path for isolation during testing.

2. ToolResult is an immutable type

The frozen=True dataclass means tools cannot modify results after run() returns. Any step in the pipeline (optimistic lock injection, byte budget truncation, etc.) produces a new object rather than modifying the original, reducing the possibility of data races.

3. Three states, not just success/failure

status=partial is for situations where "the result is usable but discounted" — for example, reading a large file and returning only the first 500 lines, or using encoding fallback. When the model sees partial it knows the result may be incomplete and can ask follow-up questions or adjust strategy; seeing success it can use the result confidently.


Summary

Step Key Points
Inherit Tool run() returns ToolResult, get_parameters() defines the parameter schema
Sandbox protection target.relative_to(project_root) is a required check for every tool involving the filesystem
Registration registry.register_tool(YourTool(project_root=...))
Prompt description is the sole basis on which the model selects a tool — clearly state "when to use it"
Tests Cover at minimum: success path, missing parameter, sandbox escape

The next article covers connecting a new LLM provider — the tool system is the agent's hands, and the LLM is the agent's brain; they are equally important extension points.


About the Source Code for This Series

All analysis in this series is based on the open source project MyCodeAgent.

The source code already has companion comments added at key locations in the order covered by this series — you can read alongside the code, or clone it directly to run, modify, and extend it to build your own agent.

git clone https://github.com/chendongqi/MyCodeAgent
cd MyCodeAgent
cp .env.example .env   # fill in your LLM API key
uv sync
uv run python main.py
Enter fullscreen mode Exit fullscreen mode

Visit PrimeSkills — a carefully curated AI Agent and skills marketplace where every piece of content is validated through real enterprise-grade workflows. No hype, only what actually works.

For more practical knowledge and interesting products, visit my personal homepage

Top comments (0)