oleg.m/waggle-os
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

moving

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

11
packages/worker/package.json
View File

@@ -2,10 +2,17 @@
"name": "@waggle/worker",
"version": "0.1.0",
"type": "module",
"exports": { ".": "./src/index.ts" },
"main": "dist/index.js",
"types": "dist/src/index.d.ts",
"exports": {
".": {
"types": "./dist/src/index.d.ts",
"import": "./dist/index.js"
}
},
"scripts": {
"dev": "tsx watch src/index.ts",
"build": "tsc",
"build": "tsc --build --force && esbuild src/index.ts --bundle --platform=node --format=esm --packages=external --sourcemap --outfile=dist/index.js",
"test": "node ../../node_modules/vitest/vitest.mjs run --root ../.. --config vitest.config.ts packages/worker/tests",
"start": "node dist/index.js"
},

69
packages/worker/src/execution-policy.ts Normal file
View File

@@ -0,0 +1,69 @@
import fs from 'node:fs';
import path from 'node:path';
import {
createSystemTools,
PermissionManager,
type ToolDefinition,
} from '@waggle/agent';
const SAFE_TEAM_ID = /^[A-Za-z0-9](?:[A-Za-z0-9_-]{0,127})$/;
export const READ_ONLY_WORKER_SYSTEM_PROMPT = [
'You are a Waggle AI agent running as a non-interactive read-only worker.',
'You may inspect the tenant workspace with read_file, search_files, and search_content, and use web_search or web_fetch for research.',
'You cannot run shell commands, execute code, or write, edit, or delete files.',
'If a task requires changes, return a proposed patch or exact instructions for a human-approved interactive run.',
'Be helpful, concise, and proactive.',
].join('\n');
export interface WorkerExecutionContext {
workspaceDir: string;
tools: ToolDefinition[];
systemPrompt: string;
}
function isContained(parent: string, candidate: string): boolean {
const relative = path.relative(parent, candidate);
return relative === '' || (
!path.isAbsolute(relative)
&& relative !== '..'
&& !relative.startsWith(`..${path.sep}`)
);
}
function canonicalChild(parent: string, childName: string): string {
const child = path.join(parent, childName);
fs.mkdirSync(child, { recursive: true });
const canonical = fs.realpathSync.native(child);
if (!isContained(parent, canonical)) {
throw new Error(`Worker tenant path escapes WAGGLE_DATA_DIR: ${childName}`);
}
if (path.relative(child, canonical) !== '') {
throw new Error(`Worker tenant path must not traverse a link or junction: ${childName}`);
}
return canonical;
}
export function createWorkerExecutionContext(teamId: string): WorkerExecutionContext {
const configuredDataDir = process.env.WAGGLE_DATA_DIR;
if (!configuredDataDir?.trim()) {
throw new Error('WAGGLE_DATA_DIR is required for non-interactive worker execution');
}
if (typeof teamId !== 'string' || !SAFE_TEAM_ID.test(teamId)) {
throw new Error('Invalid teamId for non-interactive worker execution');
}
const dataDir = path.resolve(configuredDataDir);
fs.mkdirSync(dataDir, { recursive: true });
const canonicalDataDir = fs.realpathSync.native(dataDir);
const teamsDir = canonicalChild(canonicalDataDir, 'teams');
const teamDir = canonicalChild(teamsDir, teamId);
const workspaceDir = canonicalChild(teamDir, 'files');
const tools = PermissionManager.sandbox().filterTools(createSystemTools(workspaceDir));
return {
workspaceDir,
tools,
systemPrompt: READ_ONLY_WORKER_SYSTEM_PROMPT,
};
}

15
packages/worker/src/handlers/chat-handler.ts
View File

@@ -1,25 +1,18 @@
import type { Job } from 'bullmq';
import type { JobData } from '../job-processor.js';
import type { Db } from '../../../server/src/db/connection.js';
import { runAgentLoop, createSystemTools } from '@waggle/agent';
import { runAgentLoop } from '@waggle/agent';
import { createWorkerExecutionContext } from '../execution-policy.js';
const LITELLM_URL = process.env.LITELLM_URL ?? 'http://localhost:4000/v1';
const LITELLM_API_KEY = process.env.LITELLM_API_KEY ?? process.env.LITELLM_MASTER_KEY ?? 'sk-waggle-dev';
const DEFAULT_MODEL = process.env.DEFAULT_MODEL ?? 'claude-sonnet';
export async function chatHandler(job: Job<JobData>, _db: Db): Promise<Record<string, unknown>> {
const { userId, input } = job.data;
const { teamId, userId, input } = job.data;
const message = (input as Record<string, unknown>).message as string ?? '';
const model = (input as Record<string, unknown>).model as string ?? DEFAULT_MODEL;
const workspaceDir = (input as Record<string, unknown>).workspaceDir as string ?? process.cwd();
const tools = createSystemTools(workspaceDir);
const systemPrompt = [
'You are a Waggle AI agent. You help users with tasks using your tools.',
'Use system tools (bash, read_file, write_file, edit_file, search_files, search_content) to interact with the workspace.',
'Be helpful, concise, and proactive.',
].join('\n');
const { systemPrompt, tools } = createWorkerExecutionContext(teamId);
const result = await runAgentLoop({
litellmUrl: LITELLM_URL,

21
packages/worker/src/handlers/group-handler.ts
View File

@@ -6,13 +6,14 @@ import { eq } from 'drizzle-orm';
import { executeParallel, type ExecutionDeps } from '../execution/parallel.js';
import { executeSequential } from '../execution/sequential.js';
import { executeCoordinator } from '../execution/coordinator.js';
import { runAgentLoop, createSystemTools } from '@waggle/agent';
import { runAgentLoop } from '@waggle/agent';
import { createWorkerExecutionContext } from '../execution-policy.js';
const LITELLM_URL = process.env.LITELLM_URL ?? 'http://localhost:4000/v1';
const LITELLM_API_KEY = process.env.LITELLM_API_KEY ?? process.env.LITELLM_MASTER_KEY ?? 'sk-waggle-dev';
export async function groupHandler(job: Job<JobData>, db: Db): Promise<Record<string, unknown>> {
const { input } = job.data;
const { teamId, userId, input } = job.data;
const groupId = (input as Record<string, unknown>).groupId as string;
const taskInput = (input as Record<string, unknown>).taskInput as Record<string, unknown> ?? {};
@@ -24,7 +25,7 @@ export async function groupHandler(job: Job<JobData>, db: Db): Promise<Record<st
const [group] = await db.select().from(agentGroups)
.where(eq(agentGroups.id, groupId));
if (!group) {
if (!group || group.userId !== userId) {
throw new Error(`Agent group not found: ${groupId}`);
}
@@ -40,27 +41,29 @@ export async function groupHandler(job: Job<JobData>, db: Db): Promise<Record<st
if (members.length === 0) {
throw new Error(`Agent group ${groupId} has no members`);
}
if (members.some(({ agent }) => agent.userId !== group.userId)) {
throw new Error(`Agent group not found: ${groupId}`);
}
const executionContext = createWorkerExecutionContext(teamId);
// Sort by execution order
members.sort((a: typeof members[number], b: typeof members[number]) => a.member.executionOrder - b.member.executionOrder);
// Build execution deps — wires strategies to real runAgentLoop
const workspaceDir = (taskInput.workspaceDir as string) ?? process.cwd();
const allTools = createSystemTools(workspaceDir);
const deps: ExecutionDeps = {
runAgent: async (config) => runAgentLoop({
litellmUrl: LITELLM_URL,
litellmApiKey: LITELLM_API_KEY,
model: config.model,
systemPrompt: config.systemPrompt,
systemPrompt: `${executionContext.systemPrompt}\n\n${config.systemPrompt}`,
tools: config.tools,
messages: config.messages,
maxTurns: config.maxTurns ?? 10,
}),
resolveTools: (toolNames) => {
if (toolNames.length === 0) return allTools; // No filter = all tools
return allTools.filter(t => toolNames.includes(t.name));
if (toolNames.length === 0) return executionContext.tools; // No filter = all safe tools
return executionContext.tools.filter(t => toolNames.includes(t.name));
},
};

13
packages/worker/src/handlers/task-handler.ts
View File

@@ -8,7 +8,8 @@ import type { JobData } from '../job-processor.js';
import type { Db } from '../../../server/src/db/connection.js';
import { tasks } from '../../../server/src/db/schema.js';
import { eq } from 'drizzle-orm';
import { runAgentLoop, createSystemTools } from '@waggle/agent';
import { runAgentLoop } from '@waggle/agent';
import { createWorkerExecutionContext } from '../execution-policy.js';
const LITELLM_URL = process.env.LITELLM_URL ?? 'http://localhost:4000/v1';
const LITELLM_API_KEY = process.env.LITELLM_API_KEY ?? process.env.LITELLM_MASTER_KEY ?? 'sk-waggle-dev';
@@ -34,6 +35,8 @@ export async function taskHandler(job: Job<JobData>, db: Db): Promise<Record<str
throw new Error(`Task ${taskId} does not belong to team ${teamId}`);
}
const executionContext = createWorkerExecutionContext(teamId);
// Mark task as in-progress
await db.update(tasks)
.set({ status: 'in_progress', assignedTo: userId, updatedAt: new Date() })
@@ -41,27 +44,25 @@ export async function taskHandler(job: Job<JobData>, db: Db): Promise<Record<str
// Build prompt from task context
const systemPrompt = [
executionContext.systemPrompt,
'',
'You are a Waggle AI agent executing an assigned task.',
'Complete the task described below. Be thorough but concise.',
'Use system tools (bash, read_file, write_file, edit_file, search_files, search_content) to interact with the workspace.',
'',
`Task: ${task.title}`,
task.description ? `Details: ${task.description}` : '',
task.priority ? `Priority: ${task.priority}` : '',
].filter(Boolean).join('\n');
const workspaceDir = (input as Record<string, unknown>).workspaceDir as string ?? process.cwd();
const model = (input as Record<string, unknown>).model as string ?? DEFAULT_MODEL;
try {
const tools = createSystemTools(workspaceDir);
const result = await runAgentLoop({
litellmUrl: LITELLM_URL,
litellmApiKey: LITELLM_API_KEY,
model,
systemPrompt,
tools,
tools: executionContext.tools,
messages: [{ role: 'user', content: `Execute this task: ${task.title}${task.description ? '\n\n' + task.description : ''}` }],
});

29
packages/worker/src/index.ts
View File

@@ -1,4 +1,6 @@
import { Worker } from 'bullmq';
import { resolve } from 'node:path';
import { pathToFileURL } from 'node:url';
import { createDb } from '../../server/src/db/connection.js';
import { agentJobs } from '../../server/src/db/schema.js';
import { eq } from 'drizzle-orm';
@@ -91,9 +93,32 @@ export function createWorker(redisUrl = REDIS_URL, databaseUrl?: string, queueNa
}
// Start if run directly
const isDirectRun = process.argv[1]?.replace(/\\/g, '/').includes('worker/src/index');
if (isDirectRun) {
export function isDirectModule(entryPath: string | undefined, moduleUrl: string): boolean {
return entryPath !== undefined && pathToFileURL(resolve(entryPath)).href === moduleUrl;
}
type ShutdownSignalTarget = {
once(signal: 'SIGTERM' | 'SIGINT', listener: () => void): unknown;
};
export function installShutdownHandlers(
worker: Pick<Worker, 'close'>,
signalTarget: ShutdownSignalTarget = process,
): () => Promise<void> {
let shutdownPromise: Promise<void> | undefined;
const shutdown = () => {
shutdownPromise ??= worker.close();
return shutdownPromise;
};
signalTarget.once('SIGTERM', shutdown);
signalTarget.once('SIGINT', shutdown);
return shutdown;
}
if (isDirectModule(process.argv[1], import.meta.url)) {
const { worker } = createWorker();
installShutdownHandlers(worker);
console.log('Waggle agent worker started, waiting for jobs...');
worker.on('completed', (job) => {

72
packages/worker/tests/entrypoint.test.ts Normal file
View File

@@ -0,0 +1,72 @@
import { describe, expect, it, vi } from 'vitest';
import { readFileSync } from 'node:fs';
import { resolve } from 'node:path';
import { fileURLToPath, pathToFileURL } from 'node:url';
import { installShutdownHandlers, isDirectModule } from '../src/index.js';
const REPO_ROOT = fileURLToPath(new URL('../../../', import.meta.url));
describe('cloud worker production entrypoint', () => {
it('builds a standalone compiled worker after its server dependency', () => {
const rootPackage = JSON.parse(readFileSync(resolve(REPO_ROOT, 'package.json'), 'utf8')) as {
scripts: Record<string, string>;
};
const workerPackage = JSON.parse(
readFileSync(resolve(REPO_ROOT, 'packages/worker/package.json'), 'utf8'),
) as {
main?: string;
types?: string;
exports?: Record<string, { import?: string; types?: string }>;
scripts?: Record<string, string>;
};
const build = rootPackage.scripts['build:packages'];
expect(build.indexOf('../worker')).toBeGreaterThan(build.indexOf('../server'));
expect(build).toContain('cd ../worker && npm run build');
expect(workerPackage.scripts?.build).toContain('tsc --build --force');
expect(workerPackage.scripts?.build).toContain('--bundle');
expect(workerPackage.main).toBe('dist/index.js');
expect(workerPackage.types).toBe('dist/src/index.d.ts');
expect(workerPackage.exports?.['.']).toEqual({
types: './dist/src/index.d.ts',
import: './dist/index.js',
});
});
it('runs the compiled worker as a production Compose service', () => {
const compose = readFileSync(resolve(REPO_ROOT, 'docker-compose.production.yml'), 'utf8');
expect(compose).toMatch(/\n {2}worker:\r?\n/);
expect(compose).toContain("command: ['node', 'packages/worker/dist/index.js']");
expect(compose).toContain('DATABASE_URL=postgres://waggle:${POSTGRES_PASSWORD:?set POSTGRES_PASSWORD}@postgres:5432/waggle');
expect(compose).toContain('REDIS_URL=redis://redis:6379');
});
it('recognizes the compiled entrypoint by URL identity', () => {
const distPath = resolve(REPO_ROOT, 'packages/worker/dist/index.js');
expect(isDirectModule(distPath, pathToFileURL(distPath).href)).toBe(true);
expect(isDirectModule(distPath, pathToFileURL(resolve(REPO_ROOT, 'consumer.js')).href)).toBe(false);
expect(isDirectModule(undefined, pathToFileURL(distPath).href)).toBe(false);
});
it('drains the worker once across repeated shutdown signals', async () => {
const listeners = new Map<string, () => void>();
let releaseClose!: () => void;
const close = vi.fn(() => new Promise<void>(resolve => { releaseClose = resolve; }));
const shutdown = installShutdownHandlers(
{ close },
{ once: (signal, listener) => { listeners.set(signal, listener); } },
);
expect([...listeners.keys()]).toEqual(['SIGTERM', 'SIGINT']);
listeners.get('SIGTERM')!();
listeners.get('SIGINT')!();
const draining = shutdown();
expect(shutdown()).toBe(draining);
expect(close).toHaveBeenCalledTimes(1);
releaseClose();
await draining;
});
});

92
packages/worker/tests/execution-policy.test.ts Normal file
View File

@@ -0,0 +1,92 @@
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import { createWorkerExecutionContext } from '../src/execution-policy.js';
describe('worker execution policy', () => {
let dataDir: string;
beforeEach(() => {
dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-worker-policy-'));
vi.stubEnv('WAGGLE_DATA_DIR', dataDir);
});
afterEach(() => {
vi.unstubAllEnvs();
fs.rmSync(dataDir, { recursive: true, force: true });
});
it('derives a canonical tenant workspace from WAGGLE_DATA_DIR and teamId', () => {
const context = createWorkerExecutionContext('team-123');
const expected = fs.realpathSync.native(path.join(dataDir, 'teams', 'team-123', 'files'));
expect(context.workspaceDir).toBe(expected);
expect(fs.statSync(context.workspaceDir).isDirectory()).toBe(true);
});
it('rejects a team directory link that aliases another tenant', () => {
const teamsDir = path.join(dataDir, 'teams');
const otherTeamDir = path.join(teamsDir, 'team-other');
fs.mkdirSync(path.join(otherTeamDir, 'files'), { recursive: true });
fs.symlinkSync(
otherTeamDir,
path.join(teamsDir, 'team-requested'),
process.platform === 'win32' ? 'junction' : 'dir',
);
expect(() => createWorkerExecutionContext('team-requested')).toThrow('link or junction');
});
it('rejects a files directory link that widens the tenant root', () => {
const teamDir = path.join(dataDir, 'teams', 'team-123');
fs.mkdirSync(teamDir, { recursive: true });
fs.symlinkSync(
teamDir,
path.join(teamDir, 'files'),
process.platform === 'win32' ? 'junction' : 'dir',
);
expect(() => createWorkerExecutionContext('team-123')).toThrow('link or junction');
});
it('exposes exactly the read-only system tool pool', () => {
const context = createWorkerExecutionContext('team-123');
expect(context.tools.map(tool => tool.name)).toEqual([
'read_file',
'search_files',
'search_content',
'web_search',
'web_fetch',
]);
});
it('truthfully instructs the model to propose changes instead of mutating', () => {
const context = createWorkerExecutionContext('team-123');
expect(context.systemPrompt).toContain('non-interactive read-only worker');
expect(context.systemPrompt).toContain('cannot run shell commands, execute code, or write, edit, or delete files');
expect(context.systemPrompt).toContain('proposed patch');
});
it('fails closed without WAGGLE_DATA_DIR', () => {
vi.stubEnv('WAGGLE_DATA_DIR', ' ');
expect(() => createWorkerExecutionContext('team-123')).toThrow('WAGGLE_DATA_DIR');
});
it.each([
['empty', ''],
['current directory', '.'],
['parent directory', '..'],
['forward-slash traversal', '../outside'],
['backslash traversal', '..\\outside'],
['absolute Windows path', 'C:\\outside'],
['overlong identifier', 'a'.repeat(129)],
['missing runtime identifier', undefined as unknown as string],
['null runtime identifier', null as unknown as string],
])('rejects an unsafe %s teamId', (_label, teamId) => {
expect(() => createWorkerExecutionContext(teamId)).toThrow('Invalid teamId');
});
});

83
packages/worker/tests/handlers/chat-handler.test.ts
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 () => {

318
packages/worker/tests/handlers/handlers.test.ts
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: [] },
},
];