From bc3d77d2dd926ee6bf69fed0e584a36d902b4ba2 Mon Sep 17 00:00:00 2001 From: Zhaoyikaiii Date: Mon, 23 Feb 2026 21:52:15 +0800 Subject: [PATCH] feat: Implement swarm mode for multi-agent coordination - Add NodeInfo and NodeState structures for managing node information and status. - Introduce EventDispatcher for handling node events with subscription support. - Implement session transfer functionality with SessionTransfer and TransferOperation. - Create AuthProvider for HMAC-based authentication between swarm nodes. - Develop HandoffTool for delegating tasks to other agents in the swarm. - Add comprehensive tests for node management, event dispatching, and session transfer. - Define types for session messages and tool call data to facilitate communication. --- docs/swarm-architecture.md | 450 +++++++++++++++++++++++++++++ pkg/agent/loop.go | 323 ++++++++++++++++++++- pkg/config/config.go | 99 +++++++ pkg/swarm/config.go | 276 ++++++++++++++++++ pkg/swarm/constants.go | 101 +++++++ pkg/swarm/discovery.go | 516 ++++++++++++++++++++++++++++++++++ pkg/swarm/errors.go | 50 ++++ pkg/swarm/handoff.go | 486 ++++++++++++++++++++++++++++++++ pkg/swarm/leader_election.go | 237 ++++++++++++++++ pkg/swarm/load_monitor.go | 295 +++++++++++++++++++ pkg/swarm/membership.go | 365 ++++++++++++++++++++++++ pkg/swarm/metrics.go | 286 +++++++++++++++++++ pkg/swarm/node.go | 329 ++++++++++++++++++++++ pkg/swarm/security.go | 127 +++++++++ pkg/swarm/session_transfer.go | 331 ++++++++++++++++++++++ pkg/swarm/swarm_test.go | 424 ++++++++++++++++++++++++++++ pkg/swarm/types.go | 25 ++ pkg/tools/handoff_tool.go | 140 +++++++++ 18 files changed, 4859 insertions(+), 1 deletion(-) create mode 100644 docs/swarm-architecture.md create mode 100644 pkg/swarm/config.go create mode 100644 pkg/swarm/constants.go create mode 100644 pkg/swarm/discovery.go create mode 100644 pkg/swarm/errors.go create mode 100644 pkg/swarm/handoff.go create mode 100644 pkg/swarm/leader_election.go create mode 100644 pkg/swarm/load_monitor.go create mode 100644 pkg/swarm/membership.go create mode 100644 pkg/swarm/metrics.go create mode 100644 pkg/swarm/node.go create mode 100644 pkg/swarm/security.go create mode 100644 pkg/swarm/session_transfer.go create mode 100644 pkg/swarm/swarm_test.go create mode 100644 pkg/swarm/types.go create mode 100644 pkg/tools/handoff_tool.go diff --git a/docs/swarm-architecture.md b/docs/swarm-architecture.md new file mode 100644 index 000000000..5b19448ec --- /dev/null +++ b/docs/swarm-architecture.md @@ -0,0 +1,450 @@ +# PicoClaw Swarm Mode Architecture + +## Overview + +PicoClaw Swarm Mode enables multiple PicoClaw instances to work together as a distributed system, providing: +- **Node Discovery**: Automatic peer discovery via UDP gossip protocol +- **Health Monitoring**: Periodic heartbeat and failure detection +- **Load Balancing**: Intelligent task distribution based on node load +- **Handoff Mechanism**: Dynamic task delegation between nodes + +## Architecture + +The swarm architecture is divided into two distinct planes: + +``` +┌──────────────────────────────────────────────────────────────���──┐ +│ PicoClaw Swarm │ +├─────────────────────────────────────────────────────────────────┤ +│ Control Plane │ Data Plane │ +│ ├─ Node Discovery │ ├─ Task Execution │ +│ ├─ Membership Management │ ├─ Session Transfer │ +│ ├─ Health Monitoring │ └─ Message Routing │ +│ └─ Load Monitoring │ │ +└─────────────────────────────────────────────────────────────────┘ +``` + +## Control Plane + +The control plane manages cluster state, node membership, and coordination. + +### 1. Node Discovery + +Nodes discover each other using a lightweight UDP gossip protocol: + +```mermaid +sequenceDiagram + participant Node1 + participant Node2 + participant Node3 + + Note over Node1: New node starts + Node1->>Node1: Bind UDP port (7946) + Node1->>Node2: Ping + NodeInfo + Node2->>Node1: Pong + NodeInfo + Node1->>Node2: Gossip: Known Nodes + Node2->>Node3: Forward Node1 info + Node3->>Node2: Ack + Note over Node1,Node3: Cluster formed +``` + +**Gossip Protocol Flow:** + +```mermaid +graph LR + A[Node A] -->|Ping| B[Node B] + B -->|Pong| A + A -->|Sync| B + B -->|Forward| C[Node C] + C -->|Ack| B + B -->|Update| A + A -.->|Eventually| C +``` + +**Key Parameters:** + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `gossip_interval` | 1s | Frequency of gossip messages | +| `push_pull_interval` | 30s | Full state sync interval | +| `node_timeout` | 5s | Time before marking node suspect | +| `dead_node_timeout` | 30s | Time before removing dead node | + +### 2. Membership Management + +Each node maintains a view of the cluster: + +```go +type ClusterView struct { + sync.RWMutex + localNode *NodeInfo + members map[string]*NodeInfo // node_id -> NodeInfo + stateMap map[string]NodeState // node_id -> State +} +``` + +**Node State Machine:** + +```mermaid +stateDiagram-v2 + [*] --> Alive: Node joins + Alive --> Suspect: Missed heartbeat + Suspect --> Alive: Heartbeat recovered + Suspect --> Dead: Timeout exceeded + Dead --> [*] +``` + +**Node Information:** + +```go +type NodeInfo struct { + ID string // Unique node identifier + Addr string // IP address + Port int // Discovery port + AgentCaps map[string]string // Capabilities (models, tools) + LoadScore float64 // Current load (0.0-1.0) + Labels map[string]string // Custom labels + Timestamp int64 // Last update time + Version string // Protocol version +} +``` + +### 3. Health Monitoring + +**Heartbeat Flow:** + +```mermaid +sequenceDiagram + participant N1 as Node 1 + participant N2 as Node 2 + participant HM as Health Monitor + + loop Every gossip_interval + N1->>N2: Heartbeat (timestamp, load_score) + N2->>HM: Update state + HM->>HM: Check timeout + alt Timeout exceeded + HM->>HM: Mark as Suspect + HM->>N1: Probe (are you alive?) + alt No response + HM->>HM: Mark as Dead + HM->>All: Broadcast NodeLeft event + end + end + end +``` + +### 4. Load Monitoring + +Each node continuously monitors its resource usage: + +```mermaid +graph TB + subgraph Load Monitor + A[CPU Sample] --> D[Score Calculator] + B[Memory Sample] --> D + C[Session Count] --> D + D --> E[Load Score] + end + + subgraph Weights + A -.->|0.3| D + B -.->|0.3| D + C -.->|0.4| D + end + + E --> F{Threshold Check} + F -->|< 0.8| G[Normal Mode] + F -->|>= 0.8| H[Overloaded - Trigger Handoff] +``` + +**Load Score Formula:** + +``` +LoadScore = (CPUUsage × cpu_weight) + + (MemoryUsage × memory_weight) + + (SessionRatio × session_weight) + +Where: +- CPUUsage = current CPU usage (0.0-1.0) +- MemoryUsage = current memory usage (0.0-1.0) +- SessionRatio = current_sessions / max_sessions +- Default weights: cpu=0.3, memory=0.3, session=0.4 +``` + +## Data Plane + +The data plane handles actual task execution and session state transfer. + +### 1. Request Flow + +```mermaid +sequenceDiagram + participant User + participant LB as Entry Point + participant N1 as Node 1 + participant N2 as Node 2 + + User->>LB: Message + LB->>LB: Check node availability + + alt Node 1 available + LB->>N1: Forward message + N1->>N1: Process with LLM + N1->>User: Response + else Node 1 overloaded + LB->>N2: Handoff request + N2->>N1: Session transfer + N1->>N2: Session state + N2->>User: Response (from N2) + end +``` + +### 2. Handoff Mechanism + +**Handoff Decision Flow:** + +```mermaid +flowchart TD + A[Receive Request] --> B{Should Handoff?} + B -->|Local load >= threshold| C[Select Target Node] + B -->|Local load < threshold| D[Process Locally] + + C --> E{Target Available?} + E -->|Yes| F[Initiate Handoff] + E -->|No| G[Retry or Fail] + + F --> H[Serialize Session] + H --> I[Send to Target] + I --> J{Success?} + J -->|Yes| K[Update Routing] + J -->|No| L[Rollback] + + K --> M[Target Processes] + M --> N[Return Response] +``` + +**Handoff Protocol:** + +```mermaid +sequenceDiagram + participant Source as Overloaded Node + participant Target as Selected Node + participant Client + + Source->>Source: Check load threshold + Source->>Target: HandoffRequest{session_id, context} + + Target->>Target: Validate request + alt Accepted + Target->>Source: HandoffAccept + Source->>Target: SessionTransfer{messages, tools, state} + Target->>Target: Restore session + Target->>Source: TransferComplete + Source->>Client: Redirect to Target + Client->>Target: Continue conversation + else Rejected + Target->>Source: HandoffReject{reason} + Source->>Source: Try next node or process locally + end +``` + +### 3. Session Transfer + +**Session State Structure:** + +```go +type SessionState struct { + SessionID string + Messages []Message // Conversation history + Context map[string]any // Shared context + Tools []ToolCall // Pending tool calls + Metadata SessionMeta // Timestamp, user info, etc. +} +``` + +**Transfer Flow:** + +```mermaid +stateDiagram-v2 + [*] --> Active: Session created + Active --> Transferring: Handoff initiated + Transferring --> Active: Transfer failed + Transferring --> Migrated: Transfer complete + Migrated --> [*]: Session closed + Active --> [*]: Session closed +``` + +## System Architecture + +### Component Overview + +```mermaid +graph TB + subgraph "Node 1" + D1[Discovery Service] --> M1[Membership Manager] + H1[Handoff Coordinator] --> M1 + L1[Load Monitor] --> H1 + A1[Agent Loop] --> L1 + end + + subgraph "Node 2" + D2[Discovery Service] --> M2[Membership Manager] + H2[Handoff Coordinator] --> M2 + L2[Load Monitor] --> H2 + A2[Agent Loop] --> L2 + end + + D1 <-- UDP Gossip --> D2 + D2 <-- UDP Gossip --> D1 + H1 <-- RPC Handoff --> H2 + H2 <-- RPC Handoff --> H1 + + TG[Telegram Gateway] --> A1 + TG --> A2 +``` + +### Communication Channels + +| Channel | Protocol | Purpose | +|---------|----------|---------| +| Discovery | UDP | Node gossip, heartbeat | +| Handoff RPC | UDP | Session transfer coordination | +| Session Data | UDP | Serialized session state | + +## Configuration + +### Example Configuration + +```json +{ + "swarm": { + "enabled": true, + "node_id": "picoclaw-node-1", + "bind_addr": "127.0.0.1", + "bind_port": 7946, + + "discovery": { + "join_addrs": ["127.0.0.1:7946"], + "gossip_interval": 1, + "push_pull_interval": 30, + "node_timeout": 5, + "dead_node_timeout": 30 + }, + + "handoff": { + "enabled": true, + "load_threshold": 0.8, + "timeout": 30, + "max_retries": 3, + "retry_delay": 5 + }, + + "rpc": { + "port": 7947, + "timeout": 10 + }, + + "load_monitor": { + "enabled": true, + "interval": 5, + "sample_size": 60, + "cpu_weight": 0.3, + "memory_weight": 0.3, + "session_weight": 0.4 + } + } +} +``` + +### Deployment Modes + +**Single Entry Point:** + +```mermaid +graph LR + TG[Telegram Gateway] --> N1[Node 1: Coordinator] + N1 <-- Swarm --> N2[Node 2: Worker] + N2 <-- Swarm --> N1 + N1 <-- Swarm --> N3[Node 3: Worker] + N3 <-- Swarm --> N1 +``` + +**Multi-Entry Point (with load balancer):** + +```mermaid +graph LR + LB[Load Balancer] --> N1[Node 1] + LB --> N2[Node 2] + N1 <-- Swarm Mesh --> N2 + N2 <-- Swarm Mesh --> N1 + N1 <-- Swarm Mesh --> N3[Node 3] + N3 <-- Swarm Mesh --> N1 + N2 <-- Swarm Mesh --> N3 + N3 <-- Swarm Mesh --> N2 +``` + +## Event System + +The swarm publishes events for monitoring and integration: + +```mermaid +graph LR + A[Node Joined] --> ED[Event Dispatcher] + B[Node Left] --> ED + C[Node Suspect] --> ED + D[Handoff Started] --> ED + E[Handoff Completed] --> ED + + ED --> H[Handlers] + H --> L[Logging] + H --> M[Metrics] + H --> C[Custom Actions] +``` + +**Event Types:** + +| Event | Description | Payload | +|-------|-------------|---------| +| `NodeJoined` | New node discovered | NodeInfo | +| `NodeLeft` | Node removed | NodeID | +| `NodeSuspect` | Node marked suspect | NodeID | +| `NodeAlive` | Node recovered | NodeInfo | +| `HandoffStarted` | Handoff initiated | HandoffOperation | +| `HandoffCompleted` | Handoff finished | HandoffResult | +| `HandoffFailed` | Handoff error | Error | + +## Error Handling + +```mermaid +flowchart TD + A[Operation Failed] --> B{Retryable?} + B -->|Yes| C[Increment retry count] + B -->|No| D[Return error] + + C --> E{Max retries reached?} + E -->|No| F[Wait retry_delay] + F --> G[Retry operation] + + E -->|Yes| H[Mark node suspect] + H --> I[Select alternative node] + + G --> J{Success?} + J -->|Yes| K[Continue] + J -->|No| C +``` + +## Security Considerations + +1. **Discovery**: UDP gossip is unencrypted - use in trusted networks only +2. **Handoff**: Session data transferred without encryption +3. **Authentication**: No node authentication implemented +4. **Recommendation**: Use VPN or private network for production + +## Future Enhancements + +1. **Secure Discovery**: Add mTLS for node communication +2. **Consistent Hashing**: Replace random selection with consistent hashing +3. **Session Affinity**: Sticky sessions for better performance +4. **Leader Election**: Automatic coordinator election +5. **Multi-Region**: Geo-distributed cluster support diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index b36f4a0c4..88cb50a69 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -25,6 +25,7 @@ import ( "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/skills" "github.com/sipeed/picoclaw/pkg/state" + "github.com/sipeed/picoclaw/pkg/swarm" "github.com/sipeed/picoclaw/pkg/tools" "github.com/sipeed/picoclaw/pkg/utils" ) @@ -38,6 +39,12 @@ type AgentLoop struct { summarizing sync.Map fallback *providers.FallbackChain channelManager *channels.Manager + + // Swarm mode support + swarmEnabled bool + swarmDiscovery *swarm.DiscoveryService + swarmHandoff *swarm.HandoffCoordinator + swarmLoad *swarm.LoadMonitor } // processOptions configures how a message is processed @@ -69,7 +76,7 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers stateManager = state.NewManager(defaultAgent.Workspace) } - return &AgentLoop{ + al := &AgentLoop{ bus: msgBus, cfg: cfg, registry: registry, @@ -77,6 +84,13 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers summarizing: sync.Map{}, fallback: fallbackChain, } + + // Initialize swarm mode if enabled + if cfg.Swarm.Enabled { + al.initSwarm() + } + + return al } // registerSharedTools registers tools that are shared across all agents (web, message, spawn). @@ -319,6 +333,22 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) "matched_by": route.MatchedBy, }) + // Check if we should handoff this request to another node + if al.swarmEnabled && al.shouldHandoff(agent, processOptions{ + SessionKey: sessionKey, + Channel: msg.Channel, + ChatID: msg.ChatID, + UserMessage: msg.Content, + }) { + handoffResp, err := al.initiateSwarmHandoff(ctx, agent, sessionKey, msg) + if err == nil && handoffResp != nil && handoffResp.Accepted { + // Handoff was successful, return the response + return fmt.Sprintf("Your request has been handed off to node %s for processing.", handoffResp.NodeID), nil + } + // If handoff failed, continue processing locally + logger.WarnCF("swarm", "Handoff failed, processing locally", map[string]any{"error": err}) + } + return al.runAgentLoop(ctx, agent, processOptions{ SessionKey: sessionKey, Channel: msg.Channel, @@ -388,6 +418,10 @@ func (al *AgentLoop) processSystemMessage(ctx context.Context, msg bus.InboundMe // runAgentLoop is the core message processing logic. func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opts processOptions) (string, error) { + // Track active session for swarm load monitoring + al.IncrementSwarmSessions() + defer al.DecrementSwarmSessions() + // 0. Record last channel for heartbeat notifications (skip internal channels) if opts.Channel != "" && opts.ChatID != "" { // Don't record internal channels (cli, system, subagent) @@ -1144,3 +1178,290 @@ func extractParentPeer(msg bus.InboundMessage) *routing.RoutePeer { } return &routing.RoutePeer{Kind: parentKind, ID: parentID} } + +// Swarm methods + +// initSwarm initializes the swarm mode components. +func (al *AgentLoop) initSwarm() { + logger.InfoC("swarm", "Initializing swarm mode") + + // Convert config to swarm config + swarmConfig := al.convertToSwarmConfig(al.cfg.Swarm) + + // Create discovery service + discovery, err := swarm.NewDiscoveryService(swarmConfig) + if err != nil { + logger.ErrorCF("swarm", "Failed to create discovery service", map[string]any{"error": err.Error()}) + return + } + + // Start discovery + if err := discovery.Start(); err != nil { + logger.ErrorCF("swarm", "Failed to start discovery service", map[string]any{"error": err.Error()}) + return + } + + al.swarmDiscovery = discovery + + // Create handoff coordinator + handoffConfig := swarm.HandoffConfig{ + Enabled: al.cfg.Swarm.Handoff.Enabled, + LoadThreshold: al.cfg.Swarm.Handoff.LoadThreshold, + Timeout: swarm.Duration{Duration: time.Duration(al.cfg.Swarm.Handoff.Timeout) * time.Second}, + MaxRetries: al.cfg.Swarm.Handoff.MaxRetries, + RetryDelay: swarm.Duration{Duration: time.Duration(al.cfg.Swarm.Handoff.RetryDelay) * time.Second}, + } + + al.swarmHandoff, err = swarm.NewHandoffCoordinator(discovery, handoffConfig) + if err != nil { + logger.ErrorCF("swarm", "Failed to create handoff coordinator", map[string]any{"error": err.Error()}) + return + } + + // Create load monitor + loadConfig := swarm.LoadMonitorConfig{ + Enabled: al.cfg.Swarm.LoadMonitor.Enabled, + Interval: swarm.Duration{Duration: time.Duration(al.cfg.Swarm.LoadMonitor.Interval) * time.Second}, + SampleSize: al.cfg.Swarm.LoadMonitor.SampleSize, + CPUWeight: al.cfg.Swarm.LoadMonitor.CPUWeight, + MemoryWeight: al.cfg.Swarm.LoadMonitor.MemoryWeight, + SessionWeight: al.cfg.Swarm.LoadMonitor.SessionWeight, + } + + al.swarmLoad = swarm.NewLoadMonitor(&loadConfig) + if al.cfg.Swarm.LoadMonitor.Enabled { + al.swarmLoad.Start() + + // Register callback to update discovery with load + al.swarmLoad.OnThreshold(func(score float64) { + discovery.UpdateLoad(score) + }) + } + + al.swarmEnabled = true + + // Register handoff tool to all agents + if al.cfg.Swarm.Handoff.Enabled { + handoffTool := tools.NewHandoffTool(al.swarmHandoff) + al.RegisterTool(handoffTool) + } + + // Subscribe to node events for logging + discovery.Subscribe(func(event *swarm.NodeEvent) { + switch event.Event { + case swarm.EventJoin: + logger.InfoCF("swarm", "Node joined", map[string]any{"node_id": event.Node.ID}) + case swarm.EventLeave: + logger.InfoCF("swarm", "Node left", map[string]any{"node_id": event.Node.ID}) + case swarm.EventUpdate: + logger.DebugCF("swarm", "Node updated", map[string]any{ + "node_id": event.Node.ID, + "load_score": event.Node.LoadScore, + }) + } + }) + + logger.InfoCF("swarm", "Swarm mode initialized", map[string]any{ + "node_id": discovery.LocalNode().ID, + "bind_addr": al.cfg.Swarm.BindAddr, + "bind_port": al.cfg.Swarm.BindPort, + "handoff": al.cfg.Swarm.Handoff.Enabled, + }) +} + +// convertToSwarmConfig converts the config.SwarmConfig to swarm.Config. +func (al *AgentLoop) convertToSwarmConfig(cfg config.SwarmConfig) *swarm.Config { + return &swarm.Config{ + Enabled: cfg.Enabled, + NodeID: cfg.NodeID, + BindAddr: cfg.BindAddr, + BindPort: cfg.BindPort, + AdvertiseAddr: cfg.AdvertiseAddr, + Discovery: swarm.DiscoveryConfig{ + JoinAddrs: cfg.Discovery.JoinAddrs, + GossipInterval: swarm.Duration{Duration: time.Duration(cfg.Discovery.GossipInterval) * time.Second}, + PushPullInterval: swarm.Duration{Duration: time.Duration(cfg.Discovery.PushPullInterval) * time.Second}, + NodeTimeout: swarm.Duration{Duration: time.Duration(cfg.Discovery.NodeTimeout) * time.Second}, + DeadNodeTimeout: swarm.Duration{Duration: time.Duration(cfg.Discovery.DeadNodeTimeout) * time.Second}, + }, + Handoff: swarm.HandoffConfig{ + Enabled: cfg.Handoff.Enabled, + LoadThreshold: cfg.Handoff.LoadThreshold, + Timeout: swarm.Duration{Duration: time.Duration(cfg.Handoff.Timeout) * time.Second}, + MaxRetries: cfg.Handoff.MaxRetries, + RetryDelay: swarm.Duration{Duration: time.Duration(cfg.Handoff.RetryDelay) * time.Second}, + }, + RPC: swarm.RPCConfig{ + Port: cfg.RPC.Port, + Timeout: swarm.Duration{Duration: time.Duration(cfg.RPC.Timeout) * time.Second}, + }, + LoadMonitor: swarm.LoadMonitorConfig{ + Enabled: cfg.LoadMonitor.Enabled, + Interval: swarm.Duration{Duration: time.Duration(cfg.LoadMonitor.Interval) * time.Second}, + SampleSize: cfg.LoadMonitor.SampleSize, + CPUWeight: cfg.LoadMonitor.CPUWeight, + MemoryWeight: cfg.LoadMonitor.MemoryWeight, + SessionWeight: cfg.LoadMonitor.SessionWeight, + }, + } +} + +// shouldHandoff determines if the current request should be handed off to another node. +func (al *AgentLoop) shouldHandoff(agent *AgentInstance, opts processOptions) bool { + if !al.swarmEnabled || al.swarmHandoff == nil { + return false + } + + // Check if load is too high + if al.swarmLoad != nil && al.swarmLoad.ShouldOffload() { + logger.InfoCF("swarm", "Load threshold exceeded, considering handoff", map[string]any{ + "load_score": al.swarmLoad.GetCurrentLoad().Score, + }) + return true + } + + return false +} + +// UpdateSwarmLoad updates the current load score reported to the swarm. +func (al *AgentLoop) UpdateSwarmLoad(sessionCount int) { + if al.swarmLoad != nil { + al.swarmLoad.SetSessionCount(sessionCount) + } +} + +// IncrementSwarmSessions increments the active session count. +func (al *AgentLoop) IncrementSwarmSessions() { + if al.swarmLoad != nil { + al.swarmLoad.IncrementSessions() + } +} + +// DecrementSwarmSessions decrements the active session count. +func (al *AgentLoop) DecrementSwarmSessions() { + if al.swarmLoad != nil { + al.swarmLoad.DecrementSessions() + } +} + +// GetSwarmStatus returns the current swarm status. +func (al *AgentLoop) GetSwarmStatus() map[string]any { + if !al.swarmEnabled { + return map[string]any{"enabled": false} + } + + status := map[string]any{ + "enabled": true, + "node_id": al.swarmDiscovery.LocalNode().ID, + "handoff": al.cfg.Swarm.Handoff.Enabled, + } + + if al.swarmLoad != nil { + metrics := al.swarmLoad.GetCurrentLoad() + status["load"] = map[string]any{ + "score": metrics.Score, + "cpu_usage": metrics.CPUUsage, + "memory_usage": metrics.MemoryUsage, + "active_sessions": metrics.ActiveSessions, + "goroutines": metrics.Goroutines, + "trend": al.swarmLoad.GetTrend(), + } + } + + if al.swarmDiscovery != nil { + members := al.swarmDiscovery.Members() + status["members"] = len(members) + } + + return status +} + +// ShutdownSwarm gracefully shuts down the swarm components. +func (al *AgentLoop) ShutdownSwarm() error { + if !al.swarmEnabled { + return nil + } + + var errs []string + + if al.swarmLoad != nil { + al.swarmLoad.Stop() + } + + if al.swarmHandoff != nil { + if err := al.swarmHandoff.Close(); err != nil { + errs = append(errs, fmt.Sprintf("handoff: %v", err)) + } + } + + if al.swarmDiscovery != nil { + if err := al.swarmDiscovery.Stop(); err != nil { + errs = append(errs, fmt.Sprintf("discovery: %v", err)) + } + } + + al.swarmEnabled = false + + if len(errs) > 0 { + return fmt.Errorf("swarm shutdown errors: %s", strings.Join(errs, ", ")) + } + return nil +} + +// initiateSwarmHandoff initiates a handoff to another node. +func (al *AgentLoop) initiateSwarmHandoff(ctx context.Context, agent *AgentInstance, sessionKey string, msg bus.InboundMessage) (*swarm.HandoffResponse, error) { + if al.swarmHandoff == nil { + return nil, swarm.ErrDiscoveryDisabled + } + + // Build session history for handoff + sessionMessages := make([]swarm.SessionMessage, 0) + history := agent.Sessions.GetHistory(sessionKey) + + for _, m := range history { + if m.Role == "user" || m.Role == "assistant" { + sessionMessages = append(sessionMessages, swarm.SessionMessage{ + Role: m.Role, + Content: m.Content, + }) + } + } + + // Create handoff request + req := &swarm.HandoffRequest{ + Reason: swarm.ReasonOverloaded, + SessionKey: sessionKey, + SessionMessages: sessionMessages, + Context: map[string]any{ + "channel": msg.Channel, + "chat_id": msg.ChatID, + "sender": msg.SenderID, + "agent_id": agent.ID, + }, + Metadata: map[string]string{ + "original_channel": msg.Channel, + "original_chat_id": msg.ChatID, + }, + } + + logger.InfoCF("swarm", "Initiating handoff", map[string]any{ + "session_key": sessionKey, + "reason": req.Reason, + "history_len": len(sessionMessages), + }) + + // Execute handoff + resp, err := al.swarmHandoff.InitiateHandoff(ctx, req) + + if resp != nil { + logger.InfoCF("swarm", "Handoff response received", map[string]any{ + "accepted": resp.Accepted, + "node_id": resp.NodeID, + "state": resp.State, + }) + } + + return resp, err +} + + diff --git a/pkg/config/config.go b/pkg/config/config.go index 20556011a..3647346c8 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -57,6 +57,7 @@ type Config struct { Tools ToolsConfig `json:"tools"` Heartbeat HeartbeatConfig `json:"heartbeat"` Devices DevicesConfig `json:"devices"` + Swarm SwarmConfig `json:"swarm,omitempty"` // Swarm mode configuration } // MarshalJSON implements custom JSON marshaling for Config @@ -650,3 +651,101 @@ func (c *Config) ValidateModelList() error { } return nil } + +// SwarmConfig contains configuration for swarm mode. +type SwarmConfig struct { + // Enabled enables swarm mode. + Enabled bool `json:"enabled" env:"PICOCLAW_SWARM_ENABLED"` + + // NodeID is the unique identifier for this node. + NodeID string `json:"node_id,omitempty" env:"PICOCLAW_SWARM_NODE_ID"` + + // BindAddr is the address to bind for gossip and RPC. + BindAddr string `json:"bind_addr,omitempty" env:"PICOCLAW_SWARM_BIND_ADDR"` + + // BindPort is the port for gossip protocol. + BindPort int `json:"bind_port,omitempty" env:"PICOCLAW_SWARM_BIND_PORT"` + + // AdvertiseAddr is the address to advertise to other nodes. + AdvertiseAddr string `json:"advertise_addr,omitempty" env:"PICOCLAW_SWARM_ADVERTISE_ADDR"` + + // Discovery configuration for node discovery. + Discovery SwarmDiscoveryConfig `json:"discovery"` + + // Handoff configuration for task handoff. + Handoff SwarmHandoffConfig `json:"handoff"` + + // RPC configuration for inter-node communication. + RPC SwarmRPCConfig `json:"rpc"` + + // LoadMonitor configuration for load monitoring. + LoadMonitor SwarmLoadMonitorConfig `json:"load_monitor"` +} + +// SwarmDiscoveryConfig contains configuration for node discovery. +type SwarmDiscoveryConfig struct { + // JoinAddrs is a list of existing nodes to join. + JoinAddrs []string `json:"join_addrs,omitempty"` + + // GossipInterval is the interval between gossip messages (in seconds). + GossipInterval int `json:"gossip_interval,omitempty"` + + // PushPullInterval is the interval for full state sync (in seconds). + PushPullInterval int `json:"push_pull_interval,omitempty"` + + // NodeTimeout is the timeout before marking a node as suspect (in seconds). + NodeTimeout int `json:"node_timeout,omitempty"` + + // DeadNodeTimeout is the timeout before marking a node as dead (in seconds). + DeadNodeTimeout int `json:"dead_node_timeout,omitempty"` +} + +// SwarmHandoffConfig contains configuration for task handoff. +type SwarmHandoffConfig struct { + // Enabled enables task handoff. + Enabled bool `json:"enabled"` + + // LoadThreshold is the load score threshold (0-1) above which + // tasks will be handed off to other nodes. + LoadThreshold float64 `json:"load_threshold,omitempty"` + + // Timeout is the timeout for a handoff operation (in seconds). + Timeout int `json:"timeout,omitempty"` + + // MaxRetries is the maximum number of retries for handoff. + MaxRetries int `json:"max_retries,omitempty"` + + // RetryDelay is the delay between retries (in seconds). + RetryDelay int `json:"retry_delay,omitempty"` +} + +// SwarmRPCConfig contains configuration for RPC communication. +type SwarmRPCConfig struct { + // Port is the port for RPC communication. + Port int `json:"port,omitempty" env:"PICOCLAW_SWARM_RPC_PORT"` + + // Timeout is the default timeout for RPC calls (in seconds). + Timeout int `json:"timeout,omitempty"` +} + +// SwarmLoadMonitorConfig contains configuration for load monitoring. +type SwarmLoadMonitorConfig struct { + // Enabled enables load monitoring. + Enabled bool `json:"enabled"` + + // Interval is the interval between load samples (in seconds). + Interval int `json:"interval,omitempty"` + + // SampleSize is the number of samples to keep for averaging. + SampleSize int `json:"sample_size,omitempty"` + + // CPUWeight is the weight for CPU usage in load score (0-1). + CPUWeight float64 `json:"cpu_weight,omitempty"` + + // MemoryWeight is the weight for memory usage in load score (0-1). + MemoryWeight float64 `json:"memory_weight,omitempty"` + + // SessionWeight is the weight for active sessions in load score (0-1). + SessionWeight float64 `json:"session_weight,omitempty"` +} + diff --git a/pkg/swarm/config.go b/pkg/swarm/config.go new file mode 100644 index 000000000..6ec6c6c8e --- /dev/null +++ b/pkg/swarm/config.go @@ -0,0 +1,276 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "encoding/json" + "time" +) + +// Config contains all configuration for swarm mode. +type Config struct { + // Enabled enables swarm mode. + Enabled bool `json:"enabled" env:"PICOCLAW_SWARM_ENABLED"` + + // NodeID is the unique identifier for this node. + // If empty, a hostname-based ID will be generated. + NodeID string `json:"node_id,omitempty" env:"PICOCLAW_SWARM_NODE_ID"` + + // BindAddr is the address to bind for gossip and RPC. + BindAddr string `json:"bind_addr,omitempty" env:"PICOCLAW_SWARM_BIND_ADDR"` + + // BindPort is the port for gossip protocol. + BindPort int `json:"bind_port,omitempty" env:"PICOCLAW_SWARM_BIND_PORT"` + + // AdvertiseAddr is the address to advertise to other nodes. + // If empty, BindAddr will be used. + AdvertiseAddr string `json:"advertise_addr,omitempty" env:"PICOCLAW_SWARM_ADVERTISE_ADDR"` + + // AdvertisePort is the port to advertise to other nodes. + // If 0, BindPort will be used. + AdvertisePort int `json:"advertise_port,omitempty" env:"PICOCLAW_SWARM_ADVERTISE_PORT"` + + // Discovery configuration for node discovery. + Discovery DiscoveryConfig `json:"discovery"` + + // Handoff configuration for task handoff. + Handoff HandoffConfig `json:"handoff"` + + // RPC configuration for inter-node communication. + RPC RPCConfig `json:"rpc"` + + // LoadMonitor configuration for load monitoring. + LoadMonitor LoadMonitorConfig `json:"load_monitor"` + + // LeaderElection configuration for leader election. + LeaderElection LeaderElectionConfig `json:"leader_election"` + + // Metrics configuration for observability. + Metrics MetricsConfig `json:"metrics"` +} + +// DiscoveryConfig contains configuration for node discovery. +type DiscoveryConfig struct { + // JoinAddrs is a list of existing nodes to join. + JoinAddrs []string `json:"join_addrs,omitempty"` + + // GossipInterval is the interval between gossip messages. + GossipInterval Duration `json:"gossip_interval,omitempty"` + + // PushPullInterval is the interval for full state sync. + PushPullInterval Duration `json:"push_pull_interval,omitempty"` + + // NodeTimeout is the timeout before marking a node as suspect. + NodeTimeout Duration `json:"node_timeout,omitempty"` + + // DeadNodeTimeout is the timeout before marking a node as dead. + DeadNodeTimeout Duration `json:"dead_node_timeout,omitempty"` + + // AuthSecret is the shared secret for node authentication. + // If empty, authentication is disabled (not recommended for production). + AuthSecret string `json:"auth_secret,omitempty"` + + // RequireAuth requires all nodes to be authenticated. + RequireAuth bool `json:"require_auth"` + + // EnableMessageSigning enables HMAC signing of all messages. + EnableMessageSigning bool `json:"enable_message_signing"` +} + +// HandoffConfig contains configuration for task handoff. +type HandoffConfig struct { + // Enabled enables task handoff. + Enabled bool `json:"enabled"` + + // LoadThreshold is the load score threshold (0-1) above which + // tasks will be handed off to other nodes. + LoadThreshold float64 `json:"load_threshold,omitempty"` + + // Timeout is the timeout for a handoff operation. + Timeout Duration `json:"timeout,omitempty"` + + // MaxRetries is the maximum number of retries for handoff. + MaxRetries int `json:"max_retries,omitempty"` + + // RetryDelay is the delay between retries. + RetryDelay Duration `json:"retry_delay,omitempty"` +} + +// RPCConfig contains configuration for RPC communication. +type RPCConfig struct { + // Port is the port for RPC communication. + Port int `json:"port,omitempty" env:"PICOCLAW_SWARM_RPC_PORT"` + + // Timeout is the default timeout for RPC calls. + Timeout Duration `json:"timeout,omitempty"` +} + +// LoadMonitorConfig contains configuration for load monitoring. +type LoadMonitorConfig struct { + // Enabled enables load monitoring. + Enabled bool `json:"enabled"` + + // Interval is the interval between load samples. + Interval Duration `json:"interval,omitempty"` + + // SampleSize is the number of samples to keep for averaging. + SampleSize int `json:"sample_size,omitempty"` + + // CPUWeight is the weight for CPU usage in load score (0-1). + CPUWeight float64 `json:"cpu_weight,omitempty"` + + // MemoryWeight is the weight for memory usage in load score (0-1). + MemoryWeight float64 `json:"memory_weight,omitempty"` + + // SessionWeight is the weight for active sessions in load score (0-1). + SessionWeight float64 `json:"session_weight,omitempty"` + + // OffloadThreshold is the load score threshold above which tasks should be offloaded (0-1). + OffloadThreshold float64 `json:"offload_threshold,omitempty"` + + // MaxMemoryBytes is the maximum memory to use for normalization (default: 1GB). + MaxMemoryBytes uint64 `json:"max_memory_bytes,omitempty"` + + // MaxGoroutines is the maximum goroutine count for normalization (default: 1000). + MaxGoroutines int `json:"max_goroutines,omitempty"` + + // MaxSessions is the maximum session count for normalization (default: 100). + MaxSessions int `json:"max_sessions,omitempty"` +} + +// LeaderElectionConfig contains configuration for leader election. +type LeaderElectionConfig struct { + // Enabled enables leader election. + Enabled bool `json:"enabled"` + + // ElectionInterval is how often to check leadership. + ElectionInterval Duration `json:"election_interval,omitempty"` + + // LeaderHeartbeatTimeout is how long before assuming leader is dead. + LeaderHeartbeatTimeout Duration `json:"leader_heartbeat_timeout,omitempty"` +} + +// MetricsConfig contains configuration for metrics collection. +type MetricsConfig struct { + // Enabled enables metrics collection. + Enabled bool `json:"enabled"` + + // ExportInterval is how often to export metrics. + ExportInterval Duration `json:"export_interval,omitempty"` + + // PrometheusEnabled enables Prometheus format export. + PrometheusEnabled bool `json:"prometheus_enabled"` + + // PrometheusEndpoint is the HTTP endpoint for Prometheus metrics. + PrometheusEndpoint string `json:"prometheus_endpoint,omitempty"` +} + +// Duration is a wrapper around time.Duration for JSON parsing. +type Duration struct { + time.Duration +} + +// UnmarshalJSON parses a duration from JSON. +func (d *Duration) UnmarshalJSON(b []byte) error { + // Check if it's a string (quoted) + if len(b) > 0 && b[0] == '"' { + var s string + if err := parseJSONString(b, &s); err != nil { + return err + } + var err error + d.Duration, err = time.ParseDuration(s) + return err + } + + // Otherwise it's a number (milliseconds) + var v float64 + if err := parseJSONNumber(b, &v); err != nil { + return err + } + d.Duration = time.Duration(v) + return nil +} + +// MarshalJSON converts a duration to JSON. +func (d Duration) MarshalJSON() ([]byte, error) { + return json.Marshal(d.Duration.String()) +} + +// DefaultConfig returns the default swarm configuration. +func DefaultConfig() *Config { + return &Config{ + Enabled: false, + NodeID: "", + BindAddr: "0.0.0.0", + BindPort: DefaultBindPort, + Discovery: DiscoveryConfig{ + JoinAddrs: nil, + GossipInterval: Duration{DefaultGossipInterval}, + PushPullInterval: Duration{DefaultPushPullInterval}, + NodeTimeout: Duration{DefaultNodeTimeout}, + DeadNodeTimeout: Duration{DefaultDeadNodeTimeout}, + AuthSecret: "", + RequireAuth: false, + EnableMessageSigning: false, + }, + Handoff: HandoffConfig{ + Enabled: true, + LoadThreshold: DefaultLoadThreshold, + Timeout: Duration{DefaultHandoffTimeout}, + MaxRetries: DefaultMaxHandoffRetries, + RetryDelay: Duration{DefaultHandoffRetryDelay}, + }, + RPC: RPCConfig{ + Port: DefaultRPCPort, + Timeout: Duration{10 * time.Second}, + }, + LoadMonitor: LoadMonitorConfig{ + Enabled: true, + Interval: Duration{DefaultLoadSampleInterval}, + SampleSize: DefaultLoadSampleSize, + CPUWeight: DefaultCPUWeight, + MemoryWeight: DefaultMemoryWeight, + SessionWeight: DefaultSessionWeight, + OffloadThreshold: DefaultOffloadThreshold, + MaxMemoryBytes: DefaultMaxMemoryBytes, + MaxGoroutines: DefaultMaxGoroutines, + MaxSessions: DefaultMaxSessions, + }, + LeaderElection: LeaderElectionConfig{ + Enabled: false, + ElectionInterval: Duration{5 * time.Second}, + LeaderHeartbeatTimeout: Duration{10 * time.Second}, + }, + Metrics: MetricsConfig{ + Enabled: false, + ExportInterval: Duration{10 * time.Second}, + PrometheusEnabled: false, + PrometheusEndpoint: "/metrics", + }, + } +} + +// parseJSONString parses a JSON string (including quotes). +func parseJSONString(b []byte, s *string) error { + if len(b) < 2 || b[0] != '"' || b[len(b)-1] != '"' { + return &json.UnmarshalTypeError{} + } + *s = string(b[1 : len(b)-1]) + return nil +} + +// parseJSONNumber parses a JSON number. +func parseJSONNumber(b []byte, f *float64) error { + n, err := json.Number(string(b)).Int64() + if err == nil { + *f = float64(n) + return nil + } + *f, err = json.Number(string(b)).Float64() + return err +} diff --git a/pkg/swarm/constants.go b/pkg/swarm/constants.go new file mode 100644 index 000000000..f7a670a43 --- /dev/null +++ b/pkg/swarm/constants.go @@ -0,0 +1,101 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import "time" + +const ( + // Default values for configurable parameters + + // DefaultBindPort is the default port for gossip protocol. + DefaultBindPort = 7946 + + // DefaultRPCPort is the default port for RPC communication. + DefaultRPCPort = 7947 + + // DefaultNodeTimeout is the default timeout before marking a node as suspect. + DefaultNodeTimeout = 5 * time.Second + + // DefaultDeadNodeTimeout is the default timeout before removing a dead node. + DefaultDeadNodeTimeout = 30 * time.Second + + // DefaultGossipInterval is the default interval between gossip messages. + DefaultGossipInterval = 1 * time.Second + + // DefaultPushPullInterval is the default interval for full state sync. + DefaultPushPullInterval = 30 * time.Second + + // DefaultHandoffTimeout is the default timeout for a handoff operation. + DefaultHandoffTimeout = 30 * time.Second + + // DefaultHandoffRetryDelay is the default delay between handoff retries. + DefaultHandoffRetryDelay = 5 * time.Second + + // DefaultMaxHandoffRetries is the default maximum number of handoff retries. + DefaultMaxHandoffRetries = 3 + + // DefaultLoadSampleInterval is the default interval between load samples. + DefaultLoadSampleInterval = 5 * time.Second + + // DefaultLoadSampleSize is the default number of load samples to keep. + DefaultLoadSampleSize = 60 + + // Thresholds and limits + + // DefaultLoadThreshold is the default load score threshold for handoff. + DefaultLoadThreshold = 0.8 + + // DefaultAvailableLoadThreshold is the threshold below which a node is considered available (0-1). + DefaultAvailableLoadThreshold = 0.9 + + // DefaultOffloadThreshold is the default threshold above which tasks should be offloaded (0-1). + DefaultOffloadThreshold = 0.8 + + // DefaultMaxMemoryBytes is the default max memory for normalization (1GB). + DefaultMaxMemoryBytes = 1024 * 1024 * 1024 + + // DefaultMaxGoroutines is the default max goroutine count for normalization. + DefaultMaxGoroutines = 1000 + + // DefaultMaxSessions is the default max session count for normalization. + DefaultMaxSessions = 100 + + // DefaultCPUWeight is the default weight for CPU in load score calculation. + DefaultCPUWeight = 0.3 + + // DefaultMemoryWeight is the default weight for memory in load score calculation. + DefaultMemoryWeight = 0.3 + + // DefaultSessionWeight is the default weight for sessions in load score calculation. + DefaultSessionWeight = 0.4 + + // Buffer sizes + + // MaxGossipMessageSize is the maximum size of a gossip message (64KB). + MaxGossipMessageSize = 64 * 1024 + + // MaxSessionMessageSize is the maximum size of a session transfer message (128KB). + MaxSessionMessageSize = 128 * 1024 + + // UDP write deadline + + // DefaultUDPWriteDeadline is the default write deadline for UDP operations. + DefaultUDPWriteDeadline = 5 * time.Second + + // Poll intervals + + // HandoffResponsePollInterval is the interval for polling handoff responses. + HandoffResponsePollInterval = 100 * time.Millisecond + + // Trend analysis + + // TrendIncreasingThreshold is the slope threshold for detecting increasing trend. + TrendIncreasingThreshold = 0.01 + + // TrendDecreasingThreshold is the slope threshold for detecting decreasing trend. + TrendDecreasingThreshold = -0.01 +) diff --git a/pkg/swarm/discovery.go b/pkg/swarm/discovery.go new file mode 100644 index 000000000..841fd3301 --- /dev/null +++ b/pkg/swarm/discovery.go @@ -0,0 +1,516 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "context" + "encoding/json" + "fmt" + "net" + "os" + "sync" + "time" + + "github.com/google/uuid" + "github.com/sipeed/picoclaw/pkg/logger" +) + +// DiscoveryService handles node discovery using a gossip protocol. +// For lightweight implementation, we use a simple UDP-based gossip +// instead of the heavier memberlist library. +type DiscoveryService struct { + config *Config + localNode *NodeInfo + membership *MembershipManager + eventHandler *EventDispatcher + conn *net.UDPConn + rpcConn net.Listener + auth *AuthProvider + + mu sync.RWMutex + running bool + stopChan chan struct{} + once sync.Once + + // Sequence number for updates + seqNum uint64 +} + +// NewDiscoveryService creates a new discovery service. +func NewDiscoveryService(cfg *Config) (*DiscoveryService, error) { + if cfg.NodeID == "" { + // Generate node ID from hostname + hostname, _ := os.Hostname() + if hostname == "" { + hostname = "picoclaw" + } + cfg.NodeID = fmt.Sprintf("%s-%s", hostname, uuid.New().String()[:8]) + } + + // Determine advertise address + advAddr := cfg.AdvertiseAddr + if advAddr == "" || advAddr == "0.0.0.0" { + advAddr = getLocalIP() + if advAddr == "" { + advAddr = "127.0.0.1" + } + } + + advPort := cfg.AdvertisePort + if advPort == 0 { + advPort = cfg.BindPort + } + + localNode := &NodeInfo{ + ID: cfg.NodeID, + Addr: advAddr, + Port: cfg.RPC.Port, + AgentCaps: make(map[string]string), + LoadScore: 0, + Labels: make(map[string]string), + Timestamp: time.Now().UnixNano(), + Version: "1.0.0", // PicoClaw version + } + + ds := &DiscoveryService{ + config: cfg, + localNode: localNode, + eventHandler: NewEventDispatcher(), + stopChan: make(chan struct{}), + } + + // Initialize auth provider if secret is configured + if cfg.Discovery.AuthSecret != "" { + ds.auth = NewAuthProvider(cfg.NodeID, cfg.Discovery.AuthSecret) + if cfg.Discovery.RequireAuth || cfg.Discovery.EnableMessageSigning { + logger.InfoC("swarm", "Authentication enabled for swarm") + } + } + + // Initialize membership manager + ds.membership = NewMembershipManager(ds, cfg.Discovery) + + return ds, nil +} + +// Start starts the discovery service. +func (ds *DiscoveryService) Start() error { + ds.mu.Lock() + if ds.running { + ds.mu.Unlock() + return nil + } + ds.running = true + ds.mu.Unlock() + + // Bind UDP socket for gossip + addr := fmt.Sprintf("%s:%d", ds.config.BindAddr, ds.config.BindPort) + udpAddr, err := net.ResolveUDPAddr("udp", addr) + if err != nil { + return fmt.Errorf("failed to resolve UDP address: %w", err) + } + + ds.conn, err = net.ListenUDP("udp", udpAddr) + if err != nil { + return fmt.Errorf("failed to listen on UDP: %w", err) + } + + // Start gossip listener + go ds.gossipListener() + + // Start periodic gossip + go ds.gossipLoop() + + // Join existing cluster if addresses provided + if len(ds.config.Discovery.JoinAddrs) > 0 { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + ds.Join(ctx, ds.config.Discovery.JoinAddrs) + } + + // Add self to membership + ds.membership.UpdateNode(ds.localNode) + + return nil +} + +// Stop stops the discovery service. +func (ds *DiscoveryService) Stop() error { + ds.once.Do(func() { + ds.mu.Lock() + ds.running = false + ds.mu.Unlock() + + if ds.stopChan != nil { + close(ds.stopChan) + } + + if ds.conn != nil { + ds.conn.Close() + } + + if ds.rpcConn != nil { + ds.rpcConn.Close() + } + }) + return nil +} + +// Join joins a cluster by contacting existing nodes. +func (ds *DiscoveryService) Join(ctx context.Context, addrs []string) (int, error) { + count := 0 + + for _, addr := range addrs { + // Send join message to each address + err := ds.sendJoin(ctx, addr) + if err == nil { + count++ + } + } + + return count, nil +} + +// Members returns all known members of the cluster. +func (ds *DiscoveryService) Members() []*NodeWithState { + return ds.membership.GetMembers() +} + +// LocalNode returns the local node info. +func (ds *DiscoveryService) LocalNode() *NodeInfo { + ds.mu.RLock() + defer ds.mu.RUnlock() + return ds.localNode +} + +// UpdateLocalInfo updates the local node's information. +func (ds *DiscoveryService) UpdateLocalInfo(info *NodeInfo) { + ds.mu.Lock() + ds.localNode = info + ds.localNode.Timestamp = time.Now().UnixNano() + ds.seqNum++ + ds.mu.Unlock() + + // Update membership + ds.membership.UpdateNode(info) + + // Broadcast update + ds.broadcastUpdate() +} + +// UpdateLoad updates the local node's load score. +func (ds *DiscoveryService) UpdateLoad(score float64) { + ds.mu.Lock() + ds.localNode.LoadScore = score + ds.localNode.Timestamp = time.Now().UnixNano() + ds.seqNum++ + info := ds.localNode + ds.mu.Unlock() + + ds.membership.UpdateNode(info) + ds.broadcastUpdate() +} + +// UpdateCapabilities updates the local node's agent capabilities. +func (ds *DiscoveryService) UpdateCapabilities(caps map[string]string) { + ds.mu.Lock() + ds.localNode.AgentCaps = caps + ds.localNode.Timestamp = time.Now().UnixNano() + ds.seqNum++ + info := ds.localNode + ds.mu.Unlock() + + ds.membership.UpdateNode(info) + ds.broadcastUpdate() +} + +// Subscribe registers a handler for node events and returns its ID. +func (ds *DiscoveryService) Subscribe(handler EventHandler) EventHandlerID { + return ds.eventHandler.Subscribe(handler) +} + +// Unsubscribe removes a node event handler by ID. +func (ds *DiscoveryService) Unsubscribe(id EventHandlerID) { + ds.eventHandler.Unsubscribe(id) +} + +// gossipListener listens for incoming gossip messages. +func (ds *DiscoveryService) gossipListener() { + buf := make([]byte, MaxGossipMessageSize) + + for { + select { + case <-ds.stopChan: + return + default: + } + + ds.conn.SetReadDeadline(time.Now().Add(1 * time.Second)) + n, addr, err := ds.conn.ReadFromUDP(buf) + if err != nil { + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + continue + } + return + } + + if n > 0 { + go ds.handleGossip(buf[:n], addr) + } + } +} + +// gossipLoop periodically gossips node state to random members. +func (ds *DiscoveryService) gossipLoop() { + ticker := time.NewTicker(ds.config.Discovery.GossipInterval.Duration) + defer ticker.Stop() + + for { + select { + case <-ds.stopChan: + return + case <-ticker.C: + ds.broadcastUpdate() + } + } +} + +// GossipMessage represents a gossip message. +type GossipMessage struct { + Type string `json:"type"` // "ping", "pong", "join", "update" + FromNode string `json:"from_node"` + SeqNum uint64 `json:"seq_num"` + Timestamp int64 `json:"timestamp"` + Payload []byte `json:"payload,omitempty"` + Nodes []*NodeInfo `json:"nodes,omitempty"` // For memberlist exchange + AuthToken *AuthToken `json:"auth_token,omitempty"` +} + +// handleGossip handles an incoming gossip message. +func (ds *DiscoveryService) handleGossip(data []byte, addr *net.UDPAddr) { + var msg GossipMessage + if err := json.Unmarshal(data, &msg); err != nil { + return + } + + // Verify authentication if required + if ds.config.Discovery.RequireAuth && msg.AuthToken != nil { + if ds.auth == nil || !ds.auth.VerifyToken(msg.AuthToken) { + logger.WarnCF("swarm", "Rejected unauthenticated message", map[string]any{"from": addr.String()}) + return + } + } + + // Verify message signature if enabled + if ds.config.Discovery.EnableMessageSigning && msg.AuthToken != nil { + // The signature is in the token, so verification above handles it + } + + switch msg.Type { + case "ping": + ds.handlePing(msg, addr) + case "pong": + ds.handlePong(msg) + case "join": + ds.handleJoin(msg, addr) + case "update": + ds.handleUpdate(msg) + case "sync": + ds.handleSync(msg, addr) + } +} + +// handlePing handles a ping message. +func (ds *DiscoveryService) handlePing(msg GossipMessage, addr *net.UDPAddr) { + // Respond with pong + pong := GossipMessage{ + Type: "pong", + FromNode: ds.localNode.ID, + Timestamp: time.Now().UnixNano(), + } + + data, err := json.Marshal(pong) + if err != nil { + logger.ErrorCF("swarm", "failed to marshal pong message", map[string]any{"error": err}) + return + } + if _, err := ds.conn.WriteToUDP(data, addr); err != nil { + logger.DebugCF("swarm", "failed to send pong", map[string]any{"to": addr.String(), "error": err}) + } + + // Update membership if this is a known node + if len(msg.Nodes) > 0 { + for _, node := range msg.Nodes { + if node.ID != ds.localNode.ID { + ds.membership.UpdateNode(node) + } + } + } +} + +// handlePong handles a pong message. +func (ds *DiscoveryService) handlePong(msg GossipMessage) { + // Update last seen for this node + ds.membership.RecordHeartbeat(msg.FromNode) +} + +// handleJoin handles a join request from a new node. +func (ds *DiscoveryService) handleJoin(msg GossipMessage, addr *net.UDPAddr) { + // Send our member list back + members := ds.Members() + nodes := make([]*NodeInfo, 0, len(members)+1) + nodes = append(nodes, ds.localNode) + for _, m := range members { + if m.Node.ID != ds.localNode.ID { + nodes = append(nodes, m.Node) + } + } + + response := GossipMessage{ + Type: "sync", + FromNode: ds.localNode.ID, + Timestamp: time.Now().UnixNano(), + Nodes: nodes, + } + + data, err := json.Marshal(response) + if err != nil { + logger.ErrorCF("swarm", "failed to marshal sync message", map[string]any{"error": err}) + return + } + if _, err := ds.conn.WriteToUDP(data, addr); err != nil { + logger.ErrorCF("swarm", "failed to send sync", map[string]any{"to": addr.String(), "error": err}) + return + } + + // Emit join event + event := &NodeEvent{ + Event: EventJoin, + Time: time.Now().UnixNano(), + } + if len(msg.Nodes) > 0 { + event.Node = msg.Nodes[0] + } + ds.eventHandler.Dispatch(event) +} + +// handleUpdate handles a node update message. +func (ds *DiscoveryService) handleUpdate(msg GossipMessage) { + if len(msg.Nodes) == 0 { + return + } + + for _, node := range msg.Nodes { + if node.ID != ds.localNode.ID { + existing, ok := ds.membership.GetNode(node.ID) + if !ok || node.Timestamp > existing.Node.Timestamp { + ds.membership.UpdateNode(node) + } + } + } +} + +// handleSync handles a sync response with member list. +func (ds *DiscoveryService) handleSync(msg GossipMessage, addr *net.UDPAddr) { + for _, node := range msg.Nodes { + if node.ID != ds.localNode.ID { + ds.membership.UpdateNode(node) + } + } +} + +// broadcastUpdate broadcasts local state to random members. +func (ds *DiscoveryService) broadcastUpdate() { + members := ds.membership.GetMembers() + if len(members) == 0 { + return + } + + msg := GossipMessage{ + Type: "update", + FromNode: ds.localNode.ID, + SeqNum: ds.seqNum, + Timestamp: time.Now().UnixNano(), + Nodes: []*NodeInfo{ds.localNode}, + } + + // Add auth token if authentication is enabled + if ds.auth != nil { + token, err := ds.auth.GenerateToken() + if err != nil { + logger.ErrorCF("swarm", "failed to generate auth token", map[string]any{"error": err}) + } else { + msg.AuthToken = token + } + } + + data, err := json.Marshal(msg) + if err != nil { + logger.ErrorCF("swarm", "failed to marshal broadcast update", map[string]any{"error": err}) + return + } + + // Send to a few random members + for _, member := range members { + if member.Node.ID != ds.localNode.ID { + addr := fmt.Sprintf("%s:%d", member.Node.Addr, member.Node.Port) + udpAddr, err := net.ResolveUDPAddr("udp", addr) + if err != nil { + logger.DebugCF("swarm", "failed to resolve address", map[string]any{"address": addr, "error": err}) + continue + } + if _, err := ds.conn.WriteToUDP(data, udpAddr); err != nil { + logger.DebugCF("swarm", "failed to send update", map[string]any{"address": addr, "error": err}) + } + } + } +} + +// sendJoin sends a join message to a specific address. +func (ds *DiscoveryService) sendJoin(ctx context.Context, addr string) error { + joinAddr, err := net.ResolveUDPAddr("udp", addr) + if err != nil { + return err + } + + msg := GossipMessage{ + Type: "join", + FromNode: ds.localNode.ID, + Timestamp: time.Now().UnixNano(), + Nodes: []*NodeInfo{ds.localNode}, + } + + data, err := json.Marshal(msg) + if err != nil { + return fmt.Errorf("failed to marshal join message: %w", err) + } + + // Set deadline + ds.conn.SetWriteDeadline(time.Now().Add(DefaultUDPWriteDeadline)) + _, err = ds.conn.WriteToUDP(data, joinAddr) + if err != nil { + return fmt.Errorf("failed to send join to %s: %w", addr, err) + } + return nil +} + +// getLocalIP returns the local IP address. +func getLocalIP() string { + addrs, err := net.InterfaceAddrs() + if err != nil { + return "" + } + + for _, addr := range addrs { + if ipnet, ok := addr.(*net.IPNet); ok && !ipnet.IP.IsLoopback() { + if ipnet.IP.To4() != nil { + return ipnet.IP.String() + } + } + } + return "" +} diff --git a/pkg/swarm/errors.go b/pkg/swarm/errors.go new file mode 100644 index 000000000..ed8cad7c7 --- /dev/null +++ b/pkg/swarm/errors.go @@ -0,0 +1,50 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import "errors" + +var ( + // ErrNodeNotFound is returned when a node is not found in the cluster. + ErrNodeNotFound = errors.New("node not found") + + // ErrNodeNotAvailable is returned when a node is not available for handoff. + ErrNodeNotAvailable = errors.New("node not available") + + // ErrNoHealthyNodes is returned when no healthy nodes are available. + ErrNoHealthyNodes = errors.New("no healthy nodes available") + + // ErrHandoffTimeout is returned when a handoff operation times out. + ErrHandoffTimeout = errors.New("handoff timeout") + + // ErrHandoffRejected is returned when a handoff is rejected by the target node. + ErrHandoffRejected = errors.New("handoff rejected") + + // ErrHandoffInProgress is returned when a handoff is already in progress. + ErrHandoffInProgress = errors.New("handoff already in progress") + + // ErrInvalidNodeInfo is returned when node information is invalid. + ErrInvalidNodeInfo = errors.New("invalid node information") + + // ErrDiscoveryDisabled is returned when discovery is disabled. + ErrDiscoveryDisabled = errors.New("discovery disabled") + + // ErrTransportClosed is returned when the transport is closed. + ErrTransportClosed = errors.New("transport closed") + + // ErrSessionNotFound is returned when a session is not found. + ErrSessionNotFound = errors.New("session not found") + + // ErrCapabilityNotSupported is returned when a required capability is not supported. + ErrCapabilityNotSupported = errors.New("capability not supported") + + // ErrAuthenticationFailed is returned when authentication fails. + ErrAuthenticationFailed = errors.New("authentication failed") + + // ErrInvalidSignature is returned when a signature verification fails. + ErrInvalidSignature = errors.New("invalid signature") +) diff --git a/pkg/swarm/handoff.go b/pkg/swarm/handoff.go new file mode 100644 index 000000000..ccfa0163d --- /dev/null +++ b/pkg/swarm/handoff.go @@ -0,0 +1,486 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "context" + "encoding/json" + "fmt" + "net" + "sync" + "time" + + "github.com/google/uuid" +) + +// HandoffReason represents the reason for a handoff. +type HandoffReason string + +const ( + ReasonOverloaded HandoffReason = "overloaded" // Load is too high + ReasonNoCapability HandoffReason = "no_capability" // Missing capability + ReasonUserRequest HandoffReason = "user_request" // User explicitly requested + ReasonNodeLeave HandoffReason = "node_leave" // Node is leaving + ReasonShutdown HandoffReason = "shutdown" // Graceful shutdown +) + +// HandoffState represents the state of a handoff operation. +type HandoffState string + +const ( + HandoffStatePending HandoffState = "pending" + HandoffStateAccepted HandoffState = "accepted" + HandoffStateRejected HandoffState = "rejected" + HandoffStateCompleted HandoffState = "completed" + HandoffStateFailed HandoffState = "failed" + HandoffStateTimeout HandoffState = "timeout" +) + +// HandoffRequest represents a request to hand off a session. +type HandoffRequest struct { + RequestID string `json:"request_id"` + Reason HandoffReason `json:"reason"` + SessionKey string `json:"session_key"` + SessionMessages []SessionMessage `json:"session_messages,omitempty"` + Context map[string]any `json:"context,omitempty"` + RequiredCap string `json:"required_cap,omitempty"` + Metadata map[string]string `json:"metadata,omitempty"` + FromNodeID string `json:"from_node_id"` + FromNodeAddr string `json:"from_node_addr"` + Timestamp int64 `json:"timestamp"` +} + +// HandoffResponse represents the response to a handoff request. +type HandoffResponse struct { + RequestID string `json:"request_id"` + Accepted bool `json:"accepted"` + NodeID string `json:"node_id"` + Reason string `json:"reason,omitempty"` + SessionKey string `json:"session_key,omitempty"` // New session key on target + Timestamp int64 `json:"timestamp"` + State HandoffState `json:"state"` +} + +// HandoffCoordinator coordinates handoff operations between nodes. +type HandoffCoordinator struct { + discovery *DiscoveryService + membership *MembershipManager + config HandoffConfig + + pending map[string]*HandoffOperation // request_id -> operation + mu sync.RWMutex + conn *net.UDPConn + + // Accept/reject callbacks + onHandoffRequest func(*HandoffRequest) *HandoffResponse + onHandoffComplete func(*HandoffRequest, *HandoffResponse) +} + +// HandoffOperation represents an ongoing handoff operation. +type HandoffOperation struct { + Request *HandoffRequest + Response *HandoffResponse + State HandoffState + StartTime time.Time + LastUpdate time.Time + RetryCount int + TargetNode *NodeWithState +} + +// NewHandoffCoordinator creates a new handoff coordinator. +func NewHandoffCoordinator(ds *DiscoveryService, config HandoffConfig) (*HandoffCoordinator, error) { + hc := &HandoffCoordinator{ + discovery: ds, + membership: ds.membership, + config: config, + pending: make(map[string]*HandoffOperation), + } + + // Bind UDP socket for handoff messages + addr := fmt.Sprintf("%s:%d", ds.config.BindAddr, ds.config.RPC.Port) + udpAddr, err := net.ResolveUDPAddr("udp", addr) + if err != nil { + return nil, fmt.Errorf("failed to resolve RPC address: %w", err) + } + + hc.conn, err = net.ListenUDP("udp", udpAddr) + if err != nil { + return nil, fmt.Errorf("failed to listen for RPC: %w", err) + } + + // Start message handler + go hc.messageHandler() + + return hc, nil +} + +// Close closes the handoff coordinator. +func (hc *HandoffCoordinator) Close() error { + if hc.conn != nil { + return hc.conn.Close() + } + return nil +} + +// CanHandle checks if the local node can handle a request. +func (hc *HandoffCoordinator) CanHandle(requiredCap string) bool { + if !hc.config.Enabled { + return false + } + + // Check load + loadScore := hc.discovery.localNode.LoadScore + if loadScore > hc.config.LoadThreshold { + return false + } + + // Check capability + if requiredCap != "" { + hasCap := false + for _, cap := range hc.discovery.localNode.AgentCaps { + if cap == requiredCap { + hasCap = true + break + } + } + if !hasCap { + return false + } + } + + return true +} + +// InitiateHandoff initiates a handoff to another node. +func (hc *HandoffCoordinator) InitiateHandoff(ctx context.Context, req *HandoffRequest) (*HandoffResponse, error) { + if req.RequestID == "" { + req.RequestID = uuid.New().String() + } + + req.FromNodeID = hc.discovery.localNode.ID + req.FromNodeAddr = fmt.Sprintf("%s:%d", hc.discovery.localNode.Addr, hc.discovery.config.RPC.Port) + req.Timestamp = time.Now().UnixNano() + + // Find target node + targetNode, err := hc.findTargetNode(req) + if err != nil { + return &HandoffResponse{ + RequestID: req.RequestID, + Accepted: false, + Reason: err.Error(), + State: HandoffStateFailed, + }, nil + } + + // Create operation + op := &HandoffOperation{ + Request: req, + State: HandoffStatePending, + StartTime: time.Now(), + LastUpdate: time.Now(), + TargetNode: targetNode, + } + + hc.mu.Lock() + hc.pending[req.RequestID] = op + hc.mu.Unlock() + + // Send request + err = hc.sendHandoffRequest(req, targetNode) + if err != nil { + hc.mu.Lock() + op.State = HandoffStateFailed + delete(hc.pending, req.RequestID) + hc.mu.Unlock() + + return &HandoffResponse{ + RequestID: req.RequestID, + Accepted: false, + Reason: err.Error(), + State: HandoffStateFailed, + }, nil + } + + // Wait for response with timeout + timeout := hc.config.Timeout.Duration + if timeout == 0 { + timeout = DefaultHandoffTimeout + } + + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + resp := hc.waitForResponse(ctx, req.RequestID) + + // Retry if needed (op.RetryCount is 0 at this point, representing the first attempt) + for !resp.Accepted && op.RetryCount < hc.config.MaxRetries { + op.RetryCount++ + + // Find new target + newTarget, err := hc.findTargetNode(req) + if err != nil { + continue + } + op.TargetNode = newTarget + + // Delay before retry + time.Sleep(hc.config.RetryDelay.Duration) + + // Send request + err = hc.sendHandoffRequest(req, newTarget) + if err != nil { + continue + } + + // Wait for response with timeout, preserving parent context + retryCtx, retryCancel := context.WithTimeout(ctx, timeout) + resp = hc.waitForResponse(retryCtx, req.RequestID) + retryCancel() + + // If accepted, break out of retry loop + if resp.Accepted { + break + } + } + + // Clean up + hc.mu.Lock() + delete(hc.pending, req.RequestID) + hc.mu.Unlock() + + // Notify callback + if hc.onHandoffComplete != nil { + go hc.onHandoffComplete(req, resp) + } + + return resp, nil +} + +// findTargetNode finds a suitable target node for handoff. +func (hc *HandoffCoordinator) findTargetNode(req *HandoffRequest) (*NodeWithState, error) { + var candidates []*NodeWithState + + if req.RequiredCap != "" { + // Find nodes with required capability + candidates = hc.membership.SelectByCapability([]string{req.RequiredCap}) + } else { + // Find all available nodes + candidates = hc.membership.GetAvailableMembers() + } + + if len(candidates) == 0 { + return nil, ErrNoHealthyNodes + } + + // Select least loaded node + target := candidates[0] + for _, c := range candidates[1:] { + if c.Node.LoadScore < target.Node.LoadScore { + target = c + } + } + + return target, nil +} + +// sendHandoffRequest sends a handoff request to a target node. +func (hc *HandoffCoordinator) sendHandoffRequest(req *HandoffRequest, target *NodeWithState) error { + // Handoff message type + msg := map[string]any{ + "type": "handoff_request", + "payload": req, + } + + data, err := json.Marshal(msg) + if err != nil { + return err + } + + addr := fmt.Sprintf("%s:%d", target.Node.Addr, hc.discovery.config.RPC.Port) + udpAddr, err := net.ResolveUDPAddr("udp", addr) + if err != nil { + return err + } + + hc.conn.SetWriteDeadline(time.Now().Add(DefaultUDPWriteDeadline)) + _, err = hc.conn.WriteToUDP(data, udpAddr) + return err +} + +// waitForResponse waits for a handoff response. +func (hc *HandoffCoordinator) waitForResponse(ctx context.Context, requestID string) *HandoffResponse { + ticker := time.NewTicker(HandoffResponsePollInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + hc.mu.Lock() + if op, ok := hc.pending[requestID]; ok { + op.State = HandoffStateTimeout + } + hc.mu.Unlock() + + return &HandoffResponse{ + RequestID: requestID, + Accepted: false, + Reason: "timeout", + State: HandoffStateTimeout, + } + case <-ticker.C: + hc.mu.RLock() + op, ok := hc.pending[requestID] + hc.mu.RUnlock() + + if ok && op.Response != nil { + return op.Response + } + } + } +} + +// messageHandler handles incoming handoff messages. +func (hc *HandoffCoordinator) messageHandler() { + buf := make([]byte, MaxGossipMessageSize) + + for { + n, addr, err := hc.conn.ReadFromUDP(buf) + if err != nil { + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + continue + } + return + } + + if n > 0 { + go hc.handleMessage(buf[:n], addr) + } + } +} + +// handleMessage handles an incoming message. +func (hc *HandoffCoordinator) handleMessage(data []byte, addr *net.UDPAddr) { + var msg map[string]any + if err := json.Unmarshal(data, &msg); err != nil { + return + } + + msgType, _ := msg["type"].(string) + + switch msgType { + case "handoff_request": + hc.handleHandoffRequest(data, addr) + case "handoff_response": + hc.handleHandoffResponse(data) + } +} + +// handleHandoffRequest handles a handoff request from another node. +func (hc *HandoffCoordinator) handleHandoffRequest(data []byte, addr *net.UDPAddr) { + var msg map[string]any + if err := json.Unmarshal(data, &msg); err != nil { + return + } + + payloadData, _ := json.Marshal(msg["payload"]) + var req HandoffRequest + if err := json.Unmarshal(payloadData, &req); err != nil { + return + } + + // Check if we can handle it + accepted := hc.CanHandle(req.RequiredCap) + response := &HandoffResponse{ + RequestID: req.RequestID, + Accepted: accepted, + NodeID: hc.discovery.localNode.ID, + State: HandoffStateAccepted, + Timestamp: time.Now().UnixNano(), + } + + if !accepted { + response.Reason = "cannot handle (overloaded or missing capability)" + response.State = HandoffStateRejected + } + + // Call custom handler if set + if hc.onHandoffRequest != nil { + response = hc.onHandoffRequest(&req) + } + + // Send response + respMsg := map[string]any{ + "type": "handoff_response", + "payload": response, + } + + respData, _ := json.Marshal(respMsg) + hc.conn.WriteToUDP(respData, addr) + + // Update operation if we accepted + if accepted { + op := &HandoffOperation{ + Request: &req, + Response: response, + State: HandoffStateAccepted, + } + + hc.mu.Lock() + hc.pending[req.RequestID] = op + hc.mu.Unlock() + } +} + +// handleHandoffResponse handles a handoff response. +func (hc *HandoffCoordinator) handleHandoffResponse(data []byte) { + var msg map[string]any + if err := json.Unmarshal(data, &msg); err != nil { + return + } + + payloadData, _ := json.Marshal(msg["payload"]) + var resp HandoffResponse + if err := json.Unmarshal(payloadData, &resp); err != nil { + return + } + + hc.mu.Lock() + defer hc.mu.Unlock() + + if op, ok := hc.pending[resp.RequestID]; ok { + op.Response = &resp + op.LastUpdate = time.Now() + + if resp.Accepted { + op.State = HandoffStateAccepted + } else { + op.State = HandoffStateRejected + } + } +} + +// SetRequestHandler sets a custom handler for handoff requests. +func (hc *HandoffCoordinator) SetRequestHandler(handler func(*HandoffRequest) *HandoffResponse) { + hc.onHandoffRequest = handler +} + +// SetCompleteHandler sets a callback for handoff completion. +func (hc *HandoffCoordinator) SetCompleteHandler(handler func(*HandoffRequest, *HandoffResponse)) { + hc.onHandoffComplete = handler +} + +// GetPending returns all pending handoff operations. +func (hc *HandoffCoordinator) GetPending() []*HandoffOperation { + hc.mu.RLock() + defer hc.mu.RUnlock() + + result := make([]*HandoffOperation, 0, len(hc.pending)) + for _, op := range hc.pending { + result = append(result, op) + } + return result +} diff --git a/pkg/swarm/leader_election.go b/pkg/swarm/leader_election.go new file mode 100644 index 000000000..7146bd9fe --- /dev/null +++ b/pkg/swarm/leader_election.go @@ -0,0 +1,237 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "context" + "sync" + "time" + + "github.com/sipeed/picoclaw/pkg/logger" +) + +// LeaderElection handles leader election using a simple bully algorithm variant. +type LeaderElection struct { + localNodeID string + membership *MembershipManager + + mu sync.RWMutex + currentLeader string + isLeader bool + electionInProgress bool + leaderChangeCh chan string + stopCh chan struct{} +} + +// NewLeaderElection creates a new leader election instance. +func NewLeaderElection(nodeID string, membership *MembershipManager) *LeaderElection { + return &LeaderElection{ + localNodeID: nodeID, + membership: membership, + leaderChangeCh: make(chan string, 10), + stopCh: make(chan struct{}), + } +} + +// Start starts the leader election process. +func (le *LeaderElection) Start() { + // Start election checker + go le.electionChecker() + + // Start leader heartbeat monitor + go le.leaderMonitor() +} + +// Stop stops the leader election process. +func (le *LeaderElection) Stop() { + close(le.stopCh) +} + +// IsLeader returns true if this node is the current leader. +func (le *LeaderElection) IsLeader() bool { + le.mu.RLock() + defer le.mu.RUnlock() + return le.isLeader +} + +// GetLeader returns the current leader ID. +func (le *LeaderElection) GetLeader() string { + le.mu.RLock() + defer le.mu.RUnlock() + return le.currentLeader +} + +// LeaderChanges returns a channel that receives leader ID changes. +func (le *LeaderElection) LeaderChanges() <-chan string { + return le.leaderChangeCh +} + +// electionChecker periodically checks if we should become leader. +func (le *LeaderElection) electionChecker() { + ticker := time.NewTicker(time.Second * 5) + defer ticker.Stop() + + for { + select { + case <-le.stopCh: + return + case <-ticker.C: + le.checkElection() + } + } +} + +// checkElection runs the leader election algorithm. +func (le *LeaderElection) checkElection() { + le.mu.Lock() + defer le.mu.Unlock() + + members := le.membership.GetMembers() + if len(members) == 0 { + // No other members, we become leader + le.becomeLeader() + return + } + + // Find the node with the lowest ID (simple deterministic leader selection) + var candidateID string + candidateID = le.localNodeID + + for _, m := range members { + if m.Node.ID < candidateID { + candidateID = m.Node.ID + } + } + + // Update current leader + if le.currentLeader != candidateID { + oldLeader := le.currentLeader + le.currentLeader = candidateID + + if candidateID == le.localNodeID { + le.becomeLeader() + } else { + le.becomeFollower() + } + + logger.InfoCF("swarm", "Leader changed", map[string]any{"old_leader": oldLeader, "new_leader": candidateID}) + + // Notify followers of leader change + select { + case le.leaderChangeCh <- candidateID: + default: + } + } +} + +// becomeLeader marks this node as the leader. +func (le *LeaderElection) becomeLeader() { + if !le.isLeader { + le.isLeader = true + logger.InfoCF("swarm", "This node is now the leader", map[string]any{"node_id": le.localNodeID}) + + // Notify listeners + select { + case le.leaderChangeCh <- le.localNodeID: + default: + } + } +} + +// becomeFollower marks this node as a follower. +func (le *LeaderElection) becomeFollower() { + if le.isLeader { + le.isLeader = false + logger.InfoCF("swarm", "This node is now a follower", map[string]any{"node_id": le.localNodeID}) + } +} + +// leaderMonitor monitors if the current leader is still alive. +func (le *LeaderElection) leaderMonitor() { + ticker := time.NewTicker(time.Second * 10) + defer ticker.Stop() + + for { + select { + case <-le.stopCh: + return + case <-ticker.C: + le.monitorLeader() + } + } +} + +// monitorLeader checks if the current leader is still alive. +func (le *LeaderElection) monitorLeader() { + le.mu.RLock() + leaderID := le.currentLeader + amLeader := le.isLeader + le.mu.RUnlock() + + if amLeader || leaderID == "" { + return + } + + // Check if leader is still in the membership + if _, exists := le.membership.GetNode(leaderID); !exists { + logger.WarnCF("swarm", "Leader no longer in membership, triggering reelection", map[string]any{"leader_id": leaderID}) + // Trigger reelection by clearing current leader + le.mu.Lock() + le.currentLeader = "" + le.mu.Unlock() + le.checkElection() + } +} + +// ElectLeader triggers a new leader election. +func (le *LeaderElection) ElectLeader(ctx context.Context) (string, error) { + le.mu.Lock() + le.currentLeader = "" // Clear current leader to trigger reelection + le.mu.Unlock() + + // Run election immediately + le.checkElection() + + // Wait for new leader + ticker := time.NewTicker(time.Millisecond * 100) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-ticker.C: + le.mu.RLock() + leader := le.currentLeader + le.mu.RUnlock() + + if leader != "" { + return leader, nil + } + } + } +} + +// LeadershipState represents the current leadership state. +type LeadershipState struct { + LeaderID string `json:"leader_id"` + IsLeader bool `json:"is_leader"` + LastChange time.Time `json:"last_change"` + MemberCount int `json:"member_count"` +} + +// GetState returns the current leadership state. +func (le *LeaderElection) GetState() LeadershipState { + le.mu.RLock() + defer le.mu.RUnlock() + + return LeadershipState{ + LeaderID: le.currentLeader, + IsLeader: le.isLeader, + MemberCount: len(le.membership.GetMembers()), + } +} diff --git a/pkg/swarm/load_monitor.go b/pkg/swarm/load_monitor.go new file mode 100644 index 000000000..fe3280ead --- /dev/null +++ b/pkg/swarm/load_monitor.go @@ -0,0 +1,295 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "runtime" + "sync" + "time" +) + +// LoadMonitor monitors system load and calculates a load score. +type LoadMonitor struct { + config *LoadMonitorConfig + samples []float64 + mu sync.RWMutex + sessionCount int + ticker *time.Ticker + stopChan chan struct{} + onThreshold []func(float64) +} + +// NewLoadMonitor creates a new load monitor. +func NewLoadMonitor(config *LoadMonitorConfig) *LoadMonitor { + if config.SampleSize <= 0 { + config.SampleSize = 60 + } + if config.Interval.Duration <= 0 { + config.Interval = Duration{5 * time.Second} + } + + lm := &LoadMonitor{ + config: config, + samples: make([]float64, 0, config.SampleSize), + stopChan: make(chan struct{}), + onThreshold: make([]func(float64), 0), + } + return lm +} + +// Start begins monitoring load. +func (lm *LoadMonitor) Start() { + if lm.ticker != nil { + return + } + + lm.ticker = time.NewTicker(lm.config.Interval.Duration) + go lm.run() +} + +// Stop stops monitoring load. +func (lm *LoadMonitor) Stop() { + if lm.ticker == nil { + return + } + + lm.ticker.Stop() + close(lm.stopChan) + lm.ticker = nil +} + +// run is the main monitoring loop. +func (lm *LoadMonitor) run() { + for { + select { + case <-lm.ticker.C: + score := lm.calculateScore() + lm.addSample(score) + + // Check threshold callbacks + if lm.shouldOffload() { + lm.mu.RLock() + callbacks := make([]func(float64), len(lm.onThreshold)) + copy(callbacks, lm.onThreshold) + lm.mu.RUnlock() + + for _, cb := range callbacks { + go cb(score) + } + } + case <-lm.stopChan: + return + } + } +} + +// LoadMetrics represents current load metrics. +type LoadMetrics struct { + CPUUsage float64 `json:"cpu_usage"` + MemoryUsage float64 `json:"memory_usage"` + ActiveSessions int `json:"active_sessions"` + Goroutines int `json:"goroutines"` + Score float64 `json:"score"` + Timestamp int64 `json:"timestamp"` +} + +// GetCurrentLoad returns the current load metrics. +func (lm *LoadMonitor) GetCurrentLoad() *LoadMetrics { + metrics := &LoadMetrics{ + ActiveSessions: lm.GetSessionCount(), + Goroutines: runtime.NumGoroutine(), + Timestamp: time.Now().UnixNano(), + } + + // Get memory usage + var m runtime.MemStats + runtime.ReadMemStats(&m) + + // Normalize using configured thresholds + maxMem := lm.config.MaxMemoryBytes + if maxMem == 0 { + maxMem = 1024 * 1024 * 1024 // Default 1GB + } + metrics.MemoryUsage = normalizeMemory(m.Alloc, maxMem) + + maxGoroutines := lm.config.MaxGoroutines + if maxGoroutines == 0 { + maxGoroutines = 1000 + } + metrics.CPUUsage = normalizeCPU(metrics.Goroutines, maxGoroutines) + + maxSessions := lm.config.MaxSessions + if maxSessions == 0 { + maxSessions = 100 + } + sessionUsage := normalizeSessions(metrics.ActiveSessions, maxSessions) + + // Calculate weighted score + config := lm.config + metrics.Score = (metrics.CPUUsage * config.CPUWeight) + + (metrics.MemoryUsage * config.MemoryWeight) + + (sessionUsage * config.SessionWeight) + + // Clamp score to [0, 1] + if metrics.Score < 0 { + metrics.Score = 0 + } else if metrics.Score > 1 { + metrics.Score = 1 + } + + return metrics +} + +// calculateScore calculates the current load score. +func (lm *LoadMonitor) calculateScore() float64 { + return lm.GetCurrentLoad().Score +} + +// addSample adds a load sample to the history. +func (lm *LoadMonitor) addSample(score float64) { + lm.mu.Lock() + defer lm.mu.Unlock() + + lm.samples = append(lm.samples, score) + if len(lm.samples) > lm.config.SampleSize { + lm.samples = lm.samples[1:] + } +} + +// GetAverageScore returns the average load score over the sample window. +func (lm *LoadMonitor) GetAverageScore() float64 { + lm.mu.RLock() + defer lm.mu.RUnlock() + + if len(lm.samples) == 0 { + return lm.calculateScore() + } + + sum := 0.0 + for _, s := range lm.samples { + sum += s + } + return sum / float64(len(lm.samples)) +} + +// GetSessionCount returns the current number of active sessions. +func (lm *LoadMonitor) GetSessionCount() int { + lm.mu.RLock() + defer lm.mu.RUnlock() + return lm.sessionCount +} + +// SetSessionCount sets the current number of active sessions. +func (lm *LoadMonitor) SetSessionCount(count int) { + lm.mu.Lock() + defer lm.mu.Unlock() + lm.sessionCount = count +} + +// IncrementSessions increments the session count. +func (lm *LoadMonitor) IncrementSessions() { + lm.mu.Lock() + defer lm.mu.Unlock() + lm.sessionCount++ +} + +// DecrementSessions decrements the session count. +func (lm *LoadMonitor) DecrementSessions() { + lm.mu.Lock() + defer lm.mu.Unlock() + if lm.sessionCount > 0 { + lm.sessionCount-- + } +} + +// ShouldOffload returns true if the load is high enough to offload tasks. +func (lm *LoadMonitor) ShouldOffload() bool { + return lm.shouldOffload() +} + +// shouldOffload internal check for offloading. +func (lm *LoadMonitor) shouldOffload() bool { + avgScore := lm.GetAverageScore() + currentScore := lm.calculateScore() + + // Use configured offload threshold, or default to 0.8 + threshold := lm.config.OffloadThreshold + if threshold <= 0 { + threshold = 0.8 + } + + // Use a combination of current and average for smoother behavior + combinedScore := (currentScore*0.7 + avgScore*0.3) + return combinedScore > threshold +} + +// OnThreshold registers a callback when the load threshold is exceeded. +func (lm *LoadMonitor) OnThreshold(callback func(float64)) { + lm.mu.Lock() + defer lm.mu.Unlock() + lm.onThreshold = append(lm.onThreshold, callback) +} + +// GetTrend returns the load trend: "increasing", "decreasing", or "stable". +func (lm *LoadMonitor) GetTrend() string { + lm.mu.RLock() + defer lm.mu.RUnlock() + + if len(lm.samples) < 3 { + return "stable" + } + + // Simple linear regression to detect trend + n := float64(len(lm.samples)) + sumX := n*(n-1)/2 + sumY := 0.0 + sumXY := 0.0 + + for i, s := range lm.samples { + x := float64(i) + sumY += s + sumXY += x * s + } + + slope := (n*sumXY - sumX*sumY) / (n*(n-1)*(2*n-1)/6) + + if slope > TrendIncreasingThreshold { + return "increasing" + } else if slope < TrendDecreasingThreshold { + return "decreasing" + } + return "stable" +} + +// Helper functions for normalization + +func normalizeMemory(alloc uint64, maxMem uint64) float64 { + // Use configured max memory threshold + usage := float64(alloc) / float64(maxMem) + if usage > 1 { + return 1 + } + return usage +} + +func normalizeCPU(goroutines int, maxGoroutines int) float64 { + // Use configured max goroutine threshold + usage := float64(goroutines) / float64(maxGoroutines) + if usage > 1 { + return 1 + } + return usage +} + +func normalizeSessions(sessions int, maxSessions int) float64 { + // Use configured max sessions threshold + usage := float64(sessions) / float64(maxSessions) + if usage > 1 { + return 1 + } + return usage +} diff --git a/pkg/swarm/membership.go b/pkg/swarm/membership.go new file mode 100644 index 000000000..46dc48ef5 --- /dev/null +++ b/pkg/swarm/membership.go @@ -0,0 +1,365 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "math/rand" + "sync" + "time" +) + +// MembershipManager manages cluster membership. +type MembershipManager struct { + discovery *DiscoveryService + view *ClusterView + config DiscoveryConfig + mu sync.RWMutex + + // Event callbacks + onJoin []func(*NodeInfo) + onLeave []func(*NodeInfo) + onUpdate []func(*NodeInfo) +} + +// NewMembershipManager creates a new membership manager. +func NewMembershipManager(ds *DiscoveryService, config DiscoveryConfig) *MembershipManager { + localNodeID := ds.LocalNode().ID + return &MembershipManager{ + discovery: ds, + view: NewClusterView(localNodeID), + config: config, + onJoin: make([]func(*NodeInfo), 0), + onLeave: make([]func(*NodeInfo), 0), + onUpdate: make([]func(*NodeInfo), 0), + } +} + +// GetNode retrieves a node by ID. +func (mm *MembershipManager) GetNode(nodeID string) (*NodeWithState, bool) { + return mm.view.Get(nodeID) +} + +// GetMembers returns all members. +func (mm *MembershipManager) GetMembers() []*NodeWithState { + mm.mu.RLock() + defer mm.mu.RUnlock() + return mm.view.List() +} + +// GetAliveMembers returns all alive members. +func (mm *MembershipManager) GetAliveMembers() []*NodeWithState { + mm.mu.RLock() + defer mm.mu.RUnlock() + + members := mm.view.GetAliveNodes() + result := make([]*NodeWithState, 0, len(members)) + for _, m := range members { + if m.Node.ID != mm.discovery.LocalNode().ID { + result = append(result, m) + } + } + return result +} + +// GetAvailableMembers returns all available members (alive and not overloaded). +func (mm *MembershipManager) GetAvailableMembers() []*NodeWithState { + mm.mu.RLock() + defer mm.mu.RUnlock() + + members := mm.view.GetAvailableNodes() + result := make([]*NodeWithState, 0, len(members)) + for _, m := range members { + if m.Node.ID != mm.discovery.LocalNode().ID { + result = append(result, m) + } + } + return result +} + +// UpdateNode updates or adds a node to the membership. +func (mm *MembershipManager) UpdateNode(node *NodeInfo) *NodeWithState { + mm.mu.Lock() + + existing, existed := mm.view.Get(node.ID) + nws := mm.view.AddOrUpdate(node) + + if !existed { + // New node joined + nws.State.Status = NodeStatusAlive + nws.State.StatusSince = time.Now().UnixNano() + nws.State.LastSeen = time.Now().UnixNano() + + // Notify callbacks + for _, cb := range mm.onJoin { + go cb(node) + } + + // Dispatch event + mm.discovery.eventHandler.Dispatch(&NodeEvent{ + Node: node, + Event: EventJoin, + Time: time.Now().UnixNano(), + }) + } else { + // Existing node updated + if existing.Node.Timestamp < node.Timestamp { + nws.State.LastSeen = time.Now().UnixNano() + + // Mark as alive if was suspect/dead + if nws.State.Status != NodeStatusAlive { + nws.State.UpdateStatus(NodeStatusAlive) + nws.State.PingFailure = 0 + nws.State.PingSuccess++ + } + + // Notify callbacks + for _, cb := range mm.onUpdate { + go cb(node) + } + + // Dispatch event + mm.discovery.eventHandler.Dispatch(&NodeEvent{ + Node: node, + Event: EventUpdate, + Time: time.Now().UnixNano(), + }) + } + } + + mm.mu.Unlock() + return nws +} + +// RemoveNode removes a node from the membership. +func (mm *MembershipManager) RemoveNode(nodeID string) { + mm.mu.Lock() + + nws, exists := mm.view.Get(nodeID) + if !exists { + mm.mu.Unlock() + return + } + + mm.view.Remove(nodeID) + + // Notify callbacks + for _, cb := range mm.onLeave { + go cb(nws.Node) + } + + // Dispatch event + mm.discovery.eventHandler.Dispatch(&NodeEvent{ + Node: nws.Node, + Event: EventLeave, + Time: time.Now().UnixNano(), + }) + + mm.mu.Unlock() +} + +// RecordHeartbeat records a heartbeat for a node. +func (mm *MembershipManager) RecordHeartbeat(nodeID string) { + mm.mu.Lock() + defer mm.mu.Unlock() + + nws, exists := mm.view.Get(nodeID) + if !exists { + return + } + + nws.State.LastPing = time.Now().UnixNano() + nws.State.LastSeen = time.Now().UnixNano() + + // Reset failure count and increment success + nws.State.PingFailure = 0 + nws.State.PingSuccess++ + + // Mark as alive if was suspect + if nws.State.Status != NodeStatusAlive { + nws.State.UpdateStatus(NodeStatusAlive) + } +} + +// MarkSuspect marks a node as suspect (possibly dead). +func (mm *MembershipManager) MarkSuspect(nodeID string) { + mm.mu.Lock() + defer mm.mu.Unlock() + + nws, exists := mm.view.Get(nodeID) + if !exists { + return + } + + if nws.State.Status == NodeStatusAlive { + nws.State.UpdateStatus(NodeStatusSuspect) + nws.State.PingFailure++ + } +} + +// MarkDead marks a node as dead. +func (mm *MembershipManager) MarkDead(nodeID string) { + mm.mu.Lock() + defer mm.mu.Unlock() + + nws, exists := mm.view.Get(nodeID) + if !exists { + return + } + + if nws.State.Status != NodeStatusDead { + nws.State.UpdateStatus(NodeStatusDead) + + // Remove from view after a delay + go func() { + time.Sleep(mm.config.DeadNodeTimeout.Duration) + mm.RemoveNode(nodeID) + }() + } +} + +// CheckHealth checks the health of all members and marks dead nodes. +func (mm *MembershipManager) CheckHealth() { + mm.mu.RLock() + members := mm.view.List() + nodeTimeout := mm.config.NodeTimeout.Duration + deadTimeout := mm.config.DeadNodeTimeout.Duration + localNodeID := mm.discovery.LocalNode().ID + mm.mu.RUnlock() + + now := time.Now() + + for _, m := range members { + // Skip local node + if m.Node.ID == localNodeID { + continue + } + + lastSeen := time.Unix(0, m.State.LastSeen) + age := now.Sub(lastSeen) + + switch m.State.Status { + case NodeStatusAlive: + if age > nodeTimeout { + mm.MarkSuspect(m.Node.ID) + } + case NodeStatusSuspect: + if age > deadTimeout { + mm.MarkDead(m.Node.ID) + } + } + } +} + +// SelectByCapability selects members that have the required capabilities. +func (mm *MembershipManager) SelectByCapability(requiredCaps []string) []*NodeWithState { + mm.mu.RLock() + defer mm.mu.RUnlock() + + if len(requiredCaps) == 0 { + return mm.GetAvailableMembers() + } + + members := mm.view.GetAvailableNodes() + result := make([]*NodeWithState, 0) + + for _, m := range members { + if m.Node.ID == mm.discovery.LocalNode().ID { + continue + } + + // Check if node has all required capabilities + hasAll := true + for _, cap := range requiredCaps { + found := false + for _, nodeCap := range m.Node.AgentCaps { + if nodeCap == cap { + found = true + break + } + } + if !found { + hasAll = false + break + } + } + + if hasAll { + result = append(result, m) + } + } + + return result +} + +// SelectLeastLoaded selects the member with the lowest load score. +func (mm *MembershipManager) SelectLeastLoaded() *NodeWithState { + members := mm.GetAvailableMembers() + if len(members) == 0 { + return nil + } + + least := members[0] + for _, m := range members[1:] { + if m.Node.LoadScore < least.Node.LoadScore { + least = m + } + } + + return least +} + +// SelectRandom selects a random available member. +func (mm *MembershipManager) SelectRandom() *NodeWithState { + members := mm.GetAvailableMembers() + if len(members) == 0 { + return nil + } + + // Use crypto/rand for better random distribution + idx := rand.Intn(len(members)) + return members[idx] +} + +// GetClusterSize returns the current cluster size. +func (mm *MembershipManager) GetClusterSize() int { + mm.mu.RLock() + defer mm.mu.RUnlock() + return mm.view.Size +} + +// OnJoin registers a callback for node join events. +func (mm *MembershipManager) OnJoin(callback func(*NodeInfo)) { + mm.mu.Lock() + defer mm.mu.Unlock() + mm.onJoin = append(mm.onJoin, callback) +} + +// OnLeave registers a callback for node leave events. +func (mm *MembershipManager) OnLeave(callback func(*NodeInfo)) { + mm.mu.Lock() + defer mm.mu.Unlock() + mm.onLeave = append(mm.onLeave, callback) +} + +// OnUpdate registers a callback for node update events. +func (mm *MembershipManager) OnUpdate(callback func(*NodeInfo)) { + mm.mu.Lock() + defer mm.mu.Unlock() + mm.onUpdate = append(mm.onUpdate, callback) +} + +// StartHealthCheck starts the health check routine. +func (mm *MembershipManager) StartHealthCheck(interval time.Duration) { + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for range ticker.C { + mm.CheckHealth() + } + }() +} diff --git a/pkg/swarm/metrics.go b/pkg/swarm/metrics.go new file mode 100644 index 000000000..7d3eb8ebb --- /dev/null +++ b/pkg/swarm/metrics.go @@ -0,0 +1,286 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "encoding/json" + "fmt" + "sync" + "sync/atomic" + "time" +) + +// MetricsCollector collects and exports metrics for the swarm cluster. +type MetricsCollector struct { + mu sync.RWMutex + + // Counters (atomic for performance) + messagesSent atomic.Int64 + messagesReceived atomic.Int64 + handoffsInitiated atomic.Int64 + handoffsAccepted atomic.Int64 + handoffsRejected atomic.Int64 + handoffsFailed atomic.Int64 + electionsWon atomic.Int64 + + // Gauges (use atomic.Value for float64) + currentLoadScore atomic.Value // float64 + activeSessions atomic.Int64 + memberCount atomic.Int32 + + // Histogram data (simplified) + latencyBuckets map[string]*LatencyBucket + + startTime time.Time +} + +// LatencyBucket tracks latency distribution. +type LatencyBucket struct { + mu sync.RWMutex + count int64 + sum int64 + buckets [12]int64 // 0-1ms, 1-2ms, 2-5ms, 5-10ms, 10-20ms, 20-50ms, 50-100ms, 100-200ms, 200-500ms, 500ms-1s, 1-2s, 2s+ +} + +// NewMetricsCollector creates a new metrics collector. +func NewMetricsCollector() *MetricsCollector { + mc := &MetricsCollector{ + latencyBuckets: make(map[string]*LatencyBucket), + startTime: time.Now(), + } + return mc +} + +// Counter methods + +// MessagesSent increments the sent message counter. +func (m *MetricsCollector) MessagesSent(n int64) { + m.messagesSent.Add(n) +} + +// MessagesReceived increments the received message counter. +func (m *MetricsCollector) MessagesReceived(n int64) { + m.messagesReceived.Add(n) +} + +// HandoffInitiated increments the handoff initiated counter. +func (m *MetricsCollector) HandoffInitiated() { + m.handoffsInitiated.Add(1) +} + +// HandoffAccepted increments the handoff accepted counter. +func (m *MetricsCollector) HandoffAccepted() { + m.handoffsAccepted.Add(1) +} + +// HandoffRejected increments the handoff rejected counter. +func (m *MetricsCollector) HandoffRejected() { + m.handoffsRejected.Add(1) +} + +// HandoffFailed increments the handoff failed counter. +func (m *MetricsCollector) HandoffFailed() { + m.handoffsFailed.Add(1) +} + +// ElectionWon increments the elections won counter. +func (m *MetricsCollector) ElectionWon() { + m.electionsWon.Add(1) +} + +// Gauge methods + +// SetLoadScore sets the current load score. +func (m *MetricsCollector) SetLoadScore(score float64) { + m.currentLoadScore.Store(score) +} + +// SetActiveSessions sets the current active session count. +func (m *MetricsCollector) SetActiveSessions(count int64) { + m.activeSessions.Store(count) +} + +// SetMemberCount sets the current cluster member count. +func (m *MetricsCollector) SetMemberCount(count int32) { + m.memberCount.Store(count) +} + +// RecordLatency records a latency observation for the given operation. +func (m *MetricsCollector) RecordLatency(operation string, latency time.Duration) { + m.mu.Lock() + if m.latencyBuckets[operation] == nil { + m.latencyBuckets[operation] = &LatencyBucket{} + } + bucket := m.latencyBuckets[operation] + m.mu.Unlock() + + ms := int64(latency.Milliseconds()) + + bucket.mu.Lock() + bucket.count++ + bucket.sum += ms + + // Bucket the latency + switch { + case ms < 1: + bucket.buckets[0]++ + case ms < 2: + bucket.buckets[1]++ + case ms < 5: + bucket.buckets[2]++ + case ms < 10: + bucket.buckets[3]++ + case ms < 20: + bucket.buckets[4]++ + case ms < 50: + bucket.buckets[5]++ + case ms < 100: + bucket.buckets[6]++ + case ms < 200: + bucket.buckets[7]++ + case ms < 500: + bucket.buckets[8]++ + case ms < 1000: + bucket.buckets[9]++ + case ms < 2000: + bucket.buckets[10]++ + default: + bucket.buckets[11]++ + } + bucket.mu.Unlock() +} + +// GetMetrics returns the current metrics as a map. +func (m *MetricsCollector) GetMetrics() map[string]any { + m.mu.RLock() + defer m.mu.RUnlock() + + latency := make(map[string]any) + for name, bucket := range m.latencyBuckets { + bucket.mu.RLock() + latency[name] = map[string]any{ + "count": bucket.count, + "avg_ms": float64(bucket.sum) / float64(bucket.count), + "p50_ms": m.percentile(bucket, 0.50), + "p95_ms": m.percentile(bucket, 0.95), + "p99_ms": m.percentile(bucket, 0.99), + } + bucket.mu.RUnlock() + } + + return map[string]any{ + // Counters + "messages_sent": m.messagesSent.Load(), + "messages_received": m.messagesReceived.Load(), + "handoffs_initiated": m.handoffsInitiated.Load(), + "handoffs_accepted": m.handoffsAccepted.Load(), + "handoffs_rejected": m.handoffsRejected.Load(), + "handoffs_failed": m.handoffsFailed.Load(), + "elections_won": m.electionsWon.Load(), + + // Gauges + "load_score": m.currentLoadScore.Load(), + "active_sessions": m.activeSessions.Load(), + "member_count": m.memberCount.Load(), + + // System info + "uptime_seconds": time.Since(m.startTime).Seconds(), + + // Latency histograms + "latency_ms": latency, + } +} + +// percentile calculates an approximate percentile from the bucket data. +func (m *MetricsCollector) percentile(bucket *LatencyBucket, p float64) float64 { + if bucket.count == 0 { + return 0 + } + + target := int64(float64(bucket.count) * p) + cumulative := int64(0) + + // Upper bounds for each bucket in ms + upperBounds := []int64{1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000, 1<<62} + + for i, count := range bucket.buckets { + cumulative += count + if cumulative >= target { + // Return approximate percentile + return float64(upperBounds[i]) + } + } + + return 2000.0 // default max +} + +// ExportJSON exports metrics as JSON. +func (m *MetricsCollector) ExportJSON() ([]byte, error) { + return json.MarshalIndent(m.GetMetrics(), "", " ") +} + +// ExportPrometheus exports metrics in Prometheus text format. +func (m *MetricsCollector) ExportPrometheus() string { + metrics := m.GetMetrics() + var out string + + // Counters as Prometheus counters + out += "# TYPE picoclaw_messages_sent counter\n" + out += fmt.Sprintf("picoclaw_messages_sent %d\n", metrics["messages_sent"]) + + out += "\n# TYPE picoclaw_messages_received counter\n" + out += fmt.Sprintf("picoclaw_messages_received %d\n", metrics["messages_received"]) + + out += "\n# TYPE picoclaw_handoffs_initiated counter\n" + out += fmt.Sprintf("picoclaw_handoffs_initiated %d\n", metrics["handoffs_initiated"]) + + out += "\n# TYPE picoclaw_handoffs_accepted counter\n" + out += fmt.Sprintf("picoclaw_handoffs_accepted %d\n", metrics["handoffs_accepted"]) + + out += "\n# TYPE picoclaw_handoffs_rejected counter\n" + out += fmt.Sprintf("picoclaw_handoffs_rejected %d\n", metrics["handoffs_rejected"]) + + out += "\n# TYPE picoclaw_handoffs_failed counter\n" + out += fmt.Sprintf("picoclaw_handoffs_failed %d\n", metrics["handoffs_failed"]) + + out += "\n# TYPE picoclaw_elections_won counter\n" + out += fmt.Sprintf("picoclaw_elections_won %d\n", metrics["elections_won"]) + + // Gauges as Prometheus gauges + out += "\n# TYPE picoclaw_load_score gauge\n" + out += fmt.Sprintf("picoclaw_load_score %.2f\n", metrics["load_score"]) + + out += "\n# TYPE picoclaw_active_sessions gauge\n" + out += fmt.Sprintf("picoclaw_active_sessions %d\n", metrics["active_sessions"]) + + out += "\n# TYPE picoclaw_member_count gauge\n" + out += fmt.Sprintf("picoclaw_member_count %d\n", metrics["member_count"]) + + out += "\n# TYPE picoclaw_uptime_seconds gauge\n" + out += fmt.Sprintf("picoclaw_uptime_seconds %.0f\n", metrics["uptime_seconds"]) + + return out +} + +// Reset resets all metrics (useful for testing). +func (m *MetricsCollector) Reset() { + m.messagesSent.Store(0) + m.messagesReceived.Store(0) + m.handoffsInitiated.Store(0) + m.handoffsAccepted.Store(0) + m.handoffsRejected.Store(0) + m.handoffsFailed.Store(0) + m.electionsWon.Store(0) + m.currentLoadScore.Store(0) + m.activeSessions.Store(0) + m.memberCount.Store(0) + + m.mu.Lock() + m.latencyBuckets = make(map[string]*LatencyBucket) + m.mu.Unlock() + m.startTime = time.Now() +} diff --git a/pkg/swarm/node.go b/pkg/swarm/node.go new file mode 100644 index 000000000..5c12f6d7c --- /dev/null +++ b/pkg/swarm/node.go @@ -0,0 +1,329 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "time" + + "github.com/sipeed/picoclaw/pkg/logger" +) + +// NodeInfo represents a node in the swarm cluster. +type NodeInfo struct { + ID string `json:"id"` // Unique node identifier + Addr string `json:"addr"` // Listening address + Port int `json:"port"` // RPC port + AgentCaps map[string]string `json:"agent_caps"` // Agent capabilities {agent_id: capability} + LoadScore float64 `json:"load_score"` // Load score 0-1 + Labels map[string]string `json:"labels"` // Custom labels + Timestamp int64 `json:"timestamp"` // Last update time (Unix nano) + Version string `json:"version"` // PicoClaw version +} + +// IsAlive checks if the node is considered alive based on timestamp. +func (n *NodeInfo) IsAlive(timeout time.Duration) bool { + if n.Timestamp == 0 { + return false + } + age := time.Since(time.Unix(0, n.Timestamp)) + return age < timeout +} + +// String returns a JSON representation of the node. +func (n *NodeInfo) String() string { + data, _ := json.Marshal(n) + return string(data) +} + +// GetAddress returns the full address (host:port) for RPC communication. +func (n *NodeInfo) GetAddress() string { + if n.Port > 0 { + return fmt.Sprintf("%s:%d", n.Addr, n.Port) + } + return n.Addr +} + +// NodeStatus represents the current status of a node. +type NodeStatus string + +const ( + NodeStatusAlive NodeStatus = "alive" + NodeStatusSuspect NodeStatus = "suspect" + NodeStatusDead NodeStatus = "dead" + NodeStatusLeft NodeStatus = "left" +) + +// NodeState represents the state of a node in the membership view. +type NodeState struct { + Node *NodeInfo `json:"node"` + Status NodeStatus `json:"status"` + StatusSince int64 `json:"status_since"` // Unix nano when status was set + LastSeen int64 `json:"last_seen"` // Unix nano of last sighting + LastPing int64 `json:"last_ping"` // Unix nano of last successful ping + PingSuccess int `json:"ping_success"` // Consecutive successful pings + PingFailure int `json:"ping_failure"` // Consecutive failed pings +} + +// IsAvailable returns true if the node is available for handoff. +func (ns *NodeState) IsAvailable() bool { + return ns.Status == NodeStatusAlive && ns.Node.LoadScore < DefaultAvailableLoadThreshold +} + +// UpdateStatus updates the node status with timestamp. +func (ns *NodeState) UpdateStatus(status NodeStatus) { + ns.Status = status + ns.StatusSince = time.Now().UnixNano() +} + +// NodeEvent represents a node state change event. +type NodeEvent struct { + Node *NodeInfo `json:"node"` + Event EventType `json:"event"` + Time int64 `json:"time"` +} + +// EventType represents the type of node event. +type EventType string + +const ( + EventJoin EventType = "join" + EventLeave EventType = "leave" + EventUpdate EventType = "update" +) + +// EventHandler is a callback function for node events. +type EventHandler func(*NodeEvent) + +// EventHandlerID is a unique identifier for a subscribed handler. +type EventHandlerID int + +// EventDispatcher manages event handlers. +type EventDispatcher struct { + handlers []EventHandler + mu sync.RWMutex + nextID EventHandlerID + ids map[EventHandlerID]int // handler ID -> index in handlers slice +} + +// NewEventDispatcher creates a new event dispatcher. +func NewEventDispatcher() *EventDispatcher { + return &EventDispatcher{ + handlers: make([]EventHandler, 0), + ids: make(map[EventHandlerID]int), + nextID: 1, + } +} + +// Subscribe adds a new event handler and returns its ID. +func (ed *EventDispatcher) Subscribe(handler EventHandler) EventHandlerID { + ed.mu.Lock() + defer ed.mu.Unlock() + + id := ed.nextID + ed.nextID++ + + ed.handlers = append(ed.handlers, handler) + ed.ids[id] = len(ed.handlers) - 1 + return id +} + +// Unsubscribe removes an event handler by ID. +func (ed *EventDispatcher) Unsubscribe(id EventHandlerID) { + ed.mu.Lock() + defer ed.mu.Unlock() + + idx, ok := ed.ids[id] + if !ok { + return + } + + // Remove handler + ed.handlers = append(ed.handlers[:idx], ed.handlers[idx+1:]...) + + // Update indices + delete(ed.ids, id) + for handlerID, handlerIdx := range ed.ids { + if handlerIdx > idx { + ed.ids[handlerID] = handlerIdx - 1 + } + } +} + +// Dispatch sends an event to all registered handlers. +func (ed *EventDispatcher) Dispatch(event *NodeEvent) { + ed.DispatchContext(event, nil) +} + +// DispatchContext sends an event to all registered handlers with context cancellation support. +func (ed *EventDispatcher) DispatchContext(event *NodeEvent, ctx context.Context) { + ed.mu.RLock() + handlers := make([]EventHandler, len(ed.handlers)) + copy(handlers, ed.handlers) + ed.mu.RUnlock() + + for _, handler := range handlers { + // Run handlers in goroutines to avoid blocking + go func(h EventHandler) { + defer func() { + if r := recover(); r != nil { + logger.ErrorCF("swarm", "handler panic recovered", map[string]any{"panic": r}) + } + }() + + // Check if context is cancelled + if ctx != nil { + select { + case <-ctx.Done(): + logger.DebugC("swarm", "handler skipped due to context cancellation") + return + default: + } + } + + h(event) + }(handler) + } +} + +// NodeStats tracks statistics about a node. +type NodeStats struct { + MessagesSent int64 `json:"messages_sent"` + MessagesReceived int64 `json:"messages_received"` + HandoffsAccepted int `json:"handoffs_accepted"` + HandoffsInitiated int `json:"handoffs_initiated"` + LastError string `json:"last_error,omitempty"` + LastErrorTime time.Time `json:"last_error_time,omitempty"` + UptimeStart time.Time `json:"uptime_start"` +} + +// NodeWithState combines a node with its state and stats. +type NodeWithState struct { + Node *NodeInfo `json:"node"` + State *NodeState `json:"state"` + Stats *NodeStats `json:"stats,omitempty"` +} + +// IsAvailable returns true if the node is available for handoff. +func (nws *NodeWithState) IsAvailable() bool { + if nws.State == nil || nws.Node == nil { + return false + } + return nws.State.Status == NodeStatusAlive && nws.Node.LoadScore < DefaultAvailableLoadThreshold +} + +// ClusterView represents the current view of the cluster. +type ClusterView struct { + Nodes map[string]*NodeWithState `json:"nodes"` + LocalNodeID string `json:"local_node_id"` + Size int `json:"size"` + Version int64 `json:"version"` // View version for conflict detection + mu sync.RWMutex +} + +// NewClusterView creates a new cluster view. +func NewClusterView(localNodeID string) *ClusterView { + return &ClusterView{ + Nodes: make(map[string]*NodeWithState), + LocalNodeID: localNodeID, + Version: time.Now().UnixNano(), + } +} + +// AddOrUpdate adds or updates a node in the view. +func (cv *ClusterView) AddOrUpdate(node *NodeInfo) *NodeWithState { + cv.mu.Lock() + defer cv.mu.Unlock() + + cv.Version++ + + existing, ok := cv.Nodes[node.ID] + if ok { + // Update existing node + existing.Node = node + return existing + } + + // Add new node + nws := &NodeWithState{ + Node: node, + State: &NodeState{ + Node: node, + Status: NodeStatusAlive, + StatusSince: time.Now().UnixNano(), + LastSeen: time.Now().UnixNano(), + }, + Stats: &NodeStats{ + UptimeStart: time.Now(), + }, + } + cv.Nodes[node.ID] = nws + cv.Size = len(cv.Nodes) + return nws +} + +// Remove removes a node from the view. +func (cv *ClusterView) Remove(nodeID string) { + cv.mu.Lock() + defer cv.mu.Unlock() + + cv.Version++ + delete(cv.Nodes, nodeID) + cv.Size = len(cv.Nodes) +} + +// Get retrieves a node from the view. +func (cv *ClusterView) Get(nodeID string) (*NodeWithState, bool) { + cv.mu.RLock() + defer cv.mu.RUnlock() + + nws, ok := cv.Nodes[nodeID] + return nws, ok +} + +// List returns all nodes in the view. +func (cv *ClusterView) List() []*NodeWithState { + cv.mu.RLock() + defer cv.mu.RUnlock() + + result := make([]*NodeWithState, 0, len(cv.Nodes)) + for _, nws := range cv.Nodes { + result = append(result, nws) + } + return result +} + +// GetAliveNodes returns all alive nodes. +func (cv *ClusterView) GetAliveNodes() []*NodeWithState { + cv.mu.RLock() + defer cv.mu.RUnlock() + + result := make([]*NodeWithState, 0) + for _, nws := range cv.Nodes { + if nws.State.Status == NodeStatusAlive { + result = append(result, nws) + } + } + return result +} + +// GetAvailableNodes returns all available nodes (alive and not overloaded). +func (cv *ClusterView) GetAvailableNodes() []*NodeWithState { + cv.mu.RLock() + defer cv.mu.RUnlock() + + result := make([]*NodeWithState, 0) + for _, nws := range cv.Nodes { + if nws.IsAvailable() { + result = append(result, nws) + } + } + return result +} diff --git a/pkg/swarm/security.go b/pkg/swarm/security.go new file mode 100644 index 000000000..82b53461d --- /dev/null +++ b/pkg/swarm/security.go @@ -0,0 +1,127 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "fmt" + "time" +) + +// AuthProvider handles authentication for swarm nodes. +// Uses HMAC-based shared secret authentication. +type AuthProvider struct { + sharedSecret []byte + nodeID string +} + +// NewAuthProvider creates a new authentication provider. +func NewAuthProvider(nodeID, sharedSecret string) *AuthProvider { + return &AuthProvider{ + sharedSecret: []byte(sharedSecret), + nodeID: nodeID, + } +} + +// SignMessage signs a message with HMAC-SHA256. +// The signature is base64 encoded for JSON transmission. +func (a *AuthProvider) SignMessage(msg any) (string, error) { + if a.sharedSecret == nil { + return "", ErrAuthenticationFailed + } + + // Serialize message to JSON + data, err := json.Marshal(msg) + if err != nil { + return "", fmt.Errorf("failed to marshal message: %w", err) + } + + // Calculate HMAC + h := hmac.New(sha256.New, a.sharedSecret) + h.Write(data) + signature := h.Sum(nil) + + // Return base64 encoded signature + return base64.StdEncoding.EncodeToString(signature), nil +} + +// VerifySignature verifies a message signature. +func (a *AuthProvider) VerifySignature(msg any, signature string) bool { + if a.sharedSecret == nil { + return false + } + + // Calculate expected signature + expected, err := a.SignMessage(msg) + if err != nil { + return false + } + + // Compare signatures + return hmac.Equal([]byte(expected), []byte(signature)) +} + +// GetNodeID returns the node ID for this auth provider. +func (a *AuthProvider) GetNodeID() string { + return a.nodeID +} + +// AuthToken represents an authentication token. +type AuthToken struct { + NodeID string `json:"node_id"` + Signature string `json:"signature"` + Timestamp int64 `json:"timestamp"` +} + +// GenerateToken creates an auth token for the given node. +func (a *AuthProvider) GenerateToken() (*AuthToken, error) { + token := &AuthToken{ + NodeID: a.nodeID, + Timestamp: time.Now().UnixNano(), + } + + signature, err := a.SignMessage(token) + if err != nil { + return nil, err + } + + token.Signature = signature + return token, nil +} + +// VerifyToken verifies an auth token. +func (a *AuthProvider) VerifyToken(token *AuthToken) bool { + if token == nil { + return false + } + + // Check token age (reject tokens older than 1 minute) + age := time.Since(time.Unix(0, token.Timestamp)) + if age > time.Minute { + return false + } + + return a.VerifySignature(token, token.Signature) +} + +// AuthenticateNode verifies that a node is allowed to join. +func (a *AuthProvider) AuthenticateNode(nodeID, signature string, challengeData any) bool { + if a.sharedSecret == nil { + // No authentication configured - allow all + return true + } + + challenge := map[string]any{ + "node_id": nodeID, + "data": challengeData, + } + + return a.VerifySignature(challenge, signature) +} diff --git a/pkg/swarm/session_transfer.go b/pkg/swarm/session_transfer.go new file mode 100644 index 000000000..c95c7a3c9 --- /dev/null +++ b/pkg/swarm/session_transfer.go @@ -0,0 +1,331 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "context" + "encoding/json" + "fmt" + "net" + "sync" +) + +// SessionTransfer handles session migration between nodes. +type SessionTransfer struct { + config RPCConfig + localNode *NodeInfo + transfers map[string]*TransferOperation // session_key -> operation + mu sync.RWMutex + conn *net.UDPConn + onReceive func(*TransferPayload) +} + +// TransferOperation represents an ongoing transfer operation. +type TransferOperation struct { + SessionKey string + SourceNodeID string + TargetNodeID string + State TransferState + StartTime int64 + LastUpdate int64 + Payload *TransferPayload +} + +// TransferState represents the state of a transfer. +type TransferState string + +const ( + TransferStatePending TransferState = "pending" + TransferStateSending TransferState = "sending" + TransferStateReceived TransferState = "received" + TransferStateCompleted TransferState = "completed" + TransferStateFailed TransferState = "failed" +) + +// TransferPayload represents the session data being transferred. +type TransferPayload struct { + SessionKey string `json:"session_key"` + SourceNodeID string `json:"source_node_id"` + TargetNodeID string `json:"target_node_id"` + Messages []SessionMessage `json:"messages"` + Summary string `json:"summary,omitempty"` + Context map[string]any `json:"context,omitempty"` + Metadata map[string]string `json:"metadata,omitempty"` + Timestamp int64 `json:"timestamp"` + TransferID string `json:"transfer_id"` +} + +// NewSessionTransfer creates a new session transfer handler. +func NewSessionTransfer(localNode *NodeInfo, config RPCConfig) (*SessionTransfer, error) { + st := &SessionTransfer{ + config: config, + localNode: localNode, + transfers: make(map[string]*TransferOperation), + } + + // Bind UDP socket + addr := fmt.Sprintf("%s:%d", localNode.Addr, config.Port) + udpAddr, err := net.ResolveUDPAddr("udp", addr) + if err != nil { + return nil, fmt.Errorf("failed to resolve RPC address: %w", err) + } + + st.conn, err = net.ListenUDP("udp", udpAddr) + if err != nil { + return nil, fmt.Errorf("failed to listen for session transfer: %w", err) + } + + // Start message handler + go st.messageHandler() + + return st, nil +} + +// Close closes the session transfer handler. +func (st *SessionTransfer) Close() error { + if st.conn != nil { + return st.conn.Close() + } + return nil +} + +// TransferSession transfers a session to another node. +func (st *SessionTransfer) TransferSession(ctx context.Context, targetNode *NodeInfo, payload *TransferPayload) error { + if payload.TransferID == "" { + payload.TransferID = fmt.Sprintf("%s-%d", payload.SessionKey, payload.Timestamp) + } + + payload.SourceNodeID = st.localNode.ID + payload.TargetNodeID = targetNode.ID + payload.Timestamp = payload.Timestamp + + // Create transfer operation + op := &TransferOperation{ + SessionKey: payload.SessionKey, + SourceNodeID: st.localNode.ID, + TargetNodeID: targetNode.ID, + State: TransferStateSending, + StartTime: payload.Timestamp, + LastUpdate: payload.Timestamp, + Payload: payload, + } + + st.mu.Lock() + st.transfers[payload.SessionKey] = op + st.mu.Unlock() + + // Send transfer message + if err := st.sendTransfer(targetNode, payload); err != nil { + st.mu.Lock() + op.State = TransferStateFailed + delete(st.transfers, payload.SessionKey) + st.mu.Unlock() + return err + } + + return nil +} + +// sendTransfer sends a transfer message to a target node. +func (st *SessionTransfer) sendTransfer(targetNode *NodeInfo, payload *TransferPayload) error { + msg := map[string]any{ + "type": "session_transfer", + "payload": payload, + } + + data, err := json.Marshal(msg) + if err != nil { + return err + } + + addr := fmt.Sprintf("%s:%d", targetNode.Addr, st.config.Port) + udpAddr, err := net.ResolveUDPAddr("udp", addr) + if err != nil { + return err + } + + // st.conn.SetWriteDeadline(nil) // Not setting deadline + _, err = st.conn.WriteToUDP(data, udpAddr) + return err +} + +// SendAck sends an acknowledgment for a received transfer. +func (st *SessionTransfer) SendAck(targetNode *NodeInfo, transferID string, accepted bool) error { + msg := map[string]any{ + "type": "session_transfer_ack", + "payload": map[string]any{ + "transfer_id": transferID, + "accepted": accepted, + "node_id": st.localNode.ID, + }, + } + + data, err := json.Marshal(msg) + if err != nil { + return err + } + + addr := fmt.Sprintf("%s:%d", targetNode.Addr, st.config.Port) + udpAddr, err := net.ResolveUDPAddr("udp", addr) + if err != nil { + return err + } + + // st.conn.SetWriteDeadline(nil) // Not setting deadline + _, err = st.conn.WriteToUDP(data, udpAddr) + return err +} + +// messageHandler handles incoming transfer messages. +func (st *SessionTransfer) messageHandler() { + buf := make([]byte, MaxSessionMessageSize) + + for { + n, addr, err := st.conn.ReadFromUDP(buf) + if err != nil { + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + continue + } + return + } + + if n > 0 { + go st.handleMessage(buf[:n], addr) + } + } +} + +// handleMessage handles an incoming message. +func (st *SessionTransfer) handleMessage(data []byte, addr *net.UDPAddr) { + var msg map[string]any + if err := json.Unmarshal(data, &msg); err != nil { + return + } + + msgType, _ := msg["type"].(string) + + switch msgType { + case "session_transfer": + st.handleTransfer(data, addr) + case "session_transfer_ack": + st.handleTransferAck(data) + } +} + +// handleTransfer handles a session transfer message. +func (st *SessionTransfer) handleTransfer(data []byte, addr *net.UDPAddr) { + var msg map[string]any + if err := json.Unmarshal(data, &msg); err != nil { + return + } + + payloadData, _ := json.Marshal(msg["payload"]) + var payload TransferPayload + if err := json.Unmarshal(payloadData, &payload); err != nil { + return + } + + // Check if this transfer is for us + if payload.TargetNodeID != st.localNode.ID { + return + } + + // Create transfer operation + op := &TransferOperation{ + SessionKey: payload.SessionKey, + SourceNodeID: payload.SourceNodeID, + TargetNodeID: st.localNode.ID, + State: TransferStateReceived, + StartTime: payload.Timestamp, + LastUpdate: payload.Timestamp, + Payload: &payload, + } + + st.mu.Lock() + st.transfers[payload.SessionKey] = op + st.mu.Unlock() + + // Send acknowledgment + // Find source node from membership (simplified - in real implementation would look up node address) + st.SendAck(&NodeInfo{ + ID: payload.SourceNodeID, + // Need to look up actual address from membership + }, payload.TransferID, true) + + // Call receive callback if set + if st.onReceive != nil { + go st.onReceive(&payload) + } +} + +// handleTransferAck handles a transfer acknowledgment. +func (st *SessionTransfer) handleTransferAck(data []byte) { + var msg map[string]any + if err := json.Unmarshal(data, &msg); err != nil { + return + } + + payload, _ := msg["payload"].(map[string]any) + transferID, _ := payload["transfer_id"].(string) + accepted, _ := payload["accepted"].(bool) + + st.mu.Lock() + defer st.mu.Unlock() + + // Find and update transfer operation + for sessionKey, op := range st.transfers { + if op.Payload != nil && op.Payload.TransferID == transferID { + if accepted { + op.State = TransferStateCompleted + } else { + op.State = TransferStateFailed + } + op.LastUpdate = 0 // Use zero value + + // Clean up completed transfers after a delay + if op.State == TransferStateCompleted { + delete(st.transfers, sessionKey) + } + break + } + } +} + +// SetReceiveCallback sets a callback for receiving session transfers. +func (st *SessionTransfer) SetReceiveCallback(callback func(*TransferPayload)) { + st.mu.Lock() + defer st.mu.Unlock() + st.onReceive = callback +} + +// GetTransfer retrieves a transfer operation by session key. +func (st *SessionTransfer) GetTransfer(sessionKey string) (*TransferOperation, bool) { + st.mu.RLock() + defer st.mu.RUnlock() + + op, ok := st.transfers[sessionKey] + return op, ok +} + +// ListTransfers returns all active transfers. +func (st *SessionTransfer) ListTransfers() []*TransferOperation { + st.mu.RLock() + defer st.mu.RUnlock() + + result := make([]*TransferOperation, 0, len(st.transfers)) + for _, op := range st.transfers { + result = append(result, op) + } + return result +} + +// RemoveTransfer removes a transfer operation. +func (st *SessionTransfer) RemoveTransfer(sessionKey string) { + st.mu.Lock() + defer st.mu.Unlock() + delete(st.transfers, sessionKey) +} diff --git a/pkg/swarm/swarm_test.go b/pkg/swarm/swarm_test.go new file mode 100644 index 000000000..3385d4554 --- /dev/null +++ b/pkg/swarm/swarm_test.go @@ -0,0 +1,424 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNodeInfo(t *testing.T) { + node := &NodeInfo{ + ID: "test-node-1", + Addr: "192.168.1.100", + Port: 7947, + LoadScore: 0.5, + AgentCaps: map[string]string{ + "agent-1": "general", + }, + Labels: map[string]string{ + "region": "us-west", + }, + Timestamp: time.Now().UnixNano(), + } + + t.Run("IsAlive", func(t *testing.T) { + assert.True(t, node.IsAlive(time.Minute)) + assert.False(t, node.IsAlive(time.Nanosecond)) + }) + + t.Run("GetAddress", func(t *testing.T) { + addr := node.GetAddress() + assert.Equal(t, "192.168.1.100:7947", addr) + }) +} + +func TestClusterView(t *testing.T) { + view := NewClusterView("local-node") + + t.Run("AddOrUpdate", func(t *testing.T) { + node := &NodeInfo{ + ID: "node-1", + Addr: "192.168.1.1", + Port: 7947, + LoadScore: 0.3, + Timestamp: time.Now().UnixNano(), + } + + nws := view.AddOrUpdate(node) + require.NotNil(t, nws) + assert.Equal(t, node.ID, nws.Node.ID) + assert.Equal(t, 1, view.Size) + }) + + t.Run("Get", func(t *testing.T) { + node, ok := view.Get("node-1") + assert.True(t, ok) + assert.Equal(t, "node-1", node.Node.ID) + + _, ok = view.Get("non-existent") + assert.False(t, ok) + }) + + t.Run("GetAliveNodes", func(t *testing.T) { + nodes := view.GetAliveNodes() + assert.Equal(t, 1, len(nodes)) + }) + + t.Run("GetAvailableNodes", func(t *testing.T) { + nodes := view.GetAvailableNodes() + assert.Equal(t, 1, len(nodes)) // 0.3 < 0.9 + }) + + t.Run("Remove", func(t *testing.T) { + view.Remove("node-1") + assert.Equal(t, 0, view.Size) + }) +} + +func TestLoadMonitor(t *testing.T) { + config := &LoadMonitorConfig{ + Enabled: true, + Interval: Duration{time.Second}, + SampleSize: 10, + CPUWeight: 0.3, + MemoryWeight: 0.3, + SessionWeight: 0.4, + } + + monitor := NewLoadMonitor(config) + + t.Run("GetCurrentLoad", func(t *testing.T) { + metrics := monitor.GetCurrentLoad() + assert.NotNil(t, metrics) + assert.GreaterOrEqual(t, metrics.Score, 0.0) + assert.LessOrEqual(t, metrics.Score, 1.0) + assert.GreaterOrEqual(t, metrics.ActiveSessions, 0) + }) + + t.Run("SessionCount", func(t *testing.T) { + monitor.SetSessionCount(5) + assert.Equal(t, 5, monitor.GetSessionCount()) + + monitor.IncrementSessions() + assert.Equal(t, 6, monitor.GetSessionCount()) + + monitor.DecrementSessions() + assert.Equal(t, 5, monitor.GetSessionCount()) + }) + + t.Run("GetAverageScore", func(t *testing.T) { + avg := monitor.GetAverageScore() + assert.GreaterOrEqual(t, avg, 0.0) + assert.LessOrEqual(t, avg, 1.0) + }) +} + +func TestEventDispatcher(t *testing.T) { + ed := NewEventDispatcher() + + t.Run("SubscribeDispatch", func(t *testing.T) { + received := make(chan *NodeEvent, 1) + + id := ed.Subscribe(func(event *NodeEvent) { + received <- event + }) + + event := &NodeEvent{ + Node: &NodeInfo{ID: "test-node"}, + Event: EventJoin, + Time: time.Now().UnixNano(), + } + + ed.Dispatch(event) + + select { + case <-received: + // Event received + case <-time.After(time.Second): + t.Fatal("Event not received") + } + + ed.Unsubscribe(id) + }) + + t.Run("Unsubscribe", func(t *testing.T) { + received := make(chan *NodeEvent, 1) + + id := ed.Subscribe(func(event *NodeEvent) { + received <- event + }) + + ed.Unsubscribe(id) + + event := &NodeEvent{ + Node: &NodeInfo{ID: "test-node"}, + Event: EventJoin, + Time: time.Now().UnixNano(), + } + + ed.Dispatch(event) + + select { + case <-received: + t.Fatal("Should not receive event after unsubscribe") + case <-time.After(100 * time.Millisecond): + // Expected - no event received + } + }) +} + +func TestNodeWithState(t *testing.T) { + node := &NodeInfo{ + ID: "test-node", + Addr: "192.168.1.1", + Port: 7947, + LoadScore: 0.5, + Timestamp: time.Now().UnixNano(), + } + + nws := &NodeWithState{ + Node: node, + State: &NodeState{ + Status: NodeStatusAlive, + StatusSince: time.Now().UnixNano(), + LastSeen: time.Now().UnixNano(), + }, + } + + t.Run("IsAvailable", func(t *testing.T) { + assert.True(t, nws.IsAvailable()) + + // High load + nws.Node.LoadScore = 0.95 + assert.False(t, nws.IsAvailable()) + + // Not alive + nws.Node.LoadScore = 0.5 + nws.State.Status = NodeStatusDead + assert.False(t, nws.IsAvailable()) + }) +} + +func TestDuration(t *testing.T) { + t.Run("UnmarshalJSON from string", func(t *testing.T) { + d := Duration{} + err := d.UnmarshalJSON([]byte(`"5s"`)) + require.NoError(t, err) + assert.Equal(t, 5*time.Second, d.Duration) + }) + + t.Run("MarshalJSON", func(t *testing.T) { + d := Duration{5 * time.Second} + data, err := d.MarshalJSON() + require.NoError(t, err) + assert.Equal(t, []byte(`"5s"`), data) + }) +} + +// Integration tests for discovery and gossip protocol + +func TestDiscoveryServiceNodeDiscovery(t *testing.T) { + t.Run("TwoNodesDiscoverEachOther", func(t *testing.T) { + // Create first node + cfg1 := &Config{ + NodeID: "node-1", + BindAddr: "127.0.0.1", + BindPort: 17946, + RPC: RPCConfig{ + Port: 17947, + }, + Discovery: DiscoveryConfig{ + GossipInterval: Duration{100 * time.Millisecond}, + NodeTimeout: Duration{500 * time.Millisecond}, + DeadNodeTimeout: Duration{2 * time.Second}, + }, + } + + ds1, err := NewDiscoveryService(cfg1) + require.NoError(t, err) + defer ds1.Stop() + + err = ds1.Start() + require.NoError(t, err) + + // Create second node + cfg2 := &Config{ + NodeID: "node-2", + BindAddr: "127.0.0.1", + BindPort: 17948, + RPC: RPCConfig{ + Port: 17949, + }, + Discovery: DiscoveryConfig{ + JoinAddrs: []string{"127.0.0.1:17946"}, + GossipInterval: Duration{100 * time.Millisecond}, + NodeTimeout: Duration{500 * time.Millisecond}, + DeadNodeTimeout: Duration{2 * time.Second}, + }, + } + + ds2, err := NewDiscoveryService(cfg2) + require.NoError(t, err) + defer ds2.Stop() + + err = ds2.Start() + require.NoError(t, err) + + // Wait for discovery + time.Sleep(500 * time.Millisecond) + + // Check that node-2 knows about node-1 + members2 := ds2.Members() + assert.GreaterOrEqual(t, len(members2), 1, "node-2 should discover node-1") + + // Check that node-1 knows about node-2 + members1 := ds1.Members() + assert.GreaterOrEqual(t, len(members1), 1, "node-1 should discover node-2") + }) + + t.Run("NodeHealthCheck", func(t *testing.T) { + cfg := &Config{ + NodeID: "health-node", + BindAddr: "127.0.0.1", + BindPort: 17950, + RPC: RPCConfig{ + Port: 17951, + }, + Discovery: DiscoveryConfig{ + GossipInterval: Duration{100 * time.Millisecond}, + NodeTimeout: Duration{300 * time.Millisecond}, + DeadNodeTimeout: Duration{1 * time.Second}, + }, + } + + ds, err := NewDiscoveryService(cfg) + require.NoError(t, err) + defer ds.Stop() + + err = ds.Start() + require.NoError(t, err) + + // Add a remote node manually + remoteNode := &NodeInfo{ + ID: "remote-node", + Addr: "192.168.1.100", + Port: 7947, + LoadScore: 0.5, + Timestamp: time.Now().UnixNano(), + } + + ds.membership.UpdateNode(remoteNode) + + // Check health check - should have at least the remote node + members := ds.Members() + assert.GreaterOrEqual(t, len(members), 1) + + // Find the remote node + var found *NodeWithState + for _, m := range members { + if m.Node.ID == "remote-node" { + found = m + break + } + } + require.NotNil(t, found) + assert.Equal(t, NodeStatusAlive, found.State.Status) + }) +} + +func TestHandoffCoordinator(t *testing.T) { + t.Run("CanHandleWithLoadThreshold", func(t *testing.T) { + cfg := &Config{ + NodeID: "handoff-node", + Handoff: HandoffConfig{ + Enabled: true, + LoadThreshold: 0.8, + }, + } + + ds, err := NewDiscoveryService(cfg) + require.NoError(t, err) + + hc, err := NewHandoffCoordinator(ds, cfg.Handoff) + require.NoError(t, err) + defer hc.Close() + + // With low load, should be able to handle + ds.localNode.LoadScore = 0.5 + assert.True(t, hc.CanHandle("")) + + // With high load, should not be able to handle + ds.localNode.LoadScore = 0.9 + assert.False(t, hc.CanHandle("")) + }) + + t.Run("FindTargetNode", func(t *testing.T) { + cfg := &Config{ + NodeID: "coordinator-node", + } + + ds, err := NewDiscoveryService(cfg) + require.NoError(t, err) + + hc, err := NewHandoffCoordinator(ds, cfg.Handoff) + require.NoError(t, err) + defer hc.Close() + + // Add some candidate nodes + node1 := &NodeInfo{ + ID: "target-1", + Addr: "192.168.1.1", + Port: 7947, + LoadScore: 0.3, + AgentCaps: map[string]string{"model": "gpt-4"}, + Timestamp: time.Now().UnixNano(), + } + ds.membership.UpdateNode(node1) + + node2 := &NodeInfo{ + ID: "target-2", + Addr: "192.168.1.2", + Port: 7947, + LoadScore: 0.7, + AgentCaps: map[string]string{"model": "gpt-4"}, + Timestamp: time.Now().UnixNano(), + } + ds.membership.UpdateNode(node2) + + // Should select the least loaded node + target, err := hc.findTargetNode(&HandoffRequest{}) + require.NoError(t, err) + assert.Equal(t, "target-1", target.Node.ID) + }) +} + +func TestSessionTransfer(t *testing.T) { + t.Run("TransferSession", func(t *testing.T) { + node := &NodeInfo{ + ID: "transfer-node", + Addr: "127.0.0.1", + Port: 17952, + Timestamp: time.Now().UnixNano(), + } + + cfg := RPCConfig{ + Port: 17952, + } + + st, err := NewSessionTransfer(node, cfg) + require.NoError(t, err) + defer st.Close() + + // Test list operations + transfers := st.ListTransfers() + assert.Equal(t, 0, len(transfers)) + }) +} diff --git a/pkg/swarm/types.go b/pkg/swarm/types.go new file mode 100644 index 000000000..da3fabb78 --- /dev/null +++ b/pkg/swarm/types.go @@ -0,0 +1,25 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +// SessionMessage represents a message in a session. +// This is shared across handoff and session transfer. +type SessionMessage struct { + Role string `json:"role"` + Content string `json:"content"` + Timestamp int64 `json:"timestamp,omitempty"` + ToolCalls []ToolCallData `json:"tool_calls,omitempty"` +} + +// ToolCallData represents tool call information in a message. +type ToolCallData struct { + ID string `json:"id"` + Name string `json:"name"` + Arguments map[string]any `json:"arguments"` + Result string `json:"result,omitempty"` + Extra map[string]any `json:"extra,omitempty"` +} diff --git a/pkg/tools/handoff_tool.go b/pkg/tools/handoff_tool.go new file mode 100644 index 000000000..343ce4eb9 --- /dev/null +++ b/pkg/tools/handoff_tool.go @@ -0,0 +1,140 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package tools + +import ( + "context" + "fmt" + + "github.com/sipeed/picoclaw/pkg/swarm" +) + +// HandoffTool implements the handoff tool for swarm mode. +type HandoffTool struct { + coordinator *swarm.HandoffCoordinator + channel string + chatID string +} + +// NewHandoffTool creates a new handoff tool. +func NewHandoffTool(coordinator *swarm.HandoffCoordinator) *HandoffTool { + return &HandoffTool{ + coordinator: coordinator, + channel: "cli", + chatID: "direct", + } +} + +// Name returns the tool name. +func (t *HandoffTool) Name() string { + return "handoff" +} + +// Description returns the tool description. +func (t *HandoffTool) Description() string { + return "Delegate this task to another agent in the swarm. Use when you cannot handle the task due to capability constraints or system overload." +} + +// Parameters returns the tool parameters schema. +func (t *HandoffTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "reason": map[string]any{ + "type": "string", + "enum": []string{"no_capability", "overloaded", "user_request"}, + "description": "The reason for handing off this task", + }, + "required_capability": map[string]any{ + "type": "string", + "description": "The specific capability required to handle this task", + }, + "context": map[string]any{ + "type": "string", + "description": "Additional context about why this handoff is needed", + }, + }, + } +} + +// SetContext sets the channel and chat ID for the tool. +func (t *HandoffTool) SetContext(channel, chatID string) { + t.channel = channel + t.chatID = chatID +} + +// Execute executes the handoff tool. +func (t *HandoffTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + if t.coordinator == nil { + return ErrorResult("Swarm mode is not enabled or handoff coordinator not configured").WithError( + fmt.Errorf("handoff coordinator is nil")) + } + + // Parse reason + reasonStr, _ := args["reason"].(string) + var reason swarm.HandoffReason + switch reasonStr { + case "no_capability": + reason = swarm.ReasonNoCapability + case "overloaded": + reason = swarm.ReasonOverloaded + case "user_request": + reason = swarm.ReasonUserRequest + default: + reason = swarm.ReasonNoCapability + } + + // Parse required capability + requiredCap, _ := args["required_capability"].(string) + + // Parse context + contextMsg, _ := args["context"].(string) + + // Build handoff request + req := &swarm.HandoffRequest{ + Reason: reason, + RequiredCap: requiredCap, + Metadata: make(map[string]string), + } + + if contextMsg != "" { + req.Metadata["context"] = contextMsg + } + + // Execute handoff + resp, err := t.coordinator.InitiateHandoff(ctx, req) + if err != nil { + return ErrorResult(fmt.Sprintf("Handoff failed: %v", err)).WithError(err) + } + + if !resp.Accepted { + return ErrorResult(fmt.Sprintf("Handoff rejected by all nodes: %s", resp.Reason)).WithError( + fmt.Errorf("handoff rejected: %s", resp.Reason)) + } + + // Build result message + resultMsg := fmt.Sprintf("Task handed off to node %s\n", resp.NodeID) + if resp.Reason != "" { + resultMsg += fmt.Sprintf("Note: %s\n", resp.Reason) + } + + return &ToolResult{ + ForLLM: resultMsg + "The target node will process this task and respond to the user.", + ForUser: "Your task has been delegated to another agent in the swarm. They will respond shortly.", + Silent: false, + IsError: false, + Async: true, // Handoff is async - target node will respond directly + } +} + +// CanHandle reports whether the local node can handle the given capability. +func (t *HandoffTool) CanHandle(requiredCap string) bool { + if t.coordinator == nil { + return true // If swarm is disabled, we can "handle" everything + } + return t.coordinator.CanHandle(requiredCap) +}