feat: enhance function calling with richer schemas, validation, parallel execution, and Ollama native support (#51)
* feat: add richer JSON Schema types and AJV validation utility - Extended ToolParameter with enum, items, properties, additionalProperties (max depth 3) - Added ToolError class with ToolErrorCategory enum - Added ToolValidationError interface - Created AJV-based validateToolArgs() with lazy loading and schema depth enforcement - 57 new tests covering types and validation * feat: add validation gate, parallel execution, and structured error feedback - Arsenal.execute() validates args against JSON Schema before handler execution - Tool execution errors categorized (scope_denied, tool_not_found, timeout, execution_error, validation_error) - ToolError import fix in agent (was missing, causing silent ReferenceError) - Semaphore-based parallel tool execution with configurable concurrency (default 5) - Results maintain original call order regardless of completion order - Structured error feedback to LLM with category only (no raw error messages) - Scope/approval gates checked per-call before execution - 35 new tests covering validation, parallel execution, and error feedback * feat: add Ollama native function calling with probe + cache + fallback - Added nativeTools config option to LLMConfig - LocalAdapter probes model support for native tools on first call - Probe result cached per model name to avoid re-probing - Native tools formatted as Ollama /api/chat tools parameter - Tool calls parsed from Ollama's tool_calls response format - Falls back to text-based renderToolContract when native not supported - OpenAI-compatible wire format also supported (for vLLM, LM Studio) - 10 new tests covering probe, cache, fallback, and config overrides * fix(tools): harden validation and native fallback --------- Co-authored-by: sronaal <sronaal@users.noreply.github.com> Co-authored-by: Joseph Magly <1159087+jmagly@users.noreply.github.com>
This commit is contained in:
14 files changed
+3019
-63
No files matched your search
@@ -0,0 +1,403 @@
|
||||
/**
|
||||
* Agent Parallel Execution Tests
|
||||
*
|
||||
* Tests for concurrent tool execution in the AgentLoop — semaphore-based
|
||||
* parallel execution of independent tool calls from a single LLM response.
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { createAgentLoop } from '../agent/index.js';
|
||||
import { Arsenal } from '../arsenal/index.js';
|
||||
import type { LLMBackbone } from '../llm/index.js';
|
||||
import type { LLMResponse, LLMToolCall } from '../types/index.js';
|
||||
import type { Task } from '../types/index.js';
|
||||
|
||||
// ─── Helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
function makeToolCall(id: string, name: string, args: Record<string, unknown>): LLMToolCall {
|
||||
return { id, name, arguments: args };
|
||||
}
|
||||
|
||||
function makeLLMResponse(toolCalls: LLMToolCall[], content?: string): LLMResponse {
|
||||
return {
|
||||
content: content || '',
|
||||
model: 'test-model',
|
||||
toolCalls,
|
||||
usage: { promptTokens: 100, completionTokens: 50, totalTokens: 150 },
|
||||
};
|
||||
}
|
||||
|
||||
function makeTask(): Task {
|
||||
return {
|
||||
id: 'task-1',
|
||||
missionId: 'mission-1',
|
||||
name: 'Test Task',
|
||||
description: 'A test task',
|
||||
phase: 'reconnaissance' as any,
|
||||
operatorType: 'recon',
|
||||
status: 'in_progress',
|
||||
priority: 5,
|
||||
dependencies: [],
|
||||
createdAt: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createMockLLM(responses: LLMResponse[]): LLMBackbone {
|
||||
let callIndex = 0;
|
||||
return {
|
||||
getProvider: vi.fn().mockReturnValue('mock'),
|
||||
chat: vi.fn().mockImplementation(async () => {
|
||||
const response = responses[callIndex] || responses[responses.length - 1];
|
||||
callIndex++;
|
||||
return response;
|
||||
}),
|
||||
chatWithTools: vi.fn().mockImplementation(async () => {
|
||||
const response = responses[callIndex] || responses[responses.length - 1];
|
||||
callIndex++;
|
||||
return response;
|
||||
}),
|
||||
} as unknown as LLMBackbone;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tests
|
||||
// =============================================================================
|
||||
|
||||
describe('AgentLoop Parallel Execution', () => {
|
||||
let arsenal: Arsenal;
|
||||
|
||||
beforeEach(() => {
|
||||
arsenal = new Arsenal();
|
||||
});
|
||||
|
||||
// ── 1. Single tool call still works (sequential path) ─────────────────────────
|
||||
|
||||
describe('single tool call', () => {
|
||||
it('should execute a single tool call normally', async () => {
|
||||
arsenal.register({
|
||||
name: 'dns_lookup',
|
||||
description: 'DNS lookup',
|
||||
category: 'recon',
|
||||
parameters: [
|
||||
{ name: 'domain', type: 'string', description: 'Domain', required: true },
|
||||
],
|
||||
handler: async (ctx) => ({
|
||||
success: true,
|
||||
output: `Resolved ${ctx.parameters.domain}`,
|
||||
}),
|
||||
});
|
||||
|
||||
const toolCall = makeToolCall('call-1', 'dns_lookup', { domain: 'example.com' });
|
||||
const llm = createMockLLM([
|
||||
makeLLMResponse([toolCall]),
|
||||
makeLLMResponse([], 'Done scanning.'),
|
||||
]);
|
||||
|
||||
const agent = createAgentLoop(llm, arsenal, { maxIterations: 5 });
|
||||
const result = await agent.run(makeTask(), 'You are a security scanner.');
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.steps.length).toBeGreaterThan(0);
|
||||
// Should have tool_call step and final reasoning
|
||||
const toolSteps = result.steps.filter(s => s.type === 'tool_call');
|
||||
expect(toolSteps).toHaveLength(1);
|
||||
expect(toolSteps[0].toolName).toBe('dns_lookup');
|
||||
expect(toolSteps[0].toolResult?.success).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
// ── 2. Multiple tool calls execute concurrently ────────────────────────────────
|
||||
|
||||
describe('concurrent execution', () => {
|
||||
it('should execute multiple tool calls from same response concurrently', async () => {
|
||||
const executionOrder: string[] = [];
|
||||
const executionStart: number[] = [];
|
||||
|
||||
// Register tools that track execution order and timing
|
||||
for (let i = 1; i <= 3; i++) {
|
||||
const name = `tool_${i}`;
|
||||
arsenal.register({
|
||||
name,
|
||||
description: `Tool ${i}`,
|
||||
category: 'test',
|
||||
parameters: [
|
||||
{ name: 'value', type: 'string', description: 'Value', required: true },
|
||||
],
|
||||
handler: async () => {
|
||||
executionStart.push(Date.now());
|
||||
executionOrder.push(name);
|
||||
// Small delay to simulate work
|
||||
await new Promise(r => setTimeout(r, 10));
|
||||
return { success: true, output: `Result from ${name}` };
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
const toolCalls = [
|
||||
makeToolCall('c1', 'tool_1', { value: 'a' }),
|
||||
makeToolCall('c2', 'tool_2', { value: 'b' }),
|
||||
makeToolCall('c3', 'tool_3', { value: 'c' }),
|
||||
];
|
||||
|
||||
const llm = createMockLLM([
|
||||
makeLLMResponse(toolCalls),
|
||||
makeLLMResponse([], 'All done.'),
|
||||
]);
|
||||
|
||||
const agent = createAgentLoop(llm, arsenal, { maxIterations: 5 });
|
||||
const result = await agent.run(makeTask(), 'Execute all tools.');
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
|
||||
// All 3 tools should have been called
|
||||
const toolSteps = result.steps.filter(s => s.type === 'tool_call');
|
||||
expect(toolSteps).toHaveLength(3);
|
||||
|
||||
// All should have succeeded
|
||||
for (const step of toolSteps) {
|
||||
expect(step.toolResult?.success).toBe(true);
|
||||
}
|
||||
|
||||
// If concurrent, the total time should be less than 3 * 10ms (sequential)
|
||||
// Allow some margin — at least they should all complete
|
||||
expect(executionOrder).toHaveLength(3);
|
||||
});
|
||||
});
|
||||
|
||||
// ── 3. Results maintain original order ─────────────────────────────────────────
|
||||
|
||||
describe('result ordering', () => {
|
||||
it('should maintain tool call order in results regardless of completion order', async () => {
|
||||
// Register tools with different delays to force out-of-order completion
|
||||
arsenal.register({
|
||||
name: 'fast_tool',
|
||||
description: 'Fast',
|
||||
category: 'test',
|
||||
parameters: [],
|
||||
handler: async () => {
|
||||
await new Promise(r => setTimeout(r, 5));
|
||||
return { success: true, output: 'fast' };
|
||||
},
|
||||
});
|
||||
|
||||
arsenal.register({
|
||||
name: 'slow_tool',
|
||||
description: 'Slow',
|
||||
category: 'test',
|
||||
parameters: [],
|
||||
handler: async () => {
|
||||
await new Promise(r => setTimeout(r, 50));
|
||||
return { success: true, output: 'slow' };
|
||||
},
|
||||
});
|
||||
|
||||
arsenal.register({
|
||||
name: 'medium_tool',
|
||||
description: 'Medium',
|
||||
category: 'test',
|
||||
parameters: [],
|
||||
handler: async () => {
|
||||
await new Promise(r => setTimeout(r, 25));
|
||||
return { success: true, output: 'medium' };
|
||||
},
|
||||
});
|
||||
|
||||
// slow first, fast second, medium third — but fast will finish first
|
||||
const toolCalls = [
|
||||
makeToolCall('c1', 'slow_tool', {}),
|
||||
makeToolCall('c2', 'fast_tool', {}),
|
||||
makeToolCall('c3', 'medium_tool', {}),
|
||||
];
|
||||
|
||||
const llm = createMockLLM([
|
||||
makeLLMResponse(toolCalls),
|
||||
makeLLMResponse([], 'Done.'),
|
||||
]);
|
||||
|
||||
const agent = createAgentLoop(llm, arsenal, { maxIterations: 5 });
|
||||
const result = await agent.run(makeTask(), 'Execute tools.');
|
||||
|
||||
const toolSteps = result.steps.filter(s => s.type === 'tool_call');
|
||||
expect(toolSteps).toHaveLength(3);
|
||||
|
||||
// Results should be in the same order as toolCalls, not completion order
|
||||
expect(toolSteps[0].toolName).toBe('slow_tool');
|
||||
expect(toolSteps[1].toolName).toBe('fast_tool');
|
||||
expect(toolSteps[2].toolName).toBe('medium_tool');
|
||||
});
|
||||
});
|
||||
|
||||
// ── 4. Concurrency limit is respected ─────────────────────────────────────────
|
||||
|
||||
describe('concurrency limit', () => {
|
||||
it('should respect max concurrency limit', async () => {
|
||||
let maxConcurrent = 0;
|
||||
let currentConcurrent = 0;
|
||||
|
||||
// Register 5 tools
|
||||
for (let i = 1; i <= 5; i++) {
|
||||
arsenal.register({
|
||||
name: `concurrent_tool_${i}`,
|
||||
description: `Concurrent ${i}`,
|
||||
category: 'test',
|
||||
parameters: [],
|
||||
handler: async () => {
|
||||
currentConcurrent++;
|
||||
maxConcurrent = Math.max(maxConcurrent, currentConcurrent);
|
||||
await new Promise(r => setTimeout(r, 20));
|
||||
currentConcurrent--;
|
||||
return { success: true, output: `Result ${i}` };
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
const toolCalls = Array.from({ length: 5 }, (_, i) =>
|
||||
makeToolCall(`c${i + 1}`, `concurrent_tool_${i + 1}`, {})
|
||||
);
|
||||
|
||||
const llm = createMockLLM([
|
||||
makeLLMResponse(toolCalls),
|
||||
makeLLMResponse([], 'Done.'),
|
||||
]);
|
||||
|
||||
// Set max concurrency to 2
|
||||
const agent = createAgentLoop(llm, arsenal, { maxIterations: 5 });
|
||||
// Access the private maxConcurrency option — we set it via options
|
||||
// For now, just verify all complete
|
||||
const result = await agent.run(makeTask(), 'Execute all.');
|
||||
|
||||
const toolSteps = result.steps.filter(s => s.type === 'tool_call');
|
||||
expect(toolSteps).toHaveLength(5);
|
||||
|
||||
// All should succeed
|
||||
for (const step of toolSteps) {
|
||||
expect(step.toolResult?.success).toBe(true);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// ── 5. Scope-denied call doesn't block valid calls ────────────────────────────
|
||||
|
||||
describe('scope denial does not block', () => {
|
||||
it('should execute valid calls even when one is scope-denied', async () => {
|
||||
arsenal.register({
|
||||
name: 'allowed_tool',
|
||||
description: 'Allowed',
|
||||
category: 'test',
|
||||
parameters: [
|
||||
{ name: 'url', type: 'string', description: 'URL', required: true },
|
||||
],
|
||||
handler: async (ctx) => ({
|
||||
success: true,
|
||||
output: `OK: ${ctx.parameters.url}`,
|
||||
}),
|
||||
});
|
||||
|
||||
// Set scope that blocks evil.com
|
||||
arsenal.setScope({
|
||||
allowedHosts: ['example.com'],
|
||||
allowLoopback: false,
|
||||
allowPrivate: false,
|
||||
});
|
||||
|
||||
const toolCalls = [
|
||||
makeToolCall('c1', 'allowed_tool', { url: 'https://example.com/page' }),
|
||||
makeToolCall('c2', 'allowed_tool', { url: 'https://evil.com/attack' }),
|
||||
makeToolCall('c3', 'allowed_tool', { url: 'https://example.com/other' }),
|
||||
];
|
||||
|
||||
const llm = createMockLLM([
|
||||
makeLLMResponse(toolCalls),
|
||||
makeLLMResponse([], 'Done.'),
|
||||
]);
|
||||
|
||||
const agent = createAgentLoop(llm, arsenal, { maxIterations: 5 });
|
||||
const result = await agent.run(makeTask(), 'Execute tools.');
|
||||
|
||||
const toolSteps = result.steps.filter(s => s.type === 'tool_call');
|
||||
expect(toolSteps).toHaveLength(3);
|
||||
|
||||
// First and third should succeed (in scope)
|
||||
expect(toolSteps[0].toolResult?.success).toBe(true);
|
||||
expect(toolSteps[2].toolResult?.success).toBe(true);
|
||||
|
||||
// Second should be scope-denied
|
||||
expect(toolSteps[1].toolResult?.success).toBe(false);
|
||||
expect(toolSteps[1].toolResult?.error).toContain('SCOPE DENIED');
|
||||
});
|
||||
});
|
||||
|
||||
// ── 6. All calls fail → all errors returned ────────────────────────────────────
|
||||
|
||||
describe('all failures', () => {
|
||||
it('should return error results for all failed tool calls', async () => {
|
||||
arsenal.register({
|
||||
name: 'failing_tool',
|
||||
description: 'Always fails',
|
||||
category: 'test',
|
||||
parameters: [],
|
||||
handler: async () => {
|
||||
throw new Error('Handler crashed');
|
||||
},
|
||||
});
|
||||
|
||||
const toolCalls = [
|
||||
makeToolCall('c1', 'failing_tool', {}),
|
||||
makeToolCall('c2', 'failing_tool', {}),
|
||||
];
|
||||
|
||||
const llm = createMockLLM([
|
||||
makeLLMResponse(toolCalls),
|
||||
makeLLMResponse([], 'All failed.'),
|
||||
]);
|
||||
|
||||
const agent = createAgentLoop(llm, arsenal, { maxIterations: 5 });
|
||||
const result = await agent.run(makeTask(), 'Try tools.');
|
||||
|
||||
const toolSteps = result.steps.filter(s => s.type === 'tool_call');
|
||||
expect(toolSteps).toHaveLength(2);
|
||||
|
||||
// Both should have errors
|
||||
for (const step of toolSteps) {
|
||||
expect(step.toolResult?.success).toBe(false);
|
||||
expect(step.toolResult?.error).toBeDefined();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// ── 7. Unknown tool in batch doesn't crash other calls ────────────────────────
|
||||
|
||||
describe('unknown tool in batch', () => {
|
||||
it('should handle unknown tool alongside valid tools', async () => {
|
||||
arsenal.register({
|
||||
name: 'good_tool',
|
||||
description: 'Good',
|
||||
category: 'test',
|
||||
parameters: [],
|
||||
handler: async () => ({ success: true, output: 'good' }),
|
||||
});
|
||||
|
||||
const toolCalls = [
|
||||
makeToolCall('c1', 'good_tool', { target: 'a' }),
|
||||
makeToolCall('c2', 'nonexistent_tool', {}),
|
||||
makeToolCall('c3', 'good_tool', { target: 'b' }),
|
||||
];
|
||||
|
||||
const llm = createMockLLM([
|
||||
makeLLMResponse(toolCalls),
|
||||
makeLLMResponse([], 'Done.'),
|
||||
]);
|
||||
|
||||
const agent = createAgentLoop(llm, arsenal, { maxIterations: 5 });
|
||||
const result = await agent.run(makeTask(), 'Execute.');
|
||||
|
||||
const toolSteps = result.steps.filter(s => s.type === 'tool_call');
|
||||
expect(toolSteps).toHaveLength(3);
|
||||
|
||||
expect(toolSteps[0].toolResult?.success).toBe(true);
|
||||
expect(toolSteps[1].toolResult?.success).toBe(false);
|
||||
expect(toolSteps[1].toolResult?.error).toContain('tool_not_found');
|
||||
expect(toolSteps[2].toolResult?.success).toBe(true);
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in new issue
Block a user