feat: add base multi-agent collaboration framework (#294)
Introduce the foundation for multi-agent collaboration in a new pkg/agent/multi package. This is a non-invasive addition that creates new abstractions without modifying any existing code. - Agent interface with Name, Role, SystemPrompt, Capabilities, Tools - BaseAgent embeddable struct for concrete agent implementations - SharedContext blackboard for inter-agent data sharing and event log - AgentRegistry for lifecycle management and capability-based routing - HandoffRequest/HandoffResult for task delegation between agents - Comprehensive test suite (24 tests) including concurrency and integration
This commit is contained in:
Edouard CLAUDE
parent
3390576eea
commit
7ab8431b33
4 changed files with 1161 additions and 0 deletions
124
pkg/agent/multi/agent.go
Normal file
124
pkg/agent/multi/agent.go
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
// PicoClaw - Ultra-lightweight personal AI agent
|
||||
// License: MIT
|
||||
//
|
||||
// Copyright (c) 2026 PicoClaw contributors
|
||||
|
||||
// Package multi provides the foundation for multi-agent collaboration.
|
||||
// It defines the Agent interface, shared context, and agent registry
|
||||
// that enable multiple specialized agents to work together within
|
||||
// a single PicoClaw session.
|
||||
//
|
||||
// This package is designed to be non-invasive: it introduces new
|
||||
// abstractions without modifying the existing AgentLoop or SubagentManager.
|
||||
// The existing subagent system can be gradually migrated to use these
|
||||
// interfaces.
|
||||
package multi
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
)
|
||||
|
||||
// Agent defines the interface that all agents must implement.
|
||||
// Each agent has a unique name, a role description, a system prompt,
|
||||
// and a set of capabilities that determine which tasks it can handle.
|
||||
type Agent interface {
|
||||
// Name returns the unique identifier for this agent.
|
||||
Name() string
|
||||
|
||||
// Role returns a human-readable description of what this agent does.
|
||||
Role() string
|
||||
|
||||
// SystemPrompt returns the system prompt used when this agent
|
||||
// interacts with the LLM.
|
||||
SystemPrompt() string
|
||||
|
||||
// Capabilities returns the list of capability tags this agent supports.
|
||||
// These are used by the registry to match agents to tasks.
|
||||
// Example: ["code", "search", "file_operations"]
|
||||
Capabilities() []string
|
||||
|
||||
// Tools returns the tool registry available to this agent.
|
||||
// Each agent can have a different set of tools.
|
||||
Tools() *tools.ToolRegistry
|
||||
|
||||
// Execute runs the agent on the given task within the provided context.
|
||||
// The shared context allows reading/writing data visible to other agents.
|
||||
// Returns the agent's response content and any error.
|
||||
Execute(ctx context.Context, task string, shared *SharedContext) (string, error)
|
||||
}
|
||||
|
||||
// AgentConfig holds configuration for creating a BaseAgent.
|
||||
type AgentConfig struct {
|
||||
// Name is the unique identifier for the agent.
|
||||
Name string
|
||||
|
||||
// Role describes what this agent specializes in.
|
||||
Role string
|
||||
|
||||
// SystemPrompt is the prompt sent to the LLM.
|
||||
SystemPrompt string
|
||||
|
||||
// Capabilities lists the capability tags.
|
||||
Capabilities []string
|
||||
}
|
||||
|
||||
// BaseAgent provides a minimal Agent implementation that can be embedded
|
||||
// in concrete agent types. It handles the common fields (name, role, prompt,
|
||||
// capabilities) and leaves Execute to be implemented by the concrete type.
|
||||
type BaseAgent struct {
|
||||
config AgentConfig
|
||||
tools *tools.ToolRegistry
|
||||
}
|
||||
|
||||
// NewBaseAgent creates a new BaseAgent with the given configuration.
|
||||
func NewBaseAgent(cfg AgentConfig, registry *tools.ToolRegistry) *BaseAgent {
|
||||
if registry == nil {
|
||||
registry = tools.NewToolRegistry()
|
||||
}
|
||||
return &BaseAgent{
|
||||
config: cfg,
|
||||
tools: registry,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *BaseAgent) Name() string { return a.config.Name }
|
||||
func (a *BaseAgent) Role() string { return a.config.Role }
|
||||
func (a *BaseAgent) SystemPrompt() string { return a.config.SystemPrompt }
|
||||
func (a *BaseAgent) Capabilities() []string { return a.config.Capabilities }
|
||||
func (a *BaseAgent) Tools() *tools.ToolRegistry { return a.tools }
|
||||
|
||||
// HandoffRequest represents a request to delegate a task from one agent
|
||||
// to another. It carries the task description and optional metadata
|
||||
// for routing.
|
||||
type HandoffRequest struct {
|
||||
// FromAgent is the name of the agent delegating the task.
|
||||
FromAgent string
|
||||
|
||||
// ToAgent is the name of the target agent. If empty, the registry
|
||||
// will select the best agent based on RequiredCapability.
|
||||
ToAgent string
|
||||
|
||||
// RequiredCapability is used for capability-based routing when
|
||||
// ToAgent is not specified.
|
||||
RequiredCapability string
|
||||
|
||||
// Task is the description of what needs to be done.
|
||||
Task string
|
||||
|
||||
// Context carries additional key-value data for the target agent.
|
||||
Context map[string]interface{}
|
||||
}
|
||||
|
||||
// HandoffResult contains the outcome of a hand-off operation.
|
||||
type HandoffResult struct {
|
||||
// AgentName is the name of the agent that handled the task.
|
||||
AgentName string
|
||||
|
||||
// Content is the response produced by the agent.
|
||||
Content string
|
||||
|
||||
// Err is set if the hand-off or execution failed.
|
||||
Err error
|
||||
}
|
||||
151
pkg/agent/multi/context.go
Normal file
151
pkg/agent/multi/context.go
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
// PicoClaw - Ultra-lightweight personal AI agent
|
||||
// License: MIT
|
||||
//
|
||||
// Copyright (c) 2026 PicoClaw contributors
|
||||
|
||||
package multi
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
// SharedContext implements a thread-safe blackboard pattern where multiple
|
||||
// agents can read from and write to a common session context.
|
||||
//
|
||||
// It provides:
|
||||
// - Key-value storage for arbitrary data sharing between agents
|
||||
// - An append-only event log for agent activity tracking
|
||||
// - Thread-safe access via read-write mutex
|
||||
//
|
||||
// This is intentionally simple and in-memory. Future iterations may add
|
||||
// persistence, TTL, or namespace isolation.
|
||||
type SharedContext struct {
|
||||
mu sync.RWMutex
|
||||
data map[string]interface{}
|
||||
events []Event
|
||||
}
|
||||
|
||||
// Event records an action taken by an agent within the shared context.
|
||||
// Events are append-only and provide an audit trail of agent activity.
|
||||
type Event struct {
|
||||
// Agent is the name of the agent that produced this event.
|
||||
Agent string
|
||||
|
||||
// Type categorizes the event (e.g., "handoff", "result", "error").
|
||||
Type string
|
||||
|
||||
// Content is the event payload.
|
||||
Content string
|
||||
}
|
||||
|
||||
// NewSharedContext creates a new empty SharedContext.
|
||||
func NewSharedContext() *SharedContext {
|
||||
return &SharedContext{
|
||||
data: make(map[string]interface{}),
|
||||
events: make([]Event, 0),
|
||||
}
|
||||
}
|
||||
|
||||
// Set stores a value in the shared context under the given key.
|
||||
// Overwrites any existing value for the same key.
|
||||
func (sc *SharedContext) Set(key string, value interface{}) {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
sc.data[key] = value
|
||||
}
|
||||
|
||||
// Get retrieves a value from the shared context.
|
||||
// Returns the value and true if found, nil and false otherwise.
|
||||
func (sc *SharedContext) Get(key string) (interface{}, bool) {
|
||||
sc.mu.RLock()
|
||||
defer sc.mu.RUnlock()
|
||||
v, ok := sc.data[key]
|
||||
return v, ok
|
||||
}
|
||||
|
||||
// GetString retrieves a string value from the shared context.
|
||||
// Returns empty string if the key doesn't exist or isn't a string.
|
||||
func (sc *SharedContext) GetString(key string) string {
|
||||
v, ok := sc.Get(key)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
s, _ := v.(string)
|
||||
return s
|
||||
}
|
||||
|
||||
// Delete removes a key from the shared context.
|
||||
func (sc *SharedContext) Delete(key string) {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
delete(sc.data, key)
|
||||
}
|
||||
|
||||
// Keys returns all keys currently stored in the shared context.
|
||||
func (sc *SharedContext) Keys() []string {
|
||||
sc.mu.RLock()
|
||||
defer sc.mu.RUnlock()
|
||||
keys := make([]string, 0, len(sc.data))
|
||||
for k := range sc.data {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
// AddEvent appends an event to the shared context's event log.
|
||||
func (sc *SharedContext) AddEvent(agent, eventType, content string) {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
sc.events = append(sc.events, Event{
|
||||
Agent: agent,
|
||||
Type: eventType,
|
||||
Content: content,
|
||||
})
|
||||
}
|
||||
|
||||
// Events returns a copy of all events in the shared context.
|
||||
func (sc *SharedContext) Events() []Event {
|
||||
sc.mu.RLock()
|
||||
defer sc.mu.RUnlock()
|
||||
cp := make([]Event, len(sc.events))
|
||||
copy(cp, sc.events)
|
||||
return cp
|
||||
}
|
||||
|
||||
// EventsByAgent returns all events produced by the given agent.
|
||||
func (sc *SharedContext) EventsByAgent(agent string) []Event {
|
||||
sc.mu.RLock()
|
||||
defer sc.mu.RUnlock()
|
||||
var filtered []Event
|
||||
for _, e := range sc.events {
|
||||
if e.Agent == agent {
|
||||
filtered = append(filtered, e)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
// EventsByType returns all events of the given type.
|
||||
func (sc *SharedContext) EventsByType(eventType string) []Event {
|
||||
sc.mu.RLock()
|
||||
defer sc.mu.RUnlock()
|
||||
var filtered []Event
|
||||
for _, e := range sc.events {
|
||||
if e.Type == eventType {
|
||||
filtered = append(filtered, e)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
// Snapshot returns a shallow copy of the entire data map.
|
||||
// Useful for debugging or serialization.
|
||||
func (sc *SharedContext) Snapshot() map[string]interface{} {
|
||||
sc.mu.RLock()
|
||||
defer sc.mu.RUnlock()
|
||||
snap := make(map[string]interface{}, len(sc.data))
|
||||
for k, v := range sc.data {
|
||||
snap[k] = v
|
||||
}
|
||||
return snap
|
||||
}
|
||||
633
pkg/agent/multi/multi_test.go
Normal file
633
pkg/agent/multi/multi_test.go
Normal file
|
|
@ -0,0 +1,633 @@
|
|||
// PicoClaw - Ultra-lightweight personal AI agent
|
||||
// License: MIT
|
||||
//
|
||||
// Copyright (c) 2026 PicoClaw contributors
|
||||
|
||||
package multi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
)
|
||||
|
||||
// mockAgent is a simple Agent implementation for testing.
|
||||
type mockAgent struct {
|
||||
name string
|
||||
role string
|
||||
systemPrompt string
|
||||
capabilities []string
|
||||
toolRegistry *tools.ToolRegistry
|
||||
executeFunc func(ctx context.Context, task string, shared *SharedContext) (string, error)
|
||||
}
|
||||
|
||||
func newMockAgent(name, role string, capabilities []string) *mockAgent {
|
||||
return &mockAgent{
|
||||
name: name,
|
||||
role: role,
|
||||
systemPrompt: fmt.Sprintf("You are %s, a %s agent.", name, role),
|
||||
capabilities: capabilities,
|
||||
toolRegistry: tools.NewToolRegistry(),
|
||||
executeFunc: func(ctx context.Context, task string, shared *SharedContext) (string, error) {
|
||||
return fmt.Sprintf("[%s] completed: %s", name, task), nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mockAgent) Name() string { return m.name }
|
||||
func (m *mockAgent) Role() string { return m.role }
|
||||
func (m *mockAgent) SystemPrompt() string { return m.systemPrompt }
|
||||
func (m *mockAgent) Capabilities() []string { return m.capabilities }
|
||||
func (m *mockAgent) Tools() *tools.ToolRegistry { return m.toolRegistry }
|
||||
func (m *mockAgent) Execute(ctx context.Context, task string, shared *SharedContext) (string, error) {
|
||||
return m.executeFunc(ctx, task, shared)
|
||||
}
|
||||
|
||||
// --- SharedContext Tests ---
|
||||
|
||||
func TestSharedContext_SetGet(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
|
||||
sc.Set("key1", "value1")
|
||||
sc.Set("key2", 42)
|
||||
|
||||
v1, ok := sc.Get("key1")
|
||||
if !ok || v1 != "value1" {
|
||||
t.Errorf("expected key1=value1, got %v (ok=%v)", v1, ok)
|
||||
}
|
||||
|
||||
v2, ok := sc.Get("key2")
|
||||
if !ok || v2 != 42 {
|
||||
t.Errorf("expected key2=42, got %v (ok=%v)", v2, ok)
|
||||
}
|
||||
|
||||
_, ok = sc.Get("nonexistent")
|
||||
if ok {
|
||||
t.Error("expected nonexistent key to return false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedContext_GetString(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
|
||||
sc.Set("str", "hello")
|
||||
sc.Set("num", 42)
|
||||
|
||||
if s := sc.GetString("str"); s != "hello" {
|
||||
t.Errorf("expected 'hello', got %q", s)
|
||||
}
|
||||
|
||||
if s := sc.GetString("num"); s != "" {
|
||||
t.Errorf("expected empty string for non-string, got %q", s)
|
||||
}
|
||||
|
||||
if s := sc.GetString("missing"); s != "" {
|
||||
t.Errorf("expected empty string for missing key, got %q", s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedContext_Delete(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
|
||||
sc.Set("key", "value")
|
||||
sc.Delete("key")
|
||||
|
||||
_, ok := sc.Get("key")
|
||||
if ok {
|
||||
t.Error("expected deleted key to return false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedContext_Keys(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
|
||||
sc.Set("a", 1)
|
||||
sc.Set("b", 2)
|
||||
sc.Set("c", 3)
|
||||
|
||||
keys := sc.Keys()
|
||||
if len(keys) != 3 {
|
||||
t.Errorf("expected 3 keys, got %d", len(keys))
|
||||
}
|
||||
|
||||
keySet := make(map[string]bool)
|
||||
for _, k := range keys {
|
||||
keySet[k] = true
|
||||
}
|
||||
for _, expected := range []string{"a", "b", "c"} {
|
||||
if !keySet[expected] {
|
||||
t.Errorf("expected key %q in keys", expected)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedContext_Events(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
|
||||
sc.AddEvent("agent1", "handoff", "delegating to agent2")
|
||||
sc.AddEvent("agent2", "result", "task completed")
|
||||
sc.AddEvent("agent1", "error", "something failed")
|
||||
|
||||
events := sc.Events()
|
||||
if len(events) != 3 {
|
||||
t.Fatalf("expected 3 events, got %d", len(events))
|
||||
}
|
||||
|
||||
if events[0].Agent != "agent1" || events[0].Type != "handoff" {
|
||||
t.Errorf("unexpected first event: %+v", events[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedContext_EventsByAgent(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
|
||||
sc.AddEvent("agent1", "handoff", "task1")
|
||||
sc.AddEvent("agent2", "result", "done")
|
||||
sc.AddEvent("agent1", "result", "task2")
|
||||
|
||||
agent1Events := sc.EventsByAgent("agent1")
|
||||
if len(agent1Events) != 2 {
|
||||
t.Errorf("expected 2 events for agent1, got %d", len(agent1Events))
|
||||
}
|
||||
|
||||
agent2Events := sc.EventsByAgent("agent2")
|
||||
if len(agent2Events) != 1 {
|
||||
t.Errorf("expected 1 event for agent2, got %d", len(agent2Events))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedContext_EventsByType(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
|
||||
sc.AddEvent("a1", "result", "r1")
|
||||
sc.AddEvent("a2", "error", "e1")
|
||||
sc.AddEvent("a3", "result", "r2")
|
||||
|
||||
results := sc.EventsByType("result")
|
||||
if len(results) != 2 {
|
||||
t.Errorf("expected 2 result events, got %d", len(results))
|
||||
}
|
||||
|
||||
errors := sc.EventsByType("error")
|
||||
if len(errors) != 1 {
|
||||
t.Errorf("expected 1 error event, got %d", len(errors))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedContext_Snapshot(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
|
||||
sc.Set("key", "value")
|
||||
snap := sc.Snapshot()
|
||||
|
||||
// Modify original
|
||||
sc.Set("key", "changed")
|
||||
|
||||
// Snapshot should be independent
|
||||
if snap["key"] != "value" {
|
||||
t.Errorf("snapshot should be independent, got %v", snap["key"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedContext_ConcurrentAccess(t *testing.T) {
|
||||
sc := NewSharedContext()
|
||||
var wg sync.WaitGroup
|
||||
|
||||
// Concurrent writes
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
sc.Set(fmt.Sprintf("key-%d", i), i)
|
||||
sc.AddEvent(fmt.Sprintf("agent-%d", i), "write", "data")
|
||||
}(i)
|
||||
}
|
||||
|
||||
// Concurrent reads
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
sc.Get(fmt.Sprintf("key-%d", i))
|
||||
sc.Keys()
|
||||
sc.Events()
|
||||
}(i)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
keys := sc.Keys()
|
||||
if len(keys) != 100 {
|
||||
t.Errorf("expected 100 keys after concurrent writes, got %d", len(keys))
|
||||
}
|
||||
}
|
||||
|
||||
// --- BaseAgent Tests ---
|
||||
|
||||
func TestBaseAgent_Fields(t *testing.T) {
|
||||
registry := tools.NewToolRegistry()
|
||||
agent := NewBaseAgent(AgentConfig{
|
||||
Name: "coder",
|
||||
Role: "Code generation and review",
|
||||
SystemPrompt: "You are a coding agent.",
|
||||
Capabilities: []string{"code", "review"},
|
||||
}, registry)
|
||||
|
||||
if agent.Name() != "coder" {
|
||||
t.Errorf("expected name 'coder', got %q", agent.Name())
|
||||
}
|
||||
if agent.Role() != "Code generation and review" {
|
||||
t.Errorf("unexpected role: %q", agent.Role())
|
||||
}
|
||||
if agent.SystemPrompt() != "You are a coding agent." {
|
||||
t.Errorf("unexpected system prompt: %q", agent.SystemPrompt())
|
||||
}
|
||||
caps := agent.Capabilities()
|
||||
if len(caps) != 2 || caps[0] != "code" || caps[1] != "review" {
|
||||
t.Errorf("unexpected capabilities: %v", caps)
|
||||
}
|
||||
if agent.Tools() != registry {
|
||||
t.Error("expected same tool registry")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBaseAgent_NilRegistry(t *testing.T) {
|
||||
agent := NewBaseAgent(AgentConfig{Name: "test"}, nil)
|
||||
if agent.Tools() == nil {
|
||||
t.Error("expected non-nil default tool registry")
|
||||
}
|
||||
}
|
||||
|
||||
// --- AgentRegistry Tests ---
|
||||
|
||||
func TestAgentRegistry_Register(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
agent := newMockAgent("coder", "coding", []string{"code"})
|
||||
|
||||
err := r.Register(agent)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Duplicate registration
|
||||
err = r.Register(agent)
|
||||
if err == nil {
|
||||
t.Error("expected error on duplicate registration")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_Unregister(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
agent := newMockAgent("coder", "coding", []string{"code"})
|
||||
r.Register(agent)
|
||||
|
||||
err := r.Unregister("coder")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
// Verify agent is gone
|
||||
if r.Get("coder") != nil {
|
||||
t.Error("expected nil after unregister")
|
||||
}
|
||||
|
||||
// Unregister nonexistent
|
||||
err = r.Unregister("nonexistent")
|
||||
if err == nil {
|
||||
t.Error("expected error for nonexistent agent")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_Get(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
agent := newMockAgent("searcher", "search", []string{"search"})
|
||||
r.Register(agent)
|
||||
|
||||
got := r.Get("searcher")
|
||||
if got == nil {
|
||||
t.Fatal("expected non-nil agent")
|
||||
}
|
||||
if got.Name() != "searcher" {
|
||||
t.Errorf("expected 'searcher', got %q", got.Name())
|
||||
}
|
||||
|
||||
if r.Get("nonexistent") != nil {
|
||||
t.Error("expected nil for nonexistent agent")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_List(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
r.Register(newMockAgent("a", "role-a", nil))
|
||||
r.Register(newMockAgent("b", "role-b", nil))
|
||||
r.Register(newMockAgent("c", "role-c", nil))
|
||||
|
||||
names := r.List()
|
||||
if len(names) != 3 {
|
||||
t.Errorf("expected 3 agents, got %d", len(names))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_FindByCapability(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
r.Register(newMockAgent("coder", "coding", []string{"code", "review"}))
|
||||
r.Register(newMockAgent("searcher", "searching", []string{"search", "web"}))
|
||||
r.Register(newMockAgent("reviewer", "reviewing", []string{"review"}))
|
||||
|
||||
codeAgents := r.FindByCapability("code")
|
||||
if len(codeAgents) != 1 {
|
||||
t.Errorf("expected 1 agent with 'code', got %d", len(codeAgents))
|
||||
}
|
||||
|
||||
reviewAgents := r.FindByCapability("review")
|
||||
if len(reviewAgents) != 2 {
|
||||
t.Errorf("expected 2 agents with 'review', got %d", len(reviewAgents))
|
||||
}
|
||||
|
||||
noneAgents := r.FindByCapability("nonexistent")
|
||||
if len(noneAgents) != 0 {
|
||||
t.Errorf("expected 0 agents, got %d", len(noneAgents))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_SharedContext(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
sc := r.SharedContext()
|
||||
|
||||
if sc == nil {
|
||||
t.Fatal("expected non-nil shared context")
|
||||
}
|
||||
|
||||
// Verify it's the same instance
|
||||
if r.SharedContext() != sc {
|
||||
t.Error("expected same shared context instance")
|
||||
}
|
||||
}
|
||||
|
||||
// --- Handoff Tests ---
|
||||
|
||||
func TestAgentRegistry_Handoff_DirectRouting(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
r.Register(newMockAgent("coder", "coding", []string{"code"}))
|
||||
|
||||
result := r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
ToAgent: "coder",
|
||||
Task: "write a function",
|
||||
})
|
||||
|
||||
if result.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", result.Err)
|
||||
}
|
||||
if result.AgentName != "coder" {
|
||||
t.Errorf("expected agent 'coder', got %q", result.AgentName)
|
||||
}
|
||||
if result.Content != "[coder] completed: write a function" {
|
||||
t.Errorf("unexpected content: %q", result.Content)
|
||||
}
|
||||
|
||||
// Verify events were recorded
|
||||
events := r.SharedContext().Events()
|
||||
if len(events) < 2 {
|
||||
t.Fatalf("expected at least 2 events, got %d", len(events))
|
||||
}
|
||||
if events[0].Type != "handoff" {
|
||||
t.Errorf("expected first event type 'handoff', got %q", events[0].Type)
|
||||
}
|
||||
if events[1].Type != "result" {
|
||||
t.Errorf("expected second event type 'result', got %q", events[1].Type)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_Handoff_CapabilityRouting(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
r.Register(newMockAgent("coder", "coding", []string{"code"}))
|
||||
r.Register(newMockAgent("searcher", "searching", []string{"search", "web"}))
|
||||
|
||||
result := r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
RequiredCapability: "search",
|
||||
Task: "find documentation",
|
||||
})
|
||||
|
||||
if result.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", result.Err)
|
||||
}
|
||||
if result.AgentName != "searcher" {
|
||||
t.Errorf("expected agent 'searcher', got %q", result.AgentName)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_Handoff_NotFound(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
|
||||
// Target agent not found
|
||||
result := r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
ToAgent: "nonexistent",
|
||||
Task: "do something",
|
||||
})
|
||||
if result.Err == nil {
|
||||
t.Error("expected error for nonexistent agent")
|
||||
}
|
||||
|
||||
// Capability not found
|
||||
result = r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
RequiredCapability: "nonexistent",
|
||||
Task: "do something",
|
||||
})
|
||||
if result.Err == nil {
|
||||
t.Error("expected error for nonexistent capability")
|
||||
}
|
||||
|
||||
// Neither ToAgent nor RequiredCapability
|
||||
result = r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
Task: "do something",
|
||||
})
|
||||
if result.Err == nil {
|
||||
t.Error("expected error when neither routing field is set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_Handoff_ExecutionError(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
|
||||
failAgent := newMockAgent("failer", "failing", []string{"fail"})
|
||||
failAgent.executeFunc = func(ctx context.Context, task string, shared *SharedContext) (string, error) {
|
||||
return "", fmt.Errorf("execution failed: %s", task)
|
||||
}
|
||||
r.Register(failAgent)
|
||||
|
||||
result := r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
ToAgent: "failer",
|
||||
Task: "break things",
|
||||
})
|
||||
|
||||
if result.Err == nil {
|
||||
t.Fatal("expected execution error")
|
||||
}
|
||||
if result.AgentName != "failer" {
|
||||
t.Errorf("expected agent 'failer', got %q", result.AgentName)
|
||||
}
|
||||
|
||||
// Verify error event was recorded
|
||||
errorEvents := r.SharedContext().EventsByType("error")
|
||||
if len(errorEvents) == 0 {
|
||||
t.Error("expected error event to be recorded")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_Handoff_ContextPassing(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
|
||||
// Agent that reads from shared context
|
||||
reader := newMockAgent("reader", "reading", []string{"read"})
|
||||
reader.executeFunc = func(ctx context.Context, task string, shared *SharedContext) (string, error) {
|
||||
val := shared.GetString("input_data")
|
||||
return fmt.Sprintf("read: %s", val), nil
|
||||
}
|
||||
r.Register(reader)
|
||||
|
||||
result := r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
ToAgent: "reader",
|
||||
Task: "process data",
|
||||
Context: map[string]interface{}{"input_data": "hello world"},
|
||||
})
|
||||
|
||||
if result.Err != nil {
|
||||
t.Fatalf("unexpected error: %v", result.Err)
|
||||
}
|
||||
if result.Content != "read: hello world" {
|
||||
t.Errorf("unexpected content: %q", result.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_Handoff_AgentStateTransition(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
|
||||
// Agent that checks its own state via a channel
|
||||
stateChecked := make(chan AgentState, 1)
|
||||
statefulAgent := newMockAgent("stateful", "checking", []string{"check"})
|
||||
statefulAgent.executeFunc = func(ctx context.Context, task string, shared *SharedContext) (string, error) {
|
||||
state, _ := r.GetAgentState("stateful")
|
||||
stateChecked <- state
|
||||
return "done", nil
|
||||
}
|
||||
r.Register(statefulAgent)
|
||||
|
||||
// Before hand-off: idle
|
||||
state, ok := r.GetAgentState("stateful")
|
||||
if !ok || state != AgentIdle {
|
||||
t.Errorf("expected idle state before handoff, got %v", state)
|
||||
}
|
||||
|
||||
r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
ToAgent: "stateful",
|
||||
Task: "check state",
|
||||
})
|
||||
|
||||
// During execution: should have been active
|
||||
duringState := <-stateChecked
|
||||
if duringState != AgentActive {
|
||||
t.Errorf("expected active state during execution, got %v", duringState)
|
||||
}
|
||||
|
||||
// After hand-off: idle again
|
||||
state, _ = r.GetAgentState("stateful")
|
||||
if state != AgentIdle {
|
||||
t.Errorf("expected idle state after handoff, got %v", state)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentRegistry_Handoff_ContextCancellation(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
|
||||
// Agent that respects context cancellation
|
||||
cancelAgent := newMockAgent("cancellable", "cancelling", []string{"cancel"})
|
||||
cancelAgent.executeFunc = func(ctx context.Context, task string, shared *SharedContext) (string, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
default:
|
||||
return "completed before cancel", nil
|
||||
}
|
||||
}
|
||||
r.Register(cancelAgent)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // Cancel immediately
|
||||
|
||||
result := r.Handoff(ctx, HandoffRequest{
|
||||
FromAgent: "main",
|
||||
ToAgent: "cancellable",
|
||||
Task: "should be cancelled",
|
||||
})
|
||||
|
||||
if result.Err == nil {
|
||||
// The agent might complete before checking ctx, which is fine
|
||||
// Just verify it ran
|
||||
if result.Content == "" {
|
||||
t.Error("expected some content")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Integration Test ---
|
||||
|
||||
func TestMultiAgent_Integration(t *testing.T) {
|
||||
r := NewAgentRegistry()
|
||||
|
||||
// Register a chain of agents
|
||||
analyzer := newMockAgent("analyzer", "analysis", []string{"analyze"})
|
||||
analyzer.executeFunc = func(ctx context.Context, task string, shared *SharedContext) (string, error) {
|
||||
shared.Set("analysis_result", "code needs refactoring")
|
||||
return "Analysis complete", nil
|
||||
}
|
||||
|
||||
coder := newMockAgent("coder", "coding", []string{"code"})
|
||||
coder.executeFunc = func(ctx context.Context, task string, shared *SharedContext) (string, error) {
|
||||
analysis := shared.GetString("analysis_result")
|
||||
return fmt.Sprintf("Applied fix based on: %s", analysis), nil
|
||||
}
|
||||
|
||||
r.Register(analyzer)
|
||||
r.Register(coder)
|
||||
|
||||
// Step 1: Analyze
|
||||
result1 := r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
RequiredCapability: "analyze",
|
||||
Task: "review the codebase",
|
||||
})
|
||||
if result1.Err != nil {
|
||||
t.Fatalf("analysis failed: %v", result1.Err)
|
||||
}
|
||||
|
||||
// Step 2: Code fix based on analysis result (already in shared context)
|
||||
result2 := r.Handoff(context.Background(), HandoffRequest{
|
||||
FromAgent: "main",
|
||||
RequiredCapability: "code",
|
||||
Task: "fix the issues found",
|
||||
})
|
||||
if result2.Err != nil {
|
||||
t.Fatalf("coding failed: %v", result2.Err)
|
||||
}
|
||||
|
||||
if result2.Content != "Applied fix based on: code needs refactoring" {
|
||||
t.Errorf("unexpected content: %q", result2.Content)
|
||||
}
|
||||
|
||||
// Verify full event trail
|
||||
events := r.SharedContext().Events()
|
||||
if len(events) != 4 { // 2 handoffs + 2 results
|
||||
t.Errorf("expected 4 events, got %d", len(events))
|
||||
}
|
||||
}
|
||||
253
pkg/agent/multi/registry.go
Normal file
253
pkg/agent/multi/registry.go
Normal file
|
|
@ -0,0 +1,253 @@
|
|||
// PicoClaw - Ultra-lightweight personal AI agent
|
||||
// License: MIT
|
||||
//
|
||||
// Copyright (c) 2026 PicoClaw contributors
|
||||
|
||||
package multi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
)
|
||||
|
||||
// AgentState represents the current state of an agent in the registry.
|
||||
type AgentState int
|
||||
|
||||
const (
|
||||
// AgentIdle means the agent is registered but not currently executing.
|
||||
AgentIdle AgentState = iota
|
||||
|
||||
// AgentActive means the agent is currently executing a task.
|
||||
AgentActive
|
||||
)
|
||||
|
||||
// agentEntry holds an agent and its runtime state within the registry.
|
||||
type agentEntry struct {
|
||||
agent Agent
|
||||
state AgentState
|
||||
}
|
||||
|
||||
// AgentRegistry manages the lifecycle of agents and provides capability-based
|
||||
// routing for hand-off requests. It is the central coordinator for multi-agent
|
||||
// collaboration within a session.
|
||||
//
|
||||
// Thread-safe: all operations are protected by a read-write mutex.
|
||||
type AgentRegistry struct {
|
||||
mu sync.RWMutex
|
||||
agents map[string]*agentEntry
|
||||
shared *SharedContext
|
||||
}
|
||||
|
||||
// NewAgentRegistry creates a new AgentRegistry with a fresh SharedContext.
|
||||
func NewAgentRegistry() *AgentRegistry {
|
||||
return &AgentRegistry{
|
||||
agents: make(map[string]*agentEntry),
|
||||
shared: NewSharedContext(),
|
||||
}
|
||||
}
|
||||
|
||||
// Register adds an agent to the registry. Returns an error if an agent
|
||||
// with the same name is already registered.
|
||||
func (r *AgentRegistry) Register(agent Agent) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
name := agent.Name()
|
||||
if _, exists := r.agents[name]; exists {
|
||||
return fmt.Errorf("agent %q already registered", name)
|
||||
}
|
||||
|
||||
r.agents[name] = &agentEntry{
|
||||
agent: agent,
|
||||
state: AgentIdle,
|
||||
}
|
||||
|
||||
logger.InfoCF("multi", "Agent registered",
|
||||
map[string]interface{}{
|
||||
"name": name,
|
||||
"role": agent.Role(),
|
||||
"capabilities": agent.Capabilities(),
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Unregister removes an agent from the registry.
|
||||
// Returns an error if the agent is currently active.
|
||||
func (r *AgentRegistry) Unregister(name string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
entry, exists := r.agents[name]
|
||||
if !exists {
|
||||
return fmt.Errorf("agent %q not found", name)
|
||||
}
|
||||
|
||||
if entry.state == AgentActive {
|
||||
return fmt.Errorf("cannot unregister active agent %q", name)
|
||||
}
|
||||
|
||||
delete(r.agents, name)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get returns the agent with the given name, or nil if not found.
|
||||
func (r *AgentRegistry) Get(name string) Agent {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
entry, ok := r.agents[name]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return entry.agent
|
||||
}
|
||||
|
||||
// List returns the names of all registered agents.
|
||||
func (r *AgentRegistry) List() []string {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
names := make([]string, 0, len(r.agents))
|
||||
for name := range r.agents {
|
||||
names = append(names, name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
// FindByCapability returns all agents that have the specified capability.
|
||||
func (r *AgentRegistry) FindByCapability(capability string) []Agent {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
var matches []Agent
|
||||
for _, entry := range r.agents {
|
||||
for _, cap := range entry.agent.Capabilities() {
|
||||
if cap == capability {
|
||||
matches = append(matches, entry.agent)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return matches
|
||||
}
|
||||
|
||||
// SharedContext returns the registry's shared context instance.
|
||||
func (r *AgentRegistry) SharedContext() *SharedContext {
|
||||
return r.shared
|
||||
}
|
||||
|
||||
// Handoff delegates a task to another agent based on the HandoffRequest.
|
||||
//
|
||||
// Routing logic:
|
||||
// 1. If ToAgent is specified, route directly to that agent.
|
||||
// 2. If RequiredCapability is specified, find the first idle agent
|
||||
// with that capability.
|
||||
// 3. If no suitable agent is found, return an error.
|
||||
//
|
||||
// The hand-off records events in the shared context for traceability.
|
||||
func (r *AgentRegistry) Handoff(ctx context.Context, req HandoffRequest) *HandoffResult {
|
||||
// Inject hand-off context data into shared context
|
||||
if req.Context != nil {
|
||||
for k, v := range req.Context {
|
||||
r.shared.Set(k, v)
|
||||
}
|
||||
}
|
||||
|
||||
// Record the hand-off request as an event
|
||||
r.shared.AddEvent(req.FromAgent, "handoff",
|
||||
fmt.Sprintf("delegating task to %s (capability: %s): %s",
|
||||
req.ToAgent, req.RequiredCapability, req.Task))
|
||||
|
||||
// Resolve target agent
|
||||
target, err := r.resolveTarget(req)
|
||||
if err != nil {
|
||||
r.shared.AddEvent(req.FromAgent, "error", err.Error())
|
||||
return &HandoffResult{Err: err}
|
||||
}
|
||||
|
||||
// Mark agent as active
|
||||
r.setAgentState(target.Name(), AgentActive)
|
||||
defer r.setAgentState(target.Name(), AgentIdle)
|
||||
|
||||
logger.InfoCF("multi", "Executing hand-off",
|
||||
map[string]interface{}{
|
||||
"from": req.FromAgent,
|
||||
"to": target.Name(),
|
||||
"task_len": len(req.Task),
|
||||
"capability": req.RequiredCapability,
|
||||
})
|
||||
|
||||
// Execute the target agent
|
||||
content, execErr := target.Execute(ctx, req.Task, r.shared)
|
||||
|
||||
// Record the result
|
||||
eventType := "result"
|
||||
eventContent := content
|
||||
if execErr != nil {
|
||||
eventType = "error"
|
||||
eventContent = execErr.Error()
|
||||
}
|
||||
r.shared.AddEvent(target.Name(), eventType, eventContent)
|
||||
|
||||
return &HandoffResult{
|
||||
AgentName: target.Name(),
|
||||
Content: content,
|
||||
Err: execErr,
|
||||
}
|
||||
}
|
||||
|
||||
// resolveTarget finds the appropriate agent for a hand-off request.
|
||||
func (r *AgentRegistry) resolveTarget(req HandoffRequest) (Agent, error) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
|
||||
// Direct routing by name
|
||||
if req.ToAgent != "" {
|
||||
entry, ok := r.agents[req.ToAgent]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("target agent %q not found", req.ToAgent)
|
||||
}
|
||||
return entry.agent, nil
|
||||
}
|
||||
|
||||
// Capability-based routing: find first idle agent with the capability
|
||||
if req.RequiredCapability != "" {
|
||||
for _, entry := range r.agents {
|
||||
if entry.state != AgentIdle {
|
||||
continue
|
||||
}
|
||||
for _, cap := range entry.agent.Capabilities() {
|
||||
if cap == req.RequiredCapability {
|
||||
return entry.agent, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("no idle agent found with capability %q", req.RequiredCapability)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("hand-off request must specify ToAgent or RequiredCapability")
|
||||
}
|
||||
|
||||
// setAgentState updates the state of an agent in the registry.
|
||||
func (r *AgentRegistry) setAgentState(name string, state AgentState) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if entry, ok := r.agents[name]; ok {
|
||||
entry.state = state
|
||||
}
|
||||
}
|
||||
|
||||
// GetAgentState returns the current state of an agent.
|
||||
func (r *AgentRegistry) GetAgentState(name string) (AgentState, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
entry, ok := r.agents[name]
|
||||
if !ok {
|
||||
return AgentIdle, false
|
||||
}
|
||||
return entry.state, true
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue