add some scriptsmain

1 files changed, 18 insertions, 2 deletions

diff --git a/internal/controlplane/manager.go b/internal/controlplane/manager.go
index 42f5efe..cdafc36 100644
--- a/internal/controlplane/manager.go
+++ b/internal/controlplane/manager.go
@@ -10,6 +10,7 @@ import (
 
 	"aigw/internal/auth"
 	"aigw/internal/catalog"
+	"aigw/internal/domain"
 )
 
 const broadcastQueueSize = 128
@@ -22,6 +23,10 @@ type managerStore interface {
 	RedisEnabled() bool
 }
 
+type policyReplacer interface {
+	ReplacePolicies([]domain.LimitPolicy)
+}
+
 type Manager struct {
 	store          managerStore
 	catalog        *catalog.Catalog
@@ -32,15 +37,20 @@ type Manager struct {
 	redisConnected atomic.Bool
 	reloadMu       sync.Mutex
 	broadcasts     chan ChangeEvent
+	policyTarget   policyReplacer
 }
 
-func NewManager(store managerStore, modelCatalog *catalog.Catalog, authenticator *auth.StaticAuthenticator, logger *slog.Logger, pollInterval time.Duration) *Manager {
+func NewManager(store managerStore, modelCatalog *catalog.Catalog, authenticator *auth.StaticAuthenticator, logger *slog.Logger, pollInterval time.Duration, policyTargets ...policyReplacer) *Manager {
 	if pollInterval <= 0 {
 		pollInterval = 30 * time.Second
 	}
+	var policyTarget policyReplacer
+	if len(policyTargets) > 0 {
+		policyTarget = policyTargets[0]
+	}
 	return &Manager{
 		store: store, catalog: modelCatalog, authenticator: authenticator,
-		logger: logger, pollInterval: pollInterval, broadcasts: make(chan ChangeEvent, broadcastQueueSize),
+		logger: logger, pollInterval: pollInterval, broadcasts: make(chan ChangeEvent, broadcastQueueSize), policyTarget: policyTarget,
 	}
 }
 
@@ -53,6 +63,9 @@ func (m *Manager) Reload(ctx context.Context) (int64, error) {
 	}
 	m.catalog.Replace(snapshot.Models)
 	m.authenticator.ReplaceHashed(snapshot.APIKeys)
+	if m.policyTarget != nil {
+		m.policyTarget.ReplacePolicies(snapshot.Limits)
+	}
 	m.generation.Store(snapshot.Generation)
 	m.logger.Info("control_plane_reloaded", "generation", snapshot.Generation, "models", len(snapshot.Models), "api_keys", len(snapshot.APIKeys))
 	return snapshot.Generation, nil
@@ -195,7 +208,10 @@ func (m *Manager) runBroadcasts(ctx context.Context) {
 			err := m.store.PublishChange(publishContext, event)
 			cancel()
 			if err != nil {
+				m.redisConnected.Store(false)
 				m.logger.Warn("control_plane_publish_failed", "generation", event.Generation, "resource", event.Resource, "id", event.ID, "error", err)
+			} else {
+				m.redisConnected.Store(true)
 			}
 		}
 	}