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.
This commit is contained in:
parent
b9a66248d8
commit
bc3d77d2dd
18 changed files with 4859 additions and 1 deletions
450
docs/swarm-architecture.md
Normal file
450
docs/swarm-architecture.md
Normal file
|
|
@ -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:
|
||||||
|
|
||||||
|
```
|
||||||
|
┌──────────────────────────────────────────────────────────────<EFBFBD><EFBFBD><EFBFBD>──┐
|
||||||
|
│ 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
|
||||||
|
|
@ -25,6 +25,7 @@ import (
|
||||||
"github.com/sipeed/picoclaw/pkg/routing"
|
"github.com/sipeed/picoclaw/pkg/routing"
|
||||||
"github.com/sipeed/picoclaw/pkg/skills"
|
"github.com/sipeed/picoclaw/pkg/skills"
|
||||||
"github.com/sipeed/picoclaw/pkg/state"
|
"github.com/sipeed/picoclaw/pkg/state"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/swarm"
|
||||||
"github.com/sipeed/picoclaw/pkg/tools"
|
"github.com/sipeed/picoclaw/pkg/tools"
|
||||||
"github.com/sipeed/picoclaw/pkg/utils"
|
"github.com/sipeed/picoclaw/pkg/utils"
|
||||||
)
|
)
|
||||||
|
|
@ -38,6 +39,12 @@ type AgentLoop struct {
|
||||||
summarizing sync.Map
|
summarizing sync.Map
|
||||||
fallback *providers.FallbackChain
|
fallback *providers.FallbackChain
|
||||||
channelManager *channels.Manager
|
channelManager *channels.Manager
|
||||||
|
|
||||||
|
// Swarm mode support
|
||||||
|
swarmEnabled bool
|
||||||
|
swarmDiscovery *swarm.DiscoveryService
|
||||||
|
swarmHandoff *swarm.HandoffCoordinator
|
||||||
|
swarmLoad *swarm.LoadMonitor
|
||||||
}
|
}
|
||||||
|
|
||||||
// processOptions configures how a message is processed
|
// 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)
|
stateManager = state.NewManager(defaultAgent.Workspace)
|
||||||
}
|
}
|
||||||
|
|
||||||
return &AgentLoop{
|
al := &AgentLoop{
|
||||||
bus: msgBus,
|
bus: msgBus,
|
||||||
cfg: cfg,
|
cfg: cfg,
|
||||||
registry: registry,
|
registry: registry,
|
||||||
|
|
@ -77,6 +84,13 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
|
||||||
summarizing: sync.Map{},
|
summarizing: sync.Map{},
|
||||||
fallback: fallbackChain,
|
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).
|
// 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,
|
"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{
|
return al.runAgentLoop(ctx, agent, processOptions{
|
||||||
SessionKey: sessionKey,
|
SessionKey: sessionKey,
|
||||||
Channel: msg.Channel,
|
Channel: msg.Channel,
|
||||||
|
|
@ -388,6 +418,10 @@ func (al *AgentLoop) processSystemMessage(ctx context.Context, msg bus.InboundMe
|
||||||
|
|
||||||
// runAgentLoop is the core message processing logic.
|
// runAgentLoop is the core message processing logic.
|
||||||
func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opts processOptions) (string, error) {
|
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)
|
// 0. Record last channel for heartbeat notifications (skip internal channels)
|
||||||
if opts.Channel != "" && opts.ChatID != "" {
|
if opts.Channel != "" && opts.ChatID != "" {
|
||||||
// Don't record internal channels (cli, system, subagent)
|
// 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}
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -57,6 +57,7 @@ type Config struct {
|
||||||
Tools ToolsConfig `json:"tools"`
|
Tools ToolsConfig `json:"tools"`
|
||||||
Heartbeat HeartbeatConfig `json:"heartbeat"`
|
Heartbeat HeartbeatConfig `json:"heartbeat"`
|
||||||
Devices DevicesConfig `json:"devices"`
|
Devices DevicesConfig `json:"devices"`
|
||||||
|
Swarm SwarmConfig `json:"swarm,omitempty"` // Swarm mode configuration
|
||||||
}
|
}
|
||||||
|
|
||||||
// MarshalJSON implements custom JSON marshaling for Config
|
// MarshalJSON implements custom JSON marshaling for Config
|
||||||
|
|
@ -650,3 +651,101 @@ func (c *Config) ValidateModelList() error {
|
||||||
}
|
}
|
||||||
return nil
|
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"`
|
||||||
|
}
|
||||||
|
|
||||||
|
|
|
||||||
276
pkg/swarm/config.go
Normal file
276
pkg/swarm/config.go
Normal file
|
|
@ -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
|
||||||
|
}
|
||||||
101
pkg/swarm/constants.go
Normal file
101
pkg/swarm/constants.go
Normal file
|
|
@ -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
|
||||||
|
)
|
||||||
516
pkg/swarm/discovery.go
Normal file
516
pkg/swarm/discovery.go
Normal file
|
|
@ -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 ""
|
||||||
|
}
|
||||||
50
pkg/swarm/errors.go
Normal file
50
pkg/swarm/errors.go
Normal file
|
|
@ -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")
|
||||||
|
)
|
||||||
486
pkg/swarm/handoff.go
Normal file
486
pkg/swarm/handoff.go
Normal file
|
|
@ -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
|
||||||
|
}
|
||||||
237
pkg/swarm/leader_election.go
Normal file
237
pkg/swarm/leader_election.go
Normal file
|
|
@ -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()),
|
||||||
|
}
|
||||||
|
}
|
||||||
295
pkg/swarm/load_monitor.go
Normal file
295
pkg/swarm/load_monitor.go
Normal file
|
|
@ -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
|
||||||
|
}
|
||||||
365
pkg/swarm/membership.go
Normal file
365
pkg/swarm/membership.go
Normal file
|
|
@ -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()
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
286
pkg/swarm/metrics.go
Normal file
286
pkg/swarm/metrics.go
Normal file
|
|
@ -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()
|
||||||
|
}
|
||||||
329
pkg/swarm/node.go
Normal file
329
pkg/swarm/node.go
Normal file
|
|
@ -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
|
||||||
|
}
|
||||||
127
pkg/swarm/security.go
Normal file
127
pkg/swarm/security.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
331
pkg/swarm/session_transfer.go
Normal file
331
pkg/swarm/session_transfer.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
424
pkg/swarm/swarm_test.go
Normal file
424
pkg/swarm/swarm_test.go
Normal file
|
|
@ -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))
|
||||||
|
})
|
||||||
|
}
|
||||||
25
pkg/swarm/types.go
Normal file
25
pkg/swarm/types.go
Normal file
|
|
@ -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"`
|
||||||
|
}
|
||||||
140
pkg/tools/handoff_tool.go
Normal file
140
pkg/tools/handoff_tool.go
Normal file
|
|
@ -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)
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue