How to Create Custom Tools and Extend the Tool Registry in Kimi Code

Kimi Code provides an agent-scoped tool registry that lets you register custom ExecutableTool objects via the IAgentToolRegistryService contract, enabling persistent, isolated capabilities that survive session restarts.

Creating custom tools and extending the tool registry in Kimi Code allows you to augment agent capabilities with domain-specific logic. The MoonshotAI/kimi-code repository implements a robust dependency injection framework where each agent receives its own tool registry, ensuring clean isolation between different agent instances. By implementing the ExecutableTool interface and hooking into the agent lifecycle, you can add bespoke functionality that behaves exactly like native tools.

Understanding the Tool Registry Architecture

The tool registry system centers on the IAgentToolRegistryService contract defined in packages/agent-core-v2/src/agent/toolRegistry/toolRegistry.ts. This service maintains an internal Map<string, ToolEntry> that stores all available tools for a specific agent instance.

The ExecutableTool Contract

Every tool must implement the ExecutableTool interface from the shared toolContract. The contract requires a resolveExecution method that returns an object containing:

  • An approvalRule string that determines whether user confirmation is required
  • An execute callback function containing the actual tool logic

When the agent selects a tool during a turn, the registry invokes resolveExecution with the parsed arguments, retrieves the approval rule, and then executes the callback if approved.

Agent-Scoped Isolation

Unlike global registries, Kimi Code's implementation in packages/agent-core-v2/src/agent/toolRegistry/toolRegistryService.ts creates a unique registry for each agent. This means tools registered in one agent do not pollute the namespace of another, which is critical for multi-agent scenarios involving sub-agents (Agent tool) or swarms (AgentSwarm tool). The registry also returns an IDisposable when registering a tool, enabling automatic cleanup when scopes are destroyed.

Implementing a Custom ExecutableTool

To create a custom tool, define an object that conforms to the ExecutableTool interface. Below is a complete implementation of an Echo tool that returns input text unchanged:

// src/custom-tools/echoTool.ts
import type { ExecutableTool, ToolExecutionResult } from '#/tool/toolContract';

export const EchoTool: ExecutableTool = {
  name: 'Echo',
  description: 'Returns the supplied text unchanged.',
  // Zod-compatible JSON schema of arguments
  parameters: {
    type: 'object',
    properties: {
      text: { type: 'string', description: 'The text to echo back.' },
    },
    required: ['text'],
    additionalProperties: false,
  },

  // Called by the engine when the tool is selected
  resolveExecution: (args) => ({
    // The approval rule matches the tool name; can be overridden by policy
    approvalRule: 'Echo',
    // The actual implementation – can be async
    execute: async (ctx) => {
      const { text } = args as { text: string };
      const result: ToolExecutionResult = { output: text };
      return result;
    },
  }),
};

The parameters property must provide a valid JSON Schema object describing the tool's arguments. This schema drives the agent's ability to construct proper tool calls and validates inputs before execution.

Registering Tools in the Agent Lifecycle

Tools must be registered when an agent scope is created. Kimi Code uses a dependency injection container with lifecycle management, allowing you to hook registration logic into the OnScopeCreated activation event.

// src/plugin/registerEcho.ts
import { registerScopedService } from '#/_base/di/scope';
import { LifecycleScope, ScopeActivation } from '#/_base/di/scope';
import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry';
import { EchoTool } from '../custom-tools/echoTool';

class EchoToolRegistrar {
  constructor(@IAgentToolRegistryService private readonly registry: IAgentToolRegistryService) {
    // Register the tool; the returned disposable is managed by the DI container
    this.registry.register(EchoTool);
  }
}

// Bind the registrar so it is instantiated for every new Agent
registerScopedService(
  LifecycleScope.Agent,
  IAgentToolRegistryService,
  EchoToolRegistrar,
  ScopeActivation.OnScopeCreated,
  'echoToolRegistrar',
);

The register method in toolRegistryService.ts (lines 33-44) accepts the tool definition and returns an IDisposable. When this disposable is triggered—either manually or automatically during scope teardown—the tool is removed from the registry via unregisterTool, preventing resource leaks.

Persisting User Tools Across Sessions

For tools defined by end-users rather than system plugins, Kimi Code provides the AgentUserToolService in packages/agent-core-v2/src/agent/userTool/userToolService.ts. This service bridges the persisted UserToolModel with the live registry, ensuring custom tools survive application restarts.

When registering a user tool, the service performs three operations:

  1. Dispatches a registerUserTool operation to the wire model for persistence
  2. Calls registry.register with the tool definition marked as source: 'user'
  3. Optionally adds the tool name to the agent's active-tool list via profile.addActiveTool

During session restoration, the restoreRegisteredTools method recreates live registrations from the persisted model, guaranteeing that a resumed agent sees exactly the same custom tools that existed before the restart.

Using Custom Tools in Agent Sessions

Once registered, custom tools are indistinguishable from built-in capabilities. The agent can invoke your tool using standard JSON tool calls:

{
  "tool": "Echo",
  "arguments": { "text": "Hello, world!" }
}

Because the Echo tool specifies an approvalRule matching its name, and assuming default policies treat read-only tools as auto-approved, the agent will receive the response "Hello, world!" without requiring explicit user confirmation.

Summary

  • Agent-scoped registries: Each agent maintains an isolated IAgentToolRegistryService instance, preventing namespace collisions in multi-agent scenarios.
  • ExecutableTool interface: Custom tools must implement resolveExecution to provide approval rules and execution logic, along with JSON Schema parameters.
  • Lifecycle integration: Register tools during scope creation using registerScopedService with ScopeActivation.OnScopeCreated.
  • Automatic cleanup: The register method returns an IDisposable that automatically unregisters tools when the agent scope ends.
  • Session persistence: Use AgentUserToolService to bridge user-defined tools to persistent storage and restore them across restarts via restoreRegisteredTools.

Frequently Asked Questions

What is the IAgentToolRegistryService in Kimi Code?

The IAgentToolRegistryService is the core contract for tool management in Kimi Code, defined in packages/agent-core-v2/src/agent/toolRegistry/toolRegistry.ts. It provides methods to register, unregister, and resolve ExecutableTool objects within an agent-specific scope. The concrete implementation in toolRegistryService.ts stores tools in a Map<string, ToolEntry> and handles the lifecycle of tool availability during an agent session.

How do custom tools maintain isolation between agents?

Custom tools remain isolated because Kimi Code instantiates a unique tool registry for each agent via dependency injection scoping. When you register a tool using LifecycleScope.Agent, the registration exists only within that specific agent's context. This design prevents tools registered for one agent from appearing in another agent's available tool set, which is essential when using sub-agents or agent swarms.

Can custom tools persist across session restarts?

Yes, custom tools can persist when registered through the AgentUserToolService. This service, located in packages/agent-core-v2/src/agent/userTool/userToolService.ts, dispatches registerUserTool operations to the wire model for storage. When the application restarts, the restoreRegisteredTools method rebuilds the live registry from the persisted UserToolModel, ensuring your custom capabilities remain available without re-registration.

What approval rules apply to custom tools?

The approval rule for a custom tool is returned by the resolveExecution method as the approvalRule property. This string typically matches the tool name (e.g., 'Echo') and is evaluated against the agent's policy configuration. If the policy defines the rule as auto-approved, the tool executes immediately; otherwise, the system prompts for user confirmation before invoking the execute callback.

Have a question about this repo?

These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:

Share the following with your agent to get started:
curl -s "https://instagit.com/install.md"

Works with
Claude Codex Cursor VS Code OpenClaw Any MCP Client

Maintain an open-source project? Get it listed too →