This commit is contained in:
21
packages/worker/package.json
Normal file
21
packages/worker/package.json
Normal file
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"name": "@waggle/worker",
|
||||
"version": "0.1.0",
|
||||
"type": "module",
|
||||
"exports": { ".": "./src/index.ts" },
|
||||
"scripts": {
|
||||
"dev": "tsx watch src/index.ts",
|
||||
"build": "tsc",
|
||||
"test": "node ../../node_modules/vitest/vitest.mjs run --root ../.. --config vitest.config.ts packages/worker/tests",
|
||||
"start": "node dist/index.js"
|
||||
},
|
||||
"dependencies": {
|
||||
"@waggle/shared": "*",
|
||||
"@waggle/core": "*",
|
||||
"@waggle/agent": "*",
|
||||
"bullmq": "^5.0.0",
|
||||
"drizzle-orm": "^0.45.2",
|
||||
"ioredis": "^5.4.0",
|
||||
"postgres": "^3.4.0"
|
||||
}
|
||||
}
|
||||
150
packages/worker/src/execution/coordinator.ts
Normal file
150
packages/worker/src/execution/coordinator.ts
Normal file
@@ -0,0 +1,150 @@
|
||||
import type { AgentMemberConfig, AgentResult, ExecutionDeps } from './parallel.js';
|
||||
|
||||
/**
|
||||
* Coordinator execution strategy: a lead agent plans and delegates,
|
||||
* worker agents execute subtasks, then the lead synthesizes results.
|
||||
* Best for complex tasks requiring decomposition and integration.
|
||||
*/
|
||||
export async function executeCoordinator(
|
||||
members: AgentMemberConfig[],
|
||||
taskInput: Record<string, unknown>,
|
||||
deps?: ExecutionDeps,
|
||||
): Promise<Record<string, unknown>> {
|
||||
const lead = members.find(m => m.member.roleInGroup === 'lead');
|
||||
const workers = members.filter(m => m.member.roleInGroup === 'worker');
|
||||
|
||||
if (!lead) {
|
||||
throw new Error('Coordinator strategy requires a lead agent');
|
||||
}
|
||||
|
||||
const task = (taskInput.task as string) ?? JSON.stringify(taskInput);
|
||||
|
||||
// ── Stub mode (backward compat) ──────────────────────────────────
|
||||
if (!deps) {
|
||||
const plan = {
|
||||
agentId: lead.agent.id,
|
||||
agentName: lead.agent.name,
|
||||
phase: 'planning',
|
||||
output: `[Stub] ${lead.agent.name} delegated to ${workers.length} workers`,
|
||||
subtasks: workers.map(w => ({
|
||||
assignedTo: w.agent.name,
|
||||
description: `[Stub] Subtask for ${w.agent.name}`,
|
||||
})),
|
||||
};
|
||||
|
||||
const workerResults = await Promise.all(
|
||||
workers.map(async ({ agent }) => ({
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
phase: 'execution',
|
||||
model: agent.model,
|
||||
output: `[Stub] ${agent.name} completed delegated work`,
|
||||
})),
|
||||
);
|
||||
|
||||
const synthesis = {
|
||||
agentId: lead.agent.id,
|
||||
agentName: lead.agent.name,
|
||||
phase: 'synthesis',
|
||||
output: `[Stub] ${lead.agent.name} synthesized ${workerResults.length} worker outputs`,
|
||||
};
|
||||
|
||||
return {
|
||||
strategy: 'coordinator',
|
||||
leadAgent: lead.agent.name,
|
||||
workerCount: workers.length,
|
||||
plan,
|
||||
workerResults,
|
||||
synthesis,
|
||||
finalOutput: synthesis.output,
|
||||
};
|
||||
}
|
||||
|
||||
// ── Real execution: 3-phase pattern ──────────────────────────────
|
||||
|
||||
// Phase 1: Lead agent creates plan
|
||||
const leadPrompt = lead.agent.systemPrompt ?? `You are ${lead.agent.name}, a lead coordinator.`;
|
||||
const planResult = await deps.runAgent({
|
||||
model: lead.agent.model,
|
||||
systemPrompt: `${leadPrompt}\n\nYou are the lead coordinator. Analyze the task and create a clear plan for ${workers.length} worker agent(s). Output a structured plan with one subtask per worker.`,
|
||||
tools: deps.resolveTools(lead.agent.tools),
|
||||
messages: [{ role: 'user', content: task }],
|
||||
});
|
||||
|
||||
const plan = {
|
||||
agentId: lead.agent.id,
|
||||
agentName: lead.agent.name,
|
||||
phase: 'planning',
|
||||
output: planResult.content,
|
||||
usage: planResult.usage,
|
||||
};
|
||||
|
||||
// Phase 2: Workers execute in parallel, receiving the plan as context
|
||||
const workerResults: AgentResult[] = await Promise.all(
|
||||
workers.map(async ({ agent, member }) => {
|
||||
try {
|
||||
const workerPrompt = agent.systemPrompt ?? `You are ${agent.name}. Execute your assigned subtask.`;
|
||||
const result = await deps.runAgent({
|
||||
model: agent.model,
|
||||
systemPrompt: `${workerPrompt}\n\n## Coordinator's Plan\n${planResult.content}`,
|
||||
tools: deps.resolveTools(agent.tools),
|
||||
messages: [{ role: 'user', content: task }],
|
||||
});
|
||||
|
||||
return {
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
role: member.roleInGroup,
|
||||
model: agent.model,
|
||||
output: result.content,
|
||||
toolsUsed: result.toolsUsed,
|
||||
usage: result.usage,
|
||||
};
|
||||
} catch (err: unknown) {
|
||||
return {
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
role: member.roleInGroup,
|
||||
model: agent.model,
|
||||
output: '',
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
};
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
// Phase 3: Lead synthesizes worker results (include failure notices)
|
||||
const successOutputs = workerResults
|
||||
.filter(r => !r.error)
|
||||
.map(r => `## ${r.agentName}\n${r.output}`);
|
||||
const failureNotices = workerResults
|
||||
.filter(r => r.error)
|
||||
.map(r => `## ${r.agentName} (FAILED)\nError: ${r.error}`);
|
||||
const workerOutputs = [...successOutputs, ...failureNotices].join('\n\n');
|
||||
|
||||
const synthesisResult = await deps.runAgent({
|
||||
model: lead.agent.model,
|
||||
systemPrompt: `${leadPrompt}\n\nSynthesize the following worker outputs into a comprehensive final result.`,
|
||||
tools: deps.resolveTools(lead.agent.tools),
|
||||
messages: [{ role: 'user', content: workerOutputs || 'No worker outputs available.' }],
|
||||
});
|
||||
|
||||
const synthesis = {
|
||||
agentId: lead.agent.id,
|
||||
agentName: lead.agent.name,
|
||||
phase: 'synthesis',
|
||||
output: synthesisResult.content,
|
||||
usage: synthesisResult.usage,
|
||||
};
|
||||
|
||||
return {
|
||||
strategy: 'coordinator',
|
||||
leadAgent: lead.agent.name,
|
||||
workerCount: workers.length,
|
||||
plan,
|
||||
workerResults,
|
||||
synthesis,
|
||||
finalOutput: synthesis.output,
|
||||
errors: workerResults.filter(r => r.error).map(r => ({ agent: r.agentName, error: r.error })),
|
||||
};
|
||||
}
|
||||
99
packages/worker/src/execution/parallel.ts
Normal file
99
packages/worker/src/execution/parallel.ts
Normal file
@@ -0,0 +1,99 @@
|
||||
import type { ToolDefinition } from '@waggle/agent';
|
||||
|
||||
export interface AgentMemberConfig {
|
||||
member: { roleInGroup: string; executionOrder: number };
|
||||
agent: { id: string; name: string; model: string; systemPrompt: string | null; tools: string[] };
|
||||
}
|
||||
|
||||
/** Dependencies injected for real execution (omit for stub mode) */
|
||||
export interface ExecutionDeps {
|
||||
runAgent: (config: {
|
||||
model: string;
|
||||
systemPrompt: string;
|
||||
tools: ToolDefinition[];
|
||||
messages: Array<{ role: string; content: string }>;
|
||||
maxTurns?: number;
|
||||
}) => Promise<{ content: string; toolsUsed: string[]; usage: { inputTokens: number; outputTokens: number } }>;
|
||||
/** Resolve tool names to ToolDefinition[] from available tools */
|
||||
resolveTools: (toolNames: string[]) => ToolDefinition[];
|
||||
}
|
||||
|
||||
export interface AgentResult {
|
||||
agentId: string;
|
||||
agentName: string;
|
||||
role: string;
|
||||
model: string;
|
||||
output: string;
|
||||
toolsUsed?: string[];
|
||||
usage?: { inputTokens: number; outputTokens: number };
|
||||
error?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Parallel execution strategy: all agents run concurrently, results are merged.
|
||||
* Best for independent subtasks that don't depend on each other.
|
||||
*/
|
||||
export async function executeParallel(
|
||||
members: AgentMemberConfig[],
|
||||
taskInput: Record<string, unknown>,
|
||||
deps?: ExecutionDeps,
|
||||
): Promise<Record<string, unknown>> {
|
||||
const task = (taskInput.task as string) ?? JSON.stringify(taskInput);
|
||||
|
||||
const results: AgentResult[] = await Promise.all(
|
||||
members.map(async ({ agent, member }) => {
|
||||
if (!deps) {
|
||||
// Stub mode (backward compat for existing tests without deps)
|
||||
return {
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
role: member.roleInGroup,
|
||||
model: agent.model,
|
||||
output: `[Stub] ${agent.name} processed task with model ${agent.model}`,
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
const tools = deps.resolveTools(agent.tools);
|
||||
const systemPrompt = agent.systemPrompt ?? `You are ${agent.name}. Complete the assigned task thoroughly.`;
|
||||
|
||||
const result = await deps.runAgent({
|
||||
model: agent.model,
|
||||
systemPrompt,
|
||||
tools,
|
||||
messages: [{ role: 'user', content: task }],
|
||||
});
|
||||
|
||||
return {
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
role: member.roleInGroup,
|
||||
model: agent.model,
|
||||
output: result.content,
|
||||
toolsUsed: result.toolsUsed,
|
||||
usage: result.usage,
|
||||
};
|
||||
} catch (err: unknown) {
|
||||
return {
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
role: member.roleInGroup,
|
||||
model: agent.model,
|
||||
output: '',
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
};
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
return {
|
||||
strategy: 'parallel',
|
||||
agentCount: results.length,
|
||||
results,
|
||||
mergedOutput: results
|
||||
.filter(r => !r.error)
|
||||
.map(r => r.output)
|
||||
.join('\n\n---\n\n'),
|
||||
errors: results.filter(r => r.error).map(r => ({ agent: r.agentName, error: r.error })),
|
||||
};
|
||||
}
|
||||
85
packages/worker/src/execution/sequential.ts
Normal file
85
packages/worker/src/execution/sequential.ts
Normal file
@@ -0,0 +1,85 @@
|
||||
import type { AgentMemberConfig, AgentResult, ExecutionDeps } from './parallel.js';
|
||||
|
||||
/** Stub-mode result carries an extra `inputFrom` marker for backward-compat tests. */
|
||||
type StubAgentResult = AgentResult & { inputFrom: string };
|
||||
|
||||
/**
|
||||
* Sequential execution strategy: agents run one after another,
|
||||
* each receiving the previous agent's output as context.
|
||||
* Best for pipeline-style workflows (research -> draft -> review).
|
||||
*/
|
||||
export async function executeSequential(
|
||||
members: AgentMemberConfig[],
|
||||
taskInput: Record<string, unknown>,
|
||||
deps?: ExecutionDeps,
|
||||
): Promise<Record<string, unknown>> {
|
||||
const task = (taskInput.task as string) ?? JSON.stringify(taskInput);
|
||||
const results: AgentResult[] = [];
|
||||
let previousOutput: string | null = null;
|
||||
|
||||
for (const { agent, member } of members) {
|
||||
if (!deps) {
|
||||
// Stub mode (backward compat)
|
||||
const output: StubAgentResult = {
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
role: member.roleInGroup,
|
||||
model: agent.model,
|
||||
output: `[Stub] ${agent.name} processed with input from previous step`,
|
||||
// Preserve inputFrom for backward compat with existing tests
|
||||
inputFrom: previousOutput === null ? 'taskInput' : (results[results.length - 1]?.agentName ?? 'taskInput'),
|
||||
};
|
||||
results.push(output);
|
||||
previousOutput = output.output;
|
||||
continue;
|
||||
}
|
||||
|
||||
try {
|
||||
const tools = deps.resolveTools(agent.tools);
|
||||
const basePrompt = agent.systemPrompt ?? `You are ${agent.name}. Complete the assigned task.`;
|
||||
|
||||
// Chain: inject previous agent's output as context
|
||||
const systemPrompt = previousOutput
|
||||
? `${basePrompt}\n\n## Previous Agent's Output\n${previousOutput}`
|
||||
: basePrompt;
|
||||
|
||||
const result = await deps.runAgent({
|
||||
model: agent.model,
|
||||
systemPrompt,
|
||||
tools,
|
||||
messages: [{ role: 'user', content: task }],
|
||||
});
|
||||
|
||||
const agentResult: AgentResult = {
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
role: member.roleInGroup,
|
||||
model: agent.model,
|
||||
output: result.content,
|
||||
toolsUsed: result.toolsUsed,
|
||||
usage: result.usage,
|
||||
};
|
||||
results.push(agentResult);
|
||||
previousOutput = result.content;
|
||||
} catch (err: unknown) {
|
||||
// Error in chain stops execution — subsequent agents depend on this output
|
||||
results.push({
|
||||
agentId: agent.id,
|
||||
agentName: agent.name,
|
||||
role: member.roleInGroup,
|
||||
model: agent.model,
|
||||
output: '',
|
||||
error: err instanceof Error ? err.message : String(err),
|
||||
});
|
||||
break; // Stop chain on failure
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
strategy: 'sequential',
|
||||
agentCount: results.length,
|
||||
results,
|
||||
finalOutput: results.filter(r => !r.error).at(-1)?.output ?? null,
|
||||
chainBroken: results.some(r => r.error),
|
||||
};
|
||||
}
|
||||
40
packages/worker/src/handlers/chat-handler.ts
Normal file
40
packages/worker/src/handlers/chat-handler.ts
Normal file
@@ -0,0 +1,40 @@
|
||||
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';
|
||||
|
||||
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 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 result = await runAgentLoop({
|
||||
litellmUrl: LITELLM_URL,
|
||||
litellmApiKey: LITELLM_API_KEY,
|
||||
model,
|
||||
systemPrompt,
|
||||
tools,
|
||||
messages: [{ role: 'user', content: message }],
|
||||
});
|
||||
|
||||
return {
|
||||
response: result.content,
|
||||
userId,
|
||||
model,
|
||||
toolsUsed: result.toolsUsed,
|
||||
tokensUsed: result.usage.inputTokens + result.usage.outputTokens,
|
||||
};
|
||||
}
|
||||
53
packages/worker/src/handlers/cron-handler.ts
Normal file
53
packages/worker/src/handlers/cron-handler.ts
Normal file
@@ -0,0 +1,53 @@
|
||||
import type { Job } from 'bullmq';
|
||||
import type { Db } from '../../../server/src/db/connection.js';
|
||||
import type { JobData, JobHandler } from '../job-processor.js';
|
||||
|
||||
const DELEGATABLE_JOB_TYPES = new Set(['chat', 'task', 'waggle', 'group']);
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === 'object' && value !== null && !Array.isArray(value);
|
||||
}
|
||||
|
||||
export type CronDispatch = (job: Job<JobData>, db: Db) => Promise<Record<string, unknown>>;
|
||||
|
||||
/**
|
||||
* Execute a scheduled wrapper job through one of the worker's real handlers.
|
||||
* The explicit allowlist prevents cron jobs from recursively dispatching cron
|
||||
* or reaching an unregistered handler by passing arbitrary job types.
|
||||
*/
|
||||
export function createCronHandler(dispatch: CronDispatch): JobHandler {
|
||||
return async (job, db) => {
|
||||
const input = job.data.input;
|
||||
const delegatedJobType = typeof input.jobType === 'string' ? input.jobType.trim() : '';
|
||||
|
||||
if (!delegatedJobType) {
|
||||
throw new Error('cron job requires input.jobType (chat, task, waggle, or group)');
|
||||
}
|
||||
|
||||
if (!DELEGATABLE_JOB_TYPES.has(delegatedJobType)) {
|
||||
throw new Error(
|
||||
`cron job cannot delegate to "${delegatedJobType}"; supported types: ${[...DELEGATABLE_JOB_TYPES].join(', ')}`,
|
||||
);
|
||||
}
|
||||
|
||||
const delegatedInput = input.jobConfig ?? input.input;
|
||||
if (!isRecord(delegatedInput)) {
|
||||
throw new Error('cron job requires input.jobConfig or input.input as an object');
|
||||
}
|
||||
|
||||
const delegatedJob = {
|
||||
...job,
|
||||
data: {
|
||||
...job.data,
|
||||
jobType: delegatedJobType,
|
||||
input: delegatedInput,
|
||||
},
|
||||
} as Job<JobData>;
|
||||
|
||||
const result = await dispatch(delegatedJob, db);
|
||||
return {
|
||||
...result,
|
||||
cron: { delegatedJobType },
|
||||
};
|
||||
};
|
||||
}
|
||||
78
packages/worker/src/handlers/group-handler.ts
Normal file
78
packages/worker/src/handlers/group-handler.ts
Normal file
@@ -0,0 +1,78 @@
|
||||
import type { Job } from 'bullmq';
|
||||
import type { JobData } from '../job-processor.js';
|
||||
import type { Db } from '../../../server/src/db/connection.js';
|
||||
import { agentGroups, agentGroupMembers, agents } from '../../../server/src/db/schema.js';
|
||||
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';
|
||||
|
||||
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 groupId = (input as Record<string, unknown>).groupId as string;
|
||||
const taskInput = (input as Record<string, unknown>).taskInput as Record<string, unknown> ?? {};
|
||||
|
||||
if (!groupId) {
|
||||
throw new Error('groupHandler requires input.groupId');
|
||||
}
|
||||
|
||||
// Load group config
|
||||
const [group] = await db.select().from(agentGroups)
|
||||
.where(eq(agentGroups.id, groupId));
|
||||
|
||||
if (!group) {
|
||||
throw new Error(`Agent group not found: ${groupId}`);
|
||||
}
|
||||
|
||||
// Load members with their agent configs
|
||||
const members = await db.select({
|
||||
member: agentGroupMembers,
|
||||
agent: agents,
|
||||
})
|
||||
.from(agentGroupMembers)
|
||||
.innerJoin(agents, eq(agentGroupMembers.agentId, agents.id))
|
||||
.where(eq(agentGroupMembers.groupId, groupId));
|
||||
|
||||
if (members.length === 0) {
|
||||
throw new Error(`Agent group ${groupId} has no members`);
|
||||
}
|
||||
|
||||
// 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,
|
||||
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));
|
||||
},
|
||||
};
|
||||
|
||||
// Dispatch to strategy with real execution deps
|
||||
switch (group.strategy) {
|
||||
case 'parallel':
|
||||
return executeParallel(members, taskInput, deps);
|
||||
case 'sequential':
|
||||
return executeSequential(members, taskInput, deps);
|
||||
case 'coordinator':
|
||||
return executeCoordinator(members, taskInput, deps);
|
||||
default:
|
||||
throw new Error(`Unknown execution strategy: ${group.strategy}`);
|
||||
}
|
||||
}
|
||||
90
packages/worker/src/handlers/task-handler.ts
Normal file
90
packages/worker/src/handlers/task-handler.ts
Normal file
@@ -0,0 +1,90 @@
|
||||
/**
|
||||
* Task handler — runs agent loop with task context.
|
||||
* Replaces the stub implementation with real agent invocation.
|
||||
*/
|
||||
|
||||
import type { Job } from 'bullmq';
|
||||
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';
|
||||
|
||||
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 taskHandler(job: Job<JobData>, db: Db): Promise<Record<string, unknown>> {
|
||||
const { teamId, userId, input } = job.data;
|
||||
const taskId = (input as Record<string, unknown>).taskId as string | undefined;
|
||||
|
||||
if (!taskId) {
|
||||
throw new Error('taskHandler requires input.taskId');
|
||||
}
|
||||
|
||||
// Load task from DB to verify it exists and belongs to the team
|
||||
const [task] = await db.select().from(tasks)
|
||||
.where(eq(tasks.id, taskId));
|
||||
|
||||
if (!task) {
|
||||
throw new Error(`Task not found: ${taskId}`);
|
||||
}
|
||||
|
||||
if (task.teamId !== teamId) {
|
||||
throw new Error(`Task ${taskId} does not belong to team ${teamId}`);
|
||||
}
|
||||
|
||||
// Mark task as in-progress
|
||||
await db.update(tasks)
|
||||
.set({ status: 'in_progress', assignedTo: userId, updatedAt: new Date() })
|
||||
.where(eq(tasks.id, taskId));
|
||||
|
||||
// Build prompt from task context
|
||||
const 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,
|
||||
messages: [{ role: 'user', content: `Execute this task: ${task.title}${task.description ? '\n\n' + task.description : ''}` }],
|
||||
});
|
||||
|
||||
// Mark task as completed
|
||||
await db.update(tasks)
|
||||
.set({ status: 'completed', updatedAt: new Date() })
|
||||
.where(eq(tasks.id, taskId));
|
||||
|
||||
return {
|
||||
taskId,
|
||||
taskTitle: task.title,
|
||||
status: 'completed',
|
||||
result: result.content,
|
||||
toolsUsed: result.toolsUsed,
|
||||
tokensUsed: result.usage.inputTokens + result.usage.outputTokens,
|
||||
userId,
|
||||
};
|
||||
} catch (err) {
|
||||
// Reset task to open on failure (not stuck in in_progress)
|
||||
await db.update(tasks)
|
||||
.set({ status: 'open', updatedAt: new Date() })
|
||||
.where(eq(tasks.id, taskId));
|
||||
|
||||
throw err; // Let the worker mark the job as failed
|
||||
}
|
||||
}
|
||||
111
packages/worker/src/handlers/waggle-handler.ts
Normal file
111
packages/worker/src/handlers/waggle-handler.ts
Normal file
@@ -0,0 +1,111 @@
|
||||
/**
|
||||
* Waggle handler — routes Waggle Dance protocol messages through the dispatcher,
|
||||
* with fallback to hive query for legacy topic-based calls.
|
||||
*/
|
||||
|
||||
import type { Job } from 'bullmq';
|
||||
import type { JobData } from '../job-processor.js';
|
||||
import type { Db } from '../../../server/src/db/connection.js';
|
||||
import { teamEntities, tasks } from '../../../server/src/db/schema.js';
|
||||
import { sql } from 'drizzle-orm';
|
||||
import { WaggleDanceDispatcher } from '@waggle/waggle-dance';
|
||||
import type { WaggleMessage } from '@waggle/shared';
|
||||
import { CapabilityRouter, createSystemTools } from '@waggle/agent';
|
||||
|
||||
export interface WaggleHandlerDeps {
|
||||
enqueueWorker?: (input: {
|
||||
teamId: string;
|
||||
userId: string;
|
||||
task: string;
|
||||
role: string;
|
||||
context?: string;
|
||||
}) => Promise<string>;
|
||||
}
|
||||
|
||||
export function createWaggleHandler(deps: WaggleHandlerDeps = {}) {
|
||||
const capabilityRouter = new CapabilityRouter({
|
||||
toolNames: createSystemTools(process.cwd()).map(tool => tool.name),
|
||||
skills: [],
|
||||
plugins: [],
|
||||
mcpServers: [],
|
||||
subAgentRoles: ['researcher', 'writer', 'coder', 'analyst', 'reviewer', 'planner'],
|
||||
connectors: [],
|
||||
});
|
||||
|
||||
return async function waggleHandler(job: Job<JobData>, db: Db): Promise<Record<string, unknown>> {
|
||||
const { teamId, userId, input } = job.data;
|
||||
const inputObj = input as Record<string, unknown>;
|
||||
|
||||
// ── Protocol message routing (type + subtype present) ────────────
|
||||
if (inputObj.type && inputObj.subtype) {
|
||||
const dispatcher = new WaggleDanceDispatcher({
|
||||
searchMemory: async (query: string) => {
|
||||
const searchTerm = `%${query}%`;
|
||||
const entities = await db.select().from(teamEntities)
|
||||
.where(sql`${teamEntities.teamId} = ${teamId} AND ${teamEntities.name} ILIKE ${searchTerm}`)
|
||||
.limit(10);
|
||||
return entities.map((e: typeof entities[number]) => `[${e.entityType}] ${e.name}`).join('\n') || 'No matching knowledge found.';
|
||||
},
|
||||
resolveCapability: (query: string) => {
|
||||
return capabilityRouter.resolve(query).map(route => ({
|
||||
source: route.source,
|
||||
name: route.name,
|
||||
description: route.description,
|
||||
available: route.available,
|
||||
}));
|
||||
},
|
||||
spawnWorker: async (task: string, role: string, context?: string) => {
|
||||
if (!deps.enqueueWorker) {
|
||||
throw new Error('Worker delegation is unavailable in this runtime');
|
||||
}
|
||||
|
||||
const childJobId = await deps.enqueueWorker({ teamId, userId, task, role, ...(context ? { context } : {}) });
|
||||
return `Worker job ${childJobId} queued for: ${task} (${role})`;
|
||||
},
|
||||
});
|
||||
|
||||
// `inputObj` is untrusted job-data; the dispatcher revalidates the
|
||||
// type/subtype combo at runtime before acting on it.
|
||||
const result = await dispatcher.dispatch(inputObj as unknown as WaggleMessage);
|
||||
return {
|
||||
dispatched: true,
|
||||
type: inputObj.type,
|
||||
subtype: inputObj.subtype,
|
||||
...result,
|
||||
};
|
||||
}
|
||||
|
||||
// ── Legacy fallback: topic-based hive query ──────────────────────
|
||||
const topic = (inputObj.topic as string) ?? '';
|
||||
const searchTerm = `%${topic}%`;
|
||||
|
||||
// Query existing team knowledge
|
||||
const entities = await db.select().from(teamEntities)
|
||||
.where(sql`${teamEntities.teamId} = ${teamId} AND ${teamEntities.name} ILIKE ${searchTerm}`)
|
||||
.limit(5);
|
||||
|
||||
// Query related tasks
|
||||
const relatedTasks = await db.select().from(tasks)
|
||||
.where(sql`${tasks.teamId} = ${teamId} AND ${tasks.title} ILIKE ${searchTerm}`)
|
||||
.limit(5);
|
||||
|
||||
const gaps = [
|
||||
...(entities.length === 0 ? [`No team knowledge matched "${topic}".`] : []),
|
||||
...(relatedTasks.length === 0 ? [`No team tasks matched "${topic}".`] : []),
|
||||
];
|
||||
|
||||
return {
|
||||
topic,
|
||||
existingKnowledge: entities.length,
|
||||
entities: entities.map((e: typeof entities[number]) => ({ id: e.id, name: e.name, type: e.entityType })),
|
||||
relatedTasks: relatedTasks.length,
|
||||
tasks: relatedTasks.map((t: typeof relatedTasks[number]) => ({ id: t.id, title: t.title, status: t.status })),
|
||||
gaps,
|
||||
recommendation: entities.length === 0 && relatedTasks.length === 0
|
||||
? `Create a knowledge entry or task for "${topic}" so the team can build a shared record.`
|
||||
: `Review the ${entities.length} matching knowledge entr${entities.length === 1 ? 'y' : 'ies'} and ${relatedTasks.length} matching task${relatedTasks.length === 1 ? '' : 's'} for "${topic}".`,
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
export const waggleHandler = createWaggleHandler();
|
||||
106
packages/worker/src/index.ts
Normal file
106
packages/worker/src/index.ts
Normal file
@@ -0,0 +1,106 @@
|
||||
import { Worker } from 'bullmq';
|
||||
import { createDb } from '../../server/src/db/connection.js';
|
||||
import { agentJobs } from '../../server/src/db/schema.js';
|
||||
import { eq } from 'drizzle-orm';
|
||||
import { JobProcessor, type JobData } from './job-processor.js';
|
||||
import Redis from 'ioredis';
|
||||
import { chatHandler } from './handlers/chat-handler.js';
|
||||
import { taskHandler } from './handlers/task-handler.js';
|
||||
import { groupHandler } from './handlers/group-handler.js';
|
||||
import { createCronHandler } from './handlers/cron-handler.js';
|
||||
import { createWaggleHandler } from './handlers/waggle-handler.js';
|
||||
import { JobService } from '../../server/src/services/job-service.js';
|
||||
|
||||
const REDIS_URL = process.env.REDIS_URL ?? 'redis://localhost:6381';
|
||||
|
||||
function requireDatabaseUrl(): string {
|
||||
const url = process.env.DATABASE_URL;
|
||||
if (!url) throw new Error('DATABASE_URL environment variable is required');
|
||||
return url;
|
||||
}
|
||||
|
||||
export function createWorker(redisUrl = REDIS_URL, databaseUrl?: string, queueName = 'waggle-jobs') {
|
||||
const resolvedDbUrl = databaseUrl ?? requireDatabaseUrl();
|
||||
const db = createDb(resolvedDbUrl);
|
||||
const processor = new JobProcessor();
|
||||
const redisPub = new Redis(redisUrl);
|
||||
const jobService = new JobService(db, redisUrl, queueName);
|
||||
|
||||
// Register job handlers
|
||||
processor.register('chat', chatHandler);
|
||||
processor.register('task', taskHandler);
|
||||
processor.register('waggle', createWaggleHandler({
|
||||
enqueueWorker: async ({ teamId, userId, task, role, context }) => {
|
||||
const childJob = await jobService.createJob(teamId, userId, 'chat', {
|
||||
message: task,
|
||||
role,
|
||||
...(context ? { context } : {}),
|
||||
});
|
||||
return childJob.id;
|
||||
},
|
||||
}));
|
||||
processor.register('group', groupHandler);
|
||||
processor.register('cron', createCronHandler((job, handlerDb) => processor.process(job, handlerDb)));
|
||||
|
||||
const url = new URL(redisUrl);
|
||||
const worker = new Worker<JobData>(queueName, async (job) => {
|
||||
// Update status to running
|
||||
await db.update(agentJobs)
|
||||
.set({ status: 'running', startedAt: new Date() })
|
||||
.where(eq(agentJobs.id, job.data.jobId));
|
||||
|
||||
try {
|
||||
const result = await processor.process(job, db);
|
||||
|
||||
// Update status to completed
|
||||
await db.update(agentJobs)
|
||||
.set({ status: 'completed', completedAt: new Date(), output: result })
|
||||
.where(eq(agentJobs.id, job.data.jobId));
|
||||
|
||||
// Publish progress to Redis (includes teamId for gateway routing)
|
||||
await redisPub.publish(`job:${job.data.jobId}:progress`, JSON.stringify({
|
||||
status: 'completed',
|
||||
output: result,
|
||||
userId: job.data.userId,
|
||||
teamId: job.data.teamId,
|
||||
}));
|
||||
|
||||
return result;
|
||||
} catch (error: unknown) {
|
||||
const message = error instanceof Error ? error.message : String(error);
|
||||
await db.update(agentJobs)
|
||||
.set({ status: 'failed', completedAt: new Date(), output: { error: message } })
|
||||
.where(eq(agentJobs.id, job.data.jobId));
|
||||
throw error;
|
||||
}
|
||||
}, {
|
||||
connection: {
|
||||
host: url.hostname,
|
||||
port: parseInt(url.port || '6379'),
|
||||
},
|
||||
concurrency: 5,
|
||||
});
|
||||
|
||||
// Clean up shared Redis publisher when worker closes
|
||||
worker.on('closed', () => {
|
||||
redisPub.quit().catch(() => {});
|
||||
jobService.close().catch(() => {});
|
||||
});
|
||||
|
||||
return { worker, processor, db, redisPub, jobService };
|
||||
}
|
||||
|
||||
// Start if run directly
|
||||
const isDirectRun = process.argv[1]?.replace(/\\/g, '/').includes('worker/src/index');
|
||||
if (isDirectRun) {
|
||||
const { worker } = createWorker();
|
||||
console.log('Waggle agent worker started, waiting for jobs...');
|
||||
|
||||
worker.on('completed', (job) => {
|
||||
console.log(`Job ${job.id} completed`);
|
||||
});
|
||||
|
||||
worker.on('failed', (job, err) => {
|
||||
console.error(`Job ${job?.id} failed:`, err.message);
|
||||
});
|
||||
}
|
||||
28
packages/worker/src/job-processor.ts
Normal file
28
packages/worker/src/job-processor.ts
Normal file
@@ -0,0 +1,28 @@
|
||||
import type { Job } from 'bullmq';
|
||||
import type { Db } from '../../server/src/db/connection.js';
|
||||
|
||||
export interface JobData {
|
||||
jobId: string;
|
||||
teamId: string;
|
||||
userId: string;
|
||||
jobType: string;
|
||||
input: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export type JobHandler = (job: Job<JobData>, db: Db) => Promise<Record<string, unknown>>;
|
||||
|
||||
export class JobProcessor {
|
||||
private handlers = new Map<string, JobHandler>();
|
||||
|
||||
register(jobType: string, handler: JobHandler): void {
|
||||
this.handlers.set(jobType, handler);
|
||||
}
|
||||
|
||||
async process(job: Job<JobData>, db: Db): Promise<Record<string, unknown>> {
|
||||
const handler = this.handlers.get(job.data.jobType);
|
||||
if (!handler) {
|
||||
throw new Error(`No handler registered for job type: ${job.data.jobType}`);
|
||||
}
|
||||
return handler(job, db);
|
||||
}
|
||||
}
|
||||
352
packages/worker/tests/execution/strategies.test.ts
Normal file
352
packages/worker/tests/execution/strategies.test.ts
Normal file
@@ -0,0 +1,352 @@
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { executeParallel, type AgentMemberConfig, type ExecutionDeps } from '../../src/execution/parallel.js';
|
||||
import { executeSequential } from '../../src/execution/sequential.js';
|
||||
import { executeCoordinator } from '../../src/execution/coordinator.js';
|
||||
|
||||
const mockMembers: AgentMemberConfig[] = [
|
||||
{
|
||||
member: { roleInGroup: 'lead', executionOrder: 0 },
|
||||
agent: { id: 'a1', name: 'leader', model: 'claude-sonnet', systemPrompt: 'You are a leader.', tools: ['search_memory'] },
|
||||
},
|
||||
{
|
||||
member: { roleInGroup: 'worker', executionOrder: 1 },
|
||||
agent: { id: 'a2', name: 'researcher', model: 'claude-haiku', systemPrompt: 'You are a researcher.', tools: ['web_search'] },
|
||||
},
|
||||
{
|
||||
member: { roleInGroup: 'worker', executionOrder: 2 },
|
||||
agent: { id: 'a3', name: 'writer', model: 'claude-haiku', systemPrompt: 'You are a writer.', tools: ['write_file'] },
|
||||
},
|
||||
];
|
||||
|
||||
// ── Mock ExecutionDeps ─────────────────────────────────────────────────
|
||||
|
||||
function createMockDeps(overrides?: Partial<ExecutionDeps>): ExecutionDeps {
|
||||
return {
|
||||
runAgent: vi.fn(async (config) => ({
|
||||
content: `Output from agent with prompt: ${config.systemPrompt.slice(0, 30)}...`,
|
||||
toolsUsed: ['mock_tool'],
|
||||
usage: { inputTokens: 100, outputTokens: 50 },
|
||||
})),
|
||||
resolveTools: vi.fn((names) => names.map(n => ({
|
||||
name: n,
|
||||
description: `Mock ${n}`,
|
||||
parameters: { type: 'object', properties: {} },
|
||||
execute: async () => 'ok',
|
||||
}))),
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// Stub mode (backward compat — no deps passed)
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
describe('Execution Strategies (stub mode)', () => {
|
||||
describe('parallel', () => {
|
||||
it('runs all agents and merges output', async () => {
|
||||
const result = await executeParallel(mockMembers, { task: 'test' });
|
||||
|
||||
expect(result.strategy).toBe('parallel');
|
||||
expect(result.agentCount).toBe(3);
|
||||
expect(result.results).toHaveLength(3);
|
||||
expect(result.mergedOutput).toContain('leader');
|
||||
expect(result.mergedOutput).toContain('researcher');
|
||||
expect(result.mergedOutput).toContain('writer');
|
||||
});
|
||||
|
||||
it('includes agent metadata in each result', async () => {
|
||||
const result = await executeParallel(mockMembers, { task: 'test' });
|
||||
const results = result.results as Array<Record<string, unknown>>;
|
||||
|
||||
expect(results[0]).toMatchObject({
|
||||
agentId: 'a1',
|
||||
agentName: 'leader',
|
||||
model: 'claude-sonnet',
|
||||
role: 'lead',
|
||||
});
|
||||
});
|
||||
|
||||
it('handles single agent', async () => {
|
||||
const single = [mockMembers[0]];
|
||||
const result = await executeParallel(single, {});
|
||||
|
||||
expect(result.agentCount).toBe(1);
|
||||
expect(result.results).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('sequential', () => {
|
||||
it('chains output through agents in order', async () => {
|
||||
const result = await executeSequential(mockMembers, { task: 'test' });
|
||||
|
||||
expect(result.strategy).toBe('sequential');
|
||||
expect(result.agentCount).toBe(3);
|
||||
expect(result.results).toHaveLength(3);
|
||||
expect(result.finalOutput).toContain('writer');
|
||||
});
|
||||
|
||||
it('first agent receives taskInput, subsequent receive previous output', async () => {
|
||||
const result = await executeSequential(mockMembers, { task: 'test' });
|
||||
const results = result.results as Array<Record<string, unknown>>;
|
||||
|
||||
expect(results[0].inputFrom).toBe('taskInput');
|
||||
expect(results[1].inputFrom).toBe('leader');
|
||||
expect(results[2].inputFrom).toBe('researcher');
|
||||
});
|
||||
|
||||
it('handles single agent', async () => {
|
||||
const single = [mockMembers[0]];
|
||||
const result = await executeSequential(single, { data: 'input' });
|
||||
|
||||
expect(result.agentCount).toBe(1);
|
||||
expect(result.finalOutput).toContain('leader');
|
||||
});
|
||||
|
||||
it('handles empty members list', async () => {
|
||||
const result = await executeSequential([], {});
|
||||
|
||||
expect(result.agentCount).toBe(0);
|
||||
expect(result.results).toHaveLength(0);
|
||||
expect(result.finalOutput).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe('coordinator', () => {
|
||||
it('lead delegates, workers execute, lead synthesizes', async () => {
|
||||
const result = await executeCoordinator(mockMembers, { task: 'test' });
|
||||
|
||||
expect(result.strategy).toBe('coordinator');
|
||||
expect(result.leadAgent).toBe('leader');
|
||||
expect(result.workerCount).toBe(2);
|
||||
expect(result.plan).toBeDefined();
|
||||
expect(result.workerResults).toHaveLength(2);
|
||||
expect(result.synthesis).toBeDefined();
|
||||
expect(result.finalOutput).toContain('synthesized');
|
||||
});
|
||||
|
||||
it('only workers appear in workerResults (not lead)', async () => {
|
||||
const result = await executeCoordinator(mockMembers, {});
|
||||
const workerResults = result.workerResults as Array<Record<string, unknown>>;
|
||||
|
||||
const agentNames = workerResults.map(r => r.agentName);
|
||||
expect(agentNames).toContain('researcher');
|
||||
expect(agentNames).toContain('writer');
|
||||
expect(agentNames).not.toContain('leader');
|
||||
});
|
||||
|
||||
it('plan includes subtask assignments', async () => {
|
||||
const result = await executeCoordinator(mockMembers, {});
|
||||
const plan = result.plan as Record<string, unknown>;
|
||||
|
||||
expect(plan.phase).toBe('planning');
|
||||
expect(plan.subtasks).toHaveLength(2);
|
||||
});
|
||||
|
||||
it('synthesis references worker count', async () => {
|
||||
const result = await executeCoordinator(mockMembers, {});
|
||||
const synthesis = result.synthesis as Record<string, unknown>;
|
||||
|
||||
expect(synthesis.phase).toBe('synthesis');
|
||||
expect(synthesis.output).toContain('2');
|
||||
});
|
||||
|
||||
it('throws if no lead agent', async () => {
|
||||
const noLead = mockMembers.map(m => ({
|
||||
...m,
|
||||
member: { ...m.member, roleInGroup: 'worker' },
|
||||
}));
|
||||
|
||||
await expect(executeCoordinator(noLead, {})).rejects.toThrow('requires a lead');
|
||||
});
|
||||
|
||||
it('works with lead and no workers', async () => {
|
||||
const leadOnly: AgentMemberConfig[] = [{
|
||||
member: { roleInGroup: 'lead', executionOrder: 0 },
|
||||
agent: { id: 'a1', name: 'solo-lead', model: 'claude-sonnet', systemPrompt: null, tools: [] },
|
||||
}];
|
||||
|
||||
const result = await executeCoordinator(leadOnly, {});
|
||||
|
||||
expect(result.workerCount).toBe(0);
|
||||
expect(result.workerResults).toHaveLength(0);
|
||||
expect(result.finalOutput).toContain('synthesized 0');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// Real execution (with mocked runAgent via ExecutionDeps)
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
describe('Execution Strategies (real execution)', () => {
|
||||
describe('parallel with deps', () => {
|
||||
it('runs all agents concurrently via runAgent', async () => {
|
||||
const deps = createMockDeps();
|
||||
const result = await executeParallel(mockMembers, { task: 'analyze data' }, deps);
|
||||
|
||||
expect(result.strategy).toBe('parallel');
|
||||
expect(result.agentCount).toBe(3);
|
||||
expect(deps.runAgent).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
|
||||
it('each agent gets its own systemPrompt and resolved tools', async () => {
|
||||
const deps = createMockDeps();
|
||||
await executeParallel(mockMembers, { task: 'analyze' }, deps);
|
||||
|
||||
// Check that resolveTools was called with each agent's tool list
|
||||
expect(deps.resolveTools).toHaveBeenCalledWith(['search_memory']);
|
||||
expect(deps.resolveTools).toHaveBeenCalledWith(['web_search']);
|
||||
expect(deps.resolveTools).toHaveBeenCalledWith(['write_file']);
|
||||
});
|
||||
|
||||
it('captures output, toolsUsed, and usage from each agent', async () => {
|
||||
const deps = createMockDeps();
|
||||
const result = await executeParallel(mockMembers, { task: 'test' }, deps);
|
||||
const results = result.results as Array<Record<string, unknown>>;
|
||||
|
||||
expect(results[0].toolsUsed).toEqual(['mock_tool']);
|
||||
expect(results[0].usage).toEqual({ inputTokens: 100, outputTokens: 50 });
|
||||
expect(typeof results[0].output).toBe('string');
|
||||
expect((results[0].output as string).length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it('handles one agent failing (others still complete)', async () => {
|
||||
const deps = createMockDeps({
|
||||
runAgent: vi.fn(async (config) => {
|
||||
if (config.systemPrompt.includes('researcher')) throw new Error('API timeout');
|
||||
return { content: 'OK', toolsUsed: [], usage: { inputTokens: 10, outputTokens: 5 } };
|
||||
}),
|
||||
});
|
||||
const result = await executeParallel(mockMembers, { task: 'test' }, deps);
|
||||
|
||||
// All 3 agents attempted
|
||||
expect(result.agentCount).toBe(3);
|
||||
// Error captured, not thrown
|
||||
const errors = result.errors as Array<Record<string, unknown>>;
|
||||
expect(errors).toHaveLength(1);
|
||||
expect(errors[0].agent).toBe('researcher');
|
||||
expect(errors[0].error).toBe('API timeout');
|
||||
// Other agents still succeeded
|
||||
expect(result.mergedOutput).toContain('OK');
|
||||
});
|
||||
});
|
||||
|
||||
describe('sequential with deps', () => {
|
||||
it('runs agents in order, each receives previous output as context', async () => {
|
||||
let callOrder = 0;
|
||||
const deps = createMockDeps({
|
||||
runAgent: vi.fn(async (config) => {
|
||||
callOrder++;
|
||||
return {
|
||||
content: `Step ${callOrder} result`,
|
||||
toolsUsed: [],
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
};
|
||||
}),
|
||||
});
|
||||
|
||||
const result = await executeSequential(mockMembers, { task: 'research then write' }, deps);
|
||||
|
||||
expect(result.strategy).toBe('sequential');
|
||||
expect(deps.runAgent).toHaveBeenCalledTimes(3);
|
||||
// Final output comes from last agent
|
||||
expect(result.finalOutput).toBe('Step 3 result');
|
||||
});
|
||||
|
||||
it('second agent receives first agents output in systemPrompt', async () => {
|
||||
const deps = createMockDeps({
|
||||
runAgent: vi.fn(async (config) => ({
|
||||
content: `Output: ${config.systemPrompt.includes('Previous Agent') ? 'chained' : 'first'}`,
|
||||
toolsUsed: [],
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
})),
|
||||
});
|
||||
|
||||
await executeSequential(mockMembers, { task: 'pipeline' }, deps);
|
||||
|
||||
// Second call should have "Previous Agent's Output" in the system prompt
|
||||
const calls = vi.mocked(deps.runAgent).mock.calls;
|
||||
expect(calls[0][0].systemPrompt).not.toContain('Previous Agent');
|
||||
expect(calls[1][0].systemPrompt).toContain('Previous Agent');
|
||||
expect(calls[2][0].systemPrompt).toContain('Previous Agent');
|
||||
});
|
||||
|
||||
it('error in middle agent stops the chain', async () => {
|
||||
const deps = createMockDeps({
|
||||
runAgent: vi.fn(async (config) => {
|
||||
if (config.systemPrompt.includes('researcher')) throw new Error('Crash');
|
||||
return { content: 'OK', toolsUsed: [], usage: { inputTokens: 10, outputTokens: 5 } };
|
||||
}),
|
||||
});
|
||||
|
||||
const result = await executeSequential(mockMembers, { task: 'test' }, deps);
|
||||
|
||||
// Chain stopped at researcher (agent 2), writer (agent 3) never ran
|
||||
expect(result.agentCount).toBe(2); // leader + researcher (failed)
|
||||
expect(result.chainBroken).toBe(true);
|
||||
expect(deps.runAgent).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe('coordinator with deps', () => {
|
||||
it('runs 3-phase pattern: plan → execute → synthesize', async () => {
|
||||
const deps = createMockDeps();
|
||||
const result = await executeCoordinator(mockMembers, { task: 'complex task' }, deps);
|
||||
|
||||
expect(result.strategy).toBe('coordinator');
|
||||
// Lead calls: 1 planning + 1 synthesis = 2
|
||||
// Workers: 2 parallel calls
|
||||
// Total: 4
|
||||
expect(deps.runAgent).toHaveBeenCalledTimes(4);
|
||||
|
||||
const plan = result.plan as Record<string, unknown>;
|
||||
expect(plan.phase).toBe('planning');
|
||||
expect(typeof plan.output).toBe('string');
|
||||
|
||||
const synthesis = result.synthesis as Record<string, unknown>;
|
||||
expect(synthesis.phase).toBe('synthesis');
|
||||
});
|
||||
|
||||
it('workers receive the plan as context', async () => {
|
||||
const deps = createMockDeps({
|
||||
runAgent: vi.fn(async (config) => ({
|
||||
content: config.systemPrompt.includes('Coordinator') ? 'Worker got plan' : 'Plan created',
|
||||
toolsUsed: [],
|
||||
usage: { inputTokens: 10, outputTokens: 5 },
|
||||
})),
|
||||
});
|
||||
|
||||
const result = await executeCoordinator(mockMembers, { task: 'test' }, deps);
|
||||
const workerResults = result.workerResults as Array<Record<string, unknown>>;
|
||||
|
||||
// Workers should have received the plan in their systemPrompt
|
||||
expect(workerResults[0].output).toBe('Worker got plan');
|
||||
expect(workerResults[1].output).toBe('Worker got plan');
|
||||
});
|
||||
|
||||
it('handles zero workers (lead does everything)', async () => {
|
||||
const leadOnly: AgentMemberConfig[] = [{
|
||||
member: { roleInGroup: 'lead', executionOrder: 0 },
|
||||
agent: { id: 'a1', name: 'solo', model: 'claude-sonnet', systemPrompt: 'Solo lead', tools: [] },
|
||||
}];
|
||||
|
||||
const deps = createMockDeps();
|
||||
const result = await executeCoordinator(leadOnly, { task: 'solo task' }, deps);
|
||||
|
||||
expect(result.workerCount).toBe(0);
|
||||
expect(result.workerResults).toHaveLength(0);
|
||||
// Lead still runs plan + synthesis = 2 calls
|
||||
expect(deps.runAgent).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('still throws if no lead (even with deps)', async () => {
|
||||
const noLead = mockMembers.map(m => ({
|
||||
...m,
|
||||
member: { ...m.member, roleInGroup: 'worker' },
|
||||
}));
|
||||
const deps = createMockDeps();
|
||||
|
||||
await expect(executeCoordinator(noLead, {}, deps)).rejects.toThrow('requires a lead');
|
||||
});
|
||||
});
|
||||
});
|
||||
126
packages/worker/tests/handlers/chat-handler.test.ts
Normal file
126
packages/worker/tests/handlers/chat-handler.test.ts
Normal file
@@ -0,0 +1,126 @@
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import type { Job } from 'bullmq';
|
||||
import type { Db } from '../../../server/src/db/connection.js';
|
||||
import type { JobData } from '../../src/job-processor.js';
|
||||
|
||||
/** Build a minimal Job<JobData> mock — only `.data` is exercised by handlers. */
|
||||
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(() => []),
|
||||
}));
|
||||
|
||||
import { chatHandler } from '../../src/handlers/chat-handler.js';
|
||||
import { runAgentLoop, createSystemTools } from '@waggle/agent';
|
||||
|
||||
describe('Chat Handler', () => {
|
||||
const mockDb = {} as Db;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('calls runAgentLoop and returns response with correct shape', async () => {
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j1',
|
||||
teamId: 't1',
|
||||
userId: 'u1',
|
||||
jobType: 'chat',
|
||||
input: { message: 'hello world' },
|
||||
});
|
||||
|
||||
const result = await chatHandler(mockJob, mockDb);
|
||||
|
||||
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.tokensUsed).toBe(150);
|
||||
});
|
||||
|
||||
it('passes correct config to runAgentLoop', async () => {
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j2',
|
||||
teamId: 't1',
|
||||
userId: 'u1',
|
||||
jobType: 'chat',
|
||||
input: { message: 'test message', model: 'gpt-4o' },
|
||||
});
|
||||
|
||||
await chatHandler(mockJob, mockDb);
|
||||
|
||||
expect(runAgentLoop).toHaveBeenCalledOnce();
|
||||
const callArgs = vi.mocked(runAgentLoop).mock.calls[0][0];
|
||||
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([]);
|
||||
});
|
||||
|
||||
it('creates system tools with workspaceDir from input', async () => {
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j3',
|
||||
teamId: 't1',
|
||||
userId: 'u1',
|
||||
jobType: 'chat',
|
||||
input: { message: 'test', workspaceDir: '/custom/workspace' },
|
||||
});
|
||||
|
||||
await chatHandler(mockJob, mockDb);
|
||||
|
||||
expect(createSystemTools).toHaveBeenCalledWith('/custom/workspace');
|
||||
});
|
||||
|
||||
it('uses default model when not specified in input', async () => {
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j4',
|
||||
teamId: 't1',
|
||||
userId: 'u1',
|
||||
jobType: 'chat',
|
||||
input: { message: 'hi' },
|
||||
});
|
||||
|
||||
await chatHandler(mockJob, mockDb);
|
||||
|
||||
const callArgs = vi.mocked(runAgentLoop).mock.calls[0][0];
|
||||
expect(callArgs.model).toBe('claude-sonnet');
|
||||
});
|
||||
|
||||
it('handles missing message gracefully', async () => {
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j5',
|
||||
teamId: 't1',
|
||||
userId: 'u1',
|
||||
jobType: 'chat',
|
||||
input: {},
|
||||
});
|
||||
|
||||
const result = await chatHandler(mockJob, mockDb);
|
||||
|
||||
expect(result.response).toBe('Agent response here');
|
||||
expect(result.userId).toBe('u1');
|
||||
const callArgs = vi.mocked(runAgentLoop).mock.calls[0][0];
|
||||
expect(callArgs.messages).toEqual([{ role: 'user', content: '' }]);
|
||||
});
|
||||
|
||||
it('includes userId in response', async () => {
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j6',
|
||||
teamId: 't1',
|
||||
userId: 'user-abc',
|
||||
jobType: 'chat',
|
||||
input: { message: 'test' },
|
||||
});
|
||||
|
||||
const result = await chatHandler(mockJob, mockDb);
|
||||
|
||||
expect(result.userId).toBe('user-abc');
|
||||
});
|
||||
});
|
||||
66
packages/worker/tests/handlers/cron-handler.test.ts
Normal file
66
packages/worker/tests/handlers/cron-handler.test.ts
Normal file
@@ -0,0 +1,66 @@
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
import type { Job } from 'bullmq';
|
||||
import type { Db } from '../../../server/src/db/connection.js';
|
||||
import type { JobData } from '../../src/job-processor.js';
|
||||
import { createCronHandler } from '../../src/handlers/cron-handler.js';
|
||||
|
||||
function makeJob(input: Record<string, unknown>): Job<JobData> {
|
||||
return {
|
||||
data: {
|
||||
jobId: 'cron-job-1',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'cron',
|
||||
input,
|
||||
},
|
||||
} as unknown as Job<JobData>;
|
||||
}
|
||||
|
||||
describe('Cron handler', () => {
|
||||
const db = {} as Db;
|
||||
|
||||
it('delegates a scheduled job to a real worker handler', async () => {
|
||||
const dispatch = vi.fn(async (job: Job<JobData>) => ({
|
||||
response: `ran ${job.data.jobType}`,
|
||||
input: job.data.input,
|
||||
}));
|
||||
const handler = createCronHandler(dispatch);
|
||||
|
||||
const result = await handler(makeJob({
|
||||
jobType: 'chat',
|
||||
jobConfig: { message: 'send the daily brief' },
|
||||
}), db);
|
||||
|
||||
expect(dispatch).toHaveBeenCalledOnce();
|
||||
const delegatedJob = dispatch.mock.calls[0][0];
|
||||
expect(delegatedJob.data.jobType).toBe('chat');
|
||||
expect(delegatedJob.data.input).toEqual({ message: 'send the daily brief' });
|
||||
expect(result).toEqual({
|
||||
response: 'ran chat',
|
||||
input: { message: 'send the daily brief' },
|
||||
cron: { delegatedJobType: 'chat' },
|
||||
});
|
||||
});
|
||||
|
||||
it('accepts input as the delegated payload for queue clients', async () => {
|
||||
const dispatch = vi.fn(async () => ({ ok: true }));
|
||||
const handler = createCronHandler(dispatch);
|
||||
|
||||
await handler(makeJob({
|
||||
jobType: 'task',
|
||||
input: { taskId: 'task-1' },
|
||||
}), db);
|
||||
|
||||
expect(dispatch.mock.calls[0][0].data.input).toEqual({ taskId: 'task-1' });
|
||||
});
|
||||
|
||||
it('rejects missing, unsupported, and malformed delegation targets', async () => {
|
||||
const handler = createCronHandler(vi.fn(async () => ({ ok: true })));
|
||||
|
||||
await expect(handler(makeJob({}), db)).rejects.toThrow('requires input.jobType');
|
||||
await expect(handler(makeJob({ jobType: 'cron', jobConfig: {} }), db))
|
||||
.rejects.toThrow('cannot delegate to "cron"');
|
||||
await expect(handler(makeJob({ jobType: 'chat' }), db))
|
||||
.rejects.toThrow('requires input.jobConfig or input.input');
|
||||
});
|
||||
});
|
||||
340
packages/worker/tests/handlers/handlers.test.ts
Normal file
340
packages/worker/tests/handlers/handlers.test.ts
Normal file
@@ -0,0 +1,340 @@
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import type { Job } from 'bullmq';
|
||||
import type { Db } from '../../../server/src/db/connection.js';
|
||||
import type { JobData } from '../../src/job-processor.js';
|
||||
|
||||
/** Build a minimal Job<JobData> mock — only `.data` is exercised by handlers. */
|
||||
function makeJob(data: JobData): Job<JobData> {
|
||||
return { data } as unknown as 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' },
|
||||
]),
|
||||
}));
|
||||
|
||||
// ── Mock drizzle-orm ────────────────────────────────────────────────────
|
||||
vi.mock('drizzle-orm', () => ({
|
||||
eq: vi.fn((col: unknown, val: unknown) => ({ column: col, value: val })),
|
||||
}));
|
||||
|
||||
// ── Mock server DB schema ───────────────────────────────────────────────
|
||||
vi.mock('../../../server/src/db/schema.js', () => ({
|
||||
tasks: {
|
||||
id: 'tasks.id',
|
||||
teamId: 'tasks.teamId',
|
||||
status: 'tasks.status',
|
||||
assignedTo: 'tasks.assignedTo',
|
||||
updatedAt: 'tasks.updatedAt',
|
||||
title: 'tasks.title',
|
||||
description: 'tasks.description',
|
||||
priority: 'tasks.priority',
|
||||
},
|
||||
agentGroups: { id: 'agent_groups.id' },
|
||||
agentGroupMembers: {
|
||||
groupId: 'agent_group_members.groupId',
|
||||
agentId: 'agent_group_members.agentId',
|
||||
executionOrder: 'agent_group_members.executionOrder',
|
||||
},
|
||||
agents: { id: 'agents.id' },
|
||||
}));
|
||||
|
||||
import { taskHandler } from '../../src/handlers/task-handler.js';
|
||||
import { groupHandler } from '../../src/handlers/group-handler.js';
|
||||
import { runAgentLoop } from '@waggle/agent';
|
||||
|
||||
// ── Helper: mock DB ─────────────────────────────────────────────────────
|
||||
|
||||
function createMockDb(overrides?: {
|
||||
selectResult?: unknown[];
|
||||
selectResultSecond?: unknown[];
|
||||
}): Db {
|
||||
const selectResult = overrides?.selectResult ?? [];
|
||||
const hasSecondResult = overrides?.selectResultSecond !== undefined;
|
||||
const selectResultSecond = overrides?.selectResultSecond ?? [];
|
||||
let selectCallCount = 0;
|
||||
|
||||
const mockChain = (results: unknown[]) => ({
|
||||
from: vi.fn().mockReturnValue({
|
||||
where: vi.fn().mockResolvedValue(results),
|
||||
innerJoin: vi.fn().mockReturnValue({
|
||||
where: vi.fn().mockResolvedValue(results),
|
||||
}),
|
||||
}),
|
||||
});
|
||||
|
||||
return {
|
||||
select: vi.fn(() => {
|
||||
selectCallCount++;
|
||||
if (selectCallCount === 1) return mockChain(selectResult);
|
||||
return mockChain(hasSecondResult ? selectResultSecond : selectResult);
|
||||
}),
|
||||
update: vi.fn().mockReturnValue({
|
||||
set: vi.fn().mockReturnValue({
|
||||
where: vi.fn().mockResolvedValue(undefined),
|
||||
}),
|
||||
}),
|
||||
} as unknown as Db;
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// Task Handler Tests (11F-7)
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
describe('Task Handler', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('executes a task job and returns results', async () => {
|
||||
const mockTask = {
|
||||
id: 'task-1',
|
||||
title: 'Fix login bug',
|
||||
description: 'Users cannot log in with SSO',
|
||||
teamId: 'team-1',
|
||||
priority: 'high',
|
||||
status: 'open',
|
||||
};
|
||||
|
||||
const mockDb = createMockDb({ selectResult: [mockTask] });
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j1',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'task',
|
||||
input: { taskId: 'task-1' },
|
||||
});
|
||||
|
||||
const result = await taskHandler(mockJob, mockDb);
|
||||
|
||||
expect(result.taskId).toBe('task-1');
|
||||
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.tokensUsed).toBe(300); // 200 + 100
|
||||
expect(result.userId).toBe('user-1');
|
||||
|
||||
// Verify runAgentLoop was called
|
||||
expect(runAgentLoop).toHaveBeenCalledOnce();
|
||||
|
||||
// Verify DB was updated (in-progress then completed)
|
||||
expect(mockDb.update).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('throws error and resets task to open on agent failure', async () => {
|
||||
const mockTask = {
|
||||
id: 'task-2',
|
||||
title: 'Deploy service',
|
||||
description: null,
|
||||
teamId: 'team-1',
|
||||
priority: null,
|
||||
status: 'open',
|
||||
};
|
||||
|
||||
const mockDb = createMockDb({ selectResult: [mockTask] });
|
||||
vi.mocked(runAgentLoop).mockRejectedValueOnce(new Error('LLM API timeout'));
|
||||
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j2',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'task',
|
||||
input: { taskId: 'task-2' },
|
||||
});
|
||||
|
||||
await expect(taskHandler(mockJob, mockDb)).rejects.toThrow('LLM API timeout');
|
||||
|
||||
// Verify DB was updated: first to in_progress, then back to open on error
|
||||
expect(mockDb.update).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('throws if taskId is missing from input', async () => {
|
||||
const mockDb = createMockDb();
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j3',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'task',
|
||||
input: {}, // no taskId
|
||||
});
|
||||
|
||||
await expect(taskHandler(mockJob, mockDb)).rejects.toThrow('taskHandler requires input.taskId');
|
||||
});
|
||||
|
||||
it('throws if task is not found in DB', async () => {
|
||||
const mockDb = createMockDb({ selectResult: [] }); // no task found
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j4',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'task',
|
||||
input: { taskId: 'nonexistent-task' },
|
||||
});
|
||||
|
||||
await expect(taskHandler(mockJob, mockDb)).rejects.toThrow('Task not found: nonexistent-task');
|
||||
});
|
||||
|
||||
it('throws if task belongs to a different team', async () => {
|
||||
const mockTask = {
|
||||
id: 'task-5',
|
||||
title: 'Other team task',
|
||||
description: null,
|
||||
teamId: 'team-other', // different team
|
||||
priority: null,
|
||||
status: 'open',
|
||||
};
|
||||
|
||||
const mockDb = createMockDb({ selectResult: [mockTask] });
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j5',
|
||||
teamId: 'team-1', // requesting team doesn't match
|
||||
userId: 'user-1',
|
||||
jobType: 'task',
|
||||
input: { taskId: 'task-5' },
|
||||
});
|
||||
|
||||
await expect(taskHandler(mockJob, mockDb)).rejects.toThrow(
|
||||
'Task task-5 does not belong to team team-1',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// Group Handler Tests (11F-7)
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
describe('Group Handler', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('throws if groupId is missing from input', async () => {
|
||||
const mockDb = createMockDb();
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j1',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'group',
|
||||
input: {}, // no groupId
|
||||
});
|
||||
|
||||
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow('groupHandler requires input.groupId');
|
||||
});
|
||||
|
||||
it('throws if group is not found in DB', async () => {
|
||||
const mockDb = createMockDb({ selectResult: [] });
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j2',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'group',
|
||||
input: { groupId: 'nonexistent-group' },
|
||||
});
|
||||
|
||||
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow('Agent group not found: nonexistent-group');
|
||||
});
|
||||
|
||||
it('throws if group has no members', async () => {
|
||||
const mockGroup = { id: 'group-1', strategy: 'parallel', name: 'Test Group' };
|
||||
// First select returns the group, second returns empty members
|
||||
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: [] });
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j3',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'group',
|
||||
input: { groupId: 'group-1' },
|
||||
});
|
||||
|
||||
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow('Agent group group-1 has no members');
|
||||
});
|
||||
|
||||
it('dispatches parallel strategy correctly', async () => {
|
||||
const mockGroup = { id: 'group-1', strategy: 'parallel', name: 'Parallel Group' };
|
||||
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: [] },
|
||||
},
|
||||
{
|
||||
member: { roleInGroup: 'worker', executionOrder: 1, groupId: 'group-1', agentId: 'a2' },
|
||||
agent: { id: 'a2', name: 'Agent Beta', model: 'claude-haiku', systemPrompt: 'You are Beta', tools: [] },
|
||||
},
|
||||
];
|
||||
|
||||
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: mockMembers });
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j4',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'group',
|
||||
input: { groupId: 'group-1', taskInput: { task: 'analyze data' } },
|
||||
});
|
||||
|
||||
const result = await groupHandler(mockJob, mockDb);
|
||||
|
||||
expect(result.strategy).toBe('parallel');
|
||||
expect(result.agentCount).toBe(2);
|
||||
// runAgentLoop called once per member
|
||||
expect(runAgentLoop).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('dispatches sequential strategy correctly', async () => {
|
||||
const mockGroup = { id: 'group-2', strategy: 'sequential', name: 'Sequential Group' };
|
||||
const mockMembers = [
|
||||
{
|
||||
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-2', agentId: 'a1' },
|
||||
agent: { id: 'a1', 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: [] },
|
||||
},
|
||||
];
|
||||
|
||||
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: mockMembers });
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j5',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'group',
|
||||
input: { groupId: 'group-2', taskInput: { task: 'research and draft' } },
|
||||
});
|
||||
|
||||
const result = await groupHandler(mockJob, mockDb);
|
||||
|
||||
expect(result.strategy).toBe('sequential');
|
||||
expect(result.agentCount).toBe(2);
|
||||
expect(runAgentLoop).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('throws for unknown execution strategy', async () => {
|
||||
const mockGroup = { id: 'group-3', strategy: 'unknown_strategy', name: 'Bad Group' };
|
||||
const mockMembers = [
|
||||
{
|
||||
member: { roleInGroup: 'worker', executionOrder: 0, groupId: 'group-3', agentId: 'a1' },
|
||||
agent: { id: 'a1', name: 'Agent', model: 'claude-sonnet', systemPrompt: null, tools: [] },
|
||||
},
|
||||
];
|
||||
|
||||
const mockDb = createMockDb({ selectResult: [mockGroup], selectResultSecond: mockMembers });
|
||||
const mockJob = makeJob({
|
||||
jobId: 'j6',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'group',
|
||||
input: { groupId: 'group-3', taskInput: {} },
|
||||
});
|
||||
|
||||
await expect(groupHandler(mockJob, mockDb)).rejects.toThrow('Unknown execution strategy: unknown_strategy');
|
||||
});
|
||||
});
|
||||
92
packages/worker/tests/handlers/waggle-dispatch.test.ts
Normal file
92
packages/worker/tests/handlers/waggle-dispatch.test.ts
Normal file
@@ -0,0 +1,92 @@
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
|
||||
/**
|
||||
* Tests for waggle handler dispatcher integration.
|
||||
* Since waggle-handler depends on BullMQ Job + DB, we test the dispatcher
|
||||
* integration logic directly by importing the dispatcher.
|
||||
*/
|
||||
import { WaggleDanceDispatcher } from '@waggle/waggle-dance';
|
||||
import type { DispatchDeps } from '@waggle/waggle-dance';
|
||||
|
||||
function makeDeps(overrides?: Partial<DispatchDeps>): DispatchDeps {
|
||||
return {
|
||||
searchMemory: vi.fn(async (q: string) => `Knowledge for: ${q}`),
|
||||
resolveCapability: vi.fn((q: string) => [
|
||||
{ source: 'native', name: q, description: `Capability: ${q}`, available: true },
|
||||
]),
|
||||
spawnWorker: vi.fn(async (task: string, role: string) => `Worker spawned: ${task} (${role})`),
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
describe('waggle handler dispatcher integration', () => {
|
||||
it('task_delegation routes through dispatcher to spawnWorker', async () => {
|
||||
const deps = makeDeps();
|
||||
const dispatcher = new WaggleDanceDispatcher(deps);
|
||||
|
||||
const result = await dispatcher.dispatch({
|
||||
id: 'msg-1', teamId: 't1', senderId: 'agent-1',
|
||||
type: 'request', subtype: 'task_delegation',
|
||||
content: { task: 'Analyze Q4 data', role: 'analyst' },
|
||||
referenceId: null, routing: null, createdAt: new Date(),
|
||||
});
|
||||
|
||||
expect(result.handled).toBe(true);
|
||||
expect(result.response).toContain('Analyze Q4 data');
|
||||
expect(deps.spawnWorker).toHaveBeenCalledWith('Analyze Q4 data', 'analyst', undefined);
|
||||
});
|
||||
|
||||
it('knowledge_check routes through dispatcher to searchMemory', async () => {
|
||||
const deps = makeDeps();
|
||||
const dispatcher = new WaggleDanceDispatcher(deps);
|
||||
|
||||
const result = await dispatcher.dispatch({
|
||||
id: 'msg-2', teamId: 't1', senderId: 'agent-1',
|
||||
type: 'request', subtype: 'knowledge_check',
|
||||
content: { query: 'TypeScript patterns' },
|
||||
referenceId: null, routing: null, createdAt: new Date(),
|
||||
});
|
||||
|
||||
expect(result.handled).toBe(true);
|
||||
expect(result.response).toBe('Knowledge for: TypeScript patterns');
|
||||
expect(deps.searchMemory).toHaveBeenCalledWith('TypeScript patterns');
|
||||
});
|
||||
|
||||
it('skill_request routes through dispatcher to resolveCapability', async () => {
|
||||
const deps = makeDeps();
|
||||
const dispatcher = new WaggleDanceDispatcher(deps);
|
||||
|
||||
const result = await dispatcher.dispatch({
|
||||
id: 'msg-3', teamId: 't1', senderId: 'agent-1',
|
||||
type: 'request', subtype: 'skill_request',
|
||||
content: { skill: 'web_search' },
|
||||
referenceId: null, routing: null, createdAt: new Date(),
|
||||
});
|
||||
|
||||
expect(result.handled).toBe(true);
|
||||
expect(deps.resolveCapability).toHaveBeenCalledWith('web_search');
|
||||
});
|
||||
|
||||
it('invalid message type returns error', async () => {
|
||||
const deps = makeDeps();
|
||||
const dispatcher = new WaggleDanceDispatcher(deps);
|
||||
|
||||
const result = await dispatcher.dispatch({
|
||||
id: 'msg-4', teamId: 't1', senderId: 'agent-1',
|
||||
type: 'broadcast', subtype: 'task_delegation', // invalid combo
|
||||
content: {},
|
||||
referenceId: null, routing: null, createdAt: new Date(),
|
||||
});
|
||||
|
||||
expect(result.handled).toBe(false);
|
||||
expect(result.error).toContain('Invalid message');
|
||||
});
|
||||
|
||||
it('existing dispatcher tests still pass (14 tests in dispatcher.test.ts)', async () => {
|
||||
// This test verifies the dispatcher module imports correctly
|
||||
// The actual 14 tests run in packages/waggle-dance/tests/dispatcher.test.ts
|
||||
const deps = makeDeps();
|
||||
const dispatcher = new WaggleDanceDispatcher(deps);
|
||||
expect(dispatcher).toBeDefined();
|
||||
});
|
||||
});
|
||||
77
packages/worker/tests/handlers/waggle-handler.test.ts
Normal file
77
packages/worker/tests/handlers/waggle-handler.test.ts
Normal file
@@ -0,0 +1,77 @@
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
import type { Job } from 'bullmq';
|
||||
import type { Db } from '../../../server/src/db/connection.js';
|
||||
import type { JobData } from '../../src/job-processor.js';
|
||||
import { createWaggleHandler } from '../../src/handlers/waggle-handler.js';
|
||||
|
||||
function makeJob(input: Record<string, unknown>): Job<JobData> {
|
||||
return {
|
||||
data: {
|
||||
jobId: 'waggle-job-1',
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
jobType: 'waggle',
|
||||
input,
|
||||
},
|
||||
} as unknown as Job<JobData>;
|
||||
}
|
||||
|
||||
describe('Waggle worker handler', () => {
|
||||
const db = {} as Db;
|
||||
|
||||
it('queues a real child chat job for task delegation', async () => {
|
||||
const enqueueWorker = vi.fn(async () => 'child-job-1');
|
||||
const handler = createWaggleHandler({ enqueueWorker });
|
||||
|
||||
const result = await handler(makeJob({
|
||||
type: 'request',
|
||||
subtype: 'task_delegation',
|
||||
content: { task: 'Analyze Q4 data', role: 'analyst', context: 'Use the latest report' },
|
||||
}), db);
|
||||
|
||||
expect(enqueueWorker).toHaveBeenCalledWith({
|
||||
teamId: 'team-1',
|
||||
userId: 'user-1',
|
||||
task: 'Analyze Q4 data',
|
||||
role: 'analyst',
|
||||
context: 'Use the latest report',
|
||||
});
|
||||
expect(result.response).toContain('child-job-1');
|
||||
expect(result.handled).toBe(true);
|
||||
});
|
||||
|
||||
it('reports missing capabilities as unavailable instead of fabricating success', async () => {
|
||||
const handler = createWaggleHandler();
|
||||
|
||||
const result = await handler(makeJob({
|
||||
type: 'request',
|
||||
subtype: 'skill_request',
|
||||
content: { query: 'something-that-does-not-exist' },
|
||||
}), db);
|
||||
|
||||
expect(result.response).toContain('available):');
|
||||
expect(result.response).toContain('not available');
|
||||
});
|
||||
|
||||
it('returns a concrete legacy knowledge gap instead of a stub marker', async () => {
|
||||
const mockDb = {
|
||||
select: vi.fn(() => ({
|
||||
from: vi.fn(() => ({
|
||||
where: vi.fn(() => ({
|
||||
limit: vi.fn(async () => []),
|
||||
})),
|
||||
})),
|
||||
})),
|
||||
} as unknown as Db;
|
||||
const handler = createWaggleHandler();
|
||||
|
||||
const result = await handler(makeJob({ topic: 'Q4 planning' }), mockDb);
|
||||
|
||||
expect(result.gaps).toEqual([
|
||||
'No team knowledge matched "Q4 planning".',
|
||||
'No team tasks matched "Q4 planning".',
|
||||
]);
|
||||
expect(result.recommendation).toContain('Create a knowledge entry or task');
|
||||
expect(JSON.stringify(result)).not.toContain('[Stub]');
|
||||
});
|
||||
});
|
||||
139
packages/worker/tests/job-processor.test.ts
Normal file
139
packages/worker/tests/job-processor.test.ts
Normal file
@@ -0,0 +1,139 @@
|
||||
import { describe, it, expect, beforeAll, afterAll } from 'vitest';
|
||||
import type { Job } from 'bullmq';
|
||||
import { createWorker } from '../src/index.js';
|
||||
import { JobService } from '../../server/src/services/job-service.js';
|
||||
import { createDb } from '../../server/src/db/connection.js';
|
||||
import { users, teams, teamMembers, agentJobs } from '../../server/src/db/schema.js';
|
||||
import { eq, sql } from 'drizzle-orm';
|
||||
import type { JobData } from '../src/job-processor.js';
|
||||
|
||||
const REDIS_URL = 'redis://localhost:6381';
|
||||
const DATABASE_URL = 'postgres://waggle:waggle_dev@localhost:5434/waggle';
|
||||
const QUEUE_NAME = `waggle-jobs-test-${Date.now()}`;
|
||||
|
||||
/** Poll for job status until it reaches target or timeout. */
|
||||
async function waitForJobStatus(
|
||||
jobService: JobService,
|
||||
jobId: string,
|
||||
target: string,
|
||||
timeoutMs = 10_000,
|
||||
intervalMs = 200,
|
||||
): Promise<Awaited<ReturnType<JobService['getJob']>>> {
|
||||
const start = Date.now();
|
||||
while (Date.now() - start < timeoutMs) {
|
||||
const job = await jobService.getJob(jobId);
|
||||
if (job?.status === target || job?.status === 'failed') return job;
|
||||
await new Promise(r => setTimeout(r, intervalMs));
|
||||
}
|
||||
// Return last state even on timeout
|
||||
return jobService.getJob(jobId);
|
||||
}
|
||||
|
||||
describe('BullMQ Worker', () => {
|
||||
let db: ReturnType<typeof createDb>;
|
||||
let jobService: JobService;
|
||||
let workerInstance: ReturnType<typeof createWorker>;
|
||||
let testUserId: string;
|
||||
let testTeamId: string;
|
||||
|
||||
beforeAll(async () => {
|
||||
db = createDb(DATABASE_URL);
|
||||
jobService = new JobService(db, REDIS_URL, QUEUE_NAME);
|
||||
workerInstance = createWorker(REDIS_URL, DATABASE_URL, QUEUE_NAME);
|
||||
|
||||
// Override handlers with fast mocks (real handlers call LiteLLM which isn't running in tests)
|
||||
workerInstance.processor.register('chat', async (job) => ({
|
||||
response: `Mock response for: ${String(job.data.input.message ?? '')}`,
|
||||
model: 'mock',
|
||||
}));
|
||||
workerInstance.processor.register('task', async (job) => ({
|
||||
result: `Mock task result`,
|
||||
input: job.data.input,
|
||||
}));
|
||||
|
||||
// Create test data
|
||||
const [user] = await db.insert(users).values({
|
||||
clerkId: 'worker_test_user_' + Date.now(),
|
||||
displayName: 'Worker Test',
|
||||
email: `worker_${Date.now()}@test.com`,
|
||||
}).returning();
|
||||
testUserId = user.id;
|
||||
|
||||
const [team] = await db.insert(teams).values({
|
||||
name: 'Worker Team',
|
||||
slug: 'worker-test-' + Date.now(),
|
||||
ownerId: testUserId,
|
||||
}).returning();
|
||||
testTeamId = team.id;
|
||||
|
||||
await db.insert(teamMembers).values({
|
||||
teamId: testTeamId,
|
||||
userId: testUserId,
|
||||
role: 'owner',
|
||||
});
|
||||
}, 15_000);
|
||||
|
||||
afterAll(async () => {
|
||||
try { await workerInstance.worker.close(); } catch { /* ignore */ }
|
||||
try { await jobService.close(); } catch { /* ignore */ }
|
||||
await db.execute(sql`DELETE FROM agent_jobs WHERE team_id = ${testTeamId}`);
|
||||
await db.execute(sql`DELETE FROM team_members WHERE user_id = ${testUserId}`);
|
||||
await db.execute(sql`DELETE FROM teams WHERE id = ${testTeamId}`);
|
||||
await db.execute(sql`DELETE FROM users WHERE id = ${testUserId}`);
|
||||
}, 15_000);
|
||||
|
||||
it('queues job and worker picks it up', async () => {
|
||||
const job = await jobService.createJob(testTeamId, testUserId, 'chat', { message: 'hello' });
|
||||
expect(job.status).toBe('queued');
|
||||
|
||||
const updated = await waitForJobStatus(jobService, job.id, 'completed');
|
||||
expect(updated?.status).toBe('completed');
|
||||
expect(updated?.output).toBeDefined();
|
||||
}, 15_000);
|
||||
|
||||
it('tracks job status transitions', async () => {
|
||||
const job = await jobService.createJob(testTeamId, testUserId, 'chat', { message: 'status test' });
|
||||
|
||||
const completed = await waitForJobStatus(jobService, job.id, 'completed');
|
||||
expect(completed?.startedAt).toBeDefined();
|
||||
expect(completed?.completedAt).toBeDefined();
|
||||
expect(completed?.status).toBe('completed');
|
||||
}, 15_000);
|
||||
|
||||
it('marks failed jobs', async () => {
|
||||
// Register a failing handler
|
||||
workerInstance.processor.register('fail_test', async () => {
|
||||
throw new Error('Intentional failure');
|
||||
});
|
||||
|
||||
const job = await jobService.createJob(testTeamId, testUserId, 'fail_test', {});
|
||||
|
||||
const failed = await waitForJobStatus(jobService, job.id, 'failed');
|
||||
expect(failed?.status).toBe('failed');
|
||||
expect((failed?.output as Record<string, unknown>)?.error).toContain('Intentional failure');
|
||||
}, 15_000);
|
||||
|
||||
it('executes cron wrapper jobs through the registered target handler', async () => {
|
||||
workerInstance.processor.register('chat', async (job) => ({
|
||||
response: String(job.data.input.message),
|
||||
}));
|
||||
|
||||
const result = await workerInstance.processor.process({
|
||||
data: {
|
||||
jobId: 'cron-wrapper-1',
|
||||
teamId: testTeamId,
|
||||
userId: testUserId,
|
||||
jobType: 'cron',
|
||||
input: {
|
||||
jobType: 'chat',
|
||||
jobConfig: { message: 'scheduled brief' },
|
||||
},
|
||||
},
|
||||
} as unknown as Job<JobData>, db);
|
||||
|
||||
expect(result).toEqual({
|
||||
response: 'scheduled brief',
|
||||
cron: { delegatedJobType: 'chat' },
|
||||
});
|
||||
}, 15_000);
|
||||
});
|
||||
24
packages/worker/tsconfig.json
Normal file
24
packages/worker/tsconfig.json
Normal file
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2022",
|
||||
"module": "ESNext",
|
||||
"moduleResolution": "bundler",
|
||||
"lib": ["ES2022"],
|
||||
"outDir": "dist",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true,
|
||||
"forceConsistentCasingInFileNames": true,
|
||||
"resolveJsonModule": true,
|
||||
"declaration": true,
|
||||
"declarationMap": true,
|
||||
"sourceMap": true,
|
||||
"composite": true
|
||||
},
|
||||
"references": [
|
||||
{ "path": "../shared" },
|
||||
{ "path": "../server" }
|
||||
],
|
||||
"include": ["src/**/*.ts"],
|
||||
"exclude": ["node_modules", "dist", "tests"]
|
||||
}
|
||||
Reference in New Issue
Block a user