From f1b581d67b8ad68958da5475a7ea10a811d458ea Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Mon, 1 Dec 2025 19:07:16 +0000
Subject: [PATCH] Refactor GetModelPartitions pagination logic

Co-authored-by: eric.okuma <eric.okuma@rilldata.com>
---
 runtime/server/controller.go                 | 15 +++++--
 runtime/server/controller_pagination_test.go | 44 ++++++++++++++++++++
 2 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 runtime/server/controller_pagination_test.go

diff --git a/runtime/server/controller.go b/runtime/server/controller.go
index d167b1f783f..0841d04044f 100644
--- a/runtime/server/controller.go
+++ b/runtime/server/controller.go
@@ -327,13 +327,14 @@ func (s *Server) GetModelPartitions(ctx context.Context, req *runtimev1.GetModel
 	defer release()
 
 	defaultPageSize := 100
+	pageLimit := pagination.ValidPageSize(req.PageSize, defaultPageSize)
 	opts := &drivers.FindModelPartitionsOptions{
 		ModelID:          partitionsModelID,
 		WherePending:     req.Pending,
 		WhereErrored:     req.Errored,
 		BeforeExecutedOn: beforeExecutedOn,
 		AfterKey:         afterKey,
-		Limit:            pagination.ValidPageSize(req.PageSize, defaultPageSize),
+		Limit:            pageLimit,
 	}
 
 	partitions, err := catalog.FindModelPartitions(ctx, opts)
@@ -342,9 +343,9 @@ func (s *Server) GetModelPartitions(ctx context.Context, req *runtimev1.GetModel
 	}
 
 	var nextPageToken string
-	if len(partitions) == pagination.ValidPageSize(req.PageSize, defaultPageSize) {
+	if len(partitions) == pageLimit && len(partitions) > 0 {
 		last := partitions[len(partitions)-1]
-		nextPageToken = pagination.MarshalPageToken(last.Index, last.Key)
+		nextPageToken = modelPartitionPageToken(last)
 	}
 
 	return &runtimev1.GetModelPartitionsResponse{
@@ -353,6 +354,14 @@ func (s *Server) GetModelPartitions(ctx context.Context, req *runtimev1.GetModel
 	}, nil
 }
 
+func modelPartitionPageToken(partition drivers.ModelPartition) string {
+	var executedOn time.Time
+	if partition.ExecutedOn != nil {
+		executedOn = *partition.ExecutedOn
+	}
+	return pagination.MarshalPageToken(executedOn, partition.Key)
+}
+
 // CreateTrigger implements runtimev1.RuntimeServiceServer
 func (s *Server) CreateTrigger(ctx context.Context, req *runtimev1.CreateTriggerRequest) (*runtimev1.CreateTriggerResponse, error) {
 	s.addInstanceRequestAttributes(ctx, req.InstanceId)
diff --git a/runtime/server/controller_pagination_test.go b/runtime/server/controller_pagination_test.go
new file mode 100644
index 00000000000..ece6cabc5b0
--- /dev/null
+++ b/runtime/server/controller_pagination_test.go
@@ -0,0 +1,44 @@
+package server
+
+import (
+	"testing"
+	"time"
+
+	"github.com/rilldata/rill/runtime/drivers"
+	"github.com/rilldata/rill/runtime/pkg/pagination"
+	"github.com/stretchr/testify/require"
+)
+
+func TestModelPartitionPageToken_ExecutedOn(t *testing.T) {
+	ts := time.Date(2024, time.March, 5, 12, 0, 0, 0, time.UTC)
+	partition := drivers.ModelPartition{
+		Key:        "abc123",
+		ExecutedOn: &ts,
+	}
+
+	token := modelPartitionPageToken(partition)
+	require.NotEmpty(t, token)
+
+	var decoded time.Time
+	var decodedKey string
+	err := pagination.UnmarshalPageToken(token, &decoded, &decodedKey)
+	require.NoError(t, err)
+	require.True(t, ts.Equal(decoded))
+	require.Equal(t, partition.Key, decodedKey)
+}
+
+func TestModelPartitionPageToken_NilExecutedOn(t *testing.T) {
+	partition := drivers.ModelPartition{
+		Key: "pending",
+	}
+
+	token := modelPartitionPageToken(partition)
+	require.NotEmpty(t, token)
+
+	var decoded time.Time
+	var decodedKey string
+	err := pagination.UnmarshalPageToken(token, &decoded, &decodedKey)
+	require.NoError(t, err)
+	require.True(t, decoded.IsZero())
+	require.Equal(t, partition.Key, decodedKey)
+}