Add Calibrate AI Control UI and fleet LLM backend wiring.
Some checks failed
CI Docker Mining Proof / Linux agent hashrate proof (push) Has been cancelled
Some checks failed
CI Docker Mining Proof / Linux agent hashrate proof (push) Has been cancelled
Operators toggle Logic gates vs AI Control on Settings, refresh local Ollama models, and save ai_endpoint settings via Calibrate PUT; server scheduler and agent snapshot/command paths support stateless 60s fleet decisions.
This commit is contained in:
143
server/internal/ai/client.go
Normal file
143
server/internal/ai/client.go
Normal file
@@ -0,0 +1,143 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
var defaultHTTPClient = &http.Client{Timeout: 45 * time.Second}
|
||||
|
||||
// ListModels GET {endpoint}/models — OpenAI-compatible model list.
|
||||
func ListModels(ctx context.Context, endpoint string) ([]string, error) {
|
||||
return ListModelsWithClient(ctx, endpoint, defaultHTTPClient)
|
||||
}
|
||||
|
||||
func ListModelsWithClient(ctx context.Context, endpoint string, client *http.Client) ([]string, error) {
|
||||
base := normalizeEndpoint(endpoint)
|
||||
if base == "" {
|
||||
return nil, fmt.Errorf("endpoint is required")
|
||||
}
|
||||
if client == nil {
|
||||
client = defaultHTTPClient
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+"/models", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("models: status %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
var out struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
Models []struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"models"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &out); err != nil {
|
||||
return nil, fmt.Errorf("models: parse: %w", err)
|
||||
}
|
||||
names := make([]string, 0)
|
||||
seen := map[string]bool{}
|
||||
for _, m := range out.Data {
|
||||
id := strings.TrimSpace(m.ID)
|
||||
if id != "" && !seen[id] {
|
||||
seen[id] = true
|
||||
names = append(names, id)
|
||||
}
|
||||
}
|
||||
for _, m := range out.Models {
|
||||
name := strings.TrimSpace(m.Name)
|
||||
if name != "" && !seen[name] {
|
||||
seen[name] = true
|
||||
names = append(names, name)
|
||||
}
|
||||
}
|
||||
return names, nil
|
||||
}
|
||||
|
||||
// Decide POST chat/completions — single turn, no conversation history.
|
||||
func Decide(ctx context.Context, endpoint, model, systemPrompt, userPrompt string) (string, error) {
|
||||
return DecideWithClient(ctx, endpoint, model, systemPrompt, userPrompt, defaultHTTPClient)
|
||||
}
|
||||
|
||||
func DecideWithClient(ctx context.Context, endpoint, model, systemPrompt, userPrompt string, client *http.Client) (string, error) {
|
||||
base := normalizeEndpoint(endpoint)
|
||||
if base == "" {
|
||||
return "", fmt.Errorf("endpoint is required")
|
||||
}
|
||||
if client == nil {
|
||||
client = defaultHTTPClient
|
||||
}
|
||||
if strings.TrimSpace(model) == "" {
|
||||
model = "llama3.2"
|
||||
}
|
||||
payload := map[string]interface{}{
|
||||
"model": model,
|
||||
"messages": []map[string]string{
|
||||
{"role": "system", "content": systemPrompt},
|
||||
{"role": "user", "content": userPrompt},
|
||||
},
|
||||
"stream": false,
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("completions: status %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
|
||||
}
|
||||
var completion struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &completion); err != nil {
|
||||
return "", fmt.Errorf("completions: parse: %w", err)
|
||||
}
|
||||
if completion.Error != nil && completion.Error.Message != "" {
|
||||
return "", fmt.Errorf("completions: %s", completion.Error.Message)
|
||||
}
|
||||
if len(completion.Choices) == 0 {
|
||||
return "", fmt.Errorf("completions: empty choices")
|
||||
}
|
||||
return strings.TrimSpace(completion.Choices[0].Message.Content), nil
|
||||
}
|
||||
|
||||
func normalizeEndpoint(endpoint string) string {
|
||||
endpoint = strings.TrimSpace(endpoint)
|
||||
endpoint = strings.TrimRight(endpoint, "/")
|
||||
if endpoint == "" {
|
||||
return ""
|
||||
}
|
||||
if !strings.HasSuffix(endpoint, "/v1") {
|
||||
endpoint += "/v1"
|
||||
}
|
||||
return endpoint
|
||||
}
|
||||
66
server/internal/ai/client_test.go
Normal file
66
server/internal/ai/client_test.go
Normal file
@@ -0,0 +1,66 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestListModelsOpenAI(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/models" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"data": []map[string]string{{"id": "llama3.2"}, {"id": "mistral"}},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
models, err := ListModelsWithClient(context.Background(), srv.URL+"/v1", srv.Client())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(models) != 2 || models[0] != "llama3.2" {
|
||||
t.Fatalf("models: %v", models)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecideOpenAI(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/chat/completions" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
var req map[string]interface{}
|
||||
_ = json.NewDecoder(r.Body).Decode(&req)
|
||||
msgs, _ := req["messages"].([]interface{})
|
||||
if len(msgs) != 2 {
|
||||
t.Fatalf("expected single-turn messages, got %d", len(msgs))
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]interface{}{
|
||||
"choices": []map[string]interface{}{
|
||||
{"message": map[string]string{"content": `{"commands":[{"type":"noop","args":{}}]}`}},
|
||||
},
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
out, err := DecideWithClient(context.Background(), srv.URL+"/v1", "test-model", "sys", "user", srv.Client())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cmds := ParseCommands(out)
|
||||
if len(cmds) != 1 || cmds[0].Type != CmdNoop {
|
||||
t.Fatalf("parse: %+v", cmds)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeEndpoint(t *testing.T) {
|
||||
if got := normalizeEndpoint("http://127.0.0.1:11434"); got != "http://127.0.0.1:11434/v1" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
169
server/internal/ai/commands.go
Normal file
169
server/internal/ai/commands.go
Normal file
@@ -0,0 +1,169 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
CmdBulkCommand = "bulk_command"
|
||||
CmdAgentCommand = "agent_command"
|
||||
CmdDiscoverAndJoin = "discover_and_join"
|
||||
CmdRestartMining = "restart_mining"
|
||||
CmdReorderTiers = "reorder_tiers"
|
||||
CmdSpreadNow = "spread_now"
|
||||
CmdStageFetch = "stage_fetch"
|
||||
CmdSetAgentVersion = "set_agent_version"
|
||||
CmdNoop = "noop"
|
||||
)
|
||||
|
||||
var knownCommands = map[string]bool{
|
||||
CmdBulkCommand: true,
|
||||
CmdAgentCommand: true,
|
||||
CmdDiscoverAndJoin: true,
|
||||
CmdRestartMining: true,
|
||||
CmdReorderTiers: true,
|
||||
CmdSpreadNow: true,
|
||||
CmdStageFetch: true,
|
||||
CmdSetAgentVersion: true,
|
||||
CmdNoop: true,
|
||||
}
|
||||
|
||||
var jsonBlockRe = regexp.MustCompile(`(?s)\{[\s\n]*"commands"\s*:\s*\[[\s\S]*?\]\s*\}`)
|
||||
|
||||
// ParseCommands extracts fleet commands from LLM text (JSON block, tool-call, or COMMAND: lines).
|
||||
func ParseCommands(raw string) []Command {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return nil
|
||||
}
|
||||
if cmds := parseCommandsJSON(raw); len(cmds) > 0 {
|
||||
return cmds
|
||||
}
|
||||
if block := jsonBlockRe.FindString(raw); block != "" && block != raw {
|
||||
if cmds := parseCommandsJSON(block); len(cmds) > 0 {
|
||||
return cmds
|
||||
}
|
||||
}
|
||||
if cmd := parseToolCall(raw); cmd != nil {
|
||||
return []Command{*cmd}
|
||||
}
|
||||
return parseCommandLines(raw)
|
||||
}
|
||||
|
||||
func parseCommandsJSON(raw string) []Command {
|
||||
var envelope struct {
|
||||
Commands []Command `json:"commands"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw), &envelope); err == nil && len(envelope.Commands) > 0 {
|
||||
return normalizeCommands(envelope.Commands)
|
||||
}
|
||||
// Bare array
|
||||
var arr []Command
|
||||
if err := json.Unmarshal([]byte(raw), &arr); err == nil && len(arr) > 0 {
|
||||
return normalizeCommands(arr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func parseToolCall(raw string) *Command {
|
||||
var tool struct {
|
||||
Tool string `json:"tool"`
|
||||
Args map[string]interface{} `json:"args"`
|
||||
Type string `json:"type"`
|
||||
}
|
||||
start := strings.Index(raw, "{")
|
||||
end := strings.LastIndex(raw, "}")
|
||||
if start < 0 || end <= start {
|
||||
return nil
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw[start:end+1]), &tool); err != nil {
|
||||
return nil
|
||||
}
|
||||
name := strings.TrimSpace(tool.Tool)
|
||||
if name == "" {
|
||||
name = strings.TrimSpace(tool.Type)
|
||||
}
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
name = normalizeCommandType(name)
|
||||
if !knownCommands[name] && name != "restart_agent" {
|
||||
return &Command{Type: CmdAgentCommand, Args: map[string]interface{}{"action": name, "args": tool.Args}}
|
||||
}
|
||||
if name == "restart_agent" {
|
||||
name = CmdRestartMining
|
||||
}
|
||||
return &Command{Type: name, Args: tool.Args}
|
||||
}
|
||||
|
||||
func parseCommandLines(raw string) []Command {
|
||||
var out []Command
|
||||
for _, line := range strings.Split(raw, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if !strings.HasPrefix(strings.ToUpper(line), "COMMAND:") {
|
||||
continue
|
||||
}
|
||||
rest := strings.TrimSpace(line[len("COMMAND:"):])
|
||||
if rest == "" {
|
||||
continue
|
||||
}
|
||||
parts := strings.Fields(rest)
|
||||
cmdType := normalizeCommandType(parts[0])
|
||||
args := map[string]interface{}{}
|
||||
for _, p := range parts[1:] {
|
||||
kv := strings.SplitN(p, "=", 2)
|
||||
if len(kv) == 2 {
|
||||
args[kv[0]] = kv[1]
|
||||
}
|
||||
}
|
||||
out = append(out, Command{Type: cmdType, Args: args})
|
||||
}
|
||||
return normalizeCommands(out)
|
||||
}
|
||||
|
||||
func normalizeCommands(cmds []Command) []Command {
|
||||
out := make([]Command, 0, len(cmds))
|
||||
for _, c := range cmds {
|
||||
typ := normalizeCommandType(c.Type)
|
||||
if typ == "" {
|
||||
continue
|
||||
}
|
||||
if typ == "restart_agent" {
|
||||
typ = CmdRestartMining
|
||||
}
|
||||
args := c.Args
|
||||
if args == nil {
|
||||
args = map[string]interface{}{}
|
||||
}
|
||||
out = append(out, Command{Type: typ, Args: args})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalizeCommandType(s string) string {
|
||||
s = strings.TrimSpace(strings.ToLower(s))
|
||||
s = strings.ReplaceAll(s, "-", "_")
|
||||
if s == "restart" {
|
||||
return CmdRestartMining
|
||||
}
|
||||
if knownCommands[s] {
|
||||
return s
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// FormatExecuted summarizes commands for audit log storage.
|
||||
func FormatExecuted(cmds []Command, results []string) string {
|
||||
parts := make([]string, 0, len(cmds))
|
||||
for i, c := range cmds {
|
||||
msg := c.Type
|
||||
if i < len(results) && results[i] != "" {
|
||||
msg = fmt.Sprintf("%s:%s", c.Type, results[i])
|
||||
}
|
||||
parts = append(parts, msg)
|
||||
}
|
||||
return strings.Join(parts, "; ")
|
||||
}
|
||||
39
server/internal/ai/commands_test.go
Normal file
39
server/internal/ai/commands_test.go
Normal file
@@ -0,0 +1,39 @@
|
||||
package ai
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestParseCommandsJSONBlock(t *testing.T) {
|
||||
raw := `Here is my plan:
|
||||
{"commands":[{"type":"restart_mining","args":{}}]}`
|
||||
cmds := ParseCommands(raw)
|
||||
if len(cmds) != 1 || cmds[0].Type != CmdRestartMining {
|
||||
t.Fatalf("got %+v", cmds)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseCommandsToolCall(t *testing.T) {
|
||||
raw := `{"tool":"restart_agent","args":{}}`
|
||||
cmds := ParseCommands(raw)
|
||||
if len(cmds) != 1 || cmds[0].Type != CmdRestartMining {
|
||||
t.Fatalf("got %+v", cmds)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseCommandsLineFallback(t *testing.T) {
|
||||
raw := "COMMAND: spread_now\nCOMMAND: noop"
|
||||
cmds := ParseCommands(raw)
|
||||
if len(cmds) != 2 || cmds[0].Type != CmdSpreadNow || cmds[1].Type != CmdNoop {
|
||||
t.Fatalf("got %+v", cmds)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseCommandsAgentCommand(t *testing.T) {
|
||||
raw := `{"commands":[{"type":"agent_command","args":{"action":"pause"}}]}`
|
||||
cmds := ParseCommands(raw)
|
||||
if len(cmds) != 1 || cmds[0].Type != CmdAgentCommand {
|
||||
t.Fatalf("got %+v", cmds)
|
||||
}
|
||||
if cmds[0].Args["action"] != "pause" {
|
||||
t.Fatalf("args: %+v", cmds[0].Args)
|
||||
}
|
||||
}
|
||||
135
server/internal/ai/mission_prompt.go
Normal file
135
server/internal/ai/mission_prompt.go
Normal file
@@ -0,0 +1,135 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Default spread onion tiers (14) for prompt context.
|
||||
var defaultSpreadTiers = []string{
|
||||
"vuln_recon", "docker", "wsl", "powershell", "dotnet", "bits_curl",
|
||||
"do_peer", "wsus_cache_peer", "dns_txt", "webrtc_mesh",
|
||||
"smb", "winrm", "linux", "gpo",
|
||||
}
|
||||
|
||||
// SystemPrompt returns the fleet AI system instructions.
|
||||
func SystemPrompt() string {
|
||||
return strings.TrimSpace(`You are the AetherForge fleet controller for the operator's own machines.
|
||||
Respond with short answers only. Prefer JSON: {"commands":[{"type":"noop","args":{}}]}.
|
||||
Valid command types: bulk_command, agent_command, discover_and_join, restart_mining, reorder_tiers, spread_now, stage_fetch, set_agent_version, noop.
|
||||
agent_command args: action (required), command, path, data.
|
||||
bulk_command args: agent_ids (array), action, command.
|
||||
reorder_tiers args: tier_order (array of strings), skip_tiers (optional array).
|
||||
stage_fetch args: data (JSON manifest string).
|
||||
set_agent_version args: module or build_id.
|
||||
You have complete control in AI mode. Never target third-party systems.`)
|
||||
}
|
||||
|
||||
// BuildUserPrompt renders the per-agent snapshot for one decision cycle.
|
||||
func BuildUserPrompt(s AgentSnapshot) string {
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "Agent: name=%q id=%s", s.Name, s.AgentID)
|
||||
if s.Worker != "" {
|
||||
fmt.Fprintf(&b, " worker=%s", s.Worker)
|
||||
}
|
||||
b.WriteString("\n")
|
||||
fmt.Fprintf(&b, "Platform: GOOS=%s version=%s build=%s\n", firstNonEmpty(s.GOOS, s.Platform), s.Version, s.BuildID)
|
||||
if len(s.Capabilities) > 0 {
|
||||
flags := make([]string, 0, len(s.Capabilities))
|
||||
for k, v := range s.Capabilities {
|
||||
if v {
|
||||
flags = append(flags, k)
|
||||
}
|
||||
}
|
||||
if len(flags) > 0 {
|
||||
fmt.Fprintf(&b, "Forge capabilities: %s\n", strings.Join(flags, ", "))
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(&b, "LOTL tier: %s\n", emptyDash(s.LOTLTier))
|
||||
fmt.Fprintf(&b, "Mining hashrate: %.2f H/s\n", s.MiningHashrate)
|
||||
if s.ActiveMethod != "" {
|
||||
fmt.Fprintf(&b, "Active method: %s\n", s.ActiveMethod)
|
||||
}
|
||||
if s.ChainExhausted {
|
||||
b.WriteString("Mining chain exhausted: true\n")
|
||||
}
|
||||
if len(s.ChainOrder) > 0 {
|
||||
fmt.Fprintf(&b, "Chain order: %s\n", strings.Join(s.ChainOrder, " → "))
|
||||
}
|
||||
|
||||
b.WriteString("LOTL attempts (all tiers):\n")
|
||||
attemptByTier := map[string]TierAttempt{}
|
||||
for _, a := range s.LOTLAttempts {
|
||||
attemptByTier[a.Tier] = a
|
||||
}
|
||||
for _, tier := range defaultSpreadTiers {
|
||||
if a, ok := attemptByTier[tier]; ok {
|
||||
status := "fail"
|
||||
if a.OK {
|
||||
status = "ok"
|
||||
}
|
||||
if a.Error != "" {
|
||||
fmt.Fprintf(&b, " - %s: %s (%s)\n", tier, status, a.Error)
|
||||
} else {
|
||||
fmt.Fprintf(&b, " - %s: %s\n", tier, status)
|
||||
}
|
||||
} else {
|
||||
fmt.Fprintf(&b, " - %s: pending\n", tier)
|
||||
}
|
||||
}
|
||||
for _, a := range s.LOTLAttempts {
|
||||
if _, listed := attemptByTier[a.Tier]; listed {
|
||||
continue
|
||||
}
|
||||
found := false
|
||||
for _, t := range defaultSpreadTiers {
|
||||
if t == a.Tier {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
status := "fail"
|
||||
if a.OK {
|
||||
status = "ok"
|
||||
}
|
||||
fmt.Fprintf(&b, " - %s: %s\n", a.Tier, status)
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Fprintf(&b, "Join lane: %s\n", emptyDash(s.JoinLane))
|
||||
fmt.Fprintf(&b, "Spread state: %s\n", emptyDash(s.SpreadState))
|
||||
if s.VulnRisk != nil {
|
||||
fmt.Fprintf(&b, "Vuln risk score: %d\n", *s.VulnRisk)
|
||||
} else {
|
||||
b.WriteString("Vuln risk score: n/a\n")
|
||||
}
|
||||
if s.AdaptiveSummary != "" {
|
||||
fmt.Fprintf(&b, "Adaptive strategy summary: %s\n", s.AdaptiveSummary)
|
||||
} else if s.Adaptive != nil {
|
||||
if raw, err := json.Marshal(s.Adaptive); err == nil {
|
||||
fmt.Fprintf(&b, "Adaptive strategy: %s\n", string(raw))
|
||||
}
|
||||
}
|
||||
|
||||
b.WriteString("\nIf all 14 tiers failed and hashrate=0, you MAY force restart mining chain (restart_mining).\n")
|
||||
b.WriteString("Return JSON commands array for this agent only.\n")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func firstNonEmpty(vals ...string) string {
|
||||
for _, v := range vals {
|
||||
if strings.TrimSpace(v) != "" {
|
||||
return strings.TrimSpace(v)
|
||||
}
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
func emptyDash(s string) string {
|
||||
if strings.TrimSpace(s) == "" {
|
||||
return "—"
|
||||
}
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
188
server/internal/ai/scheduler.go
Normal file
188
server/internal/ai/scheduler.go
Normal file
@@ -0,0 +1,188 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SnapshotProvider supplies live agent telemetry for decision cycles.
|
||||
type SnapshotProvider interface {
|
||||
ConnectedAgentIDs() []string
|
||||
AgentSnapshot(agentID string) (AgentSnapshot, bool)
|
||||
}
|
||||
|
||||
// CommandExecutor runs parsed fleet commands.
|
||||
type CommandExecutor interface {
|
||||
Execute(agentID string, cmd Command) (summary string, err error)
|
||||
}
|
||||
|
||||
// ConfigProvider reads current Fleet AI Control settings.
|
||||
type ConfigProvider interface {
|
||||
AIConfig() Config
|
||||
}
|
||||
|
||||
// DecisionStore persists decision audit rows.
|
||||
type DecisionStore interface {
|
||||
InsertAIDecision(agentID, promptHash, response, commandsExecuted string) error
|
||||
}
|
||||
|
||||
// Scheduler runs periodic LLM decisions for online agents.
|
||||
type Scheduler struct {
|
||||
cfg ConfigProvider
|
||||
snap SnapshotProvider
|
||||
exec CommandExecutor
|
||||
store DecisionStore
|
||||
stop chan struct{}
|
||||
wg sync.WaitGroup
|
||||
|
||||
lastRunMu sync.Mutex
|
||||
lastRun map[string]time.Time
|
||||
}
|
||||
|
||||
func NewScheduler(cfg ConfigProvider, snap SnapshotProvider, exec CommandExecutor, store DecisionStore) *Scheduler {
|
||||
return &Scheduler{
|
||||
cfg: cfg,
|
||||
snap: snap,
|
||||
exec: exec,
|
||||
store: store,
|
||||
stop: make(chan struct{}),
|
||||
lastRun: make(map[string]time.Time),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Scheduler) Start() {
|
||||
s.wg.Add(1)
|
||||
go s.loop()
|
||||
}
|
||||
|
||||
func (s *Scheduler) Stop() {
|
||||
close(s.stop)
|
||||
s.wg.Wait()
|
||||
}
|
||||
|
||||
func (s *Scheduler) loop() {
|
||||
defer s.wg.Done()
|
||||
ticker := time.NewTicker(5 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-s.stop:
|
||||
return
|
||||
case <-ticker.C:
|
||||
s.tick()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Tick runs one scheduler pass (exported for tests).
|
||||
func (s *Scheduler) Tick() {
|
||||
s.tick()
|
||||
}
|
||||
|
||||
func (s *Scheduler) tick() {
|
||||
if s.cfg == nil || s.snap == nil {
|
||||
return
|
||||
}
|
||||
cfg := s.cfg.AIConfig()
|
||||
if !cfg.Enabled {
|
||||
return
|
||||
}
|
||||
interval := time.Duration(cfg.IntervalSec) * time.Second
|
||||
if interval < time.Second {
|
||||
interval = 60 * time.Second
|
||||
}
|
||||
ids := s.snap.ConnectedAgentIDs()
|
||||
for i, agentID := range ids {
|
||||
if !s.shouldRun(agentID, interval, i) {
|
||||
continue
|
||||
}
|
||||
s.runAgent(context.Background(), agentID, cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Scheduler) shouldRun(agentID string, interval time.Duration, staggerIndex int) bool {
|
||||
s.lastRunMu.Lock()
|
||||
defer s.lastRunMu.Unlock()
|
||||
last, ok := s.lastRun[agentID]
|
||||
if !ok {
|
||||
offset := time.Duration(staggerIndex%max(1, int(interval/time.Second))) * time.Second
|
||||
if offset > 0 {
|
||||
s.lastRun[agentID] = time.Now().Add(-interval + offset)
|
||||
}
|
||||
return true
|
||||
}
|
||||
return time.Since(last) >= interval
|
||||
}
|
||||
|
||||
func (s *Scheduler) markRun(agentID string) {
|
||||
s.lastRunMu.Lock()
|
||||
s.lastRun[agentID] = time.Now()
|
||||
s.lastRunMu.Unlock()
|
||||
}
|
||||
|
||||
func (s *Scheduler) runAgent(ctx context.Context, agentID string, cfg Config) {
|
||||
snap, ok := s.snap.AgentSnapshot(agentID)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
userPrompt := BuildUserPrompt(snap)
|
||||
systemPrompt := SystemPrompt()
|
||||
promptHash := hashPrompt(userPrompt)
|
||||
|
||||
decide := Decide
|
||||
if DecideFunc != nil {
|
||||
decide = DecideFunc
|
||||
}
|
||||
response, err := decide(ctx, cfg.Endpoint, cfg.Model, systemPrompt, userPrompt)
|
||||
if err != nil {
|
||||
log.Printf("[fleet-ai] agent %s decide: %v", agentID, err)
|
||||
if s.store != nil {
|
||||
_ = s.store.InsertAIDecision(agentID, promptHash, "", "error:"+err.Error())
|
||||
}
|
||||
s.markRun(agentID)
|
||||
return
|
||||
}
|
||||
|
||||
cmds := ParseCommands(response)
|
||||
results := make([]string, 0, len(cmds))
|
||||
for _, cmd := range cmds {
|
||||
if cmd.Type == CmdNoop {
|
||||
results = append(results, "ok")
|
||||
continue
|
||||
}
|
||||
if s.exec == nil {
|
||||
results = append(results, "no executor")
|
||||
continue
|
||||
}
|
||||
sum, execErr := s.exec.Execute(agentID, cmd)
|
||||
if execErr != nil {
|
||||
results = append(results, "err:"+execErr.Error())
|
||||
} else {
|
||||
results = append(results, sum)
|
||||
}
|
||||
}
|
||||
executed := FormatExecuted(cmds, results)
|
||||
if s.store != nil {
|
||||
_ = s.store.InsertAIDecision(agentID, promptHash, response, executed)
|
||||
}
|
||||
s.markRun(agentID)
|
||||
}
|
||||
|
||||
func hashPrompt(prompt string) string {
|
||||
h := sha256.Sum256([]byte(prompt))
|
||||
return hex.EncodeToString(h[:8])
|
||||
}
|
||||
|
||||
// DecideFunc allows tests to override LLM calls.
|
||||
var DecideFunc func(ctx context.Context, endpoint, model, systemPrompt, userPrompt string) (string, error)
|
||||
|
||||
func max(a, b int) int {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
92
server/internal/ai/scheduler_test.go
Normal file
92
server/internal/ai/scheduler_test.go
Normal file
@@ -0,0 +1,92 @@
|
||||
package ai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type mockSnap struct {
|
||||
ids []string
|
||||
snap AgentSnapshot
|
||||
}
|
||||
|
||||
func (m *mockSnap) ConnectedAgentIDs() []string { return m.ids }
|
||||
func (m *mockSnap) AgentSnapshot(string) (AgentSnapshot, bool) {
|
||||
return m.snap, true
|
||||
}
|
||||
|
||||
type mockExec struct {
|
||||
mu sync.Mutex
|
||||
calls []Command
|
||||
}
|
||||
|
||||
func (m *mockExec) Execute(_ string, cmd Command) (string, error) {
|
||||
m.mu.Lock()
|
||||
m.calls = append(m.calls, cmd)
|
||||
m.mu.Unlock()
|
||||
return cmd.Type, nil
|
||||
}
|
||||
|
||||
type mockCfg struct{ cfg Config }
|
||||
|
||||
func (m *mockCfg) AIConfig() Config { return m.cfg }
|
||||
|
||||
type mockStore struct {
|
||||
mu sync.Mutex
|
||||
rows []string
|
||||
}
|
||||
|
||||
func (m *mockStore) InsertAIDecision(_, _, _, executed string) error {
|
||||
m.mu.Lock()
|
||||
m.rows = append(m.rows, executed)
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestSchedulerExecutesRestartCommand(t *testing.T) {
|
||||
old := DecideFunc
|
||||
defer func() { DecideFunc = old }()
|
||||
DecideFunc = func(_ context.Context, _, _, _, _ string) (string, error) {
|
||||
return `{"commands":[{"type":"restart_mining","args":{}}]}`, nil
|
||||
}
|
||||
|
||||
exec := &mockExec{}
|
||||
store := &mockStore{}
|
||||
sched := NewScheduler(
|
||||
&mockCfg{cfg: Config{Enabled: true, Endpoint: "http://test/v1", IntervalSec: 1}},
|
||||
&mockSnap{ids: []string{"agent-1"}, snap: AgentSnapshot{AgentID: "agent-1", Name: "host"}},
|
||||
exec,
|
||||
store,
|
||||
)
|
||||
sched.lastRun["agent-1"] = time.Now().Add(-2 * time.Minute)
|
||||
sched.Tick()
|
||||
|
||||
exec.mu.Lock()
|
||||
n := len(exec.calls)
|
||||
call := exec.calls
|
||||
exec.mu.Unlock()
|
||||
if n != 1 || call[0].Type != CmdRestartMining {
|
||||
t.Fatalf("calls: %+v", call)
|
||||
}
|
||||
store.mu.Lock()
|
||||
defer store.mu.Unlock()
|
||||
if len(store.rows) != 1 || store.rows[0] != "restart_mining:restart_mining" {
|
||||
t.Fatalf("store: %v", store.rows)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerNoOpWhenDisabled(t *testing.T) {
|
||||
exec := &mockExec{}
|
||||
sched := NewScheduler(
|
||||
&mockCfg{cfg: Config{Enabled: false}},
|
||||
&mockSnap{ids: []string{"agent-1"}},
|
||||
exec,
|
||||
nil,
|
||||
)
|
||||
sched.Tick()
|
||||
if len(exec.calls) != 0 {
|
||||
t.Fatalf("expected no calls")
|
||||
}
|
||||
}
|
||||
63
server/internal/ai/types.go
Normal file
63
server/internal/ai/types.go
Normal file
@@ -0,0 +1,63 @@
|
||||
package ai
|
||||
|
||||
import "crypto-miner-server/internal/strategy"
|
||||
|
||||
// TierAttempt mirrors agent LOTL tier attempt telemetry.
|
||||
type TierAttempt struct {
|
||||
Tier string `json:"tier"`
|
||||
OK bool `json:"ok"`
|
||||
Error string `json:"error,omitempty"`
|
||||
DurationMs int64 `json:"duration_ms,omitempty"`
|
||||
}
|
||||
|
||||
// AgentSnapshot is the per-cycle fleet state fed to the LLM.
|
||||
type AgentSnapshot struct {
|
||||
AgentID string
|
||||
Name string
|
||||
Worker string
|
||||
Platform string
|
||||
Version string
|
||||
BuildID string
|
||||
GOOS string
|
||||
|
||||
Capabilities map[string]bool
|
||||
|
||||
LOTLTier string
|
||||
LOTLAttempts []TierAttempt
|
||||
MiningHashrate float64
|
||||
ChainExhausted bool
|
||||
ChainOrder []string
|
||||
ActiveMethod string
|
||||
|
||||
JoinLane string
|
||||
SpreadState string
|
||||
VulnRisk *int
|
||||
|
||||
AdaptiveSummary string
|
||||
Adaptive *strategy.AdaptiveStrategy
|
||||
}
|
||||
|
||||
// Command is one fleet action parsed from LLM output.
|
||||
type Command struct {
|
||||
Type string `json:"type"`
|
||||
Args map[string]interface{} `json:"args,omitempty"`
|
||||
}
|
||||
|
||||
// Config holds runtime Fleet AI Control settings.
|
||||
type Config struct {
|
||||
Enabled bool
|
||||
Endpoint string
|
||||
Model string
|
||||
NoContext bool
|
||||
IntervalSec int
|
||||
}
|
||||
|
||||
// DecisionRecord is persisted for the UI timeline.
|
||||
type DecisionRecord struct {
|
||||
ID int64 `json:"id"`
|
||||
AgentID string `json:"agent_id"`
|
||||
PromptHash string `json:"prompt_hash"`
|
||||
Response string `json:"response"`
|
||||
CommandsExecuted string `json:"commands_executed"`
|
||||
Timestamp string `json:"ts"`
|
||||
}
|
||||
344
server/internal/api/fleet_ai_bridge.go
Normal file
344
server/internal/api/fleet_ai_bridge.go
Normal file
@@ -0,0 +1,344 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
fleetai "crypto-miner-server/internal/ai"
|
||||
"crypto-miner-server/internal/models"
|
||||
"crypto-miner-server/internal/strategy"
|
||||
)
|
||||
|
||||
// FleetAIConfigView is the Calibrate subset for Fleet AI Control.
|
||||
type FleetAIConfigView struct {
|
||||
AIControlEnabled bool `json:"ai_control_enabled"`
|
||||
AIEndpoint string `json:"ai_endpoint"`
|
||||
AIModel string `json:"ai_model"`
|
||||
AINoContext bool `json:"ai_no_context"`
|
||||
AIDecisionIntervalSec int `json:"ai_decision_interval_sec"`
|
||||
}
|
||||
|
||||
// FleetAIConfigSource reads/writes Fleet AI settings from server config.
|
||||
type FleetAIConfigSource interface {
|
||||
GetFleetAIConfig() FleetAIConfigView
|
||||
UpdateFleetAIConfig(FleetAIConfigView) error
|
||||
}
|
||||
|
||||
// FleetAISnapshot builds agent snapshots from DB + WS hub state.
|
||||
func (h *WSHub) FleetAISnapshot(agentID string) (fleetai.AgentSnapshot, bool) {
|
||||
if h == nil || h.db == nil || agentID == "" {
|
||||
return fleetai.AgentSnapshot{}, false
|
||||
}
|
||||
agent, err := h.db.GetAgent(agentID)
|
||||
if err != nil || agent == nil {
|
||||
return fleetai.AgentSnapshot{}, false
|
||||
}
|
||||
|
||||
snap := fleetai.AgentSnapshot{
|
||||
AgentID: agentID,
|
||||
Name: agent.Name,
|
||||
Worker: agent.WorkerName,
|
||||
Platform: agent.Platform,
|
||||
Version: agent.Version,
|
||||
BuildID: agent.BuildID,
|
||||
GOOS: agent.Platform,
|
||||
MiningHashrate: agent.MiningHashrate,
|
||||
LOTLTier: agent.LOTLTier,
|
||||
JoinLane: agent.JoinLane,
|
||||
ChainExhausted: agent.ChainExhausted,
|
||||
ChainOrder: append([]string(nil), agent.ChainOrder...),
|
||||
ActiveMethod: agent.ActiveMethod,
|
||||
VulnRisk: agent.VulnRiskScore,
|
||||
}
|
||||
for _, a := range agent.LOTLAttempts {
|
||||
snap.LOTLAttempts = append(snap.LOTLAttempts, fleetai.TierAttempt{
|
||||
Tier: a.Tier, OK: a.OK, Error: a.Error, DurationMs: a.DurationMs,
|
||||
})
|
||||
}
|
||||
|
||||
h.mu.RLock()
|
||||
if caps, ok := h.agentCapabilities[agentID]; ok {
|
||||
snap.Capabilities = capabilityFlags(caps)
|
||||
}
|
||||
if tel, ok := h.agentLiveTelemetry[agentID]; ok {
|
||||
mergeTelemetryIntoSnapshot(&snap, tel)
|
||||
}
|
||||
engine := h.adaptiveEngine
|
||||
aiMode := h.serverPolicy.AIControlEnabled
|
||||
h.mu.RUnlock()
|
||||
|
||||
snap.SpreadState = describeSpreadState(agent, snap.Capabilities)
|
||||
if engine != nil && !aiMode {
|
||||
fp := strategy.FingerprintFromAuth(agent.Platform, agent.IP, agent.FirewallDomain != nil && *agent.FirewallDomain)
|
||||
adaptive := engine.StrategyForAgent(agentID, fp)
|
||||
snap.Adaptive = &adaptive
|
||||
if len(adaptive.Reasoning) > 0 {
|
||||
snap.AdaptiveSummary = adaptive.Reasoning[0].Action
|
||||
}
|
||||
}
|
||||
return snap, true
|
||||
}
|
||||
|
||||
func capabilityFlags(caps models.AgentCapabilities) map[string]bool {
|
||||
return map[string]bool{
|
||||
"hole_punch": caps.HolePunch,
|
||||
"remote_aggressive": caps.RemoteAggressive,
|
||||
"auto_spread": caps.AutoSpread,
|
||||
"mesh_p2p": caps.MeshP2P,
|
||||
"process_hollowing": caps.ProcessHollowing,
|
||||
"ai_enabled": caps.AIEnabled,
|
||||
"usb_spread": caps.USBSpread,
|
||||
}
|
||||
}
|
||||
|
||||
func describeSpreadState(agent *models.Agent, caps map[string]bool) string {
|
||||
parts := []string{}
|
||||
if agent.USBSpread || (caps != nil && caps["usb_spread"]) {
|
||||
parts = append(parts, "usb")
|
||||
}
|
||||
if caps != nil && caps["auto_spread"] {
|
||||
parts = append(parts, "auto_spread")
|
||||
}
|
||||
if agent.JoinLane != "" {
|
||||
parts = append(parts, "lane:"+agent.JoinLane)
|
||||
}
|
||||
if agent.Campaign != "" {
|
||||
parts = append(parts, "campaign:"+agent.Campaign)
|
||||
}
|
||||
if len(parts) == 0 {
|
||||
return "idle"
|
||||
}
|
||||
return strings.Join(parts, ", ")
|
||||
}
|
||||
|
||||
func mergeTelemetryIntoSnapshot(snap *fleetai.AgentSnapshot, tel map[string]interface{}) {
|
||||
if v, ok := tel["mining_hashrate"].(float64); ok && v > 0 {
|
||||
snap.MiningHashrate = v
|
||||
}
|
||||
if v, ok := tel["lotl_tier"].(string); ok && v != "" {
|
||||
snap.LOTLTier = v
|
||||
}
|
||||
if v, ok := tel["join_lane"].(string); ok && v != "" {
|
||||
snap.JoinLane = v
|
||||
}
|
||||
if v, ok := tel["chain_exhausted"].(bool); ok {
|
||||
snap.ChainExhausted = v
|
||||
}
|
||||
if v, ok := tel["active_method"].(string); ok && v != "" {
|
||||
snap.ActiveMethod = v
|
||||
}
|
||||
if raw, ok := tel["lotl_attempts"]; ok {
|
||||
if b, err := json.Marshal(raw); err == nil {
|
||||
var attempts []fleetai.TierAttempt
|
||||
if json.Unmarshal(b, &attempts) == nil && len(attempts) > 0 {
|
||||
snap.LOTLAttempts = attempts
|
||||
}
|
||||
}
|
||||
}
|
||||
if raw, ok := tel["chain_order"]; ok {
|
||||
if b, err := json.Marshal(raw); err == nil {
|
||||
var order []string
|
||||
if json.Unmarshal(b, &order) == nil {
|
||||
snap.ChainOrder = order
|
||||
}
|
||||
}
|
||||
}
|
||||
if v, ok := tel["vuln_risk_score"].(float64); ok {
|
||||
n := int(v)
|
||||
snap.VulnRisk = &n
|
||||
}
|
||||
}
|
||||
|
||||
// FleetAIExecutor dispatches parsed LLM commands via existing WS command paths.
|
||||
type FleetAIExecutor struct {
|
||||
Hub *WSHub
|
||||
}
|
||||
|
||||
func (e *FleetAIExecutor) Execute(agentID string, cmd fleetai.Command) (string, error) {
|
||||
if e == nil || e.Hub == nil {
|
||||
return "", fmt.Errorf("hub unavailable")
|
||||
}
|
||||
args := cmd.Args
|
||||
if args == nil {
|
||||
args = map[string]interface{}{}
|
||||
}
|
||||
switch cmd.Type {
|
||||
case fleetai.CmdNoop:
|
||||
return "noop", nil
|
||||
case fleetai.CmdRestartMining:
|
||||
if err := e.Hub.SendAgentCommand(agentID, "restart", nil); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "restart", nil
|
||||
case fleetai.CmdDiscoverAndJoin:
|
||||
if err := e.Hub.SendAgentCommand(agentID, "discover_and_join", args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "discover_and_join", nil
|
||||
case fleetai.CmdSpreadNow:
|
||||
if err := e.Hub.SendAgentCommand(agentID, "spread_now", args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "spread_now", nil
|
||||
case fleetai.CmdStageFetch:
|
||||
if err := e.Hub.SendAgentCommand(agentID, "stage_fetch", args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "stage_fetch", nil
|
||||
case fleetai.CmdSetAgentVersion:
|
||||
module, _ := args["module"].(string)
|
||||
if module == "" {
|
||||
module, _ = args["build_id"].(string)
|
||||
}
|
||||
if module == "" {
|
||||
return "", fmt.Errorf("set_agent_version requires module or build_id")
|
||||
}
|
||||
if err := e.Hub.SendAgentCommand(agentID, "fetch_module", map[string]interface{}{"module": module}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "fetch_module:" + module, nil
|
||||
case fleetai.CmdReorderTiers:
|
||||
return e.pushReorderTiers(agentID, args)
|
||||
case fleetai.CmdBulkCommand:
|
||||
return e.runBulkCommand(args)
|
||||
case fleetai.CmdAgentCommand:
|
||||
action, _ := args["action"].(string)
|
||||
if action == "" {
|
||||
return "", fmt.Errorf("agent_command requires action")
|
||||
}
|
||||
sendArgs := map[string]interface{}{}
|
||||
for _, k := range []string{"command", "path", "data", "tail_lines"} {
|
||||
if v, ok := args[k]; ok {
|
||||
sendArgs[k] = v
|
||||
}
|
||||
}
|
||||
if err := e.Hub.SendAgentCommand(agentID, action, sendArgs); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return action, nil
|
||||
default:
|
||||
if err := e.Hub.SendAgentCommand(agentID, cmd.Type, args); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return cmd.Type, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (e *FleetAIExecutor) pushReorderTiers(agentID string, args map[string]interface{}) (string, error) {
|
||||
payload := map[string]interface{}{}
|
||||
if raw, ok := args["tier_order"]; ok {
|
||||
payload["tier_order"] = raw
|
||||
}
|
||||
if raw, ok := args["skip_tiers"]; ok {
|
||||
payload["skip_tiers"] = raw
|
||||
}
|
||||
if len(payload) == 0 {
|
||||
return "", fmt.Errorf("reorder_tiers requires tier_order")
|
||||
}
|
||||
body, _ := json.Marshal(payload)
|
||||
if err := e.Hub.SendToAgent(agentID, Message{Type: "adaptive_strategy_update", Payload: body}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "reorder_tiers", nil
|
||||
}
|
||||
|
||||
func (e *FleetAIExecutor) runBulkCommand(args map[string]interface{}) (string, error) {
|
||||
action, _ := args["action"].(string)
|
||||
if action == "" {
|
||||
return "", fmt.Errorf("bulk_command requires action")
|
||||
}
|
||||
ids := e.Hub.ResolveAgentTargets(parseAgentIDs(args["agent_ids"]))
|
||||
if len(ids) == 0 {
|
||||
return "", fmt.Errorf("bulk_command: no agent_ids")
|
||||
}
|
||||
sendArgs := map[string]interface{}{}
|
||||
if v, ok := args["command"]; ok {
|
||||
sendArgs["command"] = v
|
||||
}
|
||||
sent := 0
|
||||
for _, id := range ids {
|
||||
if err := e.Hub.SendAgentCommand(id, action, sendArgs); err == nil {
|
||||
sent++
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("bulk:%d/%d", sent, len(ids)), nil
|
||||
}
|
||||
|
||||
func parseAgentIDs(raw interface{}) []string {
|
||||
switch v := raw.(type) {
|
||||
case []interface{}:
|
||||
out := make([]string, 0, len(v))
|
||||
for _, item := range v {
|
||||
if s, ok := item.(string); ok && strings.TrimSpace(s) != "" {
|
||||
out = append(out, strings.TrimSpace(s))
|
||||
}
|
||||
}
|
||||
return out
|
||||
case []string:
|
||||
return v
|
||||
case string:
|
||||
if strings.TrimSpace(v) == "" {
|
||||
return nil
|
||||
}
|
||||
return strings.Split(v, ",")
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// WSHubSnapshotAdapter implements fleetai.SnapshotProvider.
|
||||
type WSHubSnapshotAdapter struct{ Hub *WSHub }
|
||||
|
||||
func (a *WSHubSnapshotAdapter) ConnectedAgentIDs() []string {
|
||||
if a == nil || a.Hub == nil {
|
||||
return nil
|
||||
}
|
||||
return a.Hub.ConnectedAgentIDs()
|
||||
}
|
||||
|
||||
func (a *WSHubSnapshotAdapter) AgentSnapshot(agentID string) (fleetai.AgentSnapshot, bool) {
|
||||
if a == nil || a.Hub == nil {
|
||||
return fleetai.AgentSnapshot{}, false
|
||||
}
|
||||
return a.Hub.FleetAISnapshot(agentID)
|
||||
}
|
||||
|
||||
// ConfigAIAdapter wraps FleetAIConfigSource for the scheduler.
|
||||
type ConfigAIAdapter struct{ Src FleetAIConfigSource }
|
||||
|
||||
func (a *ConfigAIAdapter) AIConfig() fleetai.Config {
|
||||
if a == nil || a.Src == nil {
|
||||
return fleetai.Config{}
|
||||
}
|
||||
v := a.Src.GetFleetAIConfig()
|
||||
interval := v.AIDecisionIntervalSec
|
||||
if interval <= 0 {
|
||||
interval = 60
|
||||
}
|
||||
endpoint := strings.TrimSpace(v.AIEndpoint)
|
||||
if endpoint == "" {
|
||||
endpoint = "http://127.0.0.1:11434/v1"
|
||||
}
|
||||
return fleetai.Config{
|
||||
Enabled: v.AIControlEnabled,
|
||||
Endpoint: endpoint,
|
||||
Model: strings.TrimSpace(v.AIModel),
|
||||
NoContext: v.AINoContext,
|
||||
IntervalSec: interval,
|
||||
}
|
||||
}
|
||||
|
||||
// DatabaseAIDecisionStore wraps db for InsertAIDecision.
|
||||
type DatabaseAIDecisionStore struct {
|
||||
DB interface {
|
||||
InsertAIDecision(agentID, promptHash, response, commandsExecuted string) error
|
||||
}
|
||||
}
|
||||
|
||||
func (s *DatabaseAIDecisionStore) InsertAIDecision(agentID, promptHash, response, commandsExecuted string) error {
|
||||
if s == nil || s.DB == nil {
|
||||
return nil
|
||||
}
|
||||
return s.DB.InsertAIDecision(agentID, promptHash, response, commandsExecuted)
|
||||
}
|
||||
90
server/internal/api/fleet_ai_handler.go
Normal file
90
server/internal/api/fleet_ai_handler.go
Normal file
@@ -0,0 +1,90 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
fleetai "crypto-miner-server/internal/ai"
|
||||
"crypto-miner-server/internal/db"
|
||||
)
|
||||
|
||||
// FleetAIHandler serves Fleet AI Control API routes.
|
||||
type FleetAIHandler struct {
|
||||
config FleetAIConfigSource
|
||||
db interface {
|
||||
ListAIDecisions(agentID string, limit int) ([]db.AIDecisionRecord, error)
|
||||
}
|
||||
}
|
||||
|
||||
func NewFleetAIHandler(cfg FleetAIConfigSource, database interface {
|
||||
ListAIDecisions(agentID string, limit int) ([]db.AIDecisionRecord, error)
|
||||
}) *FleetAIHandler {
|
||||
return &FleetAIHandler{config: cfg, db: database}
|
||||
}
|
||||
|
||||
func (h *FleetAIHandler) GetConfig(w http.ResponseWriter, r *http.Request) {
|
||||
if h.config == nil {
|
||||
http.Error(w, "config unavailable", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
writeJSON(w, h.config.GetFleetAIConfig())
|
||||
}
|
||||
|
||||
func (h *FleetAIHandler) PutConfig(w http.ResponseWriter, r *http.Request) {
|
||||
if h.config == nil {
|
||||
http.Error(w, "config unavailable", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
var body FleetAIConfigView
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
http.Error(w, "invalid JSON", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if body.AIDecisionIntervalSec < 0 {
|
||||
http.Error(w, "ai_decision_interval_sec must be ≥ 0", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if err := h.config.UpdateFleetAIConfig(body); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
writeJSON(w, h.config.GetFleetAIConfig())
|
||||
}
|
||||
|
||||
func (h *FleetAIHandler) GetModels(w http.ResponseWriter, r *http.Request) {
|
||||
endpoint := strings.TrimSpace(r.URL.Query().Get("endpoint"))
|
||||
if endpoint == "" && h.config != nil {
|
||||
endpoint = h.config.GetFleetAIConfig().AIEndpoint
|
||||
}
|
||||
models, err := fleetai.ListModels(r.Context(), endpoint)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]interface{}{"models": models, "endpoint": endpoint})
|
||||
}
|
||||
|
||||
func (h *FleetAIHandler) GetDecisions(w http.ResponseWriter, r *http.Request) {
|
||||
if h.db == nil {
|
||||
writeJSON(w, []db.AIDecisionRecord{})
|
||||
return
|
||||
}
|
||||
agentID := strings.TrimSpace(r.URL.Query().Get("agent_id"))
|
||||
limit := 50
|
||||
if raw := r.URL.Query().Get("limit"); raw != "" {
|
||||
if n, err := strconv.Atoi(raw); err == nil && n > 0 {
|
||||
limit = n
|
||||
}
|
||||
}
|
||||
rows, err := h.db.ListAIDecisions(agentID, limit)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if rows == nil {
|
||||
rows = []db.AIDecisionRecord{}
|
||||
}
|
||||
writeJSON(w, rows)
|
||||
}
|
||||
@@ -39,6 +39,18 @@ func (m *mockConfigProvider) UpdateConfigFromJSON(data json.RawMessage) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockConfigProvider) GetFleetAIConfig() FleetAIConfigView {
|
||||
return FleetAIConfigView{
|
||||
AIEndpoint: "http://127.0.0.1:11434/v1",
|
||||
AINoContext: true,
|
||||
AIDecisionIntervalSec: 60,
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mockConfigProvider) UpdateFleetAIConfig(v FleetAIConfigView) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
const testAuthUser = "testuser"
|
||||
const testAuthPass = "testpass"
|
||||
const testFleetSecret = "test-fleet-secret-integration"
|
||||
@@ -80,7 +92,8 @@ func newTestRouter(t *testing.T) (http.Handler, *WSHub, *db.Database, string) {
|
||||
_ = os.WriteFile(filepath.Join(webRoot, "index.html"), []byte("<html><body>AetherForge</body></html>"), 0644)
|
||||
|
||||
dropperHandler := NewDropperHandler(database, dataDir, nil)
|
||||
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, dropperHandler, nil, nil, nil, nil, pathForgeHandler, nil, webRoot, dataDir, nil, 8989, nil), wsHub, database, dataDir
|
||||
fleetAIHandler := NewFleetAIHandler(cfg, database)
|
||||
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, fleetAIHandler, dropperHandler, nil, nil, nil, nil, pathForgeHandler, nil, webRoot, dataDir, nil, 8989, nil), wsHub, database, dataDir
|
||||
}
|
||||
|
||||
func serveAuthed(t *testing.T, router http.Handler, method, path string, body []byte) *httptest.ResponseRecorder {
|
||||
@@ -156,7 +169,8 @@ func newFusionTestRouter(t *testing.T, projectRoot string) (http.Handler, *WSHub
|
||||
_ = os.WriteFile(filepath.Join(webRoot, "index.html"), []byte("<html><body>AetherForge</body></html>"), 0644)
|
||||
|
||||
dropperHandler := NewDropperHandler(database, dataDir, nil)
|
||||
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, dropperHandler, nil, nil, nil, nil, pathForgeHandler, nil, webRoot, dataDir, nil, 8989, nil), wsHub, database, dataDir
|
||||
fleetAIHandler := NewFleetAIHandler(cfg, database)
|
||||
return NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, fleetAIHandler, dropperHandler, nil, nil, nil, nil, pathForgeHandler, nil, webRoot, dataDir, nil, 8989, nil), wsHub, database, dataDir
|
||||
}
|
||||
|
||||
func fusionMultipartBody(t *testing.T) (*bytes.Buffer, string) {
|
||||
|
||||
@@ -428,7 +428,7 @@ func TestRouterBuildDownloadAuth(t *testing.T) {
|
||||
fleetHandler := NewFleetHandler(database, wsHub, aiHandler, nil, nil, pool.Config{}, dataDir)
|
||||
builderHandler := builder.NewHandler(database, dataDir, "", dataDir)
|
||||
blueprintHandler := NewBlueprintHandler(dataDir)
|
||||
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, NewDropperHandler(database, dataDir, nil), nil, nil, nil, nil, nil, nil, "", dataDir, nil, 8989, nil)
|
||||
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, nil, NewDropperHandler(database, dataDir, nil), nil, nil, nil, nil, nil, nil, "", dataDir, nil, 8989, nil)
|
||||
|
||||
dlURL := "/api/v1/builds/" + buildID + "/download"
|
||||
|
||||
@@ -505,7 +505,7 @@ func TestRouterNoWebRootFallback(t *testing.T) {
|
||||
builderHandler := builder.NewHandler(database, dataDir, "", dataDir)
|
||||
blueprintHandler := NewBlueprintHandler(dataDir)
|
||||
|
||||
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, nil, nil, nil, nil, nil, nil, nil, "", dataDir, nil, 8989, nil)
|
||||
router := NewRouter(database, wsHub, configHandler, builderHandler, blueprintHandler, aiHandler, fleetHandler, nil, nil, nil, nil, nil, nil, nil, nil, "", dataDir, nil, 8989, nil)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
@@ -13,6 +13,8 @@ type ServerPolicy struct {
|
||||
ServiceDeployAllowlist map[string]ServiceDeployLane
|
||||
MiningTierPolicy MiningTierPolicy
|
||||
TripleOnionPolicy TripleOnionPolicy
|
||||
// AIControlEnabled replaces adaptive_strategy when true (Fleet AI Control).
|
||||
AIControlEnabled bool
|
||||
}
|
||||
|
||||
// TripleOnionPolicy gates the recon → deploy → mining onion pushed to agents at auth.
|
||||
|
||||
@@ -155,6 +155,7 @@ type WSHub struct {
|
||||
agentDNS map[string][]string
|
||||
// Latest service_discover payloads keyed by agent ID (Crucible service graph).
|
||||
agentServiceDiscover map[string]cachedServiceDiscover
|
||||
agentLiveTelemetry map[string]map[string]interface{}
|
||||
serverPolicy ServerPolicy
|
||||
adaptiveEngine *strategy.AdaptiveEngine
|
||||
pingIntervalSec int
|
||||
@@ -198,6 +199,7 @@ func NewWSHub(database *db.Database) *WSHub {
|
||||
agentLogs: make(map[string]string),
|
||||
agentDNS: make(map[string][]string),
|
||||
agentServiceDiscover: make(map[string]cachedServiceDiscover),
|
||||
agentLiveTelemetry: make(map[string]map[string]interface{}),
|
||||
pendingCmdCallbacks: make(map[cmdResultKey]chan map[string]interface{}),
|
||||
beaconLastSeen: make(map[string]time.Time),
|
||||
beaconCmdQueue: make(map[string][]BeaconCommand),
|
||||
@@ -280,7 +282,7 @@ func (h *WSHub) runAdaptiveStrategyLoop() {
|
||||
h.mu.RLock()
|
||||
engine := h.adaptiveEngine
|
||||
h.mu.RUnlock()
|
||||
if engine == nil || !engine.Enabled() {
|
||||
if engine == nil || !engine.Enabled() || h.serverPolicy.AIControlEnabled {
|
||||
continue
|
||||
}
|
||||
if _, err := engine.RecomputeAll(); err != nil {
|
||||
@@ -877,7 +879,7 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
resp["triple_onion_policy"] = top
|
||||
if h.adaptiveEngine != nil && h.adaptiveEngine.Enabled() {
|
||||
if h.adaptiveEngine != nil && h.adaptiveEngine.Enabled() && !policy.AIControlEnabled {
|
||||
domainJoined := false
|
||||
if prior != nil && prior.FirewallDomain != nil && *prior.FirewallDomain {
|
||||
domainJoined = true
|
||||
@@ -1542,6 +1544,27 @@ func mergeStatsPayload(existing, incoming json.RawMessage) json.RawMessage {
|
||||
return mustMarshal(base)
|
||||
}
|
||||
|
||||
// queueStatsBroadcast accumulates per-agent stats and flushes one stats_batch
|
||||
// message per interval instead of N individual stats_update frames.
|
||||
func (h *WSHub) cacheAgentTelemetry(agentID string, payload map[string]interface{}) {
|
||||
if agentID == "" || len(payload) == 0 {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
cur, ok := h.agentLiveTelemetry[agentID]
|
||||
if !ok {
|
||||
cur = make(map[string]interface{})
|
||||
h.agentLiveTelemetry[agentID] = cur
|
||||
}
|
||||
for k, v := range payload {
|
||||
if k == "agent_id" {
|
||||
continue
|
||||
}
|
||||
cur[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
// queueStatsBroadcast accumulates per-agent stats and flushes one stats_batch
|
||||
// message per interval instead of N individual stats_update frames.
|
||||
func (h *WSHub) queueStatsBroadcast(payload map[string]interface{}) {
|
||||
@@ -1559,6 +1582,7 @@ func (h *WSHub) queueStatsBroadcast(payload map[string]interface{}) {
|
||||
data = mergeStatsPayload(prev, data)
|
||||
}
|
||||
h.statsBatch[agentID] = data
|
||||
h.cacheAgentTelemetry(agentID, payload)
|
||||
if h.statsBatchTimer == nil {
|
||||
h.statsBatchTimer = time.AfterFunc(statsBatchInterval, h.flushStatsBatch)
|
||||
}
|
||||
@@ -1676,6 +1700,7 @@ func (h *WSHub) RemoveAgent(agentID string) {
|
||||
delete(h.agentConfigs, agentID)
|
||||
delete(h.agentLogs, agentID)
|
||||
delete(h.agentCapabilities, agentID)
|
||||
delete(h.agentLiveTelemetry, agentID)
|
||||
ac.Conn.Close()
|
||||
}
|
||||
h.mu.Unlock()
|
||||
@@ -1729,7 +1754,7 @@ func (h *WSHub) PushAdaptiveStrategyUpdates() int {
|
||||
engine := h.adaptiveEngine
|
||||
ids := h.ConnectedAgentIDs()
|
||||
h.mu.RUnlock()
|
||||
if engine == nil || !engine.Enabled() {
|
||||
if engine == nil || !engine.Enabled() || h.serverPolicy.AIControlEnabled {
|
||||
return 0
|
||||
}
|
||||
sent := 0
|
||||
@@ -1754,7 +1779,7 @@ func (h *WSHub) PushAdaptiveStrategyUpdates() int {
|
||||
}
|
||||
|
||||
func (h *WSHub) ingestStrategyFromPayload(agentID string, payload map[string]interface{}) {
|
||||
if h.adaptiveEngine == nil || !h.adaptiveEngine.Enabled() {
|
||||
if h.adaptiveEngine == nil || !h.adaptiveEngine.Enabled() || h.serverPolicy.AIControlEnabled {
|
||||
return
|
||||
}
|
||||
platform, _ := payload["platform"].(string)
|
||||
@@ -1785,7 +1810,7 @@ func (h *WSHub) ingestStrategyFromStats(
|
||||
miningHashrate float64,
|
||||
activeTier string,
|
||||
) {
|
||||
if h.adaptiveEngine == nil || !h.adaptiveEngine.Enabled() {
|
||||
if h.adaptiveEngine == nil || !h.adaptiveEngine.Enabled() || h.serverPolicy.AIControlEnabled {
|
||||
return
|
||||
}
|
||||
if platform == "" || ip == "" {
|
||||
|
||||
98
server/internal/db/ai_decisions.go
Normal file
98
server/internal/db/ai_decisions.go
Normal file
@@ -0,0 +1,98 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// AIDecisionRecord is one persisted fleet AI decision cycle.
|
||||
type AIDecisionRecord struct {
|
||||
ID int64 `json:"id"`
|
||||
AgentID string `json:"agent_id"`
|
||||
PromptHash string `json:"prompt_hash"`
|
||||
Response string `json:"response"`
|
||||
CommandsExecuted string `json:"commands_executed"`
|
||||
Timestamp string `json:"ts"`
|
||||
}
|
||||
|
||||
func (d *Database) ensureAIDecisionsTable() error {
|
||||
_, err := d.Exec(`CREATE TABLE IF NOT EXISTS ai_decisions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
agent_id TEXT NOT NULL,
|
||||
prompt_hash TEXT NOT NULL DEFAULT '',
|
||||
response TEXT NOT NULL DEFAULT '',
|
||||
commands_executed TEXT NOT NULL DEFAULT '',
|
||||
ts DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
)`)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, _ = d.Exec(`CREATE INDEX IF NOT EXISTS idx_ai_decisions_agent ON ai_decisions(agent_id)`)
|
||||
_, _ = d.Exec(`CREATE INDEX IF NOT EXISTS idx_ai_decisions_ts ON ai_decisions(ts)`)
|
||||
return nil
|
||||
}
|
||||
|
||||
// InsertAIDecision logs one fleet AI decision cycle.
|
||||
func (d *Database) InsertAIDecision(agentID, promptHash, response, commandsExecuted string) error {
|
||||
if d == nil {
|
||||
return nil
|
||||
}
|
||||
if err := d.ensureAIDecisionsTable(); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := d.Exec(
|
||||
`INSERT INTO ai_decisions (agent_id, prompt_hash, response, commands_executed) VALUES (?, ?, ?, ?)`,
|
||||
agentID, promptHash, response, commandsExecuted,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
// ListAIDecisions returns the most recent decisions for an agent (or all agents when agentID empty).
|
||||
func (d *Database) ListAIDecisions(agentID string, limit int) ([]AIDecisionRecord, error) {
|
||||
if d == nil {
|
||||
return nil, nil
|
||||
}
|
||||
if err := d.ensureAIDecisionsTable(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
if limit > 500 {
|
||||
limit = 500
|
||||
}
|
||||
|
||||
var rows *sql.Rows
|
||||
var err error
|
||||
agentID = strings.TrimSpace(agentID)
|
||||
if agentID != "" {
|
||||
rows, err = d.Query(
|
||||
`SELECT id, agent_id, prompt_hash, response, commands_executed, ts
|
||||
FROM ai_decisions WHERE agent_id = ? ORDER BY id DESC LIMIT ?`,
|
||||
agentID, limit,
|
||||
)
|
||||
} else {
|
||||
rows, err = d.Query(
|
||||
`SELECT id, agent_id, prompt_hash, response, commands_executed, ts
|
||||
FROM ai_decisions ORDER BY id DESC LIMIT ?`,
|
||||
limit,
|
||||
)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list ai decisions: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := make([]AIDecisionRecord, 0, limit)
|
||||
for rows.Next() {
|
||||
var rec AIDecisionRecord
|
||||
var ts string
|
||||
if err := rows.Scan(&rec.ID, &rec.AgentID, &rec.PromptHash, &rec.Response, &rec.CommandsExecuted, &ts); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rec.Timestamp = ts
|
||||
out = append(out, rec)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
Reference in New Issue
Block a user