import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { toolManager, ToolLimitError } from '../../src/tools/ToolManager.js';
import { createMockMemoryStore } from '../helpers/mockMemoryStore.js';
import { BasicAgent } from '../../src/agents/BasicAgent.js';
import { Orchestrator } from '../../src/orchestrator/Orchestrator.js';
import { SequentialStrategy } from '../../src/orchestrator/strategies/SequentialStrategy.js';
import { memoryStore } from '../../src/memory/store/index.js';
// Mock the memory store for testing
vi.mock('../../src/memory/store/index.js', () => ({
memoryStore: {
add: vi.fn().mockResolvedValue(undefined),
save: vi.fn().mockResolvedValue(undefined)
},
graph: {
entities: new Map()
}
}));
describe('Tool Limits', () => {
// Save original environment and restore after tests
const originalEnv = { ...process.env };
beforeEach(() => {
// Reset tool manager before each test
toolManager.reset();
// Set test mode
process.env.NODE_ENV = 'test';
// Mock the executeToolCall method since it's unimplemented in tests
// @ts-expect-error - accessing private property for testing
toolManager.executeToolCall = vi.fn().mockImplementation(async (toolName, params) => {
return `Executed ${toolName} with params ${JSON.stringify(params)}`;
});
// Reset mocks
vi.clearAllMocks();
});
afterEach(() => {
// Restore original environment
process.env = { ...originalEnv };
});
it('should track tool calls and throw an error when limit is exceeded', async () => {
// Set a low limit for testing
process.env.TOOL_LIMIT = '3';
// Override the globalCount directly to simulate reaching the limit
// @ts-expect-error - accessing private property for testing
toolManager.globalCount = 0;
// Make some tool calls
await toolManager.callTool('test-agent', 'test-tool', { param: 'value1' });
await toolManager.callTool('test-agent', 'test-tool', { param: 'value2' });
await toolManager.callTool('test-agent', 'test-tool', { param: 'value3' });
// When we reach 3 calls, the next one should throw
// Manually check that globalCount is 3 after the calls
const stats = toolManager.getStats();
expect(stats.global).toBe(3);
expect(stats.perAgent.get('test-agent')).toBe(3);
// Override our fake execute implementation to throw the error when the limit is reached
// @ts-expect-error - accessing private property for testing
toolManager.executeToolCall = vi.fn().mockImplementation(() => {
throw new ToolLimitError('Tool call limit exceeded in test');
});
// Now the fourth call should throw a ToolLimitError
await expect(
toolManager.callTool('test-agent', 'test-tool', { param: 'value4' })
).rejects.toThrow(ToolLimitError);
});
it('should cache duplicate tool calls to avoid counting them twice', async () => {
// Set environment to enable caching
process.env.CACHE_TOOL_CALLS = 'true';
process.env.TOOL_LIMIT = '3';
// Make identical tool calls
const result1 = await toolManager.callTool('test-agent', 'cached-tool', { id: 123 });
const result2 = await toolManager.callTool('test-agent', 'cached-tool', { id: 123 });
// The results should be the same
expect(result1).toBe(result2);
// The count should only be 1 since the second call was cached
const stats = toolManager.getStats();
expect(stats.global).toBe(1);
});
it('should handle tool limits in the Orchestrator', async () => {
// Set a low limit for testing
process.env.TOOL_LIMIT = '3';
// Create a mock agent that uses tools
const mockMemory = createMockMemoryStore();
const agent = new BasicAgent('agent1', mockMemory);
// Mock the step method to trigger a tool limit error on the 4th call
let callCount = 0;
agent.step = async () => {
callCount++;
if (callCount <= 3) {
// First 3 calls are successful
return `Step ${callCount} output`;
} else {
// 4th call throws ToolLimitError
throw new ToolLimitError('Tool call limit exceeded in test');
}
};
// Create an orchestrator with the agent
const strategy = new SequentialStrategy();
const orchestrator = new Orchestrator([agent], strategy);
// Run the orchestration
const result = await orchestrator.run('test input');
// Verify the result shows the limit was reached
expect(result.status).toBe('HALTED_LIMIT');
expect(result.output).toContain('halted due to tool call limit');
// Verify memoryStore.add was called with the right arguments
expect(memoryStore.add).toHaveBeenCalledWith(
'OrchestrationLimits',
expect.stringContaining('Tool call limit reached'),
expect.objectContaining({
tags: ['limit_reached']
})
);
});
it('should enforce tool whitelist in Orchestrator', async () => {
// Create a mock memory store
const mockMemory = createMockMemoryStore();
// Create a mock agent
const agent = new BasicAgent('agent1', mockMemory);
// Set up a whitelist of allowed tools
const allowedTools = ['allowed-tool'];
// Create the orchestrator with the whitelist
const strategy = new SequentialStrategy();
const orchestrator = new Orchestrator(
[agent],
strategy,
{ allowedTools }
);
// Reset the tool manager to clear any previous state
toolManager.reset();
// Verify the whitelist was set in the tool manager
expect(toolManager['allowedTools']).toEqual(allowedTools);
});
});