feat: implement swarm mode for multi-agent coordination

This commit is contained in:
Zhaoyikaiii 2026-02-28 15:06:41 +08:00
parent 6f5930624b
commit 74c192bcdd
No known key found for this signature in database
GPG key ID: 14C8E2DD713BB8D6
53 changed files with 10922 additions and 44 deletions

4
.gitignore vendored
View file

@ -47,3 +47,7 @@ dist/
# Windows Application Icon/Resource
*.syso
# PM2 and process manager configs (dev-specific)
ecosystem.config.js
.pm2/

View file

@ -6,6 +6,7 @@ import (
func NewGatewayCommand() *cobra.Command {
var debug bool
var configPath string
cmd := &cobra.Command{
Use: "gateway",
@ -13,11 +14,12 @@ func NewGatewayCommand() *cobra.Command {
Short: "Start picoclaw gateway",
Args: cobra.NoArgs,
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().StringVarP(&configPath, "config", "c", "", "Path to config file")
return cmd
}

View file

@ -39,13 +39,13 @@ import (
"github.com/sipeed/picoclaw/pkg/voice"
)
func gatewayCmd(debug bool) error {
func gatewayCmd(debug bool, configPath string) error {
if debug {
logger.SetLevel(logger.DEBUG)
fmt.Println("🔍 Debug mode enabled")
}
cfg, err := internal.LoadConfig()
cfg, err := loadConfigWithPath(configPath)
if err != nil {
return fmt.Errorf("error loading config: %w", err)
}
@ -252,3 +252,12 @@ func setupCronTool(
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
View 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
View file

@ -15,6 +15,7 @@ require (
github.com/mdp/qrterminal/v3 v3.2.1
github.com/modelcontextprotocol/go-sdk v1.3.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/openai/openai-go/v3 v3.22.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/mattn/go-colorable v0.1.14 // 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/petermattis/goid v0.0.0-20260113132338-7c7de50cc741 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect

6
go.sum
View file

@ -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/mymmrac/telego v1.6.0 h1:Zc8rgyHozvd/7ZgyrigyHdAF9koHYMfilYfyB6wlFC0=
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/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=

View file

@ -30,6 +30,7 @@ import (
"github.com/sipeed/picoclaw/pkg/routing"
"github.com/sipeed/picoclaw/pkg/skills"
"github.com/sipeed/picoclaw/pkg/state"
"github.com/sipeed/picoclaw/pkg/swarm"
"github.com/sipeed/picoclaw/pkg/tools"
"github.com/sipeed/picoclaw/pkg/utils"
"github.com/sipeed/picoclaw/pkg/voice"
@ -46,6 +47,22 @@ type AgentLoop struct {
channelManager *channels.Manager
mediaStore media.MediaStore
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
@ -84,7 +101,7 @@ func NewAgentLoop(
stateManager = state.NewManager(defaultAgent.Workspace)
}
return &AgentLoop{
al := &AgentLoop{
bus: msgBus,
cfg: cfg,
registry: registry,
@ -92,6 +109,13 @@ func NewAgentLoop(
summarizing: sync.Map{},
fallback: fallbackChain,
}
// Initialize swarm mode if enabled
if cfg.Swarm.Enabled {
al.initSwarm()
}
return al
}
// registerSharedTools registers tools that are shared across all agents (web, message, spawn).
@ -208,6 +232,10 @@ func registerSharedTools(
}
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)
// Initialize MCP servers for all agents
@ -350,6 +378,12 @@ func (al *AgentLoop) Run(ctx context.Context) error {
func (al *AgentLoop) Stop() {
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) {
@ -362,6 +396,12 @@ func (al *AgentLoop) RegisterTool(tool tools.Tool) {
func (al *AgentLoop) SetChannelManager(cm *channels.Manager) {
al.channelManager = cm
al.commandRegistry = cm.CommandRegistry()
// Wire up swarm commands
if al.swarmEnabled {
al.registerSwarmCommands()
}
}
// 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)
}
// 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
if response, handled := al.handleCommand(ctx, msg); handled {
return response, nil
@ -588,6 +647,22 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
"matched_by": route.MatchedBy,
})
// Check if we should handoff this request to another node
if al.swarmEnabled && al.shouldHandoff(agent, processOptions{
SessionKey: sessionKey,
Channel: msg.Channel,
ChatID: msg.ChatID,
UserMessage: msg.Content,
}) {
handoffResp, err := al.initiateSwarmHandoff(ctx, agent, sessionKey, msg)
if err == nil && handoffResp != nil && handoffResp.Accepted {
// Handoff was successful, return the response
return fmt.Sprintf("Your request has been handed off to node %s for processing.", handoffResp.NodeID), nil
}
// If handoff failed, continue processing locally
logger.WarnCF("swarm", "Handoff failed, processing locally", map[string]any{"error": err})
}
return al.runAgentLoop(ctx, agent, processOptions{
SessionKey: sessionKey,
Channel: msg.Channel,
@ -671,6 +746,10 @@ func (al *AgentLoop) runAgentLoop(
agent *AgentInstance,
opts processOptions,
) (string, error) {
// Track active session for swarm load monitoring
al.IncrementSwarmSessions()
defer al.DecrementSwarmSessions()
// 0. Record last channel for heartbeat notifications (skip internal channels)
if opts.Channel != "" && opts.ChatID != "" {
// Don't record internal channels (cli, system, subagent)
@ -1177,6 +1256,15 @@ func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, c
go func() {
defer al.summarizing.Delete(summarizeKey)
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)
}()
}
@ -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
}

925
pkg/agent/loop_swarm.go Normal file
View 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
View 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)
}

View file

@ -79,12 +79,13 @@ type Manager struct {
bus *bus.MessageBus
config *config.Config
mediaStore media.MediaStore
commands *CommandRegistry
dispatchTask *asyncTask
mux *http.ServeMux
httpServer *http.Server
mu sync.RWMutex
placeholders sync.Map // "channel:chatID" → placeholderID (string)
typingStops sync.Map // "channel:chatID" → func()
placeholders sync.Map // "channel:chatID" → placeholderEntry
typingStops sync.Map // "channel:chatID" → typingEntry
reactionUndos sync.Map // "channel:chatID" → reactionEntry
}
@ -154,6 +155,7 @@ func NewManager(cfg *config.Config, messageBus *bus.MessageBus, store media.Medi
bus: messageBus,
config: cfg,
mediaStore: store,
commands: NewCommandRegistry(),
}
if err := m.initChannels(); err != nil {
@ -786,6 +788,12 @@ func (m *Manager) RegisterChannel(name string, channel 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) {
m.mu.Lock()
defer m.mu.Unlock()

View file

@ -56,6 +56,8 @@ type PicoChannel struct {
connCount atomic.Int32
ctx context.Context
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.
@ -403,6 +405,9 @@ func (c *PicoChannel) handleMessage(pc *picoConn, msg PicoMessage) {
case TypeMessageSend:
c.handleMessageSend(pc, msg)
case TypeNodeRequest:
c.handleNodeRequest(pc, msg)
default:
errMsg := newError("unknown_type", fmt.Sprintf("unknown message type: %s", msg.Type))
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)
}
// 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.
func truncate(s string, maxLen int) string {
runes := []rune(s)

View file

@ -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
import "time"
import (
"time"
// Protocol message types.
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"
picoproto "github.com/sipeed/picoclaw/pkg/pico/protocol"
)
// 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.
type PicoMessage 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"`
}
// This is a local type alias for the canonical protocol.Message type.
type PicoMessage = picoproto.Message
// newMessage creates a PicoMessage with the given type and payload.
func newMessage(msgType string, payload map[string]any) PicoMessage {
return PicoMessage{
Type: msgType,
Timestamp: time.Now().UnixMilli(),
Payload: payload,
}
return picoproto.NewMessage(msgType, payload)
}
// newError creates an error PicoMessage.
func newError(code, message string) PicoMessage {
return newMessage(TypeError, map[string]any{
"code": code,
"message": message,
})
return picoproto.NewError(code, 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()
}

View file

@ -58,6 +58,7 @@ type Config struct {
Tools ToolsConfig `json:"tools"`
Heartbeat HeartbeatConfig `json:"heartbeat"`
Devices DevicesConfig `json:"devices"`
Swarm SwarmConfig `json:"swarm,omitempty"` // Swarm mode configuration
}
// MarshalJSON implements custom JSON marshaling for Config
@ -900,3 +901,134 @@ func (t *ToolsConfig) IsToolEnabled(name string) bool {
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
View 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
View 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
View 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
View 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
View 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
View 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}
}

View 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
View 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
View 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
View 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
View 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
View 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
}

View 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
View 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))
}

View 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
}

View 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
View 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
View 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
View 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
}

View 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
View 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()
}

View 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
View 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
View 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
View 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()
}

View 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
View 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
View 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
}

View 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
View 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
View 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
View 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)
}
}
}

View 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
}

View 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
View 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
View 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
View 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
View 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
View 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
View 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())
}