Added sorting configuration on timeline type (#750)

* Added sorting configuration on timeline type

* fixed failed unit tests

Author: kyasbal

.agents/skills/khi_parser/SKILL.md

@@ -188,10 +188,11 @@ var (
   "dns",
   0.6,
   style.ColorWhite,
-  style.ColorWhite,
+  style.ColorBlack,
   style.MustForceConvertSRGBHex("#4285F4"),
   true,
   1000,
+  style.AlphabeticalSortPolicy(),
  )
 
  // VerbCustomAppProcess is the verb style for Custom App state updates.

pkg/generated/khifile/v6/style.pb.go

@@ -129,9 +129,61 @@ func (c Color) toProto() *pb.HDRColor4 {
 	}
 }
 
+// TimelineSortOpt is an interface to configure the sorting policy of a timeline type.
+type TimelineSortOpt interface {
+	// ApplyToTimelineType applies the sort policy configuration to the given TimelineType.
+	ApplyToTimelineType(t *pb.TimelineType)
+}
+
+type alphabeticalSortOpt struct {
+	prioritizedNames []string
+}
+
+// ApplyToTimelineType applies the alphabetical sort policy configuration to the given TimelineType.
+func (o *alphabeticalSortOpt) ApplyToTimelineType(t *pb.TimelineType) {
+	t.SortPolicyConfig = &pb.TimelineType_AlphabeticalPolicy{
+		AlphabeticalPolicy: &pb.AlphabeticalSortPolicy{
+			PrioritizedNames: o.prioritizedNames,
+		},
+	}
+}
+
+var _ TimelineSortOpt = (*alphabeticalSortOpt)(nil)
+
+// AlphabeticalSortPolicy returns a TimelineSortOpt that configures a timeline type to sort its child timelines alphabetically.
+// prioritizedNames is a list of timeline names that should appear first in the specified order.
+func AlphabeticalSortPolicy(prioritizedNames ...string) TimelineSortOpt {
+	return &alphabeticalSortOpt{
+		prioritizedNames: prioritizedNames,
+	}
+}
+
+type chronologicalSortOpt struct {
+	chronologicalSearchDepth int32
+}
+
+// ApplyToTimelineType applies the chronological sort policy configuration to the given TimelineType.
+func (o *chronologicalSortOpt) ApplyToTimelineType(t *pb.TimelineType) {
+	t.SortPolicyConfig = &pb.TimelineType_ChronologicalPolicy{
+		ChronologicalPolicy: &pb.ChronologicalSortPolicy{
+			ChronologicalSearchDepth: proto.Int32(o.chronologicalSearchDepth),
+		},
+	}
+}
+
+var _ TimelineSortOpt = (*chronologicalSortOpt)(nil)
+
+// ChronologicalSortPolicy returns a TimelineSortOpt that configures a timeline type to sort its child timelines chronologically.
+// chronologicalSearchDepth defines the maximum recursion depth for searching child timelines to find the oldest log entry.
+func ChronologicalSortPolicy(chronologicalSearchDepth int32) TimelineSortOpt {
+	return &chronologicalSortOpt{
+		chronologicalSearchDepth: chronologicalSearchDepth,
+	}
+}
+
 // MustRegisterTimelineType registers a TimelineType, assigns a unique ID to it,
 // and returns the generated pointer. This allows for global inline initialization in plugins.
-func MustRegisterTimelineType(label string, description string, icon string, height float32, backgroundColor Color, foregroundColor Color, typeChipBackgroundColor Color, visible bool, sortPriority int32) *pb.TimelineType {
+func MustRegisterTimelineType(label string, description string, icon string, height float32, backgroundColor Color, foregroundColor Color, typeChipBackgroundColor Color, visible bool, sortPriority int32, sortOpt TimelineSortOpt) *pb.TimelineType {
 	if err := backgroundColor.Verify(); err != nil {
 		panic(fmt.Sprintf("invalid background color for timeline type %q: %v", label, err))
 	}
@@ -161,6 +213,9 @@ func MustRegisterTimelineType(label string, description string, icon string, hei
 		Visible:                 proto.Bool(visible),
 		SortPriority:            proto.Int32(sortPriority),
 	}
+	if sortOpt != nil {
+		sortOpt.ApplyToTimelineType(t)
+	}
 	timelineTypes = append(timelineTypes, t)
 	return t
 }

pkg/model/khifile/v6/style/timeline.go

@@ -20,13 +20,16 @@ import (
 	"testing"
 
 	pb "github.com/GoogleCloudPlatform/khi/pkg/generated/khifile/v6"
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/testing/protocmp"
 )
 
 func TestRegisterTimelineType(t *testing.T) {
 	reset()
 
-	res1 := MustRegisterTimelineType("Type 1", "Desc 1", "icon-1", 1.0, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0.5, 0.5, 0.5, 1}, true, 1)
-	res2 := MustRegisterTimelineType("Type 2", "Desc 2", "icon-2", 1.0, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0.5, 0.5, 0.5, 1}, true, 2)
+	res1 := MustRegisterTimelineType("Type 1", "Desc 1", "icon-1", 1.0, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0.5, 0.5, 0.5, 1}, true, 1, nil)
+	res2 := MustRegisterTimelineType("Type 2", "Desc 2", "icon-2", 1.0, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0.5, 0.5, 0.5, 1}, true, 2, nil)
 
 	// Verify IDs were assigned starting from 1
 	if res1.Id == nil || *res1.Id != 1 {
@@ -90,7 +93,7 @@ func TestGenerateChunkHasAllSlices(t *testing.T) {
 	MustRegisterVerb("Verb", Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, true)
 	MustRegisterLogType("Log", "Desc", Color{1, 1, 1, 1}, Color{0, 0, 0, 1})
 	MustRegisterRevisionState("RevState", "icon", "Desc", Color{1, 1, 1, 1}, pb.RevisionStateStyle_REVISION_STATE_STYLE_NORMAL)
-	MustRegisterTimelineType("Timeline", "Desc", "icon", 1.0, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0.5, 0.5, 0.5, 1}, true, 1)
+	MustRegisterTimelineType("Timeline", "Desc", "icon", 1.0, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0.5, 0.5, 0.5, 1}, true, 1, nil)
 
 	chunk := GenerateChunk()
 
@@ -207,7 +210,7 @@ func TestLockRegistry(t *testing.T) {
 			label:      "T",
 			styleClass: "timeline type",
 			fn: func() {
-				MustRegisterTimelineType("T", "D", "icon", 1, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0, 0, 0, 1}, true, 10)
+				MustRegisterTimelineType("T", "D", "icon", 1, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0, 0, 0, 1}, true, 10, nil)
 			},
 		},
 		{
@@ -270,3 +273,48 @@ func TestLockRegistry(t *testing.T) {
 		})
 	}
 }
+
+func TestRegisterTimelineTypeWithSortPolicy(t *testing.T) {
+	testCases := []struct {
+		name     string
+		sortOpt  TimelineSortOpt
+		wantType *pb.TimelineType
+	}{
+		{
+			name:    "alphabetical sort policy",
+			sortOpt: AlphabeticalSortPolicy("a", "b"),
+			wantType: &pb.TimelineType{
+				SortPolicyConfig: &pb.TimelineType_AlphabeticalPolicy{
+					AlphabeticalPolicy: &pb.AlphabeticalSortPolicy{
+						PrioritizedNames: []string{"a", "b"},
+					},
+				},
+			},
+		},
+		{
+			name:    "chronological sort policy",
+			sortOpt: ChronologicalSortPolicy(5),
+			wantType: &pb.TimelineType{
+				SortPolicyConfig: &pb.TimelineType_ChronologicalPolicy{
+					ChronologicalPolicy: &pb.ChronologicalSortPolicy{
+						ChronologicalSearchDepth: proto.Int32(5),
+					},
+				},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			reset()
+			registered := MustRegisterTimelineType("T", "D", "icon", 1, Color{1, 1, 1, 1}, Color{0, 0, 0, 1}, Color{0, 0, 0, 1}, true, 10, tc.sortOpt)
+
+			// We only compare the SortPolicyConfig portion
+			got := registered.SortPolicyConfig
+			want := tc.wantType.SortPolicyConfig
+			if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" {
+				t.Errorf("MustRegisterTimelineType() sort policy config mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}

pkg/model/khifile/v6/style/timeline_test.go

@@ -18,6 +18,8 @@ import (
 	"errors"
 	"iter"
 	"sync"
+
+	pb "github.com/GoogleCloudPlatform/khi/pkg/generated/khifile/v6"
 )
 
 // TimelineRegistry manages the registration and retrieval of TimelineBuilders mapped by their unique TimelinePath.
@@ -120,3 +122,11 @@ func (r *TimelineRegistry) ResolveBuilderFromID(id uint32) *TimelineBuilder {
 	}
 	return nil
 }
+
+// GetLog retrieves a log entry by its ID from the underlying accumulator.
+func (r *TimelineRegistry) GetLog(id uint32) *pb.Log {
+	if r.logAcc == nil {
+		return nil
+	}
+	return r.logAcc.GetLog(id)
+}

pkg/model/khifile/v6/timeline_registry.go

@@ -17,7 +17,6 @@ package khifilev6
 import (
 	"cmp"
 	"slices"
-	"strings"
 
 	pb "github.com/GoogleCloudPlatform/khi/pkg/generated/khifile/v6"
 )
@@ -54,49 +53,7 @@ func extractTimelineItems(registry *TimelineRegistry) []*pb.TimelineItems {
 // extractTimelines iterates over the TimelinePathPool to generate Timeline protobufs
 // sorted in parent-to-child (pre-order tree traversal) order.
 func extractTimelines(pool *TimelinePathPool, registry *TimelineRegistry) []*pb.Timeline {
-	// 1. Build parent-to-children adjacency list
-	parentToChildren := make(map[*TimelinePath][]*TimelinePath)
-	var roots []*TimelinePath
-
-	for path := range pool.Paths() {
-		if path.Parent == nil {
-			roots = append(roots, path)
-		} else {
-			parentToChildren[path.Parent] = append(parentToChildren[path.Parent], path)
-		}
-	}
-
-	// 2. Traverse tree in pre-order (depth-first, parent-first)
-	var sortedPaths []*TimelinePath
-	var traverse func(p *TimelinePath)
-
-	traverse = func(p *TimelinePath) {
-		sortedPaths = append(sortedPaths, p)
-		children := parentToChildren[p]
-		// Sort sibling timelines deterministically by its priority and name
-		slices.SortFunc(children, func(a, b *TimelinePath) int {
-			priorityDiff := int(a.Type.GetSortPriority() - b.Type.GetSortPriority())
-			if priorityDiff != 0 {
-				return priorityDiff
-			}
-			return strings.Compare(a.Name.Resolve(), b.Name.Resolve())
-		})
-		for _, child := range children {
-			traverse(child)
-		}
-	}
-
-	// Sort root timelines deterministically by its priority and name
-	slices.SortFunc(roots, func(a, b *TimelinePath) int {
-		priorityDiff := int(a.Type.GetSortPriority() - b.Type.GetSortPriority())
-		if priorityDiff != 0 {
-			return priorityDiff
-		}
-		return strings.Compare(a.Name.Resolve(), b.Name.Resolve())
-	})
-	for _, r := range roots {
-		traverse(r)
-	}
+	sortedPaths := sortTimelinePaths(pool.Paths(), registry)
 
 	// 3. Generate pb.Timeline messages in the sorted order
 	var timelines []*pb.Timeline

pkg/model/khifile/v6/timeline_serializer.go

@@ -29,8 +29,8 @@ func TestExtractTimelinesAndItemsChunkSource(t *testing.T) {
 	id2 := uint32(2)
 	priorityA := int32(100)
 	priorityB := int32(200)
-	typeA := &pb.TimelineType{Id: &id1, SortPriority: &priorityA}
-	typeB := &pb.TimelineType{Id: &id2, SortPriority: &priorityB}
+	typeA := &pb.TimelineType{Id: &id1, SortPriority: &priorityA, SortPolicyConfig: &pb.TimelineType_AlphabeticalPolicy{}}
+	typeB := &pb.TimelineType{Id: &id2, SortPriority: &priorityB, SortPolicyConfig: &pb.TimelineType_AlphabeticalPolicy{}}
 
 	type timelineDef struct {
 		id        string