chore: add offset-based pagination support to aibridge list endpoint (#20393)

Necessary for the frontend to be able to paginate easily. Cursor pagination is good for fetching all events, but doesn't play very well when a pagination component gets involved. Adds support for `?offset=x` to the existing endpoint. The cursor-based pagination (`?after_id=x`) is still supported. The two pagination modes are mutually exclusive, and are documented as such. If both are supplied, the request will be rejected. Also adds a `total` property to the response that contains the full count of items matching the filter. We already have indices in place so I don't think this will impact performance (or we can revisit it before GA).
2026-09-24 15:04:27 +08:00 · 2025-10-21 11:50:00 +00:00
parent 01ff28db11
commit ea261a1f7c
17 changed files with 328 additions and 48 deletions
@@ -33,7 +33,8 @@ const (
 // @Tags AIBridge
 // @Param q query string false "Search query in the format `key:value`. Available keys are: initiator, provider, model, started_after, started_before."
 // @Param limit query int false "Page limit"
-// @Param after_id query string false "Cursor pagination after ID"
+// @Param after_id query string false "Cursor pagination after ID (cannot be used with offset)"
+// @Param offset query int false "Offset pagination (cannot be used with after_id)"
 // @Success 200 {object} codersdk.AIBridgeListInterceptionsResponse
 // @Router /api/experimental/aibridge/interceptions [get]
 func (api *API) aiBridgeListInterceptions(rw http.ResponseWriter, r *http.Request) {
@@ -44,10 +45,10 @@ func (api *API) aiBridgeListInterceptions(rw http.ResponseWriter, r *http.Reques
 	if !ok {
 		return
 	}
-	if page.Offset != 0 {
+	if page.AfterID != uuid.Nil && page.Offset != 0 {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
-			Message: "Offset pagination is not supported.",
-			Detail:  "Offset pagination is not supported for AIBridge interceptions. Use cursor pagination instead with after_id.",
+			Message: "Query parameters have invalid values.",
+			Detail:  "Cannot use both after_id and offset pagination in the same request.",
 		})
 		return
 	}
@@ -72,7 +73,10 @@ func (api *API) aiBridgeListInterceptions(rw http.ResponseWriter, r *http.Reques
 		return
 	}

-	var rows []database.AIBridgeInterception
+	var (
+		count int64
+		rows  []database.AIBridgeInterception
+	)
 	err := api.Database.InTx(func(db database.Store) error {
 		// Ensure the after_id interception exists and is visible to the user.
 		if page.AfterID != uuid.Nil {
@@ -83,6 +87,19 @@ func (api *API) aiBridgeListInterceptions(rw http.ResponseWriter, r *http.Reques
 		}

 		var err error
+		// Get the full count of authorized interceptions matching the filter
+		// for pagination purposes.
+		count, err = db.CountAIBridgeInterceptions(ctx, database.CountAIBridgeInterceptionsParams{
+			StartedAfter:  filter.StartedAfter,
+			StartedBefore: filter.StartedBefore,
+			InitiatorID:   filter.InitiatorID,
+			Provider:      filter.Provider,
+			Model:         filter.Model,
+		})
+		if err != nil {
+			return xerrors.Errorf("count authorized aibridge interceptions: %w", err)
+		}
+
 		// This only returns authorized interceptions (when using dbauthz).
 		rows, err = db.ListAIBridgeInterceptions(ctx, filter)
 		if err != nil {
@@ -110,6 +127,7 @@ func (api *API) aiBridgeListInterceptions(rw http.ResponseWriter, r *http.Reques
 	}

 	httpapi.Write(ctx, rw, http.StatusOK, codersdk.AIBridgeListInterceptionsResponse{
+		Total:   count,
 		Results: items,
 	})
 }
@@ -196,42 +196,7 @@ func TestAIBridgeListInterceptions(t *testing.T) {
 			allInterceptionIDs = append(allInterceptionIDs, interception.ID)
 		}

-		// Get all interceptions one by one from the API using cursor
-		// pagination.
-		getAllInterceptionsOneByOne := func() []uuid.UUID {
-			interceptionIDs := []uuid.UUID{}
-			for {
-				afterID := uuid.Nil
-				if len(interceptionIDs) > 0 {
-					afterID = interceptionIDs[len(interceptionIDs)-1]
-				}
-				res, err := experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{
-					Pagination: codersdk.Pagination{
-						AfterID: afterID,
-						Limit:   1,
-					},
-				})
-				require.NoError(t, err)
-				if len(res.Results) == 0 {
-					break
-				}
-				require.Len(t, res.Results, 1)
-				interceptionIDs = append(interceptionIDs, res.Results[0].ID)
-			}
-			return interceptionIDs
-		}
-
-		// First attempt: get all interceptions one by one.
-		gotInterceptionIDs1 := getAllInterceptionsOneByOne()
-		// We should have all of the interceptions returned:
-		require.ElementsMatch(t, allInterceptionIDs, gotInterceptionIDs1)
-
-		// Second attempt: get all interceptions one by one again.
-		gotInterceptionIDs2 := getAllInterceptionsOneByOne()
-		// They should be returned in the exact same order.
-		require.Equal(t, gotInterceptionIDs1, gotInterceptionIDs2)
-
-		// Try to get an invalid limit.
+		// Try to fetch with an invalid limit.
 		res, err := experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{
 			Pagination: codersdk.Pagination{
 				Limit: 1001,
@@ -241,6 +206,66 @@ func TestAIBridgeListInterceptions(t *testing.T) {
 		require.ErrorAs(t, err, &sdkErr)
 		require.Contains(t, sdkErr.Message, "Invalid pagination limit value.")
 		require.Empty(t, res.Results)
+
+		// Try to fetch with both after_id and offset pagination.
+		res, err = experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{
+			Pagination: codersdk.Pagination{
+				AfterID: allInterceptionIDs[0],
+				Offset:  1,
+			},
+		})
+		require.ErrorAs(t, err, &sdkErr)
+		require.Contains(t, sdkErr.Message, "Query parameters have invalid values")
+		require.Contains(t, sdkErr.Detail, "Cannot use both after_id and offset pagination in the same request.")
+
+		// Iterate over all interceptions using both cursor and offset
+		// pagination modes.
+		for _, paginationMode := range []string{"after_id", "offset"} {
+			t.Run(paginationMode, func(t *testing.T) {
+				t.Parallel()
+
+				ctx := testutil.Context(t, testutil.WaitLong)
+
+				// Get all interceptions one by one using the given pagination
+				// mode.
+				getAllInterceptionsOneByOne := func() []uuid.UUID {
+					interceptionIDs := []uuid.UUID{}
+					for {
+						pagination := codersdk.Pagination{
+							Limit: 1,
+						}
+						if paginationMode == "after_id" {
+							if len(interceptionIDs) > 0 {
+								pagination.AfterID = interceptionIDs[len(interceptionIDs)-1]
+							}
+						} else {
+							pagination.Offset = len(interceptionIDs)
+						}
+						res, err := experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{
+							Pagination: pagination,
+						})
+						require.NoError(t, err)
+						if len(res.Results) == 0 {
+							break
+						}
+						require.EqualValues(t, len(allInterceptionIDs), res.Total)
+						require.Len(t, res.Results, 1)
+						interceptionIDs = append(interceptionIDs, res.Results[0].ID)
+					}
+					return interceptionIDs
+				}
+
+				// First attempt: get all interceptions one by one.
+				gotInterceptionIDs1 := getAllInterceptionsOneByOne()
+				// We should have all of the interceptions returned:
+				require.ElementsMatch(t, allInterceptionIDs, gotInterceptionIDs1)
+
+				// Second attempt: get all interceptions one by one again.
+				gotInterceptionIDs2 := getAllInterceptionsOneByOne()
+				// They should be returned in the exact same order.
+				require.Equal(t, gotInterceptionIDs1, gotInterceptionIDs2)
+			})
+		}
 	})

 	t.Run("Authorized", func(t *testing.T) {
@@ -276,6 +301,7 @@ func TestAIBridgeListInterceptions(t *testing.T) {
 		// Admin can see all interceptions.
 		res, err := adminExperimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{})
 		require.NoError(t, err)
+		require.EqualValues(t, 2, res.Total)
 		require.Len(t, res.Results, 2)
 		require.Equal(t, i1.ID, res.Results[0].ID)
 		require.Equal(t, i2.ID, res.Results[1].ID)
@@ -283,6 +309,7 @@ func TestAIBridgeListInterceptions(t *testing.T) {
 		// Second user can only see their own interceptions.
 		res, err = secondUserExperimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{})
 		require.NoError(t, err)
+		require.EqualValues(t, 1, res.Total)
 		require.Len(t, res.Results, 1)
 		require.Equal(t, i2.ID, res.Results[0].ID)
 	})
@@ -436,6 +463,7 @@ func TestAIBridgeListInterceptions(t *testing.T) {
 				ctx := testutil.Context(t, testutil.WaitLong)
 				res, err := experimentalClient.AIBridgeListInterceptions(ctx, tc.filter)
 				require.NoError(t, err)
+				require.EqualValues(t, len(tc.want), res.Total)
 				// We just compare UUID strings for the sake of this test.
 				wantIDs := make([]string, len(tc.want))
 				for i, r := range tc.want {