@johpaz/hive-core 1.0.8 → 1.0.10

package/src/agents/registry.test.ts

@@ -0,0 +1,275 @@
+import { describe, it, expect, beforeEach } from "bun:test";
+import { AgentRegistry, type AgentDefinition, type Task } from "./registry";
+
+describe("AgentRegistry", () => {
+  let registry: AgentRegistry;
+
+  const createAgent = (id: string, capabilities: string[], maxTasks = 3): AgentDefinition => ({
+    id,
+    name: `Agent ${id}`,
+    description: `Test agent ${id}`,
+    capabilities,
+    modelConfig: {
+      provider: "openai",
+      model: "gpt-4",
+    },
+    systemPrompt: "You are a helpful assistant",
+    maxConcurrentTasks: maxTasks,
+    currentLoad: 0,
+  });
+
+  const createTask = (id: string, capabilities: string[], priority: Task["priority"] = "normal"): Task => ({
+    id,
+    description: `Task ${id}`,
+    requiredCapabilities: capabilities,
+    priority,
+  });
+
+  beforeEach(() => {
+    registry = new AgentRegistry();
+  });
+
+  describe("register", () => {
+    it("should register an agent", () => {
+      const agent = createAgent("agent-1", ["code", "test"]);
+      registry.register(agent);
+
+      expect(registry.get("agent-1")).toBeDefined();
+      expect(registry.get("agent-1")?.name).toBe("Agent agent-1");
+    });
+
+    it("should update existing agent", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.register(createAgent("agent-1", ["code", "test", "deploy"]));
+
+      const agent = registry.get("agent-1");
+      expect(agent?.capabilities.length).toBe(3);
+    });
+
+    it("should index capabilities", () => {
+      registry.register(createAgent("agent-1", ["code", "test"]));
+      registry.register(createAgent("agent-2", ["test", "deploy"]));
+
+      const codeAgents = registry.getByCapability("code");
+      const testAgents = registry.getByCapability("test");
+
+      expect(codeAgents.length).toBe(1);
+      expect(testAgents.length).toBe(2);
+    });
+  });
+
+  describe("unregister", () => {
+    it("should unregister an agent", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.unregister("agent-1");
+
+      expect(registry.get("agent-1")).toBeUndefined();
+    });
+
+    it("should clean up capability index", () => {
+      registry.register(createAgent("agent-1", ["code", "test"]));
+      registry.register(createAgent("agent-2", ["test"]));
+      registry.unregister("agent-1");
+
+      const testAgents = registry.getByCapability("test");
+      expect(testAgents.length).toBe(1);
+      expect(testAgents[0]?.id).toBe("agent-2");
+    });
+
+    it("should handle non-existent agent", () => {
+      expect(() => registry.unregister("non-existent")).not.toThrow();
+    });
+  });
+
+  describe("getByCapability", () => {
+    it("should return agents with capability", () => {
+      registry.register(createAgent("agent-1", ["code", "test"]));
+      registry.register(createAgent("agent-2", ["test", "deploy"]));
+      registry.register(createAgent("agent-3", ["code"]));
+
+      const testAgents = registry.getByCapability("test");
+      expect(testAgents.length).toBe(2);
+    });
+
+    it("should return empty array for unknown capability", () => {
+      expect(registry.getByCapability("unknown")).toEqual([]);
+    });
+  });
+
+  describe("findBestAgentForTask", () => {
+    it("should find agent with matching capabilities", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.register(createAgent("agent-2", ["test", "deploy"]));
+
+      const task = createTask("task-1", ["test"]);
+      const agent = registry.findBestAgentForTask(task);
+
+      expect(agent?.id).toBe("agent-2");
+    });
+
+    it("should prefer agent with all capabilities", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.register(createAgent("agent-2", ["code", "test", "deploy"]));
+
+      const task = createTask("task-1", ["code", "test"]);
+      const agent = registry.findBestAgentForTask(task);
+
+      expect(agent?.id).toBe("agent-2");
+    });
+
+    it("should consider load when selecting", () => {
+      registry.register(createAgent("agent-1", ["code"], 3));
+      registry.register(createAgent("agent-2", ["code"], 3));
+
+      registry.setLoad("agent-1", 2);
+      registry.setLoad("agent-2", 0);
+
+      const task = createTask("task-1", ["code"]);
+      const agent = registry.findBestAgentForTask(task);
+
+      expect(agent?.id).toBe("agent-2");
+    });
+
+    it("should return least loaded when all at capacity", () => {
+      registry.register(createAgent("agent-1", ["code"], 3));
+      registry.register(createAgent("agent-2", ["code"], 3));
+
+      registry.setLoad("agent-1", 3);
+      registry.setLoad("agent-2", 3);
+
+      const task = createTask("task-1", ["code"]);
+      const agent = registry.findBestAgentForTask(task);
+
+      expect(agent).toBeDefined();
+    });
+
+    it("should return undefined when no agents available", () => {
+      const task = createTask("task-1", ["unknown"]);
+      expect(registry.findBestAgentForTask(task)).toBeUndefined();
+    });
+
+    it("should handle empty capabilities", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      const task = createTask("task-1", []);
+
+      const agent = registry.findBestAgentForTask(task);
+      expect(agent).toBeDefined();
+    });
+  });
+
+  describe("findTeamForTask", () => {
+    it("should find team for multiple capabilities", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.register(createAgent("agent-2", ["test"]));
+      registry.register(createAgent("agent-3", ["deploy"]));
+
+      const task = createTask("task-1", ["code", "test", "deploy"]);
+      const team = registry.findTeamForTask(task);
+
+      expect(team.length).toBe(3);
+      expect(team.map(a => a.id)).toContain("agent-1");
+      expect(team.map(a => a.id)).toContain("agent-2");
+      expect(team.map(a => a.id)).toContain("agent-3");
+    });
+
+    it("should not duplicate agents", () => {
+      registry.register(createAgent("agent-1", ["code", "test", "deploy"]));
+
+      const task = createTask("task-1", ["code", "test"]);
+      const team = registry.findTeamForTask(task);
+
+      expect(team.length).toBe(1);
+    });
+
+    it("should prefer available agents", () => {
+      registry.register(createAgent("agent-1", ["code"], 1));
+      registry.register(createAgent("agent-2", ["code"], 3));
+
+      registry.setLoad("agent-1", 1);
+
+      const task = createTask("task-1", ["code"]);
+      const team = registry.findTeamForTask(task);
+
+      expect(team[0]?.id).toBe("agent-2");
+    });
+
+    it("should skip capabilities with no agents", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+
+      const task = createTask("task-1", ["code", "nonexistent"]);
+      const team = registry.findTeamForTask(task);
+
+      expect(team.length).toBe(1);
+      expect(team[0]?.id).toBe("agent-1");
+    });
+  });
+
+  describe("load management", () => {
+    it("should increment load", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.incrementLoad("agent-1");
+
+      expect(registry.get("agent-1")?.currentLoad).toBe(1);
+    });
+
+    it("should decrement load", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.setLoad("agent-1", 2);
+      registry.decrementLoad("agent-1");
+
+      expect(registry.get("agent-1")?.currentLoad).toBe(1);
+    });
+
+    it("should not go below zero", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.decrementLoad("agent-1");
+
+      expect(registry.get("agent-1")?.currentLoad).toBe(0);
+    });
+
+    it("should set load", () => {
+      registry.register(createAgent("agent-1", ["code"], 5));
+      registry.setLoad("agent-1", 3);
+
+      expect(registry.get("agent-1")?.currentLoad).toBe(3);
+    });
+
+    it("should clamp load to valid range", () => {
+      registry.register(createAgent("agent-1", ["code"], 5));
+      registry.setLoad("agent-1", 10);
+
+      expect(registry.get("agent-1")?.currentLoad).toBe(5);
+
+      registry.setLoad("agent-1", -1);
+      expect(registry.get("agent-1")?.currentLoad).toBe(0);
+    });
+  });
+
+  describe("getStats", () => {
+    it("should return correct stats", () => {
+      registry.register(createAgent("agent-1", ["code", "test"], 3));
+      registry.register(createAgent("agent-2", ["test"], 2));
+      registry.setLoad("agent-1", 1);
+
+      const stats = registry.getStats();
+
+      expect(stats.totalAgents).toBe(2);
+      expect(stats.byCapability["code"]).toBe(1);
+      expect(stats.byCapability["test"]).toBe(2);
+      expect(stats.totalCapacity).toBe(5);
+      expect(stats.currentLoad).toBe(1);
+      expect(stats.availableSlots).toBe(4);
+    });
+  });
+
+  describe("clear", () => {
+    it("should clear all agents", () => {
+      registry.register(createAgent("agent-1", ["code"]));
+      registry.register(createAgent("agent-2", ["test"]));
+      registry.clear();
+
+      expect(registry.getAll().length).toBe(0);
+      expect(registry.getByCapability("code").length).toBe(0);
+    });
+  });
+});

package/src/agents/registry.ts

@@ -0,0 +1,273 @@
+import { logger } from "../utils/logger.ts";
+
+export interface ModelConfig {
+  provider: string;
+  model: string;
+  temperature?: number;
+  maxTokens?: number;
+  baseUrl?: string;
+}
+
+export interface AgentDefinition {
+  id: string;
+  name: string;
+  description: string;
+  capabilities: string[];
+  modelConfig: ModelConfig;
+  systemPrompt: string;
+  maxConcurrentTasks: number;
+  currentLoad: number;
+  priority?: number;
+  tags?: string[];
+  metadata?: Record<string, unknown>;
+}
+
+export interface Task {
+  id: string;
+  description: string;
+  requiredCapabilities: string[];
+  priority: "low" | "normal" | "high" | "critical";
+  payload?: Record<string, unknown>;
+  requiresTeam?: boolean;
+  maxAgents?: number;
+  deadline?: number;
+  parentTaskId?: string;
+}
+
+export interface TaskResult {
+  taskId: string;
+  agentId: string;
+  status: "success" | "partial" | "failed";
+  output: unknown;
+  duration: number;
+  subtasks?: TaskResult[];
+}
+
+export interface AgentStats {
+  totalAgents: number;
+  byCapability: Record<string, number>;
+  totalCapacity: number;
+  currentLoad: number;
+  availableSlots: number;
+}
+
+export class AgentRegistry {
+  private agents: Map<string, AgentDefinition> = new Map();
+  private capabilities: Map<string, Set<string>> = new Map();
+  private log = logger.child("agent-registry");
+
+  register(agent: AgentDefinition): void {
+    if (this.agents.has(agent.id)) {
+      this.log.warn(`Agent ${agent.id} already registered, updating`);
+    }
+
+    this.agents.set(agent.id, { ...agent, currentLoad: agent.currentLoad ?? 0 });
+
+    for (const cap of agent.capabilities) {
+      if (!this.capabilities.has(cap)) {
+        this.capabilities.set(cap, new Set());
+      }
+      this.capabilities.get(cap)!.add(agent.id);
+    }
+
+    this.log.info(`Agent registered: ${agent.id}`, {
+      name: agent.name,
+      capabilities: agent.capabilities,
+      maxTasks: agent.maxConcurrentTasks,
+    });
+  }
+
+  unregister(agentId: string): void {
+    const agent = this.agents.get(agentId);
+    if (!agent) {
+      this.log.warn(`Agent ${agentId} not found`);
+      return;
+    }
+
+    for (const cap of agent.capabilities) {
+      this.capabilities.get(cap)?.delete(agentId);
+      if (this.capabilities.get(cap)?.size === 0) {
+        this.capabilities.delete(cap);
+      }
+    }
+
+    this.agents.delete(agentId);
+    this.log.info(`Agent unregistered: ${agentId}`);
+  }
+
+  get(agentId: string): AgentDefinition | undefined {
+    return this.agents.get(agentId);
+  }
+
+  getAll(): AgentDefinition[] {
+    return Array.from(this.agents.values());
+  }
+
+  getByCapability(capability: string): AgentDefinition[] {
+    const agentIds = this.capabilities.get(capability);
+    if (!agentIds) return [];
+
+    return Array.from(agentIds)
+      .map((id) => this.agents.get(id))
+      .filter((a): a is AgentDefinition => !!a);
+  }
+
+  findBestAgentForTask(task: Task): AgentDefinition | undefined {
+    const candidates = this.findCandidates(task.requiredCapabilities);
+
+    if (candidates.length === 0) {
+      this.log.warn(`No agents found for capabilities: ${task.requiredCapabilities.join(", ")}`);
+      return undefined;
+    }
+
+    const available = candidates.filter(
+      (a) => a.currentLoad < a.maxConcurrentTasks
+    );
+
+    if (available.length === 0) {
+      this.log.warn(`All agents at capacity for task ${task.id}`);
+      return candidates.reduce((best, curr) =>
+        curr.currentLoad < best.currentLoad ? curr : best
+      );
+    }
+
+    const scored = available
+      .map((agent) => ({
+        agent,
+        score: this.calculateScore(agent, task),
+      }))
+      .sort((a, b) => {
+        if (b.score !== a.score) return b.score - a.score;
+        return a.agent.currentLoad - b.agent.currentLoad;
+      });
+
+    return scored[0]?.agent;
+  }
+
+  findTeamForTask(task: Task): AgentDefinition[] {
+    const uniqueCapabilities = [...new Set(task.requiredCapabilities)];
+    const team: AgentDefinition[] = [];
+    const usedAgents = new Set<string>();
+
+    for (const capability of uniqueCapabilities) {
+      const candidates = this.getByCapability(capability)
+        .filter((a) => !usedAgents.has(a.id) && a.currentLoad < a.maxConcurrentTasks);
+
+      if (candidates.length === 0) {
+        const anyAvailable = this.getByCapability(capability)
+          .filter((a) => !usedAgents.has(a.id));
+        if (anyAvailable.length === 0) {
+          this.log.warn(`No agent available for capability: ${capability}`);
+          continue;
+        }
+        const best = anyAvailable.reduce((best, curr) =>
+          curr.currentLoad < best.currentLoad ? curr : best
+        );
+        team.push(best);
+        usedAgents.add(best.id);
+      } else {
+        const best = candidates.reduce((best, curr) =>
+          curr.currentLoad < best.currentLoad ? curr : best
+        );
+        team.push(best);
+        usedAgents.add(best.id);
+      }
+    }
+
+    return team;
+  }
+
+  incrementLoad(agentId: string): void {
+    const agent = this.agents.get(agentId);
+    if (agent) {
+      agent.currentLoad++;
+      this.log.debug(`Agent ${agentId} load increased to ${agent.currentLoad}`);
+    }
+  }
+
+  decrementLoad(agentId: string): void {
+    const agent = this.agents.get(agentId);
+    if (agent && agent.currentLoad > 0) {
+      agent.currentLoad--;
+      this.log.debug(`Agent ${agentId} load decreased to ${agent.currentLoad}`);
+    }
+  }
+
+  setLoad(agentId: string, load: number): void {
+    const agent = this.agents.get(agentId);
+    if (agent) {
+      agent.currentLoad = Math.max(0, Math.min(load, agent.maxConcurrentTasks));
+    }
+  }
+
+  getStats(): AgentStats {
+    const byCapability: Record<string, number> = {};
+    
+    for (const [cap, agents] of this.capabilities) {
+      byCapability[cap] = agents.size;
+    }
+
+    const agents = Array.from(this.agents.values());
+    const totalCapacity = agents.reduce((sum, a) => sum + a.maxConcurrentTasks, 0);
+    const currentLoad = agents.reduce((sum, a) => sum + a.currentLoad, 0);
+
+    return {
+      totalAgents: this.agents.size,
+      byCapability,
+      totalCapacity,
+      currentLoad,
+      availableSlots: totalCapacity - currentLoad,
+    };
+  }
+
+  clear(): void {
+    this.agents.clear();
+    this.capabilities.clear();
+    this.log.info("All agents cleared");
+  }
+
+  private findCandidates(capabilities: string[]): AgentDefinition[] {
+    if (capabilities.length === 0) {
+      return this.getAll();
+    }
+
+    const agentMatchCount = new Map<string, number>();
+
+    for (const cap of capabilities) {
+      const agentIds = this.capabilities.get(cap);
+      if (agentIds) {
+        for (const id of agentIds) {
+          agentMatchCount.set(id, (agentMatchCount.get(id) ?? 0) + 1);
+        }
+      }
+    }
+
+    const sortedIds = Array.from(agentMatchCount.entries())
+      .sort((a, b) => b[1] - a[1])
+      .map(([id]) => id);
+
+    return sortedIds
+      .map((id) => this.agents.get(id))
+      .filter((a): a is AgentDefinition => !!a);
+  }
+
+  private calculateScore(agent: AgentDefinition, task: Task): number {
+    const matchedCapabilities = task.requiredCapabilities.filter((c) =>
+      agent.capabilities.includes(c)
+    ).length;
+
+    const capabilityScore = (matchedCapabilities / task.requiredCapabilities.length) * 100;
+
+    const loadScore = (1 - agent.currentLoad / agent.maxConcurrentTasks) * 50;
+
+    const priorityScore = task.priority === "critical" ? 20 :
+      task.priority === "high" ? 10 :
+      task.priority === "normal" ? 5 : 0;
+
+    const agentPriorityScore = (agent.priority ?? 5) * 2;
+
+    return capabilityScore + loadScore + priorityScore + agentPriorityScore;
+  }
+}
+
+export const agentRegistry = new AgentRegistry();