Add Peering Support (#296)

Makefile

@@ -24,10 +24,11 @@ proxy/ui_dist/placeholder.txt:
 	touch $@
 
 test: proxy/ui_dist/placeholder.txt
-	go test -short -v -count=1 ./proxy
+	go test -short ./proxy/...
 
+# for CI - full test (takes longer)
 test-all: proxy/ui_dist/placeholder.txt
-	go test -v -count=1 ./proxy
+	go test -count=1 ./proxy/...
 
 ui/node_modules:
 	cd ui && npm install
@@ -81,4 +82,4 @@ release:
 	git tag "$$new_tag";
 
 # Phony targets
-.PHONY: all clean ui mac linux windows simple-responder
+.PHONY: all clean ui mac linux windows simple-responder test test-all

llama-swap.go

@@ -16,6 +16,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/mostlygeek/llama-swap/event"
 	"github.com/mostlygeek/llama-swap/proxy"
+	"github.com/mostlygeek/llama-swap/proxy/config"
 )
 
 var (
@@ -38,13 +39,13 @@ func main() {
 		os.Exit(0)
 	}
 
-	config, err := proxy.LoadConfig(*configPath)
+	conf, err := config.LoadConfig(*configPath)
 	if err != nil {
 		fmt.Printf("Error loading config: %v\n", err)
 		os.Exit(1)
 	}
 
-	if len(config.Profiles) > 0 {
+	if len(conf.Profiles) > 0 {
 		fmt.Println("WARNING: Profile functionality has been removed in favor of Groups. See the README for more information.")
 	}
 
@@ -67,15 +68,15 @@ func main() {
 	// Support for watching config and reloading when it changes
 	reloadProxyManager := func() {
 		if currentPM, ok := srv.Handler.(*proxy.ProxyManager); ok {
-			config, err = proxy.LoadConfig(*configPath)
+			conf, err = config.LoadConfig(*configPath)
 			if err != nil {
 				fmt.Printf("Warning, unable to reload configuration: %v\n", err)
 				return
 			}
 
 			fmt.Println("Configuration Changed")
 			currentPM.Shutdown()
-			srv.Handler = proxy.New(config)
+			srv.Handler = proxy.New(conf)
 			fmt.Println("Configuration Reloaded")
 
 			// wait a few seconds and tell any UI to reload
@@ -85,12 +86,12 @@ func main() {
 				})
 			})
 		} else {
-			config, err = proxy.LoadConfig(*configPath)
+			conf, err = config.LoadConfig(*configPath)
 			if err != nil {
 				fmt.Printf("Error, unable to load configuration: %v\n", err)
 				os.Exit(1)
 			}
-			srv.Handler = proxy.New(config)
+			srv.Handler = proxy.New(conf)
 		}
 	}
 

proxy/config.go → proxy/config/config.go

@@ -1,4 +1,4 @@
-package proxy
+package config
 
 import (
 	"fmt"
@@ -154,6 +154,7 @@ type Config struct {
 	Models             map[string]ModelConfig `yaml:"models"` /* key is model ID */
 	Profiles           map[string][]string    `yaml:"profiles"`
 	Groups             map[string]GroupConfig `yaml:"groups"` /* key is group ID */
+	Peers              map[string]PeerConfig  `yaml:"peers"`  /* key is peer ID */
 
 	// for key/value replacements in model's cmd, cmdStop, proxy, checkEndPoint
 	Macros map[string]string `yaml:"macros"`

proxy/config_posix_test.go → proxy/config/config_posix_test.go

@@ -1,10 +1,11 @@
 //go:build !windows
 
-package proxy
+package config
 
 import (
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"testing"
 
@@ -148,6 +149,14 @@ groups:
     persistent: true
     members:
       - "model4"
+peers:
+  desktop:
+    name: "Desktop"
+    description: "runs Linux"
+    baseURL: "http://10.0.4.11:8080"
+    apikey: "secret-key"
+    priority: 10
+    filters: []
 `
 
 	if err := os.WriteFile(tempFile, []byte(content), 0644); err != nil {
@@ -232,6 +241,17 @@ groups:
 				Members:    []string{"model4"},
 			},
 		},
+		Peers: map[string]PeerConfig{
+			"desktop": {
+				Name:        "Desktop",
+				Description: "runs Linux",
+				BaseURL:     "http://10.0.4.11:8080",
+				ApiKey:      "secret-key",
+				Priority:    10,
+				Filters:     []string{},
+				reFilters:   []*regexp.Regexp{}, /* leave blank, test in peer_test.go */
+			},
+		},
 	}
 
 	assert.Equal(t, expected, config)

proxy/config_test.go → proxy/config/config_test.go

@@ -1,4 +1,4 @@
-package proxy
+package config
 
 import (
 	"slices"

proxy/config_windows_test.go → proxy/config/config_windows_test.go

@@ -1,6 +1,6 @@
 //go:build windows
 
-package proxy
+package config
 
 import (
 	"os"
@@ -221,6 +221,7 @@ groups:
 				Members:    []string{"model4"},
 			},
 		},
+		Peers: nil, // empty here, see config_posix_test.go
 	}
 
 	assert.Equal(t, expected, config)

proxy/config/peer.go

@@ -0,0 +1,46 @@
+package config
+
+import (
+	"fmt"
+	"regexp"
+)
+
+type PeerConfig struct {
+	Name        string           `yaml:"name"`
+	Description string           `yaml:"description"`
+	BaseURL     string           `yaml:"baseURL"`
+	ApiKey      string           `yaml:"apikey"`
+	Priority    int              `yaml:"priority"`
+	Filters     []string         `yaml:"filters"`
+	reFilters   []*regexp.Regexp `yaml:"-"`
+}
+
+// set default values for GroupConfig
+func (c *PeerConfig) UnmarshalYAML(unmarshal func(any) error) error {
+	type rawConfig PeerConfig
+	defaults := rawConfig{
+		Name:        "",
+		Description: "",
+		BaseURL:     "",
+		ApiKey:      "",
+		Priority:    0,
+		Filters:     []string{},
+		reFilters:   []*regexp.Regexp{},
+	}
+
+	if err := unmarshal(&defaults); err != nil {
+		return err
+	}
+
+	// compile regex filters and store compiled patterns in reFilters
+	for _, pat := range defaults.Filters {
+		r, err := regexp.Compile(pat)
+		if err != nil {
+			return fmt.Errorf("failed to compile peer filter %q: %w", pat, err)
+		}
+		defaults.reFilters = append(defaults.reFilters, r)
+	}
+
+	*c = PeerConfig(defaults)
+	return nil
+}

proxy/config/peer_test.go

@@ -0,0 +1,99 @@
+package config
+
+import (
+	"strings"
+	"testing"
+
+	"gopkg.in/yaml.v3"
+)
+
+// Tests that defaults are set when unmarshaling an empty/minimal YAML.
+func TestPeerConfig_Defaults(t *testing.T) {
+	var pc PeerConfig
+	data := `{}`
+
+	if err := yaml.Unmarshal([]byte(data), &pc); err != nil {
+		t.Fatalf("unexpected unmarshal error: %v", err)
+	}
+
+	if pc.Name != "" {
+		t.Errorf("Name expected %q, got %q", "", pc.Name)
+	}
+	if pc.Description != "" {
+		t.Errorf("Description expected %q, got %q", "", pc.Description)
+	}
+	if pc.BaseURL != "" {
+		t.Errorf("BaseURL expected %q, got %q", "", pc.BaseURL)
+	}
+	if pc.ApiKey != "" {
+		t.Errorf("ApiKey expected %q, got %q", "", pc.ApiKey)
+	}
+	if pc.Priority != 0 {
+		t.Errorf("Priority expected %d, got %d", 0, pc.Priority)
+	}
+	if len(pc.Filters) != 0 {
+		t.Errorf("Filters expected length %d, got %d", 0, len(pc.Filters))
+	}
+	if len(pc.reFilters) != 0 {
+		t.Errorf("reFilters expected length %d, got %d", 0, len(pc.reFilters))
+	}
+}
+
+// Tests that valid regex patterns in Filters are compiled into reFilters and work as expected.
+func TestPeerConfig_RegexCompileSuccess(t *testing.T) {
+	var pc PeerConfig
+	data := `
+filters:
+  - "^foo.*"
+  - "ba[rz]$"
+`
+
+	if err := yaml.Unmarshal([]byte(data), &pc); err != nil {
+		t.Fatalf("unexpected unmarshal error: %v", err)
+	}
+
+	if len(pc.Filters) != 2 {
+		t.Fatalf("expected Filters length 2, got %d", len(pc.Filters))
+	}
+	if len(pc.reFilters) != 2 {
+		t.Fatalf("expected reFilters length 2, got %d", len(pc.reFilters))
+	}
+
+	// first pattern ^foo.*
+	if !pc.reFilters[0].MatchString("foobar") {
+		t.Errorf("expected pattern %q to match %q", pc.Filters[0], "foobar")
+	}
+	if pc.reFilters[0].MatchString("barfoo") {
+		t.Errorf("expected pattern %q NOT to match %q", pc.Filters[0], "barfoo")
+	}
+
+	// second pattern ba[rz]$
+	if !pc.reFilters[1].MatchString("bar") {
+		t.Errorf("expected pattern %q to match %q", pc.Filters[1], "bar")
+	}
+	if !pc.reFilters[1].MatchString("baz") {
+		t.Errorf("expected pattern %q to match %q", pc.Filters[1], "baz")
+	}
+	if pc.reFilters[1].MatchString("bax") {
+		t.Errorf("expected pattern %q NOT to match %q", pc.Filters[1], "bax")
+	}
+}
+
+// Tests that an invalid regex produces an error during Unmarshal.
+func TestPeerConfig_RegexCompileFailure(t *testing.T) {
+	var pc PeerConfig
+	data := `
+filters:
+  - "("
+`
+
+	err := yaml.Unmarshal([]byte(data), &pc)
+	if err == nil {
+		t.Fatalf("expected error compiling invalid regex, got nil")
+	}
+	// Optionally ensure our error message path was used
+	if !strings.Contains(err.Error(), "failed to compile peer filter") {
+		t.Logf("warning: error did not contain expected text; full error: %v", err)
+		t.Fail()
+	}
+}

proxy/helpers_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/gin-gonic/gin"
+	"github.com/mostlygeek/llama-swap/proxy/config"
 	"gopkg.in/yaml.v3"
 )
 
@@ -65,18 +66,18 @@ func getTestPort() int {
 	return port
 }
 
-func getTestSimpleResponderConfig(expectedMessage string) ModelConfig {
+func getTestSimpleResponderConfig(expectedMessage string) config.ModelConfig {
 	return getTestSimpleResponderConfigPort(expectedMessage, getTestPort())
 }
 
-func getTestSimpleResponderConfigPort(expectedMessage string, port int) ModelConfig {
+func getTestSimpleResponderConfigPort(expectedMessage string, port int) config.ModelConfig {
 	// Create a YAML string with just the values we want to set
 	yamlStr := fmt.Sprintf(`
 cmd: '%s --port %d --silent --respond %s'
 proxy: "http://127.0.0.1:%d"
 `, simpleResponderPath, port, expectedMessage, port)
 
-	var cfg ModelConfig
+	var cfg config.ModelConfig
 	if err := yaml.Unmarshal([]byte(yamlStr), &cfg); err != nil {
 		panic(fmt.Sprintf("failed to unmarshal test config: %v in [%s]", err, yamlStr))
 	}

proxy/metrics_monitor.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/mostlygeek/llama-swap/event"
+	"github.com/mostlygeek/llama-swap/proxy/config"
 )
 
 // TokenMetrics represents parsed token statistics from llama-server logs
@@ -38,7 +39,7 @@ type MetricsMonitor struct {
 	nextID     int
 }
 
-func NewMetricsMonitor(config *Config) *MetricsMonitor {
+func NewMetricsMonitor(config *config.Config) *MetricsMonitor {
 	maxMetrics := config.MetricsMaxInMemory
 	if maxMetrics <= 0 {
 		maxMetrics = 1000 // Default fallback

proxy/process.go

@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"github.com/mostlygeek/llama-swap/event"
+	"github.com/mostlygeek/llama-swap/proxy/config"
 )
 
 type ProcessState string
@@ -39,7 +40,7 @@ const (
 
 type Process struct {
 	ID     string
-	config ModelConfig
+	config config.ModelConfig
 	cmd    *exec.Cmd
 
 	// PR #155 called to cancel the upstream process
@@ -74,7 +75,7 @@ type Process struct {
 	failedStartCount int
 }
 
-func NewProcess(ID string, healthCheckTimeout int, config ModelConfig, processLogger *LogMonitor, proxyLogger *LogMonitor) *Process {
+func NewProcess(ID string, healthCheckTimeout int, config config.ModelConfig, processLogger *LogMonitor, proxyLogger *LogMonitor) *Process {
 	concurrentLimit := 10
 	if config.ConcurrencyLimit > 0 {
 		concurrentLimit = config.ConcurrencyLimit
@@ -539,7 +540,7 @@ func (p *Process) cmdStopUpstreamProcess() error {
 
 	if p.config.CmdStop != "" {
 		// replace ${PID} with the pid of the process
-		stopArgs, err := SanitizeCommand(strings.ReplaceAll(p.config.CmdStop, "${PID}", fmt.Sprintf("%d", p.cmd.Process.Pid)))
+		stopArgs, err := config.SanitizeCommand(strings.ReplaceAll(p.config.CmdStop, "${PID}", fmt.Sprintf("%d", p.cmd.Process.Pid)))
 		if err != nil {
 			p.proxyLogger.Errorf("<%s> Failed to sanitize stop command: %v", p.ID, err)
 			return err