How to Extend the Osmosis Agent Toolkit with Custom Tools: A Complete Developer's Guide

You can extend the Osmosis Agent Toolkit by implementing the Tool<I, O> interface, registering the class in OsmosisAgentToolkit, and exposing it through the AI SDK wrapper—no core modifications required.

The jonator/osmosis-agent-toolkit provides a minimal, type-safe plugin architecture that lets you add bespoke functionality beyond the built-in swap and account tools. Because every tool is simply a class implementing a standard interface, you can inject custom analytics, governance actions, or proprietary data sources while maintaining full compatibility with LLM-driven agents.

Understanding the Plugin Architecture

The toolkit follows a facade pattern built around three core components:

• Tool<I, O> – Defined in packages/core/src/tools/tool.ts, this interface requires a name, description, optional Zod parameters schema, and an async call method.
• OsmosisAgentToolkit – Located in packages/core/src/toolkit.ts, this class constructs built-in tools and exposes them via getters.
• AI SDK Wrapper – Found in packages/ai-sdk/src/toolkit.ts, this converts core Tool instances into AI SDK-compatible functions.

Because dependencies are injected through constructors rather than hard-coded, you can add unlimited custom tools without forking the repository.

Step 1: Create a Custom Tool Class

Implement the Tool interface in a new file under packages/core/src/tools/. Use Zod to define parameter schemas for automatic JSON Schema generation.

// packages/core/src/tools/custom-price-history.ts
import { z } from 'zod';
import type { Tool } from './tool.js';
import type { OsmosisSqsQueryClient } from '../queries/sqs/client.js';

export class PriceHistoryTool
  implements
    Tool<
      {
        tickers: string[];
        days?: number;
      },
      Record<string, number[]>
    >
{
  readonly name = 'getPriceHistory';
  readonly description = 'Fetch historic price points for the supplied tickers.';
  readonly parameters = z.object({
    tickers: z.array(z.string()).describe('Asset symbols, e.g. ["OSMO","ATOM"]'),
    days: z
      .number()
      .int()
      .min(1)
      .max(365)
      .default(7)
      .describe('Number of days to look back'),
  });

  constructor(private readonly sqsClient: OsmosisSqsQueryClient) {}

  async call({
    tickers,
    days = 7,
  }: {
    tickers: string[];
    days?: number;
  }): Promise<Record<string, number[]>> {
    const history = await this.sqsClient.getPriceHistory(tickers, days);
    return history;
  }
}

Key implementation details:

• Generic types – I defines the input shape; O defines the return type.
• Dependency injection – Pass external services like OsmosisSqsQueryClient via the constructor to keep the tool testable and pure.
• Zod schemas – The optional parameters field enables automatic OpenAPI-compatible schema generation for LLM function calling.

Step 2: Register the Tool in the Core Toolkit

Modify packages/core/src/toolkit.ts to instantiate your tool and expose it via a public getter. Follow the existing pattern used for built-in tools like the account tool.

// packages/core/src/toolkit.ts
import { PriceHistoryTool } from './tools/custom-price-history.js';

export class OsmosisAgentToolkit {
  // Existing tools...
  protected readonly _priceHistoryTool = new PriceHistoryTool(this._sqsClient);

  get priceHistoryTool() {
    return this._priceHistoryTool;
  }

  // Optional: expose all tools as an array for easy iteration
  get allTools(): Tool<any, any>[] {
    return [
      this._accountTool,
      this._swapQuoteInGivenOutTool,
      this._swapQuoteOutGivenInTool,
      this._sendSwapInGivenOutQuoteTxTool,
      this._sendSwapOutGivenInQuoteTxTool,
      this._priceHistoryTool,
    ];
  }
}

The getter pattern (get toolName()) maintains a stable public API and allows downstream wrappers to discover tools dynamically via reflection.

Step 3: Expose the Tool to the AI SDK

If you are using the AI SDK package, extend the wrapper in packages/ai-sdk/src/toolkit.ts to include your custom tool in the exported tools map.

// packages/ai-sdk/src/toolkit.ts
import { PriceHistoryTool } from '@osmosis-agent-toolkit/core/src/tools/custom-price-history.js';
import { makeAiSdkTool } from './utils.js'; // hypothetical utility

export class OsmosisAgentToolkit extends CoreOsmosisAgentToolkit {
  get tools() {
    return {
      ...super.tools,
      priceHistoryTool: makeAiSdkTool(super['priceHistoryTool']),
    };
  }
}

The wrapper converts the core Tool instance into an AI SDK tool object that LLMs can invoke via function calling. Accessing the tool through super['priceHistoryTool'] ensures you retrieve the instance from the parent class's getter.

Step 4: Use the Custom Tool

Once registered, use your tool both programmatically and within LLM contexts.

Direct programmatic usage:

import { OsmosisAgentToolkit } from '@osmosis-agent-toolkit/ai-sdk';

const toolkit = new OsmosisAgentToolkit('your mnemonic phrase here');

const history = await toolkit.priceHistoryTool.call({
  tickers: ['OSMO', 'ATOM'],
  days: 14,
});
console.log(history); // { OSMO: [0.98, 1.02, ...], ATOM: [9.5, 9.7, ...] }

LLM-driven usage:

const response = await openai.chat.completions.create({
  model: 'gpt-4',
  messages: [{ role: 'user', content: 'What was OSMO price last week?' }],
  tools: Object.values(toolkit.tools),
});

Testing Your Custom Tool

Add unit tests under packages/core/tests/ to verify logic and dependency injection. The following Vitest example demonstrates mocking the SQS client:

// packages/core/tests/price-history.test.ts
import { describe, it, expect, vi } from 'vitest';
import { PriceHistoryTool } from '../src/tools/custom-price-history.js';

describe('PriceHistoryTool', () => {
  it('returns mocked price history', async () => {
    const mockClient = {
      getPriceHistory: vi.fn().mockResolvedValue({
        OSMO: [1.0, 0.98, 1.02],
      }),
    };

    const tool = new PriceHistoryTool(mockClient as any);
    const result = await tool.call({ tickers: ['OSMO'], days: 3 });

    expect(result.OSMO).toEqual([1.0, 0.98, 1.02]);
    expect(mockClient.getPriceHistory).toHaveBeenCalledWith(['OSMO'], 3);
  });
});

Isolating external dependencies ensures your custom tool logic remains robust regardless of network conditions or API changes.

Summary

• Implement the Tool<I, O> interface with strongly typed inputs and outputs.
• Inject dependencies via the constructor to maintain testability.
• Register the tool in packages/core/src/toolkit.ts by adding a protected field and public getter.
• Expose the tool through the AI SDK wrapper by extending the tools object in packages/ai-sdk/src/toolkit.ts.
• Test thoroughly using mocked dependencies to ensure reliability.

Following this pattern allows you to extend the Osmosis Agent Toolkit indefinitely with custom analytics, governance features, or proprietary data sources while maintaining full type safety and LLM compatibility.

Frequently Asked Questions

Do I need to modify the core repository to add custom tools?

No. Because the toolkit uses a pure interface-based plugin system, you can implement the Tool interface in your own codebase and pass instances to the toolkit's constructor or extend the class. However, for full integration with the AI SDK wrapper and the allTools getter, modifying the core files as shown above provides the cleanest developer experience.

What schema library should I use for parameter validation?

The toolkit uses Zod for runtime validation and type inference. Define your parameters field as a z.ZodType to enable automatic JSON Schema generation for OpenAI function calling. Other schema libraries will work if they can produce compatible JSON Schema descriptions, but Zod is the tested and documented standard in packages/core/src/tools/tool.ts.

How do I handle errors inside a custom tool's call method?

Throw standard JavaScript Error objects with descriptive messages. The AI SDK wrapper catches these and returns them as tool error results to the LLM, allowing the agent to retry or request clarification. For recoverable errors, consider returning a specific error shape in your output type O rather than throwing, so the LLM receives structured data for decision-making.

Can custom tools access the wallet or account state?

Yes. Follow the dependency injection pattern shown in packages/core/src/tools/account.ts. Accept the Account class or wallet client in your constructor, then use those methods inside call(). This keeps your tool decoupled from specific wallet implementations while allowing on-chain transactions or balance queries.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}