From 830f818c58933ca40fde1084a4c72f331ca65dd3 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@localai.io>
Date: Mon, 1 Jun 2026 07:54:42 +0000
Subject: [PATCH] feat(p2p): gossip free VRAM per node and add testable online
 check

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
---
 core/p2p/p2p.go              |  8 +++++---
 core/schema/localai.go       | 20 +++++++++++++++++---
 core/schema/nodedata_test.go | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 6 deletions(-)
 create mode 100644 core/schema/nodedata_test.go

diff --git a/core/p2p/p2p.go b/core/p2p/p2p.go
index d03a9c50c..a108147bf 100644
--- a/core/p2p/p2p.go
+++ b/core/p2p/p2p.go
@@ -15,6 +15,7 @@ import (
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/mudler/LocalAI/core/schema"
 	"github.com/mudler/LocalAI/pkg/utils"
+	"github.com/mudler/LocalAI/pkg/xsysinfo"
 	"github.com/mudler/edgevpn/pkg/config"
 	"github.com/mudler/edgevpn/pkg/node"
 	"github.com/mudler/edgevpn/pkg/protocol"
@@ -348,9 +349,10 @@ func ExposeService(ctx context.Context, host, port, token, servicesID string) (*
 		func() {
 			updatedMap := map[string]any{}
 			updatedMap[name] = &schema.NodeData{
-				Name:     name,
-				LastSeen: time.Now(),
-				ID:       nodeID(name),
+				Name:          name,
+				LastSeen:      time.Now(),
+				ID:            nodeID(name),
+				AvailableVRAM: xsysinfo.GetGPUAggregateInfo().FreeVRAM,
 			}
 			ledger.Add(servicesID, updatedMap)
 		},
diff --git a/core/schema/localai.go b/core/schema/localai.go
index 8704f8ad8..9a2b1bc39 100644
--- a/core/schema/localai.go
+++ b/core/schema/localai.go
@@ -121,18 +121,32 @@ type StoresFindResponse struct {
 	Similarities []float32   `json:"similarities" yaml:"similarities"`
 }
 
+// NodeOnlineWindow is how long after its last announce a node still counts as
+// online. Nodes re-announce into the edgevpn ledger every 20s (see core/p2p
+// ExposeService), so 40s tolerates a single missed announce.
+const NodeOnlineWindow = 40 * time.Second
+
 type NodeData struct {
 	Name          string
 	ID            string
 	TunnelAddress string
 	ServiceID     string
 	LastSeen      time.Time
+	// AvailableVRAM is the node's free GPU VRAM in bytes at its last announce,
+	// gossiped so federation selection can prefer peers with more headroom.
+	// Zero for CPU-only nodes and for peers on an older version that does not
+	// publish it; the routing policy treats zero as the lowest VRAM tier.
+	AvailableVRAM uint64
 }
 
 func (d NodeData) IsOnline() bool {
-	now := time.Now()
-	// if the node was seen in the last 40 seconds, it's online
-	return now.Sub(d.LastSeen) < 40*time.Second
+	return d.IsOnlineAt(time.Now())
+}
+
+// IsOnlineAt reports whether the node counts as online relative to now. It is
+// split from IsOnline so selection logic can be exercised with a fixed clock.
+func (d NodeData) IsOnlineAt(now time.Time) bool {
+	return now.Sub(d.LastSeen) < NodeOnlineWindow
 }
 
 type P2PNodesResponse struct {
diff --git a/core/schema/nodedata_test.go b/core/schema/nodedata_test.go
new file mode 100644
index 000000000..61a21ed21
--- /dev/null
+++ b/core/schema/nodedata_test.go
@@ -0,0 +1,34 @@
+package schema_test
+
+import (
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/schema"
+)
+
+var _ = Describe("NodeData", func() {
+	ref := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+
+	It("is online within the online window", func() {
+		nd := schema.NodeData{LastSeen: ref.Add(-30 * time.Second)}
+		Expect(nd.IsOnlineAt(ref)).To(BeTrue())
+	})
+
+	It("is offline once the online window has elapsed", func() {
+		nd := schema.NodeData{LastSeen: ref.Add(-50 * time.Second)}
+		Expect(nd.IsOnlineAt(ref)).To(BeFalse())
+	})
+
+	It("treats exactly the window boundary as offline (strict less-than)", func() {
+		nd := schema.NodeData{LastSeen: ref.Add(-schema.NodeOnlineWindow)}
+		Expect(nd.IsOnlineAt(ref)).To(BeFalse())
+	})
+
+	It("carries AvailableVRAM in bytes", func() {
+		nd := schema.NodeData{AvailableVRAM: 8_000_000_000}
+		Expect(nd.AvailableVRAM).To(Equal(uint64(8_000_000_000)))
+	})
+})