diff --git a/pkg/agent/multi/agent.go b/pkg/agent/multi/agent.go new file mode 100644 index 0000000000..a61602e72e --- /dev/null +++ b/pkg/agent/multi/agent.go @@ -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 +} diff --git a/pkg/agent/multi/context.go b/pkg/agent/multi/context.go new file mode 100644 index 0000000000..0ec090e5b8 --- /dev/null +++ b/pkg/agent/multi/context.go @@ -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 +} diff --git a/pkg/agent/multi/multi_test.go b/pkg/agent/multi/multi_test.go new file mode 100644 index 0000000000..f8626bbafd --- /dev/null +++ b/pkg/agent/multi/multi_test.go @@ -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)) + } +} diff --git a/pkg/agent/multi/registry.go b/pkg/agent/multi/registry.go new file mode 100644 index 0000000000..7e02f8e7ee --- /dev/null +++ b/pkg/agent/multi/registry.go @@ -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 +}