[Backport 5.3.9104] Context: fix bug with result limits (#62054)

Context: fix bug with result limits (#61806) While writing eval scripts for Cody context, I noticed that when searching multiple repos, we sometimes we return greater than the intended number of results. This PR makes a small fix so we always respect the result limits. (cherry picked from commit 6a4ac36) Co-authored-by: Julie Tibshirani <[email protected]>
sourcegraph · Apr 22, 2024 · ece951a · ece951a
1 parent e2cf4ee
commit ece951a
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/internal/codycontext/context.go b/internal/codycontext/context.go
@@ -302,6 +302,12 @@ func (c *CodyContextClient) getKeywordContext(ctx context.Context, args GetConte
 			mu.Lock()
 			defer mu.Unlock()
 
+			// Another caller may have already hit the limit, but we haven't yet responded
+			// to the cancellation. Return immediately in this case.
+			if len(collected) >= limit {
+				return
+			}
+
 			for _, res := range e.Results {
 				if fm, ok := res.(*result.FileMatch); ok {
 					collected = append(collected, filter(fileMatchToContextMatches(fm))...)