Add additional steelthread tests for LookupResources and fix issue in…

… LR2

internal/graph/lookupresources2.go

@@ -2,6 +2,7 @@ package graph
 
 import (
 	"context"
+	"errors"
 	"slices"
 	"sort"
 
@@ -583,27 +584,55 @@ func (crr *CursoredLookupResources2) redispatchOrReport(
 			// The stream that collects the results of the dispatch will add metadata to the response,
 			// map the results found based on the mapping data in the results and, if the entrypoint is not
 			// direct, issue a check to further filter the results.
-			stream, completed := lookupResourcesDispatchStreamForEntrypoint(ctx, foundResources, parentStream, entrypoint, ci, parentRequest, crr.dc)
-
-			// Dispatch the found resources as the subjects for the next call, to continue the
-			// resolution.
-			err = crr.dl.DispatchLookupResources2(&v1.DispatchLookupResources2Request{
-				ResourceRelation: parentRequest.ResourceRelation,
-				SubjectRelation:  newSubjectType,
-				SubjectIds:       filteredSubjectIDs,
-				TerminalSubject:  parentRequest.TerminalSubject,
-				Metadata: &v1.ResolverMeta{
-					AtRevision:     parentRequest.Revision.String(),
-					DepthRemaining: parentRequest.Metadata.DepthRemaining - 1,
-				},
-				OptionalCursor: ci.currentCursor,
-				OptionalLimit:  ci.limits.currentLimit,
-			}, stream)
-			if err != nil {
-				return err
+			currentCursor := ci.currentCursor
+
+			// Loop until we've produced enough results to satisfy the limit. This is necessary because
+			// the dispatch may return a set of results that, after checking, is less than the limit.
+			for {
+				stream, completed := lookupResourcesDispatchStreamForEntrypoint(ctx, foundResources, parentStream, entrypoint, ci, parentRequest, crr.dc)
+
+				// NOTE: if the entrypoint is a direct result, then all results returned by the dispatch will, themselves,
+				// be direct results. In this case, we can request the full limit of results. If the entrypoint is not a
+				// direct result, then we must request more than the limit in the hope that we get enough results to satisfy the
+				// limit after filtering.
+				var limit uint32 = uint32(datastore.FilterMaximumIDCount)
+				if entrypoint.IsDirectResult() {
+					limit = parentRequest.OptionalLimit
+				}
+
+				// Dispatch the found resources as the subjects for the next call, to continue the
+				// resolution.
+				err = crr.dl.DispatchLookupResources2(&v1.DispatchLookupResources2Request{
+					ResourceRelation: parentRequest.ResourceRelation,
+					SubjectRelation:  newSubjectType,
+					SubjectIds:       filteredSubjectIDs,
+					TerminalSubject:  parentRequest.TerminalSubject,
+					Metadata: &v1.ResolverMeta{
+						AtRevision:     parentRequest.Revision.String(),
+						DepthRemaining: parentRequest.Metadata.DepthRemaining - 1,
+					},
+					OptionalCursor: currentCursor,
+					OptionalLimit:  limit, // Request more than the limit to hopefully get enough results.
+				}, stream)
+				if err != nil {
+					// If the dispatch was canceled due to the limit, do not treat it as an error.
+					if errors.Is(err, errCanceledBecauseLimitReached) {
+						return err
+					}
+				}
+
+				nextCursor, err := completed()
+				if err != nil {
+					return err
+				}
+
+				if nextCursor == nil || ci.limits.hasExhaustedLimit() {
+					break
+				}
+				currentCursor = nextCursor
 			}
 
-			return completed()
+			return nil
 		})
 }
 
@@ -615,25 +644,20 @@ func lookupResourcesDispatchStreamForEntrypoint(
 	ci cursorInformation,
 	parentRequest ValidatedLookupResources2Request,
 	dc dispatch.Check,
-) (dispatch.LookupResources2Stream, func() error) {
+) (dispatch.LookupResources2Stream, func() (*v1.Cursor, error)) {
 	// Branch the context so that the dispatch can be canceled without canceling the parent
 	// call.
 	sctx, cancelDispatch := branchContext(ctx)
 
 	needsCallAddedToMetadata := true
 	resultsToCheck := make([]*v1.DispatchLookupResources2Response, 0, int(datastore.FilterMaximumIDCount))
+	var nextCursor *v1.Cursor
 
 	publishResultToParentStream := func(
 		result *v1.DispatchLookupResources2Response,
 		additionalMissingContext []string,
 		additionalMetadata *v1.ResponseMeta,
 	) error {
-		// If we've exhausted the limit of resources to be returned, nothing more to do.
-		if ci.limits.hasExhaustedLimit() {
-			cancelDispatch(errCanceledBecauseLimitReached)
-			return nil
-		}
-
 		// Map the found resources via the subject+resources used for dispatching, to determine
 		// if any need to be made conditional due to caveats.
 		mappedResource, err := foundResources.mapPossibleResource(result.Resource)
@@ -666,7 +690,7 @@ func lookupResourcesDispatchStreamForEntrypoint(
 			metadata = addAdditionalDepthRequired(metadata)
 		}
 
-		missingContextParameters := mapz.NewSet[string](mappedResource.MissingContextParams...)
+		missingContextParameters := mapz.NewSet(mappedResource.MissingContextParams...)
 		missingContextParameters.Extend(result.Resource.MissingContextParams)
 		missingContextParameters.Extend(additionalMissingContext)
 
@@ -766,6 +790,8 @@ func lookupResourcesDispatchStreamForEntrypoint(
 		default:
 		}
 
+		nextCursor = result.AfterResponseCursor
+
 		// If the entrypoint is a direct result, simply publish the found resource.
 		if entrypoint.IsDirectResult() {
 			return publishResultToParentStream(result, nil, emptyMetadata)
@@ -775,9 +801,8 @@ func lookupResourcesDispatchStreamForEntrypoint(
 		return batchCheckAndPublishIfNecessary(result)
 	})
 
-	return wrappedStream, func() error {
+	return wrappedStream, func() (*v1.Cursor, error) {
 		defer cancelDispatch(nil)
-
-		return batchCheckAndPublishIfNecessary(nil)
+		return nextCursor, batchCheckAndPublishIfNecessary(nil)
 	}
 }

internal/services/integrationtesting/testconfigs/directandindirect.yaml

@@ -0,0 +1,45 @@
+---
+schema: |+
+  definition user {}
+
+  definition organization {
+    relation admin: user
+    relation member: user
+    permission is_member = member + admin
+  }
+
+  definition document {
+    relation banned: user
+    relation viewer: user | user:*
+    relation org: organization
+
+    permission edit = org->is_member
+    permission view = (viewer - banned) + edit
+    permission vsb = viewer - banned
+    permission indirect_view = vsb + edit
+  }
+
+relationships: |
+  document:firstdoc#viewer@user:tom
+  document:firstdoc#banned@user:fred
+  document:firstdoc#org@organization:someorg
+  organization:someorg#member@user:fred
+
+  document:publicdoc#viewer@user:*
+  document:publicdoc#banned@user:fred
+  document:publicdoc#org@organization:someorg
+assertions:
+  assertTrue:
+    - "document:firstdoc#view@user:tom"
+    - "document:firstdoc#view@user:fred"
+    - "document:firstdoc#indirect_view@user:tom"
+    - "document:firstdoc#indirect_view@user:fred"
+    - "document:firstdoc#vsb@user:tom"
+    - "document:publicdoc#view@user:tom"
+    - "document:publicdoc#view@user:fred"
+    - "document:publicdoc#indirect_view@user:tom"
+    - "document:publicdoc#indirect_view@user:fred"
+    - "document:publicdoc#vsb@user:tom"
+  assertFalse:
+    - "document:firstdoc#vsb@user:fred"
+    - "document:publicdoc#vsb@user:fred"

internal/services/steelthreadtesting/definitions.go

@@ -209,6 +209,54 @@ var steelThreadTestCases = []steelThreadTestCase{
 				},
 				resultsFileName: "basic-lookup-resources-cursored-lookup-resources-for-fred-page-size-100-results.yaml",
 			},
+			{
+				name:          "indirect without other permission, page size 5",
+				operationName: "cursoredLookupResources",
+				arguments: map[string]any{
+					"resource_type":     "document",
+					"permission":        "vsb",
+					"subject_type":      "user",
+					"subject_object_id": "fred",
+					"page_size":         5,
+				},
+				resultsFileName: "basic-lookup-resources-indirect-without-other-permission-page-size-5-results.yaml",
+			},
+			{
+				name:          "indirect without other permission, page size 16",
+				operationName: "cursoredLookupResources",
+				arguments: map[string]any{
+					"resource_type":     "document",
+					"permission":        "vsb",
+					"subject_type":      "user",
+					"subject_object_id": "fred",
+					"page_size":         16,
+				},
+				resultsFileName: "basic-lookup-resources-indirect-without-other-permission-page-size-16-results.yaml",
+			},
+			{
+				name:          "vsb_plus_nil, page size 16",
+				operationName: "cursoredLookupResources",
+				arguments: map[string]any{
+					"resource_type":     "document",
+					"permission":        "vsb_plus_nil",
+					"subject_type":      "user",
+					"subject_object_id": "fred",
+					"page_size":         16,
+				},
+				resultsFileName: "basic-lookup-resources-vsb-plus-nil-page-size-16-results.yaml",
+			},
+			{
+				name:          "edit, page size 16",
+				operationName: "cursoredLookupResources",
+				arguments: map[string]any{
+					"resource_type":     "document",
+					"permission":        "edit",
+					"subject_type":      "user",
+					"subject_object_id": "fred",
+					"page_size":         16,
+				},
+				resultsFileName: "basic-lookup-resources-edit-page-size-16-results.yaml",
+			},
 		},
 	},
 	{

internal/services/steelthreadtesting/operations.go

@@ -6,6 +6,7 @@ package steelthreadtesting
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"sort"
 	"strings"
@@ -147,6 +148,7 @@ func cursoredLookupResources(parameters map[string]any, client v1.PermissionsSer
 
 	var currentCursor *v1.Cursor
 	nodeSets := make([][]yaml.Node, 0)
+	resultCounts := make([]int, 0)
 	for {
 		r, err := client.LookupResources(ctx, &v1.LookupResourcesRequest{
 			ResourceObjectType: parameters["resource_type"].(string),
@@ -171,6 +173,7 @@ func cursoredLookupResources(parameters map[string]any, client v1.PermissionsSer
 		}
 
 		foundResources := mapz.NewSet[string]()
+		resultCount := 0
 		for {
 			resp, err := r.Recv()
 			if err != nil {
@@ -183,12 +186,15 @@ func cursoredLookupResources(parameters map[string]any, client v1.PermissionsSer
 
 			foundResources.Add(formatResolvedResource(resp))
 			currentCursor = resp.AfterResultCursor
+			resultCount++
 		}
 
 		if foundResources.IsEmpty() {
 			break
 		}
 
+		resultCounts = append(resultCounts, resultCount)
+
 		foundResourcesSlice := foundResources.AsSlice()
 		sort.Strings(foundResourcesSlice)
 
@@ -204,6 +210,16 @@ func cursoredLookupResources(parameters map[string]any, client v1.PermissionsSer
 		nodeSets = append(nodeSets, yamlNodes)
 	}
 
+	for index, count := range resultCounts {
+		if index == len(resultCounts)-1 {
+			continue
+		}
+
+		if count != parameters["page_size"].(int) {
+			return nil, fmt.Errorf("expected full page size of %d for page #%d, got %d", parameters["page_size"].(int), index, count)
+		}
+	}
+
 	return nodeSets, nil
 }
 

internal/services/steelthreadtesting/steelresults/basic-lookup-resources-edit-page-size-16-results.yaml

@@ -0,0 +1,11 @@
+---
+- - 'doc-90'
+  - 'doc-91'
+  - 'doc-92'
+  - 'doc-93'
+  - 'doc-94'
+  - 'doc-95'
+  - 'doc-96'
+  - 'doc-97'
+  - 'doc-98'
+  - 'doc-99'

internal/services/steelthreadtesting/steelresults/basic-lookup-resources-indirect-without-other-permission-page-size-16-results.yaml

@@ -0,0 +1,87 @@
+---
+- - 'doc-0'
+  - 'doc-1'
+  - 'doc-10'
+  - 'doc-11'
+  - 'doc-12'
+  - 'doc-13'
+  - 'doc-14'
+  - 'doc-15'
+  - 'doc-16'
+  - 'doc-17'
+  - 'doc-18'
+  - 'doc-19'
+  - 'doc-2'
+  - 'doc-20'
+  - 'doc-21'
+  - 'doc-22'
+- - 'doc-23'
+  - 'doc-24'
+  - 'doc-25'
+  - 'doc-26'
+  - 'doc-27'
+  - 'doc-28'
+  - 'doc-29'
+  - 'doc-3'
+  - 'doc-30'
+  - 'doc-31'
+  - 'doc-32'
+  - 'doc-33'
+  - 'doc-34'
+  - 'doc-35'
+  - 'doc-36'
+  - 'doc-37'
+- - 'doc-38'
+  - 'doc-39'
+  - 'doc-4'
+  - 'doc-40'
+  - 'doc-41'
+  - 'doc-42'
+  - 'doc-43'
+  - 'doc-44'
+  - 'doc-45'
+  - 'doc-46'
+  - 'doc-47'
+  - 'doc-48'
+  - 'doc-49'
+  - 'doc-5'
+  - 'doc-50'
+  - 'doc-51'
+- - 'doc-52'
+  - 'doc-53'
+  - 'doc-54'
+  - 'doc-55'
+  - 'doc-56'
+  - 'doc-57'
+  - 'doc-58'
+  - 'doc-59'
+  - 'doc-6'
+  - 'doc-60'
+  - 'doc-61'
+  - 'doc-62'
+  - 'doc-63'
+  - 'doc-64'
+  - 'doc-65'
+  - 'doc-66'
+- - 'doc-67'
+  - 'doc-68'
+  - 'doc-69'
+  - 'doc-7'
+  - 'doc-70'
+  - 'doc-71'
+  - 'doc-72'
+  - 'doc-73'
+  - 'doc-74'
+  - 'doc-75'
+  - 'doc-76'
+  - 'doc-77'
+  - 'doc-78'
+  - 'doc-79'
+  - 'doc-8'
+  - 'doc-9'
+- - 'doc-9'
+  - 'public-doc-0'
+  - 'public-doc-1'
+  - 'public-doc-3'
+- - 'public-doc-1'
+  - 'public-doc-3'