diff --git a/.gitignore b/.gitignore index 02ef18d1f..03b01e38e 100644 --- a/.gitignore +++ b/.gitignore @@ -47,3 +47,7 @@ dist/ # Windows Application Icon/Resource *.syso + +# PM2 and process manager configs (dev-specific) +ecosystem.config.js +.pm2/ diff --git a/cmd/picoclaw/internal/gateway/command.go b/cmd/picoclaw/internal/gateway/command.go index 66a56f9ce..a4cf111a4 100644 --- a/cmd/picoclaw/internal/gateway/command.go +++ b/cmd/picoclaw/internal/gateway/command.go @@ -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 } diff --git a/cmd/picoclaw/internal/gateway/helpers.go b/cmd/picoclaw/internal/gateway/helpers.go index 174f5db62..bd7e66ccf 100644 --- a/cmd/picoclaw/internal/gateway/helpers.go +++ b/cmd/picoclaw/internal/gateway/helpers.go @@ -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) +} diff --git a/docs/swarm-architecture.md b/docs/swarm-architecture.md new file mode 100644 index 000000000..841a1d651 --- /dev/null +++ b/docs/swarm-architecture.md @@ -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.` | Heartbeat / KV renewal | +| `picoclaw.swarm.node..msg` | Direct node messaging (request-reply) | +| `picoclaw.swarm.handoff.` | Targeted handoff requests (request-reply) | +| `picoclaw.swarm.leader` | Leader election announcements | +| `picoclaw.swarm.metrics.` | Metrics publishing | + +## JetStream KV Buckets + +| Bucket | Key Pattern | TTL | Purpose | +|--------|-------------|-----|---------| +| `swarm_members` | `node.` | 15s (configurable) | Membership liveness via TTL expiry | +| `swarm_status` | `status.` | 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., 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. + 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..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) diff --git a/go.mod b/go.mod index 238bd405c..cdfdbd989 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 060594d06..9610aa5f9 100644 --- a/go.sum +++ b/go.sum @@ -138,6 +138,12 @@ github.com/modelcontextprotocol/go-sdk v1.3.0 h1:gMfZkv3DzQF5q/DcQePo5rahEY+sguy github.com/modelcontextprotocol/go-sdk v1.3.0/go.mod h1:AnQ//Qc6+4nIyyrB4cxBU7UW9VibK4iOZBeyP/rF1IE= github.com/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= diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 1ab79f3ca..088151385 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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 } diff --git a/pkg/agent/loop_swarm.go b/pkg/agent/loop_swarm.go new file mode 100644 index 000000000..a5e4d0092 --- /dev/null +++ b/pkg/agent/loop_swarm.go @@ -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 +} diff --git a/pkg/channels/commands.go b/pkg/channels/commands.go new file mode 100644 index 000000000..e0ebbc11a --- /dev/null +++ b/pkg/channels/commands.go @@ -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) +} diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index fdd6d0c1f..92daf73e9 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -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() diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index 8d8b62a67..c93e11edf 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -50,12 +50,14 @@ func (pc *picoConn) close() { // It serves as the reference implementation for all optional capability interfaces. type PicoChannel struct { *channels.BaseChannel - config config.PicoConfig - upgrader websocket.Upgrader - connections sync.Map // connID โ†’ *picoConn - connCount atomic.Int32 - ctx context.Context - cancel context.CancelFunc + config config.PicoConfig + upgrader websocket.Upgrader + connections sync.Map // connID โ†’ *picoConn + 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) diff --git a/pkg/channels/pico/protocol.go b/pkg/channels/pico/protocol.go index 0a630e193..4620cc7d7 100644 --- a/pkg/channels/pico/protocol.go +++ b/pkg/channels/pico/protocol.go @@ -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() } diff --git a/pkg/config/config.go b/pkg/config/config.go index 0ee3acfe0..a6ac517c9 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -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"` +} diff --git a/pkg/kv/bucket.go b/pkg/kv/bucket.go new file mode 100644 index 000000000..c7e2a6c9c --- /dev/null +++ b/pkg/kv/bucket.go @@ -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 +} diff --git a/pkg/kv/nats.go b/pkg/kv/nats.go new file mode 100644 index 000000000..27eb3a1a1 --- /dev/null +++ b/pkg/kv/nats.go @@ -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) + } + }) +} diff --git a/pkg/pico/addr.go b/pkg/pico/addr.go new file mode 100644 index 000000000..185d09718 --- /dev/null +++ b/pkg/pico/addr.go @@ -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) +} diff --git a/pkg/pico/client.go b/pkg/pico/client.go new file mode 100644 index 000000000..d3d27b42e --- /dev/null +++ b/pkg/pico/client.go @@ -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 + } +} diff --git a/pkg/pico/node_client.go b/pkg/pico/node_client.go new file mode 100644 index 000000000..8c7f0a980 --- /dev/null +++ b/pkg/pico/node_client.go @@ -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 +} diff --git a/pkg/pico/node_payload.go b/pkg/pico/node_payload.go new file mode 100644 index 000000000..04b78c996 --- /dev/null +++ b/pkg/pico/node_payload.go @@ -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} +} diff --git a/pkg/pico/protocol/protocol.go b/pkg/pico/protocol/protocol.go new file mode 100644 index 000000000..2f5511bd5 --- /dev/null +++ b/pkg/pico/protocol/protocol.go @@ -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, + }) +} diff --git a/pkg/pico/types.go b/pkg/pico/types.go new file mode 100644 index 000000000..492ea17ef --- /dev/null +++ b/pkg/pico/types.go @@ -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"` +} diff --git a/pkg/swarm/bucket_config.go b/pkg/swarm/bucket_config.go new file mode 100644 index 000000000..0340b31d9 --- /dev/null +++ b/pkg/swarm/bucket_config.go @@ -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() +} diff --git a/pkg/swarm/config.go b/pkg/swarm/config.go new file mode 100644 index 000000000..e9afaebce --- /dev/null +++ b/pkg/swarm/config.go @@ -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", + }, + } +} diff --git a/pkg/swarm/constants.go b/pkg/swarm/constants.go new file mode 100644 index 000000000..1d17017af --- /dev/null +++ b/pkg/swarm/constants.go @@ -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. + SubjectHeartbeat = SubjectPrefix + ".heartbeat" + + // SubjectMetrics is the prefix for metrics messages: picoclaw.swarm.metrics. + SubjectMetrics = SubjectPrefix + ".metrics" + + // SubjectNodeMsg is the prefix for direct node messages: picoclaw.swarm.node..msg + SubjectNodeMsg = SubjectPrefix + ".node" + + // SubjectHandoff is the prefix for targeted handoff: picoclaw.swarm.handoff. + 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 +) diff --git a/pkg/swarm/coordinator.go b/pkg/swarm/coordinator.go new file mode 100644 index 000000000..5431c59c6 --- /dev/null +++ b/pkg/swarm/coordinator.go @@ -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 +} diff --git a/pkg/swarm/coordinator_interface.go b/pkg/swarm/coordinator_interface.go new file mode 100644 index 000000000..6c879b7cf --- /dev/null +++ b/pkg/swarm/coordinator_interface.go @@ -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 +} diff --git a/pkg/swarm/crypto.go b/pkg/swarm/crypto.go new file mode 100644 index 000000000..92ec6c50c --- /dev/null +++ b/pkg/swarm/crypto.go @@ -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)) +} diff --git a/pkg/swarm/discovery_adapter.go b/pkg/swarm/discovery_adapter.go new file mode 100644 index 000000000..1dd5a84a8 --- /dev/null +++ b/pkg/swarm/discovery_adapter.go @@ -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 +} diff --git a/pkg/swarm/discovery_interface.go b/pkg/swarm/discovery_interface.go new file mode 100644 index 000000000..96ec058da --- /dev/null +++ b/pkg/swarm/discovery_interface.go @@ -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 +} diff --git a/pkg/swarm/discovery_nats.go b/pkg/swarm/discovery_nats.go new file mode 100644 index 000000000..d492524a2 --- /dev/null +++ b/pkg/swarm/discovery_nats.go @@ -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() +} diff --git a/pkg/swarm/errors.go b/pkg/swarm/errors.go new file mode 100644 index 000000000..505625c93 --- /dev/null +++ b/pkg/swarm/errors.go @@ -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) +} diff --git a/pkg/swarm/handoff.go b/pkg/swarm/handoff.go new file mode 100644 index 000000000..2e7e9867c --- /dev/null +++ b/pkg/swarm/handoff.go @@ -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. +// +// 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 +} diff --git a/pkg/swarm/leader_election.go b/pkg/swarm/leader_election.go new file mode 100644 index 000000000..01a667076 --- /dev/null +++ b/pkg/swarm/leader_election.go @@ -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, + } +} diff --git a/pkg/swarm/load_monitor.go b/pkg/swarm/load_monitor.go new file mode 100644 index 000000000..52148c5e5 --- /dev/null +++ b/pkg/swarm/load_monitor.go @@ -0,0 +1,331 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package swarm + +import ( + "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() +} diff --git a/pkg/swarm/load_reporter_interface.go b/pkg/swarm/load_reporter_interface.go new file mode 100644 index 000000000..77c92e7c3 --- /dev/null +++ b/pkg/swarm/load_reporter_interface.go @@ -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 +} diff --git a/pkg/swarm/manager.go b/pkg/swarm/manager.go new file mode 100644 index 000000000..3a55afe70 --- /dev/null +++ b/pkg/swarm/manager.go @@ -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() +} diff --git a/pkg/swarm/membership.go b/pkg/swarm/membership.go new file mode 100644 index 000000000..1d7e2ac36 --- /dev/null +++ b/pkg/swarm/membership.go @@ -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() + } + }() +} diff --git a/pkg/swarm/metrics.go b/pkg/swarm/metrics.go new file mode 100644 index 000000000..67fb9489e --- /dev/null +++ b/pkg/swarm/metrics.go @@ -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() +} diff --git a/pkg/swarm/nats_membership.go b/pkg/swarm/nats_membership.go new file mode 100644 index 000000000..55e4c8852 --- /dev/null +++ b/pkg/swarm/nats_membership.go @@ -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) + } + } +} diff --git a/pkg/swarm/nats_status.go b/pkg/swarm/nats_status.go new file mode 100644 index 000000000..9578e80e8 --- /dev/null +++ b/pkg/swarm/nats_status.go @@ -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 +} diff --git a/pkg/swarm/node.go b/pkg/swarm/node.go new file mode 100644 index 000000000..6bc63879f --- /dev/null +++ b/pkg/swarm/node.go @@ -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 +} diff --git a/pkg/swarm/node_discovery_interface.go b/pkg/swarm/node_discovery_interface.go new file mode 100644 index 000000000..944843dfc --- /dev/null +++ b/pkg/swarm/node_discovery_interface.go @@ -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) +} diff --git a/pkg/swarm/node_status.go b/pkg/swarm/node_status.go new file mode 100644 index 000000000..c4d3911bb --- /dev/null +++ b/pkg/swarm/node_status.go @@ -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 +} diff --git a/pkg/swarm/nonce.go b/pkg/swarm/nonce.go new file mode 100644 index 000000000..5b13a75fb --- /dev/null +++ b/pkg/swarm/nonce.go @@ -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(), + } +} diff --git a/pkg/swarm/policy.go b/pkg/swarm/policy.go new file mode 100644 index 000000000..64147fc3d --- /dev/null +++ b/pkg/swarm/policy.go @@ -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) + } + } +} diff --git a/pkg/swarm/router_default.go b/pkg/swarm/router_default.go new file mode 100644 index 000000000..6a3376d07 --- /dev/null +++ b/pkg/swarm/router_default.go @@ -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 +} diff --git a/pkg/swarm/router_interface.go b/pkg/swarm/router_interface.go new file mode 100644 index 000000000..ab5752220 --- /dev/null +++ b/pkg/swarm/router_interface.go @@ -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 +} diff --git a/pkg/swarm/swarm_test.go b/pkg/swarm/swarm_test.go new file mode 100644 index 000000000..c244b0f2e --- /dev/null +++ b/pkg/swarm/swarm_test.go @@ -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") +} diff --git a/pkg/swarm/todo.go b/pkg/swarm/todo.go new file mode 100644 index 000000000..78eb7997d --- /dev/null +++ b/pkg/swarm/todo.go @@ -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()) +} diff --git a/pkg/swarm/view.go b/pkg/swarm/view.go new file mode 100644 index 000000000..c07dabc15 --- /dev/null +++ b/pkg/swarm/view.go @@ -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.) +} diff --git a/pkg/tools/handoff_tool.go b/pkg/tools/handoff_tool.go new file mode 100644 index 000000000..343ce4eb9 --- /dev/null +++ b/pkg/tools/handoff_tool.go @@ -0,0 +1,140 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// Swarm mode support for multi-agent coordination +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package tools + +import ( + "context" + "fmt" + + "github.com/sipeed/picoclaw/pkg/swarm" +) + +// HandoffTool implements the handoff tool for swarm mode. +type HandoffTool struct { + coordinator *swarm.HandoffCoordinator + channel string + chatID string +} + +// NewHandoffTool creates a new handoff tool. +func NewHandoffTool(coordinator *swarm.HandoffCoordinator) *HandoffTool { + return &HandoffTool{ + coordinator: coordinator, + channel: "cli", + chatID: "direct", + } +} + +// Name returns the tool name. +func (t *HandoffTool) Name() string { + return "handoff" +} + +// Description returns the tool description. +func (t *HandoffTool) Description() string { + return "Delegate this task to another agent in the swarm. Use when you cannot handle the task due to capability constraints or system overload." +} + +// Parameters returns the tool parameters schema. +func (t *HandoffTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "reason": map[string]any{ + "type": "string", + "enum": []string{"no_capability", "overloaded", "user_request"}, + "description": "The reason for handing off this task", + }, + "required_capability": map[string]any{ + "type": "string", + "description": "The specific capability required to handle this task", + }, + "context": map[string]any{ + "type": "string", + "description": "Additional context about why this handoff is needed", + }, + }, + } +} + +// SetContext sets the channel and chat ID for the tool. +func (t *HandoffTool) SetContext(channel, chatID string) { + t.channel = channel + t.chatID = chatID +} + +// Execute executes the handoff tool. +func (t *HandoffTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + if t.coordinator == nil { + return ErrorResult("Swarm mode is not enabled or handoff coordinator not configured").WithError( + fmt.Errorf("handoff coordinator is nil")) + } + + // Parse reason + reasonStr, _ := args["reason"].(string) + var reason swarm.HandoffReason + switch reasonStr { + case "no_capability": + reason = swarm.ReasonNoCapability + case "overloaded": + reason = swarm.ReasonOverloaded + case "user_request": + reason = swarm.ReasonUserRequest + default: + reason = swarm.ReasonNoCapability + } + + // Parse required capability + requiredCap, _ := args["required_capability"].(string) + + // Parse context + contextMsg, _ := args["context"].(string) + + // Build handoff request + req := &swarm.HandoffRequest{ + Reason: reason, + RequiredCap: requiredCap, + Metadata: make(map[string]string), + } + + if contextMsg != "" { + req.Metadata["context"] = contextMsg + } + + // Execute handoff + resp, err := t.coordinator.InitiateHandoff(ctx, req) + if err != nil { + return ErrorResult(fmt.Sprintf("Handoff failed: %v", err)).WithError(err) + } + + if !resp.Accepted { + return ErrorResult(fmt.Sprintf("Handoff rejected by all nodes: %s", resp.Reason)).WithError( + fmt.Errorf("handoff rejected: %s", resp.Reason)) + } + + // Build result message + resultMsg := fmt.Sprintf("Task handed off to node %s\n", resp.NodeID) + if resp.Reason != "" { + resultMsg += fmt.Sprintf("Note: %s\n", resp.Reason) + } + + return &ToolResult{ + ForLLM: resultMsg + "The target node will process this task and respond to the user.", + ForUser: "Your task has been delegated to another agent in the swarm. They will respond shortly.", + Silent: false, + IsError: false, + Async: true, // Handoff is async - target node will respond directly + } +} + +// CanHandle reports whether the local node can handle the given capability. +func (t *HandoffTool) CanHandle(requiredCap string) bool { + if t.coordinator == nil { + return true // If swarm is disabled, we can "handle" everything + } + return t.coordinator.CanHandle(requiredCap) +} diff --git a/pkg/tools/swarm_batch.go b/pkg/tools/swarm_batch.go new file mode 100644 index 000000000..fc9beda42 --- /dev/null +++ b/pkg/tools/swarm_batch.go @@ -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, + } +} diff --git a/pkg/tools/swarm_nodes.go b/pkg/tools/swarm_nodes.go new file mode 100644 index 000000000..adc14d9e8 --- /dev/null +++ b/pkg/tools/swarm_nodes.go @@ -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()) +}