feat: implement swarm mode for multi-agent coordination
This commit is contained in:
parent
6f5930624b
commit
74c192bcdd
53 changed files with 10922 additions and 44 deletions
4
.gitignore
vendored
4
.gitignore
vendored
|
|
@ -47,3 +47,7 @@ dist/
|
||||||
|
|
||||||
# Windows Application Icon/Resource
|
# Windows Application Icon/Resource
|
||||||
*.syso
|
*.syso
|
||||||
|
|
||||||
|
# PM2 and process manager configs (dev-specific)
|
||||||
|
ecosystem.config.js
|
||||||
|
.pm2/
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,7 @@ import (
|
||||||
|
|
||||||
func NewGatewayCommand() *cobra.Command {
|
func NewGatewayCommand() *cobra.Command {
|
||||||
var debug bool
|
var debug bool
|
||||||
|
var configPath string
|
||||||
|
|
||||||
cmd := &cobra.Command{
|
cmd := &cobra.Command{
|
||||||
Use: "gateway",
|
Use: "gateway",
|
||||||
|
|
@ -13,11 +14,12 @@ func NewGatewayCommand() *cobra.Command {
|
||||||
Short: "Start picoclaw gateway",
|
Short: "Start picoclaw gateway",
|
||||||
Args: cobra.NoArgs,
|
Args: cobra.NoArgs,
|
||||||
RunE: func(_ *cobra.Command, _ []string) error {
|
RunE: func(_ *cobra.Command, _ []string) error {
|
||||||
return gatewayCmd(debug)
|
return gatewayCmd(debug, configPath)
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
cmd.Flags().BoolVarP(&debug, "debug", "d", false, "Enable debug logging")
|
cmd.Flags().BoolVarP(&debug, "debug", "d", false, "Enable debug logging")
|
||||||
|
cmd.Flags().StringVarP(&configPath, "config", "c", "", "Path to config file")
|
||||||
|
|
||||||
return cmd
|
return cmd
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -39,13 +39,13 @@ import (
|
||||||
"github.com/sipeed/picoclaw/pkg/voice"
|
"github.com/sipeed/picoclaw/pkg/voice"
|
||||||
)
|
)
|
||||||
|
|
||||||
func gatewayCmd(debug bool) error {
|
func gatewayCmd(debug bool, configPath string) error {
|
||||||
if debug {
|
if debug {
|
||||||
logger.SetLevel(logger.DEBUG)
|
logger.SetLevel(logger.DEBUG)
|
||||||
fmt.Println("🔍 Debug mode enabled")
|
fmt.Println("🔍 Debug mode enabled")
|
||||||
}
|
}
|
||||||
|
|
||||||
cfg, err := internal.LoadConfig()
|
cfg, err := loadConfigWithPath(configPath)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("error loading config: %w", err)
|
return fmt.Errorf("error loading config: %w", err)
|
||||||
}
|
}
|
||||||
|
|
@ -252,3 +252,12 @@ func setupCronTool(
|
||||||
|
|
||||||
return cronService
|
return cronService
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func loadConfigWithPath(configPath string) (*config.Config, error) {
|
||||||
|
path := configPath
|
||||||
|
if path == "" {
|
||||||
|
path = internal.GetConfigPath()
|
||||||
|
}
|
||||||
|
fmt.Printf("📁 Loading config from: %s\n", path)
|
||||||
|
return config.LoadConfig(path)
|
||||||
|
}
|
||||||
|
|
|
||||||
572
docs/swarm-architecture.md
Normal file
572
docs/swarm-architecture.md
Normal file
|
|
@ -0,0 +1,572 @@
|
||||||
|
# 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 NATS JetStream KV with TTL-based liveness
|
||||||
|
- **Health Monitoring**: Heartbeat renewal with configurable TTL and failure detection
|
||||||
|
- **Load Balancing**: Intelligent task distribution based on node load scoring
|
||||||
|
- **Handoff Mechanism**: Dynamic task delegation via NATS request-reply
|
||||||
|
- **Leader Election**: Distributed leader election via KV CAS (Compare-And-Swap) lock
|
||||||
|
|
||||||
|
**Transport**: All swarm communication uses NATS as the sole transport layer. There is no UDP gossip, no custom RPC, and no unencrypted channels.
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
```
|
||||||
|
┌─────────────────────────────────────────────────────────────────┐
|
||||||
|
│ PicoClaw Swarm │
|
||||||
|
│ (NATS as sole transport) │
|
||||||
|
├─────────────────────────────────────────────────────────────────┤
|
||||||
|
│ Control Plane │ Data Plane │
|
||||||
|
│ ├─ Node Discovery (KV + TTL) │ ├─ Task Execution │
|
||||||
|
│ ├─ Membership Management │ ├─ Session Transfer │
|
||||||
|
│ ├─ Health Monitoring │ │ (NATS request-reply) │
|
||||||
|
│ ├─ Load Monitoring │ └─ Message Routing │
|
||||||
|
│ └─ Leader Election (KV CAS) │ (targeted subjects) │
|
||||||
|
└─────────────────────────────────────────────────────────────────┘
|
||||||
|
```
|
||||||
|
|
||||||
|
## NATS Subject Namespace
|
||||||
|
|
||||||
|
All swarm messages are published under a common subject prefix. NATS ACL rules should restrict `picoclaw.swarm.>` to authenticated swarm node identities only.
|
||||||
|
|
||||||
|
| Subject Pattern | Purpose |
|
||||||
|
|-----------------|---------|
|
||||||
|
| `picoclaw.swarm.heartbeat.<nodeID>` | Heartbeat / KV renewal |
|
||||||
|
| `picoclaw.swarm.node.<nodeID>.msg` | Direct node messaging (request-reply) |
|
||||||
|
| `picoclaw.swarm.handoff.<nodeID>` | Targeted handoff requests (request-reply) |
|
||||||
|
| `picoclaw.swarm.leader` | Leader election announcements |
|
||||||
|
| `picoclaw.swarm.metrics.<nodeID>` | Metrics publishing |
|
||||||
|
|
||||||
|
## JetStream KV Buckets
|
||||||
|
|
||||||
|
| Bucket | Key Pattern | TTL | Purpose |
|
||||||
|
|--------|-------------|-----|---------|
|
||||||
|
| `swarm_members` | `node.<nodeID>` | 15s (configurable) | Membership liveness via TTL expiry |
|
||||||
|
| `swarm_status` | `status.<nodeID>` | 30s | Detailed node status cache (CPU, mem, tasks) |
|
||||||
|
| `swarm_leader` | `leader` | 10s | Leader election CAS lock |
|
||||||
|
|
||||||
|
## Control Plane
|
||||||
|
|
||||||
|
### 1. Node Discovery
|
||||||
|
|
||||||
|
Nodes discover each other by writing entries to the `swarm_members` KV bucket and watching for changes. Liveness is determined by TTL — if a node stops renewing its entry, the key expires and the node is considered dead.
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
sequenceDiagram
|
||||||
|
participant Node1
|
||||||
|
participant NATS as NATS JetStream KV
|
||||||
|
participant Node2
|
||||||
|
|
||||||
|
Note over Node1: New node starts
|
||||||
|
Node1->>NATS: KV Put (node.node-1, NodeInfo, TTL=15s)
|
||||||
|
Node1->>NATS: KV Watch (swarm_members.>)
|
||||||
|
Node2->>NATS: KV Put (node.node-2, NodeInfo, TTL=15s)
|
||||||
|
NATS-->>Node1: Watch event: node.node-2 created
|
||||||
|
NATS-->>Node2: Watch event: node.node-1 exists
|
||||||
|
Note over Node1,Node2: Cluster formed
|
||||||
|
|
||||||
|
loop Every heartbeat_interval (5s)
|
||||||
|
Node1->>NATS: KV Put (renew TTL)
|
||||||
|
Node2->>NATS: KV Put (renew TTL)
|
||||||
|
end
|
||||||
|
|
||||||
|
Note over Node1: Node1 crashes
|
||||||
|
NATS-->>Node2: Watch event: node.node-1 expired (TTL)
|
||||||
|
Note over Node2: Mark Node1 as dead
|
||||||
|
```
|
||||||
|
|
||||||
|
**Key Parameters:**
|
||||||
|
|
||||||
|
| Parameter | Default | Description |
|
||||||
|
|-----------|---------|-------------|
|
||||||
|
| `heartbeat_interval` | 5s | Frequency of KV entry renewal |
|
||||||
|
| `member_ttl` | 15s | TTL for member entries (must be > heartbeat_interval) |
|
||||||
|
| `node_timeout` | 5s | Time before marking node as suspect |
|
||||||
|
| `dead_node_timeout` | 30s | Time before removing dead node from view |
|
||||||
|
|
||||||
|
### 2. Membership Management
|
||||||
|
|
||||||
|
Each node maintains a local `ClusterView` populated by KV watch events:
|
||||||
|
|
||||||
|
```go
|
||||||
|
type ClusterView struct {
|
||||||
|
sync.RWMutex
|
||||||
|
localNode *NodeInfo
|
||||||
|
members map[string]*NodeWithState // node_id -> NodeWithState
|
||||||
|
Size int
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
**Node State Machine:**
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
stateDiagram-v2
|
||||||
|
[*] --> Alive: KV entry created
|
||||||
|
Alive --> Suspect: Heartbeat not renewed within node_timeout
|
||||||
|
Suspect --> Alive: KV entry renewed
|
||||||
|
Suspect --> Dead: KV entry expired (TTL)
|
||||||
|
Dead --> [*]: Removed from view
|
||||||
|
Alive --> Left: Node gracefully stops (deletes KV entry)
|
||||||
|
Left --> [*]
|
||||||
|
```
|
||||||
|
|
||||||
|
**Node Information:**
|
||||||
|
|
||||||
|
```go
|
||||||
|
type NodeInfo struct {
|
||||||
|
ID string // Unique node identifier
|
||||||
|
Addr string // IP address
|
||||||
|
Port int // Service 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 (UnixNano)
|
||||||
|
Version string // Protocol version
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. Health Monitoring
|
||||||
|
|
||||||
|
Health is driven by KV TTL — no separate probe/ping mechanism is needed:
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
sequenceDiagram
|
||||||
|
participant Node as Node
|
||||||
|
participant KV as NATS KV (swarm_members)
|
||||||
|
participant Watcher as Other Nodes (KV Watch)
|
||||||
|
|
||||||
|
loop Every heartbeat_interval
|
||||||
|
Node->>KV: Put (node.<id>, NodeInfo, TTL=15s)
|
||||||
|
KV-->>Watcher: Watch event: key updated
|
||||||
|
Watcher->>Watcher: Update LastSeen, mark Alive
|
||||||
|
end
|
||||||
|
|
||||||
|
Note over Node: Node stops renewing
|
||||||
|
KV->>KV: TTL expires, key deleted
|
||||||
|
KV-->>Watcher: Watch event: key deleted
|
||||||
|
Watcher->>Watcher: Mark as Dead, dispatch NodeLeft event
|
||||||
|
```
|
||||||
|
|
||||||
|
### 4. Load Monitoring
|
||||||
|
|
||||||
|
Each node continuously monitors its resource usage (transport-agnostic):
|
||||||
|
|
||||||
|
```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
|
||||||
|
```
|
||||||
|
|
||||||
|
### 5. Leader Election
|
||||||
|
|
||||||
|
Leader election uses a CAS (Compare-And-Swap) lock on the `swarm_leader` KV bucket:
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
sequenceDiagram
|
||||||
|
participant NodeA
|
||||||
|
participant KV as NATS KV (swarm_leader)
|
||||||
|
participant NodeB
|
||||||
|
|
||||||
|
NodeA->>KV: Create("leader", {nodeID: A}, TTL=10s)
|
||||||
|
KV-->>NodeA: OK (revision=1)
|
||||||
|
Note over NodeA: Becomes leader
|
||||||
|
|
||||||
|
NodeB->>KV: Create("leader", {nodeID: B}, TTL=10s)
|
||||||
|
KV-->>NodeB: Key exists (conflict)
|
||||||
|
NodeB->>KV: Watch("leader")
|
||||||
|
Note over NodeB: Becomes follower, watches for changes
|
||||||
|
|
||||||
|
loop Every renewal_interval (3s)
|
||||||
|
NodeA->>KV: Update("leader", {nodeID: A}, revision=N)
|
||||||
|
KV-->>NodeA: OK (revision=N+1)
|
||||||
|
end
|
||||||
|
|
||||||
|
Note over NodeA: NodeA crashes, stops renewing
|
||||||
|
KV->>KV: TTL expires, key deleted
|
||||||
|
KV-->>NodeB: Watch event: "leader" deleted
|
||||||
|
NodeB->>KV: Create("leader", {nodeID: B}, TTL=10s)
|
||||||
|
KV-->>NodeB: OK
|
||||||
|
Note over NodeB: Becomes new leader
|
||||||
|
```
|
||||||
|
|
||||||
|
## Data Plane
|
||||||
|
|
||||||
|
### 1. Request Routing
|
||||||
|
|
||||||
|
Inter-node messages use NATS request-reply on targeted subjects:
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
sequenceDiagram
|
||||||
|
participant User
|
||||||
|
participant N1 as Node 1
|
||||||
|
participant NATS
|
||||||
|
participant N2 as Node 2
|
||||||
|
|
||||||
|
User->>N1: Message
|
||||||
|
N1->>N1: Check local load
|
||||||
|
|
||||||
|
alt Load OK
|
||||||
|
N1->>N1: Process locally
|
||||||
|
N1->>User: Response
|
||||||
|
else Overloaded
|
||||||
|
N1->>NATS: Request on picoclaw.swarm.node.N2.msg
|
||||||
|
NATS->>N2: Deliver request
|
||||||
|
N2->>NATS: Reply with response
|
||||||
|
NATS->>N1: Deliver reply
|
||||||
|
N1->>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[NATS Request to target]
|
||||||
|
E -->|No| G[Retry or Fail]
|
||||||
|
|
||||||
|
F --> H[Target validates & processes]
|
||||||
|
H --> I{Reply received?}
|
||||||
|
I -->|Yes, accepted| J[Return target's response]
|
||||||
|
I -->|Timeout| K[Retry next node or fail]
|
||||||
|
```
|
||||||
|
|
||||||
|
**Handoff Protocol (NATS request-reply):**
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
sequenceDiagram
|
||||||
|
participant Source as Overloaded Node
|
||||||
|
participant NATS
|
||||||
|
participant Target as Selected Node
|
||||||
|
|
||||||
|
Source->>Source: Check load threshold
|
||||||
|
Source->>NATS: Request on picoclaw.swarm.handoff.<targetID>
|
||||||
|
NATS->>Target: Deliver HandoffRequest
|
||||||
|
|
||||||
|
Target->>Target: Validate (load OK? can handle?)
|
||||||
|
alt Accepted
|
||||||
|
Target->>Target: Process request
|
||||||
|
Target->>NATS: Reply: HandoffResponse{accepted, result}
|
||||||
|
NATS->>Source: Deliver reply
|
||||||
|
Source->>Source: Return result to caller
|
||||||
|
else Rejected
|
||||||
|
Target->>NATS: Reply: HandoffResponse{rejected, reason}
|
||||||
|
NATS->>Source: Deliver rejection
|
||||||
|
Source->>Source: Try next node or process locally
|
||||||
|
end
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. Direct Node Messaging
|
||||||
|
|
||||||
|
Nodes can send arbitrary messages to specific peers using NATS request-reply:
|
||||||
|
|
||||||
|
```
|
||||||
|
Subject: picoclaw.swarm.node.<targetNodeID>.msg
|
||||||
|
Payload: JSON {action, message, channel, chat_id, sender_id, trace_id}
|
||||||
|
Reply: JSON {response} or {error}
|
||||||
|
```
|
||||||
|
|
||||||
|
Context parameters (`channel`, `chat_id`, `sender_id`, `trace_id`) enable cross-node audit trails and conversation continuity.
|
||||||
|
|
||||||
|
## System Architecture
|
||||||
|
|
||||||
|
### Component Overview
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
graph TB
|
||||||
|
subgraph "Node 1"
|
||||||
|
D1[Discovery Service] --> M1[Membership Manager]
|
||||||
|
H1[Handoff Coordinator] --> M1
|
||||||
|
L1[Load Monitor] --> H1
|
||||||
|
LE1[Leader Election] --> D1
|
||||||
|
A1[Agent Loop] --> L1
|
||||||
|
end
|
||||||
|
|
||||||
|
subgraph "Node 2"
|
||||||
|
D2[Discovery Service] --> M2[Membership Manager]
|
||||||
|
H2[Handoff Coordinator] --> M2
|
||||||
|
L2[Load Monitor] --> H2
|
||||||
|
LE2[Leader Election] --> D2
|
||||||
|
A2[Agent Loop] --> L2
|
||||||
|
end
|
||||||
|
|
||||||
|
subgraph NATS["NATS Server (JetStream)"]
|
||||||
|
KV1[KV: swarm_members]
|
||||||
|
KV2[KV: swarm_status]
|
||||||
|
KV3[KV: swarm_leader]
|
||||||
|
SUB[Subjects: picoclaw.swarm.*]
|
||||||
|
end
|
||||||
|
|
||||||
|
D1 <-->|KV Watch + Put| KV1
|
||||||
|
D2 <-->|KV Watch + Put| KV1
|
||||||
|
H1 <-->|Request-Reply| SUB
|
||||||
|
H2 <-->|Request-Reply| SUB
|
||||||
|
LE1 <-->|CAS Lock| KV3
|
||||||
|
LE2 <-->|CAS Lock| KV3
|
||||||
|
```
|
||||||
|
|
||||||
|
### Communication Channels
|
||||||
|
|
||||||
|
| Channel | Protocol | Purpose |
|
||||||
|
|---------|----------|---------|
|
||||||
|
| Discovery | NATS KV (JetStream) | Membership liveness via TTL |
|
||||||
|
| Status Cache | NATS KV (JetStream) | Detailed node status (CPU, mem, tasks) |
|
||||||
|
| Leader Election | NATS KV CAS (JetStream) | Distributed leader lock |
|
||||||
|
| Handoff | NATS request-reply | Session transfer coordination |
|
||||||
|
| Node Messaging | NATS request-reply | Direct inter-node messages |
|
||||||
|
| Metrics | NATS publish | Observability data |
|
||||||
|
|
||||||
|
## Configuration
|
||||||
|
|
||||||
|
### Example Configuration
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"swarm": {
|
||||||
|
"enabled": true,
|
||||||
|
"node_id": "picoclaw-node-1",
|
||||||
|
|
||||||
|
"discovery": {
|
||||||
|
"nats_url": "nats://nats.example.com:4222",
|
||||||
|
"nats_creds_file": "/etc/picoclaw/nats.creds",
|
||||||
|
"nats_tls_cert": "/etc/picoclaw/client-cert.pem",
|
||||||
|
"nats_tls_key": "/etc/picoclaw/client-key.pem",
|
||||||
|
"nats_tls_ca_cert": "/etc/picoclaw/ca.pem",
|
||||||
|
"heartbeat_interval": "5s",
|
||||||
|
"member_ttl": "15s",
|
||||||
|
"node_timeout": "5s",
|
||||||
|
"dead_node_timeout": "30s"
|
||||||
|
},
|
||||||
|
|
||||||
|
"handoff": {
|
||||||
|
"enabled": true,
|
||||||
|
"load_threshold": 0.8,
|
||||||
|
"timeout": "30s",
|
||||||
|
"max_retries": 3,
|
||||||
|
"retry_delay": "5s",
|
||||||
|
"request_timeout": "10s"
|
||||||
|
},
|
||||||
|
|
||||||
|
"load_monitor": {
|
||||||
|
"enabled": true,
|
||||||
|
"interval": "5s",
|
||||||
|
"sample_size": 60,
|
||||||
|
"cpu_weight": 0.3,
|
||||||
|
"memory_weight": 0.3,
|
||||||
|
"session_weight": 0.4
|
||||||
|
},
|
||||||
|
|
||||||
|
"leader_election": {
|
||||||
|
"enabled": false,
|
||||||
|
"lock_ttl": "10s",
|
||||||
|
"renewal_interval": "3s"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### Minimal Two-Node Setup
|
||||||
|
|
||||||
|
**1. Start NATS with JetStream:**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
docker run -d --name nats \
|
||||||
|
-p 4222:4222 \
|
||||||
|
nats:latest -js
|
||||||
|
```
|
||||||
|
|
||||||
|
**2. Node 1 config (`node1.json`):**
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"swarm": {
|
||||||
|
"enabled": true,
|
||||||
|
"node_id": "node-1",
|
||||||
|
"discovery": {
|
||||||
|
"nats_url": "nats://localhost:4222"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
**3. Node 2 config (`node2.json`):**
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"swarm": {
|
||||||
|
"enabled": true,
|
||||||
|
"node_id": "node-2",
|
||||||
|
"discovery": {
|
||||||
|
"nats_url": "nats://localhost:4222"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
**4. Verify cluster:**
|
||||||
|
|
||||||
|
Once both nodes start, each should discover the other via KV watch within one `heartbeat_interval` (5s). Use the `swarm_nodes` tool or `/nodes` command to verify membership.
|
||||||
|
|
||||||
|
### Deployment Modes
|
||||||
|
|
||||||
|
**Single Entry Point:**
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
graph LR
|
||||||
|
TG[Telegram Gateway] --> N1[Node 1: Leader]
|
||||||
|
N1 <-->|NATS| N2[Node 2: Worker]
|
||||||
|
N1 <-->|NATS| N3[Node 3: Worker]
|
||||||
|
```
|
||||||
|
|
||||||
|
**Multi-Entry Point (with load balancer):**
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
graph LR
|
||||||
|
LB[Load Balancer] --> N1[Node 1]
|
||||||
|
LB --> N2[Node 2]
|
||||||
|
N1 <-->|NATS Mesh| N2
|
||||||
|
N1 <-->|NATS Mesh| N3[Node 3]
|
||||||
|
N2 <-->|NATS Mesh| N3
|
||||||
|
```
|
||||||
|
|
||||||
|
## 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[Leader Changed] --> ED
|
||||||
|
E[Handoff Started] --> ED
|
||||||
|
F[Handoff Completed] --> ED
|
||||||
|
|
||||||
|
ED --> H[Handlers]
|
||||||
|
H --> L[Logging]
|
||||||
|
H --> M[Metrics]
|
||||||
|
H --> U[Custom Actions]
|
||||||
|
```
|
||||||
|
|
||||||
|
**Event Types:**
|
||||||
|
|
||||||
|
| Event | Description | Payload |
|
||||||
|
|-------|-------------|---------|
|
||||||
|
| `NodeJoined` | New node discovered via KV watch | NodeInfo |
|
||||||
|
| `NodeLeft` | Node KV entry expired or deleted | NodeID |
|
||||||
|
| `NodeSuspect` | Node heartbeat overdue | NodeID |
|
||||||
|
| `NodeAlive` | Node recovered (KV entry renewed) | NodeInfo |
|
||||||
|
| `LeaderChanged` | Leader election result changed | LeaderID |
|
||||||
|
| `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 immediately]
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
### Threat Model
|
||||||
|
|
||||||
|
PicoClaw Swarm assumes nodes are deployed in a **trusted or semi-trusted network**. The NATS server is the single trust boundary.
|
||||||
|
|
||||||
|
| Concern | Mitigation |
|
||||||
|
|---------|------------|
|
||||||
|
| **Transport encryption** | NATS TLS (configure `nats_tls_*` options) |
|
||||||
|
| **Node authentication** | NATS credentials file (NKeys/JWT) or mTLS client certs |
|
||||||
|
| **Subject authorization** | NATS ACL: restrict `picoclaw.swarm.>` to swarm node identities |
|
||||||
|
| **Inter-node message integrity** | NATS connection-level TLS ensures no tampering |
|
||||||
|
| **Unauthorized KV access** | JetStream permissions: only swarm nodes can read/write swarm KV buckets |
|
||||||
|
|
||||||
|
### Production Recommendations
|
||||||
|
|
||||||
|
1. **Always enable TLS** — set `nats_tls_cert`, `nats_tls_key`, `nats_tls_ca_cert`
|
||||||
|
2. **Use NKeys or JWT credentials** — set `nats_creds_file` for per-node identity
|
||||||
|
3. **Configure NATS ACLs** — restrict `picoclaw.swarm.>` publish/subscribe to swarm accounts
|
||||||
|
4. **Network isolation** — NATS port (4222) should not be exposed to public internet
|
||||||
|
5. **Rotate credentials** — use short-lived JWTs with NATS account server for production
|
||||||
|
|
||||||
|
### Minimal Secure NATS Configuration
|
||||||
|
|
||||||
|
```conf
|
||||||
|
# nats-server.conf
|
||||||
|
listen: 0.0.0.0:4222
|
||||||
|
jetstream: enabled
|
||||||
|
|
||||||
|
tls {
|
||||||
|
cert_file: "/etc/nats/server-cert.pem"
|
||||||
|
key_file: "/etc/nats/server-key.pem"
|
||||||
|
ca_file: "/etc/nats/ca.pem"
|
||||||
|
verify: true # require client certs (mTLS)
|
||||||
|
}
|
||||||
|
|
||||||
|
authorization {
|
||||||
|
swarm_user = {
|
||||||
|
publish = ["picoclaw.swarm.>", "$JS.API.>"]
|
||||||
|
subscribe = ["picoclaw.swarm.>", "_INBOX.>"]
|
||||||
|
}
|
||||||
|
users = [
|
||||||
|
{ nkey: "UABC...", permissions: $swarm_user }
|
||||||
|
]
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## Future Enhancements
|
||||||
|
|
||||||
|
1. **Consistent Hashing**: Route requests to stable node assignments for session affinity
|
||||||
|
2. **Multi-Region**: Geo-distributed clusters with region-aware subject prefixes
|
||||||
|
3. **Graceful Draining**: Leader-coordinated node shutdown with session migration
|
||||||
|
4. **Observability**: Prometheus metrics endpoint for swarm health (connection count, handoff latency, leader changes)
|
||||||
3
go.mod
3
go.mod
|
|
@ -15,6 +15,7 @@ require (
|
||||||
github.com/mdp/qrterminal/v3 v3.2.1
|
github.com/mdp/qrterminal/v3 v3.2.1
|
||||||
github.com/modelcontextprotocol/go-sdk v1.3.0
|
github.com/modelcontextprotocol/go-sdk v1.3.0
|
||||||
github.com/mymmrac/telego v1.6.0
|
github.com/mymmrac/telego v1.6.0
|
||||||
|
github.com/nats-io/nats.go v1.49.0
|
||||||
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
|
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
|
||||||
github.com/openai/openai-go/v3 v3.22.0
|
github.com/openai/openai-go/v3 v3.22.0
|
||||||
github.com/rivo/tview v0.42.0
|
github.com/rivo/tview v0.42.0
|
||||||
|
|
@ -42,6 +43,8 @@ require (
|
||||||
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
|
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
|
||||||
github.com/mattn/go-colorable v0.1.14 // indirect
|
github.com/mattn/go-colorable v0.1.14 // indirect
|
||||||
github.com/mattn/go-isatty v0.0.20 // indirect
|
github.com/mattn/go-isatty v0.0.20 // indirect
|
||||||
|
github.com/nats-io/nkeys v0.4.12 // indirect
|
||||||
|
github.com/nats-io/nuid v1.0.1 // indirect
|
||||||
github.com/ncruces/go-strftime v1.0.0 // indirect
|
github.com/ncruces/go-strftime v1.0.0 // indirect
|
||||||
github.com/petermattis/goid v0.0.0-20260113132338-7c7de50cc741 // indirect
|
github.com/petermattis/goid v0.0.0-20260113132338-7c7de50cc741 // indirect
|
||||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||||
|
|
|
||||||
6
go.sum
6
go.sum
|
|
@ -138,6 +138,12 @@ github.com/modelcontextprotocol/go-sdk v1.3.0 h1:gMfZkv3DzQF5q/DcQePo5rahEY+sguy
|
||||||
github.com/modelcontextprotocol/go-sdk v1.3.0/go.mod h1:AnQ//Qc6+4nIyyrB4cxBU7UW9VibK4iOZBeyP/rF1IE=
|
github.com/modelcontextprotocol/go-sdk v1.3.0/go.mod h1:AnQ//Qc6+4nIyyrB4cxBU7UW9VibK4iOZBeyP/rF1IE=
|
||||||
github.com/mymmrac/telego v1.6.0 h1:Zc8rgyHozvd/7ZgyrigyHdAF9koHYMfilYfyB6wlFC0=
|
github.com/mymmrac/telego v1.6.0 h1:Zc8rgyHozvd/7ZgyrigyHdAF9koHYMfilYfyB6wlFC0=
|
||||||
github.com/mymmrac/telego v1.6.0/go.mod h1:xt6ZWA8zi8KmuzryE1ImEdl9JSwjHNpM4yhC7D8hU4Y=
|
github.com/mymmrac/telego v1.6.0/go.mod h1:xt6ZWA8zi8KmuzryE1ImEdl9JSwjHNpM4yhC7D8hU4Y=
|
||||||
|
github.com/nats-io/nats.go v1.49.0 h1:yh/WvY59gXqYpgl33ZI+XoVPKyut/IcEaqtsiuTJpoE=
|
||||||
|
github.com/nats-io/nats.go v1.49.0/go.mod h1:fDCn3mN5cY8HooHwE2ukiLb4p4G4ImmzvXyJt+tGwdw=
|
||||||
|
github.com/nats-io/nkeys v0.4.12 h1:nssm7JKOG9/x4J8II47VWCL1Ds29avyiQDRn0ckMvDc=
|
||||||
|
github.com/nats-io/nkeys v0.4.12/go.mod h1:MT59A1HYcjIcyQDJStTfaOY6vhy9XTUjOFo+SVsvpBg=
|
||||||
|
github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw=
|
||||||
|
github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
|
||||||
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||||
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
|
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
|
||||||
|
|
|
||||||
|
|
@ -30,6 +30,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"
|
||||||
"github.com/sipeed/picoclaw/pkg/voice"
|
"github.com/sipeed/picoclaw/pkg/voice"
|
||||||
|
|
@ -46,6 +47,22 @@ type AgentLoop struct {
|
||||||
channelManager *channels.Manager
|
channelManager *channels.Manager
|
||||||
mediaStore media.MediaStore
|
mediaStore media.MediaStore
|
||||||
transcriber voice.Transcriber
|
transcriber voice.Transcriber
|
||||||
|
|
||||||
|
// Lifecycle context for proper shutdown propagation
|
||||||
|
ctx context.Context
|
||||||
|
cancel context.CancelFunc
|
||||||
|
|
||||||
|
// Swarm mode support
|
||||||
|
swarmEnabled bool
|
||||||
|
swarmInitError error // Set if swarm initialization fails
|
||||||
|
swarmDiscovery swarm.Discovery
|
||||||
|
swarmHandoff *swarm.HandoffCoordinator
|
||||||
|
swarmLoad *swarm.LoadMonitor
|
||||||
|
swarmLeaderElection *swarm.LeaderElection
|
||||||
|
swarmSigner *swarm.Signer // Message-layer signature verification for node-to-node messages
|
||||||
|
|
||||||
|
// Dynamic command registry (wired from channel manager)
|
||||||
|
commandRegistry *channels.CommandRegistry
|
||||||
}
|
}
|
||||||
|
|
||||||
// processOptions configures how a message is processed
|
// processOptions configures how a message is processed
|
||||||
|
|
@ -84,7 +101,7 @@ func NewAgentLoop(
|
||||||
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,
|
||||||
|
|
@ -92,6 +109,13 @@ func NewAgentLoop(
|
||||||
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).
|
||||||
|
|
@ -208,6 +232,10 @@ func registerSharedTools(
|
||||||
}
|
}
|
||||||
|
|
||||||
func (al *AgentLoop) Run(ctx context.Context) error {
|
func (al *AgentLoop) Run(ctx context.Context) error {
|
||||||
|
// Store context for proper cancellation propagation
|
||||||
|
al.ctx, al.cancel = context.WithCancel(ctx)
|
||||||
|
defer al.cancel()
|
||||||
|
|
||||||
al.running.Store(true)
|
al.running.Store(true)
|
||||||
|
|
||||||
// Initialize MCP servers for all agents
|
// Initialize MCP servers for all agents
|
||||||
|
|
@ -350,6 +378,12 @@ func (al *AgentLoop) Run(ctx context.Context) error {
|
||||||
|
|
||||||
func (al *AgentLoop) Stop() {
|
func (al *AgentLoop) Stop() {
|
||||||
al.running.Store(false)
|
al.running.Store(false)
|
||||||
|
// Gracefully shutdown swarm mode components if enabled
|
||||||
|
if al.swarmEnabled {
|
||||||
|
if err := al.ShutdownSwarm(); err != nil {
|
||||||
|
logger.WarnCF("swarm", "Error during swarm shutdown", map[string]any{"error": err})
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (al *AgentLoop) RegisterTool(tool tools.Tool) {
|
func (al *AgentLoop) RegisterTool(tool tools.Tool) {
|
||||||
|
|
@ -362,6 +396,12 @@ func (al *AgentLoop) RegisterTool(tool tools.Tool) {
|
||||||
|
|
||||||
func (al *AgentLoop) SetChannelManager(cm *channels.Manager) {
|
func (al *AgentLoop) SetChannelManager(cm *channels.Manager) {
|
||||||
al.channelManager = cm
|
al.channelManager = cm
|
||||||
|
al.commandRegistry = cm.CommandRegistry()
|
||||||
|
|
||||||
|
// Wire up swarm commands
|
||||||
|
if al.swarmEnabled {
|
||||||
|
al.registerSwarmCommands()
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// SetMediaStore injects a MediaStore for media lifecycle management.
|
// SetMediaStore injects a MediaStore for media lifecycle management.
|
||||||
|
|
@ -545,6 +585,25 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
|
||||||
return al.processSystemMessage(ctx, msg)
|
return al.processSystemMessage(ctx, msg)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Check for @node-id routing syntax
|
||||||
|
if al.swarmEnabled {
|
||||||
|
if targetNodeID, content := tools.ParseNodeMention(msg.Content); targetNodeID != "" {
|
||||||
|
logger.InfoCF("swarm", "Node mention detected", map[string]any{
|
||||||
|
"target": targetNodeID,
|
||||||
|
"content": utils.Truncate(content, 50),
|
||||||
|
})
|
||||||
|
return al.handleNodeRouting(ctx, msg, targetNodeID, content)
|
||||||
|
}
|
||||||
|
|
||||||
|
// For non-mentioned messages in swarm mode, only leader processes
|
||||||
|
if al.swarmLeaderElection != nil && !al.swarmLeaderElection.IsLeader() {
|
||||||
|
logger.DebugCF("swarm", "Ignoring message (not leader)", map[string]any{
|
||||||
|
"content": utils.Truncate(msg.Content, 50),
|
||||||
|
})
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Check for commands
|
// Check for commands
|
||||||
if response, handled := al.handleCommand(ctx, msg); handled {
|
if response, handled := al.handleCommand(ctx, msg); handled {
|
||||||
return response, nil
|
return response, nil
|
||||||
|
|
@ -588,6 +647,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,
|
||||||
|
|
@ -671,6 +746,10 @@ func (al *AgentLoop) runAgentLoop(
|
||||||
agent *AgentInstance,
|
agent *AgentInstance,
|
||||||
opts processOptions,
|
opts processOptions,
|
||||||
) (string, error) {
|
) (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)
|
||||||
|
|
@ -1177,6 +1256,15 @@ func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, c
|
||||||
go func() {
|
go func() {
|
||||||
defer al.summarizing.Delete(summarizeKey)
|
defer al.summarizing.Delete(summarizeKey)
|
||||||
logger.Debug("Memory threshold reached. Optimizing conversation history...")
|
logger.Debug("Memory threshold reached. Optimizing conversation history...")
|
||||||
|
if !constants.IsInternalChannel(channel) {
|
||||||
|
pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||||
|
defer pubCancel()
|
||||||
|
al.bus.PublishOutbound(pubCtx, bus.OutboundMessage{
|
||||||
|
Channel: channel,
|
||||||
|
ChatID: chatID,
|
||||||
|
Content: "Memory threshold reached. Optimizing conversation history...",
|
||||||
|
})
|
||||||
|
}
|
||||||
al.summarizeSession(agent, sessionKey)
|
al.summarizeSession(agent, sessionKey)
|
||||||
}()
|
}()
|
||||||
}
|
}
|
||||||
|
|
@ -1543,6 +1631,19 @@ func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Check dynamic command registry
|
||||||
|
if al.commandRegistry != nil {
|
||||||
|
cmdName := strings.TrimPrefix(parts[0], "/")
|
||||||
|
if entry, ok := al.commandRegistry.Get(cmdName); ok {
|
||||||
|
argsStr := strings.Join(args, " ")
|
||||||
|
response, err := entry.Handler(ctx, argsStr, msg)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Sprintf("Command error: %v", err), true
|
||||||
|
}
|
||||||
|
return response, true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
return "", false
|
return "", false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
925
pkg/agent/loop_swarm.go
Normal file
925
pkg/agent/loop_swarm.go
Normal file
|
|
@ -0,0 +1,925 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package agent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/nats-io/nats.go"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/swarm"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/tools"
|
||||||
|
)
|
||||||
|
|
||||||
|
// handleNodeRouting handles routing a message to a specific node in the swarm via NATS.
|
||||||
|
func (al *AgentLoop) handleNodeRouting(
|
||||||
|
ctx context.Context,
|
||||||
|
msg bus.InboundMessage,
|
||||||
|
targetNodeID, content string,
|
||||||
|
) (string, error) {
|
||||||
|
// Check if target is this node
|
||||||
|
if targetNodeID == al.swarmDiscovery.LocalNode().ID {
|
||||||
|
logger.InfoCF("swarm", "Target is this node, processing locally", map[string]any{"node_id": targetNodeID})
|
||||||
|
msg.Content = content
|
||||||
|
return al.processMessage(ctx, msg)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Find target node
|
||||||
|
members := al.swarmDiscovery.Members()
|
||||||
|
var target *swarm.NodeWithState
|
||||||
|
for _, m := range members {
|
||||||
|
if m.Node.ID == targetNodeID {
|
||||||
|
target = m
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if target == nil {
|
||||||
|
return fmt.Sprintf("Node '%s' not found in cluster. Use /nodes to see available nodes.", targetNodeID), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if target is available
|
||||||
|
if target.State.Status != swarm.NodeStatusAlive {
|
||||||
|
return fmt.Sprintf("Node '%s' is not alive (status: %s)", targetNodeID, target.State.Status), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if target is overloaded using configurable threshold with hysteresis
|
||||||
|
rejectThreshold := al.cfg.Swarm.LoadMonitor.RoutingRejectThreshold
|
||||||
|
if rejectThreshold == 0 {
|
||||||
|
rejectThreshold = swarm.DefaultRoutingRejectThreshold // Default to 0.9
|
||||||
|
}
|
||||||
|
if target.Node.LoadScore > rejectThreshold {
|
||||||
|
return fmt.Sprintf("Node '%s' is overloaded (load: %.0f%%, threshold: %.0f%%)",
|
||||||
|
targetNodeID, target.Node.LoadScore*100, rejectThreshold*100), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
nc := al.swarmDiscovery.NATSConn()
|
||||||
|
if nc == nil {
|
||||||
|
return "Node-to-node communication not initialized (NATS not connected)", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Send message via NATS to target node's message subject
|
||||||
|
subject := swarm.SubjectNodeMsg + "." + targetNodeID + ".msg"
|
||||||
|
|
||||||
|
payload := map[string]any{
|
||||||
|
"content": content,
|
||||||
|
"channel": msg.Channel,
|
||||||
|
"chat_id": msg.ChatID,
|
||||||
|
"sender_id": msg.SenderID,
|
||||||
|
"from_node": al.swarmDiscovery.LocalNode().ID,
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Sprintf("Failed to marshal message: %v", err), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Sending message to remote node via NATS", map[string]any{
|
||||||
|
"target": targetNodeID,
|
||||||
|
"subject": subject,
|
||||||
|
"load": target.Node.LoadScore,
|
||||||
|
})
|
||||||
|
|
||||||
|
// Use NATS request-reply for synchronous response
|
||||||
|
reply, err := nc.RequestWithContext(ctx, subject, data)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Sprintf("Failed to send message to node '%s': %v", targetNodeID, err), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parse response
|
||||||
|
var response map[string]any
|
||||||
|
if err := json.Unmarshal(reply.Data, &response); err != nil {
|
||||||
|
return string(reply.Data), nil //nolint:nilerr // Fallback to raw response on parse failure
|
||||||
|
}
|
||||||
|
|
||||||
|
if errMsg, ok := response["error"].(string); ok && errMsg != "" {
|
||||||
|
return fmt.Sprintf("Node '%s' error: %s", targetNodeID, errMsg), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp, ok := response["response"].(string); ok {
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return string(reply.Data), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// initSwarm initializes the swarm mode components.
|
||||||
|
func (al *AgentLoop) initSwarm() {
|
||||||
|
logger.InfoC("swarm", "Initializing swarm mode")
|
||||||
|
|
||||||
|
// Convert config and create discovery service directly (NATS-only)
|
||||||
|
swarmConfig := al.convertToSwarmConfig(al.cfg.Swarm)
|
||||||
|
discovery, err := swarm.NewDiscoveryService(swarmConfig)
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF(
|
||||||
|
"swarm",
|
||||||
|
"SWARM INIT FAILED: discovery service creation failed — swarm mode will be disabled",
|
||||||
|
map[string]any{"error": err.Error()},
|
||||||
|
)
|
||||||
|
al.swarmInitError = fmt.Errorf("discovery service creation failed: %w", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := discovery.Start(); err != nil {
|
||||||
|
logger.ErrorCF(
|
||||||
|
"swarm",
|
||||||
|
"SWARM INIT FAILED: discovery service start failed — swarm mode will be disabled",
|
||||||
|
map[string]any{"error": err.Error()},
|
||||||
|
)
|
||||||
|
al.swarmInitError = fmt.Errorf("discovery service start failed: %w", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
al.swarmDiscovery = discovery
|
||||||
|
|
||||||
|
// Create handoff coordinator using config from convertToSwarmConfig (with default duration protection)
|
||||||
|
al.swarmHandoff = swarm.NewHandoffCoordinator(discovery, swarmConfig.Handoff)
|
||||||
|
|
||||||
|
// Configure crypto settings if provided
|
||||||
|
if al.cfg.Swarm.Handoff.SharedSecret != "" || al.cfg.Swarm.Handoff.EncryptionKey != "" {
|
||||||
|
cryptoCfg := swarm.CryptoConfig{
|
||||||
|
SharedSecret: al.cfg.Swarm.Handoff.SharedSecret,
|
||||||
|
EncryptionKey: al.cfg.Swarm.Handoff.EncryptionKey,
|
||||||
|
RequireAuth: al.cfg.Swarm.Handoff.RequireAuth,
|
||||||
|
RequireEncryption: al.cfg.Swarm.Handoff.RequireEncryption,
|
||||||
|
}
|
||||||
|
if err := al.swarmHandoff.SetCryptoConfig(cryptoCfg); err != nil {
|
||||||
|
logger.ErrorCF(
|
||||||
|
"swarm",
|
||||||
|
"Failed to configure handoff crypto — handoff disabled",
|
||||||
|
map[string]any{"error": err.Error()},
|
||||||
|
)
|
||||||
|
al.swarmHandoff.Close()
|
||||||
|
al.swarmHandoff = nil
|
||||||
|
} else {
|
||||||
|
logger.InfoCF("swarm", "Handoff crypto configured", map[string]any{
|
||||||
|
"auth_enabled": al.cfg.Swarm.Handoff.SharedSecret != "",
|
||||||
|
"encryption_enabled": al.cfg.Swarm.Handoff.EncryptionKey != "",
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Initialize signer for node-to-node message authentication
|
||||||
|
// Uses the same SharedSecret as handoff for consistency
|
||||||
|
if al.cfg.Swarm.Handoff.SharedSecret != "" {
|
||||||
|
al.swarmSigner = swarm.NewSigner(al.cfg.Swarm.Handoff.SharedSecret)
|
||||||
|
logger.InfoCF("swarm", "Node message signing enabled", map[string]any{
|
||||||
|
"require_auth": al.cfg.Swarm.Handoff.RequireAuth,
|
||||||
|
})
|
||||||
|
} else if al.cfg.Swarm.Handoff.RequireAuth {
|
||||||
|
logger.WarnCF(
|
||||||
|
"swarm",
|
||||||
|
"RequireAuth is set but no SharedSecret provided - node messages will not be signed or verified",
|
||||||
|
nil,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
if al.swarmHandoff != nil {
|
||||||
|
if err := al.swarmHandoff.Start(); err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to start handoff coordinator", map[string]any{"error": err.Error()})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create and start load monitor using config from convertToSwarmConfig (with default duration protection)
|
||||||
|
al.swarmLoad = swarm.NewLoadMonitor(&swarmConfig.LoadMonitor)
|
||||||
|
if al.cfg.Swarm.LoadMonitor.Enabled {
|
||||||
|
al.swarmLoad.Start()
|
||||||
|
al.swarmLoad.OnThreshold(func(score float64) {
|
||||||
|
discovery.UpdateLoad(score)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wire load monitor to discovery service for detailed status publishing
|
||||||
|
discovery.SetLoadMonitor(al.swarmLoad)
|
||||||
|
|
||||||
|
al.swarmEnabled = true
|
||||||
|
|
||||||
|
// Initialize leader election if enabled
|
||||||
|
if al.cfg.Swarm.LeaderElection.Enabled {
|
||||||
|
al.initSwarmLeaderElection(discovery)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Set up NATS subscription for incoming directed node messages
|
||||||
|
al.setupNodeMessageSubscription(discovery)
|
||||||
|
|
||||||
|
// Register swarm tools
|
||||||
|
al.registerSwarmTools(discovery)
|
||||||
|
|
||||||
|
// Subscribe to node events for logging
|
||||||
|
al.subscribeSwarmEvents(discovery)
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Swarm mode initialized", map[string]any{
|
||||||
|
"node_id": discovery.LocalNode().ID,
|
||||||
|
"nats_url": al.cfg.Swarm.Discovery.NATSURL,
|
||||||
|
"handoff": al.cfg.Swarm.Handoff.Enabled,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// setupNodeMessageSubscription subscribes to incoming directed messages via NATS.
|
||||||
|
func (al *AgentLoop) setupNodeMessageSubscription(discovery swarm.Discovery) {
|
||||||
|
nc := discovery.NATSConn()
|
||||||
|
if nc == nil {
|
||||||
|
logger.WarnC("swarm", "NATS connection not available, node message subscription skipped")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
localNodeID := discovery.LocalNode().ID
|
||||||
|
subject := swarm.SubjectNodeMsg + "." + localNodeID + ".msg"
|
||||||
|
|
||||||
|
_, err := nc.Subscribe(subject, func(msg *nats.Msg) {
|
||||||
|
al.handleIncomingNATSNodeMessage(msg)
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to subscribe to node message subject", map[string]any{
|
||||||
|
"subject": subject,
|
||||||
|
"error": err.Error(),
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Subscribed to node messages", map[string]any{"subject": subject})
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleIncomingNATSNodeMessage handles a message received from another node via NATS.
|
||||||
|
func (al *AgentLoop) handleIncomingNATSNodeMessage(msg *nats.Msg) {
|
||||||
|
var payload map[string]any
|
||||||
|
if err := json.Unmarshal(msg.Data, &payload); err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to decode incoming node message", map[string]any{"error": err.Error()})
|
||||||
|
al.sendNodeReply(msg, "invalid message format", nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify sender is a known swarm member (routing authorization)
|
||||||
|
fromNode, _ := payload["from_node"].(string)
|
||||||
|
if fromNode != "" && al.swarmDiscovery != nil {
|
||||||
|
members := al.swarmDiscovery.Members()
|
||||||
|
isKnownMember := false
|
||||||
|
for _, m := range members {
|
||||||
|
if m.Node.ID == fromNode {
|
||||||
|
isKnownMember = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !isKnownMember {
|
||||||
|
logger.WarnCF("swarm", "Rejected node message from unknown member", map[string]any{
|
||||||
|
"from_node": fromNode,
|
||||||
|
"known_members": len(members),
|
||||||
|
})
|
||||||
|
al.sendNodeReply(msg, "unknown node", nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify message signature if signer is configured
|
||||||
|
signature, hasSignature := payload["signature"].(string)
|
||||||
|
if hasSignature && al.swarmSigner != nil {
|
||||||
|
// Create a copy of payload without signature for verification
|
||||||
|
payloadToVerify := make(map[string]any)
|
||||||
|
for k, v := range payload {
|
||||||
|
if k != "signature" {
|
||||||
|
payloadToVerify[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !al.swarmSigner.Verify(payloadToVerify, signature) {
|
||||||
|
logger.WarnCF("swarm", "Rejected node message with invalid signature", map[string]any{
|
||||||
|
"from_node": fromNode,
|
||||||
|
})
|
||||||
|
al.sendNodeReply(msg, "invalid signature", nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
} else if al.cfg.Swarm.Handoff.RequireAuth && al.swarmSigner == nil {
|
||||||
|
// RequireAuth is enabled but no signer configured - reject for safety
|
||||||
|
logger.WarnCF(
|
||||||
|
"swarm",
|
||||||
|
"Rejected node message - RequireAuth enabled but no signature verification configured",
|
||||||
|
map[string]any{
|
||||||
|
"from_node": fromNode,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
al.sendNodeReply(msg, "signature required but not configured", nil)
|
||||||
|
return
|
||||||
|
} else if al.cfg.Swarm.Handoff.RequireAuth && !hasSignature {
|
||||||
|
// RequireAuth is enabled but message has no signature
|
||||||
|
logger.WarnCF("swarm", "Rejected unsigned node message - RequireAuth enabled", map[string]any{
|
||||||
|
"from_node": fromNode,
|
||||||
|
})
|
||||||
|
al.sendNodeReply(msg, "signature required", nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check for action-based request (lightweight, no LLM processing)
|
||||||
|
action, hasAction := payload["action"].(string)
|
||||||
|
if hasAction {
|
||||||
|
al.handleNodeActionRequest(msg, action, payload, fromNode)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Content-based message (full LLM processing)
|
||||||
|
content, _ := payload["content"].(string)
|
||||||
|
channel, _ := payload["channel"].(string)
|
||||||
|
chatID, _ := payload["chat_id"].(string)
|
||||||
|
senderID, _ := payload["sender_id"].(string)
|
||||||
|
|
||||||
|
// Require content for message processing
|
||||||
|
if content == "" {
|
||||||
|
al.sendNodeReply(msg, "empty content", nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Processing incoming node message", map[string]any{
|
||||||
|
"from": fromNode,
|
||||||
|
"content": content[:min(50, len(content))],
|
||||||
|
})
|
||||||
|
|
||||||
|
// Create an inbound message and process with a timeout context
|
||||||
|
inboundMsg := bus.InboundMessage{
|
||||||
|
Content: content,
|
||||||
|
Channel: channel,
|
||||||
|
ChatID: chatID,
|
||||||
|
SenderID: senderID,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Use parent context from AgentLoop for proper cancellation propagation
|
||||||
|
// If al.ctx is not set (e.g., during initialization), fall back to background context
|
||||||
|
const nodeMessageTimeout = 120 * time.Second
|
||||||
|
parentCtx := al.ctx
|
||||||
|
if parentCtx == nil {
|
||||||
|
parentCtx = context.Background()
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithTimeout(parentCtx, nodeMessageTimeout)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
response, err := al.processMessage(ctx, inboundMsg)
|
||||||
|
|
||||||
|
// Send reply if request-reply pattern
|
||||||
|
if msg.Reply != "" {
|
||||||
|
var respData []byte
|
||||||
|
if err != nil {
|
||||||
|
respData, _ = json.Marshal(map[string]any{"error": err.Error()})
|
||||||
|
} else {
|
||||||
|
respData, _ = json.Marshal(map[string]any{"response": response})
|
||||||
|
}
|
||||||
|
msg.Respond(respData)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleNodeActionRequest handles lightweight action requests (status, ping, health, etc.)
|
||||||
|
// These actions bypass LLM processing and return structured responses directly.
|
||||||
|
func (al *AgentLoop) handleNodeActionRequest(msg *nats.Msg, action string, payload map[string]any, fromNode string) {
|
||||||
|
logger.InfoCF("swarm", "Processing node action request", map[string]any{
|
||||||
|
"from": fromNode,
|
||||||
|
"action": action,
|
||||||
|
})
|
||||||
|
|
||||||
|
var response any
|
||||||
|
var err error
|
||||||
|
|
||||||
|
switch action {
|
||||||
|
case swarm.NodeActionPing:
|
||||||
|
response = map[string]any{
|
||||||
|
"status": "pong",
|
||||||
|
"node_id": al.swarmDiscovery.LocalNode().ID,
|
||||||
|
}
|
||||||
|
|
||||||
|
case swarm.NodeActionStatus:
|
||||||
|
response = al.GetSwarmStatus()
|
||||||
|
|
||||||
|
case swarm.NodeActionHealth:
|
||||||
|
response = map[string]any{
|
||||||
|
"status": "healthy",
|
||||||
|
"node_id": al.swarmDiscovery.LocalNode().ID,
|
||||||
|
}
|
||||||
|
|
||||||
|
case swarm.NodeActionLeader:
|
||||||
|
if al.swarmLeaderElection != nil {
|
||||||
|
response = map[string]any{
|
||||||
|
"leader_id": al.swarmLeaderElection.GetLeader(),
|
||||||
|
"is_leader": al.swarmLeaderElection.IsLeader(),
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
response = map[string]any{
|
||||||
|
"leader_id": nil,
|
||||||
|
"message": "leader election not enabled",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
case swarm.NodeActionMetrics:
|
||||||
|
response = al.GetSwarmStatus()
|
||||||
|
|
||||||
|
default:
|
||||||
|
err = fmt.Errorf("unknown action: %s", action)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Send reply
|
||||||
|
if msg.Reply != "" {
|
||||||
|
respData, _ := json.Marshal(map[string]any{
|
||||||
|
"response": response,
|
||||||
|
"error": err,
|
||||||
|
})
|
||||||
|
msg.Respond(respData)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// sendNodeReply sends a standardized error reply to a node message.
|
||||||
|
func (al *AgentLoop) sendNodeReply(msg *nats.Msg, errMsg string, payload map[string]any) {
|
||||||
|
if msg.Reply == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
reply := map[string]any{"error": errMsg}
|
||||||
|
if payload != nil {
|
||||||
|
for k, v := range payload {
|
||||||
|
reply[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
respData, _ := json.Marshal(reply)
|
||||||
|
msg.Respond(respData)
|
||||||
|
}
|
||||||
|
|
||||||
|
// initSwarmLeaderElection initializes the leader election module using NATS KV CAS lock.
|
||||||
|
func (al *AgentLoop) initSwarmLeaderElection(discovery swarm.Discovery) {
|
||||||
|
js := discovery.JetStream()
|
||||||
|
if js == nil {
|
||||||
|
logger.WarnCF("swarm", "JetStream not available, leader election disabled", nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert config
|
||||||
|
leaderElectionConfig := swarm.LeaderElectionConfig{
|
||||||
|
Enabled: al.cfg.Swarm.LeaderElection.Enabled,
|
||||||
|
LockTTL: swarm.Duration{
|
||||||
|
Duration: time.Duration(al.cfg.Swarm.LeaderElection.LockTTL) * time.Second,
|
||||||
|
},
|
||||||
|
RenewalInterval: swarm.Duration{
|
||||||
|
Duration: time.Duration(al.cfg.Swarm.LeaderElection.RenewalInterval) * time.Second,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create leader election instance with JetStream
|
||||||
|
buckets := discovery.Buckets()
|
||||||
|
if buckets == nil || buckets.Leader == nil {
|
||||||
|
logger.WarnCF("swarm", "Leader bucket not available, leader election disabled", nil)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
leaderElection, err := swarm.NewLeaderElection(
|
||||||
|
discovery.LocalNode().ID,
|
||||||
|
buckets.Leader,
|
||||||
|
leaderElectionConfig,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to create leader election", map[string]any{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start leader election
|
||||||
|
if err := leaderElection.Start(); err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to start leader election", map[string]any{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
al.swarmLeaderElection = leaderElection
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Leader election initialized (KV CAS lock)", map[string]any{
|
||||||
|
"node_id": discovery.LocalNode().ID,
|
||||||
|
"enabled": al.cfg.Swarm.LeaderElection.Enabled,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// sendMessageToNode sends a message to a target node via NATS.
|
||||||
|
func (al *AgentLoop) sendMessageToNode(
|
||||||
|
ctx context.Context,
|
||||||
|
targetNodeID, content, channel, chatID, senderID string,
|
||||||
|
) (string, error) {
|
||||||
|
nc := al.swarmDiscovery.NATSConn()
|
||||||
|
if nc == nil {
|
||||||
|
return "", fmt.Errorf("NATS connection not available")
|
||||||
|
}
|
||||||
|
|
||||||
|
subject := swarm.SubjectNodeMsg + "." + targetNodeID + ".msg"
|
||||||
|
|
||||||
|
payload := map[string]any{
|
||||||
|
"content": content,
|
||||||
|
"channel": channel,
|
||||||
|
"chat_id": chatID,
|
||||||
|
"sender_id": senderID,
|
||||||
|
"from_node": al.swarmDiscovery.LocalNode().ID,
|
||||||
|
"timestamp": time.Now().UnixNano(),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add signature if signer is configured
|
||||||
|
if al.swarmSigner != nil {
|
||||||
|
signature, err := al.swarmSigner.Sign(payload)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to sign message: %w", err)
|
||||||
|
}
|
||||||
|
payload["signature"] = signature
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := json.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to marshal message: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
reply, err := nc.RequestWithContext(ctx, subject, data)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to send message to node %s: %w", targetNodeID, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var response map[string]any
|
||||||
|
if err := json.Unmarshal(reply.Data, &response); err != nil {
|
||||||
|
return string(reply.Data), nil //nolint:nilerr // Fallback to raw response on parse failure
|
||||||
|
}
|
||||||
|
|
||||||
|
if errMsg, ok := response["error"].(string); ok && errMsg != "" {
|
||||||
|
return "", fmt.Errorf("node %s error: %s", targetNodeID, errMsg)
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp, ok := response["response"].(string); ok {
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return string(reply.Data), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// sendActionToNode sends an action-based request to a target node via NATS.
|
||||||
|
// Used for fast operations like 'status' that don't require LLM processing.
|
||||||
|
func (al *AgentLoop) sendActionToNode(
|
||||||
|
ctx context.Context,
|
||||||
|
targetNodeID, action string,
|
||||||
|
) (string, error) {
|
||||||
|
nc := al.swarmDiscovery.NATSConn()
|
||||||
|
if nc == nil {
|
||||||
|
return "", fmt.Errorf("NATS connection not available")
|
||||||
|
}
|
||||||
|
|
||||||
|
subject := swarm.SubjectNodeMsg + "." + targetNodeID + ".msg"
|
||||||
|
|
||||||
|
payload := map[string]any{
|
||||||
|
"action": action,
|
||||||
|
"from_node": al.swarmDiscovery.LocalNode().ID,
|
||||||
|
"timestamp": time.Now().UnixNano(),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add signature if signer is configured
|
||||||
|
if al.swarmSigner != nil {
|
||||||
|
signature, err := al.swarmSigner.Sign(payload)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to sign action: %w", err)
|
||||||
|
}
|
||||||
|
payload["signature"] = signature
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := json.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to marshal action: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
reply, err := nc.RequestWithContext(ctx, subject, data)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to send action to node %s: %w", targetNodeID, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var response map[string]any
|
||||||
|
if err := json.Unmarshal(reply.Data, &response); err != nil {
|
||||||
|
return string(reply.Data), nil //nolint:nilerr // Fallback to raw response on parse failure
|
||||||
|
}
|
||||||
|
|
||||||
|
if errMsg, ok := response["error"].(string); ok && errMsg != "" {
|
||||||
|
return "", fmt.Errorf("node error: %s", errMsg)
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp, ok := response["response"].(string); ok {
|
||||||
|
return resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return string(reply.Data), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// registerSwarmTools registers swarm-related tools (handoff, routing).
|
||||||
|
func (al *AgentLoop) registerSwarmTools(discovery swarm.Discovery) {
|
||||||
|
localNodeID := discovery.LocalNode().ID
|
||||||
|
|
||||||
|
swarmTool := tools.NewSwarmTool(discovery, al.swarmLoad, localNodeID)
|
||||||
|
swarmTool.SetSendMessageFn(al.sendMessageToNode)
|
||||||
|
al.RegisterTool(swarmTool)
|
||||||
|
|
||||||
|
nodesTool := tools.NewSwarmNodesTool(discovery, al.swarmLoad, localNodeID)
|
||||||
|
al.RegisterTool(nodesTool)
|
||||||
|
|
||||||
|
routeTool := tools.NewSwarmRouteTool(discovery, al.swarmHandoff, localNodeID)
|
||||||
|
routeTool.SetSendMessageFn(al.sendMessageToNode)
|
||||||
|
routeTool.SetSendActionFn(al.sendActionToNode)
|
||||||
|
al.RegisterTool(routeTool)
|
||||||
|
|
||||||
|
// Register batch query tool for parallel node queries
|
||||||
|
batchTool := tools.NewSwarmBatchTool(discovery, localNodeID)
|
||||||
|
batchTool.SetSendActionFn(al.sendActionToNode)
|
||||||
|
al.RegisterTool(batchTool)
|
||||||
|
}
|
||||||
|
|
||||||
|
// subscribeSwarmEvents subscribes to node join/leave events.
|
||||||
|
func (al *AgentLoop) subscribeSwarmEvents(discovery swarm.Discovery) {
|
||||||
|
_ = 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})
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// registerSwarmCommands registers swarm slash commands on the channel command registry.
|
||||||
|
func (al *AgentLoop) registerSwarmCommands() {
|
||||||
|
if al.commandRegistry == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// If swarm init failed, register a /nodes command that shows the error.
|
||||||
|
if al.swarmDiscovery == nil {
|
||||||
|
if al.swarmInitError != nil {
|
||||||
|
al.commandRegistry.Register("nodes", "List swarm cluster nodes", func(
|
||||||
|
ctx context.Context, args string, msg bus.InboundMessage,
|
||||||
|
) (string, error) {
|
||||||
|
return fmt.Sprintf("Swarm mode failed to initialize: %s", al.swarmInitError), nil
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
localNodeID := al.swarmDiscovery.LocalNode().ID
|
||||||
|
|
||||||
|
al.commandRegistry.Register("nodes", "List swarm cluster nodes", func(
|
||||||
|
ctx context.Context, args string, msg bus.InboundMessage,
|
||||||
|
) (string, error) {
|
||||||
|
verbose := strings.Contains(args, "verbose") || strings.Contains(args, "-v")
|
||||||
|
return tools.FormatClusterStatus(al.swarmDiscovery, al.swarmLoad, localNodeID, verbose), nil
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// convertToSwarmConfig converts the config.SwarmConfig to swarm.Config.
|
||||||
|
func (al *AgentLoop) convertToSwarmConfig(cfg config.SwarmConfig) *swarm.Config {
|
||||||
|
applyDurationDefault := func(val int, defaultVal time.Duration) time.Duration {
|
||||||
|
if val <= 0 {
|
||||||
|
return defaultVal
|
||||||
|
}
|
||||||
|
return time.Duration(val) * time.Second
|
||||||
|
}
|
||||||
|
|
||||||
|
return &swarm.Config{
|
||||||
|
Enabled: cfg.Enabled,
|
||||||
|
NodeID: cfg.NodeID,
|
||||||
|
Discovery: swarm.DiscoveryConfig{
|
||||||
|
NATSURL: cfg.Discovery.NATSURL,
|
||||||
|
NATSCredsFile: cfg.Discovery.NATSCredsFile,
|
||||||
|
NATSTLSCert: cfg.Discovery.NATSTLSCert,
|
||||||
|
NATSTLSKey: cfg.Discovery.NATSTLSKey,
|
||||||
|
NATSTLSCACert: cfg.Discovery.NATSTLSCACert,
|
||||||
|
HeartbeatInterval: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.Discovery.HeartbeatInterval, swarm.DefaultHeartbeatInterval),
|
||||||
|
},
|
||||||
|
MemberTTL: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.Discovery.MemberTTL, swarm.DefaultMemberTTL),
|
||||||
|
},
|
||||||
|
SubjectPrefix: cfg.Discovery.SubjectPrefix,
|
||||||
|
NodeTimeout: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.Discovery.NodeTimeout, swarm.DefaultNodeTimeout),
|
||||||
|
},
|
||||||
|
DeadNodeTimeout: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.Discovery.DeadNodeTimeout, swarm.DefaultDeadNodeTimeout),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Handoff: swarm.HandoffConfig{
|
||||||
|
Enabled: cfg.Handoff.Enabled,
|
||||||
|
LoadThreshold: cfg.Handoff.LoadThreshold,
|
||||||
|
Timeout: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.Handoff.Timeout, swarm.DefaultHandoffTimeout),
|
||||||
|
},
|
||||||
|
MaxRetries: cfg.Handoff.MaxRetries,
|
||||||
|
RetryDelay: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.Handoff.RetryDelay, swarm.DefaultHandoffRetryDelay),
|
||||||
|
},
|
||||||
|
RequestTimeout: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.Handoff.RequestTimeout, swarm.DefaultHandoffRequestTimeout),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
LoadMonitor: swarm.LoadMonitorConfig{
|
||||||
|
Enabled: cfg.LoadMonitor.Enabled,
|
||||||
|
Interval: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.LoadMonitor.Interval, swarm.DefaultLoadSampleInterval),
|
||||||
|
},
|
||||||
|
SampleSize: cfg.LoadMonitor.SampleSize,
|
||||||
|
CPUWeight: cfg.LoadMonitor.CPUWeight,
|
||||||
|
MemoryWeight: cfg.LoadMonitor.MemoryWeight,
|
||||||
|
SessionWeight: cfg.LoadMonitor.SessionWeight,
|
||||||
|
},
|
||||||
|
LeaderElection: swarm.LeaderElectionConfig{
|
||||||
|
Enabled: cfg.LeaderElection.Enabled,
|
||||||
|
LockTTL: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.LeaderElection.LockTTL, swarm.DefaultLeaderLockTTL),
|
||||||
|
},
|
||||||
|
RenewalInterval: swarm.Duration{
|
||||||
|
Duration: applyDurationDefault(cfg.LeaderElection.RenewalInterval, swarm.DefaultLeaderRenewalInterval),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
}
|
||||||
|
|
||||||
|
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 with metrics and component health.
|
||||||
|
func (al *AgentLoop) GetSwarmStatus() map[string]any {
|
||||||
|
if !al.swarmEnabled {
|
||||||
|
if al.swarmInitError != nil {
|
||||||
|
return map[string]any{
|
||||||
|
"enabled": false,
|
||||||
|
"error": al.swarmInitError.Error(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return map[string]any{"enabled": false}
|
||||||
|
}
|
||||||
|
|
||||||
|
status := map[string]any{
|
||||||
|
"enabled": true,
|
||||||
|
"node_id": al.swarmDiscovery.LocalNode().ID,
|
||||||
|
"handoff": al.cfg.Swarm.Handoff.Enabled,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Component health status
|
||||||
|
health := map[string]bool{
|
||||||
|
"discovery": al.swarmDiscovery != nil,
|
||||||
|
"handoff": al.swarmHandoff != nil,
|
||||||
|
"load_monitor": al.swarmLoad != nil,
|
||||||
|
"leader_election": al.swarmLeaderElection != nil,
|
||||||
|
}
|
||||||
|
if al.swarmHandoff != nil {
|
||||||
|
for k, v := range al.swarmHandoff.ComponentHealth() {
|
||||||
|
health[k] = v
|
||||||
|
}
|
||||||
|
circuitBreakerStatus := al.swarmHandoff.GetCircuitBreakerStatus()
|
||||||
|
if len(circuitBreakerStatus) > 0 {
|
||||||
|
status["circuit_breaker"] = circuitBreakerStatus
|
||||||
|
}
|
||||||
|
}
|
||||||
|
status["components"] = health
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|
||||||
|
if al.swarmHandoff != nil {
|
||||||
|
status["handoff_metrics"] = al.swarmHandoff.GetMetrics().Snapshot()
|
||||||
|
}
|
||||||
|
|
||||||
|
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.swarmLeaderElection != nil {
|
||||||
|
al.swarmLeaderElection.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),
|
||||||
|
})
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
77
pkg/channels/commands.go
Normal file
77
pkg/channels/commands.go
Normal file
|
|
@ -0,0 +1,77 @@
|
||||||
|
package channels
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"sort"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/bus"
|
||||||
|
)
|
||||||
|
|
||||||
|
// CommandHandler processes a slash command and returns a text response.
|
||||||
|
// args contains everything after the command name (e.g. for "/nodes verbose", args = "verbose").
|
||||||
|
// msg provides the full inbound message context (channel, sender, chat, etc.).
|
||||||
|
type CommandHandler func(ctx context.Context, args string, msg bus.InboundMessage) (string, error)
|
||||||
|
|
||||||
|
// CommandEntry holds a registered command.
|
||||||
|
type CommandEntry struct {
|
||||||
|
Name string // command name without leading slash, e.g. "nodes"
|
||||||
|
Description string // human-readable description, e.g. "List swarm cluster nodes"
|
||||||
|
Handler CommandHandler // function to execute
|
||||||
|
}
|
||||||
|
|
||||||
|
// CommandRegistry is a thread-safe registry of slash commands.
|
||||||
|
// External modules register commands here; the agent loop checks it when
|
||||||
|
// processing inbound messages that start with "/".
|
||||||
|
type CommandRegistry struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
commands map[string]*CommandEntry
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewCommandRegistry creates an empty command registry.
|
||||||
|
func NewCommandRegistry() *CommandRegistry {
|
||||||
|
return &CommandRegistry{
|
||||||
|
commands: make(map[string]*CommandEntry),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Register adds or replaces a command in the registry.
|
||||||
|
// name should NOT include the leading slash.
|
||||||
|
func (r *CommandRegistry) Register(name, description string, handler CommandHandler) {
|
||||||
|
r.mu.Lock()
|
||||||
|
defer r.mu.Unlock()
|
||||||
|
r.commands[name] = &CommandEntry{
|
||||||
|
Name: name,
|
||||||
|
Description: description,
|
||||||
|
Handler: handler,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get looks up a command by name. Returns nil, false if not found.
|
||||||
|
func (r *CommandRegistry) Get(name string) (*CommandEntry, bool) {
|
||||||
|
r.mu.RLock()
|
||||||
|
defer r.mu.RUnlock()
|
||||||
|
entry, ok := r.commands[name]
|
||||||
|
return entry, ok
|
||||||
|
}
|
||||||
|
|
||||||
|
// List returns all registered commands sorted by name.
|
||||||
|
func (r *CommandRegistry) List() []*CommandEntry {
|
||||||
|
r.mu.RLock()
|
||||||
|
defer r.mu.RUnlock()
|
||||||
|
entries := make([]*CommandEntry, 0, len(r.commands))
|
||||||
|
for _, e := range r.commands {
|
||||||
|
entries = append(entries, e)
|
||||||
|
}
|
||||||
|
sort.Slice(entries, func(i, j int) bool {
|
||||||
|
return entries[i].Name < entries[j].Name
|
||||||
|
})
|
||||||
|
return entries
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove unregisters a command by name. No-op if not found.
|
||||||
|
func (r *CommandRegistry) Remove(name string) {
|
||||||
|
r.mu.Lock()
|
||||||
|
defer r.mu.Unlock()
|
||||||
|
delete(r.commands, name)
|
||||||
|
}
|
||||||
|
|
@ -79,12 +79,13 @@ type Manager struct {
|
||||||
bus *bus.MessageBus
|
bus *bus.MessageBus
|
||||||
config *config.Config
|
config *config.Config
|
||||||
mediaStore media.MediaStore
|
mediaStore media.MediaStore
|
||||||
|
commands *CommandRegistry
|
||||||
dispatchTask *asyncTask
|
dispatchTask *asyncTask
|
||||||
mux *http.ServeMux
|
mux *http.ServeMux
|
||||||
httpServer *http.Server
|
httpServer *http.Server
|
||||||
mu sync.RWMutex
|
mu sync.RWMutex
|
||||||
placeholders sync.Map // "channel:chatID" → placeholderID (string)
|
placeholders sync.Map // "channel:chatID" → placeholderEntry
|
||||||
typingStops sync.Map // "channel:chatID" → func()
|
typingStops sync.Map // "channel:chatID" → typingEntry
|
||||||
reactionUndos sync.Map // "channel:chatID" → reactionEntry
|
reactionUndos sync.Map // "channel:chatID" → reactionEntry
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -154,6 +155,7 @@ func NewManager(cfg *config.Config, messageBus *bus.MessageBus, store media.Medi
|
||||||
bus: messageBus,
|
bus: messageBus,
|
||||||
config: cfg,
|
config: cfg,
|
||||||
mediaStore: store,
|
mediaStore: store,
|
||||||
|
commands: NewCommandRegistry(),
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := m.initChannels(); err != nil {
|
if err := m.initChannels(); err != nil {
|
||||||
|
|
@ -786,6 +788,12 @@ func (m *Manager) RegisterChannel(name string, channel Channel) {
|
||||||
m.channels[name] = channel
|
m.channels[name] = channel
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// CommandRegistry returns the shared command registry.
|
||||||
|
// External modules use this to register slash commands that work across all channels.
|
||||||
|
func (m *Manager) CommandRegistry() *CommandRegistry {
|
||||||
|
return m.commands
|
||||||
|
}
|
||||||
|
|
||||||
func (m *Manager) UnregisterChannel(name string) {
|
func (m *Manager) UnregisterChannel(name string) {
|
||||||
m.mu.Lock()
|
m.mu.Lock()
|
||||||
defer m.mu.Unlock()
|
defer m.mu.Unlock()
|
||||||
|
|
|
||||||
|
|
@ -56,6 +56,8 @@ type PicoChannel struct {
|
||||||
connCount atomic.Int32
|
connCount atomic.Int32
|
||||||
ctx context.Context
|
ctx context.Context
|
||||||
cancel context.CancelFunc
|
cancel context.CancelFunc
|
||||||
|
nodeRequestHandler func(payload map[string]any) (map[string]any, error)
|
||||||
|
nodeValidator func(sourceNodeID string) bool // validates source is a known swarm member
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewPicoChannel creates a new Pico Protocol channel.
|
// NewPicoChannel creates a new Pico Protocol channel.
|
||||||
|
|
@ -403,6 +405,9 @@ func (c *PicoChannel) handleMessage(pc *picoConn, msg PicoMessage) {
|
||||||
case TypeMessageSend:
|
case TypeMessageSend:
|
||||||
c.handleMessageSend(pc, msg)
|
c.handleMessageSend(pc, msg)
|
||||||
|
|
||||||
|
case TypeNodeRequest:
|
||||||
|
c.handleNodeRequest(pc, msg)
|
||||||
|
|
||||||
default:
|
default:
|
||||||
errMsg := newError("unknown_type", fmt.Sprintf("unknown message type: %s", msg.Type))
|
errMsg := newError("unknown_type", fmt.Sprintf("unknown message type: %s", msg.Type))
|
||||||
pc.writeJSON(errMsg)
|
pc.writeJSON(errMsg)
|
||||||
|
|
@ -452,6 +457,119 @@ func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) {
|
||||||
c.HandleMessage(c.ctx, peer, msg.ID, senderID, chatID, content, nil, metadata, sender)
|
c.HandleMessage(c.ctx, peer, msg.ID, senderID, chatID, content, nil, metadata, sender)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetNodeRequestHandler sets the handler for incoming inter-node request messages.
|
||||||
|
// The handler receives the full payload map and returns a reply payload map.
|
||||||
|
func (c *PicoChannel) SetNodeRequestHandler(handler func(payload map[string]any) (map[string]any, error)) {
|
||||||
|
c.nodeRequestHandler = handler
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetNodeValidator sets a function that validates whether a source_node_id
|
||||||
|
// is a known, alive member of the swarm cluster. This prevents arbitrary
|
||||||
|
// WebSocket clients from issuing node.request messages even if they hold
|
||||||
|
// a valid Pico token.
|
||||||
|
func (c *PicoChannel) SetNodeValidator(validator func(sourceNodeID string) bool) {
|
||||||
|
c.nodeValidator = validator
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleNodeRequest processes an incoming node.request message from a swarm peer.
|
||||||
|
//
|
||||||
|
// Security model:
|
||||||
|
// - The WebSocket connection is already authenticated via Pico token.
|
||||||
|
// - If a nodeValidator is set, source_node_id MUST be a known alive swarm member.
|
||||||
|
// This prevents non-swarm WS clients from reaching inter-node control plane.
|
||||||
|
// - For production, also configure NATS-level ACLs as the primary trust boundary;
|
||||||
|
// this Pico path is a secondary/legacy channel.
|
||||||
|
func (c *PicoChannel) handleNodeRequest(pc *picoConn, msg PicoMessage) {
|
||||||
|
requestID, _ := msg.Payload["request_id"].(string)
|
||||||
|
sourceNodeID, _ := msg.Payload["source_node_id"].(string)
|
||||||
|
action, _ := msg.Payload["action"].(string)
|
||||||
|
|
||||||
|
logger.InfoCF("pico", "Received node request", map[string]any{
|
||||||
|
"request_id": requestID,
|
||||||
|
"source_node_id": sourceNodeID,
|
||||||
|
"action": action,
|
||||||
|
})
|
||||||
|
|
||||||
|
// Validate source node identity against swarm membership.
|
||||||
|
if c.nodeValidator != nil && !c.nodeValidator(sourceNodeID) {
|
||||||
|
logger.WarnCF("pico", "Rejected node request from unknown/dead node", map[string]any{
|
||||||
|
"request_id": requestID,
|
||||||
|
"source_node_id": sourceNodeID,
|
||||||
|
})
|
||||||
|
reply := newMessage(TypeNodeReply, map[string]any{
|
||||||
|
"request_id": requestID,
|
||||||
|
"error": "source node not recognized as a cluster member",
|
||||||
|
})
|
||||||
|
reply.ID = msg.ID
|
||||||
|
_ = pc.writeJSON(reply)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var replyPayload map[string]any
|
||||||
|
|
||||||
|
if c.nodeRequestHandler == nil {
|
||||||
|
replyPayload = map[string]any{
|
||||||
|
"request_id": requestID,
|
||||||
|
"error": "no node request handler registered",
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Start watchdog goroutine to send processing heartbeats
|
||||||
|
stopHeartbeat := make(chan struct{})
|
||||||
|
go c.watchDog(pc, msg.ID, requestID, stopHeartbeat)
|
||||||
|
|
||||||
|
// Call the handler and stop heartbeat when done
|
||||||
|
var err error
|
||||||
|
replyPayload, err = c.nodeRequestHandler(msg.Payload)
|
||||||
|
close(stopHeartbeat)
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
replyPayload = map[string]any{
|
||||||
|
"request_id": requestID,
|
||||||
|
"error": err.Error(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Ensure request_id is always in the reply
|
||||||
|
if replyPayload != nil {
|
||||||
|
replyPayload["request_id"] = requestID
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
reply := newMessage(TypeNodeReply, replyPayload)
|
||||||
|
reply.ID = msg.ID
|
||||||
|
if err := pc.writeJSON(reply); err != nil {
|
||||||
|
logger.ErrorCF("pico", "Failed to send node reply", map[string]any{
|
||||||
|
"error": err.Error(),
|
||||||
|
"request_id": requestID,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// watchDog sends periodic processing heartbeats to keep the connection alive
|
||||||
|
// while a long-running node request is being handled. It stops when the
|
||||||
|
// stopHeartbeat channel is closed.
|
||||||
|
func (c *PicoChannel) watchDog(pc *picoConn, msgID, requestID string, stop <-chan struct{}) {
|
||||||
|
ticker := time.NewTicker(nodeHeartbeatInterval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-stop:
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
heartbeat := newMessage(TypeNodeProcessing, map[string]any{
|
||||||
|
"request_id": requestID,
|
||||||
|
})
|
||||||
|
heartbeat.ID = msgID
|
||||||
|
if err := pc.writeJSON(heartbeat); err != nil {
|
||||||
|
logger.DebugCF("pico", "Failed to send processing heartbeat", map[string]any{
|
||||||
|
"error": err.Error(),
|
||||||
|
"request_id": requestID,
|
||||||
|
})
|
||||||
|
// Don't return - keep trying until stop is closed
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// truncate truncates a string to maxLen runes.
|
// truncate truncates a string to maxLen runes.
|
||||||
func truncate(s string, maxLen int) string {
|
func truncate(s string, maxLen int) string {
|
||||||
runes := []rune(s)
|
runes := []rune(s)
|
||||||
|
|
|
||||||
|
|
@ -1,46 +1,59 @@
|
||||||
|
// Package pico provides the Pico Protocol WebSocket channel implementation.
|
||||||
|
// This file contains the local PicoMessage type wrapper for compatibility with
|
||||||
|
// the canonical protocol definitions in pkg/pico/protocol.
|
||||||
package pico
|
package pico
|
||||||
|
|
||||||
import "time"
|
import (
|
||||||
|
"time"
|
||||||
|
|
||||||
// Protocol message types.
|
picoproto "github.com/sipeed/picoclaw/pkg/pico/protocol"
|
||||||
const (
|
|
||||||
// TypeMessageSend is sent from client to server.
|
|
||||||
TypeMessageSend = "message.send"
|
|
||||||
TypeMediaSend = "media.send"
|
|
||||||
TypePing = "ping"
|
|
||||||
|
|
||||||
// TypeMessageCreate is sent from server to client.
|
|
||||||
TypeMessageCreate = "message.create"
|
|
||||||
TypeMessageUpdate = "message.update"
|
|
||||||
TypeMediaCreate = "media.create"
|
|
||||||
TypeTypingStart = "typing.start"
|
|
||||||
TypeTypingStop = "typing.stop"
|
|
||||||
TypeError = "error"
|
|
||||||
TypePong = "pong"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// nodeHeartbeatInterval is the interval between processing heartbeats sent
|
||||||
|
// during long-running node requests to keep the connection alive.
|
||||||
|
const nodeHeartbeatInterval = 15 * time.Second
|
||||||
|
|
||||||
// PicoMessage is the wire format for all Pico Protocol messages.
|
// PicoMessage is the wire format for all Pico Protocol messages.
|
||||||
type PicoMessage struct {
|
// This is a local type alias for the canonical protocol.Message type.
|
||||||
Type string `json:"type"`
|
type PicoMessage = picoproto.Message
|
||||||
ID string `json:"id,omitempty"`
|
|
||||||
SessionID string `json:"session_id,omitempty"`
|
|
||||||
Timestamp int64 `json:"timestamp,omitempty"`
|
|
||||||
Payload map[string]any `json:"payload,omitempty"`
|
|
||||||
}
|
|
||||||
|
|
||||||
// newMessage creates a PicoMessage with the given type and payload.
|
// newMessage creates a PicoMessage with the given type and payload.
|
||||||
func newMessage(msgType string, payload map[string]any) PicoMessage {
|
func newMessage(msgType string, payload map[string]any) PicoMessage {
|
||||||
return PicoMessage{
|
return picoproto.NewMessage(msgType, payload)
|
||||||
Type: msgType,
|
|
||||||
Timestamp: time.Now().UnixMilli(),
|
|
||||||
Payload: payload,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// newError creates an error PicoMessage.
|
// newError creates an error PicoMessage.
|
||||||
func newError(code, message string) PicoMessage {
|
func newError(code, message string) PicoMessage {
|
||||||
return newMessage(TypeError, map[string]any{
|
return picoproto.NewError(code, message)
|
||||||
"code": code,
|
}
|
||||||
"message": message,
|
|
||||||
})
|
// Message type constants re-exported from pkg/pico/protocol for convenience.
|
||||||
|
// The canonical definitions are in pkg/pico/protocol/protocol.go.
|
||||||
|
const (
|
||||||
|
TypeMessageSend = picoproto.TypeMessageSend
|
||||||
|
TypeMediaSend = picoproto.TypeMediaSend
|
||||||
|
TypePing = picoproto.TypePing
|
||||||
|
TypeMessageCreate = picoproto.TypeMessageCreate
|
||||||
|
TypeMessageUpdate = picoproto.TypeMessageUpdate
|
||||||
|
TypeMediaCreate = picoproto.TypeMediaCreate
|
||||||
|
TypeTypingStart = picoproto.TypeTypingStart
|
||||||
|
TypeTypingStop = picoproto.TypeTypingStop
|
||||||
|
TypeError = picoproto.TypeError
|
||||||
|
TypePong = picoproto.TypePong
|
||||||
|
TypeNodeRequest = picoproto.TypeNodeRequest
|
||||||
|
TypeNodeReply = picoproto.TypeNodeReply
|
||||||
|
TypeNodeProcessing = picoproto.TypeNodeProcessing
|
||||||
|
)
|
||||||
|
|
||||||
|
// Ensure PicoMessage matches protocol.Message at compile time.
|
||||||
|
var _ = func() struct{} {
|
||||||
|
// Type alias compatibility check
|
||||||
|
type _ = picoproto.Message
|
||||||
|
type _ = PicoMessage
|
||||||
|
return struct{}{}
|
||||||
|
}()
|
||||||
|
|
||||||
|
// PicoMessageTimestamp returns the current timestamp for Pico messages.
|
||||||
|
func PicoMessageTimestamp() int64 {
|
||||||
|
return time.Now().UnixMilli()
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -58,6 +58,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
|
||||||
|
|
@ -900,3 +901,134 @@ func (t *ToolsConfig) IsToolEnabled(name string) bool {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 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"`
|
||||||
|
|
||||||
|
// Discovery configuration for NATS-based node discovery.
|
||||||
|
Discovery SwarmDiscoveryConfig `json:"discovery"`
|
||||||
|
|
||||||
|
// Handoff configuration for task handoff.
|
||||||
|
Handoff SwarmHandoffConfig `json:"handoff"`
|
||||||
|
|
||||||
|
// LoadMonitor configuration for load monitoring.
|
||||||
|
LoadMonitor SwarmLoadMonitorConfig `json:"load_monitor"`
|
||||||
|
|
||||||
|
// LeaderElection configuration for leader election.
|
||||||
|
LeaderElection SwarmLeaderElectionConfig `json:"leader_election"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// SwarmDiscoveryConfig contains configuration for NATS-based node discovery.
|
||||||
|
type SwarmDiscoveryConfig struct {
|
||||||
|
// NATSURL is the NATS server URL (required).
|
||||||
|
NATSURL string `json:"nats_url,omitempty" env:"PICOCLAW_SWARM_NATS_URL"`
|
||||||
|
|
||||||
|
// NATSCredsFile is the path to NATS credentials file (NKeys/JWT auth).
|
||||||
|
NATSCredsFile string `json:"nats_creds_file,omitempty"`
|
||||||
|
|
||||||
|
// NATSTLSCert is the path to the client TLS certificate for mTLS.
|
||||||
|
NATSTLSCert string `json:"nats_tls_cert,omitempty"`
|
||||||
|
|
||||||
|
// NATSTLSKey is the path to the client TLS key for mTLS.
|
||||||
|
NATSTLSKey string `json:"nats_tls_key,omitempty"`
|
||||||
|
|
||||||
|
// NATSTLSCACert is the path to the CA certificate for TLS verification.
|
||||||
|
NATSTLSCACert string `json:"nats_tls_ca_cert,omitempty"`
|
||||||
|
|
||||||
|
// HeartbeatInterval controls how often the node renews its KV entry (in seconds).
|
||||||
|
HeartbeatInterval int `json:"heartbeat_interval,omitempty"`
|
||||||
|
|
||||||
|
// MemberTTL is the TTL for member entries in the KV bucket (in seconds).
|
||||||
|
MemberTTL int `json:"member_ttl,omitempty"`
|
||||||
|
|
||||||
|
// SubjectPrefix overrides the default "picoclaw.swarm" subject prefix.
|
||||||
|
SubjectPrefix string `json:"subject_prefix,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"`
|
||||||
|
|
||||||
|
// RequestTimeout is the timeout for a single NATS handoff request-reply (in seconds).
|
||||||
|
RequestTimeout int `json:"request_timeout,omitempty"`
|
||||||
|
|
||||||
|
// SharedSecret is the HMAC secret for message signing (base64 encoded or raw string).
|
||||||
|
// If empty, signing is disabled (NOT RECOMMENDED for production).
|
||||||
|
SharedSecret string `json:"shared_secret,omitempty" env:"PICOCLAW_SWARM_SHARED_SECRET"`
|
||||||
|
|
||||||
|
// EncryptionKey is the AES-256 key for encrypting session data (base64 encoded or raw string).
|
||||||
|
// Must be 32 bytes when decoded. If empty, encryption is disabled.
|
||||||
|
EncryptionKey string `json:"encryption_key,omitempty" env:"PICOCLAW_SWARM_ENCRYPTION_KEY"`
|
||||||
|
|
||||||
|
// RequireAuth enables strict authentication - reject unsigned messages.
|
||||||
|
RequireAuth bool `json:"require_auth"`
|
||||||
|
|
||||||
|
// RequireEncryption enables strict encryption - reject unencrypted session data.
|
||||||
|
RequireEncryption bool `json:"require_encryption"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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"`
|
||||||
|
|
||||||
|
// RoutingRejectThreshold is the load score above which routing requests are rejected (0-1).
|
||||||
|
// Provides hysteresis to prevent oscillation. Default: 0.9
|
||||||
|
RoutingRejectThreshold float64 `json:"routing_reject_threshold,omitempty"`
|
||||||
|
|
||||||
|
// RoutingAcceptThreshold is the load score below which routing requests are accepted (0-1).
|
||||||
|
// Once rejected, load must drop below this to accept again. Default: 0.75
|
||||||
|
RoutingAcceptThreshold float64 `json:"routing_accept_threshold,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// SwarmLeaderElectionConfig contains configuration for leader election.
|
||||||
|
type SwarmLeaderElectionConfig struct {
|
||||||
|
// Enabled enables leader election.
|
||||||
|
Enabled bool `json:"enabled"`
|
||||||
|
|
||||||
|
// LockTTL is how long the leader lock lives without renewal (in seconds).
|
||||||
|
LockTTL int `json:"lock_ttl,omitempty"`
|
||||||
|
|
||||||
|
// RenewalInterval is how often the leader renews its lock (in seconds).
|
||||||
|
RenewalInterval int `json:"renewal_interval,omitempty"`
|
||||||
|
}
|
||||||
|
|
|
||||||
115
pkg/kv/bucket.go
Normal file
115
pkg/kv/bucket.go
Normal file
|
|
@ -0,0 +1,115 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Generic Key-Value storage abstraction
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
// Package kv provides a generic key-value storage abstraction.
|
||||||
|
// It supports multiple backends (NATS JetStream, etcd, Consul, etc.)
|
||||||
|
// with CAS (Compare-And-Swap) operations and change watching.
|
||||||
|
package kv
|
||||||
|
|
||||||
|
import "time"
|
||||||
|
|
||||||
|
// BucketSpec defines the configuration for a bucket.
|
||||||
|
type BucketSpec struct {
|
||||||
|
Name string // Bucket name
|
||||||
|
Description string // Human-readable description
|
||||||
|
TTL time.Duration // TTL for entries
|
||||||
|
MaxValueSize int32 // Maximum value size
|
||||||
|
Storage StorageType // Storage backend
|
||||||
|
}
|
||||||
|
|
||||||
|
// StorageType defines the storage backend for a bucket.
|
||||||
|
type StorageType string
|
||||||
|
|
||||||
|
const (
|
||||||
|
StorageFile StorageType = "file"
|
||||||
|
StorageMemory StorageType = "memory"
|
||||||
|
)
|
||||||
|
|
||||||
|
// BucketEntry represents a key-value entry in a bucket.
|
||||||
|
type BucketEntry struct {
|
||||||
|
Key string // The key
|
||||||
|
Value []byte // The value
|
||||||
|
Revision uint64 // Revision number (for CAS)
|
||||||
|
Operation EntryOperation // The operation that triggered this entry
|
||||||
|
}
|
||||||
|
|
||||||
|
// EntryOperation represents the type of operation on an entry.
|
||||||
|
type EntryOperation string
|
||||||
|
|
||||||
|
const (
|
||||||
|
EntryOperationPut EntryOperation = "put"
|
||||||
|
EntryOperationDelete EntryOperation = "delete"
|
||||||
|
EntryOperationPurge EntryOperation = "purge"
|
||||||
|
)
|
||||||
|
|
||||||
|
// BucketWatcher watches for changes in a bucket.
|
||||||
|
type BucketWatcher interface {
|
||||||
|
// Updates returns a channel of bucket entry updates.
|
||||||
|
// The channel is closed when Stop() is called.
|
||||||
|
Updates() <-chan *BucketEntry
|
||||||
|
|
||||||
|
// Stop stops watching and closes the updates channel.
|
||||||
|
Stop()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Bucket provides key-value storage with watch capabilities.
|
||||||
|
//
|
||||||
|
// This interface abstracts KV operations for different storage backends.
|
||||||
|
// Implementations can use NATS JetStream KV, etcd, Consul, Redis, etc.
|
||||||
|
type Bucket interface {
|
||||||
|
// Name returns the bucket name.
|
||||||
|
Name() string
|
||||||
|
|
||||||
|
// Get retrieves a value by key.
|
||||||
|
// Returns nil if key does not exist.
|
||||||
|
Get(key string) (*BucketEntry, error)
|
||||||
|
|
||||||
|
// Put stores a value under the given key (creates or updates).
|
||||||
|
Put(key string, value []byte) error
|
||||||
|
|
||||||
|
// Create stores a value only if the key does not exist.
|
||||||
|
// Returns the revision number of the new entry.
|
||||||
|
// Returns an error if the key already exists.
|
||||||
|
Create(key string, value []byte) (uint64, error)
|
||||||
|
|
||||||
|
// Update updates a value with CAS (Compare-And-Swap) semantics.
|
||||||
|
// The operation only succeeds if the entry's revision matches the expected revision.
|
||||||
|
// Returns the new revision number.
|
||||||
|
// Returns an error if the revision doesn't match (CAS failure).
|
||||||
|
Update(key string, value []byte, revision uint64) (uint64, error)
|
||||||
|
|
||||||
|
// Delete removes a key.
|
||||||
|
Delete(key string) error
|
||||||
|
|
||||||
|
// Keys returns all keys in the bucket.
|
||||||
|
Keys() ([]string, error)
|
||||||
|
|
||||||
|
// Watch starts watching for changes in the bucket.
|
||||||
|
// If keyPrefix is provided, only watch keys with that prefix.
|
||||||
|
// Returns a BucketWatcher that can be stopped.
|
||||||
|
Watch(keyPrefix string) (BucketWatcher, error)
|
||||||
|
|
||||||
|
// WatchAll watches all keys in the bucket.
|
||||||
|
WatchAll() (BucketWatcher, error)
|
||||||
|
|
||||||
|
// Close closes the bucket and releases resources.
|
||||||
|
Close() error
|
||||||
|
}
|
||||||
|
|
||||||
|
// BucketFactory creates buckets of different types.
|
||||||
|
type BucketFactory interface {
|
||||||
|
// CreateBucket creates a new bucket with the given specification.
|
||||||
|
CreateBucket(spec BucketSpec) (Bucket, error)
|
||||||
|
|
||||||
|
// GetBucket returns an existing bucket by name.
|
||||||
|
GetBucket(name string) (Bucket, error)
|
||||||
|
|
||||||
|
// CreateIfNotExists creates a bucket if it doesn't exist, or returns the existing one.
|
||||||
|
CreateIfNotExists(spec BucketSpec) (Bucket, error)
|
||||||
|
|
||||||
|
// Close closes all buckets and releases resources.
|
||||||
|
Close() error
|
||||||
|
}
|
||||||
253
pkg/kv/nats.go
Normal file
253
pkg/kv/nats.go
Normal file
|
|
@ -0,0 +1,253 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Generic Key-Value storage abstraction - NATS JetStream implementation
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package kv
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"github.com/nats-io/nats.go"
|
||||||
|
)
|
||||||
|
|
||||||
|
// NATSBucketFactory implements BucketFactory using NATS JetStream KV.
|
||||||
|
type NATSBucketFactory struct {
|
||||||
|
js nats.JetStreamContext
|
||||||
|
buckets map[string]Bucket
|
||||||
|
mu sync.RWMutex
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewNATSBucketFactory creates a new NATS-backed bucket factory.
|
||||||
|
func NewNATSBucketFactory(js nats.JetStreamContext) *NATSBucketFactory {
|
||||||
|
return &NATSBucketFactory{
|
||||||
|
js: js,
|
||||||
|
buckets: make(map[string]Bucket),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateBucket creates a new bucket with the given specification.
|
||||||
|
func (nf *NATSBucketFactory) CreateBucket(spec BucketSpec) (Bucket, error) {
|
||||||
|
nf.mu.Lock()
|
||||||
|
defer nf.mu.Unlock()
|
||||||
|
|
||||||
|
// Check if already exists
|
||||||
|
if _, exists := nf.buckets[spec.Name]; exists {
|
||||||
|
return nil, fmt.Errorf("bucket %s already exists", spec.Name)
|
||||||
|
}
|
||||||
|
|
||||||
|
storage := nats.FileStorage
|
||||||
|
if spec.Storage == StorageMemory {
|
||||||
|
storage = nats.MemoryStorage
|
||||||
|
}
|
||||||
|
|
||||||
|
kv, err := nf.js.CreateKeyValue(&nats.KeyValueConfig{
|
||||||
|
Bucket: spec.Name,
|
||||||
|
Description: spec.Description,
|
||||||
|
MaxValueSize: spec.MaxValueSize,
|
||||||
|
TTL: spec.TTL,
|
||||||
|
Storage: storage,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create KV bucket %s: %w", spec.Name, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
bucket := &natsBucket{
|
||||||
|
kv: kv,
|
||||||
|
spec: spec,
|
||||||
|
}
|
||||||
|
nf.buckets[spec.Name] = bucket
|
||||||
|
|
||||||
|
return bucket, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetBucket returns an existing bucket by name.
|
||||||
|
func (nf *NATSBucketFactory) GetBucket(name string) (Bucket, error) {
|
||||||
|
nf.mu.RLock()
|
||||||
|
defer nf.mu.RUnlock()
|
||||||
|
|
||||||
|
if b, exists := nf.buckets[name]; exists {
|
||||||
|
return b, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to load from NATS
|
||||||
|
kv, err := nf.js.KeyValue(name)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("bucket %s not found: %w", name, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
bucket := &natsBucket{
|
||||||
|
kv: kv,
|
||||||
|
spec: BucketSpec{Name: name},
|
||||||
|
}
|
||||||
|
nf.buckets[name] = bucket
|
||||||
|
|
||||||
|
return bucket, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// CreateIfNotExists creates a bucket if it doesn't exist.
|
||||||
|
func (nf *NATSBucketFactory) CreateIfNotExists(spec BucketSpec) (Bucket, error) {
|
||||||
|
b, err := nf.GetBucket(spec.Name)
|
||||||
|
if err == nil {
|
||||||
|
return b, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return nf.CreateBucket(spec)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close closes all buckets.
|
||||||
|
func (nf *NATSBucketFactory) Close() error {
|
||||||
|
nf.mu.Lock()
|
||||||
|
defer nf.mu.Unlock()
|
||||||
|
|
||||||
|
var lastErr error
|
||||||
|
for name := range nf.buckets {
|
||||||
|
delete(nf.buckets, name)
|
||||||
|
}
|
||||||
|
|
||||||
|
return lastErr
|
||||||
|
}
|
||||||
|
|
||||||
|
// natsBucket implements Bucket using NATS JetStream KV.
|
||||||
|
type natsBucket struct {
|
||||||
|
kv nats.KeyValue
|
||||||
|
spec BucketSpec
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name returns the bucket name.
|
||||||
|
func (b *natsBucket) Name() string {
|
||||||
|
return b.spec.Name
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get retrieves a value by key.
|
||||||
|
func (b *natsBucket) Get(key string) (*BucketEntry, error) {
|
||||||
|
entry, err := b.kv.Get(key)
|
||||||
|
if err != nil {
|
||||||
|
if err == nats.ErrKeyNotFound {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return &BucketEntry{
|
||||||
|
Key: entry.Key(),
|
||||||
|
Value: entry.Value(),
|
||||||
|
Revision: entry.Revision(),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Put stores a value under the given key.
|
||||||
|
func (b *natsBucket) Put(key string, value []byte) error {
|
||||||
|
_, err := b.kv.Put(key, value)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create puts a key only if it doesn't exist (Create operation).
|
||||||
|
func (b *natsBucket) Create(key string, value []byte) (uint64, error) {
|
||||||
|
return b.kv.Create(key, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update updates a key with CAS (Compare-And-Swap).
|
||||||
|
func (b *natsBucket) Update(key string, value []byte, revision uint64) (uint64, error) {
|
||||||
|
return b.kv.Update(key, value, revision)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Delete removes a key.
|
||||||
|
func (b *natsBucket) Delete(key string) error {
|
||||||
|
return b.kv.Delete(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Keys returns all keys in the bucket.
|
||||||
|
func (b *natsBucket) Keys() ([]string, error) {
|
||||||
|
return b.kv.Keys()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Watch starts watching for changes in the bucket.
|
||||||
|
func (b *natsBucket) Watch(keyPrefix string) (BucketWatcher, error) {
|
||||||
|
watcher, err := b.kv.Watch(keyPrefix)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return &natsBucketWatcher{watcher: watcher}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// WatchAll watches all keys in the bucket.
|
||||||
|
func (b *natsBucket) WatchAll() (BucketWatcher, error) {
|
||||||
|
watcher, err := b.kv.WatchAll()
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return &natsBucketWatcher{watcher: watcher}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close closes the bucket.
|
||||||
|
func (b *natsBucket) Close() error {
|
||||||
|
// NATS KV buckets are managed by JetStream, not by individual holders
|
||||||
|
// We don't actually delete the bucket, just release our reference
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// natsBucketWatcher implements BucketWatcher using NATS KV watcher.
|
||||||
|
type natsBucketWatcher struct {
|
||||||
|
watcher nats.KeyWatcher
|
||||||
|
updates chan *BucketEntry
|
||||||
|
once sync.Once
|
||||||
|
}
|
||||||
|
|
||||||
|
// Updates returns a channel of bucket entry updates.
|
||||||
|
func (w *natsBucketWatcher) Updates() <-chan *BucketEntry {
|
||||||
|
if w.updates == nil {
|
||||||
|
w.updates = make(chan *BucketEntry, 32)
|
||||||
|
go w.forwardUpdates()
|
||||||
|
}
|
||||||
|
return w.updates
|
||||||
|
}
|
||||||
|
|
||||||
|
// forwardUpdates forwards NATS watcher updates to our channel.
|
||||||
|
func (w *natsBucketWatcher) forwardUpdates() {
|
||||||
|
for entry := range w.watcher.Updates() {
|
||||||
|
if entry == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var op EntryOperation
|
||||||
|
switch entry.Operation() {
|
||||||
|
case nats.KeyValuePut:
|
||||||
|
op = EntryOperationPut
|
||||||
|
case nats.KeyValueDelete:
|
||||||
|
op = EntryOperationDelete
|
||||||
|
case nats.KeyValuePurge:
|
||||||
|
op = EntryOperationPurge
|
||||||
|
default:
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
bEntry := &BucketEntry{
|
||||||
|
Key: entry.Key(),
|
||||||
|
Value: entry.Value(),
|
||||||
|
Revision: entry.Revision(),
|
||||||
|
Operation: op,
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case w.updates <- bEntry:
|
||||||
|
default:
|
||||||
|
// Channel full, drop the update
|
||||||
|
// In production, you might want to handle this differently
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop stops watching.
|
||||||
|
func (w *natsBucketWatcher) Stop() {
|
||||||
|
w.once.Do(func() {
|
||||||
|
w.watcher.Stop()
|
||||||
|
if w.updates != nil {
|
||||||
|
close(w.updates)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
22
pkg/pico/addr.go
Normal file
22
pkg/pico/addr.go
Normal file
|
|
@ -0,0 +1,22 @@
|
||||||
|
// Package pico provides a reusable WebSocket client for the Pico Protocol.
|
||||||
|
//
|
||||||
|
// In addition to the low-level Client, this package defines shared types and
|
||||||
|
// helpers used by both the Pico channel (server) and the swarm subsystem
|
||||||
|
// (client) for inter-node communication.
|
||||||
|
//
|
||||||
|
// The Client type encapsulates the connect → send → receive → close lifecycle
|
||||||
|
// for a single request-reply exchange with a Pico WebSocket endpoint. It is
|
||||||
|
// intentionally stateless: each call to SendRequest opens a new connection,
|
||||||
|
// performs the exchange, and closes the connection.
|
||||||
|
//
|
||||||
|
// This package depends only on pkg/pico/protocol and gorilla/websocket;
|
||||||
|
// it has no knowledge of swarm, channels, or any other higher-level construct.
|
||||||
|
package pico
|
||||||
|
|
||||||
|
import "fmt"
|
||||||
|
|
||||||
|
// BuildNodeAddr constructs a host:port address string for inter-node
|
||||||
|
// communication. It replaces the scattered fmt.Sprintf("%s:%d", ...) calls.
|
||||||
|
func BuildNodeAddr(addr string, port int) string {
|
||||||
|
return fmt.Sprintf("%s:%d", addr, port)
|
||||||
|
}
|
||||||
107
pkg/pico/client.go
Normal file
107
pkg/pico/client.go
Normal file
|
|
@ -0,0 +1,107 @@
|
||||||
|
package pico
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/gorilla/websocket"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/pico/protocol"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DefaultReadTimeout is the maximum time to wait for a reply after sending.
|
||||||
|
// For inter-node communication, LLM responses may take longer than typical HTTP requests.
|
||||||
|
const DefaultReadTimeout = 2 * time.Minute
|
||||||
|
|
||||||
|
// Client is a lightweight, stateless Pico WebSocket client.
|
||||||
|
// Each SendRequest call dials a new connection, performs a single
|
||||||
|
// request-reply exchange, and closes the connection.
|
||||||
|
type Client struct {
|
||||||
|
token string
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewClient creates a new Pico WebSocket client.
|
||||||
|
// If token is non-empty it is sent as a Bearer token in the upgrade request.
|
||||||
|
func NewClient(token string) *Client {
|
||||||
|
return &Client{token: token}
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildWSURL constructs the canonical Pico WebSocket URL for a given
|
||||||
|
// host address and session ID.
|
||||||
|
func BuildWSURL(addr, sessionID string) string {
|
||||||
|
return fmt.Sprintf("ws://%s/pico/ws?session_id=%s", addr, sessionID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// SendRequest dials the target Pico WebSocket endpoint, sends msg, and blocks
|
||||||
|
// until a single reply message is received (or the context / read timeout fires).
|
||||||
|
//
|
||||||
|
// The caller is responsible for constructing the outbound protocol.Message
|
||||||
|
// (including Type, ID, Payload, etc.) and for interpreting the reply.
|
||||||
|
//
|
||||||
|
// For long-running requests, the server may send TypeNodeProcessing messages
|
||||||
|
// to keep the connection alive. The client resets the read timeout on each
|
||||||
|
// processing message until the final reply is received.
|
||||||
|
func (c *Client) SendRequest(
|
||||||
|
ctx context.Context,
|
||||||
|
addr, sessionID string,
|
||||||
|
msg protocol.Message,
|
||||||
|
) (protocol.Message, error) {
|
||||||
|
wsURL := BuildWSURL(addr, sessionID)
|
||||||
|
|
||||||
|
header := http.Header{}
|
||||||
|
if c.token != "" {
|
||||||
|
header.Set("Authorization", "Bearer "+c.token)
|
||||||
|
}
|
||||||
|
|
||||||
|
dialer := websocket.Dialer{
|
||||||
|
ReadBufferSize: 1024,
|
||||||
|
WriteBufferSize: 1024,
|
||||||
|
}
|
||||||
|
|
||||||
|
conn, resp, err := dialer.DialContext(ctx, wsURL, header)
|
||||||
|
if err != nil {
|
||||||
|
if resp != nil && resp.Body != nil {
|
||||||
|
resp.Body.Close()
|
||||||
|
return protocol.Message{}, fmt.Errorf(
|
||||||
|
"pico WebSocket dial failed: %s (status: %d)",
|
||||||
|
err.Error(), resp.StatusCode,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
return protocol.Message{}, fmt.Errorf("pico WebSocket dial failed: %w", err)
|
||||||
|
}
|
||||||
|
if resp != nil && resp.Body != nil {
|
||||||
|
resp.Body.Close()
|
||||||
|
}
|
||||||
|
defer conn.Close()
|
||||||
|
|
||||||
|
if writeErr := conn.WriteJSON(msg); writeErr != nil {
|
||||||
|
return protocol.Message{}, fmt.Errorf("failed to send pico request: %w", writeErr)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Read loop that handles processing heartbeat messages
|
||||||
|
for {
|
||||||
|
// Set read deadline before each read
|
||||||
|
conn.SetReadDeadline(time.Now().Add(DefaultReadTimeout))
|
||||||
|
|
||||||
|
_, rawMsg, err := conn.ReadMessage()
|
||||||
|
if err != nil {
|
||||||
|
return protocol.Message{}, fmt.Errorf("failed to read pico reply: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var reply protocol.Message
|
||||||
|
if err := json.Unmarshal(rawMsg, &reply); err != nil {
|
||||||
|
return protocol.Message{}, fmt.Errorf("failed to parse pico reply: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Handle processing heartbeat - reset timeout and continue waiting
|
||||||
|
if reply.Type == protocol.TypeNodeProcessing {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Return on final reply or error
|
||||||
|
return reply, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
123
pkg/pico/node_client.go
Normal file
123
pkg/pico/node_client.go
Normal file
|
|
@ -0,0 +1,123 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package pico
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/pico/protocol"
|
||||||
|
)
|
||||||
|
|
||||||
|
// PicoNodeClient sends inter-node messages by connecting to a target node's
|
||||||
|
// Pico WebSocket endpoint. It delegates the low-level WebSocket lifecycle to
|
||||||
|
// Client and focuses on swarm-specific payload construction.
|
||||||
|
type PicoNodeClient struct {
|
||||||
|
sourceNodeID string
|
||||||
|
client *Client
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewPicoNodeClient creates a new Pico-based inter-node client.
|
||||||
|
func NewPicoNodeClient(sourceNodeID, token string) *PicoNodeClient {
|
||||||
|
return &PicoNodeClient{
|
||||||
|
sourceNodeID: sourceNodeID,
|
||||||
|
client: NewClient(token),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SendMessage connects to the target node's Pico WebSocket, sends a node.request
|
||||||
|
// with action "message", and blocks until a node.reply is received (or timeout).
|
||||||
|
func (c *PicoNodeClient) SendMessage(
|
||||||
|
ctx context.Context,
|
||||||
|
targetAddr, targetNodeID, content, channel, chatID, senderID string,
|
||||||
|
) (string, error) {
|
||||||
|
payload := NewMessagePayload(c.sourceNodeID, content, channel, chatID, senderID)
|
||||||
|
return c.sendRequest(ctx, targetAddr, targetNodeID, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// SendNodeAction sends an action-based request to a target node via Pico.
|
||||||
|
// The payload must contain an "action" key. Returns the raw reply payload.
|
||||||
|
func (c *PicoNodeClient) SendNodeAction(
|
||||||
|
ctx context.Context,
|
||||||
|
targetAddr string,
|
||||||
|
payload NodePayload,
|
||||||
|
) (NodePayload, error) {
|
||||||
|
requestID := uuid.New().String()
|
||||||
|
payload[PayloadKeySourceNodeID] = c.sourceNodeID
|
||||||
|
payload[PayloadKeyRequestID] = requestID
|
||||||
|
|
||||||
|
logger.InfoCF("pico", "Sending node action via Pico", map[string]any{
|
||||||
|
"action": payload.Action(),
|
||||||
|
"target_addr": targetAddr,
|
||||||
|
"request_id": requestID,
|
||||||
|
})
|
||||||
|
|
||||||
|
return c.doSend(ctx, targetAddr, requestID, payload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// sendRequest is the internal method for sending a request and returning the "response" string.
|
||||||
|
func (c *PicoNodeClient) sendRequest(
|
||||||
|
ctx context.Context,
|
||||||
|
targetAddr, targetNodeID string,
|
||||||
|
payload NodePayload,
|
||||||
|
) (string, error) {
|
||||||
|
requestID := uuid.New().String()
|
||||||
|
payload[PayloadKeyRequestID] = requestID
|
||||||
|
|
||||||
|
logger.InfoCF("pico", "Sending node request via Pico", map[string]any{
|
||||||
|
"action": payload.Action(),
|
||||||
|
"target_node_id": targetNodeID,
|
||||||
|
"target_addr": targetAddr,
|
||||||
|
"request_id": requestID,
|
||||||
|
})
|
||||||
|
|
||||||
|
replyPayload, err := c.doSend(ctx, targetAddr, requestID, payload)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
if errStr := replyPayload.ErrorMsg(); errStr != "" {
|
||||||
|
return "", fmt.Errorf("node error: %s", errStr)
|
||||||
|
}
|
||||||
|
|
||||||
|
return replyPayload.Response(), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// doSend builds a protocol.Message, delegates to Client.SendRequest,
|
||||||
|
// and validates the reply envelope before returning the payload.
|
||||||
|
func (c *PicoNodeClient) doSend(
|
||||||
|
ctx context.Context,
|
||||||
|
targetAddr, requestID string,
|
||||||
|
payload NodePayload,
|
||||||
|
) (NodePayload, error) {
|
||||||
|
reqMsg := protocol.Message{
|
||||||
|
Type: protocol.TypeNodeRequest,
|
||||||
|
ID: requestID,
|
||||||
|
SessionID: c.sourceNodeID,
|
||||||
|
Timestamp: time.Now().UnixMilli(),
|
||||||
|
Payload: payload,
|
||||||
|
}
|
||||||
|
|
||||||
|
reply, err := c.client.SendRequest(ctx, targetAddr, c.sourceNodeID, reqMsg)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if reply.Type != protocol.TypeNodeReply {
|
||||||
|
return nil, fmt.Errorf("unexpected reply type: %s (expected %s)", reply.Type, protocol.TypeNodeReply)
|
||||||
|
}
|
||||||
|
|
||||||
|
replyPayload := NodePayload(reply.Payload)
|
||||||
|
if replyPayload.RequestID() != requestID {
|
||||||
|
return nil, fmt.Errorf("request ID mismatch: got %s, want %s", replyPayload.RequestID(), requestID)
|
||||||
|
}
|
||||||
|
|
||||||
|
return replyPayload, nil
|
||||||
|
}
|
||||||
132
pkg/pico/node_payload.go
Normal file
132
pkg/pico/node_payload.go
Normal file
|
|
@ -0,0 +1,132 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package pico
|
||||||
|
|
||||||
|
// NodePayload is the typed payload exchanged between nodes via the
|
||||||
|
// Pico channel node.request / node.reply protocol. It provides named
|
||||||
|
// constants for field keys and typed accessor methods so that callers
|
||||||
|
// never need to use raw string literals.
|
||||||
|
type NodePayload map[string]any
|
||||||
|
|
||||||
|
// Payload field key constants.
|
||||||
|
const (
|
||||||
|
PayloadKeyAction = "action"
|
||||||
|
PayloadKeyRequestID = "request_id"
|
||||||
|
PayloadKeySourceNodeID = "source_node_id"
|
||||||
|
PayloadKeyContent = "content"
|
||||||
|
PayloadKeyChannel = "channel"
|
||||||
|
PayloadKeyChatID = "chat_id"
|
||||||
|
PayloadKeySenderID = "sender_id"
|
||||||
|
PayloadKeyMetadata = "metadata"
|
||||||
|
PayloadKeyError = "error"
|
||||||
|
PayloadKeyResponse = "response"
|
||||||
|
PayloadKeyRequest = "request" // nested handoff request object
|
||||||
|
PayloadKeyHandoffResp = "handoff_response" // nested handoff response object
|
||||||
|
)
|
||||||
|
|
||||||
|
// Node action constants used for action-based routing over Pico.
|
||||||
|
const (
|
||||||
|
NodeActionMessage = "message"
|
||||||
|
NodeActionStatus = "status"
|
||||||
|
NodeActionHandoffRequest = "handoff_request"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Accessors
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
// str is a small helper that extracts a string value from the payload.
|
||||||
|
func (p NodePayload) str(key string) string {
|
||||||
|
v, _ := p[key].(string)
|
||||||
|
return v
|
||||||
|
}
|
||||||
|
|
||||||
|
// Action returns the action field (e.g. "message", "handoff_request").
|
||||||
|
func (p NodePayload) Action() string { return p.str(PayloadKeyAction) }
|
||||||
|
|
||||||
|
// RequestID returns the request_id field.
|
||||||
|
func (p NodePayload) RequestID() string { return p.str(PayloadKeyRequestID) }
|
||||||
|
|
||||||
|
// SourceNodeID returns the source_node_id field.
|
||||||
|
func (p NodePayload) SourceNodeID() string { return p.str(PayloadKeySourceNodeID) }
|
||||||
|
|
||||||
|
// Content returns the content field.
|
||||||
|
func (p NodePayload) Content() string { return p.str(PayloadKeyContent) }
|
||||||
|
|
||||||
|
// Channel returns the channel field.
|
||||||
|
func (p NodePayload) Channel() string { return p.str(PayloadKeyChannel) }
|
||||||
|
|
||||||
|
// ChatID returns the chat_id field.
|
||||||
|
func (p NodePayload) ChatID() string { return p.str(PayloadKeyChatID) }
|
||||||
|
|
||||||
|
// SenderID returns the sender_id field.
|
||||||
|
func (p NodePayload) SenderID() string { return p.str(PayloadKeySenderID) }
|
||||||
|
|
||||||
|
// ErrorMsg returns the error field.
|
||||||
|
func (p NodePayload) ErrorMsg() string { return p.str(PayloadKeyError) }
|
||||||
|
|
||||||
|
// Response returns the response field.
|
||||||
|
func (p NodePayload) Response() string { return p.str(PayloadKeyResponse) }
|
||||||
|
|
||||||
|
// Metadata extracts the metadata map, converting map[string]any to map[string]string.
|
||||||
|
func (p NodePayload) Metadata() map[string]string {
|
||||||
|
raw, ok := p[PayloadKeyMetadata]
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
m, ok := raw.(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
result := make(map[string]string, len(m))
|
||||||
|
for k, v := range m {
|
||||||
|
if s, ok := v.(string); ok {
|
||||||
|
result[k] = s
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// RawValue returns the raw value for an arbitrary key.
|
||||||
|
func (p NodePayload) RawValue(key string) (any, bool) {
|
||||||
|
v, ok := p[key]
|
||||||
|
return v, ok
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Builder helpers
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
// NewNodePayload creates an empty NodePayload.
|
||||||
|
func NewNodePayload() NodePayload {
|
||||||
|
return make(NodePayload)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewMessagePayload creates a NodePayload pre-filled for a "message" action.
|
||||||
|
func NewMessagePayload(sourceNodeID, content, channel, chatID, senderID string) NodePayload {
|
||||||
|
return NodePayload{
|
||||||
|
PayloadKeyAction: NodeActionMessage,
|
||||||
|
PayloadKeySourceNodeID: sourceNodeID,
|
||||||
|
PayloadKeyContent: content,
|
||||||
|
PayloadKeyChannel: channel,
|
||||||
|
PayloadKeyChatID: chatID,
|
||||||
|
PayloadKeySenderID: senderID,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Reply constructors
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
// ErrorReply creates a reply payload carrying an error message.
|
||||||
|
func ErrorReply(msg string) NodePayload {
|
||||||
|
return NodePayload{PayloadKeyError: msg}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ResponseReply creates a reply payload carrying a string response.
|
||||||
|
func ResponseReply(response string) NodePayload {
|
||||||
|
return NodePayload{PayloadKeyResponse: response}
|
||||||
|
}
|
||||||
65
pkg/pico/protocol/protocol.go
Normal file
65
pkg/pico/protocol/protocol.go
Normal file
|
|
@ -0,0 +1,65 @@
|
||||||
|
// Package protocol defines the Pico Protocol wire format shared by
|
||||||
|
// the Pico channel (server) and the swarm PicoNodeClient (client).
|
||||||
|
// This package has zero internal dependencies to stay at the bottom
|
||||||
|
// of the dependency graph.
|
||||||
|
package protocol
|
||||||
|
|
||||||
|
import "time"
|
||||||
|
|
||||||
|
// Message type constants for the Pico Protocol.
|
||||||
|
const (
|
||||||
|
// TypeMessageSend is sent from client to server.
|
||||||
|
TypeMessageSend = "message.send"
|
||||||
|
// TypeMediaSend is sent from client to server for media.
|
||||||
|
TypeMediaSend = "media.send"
|
||||||
|
// TypePing is a client ping.
|
||||||
|
TypePing = "ping"
|
||||||
|
|
||||||
|
// TypeMessageCreate is sent from server to client.
|
||||||
|
TypeMessageCreate = "message.create"
|
||||||
|
// TypeMessageUpdate is sent from server to client for message updates.
|
||||||
|
TypeMessageUpdate = "message.update"
|
||||||
|
// TypeMediaCreate is sent from server to client for media.
|
||||||
|
TypeMediaCreate = "media.create"
|
||||||
|
// TypeTypingStart indicates typing has started.
|
||||||
|
TypeTypingStart = "typing.start"
|
||||||
|
// TypeTypingStop indicates typing has stopped.
|
||||||
|
TypeTypingStop = "typing.stop"
|
||||||
|
// TypeError is an error message.
|
||||||
|
TypeError = "error"
|
||||||
|
// TypePong is a server pong reply.
|
||||||
|
TypePong = "pong"
|
||||||
|
|
||||||
|
// TypeNodeRequest is for inter-node swarm communication.
|
||||||
|
TypeNodeRequest = "node.request"
|
||||||
|
// TypeNodeReply is the reply for inter-node swarm communication.
|
||||||
|
TypeNodeReply = "node.reply"
|
||||||
|
// TypeNodeProcessing is sent periodically while processing a long-running request.
|
||||||
|
TypeNodeProcessing = "node.processing"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Message is the wire format for all Pico Protocol messages.
|
||||||
|
type Message struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
ID string `json:"id,omitempty"`
|
||||||
|
SessionID string `json:"session_id,omitempty"`
|
||||||
|
Timestamp int64 `json:"timestamp,omitempty"`
|
||||||
|
Payload map[string]any `json:"payload,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewMessage creates a Message with the given type, payload, and current timestamp.
|
||||||
|
func NewMessage(msgType string, payload map[string]any) Message {
|
||||||
|
return Message{
|
||||||
|
Type: msgType,
|
||||||
|
Timestamp: time.Now().UnixMilli(),
|
||||||
|
Payload: payload,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewError creates an error Message with code and human-readable message.
|
||||||
|
func NewError(code, message string) Message {
|
||||||
|
return NewMessage(TypeError, map[string]any{
|
||||||
|
"code": code,
|
||||||
|
"message": message,
|
||||||
|
})
|
||||||
|
}
|
||||||
45
pkg/pico/types.go
Normal file
45
pkg/pico/types.go
Normal file
|
|
@ -0,0 +1,45 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package pico
|
||||||
|
|
||||||
|
// SessionMessage represents a message in a session.
|
||||||
|
// This is shared across handoff, session transfer, and inter-node communication.
|
||||||
|
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"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// DirectMessage represents a direct message sent to another node.
|
||||||
|
type DirectMessage struct {
|
||||||
|
MessageID string `json:"message_id"`
|
||||||
|
SourceNodeID string `json:"source_node_id"`
|
||||||
|
TargetNodeID string `json:"target_node_id"`
|
||||||
|
Action string `json:"action,omitempty"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
Channel string `json:"channel"`
|
||||||
|
ChatID string `json:"chat_id"`
|
||||||
|
SenderID string `json:"sender_id"`
|
||||||
|
Metadata map[string]string `json:"metadata,omitempty"`
|
||||||
|
Timestamp int64 `json:"timestamp"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// DirectMessageResponse represents a response to a direct message.
|
||||||
|
type DirectMessageResponse struct {
|
||||||
|
MessageID string `json:"message_id"`
|
||||||
|
Response string `json:"response"`
|
||||||
|
Error string `json:"error,omitempty"`
|
||||||
|
}
|
||||||
103
pkg/swarm/bucket_config.go
Normal file
103
pkg/swarm/bucket_config.go
Normal file
|
|
@ -0,0 +1,103 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"github.com/sipeed/picoclaw/pkg/kv"
|
||||||
|
)
|
||||||
|
|
||||||
|
// BucketType identifies the type of bucket.
|
||||||
|
type BucketType string
|
||||||
|
|
||||||
|
const (
|
||||||
|
BucketTypeMembers BucketType = "members"
|
||||||
|
BucketTypeStatus BucketType = "status"
|
||||||
|
BucketTypeLeader BucketType = "leader"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DefaultBucketSpecs returns the default bucket specifications for all bucket types.
|
||||||
|
func DefaultBucketSpecs() map[BucketType]kv.BucketSpec {
|
||||||
|
return map[BucketType]kv.BucketSpec{
|
||||||
|
BucketTypeMembers: {
|
||||||
|
Name: "swarm_members",
|
||||||
|
Description: "PicoClaw Swarm Node Membership",
|
||||||
|
TTL: DefaultMemberTTL,
|
||||||
|
MaxValueSize: 1024,
|
||||||
|
Storage: kv.StorageFile,
|
||||||
|
},
|
||||||
|
BucketTypeStatus: {
|
||||||
|
Name: "swarm_status",
|
||||||
|
Description: "PicoClaw Swarm Node Detailed Status",
|
||||||
|
TTL: DefaultDetailedStatusTTL,
|
||||||
|
MaxValueSize: 4096,
|
||||||
|
Storage: kv.StorageMemory,
|
||||||
|
},
|
||||||
|
BucketTypeLeader: {
|
||||||
|
Name: "swarm_leader",
|
||||||
|
Description: "PicoClaw Swarm Leader Election Lock",
|
||||||
|
TTL: DefaultLeaderLockTTL,
|
||||||
|
MaxValueSize: 512,
|
||||||
|
Storage: kv.StorageMemory,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MemberBucketSpec returns the specification for the members bucket.
|
||||||
|
func MemberBucketSpec() kv.BucketSpec {
|
||||||
|
return DefaultBucketSpecs()[BucketTypeMembers]
|
||||||
|
}
|
||||||
|
|
||||||
|
// StatusBucketSpec returns the specification for the status bucket.
|
||||||
|
func StatusBucketSpec() kv.BucketSpec {
|
||||||
|
return DefaultBucketSpecs()[BucketTypeStatus]
|
||||||
|
}
|
||||||
|
|
||||||
|
// LeaderBucketSpec returns the specification for the leader bucket.
|
||||||
|
func LeaderBucketSpec() kv.BucketSpec {
|
||||||
|
return DefaultBucketSpecs()[BucketTypeLeader]
|
||||||
|
}
|
||||||
|
|
||||||
|
// BucketSet provides access to all buckets used by the swarm.
|
||||||
|
type BucketSet struct {
|
||||||
|
Members kv.Bucket
|
||||||
|
Status kv.Bucket
|
||||||
|
Leader kv.Bucket
|
||||||
|
factory kv.BucketFactory
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewBucketSet creates a new bucket set using the given factory.
|
||||||
|
func NewBucketSet(factory kv.BucketFactory) (*BucketSet, error) {
|
||||||
|
bs := &BucketSet{factory: factory}
|
||||||
|
|
||||||
|
var err error
|
||||||
|
specs := DefaultBucketSpecs()
|
||||||
|
|
||||||
|
// Create or get members bucket
|
||||||
|
bs.Members, err = factory.CreateIfNotExists(specs[BucketTypeMembers])
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create or get status bucket
|
||||||
|
bs.Status, err = factory.CreateIfNotExists(specs[BucketTypeStatus])
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create or get leader bucket
|
||||||
|
bs.Leader, err = factory.CreateIfNotExists(specs[BucketTypeLeader])
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
return bs, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close closes all buckets in the set.
|
||||||
|
func (bs *BucketSet) Close() error {
|
||||||
|
return bs.factory.Close()
|
||||||
|
}
|
||||||
269
pkg/swarm/config.go
Normal file
269
pkg/swarm/config.go
Normal file
|
|
@ -0,0 +1,269 @@
|
||||||
|
// 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"`
|
||||||
|
|
||||||
|
// Discovery configuration for node discovery via NATS.
|
||||||
|
Discovery DiscoveryConfig `json:"discovery"`
|
||||||
|
|
||||||
|
// Handoff configuration for task handoff.
|
||||||
|
Handoff HandoffConfig `json:"handoff"`
|
||||||
|
|
||||||
|
// 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 NATS-based node discovery.
|
||||||
|
type DiscoveryConfig struct {
|
||||||
|
// NATSURL is the NATS server URL (required).
|
||||||
|
NATSURL string `json:"nats_url,omitempty" env:"PICOCLAW_SWARM_NATS_URL"`
|
||||||
|
|
||||||
|
// NATSCredsFile is the path to NATS credentials file (NKeys/JWT auth).
|
||||||
|
NATSCredsFile string `json:"nats_creds_file,omitempty" env:"PICOCLAW_SWARM_NATS_CREDS_FILE"`
|
||||||
|
|
||||||
|
// NATSTLSCert is the path to the client TLS certificate for mTLS.
|
||||||
|
NATSTLSCert string `json:"nats_tls_cert,omitempty"`
|
||||||
|
|
||||||
|
// NATSTLSKey is the path to the client TLS key for mTLS.
|
||||||
|
NATSTLSKey string `json:"nats_tls_key,omitempty"`
|
||||||
|
|
||||||
|
// NATSTLSCACert is the path to the CA certificate for TLS verification.
|
||||||
|
NATSTLSCACert string `json:"nats_tls_ca_cert,omitempty"`
|
||||||
|
|
||||||
|
// HeartbeatInterval controls how often the node renews its KV entry.
|
||||||
|
HeartbeatInterval Duration `json:"heartbeat_interval,omitempty"`
|
||||||
|
|
||||||
|
// MemberTTL is the TTL for member entries in the KV bucket.
|
||||||
|
MemberTTL Duration `json:"member_ttl,omitempty"`
|
||||||
|
|
||||||
|
// SubjectPrefix overrides the default "picoclaw.swarm" subject prefix.
|
||||||
|
SubjectPrefix string `json:"subject_prefix,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"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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"`
|
||||||
|
|
||||||
|
// RequestTimeout is the timeout for a single NATS request-reply handoff.
|
||||||
|
RequestTimeout Duration `json:"request_timeout,omitempty"`
|
||||||
|
|
||||||
|
// MaxSessionHistory is the maximum number of recent session messages to include in handoff.
|
||||||
|
// Set to 0 to send all messages (not recommended for production). Default: 10
|
||||||
|
MaxSessionHistory int `json:"max_session_history,omitempty"`
|
||||||
|
|
||||||
|
// Crypto contains security configuration for handoff.
|
||||||
|
Crypto CryptoConfig `json:"crypto"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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"`
|
||||||
|
|
||||||
|
// RoutingRejectThreshold is the load score above which routing requests are rejected (0-1).
|
||||||
|
// This provides hysteresis to prevent oscillation when load is near threshold.
|
||||||
|
RoutingRejectThreshold float64 `json:"routing_reject_threshold,omitempty"`
|
||||||
|
|
||||||
|
// RoutingAcceptThreshold is the load score below which routing requests are accepted (0-1).
|
||||||
|
// Provides hysteresis - once rejected, load must drop below this to accept again.
|
||||||
|
RoutingAcceptThreshold float64 `json:"routing_accept_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"`
|
||||||
|
|
||||||
|
// LockTTL is how long the leader lock lives without renewal.
|
||||||
|
LockTTL Duration `json:"lock_ttl,omitempty"`
|
||||||
|
|
||||||
|
// RenewalInterval is how often the leader renews its lock.
|
||||||
|
// Must be less than LockTTL.
|
||||||
|
RenewalInterval Duration `json:"renewal_interval,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// CryptoConfig contains cryptographic settings for secure swarm communication.
|
||||||
|
type CryptoConfig struct {
|
||||||
|
// SharedSecret is the HMAC secret for message signing (base64 encoded or raw string).
|
||||||
|
// If empty, signing is disabled (NOT RECOMMENDED for production).
|
||||||
|
SharedSecret string `json:"shared_secret,omitempty" env:"PICOCLAW_SWARM_SHARED_SECRET"`
|
||||||
|
|
||||||
|
// EncryptionKey is the AES-256 key for encrypting session data (base64 encoded or raw string).
|
||||||
|
// Must be 32 bytes when decoded. If empty, encryption is disabled.
|
||||||
|
EncryptionKey string `json:"encryption_key,omitempty" env:"PICOCLAW_SWARM_ENCRYPTION_KEY"`
|
||||||
|
|
||||||
|
// RequireAuth enables strict authentication - reject unsigned messages.
|
||||||
|
RequireAuth bool `json:"require_auth"`
|
||||||
|
|
||||||
|
// RequireEncryption enables strict encryption - reject unencrypted session data.
|
||||||
|
RequireEncryption bool `json:"require_encryption"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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.
|
||||||
|
// Supports:
|
||||||
|
// - String: Go duration format, e.g. "5s", "100ms", "2m30s"
|
||||||
|
// - Number: interpreted as seconds (e.g. 5 means 5s, 0.5 means 500ms)
|
||||||
|
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 := json.Unmarshal(b, &s); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
var err error
|
||||||
|
d.Duration, err = time.ParseDuration(s)
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Otherwise it's a number — interpret as seconds for human-friendly config
|
||||||
|
var v float64
|
||||||
|
if err := json.Unmarshal(b, &v); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
d.Duration = time.Duration(v * float64(time.Second))
|
||||||
|
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: "",
|
||||||
|
Discovery: DiscoveryConfig{
|
||||||
|
NATSURL: DefaultNATSURL,
|
||||||
|
HeartbeatInterval: Duration{DefaultHeartbeatInterval},
|
||||||
|
MemberTTL: Duration{DefaultMemberTTL},
|
||||||
|
NodeTimeout: Duration{DefaultNodeTimeout},
|
||||||
|
DeadNodeTimeout: Duration{DefaultDeadNodeTimeout},
|
||||||
|
},
|
||||||
|
Handoff: HandoffConfig{
|
||||||
|
Enabled: true,
|
||||||
|
LoadThreshold: DefaultLoadThreshold,
|
||||||
|
Timeout: Duration{DefaultHandoffTimeout},
|
||||||
|
MaxRetries: DefaultMaxHandoffRetries,
|
||||||
|
RetryDelay: Duration{DefaultHandoffRetryDelay},
|
||||||
|
RequestTimeout: Duration{DefaultHandoffRequestTimeout},
|
||||||
|
MaxSessionHistory: DefaultMaxSessionHistory,
|
||||||
|
},
|
||||||
|
LoadMonitor: LoadMonitorConfig{
|
||||||
|
Enabled: true,
|
||||||
|
Interval: Duration{DefaultLoadSampleInterval},
|
||||||
|
SampleSize: DefaultLoadSampleSize,
|
||||||
|
CPUWeight: DefaultCPUWeight,
|
||||||
|
MemoryWeight: DefaultMemoryWeight,
|
||||||
|
SessionWeight: DefaultSessionWeight,
|
||||||
|
OffloadThreshold: DefaultOffloadThreshold,
|
||||||
|
RoutingRejectThreshold: DefaultRoutingRejectThreshold,
|
||||||
|
RoutingAcceptThreshold: DefaultRoutingAcceptThreshold,
|
||||||
|
MaxMemoryBytes: DefaultMaxMemoryBytes,
|
||||||
|
MaxGoroutines: DefaultMaxGoroutines,
|
||||||
|
MaxSessions: DefaultMaxSessions,
|
||||||
|
},
|
||||||
|
LeaderElection: LeaderElectionConfig{
|
||||||
|
Enabled: false,
|
||||||
|
LockTTL: Duration{DefaultLeaderLockTTL},
|
||||||
|
RenewalInterval: Duration{DefaultLeaderRenewalInterval},
|
||||||
|
},
|
||||||
|
Metrics: MetricsConfig{
|
||||||
|
Enabled: false,
|
||||||
|
ExportInterval: Duration{10 * time.Second},
|
||||||
|
PrometheusEnabled: false,
|
||||||
|
PrometheusEndpoint: "/metrics",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
199
pkg/swarm/constants.go
Normal file
199
pkg/swarm/constants.go
Normal file
|
|
@ -0,0 +1,199 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import "time"
|
||||||
|
|
||||||
|
// LoadTrend represents the direction of load change over time.
|
||||||
|
type LoadTrend string
|
||||||
|
|
||||||
|
const (
|
||||||
|
LoadTrendIncreasing LoadTrend = "increasing"
|
||||||
|
LoadTrendDecreasing LoadTrend = "decreasing"
|
||||||
|
LoadTrendStable LoadTrend = "stable"
|
||||||
|
)
|
||||||
|
|
||||||
|
// NATS Subject namespace.
|
||||||
|
// All swarm communication uses these subject prefixes.
|
||||||
|
// ACL rule: only swarm node identities may publish/subscribe picoclaw.swarm.*
|
||||||
|
const (
|
||||||
|
// SubjectPrefix is the root namespace for all swarm subjects.
|
||||||
|
SubjectPrefix = "picoclaw.swarm"
|
||||||
|
|
||||||
|
// SubjectHeartbeat is the prefix for heartbeat messages: picoclaw.swarm.heartbeat.<nodeID>
|
||||||
|
SubjectHeartbeat = SubjectPrefix + ".heartbeat"
|
||||||
|
|
||||||
|
// SubjectMetrics is the prefix for metrics messages: picoclaw.swarm.metrics.<nodeID>
|
||||||
|
SubjectMetrics = SubjectPrefix + ".metrics"
|
||||||
|
|
||||||
|
// SubjectNodeMsg is the prefix for direct node messages: picoclaw.swarm.node.<nodeID>.msg
|
||||||
|
SubjectNodeMsg = SubjectPrefix + ".node"
|
||||||
|
|
||||||
|
// SubjectHandoff is the prefix for targeted handoff: picoclaw.swarm.handoff.<targetNodeID>
|
||||||
|
SubjectHandoff = SubjectPrefix + ".handoff"
|
||||||
|
|
||||||
|
// SubjectLeader is the subject for leader election announcements.
|
||||||
|
SubjectLeader = SubjectPrefix + ".leader"
|
||||||
|
)
|
||||||
|
|
||||||
|
// NATS JetStream KV Bucket names.
|
||||||
|
const (
|
||||||
|
// KVBucketMembers stores node membership with TTL-based liveness.
|
||||||
|
KVBucketMembers = "swarm_members"
|
||||||
|
|
||||||
|
// KVBucketStatus stores detailed node status (CPU, memory, tasks, etc.).
|
||||||
|
KVBucketStatus = "swarm_status"
|
||||||
|
|
||||||
|
// KVBucketLeader stores the leader election CAS lock.
|
||||||
|
KVBucketLeader = "swarm_leader"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Default NATS connection settings.
|
||||||
|
const (
|
||||||
|
// DefaultNATSURL is the default NATS server URL.
|
||||||
|
DefaultNATSURL = "nats://localhost:4222"
|
||||||
|
|
||||||
|
// DefaultNATSPrefix is the default prefix for NATS KV buckets.
|
||||||
|
DefaultNATSPrefix = "picoclaw_swarm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Default intervals for NATS-based operations.
|
||||||
|
const (
|
||||||
|
// DefaultHeartbeatInterval is how often a node renews its KV entry.
|
||||||
|
DefaultHeartbeatInterval = 5 * time.Second
|
||||||
|
|
||||||
|
// DefaultMemberTTL is the TTL for member entries in the KV bucket.
|
||||||
|
DefaultMemberTTL = 15 * time.Second
|
||||||
|
|
||||||
|
// DefaultDetailedStatusInterval is the interval for publishing detailed status.
|
||||||
|
DefaultDetailedStatusInterval = 10 * time.Second
|
||||||
|
|
||||||
|
// DefaultDetailedStatusTTL is the TTL for detailed status entries.
|
||||||
|
DefaultDetailedStatusTTL = 30 * time.Second
|
||||||
|
|
||||||
|
// DefaultLeaderLockTTL is how long the leader lock lives without renewal.
|
||||||
|
DefaultLeaderLockTTL = 10 * time.Second
|
||||||
|
|
||||||
|
// DefaultLeaderRenewalInterval is how often the leader renews its lock.
|
||||||
|
// Must be less than DefaultLeaderLockTTL.
|
||||||
|
DefaultLeaderRenewalInterval = 3 * time.Second
|
||||||
|
|
||||||
|
// DefaultHandoffRequestTimeout is the timeout for a single NATS handoff request-reply.
|
||||||
|
DefaultHandoffRequestTimeout = 10 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// Default timeouts for node health.
|
||||||
|
const (
|
||||||
|
// DefaultNodeTimeout is the timeout before marking a node as suspect.
|
||||||
|
DefaultNodeTimeout = 5 * time.Second
|
||||||
|
|
||||||
|
// DefaultDeadNodeTimeout is the timeout before removing a dead node.
|
||||||
|
DefaultDeadNodeTimeout = 30 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// Default handoff settings.
|
||||||
|
const (
|
||||||
|
// 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
|
||||||
|
|
||||||
|
// DefaultMaxHandoffPayloadBytes is the max size of a handoff NATS message.
|
||||||
|
// NATS default max_payload is 1MB; we leave headroom for envelope overhead.
|
||||||
|
DefaultMaxHandoffPayloadBytes = 900 * 1024 // 900KB
|
||||||
|
)
|
||||||
|
|
||||||
|
// Default load monitor settings.
|
||||||
|
const (
|
||||||
|
// DefaultLoadSampleInterval is the default interval between load samples.
|
||||||
|
DefaultLoadSampleInterval = 5 * time.Second
|
||||||
|
|
||||||
|
// DefaultLoadSampleSize is the default number of load samples to keep.
|
||||||
|
DefaultLoadSampleSize = 60
|
||||||
|
)
|
||||||
|
|
||||||
|
// Load thresholds and limits.
|
||||||
|
const (
|
||||||
|
// 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
|
||||||
|
|
||||||
|
// DefaultRoutingRejectThreshold is the load threshold above which routing requests are rejected.
|
||||||
|
// Provides hysteresis to prevent oscillation. Default: 0.9
|
||||||
|
DefaultRoutingRejectThreshold = 0.9
|
||||||
|
|
||||||
|
// DefaultRoutingAcceptThreshold is the load threshold below which routing requests are accepted.
|
||||||
|
// Once rejected, load must drop below this to accept again. Default: 0.75
|
||||||
|
DefaultRoutingAcceptThreshold = 0.75
|
||||||
|
|
||||||
|
// 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
|
||||||
|
|
||||||
|
// DefaultMaxSessionHistory is the default number of recent messages to include in handoff.
|
||||||
|
DefaultMaxSessionHistory = 10
|
||||||
|
)
|
||||||
|
|
||||||
|
// Load score weights.
|
||||||
|
const (
|
||||||
|
// 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
|
||||||
|
)
|
||||||
|
|
||||||
|
// Node action constants for action-based node requests.
|
||||||
|
// These are lightweight operations that don't trigger LLM processing.
|
||||||
|
const (
|
||||||
|
// NodeActionStatus requests node status information.
|
||||||
|
NodeActionStatus = "status"
|
||||||
|
|
||||||
|
// NodeActionPing checks if the node is responsive.
|
||||||
|
NodeActionPing = "ping"
|
||||||
|
|
||||||
|
// NodeActionHealth checks node health metrics.
|
||||||
|
NodeActionHealth = "health"
|
||||||
|
|
||||||
|
// NodeActionLeader queries current leader in leader election.
|
||||||
|
NodeActionLeader = "leader"
|
||||||
|
|
||||||
|
// NodeActionMetrics requests detailed metrics.
|
||||||
|
NodeActionMetrics = "metrics"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Default timeout for action-based requests (shorter than message processing).
|
||||||
|
const (
|
||||||
|
// DefaultNodeActionTimeout is the timeout for lightweight node actions.
|
||||||
|
DefaultNodeActionTimeout = 10 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// Trend analysis thresholds.
|
||||||
|
const (
|
||||||
|
// TrendIncreasingThreshold is the slope threshold for detecting increasing trend.
|
||||||
|
TrendIncreasingThreshold = 0.01
|
||||||
|
|
||||||
|
// TrendDecreasingThreshold is the slope threshold for detecting decreasing trend.
|
||||||
|
TrendDecreasingThreshold = -0.01
|
||||||
|
)
|
||||||
64
pkg/swarm/coordinator.go
Normal file
64
pkg/swarm/coordinator.go
Normal file
|
|
@ -0,0 +1,64 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
)
|
||||||
|
|
||||||
|
// CoordinatorAdapter adapts LeaderElection to the Coordinator interface.
|
||||||
|
type CoordinatorAdapter struct {
|
||||||
|
le *LeaderElection
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewCoordinatorAdapter creates a coordinator adapter from a BucketSet.
|
||||||
|
func NewCoordinatorAdapter(
|
||||||
|
buckets *BucketSet,
|
||||||
|
nodeID string,
|
||||||
|
config LeaderElectionConfig,
|
||||||
|
) (*CoordinatorAdapter, error) {
|
||||||
|
if buckets == nil {
|
||||||
|
return nil, fmt.Errorf("buckets cannot be nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
le, err := NewLeaderElection(nodeID, buckets.Leader, config)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create leader election: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return &CoordinatorAdapter{le: le}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start starts the coordinator.
|
||||||
|
func (ca *CoordinatorAdapter) Start() error {
|
||||||
|
return ca.le.Start()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop stops the coordinator.
|
||||||
|
func (ca *CoordinatorAdapter) Stop() {
|
||||||
|
ca.le.Stop()
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsLeader returns true if this node is the current leader.
|
||||||
|
func (ca *CoordinatorAdapter) IsLeader() bool {
|
||||||
|
return ca.le.IsLeader()
|
||||||
|
}
|
||||||
|
|
||||||
|
// LeaderID returns the current leader's node ID.
|
||||||
|
func (ca *CoordinatorAdapter) LeaderID() string {
|
||||||
|
return ca.le.GetLeader()
|
||||||
|
}
|
||||||
|
|
||||||
|
// LeaderChanges returns a channel that receives leader ID changes.
|
||||||
|
func (ca *CoordinatorAdapter) LeaderChanges() <-chan string {
|
||||||
|
return ca.le.LeaderChanges()
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetLeaderElection returns the underlying LeaderElection for advanced usage.
|
||||||
|
func (ca *CoordinatorAdapter) GetLeaderElection() *LeaderElection {
|
||||||
|
return ca.le
|
||||||
|
}
|
||||||
30
pkg/swarm/coordinator_interface.go
Normal file
30
pkg/swarm/coordinator_interface.go
Normal file
|
|
@ -0,0 +1,30 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
// Coordinator provides leader election functionality.
|
||||||
|
//
|
||||||
|
// The coordinator ensures exactly one node in the cluster acts as leader
|
||||||
|
// for operations that require cluster-wide coordination.
|
||||||
|
type Coordinator interface {
|
||||||
|
// Start starts the coordinator.
|
||||||
|
Start() error
|
||||||
|
|
||||||
|
// Stop stops the coordinator.
|
||||||
|
Stop()
|
||||||
|
|
||||||
|
// IsLeader returns true if this node is the current leader.
|
||||||
|
IsLeader() bool
|
||||||
|
|
||||||
|
// LeaderID returns the current leader's node ID.
|
||||||
|
// Returns empty string if no leader is elected.
|
||||||
|
LeaderID() string
|
||||||
|
|
||||||
|
// LeaderChanges returns a channel that receives leader ID changes.
|
||||||
|
// The channel is closed when Stop() is called.
|
||||||
|
LeaderChanges() <-chan string
|
||||||
|
}
|
||||||
276
pkg/swarm/crypto.go
Normal file
276
pkg/swarm/crypto.go
Normal file
|
|
@ -0,0 +1,276 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm cryptographic utilities for secure inter-node communication
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/aes"
|
||||||
|
"crypto/cipher"
|
||||||
|
"crypto/hmac"
|
||||||
|
"crypto/rand"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/base64"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Signer creates and verifies HMAC-SHA256 signatures.
|
||||||
|
type Signer struct {
|
||||||
|
secret []byte
|
||||||
|
enabled bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSigner creates a new HMAC signer.
|
||||||
|
func NewSigner(secret string) *Signer {
|
||||||
|
s := &Signer{enabled: secret != ""}
|
||||||
|
if s.enabled {
|
||||||
|
// Support both base64-encoded and raw secrets
|
||||||
|
if decoded, err := base64.StdEncoding.DecodeString(secret); err == nil && len(decoded) > 0 {
|
||||||
|
s.secret = decoded
|
||||||
|
} else {
|
||||||
|
s.secret = []byte(secret)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sign computes HMAC-SHA256 of the canonical JSON representation of data.
|
||||||
|
func (s *Signer) Sign(data any) (string, error) {
|
||||||
|
if !s.enabled {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create canonical JSON (sorted keys)
|
||||||
|
canonical, err := canonicalJSON(data)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to canonicalize: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
h := hmac.New(sha256.New, s.secret)
|
||||||
|
h.Write(canonical)
|
||||||
|
sig := h.Sum(nil)
|
||||||
|
return base64.StdEncoding.EncodeToString(sig), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify verifies the HMAC-SHA256 signature of data.
|
||||||
|
func (s *Signer) Verify(data any, signature string) bool {
|
||||||
|
if !s.enabled {
|
||||||
|
return true // Disabled means accept all
|
||||||
|
}
|
||||||
|
|
||||||
|
if signature == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
expected, err := s.Sign(data)
|
||||||
|
if err != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
return hmac.Equal([]byte(signature), []byte(expected))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Encryptor encrypts/decrypts session data using AES-GCM.
|
||||||
|
type Encryptor struct {
|
||||||
|
key []byte
|
||||||
|
enabled bool
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewEncryptor creates a new AES-GCM encryptor.
|
||||||
|
func NewEncryptor(key string) (*Encryptor, error) {
|
||||||
|
if key == "" {
|
||||||
|
return &Encryptor{enabled: false}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Support both base64-encoded and raw keys
|
||||||
|
var keyBytes []byte
|
||||||
|
if decoded, err := base64.StdEncoding.DecodeString(key); err == nil && len(decoded) > 0 {
|
||||||
|
keyBytes = decoded
|
||||||
|
} else {
|
||||||
|
keyBytes = []byte(key)
|
||||||
|
}
|
||||||
|
|
||||||
|
// AES-256 requires 32-byte key
|
||||||
|
if len(keyBytes) != 32 {
|
||||||
|
return nil, fmt.Errorf("encryption key must be 32 bytes for AES-256, got %d bytes", len(keyBytes))
|
||||||
|
}
|
||||||
|
|
||||||
|
return &Encryptor{
|
||||||
|
key: keyBytes,
|
||||||
|
enabled: true,
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// EncryptedData wraps encrypted content with nonce.
|
||||||
|
type EncryptedData struct {
|
||||||
|
Nonce string `json:"nonce"`
|
||||||
|
Cipher string `json:"cipher"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Encrypt encrypts the provided data using AES-GCM.
|
||||||
|
func (e *Encryptor) Encrypt(data []byte) (*EncryptedData, error) {
|
||||||
|
if !e.enabled {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
block, err := aes.NewCipher(e.key)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create cipher: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
gcm, err := cipher.NewGCM(block)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create GCM: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
nonce := make([]byte, gcm.NonceSize())
|
||||||
|
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to generate nonce: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cipherText := gcm.Seal(nonce, nonce, data, nil)
|
||||||
|
|
||||||
|
return &EncryptedData{
|
||||||
|
Nonce: base64.StdEncoding.EncodeToString(nonce),
|
||||||
|
Cipher: base64.StdEncoding.EncodeToString(cipherText),
|
||||||
|
}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Decrypt decrypts the provided data using AES-GCM.
|
||||||
|
func (e *Encryptor) Decrypt(enc *EncryptedData) ([]byte, error) {
|
||||||
|
if !e.enabled || enc == nil {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
nonce, err := base64.StdEncoding.DecodeString(enc.Nonce)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("invalid nonce encoding: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cipherText, err := base64.StdEncoding.DecodeString(enc.Cipher)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("invalid cipher encoding: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
block, err := aes.NewCipher(e.key)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create cipher: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
gcm, err := cipher.NewGCM(block)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create GCM: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(cipherText) < gcm.NonceSize() {
|
||||||
|
return nil, fmt.Errorf("cipher text too short")
|
||||||
|
}
|
||||||
|
|
||||||
|
plaintext, err := gcm.Open(nil, nonce, cipherText, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("decryption failed: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return plaintext, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// canonicalJSON creates a canonical JSON representation with sorted keys.
|
||||||
|
func canonicalJSON(data any) ([]byte, error) {
|
||||||
|
// Marshal to JSON then unmarshal into map[string]any for sorting
|
||||||
|
j, err := json.Marshal(data)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
var raw map[string]any
|
||||||
|
if json.Unmarshal(j, &raw) == nil {
|
||||||
|
// Successfully parsed as object, sort keys
|
||||||
|
return json.Marshal(sortMap(raw))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Not an object, return as-is
|
||||||
|
return j, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func sortMap(m map[string]any) map[string]any {
|
||||||
|
sorted := make(map[string]any)
|
||||||
|
for k, v := range m {
|
||||||
|
switch val := v.(type) {
|
||||||
|
case map[string]any:
|
||||||
|
sorted[k] = sortMap(val)
|
||||||
|
case []any:
|
||||||
|
sorted[k] = sortSlice(val)
|
||||||
|
default:
|
||||||
|
sorted[k] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return sorted
|
||||||
|
}
|
||||||
|
|
||||||
|
func sortSlice(s []any) []any {
|
||||||
|
sorted := make([]any, len(s))
|
||||||
|
for i, v := range s {
|
||||||
|
switch val := v.(type) {
|
||||||
|
case map[string]any:
|
||||||
|
sorted[i] = sortMap(val)
|
||||||
|
case []any:
|
||||||
|
sorted[i] = sortSlice(val)
|
||||||
|
default:
|
||||||
|
sorted[i] = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return sorted
|
||||||
|
}
|
||||||
|
|
||||||
|
// GenerateKey generates a random 32-byte key for AES-256 encryption.
|
||||||
|
func GenerateKey() (string, error) {
|
||||||
|
key := make([]byte, 32)
|
||||||
|
if _, err := io.ReadFull(rand.Reader, key); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return base64.StdEncoding.EncodeToString(key), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GenerateSecret generates a random HMAC secret.
|
||||||
|
func GenerateSecret() (string, error) {
|
||||||
|
secret := make([]byte, 32)
|
||||||
|
if _, err := io.ReadFull(rand.Reader, secret); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return base64.StdEncoding.EncodeToString(secret), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// NormalizeKey ensures a key is exactly 32 bytes by hashing or truncating.
|
||||||
|
func NormalizeKey(key string) string {
|
||||||
|
if len(key) >= 32 {
|
||||||
|
return key[:32]
|
||||||
|
}
|
||||||
|
h := sha256.Sum256([]byte(key))
|
||||||
|
return string(h[:])
|
||||||
|
}
|
||||||
|
|
||||||
|
// ComputeMessageSignature creates a signature from key components for manual signing.
|
||||||
|
func ComputeMessageSignature(secret string, parts ...string) string {
|
||||||
|
var s strings.Builder
|
||||||
|
for i, p := range parts {
|
||||||
|
if i > 0 {
|
||||||
|
s.WriteByte('|')
|
||||||
|
}
|
||||||
|
s.WriteString(p)
|
||||||
|
}
|
||||||
|
|
||||||
|
h := hmac.New(sha256.New, []byte(secret))
|
||||||
|
h.Write([]byte(s.String()))
|
||||||
|
return base64.StdEncoding.EncodeToString(h.Sum(nil))
|
||||||
|
}
|
||||||
|
|
||||||
|
// VerifyMessageSignature verifies a manually created signature.
|
||||||
|
func VerifyMessageSignature(secret, signature string, parts ...string) bool {
|
||||||
|
expected := ComputeMessageSignature(secret, parts...)
|
||||||
|
return hmac.Equal([]byte(signature), []byte(expected))
|
||||||
|
}
|
||||||
143
pkg/swarm/discovery_adapter.go
Normal file
143
pkg/swarm/discovery_adapter.go
Normal file
|
|
@ -0,0 +1,143 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DiscoveryAdapter adapts the existing DiscoveryService to the NodeDiscovery interface.
|
||||||
|
type DiscoveryAdapter struct {
|
||||||
|
ds Discovery
|
||||||
|
watchers []func(View)
|
||||||
|
watcherIDs []int // Track watcher IDs for removal
|
||||||
|
nextID int
|
||||||
|
mu sync.RWMutex
|
||||||
|
cancel chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewDiscoveryAdapter creates a new discovery adapter.
|
||||||
|
func NewDiscoveryAdapter(ds Discovery) *DiscoveryAdapter {
|
||||||
|
da := &DiscoveryAdapter{
|
||||||
|
ds: ds,
|
||||||
|
watchers: make([]func(View), 0),
|
||||||
|
watcherIDs: make([]int, 0),
|
||||||
|
nextID: 1,
|
||||||
|
cancel: make(chan struct{}),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Subscribe to node events to update watchers
|
||||||
|
ds.Subscribe(func(event *NodeEvent) {
|
||||||
|
da.notifyWatchers()
|
||||||
|
})
|
||||||
|
|
||||||
|
return da
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start starts the discovery service.
|
||||||
|
func (da *DiscoveryAdapter) Start() error {
|
||||||
|
return da.ds.Start()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop stops the discovery service.
|
||||||
|
func (da *DiscoveryAdapter) Stop() error {
|
||||||
|
close(da.cancel)
|
||||||
|
return da.ds.Stop()
|
||||||
|
}
|
||||||
|
|
||||||
|
// AliveNodes returns the current list of alive nodes.
|
||||||
|
func (da *DiscoveryAdapter) AliveNodes() []*NodeInfo {
|
||||||
|
members := da.ds.Members()
|
||||||
|
result := make([]*NodeInfo, 0)
|
||||||
|
for _, m := range members {
|
||||||
|
if m.State.Status == NodeStatusAlive && m.Node.ID != da.ds.LocalNode().ID {
|
||||||
|
result = append(result, m.Node)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// WatchNodes registers a callback for cluster view changes.
|
||||||
|
func (da *DiscoveryAdapter) WatchNodes(callback func(View)) func() {
|
||||||
|
da.mu.Lock()
|
||||||
|
id := da.nextID
|
||||||
|
da.nextID++
|
||||||
|
da.watchers = append(da.watchers, callback)
|
||||||
|
da.watcherIDs = append(da.watcherIDs, id)
|
||||||
|
da.mu.Unlock()
|
||||||
|
|
||||||
|
// Send initial view
|
||||||
|
callback(da.buildView())
|
||||||
|
|
||||||
|
// Return cancel function
|
||||||
|
closed := false
|
||||||
|
return func() {
|
||||||
|
if closed {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
closed = true
|
||||||
|
|
||||||
|
da.mu.Lock()
|
||||||
|
defer da.mu.Unlock()
|
||||||
|
|
||||||
|
// Find and remove the watcher
|
||||||
|
for i, watcherID := range da.watcherIDs {
|
||||||
|
if watcherID == id {
|
||||||
|
// Remove by swapping with last element
|
||||||
|
lastIdx := len(da.watchers) - 1
|
||||||
|
da.watchers[i] = da.watchers[lastIdx]
|
||||||
|
da.watcherIDs[i] = da.watcherIDs[lastIdx]
|
||||||
|
da.watchers = da.watchers[:lastIdx]
|
||||||
|
da.watcherIDs = da.watcherIDs[:lastIdx]
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// notifyWatchers notifies all registered watchers of view changes.
|
||||||
|
func (da *DiscoveryAdapter) notifyWatchers() {
|
||||||
|
view := da.buildView()
|
||||||
|
|
||||||
|
da.mu.RLock()
|
||||||
|
watchers := make([]func(View), len(da.watchers))
|
||||||
|
copy(watchers, da.watchers)
|
||||||
|
da.mu.RUnlock()
|
||||||
|
|
||||||
|
for _, w := range watchers {
|
||||||
|
go func(cb func(View)) {
|
||||||
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
// Log panic but don't crash
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
cb(view)
|
||||||
|
}(w)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildView builds a View from the current membership state.
|
||||||
|
func (da *DiscoveryAdapter) buildView() View {
|
||||||
|
members := da.ds.Members()
|
||||||
|
nodes := make([]*NodeInfo, 0, len(members))
|
||||||
|
for _, m := range members {
|
||||||
|
nodes = append(nodes, m.Node)
|
||||||
|
}
|
||||||
|
|
||||||
|
return View{
|
||||||
|
Nodes: nodes,
|
||||||
|
LocalNodeID: da.ds.LocalNode().ID,
|
||||||
|
Version: time.Now().UnixNano(),
|
||||||
|
Timestamp: time.Now().UnixNano(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetDiscovery returns the underlying Discovery service.
|
||||||
|
func (da *DiscoveryAdapter) GetDiscovery() Discovery {
|
||||||
|
return da.ds
|
||||||
|
}
|
||||||
74
pkg/swarm/discovery_interface.go
Normal file
74
pkg/swarm/discovery_interface.go
Normal file
|
|
@ -0,0 +1,74 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import "github.com/nats-io/nats.go"
|
||||||
|
|
||||||
|
// Discovery is the interface for node discovery implementations.
|
||||||
|
// In the current architecture, NATS is the sole transport.
|
||||||
|
//
|
||||||
|
//nolint:interfacebloat // Single concrete implementation; splitting would add indirection without benefit.
|
||||||
|
type Discovery interface {
|
||||||
|
// Start starts the discovery service.
|
||||||
|
Start() error
|
||||||
|
|
||||||
|
// Stop stops the discovery service.
|
||||||
|
Stop() error
|
||||||
|
|
||||||
|
// LocalNode returns the local node info.
|
||||||
|
LocalNode() *NodeInfo
|
||||||
|
|
||||||
|
// UpdateLocalInfo updates the local node's information.
|
||||||
|
UpdateLocalInfo(info *NodeInfo)
|
||||||
|
|
||||||
|
// UpdateLoad updates the local node's load score.
|
||||||
|
UpdateLoad(score float64)
|
||||||
|
|
||||||
|
// UpdateCapabilities updates the local node's agent capabilities.
|
||||||
|
UpdateCapabilities(caps map[string]string)
|
||||||
|
|
||||||
|
// Members returns all known members.
|
||||||
|
Members() []*NodeWithState
|
||||||
|
|
||||||
|
// GetNode returns a node by ID.
|
||||||
|
GetNode(nodeID string) (*NodeWithState, bool)
|
||||||
|
|
||||||
|
// Subscribe registers a handler for node events.
|
||||||
|
Subscribe(handler EventHandler) EventHandlerID
|
||||||
|
|
||||||
|
// Unsubscribe removes a node event handler.
|
||||||
|
Unsubscribe(id EventHandlerID)
|
||||||
|
|
||||||
|
// PublishLocalUpdate publishes local node state.
|
||||||
|
PublishLocalUpdate()
|
||||||
|
|
||||||
|
// DispatchEvent dispatches a node event to registered handlers.
|
||||||
|
DispatchEvent(event *NodeEvent)
|
||||||
|
|
||||||
|
// GetMembershipManager returns the membership manager.
|
||||||
|
GetMembershipManager() *MembershipManager
|
||||||
|
|
||||||
|
// NATSConn returns the underlying NATS connection.
|
||||||
|
// Used by handoff, leader election, and messaging subsystems.
|
||||||
|
NATSConn() *nats.Conn
|
||||||
|
|
||||||
|
// JetStream returns the JetStream context.
|
||||||
|
// Used for KV operations (membership, leader election).
|
||||||
|
JetStream() nats.JetStreamContext
|
||||||
|
|
||||||
|
// Buckets returns the KV bucket set.
|
||||||
|
// Used for leader election and other distributed state operations.
|
||||||
|
Buckets() *BucketSet
|
||||||
|
}
|
||||||
|
|
||||||
|
// DetailedStatusProvider is an optional capability interface for discovery
|
||||||
|
// implementations that maintain a local cache of detailed node status.
|
||||||
|
// Tools should use this interface via type assertion instead of depending
|
||||||
|
// on a concrete discovery type.
|
||||||
|
type DetailedStatusProvider interface {
|
||||||
|
GetDetailedStatus() *NodeDetailedStatus
|
||||||
|
}
|
||||||
456
pkg/swarm/discovery_nats.go
Normal file
456
pkg/swarm/discovery_nats.go
Normal file
|
|
@ -0,0 +1,456 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"net"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/nats-io/nats.go"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/kv"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DiscoveryService implements node discovery using NATS JetStream KeyValue store.
|
||||||
|
// This is the sole discovery implementation — all swarm communication goes through NATS.
|
||||||
|
type DiscoveryService struct {
|
||||||
|
config *Config
|
||||||
|
localNode *NodeInfo
|
||||||
|
membership *MembershipManager
|
||||||
|
eventHandler *EventDispatcher
|
||||||
|
nc *nats.Conn
|
||||||
|
js nats.JetStreamContext
|
||||||
|
buckets *BucketSet
|
||||||
|
|
||||||
|
detailedStatus *NodeDetailedStatus
|
||||||
|
loadMonitor *LoadMonitor
|
||||||
|
|
||||||
|
mu sync.RWMutex
|
||||||
|
running bool
|
||||||
|
stopChan chan struct{}
|
||||||
|
once sync.Once
|
||||||
|
|
||||||
|
seqNum uint64
|
||||||
|
startTime time.Time
|
||||||
|
todoList []string
|
||||||
|
activeTasks int
|
||||||
|
lastError string
|
||||||
|
version string
|
||||||
|
}
|
||||||
|
|
||||||
|
// generateStableNodeID generates a stable node ID that persists across restarts.
|
||||||
|
// It uses hostname as the base, optionally loading a previously generated ID from a file.
|
||||||
|
func generateStableNodeID() string {
|
||||||
|
hostname, _ := os.Hostname()
|
||||||
|
if hostname == "" {
|
||||||
|
hostname = "picoclaw"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to load a previously generated stable ID from a file
|
||||||
|
// This ensures the same ID is used across restarts even if hostname changes
|
||||||
|
dataDir := os.Getenv("PICOCLAW_DATA_DIR")
|
||||||
|
if dataDir == "" {
|
||||||
|
// Default to current directory
|
||||||
|
dataDir = "."
|
||||||
|
}
|
||||||
|
nodeIDFile := fmt.Sprintf("%s/.swarm_node_id", dataDir)
|
||||||
|
|
||||||
|
if idBytes, err := os.ReadFile(nodeIDFile); err == nil {
|
||||||
|
storedID := strings.TrimSpace(string(idBytes))
|
||||||
|
if storedID != "" {
|
||||||
|
logger.InfoCF("swarm", "Loaded stable node ID from file", map[string]any{"node_id": storedID})
|
||||||
|
return storedID
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Generate a stable ID based on hostname
|
||||||
|
// Using just hostname ensures stability across restarts
|
||||||
|
// For multiple instances on the same host, users should explicitly set PICOCLAW_SWARM_NODE_ID
|
||||||
|
stableID := hostname
|
||||||
|
|
||||||
|
// Optionally persist this ID for future use
|
||||||
|
_ = os.WriteFile(nodeIDFile, []byte(stableID), 0o644)
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Generated stable node ID", map[string]any{
|
||||||
|
"node_id": stableID,
|
||||||
|
"node_id_file": nodeIDFile,
|
||||||
|
"env_override": "PICOCLAW_SWARM_NODE_ID",
|
||||||
|
})
|
||||||
|
return stableID
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 ""
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewDiscoveryService creates a new NATS-based discovery service.
|
||||||
|
func NewDiscoveryService(cfg *Config) (*DiscoveryService, error) {
|
||||||
|
if cfg.NodeID == "" {
|
||||||
|
cfg.NodeID = generateStableNodeID()
|
||||||
|
logger.WarnCF(
|
||||||
|
"swarm",
|
||||||
|
"NodeID not configured, using auto-generated stable ID. For production, explicitly set PICOCLAW_SWARM_NODE_ID environment variable",
|
||||||
|
map[string]any{
|
||||||
|
"node_id": cfg.NodeID,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Determine advertise address
|
||||||
|
advAddr := getLocalIP()
|
||||||
|
if advAddr == "" {
|
||||||
|
advAddr = "127.0.0.1"
|
||||||
|
}
|
||||||
|
|
||||||
|
localNode := &NodeInfo{
|
||||||
|
ID: cfg.NodeID,
|
||||||
|
Addr: advAddr,
|
||||||
|
Port: 0, // No direct RPC port; all communication via NATS
|
||||||
|
AgentCaps: make(map[string]string),
|
||||||
|
LoadScore: 0,
|
||||||
|
Labels: make(map[string]string),
|
||||||
|
Timestamp: time.Now().UnixNano(),
|
||||||
|
Version: "1.0.0",
|
||||||
|
}
|
||||||
|
|
||||||
|
ds := &DiscoveryService{
|
||||||
|
config: cfg,
|
||||||
|
localNode: localNode,
|
||||||
|
eventHandler: NewEventDispatcher(),
|
||||||
|
stopChan: make(chan struct{}),
|
||||||
|
detailedStatus: NewNodeDetailedStatus(),
|
||||||
|
startTime: time.Now(),
|
||||||
|
todoList: make([]string, 0),
|
||||||
|
activeTasks: 0,
|
||||||
|
version: "1.0.0",
|
||||||
|
}
|
||||||
|
|
||||||
|
// Initialize membership manager
|
||||||
|
ds.membership = NewMembershipManager(ds, cfg.Discovery)
|
||||||
|
|
||||||
|
return ds, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildNATSOptions constructs NATS connection options from config.
|
||||||
|
func buildNATSOptions(cfg DiscoveryConfig, nodeID string) []nats.Option {
|
||||||
|
var opts []nats.Option
|
||||||
|
|
||||||
|
opts = append(opts, nats.Name("picoclaw-"+nodeID))
|
||||||
|
|
||||||
|
// NKeys / JWT credentials
|
||||||
|
if cfg.NATSCredsFile != "" {
|
||||||
|
opts = append(opts, nats.UserCredentials(cfg.NATSCredsFile))
|
||||||
|
}
|
||||||
|
|
||||||
|
// mTLS client certificate
|
||||||
|
if cfg.NATSTLSCert != "" && cfg.NATSTLSKey != "" {
|
||||||
|
opts = append(opts, nats.ClientCert(cfg.NATSTLSCert, cfg.NATSTLSKey))
|
||||||
|
}
|
||||||
|
|
||||||
|
// CA certificate for server verification
|
||||||
|
if cfg.NATSTLSCACert != "" {
|
||||||
|
opts = append(opts, nats.RootCAs(cfg.NATSTLSCACert))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Auto-reconnect
|
||||||
|
opts = append(opts, nats.MaxReconnects(-1))
|
||||||
|
opts = append(opts, nats.ReconnectWait(2*time.Second))
|
||||||
|
|
||||||
|
return opts
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start starts the NATS discovery service.
|
||||||
|
func (ds *DiscoveryService) Start() error {
|
||||||
|
ds.mu.Lock()
|
||||||
|
if ds.running {
|
||||||
|
ds.mu.Unlock()
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
ds.running = true
|
||||||
|
ds.mu.Unlock()
|
||||||
|
|
||||||
|
// Get NATS URL from config or use default
|
||||||
|
natsURL := ds.config.Discovery.NATSURL
|
||||||
|
if natsURL == "" {
|
||||||
|
natsURL = DefaultNATSURL
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build connection options (TLS, creds, reconnect)
|
||||||
|
opts := buildNATSOptions(ds.config.Discovery, ds.config.NodeID)
|
||||||
|
|
||||||
|
// Connect to NATS
|
||||||
|
var err error
|
||||||
|
ds.nc, err = nats.Connect(natsURL, opts...)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to connect to NATS: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get JetStream context
|
||||||
|
ds.js, err = ds.nc.JetStream()
|
||||||
|
if err != nil {
|
||||||
|
ds.nc.Close()
|
||||||
|
return fmt.Errorf("failed to get JetStream context: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create bucket factory and initialize buckets
|
||||||
|
factory := kv.NewNATSBucketFactory(ds.js)
|
||||||
|
ds.buckets, err = NewBucketSet(factory)
|
||||||
|
if err != nil {
|
||||||
|
ds.nc.Close()
|
||||||
|
return fmt.Errorf("failed to create buckets: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "NATS discovery service started", map[string]any{
|
||||||
|
"nats_url": natsURL,
|
||||||
|
"node_id": ds.localNode.ID,
|
||||||
|
})
|
||||||
|
|
||||||
|
// Start heartbeat routine
|
||||||
|
go ds.heartbeatLoop()
|
||||||
|
|
||||||
|
// Start detailed status publisher
|
||||||
|
go ds.detailedStatusLoop()
|
||||||
|
|
||||||
|
// Start watch for other nodes
|
||||||
|
go ds.watchNodes()
|
||||||
|
|
||||||
|
// Start watch for detailed status from other nodes
|
||||||
|
go ds.watchDetailedStatus()
|
||||||
|
|
||||||
|
// Add self to membership
|
||||||
|
ds.membership.UpdateNode(ds.localNode)
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop stops the NATS 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Delete our entry from members bucket
|
||||||
|
if ds.buckets != nil && ds.localNode != nil {
|
||||||
|
_ = ds.buckets.Members.Delete(ds.localNode.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close buckets
|
||||||
|
if ds.buckets != nil {
|
||||||
|
_ = ds.buckets.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
if ds.nc != nil {
|
||||||
|
ds.nc.Close()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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++
|
||||||
|
info = ds.localNode
|
||||||
|
ds.mu.Unlock()
|
||||||
|
|
||||||
|
ds.membership.UpdateNode(info)
|
||||||
|
ds.publishNode()
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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.publishNode()
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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.publishNode()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Members returns all known members.
|
||||||
|
func (ds *DiscoveryService) Members() []*NodeWithState {
|
||||||
|
return ds.membership.GetMembers()
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetNode returns a node by ID.
|
||||||
|
func (ds *DiscoveryService) GetNode(nodeID string) (*NodeWithState, bool) {
|
||||||
|
return ds.membership.GetNode(nodeID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Subscribe registers a handler for node events.
|
||||||
|
func (ds *DiscoveryService) Subscribe(handler EventHandler) EventHandlerID {
|
||||||
|
return ds.eventHandler.Subscribe(handler)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Unsubscribe removes a node event handler.
|
||||||
|
func (ds *DiscoveryService) Unsubscribe(id EventHandlerID) {
|
||||||
|
ds.eventHandler.Unsubscribe(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
// PublishLocalUpdate publishes local node state.
|
||||||
|
func (ds *DiscoveryService) PublishLocalUpdate() {
|
||||||
|
ds.publishNode()
|
||||||
|
}
|
||||||
|
|
||||||
|
// DispatchEvent dispatches a node event to registered handlers.
|
||||||
|
func (ds *DiscoveryService) DispatchEvent(event *NodeEvent) {
|
||||||
|
ds.eventHandler.Dispatch(event)
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetMembershipManager returns the membership manager.
|
||||||
|
func (ds *DiscoveryService) GetMembershipManager() *MembershipManager {
|
||||||
|
return ds.membership
|
||||||
|
}
|
||||||
|
|
||||||
|
// NATSConn returns the underlying NATS connection.
|
||||||
|
func (ds *DiscoveryService) NATSConn() *nats.Conn {
|
||||||
|
return ds.nc
|
||||||
|
}
|
||||||
|
|
||||||
|
// JetStream returns the JetStream context.
|
||||||
|
func (ds *DiscoveryService) JetStream() nats.JetStreamContext {
|
||||||
|
return ds.js
|
||||||
|
}
|
||||||
|
|
||||||
|
// Buckets returns the bucket set.
|
||||||
|
func (ds *DiscoveryService) Buckets() *BucketSet {
|
||||||
|
return ds.buckets
|
||||||
|
}
|
||||||
|
|
||||||
|
// Helper methods for todo list, active tasks, and error tracking.
|
||||||
|
|
||||||
|
// getTodoList returns the current todo list.
|
||||||
|
func (ds *DiscoveryService) getTodoList() []string {
|
||||||
|
ds.mu.RLock()
|
||||||
|
defer ds.mu.RUnlock()
|
||||||
|
return ds.todoList
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetTodoList sets the todo list for detailed status reporting.
|
||||||
|
func (ds *DiscoveryService) SetTodoList(todos []string) {
|
||||||
|
ds.mu.Lock()
|
||||||
|
ds.todoList = todos
|
||||||
|
ds.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// AddTodo adds a todo item to the todo list.
|
||||||
|
func (ds *DiscoveryService) AddTodo(todo string) {
|
||||||
|
ds.mu.Lock()
|
||||||
|
defer ds.mu.Unlock()
|
||||||
|
ds.todoList = append(ds.todoList, todo)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RemoveTodo removes a todo item from the todo list.
|
||||||
|
func (ds *DiscoveryService) RemoveTodo(todo string) {
|
||||||
|
ds.mu.Lock()
|
||||||
|
defer ds.mu.Unlock()
|
||||||
|
for i, t := range ds.todoList {
|
||||||
|
if t == todo {
|
||||||
|
ds.todoList = append(ds.todoList[:i], ds.todoList[i+1:]...)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ClearTodoList clears the todo list.
|
||||||
|
func (ds *DiscoveryService) ClearTodoList() {
|
||||||
|
ds.mu.Lock()
|
||||||
|
ds.todoList = make([]string, 0)
|
||||||
|
ds.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// getActiveTasks returns the current active tasks count.
|
||||||
|
func (ds *DiscoveryService) getActiveTasks() int {
|
||||||
|
ds.mu.RLock()
|
||||||
|
defer ds.mu.RUnlock()
|
||||||
|
return ds.activeTasks
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetActiveTasks sets the active tasks count.
|
||||||
|
func (ds *DiscoveryService) SetActiveTasks(count int) {
|
||||||
|
ds.mu.Lock()
|
||||||
|
ds.activeTasks = count
|
||||||
|
ds.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// IncrementActiveTasks increments the active tasks count.
|
||||||
|
func (ds *DiscoveryService) IncrementActiveTasks() {
|
||||||
|
ds.mu.Lock()
|
||||||
|
ds.activeTasks++
|
||||||
|
ds.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// DecrementActiveTasks decrements the active tasks count.
|
||||||
|
func (ds *DiscoveryService) DecrementActiveTasks() {
|
||||||
|
ds.mu.Lock()
|
||||||
|
if ds.activeTasks > 0 {
|
||||||
|
ds.activeTasks--
|
||||||
|
}
|
||||||
|
ds.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// getLastError returns the last error.
|
||||||
|
func (ds *DiscoveryService) getLastError() string {
|
||||||
|
ds.mu.RLock()
|
||||||
|
defer ds.mu.RUnlock()
|
||||||
|
return ds.lastError
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetLastError sets the last error for detailed status reporting.
|
||||||
|
func (ds *DiscoveryService) SetLastError(err string) {
|
||||||
|
ds.mu.Lock()
|
||||||
|
ds.lastError = err
|
||||||
|
ds.mu.Unlock()
|
||||||
|
}
|
||||||
81
pkg/swarm/errors.go
Normal file
81
pkg/swarm/errors.go
Normal file
|
|
@ -0,0 +1,81 @@
|
||||||
|
// 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")
|
||||||
|
|
||||||
|
// ErrHandoffNoResponders is returned when no nodes responded to a handoff request.
|
||||||
|
ErrHandoffNoResponders = errors.New("no nodes responded to handoff request")
|
||||||
|
|
||||||
|
// 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")
|
||||||
|
|
||||||
|
// 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")
|
||||||
|
|
||||||
|
// ErrNATSNotConnected is returned when the NATS connection is not established.
|
||||||
|
ErrNATSNotConnected = errors.New("NATS not connected")
|
||||||
|
|
||||||
|
// ErrLeaderLockFailed is returned when leader lock acquisition fails.
|
||||||
|
ErrLeaderLockFailed = errors.New("leader lock acquisition failed")
|
||||||
|
|
||||||
|
// ErrLeaderLockLost is returned when the leader lock is lost (CAS renewal failed).
|
||||||
|
ErrLeaderLockLost = errors.New("leader lock lost")
|
||||||
|
|
||||||
|
// ErrKVBucketUnavailable is returned when a required KV bucket cannot be accessed.
|
||||||
|
ErrKVBucketUnavailable = errors.New("KV bucket unavailable")
|
||||||
|
)
|
||||||
|
|
||||||
|
// IsBusinessRejection returns true if the error represents a legitimate business rejection
|
||||||
|
// (e.g., no nodes available, node overloaded) rather than a system failure.
|
||||||
|
// Business rejections should be communicated via HandoffResponse.Accepted=false.
|
||||||
|
// System failures should be returned as Go errors for proper handling.
|
||||||
|
func IsBusinessRejection(err error) bool {
|
||||||
|
if err == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return err == ErrNoHealthyNodes ||
|
||||||
|
err == ErrNodeNotAvailable ||
|
||||||
|
err == ErrHandoffRejected ||
|
||||||
|
err == ErrHandoffInProgress ||
|
||||||
|
err == ErrCapabilityNotSupported
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsSystemError returns true if the error represents a system failure
|
||||||
|
// (e.g., NATS unavailable, timeout, serialization failure).
|
||||||
|
// System errors should propagate up the call stack as Go errors.
|
||||||
|
func IsSystemError(err error) bool {
|
||||||
|
if err == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return !IsBusinessRejection(err)
|
||||||
|
}
|
||||||
956
pkg/swarm/handoff.go
Normal file
956
pkg/swarm/handoff.go
Normal file
|
|
@ -0,0 +1,956 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
"github.com/nats-io/nats.go"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// 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"
|
||||||
|
HandoffStateDuplicate HandoffState = "duplicate" // Replay attack detected
|
||||||
|
HandoffStateExpired HandoffState = "expired" // Timestamp outside valid window
|
||||||
|
)
|
||||||
|
|
||||||
|
// Security constants for replay attack prevention.
|
||||||
|
const (
|
||||||
|
// DefaultTimestampWindow is the acceptable timestamp skew (±60 seconds).
|
||||||
|
DefaultTimestampWindow = 60 * time.Second
|
||||||
|
|
||||||
|
// DefaultNonceTTL is how long nonces are kept in cache.
|
||||||
|
DefaultNonceTTL = 5 * time.Minute
|
||||||
|
|
||||||
|
// DefaultNonceCacheSize is the ring buffer size for nonce tracking.
|
||||||
|
DefaultNonceCacheSize = 1000
|
||||||
|
)
|
||||||
|
|
||||||
|
// SessionMessage represents a message in a session (local type replacing picolib.SessionMessage).
|
||||||
|
type SessionMessage struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// HandoffRequest represents a request to hand off a session.
|
||||||
|
type HandoffRequest struct {
|
||||||
|
// Core identification
|
||||||
|
RequestID string `json:"request_id"` // Unique handoff ID for idempotency
|
||||||
|
|
||||||
|
// Handoff details
|
||||||
|
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"`
|
||||||
|
|
||||||
|
// Routing information
|
||||||
|
FromNodeID string `json:"from_node_id"`
|
||||||
|
FromNodeAddr string `json:"from_node_addr"`
|
||||||
|
TargetNodeID string `json:"target_node_id,omitempty"`
|
||||||
|
|
||||||
|
// Timestamp for replay protection (Unix nanoseconds)
|
||||||
|
Timestamp int64 `json:"timestamp"`
|
||||||
|
|
||||||
|
// Security fields
|
||||||
|
Nonce string `json:"nonce,omitempty"` // Unique value for replay protection
|
||||||
|
Signature string `json:"signature,omitempty"` // HMAC-SHA256 signature
|
||||||
|
EncryptedData *EncryptedData `json:"encrypted,omitempty"` // AES-GCM encrypted session/context
|
||||||
|
}
|
||||||
|
|
||||||
|
// signingPayload returns the canonical payload for signature calculation.
|
||||||
|
// Covers all routing fields to prevent message redirection attacks.
|
||||||
|
func (r *HandoffRequest) signingPayload() map[string]any {
|
||||||
|
// Include payload hash for large payloads
|
||||||
|
payloadHash := ""
|
||||||
|
if len(r.SessionMessages) > 0 || len(r.Context) > 0 {
|
||||||
|
payloadData := map[string]any{
|
||||||
|
"session_messages": r.SessionMessages,
|
||||||
|
"context": r.Context,
|
||||||
|
}
|
||||||
|
if j, err := json.Marshal(payloadData); err == nil {
|
||||||
|
hash := sha256.Sum256(j)
|
||||||
|
payloadHash = fmt.Sprintf("%x", hash)[:16]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return map[string]any{
|
||||||
|
"request_id": r.RequestID,
|
||||||
|
"nonce": r.Nonce,
|
||||||
|
"from_node_id": r.FromNodeID,
|
||||||
|
"to_node_id": r.TargetNodeID, // Prevents redirection attacks
|
||||||
|
"session_key": r.SessionKey,
|
||||||
|
"required_cap": r.RequiredCap,
|
||||||
|
"timestamp": r.Timestamp,
|
||||||
|
"payload_hash": payloadHash,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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"`
|
||||||
|
|
||||||
|
// TruncationInfo indicates if session history was truncated during handoff.
|
||||||
|
TruncationInfo *TruncationInfo `json:"truncation_info,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// TruncationInfo contains details about message truncation.
|
||||||
|
type TruncationInfo struct {
|
||||||
|
OriginalCount int `json:"original_count"`
|
||||||
|
RemainingCount int `json:"remaining_count"`
|
||||||
|
BytesRemoved int `json:"bytes_removed"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// HandoffCoordinator coordinates handoff operations between nodes via NATS request-reply.
|
||||||
|
//
|
||||||
|
// Each node subscribes to its own targeted handoff subject:
|
||||||
|
//
|
||||||
|
// picoclaw.swarm.handoff.<localNodeID>
|
||||||
|
//
|
||||||
|
// When initiating a handoff, the coordinator sends a NATS request to the target node's
|
||||||
|
// handoff subject and waits for a synchronous reply.
|
||||||
|
type HandoffCoordinator struct {
|
||||||
|
discovery Discovery
|
||||||
|
membership *MembershipManager
|
||||||
|
config HandoffConfig
|
||||||
|
nc *nats.Conn
|
||||||
|
|
||||||
|
pending map[string]*HandoffOperation // request_id -> operation
|
||||||
|
mu sync.RWMutex
|
||||||
|
|
||||||
|
sub *nats.Subscription // Subscription for incoming handoff requests
|
||||||
|
|
||||||
|
// Security
|
||||||
|
signer *Signer
|
||||||
|
encryptor *Encryptor
|
||||||
|
cryptoConfig CryptoConfig
|
||||||
|
nonceCache *NonceCache
|
||||||
|
timestampWindow time.Duration
|
||||||
|
|
||||||
|
// Idempotency: track processed handoff IDs
|
||||||
|
processed map[string]*HandoffResponse // request_id -> cached response
|
||||||
|
processedMu sync.RWMutex
|
||||||
|
|
||||||
|
// Failure policies
|
||||||
|
circuitBreaker *CircuitBreaker
|
||||||
|
nodeCooldown *NodeCooldown
|
||||||
|
failurePolicy FailurePolicy
|
||||||
|
|
||||||
|
// Accept/reject callbacks
|
||||||
|
onHandoffRequest func(*HandoffRequest) *HandoffResponse
|
||||||
|
onHandoffComplete func(*HandoffRequest, *HandoffResponse)
|
||||||
|
|
||||||
|
// Metrics
|
||||||
|
metrics *HandoffMetrics
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 Discovery, config HandoffConfig) *HandoffCoordinator {
|
||||||
|
policy := DefaultFailurePolicy()
|
||||||
|
return &HandoffCoordinator{
|
||||||
|
discovery: ds,
|
||||||
|
membership: ds.GetMembershipManager(),
|
||||||
|
config: config,
|
||||||
|
nc: ds.NATSConn(),
|
||||||
|
pending: make(map[string]*HandoffOperation),
|
||||||
|
processed: make(map[string]*HandoffResponse),
|
||||||
|
metrics: NewHandoffMetrics(),
|
||||||
|
nonceCache: NewNonceCache(DefaultNonceTTL, DefaultNonceCacheSize),
|
||||||
|
timestampWindow: DefaultTimestampWindow,
|
||||||
|
circuitBreaker: NewCircuitBreaker(policy.MaxConsecutiveFailures, policy.CircuitBreakerCooldown.Duration),
|
||||||
|
nodeCooldown: NewNodeCooldown(policy.OverloadCooldown.Duration),
|
||||||
|
failurePolicy: *policy,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetCryptoConfig configures the cryptographic components for secure handoff.
|
||||||
|
// Returns error if RequireAuth or RequireEncryption is enabled without proper credentials.
|
||||||
|
func (hc *HandoffCoordinator) SetCryptoConfig(cfg CryptoConfig) error {
|
||||||
|
hc.cryptoConfig = cfg
|
||||||
|
hc.signer = NewSigner(cfg.SharedSecret)
|
||||||
|
|
||||||
|
if cfg.EncryptionKey != "" {
|
||||||
|
enc, err := NewEncryptor(cfg.EncryptionKey)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to initialize encryptor: %w", err)
|
||||||
|
}
|
||||||
|
hc.encryptor = enc
|
||||||
|
}
|
||||||
|
|
||||||
|
// Strict validation: error out if auth/encryption required but credentials missing
|
||||||
|
if cfg.RequireAuth && cfg.SharedSecret == "" {
|
||||||
|
return fmt.Errorf("RequireAuth=true but SharedSecret is empty; authentication cannot be enforced")
|
||||||
|
}
|
||||||
|
|
||||||
|
if cfg.RequireEncryption && cfg.EncryptionKey == "" {
|
||||||
|
return fmt.Errorf("RequireEncryption=true but EncryptionKey is empty; encryption cannot be enforced")
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start begins listening for incoming handoff requests on the targeted NATS subject.
|
||||||
|
func (hc *HandoffCoordinator) Start() error {
|
||||||
|
if hc.nc == nil {
|
||||||
|
return ErrNATSNotConnected
|
||||||
|
}
|
||||||
|
|
||||||
|
localNodeID := hc.discovery.LocalNode().ID
|
||||||
|
subject := SubjectHandoff + "." + localNodeID
|
||||||
|
|
||||||
|
var err error
|
||||||
|
hc.sub, err = hc.nc.Subscribe(subject, hc.handleIncomingNATSMsg)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to subscribe to handoff subject %s: %w", subject, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Handoff coordinator started", map[string]any{
|
||||||
|
"subject": subject,
|
||||||
|
})
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Close cleans up the handoff coordinator.
|
||||||
|
func (hc *HandoffCoordinator) Close() error {
|
||||||
|
if hc.sub != nil {
|
||||||
|
return hc.sub.Unsubscribe()
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleIncomingNATSMsg handles a NATS message containing a HandoffRequest.
|
||||||
|
func (hc *HandoffCoordinator) handleIncomingNATSMsg(msg *nats.Msg) {
|
||||||
|
var req HandoffRequest
|
||||||
|
if err := json.Unmarshal(msg.Data, &req); err != nil {
|
||||||
|
hc.metrics.RecordFailed()
|
||||||
|
logger.ErrorCF("swarm", "Failed to decode handoff request", map[string]any{"error": err})
|
||||||
|
hc.rejectHandoff(msg, "", "invalid_format", "")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate required fields
|
||||||
|
if req.RequestID == "" || req.FromNodeID == "" {
|
||||||
|
hc.metrics.RecordFailed()
|
||||||
|
hc.rejectHandoff(msg, req.RequestID, "missing_required_fields", "")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check idempotency: was this handoff already processed?
|
||||||
|
hc.processedMu.RLock()
|
||||||
|
if cachedResp, exists := hc.processed[req.RequestID]; exists {
|
||||||
|
hc.processedMu.RUnlock()
|
||||||
|
logger.InfoCF("swarm", "Returning cached handoff response (duplicate)", map[string]any{
|
||||||
|
"request_id": req.RequestID,
|
||||||
|
"from_node": req.FromNodeID,
|
||||||
|
})
|
||||||
|
respData, _ := json.Marshal(cachedResp)
|
||||||
|
msg.Respond(respData)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
hc.processedMu.RUnlock()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
// Timestamp validation: check for replay attacks
|
||||||
|
if req.Timestamp > 0 {
|
||||||
|
reqTime := time.Unix(0, req.Timestamp)
|
||||||
|
skew := reqTime.Sub(now)
|
||||||
|
|
||||||
|
// Log clock skew for debugging (even if within window)
|
||||||
|
if skew.Abs() > time.Second {
|
||||||
|
logger.DebugCF("swarm", "Handoff timestamp skew", map[string]any{
|
||||||
|
"skew_ms": skew.Milliseconds(),
|
||||||
|
"from_node": req.FromNodeID,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
if skew.Abs() > hc.timestampWindow {
|
||||||
|
hc.metrics.RecordAuthFailure()
|
||||||
|
logger.WarnCF("swarm", "Handoff request timestamp outside valid window", map[string]any{
|
||||||
|
"skew_seconds": skew.Seconds(),
|
||||||
|
"window_seconds": hc.timestampWindow.Seconds(),
|
||||||
|
"from_node": req.FromNodeID,
|
||||||
|
"request_id": req.RequestID,
|
||||||
|
})
|
||||||
|
resp := &HandoffResponse{
|
||||||
|
RequestID: req.RequestID,
|
||||||
|
Accepted: false,
|
||||||
|
Reason: "unauthorized",
|
||||||
|
State: HandoffStateExpired,
|
||||||
|
Timestamp: now.UnixNano(),
|
||||||
|
}
|
||||||
|
hc.cacheAndRespond(msg, req.RequestID, resp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Nonce validation: check for replay attacks
|
||||||
|
if req.Nonce != "" && hc.cryptoConfig.RequireAuth {
|
||||||
|
if !hc.nonceCache.CheckAndAdd(req.Nonce, req.FromNodeID) {
|
||||||
|
hc.metrics.RecordAuthFailure()
|
||||||
|
logger.WarnCF("swarm", "Handoff request with duplicate nonce rejected (replay attack)", map[string]any{
|
||||||
|
"from_node": req.FromNodeID,
|
||||||
|
"request_id": req.RequestID,
|
||||||
|
})
|
||||||
|
resp := &HandoffResponse{
|
||||||
|
RequestID: req.RequestID,
|
||||||
|
Accepted: false,
|
||||||
|
Reason: "unauthorized",
|
||||||
|
State: HandoffStateDuplicate,
|
||||||
|
Timestamp: now.UnixNano(),
|
||||||
|
}
|
||||||
|
hc.cacheAndRespond(msg, req.RequestID, resp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify signature if authentication is required
|
||||||
|
if hc.signer != nil && hc.cryptoConfig.RequireAuth {
|
||||||
|
if !hc.signer.Verify(req.signingPayload(), req.Signature) {
|
||||||
|
hc.metrics.RecordAuthFailure()
|
||||||
|
logger.WarnCF("swarm", "Handoff request signature verification failed", map[string]any{
|
||||||
|
"from_node": req.FromNodeID,
|
||||||
|
"request_id": req.RequestID,
|
||||||
|
})
|
||||||
|
resp := &HandoffResponse{
|
||||||
|
RequestID: req.RequestID,
|
||||||
|
Accepted: false,
|
||||||
|
Reason: "unauthorized",
|
||||||
|
State: HandoffStateRejected,
|
||||||
|
Timestamp: now.UnixNano(),
|
||||||
|
}
|
||||||
|
hc.cacheAndRespond(msg, req.RequestID, resp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Don't self-handoff
|
||||||
|
localNodeID := hc.discovery.LocalNode().ID
|
||||||
|
if req.FromNodeID == localNodeID {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Decrypt session data if encrypted
|
||||||
|
if req.EncryptedData != nil && hc.encryptor != nil {
|
||||||
|
decrypted, err := hc.encryptor.Decrypt(req.EncryptedData)
|
||||||
|
if err != nil {
|
||||||
|
hc.metrics.RecordDecryptError()
|
||||||
|
logger.WarnCF("swarm", "Handoff data decryption failed", map[string]any{
|
||||||
|
"from_node": req.FromNodeID,
|
||||||
|
"request_id": req.RequestID,
|
||||||
|
})
|
||||||
|
resp := &HandoffResponse{
|
||||||
|
RequestID: req.RequestID,
|
||||||
|
Accepted: false,
|
||||||
|
Reason: "unauthorized",
|
||||||
|
State: HandoffStateRejected,
|
||||||
|
Timestamp: now.UnixNano(),
|
||||||
|
}
|
||||||
|
hc.cacheAndRespond(msg, req.RequestID, resp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Unmarshal the decrypted data into a temporary struct to extract session/context
|
||||||
|
var decryptedData struct {
|
||||||
|
SessionMessages []SessionMessage `json:"session_messages"`
|
||||||
|
Context map[string]any `json:"context"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(decrypted, &decryptedData); err == nil {
|
||||||
|
req.SessionMessages = decryptedData.SessionMessages
|
||||||
|
req.Context = decryptedData.Context
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Process the request
|
||||||
|
resp := hc.HandleIncomingHandoff(&req)
|
||||||
|
resp.Timestamp = now.UnixNano()
|
||||||
|
|
||||||
|
// Cache response for idempotency
|
||||||
|
hc.cacheAndRespond(msg, req.RequestID, resp)
|
||||||
|
}
|
||||||
|
|
||||||
|
// rejectHandoff sends a generic rejection response without caching.
|
||||||
|
func (hc *HandoffCoordinator) rejectHandoff(msg *nats.Msg, requestID, reason, state string) {
|
||||||
|
resp := &HandoffResponse{
|
||||||
|
RequestID: requestID,
|
||||||
|
Accepted: false,
|
||||||
|
Reason: "unauthorized", // Always return generic reason to caller
|
||||||
|
State: HandoffStateRejected,
|
||||||
|
Timestamp: time.Now().UnixNano(),
|
||||||
|
}
|
||||||
|
if state != "" {
|
||||||
|
resp.State = HandoffState(state)
|
||||||
|
}
|
||||||
|
respData, _ := json.Marshal(resp)
|
||||||
|
msg.Respond(respData)
|
||||||
|
|
||||||
|
// Log the real reason internally
|
||||||
|
logger.DebugCF("swarm", "Handoff rejected", map[string]any{
|
||||||
|
"request_id": requestID,
|
||||||
|
"reason": reason,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// cacheAndRespond caches the response and sends it.
|
||||||
|
func (hc *HandoffCoordinator) cacheAndRespond(msg *nats.Msg, requestID string, resp *HandoffResponse) {
|
||||||
|
hc.processedMu.Lock()
|
||||||
|
hc.processed[requestID] = resp
|
||||||
|
hc.processedMu.Unlock()
|
||||||
|
|
||||||
|
respData, err := json.Marshal(resp)
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to encode handoff response", map[string]any{"error": err})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if err := msg.Respond(respData); err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to send handoff response", map[string]any{"error": err})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
localNode := hc.discovery.LocalNode()
|
||||||
|
loadScore := localNode.LoadScore
|
||||||
|
if loadScore > hc.config.LoadThreshold {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check capability
|
||||||
|
if requiredCap != "" {
|
||||||
|
hasCap := false
|
||||||
|
for _, cap := range localNode.AgentCaps {
|
||||||
|
if cap == requiredCap {
|
||||||
|
hasCap = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !hasCap {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// InitiateHandoff initiates a handoff to another node via NATS request-reply.
|
||||||
|
func (hc *HandoffCoordinator) InitiateHandoff(ctx context.Context, req *HandoffRequest) (*HandoffResponse, error) {
|
||||||
|
if hc.nc == nil {
|
||||||
|
return nil, ErrNATSNotConnected
|
||||||
|
}
|
||||||
|
|
||||||
|
if req.RequestID == "" {
|
||||||
|
req.RequestID = uuid.New().String()
|
||||||
|
}
|
||||||
|
|
||||||
|
localNode := hc.discovery.LocalNode()
|
||||||
|
req.FromNodeID = localNode.ID
|
||||||
|
req.FromNodeAddr = localNode.Addr
|
||||||
|
req.Timestamp = time.Now().UnixNano()
|
||||||
|
|
||||||
|
// Generate nonce for replay protection if auth is enabled
|
||||||
|
if hc.signer != nil && hc.cryptoConfig.RequireAuth && req.Nonce == "" {
|
||||||
|
req.Nonce = uuid.New().String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Find target node
|
||||||
|
targetNode, findErr := hc.findTargetNode(req)
|
||||||
|
if findErr != nil {
|
||||||
|
//nolint:nilerr // Failure is communicated via HandoffResponse, not Go error
|
||||||
|
return &HandoffResponse{
|
||||||
|
RequestID: req.RequestID,
|
||||||
|
Accepted: false,
|
||||||
|
Reason: findErr.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 via NATS request-reply to targeted subject
|
||||||
|
req.TargetNodeID = targetNode.Node.ID
|
||||||
|
resp, err := hc.sendHandoffRequest(ctx, targetNode.Node.ID, req)
|
||||||
|
if err != nil {
|
||||||
|
// System error - return immediately, don't retry
|
||||||
|
hc.mu.Lock()
|
||||||
|
delete(hc.pending, req.RequestID)
|
||||||
|
hc.mu.Unlock()
|
||||||
|
return nil, fmt.Errorf("handoff request failed: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Retry if needed
|
||||||
|
retryLoop:
|
||||||
|
for !resp.Accepted && op.RetryCount < hc.config.MaxRetries {
|
||||||
|
op.RetryCount++
|
||||||
|
hc.metrics.RecordRetry()
|
||||||
|
|
||||||
|
// Find new target
|
||||||
|
newTarget, findErr := hc.findTargetNode(req)
|
||||||
|
if findErr != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
op.TargetNode = newTarget
|
||||||
|
|
||||||
|
// Delay before retry
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
resp = &HandoffResponse{
|
||||||
|
RequestID: req.RequestID,
|
||||||
|
Accepted: false,
|
||||||
|
Reason: "context canceled",
|
||||||
|
State: HandoffStateFailed,
|
||||||
|
}
|
||||||
|
|
||||||
|
break retryLoop
|
||||||
|
case <-time.After(hc.config.RetryDelay.Duration):
|
||||||
|
}
|
||||||
|
|
||||||
|
req.TargetNodeID = newTarget.Node.ID
|
||||||
|
var retryErr error
|
||||||
|
resp, retryErr = hc.sendHandoffRequest(ctx, newTarget.Node.ID, req)
|
||||||
|
if retryErr != nil {
|
||||||
|
// System error during retry - abort
|
||||||
|
hc.mu.Lock()
|
||||||
|
delete(hc.pending, req.RequestID)
|
||||||
|
hc.mu.Unlock()
|
||||||
|
return nil, fmt.Errorf("handoff retry failed: %w", retryErr)
|
||||||
|
}
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
// sendHandoffRequest sends a handoff request to a specific target node via NATS request-reply.
|
||||||
|
// Returns error for system failures (NATS timeout, network error, serialization failure).
|
||||||
|
// Returns HandoffResponse for business rejections (target node rejected).
|
||||||
|
func (hc *HandoffCoordinator) sendHandoffRequest(
|
||||||
|
ctx context.Context,
|
||||||
|
targetNodeID string,
|
||||||
|
req *HandoffRequest,
|
||||||
|
) (*HandoffResponse, error) {
|
||||||
|
startTime := time.Now()
|
||||||
|
hc.metrics.RecordRequest()
|
||||||
|
|
||||||
|
subject := SubjectHandoff + "." + targetNodeID
|
||||||
|
|
||||||
|
// Prepare the request for transmission
|
||||||
|
transmitReq := *req
|
||||||
|
var truncInfo *TruncationInfo
|
||||||
|
|
||||||
|
// Data minimization: limit session history to configured max
|
||||||
|
// This reduces payload size and avoids sending unnecessary historical data
|
||||||
|
if hc.config.MaxSessionHistory > 0 && len(transmitReq.SessionMessages) > hc.config.MaxSessionHistory {
|
||||||
|
originalCount := len(transmitReq.SessionMessages)
|
||||||
|
|
||||||
|
// Preserve system/developer messages, then take most recent conversation messages
|
||||||
|
var systemMessages []SessionMessage
|
||||||
|
var conversationMessages []SessionMessage
|
||||||
|
|
||||||
|
for _, msg := range transmitReq.SessionMessages {
|
||||||
|
if msg.Role == "system" || msg.Role == "developer" {
|
||||||
|
systemMessages = append(systemMessages, msg)
|
||||||
|
} else {
|
||||||
|
conversationMessages = append(conversationMessages, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Keep all system messages, but limit conversation to MaxSessionHistory
|
||||||
|
availableSlots := hc.config.MaxSessionHistory - len(systemMessages)
|
||||||
|
if availableSlots < 0 {
|
||||||
|
availableSlots = 0
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(conversationMessages) > availableSlots {
|
||||||
|
// Take the most recent messages (from the end)
|
||||||
|
startIdx := len(conversationMessages) - availableSlots
|
||||||
|
conversationMessages = conversationMessages[startIdx:]
|
||||||
|
}
|
||||||
|
|
||||||
|
transmitReq.SessionMessages = append(systemMessages, conversationMessages...)
|
||||||
|
|
||||||
|
if len(transmitReq.SessionMessages) < originalCount {
|
||||||
|
logger.InfoCF("swarm", "Applied data minimization for handoff", map[string]any{
|
||||||
|
"original_count": originalCount,
|
||||||
|
"transmitted_count": len(transmitReq.SessionMessages),
|
||||||
|
"max_history": hc.config.MaxSessionHistory,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check message size against NATS max payload.
|
||||||
|
// If too large, truncate session history with smart rules.
|
||||||
|
data, err := json.Marshal(transmitReq)
|
||||||
|
if err != nil {
|
||||||
|
hc.metrics.RecordFailed()
|
||||||
|
return nil, fmt.Errorf("failed to marshal handoff request: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
originalSize := len(data)
|
||||||
|
originalCount := len(transmitReq.SessionMessages)
|
||||||
|
|
||||||
|
if len(data) > DefaultMaxHandoffPayloadBytes && len(transmitReq.SessionMessages) > 0 {
|
||||||
|
logger.WarnCF("swarm", "Handoff payload exceeds max size, truncating session history", map[string]any{
|
||||||
|
"original_size": originalSize,
|
||||||
|
"max_size": DefaultMaxHandoffPayloadBytes,
|
||||||
|
"original_history": originalCount,
|
||||||
|
})
|
||||||
|
|
||||||
|
transmitReq, data, truncInfo = hc.truncateSessionData(transmitReq, data)
|
||||||
|
hc.metrics.RecordTruncation()
|
||||||
|
|
||||||
|
if data == nil {
|
||||||
|
hc.metrics.RecordFailed()
|
||||||
|
return nil, fmt.Errorf("handoff payload too large even after truncation")
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Session history truncated for handoff", map[string]any{
|
||||||
|
"original_count": originalCount,
|
||||||
|
"remaining_count": len(transmitReq.SessionMessages),
|
||||||
|
"bytes_removed": originalSize - len(data),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Encrypt session data if configured
|
||||||
|
if hc.encryptor != nil && hc.cryptoConfig.RequireEncryption {
|
||||||
|
sessionData := map[string]any{
|
||||||
|
"session_messages": transmitReq.SessionMessages,
|
||||||
|
"context": transmitReq.Context,
|
||||||
|
}
|
||||||
|
plainData, _ := json.Marshal(sessionData)
|
||||||
|
encrypted, encryptErr := hc.encryptor.Encrypt(plainData)
|
||||||
|
if encryptErr == nil {
|
||||||
|
transmitReq.EncryptedData = encrypted
|
||||||
|
transmitReq.SessionMessages = nil // Clear plaintext
|
||||||
|
transmitReq.Context = nil
|
||||||
|
}
|
||||||
|
data, err = json.Marshal(transmitReq)
|
||||||
|
if err != nil {
|
||||||
|
hc.metrics.RecordFailed()
|
||||||
|
return nil, fmt.Errorf("failed to marshal handoff request after encryption: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sign the request
|
||||||
|
if hc.signer != nil {
|
||||||
|
signature, signErr := hc.signer.Sign(transmitReq.signingPayload())
|
||||||
|
if signErr != nil {
|
||||||
|
logger.WarnCF("swarm", "Failed to sign handoff request", map[string]any{"error": signErr})
|
||||||
|
} else {
|
||||||
|
transmitReq.Signature = signature
|
||||||
|
data, _ = json.Marshal(transmitReq)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Use NATS request-reply (timeout controlled by parent context)
|
||||||
|
msg, err := hc.nc.RequestWithContext(ctx, subject, data)
|
||||||
|
if err != nil {
|
||||||
|
hc.metrics.RecordTimeout()
|
||||||
|
hc.circuitBreaker.RecordFailure(targetNodeID)
|
||||||
|
// Return system error for network/timeout failures
|
||||||
|
return nil, fmt.Errorf("NATS request failed: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var resp HandoffResponse
|
||||||
|
if err := json.Unmarshal(msg.Data, &resp); err != nil {
|
||||||
|
hc.metrics.RecordFailed()
|
||||||
|
hc.circuitBreaker.RecordFailure(targetNodeID)
|
||||||
|
// Return system error for data corruption
|
||||||
|
return nil, fmt.Errorf("failed to decode handoff response: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Attach truncation info to response
|
||||||
|
resp.TruncationInfo = truncInfo
|
||||||
|
|
||||||
|
if resp.Accepted {
|
||||||
|
hc.metrics.RecordAccepted()
|
||||||
|
hc.circuitBreaker.RecordSuccess(targetNodeID)
|
||||||
|
hc.nodeCooldown.Remove(targetNodeID)
|
||||||
|
} else {
|
||||||
|
hc.metrics.RecordRejected()
|
||||||
|
hc.circuitBreaker.RecordFailure(targetNodeID)
|
||||||
|
// If rejected due to overload, add to cooldown
|
||||||
|
if resp.State == HandoffStateRejected && (resp.Reason == "cannot handle" || resp.Reason == "overloaded") {
|
||||||
|
hc.nodeCooldown.Add(targetNodeID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
hc.metrics.RecordLatency(time.Since(startTime))
|
||||||
|
|
||||||
|
return &resp, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// truncateSessionData smartly truncates session history, preserving important messages.
|
||||||
|
func (hc *HandoffCoordinator) truncateSessionData(
|
||||||
|
req HandoffRequest,
|
||||||
|
data []byte,
|
||||||
|
) (HandoffRequest, []byte, *TruncationInfo) {
|
||||||
|
truncated := req
|
||||||
|
originalCount := len(truncated.SessionMessages)
|
||||||
|
bytesRemoved := 0
|
||||||
|
|
||||||
|
// First pass: Remove oldest user/assistant messages, preserve system messages
|
||||||
|
var systemMessages []SessionMessage
|
||||||
|
var conversationMessages []SessionMessage
|
||||||
|
|
||||||
|
for _, msg := range truncated.SessionMessages {
|
||||||
|
if msg.Role == "system" || msg.Role == "developer" {
|
||||||
|
systemMessages = append(systemMessages, msg)
|
||||||
|
} else {
|
||||||
|
conversationMessages = append(conversationMessages, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Keep system messages, remove from oldest conversation messages first
|
||||||
|
for len(data) > DefaultMaxHandoffPayloadBytes && len(conversationMessages) > 0 {
|
||||||
|
removed := conversationMessages[0]
|
||||||
|
bytesRemoved += len(removed.Role) + len(removed.Content) + 20 // Approximate JSON overhead
|
||||||
|
conversationMessages = conversationMessages[1:]
|
||||||
|
|
||||||
|
truncated.SessionMessages = append(systemMessages, conversationMessages...)
|
||||||
|
var err error
|
||||||
|
data, err = json.Marshal(truncated)
|
||||||
|
if err != nil {
|
||||||
|
return req, nil, &TruncationInfo{
|
||||||
|
OriginalCount: originalCount,
|
||||||
|
RemainingCount: len(truncated.SessionMessages),
|
||||||
|
BytesRemoved: bytesRemoved,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Second pass: If still too large, start removing from the end (keep recent)
|
||||||
|
systemMsgCount := len(systemMessages)
|
||||||
|
for len(data) > DefaultMaxHandoffPayloadBytes && len(truncated.SessionMessages) > systemMsgCount {
|
||||||
|
removed := truncated.SessionMessages[len(truncated.SessionMessages)-1]
|
||||||
|
bytesRemoved += len(removed.Role) + len(removed.Content) + 20
|
||||||
|
truncated.SessionMessages = truncated.SessionMessages[:len(truncated.SessionMessages)-1]
|
||||||
|
|
||||||
|
var err error
|
||||||
|
data, err = json.Marshal(truncated)
|
||||||
|
if err != nil {
|
||||||
|
return req, nil, &TruncationInfo{
|
||||||
|
OriginalCount: originalCount,
|
||||||
|
RemainingCount: len(truncated.SessionMessages),
|
||||||
|
BytesRemoved: bytesRemoved,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(data) > DefaultMaxHandoffPayloadBytes {
|
||||||
|
return req, nil, &TruncationInfo{
|
||||||
|
OriginalCount: originalCount,
|
||||||
|
RemainingCount: len(truncated.SessionMessages),
|
||||||
|
BytesRemoved: bytesRemoved,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return truncated, data, &TruncationInfo{
|
||||||
|
OriginalCount: originalCount,
|
||||||
|
RemainingCount: len(truncated.SessionMessages),
|
||||||
|
BytesRemoved: bytesRemoved,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// HandleIncomingHandoff handles a handoff request received from another node.
|
||||||
|
func (hc *HandoffCoordinator) HandleIncomingHandoff(req *HandoffRequest) *HandoffResponse {
|
||||||
|
// Check if we can handle it
|
||||||
|
accepted := hc.CanHandle(req.RequiredCap)
|
||||||
|
localNode := hc.discovery.LocalNode()
|
||||||
|
response := &HandoffResponse{
|
||||||
|
RequestID: req.RequestID,
|
||||||
|
Accepted: accepted,
|
||||||
|
NodeID: 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
return response
|
||||||
|
}
|
||||||
|
|
||||||
|
// findTargetNode finds a suitable target node for handoff.
|
||||||
|
func (hc *HandoffCoordinator) findTargetNode(req *HandoffRequest) (*NodeWithState, error) {
|
||||||
|
var candidates []*NodeWithState
|
||||||
|
|
||||||
|
if req.RequiredCap != "" {
|
||||||
|
candidates = hc.membership.SelectByCapability([]string{req.RequiredCap})
|
||||||
|
} else {
|
||||||
|
candidates = hc.membership.GetAvailableMembers()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Filter out nodes that are in circuit breaker open state or cooldown
|
||||||
|
var healthyCandidates []*NodeWithState
|
||||||
|
for _, c := range candidates {
|
||||||
|
nodeID := c.Node.ID
|
||||||
|
// Skip if circuit is open
|
||||||
|
if hc.circuitBreaker.GetState(nodeID) == CircuitOpen {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
// Skip if in overload cooldown
|
||||||
|
if hc.nodeCooldown.IsCooled(nodeID) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
healthyCandidates = append(healthyCandidates, c)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(healthyCandidates) == 0 {
|
||||||
|
return nil, ErrNoHealthyNodes
|
||||||
|
}
|
||||||
|
|
||||||
|
// Select least loaded node
|
||||||
|
target := healthyCandidates[0]
|
||||||
|
for _, c := range healthyCandidates[1:] {
|
||||||
|
if c.Node.LoadScore < target.Node.LoadScore {
|
||||||
|
target = c
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return target, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetMetrics returns the current handoff metrics.
|
||||||
|
func (hc *HandoffCoordinator) GetMetrics() *HandoffMetrics {
|
||||||
|
return hc.metrics
|
||||||
|
}
|
||||||
|
|
||||||
|
// ComponentHealth returns the health status of handoff components.
|
||||||
|
func (hc *HandoffCoordinator) ComponentHealth() map[string]bool {
|
||||||
|
health := make(map[string]bool)
|
||||||
|
health["coordinator"] = hc.sub != nil
|
||||||
|
health["nats"] = hc.nc != nil
|
||||||
|
health["signing"] = hc.signer != nil && hc.signer.enabled
|
||||||
|
health["encryption"] = hc.encryptor != nil && hc.encryptor.enabled
|
||||||
|
health["discovery"] = hc.discovery != nil
|
||||||
|
health["membership"] = hc.membership != nil
|
||||||
|
return health
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetCircuitBreakerStatus returns the circuit breaker state for all tracked nodes.
|
||||||
|
func (hc *HandoffCoordinator) GetCircuitBreakerStatus() map[string]string {
|
||||||
|
status := make(map[string]string)
|
||||||
|
if hc.circuitBreaker == nil {
|
||||||
|
return status
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get all members
|
||||||
|
if hc.membership != nil {
|
||||||
|
members := hc.membership.GetAvailableMembers()
|
||||||
|
for _, m := range members {
|
||||||
|
nodeID := m.Node.ID
|
||||||
|
state := hc.circuitBreaker.GetState(nodeID)
|
||||||
|
if state != "" {
|
||||||
|
status[nodeID] = string(state)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return status
|
||||||
|
}
|
||||||
396
pkg/swarm/leader_election.go
Normal file
396
pkg/swarm/leader_election.go
Normal file
|
|
@ -0,0 +1,396 @@
|
||||||
|
// 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/kv"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// LeaderLockValue represents the value stored in the leader KV lock.
|
||||||
|
type LeaderLockValue struct {
|
||||||
|
NodeID string `json:"node_id"`
|
||||||
|
Timestamp int64 `json:"timestamp"` // Unix nano when lock was acquired/renewed
|
||||||
|
TTLMs int64 `json:"ttl_ms"` // TTL in milliseconds
|
||||||
|
}
|
||||||
|
|
||||||
|
// LeaderKey is the key used for leader election in the leader bucket.
|
||||||
|
const LeaderKey = "leader"
|
||||||
|
|
||||||
|
// LeaderElection handles leader election using Bucket with CAS (Compare-And-Swap).
|
||||||
|
//
|
||||||
|
// Flow:
|
||||||
|
// - On Start(), attempt to acquire the leader lock via bucket.Create (succeeds if key absent).
|
||||||
|
// - If key exists, check if TTL expired; if so, attempt CAS update to claim.
|
||||||
|
// - Leader renews the lock periodically (RenewalInterval < LockTTL).
|
||||||
|
// - Followers watch the key for changes; on delete/expire, they attempt acquisition.
|
||||||
|
// - On graceful shutdown, the leader deletes the key to trigger re-election.
|
||||||
|
type LeaderElection struct {
|
||||||
|
localNodeID string
|
||||||
|
config LeaderElectionConfig
|
||||||
|
bucket kv.Bucket
|
||||||
|
|
||||||
|
mu sync.RWMutex
|
||||||
|
currentLeader string
|
||||||
|
isLeader bool
|
||||||
|
lastRevision uint64 // For CAS operations
|
||||||
|
leaderChangeCh chan string
|
||||||
|
stopCh chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewLeaderElection creates a new leader election instance using the leader bucket.
|
||||||
|
func NewLeaderElection(nodeID string, bucket kv.Bucket, config LeaderElectionConfig) (*LeaderElection, error) {
|
||||||
|
if bucket == nil {
|
||||||
|
return nil, fmt.Errorf("bucket cannot be nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
le := &LeaderElection{
|
||||||
|
localNodeID: nodeID,
|
||||||
|
config: config,
|
||||||
|
bucket: bucket,
|
||||||
|
leaderChangeCh: make(chan string, 10),
|
||||||
|
stopCh: make(chan struct{}),
|
||||||
|
}
|
||||||
|
|
||||||
|
return le, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start starts the leader election process.
|
||||||
|
func (le *LeaderElection) Start() error {
|
||||||
|
// Attempt initial acquisition
|
||||||
|
le.tryAcquire()
|
||||||
|
|
||||||
|
// Start renewal or watch based on current role
|
||||||
|
go le.mainLoop()
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop stops the leader election process.
|
||||||
|
func (le *LeaderElection) Stop() {
|
||||||
|
le.mu.Lock()
|
||||||
|
wasLeader := le.isLeader
|
||||||
|
le.mu.Unlock()
|
||||||
|
|
||||||
|
// If we are leader, delete the key to trigger re-election
|
||||||
|
if wasLeader && le.bucket != nil {
|
||||||
|
if err := le.bucket.Delete(LeaderKey); err != nil {
|
||||||
|
logger.WarnCF("swarm", "Failed to delete leader key on shutdown", map[string]any{"error": err})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
// ElectLeader triggers a new leader election by deleting the current lock.
|
||||||
|
func (le *LeaderElection) ElectLeader(ctx context.Context) (string, error) {
|
||||||
|
// Delete current lock to force re-election
|
||||||
|
if le.bucket != nil {
|
||||||
|
_ = le.bucket.Delete(LeaderKey)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait for new leader
|
||||||
|
ticker := time.NewTicker(100 * time.Millisecond)
|
||||||
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// tryAcquire attempts to acquire the leader lock.
|
||||||
|
func (le *LeaderElection) tryAcquire() {
|
||||||
|
lockTTL := le.config.LockTTL.Duration
|
||||||
|
if lockTTL <= 0 {
|
||||||
|
lockTTL = DefaultLeaderLockTTL
|
||||||
|
}
|
||||||
|
|
||||||
|
value := LeaderLockValue{
|
||||||
|
NodeID: le.localNodeID,
|
||||||
|
Timestamp: time.Now().UnixNano(),
|
||||||
|
TTLMs: lockTTL.Milliseconds(),
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(value)
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to marshal leader lock value", map[string]any{"error": err})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try bucket.Create — succeeds only if key does not exist
|
||||||
|
rev, err := le.bucket.Create(LeaderKey, data)
|
||||||
|
if err == nil {
|
||||||
|
// Successfully acquired lock
|
||||||
|
le.mu.Lock()
|
||||||
|
le.lastRevision = rev
|
||||||
|
le.becomeLeader()
|
||||||
|
le.mu.Unlock()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Key exists — check if current leader's lock is expired
|
||||||
|
entry, err := le.bucket.Get(LeaderKey)
|
||||||
|
if err != nil {
|
||||||
|
// Key might have been deleted between Create and Get
|
||||||
|
// Try create again
|
||||||
|
rev, err = le.bucket.Create(LeaderKey, data)
|
||||||
|
if err == nil {
|
||||||
|
le.mu.Lock()
|
||||||
|
le.lastRevision = rev
|
||||||
|
le.becomeLeader()
|
||||||
|
le.mu.Unlock()
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var existing LeaderLockValue
|
||||||
|
if err = json.Unmarshal(entry.Value, &existing); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if existing lock is expired
|
||||||
|
lockAge := time.Since(time.Unix(0, existing.Timestamp))
|
||||||
|
existingTTL := time.Duration(existing.TTLMs) * time.Millisecond
|
||||||
|
if lockAge > existingTTL {
|
||||||
|
// Lock expired — attempt CAS update to claim it
|
||||||
|
rev, err = le.bucket.Update(LeaderKey, data, entry.Revision)
|
||||||
|
if err == nil {
|
||||||
|
le.mu.Lock()
|
||||||
|
le.lastRevision = rev
|
||||||
|
le.becomeLeader()
|
||||||
|
le.mu.Unlock()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// CAS failed — someone else claimed it
|
||||||
|
}
|
||||||
|
|
||||||
|
// We are a follower — update current leader info
|
||||||
|
le.mu.Lock()
|
||||||
|
le.setFollower(existing.NodeID)
|
||||||
|
le.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// mainLoop runs the renewal (if leader) or watch (if follower) loop.
|
||||||
|
func (le *LeaderElection) mainLoop() {
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-le.stopCh:
|
||||||
|
return
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
|
le.mu.RLock()
|
||||||
|
amLeader := le.isLeader
|
||||||
|
le.mu.RUnlock()
|
||||||
|
|
||||||
|
if amLeader {
|
||||||
|
le.renewalLoop()
|
||||||
|
} else {
|
||||||
|
le.watchLoop()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Backoff before re-entering loop after a role transition or error.
|
||||||
|
// Without this, a persistent issue would cause a tight spin.
|
||||||
|
select {
|
||||||
|
case <-le.stopCh:
|
||||||
|
return
|
||||||
|
case <-time.After(500 * time.Millisecond):
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// renewalLoop periodically renews the leader lock.
|
||||||
|
func (le *LeaderElection) renewalLoop() {
|
||||||
|
interval := le.config.RenewalInterval.Duration
|
||||||
|
if interval <= 0 {
|
||||||
|
interval = DefaultLeaderRenewalInterval
|
||||||
|
}
|
||||||
|
|
||||||
|
lockTTL := le.config.LockTTL.Duration
|
||||||
|
if lockTTL <= 0 {
|
||||||
|
lockTTL = DefaultLeaderLockTTL
|
||||||
|
}
|
||||||
|
|
||||||
|
ticker := time.NewTicker(interval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-le.stopCh:
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
value := LeaderLockValue{
|
||||||
|
NodeID: le.localNodeID,
|
||||||
|
Timestamp: time.Now().UnixNano(),
|
||||||
|
TTLMs: lockTTL.Milliseconds(),
|
||||||
|
}
|
||||||
|
data, err := json.Marshal(value)
|
||||||
|
if err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
le.mu.Lock()
|
||||||
|
rev, err := le.bucket.Update(LeaderKey, data, le.lastRevision)
|
||||||
|
if err != nil {
|
||||||
|
// CAS failed — lost leadership
|
||||||
|
logger.WarnCF("swarm", "Leader lock renewal failed, stepping down", map[string]any{
|
||||||
|
"node_id": le.localNodeID,
|
||||||
|
"error": err,
|
||||||
|
})
|
||||||
|
le.isLeader = false
|
||||||
|
le.currentLeader = ""
|
||||||
|
le.mu.Unlock()
|
||||||
|
return // Exit renewal loop, mainLoop will start watchLoop
|
||||||
|
}
|
||||||
|
le.lastRevision = rev
|
||||||
|
le.mu.Unlock()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// watchLoop watches the leader key for changes and attempts acquisition on delete/expire.
|
||||||
|
func (le *LeaderElection) watchLoop() {
|
||||||
|
watcher, err := le.bucket.Watch(LeaderKey)
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "Failed to watch leader key", map[string]any{"error": err})
|
||||||
|
// Retry after delay
|
||||||
|
select {
|
||||||
|
case <-le.stopCh:
|
||||||
|
return
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
return // mainLoop will retry
|
||||||
|
}
|
||||||
|
}
|
||||||
|
defer watcher.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-le.stopCh:
|
||||||
|
return
|
||||||
|
case entry := <-watcher.Updates():
|
||||||
|
if entry == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if entry.Operation == kv.EntryOperationDelete || entry.Operation == kv.EntryOperationPurge {
|
||||||
|
// Leader key deleted or expired — try to acquire
|
||||||
|
logger.InfoCF("swarm", "Leader key deleted/expired, attempting acquisition", nil)
|
||||||
|
le.tryAcquire()
|
||||||
|
|
||||||
|
le.mu.RLock()
|
||||||
|
amLeader := le.isLeader
|
||||||
|
le.mu.RUnlock()
|
||||||
|
if amLeader {
|
||||||
|
return // Exit watch loop, mainLoop will start renewalLoop
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Leader key updated — check who the leader is
|
||||||
|
var lockVal LeaderLockValue
|
||||||
|
if err := json.Unmarshal(entry.Value, &lockVal); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
le.mu.Lock()
|
||||||
|
if le.currentLeader != lockVal.NodeID {
|
||||||
|
le.setFollower(lockVal.NodeID)
|
||||||
|
}
|
||||||
|
le.mu.Unlock()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// becomeLeader marks this node as the leader. Must be called with mu held.
|
||||||
|
func (le *LeaderElection) becomeLeader() {
|
||||||
|
if !le.isLeader {
|
||||||
|
le.isLeader = true
|
||||||
|
le.currentLeader = le.localNodeID
|
||||||
|
logger.InfoCF("swarm", "This node is now the leader", map[string]any{"node_id": le.localNodeID})
|
||||||
|
|
||||||
|
select {
|
||||||
|
case le.leaderChangeCh <- le.localNodeID:
|
||||||
|
default:
|
||||||
|
logger.WarnC("swarm", "Leader change notification dropped, channel full")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// setFollower marks this node as a follower with the given leader. Must be called with mu held.
|
||||||
|
func (le *LeaderElection) setFollower(leaderID string) {
|
||||||
|
wasLeader := le.isLeader
|
||||||
|
le.isLeader = false
|
||||||
|
oldLeader := le.currentLeader
|
||||||
|
le.currentLeader = leaderID
|
||||||
|
|
||||||
|
if wasLeader || oldLeader != leaderID {
|
||||||
|
if wasLeader {
|
||||||
|
logger.InfoCF("swarm", "This node is now a follower", map[string]any{
|
||||||
|
"node_id": le.localNodeID,
|
||||||
|
"new_leader": leaderID,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case le.leaderChangeCh <- leaderID:
|
||||||
|
default:
|
||||||
|
logger.WarnC("swarm", "Leader change notification dropped, channel full")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// LeadershipState represents the current leadership state.
|
||||||
|
type LeadershipState struct {
|
||||||
|
LeaderID string `json:"leader_id"`
|
||||||
|
IsLeader bool `json:"is_leader"`
|
||||||
|
LastChange int64 `json:"last_change"` // Unix nano
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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,
|
||||||
|
}
|
||||||
|
}
|
||||||
331
pkg/swarm/load_monitor.go
Normal file
331
pkg/swarm/load_monitor.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 (
|
||||||
|
"runtime"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// 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{}
|
||||||
|
stopOnce sync.Once
|
||||||
|
started bool
|
||||||
|
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() {
|
||||||
|
lm.mu.Lock()
|
||||||
|
defer lm.mu.Unlock()
|
||||||
|
|
||||||
|
if lm.started {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
lm.stopChan = make(chan struct{})
|
||||||
|
lm.stopOnce = sync.Once{}
|
||||||
|
lm.ticker = time.NewTicker(lm.config.Interval.Duration)
|
||||||
|
lm.started = true
|
||||||
|
go lm.run()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop stops monitoring load. Safe to call multiple times.
|
||||||
|
func (lm *LoadMonitor) Stop() {
|
||||||
|
lm.stopOnce.Do(func() {
|
||||||
|
lm.mu.Lock()
|
||||||
|
defer lm.mu.Unlock()
|
||||||
|
|
||||||
|
if !lm.started {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if lm.ticker != nil {
|
||||||
|
lm.ticker.Stop()
|
||||||
|
lm.ticker = nil
|
||||||
|
}
|
||||||
|
close(lm.stopChan)
|
||||||
|
lm.started = false
|
||||||
|
|
||||||
|
logger.DebugC("swarm", "Load monitor stopped")
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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"`
|
||||||
|
MemoryBytes uint64 `json:"memory_bytes"`
|
||||||
|
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 {
|
||||||
|
// Get memory usage
|
||||||
|
var m runtime.MemStats
|
||||||
|
runtime.ReadMemStats(&m)
|
||||||
|
|
||||||
|
metrics := &LoadMetrics{
|
||||||
|
ActiveSessions: lm.GetSessionCount(),
|
||||||
|
Goroutines: runtime.NumGoroutine(),
|
||||||
|
MemoryBytes: m.Alloc,
|
||||||
|
Timestamp: time.Now().UnixNano(),
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
}
|
||||||
|
// Note: CPUUsage is approximated from goroutine count, not actual CPU sampling.
|
||||||
|
// For production accuracy, consider using OS-level CPU metrics.
|
||||||
|
metrics.CPUUsage = normalizeGoroutines(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.
|
||||||
|
func (lm *LoadMonitor) GetTrend() LoadTrend {
|
||||||
|
lm.mu.RLock()
|
||||||
|
defer lm.mu.RUnlock()
|
||||||
|
|
||||||
|
if len(lm.samples) < 3 {
|
||||||
|
return LoadTrendStable
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 LoadTrendIncreasing
|
||||||
|
} else if slope < TrendDecreasingThreshold {
|
||||||
|
return LoadTrendDecreasing
|
||||||
|
}
|
||||||
|
return LoadTrendStable
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
}
|
||||||
|
|
||||||
|
// normalizeGoroutines normalizes goroutine count to a 0-1 score.
|
||||||
|
// This is used as a CPU usage proxy; goroutine count correlates loosely
|
||||||
|
// with CPU load but does not measure actual CPU utilization.
|
||||||
|
func normalizeGoroutines(goroutines int, maxGoroutines int) float64 {
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetLoadScore returns the current load score for the LoadReporter interface.
|
||||||
|
func (lm *LoadMonitor) GetLoadScore() float64 {
|
||||||
|
return lm.GetAverageScore()
|
||||||
|
}
|
||||||
|
|
||||||
|
// LoadScore returns the current load score (alias for GetLoadScore).
|
||||||
|
func (lm *LoadMonitor) LoadScore() float64 {
|
||||||
|
return lm.GetLoadScore()
|
||||||
|
}
|
||||||
16
pkg/swarm/load_reporter_interface.go
Normal file
16
pkg/swarm/load_reporter_interface.go
Normal file
|
|
@ -0,0 +1,16 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
// LoadReporter provides load information for routing decisions.
|
||||||
|
type LoadReporter interface {
|
||||||
|
// GetLoadScore returns the current load score (0-1).
|
||||||
|
GetLoadScore() float64
|
||||||
|
|
||||||
|
// ShouldOffload returns true if the load is high enough to offload tasks.
|
||||||
|
ShouldOffload() bool
|
||||||
|
}
|
||||||
335
pkg/swarm/manager.go
Normal file
335
pkg/swarm/manager.go
Normal file
|
|
@ -0,0 +1,335 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) icoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// SwarmManager manages the swarm cluster by coordinating discovery,
|
||||||
|
// leadership, routing, load monitoring, and handoff.
|
||||||
|
//
|
||||||
|
// This is the main entry point for swarm mode. It composes the various
|
||||||
|
// swarm subsystems and provides a unified API.
|
||||||
|
type SwarmManager struct {
|
||||||
|
config *Config
|
||||||
|
|
||||||
|
// Core components
|
||||||
|
discovery NodeDiscovery
|
||||||
|
rawDS *DiscoveryService // Keep reference to access BucketSet
|
||||||
|
coordinator Coordinator
|
||||||
|
router Router
|
||||||
|
loadMonitor LoadReporter
|
||||||
|
handoff *HandoffCoordinator
|
||||||
|
|
||||||
|
// Local node info
|
||||||
|
localNode *NodeInfo
|
||||||
|
|
||||||
|
// Current view (cached for fast access)
|
||||||
|
currentView View
|
||||||
|
viewMu sync.RWMutex
|
||||||
|
|
||||||
|
// Lifecycle
|
||||||
|
running bool
|
||||||
|
mu sync.RWMutex
|
||||||
|
stopCh chan struct{}
|
||||||
|
|
||||||
|
// Callbacks
|
||||||
|
onViewChange []func(View)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSwarmManager creates a new swarm manager.
|
||||||
|
func NewSwarmManager(cfg *Config) (*SwarmManager, error) {
|
||||||
|
if !cfg.Enabled {
|
||||||
|
return nil, ErrDiscoveryDisabled
|
||||||
|
}
|
||||||
|
|
||||||
|
sm := &SwarmManager{
|
||||||
|
config: cfg,
|
||||||
|
currentView: View{LocalNodeID: cfg.NodeID},
|
||||||
|
stopCh: make(chan struct{}),
|
||||||
|
onViewChange: make([]func(View), 0),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create NATS discovery service
|
||||||
|
rawDS, err := NewDiscoveryService(cfg)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to create discovery service: %w", err)
|
||||||
|
}
|
||||||
|
sm.rawDS = rawDS
|
||||||
|
|
||||||
|
// Wrap in adapter to implement NodeDiscovery interface
|
||||||
|
ds := NewDiscoveryAdapter(rawDS)
|
||||||
|
sm.discovery = ds
|
||||||
|
sm.localNode = rawDS.LocalNode()
|
||||||
|
sm.currentView.LocalNodeID = sm.localNode.ID
|
||||||
|
|
||||||
|
// Create router
|
||||||
|
sm.router = NewDefaultRouter(cfg.Handoff)
|
||||||
|
|
||||||
|
// Create load monitor
|
||||||
|
lm := NewLoadMonitor(&cfg.LoadMonitor)
|
||||||
|
sm.loadMonitor = lm
|
||||||
|
|
||||||
|
// Create handoff coordinator
|
||||||
|
sm.handoff = NewHandoffCoordinator(rawDS, cfg.Handoff)
|
||||||
|
if err := sm.handoff.SetCryptoConfig(cfg.Handoff.Crypto); err != nil {
|
||||||
|
logger.WarnCF("swarm", "Handoff crypto configuration warning", map[string]any{"error": err})
|
||||||
|
}
|
||||||
|
|
||||||
|
return sm, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start starts the swarm manager and all its subsystems.
|
||||||
|
func (sm *SwarmManager) Start() error {
|
||||||
|
sm.mu.Lock()
|
||||||
|
defer sm.mu.Unlock()
|
||||||
|
|
||||||
|
if sm.running {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start discovery (this also creates buckets)
|
||||||
|
if err := sm.discovery.Start(); err != nil {
|
||||||
|
return fmt.Errorf("failed to start discovery: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create leader election coordinator (if enabled) - must be after discovery starts
|
||||||
|
if sm.config.LeaderElection.Enabled {
|
||||||
|
buckets := sm.rawDS.Buckets()
|
||||||
|
if buckets != nil {
|
||||||
|
coordinator, err := NewCoordinatorAdapter(buckets, sm.config.NodeID, sm.config.LeaderElection)
|
||||||
|
if err != nil {
|
||||||
|
sm.discovery.Stop()
|
||||||
|
return fmt.Errorf("failed to create coordinator: %w", err)
|
||||||
|
}
|
||||||
|
sm.coordinator = coordinator
|
||||||
|
|
||||||
|
// Start coordinator
|
||||||
|
if err := sm.coordinator.Start(); err != nil {
|
||||||
|
sm.discovery.Stop()
|
||||||
|
return fmt.Errorf("failed to start coordinator: %w", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start load monitor
|
||||||
|
if lm, ok := sm.loadMonitor.(*LoadMonitor); ok {
|
||||||
|
lm.Start()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start handoff coordinator
|
||||||
|
if sm.handoff != nil {
|
||||||
|
if err := sm.handoff.Start(); err != nil {
|
||||||
|
logger.WarnCF("swarm", "Failed to start handoff coordinator", map[string]any{"error": err})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Watch for node changes
|
||||||
|
cancelWatch := sm.discovery.WatchNodes(sm.onViewUpdate)
|
||||||
|
go func() {
|
||||||
|
<-sm.stopCh
|
||||||
|
cancelWatch()
|
||||||
|
}()
|
||||||
|
|
||||||
|
sm.running = true
|
||||||
|
logger.InfoCF("swarm", "Swarm manager started", map[string]any{
|
||||||
|
"node_id": sm.localNode.ID,
|
||||||
|
})
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop stops the swarm manager and all its subsystems.
|
||||||
|
func (sm *SwarmManager) Stop() error {
|
||||||
|
sm.mu.Lock()
|
||||||
|
defer sm.mu.Unlock()
|
||||||
|
|
||||||
|
if !sm.running {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
close(sm.stopCh)
|
||||||
|
|
||||||
|
// Stop handoff
|
||||||
|
if sm.handoff != nil {
|
||||||
|
_ = sm.handoff.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop load monitor
|
||||||
|
if lm, ok := sm.loadMonitor.(*LoadMonitor); ok {
|
||||||
|
lm.Stop()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop coordinator
|
||||||
|
if sm.coordinator != nil {
|
||||||
|
sm.coordinator.Stop()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop discovery
|
||||||
|
_ = sm.discovery.Stop()
|
||||||
|
|
||||||
|
sm.running = false
|
||||||
|
logger.InfoCF("swarm", "Swarm manager stopped", nil)
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// onViewUpdate is called when the cluster view changes.
|
||||||
|
func (sm *SwarmManager) onViewUpdate(view View) {
|
||||||
|
sm.viewMu.Lock()
|
||||||
|
sm.currentView = view
|
||||||
|
sm.viewMu.Unlock()
|
||||||
|
|
||||||
|
// Notify callbacks
|
||||||
|
for _, cb := range sm.onViewChange {
|
||||||
|
go func(f func(View)) {
|
||||||
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
logger.ErrorCF("swarm", "View change callback panic", map[string]any{"panic": r})
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
f(view)
|
||||||
|
}(cb)
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.DebugCF("swarm", "Cluster view updated", map[string]any{
|
||||||
|
"node_count": len(view.Nodes),
|
||||||
|
"alive_count": len(view.AliveNodes()),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetView returns the current cluster view.
|
||||||
|
func (sm *SwarmManager) GetView() View {
|
||||||
|
sm.viewMu.RLock()
|
||||||
|
defer sm.viewMu.RUnlock()
|
||||||
|
return sm.currentView
|
||||||
|
}
|
||||||
|
|
||||||
|
// AliveNodes returns the current list of alive nodes.
|
||||||
|
func (sm *SwarmManager) AliveNodes() []*NodeInfo {
|
||||||
|
return sm.discovery.AliveNodes()
|
||||||
|
}
|
||||||
|
|
||||||
|
// LocalNode returns the local node info.
|
||||||
|
func (sm *SwarmManager) LocalNode() *NodeInfo {
|
||||||
|
return sm.localNode
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsLeader returns true if this node is the leader.
|
||||||
|
func (sm *SwarmManager) IsLeader() bool {
|
||||||
|
if sm.coordinator == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return sm.coordinator.IsLeader()
|
||||||
|
}
|
||||||
|
|
||||||
|
// LeaderID returns the current leader ID.
|
||||||
|
func (sm *SwarmManager) LeaderID() string {
|
||||||
|
if sm.coordinator == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return sm.coordinator.LeaderID()
|
||||||
|
}
|
||||||
|
|
||||||
|
// PickNode selects a node for the given task.
|
||||||
|
func (sm *SwarmManager) PickNode(task Task) RoutingDecision {
|
||||||
|
view := sm.GetView()
|
||||||
|
return sm.router.PickNode(task, view)
|
||||||
|
}
|
||||||
|
|
||||||
|
// CanHandleLocally checks if the local node can handle the given task.
|
||||||
|
func (sm *SwarmManager) CanHandleLocally(task Task) bool {
|
||||||
|
return sm.router.CanHandleLocally(task, sm.loadMonitor.GetLoadScore())
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetLoad returns the current load score.
|
||||||
|
func (sm *SwarmManager) GetLoad() float64 {
|
||||||
|
return sm.loadMonitor.GetLoadScore()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ShouldOffload returns true if the load is high enough to offload tasks.
|
||||||
|
func (sm *SwarmManager) ShouldOffload() bool {
|
||||||
|
return sm.loadMonitor.ShouldOffload()
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateLoad updates the local node's load score.
|
||||||
|
func (sm *SwarmManager) UpdateLoad(score float64) {
|
||||||
|
if updater, ok := sm.discovery.(NodeInfoUpdater); ok {
|
||||||
|
updater.UpdateLoad(score)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateCapabilities updates the local node's agent capabilities.
|
||||||
|
func (sm *SwarmManager) UpdateCapabilities(caps map[string]string) {
|
||||||
|
if updater, ok := sm.discovery.(NodeInfoUpdater); ok {
|
||||||
|
updater.UpdateCapabilities(caps)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Handoff initiates a handoff to another node.
|
||||||
|
func (sm *SwarmManager) Handoff(ctx context.Context, req *HandoffRequest) (*HandoffResponse, error) {
|
||||||
|
if sm.handoff == nil {
|
||||||
|
return nil, fmt.Errorf("handoff not enabled")
|
||||||
|
}
|
||||||
|
return sm.handoff.InitiateHandoff(ctx, req)
|
||||||
|
}
|
||||||
|
|
||||||
|
// OnViewChange registers a callback for cluster view changes.
|
||||||
|
func (sm *SwarmManager) OnViewChange(callback func(View)) {
|
||||||
|
sm.mu.Lock()
|
||||||
|
defer sm.mu.Unlock()
|
||||||
|
sm.onViewChange = append(sm.onViewChange, callback)
|
||||||
|
}
|
||||||
|
|
||||||
|
// OnLeaderChange registers a callback for leader changes.
|
||||||
|
func (sm *SwarmManager) OnLeaderChange(callback func(leaderID string)) {
|
||||||
|
if sm.coordinator == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
ch := sm.coordinator.LeaderChanges()
|
||||||
|
go func() {
|
||||||
|
for leaderID := range ch {
|
||||||
|
callback(leaderID)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
// OnLoadThreshold registers a callback when load threshold is exceeded.
|
||||||
|
func (sm *SwarmManager) OnLoadThreshold(callback func(float64)) {
|
||||||
|
if lm, ok := sm.loadMonitor.(*LoadMonitor); ok {
|
||||||
|
lm.OnThreshold(callback)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetDiscovery returns the underlying discovery service (for compatibility).
|
||||||
|
func (sm *SwarmManager) GetDiscovery() NodeDiscovery {
|
||||||
|
return sm.discovery
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetHandoffCoordinator returns the handoff coordinator.
|
||||||
|
func (sm *SwarmManager) GetHandoffCoordinator() *HandoffCoordinator {
|
||||||
|
return sm.handoff
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetLoadMonitor returns the load monitor.
|
||||||
|
func (sm *SwarmManager) GetLoadMonitor() *LoadMonitor {
|
||||||
|
if lm, ok := sm.loadMonitor.(*LoadMonitor); ok {
|
||||||
|
return lm
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetBuckets returns the bucket set.
|
||||||
|
func (sm *SwarmManager) GetBuckets() *BucketSet {
|
||||||
|
return sm.rawDS.Buckets()
|
||||||
|
}
|
||||||
385
pkg/swarm/membership.go
Normal file
385
pkg/swarm/membership.go
Normal file
|
|
@ -0,0 +1,385 @@
|
||||||
|
// 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 Discovery
|
||||||
|
view *ClusterView
|
||||||
|
config DiscoveryConfig
|
||||||
|
mu sync.RWMutex
|
||||||
|
|
||||||
|
// Tracks nodes with pending dead-node removal goroutines
|
||||||
|
pendingRemoval map[string]struct{}
|
||||||
|
|
||||||
|
// Event callbacks
|
||||||
|
onJoin []func(*NodeInfo)
|
||||||
|
onLeave []func(*NodeInfo)
|
||||||
|
onUpdate []func(*NodeInfo)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewMembershipManager creates a new membership manager.
|
||||||
|
func NewMembershipManager(ds Discovery, config DiscoveryConfig) *MembershipManager {
|
||||||
|
localNodeID := ds.LocalNode().ID
|
||||||
|
return &MembershipManager{
|
||||||
|
discovery: ds,
|
||||||
|
view: NewClusterView(localNodeID),
|
||||||
|
config: config,
|
||||||
|
pendingRemoval: make(map[string]struct{}),
|
||||||
|
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)
|
||||||
|
|
||||||
|
// Determine what event to fire and collect callbacks while under lock
|
||||||
|
var eventType EventType
|
||||||
|
var callbacks []func(*NodeInfo)
|
||||||
|
|
||||||
|
if !existed {
|
||||||
|
// New node joined
|
||||||
|
nws.State.Status = NodeStatusAlive
|
||||||
|
nws.State.StatusSince = time.Now().UnixNano()
|
||||||
|
nws.State.LastSeen = time.Now().UnixNano()
|
||||||
|
|
||||||
|
eventType = EventJoin
|
||||||
|
callbacks = make([]func(*NodeInfo), len(mm.onJoin))
|
||||||
|
copy(callbacks, mm.onJoin)
|
||||||
|
} else if existing.Node.Timestamp < node.Timestamp {
|
||||||
|
// Existing node updated
|
||||||
|
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++
|
||||||
|
}
|
||||||
|
|
||||||
|
eventType = EventUpdate
|
||||||
|
callbacks = make([]func(*NodeInfo), len(mm.onUpdate))
|
||||||
|
copy(callbacks, mm.onUpdate)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Copy node info for safe use outside lock
|
||||||
|
nodeCopy := *node
|
||||||
|
mm.mu.Unlock()
|
||||||
|
|
||||||
|
// Fire callbacks and dispatch event outside the lock
|
||||||
|
if eventType != "" {
|
||||||
|
for _, cb := range callbacks {
|
||||||
|
go cb(&nodeCopy)
|
||||||
|
}
|
||||||
|
mm.discovery.DispatchEvent(&NodeEvent{
|
||||||
|
Node: &nodeCopy,
|
||||||
|
Event: eventType,
|
||||||
|
Time: time.Now().UnixNano(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
// Copy node info before releasing lock
|
||||||
|
nodeCopy := *nws.Node
|
||||||
|
callbacks := make([]func(*NodeInfo), len(mm.onLeave))
|
||||||
|
copy(callbacks, mm.onLeave)
|
||||||
|
|
||||||
|
mm.view.Remove(nodeID)
|
||||||
|
mm.mu.Unlock()
|
||||||
|
|
||||||
|
// Notify callbacks and dispatch event outside the lock
|
||||||
|
for _, cb := range callbacks {
|
||||||
|
go cb(&nodeCopy)
|
||||||
|
}
|
||||||
|
|
||||||
|
mm.discovery.DispatchEvent(&NodeEvent{
|
||||||
|
Node: &nodeCopy,
|
||||||
|
Event: EventLeave,
|
||||||
|
Time: time.Now().UnixNano(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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)
|
||||||
|
|
||||||
|
// Only spawn a removal goroutine if one isn't already pending
|
||||||
|
if _, pending := mm.pendingRemoval[nodeID]; !pending {
|
||||||
|
mm.pendingRemoval[nodeID] = struct{}{}
|
||||||
|
go func() {
|
||||||
|
time.Sleep(mm.config.DeadNodeTimeout.Duration)
|
||||||
|
mm.mu.Lock()
|
||||||
|
delete(mm.pendingRemoval, nodeID)
|
||||||
|
mm.mu.Unlock()
|
||||||
|
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()
|
||||||
|
|
||||||
|
localNodeID := mm.discovery.LocalNode().ID
|
||||||
|
members := mm.view.GetAvailableNodes()
|
||||||
|
result := make([]*NodeWithState, 0)
|
||||||
|
|
||||||
|
for _, m := range members {
|
||||||
|
if m.Node.ID == localNodeID {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// If no capabilities required, include all available nodes
|
||||||
|
if len(requiredCaps) == 0 {
|
||||||
|
result = append(result, m)
|
||||||
|
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.
|
||||||
|
// Note: When using NATS discovery, health checks are driven by KV TTL
|
||||||
|
// and the stale-node checker in DiscoveryService. This method is kept
|
||||||
|
// for API compatibility but can be called optionally for additional checks.
|
||||||
|
func (mm *MembershipManager) StartHealthCheck(interval time.Duration) {
|
||||||
|
go func() {
|
||||||
|
ticker := time.NewTicker(interval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
for range ticker.C {
|
||||||
|
mm.CheckHealth()
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
624
pkg/swarm/metrics.go
Normal file
624
pkg/swarm/metrics.go
Normal file
|
|
@ -0,0 +1,624 @@
|
||||||
|
// 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"
|
||||||
|
)
|
||||||
|
|
||||||
|
// HandoffMetrics tracks detailed handoff operation statistics.
|
||||||
|
// Separate from the general MetricsCollector for finer-grained tracking.
|
||||||
|
type HandoffMetrics struct {
|
||||||
|
// Counters
|
||||||
|
requests atomic.Int64
|
||||||
|
accepted atomic.Int64
|
||||||
|
rejected atomic.Int64
|
||||||
|
failed atomic.Int64
|
||||||
|
timeouts atomic.Int64
|
||||||
|
retries atomic.Int64
|
||||||
|
truncations atomic.Int64
|
||||||
|
|
||||||
|
// Result-labeled counters (for detailed observability)
|
||||||
|
resultsByReason map[string]*atomic.Int64 // result_reason -> counter
|
||||||
|
|
||||||
|
// Payload size tracking (bytes)
|
||||||
|
totalPayloadBytes atomic.Int64
|
||||||
|
maxPayloadBytes atomic.Int64
|
||||||
|
|
||||||
|
// Latency tracking (in nanoseconds)
|
||||||
|
totalLatency atomic.Int64
|
||||||
|
latencyCount atomic.Int64
|
||||||
|
|
||||||
|
// Security events
|
||||||
|
authFailures atomic.Int64
|
||||||
|
decryptErrors atomic.Int64
|
||||||
|
duplicates atomic.Int64 // Replay attacks detected
|
||||||
|
expired atomic.Int64 // Timestamp window violations
|
||||||
|
|
||||||
|
mu sync.RWMutex
|
||||||
|
lastReset time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewHandoffMetrics creates a new handoff-specific metrics collector.
|
||||||
|
func NewHandoffMetrics() *HandoffMetrics {
|
||||||
|
return &HandoffMetrics{
|
||||||
|
resultsByReason: make(map[string]*atomic.Int64),
|
||||||
|
lastReset: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordRequest records a handoff request.
|
||||||
|
func (m *HandoffMetrics) RecordRequest() {
|
||||||
|
m.requests.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordAccepted records an accepted handoff.
|
||||||
|
func (m *HandoffMetrics) RecordAccepted() {
|
||||||
|
m.accepted.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordRejected records a rejected handoff.
|
||||||
|
func (m *HandoffMetrics) RecordRejected() {
|
||||||
|
m.rejected.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordFailed records a failed handoff.
|
||||||
|
func (m *HandoffMetrics) RecordFailed() {
|
||||||
|
m.failed.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordTimeout records a timeout.
|
||||||
|
func (m *HandoffMetrics) RecordTimeout() {
|
||||||
|
m.timeouts.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordRetry records a retry attempt.
|
||||||
|
func (m *HandoffMetrics) RecordRetry() {
|
||||||
|
m.retries.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordTruncation records a message truncation event.
|
||||||
|
func (m *HandoffMetrics) RecordTruncation() {
|
||||||
|
m.truncations.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordLatency records a handoff operation latency.
|
||||||
|
func (m *HandoffMetrics) RecordLatency(d time.Duration) {
|
||||||
|
m.totalLatency.Add(d.Nanoseconds())
|
||||||
|
m.latencyCount.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordAuthFailure records an authentication failure.
|
||||||
|
func (m *HandoffMetrics) RecordAuthFailure() {
|
||||||
|
m.authFailures.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordDecryptError records a decryption error.
|
||||||
|
func (m *HandoffMetrics) RecordDecryptError() {
|
||||||
|
m.decryptErrors.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordDuplicate records a replay attack detection.
|
||||||
|
func (m *HandoffMetrics) RecordDuplicate() {
|
||||||
|
m.duplicates.Add(1)
|
||||||
|
m.recordResult("duplicate")
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordExpired records a timestamp window violation.
|
||||||
|
func (m *HandoffMetrics) RecordExpired() {
|
||||||
|
m.expired.Add(1)
|
||||||
|
m.recordResult("expired")
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordResult records a result with a specific reason label.
|
||||||
|
func (m *HandoffMetrics) RecordResult(reason string) {
|
||||||
|
m.recordResult(reason)
|
||||||
|
}
|
||||||
|
|
||||||
|
// recordResult is an internal method for recording labeled results.
|
||||||
|
func (m *HandoffMetrics) recordResult(reason string) {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
|
||||||
|
if _, exists := m.resultsByReason[reason]; !exists {
|
||||||
|
m.resultsByReason[reason] = &atomic.Int64{}
|
||||||
|
}
|
||||||
|
m.resultsByReason[reason].Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordPayloadSize records the size of a handoff payload.
|
||||||
|
func (m *HandoffMetrics) RecordPayloadSize(bytes int) {
|
||||||
|
m.totalPayloadBytes.Add(int64(bytes))
|
||||||
|
|
||||||
|
// Track max
|
||||||
|
for {
|
||||||
|
curMax := m.maxPayloadBytes.Load()
|
||||||
|
if int64(bytes) <= curMax {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if m.maxPayloadBytes.CompareAndSwap(curMax, int64(bytes)) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Snapshot returns a snapshot of current handoff metrics.
|
||||||
|
func (m *HandoffMetrics) Snapshot() map[string]any {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
|
||||||
|
result := map[string]any{
|
||||||
|
"requests": m.requests.Load(),
|
||||||
|
"accepted": m.accepted.Load(),
|
||||||
|
"rejected": m.rejected.Load(),
|
||||||
|
"failed": m.failed.Load(),
|
||||||
|
"timeouts": m.timeouts.Load(),
|
||||||
|
"retries": m.retries.Load(),
|
||||||
|
"truncations": m.truncations.Load(),
|
||||||
|
"auth_failures": m.authFailures.Load(),
|
||||||
|
"decrypt_errors": m.decryptErrors.Load(),
|
||||||
|
"duplicates": m.duplicates.Load(),
|
||||||
|
"expired": m.expired.Load(),
|
||||||
|
"avg_latency_ms": m.avgLatencyMs(),
|
||||||
|
"payload_total_bytes": m.totalPayloadBytes.Load(),
|
||||||
|
"payload_max_bytes": m.maxPayloadBytes.Load(),
|
||||||
|
"uptime_seconds": time.Since(m.lastReset).Seconds(),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add result breakdown by reason
|
||||||
|
resultBreakdown := make(map[string]int64)
|
||||||
|
for reason, counter := range m.resultsByReason {
|
||||||
|
resultBreakdown[reason] = counter.Load()
|
||||||
|
}
|
||||||
|
if len(resultBreakdown) > 0 {
|
||||||
|
result["by_reason"] = resultBreakdown
|
||||||
|
}
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// avgLatencyMs returns the average latency in milliseconds.
|
||||||
|
func (m *HandoffMetrics) avgLatencyMs() float64 {
|
||||||
|
count := m.latencyCount.Load()
|
||||||
|
if count == 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
total := m.totalLatency.Load()
|
||||||
|
return float64(total) / float64(count) / 1e6
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reset clears all handoff metrics.
|
||||||
|
func (m *HandoffMetrics) Reset() {
|
||||||
|
m.requests.Store(0)
|
||||||
|
m.accepted.Store(0)
|
||||||
|
m.rejected.Store(0)
|
||||||
|
m.failed.Store(0)
|
||||||
|
m.timeouts.Store(0)
|
||||||
|
m.retries.Store(0)
|
||||||
|
m.truncations.Store(0)
|
||||||
|
m.authFailures.Store(0)
|
||||||
|
m.decryptErrors.Store(0)
|
||||||
|
m.duplicates.Store(0)
|
||||||
|
m.expired.Store(0)
|
||||||
|
m.totalLatency.Store(0)
|
||||||
|
m.latencyCount.Store(0)
|
||||||
|
m.totalPayloadBytes.Store(0)
|
||||||
|
m.maxPayloadBytes.Store(0)
|
||||||
|
|
||||||
|
m.mu.Lock()
|
||||||
|
m.resultsByReason = make(map[string]*atomic.Int64)
|
||||||
|
m.mu.Unlock()
|
||||||
|
|
||||||
|
m.lastReset = time.Now()
|
||||||
|
}
|
||||||
|
|
||||||
|
// RoutingMetrics tracks node-to-node routing statistics.
|
||||||
|
type RoutingMetrics struct {
|
||||||
|
// Counters
|
||||||
|
requestsTotal atomic.Int64
|
||||||
|
requestsLocal atomic.Int64
|
||||||
|
requestsRemote atomic.Int64
|
||||||
|
|
||||||
|
// Result-labeled counters
|
||||||
|
resultsByResult map[string]*atomic.Int64 // ok, timeout, unauthorized, no_target, etc.
|
||||||
|
|
||||||
|
// Latency tracking (in nanoseconds)
|
||||||
|
totalLatency atomic.Int64
|
||||||
|
latencyCount atomic.Int64
|
||||||
|
|
||||||
|
// Fallback tracking
|
||||||
|
fallbackCount atomic.Int64
|
||||||
|
|
||||||
|
mu sync.RWMutex
|
||||||
|
lastReset time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewRoutingMetrics creates a new routing metrics collector.
|
||||||
|
func NewRoutingMetrics() *RoutingMetrics {
|
||||||
|
return &RoutingMetrics{
|
||||||
|
resultsByResult: make(map[string]*atomic.Int64),
|
||||||
|
lastReset: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordRequest records a routing request.
|
||||||
|
func (m *RoutingMetrics) RecordRequest(isLocal bool) {
|
||||||
|
m.requestsTotal.Add(1)
|
||||||
|
if isLocal {
|
||||||
|
m.requestsLocal.Add(1)
|
||||||
|
} else {
|
||||||
|
m.requestsRemote.Add(1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordResult records a routing result with a label.
|
||||||
|
func (m *RoutingMetrics) RecordResult(result string) {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
|
||||||
|
if _, exists := m.resultsByResult[result]; !exists {
|
||||||
|
m.resultsByResult[result] = &atomic.Int64{}
|
||||||
|
}
|
||||||
|
m.resultsByResult[result].Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordFallback records a fallback to local processing.
|
||||||
|
func (m *RoutingMetrics) RecordFallback() {
|
||||||
|
m.fallbackCount.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordLatency records a routing operation latency.
|
||||||
|
func (m *RoutingMetrics) RecordLatency(d time.Duration) {
|
||||||
|
m.totalLatency.Add(d.Nanoseconds())
|
||||||
|
m.latencyCount.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Snapshot returns a snapshot of current routing metrics.
|
||||||
|
func (m *RoutingMetrics) Snapshot() map[string]any {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
|
||||||
|
result := map[string]any{
|
||||||
|
"requests_total": m.requestsTotal.Load(),
|
||||||
|
"requests_local": m.requestsLocal.Load(),
|
||||||
|
"requests_remote": m.requestsRemote.Load(),
|
||||||
|
"fallback_count": m.fallbackCount.Load(),
|
||||||
|
"avg_latency_ms": m.avgLatencyMs(),
|
||||||
|
"uptime_seconds": time.Since(m.lastReset).Seconds(),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add result breakdown
|
||||||
|
resultBreakdown := make(map[string]int64)
|
||||||
|
for label, counter := range m.resultsByResult {
|
||||||
|
resultBreakdown[label] = counter.Load()
|
||||||
|
}
|
||||||
|
if len(resultBreakdown) > 0 {
|
||||||
|
result["by_result"] = resultBreakdown
|
||||||
|
}
|
||||||
|
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *RoutingMetrics) avgLatencyMs() float64 {
|
||||||
|
count := m.latencyCount.Load()
|
||||||
|
if count == 0 {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
total := m.totalLatency.Load()
|
||||||
|
return float64(total) / float64(count) / 1e6
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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
|
||||||
|
|
||||||
|
// NATS-specific counters
|
||||||
|
natsPublished atomic.Int64
|
||||||
|
natsReceived atomic.Int64
|
||||||
|
leaderAcquired atomic.Int64
|
||||||
|
leaderLost 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)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NATSPublished increments the NATS published message counter.
|
||||||
|
func (m *MetricsCollector) NATSPublished(n int64) {
|
||||||
|
m.natsPublished.Add(n)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NATSReceived increments the NATS received message counter.
|
||||||
|
func (m *MetricsCollector) NATSReceived(n int64) {
|
||||||
|
m.natsReceived.Add(n)
|
||||||
|
}
|
||||||
|
|
||||||
|
// LeaderAcquired increments the leader lock acquired counter.
|
||||||
|
func (m *MetricsCollector) LeaderAcquired() {
|
||||||
|
m.leaderAcquired.Add(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// LeaderLost increments the leader lock lost counter.
|
||||||
|
func (m *MetricsCollector) LeaderLost() {
|
||||||
|
m.leaderLost.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 := 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(),
|
||||||
|
|
||||||
|
// NATS counters
|
||||||
|
"nats_published": m.natsPublished.Load(),
|
||||||
|
"nats_received": m.natsReceived.Load(),
|
||||||
|
"leader_acquired": m.leaderAcquired.Load(),
|
||||||
|
"leader_lost": m.leaderLost.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.natsPublished.Store(0)
|
||||||
|
m.natsReceived.Store(0)
|
||||||
|
m.leaderAcquired.Store(0)
|
||||||
|
m.leaderLost.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()
|
||||||
|
}
|
||||||
255
pkg/swarm/nats_membership.go
Normal file
255
pkg/swarm/nats_membership.go
Normal file
|
|
@ -0,0 +1,255 @@
|
||||||
|
// 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"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/kv"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// heartbeatLoop periodically publishes heartbeat and checks for stale nodes.
|
||||||
|
func (ds *DiscoveryService) heartbeatLoop() {
|
||||||
|
interval := ds.config.Discovery.HeartbeatInterval.Duration
|
||||||
|
if interval <= 0 {
|
||||||
|
interval = DefaultHeartbeatInterval
|
||||||
|
}
|
||||||
|
|
||||||
|
ticker := time.NewTicker(interval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
// Initial publish
|
||||||
|
ds.publishNode()
|
||||||
|
|
||||||
|
// Start stale node checker
|
||||||
|
go ds.staleNodeChecker(interval)
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ds.stopChan:
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
ds.publishNode()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// publishNode publishes the local node state to NATS KV.
|
||||||
|
func (ds *DiscoveryService) publishNode() {
|
||||||
|
if ds.buckets == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
ds.mu.Lock()
|
||||||
|
now := time.Now().UnixNano()
|
||||||
|
ds.localNode.LastHeartbeat = now
|
||||||
|
ds.localNode.Timestamp = now
|
||||||
|
if ds.localNode.FirstSeen == 0 {
|
||||||
|
ds.localNode.FirstSeen = now
|
||||||
|
}
|
||||||
|
ds.localNode.Status = string(NodeStatusAlive)
|
||||||
|
node := *ds.localNode // Copy to avoid holding lock
|
||||||
|
ds.mu.Unlock()
|
||||||
|
|
||||||
|
data, err := json.Marshal(node)
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "failed to marshal node info", map[string]any{"error": err})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := ds.buckets.Members.Put(node.ID, data); err != nil {
|
||||||
|
logger.ErrorCF("swarm", "failed to publish node info", map[string]any{"error": err})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// staleNodeChecker actively checks for stale nodes by reading from KV.
|
||||||
|
func (ds *DiscoveryService) staleNodeChecker(interval time.Duration) {
|
||||||
|
checkInterval := interval
|
||||||
|
if checkInterval < 5*time.Second {
|
||||||
|
checkInterval = 5 * time.Second
|
||||||
|
}
|
||||||
|
|
||||||
|
ticker := time.NewTicker(checkInterval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ds.stopChan:
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
ds.checkStaleNodes()
|
||||||
|
ds.checkStaleStatusNodes()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// checkStaleNodes reads all nodes from KV and removes stale ones.
|
||||||
|
func (ds *DiscoveryService) checkStaleNodes() {
|
||||||
|
if ds.buckets == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
keys, err := ds.buckets.Members.Keys()
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
memberTTL := ds.config.Discovery.MemberTTL.Duration
|
||||||
|
if memberTTL <= 0 {
|
||||||
|
memberTTL = DefaultMemberTTL
|
||||||
|
}
|
||||||
|
timeout := memberTTL * 3 // Use 3x TTL as safety margin
|
||||||
|
|
||||||
|
for _, nodeID := range keys {
|
||||||
|
entry, err := ds.buckets.Members.Get(nodeID)
|
||||||
|
if err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if entry == nil {
|
||||||
|
if nws, ok := ds.membership.GetNode(nodeID); ok {
|
||||||
|
ds.membership.RemoveNode(nodeID)
|
||||||
|
ds.detailedStatus.Remove(nodeID)
|
||||||
|
logger.InfoCF("swarm", "Node removed (stale entry)", map[string]any{"node_id": nodeID})
|
||||||
|
ds.eventHandler.Dispatch(&NodeEvent{
|
||||||
|
Node: nws.Node,
|
||||||
|
Event: EventLeave,
|
||||||
|
Time: now.UnixNano(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var node NodeInfo
|
||||||
|
if err := json.Unmarshal(entry.Value, &node); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Skip local node
|
||||||
|
if node.ID == ds.localNode.ID {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
lastSeen := node.LastHeartbeat
|
||||||
|
if lastSeen == 0 {
|
||||||
|
lastSeen = node.Timestamp
|
||||||
|
}
|
||||||
|
|
||||||
|
if lastSeen > 0 {
|
||||||
|
age := now.Sub(time.Unix(0, lastSeen))
|
||||||
|
if age > timeout {
|
||||||
|
if nws, ok := ds.membership.GetNode(node.ID); ok {
|
||||||
|
ds.membership.RemoveNode(node.ID)
|
||||||
|
ds.detailedStatus.Remove(node.ID)
|
||||||
|
logger.InfoCF("swarm", "Node marked stale (no heartbeat)", map[string]any{
|
||||||
|
"node_id": node.ID,
|
||||||
|
"age": age.Seconds(),
|
||||||
|
})
|
||||||
|
ds.eventHandler.Dispatch(&NodeEvent{
|
||||||
|
Node: nws.Node,
|
||||||
|
Event: EventLeave,
|
||||||
|
Time: now.UnixNano(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// checkStaleStatusNodes reads all status entries from KV and removes stale ones.
|
||||||
|
func (ds *DiscoveryService) checkStaleStatusNodes() {
|
||||||
|
if ds.buckets == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
keys, err := ds.buckets.Status.Keys()
|
||||||
|
if err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
timeout := DefaultDetailedStatusTTL * 2
|
||||||
|
|
||||||
|
for _, nodeID := range keys {
|
||||||
|
entry, err := ds.buckets.Status.Get(nodeID)
|
||||||
|
if err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if entry == nil {
|
||||||
|
ds.detailedStatus.Remove(nodeID)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var status NodeDetailedInfo
|
||||||
|
if err := json.Unmarshal(entry.Value, &status); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if status.Timestamp > 0 {
|
||||||
|
age := now.Sub(time.Unix(0, status.Timestamp))
|
||||||
|
if age > timeout {
|
||||||
|
ds.detailedStatus.Remove(nodeID)
|
||||||
|
logger.DebugCF("swarm", "Status entry removed (stale)", map[string]any{
|
||||||
|
"node_id": nodeID,
|
||||||
|
"age": age.Seconds(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// watchNodes watches for changes in the membership bucket.
|
||||||
|
func (ds *DiscoveryService) watchNodes() {
|
||||||
|
watcher, err := ds.buckets.Members.WatchAll()
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "failed to watch members bucket", map[string]any{"error": err})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer watcher.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ds.stopChan:
|
||||||
|
return
|
||||||
|
case entry := <-watcher.Updates():
|
||||||
|
if entry == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Bucket isolation: swarm_members bucket only stores node keys
|
||||||
|
// Key format is nodeID directly (no prefix needed since bucket is isolated)
|
||||||
|
nodeID := entry.Key
|
||||||
|
|
||||||
|
// Skip our own entry
|
||||||
|
if nodeID == ds.localNode.ID {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if entry.Operation == kv.EntryOperationDelete {
|
||||||
|
if nws, ok := ds.membership.GetNode(nodeID); ok {
|
||||||
|
ds.membership.RemoveNode(nodeID)
|
||||||
|
logger.InfoCF("swarm", "Node left (member deleted)", map[string]any{"node_id": nodeID})
|
||||||
|
ds.eventHandler.Dispatch(&NodeEvent{
|
||||||
|
Node: nws.Node,
|
||||||
|
Event: EventLeave,
|
||||||
|
Time: time.Now().UnixNano(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var node NodeInfo
|
||||||
|
if err := json.Unmarshal(entry.Value, &node); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
ds.membership.UpdateNode(&node)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
155
pkg/swarm/nats_status.go
Normal file
155
pkg/swarm/nats_status.go
Normal file
|
|
@ -0,0 +1,155 @@
|
||||||
|
// 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"
|
||||||
|
"runtime"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/kv"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
// detailedStatusLoop periodically publishes detailed node status.
|
||||||
|
func (ds *DiscoveryService) detailedStatusLoop() {
|
||||||
|
interval := DefaultDetailedStatusInterval
|
||||||
|
hbInterval := ds.config.Discovery.HeartbeatInterval.Duration
|
||||||
|
if hbInterval > 0 {
|
||||||
|
interval = hbInterval * 2
|
||||||
|
if interval < DefaultDetailedStatusInterval {
|
||||||
|
interval = DefaultDetailedStatusInterval
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
ticker := time.NewTicker(interval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
// Initial publish
|
||||||
|
ds.publishDetailedStatus()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ds.stopChan:
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
ds.publishDetailedStatus()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// publishDetailedStatus publishes the detailed node status to NATS KV.
|
||||||
|
func (ds *DiscoveryService) publishDetailedStatus() {
|
||||||
|
if ds.buckets == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var loadMetrics *LoadMetrics
|
||||||
|
if ds.loadMonitor != nil {
|
||||||
|
loadMetrics = ds.loadMonitor.GetCurrentLoad()
|
||||||
|
} else {
|
||||||
|
loadMetrics = &LoadMetrics{
|
||||||
|
CPUUsage: 0,
|
||||||
|
MemoryUsage: 0,
|
||||||
|
MemoryBytes: 0,
|
||||||
|
Goroutines: runtime.NumGoroutine(),
|
||||||
|
ActiveSessions: 0,
|
||||||
|
Score: 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
status := BuildDetailedStatus(
|
||||||
|
ds.localNode.ID,
|
||||||
|
ds.localNode.Addr,
|
||||||
|
ds.localNode.Port,
|
||||||
|
0, // No HTTP port; all communication via NATS
|
||||||
|
ds.localNode.LoadScore,
|
||||||
|
loadMetrics,
|
||||||
|
ds.getTodoList(),
|
||||||
|
ds.getActiveTasks(),
|
||||||
|
ds.startTime,
|
||||||
|
ds.getLastError(),
|
||||||
|
ds.version,
|
||||||
|
)
|
||||||
|
|
||||||
|
// Update local registry
|
||||||
|
ds.detailedStatus.Update(status)
|
||||||
|
|
||||||
|
// Publish to NATS KV
|
||||||
|
data, err := json.Marshal(status)
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "failed to marshal detailed status", map[string]any{"error": err})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := ds.buckets.Status.Put(status.ID, data); err != nil {
|
||||||
|
logger.ErrorCF("swarm", "failed to publish detailed status", map[string]any{"error": err})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// watchDetailedStatus watches for detailed status updates from other nodes.
|
||||||
|
func (ds *DiscoveryService) watchDetailedStatus() {
|
||||||
|
if ds.buckets == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
watcher, err := ds.buckets.Status.WatchAll()
|
||||||
|
if err != nil {
|
||||||
|
logger.ErrorCF("swarm", "failed to watch status bucket", map[string]any{"error": err})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer watcher.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ds.stopChan:
|
||||||
|
return
|
||||||
|
case entry := <-watcher.Updates():
|
||||||
|
if entry == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Bucket isolation: swarm_status bucket only stores status keys
|
||||||
|
// Key format is nodeID directly (no prefix needed since bucket is isolated)
|
||||||
|
nodeID := entry.Key
|
||||||
|
|
||||||
|
// Skip our own entry
|
||||||
|
if nodeID == ds.localNode.ID {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
if entry.Operation == kv.EntryOperationDelete {
|
||||||
|
ds.detailedStatus.Remove(nodeID)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var status NodeDetailedInfo
|
||||||
|
if err := json.Unmarshal(entry.Value, &status); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
ds.detailedStatus.Update(&status)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetLoadMonitor sets the load monitor for detailed status reporting.
|
||||||
|
func (ds *DiscoveryService) SetLoadMonitor(lm *LoadMonitor) {
|
||||||
|
ds.loadMonitor = lm
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetVersion sets the version string for detailed status.
|
||||||
|
func (ds *DiscoveryService) SetVersion(version string) {
|
||||||
|
ds.mu.Lock()
|
||||||
|
ds.version = version
|
||||||
|
ds.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetDetailedStatus returns the detailed status registry.
|
||||||
|
func (ds *DiscoveryService) GetDetailedStatus() *NodeDetailedStatus {
|
||||||
|
return ds.detailedStatus
|
||||||
|
}
|
||||||
350
pkg/swarm/node.go
Normal file
350
pkg/swarm/node.go
Normal file
|
|
@ -0,0 +1,350 @@
|
||||||
|
// 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
|
||||||
|
Version string `json:"version"` // PicoClaw version
|
||||||
|
|
||||||
|
// Heartbeat fields for liveness detection
|
||||||
|
LastHeartbeat int64 `json:"last_heartbeat"` // Last heartbeat time (Unix nano)
|
||||||
|
FirstSeen int64 `json:"first_seen"` // When this node was first discovered
|
||||||
|
Status string `json:"status"` // Current status: "alive", "suspect", "dead", "left"
|
||||||
|
|
||||||
|
// Deprecated: Use LastHeartbeat instead
|
||||||
|
Timestamp int64 `json:"timestamp"` // Last update time (Unix nano) - for backward compatibility
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsAlive checks if the node is considered alive based on heartbeat.
|
||||||
|
func (n *NodeInfo) IsAlive(timeout time.Duration) bool {
|
||||||
|
// Prefer LastHeartbeat, fall back to Timestamp for backward compatibility
|
||||||
|
lastSeen := n.LastHeartbeat
|
||||||
|
if lastSeen == 0 {
|
||||||
|
lastSeen = n.Timestamp
|
||||||
|
}
|
||||||
|
if lastSeen == 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
age := time.Since(time.Unix(0, lastSeen))
|
||||||
|
return age < timeout
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetLastSeen returns the last seen time (heartbeat or timestamp).
|
||||||
|
func (n *NodeInfo) GetLastSeen() time.Time {
|
||||||
|
lastSeen := n.LastHeartbeat
|
||||||
|
if lastSeen == 0 {
|
||||||
|
lastSeen = n.Timestamp
|
||||||
|
}
|
||||||
|
return time.Unix(0, lastSeen)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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 canceled
|
||||||
|
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
|
||||||
|
}
|
||||||
41
pkg/swarm/node_discovery_interface.go
Normal file
41
pkg/swarm/node_discovery_interface.go
Normal file
|
|
@ -0,0 +1,41 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
// NodeDiscovery provides node discovery and membership tracking.
|
||||||
|
//
|
||||||
|
// This is the minimal interface for discovering other nodes in the swarm.
|
||||||
|
// Implementations can use NATS KV, etcd, Consul, or any other coordination service.
|
||||||
|
type NodeDiscovery interface {
|
||||||
|
// Start starts the discovery service.
|
||||||
|
Start() error
|
||||||
|
|
||||||
|
// Stop stops the discovery service.
|
||||||
|
Stop() error
|
||||||
|
|
||||||
|
// AliveNodes returns the current list of alive nodes.
|
||||||
|
// This is a point-in-time snapshot.
|
||||||
|
AliveNodes() []*NodeInfo
|
||||||
|
|
||||||
|
// WatchNodes registers a callback for cluster view changes.
|
||||||
|
// The callback receives the updated view and should return quickly.
|
||||||
|
// Returns a function that can be called to cancel the watch.
|
||||||
|
WatchNodes(callback func(View)) func()
|
||||||
|
}
|
||||||
|
|
||||||
|
// NodeInfoUpdater allows updating local node information.
|
||||||
|
// This is an optional interface that NodeDiscovery implementations may implement.
|
||||||
|
type NodeInfoUpdater interface {
|
||||||
|
// UpdateLocalInfo updates the local node's information.
|
||||||
|
UpdateLocalInfo(info *NodeInfo)
|
||||||
|
|
||||||
|
// UpdateLoad updates the local node's load score.
|
||||||
|
UpdateLoad(score float64)
|
||||||
|
|
||||||
|
// UpdateCapabilities updates the local node's agent capabilities.
|
||||||
|
UpdateCapabilities(caps map[string]string)
|
||||||
|
}
|
||||||
182
pkg/swarm/node_status.go
Normal file
182
pkg/swarm/node_status.go
Normal file
|
|
@ -0,0 +1,182 @@
|
||||||
|
// 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"
|
||||||
|
"runtime"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// NodeDetailedInfo represents detailed status information about a node.
|
||||||
|
// This is periodically written to NATS KV for fast status queries without
|
||||||
|
// needing to send requests to other nodes.
|
||||||
|
type NodeDetailedInfo struct {
|
||||||
|
// Basic node info
|
||||||
|
ID string `json:"id"`
|
||||||
|
Addr string `json:"addr"`
|
||||||
|
Port int `json:"port"`
|
||||||
|
HTTPPort int `json:"http_port"`
|
||||||
|
LoadScore float64 `json:"load_score"`
|
||||||
|
Timestamp int64 `json:"timestamp"`
|
||||||
|
|
||||||
|
// Load metrics
|
||||||
|
CPUUsage float64 `json:"cpu_usage"`
|
||||||
|
MemoryUsage float64 `json:"memory_usage"`
|
||||||
|
MemoryBytes uint64 `json:"memory_bytes"`
|
||||||
|
Goroutines int `json:"goroutines"`
|
||||||
|
ActiveSessions int `json:"active_sessions"`
|
||||||
|
|
||||||
|
// Task information
|
||||||
|
TodoList []string `json:"todo_list,omitempty"`
|
||||||
|
ActiveTasks int `json:"active_tasks"`
|
||||||
|
|
||||||
|
// System information
|
||||||
|
DiskUsagePercent float64 `json:"disk_usage_percent,omitempty"`
|
||||||
|
DiskUsedBytes uint64 `json:"disk_used_bytes,omitempty"`
|
||||||
|
DiskTotalBytes uint64 `json:"disk_total_bytes,omitempty"`
|
||||||
|
|
||||||
|
// Status
|
||||||
|
LastError string `json:"last_error,omitempty"`
|
||||||
|
Uptime int64 `json:"uptime"` // Uptime in seconds
|
||||||
|
Status string `json:"status"` // "alive", "suspect", "dead"
|
||||||
|
Version string `json:"version"` // PicoClaw version
|
||||||
|
GoVersion string `json:"go_version"` // Go runtime version
|
||||||
|
OS string `json:"os"` // Operating system
|
||||||
|
Arch string `json:"arch"` // Architecture
|
||||||
|
}
|
||||||
|
|
||||||
|
// NodeDetailedStatus is a registry for node detailed status from NATS KV.
|
||||||
|
// It caches the latest status from all nodes for fast queries.
|
||||||
|
type NodeDetailedStatus struct {
|
||||||
|
statusMap map[string]*NodeDetailedInfo // node_id -> status
|
||||||
|
mu sync.RWMutex
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewNodeDetailedStatus creates a new node detailed status registry.
|
||||||
|
func NewNodeDetailedStatus() *NodeDetailedStatus {
|
||||||
|
return &NodeDetailedStatus{
|
||||||
|
statusMap: make(map[string]*NodeDetailedInfo),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update updates the status for a node.
|
||||||
|
func (nds *NodeDetailedStatus) Update(status *NodeDetailedInfo) {
|
||||||
|
nds.mu.Lock()
|
||||||
|
defer nds.mu.Unlock()
|
||||||
|
nds.statusMap[status.ID] = status
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get returns the status for a node.
|
||||||
|
func (nds *NodeDetailedStatus) Get(nodeID string) (*NodeDetailedInfo, bool) {
|
||||||
|
nds.mu.RLock()
|
||||||
|
defer nds.mu.RUnlock()
|
||||||
|
s, ok := nds.statusMap[nodeID]
|
||||||
|
return s, ok
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetAll returns all node statuses.
|
||||||
|
func (nds *NodeDetailedStatus) GetAll() []*NodeDetailedInfo {
|
||||||
|
nds.mu.RLock()
|
||||||
|
defer nds.mu.RUnlock()
|
||||||
|
|
||||||
|
result := make([]*NodeDetailedInfo, 0, len(nds.statusMap))
|
||||||
|
for _, s := range nds.statusMap {
|
||||||
|
result = append(result, s)
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove removes a node from the registry.
|
||||||
|
func (nds *NodeDetailedStatus) Remove(nodeID string) {
|
||||||
|
nds.mu.Lock()
|
||||||
|
defer nds.mu.Unlock()
|
||||||
|
delete(nds.statusMap, nodeID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// CleanExpired removes statuses older than the given TTL.
|
||||||
|
func (nds *NodeDetailedStatus) CleanExpired(ttl time.Duration) {
|
||||||
|
nds.mu.Lock()
|
||||||
|
defer nds.mu.Unlock()
|
||||||
|
|
||||||
|
cutoff := time.Now().Add(-ttl).UnixNano()
|
||||||
|
for id, s := range nds.statusMap {
|
||||||
|
if s.Timestamp < cutoff {
|
||||||
|
delete(nds.statusMap, id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// AsMap returns the status map for external access.
|
||||||
|
func (nds *NodeDetailedStatus) AsMap() map[string]*NodeDetailedInfo {
|
||||||
|
nds.mu.RLock()
|
||||||
|
defer nds.mu.RUnlock()
|
||||||
|
|
||||||
|
result := make(map[string]*NodeDetailedInfo, len(nds.statusMap))
|
||||||
|
for k, v := range nds.statusMap {
|
||||||
|
result[k] = v
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// MarshalJSON implements json.Marshaler.
|
||||||
|
func (nds *NodeDetailedStatus) MarshalJSON() ([]byte, error) {
|
||||||
|
nds.mu.RLock()
|
||||||
|
defer nds.mu.RUnlock()
|
||||||
|
return json.Marshal(nds.statusMap)
|
||||||
|
}
|
||||||
|
|
||||||
|
// BuildDetailedStatus builds a NodeDetailedInfo from the current node state.
|
||||||
|
func BuildDetailedStatus(
|
||||||
|
nodeID, addr string,
|
||||||
|
port, httpPort int,
|
||||||
|
loadScore float64,
|
||||||
|
loadMetrics *LoadMetrics,
|
||||||
|
todoList []string,
|
||||||
|
activeTasks int,
|
||||||
|
startTime time.Time,
|
||||||
|
lastError string,
|
||||||
|
version string,
|
||||||
|
) *NodeDetailedInfo {
|
||||||
|
uptime := time.Since(startTime)
|
||||||
|
|
||||||
|
status := &NodeDetailedInfo{
|
||||||
|
ID: nodeID,
|
||||||
|
Addr: addr,
|
||||||
|
Port: port,
|
||||||
|
HTTPPort: httpPort,
|
||||||
|
LoadScore: loadScore,
|
||||||
|
Timestamp: time.Now().UnixNano(),
|
||||||
|
|
||||||
|
CPUUsage: loadMetrics.CPUUsage,
|
||||||
|
MemoryUsage: loadMetrics.MemoryUsage,
|
||||||
|
MemoryBytes: loadMetrics.MemoryBytes,
|
||||||
|
Goroutines: loadMetrics.Goroutines,
|
||||||
|
ActiveSessions: loadMetrics.ActiveSessions,
|
||||||
|
|
||||||
|
TodoList: todoList,
|
||||||
|
ActiveTasks: activeTasks,
|
||||||
|
|
||||||
|
LastError: lastError,
|
||||||
|
Uptime: int64(uptime.Seconds()),
|
||||||
|
Status: string(NodeStatusAlive),
|
||||||
|
Version: version,
|
||||||
|
GoVersion: runtime.Version(),
|
||||||
|
OS: runtime.GOOS,
|
||||||
|
Arch: runtime.GOARCH,
|
||||||
|
}
|
||||||
|
|
||||||
|
return status
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetMemoryBytes returns the current memory allocation in bytes.
|
||||||
|
func (lm *LoadMetrics) GetMemoryBytes() uint64 {
|
||||||
|
var m runtime.MemStats
|
||||||
|
runtime.ReadMemStats(&m)
|
||||||
|
return m.Alloc
|
||||||
|
}
|
||||||
113
pkg/swarm/nonce.go
Normal file
113
pkg/swarm/nonce.go
Normal file
|
|
@ -0,0 +1,113 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm nonce cache for replay attack prevention
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// NonceCache tracks seen nonces for replay attack prevention.
|
||||||
|
// Uses a simple ring buffer with TTL-based expiration.
|
||||||
|
type NonceCache struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
entries map[string]*nonceEntry
|
||||||
|
ring []*nonceEntry
|
||||||
|
ringIdx int
|
||||||
|
ringSize int
|
||||||
|
ttl time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
type nonceEntry struct {
|
||||||
|
nonce string
|
||||||
|
timestamp time.Time
|
||||||
|
nodeID string
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewNonceCache creates a new nonce cache.
|
||||||
|
func NewNonceCache(ttl time.Duration, ringSize int) *NonceCache {
|
||||||
|
if ringSize <= 0 {
|
||||||
|
ringSize = 1000 // Default ring size
|
||||||
|
}
|
||||||
|
return &NonceCache{
|
||||||
|
entries: make(map[string]*nonceEntry),
|
||||||
|
ring: make([]*nonceEntry, ringSize),
|
||||||
|
ringSize: ringSize,
|
||||||
|
ringIdx: 0,
|
||||||
|
ttl: ttl,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// CheckAndAdd returns true if the nonce is new (adds it), false if already seen.
|
||||||
|
func (nc *NonceCache) CheckAndAdd(nonce, nodeID string) bool {
|
||||||
|
if nonce == "" {
|
||||||
|
return false // Empty nonce is invalid
|
||||||
|
}
|
||||||
|
|
||||||
|
key := nodeID + ":" + nonce
|
||||||
|
|
||||||
|
// Fast path: check map
|
||||||
|
nc.mu.RLock()
|
||||||
|
if _, exists := nc.entries[key]; exists {
|
||||||
|
nc.mu.RUnlock()
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
nc.mu.RUnlock()
|
||||||
|
|
||||||
|
// Slow path: add to cache
|
||||||
|
nc.mu.Lock()
|
||||||
|
defer nc.mu.Unlock()
|
||||||
|
|
||||||
|
// Double-check after acquiring write lock
|
||||||
|
if _, exists := nc.entries[key]; exists {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clean up old entry at current ring position
|
||||||
|
if oldEntry := nc.ring[nc.ringIdx]; oldEntry != nil {
|
||||||
|
oldKey := oldEntry.nodeID + ":" + oldEntry.nonce
|
||||||
|
delete(nc.entries, oldKey)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add new entry
|
||||||
|
entry := &nonceEntry{
|
||||||
|
nonce: nonce,
|
||||||
|
timestamp: time.Now(),
|
||||||
|
nodeID: nodeID,
|
||||||
|
}
|
||||||
|
nc.entries[key] = entry
|
||||||
|
nc.ring[nc.ringIdx] = entry
|
||||||
|
nc.ringIdx = (nc.ringIdx + 1) % nc.ringSize
|
||||||
|
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cleanup removes expired entries older than TTL.
|
||||||
|
// Should be called periodically (e.g., every minute).
|
||||||
|
func (nc *NonceCache) Cleanup() {
|
||||||
|
nc.mu.Lock()
|
||||||
|
defer nc.mu.Unlock()
|
||||||
|
|
||||||
|
cutoff := time.Now().Add(-nc.ttl)
|
||||||
|
for key, entry := range nc.entries {
|
||||||
|
if entry.timestamp.Before(cutoff) {
|
||||||
|
delete(nc.entries, key)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stats returns cache statistics.
|
||||||
|
func (nc *NonceCache) Stats() map[string]any {
|
||||||
|
nc.mu.RLock()
|
||||||
|
defer nc.mu.RUnlock()
|
||||||
|
|
||||||
|
return map[string]any{
|
||||||
|
"size": len(nc.entries),
|
||||||
|
"ring_size": nc.ringSize,
|
||||||
|
"ttl_seconds": nc.ttl.Seconds(),
|
||||||
|
}
|
||||||
|
}
|
||||||
203
pkg/swarm/policy.go
Normal file
203
pkg/swarm/policy.go
Normal file
|
|
@ -0,0 +1,203 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm failure mode policies for graceful degradation
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// FailurePolicy defines how the system behaves when failures occur.
|
||||||
|
type FailurePolicy struct {
|
||||||
|
// TimeoutStrategy defines what to do on request-reply timeout.
|
||||||
|
TimeoutStrategy TimeoutStrategy `json:"timeout_strategy"`
|
||||||
|
|
||||||
|
// OverloadCooldown is how long to avoid a node after it's overloaded.
|
||||||
|
OverloadCooldown Duration `json:"overload_cooldown"`
|
||||||
|
|
||||||
|
// MaxConsecutiveFailures before circuit breaker opens.
|
||||||
|
MaxConsecutiveFailures int `json:"max_consecutive_failures"`
|
||||||
|
|
||||||
|
// CircuitBreakerCooldown is how long to wait before retrying a failed node.
|
||||||
|
CircuitBreakerCooldown Duration `json:"circuit_breaker_cooldown"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// TimeoutStrategy defines timeout fallback behavior.
|
||||||
|
type TimeoutStrategy string
|
||||||
|
|
||||||
|
const (
|
||||||
|
// TimeoutFallbackLocal: fall back to local processing on timeout.
|
||||||
|
TimeoutFallbackLocal TimeoutStrategy = "fallback_local"
|
||||||
|
// TimeoutRetry: retry with another node on timeout.
|
||||||
|
TimeoutRetry TimeoutStrategy = "retry"
|
||||||
|
// TimeoutFail: return error to user on timeout.
|
||||||
|
TimeoutFail TimeoutStrategy = "fail"
|
||||||
|
)
|
||||||
|
|
||||||
|
// DefaultFailurePolicy returns the default failure policy.
|
||||||
|
func DefaultFailurePolicy() *FailurePolicy {
|
||||||
|
return &FailurePolicy{
|
||||||
|
TimeoutStrategy: TimeoutFallbackLocal,
|
||||||
|
OverloadCooldown: Duration{30 * time.Second},
|
||||||
|
MaxConsecutiveFailures: 3,
|
||||||
|
CircuitBreakerCooldown: Duration{60 * time.Second},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// CircuitBreaker tracks node health and opens/closes circuits.
|
||||||
|
type CircuitBreaker struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
nodes map[string]*nodeState
|
||||||
|
cooldown time.Duration
|
||||||
|
maxFails int
|
||||||
|
}
|
||||||
|
|
||||||
|
type nodeState struct {
|
||||||
|
consecutiveFails int
|
||||||
|
lastFailTime time.Time
|
||||||
|
state CircuitState
|
||||||
|
}
|
||||||
|
|
||||||
|
// CircuitState represents the circuit breaker state.
|
||||||
|
type CircuitState string
|
||||||
|
|
||||||
|
const (
|
||||||
|
CircuitClosed CircuitState = "closed" // Normal operation
|
||||||
|
CircuitOpen CircuitState = "open" // Failing, stop sending
|
||||||
|
CircuitHalfOpen CircuitState = "half-open" // Testing if recovered
|
||||||
|
)
|
||||||
|
|
||||||
|
// NewCircuitBreaker creates a new circuit breaker.
|
||||||
|
func NewCircuitBreaker(maxFails int, cooldown time.Duration) *CircuitBreaker {
|
||||||
|
return &CircuitBreaker{
|
||||||
|
nodes: make(map[string]*nodeState),
|
||||||
|
cooldown: cooldown,
|
||||||
|
maxFails: maxFails,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// CanProceed returns true if requests can proceed to the node.
|
||||||
|
func (cb *CircuitBreaker) CanProceed(nodeID string) bool {
|
||||||
|
cb.mu.RLock()
|
||||||
|
defer cb.mu.RUnlock()
|
||||||
|
|
||||||
|
state, exists := cb.nodes[nodeID]
|
||||||
|
if !exists {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if circuit should be half-open (cooldown expired)
|
||||||
|
if state.state == CircuitOpen && time.Since(state.lastFailTime) > cb.cooldown {
|
||||||
|
return true // Allow one request to test
|
||||||
|
}
|
||||||
|
|
||||||
|
return state.state != CircuitOpen
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordSuccess records a successful request.
|
||||||
|
func (cb *CircuitBreaker) RecordSuccess(nodeID string) {
|
||||||
|
cb.mu.Lock()
|
||||||
|
defer cb.mu.Unlock()
|
||||||
|
|
||||||
|
if state, exists := cb.nodes[nodeID]; exists {
|
||||||
|
state.consecutiveFails = 0
|
||||||
|
if state.state == CircuitHalfOpen {
|
||||||
|
state.state = CircuitClosed
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// RecordFailure records a failed request.
|
||||||
|
func (cb *CircuitBreaker) RecordFailure(nodeID string) {
|
||||||
|
cb.mu.Lock()
|
||||||
|
defer cb.mu.Unlock()
|
||||||
|
|
||||||
|
state := cb.nodes[nodeID]
|
||||||
|
if state == nil {
|
||||||
|
state = &nodeState{}
|
||||||
|
cb.nodes[nodeID] = state
|
||||||
|
}
|
||||||
|
|
||||||
|
state.consecutiveFails++
|
||||||
|
state.lastFailTime = time.Now()
|
||||||
|
|
||||||
|
if state.consecutiveFails >= cb.maxFails {
|
||||||
|
state.state = CircuitOpen
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetState returns the current state of a node.
|
||||||
|
func (cb *CircuitBreaker) GetState(nodeID string) CircuitState {
|
||||||
|
cb.mu.RLock()
|
||||||
|
defer cb.mu.RUnlock()
|
||||||
|
|
||||||
|
if state, exists := cb.nodes[nodeID]; exists {
|
||||||
|
// Check if cooldown expired
|
||||||
|
if state.state == CircuitOpen && time.Since(state.lastFailTime) > cb.cooldown {
|
||||||
|
return CircuitHalfOpen
|
||||||
|
}
|
||||||
|
return state.state
|
||||||
|
}
|
||||||
|
return CircuitClosed
|
||||||
|
}
|
||||||
|
|
||||||
|
// NodeCooldown tracks overloaded nodes to prevent oscillation.
|
||||||
|
type NodeCooldown struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
cooled map[string]time.Time
|
||||||
|
ttl time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewNodeCooldown creates a new cooldown tracker.
|
||||||
|
func NewNodeCooldown(ttl time.Duration) *NodeCooldown {
|
||||||
|
return &NodeCooldown{
|
||||||
|
cooled: make(map[string]time.Time),
|
||||||
|
ttl: ttl,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsCooled returns true if the node is in cooldown period.
|
||||||
|
func (nc *NodeCooldown) IsCooled(nodeID string) bool {
|
||||||
|
nc.mu.RLock()
|
||||||
|
defer nc.mu.RUnlock()
|
||||||
|
|
||||||
|
if t, exists := nc.cooled[nodeID]; exists {
|
||||||
|
if time.Since(t) < nc.ttl {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add puts a node into cooldown.
|
||||||
|
func (nc *NodeCooldown) Add(nodeID string) {
|
||||||
|
nc.mu.Lock()
|
||||||
|
defer nc.mu.Unlock()
|
||||||
|
|
||||||
|
nc.cooled[nodeID] = time.Now()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove removes a node from cooldown (e.g., after successful request).
|
||||||
|
func (nc *NodeCooldown) Remove(nodeID string) {
|
||||||
|
nc.mu.Lock()
|
||||||
|
defer nc.mu.Unlock()
|
||||||
|
|
||||||
|
delete(nc.cooled, nodeID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cleanup removes expired entries.
|
||||||
|
func (nc *NodeCooldown) Cleanup() {
|
||||||
|
nc.mu.Lock()
|
||||||
|
defer nc.mu.Unlock()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
for nodeID, t := range nc.cooled {
|
||||||
|
if now.Sub(t) >= nc.ttl {
|
||||||
|
delete(nc.cooled, nodeID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
88
pkg/swarm/router_default.go
Normal file
88
pkg/swarm/router_default.go
Normal file
|
|
@ -0,0 +1,88 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import "fmt"
|
||||||
|
|
||||||
|
// DefaultRouter implements the Router interface with load-based routing.
|
||||||
|
type DefaultRouter struct {
|
||||||
|
config HandoffConfig
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewDefaultRouter creates a new default router.
|
||||||
|
func NewDefaultRouter(config HandoffConfig) *DefaultRouter {
|
||||||
|
return &DefaultRouter{config: config}
|
||||||
|
}
|
||||||
|
|
||||||
|
// PickNode selects a node for the given task.
|
||||||
|
func (r *DefaultRouter) PickNode(task Task, view View) RoutingDecision {
|
||||||
|
// Filter available nodes based on capabilities
|
||||||
|
var candidates []*NodeInfo
|
||||||
|
if len(task.RequiredCaps) > 0 {
|
||||||
|
// Need to filter by capabilities
|
||||||
|
candidates = r.filterByCapability(view.AliveNodes(), task.RequiredCaps)
|
||||||
|
} else {
|
||||||
|
candidates = view.AliveNodes()
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(candidates) == 0 {
|
||||||
|
return RoutingDecision{
|
||||||
|
Reject: true,
|
||||||
|
Reason: "no available nodes with required capabilities",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Select least loaded node
|
||||||
|
selected := candidates[0]
|
||||||
|
for _, n := range candidates[1:] {
|
||||||
|
if n.LoadScore < selected.LoadScore {
|
||||||
|
selected = n
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return RoutingDecision{
|
||||||
|
NodeID: selected.ID,
|
||||||
|
Reason: fmt.Sprintf("least loaded node (load: %.2f)", selected.LoadScore),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// CanHandleLocally checks if the local node can handle the given task.
|
||||||
|
func (r *DefaultRouter) CanHandleLocally(task Task, localLoad float64) bool {
|
||||||
|
// Check load threshold
|
||||||
|
if localLoad > r.config.LoadThreshold {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// For capability checking, we'd need to access local node caps
|
||||||
|
// This is handled by the caller in most cases
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// filterByCapability filters nodes that have all required capabilities.
|
||||||
|
func (r *DefaultRouter) filterByCapability(nodes []*NodeInfo, requiredCaps []string) []*NodeInfo {
|
||||||
|
result := make([]*NodeInfo, 0)
|
||||||
|
for _, n := range nodes {
|
||||||
|
hasAll := true
|
||||||
|
for _, cap := range requiredCaps {
|
||||||
|
found := false
|
||||||
|
for _, nodeCap := range n.AgentCaps {
|
||||||
|
if nodeCap == cap {
|
||||||
|
found = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !found {
|
||||||
|
hasAll = false
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if hasAll {
|
||||||
|
result = append(result, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
39
pkg/swarm/router_interface.go
Normal file
39
pkg/swarm/router_interface.go
Normal file
|
|
@ -0,0 +1,39 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
// Task represents a task that needs to be routed to a node.
|
||||||
|
type Task struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Type string `json:"type"`
|
||||||
|
RequiredCaps []string `json:"required_caps,omitempty"` // Required agent capabilities
|
||||||
|
Priority int `json:"priority,omitempty"`
|
||||||
|
Metadata map[string]any `json:"metadata,omitempty"`
|
||||||
|
Context map[string]string `json:"context,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// RoutingDecision is the result of routing a task.
|
||||||
|
type RoutingDecision struct {
|
||||||
|
NodeID string `json:"node_id"` // Selected node ID
|
||||||
|
Reason string `json:"reason,omitempty"` // Why this node was selected
|
||||||
|
Reject bool `json:"reject,omitempty"` // true if request should be rejected (no nodes available)
|
||||||
|
LoadTooHigh bool `json:"load_too_high,omitempty"` // true if local load is too high to accept
|
||||||
|
}
|
||||||
|
|
||||||
|
// Router provides task routing decisions based on cluster view.
|
||||||
|
//
|
||||||
|
// The router selects the best node for a task based on load,
|
||||||
|
// capabilities, and other routing policies.
|
||||||
|
type Router interface {
|
||||||
|
// PickNode selects a node for the given task.
|
||||||
|
// Returns a RoutingDecision with the selected node and reason.
|
||||||
|
PickNode(task Task, view View) RoutingDecision
|
||||||
|
|
||||||
|
// CanHandleLocally checks if the local node can handle the given task.
|
||||||
|
// Returns false if the local node is overloaded or missing capabilities.
|
||||||
|
CanHandleLocally(task Task, localLoad float64) bool
|
||||||
|
}
|
||||||
315
pkg/swarm/swarm_test.go
Normal file
315
pkg/swarm/swarm_test.go
Normal file
|
|
@ -0,0 +1,315 @@
|
||||||
|
// 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("UnmarshalJSON from number (seconds)", func(t *testing.T) {
|
||||||
|
d := Duration{}
|
||||||
|
err := d.UnmarshalJSON([]byte(`10`))
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 10*time.Second, d.Duration)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UnmarshalJSON from fractional number", func(t *testing.T) {
|
||||||
|
d := Duration{}
|
||||||
|
err := d.UnmarshalJSON([]byte(`0.5`))
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, 500*time.Millisecond, 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 NATS-based discovery, handoff, and leader election.
|
||||||
|
// These tests require a running NATS server with JetStream enabled
|
||||||
|
// and are skipped by default. Run with:
|
||||||
|
// NATS_URL=nats://localhost:4222 go test -run TestDiscovery -v
|
||||||
|
|
||||||
|
func TestDiscoveryServiceNodeDiscovery(t *testing.T) {
|
||||||
|
t.Skip("Requires running NATS server with JetStream; set NATS_URL to enable")
|
||||||
|
}
|
||||||
|
|
||||||
|
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 := NewHandoffCoordinator(ds, cfg.Handoff)
|
||||||
|
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 := NewHandoffCoordinator(ds, cfg.Handoff)
|
||||||
|
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 TestLeaderElection(t *testing.T) {
|
||||||
|
t.Skip("Requires running NATS server with JetStream; set NATS_URL to enable")
|
||||||
|
}
|
||||||
312
pkg/swarm/todo.go
Normal file
312
pkg/swarm/todo.go
Normal file
|
|
@ -0,0 +1,312 @@
|
||||||
|
// 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"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TodoItem represents a single todo item.
|
||||||
|
type TodoItem struct {
|
||||||
|
ID string `json:"id"`
|
||||||
|
Title string `json:"title"`
|
||||||
|
Desc string `json:"description,omitempty"`
|
||||||
|
Priority int `json:"priority"` // 0=low, 1=medium, 2=high
|
||||||
|
Status string `json:"status"` // "pending", "in_progress", "completed", "canceled"
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
|
DueAt time.Time `json:"due_at,omitempty"`
|
||||||
|
CompletedAt time.Time `json:"completed_at,omitempty"`
|
||||||
|
Tags []string `json:"tags,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// TodoList manages a list of todo items for a node.
|
||||||
|
type TodoList struct {
|
||||||
|
items []*TodoItem
|
||||||
|
mu sync.RWMutex
|
||||||
|
nodeID string
|
||||||
|
discovery *DiscoveryService // For publishing updates
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewTodoList creates a new todo list for a node.
|
||||||
|
func NewTodoList(nodeID string) *TodoList {
|
||||||
|
return &TodoList{
|
||||||
|
items: make([]*TodoItem, 0),
|
||||||
|
nodeID: nodeID,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetDiscovery sets the discovery service for publishing updates.
|
||||||
|
func (tl *TodoList) SetDiscovery(ds *DiscoveryService) {
|
||||||
|
tl.discovery = ds
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add adds a new todo item.
|
||||||
|
func (tl *TodoList) Add(title string, priority int) *TodoItem {
|
||||||
|
tl.mu.Lock()
|
||||||
|
defer tl.mu.Unlock()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
item := &TodoItem{
|
||||||
|
ID: generateTodoID(),
|
||||||
|
Title: title,
|
||||||
|
Priority: priority,
|
||||||
|
Status: "pending",
|
||||||
|
CreatedAt: now,
|
||||||
|
UpdatedAt: now,
|
||||||
|
}
|
||||||
|
|
||||||
|
tl.items = append(tl.items, item)
|
||||||
|
tl.publishIfNeeded()
|
||||||
|
return item
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update updates an existing todo item.
|
||||||
|
func (tl *TodoList) Update(id string, updates func(*TodoItem)) bool {
|
||||||
|
tl.mu.Lock()
|
||||||
|
defer tl.mu.Unlock()
|
||||||
|
|
||||||
|
for _, item := range tl.items {
|
||||||
|
if item.ID == id {
|
||||||
|
updates(item)
|
||||||
|
item.UpdatedAt = time.Now()
|
||||||
|
tl.publishIfNeeded()
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remove removes a todo item by ID.
|
||||||
|
func (tl *TodoList) Remove(id string) bool {
|
||||||
|
tl.mu.Lock()
|
||||||
|
defer tl.mu.Unlock()
|
||||||
|
|
||||||
|
for i, item := range tl.items {
|
||||||
|
if item.ID == id {
|
||||||
|
tl.items = append(tl.items[:i], tl.items[i+1:]...)
|
||||||
|
tl.publishIfNeeded()
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get retrieves a todo item by ID.
|
||||||
|
func (tl *TodoList) Get(id string) *TodoItem {
|
||||||
|
tl.mu.RLock()
|
||||||
|
defer tl.mu.RUnlock()
|
||||||
|
|
||||||
|
for _, item := range tl.items {
|
||||||
|
if item.ID == id {
|
||||||
|
return item
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// List returns all todo items.
|
||||||
|
func (tl *TodoList) List() []*TodoItem {
|
||||||
|
tl.mu.RLock()
|
||||||
|
defer tl.mu.RUnlock()
|
||||||
|
|
||||||
|
result := make([]*TodoItem, len(tl.items))
|
||||||
|
copy(result, tl.items)
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListByStatus returns todo items filtered by status.
|
||||||
|
func (tl *TodoList) ListByStatus(status string) []*TodoItem {
|
||||||
|
tl.mu.RLock()
|
||||||
|
defer tl.mu.RUnlock()
|
||||||
|
|
||||||
|
result := make([]*TodoItem, 0)
|
||||||
|
for _, item := range tl.items {
|
||||||
|
if item.Status == status {
|
||||||
|
result = append(result, item)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListByPriority returns todo items sorted by priority (highest first).
|
||||||
|
func (tl *TodoList) ListByPriority() []*TodoItem {
|
||||||
|
tl.mu.RLock()
|
||||||
|
defer tl.mu.RUnlock()
|
||||||
|
|
||||||
|
result := make([]*TodoItem, len(tl.items))
|
||||||
|
copy(result, tl.items)
|
||||||
|
|
||||||
|
// Simple bubble sort by priority (descending)
|
||||||
|
for i := 0; i < len(result)-1; i++ {
|
||||||
|
for j := 0; j < len(result)-i-1; j++ {
|
||||||
|
if result[j].Priority < result[j+1].Priority {
|
||||||
|
result[j], result[j+1] = result[j+1], result[j]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// Complete marks a todo item as completed.
|
||||||
|
func (tl *TodoList) Complete(id string) bool {
|
||||||
|
return tl.Update(id, func(item *TodoItem) {
|
||||||
|
item.Status = "completed"
|
||||||
|
item.CompletedAt = time.Now()
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start marks a todo item as in progress.
|
||||||
|
func (tl *TodoList) Start(id string) bool {
|
||||||
|
return tl.Update(id, func(item *TodoItem) {
|
||||||
|
item.Status = "in_progress"
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cancel marks a todo item as canceled.
|
||||||
|
func (tl *TodoList) Cancel(id string) bool {
|
||||||
|
return tl.Update(id, func(item *TodoItem) {
|
||||||
|
item.Status = "canceled"
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clear removes all completed or canceled items.
|
||||||
|
func (tl *TodoList) Clear() int {
|
||||||
|
tl.mu.Lock()
|
||||||
|
defer tl.mu.Unlock()
|
||||||
|
|
||||||
|
count := 0
|
||||||
|
newItems := make([]*TodoItem, 0, len(tl.items))
|
||||||
|
for _, item := range tl.items {
|
||||||
|
if item.Status != "completed" && item.Status != "canceled" {
|
||||||
|
newItems = append(newItems, item)
|
||||||
|
} else {
|
||||||
|
count++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
tl.items = newItems
|
||||||
|
|
||||||
|
if count > 0 {
|
||||||
|
tl.publishIfNeeded()
|
||||||
|
}
|
||||||
|
return count
|
||||||
|
}
|
||||||
|
|
||||||
|
// Count returns the total number of todo items.
|
||||||
|
func (tl *TodoList) Count() int {
|
||||||
|
tl.mu.RLock()
|
||||||
|
defer tl.mu.RUnlock()
|
||||||
|
return len(tl.items)
|
||||||
|
}
|
||||||
|
|
||||||
|
// CountByStatus returns the count of items by status.
|
||||||
|
func (tl *TodoList) CountByStatus() map[string]int {
|
||||||
|
tl.mu.RLock()
|
||||||
|
defer tl.mu.RUnlock()
|
||||||
|
|
||||||
|
counts := make(map[string]int)
|
||||||
|
for _, item := range tl.items {
|
||||||
|
counts[item.Status]++
|
||||||
|
}
|
||||||
|
return counts
|
||||||
|
}
|
||||||
|
|
||||||
|
// ToStrings converts the todo list to a slice of strings for status reporting.
|
||||||
|
func (tl *TodoList) ToStrings() []string {
|
||||||
|
tl.mu.RLock()
|
||||||
|
defer tl.mu.RUnlock()
|
||||||
|
|
||||||
|
result := make([]string, 0, len(tl.items))
|
||||||
|
for _, item := range tl.items {
|
||||||
|
prefix := ""
|
||||||
|
switch item.Status {
|
||||||
|
case "completed":
|
||||||
|
prefix = "✅ "
|
||||||
|
case "in_progress":
|
||||||
|
prefix = "🔄 "
|
||||||
|
case "canceled":
|
||||||
|
prefix = "❌ "
|
||||||
|
case "pending":
|
||||||
|
if item.Priority == 2 {
|
||||||
|
prefix = "🔴 "
|
||||||
|
} else if item.Priority == 1 {
|
||||||
|
prefix = "🟡 "
|
||||||
|
} else {
|
||||||
|
prefix = "⚪ "
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result = append(result, prefix+item.Title)
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// publishIfNeeded publishes the todo list to the discovery service if available.
|
||||||
|
func (tl *TodoList) publishIfNeeded() {
|
||||||
|
if tl.discovery == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
todoStrings := tl.ToStrings()
|
||||||
|
tl.discovery.SetTodoList(todoStrings)
|
||||||
|
}
|
||||||
|
|
||||||
|
// generateTodoID generates a unique ID for a todo item.
|
||||||
|
func generateTodoID() string {
|
||||||
|
return time.Now().Format("20060102-150405") + "-" + randomString(4)
|
||||||
|
}
|
||||||
|
|
||||||
|
// randomString generates a random string of given length.
|
||||||
|
func randomString(n int) string {
|
||||||
|
const letters = "abcdefghijklmnopqrstuvwxyz0123456789"
|
||||||
|
b := make([]byte, n)
|
||||||
|
for i := range b {
|
||||||
|
b[i] = letters[time.Now().UnixNano()%int64(len(letters))]
|
||||||
|
time.Sleep(time.Nanosecond)
|
||||||
|
}
|
||||||
|
return string(b)
|
||||||
|
}
|
||||||
|
|
||||||
|
// TodoListSummary represents a summary of a node's todo list.
|
||||||
|
type TodoListSummary struct {
|
||||||
|
NodeID string `json:"node_id"`
|
||||||
|
Total int `json:"total"`
|
||||||
|
Pending int `json:"pending"`
|
||||||
|
InProgress int `json:"in_progress"`
|
||||||
|
Completed int `json:"completed"`
|
||||||
|
Canceled int `json:"canceled"`
|
||||||
|
HighPriority int `json:"high_priority"`
|
||||||
|
Timestamp int64 `json:"timestamp"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// Summary returns a summary of the todo list.
|
||||||
|
func (tl *TodoList) Summary() *TodoListSummary {
|
||||||
|
counts := tl.CountByStatus()
|
||||||
|
highPriority := 0
|
||||||
|
for _, item := range tl.items {
|
||||||
|
if item.Priority == 2 && item.Status != "completed" && item.Status != "canceled" {
|
||||||
|
highPriority++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return &TodoListSummary{
|
||||||
|
NodeID: tl.nodeID,
|
||||||
|
Total: tl.Count(),
|
||||||
|
Pending: counts["pending"],
|
||||||
|
InProgress: counts["in_progress"],
|
||||||
|
Completed: counts["completed"],
|
||||||
|
Canceled: counts["canceled"],
|
||||||
|
HighPriority: highPriority,
|
||||||
|
Timestamp: time.Now().UnixNano(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// MarshalJSON implements json.Marshaler.
|
||||||
|
func (tl *TodoList) MarshalJSON() ([]byte, error) {
|
||||||
|
return json.Marshal(tl.List())
|
||||||
|
}
|
||||||
100
pkg/swarm/view.go
Normal file
100
pkg/swarm/view.go
Normal file
|
|
@ -0,0 +1,100 @@
|
||||||
|
// PicoClaw - Ultra-lightweight personal AI agent
|
||||||
|
// Swarm mode support for multi-agent coordination
|
||||||
|
// License: MIT
|
||||||
|
//
|
||||||
|
// Copyright (c) 2026 PicoClaw contributors
|
||||||
|
|
||||||
|
package swarm
|
||||||
|
|
||||||
|
import "time"
|
||||||
|
|
||||||
|
// View represents a snapshot of the cluster state.
|
||||||
|
// It is an immutable snapshot that can be safely passed between components.
|
||||||
|
type View struct {
|
||||||
|
Nodes []*NodeInfo `json:"nodes"`
|
||||||
|
LocalNodeID string `json:"local_node_id"`
|
||||||
|
Version int64 `json:"version"` // View version for change detection
|
||||||
|
Timestamp int64 `json:"timestamp"` // When this view was captured (Unix nano)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NodeInfoInContext extends NodeInfo with runtime state from the view.
|
||||||
|
type NodeInfoInContext struct {
|
||||||
|
Node *NodeInfo
|
||||||
|
IsAlive bool
|
||||||
|
IsLocal bool
|
||||||
|
LastSeen time.Time
|
||||||
|
LoadScore float64
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get returns node info by ID from the view.
|
||||||
|
func (v *View) Get(nodeID string) *NodeInfoInContext {
|
||||||
|
for _, n := range v.Nodes {
|
||||||
|
if n.ID == nodeID {
|
||||||
|
return &NodeInfoInContext{
|
||||||
|
Node: n,
|
||||||
|
IsAlive: n.Status == "alive" || n.Status == "",
|
||||||
|
IsLocal: n.ID == v.LocalNodeID,
|
||||||
|
LastSeen: n.GetLastSeen(),
|
||||||
|
LoadScore: n.LoadScore,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// AliveNodes returns all alive nodes excluding the local node.
|
||||||
|
func (v *View) AliveNodes() []*NodeInfo {
|
||||||
|
result := make([]*NodeInfo, 0)
|
||||||
|
for _, n := range v.Nodes {
|
||||||
|
if n.ID != v.LocalNodeID && (n.Status == "alive" || n.Status == "") {
|
||||||
|
result = append(result, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// AvailableNodes returns all alive nodes with available capacity.
|
||||||
|
func (v *View) AvailableNodes(loadThreshold float64) []*NodeInfo {
|
||||||
|
result := make([]*NodeInfo, 0)
|
||||||
|
for _, n := range v.Nodes {
|
||||||
|
if n.ID != v.LocalNodeID && (n.Status == "alive" || n.Status == "") {
|
||||||
|
if n.LoadScore < loadThreshold {
|
||||||
|
result = append(result, n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// NodeCount returns the total number of nodes in the view.
|
||||||
|
func (v *View) NodeCount() int {
|
||||||
|
return len(v.Nodes)
|
||||||
|
}
|
||||||
|
|
||||||
|
// AliveCount returns the number of alive nodes in the view.
|
||||||
|
func (v *View) AliveCount() int {
|
||||||
|
count := 0
|
||||||
|
for _, n := range v.Nodes {
|
||||||
|
if n.Status == "alive" || n.Status == "" {
|
||||||
|
count++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return count
|
||||||
|
}
|
||||||
|
|
||||||
|
// HasNode returns true if the view contains the given node ID.
|
||||||
|
func (v *View) HasNode(nodeID string) bool {
|
||||||
|
for _, n := range v.Nodes {
|
||||||
|
if n.ID == nodeID {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// NodeChange represents a change in the cluster view.
|
||||||
|
type NodeChange struct {
|
||||||
|
Node *NodeInfo
|
||||||
|
Join bool // true if node joined, false if left/updated
|
||||||
|
Updated bool // true if node was updated (status changed, etc.)
|
||||||
|
}
|
||||||
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)
|
||||||
|
}
|
||||||
267
pkg/tools/swarm_batch.go
Normal file
267
pkg/tools/swarm_batch.go
Normal file
|
|
@ -0,0 +1,267 @@
|
||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/swarm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// SwarmBatchTool queries multiple nodes in parallel.
|
||||||
|
type SwarmBatchTool struct {
|
||||||
|
discovery swarm.Discovery
|
||||||
|
localID string
|
||||||
|
sendActionFn func(ctx context.Context, targetNodeID, action string) (string, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSwarmBatchTool creates a new swarm batch query tool.
|
||||||
|
func NewSwarmBatchTool(
|
||||||
|
discovery swarm.Discovery,
|
||||||
|
localID string,
|
||||||
|
) *SwarmBatchTool {
|
||||||
|
return &SwarmBatchTool{
|
||||||
|
discovery: discovery,
|
||||||
|
localID: localID,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetSendActionFn sets the function to send actions to other nodes.
|
||||||
|
func (t *SwarmBatchTool) SetSendActionFn(
|
||||||
|
fn func(ctx context.Context, targetNodeID, action string) (string, error),
|
||||||
|
) {
|
||||||
|
t.sendActionFn = fn
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name returns the tool name.
|
||||||
|
func (t *SwarmBatchTool) Name() string {
|
||||||
|
return "swarm_batch"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Description returns the tool description.
|
||||||
|
func (t *SwarmBatchTool) Description() string {
|
||||||
|
//nolint:gosmopolitan // Intentional Chinese examples for bilingual LLM tool descriptions
|
||||||
|
return "Broadcast query to ALL nodes in parallel (fast, <10ms). " +
|
||||||
|
"IMPORTANT: Use this FIRST when asking about multiple nodes (e.g., '所有节点状态', 'node-01和node-02的状态'). " +
|
||||||
|
"DO NOT use swarm_route for each node - use this ONCE to get all data. " +
|
||||||
|
"For single node specific tasks, use swarm_route with action='message'."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters returns the tool parameters schema.
|
||||||
|
func (t *SwarmBatchTool) Parameters() map[string]any {
|
||||||
|
return map[string]any{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]any{
|
||||||
|
"node_ids": map[string]any{
|
||||||
|
"type": "array",
|
||||||
|
"items": map[string]any{"type": "string"},
|
||||||
|
"description": "List of node IDs to query (optional, defaults to all nodes)",
|
||||||
|
},
|
||||||
|
"action": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"enum": []string{"status", "message"},
|
||||||
|
"default": "status",
|
||||||
|
"description": "Action type: 'status' (default, fast) returns structured data; 'message' requires a message parameter",
|
||||||
|
},
|
||||||
|
"message": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Message to send (only for action='message')",
|
||||||
|
},
|
||||||
|
"verbose": map[string]any{
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "Include detailed information",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute executes the batch query in parallel.
|
||||||
|
func (t *SwarmBatchTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
||||||
|
if t.discovery == nil {
|
||||||
|
return ErrorResult("Swarm mode is not enabled")
|
||||||
|
}
|
||||||
|
|
||||||
|
action, _ := args["action"].(string)
|
||||||
|
if action == "" {
|
||||||
|
action = "status"
|
||||||
|
}
|
||||||
|
message, _ := args["message"].(string)
|
||||||
|
verbose := false
|
||||||
|
if v, ok := args["verbose"].(bool); ok {
|
||||||
|
verbose = v
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get target nodes
|
||||||
|
var targetIDs []string
|
||||||
|
if nodeIDsRaw, ok := args["node_ids"].([]any); ok {
|
||||||
|
for _, id := range nodeIDsRaw {
|
||||||
|
if idStr, ok := id.(string); ok {
|
||||||
|
targetIDs = append(targetIDs, idStr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// If no nodes specified, query all nodes except local
|
||||||
|
if len(targetIDs) == 0 {
|
||||||
|
members := t.discovery.Members()
|
||||||
|
for _, m := range members {
|
||||||
|
if m.Node.ID != t.localID {
|
||||||
|
targetIDs = append(targetIDs, m.Node.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(targetIDs) == 0 {
|
||||||
|
return ErrorResult("No nodes to query")
|
||||||
|
}
|
||||||
|
|
||||||
|
// For status action with detailed status available, use fast path
|
||||||
|
if action == "status" {
|
||||||
|
if provider, ok := t.discovery.(swarm.DetailedStatusProvider); ok {
|
||||||
|
if statusRegistry := provider.GetDetailedStatus(); statusRegistry != nil {
|
||||||
|
return t.executeFastStatusQuery(statusRegistry, targetIDs, verbose)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Otherwise, use parallel action queries
|
||||||
|
return t.executeParallelQueries(ctx, targetIDs, action, message)
|
||||||
|
}
|
||||||
|
|
||||||
|
// executeFastStatusQuery queries all nodes from local status cache (fastest).
|
||||||
|
func (t *SwarmBatchTool) executeFastStatusQuery(
|
||||||
|
statusRegistry *swarm.NodeDetailedStatus,
|
||||||
|
nodeIDs []string,
|
||||||
|
verbose bool,
|
||||||
|
) *ToolResult {
|
||||||
|
var sb strings.Builder
|
||||||
|
sb.WriteString(fmt.Sprintf("📊 Batch Status Query (%d nodes):\n\n", len(nodeIDs)))
|
||||||
|
|
||||||
|
for _, nodeID := range nodeIDs {
|
||||||
|
status, ok := statusRegistry.Get(nodeID)
|
||||||
|
if !ok {
|
||||||
|
sb.WriteString(fmt.Sprintf("**%s**: ❌ No status data available\n", nodeID))
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
sb.WriteString(fmt.Sprintf("**%s**", nodeID))
|
||||||
|
|
||||||
|
// Format uptime
|
||||||
|
uptime := fmt.Sprintf("%.1fm", float64(status.Uptime)/60)
|
||||||
|
sb.WriteString(fmt.Sprintf(" (up: %s)", uptime))
|
||||||
|
|
||||||
|
if verbose {
|
||||||
|
sb.WriteString(fmt.Sprintf("\n Address: %s:%d\n", status.Addr, status.HTTPPort))
|
||||||
|
sb.WriteString(fmt.Sprintf(" Load: %.1f%% | CPU: %.1f%% | Mem: %.1f%% (%.1f MB)\n",
|
||||||
|
status.LoadScore*100, status.CPUUsage*100, status.MemoryUsage*100,
|
||||||
|
float64(status.MemoryBytes)/(1024*1024)))
|
||||||
|
sb.WriteString(fmt.Sprintf(" Sessions: %d | Goroutines: %d\n",
|
||||||
|
status.ActiveSessions, status.Goroutines))
|
||||||
|
|
||||||
|
if len(status.TodoList) > 0 {
|
||||||
|
sb.WriteString(" Todo:\n")
|
||||||
|
for i, todo := range status.TodoList {
|
||||||
|
if i >= 3 {
|
||||||
|
sb.WriteString(fmt.Sprintf(" ... and %d more\n", len(status.TodoList)-3))
|
||||||
|
break
|
||||||
|
}
|
||||||
|
sb.WriteString(fmt.Sprintf(" - %s\n", todo))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if status.ActiveTasks > 0 {
|
||||||
|
sb.WriteString(fmt.Sprintf(" Active Tasks: %d\n", status.ActiveTasks))
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
sb.WriteString(fmt.Sprintf(" — Load: %.0f%% | Mem: %.0f%%",
|
||||||
|
status.LoadScore*100, status.MemoryUsage*100))
|
||||||
|
|
||||||
|
if len(status.TodoList) > 0 {
|
||||||
|
sb.WriteString(fmt.Sprintf(" | Todo: %d", len(status.TodoList)))
|
||||||
|
}
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: sb.String(),
|
||||||
|
ForUser: sb.String(),
|
||||||
|
IsError: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// executeParallelQueries sends parallel requests to all nodes.
|
||||||
|
func (t *SwarmBatchTool) executeParallelQueries(
|
||||||
|
ctx context.Context,
|
||||||
|
nodeIDs []string,
|
||||||
|
action, message string,
|
||||||
|
) *ToolResult {
|
||||||
|
if t.sendActionFn == nil {
|
||||||
|
return ErrorResult("Inter-node communication not available")
|
||||||
|
}
|
||||||
|
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
results := make(map[string]string)
|
||||||
|
errors := make(map[string]string)
|
||||||
|
var mu sync.Mutex
|
||||||
|
|
||||||
|
members := t.discovery.Members()
|
||||||
|
nodeMap := make(map[string]*swarm.NodeWithState)
|
||||||
|
for _, m := range members {
|
||||||
|
nodeMap[m.Node.ID] = m
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, nodeID := range nodeIDs {
|
||||||
|
wg.Add(1)
|
||||||
|
go func(nid string) {
|
||||||
|
defer wg.Done()
|
||||||
|
|
||||||
|
target := nodeMap[nid]
|
||||||
|
if target == nil {
|
||||||
|
mu.Lock()
|
||||||
|
errors[nid] = "Node not found"
|
||||||
|
mu.Unlock()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if target.State.Status != swarm.NodeStatusAlive {
|
||||||
|
mu.Lock()
|
||||||
|
errors[nid] = fmt.Sprintf("Node not alive (status: %s)", target.State.Status)
|
||||||
|
mu.Unlock()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
result, err := t.sendActionFn(ctx, nid, action)
|
||||||
|
mu.Lock()
|
||||||
|
defer mu.Unlock()
|
||||||
|
if err != nil {
|
||||||
|
errors[nid] = err.Error()
|
||||||
|
} else {
|
||||||
|
results[nid] = result
|
||||||
|
}
|
||||||
|
}(nodeID)
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
var sb strings.Builder
|
||||||
|
sb.WriteString(fmt.Sprintf("📊 Batch Query Results (%d nodes):\n\n", len(nodeIDs)))
|
||||||
|
|
||||||
|
for _, nodeID := range nodeIDs {
|
||||||
|
sb.WriteString(fmt.Sprintf("**%s**:\n", nodeID))
|
||||||
|
if result, ok := results[nodeID]; ok {
|
||||||
|
sb.WriteString(result)
|
||||||
|
} else if err, ok := errors[nodeID]; ok {
|
||||||
|
sb.WriteString(fmt.Sprintf("❌ Error: %s\n", err))
|
||||||
|
}
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: sb.String(),
|
||||||
|
ForUser: sb.String(),
|
||||||
|
IsError: len(errors) == 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
748
pkg/tools/swarm_nodes.go
Normal file
748
pkg/tools/swarm_nodes.go
Normal file
|
|
@ -0,0 +1,748 @@
|
||||||
|
package tools
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"sort"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/swarm"
|
||||||
|
)
|
||||||
|
|
||||||
|
// FormatClusterStatus returns a human-readable cluster status string.
|
||||||
|
// Used by the /nodes command handler.
|
||||||
|
// Now uses detailed status from NATS KV for fast queries without sending requests to other nodes.
|
||||||
|
func FormatClusterStatus(
|
||||||
|
discovery swarm.Discovery,
|
||||||
|
load *swarm.LoadMonitor,
|
||||||
|
localID string,
|
||||||
|
verbose bool,
|
||||||
|
) string {
|
||||||
|
if discovery == nil {
|
||||||
|
return "Swarm mode is not enabled."
|
||||||
|
}
|
||||||
|
|
||||||
|
members := discovery.Members()
|
||||||
|
if len(members) == 0 {
|
||||||
|
return "No nodes found in the swarm cluster."
|
||||||
|
}
|
||||||
|
|
||||||
|
var sb strings.Builder
|
||||||
|
sb.WriteString(fmt.Sprintf("Swarm Cluster Status (%d node%s):\n\n", len(members), plural(len(members))))
|
||||||
|
|
||||||
|
// Sort nodes: local first, then by ID
|
||||||
|
sortedMembers := make([]*swarm.NodeWithState, len(members))
|
||||||
|
copy(sortedMembers, members)
|
||||||
|
sort.Slice(sortedMembers, func(i, j int) bool {
|
||||||
|
// Local node always comes first
|
||||||
|
if sortedMembers[i].Node.ID == localID {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if sortedMembers[j].Node.ID == localID {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return sortedMembers[i].Node.ID < sortedMembers[j].Node.ID
|
||||||
|
})
|
||||||
|
|
||||||
|
// Try to get detailed status from discovery (if it supports it)
|
||||||
|
var detailedStatuses map[string]*swarm.NodeDetailedInfo
|
||||||
|
if provider, ok := discovery.(swarm.DetailedStatusProvider); ok {
|
||||||
|
statusRegistry := provider.GetDetailedStatus()
|
||||||
|
if statusRegistry != nil {
|
||||||
|
detailedStatuses = make(map[string]*swarm.NodeDetailedInfo)
|
||||||
|
for _, status := range statusRegistry.GetAll() {
|
||||||
|
detailedStatuses[status.ID] = status
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, m := range sortedMembers {
|
||||||
|
node := m.Node
|
||||||
|
state := m.State
|
||||||
|
|
||||||
|
localMark := " "
|
||||||
|
if node.ID == localID {
|
||||||
|
localMark = "*"
|
||||||
|
}
|
||||||
|
|
||||||
|
sb.WriteString(fmt.Sprintf("%s **%s**", localMark, node.ID))
|
||||||
|
|
||||||
|
// Check if we have detailed status from KV
|
||||||
|
detailed := detailedStatuses[node.ID]
|
||||||
|
if detailed != nil {
|
||||||
|
// Use detailed status from KV (fast path)
|
||||||
|
sb.WriteString(formatDetailedNodeStatus(detailed, verbose))
|
||||||
|
} else {
|
||||||
|
// Fallback to basic status
|
||||||
|
sb.WriteString(formatBasicNodeStatus(node, state, verbose))
|
||||||
|
}
|
||||||
|
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
sb.WriteString("* = this node\n")
|
||||||
|
sb.WriteString("\nTip: Use `@node-id: message` to route a request to a specific node.")
|
||||||
|
|
||||||
|
return sb.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatDetailedNodeStatus formats node status from detailed info (from KV).
|
||||||
|
func formatDetailedNodeStatus(node *swarm.NodeDetailedInfo, verbose bool) string {
|
||||||
|
var result strings.Builder
|
||||||
|
|
||||||
|
// Add uptime
|
||||||
|
uptime := time.Duration(node.Uptime) * time.Second
|
||||||
|
result.WriteString(fmt.Sprintf(" (uptime: %s)", formatDuration(uptime)))
|
||||||
|
|
||||||
|
if node.Status != "" {
|
||||||
|
result.WriteString(fmt.Sprintf(" [%s]", node.Status))
|
||||||
|
}
|
||||||
|
result.WriteString("\n")
|
||||||
|
|
||||||
|
if verbose {
|
||||||
|
result.WriteString(fmt.Sprintf(" Address: %s:%d (HTTP: %d)\n", node.Addr, node.Port, node.HTTPPort))
|
||||||
|
|
||||||
|
// Load metrics with emoji
|
||||||
|
result.WriteString(fmt.Sprintf(" Load: %.1f%%", node.LoadScore*100))
|
||||||
|
if node.LoadScore < 0.5 {
|
||||||
|
result.WriteString(" 🟩")
|
||||||
|
} else if node.LoadScore < 0.8 {
|
||||||
|
result.WriteString(" 🟨")
|
||||||
|
} else {
|
||||||
|
result.WriteString(" 🟥")
|
||||||
|
}
|
||||||
|
result.WriteString(fmt.Sprintf(" | CPU: %.1f%% | Mem: %.1f%% (%.1f MB)\n",
|
||||||
|
node.CPUUsage*100, node.MemoryUsage*100, float64(node.MemoryBytes)/(1024*1024)))
|
||||||
|
|
||||||
|
// Active sessions and goroutines
|
||||||
|
result.WriteString(fmt.Sprintf(" Sessions: %d | Goroutines: %d\n", node.ActiveSessions, node.Goroutines))
|
||||||
|
|
||||||
|
// Todo list
|
||||||
|
if len(node.TodoList) > 0 {
|
||||||
|
result.WriteString(" Todo:\n")
|
||||||
|
for i, todo := range node.TodoList {
|
||||||
|
if i >= 5 { // Limit to 5 items
|
||||||
|
result.WriteString(fmt.Sprintf(" ... and %d more\n", len(node.TodoList)-5))
|
||||||
|
break
|
||||||
|
}
|
||||||
|
result.WriteString(fmt.Sprintf(" - %s\n", todo))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Active tasks
|
||||||
|
if node.ActiveTasks > 0 {
|
||||||
|
result.WriteString(fmt.Sprintf(" Active Tasks: %d\n", node.ActiveTasks))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Disk usage
|
||||||
|
if node.DiskUsagePercent > 0 {
|
||||||
|
result.WriteString(fmt.Sprintf(" Disk: %.1f%% used", node.DiskUsagePercent))
|
||||||
|
if node.DiskUsedBytes > 0 && node.DiskTotalBytes > 0 {
|
||||||
|
result.WriteString(fmt.Sprintf(" (%.1f GB / %.1f GB)",
|
||||||
|
float64(node.DiskUsedBytes)/(1024*1024*1024),
|
||||||
|
float64(node.DiskTotalBytes)/(1024*1024*1024)))
|
||||||
|
}
|
||||||
|
result.WriteString("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Last error
|
||||||
|
if node.LastError != "" {
|
||||||
|
result.WriteString(fmt.Sprintf(" Last Error: %s\n", node.LastError))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Version info
|
||||||
|
if node.Version != "" {
|
||||||
|
result.WriteString(fmt.Sprintf(" Version: %s (%s/%s)\n", node.Version, node.OS, node.Arch))
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Compact format
|
||||||
|
result.WriteString(fmt.Sprintf(" — Load: %.0f%% | Mem: %.0f%%",
|
||||||
|
node.LoadScore*100, node.MemoryUsage*100))
|
||||||
|
|
||||||
|
if len(node.TodoList) > 0 {
|
||||||
|
result.WriteString(fmt.Sprintf(" | Todo: %d", len(node.TodoList)))
|
||||||
|
}
|
||||||
|
|
||||||
|
if node.ActiveTasks > 0 {
|
||||||
|
result.WriteString(fmt.Sprintf(" | Tasks: %d", node.ActiveTasks))
|
||||||
|
}
|
||||||
|
|
||||||
|
result.WriteString("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
return result.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// formatBasicNodeStatus formats node status from basic info (fallback).
|
||||||
|
func formatBasicNodeStatus(node *swarm.NodeInfo, state *swarm.NodeState, verbose bool) string {
|
||||||
|
var result strings.Builder
|
||||||
|
|
||||||
|
if state != nil {
|
||||||
|
writeStatusIndicator(&result, state.Status)
|
||||||
|
result.WriteString(fmt.Sprintf(" Load: %.0f%%", node.LoadScore*100))
|
||||||
|
}
|
||||||
|
|
||||||
|
result.WriteString(fmt.Sprintf(" @ %s:%d\n", node.Addr, node.Port))
|
||||||
|
|
||||||
|
if verbose && state != nil {
|
||||||
|
result.WriteString(fmt.Sprintf(" Status: %s\n", state.Status))
|
||||||
|
|
||||||
|
loadPercent := int(node.LoadScore * 100)
|
||||||
|
result.WriteString(fmt.Sprintf(" Load: %.2f%% ", node.LoadScore*100))
|
||||||
|
if loadPercent < 50 {
|
||||||
|
result.WriteString("🟩")
|
||||||
|
} else if loadPercent < 80 {
|
||||||
|
result.WriteString("🟨")
|
||||||
|
} else {
|
||||||
|
result.WriteString("🟥")
|
||||||
|
}
|
||||||
|
result.WriteString("\n")
|
||||||
|
|
||||||
|
lastSeen := time.Unix(0, state.LastSeen)
|
||||||
|
age := time.Since(lastSeen)
|
||||||
|
result.WriteString(fmt.Sprintf(" Last seen: %s ago\n", formatDuration(age)))
|
||||||
|
|
||||||
|
if len(node.AgentCaps) > 0 {
|
||||||
|
result.WriteString(" Capabilities:\n")
|
||||||
|
for k, v := range node.AgentCaps {
|
||||||
|
result.WriteString(fmt.Sprintf(" - %s: %s\n", k, v))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(node.Labels) > 0 {
|
||||||
|
result.WriteString(" Labels:\n")
|
||||||
|
for k, v := range node.Labels {
|
||||||
|
result.WriteString(fmt.Sprintf(" - %s: %s\n", k, v))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return result.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// writeStatusIndicator writes a status indicator to the builder.
|
||||||
|
func writeStatusIndicator(sb *strings.Builder, status swarm.NodeStatus) {
|
||||||
|
switch status {
|
||||||
|
case swarm.NodeStatusAlive:
|
||||||
|
sb.WriteString(" 🟢")
|
||||||
|
case swarm.NodeStatusSuspect:
|
||||||
|
sb.WriteString(" 🟡")
|
||||||
|
case swarm.NodeStatusDead:
|
||||||
|
sb.WriteString(" 🔴")
|
||||||
|
case swarm.NodeStatusLeft:
|
||||||
|
sb.WriteString(" ⚪")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SwarmNodesTool returns current swarm cluster node status.
|
||||||
|
type SwarmNodesTool struct {
|
||||||
|
discovery swarm.Discovery
|
||||||
|
load *swarm.LoadMonitor
|
||||||
|
localID string
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSwarmNodesTool creates a new swarm nodes status tool.
|
||||||
|
func NewSwarmNodesTool(
|
||||||
|
discovery swarm.Discovery,
|
||||||
|
load *swarm.LoadMonitor,
|
||||||
|
localID string,
|
||||||
|
) *SwarmNodesTool {
|
||||||
|
return &SwarmNodesTool{
|
||||||
|
discovery: discovery,
|
||||||
|
load: load,
|
||||||
|
localID: localID,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name returns the tool name.
|
||||||
|
func (t *SwarmNodesTool) Name() string {
|
||||||
|
return "swarm_nodes"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Description returns the tool description.
|
||||||
|
func (t *SwarmNodesTool) Description() string {
|
||||||
|
//nolint:gosmopolitan // Intentional Chinese examples for bilingual LLM tool descriptions
|
||||||
|
return "List ALL swarm cluster nodes and their detailed status in ONE call (fast, parallel, reads from local cache). " +
|
||||||
|
"IMPORTANT: This already returns status for ALL nodes - do NOT call swarm_route for each node individually. " +
|
||||||
|
"Use for questions like '查看所有节点状态', 'cluster status', '节点列表', etc."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters returns the tool parameters schema.
|
||||||
|
func (t *SwarmNodesTool) Parameters() map[string]any {
|
||||||
|
return map[string]any{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]any{
|
||||||
|
"verbose": map[string]any{
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "Whether to include detailed status for each node",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute executes the swarm nodes tool.
|
||||||
|
func (t *SwarmNodesTool) Execute(_ context.Context, args map[string]any) *ToolResult {
|
||||||
|
if t.discovery == nil {
|
||||||
|
return ErrorResult("Swarm mode is not enabled")
|
||||||
|
}
|
||||||
|
|
||||||
|
verbose := false
|
||||||
|
if rawVerbose, ok := args["verbose"]; ok {
|
||||||
|
switch v := rawVerbose.(type) {
|
||||||
|
case bool:
|
||||||
|
verbose = v
|
||||||
|
case string:
|
||||||
|
lower := strings.ToLower(strings.TrimSpace(v))
|
||||||
|
verbose = lower == "true" || lower == "1" || lower == "yes" || lower == "y"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
status := FormatClusterStatus(t.discovery, t.load, t.localID, verbose)
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: status,
|
||||||
|
ForUser: status,
|
||||||
|
IsError: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SwarmTool provides a generalized swarm entrypoint for status/query and routing.
|
||||||
|
type SwarmTool struct {
|
||||||
|
discovery swarm.Discovery
|
||||||
|
load *swarm.LoadMonitor
|
||||||
|
localID string
|
||||||
|
sendMessageFn func(ctx context.Context, targetNodeID, content, channel, chatID, senderID string) (string, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSwarmTool creates a new generalized swarm tool.
|
||||||
|
func NewSwarmTool(
|
||||||
|
discovery swarm.Discovery,
|
||||||
|
load *swarm.LoadMonitor,
|
||||||
|
localID string,
|
||||||
|
) *SwarmTool {
|
||||||
|
return &SwarmTool{
|
||||||
|
discovery: discovery,
|
||||||
|
load: load,
|
||||||
|
localID: localID,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetSendMessageFn sets the function used to send messages to other nodes.
|
||||||
|
func (t *SwarmTool) SetSendMessageFn(
|
||||||
|
fn func(ctx context.Context, targetNodeID, content, channel, chatID, senderID string) (string, error),
|
||||||
|
) {
|
||||||
|
t.sendMessageFn = fn
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name returns the tool name.
|
||||||
|
func (t *SwarmTool) Name() string {
|
||||||
|
return "swarm_tool"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Description returns the tool description.
|
||||||
|
func (t *SwarmTool) Description() string {
|
||||||
|
return "General swarm tool: list swarm nodes/status and route a message to a specific node."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters returns the tool parameters schema.
|
||||||
|
func (t *SwarmTool) Parameters() map[string]any {
|
||||||
|
return map[string]any{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]any{
|
||||||
|
"action": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"enum": []string{"list_nodes", "route_message"},
|
||||||
|
"description": "Operation type: list_nodes to query node status, route_message to send to another node",
|
||||||
|
},
|
||||||
|
"verbose": map[string]any{
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "When action is list_nodes, include detailed fields",
|
||||||
|
},
|
||||||
|
"node_id": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "When action is route_message, target node ID",
|
||||||
|
},
|
||||||
|
"message": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "When action is route_message, message content to send",
|
||||||
|
},
|
||||||
|
"channel": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Source channel identifier for routing context (optional, auto-filled from context when available)",
|
||||||
|
},
|
||||||
|
"chat_id": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Source chat/conversation ID for continuity (optional, auto-filled from context when available)",
|
||||||
|
},
|
||||||
|
"sender_id": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Sender identity for audit and reply routing (optional, auto-filled from context when available)",
|
||||||
|
},
|
||||||
|
"trace_id": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Distributed trace ID for cross-node observability (optional)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": []string{"action"},
|
||||||
|
"oneOf": []any{
|
||||||
|
map[string]any{
|
||||||
|
"properties": map[string]any{
|
||||||
|
"action": map[string]any{"const": "list_nodes"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
map[string]any{
|
||||||
|
"properties": map[string]any{
|
||||||
|
"action": map[string]any{"const": "route_message"},
|
||||||
|
},
|
||||||
|
"required": []string{"node_id", "message"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute executes the generalized swarm tool.
|
||||||
|
func (t *SwarmTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
||||||
|
if t.discovery == nil {
|
||||||
|
return ErrorResult("Swarm mode is not enabled")
|
||||||
|
}
|
||||||
|
|
||||||
|
action, _ := args["action"].(string)
|
||||||
|
action = strings.ToLower(strings.TrimSpace(action))
|
||||||
|
|
||||||
|
// Fallback inference when model omits action.
|
||||||
|
if action == "" {
|
||||||
|
if nodeID, _ := args["node_id"].(string); strings.TrimSpace(nodeID) != "" {
|
||||||
|
action = "route_message"
|
||||||
|
} else {
|
||||||
|
action = "list_nodes"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
switch action {
|
||||||
|
case "list_nodes":
|
||||||
|
verbose := false
|
||||||
|
if rawVerbose, ok := args["verbose"]; ok {
|
||||||
|
switch v := rawVerbose.(type) {
|
||||||
|
case bool:
|
||||||
|
verbose = v
|
||||||
|
case string:
|
||||||
|
lower := strings.ToLower(strings.TrimSpace(v))
|
||||||
|
verbose = lower == "true" || lower == "1" || lower == "yes" || lower == "y"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
status := FormatClusterStatus(t.discovery, t.load, t.localID, verbose)
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: status,
|
||||||
|
ForUser: status,
|
||||||
|
IsError: false,
|
||||||
|
}
|
||||||
|
|
||||||
|
case "route_message":
|
||||||
|
nodeID, _ := args["node_id"].(string)
|
||||||
|
nodeID = strings.TrimSpace(nodeID)
|
||||||
|
message, _ := args["message"].(string)
|
||||||
|
|
||||||
|
if nodeID == "" {
|
||||||
|
return ErrorResult("node_id is required when action=route_message")
|
||||||
|
}
|
||||||
|
if message == "" {
|
||||||
|
return ErrorResult("message is required when action=route_message")
|
||||||
|
}
|
||||||
|
if nodeID == t.localID {
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: fmt.Sprintf(
|
||||||
|
"Target node %s is this node. Process locally instead of routing.",
|
||||||
|
nodeID,
|
||||||
|
),
|
||||||
|
//nolint:gosmopolitan // Intentional Chinese user-facing message
|
||||||
|
ForUser: fmt.Sprintf("目标节点 %s 就是当前节点,已建议本地处理。", nodeID),
|
||||||
|
IsError: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if t.sendMessageFn == nil {
|
||||||
|
return ErrorResult("inter-node messaging is not available")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Extract optional context parameters for cross-node traceability.
|
||||||
|
channel, _ := args["channel"].(string)
|
||||||
|
chatID, _ := args["chat_id"].(string)
|
||||||
|
senderID, _ := args["sender_id"].(string)
|
||||||
|
|
||||||
|
response, err := t.sendMessageFn(ctx, nodeID, message, channel, chatID, senderID)
|
||||||
|
if err != nil {
|
||||||
|
return ErrorResult(fmt.Sprintf("failed to send message to node %s: %v", nodeID, err))
|
||||||
|
}
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: response,
|
||||||
|
ForUser: response,
|
||||||
|
IsError: false,
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
return ErrorResult("invalid action, supported: list_nodes, route_message")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SwarmRouteTool routes a request to a specific node in the swarm.
|
||||||
|
type SwarmRouteTool struct {
|
||||||
|
discovery swarm.Discovery
|
||||||
|
handoff *swarm.HandoffCoordinator
|
||||||
|
localID string
|
||||||
|
sendMessageFn func(ctx context.Context, targetNodeID, content, channel, chatID, senderID string) (string, error)
|
||||||
|
sendActionFn func(ctx context.Context, targetNodeID, action string) (string, error)
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSwarmRouteTool creates a new swarm route tool.
|
||||||
|
func NewSwarmRouteTool(
|
||||||
|
discovery swarm.Discovery,
|
||||||
|
handoff *swarm.HandoffCoordinator,
|
||||||
|
localID string,
|
||||||
|
) *SwarmRouteTool {
|
||||||
|
return &SwarmRouteTool{
|
||||||
|
discovery: discovery,
|
||||||
|
handoff: handoff,
|
||||||
|
localID: localID,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetSendMessageFn sets the function to send messages to other nodes.
|
||||||
|
func (t *SwarmRouteTool) SetSendMessageFn(
|
||||||
|
fn func(ctx context.Context, targetNodeID, content, channel, chatID, senderID string) (string, error),
|
||||||
|
) {
|
||||||
|
t.sendMessageFn = fn
|
||||||
|
}
|
||||||
|
|
||||||
|
// SetSendActionFn sets the function to send actions to other nodes.
|
||||||
|
func (t *SwarmRouteTool) SetSendActionFn(
|
||||||
|
fn func(ctx context.Context, targetNodeID, action string) (string, error),
|
||||||
|
) {
|
||||||
|
t.sendActionFn = fn
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name returns the tool name.
|
||||||
|
func (t *SwarmRouteTool) Name() string {
|
||||||
|
return "swarm_route"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Description returns the tool description.
|
||||||
|
func (t *SwarmRouteTool) Description() string {
|
||||||
|
return "Send a request to a SINGLE node. " +
|
||||||
|
"WARNING: Use swarm_batch or swarm_nodes FIRST for multi-node queries - they are MUCH FASTER. " +
|
||||||
|
"Only use this when you need ONE specific node to perform a task. " +
|
||||||
|
"For status, use action='status' (<10ms); for LLM tasks, use action='message' with a message."
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parameters returns the tool parameters schema.
|
||||||
|
func (t *SwarmRouteTool) Parameters() map[string]any {
|
||||||
|
return map[string]any{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]any{
|
||||||
|
"node_id": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "The target node ID to route the request to",
|
||||||
|
},
|
||||||
|
"action": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"enum": []string{"status", "message"},
|
||||||
|
"default": "status",
|
||||||
|
"description": "Action type: 'status' (default, fast, <10ms) returns structured data; 'message' sends to LLM for processing (slow, ~60s+)",
|
||||||
|
},
|
||||||
|
"message": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "The message to send to the target node (only for action='message')",
|
||||||
|
},
|
||||||
|
"channel": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Source channel identifier for routing context (optional)",
|
||||||
|
},
|
||||||
|
"chat_id": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Source chat/conversation ID for continuity (optional)",
|
||||||
|
},
|
||||||
|
"sender_id": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Sender identity for audit and reply routing (optional)",
|
||||||
|
},
|
||||||
|
"trace_id": map[string]any{
|
||||||
|
"type": "string",
|
||||||
|
"description": "Distributed trace ID for cross-node observability (optional)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": []string{"node_id"},
|
||||||
|
"oneOf": []any{
|
||||||
|
map[string]any{
|
||||||
|
"properties": map[string]any{
|
||||||
|
"action": map[string]any{"const": "status"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
map[string]any{
|
||||||
|
"properties": map[string]any{
|
||||||
|
"action": map[string]any{"const": "message"},
|
||||||
|
},
|
||||||
|
"required": []string{"message"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Execute executes the swarm route tool.
|
||||||
|
func (t *SwarmRouteTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
|
||||||
|
if t.discovery == nil {
|
||||||
|
return ErrorResult("Swarm mode is not enabled")
|
||||||
|
}
|
||||||
|
|
||||||
|
nodeID, _ := args["node_id"].(string)
|
||||||
|
action, _ := args["action"].(string)
|
||||||
|
message, _ := args["message"].(string)
|
||||||
|
|
||||||
|
// Smart action inference:
|
||||||
|
// - If action is explicitly set, use it
|
||||||
|
// - If message is provided, use "message"
|
||||||
|
// - Otherwise, default to "status" (fast path)
|
||||||
|
if action == "" {
|
||||||
|
if message != "" {
|
||||||
|
action = "message"
|
||||||
|
} else {
|
||||||
|
action = "status"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if nodeID == "" {
|
||||||
|
return ErrorResult("node_id is required")
|
||||||
|
}
|
||||||
|
|
||||||
|
// For "message" action, content is required
|
||||||
|
if action == "message" && message == "" {
|
||||||
|
return ErrorResult(
|
||||||
|
"message is required when action='message'. For status queries, use action='status' instead.",
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if target is this node
|
||||||
|
if nodeID == t.localID {
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: fmt.Sprintf("Target node %s is this node. Processing locally.", nodeID),
|
||||||
|
ForUser: fmt.Sprintf("This request is already on node %s.", nodeID),
|
||||||
|
IsError: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Find target node
|
||||||
|
members := t.discovery.Members()
|
||||||
|
var target *swarm.NodeWithState
|
||||||
|
for _, m := range members {
|
||||||
|
if m.Node.ID == nodeID {
|
||||||
|
target = m
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if target == nil {
|
||||||
|
return ErrorResult(fmt.Sprintf("Node %s not found in cluster", nodeID))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if target is available
|
||||||
|
if target.State.Status != swarm.NodeStatusAlive {
|
||||||
|
return ErrorResult(fmt.Sprintf("Node %s is not alive (status: %s)", nodeID, target.State.Status))
|
||||||
|
}
|
||||||
|
|
||||||
|
if target.Node.LoadScore > swarm.DefaultAvailableLoadThreshold {
|
||||||
|
return ErrorResult(fmt.Sprintf("Node %s is overloaded (load: %.0f%%)", nodeID, target.Node.LoadScore*100))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fast path: "status" action returns structured data directly
|
||||||
|
if action == "status" {
|
||||||
|
return t.executeStatusRequest(ctx, nodeID, target)
|
||||||
|
}
|
||||||
|
|
||||||
|
if t.sendMessageFn == nil {
|
||||||
|
return ErrorResult("inter-node messaging is not available")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Extract optional context parameters for cross-node traceability.
|
||||||
|
channel, _ := args["channel"].(string)
|
||||||
|
chatID, _ := args["chat_id"].(string)
|
||||||
|
senderID, _ := args["sender_id"].(string)
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Sending message to node", map[string]any{
|
||||||
|
"target": nodeID,
|
||||||
|
"message": message,
|
||||||
|
})
|
||||||
|
|
||||||
|
response, err := t.sendMessageFn(ctx, nodeID, message, channel, chatID, senderID)
|
||||||
|
if err != nil {
|
||||||
|
return ErrorResult(fmt.Sprintf("Failed to send message to node %s: %v", nodeID, err))
|
||||||
|
}
|
||||||
|
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: response,
|
||||||
|
ForUser: response,
|
||||||
|
IsError: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// executeStatusRequest sends a status action request to a node.
|
||||||
|
func (t *SwarmRouteTool) executeStatusRequest(
|
||||||
|
ctx context.Context,
|
||||||
|
nodeID string,
|
||||||
|
target *swarm.NodeWithState,
|
||||||
|
) *ToolResult {
|
||||||
|
if t.sendActionFn == nil {
|
||||||
|
return ErrorResult("inter-node actions not available")
|
||||||
|
}
|
||||||
|
|
||||||
|
logger.InfoCF("swarm", "Requesting status from node", map[string]any{
|
||||||
|
"target": nodeID,
|
||||||
|
})
|
||||||
|
|
||||||
|
response, err := t.sendActionFn(ctx, nodeID, "status")
|
||||||
|
if err != nil {
|
||||||
|
return ErrorResult(fmt.Sprintf("Failed to get status from node %s: %v", nodeID, err))
|
||||||
|
}
|
||||||
|
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: fmt.Sprintf("Node %s status:\n%s", nodeID, response),
|
||||||
|
ForUser: fmt.Sprintf("Node %s status:\n%s", nodeID, response),
|
||||||
|
IsError: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParseNodeMention extracts node ID from a message like "@node-id: message"
|
||||||
|
func ParseNodeMention(message string) (nodeID, content string) {
|
||||||
|
message = strings.TrimSpace(message)
|
||||||
|
|
||||||
|
// Check for @node-id: pattern
|
||||||
|
if strings.HasPrefix(message, "@") {
|
||||||
|
rest := message[1:]
|
||||||
|
idx := strings.IndexAny(rest, ": \n")
|
||||||
|
if idx > 0 && rest[idx] == ':' {
|
||||||
|
nodeID = rest[:idx]
|
||||||
|
content = strings.TrimSpace(rest[idx+1:])
|
||||||
|
return nodeID, content
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return "", message
|
||||||
|
}
|
||||||
|
|
||||||
|
// Helper functions
|
||||||
|
|
||||||
|
func plural(n int) string {
|
||||||
|
if n == 1 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return "s"
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatDuration(d time.Duration) string {
|
||||||
|
if d < time.Second {
|
||||||
|
return fmt.Sprintf("%dms", d.Milliseconds())
|
||||||
|
}
|
||||||
|
if d < time.Minute {
|
||||||
|
return fmt.Sprintf("%.1fs", d.Seconds())
|
||||||
|
}
|
||||||
|
if d < time.Hour {
|
||||||
|
return fmt.Sprintf("%.1fm", d.Minutes())
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%.1fh", d.Hours())
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue