This commit is contained in:
Oleg Maslov
2026-09-02 10:14:22 +02:00
parent 0c3e2ead3b
commit b20b138fe4
771 changed files with 161561 additions and 9027 deletions

View File

@@ -1,4 +1,7 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import type { Job } from 'bullmq';
import type { Db } from '../../../server/src/db/connection.js';
import type { JobData } from '../../src/job-processor.js';
@@ -8,23 +11,34 @@ function makeJob(data: JobData): Job<JobData> {
return { data } as unknown as Job<JobData>;
}
vi.mock('@waggle/agent', () => ({
runAgentLoop: vi.fn(async () => ({
content: 'Agent response here',
toolsUsed: ['search_memory'],
usage: { inputTokens: 100, outputTokens: 50 },
})),
createSystemTools: vi.fn(() => []),
}));
vi.mock('@waggle/agent', async (importOriginal) => {
const actual = await importOriginal<typeof import('@waggle/agent')>();
return {
...actual,
runAgentLoop: vi.fn(async () => ({
content: 'Agent response here',
toolsUsed: ['read_file'],
usage: { inputTokens: 100, outputTokens: 50 },
})),
};
});
import { chatHandler } from '../../src/handlers/chat-handler.js';
import { runAgentLoop, createSystemTools } from '@waggle/agent';
import { runAgentLoop } from '@waggle/agent';
describe('Chat Handler', () => {
const mockDb = {} as Db;
let dataDir: string;
beforeEach(() => {
vi.clearAllMocks();
dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-worker-chat-'));
vi.stubEnv('WAGGLE_DATA_DIR', dataDir);
});
afterEach(() => {
vi.unstubAllEnvs();
fs.rmSync(dataDir, { recursive: true, force: true });
});
it('calls runAgentLoop and returns response with correct shape', async () => {
@@ -41,7 +55,7 @@ describe('Chat Handler', () => {
expect(result.response).toBe('Agent response here');
expect(result.userId).toBe('u1');
expect(result.model).toBe('claude-sonnet');
expect(result.toolsUsed).toEqual(['search_memory']);
expect(result.toolsUsed).toEqual(['read_file']);
expect(result.tokensUsed).toBe(150);
});
@@ -61,21 +75,56 @@ describe('Chat Handler', () => {
expect(callArgs.model).toBe('gpt-4o');
expect(callArgs.messages).toEqual([{ role: 'user', content: 'test message' }]);
expect(callArgs.systemPrompt).toContain('Waggle AI agent');
expect(callArgs.tools).toEqual([]);
expect(callArgs.systemPrompt).toContain('read-only');
expect(callArgs.systemPrompt).toContain('cannot run shell commands, execute code, or write, edit, or delete files');
expect(callArgs.tools.map(tool => tool.name)).toEqual([
'read_file',
'search_files',
'search_content',
'web_search',
'web_fetch',
]);
});
it('creates system tools with workspaceDir from input', async () => {
it('ignores caller workspaceDir and roots tools in the team data directory', async () => {
const callerWorkspace = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-attacker-root-'));
const tenantRoot = path.join(dataDir, 'teams', 'team-1', 'files');
fs.mkdirSync(tenantRoot, { recursive: true });
fs.writeFileSync(path.join(tenantRoot, 'tenant.txt'), 'tenant-root-content');
fs.writeFileSync(path.join(callerWorkspace, 'tenant.txt'), 'caller-root-content');
const mockJob = makeJob({
jobId: 'j3',
teamId: 't1',
teamId: 'team-1',
userId: 'u1',
jobType: 'chat',
input: { message: 'test', workspaceDir: '/custom/workspace' },
input: { message: 'test', workspaceDir: callerWorkspace },
});
await chatHandler(mockJob, mockDb);
try {
await chatHandler(mockJob, mockDb);
expect(createSystemTools).toHaveBeenCalledWith('/custom/workspace');
const callArgs = vi.mocked(runAgentLoop).mock.calls[0][0];
const readFile = callArgs.tools.find(tool => tool.name === 'read_file');
expect(readFile).toBeDefined();
await expect(readFile!.execute({ path: 'tenant.txt' })).resolves.toContain('tenant-root-content');
} finally {
fs.rmSync(callerWorkspace, { recursive: true, force: true });
}
});
it('fails closed when WAGGLE_DATA_DIR is unavailable', async () => {
vi.stubEnv('WAGGLE_DATA_DIR', '');
const mockJob = makeJob({
jobId: 'j-no-root',
teamId: 'team-1',
userId: 'u1',
jobType: 'chat',
input: { message: 'test' },
});
await expect(chatHandler(mockJob, mockDb)).rejects.toThrow('WAGGLE_DATA_DIR');
expect(runAgentLoop).not.toHaveBeenCalled();
});
it('uses default model when not specified in input', async () => {

View File

@@ -1,4 +1,7 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import type { Job } from 'bullmq';
import type { Db } from '../../../server/src/db/connection.js';
import type { JobData } from '../../src/job-processor.js';
@@ -9,17 +12,17 @@ function makeJob(data: JobData): Job<JobData> {
}
// ── Mock @waggle/agent before importing handlers ────────────────────────
vi.mock('@waggle/agent', () => ({
runAgentLoop: vi.fn(async () => ({
content: 'Task completed successfully',
toolsUsed: ['bash', 'read_file'],
usage: { inputTokens: 200, outputTokens: 100 },
})),
createSystemTools: vi.fn(() => [
{ name: 'bash', description: 'Run shell commands', parameters: { type: 'object', properties: {} }, execute: async () => 'ok' },
{ name: 'read_file', description: 'Read a file', parameters: { type: 'object', properties: {} }, execute: async () => 'ok' },
]),
}));
vi.mock('@waggle/agent', async (importOriginal) => {
const actual = await importOriginal<typeof import('@waggle/agent')>();
return {
...actual,
runAgentLoop: vi.fn(async () => ({
content: 'Task completed successfully',
toolsUsed: ['read_file'],
usage: { inputTokens: 200, outputTokens: 100 },
})),
};
});
// ── Mock drizzle-orm ────────────────────────────────────────────────────
vi.mock('drizzle-orm', () => ({
@@ -51,6 +54,18 @@ import { taskHandler } from '../../src/handlers/task-handler.js';
import { groupHandler } from '../../src/handlers/group-handler.js';
import { runAgentLoop } from '@waggle/agent';
let dataDir: string;
beforeEach(() => {
dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-worker-handlers-'));
vi.stubEnv('WAGGLE_DATA_DIR', dataDir);
});
afterEach(() => {
vi.unstubAllEnvs();
fs.rmSync(dataDir, { recursive: true, force: true });
});
// ── Helper: mock DB ─────────────────────────────────────────────────────
function createMockDb(overrides?: {
@@ -119,7 +134,7 @@ describe('Task Handler', () => {
expect(result.taskTitle).toBe('Fix login bug');
expect(result.status).toBe('completed');
expect(result.result).toBe('Task completed successfully');
expect(result.toolsUsed).toEqual(['bash', 'read_file']);
expect(result.toolsUsed).toEqual(['read_file']);
expect(result.tokensUsed).toBe(300); // 200 + 100
expect(result.userId).toBe('user-1');
@@ -130,6 +145,101 @@ describe('Task Handler', () => {
expect(mockDb.update).toHaveBeenCalledTimes(2);
});
it('uses the trusted team root, read-only tools, and truthful task prompt', async () => {
const mockTask = {
id: 'task-policy',
title: 'Inspect login bug',
description: 'Explain the likely cause without changing files',
teamId: 'team-1',
priority: 'high',
status: 'open',
};
const tenantRoot = path.join(dataDir, 'teams', 'team-1', 'files');
const callerWorkspace = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-task-caller-'));
fs.mkdirSync(tenantRoot, { recursive: true });
fs.writeFileSync(path.join(tenantRoot, 'marker.txt'), 'tenant-root-content');
fs.writeFileSync(path.join(callerWorkspace, 'marker.txt'), 'caller-root-content');
const mockDb = createMockDb({ selectResult: [mockTask] });
const mockJob = makeJob({
jobId: 'j-policy',
teamId: 'team-1',
userId: 'user-1',
jobType: 'task',
input: { taskId: 'task-policy', model: 'gpt-4o', workspaceDir: callerWorkspace },
});
try {
await taskHandler(mockJob, mockDb);
const call = vi.mocked(runAgentLoop).mock.calls[0][0];
expect(call.model).toBe('gpt-4o');
expect(call.messages).toEqual([{
role: 'user',
content: 'Execute this task: Inspect login bug\n\nExplain the likely cause without changing files',
}]);
expect(call.systemPrompt).toContain('non-interactive read-only worker');
expect(call.systemPrompt).toContain('cannot run shell commands, execute code, or write, edit, or delete files');
expect(call.systemPrompt).toContain('Task: Inspect login bug');
expect(call.tools.map(tool => tool.name)).toEqual([
'read_file',
'search_files',
'search_content',
'web_search',
'web_fetch',
]);
const readFile = call.tools.find(tool => tool.name === 'read_file');
await expect(readFile!.execute({ path: 'marker.txt' })).resolves.toContain('tenant-root-content');
} finally {
fs.rmSync(callerWorkspace, { recursive: true, force: true });
}
});
it('fails closed before agent execution when worker data configuration is missing', async () => {
vi.stubEnv('WAGGLE_DATA_DIR', '');
const mockTask = {
id: 'task-no-root',
title: 'Inspect task',
description: null,
teamId: 'team-1',
priority: null,
status: 'open',
};
const mockDb = createMockDb({ selectResult: [mockTask] });
const mockJob = makeJob({
jobId: 'j-no-root',
teamId: 'team-1',
userId: 'user-1',
jobType: 'task',
input: { taskId: 'task-no-root' },
});
await expect(taskHandler(mockJob, mockDb)).rejects.toThrow('WAGGLE_DATA_DIR');
expect(runAgentLoop).not.toHaveBeenCalled();
});
it('fails closed before agent execution when the runtime teamId is unsafe', async () => {
const mockTask = {
id: 'task-bad-team',
title: 'Inspect task',
description: null,
teamId: '../other',
priority: null,
status: 'open',
};
const mockDb = createMockDb({ selectResult: [mockTask] });
const mockJob = makeJob({
jobId: 'j-bad-team',
teamId: '../other',
userId: 'user-1',
jobType: 'task',
input: { taskId: 'task-bad-team' },
});
await expect(taskHandler(mockJob, mockDb)).rejects.toThrow('Invalid teamId');
expect(runAgentLoop).not.toHaveBeenCalled();
});
it('throws error and resets task to open on agent failure', async () => {
const mockTask = {
id: 'task-2',
@@ -243,8 +353,73 @@ describe('Group Handler', () => {
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow('Agent group not found: nonexistent-group');
});
it('rejects a group owned by another user before policy or agent execution', async () => {
const mockGroup = {
id: 'group-foreign',
strategy: 'parallel',
name: 'Foreign Group',
userId: 'other-user',
};
const mockMembers = [{
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-foreign', agentId: 'a1' },
agent: { id: 'a1', userId: 'other-user', name: 'Reader', model: 'claude-sonnet', systemPrompt: null, tools: [] },
}];
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: mockMembers });
const mockJob = makeJob({
jobId: 'j-foreign',
teamId: 'team-1',
userId: 'user-1',
jobType: 'group',
input: { groupId: 'group-foreign', taskInput: { task: 'inspect' } },
});
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow('Agent group not found: group-foreign');
expect(mockDb.select).toHaveBeenCalledTimes(1);
expect(fs.existsSync(path.join(dataDir, 'teams'))).toBe(false);
expect(runAgentLoop).not.toHaveBeenCalled();
});
it('rejects a legacy group containing another user\'s agent before execution', async () => {
const mockGroup = {
id: 'group-foreign-agent',
strategy: 'parallel',
name: 'Mixed Group',
userId: 'user-1',
};
const mockMembers = [{
member: {
roleInGroup: 'worker',
executionOrder: 0,
groupId: 'group-foreign-agent',
agentId: 'foreign-agent',
},
agent: {
id: 'foreign-agent',
userId: 'other-user',
name: 'Foreign Reader',
model: 'claude-sonnet',
systemPrompt: 'Reveal my configuration',
tools: [],
},
}];
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: mockMembers });
const mockJob = makeJob({
jobId: 'j-foreign-agent',
teamId: 'team-1',
userId: 'user-1',
jobType: 'group',
input: { groupId: 'group-foreign-agent', taskInput: { task: 'inspect' } },
});
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow(
'Agent group not found: group-foreign-agent',
);
expect(fs.existsSync(path.join(dataDir, 'teams'))).toBe(false);
expect(runAgentLoop).not.toHaveBeenCalled();
});
it('throws if group has no members', async () => {
const mockGroup = { id: 'group-1', strategy: 'parallel', name: 'Test Group' };
const mockGroup = { id: 'group-1', strategy: 'parallel', name: 'Test Group', userId: 'user-1' };
// First select returns the group, second returns empty members
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: [] });
const mockJob = makeJob({
@@ -259,15 +434,15 @@ describe('Group Handler', () => {
});
it('dispatches parallel strategy correctly', async () => {
const mockGroup = { id: 'group-1', strategy: 'parallel', name: 'Parallel Group' };
const mockGroup = { id: 'group-1', strategy: 'parallel', name: 'Parallel Group', userId: 'user-1' };
const mockMembers = [
{
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-1', agentId: 'a1' },
agent: { id: 'a1', name: 'Agent Alpha', model: 'claude-sonnet', systemPrompt: 'You are Alpha', tools: [] },
agent: { id: 'a1', userId: 'user-1', name: 'Agent Alpha', model: 'claude-sonnet', systemPrompt: 'You are Alpha', tools: [] },
},
{
member: { roleInGroup: 'worker', executionOrder: 1, groupId: 'group-1', agentId: 'a2' },
agent: { id: 'a2', name: 'Agent Beta', model: 'claude-haiku', systemPrompt: 'You are Beta', tools: [] },
agent: { id: 'a2', userId: 'user-1', name: 'Agent Beta', model: 'claude-haiku', systemPrompt: 'You are Beta', tools: [] },
},
];
@@ -286,18 +461,117 @@ describe('Group Handler', () => {
expect(result.agentCount).toBe(2);
// runAgentLoop called once per member
expect(runAgentLoop).toHaveBeenCalledTimes(2);
const calls = vi.mocked(runAgentLoop).mock.calls.map(([config]) => config);
expect(calls[0].model).toBe('claude-sonnet');
expect(calls[0].messages).toEqual([{ role: 'user', content: 'analyze data' }]);
expect(calls[0].systemPrompt).toContain('non-interactive read-only worker');
expect(calls[0].systemPrompt).toContain('You are Alpha');
expect(calls[1].systemPrompt).toContain('You are Beta');
for (const call of calls) {
expect(call.tools.map(tool => tool.name)).toEqual([
'read_file',
'search_files',
'search_content',
'web_search',
'web_fetch',
]);
}
});
it('ignores taskInput.workspaceDir and filters requested mutating tools', async () => {
const mockGroup = { id: 'group-policy', strategy: 'parallel', name: 'Policy Group', userId: 'user-1' };
const mockMembers = [{
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-policy', agentId: 'a1' },
agent: {
id: 'a1',
userId: 'user-1',
name: 'Reader',
model: 'claude-sonnet',
systemPrompt: 'Inspect the workspace',
tools: ['bash', 'write_file', 'read_file'],
},
}];
const tenantRoot = path.join(dataDir, 'teams', 'team-1', 'files');
const callerWorkspace = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-group-caller-'));
fs.mkdirSync(tenantRoot, { recursive: true });
fs.writeFileSync(path.join(tenantRoot, 'marker.txt'), 'tenant-root-content');
fs.writeFileSync(path.join(callerWorkspace, 'marker.txt'), 'caller-root-content');
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: mockMembers });
const mockJob = makeJob({
jobId: 'j-group-policy',
teamId: 'team-1',
userId: 'user-1',
jobType: 'group',
input: {
groupId: 'group-policy',
taskInput: { task: 'inspect', workspaceDir: callerWorkspace },
},
});
try {
await groupHandler(mockJob, mockDb);
const call = vi.mocked(runAgentLoop).mock.calls[0][0];
expect(call.tools.map(tool => tool.name)).toEqual(['read_file']);
expect(call.systemPrompt).toContain('non-interactive read-only worker');
expect(call.systemPrompt).toContain('Inspect the workspace');
const readFile = call.tools.find(tool => tool.name === 'read_file');
await expect(readFile!.execute({ path: 'marker.txt' })).resolves.toContain('tenant-root-content');
} finally {
fs.rmSync(callerWorkspace, { recursive: true, force: true });
}
});
it('fails closed before group execution when worker data configuration is missing', async () => {
vi.stubEnv('WAGGLE_DATA_DIR', '');
const mockGroup = { id: 'group-no-root', strategy: 'parallel', name: 'Group', userId: 'user-1' };
const mockMembers = [{
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-no-root', agentId: 'a1' },
agent: { id: 'a1', userId: 'user-1', name: 'Reader', model: 'claude-sonnet', systemPrompt: null, tools: [] },
}];
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: mockMembers });
const mockJob = makeJob({
jobId: 'j-group-no-root',
teamId: 'team-1',
userId: 'user-1',
jobType: 'group',
input: { groupId: 'group-no-root', taskInput: { task: 'inspect' } },
});
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow('WAGGLE_DATA_DIR');
expect(runAgentLoop).not.toHaveBeenCalled();
});
it('fails closed before group execution when the runtime teamId is unsafe', async () => {
const mockGroup = { id: 'group-bad-team', strategy: 'parallel', name: 'Group', userId: 'user-1' };
const mockMembers = [{
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-bad-team', agentId: 'a1' },
agent: { id: 'a1', userId: 'user-1', name: 'Reader', model: 'claude-sonnet', systemPrompt: null, tools: [] },
}];
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: mockMembers });
const mockJob = makeJob({
jobId: 'j-group-bad-team',
teamId: '../other',
userId: 'user-1',
jobType: 'group',
input: { groupId: 'group-bad-team', taskInput: { task: 'inspect' } },
});
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow('Invalid teamId');
expect(runAgentLoop).not.toHaveBeenCalled();
});
it('dispatches sequential strategy correctly', async () => {
const mockGroup = { id: 'group-2', strategy: 'sequential', name: 'Sequential Group' };
const mockGroup = { id: 'group-2', strategy: 'sequential', name: 'Sequential Group', userId: 'user-1' };
const mockMembers = [
{
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-2', agentId: 'a1' },
agent: { id: 'a1', name: 'Researcher', model: 'claude-sonnet', systemPrompt: 'Research first', tools: [] },
agent: { id: 'a1', userId: 'user-1', name: 'Researcher', model: 'claude-sonnet', systemPrompt: 'Research first', tools: [] },
},
{
member: { roleInGroup: 'worker', executionOrder: 1, groupId: 'group-2', agentId: 'a2' },
agent: { id: 'a2', name: 'Writer', model: 'claude-haiku', systemPrompt: 'Write based on research', tools: [] },
agent: { id: 'a2', userId: 'user-1', name: 'Writer', model: 'claude-haiku', systemPrompt: 'Write based on research', tools: [] },
},
];
@@ -318,11 +592,11 @@ describe('Group Handler', () => {
});
it('throws for unknown execution strategy', async () => {
const mockGroup = { id: 'group-3', strategy: 'unknown_strategy', name: 'Bad Group' };
const mockGroup = { id: 'group-3', strategy: 'unknown_strategy', name: 'Bad Group', userId: 'user-1' };
const mockMembers = [
{
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-3', agentId: 'a1' },
agent: { id: 'a1', name: 'Agent', model: 'claude-sonnet', systemPrompt: null, tools: [] },
agent: { id: 'a1', userId: 'user-1', name: 'Agent', model: 'claude-sonnet', systemPrompt: null, tools: [] },
},
];