feat: QualityTracker.Snapshot/Restore + Router.QualityTracker() for cross-session persistence

This commit is contained in:
2026-04-05 23:40:19 +02:00
parent 20fb045cba
commit ca163dc8d4
3 changed files with 149 additions and 0 deletions

View File

@@ -0,0 +1,70 @@
package router
// QualitySnapshot is the serializable form of QualityTracker EMA data.
type QualitySnapshot struct {
Scores map[string]map[string]*EMAScore `json:"scores"` // ArmID -> TaskType.String() -> score
}
// Snapshot returns a serializable copy of current quality data.
func (qt *QualityTracker) Snapshot() QualitySnapshot {
qt.mu.RLock()
defer qt.mu.RUnlock()
snap := QualitySnapshot{
Scores: make(map[string]map[string]*EMAScore),
}
for armID, tasks := range qt.scores {
snap.Scores[string(armID)] = make(map[string]*EMAScore)
for taskType, score := range tasks {
snap.Scores[string(armID)][taskType.String()] = &EMAScore{
Value: score.Value,
Count: score.Count,
}
}
}
return snap
}
// Restore loads previously persisted quality data.
// Replaces all current scores.
func (qt *QualityTracker) Restore(snap QualitySnapshot) {
qt.mu.Lock()
defer qt.mu.Unlock()
qt.scores = make(map[ArmID]map[TaskType]*EMAScore)
for armID, tasks := range snap.Scores {
qt.scores[ArmID(armID)] = make(map[TaskType]*EMAScore)
for taskTypeStr, score := range tasks {
taskType, ok := parseTaskType(taskTypeStr)
if !ok {
continue // skip unrecognized task types
}
qt.scores[ArmID(armID)][taskType] = &EMAScore{
Value: score.Value,
Count: score.Count,
}
}
}
}
// parseTaskType maps a TaskType string representation back to its constant.
func parseTaskType(s string) (TaskType, bool) {
for _, tt := range allTaskTypes {
if tt.String() == s {
return tt, true
}
}
return 0, false
}
// allTaskTypes lists every known TaskType constant for reverse lookup.
var allTaskTypes = []TaskType{
TaskBoilerplate,
TaskGeneration,
TaskRefactor,
TaskReview,
TaskUnitTest,
TaskPlanning,
TaskOrchestration,
TaskSecurityReview,
TaskDebug,
TaskExplain,
}

View File

@@ -0,0 +1,74 @@
package router_test
import (
"encoding/json"
"testing"
"somegit.dev/Owlibou/gnoma/internal/router"
)
func TestQualityTracker_SnapshotRestore_RoundTrip(t *testing.T) {
qt := router.NewQualityTracker()
// Record some outcomes
qt.Record("anthropic/claude-3-5-sonnet", router.TaskGeneration, true)
qt.Record("anthropic/claude-3-5-sonnet", router.TaskGeneration, true)
qt.Record("anthropic/claude-3-5-sonnet", router.TaskGeneration, false)
qt.Record("ollama/gemma3", router.TaskBoilerplate, true)
snap := qt.Snapshot()
// Verify snapshot has the data
if len(snap.Scores) == 0 {
t.Fatal("snapshot scores should not be empty")
}
// Marshal and unmarshal to simulate disk persistence
data, err := json.Marshal(snap)
if err != nil {
t.Fatal(err)
}
var restored router.QualitySnapshot
if err := json.Unmarshal(data, &restored); err != nil {
t.Fatal(err)
}
// Restore into a fresh tracker
qt2 := router.NewQualityTracker()
qt2.Restore(restored)
// After restore, Quality() should return data (Count >= minObservations=3)
score, hasData := qt2.Quality("anthropic/claude-3-5-sonnet", router.TaskGeneration)
if !hasData {
t.Error("expected quality data after restore")
}
if score <= 0 {
t.Errorf("expected positive score, got %f", score)
}
}
func TestQualityTracker_Snapshot_Empty(t *testing.T) {
qt := router.NewQualityTracker()
snap := qt.Snapshot()
if snap.Scores == nil {
t.Error("scores map should be initialized (not nil)")
}
if len(snap.Scores) != 0 {
t.Errorf("expected empty scores, got %d", len(snap.Scores))
}
}
func TestQualityTracker_Restore_Replaces(t *testing.T) {
qt := router.NewQualityTracker()
qt.Record("arm-a", router.TaskDebug, true)
qt.Record("arm-a", router.TaskDebug, true)
qt.Record("arm-a", router.TaskDebug, true)
// Restore with different data — old data should be gone
empty := router.QualitySnapshot{Scores: make(map[string]map[string]*router.EMAScore)}
qt.Restore(empty)
_, hasData := qt.Quality("arm-a", router.TaskDebug)
if hasData {
t.Error("old data should be gone after restore with empty snapshot")
}
}

View File

@@ -172,6 +172,11 @@ func (r *Router) LookupArm(id ArmID) (*Arm, bool) {
return arm, ok
}
// QualityTracker returns the router's quality tracker for persistence.
func (r *Router) QualityTracker() *QualityTracker {
return r.quality
}
// Arms returns all registered arms.
func (r *Router) Arms() []*Arm {
r.mu.RLock()