feat(github): rate-limit awareness, poller throttling, and GraphQL batching (#821)
This commit is contained in:
Carlos Florêncio
@@ -654,16 +654,44 @@ func registerSecondaryRoutes(
|
||||
// registerHealthRoutes sets up the system health endpoint with all health checkers.
|
||||
func registerHealthRoutes(p routeParams) {
|
||||
var githubProvider health.GitHubStatusProvider
|
||||
var githubRateProvider health.GitHubRateLimitProvider
|
||||
if p.services.GitHub != nil {
|
||||
githubProvider = p.services.GitHub
|
||||
githubRateProvider = githubRateLimitAdapter{svc: p.services.GitHub}
|
||||
}
|
||||
githubChecker := health.NewGitHubChecker(githubProvider)
|
||||
if githubRateProvider != nil {
|
||||
githubChecker.WithRateLimitProvider(githubRateProvider)
|
||||
}
|
||||
healthSvc := health.NewService(p.log,
|
||||
health.NewGitHubChecker(githubProvider),
|
||||
githubChecker,
|
||||
health.NewAgentChecker(p.agentSettingsController),
|
||||
)
|
||||
health.RegisterRoutes(p.router, healthSvc, p.log)
|
||||
}
|
||||
|
||||
// githubRateLimitAdapter bridges the github.Service's per-resource exhaustion
|
||||
// snapshot to the structural shape consumed by the health package without
|
||||
// importing health into github (cycle).
|
||||
type githubRateLimitAdapter struct {
|
||||
svc *github.Service
|
||||
}
|
||||
|
||||
func (a githubRateLimitAdapter) ExhaustedRateLimits() []health.GitHubRateLimitStatus {
|
||||
if a.svc == nil {
|
||||
return nil
|
||||
}
|
||||
src := a.svc.ExhaustedRateLimits()
|
||||
if len(src) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]health.GitHubRateLimitStatus, len(src))
|
||||
for i, s := range src {
|
||||
out[i] = health.GitHubRateLimitStatus{Resource: s.Resource, ResetAt: s.ResetAt}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// registerMCPAndDebugRoutes registers MCP and debug routes and wires the MCP handler.
|
||||
func registerMCPAndDebugRoutes(
|
||||
p routeParams,
|
||||
|
||||
@@ -34,7 +34,7 @@ func TestComputeMergeBase_PrefersOriginOverStaleLocalBranch(t *testing.T) {
|
||||
advancedMain := strings.TrimSpace(runGitAPI(t, repoDir, "rev-parse", "HEAD"))
|
||||
runGitAPI(t, repoDir, "push", "origin", "main")
|
||||
runGitAPI(t, repoDir, "reset", "--hard", staleLocalMain)
|
||||
runGitAPI(t, repoDir, "fetch", "origin", "main:refs/remotes/origin/main")
|
||||
runGitAPI(t, repoDir, "fetch", "origin")
|
||||
|
||||
// Branch from the (advanced) origin/main and add one commit.
|
||||
runGitAPI(t, repoDir, "checkout", "-b", "feature/x", advancedMain)
|
||||
|
||||
@@ -79,7 +79,7 @@ func TestGetLog_StaleLocalBranchScenario(t *testing.T) {
|
||||
// Reset local main back to staleLocalMain (simulates "didn't fetch").
|
||||
runGit(t, repoDir, "reset", "--hard", staleLocalMain)
|
||||
// Now refetch so origin/main is ahead but local main isn't.
|
||||
runGit(t, repoDir, "fetch", "origin", "main:refs/remotes/origin/main")
|
||||
runGit(t, repoDir, "fetch", "origin")
|
||||
|
||||
// Branch from origin/main (the current upstream tip).
|
||||
runGit(t, repoDir, "checkout", "-b", "feature/x", advancedMain)
|
||||
|
||||
@@ -198,12 +198,13 @@ const (
|
||||
|
||||
// Event types for GitHub integration
|
||||
const (
|
||||
GitHubPRFeedback = "github.pr_feedback" // PR has new feedback (UI notification only)
|
||||
GitHubPRStateChanged = "github.pr_state_changed" // PR state changed (merged, closed, etc.)
|
||||
GitHubNewReviewPR = "github.new_pr_to_review" // New PR found needing review
|
||||
GitHubNewIssue = "github.new_issue" // New issue found matching issue watch
|
||||
GitHubTaskPRUpdated = "github.task_pr.updated" // TaskPR record updated (for UI refresh)
|
||||
GitHubWatchEvent = "github.watch.event" // Watch created/deleted
|
||||
GitHubPRFeedback = "github.pr_feedback" // PR has new feedback (UI notification only)
|
||||
GitHubPRStateChanged = "github.pr_state_changed" // PR state changed (merged, closed, etc.)
|
||||
GitHubNewReviewPR = "github.new_pr_to_review" // New PR found needing review
|
||||
GitHubNewIssue = "github.new_issue" // New issue found matching issue watch
|
||||
GitHubTaskPRUpdated = "github.task_pr.updated" // TaskPR record updated (for UI refresh)
|
||||
GitHubWatchEvent = "github.watch.event" // Watch created/deleted
|
||||
GitHubRateLimitUpdated = "github.rate_limit.updated" // GitHub API rate-limit snapshot changed
|
||||
)
|
||||
|
||||
// Event types for Jira integration
|
||||
|
||||
@@ -73,6 +73,7 @@ func RegisterTaskNotifications(ctx context.Context, eventBus bus.EventBus, hub *
|
||||
b.subscribe(eventBus, events.TurnCompleted, ws.ActionSessionTurnCompleted)
|
||||
b.subscribe(eventBus, events.MessageQueueStatusChanged, ws.ActionMessageQueueStatusChanged)
|
||||
b.subscribe(eventBus, events.GitHubTaskPRUpdated, ws.ActionGitHubTaskPRUpdated)
|
||||
b.subscribe(eventBus, events.GitHubRateLimitUpdated, ws.ActionGitHubRateLimitUpdated)
|
||||
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
|
||||
@@ -39,7 +39,7 @@ func TestTaskEventBroadcaster_NoDuplicateSubscriptions(t *testing.T) {
|
||||
//
|
||||
// Update this number when adding or removing event subscriptions in
|
||||
// RegisterTaskNotifications — it is intentionally exact.
|
||||
const wantSubscriptions = 48
|
||||
const wantSubscriptions = 49
|
||||
if got := len(b.subscriptions); got != wantSubscriptions {
|
||||
t.Errorf("RegisterTaskNotifications created %d subscriptions, want %d — "+
|
||||
"did an event get subscribed twice?", got, wantSubscriptions)
|
||||
|
||||
@@ -13,13 +13,79 @@ import (
|
||||
)
|
||||
|
||||
// GHClient implements Client using the gh CLI.
|
||||
type GHClient struct{}
|
||||
type GHClient struct {
|
||||
rateTracker *RateTracker
|
||||
}
|
||||
|
||||
// NewGHClient creates a new gh CLI-based client.
|
||||
func NewGHClient() *GHClient {
|
||||
return &GHClient{}
|
||||
}
|
||||
|
||||
// WithRateTracker attaches a rate tracker so stderr output is inspected for
|
||||
// rate-limit signals after every gh invocation. Returns the client for
|
||||
// chaining; safe to call before any commands run.
|
||||
func (c *GHClient) WithRateTracker(t *RateTracker) *GHClient {
|
||||
c.rateTracker = t
|
||||
return c
|
||||
}
|
||||
|
||||
// reGHRateLimit matches the prose gh prints when a request hit a primary or
|
||||
// secondary rate limit. The exact text varies by gh version and locale, but
|
||||
// "rate limit" / "API rate limit" appear consistently.
|
||||
var ghRateLimitMarkers = []string{"rate limit", "abuse detection"}
|
||||
|
||||
func ghStderrIndicatesRateLimit(stderr string) bool {
|
||||
if stderr == "" {
|
||||
return false
|
||||
}
|
||||
lower := strings.ToLower(stderr)
|
||||
for _, marker := range ghRateLimitMarkers {
|
||||
if strings.Contains(lower, marker) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// inspectRateStderr inspects the stderr from a failed `gh` invocation and
|
||||
// flags the appropriate resource bucket as exhausted. `gh api <path>` is REST
|
||||
// (Core) by default, with `api graphql` and `api search/*` as the documented
|
||||
// exceptions; non-`api` subcommands like `pr`, `issue`, and `repo` are
|
||||
// implemented against GraphQL by gh itself, so they map to the GraphQL bucket.
|
||||
func (c *GHClient) inspectRateStderr(args []string, stderr string) {
|
||||
if c.rateTracker == nil {
|
||||
return
|
||||
}
|
||||
if !ghStderrIndicatesRateLimit(stderr) {
|
||||
return
|
||||
}
|
||||
c.rateTracker.markRateExhausted(resourceForGHArgs(args), time.Time{})
|
||||
}
|
||||
|
||||
// resourceForGHArgs maps a `gh` argv to the rate-limit bucket the call hits.
|
||||
// Exposed at package scope so the table-driven test can pin every branch.
|
||||
func resourceForGHArgs(args []string) Resource {
|
||||
if len(args) == 0 {
|
||||
return ResourceCore
|
||||
}
|
||||
if args[0] != "api" {
|
||||
// `gh pr`, `gh issue`, `gh repo`, etc. are all GraphQL under the hood.
|
||||
return ResourceGraphQL
|
||||
}
|
||||
if len(args) < 2 {
|
||||
return ResourceCore
|
||||
}
|
||||
switch {
|
||||
case args[1] == "graphql":
|
||||
return ResourceGraphQL
|
||||
case strings.HasPrefix(args[1], "search/"):
|
||||
return ResourceSearch
|
||||
default:
|
||||
return ResourceCore
|
||||
}
|
||||
}
|
||||
|
||||
// GHAvailable checks if the gh CLI is installed and accessible.
|
||||
func GHAvailable() bool {
|
||||
_, err := exec.LookPath("gh")
|
||||
@@ -516,11 +582,62 @@ func (c *GHClient) run(ctx context.Context, args ...string) (string, error) {
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
if err := cmd.Run(); err != nil {
|
||||
c.inspectRateStderr(args, stderr.String())
|
||||
return stdout.String(), fmt.Errorf("gh %s: %w: %s", args[0], err, stderr.String())
|
||||
}
|
||||
return stdout.String(), nil
|
||||
}
|
||||
|
||||
// ghRateLimitResponse mirrors the GET /rate_limit JSON shape so we can seed
|
||||
// the tracker on startup and after CLI failures without parsing prose.
|
||||
type ghRateLimitResponse struct {
|
||||
Resources struct {
|
||||
Core ghRateLimitBucket `json:"core"`
|
||||
GraphQL ghRateLimitBucket `json:"graphql"`
|
||||
Search ghRateLimitBucket `json:"search"`
|
||||
} `json:"resources"`
|
||||
}
|
||||
|
||||
type ghRateLimitBucket struct {
|
||||
Limit int `json:"limit"`
|
||||
Remaining int `json:"remaining"`
|
||||
Reset int64 `json:"reset"`
|
||||
}
|
||||
|
||||
// FetchRateLimit calls `gh api rate_limit` and seeds the tracker with the
|
||||
// returned snapshots. Best-effort: a failure (e.g. CLI absent, network) is
|
||||
// logged and ignored.
|
||||
func (c *GHClient) FetchRateLimit(ctx context.Context) error {
|
||||
if c.rateTracker == nil {
|
||||
return nil
|
||||
}
|
||||
out, err := c.run(ctx, "api", "rate_limit")
|
||||
if err != nil {
|
||||
return fmt.Errorf("fetch rate limit: %w", err)
|
||||
}
|
||||
var raw ghRateLimitResponse
|
||||
if err := json.Unmarshal([]byte(out), &raw); err != nil {
|
||||
return fmt.Errorf("parse rate_limit response: %w", err)
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
record := func(resource Resource, b ghRateLimitBucket) {
|
||||
if b.Limit == 0 && b.Remaining == 0 && b.Reset == 0 {
|
||||
return
|
||||
}
|
||||
c.rateTracker.Record(RateSnapshot{
|
||||
Resource: resource,
|
||||
Limit: b.Limit,
|
||||
Remaining: b.Remaining,
|
||||
ResetAt: time.Unix(b.Reset, 0).UTC(),
|
||||
UpdatedAt: now,
|
||||
})
|
||||
}
|
||||
record(ResourceCore, raw.Resources.Core)
|
||||
record(ResourceGraphQL, raw.Resources.GraphQL)
|
||||
record(ResourceSearch, raw.Resources.Search)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *GHClient) parsePRList(data string, owner, repo string) ([]*PR, error) {
|
||||
var raw []ghPR
|
||||
if err := json.Unmarshal([]byte(data), &raw); err != nil {
|
||||
|
||||
@@ -237,3 +237,90 @@ func TestConvertGHRequestedReviewers(t *testing.T) {
|
||||
t.Errorf("unexpected third reviewer: %#v", got[2])
|
||||
}
|
||||
}
|
||||
|
||||
func TestGHStderrIndicatesRateLimit(t *testing.T) {
|
||||
cases := []struct {
|
||||
stderr string
|
||||
want bool
|
||||
}{
|
||||
{"GraphQL: API rate limit already exceeded for user ID 12345.", true},
|
||||
{"You have exceeded a secondary rate limit.", true},
|
||||
{"abuse detection mechanism triggered", true},
|
||||
{"network: connection refused", false},
|
||||
{"", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := ghStderrIndicatesRateLimit(c.stderr); got != c.want {
|
||||
t.Errorf("ghStderrIndicatesRateLimit(%q) = %v, want %v", c.stderr, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGHClient_InspectRateStderr_MarksGraphQL(t *testing.T) {
|
||||
tracker := NewRateTracker(nil, nil)
|
||||
c := NewGHClient().WithRateTracker(tracker)
|
||||
c.inspectRateStderr([]string{"pr", "view", "1"}, "GraphQL: API rate limit already exceeded")
|
||||
if !tracker.IsExhausted(ResourceGraphQL) {
|
||||
t.Errorf("expected graphql exhausted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGHClient_InspectRateStderr_MarksSearchForSearchEndpoints(t *testing.T) {
|
||||
tracker := NewRateTracker(nil, nil)
|
||||
c := NewGHClient().WithRateTracker(tracker)
|
||||
c.inspectRateStderr([]string{"api", "search/issues"}, "API rate limit exceeded")
|
||||
if !tracker.IsExhausted(ResourceSearch) {
|
||||
t.Errorf("expected search exhausted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResourceForGHArgs(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
args []string
|
||||
want Resource
|
||||
}{
|
||||
// REST endpoints under `gh api <path>` go to Core, not GraphQL.
|
||||
// This is the regression: previously every gh api failure was
|
||||
// attributed to GraphQL unless the path started with search/.
|
||||
{"api repos REST", []string{"api", "repos/o/r/pulls/1"}, ResourceCore},
|
||||
{"api user REST", []string{"api", "user"}, ResourceCore},
|
||||
{"api rate_limit REST", []string{"api", "rate_limit"}, ResourceCore},
|
||||
|
||||
// Documented exceptions still resolve to their dedicated buckets.
|
||||
{"api graphql", []string{"api", "graphql"}, ResourceGraphQL},
|
||||
{"api search/issues", []string{"api", "search/issues"}, ResourceSearch},
|
||||
{"api search/repositories", []string{"api", "search/repositories"}, ResourceSearch},
|
||||
|
||||
// Non-`api` subcommands are GraphQL — gh implements pr/issue/repo
|
||||
// against the GraphQL API.
|
||||
{"pr view", []string{"pr", "view", "1"}, ResourceGraphQL},
|
||||
{"issue list", []string{"issue", "list"}, ResourceGraphQL},
|
||||
{"repo view", []string{"repo", "view"}, ResourceGraphQL},
|
||||
|
||||
// Defensive defaults for malformed argv.
|
||||
{"empty", nil, ResourceCore},
|
||||
{"api only", []string{"api"}, ResourceCore},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := resourceForGHArgs(tc.args); got != tc.want {
|
||||
t.Errorf("resourceForGHArgs(%v) = %s, want %s", tc.args, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Regression: a 429 on `gh api repos/...` (REST) must mark Core, not GraphQL.
|
||||
// Previously this would pause the GraphQL PR monitor incorrectly.
|
||||
func TestGHClient_InspectRateStderr_RestEndpointMarksCore(t *testing.T) {
|
||||
tracker := NewRateTracker(nil, nil)
|
||||
c := NewGHClient().WithRateTracker(tracker)
|
||||
c.inspectRateStderr([]string{"api", "repos/o/r/pulls/1"}, "API rate limit exceeded")
|
||||
if !tracker.IsExhausted(ResourceCore) {
|
||||
t.Errorf("expected core bucket exhausted for REST endpoint")
|
||||
}
|
||||
if tracker.IsExhausted(ResourceGraphQL) {
|
||||
t.Errorf("REST 429 must not pause graphql bucket")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,540 @@
|
||||
package github
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// GraphQLExecutor runs a single GraphQL query against the GitHub API. Both
|
||||
// PATClient (direct HTTP) and GHClient (`gh api graphql`) implement it so the
|
||||
// batched poller works regardless of auth method.
|
||||
type GraphQLExecutor interface {
|
||||
ExecuteGraphQL(ctx context.Context, query string, variables map[string]any, out any) error
|
||||
}
|
||||
|
||||
// graphQLPRBatchAlias is the per-PR alias prefix; the index makes aliases
|
||||
// unique within a single repository block.
|
||||
const graphQLPRBatchAlias = "pr"
|
||||
|
||||
// graphQLBatchChunkSize bounds the number of aliased fields per request to
|
||||
// stay well under GitHub's per-query node-count limit (500). 50 PRs × ~10
|
||||
// fields = ~500 nodes, leaving headroom.
|
||||
const graphQLBatchChunkSize = 50
|
||||
|
||||
// reviewNode is one PR review entry from the batched GraphQL query.
|
||||
type reviewNode struct {
|
||||
State string `json:"state"`
|
||||
Author struct {
|
||||
Login string `json:"login"`
|
||||
} `json:"author"`
|
||||
SubmittedAt time.Time `json:"submittedAt"`
|
||||
}
|
||||
|
||||
// graphQLPRRef is one entry in a batched PR-status request.
|
||||
type graphQLPRRef struct {
|
||||
Owner string
|
||||
Repo string
|
||||
Number int
|
||||
}
|
||||
|
||||
// graphQLBranchRef is one entry in a batched branch-lookup request.
|
||||
type graphQLBranchRef struct {
|
||||
Owner string
|
||||
Repo string
|
||||
Branch string
|
||||
}
|
||||
|
||||
// chunkedRefs splits refs into chunks of at most graphQLBatchChunkSize so
|
||||
// callers can keep individual GraphQL queries under the node-count limit.
|
||||
func chunkedRefs[T any](refs []T) [][]T {
|
||||
if len(refs) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([][]T, 0, (len(refs)+graphQLBatchChunkSize-1)/graphQLBatchChunkSize)
|
||||
for i := 0; i < len(refs); i += graphQLBatchChunkSize {
|
||||
end := i + graphQLBatchChunkSize
|
||||
if end > len(refs) {
|
||||
end = len(refs)
|
||||
}
|
||||
out = append(out, refs[i:end])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// batchedPRResult is the decoded shape of one aliased pullRequest block.
|
||||
type batchedPRResult struct {
|
||||
State string `json:"state"`
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
IsDraft bool `json:"isDraft"`
|
||||
Mergeable string `json:"mergeable"`
|
||||
MergeStatus string `json:"mergeStateStatus"`
|
||||
HeadRefName string `json:"headRefName"`
|
||||
BaseRefName string `json:"baseRefName"`
|
||||
HeadRefOid string `json:"headRefOid"`
|
||||
Additions int `json:"additions"`
|
||||
Deletions int `json:"deletions"`
|
||||
Author struct {
|
||||
Login string `json:"login"`
|
||||
} `json:"author"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
MergedAt string `json:"mergedAt"`
|
||||
ClosedAt string `json:"closedAt"`
|
||||
Reviews struct {
|
||||
Nodes []reviewNode `json:"nodes"`
|
||||
} `json:"reviews"`
|
||||
ReviewRequests struct {
|
||||
TotalCount int `json:"totalCount"`
|
||||
} `json:"reviewRequests"`
|
||||
Commits struct {
|
||||
Nodes []struct {
|
||||
Commit struct {
|
||||
StatusCheckRollup *struct {
|
||||
State string `json:"state"`
|
||||
} `json:"statusCheckRollup"`
|
||||
} `json:"commit"`
|
||||
} `json:"nodes"`
|
||||
} `json:"commits"`
|
||||
}
|
||||
|
||||
// graphQLError mirrors a single entry in a GraphQL response's "errors" array.
|
||||
type graphQLError struct {
|
||||
Message string `json:"message"`
|
||||
Type string `json:"type"`
|
||||
Path []string `json:"path"`
|
||||
}
|
||||
|
||||
// graphQLErrorsToErr returns a non-nil error when the GraphQL response carries
|
||||
// a top-level "errors" array. The first message is included so logs are
|
||||
// actionable; the count is appended when there are multiple.
|
||||
func graphQLErrorsToErr(errs []graphQLError) error {
|
||||
if len(errs) == 0 {
|
||||
return nil
|
||||
}
|
||||
if len(errs) == 1 {
|
||||
return fmt.Errorf("graphql error: %s", errs[0].Message)
|
||||
}
|
||||
return fmt.Errorf("graphql error: %s (and %d more)", errs[0].Message, len(errs)-1)
|
||||
}
|
||||
|
||||
// graphQLRateLimit mirrors GitHub's top-level rateLimit field, the most
|
||||
// accurate source of GraphQL quota since limits are point-cost based.
|
||||
type graphQLRateLimit struct {
|
||||
Limit int `json:"limit"`
|
||||
Remaining int `json:"remaining"`
|
||||
ResetAt time.Time `json:"resetAt"`
|
||||
Cost int `json:"cost"`
|
||||
}
|
||||
|
||||
// buildBatchedPRQuery groups refs by (owner, repo) and emits one aliased
|
||||
// repository block per group, with one aliased pullRequest per number inside.
|
||||
// The shape mirrors gh pr view fields used by the existing converter so we
|
||||
// can reuse the conversion logic without reshaping callers.
|
||||
func buildBatchedPRQuery(refs []graphQLPRRef) (string, map[string]any) {
|
||||
// Group refs by (owner, repo) and assign deterministic indices so aliases
|
||||
// stay stable across runs (helpful for tests and snapshot debugging).
|
||||
type group struct {
|
||||
owner string
|
||||
repo string
|
||||
numbers []int
|
||||
}
|
||||
byKey := map[string]*group{}
|
||||
keys := []string{}
|
||||
for _, r := range refs {
|
||||
key := r.Owner + "/" + r.Repo
|
||||
g, ok := byKey[key]
|
||||
if !ok {
|
||||
g = &group{owner: r.Owner, repo: r.Repo}
|
||||
byKey[key] = g
|
||||
keys = append(keys, key)
|
||||
}
|
||||
g.numbers = append(g.numbers, r.Number)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
|
||||
var b strings.Builder
|
||||
b.WriteString("query Batch { ")
|
||||
for repoIdx, key := range keys {
|
||||
g := byKey[key]
|
||||
fmt.Fprintf(&b, `repo%d: repository(owner: %q, name: %q) { `, repoIdx, g.owner, g.repo)
|
||||
for prIdx, n := range g.numbers {
|
||||
fmt.Fprintf(&b, `%s%d: pullRequest(number: %d) { %s } `, graphQLPRBatchAlias, prIdx, n, prFieldsBlock())
|
||||
}
|
||||
b.WriteString(`} `)
|
||||
}
|
||||
b.WriteString(`rateLimit { limit remaining resetAt cost } `)
|
||||
b.WriteString(`}`)
|
||||
return b.String(), nil
|
||||
}
|
||||
|
||||
// prFieldsBlock is the GraphQL field selection used in every batched PR
|
||||
// query. Kept as a constant to make snapshot assertions stable and to keep
|
||||
// the batched and single-PR paths returning the same data.
|
||||
func prFieldsBlock() string {
|
||||
return `state title url isDraft mergeable mergeStateStatus ` +
|
||||
`headRefName baseRefName headRefOid additions deletions ` +
|
||||
`author { login } createdAt updatedAt mergedAt closedAt ` +
|
||||
`reviews(last: 100) { nodes { state author { login } submittedAt } } ` +
|
||||
`reviewRequests(first: 0) { totalCount } ` +
|
||||
`commits(last: 1) { nodes { commit { statusCheckRollup { state } } } }`
|
||||
}
|
||||
|
||||
// buildBatchedBranchQuery emits one aliased associatedPullRequests block per
|
||||
// (owner, repo, branch) ref. Used to batch the "find PR by branch" path.
|
||||
func buildBatchedBranchQuery(refs []graphQLBranchRef) (string, map[string]any) {
|
||||
var b strings.Builder
|
||||
b.WriteString("query Branches { ")
|
||||
for i, r := range refs {
|
||||
fmt.Fprintf(&b, `b%d: repository(owner: %q, name: %q) { ref(qualifiedName: %q) { associatedPullRequests(first: 1, states: OPEN) { nodes { number %s } } } } `,
|
||||
i, r.Owner, r.Repo, "refs/heads/"+r.Branch, prFieldsBlock())
|
||||
}
|
||||
b.WriteString(`rateLimit { limit remaining resetAt cost } `)
|
||||
b.WriteString(`}`)
|
||||
return b.String(), nil
|
||||
}
|
||||
|
||||
// convertBatchedPRResult turns a batched GraphQL row into the (PR, PRStatus)
|
||||
// pair the existing poller code uses.
|
||||
func convertBatchedPRResult(raw *batchedPRResult, owner, repo string, number int) *PRStatus {
|
||||
state := strings.ToLower(raw.State)
|
||||
if raw.MergedAt != "" {
|
||||
state = prStateMerged
|
||||
}
|
||||
pr := &PR{
|
||||
Number: number,
|
||||
Title: raw.Title,
|
||||
URL: raw.URL,
|
||||
HTMLURL: raw.URL,
|
||||
State: state,
|
||||
HeadBranch: raw.HeadRefName,
|
||||
HeadSHA: raw.HeadRefOid,
|
||||
BaseBranch: raw.BaseRefName,
|
||||
AuthorLogin: raw.Author.Login,
|
||||
RepoOwner: owner,
|
||||
RepoName: repo,
|
||||
Draft: raw.IsDraft,
|
||||
Mergeable: raw.Mergeable == "MERGEABLE",
|
||||
MergeableState: strings.ToLower(raw.MergeStatus),
|
||||
Additions: raw.Additions,
|
||||
Deletions: raw.Deletions,
|
||||
CreatedAt: raw.CreatedAt,
|
||||
UpdatedAt: raw.UpdatedAt,
|
||||
MergedAt: parseTimePtr(raw.MergedAt),
|
||||
ClosedAt: parseTimePtr(raw.ClosedAt),
|
||||
}
|
||||
|
||||
reviewState := summarizeReviewState(raw.Reviews.Nodes)
|
||||
checksState := ""
|
||||
if len(raw.Commits.Nodes) > 0 && raw.Commits.Nodes[0].Commit.StatusCheckRollup != nil {
|
||||
checksState = strings.ToLower(raw.Commits.Nodes[0].Commit.StatusCheckRollup.State)
|
||||
}
|
||||
return &PRStatus{
|
||||
PR: pr,
|
||||
ReviewState: reviewState,
|
||||
ChecksState: checksState,
|
||||
MergeableState: pr.MergeableState,
|
||||
}
|
||||
}
|
||||
|
||||
// summarizeReviewState collapses the review history to a single
|
||||
// "approved"/"changes_requested"/"" value matching the existing
|
||||
// service-side summary. Per-reviewer dedup: each reviewer's most-recent
|
||||
// review wins, so a CHANGES_REQUESTED followed by an APPROVED from the
|
||||
// same author resolves to APPROVED. CHANGES_REQUESTED beats APPROVED across
|
||||
// distinct reviewers.
|
||||
func summarizeReviewState(nodes []reviewNode) string {
|
||||
type latest struct {
|
||||
state string
|
||||
at time.Time
|
||||
}
|
||||
byAuthor := map[string]latest{}
|
||||
for _, n := range nodes {
|
||||
// Anonymous reviewers (deleted users) get a unique bucket so each
|
||||
// review still counts independently.
|
||||
key := n.Author.Login
|
||||
if key == "" {
|
||||
key = "<anon>:" + n.SubmittedAt.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
state := strings.ToUpper(n.State)
|
||||
// COMMENTED and PENDING are non-binding states and shouldn't replace
|
||||
// a prior APPROVED / CHANGES_REQUESTED for the same reviewer.
|
||||
if state != reviewStateApproved && state != reviewStateChangesRequested {
|
||||
if _, ok := byAuthor[key]; ok {
|
||||
continue
|
||||
}
|
||||
}
|
||||
prev, ok := byAuthor[key]
|
||||
if !ok || !n.SubmittedAt.Before(prev.at) {
|
||||
byAuthor[key] = latest{state: state, at: n.SubmittedAt}
|
||||
}
|
||||
}
|
||||
hasApproval := false
|
||||
for _, l := range byAuthor {
|
||||
switch l.state {
|
||||
case reviewStateChangesRequested:
|
||||
return computedReviewStateChangesRequested
|
||||
case reviewStateApproved:
|
||||
hasApproval = true
|
||||
}
|
||||
}
|
||||
if hasApproval {
|
||||
return computedReviewStateApproved
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// PATClient.ExecuteGraphQL satisfies GraphQLExecutor by POSTing to /graphql.
|
||||
func (c *PATClient) ExecuteGraphQL(ctx context.Context, query string, variables map[string]any, out any) error {
|
||||
body, err := json.Marshal(map[string]any{"query": query, "variables": variables})
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal graphql: %w", err)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, githubAPIBase+"/graphql", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.setGitHubHeaders(req)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return fmt.Errorf("graphql request: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
c.recordRateHeaders(resp, "/graphql")
|
||||
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 4*1024*1024))
|
||||
if resp.StatusCode >= 400 {
|
||||
c.maybeMarkRateExhaustedFromBody("/graphql", resp.StatusCode, respBody)
|
||||
return &GitHubAPIError{StatusCode: resp.StatusCode, Endpoint: "/graphql", Body: string(respBody)}
|
||||
}
|
||||
if err := json.Unmarshal(respBody, out); err != nil {
|
||||
return fmt.Errorf("decode graphql response: %w", err)
|
||||
}
|
||||
recordGraphQLRateFromPayload(c.rateTracker, respBody)
|
||||
return nil
|
||||
}
|
||||
|
||||
// recordGraphQLRateFromPayload extracts data.rateLimit from a GraphQL
|
||||
// response body and feeds it to the tracker. The GraphQL rate-limit field is
|
||||
// point-cost based and more accurate than HTTP headers.
|
||||
func recordGraphQLRateFromPayload(tracker *RateTracker, body []byte) {
|
||||
if tracker == nil {
|
||||
return
|
||||
}
|
||||
var probe struct {
|
||||
Data map[string]json.RawMessage `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &probe); err != nil {
|
||||
return
|
||||
}
|
||||
raw, ok := probe.Data["rateLimit"]
|
||||
if !ok || len(raw) == 0 {
|
||||
return
|
||||
}
|
||||
var rl graphQLRateLimit
|
||||
if err := json.Unmarshal(raw, &rl); err != nil {
|
||||
return
|
||||
}
|
||||
tracker.Record(RateSnapshot{
|
||||
Resource: ResourceGraphQL,
|
||||
Limit: rl.Limit,
|
||||
Remaining: rl.Remaining,
|
||||
ResetAt: rl.ResetAt,
|
||||
UpdatedAt: time.Now().UTC(),
|
||||
})
|
||||
}
|
||||
|
||||
// GHClient.ExecuteGraphQL satisfies GraphQLExecutor via `gh api graphql -f query=...`.
|
||||
// Variables are passed as -F field_name=value entries; they are ignored when
|
||||
// the query has no $variables.
|
||||
func (c *GHClient) ExecuteGraphQL(ctx context.Context, query string, variables map[string]any, out any) error {
|
||||
args := []string{"api", "graphql", "-f", "query=" + query}
|
||||
for k, v := range variables {
|
||||
args = append(args, "-F", fmt.Sprintf("%s=%v", k, v))
|
||||
}
|
||||
stdout, err := c.run(ctx, args...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("gh graphql: %w", err)
|
||||
}
|
||||
if err := json.Unmarshal([]byte(stdout), out); err != nil {
|
||||
return fmt.Errorf("decode graphql response: %w", err)
|
||||
}
|
||||
recordGraphQLRateFromPayload(c.rateTracker, []byte(stdout))
|
||||
return nil
|
||||
}
|
||||
|
||||
// noopGraphQLExecutorErr is returned when the active client is the noop
|
||||
// fallback (no auth). Caller paths must handle this gracefully — typically
|
||||
// by skipping the batched call entirely.
|
||||
var errGraphQLUnsupported = fmt.Errorf("github client does not support GraphQL")
|
||||
|
||||
// graphQLExecutorFor returns an executor for the given client, or
|
||||
// errGraphQLUnsupported when the client is nil/noop.
|
||||
func graphQLExecutorFor(client Client) (GraphQLExecutor, error) {
|
||||
if exec, ok := client.(GraphQLExecutor); ok && client != nil {
|
||||
return exec, nil
|
||||
}
|
||||
return nil, errGraphQLUnsupported
|
||||
}
|
||||
|
||||
// runBatchedPRQuery executes the batched query in chunks and merges the
|
||||
// results into a single map keyed by prStatusCacheKey(owner, repo, number).
|
||||
func runBatchedPRQuery(ctx context.Context, exec GraphQLExecutor, refs []graphQLPRRef) (map[string]*PRStatus, error) {
|
||||
if exec == nil || len(refs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
result := make(map[string]*PRStatus, len(refs))
|
||||
for _, chunk := range chunkedRefs(refs) {
|
||||
query, vars := buildBatchedPRQuery(chunk)
|
||||
var resp struct {
|
||||
Data map[string]json.RawMessage `json:"data"`
|
||||
Errors []graphQLError `json:"errors"`
|
||||
}
|
||||
if err := exec.ExecuteGraphQL(ctx, query, vars, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// GitHub returns HTTP 200 with a top-level "errors" array for partial
|
||||
// auth failures, schema mismatches, or per-alias errors. Surface them
|
||||
// as an error so the caller falls back to per-watch checks instead of
|
||||
// silently absorbing the failure.
|
||||
if err := graphQLErrorsToErr(resp.Errors); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := decodeBatchedPRChunk(chunk, resp.Data, result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// decodeBatchedPRChunk maps the aliased response back to the input refs and
|
||||
// fills in the result map. Refs whose alias is missing or null are skipped
|
||||
// (e.g. PR was deleted upstream); the next poller tick will retry.
|
||||
func decodeBatchedPRChunk(refs []graphQLPRRef, data map[string]json.RawMessage, result map[string]*PRStatus) error {
|
||||
type group struct {
|
||||
idx int
|
||||
owner string
|
||||
repo string
|
||||
prRefs []graphQLPRRef
|
||||
}
|
||||
groups := []*group{}
|
||||
byKey := map[string]*group{}
|
||||
for _, r := range refs {
|
||||
k := r.Owner + "/" + r.Repo
|
||||
g, ok := byKey[k]
|
||||
if !ok {
|
||||
g = &group{idx: len(groups), owner: r.Owner, repo: r.Repo}
|
||||
byKey[k] = g
|
||||
groups = append(groups, g)
|
||||
}
|
||||
g.prRefs = append(g.prRefs, r)
|
||||
}
|
||||
keys := make([]string, 0, len(byKey))
|
||||
for k := range byKey {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
// Reassign group indices in sorted order to match buildBatchedPRQuery.
|
||||
for i, k := range keys {
|
||||
byKey[k].idx = i
|
||||
}
|
||||
for _, k := range keys {
|
||||
g := byKey[k]
|
||||
raw, ok := data[fmt.Sprintf("repo%d", g.idx)]
|
||||
if !ok || len(raw) == 0 {
|
||||
continue
|
||||
}
|
||||
repoBlock := map[string]json.RawMessage{}
|
||||
if err := json.Unmarshal(raw, &repoBlock); err != nil {
|
||||
return fmt.Errorf("decode repo block: %w", err)
|
||||
}
|
||||
for prIdx, ref := range g.prRefs {
|
||||
alias := fmt.Sprintf("%s%d", graphQLPRBatchAlias, prIdx)
|
||||
rawPR, ok := repoBlock[alias]
|
||||
if !ok || len(rawPR) == 0 || string(rawPR) == "null" {
|
||||
continue
|
||||
}
|
||||
var raw batchedPRResult
|
||||
if err := json.Unmarshal(rawPR, &raw); err != nil {
|
||||
return fmt.Errorf("decode pr alias %s: %w", alias, err)
|
||||
}
|
||||
result[prStatusCacheKey(ref.Owner, ref.Repo, ref.Number)] = convertBatchedPRResult(&raw, ref.Owner, ref.Repo, ref.Number)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// runBatchedBranchQuery executes the branch-lookup query in chunks and maps
|
||||
// each branch ref to its first OPEN associated PR (if any). Result keys are
|
||||
// "owner/repo/branch" so callers can index by their input refs.
|
||||
func runBatchedBranchQuery(ctx context.Context, exec GraphQLExecutor, refs []graphQLBranchRef) (map[string]*PRStatus, error) {
|
||||
if exec == nil || len(refs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
result := make(map[string]*PRStatus, len(refs))
|
||||
for _, chunk := range chunkedRefs(refs) {
|
||||
query, vars := buildBatchedBranchQuery(chunk)
|
||||
var resp struct {
|
||||
Data map[string]json.RawMessage `json:"data"`
|
||||
Errors []graphQLError `json:"errors"`
|
||||
}
|
||||
if err := exec.ExecuteGraphQL(ctx, query, vars, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := graphQLErrorsToErr(resp.Errors); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
decodeBatchedBranchChunk(chunk, resp.Data, result)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func decodeBatchedBranchChunk(refs []graphQLBranchRef, data map[string]json.RawMessage, result map[string]*PRStatus) {
|
||||
for i, ref := range refs {
|
||||
raw, ok := data[fmt.Sprintf("b%d", i)]
|
||||
if !ok || len(raw) == 0 || string(raw) == "null" {
|
||||
continue
|
||||
}
|
||||
var inner struct {
|
||||
Ref *struct {
|
||||
AssociatedPullRequests struct {
|
||||
Nodes []struct {
|
||||
Number int `json:"number"`
|
||||
batchedPRResult
|
||||
} `json:"nodes"`
|
||||
} `json:"associatedPullRequests"`
|
||||
} `json:"ref"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &inner); err != nil {
|
||||
continue
|
||||
}
|
||||
if inner.Ref == nil || len(inner.Ref.AssociatedPullRequests.Nodes) == 0 {
|
||||
continue
|
||||
}
|
||||
node := inner.Ref.AssociatedPullRequests.Nodes[0]
|
||||
status := convertBatchedPRResult(&node.batchedPRResult, ref.Owner, ref.Repo, node.Number)
|
||||
result[graphqlBranchKey(ref.Owner, ref.Repo, ref.Branch)] = status
|
||||
}
|
||||
}
|
||||
|
||||
// graphqlBranchKey is the lookup key used in batched-branch result maps.
|
||||
// Named graphql* to avoid collision with the mock_client.go branchKey type.
|
||||
func graphqlBranchKey(owner, repo, branch string) string {
|
||||
return owner + "/" + repo + "/" + branch
|
||||
}
|
||||
|
||||
// Compile-time assertions that the existing CLI/PAT clients satisfy
|
||||
// GraphQLExecutor. If a client stops implementing it the build fails here.
|
||||
var (
|
||||
_ GraphQLExecutor = (*PATClient)(nil)
|
||||
_ GraphQLExecutor = (*GHClient)(nil)
|
||||
)
|
||||
@@ -0,0 +1,296 @@
|
||||
package github
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// stubGraphQLExecutor lets tests inspect/return canned responses without
|
||||
// touching HTTP or the gh CLI.
|
||||
type stubGraphQLExecutor struct {
|
||||
queries []string
|
||||
response string
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *stubGraphQLExecutor) ExecuteGraphQL(_ context.Context, query string, _ map[string]any, out any) error {
|
||||
s.queries = append(s.queries, query)
|
||||
if s.err != nil {
|
||||
return s.err
|
||||
}
|
||||
return json.Unmarshal([]byte(s.response), out)
|
||||
}
|
||||
|
||||
func TestBuildBatchedPRQuery_GroupsByRepo(t *testing.T) {
|
||||
q, _ := buildBatchedPRQuery([]graphQLPRRef{
|
||||
{Owner: "octo", Repo: "alpha", Number: 1},
|
||||
{Owner: "octo", Repo: "alpha", Number: 2},
|
||||
{Owner: "octo", Repo: "beta", Number: 9},
|
||||
})
|
||||
if !strings.Contains(q, `repo0: repository(owner: "octo", name: "alpha")`) {
|
||||
t.Errorf("expected repo0 alias for octo/alpha, got: %s", q)
|
||||
}
|
||||
if !strings.Contains(q, `repo1: repository(owner: "octo", name: "beta")`) {
|
||||
t.Errorf("expected repo1 alias for octo/beta, got: %s", q)
|
||||
}
|
||||
if !strings.Contains(q, `pr0: pullRequest(number: 1)`) ||
|
||||
!strings.Contains(q, `pr1: pullRequest(number: 2)`) {
|
||||
t.Errorf("expected aliased pullRequests inside repo0: %s", q)
|
||||
}
|
||||
if !strings.Contains(q, "rateLimit") {
|
||||
t.Errorf("expected rateLimit field in query")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildBatchedBranchQuery_AliasesAllBranches(t *testing.T) {
|
||||
q, _ := buildBatchedBranchQuery([]graphQLBranchRef{
|
||||
{Owner: "o", Repo: "r", Branch: "feat-1"},
|
||||
{Owner: "o", Repo: "r", Branch: "feat-2"},
|
||||
})
|
||||
if !strings.Contains(q, `b0: repository`) || !strings.Contains(q, `b1: repository`) {
|
||||
t.Errorf("expected b0/b1 aliases: %s", q)
|
||||
}
|
||||
if !strings.Contains(q, `qualifiedName: "refs/heads/feat-1"`) {
|
||||
t.Errorf("expected qualifiedName for feat-1: %s", q)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunBatchedPRQuery_DecodesAliasesBackToRefs(t *testing.T) {
|
||||
exec := &stubGraphQLExecutor{
|
||||
response: `{
|
||||
"data": {
|
||||
"repo0": {
|
||||
"pr0": {
|
||||
"state": "OPEN", "title": "PR A", "url": "https://x/1",
|
||||
"isDraft": false, "mergeable": "MERGEABLE", "mergeStateStatus": "CLEAN",
|
||||
"headRefName": "h1", "baseRefName": "main", "headRefOid": "abc",
|
||||
"author": {"login":"alice"}, "createdAt": "2026-01-01T00:00:00Z", "updatedAt": "2026-01-02T00:00:00Z",
|
||||
"reviews": {"nodes": [{"state": "APPROVED"}]},
|
||||
"reviewRequests": {"totalCount": 0},
|
||||
"commits": {"nodes": [{"commit": {"statusCheckRollup": {"state": "SUCCESS"}}}]}
|
||||
}
|
||||
},
|
||||
"rateLimit": {"limit":5000, "remaining":4999, "resetAt":"2030-01-01T00:00:00Z", "cost":1}
|
||||
}
|
||||
}`,
|
||||
}
|
||||
got, err := runBatchedPRQuery(context.Background(), exec, []graphQLPRRef{
|
||||
{Owner: "o", Repo: "r", Number: 42},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
status, ok := got[prStatusCacheKey("o", "r", 42)]
|
||||
if !ok {
|
||||
t.Fatalf("expected status for o/r#42, got keys: %v", keysOf(got))
|
||||
}
|
||||
if status.ReviewState != "approved" {
|
||||
t.Errorf("ReviewState = %q, want approved", status.ReviewState)
|
||||
}
|
||||
if status.ChecksState != "success" {
|
||||
t.Errorf("ChecksState = %q, want success", status.ChecksState)
|
||||
}
|
||||
if status.PR == nil || status.PR.Title != "PR A" {
|
||||
t.Errorf("PR.Title mismatch: %#v", status.PR)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunBatchedPRQuery_DecodesMultipleReposIntoCorrectKeys(t *testing.T) {
|
||||
// Two repos in one batch: octo/alpha#1 and octo/beta#9. The decoder sorts
|
||||
// repo group keys; this test guards against future drift between
|
||||
// buildBatchedPRQuery's sort and decodeBatchedPRChunk's sort by checking
|
||||
// each PR lands in its correct key slot.
|
||||
exec := &stubGraphQLExecutor{
|
||||
response: `{
|
||||
"data": {
|
||||
"repo0": {
|
||||
"pr0": {
|
||||
"state": "OPEN", "title": "alpha PR", "url": "https://x/a/1",
|
||||
"isDraft": false, "mergeable": "MERGEABLE", "mergeStateStatus": "CLEAN",
|
||||
"headRefName": "h1", "baseRefName": "main", "headRefOid": "aaa",
|
||||
"author": {"login":"alice"}, "createdAt": "2026-01-01T00:00:00Z", "updatedAt": "2026-01-01T00:00:00Z",
|
||||
"reviews": {"nodes": []}, "reviewRequests": {"totalCount": 0},
|
||||
"commits": {"nodes": []}
|
||||
}
|
||||
},
|
||||
"repo1": {
|
||||
"pr0": {
|
||||
"state": "OPEN", "title": "beta PR", "url": "https://x/b/9",
|
||||
"isDraft": false, "mergeable": "MERGEABLE", "mergeStateStatus": "CLEAN",
|
||||
"headRefName": "h2", "baseRefName": "main", "headRefOid": "bbb",
|
||||
"author": {"login":"alice"}, "createdAt": "2026-01-01T00:00:00Z", "updatedAt": "2026-01-01T00:00:00Z",
|
||||
"reviews": {"nodes": []}, "reviewRequests": {"totalCount": 0},
|
||||
"commits": {"nodes": []}
|
||||
}
|
||||
}
|
||||
}
|
||||
}`,
|
||||
}
|
||||
got, err := runBatchedPRQuery(context.Background(), exec, []graphQLPRRef{
|
||||
{Owner: "octo", Repo: "alpha", Number: 1},
|
||||
{Owner: "octo", Repo: "beta", Number: 9},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
if s, ok := got[prStatusCacheKey("octo", "alpha", 1)]; !ok || s.PR == nil || s.PR.Title != "alpha PR" {
|
||||
t.Errorf("alpha#1 not decoded into correct slot: %#v", s)
|
||||
}
|
||||
if s, ok := got[prStatusCacheKey("octo", "beta", 9)]; !ok || s.PR == nil || s.PR.Title != "beta PR" {
|
||||
t.Errorf("beta#9 not decoded into correct slot: %#v", s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunBatchedPRQuery_PropagatesError(t *testing.T) {
|
||||
exec := &stubGraphQLExecutor{err: errors.New("graphql 500")}
|
||||
_, err := runBatchedPRQuery(context.Background(), exec, []graphQLPRRef{{Owner: "o", Repo: "r", Number: 1}})
|
||||
if err == nil {
|
||||
t.Fatalf("expected error to propagate")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunBatchedPRQuery_SurfacesGraphQLErrors(t *testing.T) {
|
||||
exec := &stubGraphQLExecutor{
|
||||
response: `{"data": {}, "errors": [{"message": "Field 'foo' doesn't exist on type 'Repository'"}]}`,
|
||||
}
|
||||
_, err := runBatchedPRQuery(context.Background(), exec, []graphQLPRRef{{Owner: "o", Repo: "r", Number: 1}})
|
||||
if err == nil {
|
||||
t.Fatalf("expected error from non-empty errors array")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "Field 'foo'") {
|
||||
t.Errorf("expected error to include first message, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunBatchedBranchQuery_SurfacesGraphQLErrors(t *testing.T) {
|
||||
exec := &stubGraphQLExecutor{
|
||||
response: `{"data": {}, "errors": [{"message": "rate limited"}, {"message": "secondary"}]}`,
|
||||
}
|
||||
_, err := runBatchedBranchQuery(context.Background(), exec, []graphQLBranchRef{{Owner: "o", Repo: "r", Branch: "main"}})
|
||||
if err == nil {
|
||||
t.Fatalf("expected error from non-empty errors array")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "rate limited") || !strings.Contains(err.Error(), "and 1 more") {
|
||||
t.Errorf("expected error to include first message + count, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunBatchedBranchQuery_DecodesPRNode(t *testing.T) {
|
||||
exec := &stubGraphQLExecutor{
|
||||
response: `{
|
||||
"data": {
|
||||
"b0": {
|
||||
"ref": {
|
||||
"associatedPullRequests": {
|
||||
"nodes": [{
|
||||
"number": 7,
|
||||
"state": "OPEN", "title": "branch PR", "url": "https://x/7",
|
||||
"isDraft": false, "mergeable": "MERGEABLE", "mergeStateStatus": "CLEAN",
|
||||
"headRefName": "feat", "baseRefName": "main", "headRefOid": "deadbeef",
|
||||
"author": {"login":"alice"},
|
||||
"createdAt": "2026-01-01T00:00:00Z", "updatedAt": "2026-01-01T00:00:00Z",
|
||||
"reviews": {"nodes": []}, "reviewRequests": {"totalCount": 0},
|
||||
"commits": {"nodes": []}
|
||||
}]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}`,
|
||||
}
|
||||
got, err := runBatchedBranchQuery(context.Background(), exec, []graphQLBranchRef{
|
||||
{Owner: "o", Repo: "r", Branch: "feat"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
status, ok := got[graphqlBranchKey("o", "r", "feat")]
|
||||
if !ok {
|
||||
t.Fatalf("expected branch result, got keys: %v", keysOf(got))
|
||||
}
|
||||
if status.PR == nil || status.PR.Number != 7 {
|
||||
t.Errorf("expected PR number 7, got %#v", status.PR)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSummarizeReviewState_PrefersChangesRequested(t *testing.T) {
|
||||
mk := func(login, state string, sec int) reviewNode {
|
||||
var n reviewNode
|
||||
n.Author.Login = login
|
||||
n.State = state
|
||||
n.SubmittedAt = time.Unix(int64(sec), 0).UTC()
|
||||
return n
|
||||
}
|
||||
// Different reviewers: CHANGES_REQUESTED beats APPROVED.
|
||||
if got := summarizeReviewState([]reviewNode{
|
||||
mk("alice", "APPROVED", 1),
|
||||
mk("bob", "CHANGES_REQUESTED", 2),
|
||||
}); got != "changes_requested" {
|
||||
t.Errorf("got %q", got)
|
||||
}
|
||||
// Different reviewers: APPROVED + COMMENTED -> approved.
|
||||
if got := summarizeReviewState([]reviewNode{
|
||||
mk("alice", "APPROVED", 1),
|
||||
mk("bob", "COMMENTED", 2),
|
||||
}); got != "approved" {
|
||||
t.Errorf("got %q", got)
|
||||
}
|
||||
// Single COMMENTED -> empty.
|
||||
if got := summarizeReviewState([]reviewNode{
|
||||
mk("alice", "COMMENTED", 1),
|
||||
}); got != "" {
|
||||
t.Errorf("got %q", got)
|
||||
}
|
||||
// Same reviewer: CHANGES_REQUESTED then APPROVED -> approved (Greptile P1).
|
||||
if got := summarizeReviewState([]reviewNode{
|
||||
mk("alice", "CHANGES_REQUESTED", 1),
|
||||
mk("alice", "APPROVED", 2),
|
||||
}); got != "approved" {
|
||||
t.Errorf("dedup: got %q, want approved", got)
|
||||
}
|
||||
// Same reviewer: APPROVED then CHANGES_REQUESTED -> changes_requested.
|
||||
if got := summarizeReviewState([]reviewNode{
|
||||
mk("alice", "APPROVED", 1),
|
||||
mk("alice", "CHANGES_REQUESTED", 2),
|
||||
}); got != "changes_requested" {
|
||||
t.Errorf("dedup reverse: got %q, want changes_requested", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkedRefs_RespectsBatchSize(t *testing.T) {
|
||||
refs := make([]graphQLPRRef, graphQLBatchChunkSize+5)
|
||||
chunks := chunkedRefs(refs)
|
||||
if len(chunks) != 2 {
|
||||
t.Fatalf("expected 2 chunks for %d refs, got %d", len(refs), len(chunks))
|
||||
}
|
||||
if len(chunks[0]) != graphQLBatchChunkSize {
|
||||
t.Errorf("first chunk size = %d, want %d", len(chunks[0]), graphQLBatchChunkSize)
|
||||
}
|
||||
if len(chunks[1]) != 5 {
|
||||
t.Errorf("second chunk size = %d, want 5", len(chunks[1]))
|
||||
}
|
||||
}
|
||||
|
||||
func keysOf[K comparable, V any](m map[K]V) []K {
|
||||
out := make([]K, 0, len(m))
|
||||
for k := range m {
|
||||
out = append(out, k)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestGraphQLExecutorFor_NoopReturnsError(t *testing.T) {
|
||||
if _, err := graphQLExecutorFor(&NoopClient{}); err == nil {
|
||||
t.Fatalf("expected unsupported error for NoopClient")
|
||||
}
|
||||
if _, err := graphQLExecutorFor(NewPATClient("token")); err != nil {
|
||||
t.Fatalf("PATClient should satisfy GraphQLExecutor: %v", err)
|
||||
}
|
||||
if _, err := graphQLExecutorFor(NewGHClient()); err != nil {
|
||||
t.Fatalf("GHClient should satisfy GraphQLExecutor: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -234,13 +234,22 @@ type RepoBranch struct {
|
||||
|
||||
// GitHubStatus represents GitHub connection status.
|
||||
type GitHubStatus struct {
|
||||
Authenticated bool `json:"authenticated"`
|
||||
Username string `json:"username"`
|
||||
AuthMethod string `json:"auth_method"` // "gh_cli", "pat", "none"
|
||||
TokenConfigured bool `json:"token_configured"`
|
||||
TokenSecretID string `json:"token_secret_id,omitempty"`
|
||||
RequiredScopes []string `json:"required_scopes"`
|
||||
Diagnostics *AuthDiagnostics `json:"diagnostics,omitempty"`
|
||||
Authenticated bool `json:"authenticated"`
|
||||
Username string `json:"username"`
|
||||
AuthMethod string `json:"auth_method"` // "gh_cli", "pat", "none"
|
||||
TokenConfigured bool `json:"token_configured"`
|
||||
TokenSecretID string `json:"token_secret_id,omitempty"`
|
||||
RequiredScopes []string `json:"required_scopes"`
|
||||
Diagnostics *AuthDiagnostics `json:"diagnostics,omitempty"`
|
||||
RateLimit *GitHubRateLimitInfo `json:"rate_limit,omitempty"`
|
||||
}
|
||||
|
||||
// GitHubRateLimitInfo bundles known rate-limit snapshots per resource bucket
|
||||
// for surfacing in the UI.
|
||||
type GitHubRateLimitInfo struct {
|
||||
Core *RateSnapshot `json:"core,omitempty"`
|
||||
GraphQL *RateSnapshot `json:"graphql,omitempty"`
|
||||
Search *RateSnapshot `json:"search,omitempty"`
|
||||
}
|
||||
|
||||
// ConfigureTokenRequest is the request body for configuring a GitHub token.
|
||||
|
||||
@@ -31,9 +31,10 @@ func (e *GitHubAPIError) Error() string {
|
||||
|
||||
// PATClient implements Client using a GitHub Personal Access Token.
|
||||
type PATClient struct {
|
||||
token string
|
||||
httpClient *http.Client
|
||||
username string // cached after first GetAuthenticatedUser call
|
||||
token string
|
||||
httpClient *http.Client
|
||||
username string // cached after first GetAuthenticatedUser call
|
||||
rateTracker *RateTracker
|
||||
}
|
||||
|
||||
// NewPATClient creates a new PAT-based GitHub client.
|
||||
@@ -46,6 +47,81 @@ func NewPATClient(token string) *PATClient {
|
||||
}
|
||||
}
|
||||
|
||||
// WithRateTracker attaches a rate tracker so response headers are recorded.
|
||||
// Returns the client for chaining; safe to call before any requests are made.
|
||||
func (c *PATClient) WithRateTracker(t *RateTracker) *PATClient {
|
||||
c.rateTracker = t
|
||||
return c
|
||||
}
|
||||
|
||||
// recordRateHeaders feeds rate-limit data from a response into the tracker.
|
||||
// endpoint is used to pick the default resource bucket when the response
|
||||
// omits the X-RateLimit-Resource header.
|
||||
func (c *PATClient) recordRateHeaders(resp *http.Response, endpoint string) {
|
||||
if c.rateTracker == nil || resp == nil {
|
||||
return
|
||||
}
|
||||
defaultResource := ResourceCore
|
||||
if strings.HasPrefix(endpoint, "/search/") {
|
||||
defaultResource = ResourceSearch
|
||||
} else if strings.HasPrefix(endpoint, "/graphql") {
|
||||
defaultResource = ResourceGraphQL
|
||||
}
|
||||
snap, headersOK := parseRateHeaders(resp, defaultResource)
|
||||
if headersOK {
|
||||
c.rateTracker.Record(snap)
|
||||
}
|
||||
if !isRateLimitStatus(resp.StatusCode) {
|
||||
return
|
||||
}
|
||||
// On a 429, prefer the real X-RateLimit-Reset over the synthetic 1h
|
||||
// fallback. parseRateHeaders may return ok=true with Remaining>0 if the
|
||||
// secondary-limit response carried stale headers — only skip the
|
||||
// fallback when the headers themselves report exhaustion (Remaining<=0
|
||||
// + future ResetAt). Otherwise fall through to the conservative pause.
|
||||
if headersOK && snap.Exhausted() {
|
||||
return
|
||||
}
|
||||
c.rateTracker.markRateExhausted(defaultResource, time.Time{})
|
||||
}
|
||||
|
||||
// isRateLimitStatus returns true for status codes GitHub uses to signal
|
||||
// primary or secondary rate-limit exhaustion. 403 is documented for both
|
||||
// abuse-detection and primary limits when the body indicates so; 429 is
|
||||
// secondary limits.
|
||||
func isRateLimitStatus(status int) bool {
|
||||
return status == http.StatusTooManyRequests
|
||||
}
|
||||
|
||||
// maybeMarkRateExhaustedFromBody flags a rate-limit hit when GitHub returned
|
||||
// a 403/429 whose body contains the rate-limit prose. The headers may be
|
||||
// missing on these responses (esp. secondary limits), so the body is the
|
||||
// authoritative signal.
|
||||
func (c *PATClient) maybeMarkRateExhaustedFromBody(endpoint string, status int, body []byte) {
|
||||
if c.rateTracker == nil {
|
||||
return
|
||||
}
|
||||
if status != http.StatusForbidden && status != http.StatusTooManyRequests {
|
||||
return
|
||||
}
|
||||
lower := strings.ToLower(string(body))
|
||||
if !strings.Contains(lower, "rate limit") && !strings.Contains(lower, "abuse detection") {
|
||||
return
|
||||
}
|
||||
resource := ResourceCore
|
||||
if strings.HasPrefix(endpoint, "/search/") {
|
||||
resource = ResourceSearch
|
||||
} else if strings.HasPrefix(endpoint, "/graphql") {
|
||||
resource = ResourceGraphQL
|
||||
}
|
||||
// If recordRateHeaders already captured a real reset for this bucket on
|
||||
// the same response, don't clobber it with the synthetic 1h fallback.
|
||||
if existing, ok := c.rateTracker.Snapshot(resource); ok && existing.Exhausted() {
|
||||
return
|
||||
}
|
||||
c.rateTracker.markRateExhausted(resource, time.Time{})
|
||||
}
|
||||
|
||||
// setGitHubHeaders sets the common Authorization, Accept, and API version headers.
|
||||
func (c *PATClient) setGitHubHeaders(req *http.Request) {
|
||||
req.Header.Set("Authorization", "token "+c.token)
|
||||
@@ -390,9 +466,11 @@ func (c *PATClient) post(ctx context.Context, endpoint string, body []byte) erro
|
||||
return fmt.Errorf("request POST %s: %w", endpoint, err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
c.recordRateHeaders(resp, endpoint)
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
c.maybeMarkRateExhaustedFromBody(endpoint, resp.StatusCode, respBody)
|
||||
return fmt.Errorf("GitHub API POST %s returned %d: %s", endpoint, resp.StatusCode, string(respBody))
|
||||
}
|
||||
return nil
|
||||
@@ -411,9 +489,11 @@ func (c *PATClient) get(ctx context.Context, endpoint string, result interface{}
|
||||
return fmt.Errorf("request %s: %w", endpoint, err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
c.recordRateHeaders(resp, endpoint)
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
c.maybeMarkRateExhaustedFromBody(endpoint, resp.StatusCode, body)
|
||||
return &GitHubAPIError{StatusCode: resp.StatusCode, Endpoint: endpoint, Body: string(body)}
|
||||
}
|
||||
return json.NewDecoder(resp.Body).Decode(result)
|
||||
@@ -433,9 +513,11 @@ func (c *PATClient) getPaginated(ctx context.Context, endpoint string, result in
|
||||
return "", fmt.Errorf("request %s: %w", endpoint, err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
c.recordRateHeaders(resp, endpoint)
|
||||
|
||||
if resp.StatusCode >= 400 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
c.maybeMarkRateExhaustedFromBody(endpoint, resp.StatusCode, body)
|
||||
return "", &GitHubAPIError{StatusCode: resp.StatusCode, Endpoint: endpoint, Body: string(body)}
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(result); err != nil {
|
||||
|
||||
@@ -1,6 +1,11 @@
|
||||
package github
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
@@ -180,3 +185,138 @@ func TestConvertPatRequestedReviewers(t *testing.T) {
|
||||
t.Errorf("unexpected third reviewer: %#v", got[2])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPATClient_RecordsRateHeadersOnSuccess(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("X-RateLimit-Limit", "5000")
|
||||
w.Header().Set("X-RateLimit-Remaining", "4998")
|
||||
w.Header().Set("X-RateLimit-Reset", "2000000000")
|
||||
w.Header().Set("X-RateLimit-Resource", "core")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte(`{"login":"octocat"}`))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
c := newPATClientPointingAt(t, srv.URL)
|
||||
tracker := NewRateTracker(nil, nil)
|
||||
c.WithRateTracker(tracker)
|
||||
|
||||
var out struct {
|
||||
Login string `json:"login"`
|
||||
}
|
||||
if err := c.get(context.Background(), "/user", &out); err != nil {
|
||||
t.Fatalf("get: %v", err)
|
||||
}
|
||||
snap, ok := tracker.Snapshot(ResourceCore)
|
||||
if !ok {
|
||||
t.Fatalf("expected core snapshot")
|
||||
}
|
||||
if snap.Remaining != 4998 || snap.Limit != 5000 {
|
||||
t.Fatalf("snap = %+v", snap)
|
||||
}
|
||||
}
|
||||
|
||||
// Regression: when a 429 carries valid X-RateLimit-Reset headers, the reset
|
||||
// time from the headers must win — not the synthetic +1h fallback in
|
||||
// markRateExhausted. Previously, recordRateHeaders called Record(snap)
|
||||
// followed by markRateExhausted(time.Time{}), and the second call clobbered
|
||||
// the real reset with a 1-hour pause that could over-throttle the poller.
|
||||
func TestPATClient_RateLimit429_PreservesRealReset(t *testing.T) {
|
||||
realReset := time.Now().Add(7 * time.Minute).UTC().Truncate(time.Second)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("X-RateLimit-Limit", "5000")
|
||||
w.Header().Set("X-RateLimit-Remaining", "0")
|
||||
w.Header().Set("X-RateLimit-Reset", strconv.FormatInt(realReset.Unix(), 10))
|
||||
w.Header().Set("X-RateLimit-Resource", "core")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_, _ = w.Write([]byte(`{"message":"API rate limit exceeded"}`))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
c := newPATClientPointingAt(t, srv.URL)
|
||||
tracker := NewRateTracker(nil, nil)
|
||||
c.WithRateTracker(tracker)
|
||||
|
||||
var out struct{}
|
||||
if err := c.get(context.Background(), "/repos/o/r/pulls/1", &out); err == nil {
|
||||
t.Fatalf("expected error from 429")
|
||||
}
|
||||
snap, ok := tracker.Snapshot(ResourceCore)
|
||||
if !ok {
|
||||
t.Fatalf("expected core snapshot")
|
||||
}
|
||||
if !snap.ResetAt.Equal(realReset) {
|
||||
t.Errorf("expected reset_at preserved from headers (%s), got %s (off by %s)",
|
||||
realReset, snap.ResetAt, snap.ResetAt.Sub(realReset))
|
||||
}
|
||||
if !tracker.IsExhausted(ResourceCore) {
|
||||
t.Errorf("expected core to be exhausted")
|
||||
}
|
||||
}
|
||||
|
||||
// When a 429 has no rate-limit headers, the synthetic 1h fallback still
|
||||
// applies so the poller pauses instead of hammering the secondary limit.
|
||||
func TestPATClient_RateLimit429_NoHeaders_UsesSyntheticReset(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_, _ = w.Write([]byte(`{"message":"abuse detection"}`))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
c := newPATClientPointingAt(t, srv.URL)
|
||||
tracker := NewRateTracker(nil, nil)
|
||||
c.WithRateTracker(tracker)
|
||||
|
||||
var out struct{}
|
||||
if err := c.get(context.Background(), "/repos/o/r/pulls/1", &out); err == nil {
|
||||
t.Fatalf("expected error from 429")
|
||||
}
|
||||
if !tracker.IsExhausted(ResourceCore) {
|
||||
t.Fatalf("expected core exhausted via synthetic fallback")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPATClient_MarksExhaustedFromRateLimitBody(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// No headers — secondary limits sometimes omit them entirely.
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
_, _ = w.Write([]byte(`{"message":"API rate limit exceeded for user."}`))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
c := newPATClientPointingAt(t, srv.URL)
|
||||
tracker := NewRateTracker(nil, nil)
|
||||
c.WithRateTracker(tracker)
|
||||
|
||||
var out struct{}
|
||||
if err := c.get(context.Background(), "/repos/o/r/pulls/1", &out); err == nil {
|
||||
t.Fatalf("expected error from 403")
|
||||
}
|
||||
if !tracker.IsExhausted(ResourceCore) {
|
||||
t.Fatalf("expected core exhausted from body parse")
|
||||
}
|
||||
}
|
||||
|
||||
// newPATClientPointingAt builds a PATClient whose underlying HTTP client
|
||||
// reroutes any github API URL to the given test server.
|
||||
func newPATClientPointingAt(t *testing.T, baseURL string) *PATClient {
|
||||
t.Helper()
|
||||
c := NewPATClient("test-token")
|
||||
c.httpClient = &http.Client{
|
||||
Transport: &rewriteTransport{base: baseURL},
|
||||
Timeout: 2 * time.Second,
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
type rewriteTransport struct{ base string }
|
||||
|
||||
func (rt *rewriteTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
rewritten := strings.Replace(req.URL.String(), githubAPIBase, rt.base, 1)
|
||||
req2, err := http.NewRequestWithContext(req.Context(), req.Method, rewritten, req.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req2.Header = req.Header.Clone()
|
||||
return http.DefaultTransport.RoundTrip(req2)
|
||||
}
|
||||
|
||||
@@ -13,11 +13,59 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
defaultPRPollInterval = 30 * time.Second
|
||||
defaultPRPollInterval = 1 * time.Minute
|
||||
defaultReviewPollInterval = 5 * time.Minute
|
||||
defaultIssuePollInterval = 5 * time.Minute
|
||||
// rateLimitedSleepCap bounds the rate-limit sleep so a misreported reset
|
||||
// (e.g. far-future reset_at after a CLI failure) cannot wedge the loop.
|
||||
rateLimitedSleepCap = 10 * time.Minute
|
||||
)
|
||||
|
||||
// searchBucketExhausted reports whether the search bucket is currently
|
||||
// exhausted. Used by per-watch loops to bail out mid-cycle once a single
|
||||
// search request has tripped the limit, so the remaining watches don't issue
|
||||
// doomed searches that worsen secondary-limit penalties. Returns false when
|
||||
// the tracker is unavailable (tests, disabled feature).
|
||||
func (p *Poller) searchBucketExhausted(loop string) bool {
|
||||
if p.service == nil || p.service.rateTracker == nil {
|
||||
return false
|
||||
}
|
||||
if p.service.rateTracker.WaitDuration(ResourceSearch) <= 0 {
|
||||
return false
|
||||
}
|
||||
p.logger.Debug("search bucket exhausted; stopping watch loop early",
|
||||
zap.String("loop", loop))
|
||||
return true
|
||||
}
|
||||
|
||||
// waitForRateLimit returns true after sleeping the loop until the relevant
|
||||
// resource bucket has quota again, or false if ctx was cancelled. When the
|
||||
// bucket is healthy the call is a no-op and returns true immediately.
|
||||
func (p *Poller) waitForRateLimit(ctx context.Context, resource Resource, loop string) bool {
|
||||
if p.service == nil || p.service.rateTracker == nil {
|
||||
return true
|
||||
}
|
||||
d := p.service.rateTracker.WaitDuration(resource)
|
||||
if d <= 0 {
|
||||
return true
|
||||
}
|
||||
if d > rateLimitedSleepCap {
|
||||
d = rateLimitedSleepCap
|
||||
}
|
||||
p.logger.Info("github rate-limited; pausing poller loop",
|
||||
zap.String("loop", loop),
|
||||
zap.String("resource", string(resource)),
|
||||
zap.Duration("sleep", d))
|
||||
timer := time.NewTimer(d)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case <-timer.C:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// TaskBranchInfo describes a task+session that may need a PR watch.
|
||||
// RepositoryID scopes multi-repo tasks: each repo is a separate watch row.
|
||||
// Empty for legacy single-repo tasks.
|
||||
@@ -94,7 +142,9 @@ func (p *Poller) prMonitorLoop(ctx context.Context) {
|
||||
defer p.wg.Done()
|
||||
|
||||
// Run an initial check immediately so existing watches are evaluated on startup.
|
||||
p.checkPRWatches(ctx)
|
||||
if p.waitForRateLimit(ctx, ResourceGraphQL, "pr_monitor") {
|
||||
p.checkPRWatches(ctx)
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(defaultPRPollInterval)
|
||||
defer ticker.Stop()
|
||||
@@ -104,6 +154,9 @@ func (p *Poller) prMonitorLoop(ctx context.Context) {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if !p.waitForRateLimit(ctx, ResourceGraphQL, "pr_monitor") {
|
||||
return
|
||||
}
|
||||
p.checkPRWatches(ctx)
|
||||
}
|
||||
}
|
||||
@@ -117,11 +170,168 @@ func (p *Poller) checkPRWatches(ctx context.Context) {
|
||||
p.logger.Error("failed to list PR watches", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(watches) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Try the batched GraphQL path first — collapses N HTTP requests into a
|
||||
// handful of multi-aliased GraphQL calls. On error or unsupported client
|
||||
// (e.g. NoopClient), fall back to per-watch checks so one bad poll cycle
|
||||
// doesn't lose state.
|
||||
if p.tryBatchedPRWatchCheck(ctx, watches) {
|
||||
return
|
||||
}
|
||||
for _, watch := range watches {
|
||||
p.checkSinglePRWatch(ctx, watch)
|
||||
}
|
||||
}
|
||||
|
||||
// tryBatchedPRWatchCheck runs the batched GraphQL flow for the supplied
|
||||
// watches. Returns true when the path succeeded so the caller can skip the
|
||||
// per-watch fallback.
|
||||
func (p *Poller) tryBatchedPRWatchCheck(ctx context.Context, watches []*PRWatch) bool {
|
||||
exec, err := graphQLExecutorFor(p.service.client)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
numbered, searching := splitPRWatches(watches)
|
||||
statusByKey, ok := p.fetchBatchedStatuses(ctx, exec, numbered, searching)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
for _, w := range numbered {
|
||||
key := prStatusCacheKey(w.Owner, w.Repo, w.PRNumber)
|
||||
status, ok := statusByKey[key]
|
||||
if !ok {
|
||||
// Alias missing — PR may have been deleted; fall through to the
|
||||
// per-watch path which already handles the not-found case.
|
||||
p.checkSinglePRWatch(ctx, w)
|
||||
continue
|
||||
}
|
||||
p.applyPRStatus(ctx, w, status)
|
||||
}
|
||||
for _, w := range searching {
|
||||
key := graphqlBranchKey(w.Owner, w.Repo, w.Branch)
|
||||
status, ok := statusByKey[key]
|
||||
if !ok || status == nil || status.PR == nil {
|
||||
// No PR yet for this branch — record the timestamp like the
|
||||
// per-watch path so liveness updates.
|
||||
now := time.Now().UTC()
|
||||
_ = p.service.store.UpdatePRWatchTimestamps(ctx, w.ID, now, nil, "", "")
|
||||
continue
|
||||
}
|
||||
p.applyDetectedPR(ctx, w, status.PR)
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// fetchBatchedStatuses runs both the numbered-PR query and the branch query
|
||||
// against GraphQL. The combined map keys numbered watches by
|
||||
// prStatusCacheKey and searching watches by graphqlBranchKey, so callers can
|
||||
// look up either kind in one place. Returns (nil, false) when any required
|
||||
// query failed so the caller falls back to per-watch checks rather than
|
||||
// silently absorbing the failure as "no result".
|
||||
func (p *Poller) fetchBatchedStatuses(
|
||||
ctx context.Context, exec GraphQLExecutor, numbered, searching []*PRWatch,
|
||||
) (map[string]*PRStatus, bool) {
|
||||
combined := make(map[string]*PRStatus)
|
||||
|
||||
if len(numbered) > 0 {
|
||||
refs := make([]graphQLPRRef, 0, len(numbered))
|
||||
for _, w := range numbered {
|
||||
refs = append(refs, graphQLPRRef{Owner: w.Owner, Repo: w.Repo, Number: w.PRNumber})
|
||||
}
|
||||
out, err := runBatchedPRQuery(ctx, exec, refs)
|
||||
if err != nil {
|
||||
p.logger.Debug("batched PR status query failed", zap.Error(err))
|
||||
return nil, false
|
||||
}
|
||||
for k, v := range out {
|
||||
combined[k] = v
|
||||
}
|
||||
}
|
||||
if len(searching) > 0 {
|
||||
refs := make([]graphQLBranchRef, 0, len(searching))
|
||||
for _, w := range searching {
|
||||
refs = append(refs, graphQLBranchRef{Owner: w.Owner, Repo: w.Repo, Branch: w.Branch})
|
||||
}
|
||||
out, err := runBatchedBranchQuery(ctx, exec, refs)
|
||||
if err != nil {
|
||||
p.logger.Debug("batched branch query failed", zap.Error(err))
|
||||
return nil, false
|
||||
}
|
||||
for k, v := range out {
|
||||
combined[k] = v
|
||||
}
|
||||
}
|
||||
return combined, true
|
||||
}
|
||||
|
||||
// splitPRWatches partitions watches into "we know the PR number" and "still
|
||||
// searching for one on this branch" buckets so each gets its own batched
|
||||
// query shape.
|
||||
func splitPRWatches(watches []*PRWatch) (numbered, searching []*PRWatch) {
|
||||
for _, w := range watches {
|
||||
if w.PRNumber == 0 {
|
||||
searching = append(searching, w)
|
||||
} else {
|
||||
numbered = append(numbered, w)
|
||||
}
|
||||
}
|
||||
return numbered, searching
|
||||
}
|
||||
|
||||
// applyPRStatus is the post-fetch processing extracted from
|
||||
// checkSinglePRWatch so the batched path can reuse it.
|
||||
func (p *Poller) applyPRStatus(ctx context.Context, watch *PRWatch, status *PRStatus) {
|
||||
if status == nil {
|
||||
return
|
||||
}
|
||||
hasNew := status.ChecksState != watch.LastCheckStatus || status.ReviewState != watch.LastReviewState
|
||||
now := time.Now().UTC()
|
||||
if err := p.service.store.UpdatePRWatchTimestamps(ctx, watch.ID, now, nil, status.ChecksState, status.ReviewState); err != nil {
|
||||
p.logger.Error("failed to update PR watch timestamps", zap.String("id", watch.ID), zap.Error(err))
|
||||
}
|
||||
if syncErr := p.service.SyncTaskPR(ctx, watch.TaskID, status); syncErr != nil {
|
||||
p.logger.Error("failed to sync task PR",
|
||||
zap.String("task_id", watch.TaskID), zap.Error(syncErr))
|
||||
return
|
||||
}
|
||||
if status.PR != nil && (status.PR.State == prStateMerged || status.PR.State == prStateClosed) {
|
||||
p.publishPRStatusEvent(ctx, watch, status)
|
||||
if resetErr := p.service.store.UpdatePRWatchPRNumber(ctx, watch.ID, 0); resetErr != nil {
|
||||
p.logger.Error("failed to reset completed PR watch",
|
||||
zap.String("id", watch.ID), zap.Error(resetErr))
|
||||
}
|
||||
return
|
||||
}
|
||||
if hasNew {
|
||||
p.publishPRStatusEvent(ctx, watch, status)
|
||||
}
|
||||
}
|
||||
|
||||
// applyDetectedPR is the searching-watch counterpart to applyPRStatus —
|
||||
// promotes a PRWatch from "searching" to a known PR number and records the
|
||||
// task association.
|
||||
func (p *Poller) applyDetectedPR(ctx context.Context, watch *PRWatch, pr *PR) {
|
||||
now := time.Now().UTC()
|
||||
_ = p.service.store.UpdatePRWatchTimestamps(ctx, watch.ID, now, nil, "", "")
|
||||
if err := p.service.store.UpdatePRWatchPRNumber(ctx, watch.ID, pr.Number); err != nil {
|
||||
p.logger.Error("failed to update PR watch with detected PR",
|
||||
zap.String("watch_id", watch.ID), zap.Int("pr_number", pr.Number), zap.Error(err))
|
||||
return
|
||||
}
|
||||
if _, err := p.service.AssociatePRWithTask(ctx, watch.TaskID, watch.RepositoryID, pr); err != nil {
|
||||
p.logger.Error("failed to associate detected PR with task",
|
||||
zap.String("task_id", watch.TaskID), zap.Int("pr_number", pr.Number), zap.Error(err))
|
||||
return
|
||||
}
|
||||
p.logger.Info("detected PR for session branch (batched)",
|
||||
zap.String("watch_id", watch.ID),
|
||||
zap.String("branch", watch.Branch),
|
||||
zap.Int("pr_number", pr.Number))
|
||||
}
|
||||
|
||||
func (p *Poller) checkSinglePRWatch(ctx context.Context, watch *PRWatch) {
|
||||
// PRWatch with pr_number=0 means we're still searching for a PR on this branch.
|
||||
if watch.PRNumber == 0 {
|
||||
@@ -302,8 +512,9 @@ func (p *Poller) refreshStaleBranches(ctx context.Context) {
|
||||
func (p *Poller) reviewQueueLoop(ctx context.Context) {
|
||||
defer p.wg.Done()
|
||||
|
||||
// Run an initial check immediately so existing watches are evaluated on startup.
|
||||
p.checkReviewWatches(ctx)
|
||||
if p.waitForRateLimit(ctx, ResourceSearch, "review_queue") {
|
||||
p.checkReviewWatches(ctx)
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(defaultReviewPollInterval)
|
||||
defer ticker.Stop()
|
||||
@@ -313,6 +524,9 @@ func (p *Poller) reviewQueueLoop(ctx context.Context) {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if !p.waitForRateLimit(ctx, ResourceSearch, "review_queue") {
|
||||
return
|
||||
}
|
||||
p.checkReviewWatches(ctx)
|
||||
}
|
||||
}
|
||||
@@ -329,6 +543,12 @@ func (p *Poller) checkReviewWatches(ctx context.Context) {
|
||||
}
|
||||
p.logger.Debug("checking review watches", zap.Int("count", len(watches)))
|
||||
for _, watch := range watches {
|
||||
// A previous iteration in this same cycle may have exhausted the
|
||||
// search bucket. Stop early instead of issuing another doomed search
|
||||
// that would only deepen the secondary-limit penalty.
|
||||
if p.searchBucketExhausted("review_watch") {
|
||||
return
|
||||
}
|
||||
p.logger.Debug("polling review watch",
|
||||
zap.String("watch_id", watch.ID),
|
||||
zap.String("workspace_id", watch.WorkspaceID),
|
||||
@@ -368,8 +588,9 @@ func (p *Poller) checkReviewWatches(ctx context.Context) {
|
||||
func (p *Poller) issueWatchLoop(ctx context.Context) {
|
||||
defer p.wg.Done()
|
||||
|
||||
// Run an initial check immediately so existing watches are evaluated on startup.
|
||||
p.checkIssueWatches(ctx)
|
||||
if p.waitForRateLimit(ctx, ResourceSearch, "issue_watch") {
|
||||
p.checkIssueWatches(ctx)
|
||||
}
|
||||
|
||||
ticker := time.NewTicker(defaultIssuePollInterval)
|
||||
defer ticker.Stop()
|
||||
@@ -379,6 +600,9 @@ func (p *Poller) issueWatchLoop(ctx context.Context) {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if !p.waitForRateLimit(ctx, ResourceSearch, "issue_watch") {
|
||||
return
|
||||
}
|
||||
p.checkIssueWatches(ctx)
|
||||
}
|
||||
}
|
||||
@@ -395,6 +619,9 @@ func (p *Poller) checkIssueWatches(ctx context.Context) {
|
||||
}
|
||||
p.logger.Debug("checking issue watches", zap.Int("count", len(watches)))
|
||||
for _, watch := range watches {
|
||||
if p.searchBucketExhausted("issue_watch") {
|
||||
return
|
||||
}
|
||||
p.logger.Debug("polling issue watch",
|
||||
zap.String("watch_id", watch.ID),
|
||||
zap.String("workspace_id", watch.WorkspaceID),
|
||||
|
||||
@@ -3,6 +3,9 @@ package github
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -533,3 +536,336 @@ func TestRefreshStaleBranches_SkipsWhenResolverReturnsEmpty(t *testing.T) {
|
||||
t.Errorf("expected branch unchanged when resolver returns empty, got %q", updated.Branch)
|
||||
}
|
||||
}
|
||||
|
||||
// graphQLMockClient wraps MockClient with a canned GraphQL executor so tests
|
||||
// exercise the batched poll path without hitting the network.
|
||||
type graphQLMockClient struct {
|
||||
*MockClient
|
||||
prResponses []string // FIFO; one entry consumed per ExecuteGraphQL call carrying "Batch"
|
||||
branchResponses []string // FIFO; consumed for "Branches" queries
|
||||
prErr error // returned for the next "Batch" call
|
||||
branchErr error // returned for the next "Branches" call
|
||||
prQueries []string
|
||||
branchQueries []string
|
||||
}
|
||||
|
||||
func (m *graphQLMockClient) ExecuteGraphQL(_ context.Context, query string, _ map[string]any, out any) error {
|
||||
if strings.Contains(query, "query Branches") {
|
||||
m.branchQueries = append(m.branchQueries, query)
|
||||
if m.branchErr != nil {
|
||||
return m.branchErr
|
||||
}
|
||||
if len(m.branchResponses) == 0 {
|
||||
return errors.New("no canned branch response")
|
||||
}
|
||||
resp := m.branchResponses[0]
|
||||
m.branchResponses = m.branchResponses[1:]
|
||||
return json.Unmarshal([]byte(resp), out)
|
||||
}
|
||||
m.prQueries = append(m.prQueries, query)
|
||||
if m.prErr != nil {
|
||||
return m.prErr
|
||||
}
|
||||
if len(m.prResponses) == 0 {
|
||||
return errors.New("no canned PR response")
|
||||
}
|
||||
resp := m.prResponses[0]
|
||||
m.prResponses = m.prResponses[1:]
|
||||
return json.Unmarshal([]byte(resp), out)
|
||||
}
|
||||
|
||||
func setupBatchedPollerTest(t *testing.T) (*Poller, *Service, *graphQLMockClient, *Store) {
|
||||
t.Helper()
|
||||
poller, svc, mockClient, store := setupPollerTest(t)
|
||||
wrapped := &graphQLMockClient{MockClient: mockClient}
|
||||
svc.client = wrapped
|
||||
return poller, svc, wrapped, store
|
||||
}
|
||||
|
||||
func TestTryBatchedPRWatchCheck_NumberedWatch_AppliesStatus(t *testing.T) {
|
||||
poller, _, gh, store := setupBatchedPollerTest(t)
|
||||
ctx := context.Background()
|
||||
|
||||
gh.prResponses = []string{`{
|
||||
"data": {
|
||||
"repo0": {
|
||||
"pr0": {
|
||||
"state": "OPEN", "title": "Test PR", "url": "https://x/1",
|
||||
"isDraft": false, "mergeable": "MERGEABLE", "mergeStateStatus": "CLEAN",
|
||||
"headRefName": "feat", "baseRefName": "main", "headRefOid": "abc",
|
||||
"author": {"login": "alice"},
|
||||
"createdAt": "2026-01-01T00:00:00Z", "updatedAt": "2026-01-02T00:00:00Z",
|
||||
"reviews": {"nodes": [{"state": "APPROVED", "author": {"login": "bob"}, "submittedAt": "2026-01-02T00:00:00Z"}]},
|
||||
"reviewRequests": {"totalCount": 0},
|
||||
"commits": {"nodes": [{"commit": {"statusCheckRollup": {"state": "SUCCESS"}}}]}
|
||||
}
|
||||
}
|
||||
}
|
||||
}`}
|
||||
|
||||
watch := &PRWatch{
|
||||
SessionID: "s1", TaskID: "t1", Owner: "o", Repo: "r", PRNumber: 42, Branch: "feat",
|
||||
LastCheckStatus: "pending", LastReviewState: "",
|
||||
}
|
||||
if err := store.CreatePRWatch(ctx, watch); err != nil {
|
||||
t.Fatalf("create PR watch: %v", err)
|
||||
}
|
||||
taskPR := &TaskPR{
|
||||
TaskID: "t1", Owner: "o", Repo: "r", PRNumber: 42,
|
||||
PRURL: "https://x/1", PRTitle: "Test PR", HeadBranch: "feat", BaseBranch: "main", State: "open",
|
||||
}
|
||||
if err := store.CreateTaskPR(ctx, taskPR); err != nil {
|
||||
t.Fatalf("create task PR: %v", err)
|
||||
}
|
||||
|
||||
if !poller.tryBatchedPRWatchCheck(ctx, []*PRWatch{watch}) {
|
||||
t.Fatalf("expected batched path to succeed")
|
||||
}
|
||||
|
||||
updated, err := store.GetTaskPR(ctx, "t1")
|
||||
if err != nil || updated == nil {
|
||||
t.Fatalf("get task PR: err=%v, pr=%v", err, updated)
|
||||
}
|
||||
if updated.ChecksState != "success" {
|
||||
t.Errorf("ChecksState = %q, want success", updated.ChecksState)
|
||||
}
|
||||
if updated.ReviewState != "approved" {
|
||||
t.Errorf("ReviewState = %q, want approved", updated.ReviewState)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTryBatchedPRWatchCheck_SearchingWatch_DetectsPR(t *testing.T) {
|
||||
poller, _, gh, store := setupBatchedPollerTest(t)
|
||||
ctx := context.Background()
|
||||
|
||||
gh.branchResponses = []string{`{
|
||||
"data": {
|
||||
"b0": {
|
||||
"ref": {
|
||||
"associatedPullRequests": {
|
||||
"nodes": [{
|
||||
"number": 7,
|
||||
"state": "OPEN", "title": "branch PR", "url": "https://x/7",
|
||||
"isDraft": false, "mergeable": "MERGEABLE", "mergeStateStatus": "CLEAN",
|
||||
"headRefName": "feat", "baseRefName": "main", "headRefOid": "deadbeef",
|
||||
"author": {"login": "alice"},
|
||||
"createdAt": "2026-01-01T00:00:00Z", "updatedAt": "2026-01-01T00:00:00Z",
|
||||
"reviews": {"nodes": []}, "reviewRequests": {"totalCount": 0},
|
||||
"commits": {"nodes": []}
|
||||
}]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}`}
|
||||
|
||||
watch := &PRWatch{
|
||||
SessionID: "s2", TaskID: "t2", Owner: "o", Repo: "r", PRNumber: 0, Branch: "feat",
|
||||
}
|
||||
if err := store.CreatePRWatch(ctx, watch); err != nil {
|
||||
t.Fatalf("create PR watch: %v", err)
|
||||
}
|
||||
|
||||
if !poller.tryBatchedPRWatchCheck(ctx, []*PRWatch{watch}) {
|
||||
t.Fatalf("expected batched path to succeed")
|
||||
}
|
||||
|
||||
updated, err := store.GetPRWatchBySession(ctx, "s2")
|
||||
if err != nil || updated == nil {
|
||||
t.Fatalf("get PR watch: err=%v, w=%v", err, updated)
|
||||
}
|
||||
if updated.PRNumber != 7 {
|
||||
t.Errorf("PRNumber = %d, want 7 (PR should be detected and recorded)", updated.PRNumber)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTryBatchedPRWatchCheck_FallsBackOnUnsupportedClient(t *testing.T) {
|
||||
// MockClient does not implement GraphQLExecutor, so the batched path must
|
||||
// return false to trigger per-watch fallback.
|
||||
poller, _, _, store := setupPollerTest(t)
|
||||
ctx := context.Background()
|
||||
|
||||
watch := &PRWatch{SessionID: "s1", TaskID: "t1", Owner: "o", Repo: "r", PRNumber: 1, Branch: "feat"}
|
||||
if err := store.CreatePRWatch(ctx, watch); err != nil {
|
||||
t.Fatalf("create PR watch: %v", err)
|
||||
}
|
||||
if poller.tryBatchedPRWatchCheck(ctx, []*PRWatch{watch}) {
|
||||
t.Errorf("expected false when client does not implement GraphQLExecutor")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchBatchedStatuses_NumberedQueryError_ReturnsNilFalse(t *testing.T) {
|
||||
poller, _, gh, _ := setupBatchedPollerTest(t)
|
||||
ctx := context.Background()
|
||||
gh.prErr = errors.New("graphql 500")
|
||||
|
||||
numbered := []*PRWatch{{Owner: "o", Repo: "r", PRNumber: 1}}
|
||||
got, ok := poller.fetchBatchedStatuses(ctx, gh, numbered, nil)
|
||||
if ok || got != nil {
|
||||
t.Errorf("expected (nil, false) on numbered query error, got (%v, %v)", got, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchBatchedStatuses_BranchQueryError_TriggersFallback(t *testing.T) {
|
||||
// Greptile P2: when only branch query fails (no numbered watches), the
|
||||
// previous return of an empty non-nil map silently absorbed the failure.
|
||||
// The fix returns (nil, false) so the caller falls back per-watch.
|
||||
poller, _, gh, _ := setupBatchedPollerTest(t)
|
||||
ctx := context.Background()
|
||||
gh.branchErr = errors.New("graphql 500")
|
||||
|
||||
searching := []*PRWatch{{Owner: "o", Repo: "r", PRNumber: 0, Branch: "feat"}}
|
||||
got, ok := poller.fetchBatchedStatuses(ctx, gh, nil, searching)
|
||||
if ok || got != nil {
|
||||
t.Errorf("expected (nil, false) on branch query error, got (%v, %v)", got, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyPRStatus_MergedPR_ResetsWatch(t *testing.T) {
|
||||
poller, _, _, store := setupPollerTest(t)
|
||||
ctx := context.Background()
|
||||
|
||||
mergedAt := time.Now().UTC().Add(-time.Hour)
|
||||
watch := &PRWatch{
|
||||
SessionID: "s1", TaskID: "t1", Owner: "o", Repo: "r", PRNumber: 99, Branch: "feat",
|
||||
}
|
||||
if err := store.CreatePRWatch(ctx, watch); err != nil {
|
||||
t.Fatalf("create PR watch: %v", err)
|
||||
}
|
||||
taskPR := &TaskPR{
|
||||
TaskID: "t1", Owner: "o", Repo: "r", PRNumber: 99,
|
||||
PRURL: "https://x/99", PRTitle: "Merged PR", HeadBranch: "feat", BaseBranch: "main", State: "open",
|
||||
}
|
||||
if err := store.CreateTaskPR(ctx, taskPR); err != nil {
|
||||
t.Fatalf("create task PR: %v", err)
|
||||
}
|
||||
|
||||
status := &PRStatus{
|
||||
PR: &PR{
|
||||
Number: 99, State: prStateMerged, RepoOwner: "o", RepoName: "r",
|
||||
HeadBranch: "feat", BaseBranch: "main", MergedAt: &mergedAt, URL: "https://x/99", Title: "Merged PR",
|
||||
},
|
||||
}
|
||||
poller.applyPRStatus(ctx, watch, status)
|
||||
|
||||
updated, err := store.GetPRWatchBySession(ctx, "s1")
|
||||
if err != nil || updated == nil {
|
||||
t.Fatalf("get PR watch: err=%v, w=%v", err, updated)
|
||||
}
|
||||
if updated.PRNumber != 0 {
|
||||
t.Errorf("expected watch reset to PRNumber=0 after merge, got %d", updated.PRNumber)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitForRateLimit_SkipsWhenHealthy(t *testing.T) {
|
||||
poller, _, _, _ := setupPollerTest(t)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
if !poller.waitForRateLimit(ctx, ResourceGraphQL, "test") {
|
||||
t.Errorf("expected true when no exhaustion recorded")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchBucketExhausted(t *testing.T) {
|
||||
poller, svc, _, _ := setupPollerTest(t)
|
||||
|
||||
if poller.searchBucketExhausted("test") {
|
||||
t.Errorf("expected false when no exhaustion recorded")
|
||||
}
|
||||
|
||||
// Mark the search bucket exhausted with a future reset.
|
||||
svc.RateTracker().Record(RateSnapshot{
|
||||
Resource: ResourceSearch,
|
||||
Limit: 30,
|
||||
Remaining: 0,
|
||||
ResetAt: time.Now().Add(5 * time.Minute),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
|
||||
if !poller.searchBucketExhausted("test") {
|
||||
t.Errorf("expected true when search bucket exhausted")
|
||||
}
|
||||
|
||||
// Other buckets exhausting must not trip the search guard.
|
||||
poller2, svc2, _, _ := setupPollerTest(t)
|
||||
svc2.RateTracker().Record(RateSnapshot{
|
||||
Resource: ResourceCore,
|
||||
Remaining: 0,
|
||||
ResetAt: time.Now().Add(5 * time.Minute),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
if poller2.searchBucketExhausted("test") {
|
||||
t.Errorf("expected false when only core (not search) is exhausted")
|
||||
}
|
||||
}
|
||||
|
||||
// Regression: when the search bucket trips mid-cycle, checkReviewWatches must
|
||||
// stop iterating so the remaining watches don't issue doomed search requests
|
||||
// that deepen the secondary-limit penalty.
|
||||
func TestCheckReviewWatches_BailsOutWhenSearchExhaustedMidCycle(t *testing.T) {
|
||||
poller, svc, _, store := setupPollerTest(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Pre-exhaust the search bucket so the very first iteration short-circuits.
|
||||
svc.RateTracker().Record(RateSnapshot{
|
||||
Resource: ResourceSearch,
|
||||
Remaining: 0,
|
||||
ResetAt: time.Now().Add(5 * time.Minute),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
|
||||
for _, id := range []string{"rw1", "rw2", "rw3"} {
|
||||
if err := store.CreateReviewWatch(ctx, &ReviewWatch{
|
||||
ID: id,
|
||||
WorkspaceID: "ws1",
|
||||
Enabled: true,
|
||||
}); err != nil {
|
||||
t.Fatalf("seed review watch %s: %v", id, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Should return cleanly without panic / without errors despite N watches
|
||||
// being enabled.
|
||||
poller.checkReviewWatches(ctx)
|
||||
}
|
||||
|
||||
func TestCheckIssueWatches_BailsOutWhenSearchExhaustedMidCycle(t *testing.T) {
|
||||
poller, svc, _, store := setupPollerTest(t)
|
||||
ctx := context.Background()
|
||||
|
||||
svc.RateTracker().Record(RateSnapshot{
|
||||
Resource: ResourceSearch,
|
||||
Remaining: 0,
|
||||
ResetAt: time.Now().Add(5 * time.Minute),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
|
||||
for _, id := range []string{"iw1", "iw2"} {
|
||||
if err := store.CreateIssueWatch(ctx, &IssueWatch{
|
||||
ID: id,
|
||||
WorkspaceID: "ws1",
|
||||
Enabled: true,
|
||||
}); err != nil {
|
||||
t.Fatalf("seed issue watch %s: %v", id, err)
|
||||
}
|
||||
}
|
||||
|
||||
poller.checkIssueWatches(ctx)
|
||||
}
|
||||
|
||||
func TestWaitForRateLimit_ReturnsFalseOnContextCancel(t *testing.T) {
|
||||
poller, svc, _, _ := setupPollerTest(t)
|
||||
// Force exhaustion with a far-future reset.
|
||||
svc.RateTracker().Record(RateSnapshot{
|
||||
Resource: ResourceGraphQL,
|
||||
Limit: 5000,
|
||||
Remaining: 0,
|
||||
ResetAt: time.Now().Add(time.Hour),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel() // cancel immediately so the timer select returns ctx.Done first
|
||||
if poller.waitForRateLimit(ctx, ResourceGraphQL, "test") {
|
||||
t.Errorf("expected false when context cancelled during sleep")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,13 +2,36 @@ package github
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/jmoiron/sqlx"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/kandev/kandev/internal/common/logger"
|
||||
"github.com/kandev/kandev/internal/events/bus"
|
||||
)
|
||||
|
||||
// rateLimitSeedTimeout caps the startup `gh api rate_limit` probe so a slow
|
||||
// CLI cannot delay backend boot.
|
||||
const rateLimitSeedTimeout = 5 * time.Second
|
||||
|
||||
// attachRateTracker wires the rate tracker into whichever client
|
||||
// implementation we ended up with, then seeds initial snapshots. The seed is
|
||||
// best-effort — a missing CLI or network blip should not block startup.
|
||||
func attachRateTracker(client Client, tracker *RateTracker, log *logger.Logger) {
|
||||
switch c := client.(type) {
|
||||
case *PATClient:
|
||||
c.WithRateTracker(tracker)
|
||||
case *GHClient:
|
||||
c.WithRateTracker(tracker)
|
||||
seedCtx, cancel := context.WithTimeout(context.Background(), rateLimitSeedTimeout)
|
||||
defer cancel()
|
||||
if err := c.FetchRateLimit(seedCtx); err != nil {
|
||||
log.Debug("seed gh rate limit failed", zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Provide creates the full GitHub integration stack: store, client, and service.
|
||||
// Returns the service, a cleanup function, and any error.
|
||||
func Provide(
|
||||
@@ -30,6 +53,7 @@ func Provide(
|
||||
}
|
||||
|
||||
svc := NewService(client, authMethod, secrets, store, eventBus, log)
|
||||
attachRateTracker(client, svc.RateTracker(), log)
|
||||
svc.subscribeTaskEvents()
|
||||
|
||||
cleanup := func() error {
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
package github
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/kandev/kandev/internal/common/logger"
|
||||
"github.com/kandev/kandev/internal/events"
|
||||
"github.com/kandev/kandev/internal/events/bus"
|
||||
)
|
||||
|
||||
// Resource identifies a GitHub API rate-limit bucket.
|
||||
//
|
||||
// GitHub maintains separate counters for REST core, GraphQL, and search; the
|
||||
// `gh` CLI hits GraphQL even for `pr view`, so exhausting one bucket does not
|
||||
// imply the others are exhausted.
|
||||
type Resource string
|
||||
|
||||
const (
|
||||
ResourceCore Resource = "core"
|
||||
ResourceGraphQL Resource = "graphql"
|
||||
ResourceSearch Resource = "search"
|
||||
)
|
||||
|
||||
// rateUpdateDebounce throttles non-transition update events so settings-page
|
||||
// counters update at human speed rather than once per request.
|
||||
const rateUpdateDebounce = 5 * time.Second
|
||||
|
||||
// RateSnapshot captures the rate-limit state for one bucket at a point in time.
|
||||
type RateSnapshot struct {
|
||||
Resource Resource `json:"resource"`
|
||||
Remaining int `json:"remaining"`
|
||||
Limit int `json:"limit"`
|
||||
ResetAt time.Time `json:"reset_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Exhausted returns true when the bucket has no quota left and ResetAt is in
|
||||
// the future. Limit may be unknown (0) when the snapshot was synthesized from
|
||||
// an out-of-band signal (e.g. `gh` stderr).
|
||||
func (s RateSnapshot) Exhausted() bool {
|
||||
return s.Remaining <= 0 && s.ResetAt.After(time.Now())
|
||||
}
|
||||
|
||||
// RateLimitUpdatedEvent is published when a snapshot changes, either because
|
||||
// new headers arrived or because exhaustion state flipped.
|
||||
//
|
||||
// Carries all known buckets so subscribers can surface the full picture
|
||||
// without follow-up queries.
|
||||
type RateLimitUpdatedEvent struct {
|
||||
Snapshots []RateSnapshot `json:"snapshots"`
|
||||
// Trigger names the bucket whose update produced this event; useful for
|
||||
// downstream filtering/debug.
|
||||
Trigger Resource `json:"trigger"`
|
||||
// ExhaustionTransition is non-empty on the tick where the trigger bucket
|
||||
// flipped between exhausted and recovered ("exhausted" or "recovered").
|
||||
ExhaustionTransition string `json:"exhaustion_transition,omitempty"`
|
||||
}
|
||||
|
||||
// RateTracker accumulates rate-limit snapshots from the GitHub clients and
|
||||
// publishes events on change.
|
||||
type RateTracker struct {
|
||||
mu sync.RWMutex
|
||||
snapshots map[Resource]RateSnapshot
|
||||
exhausted map[Resource]bool
|
||||
lastEmitted map[Resource]time.Time
|
||||
bus bus.EventBus
|
||||
log *logger.Logger
|
||||
}
|
||||
|
||||
// NewRateTracker constructs a tracker. The event bus may be nil (tests).
|
||||
func NewRateTracker(eventBus bus.EventBus, log *logger.Logger) *RateTracker {
|
||||
return &RateTracker{
|
||||
snapshots: make(map[Resource]RateSnapshot),
|
||||
exhausted: make(map[Resource]bool),
|
||||
lastEmitted: make(map[Resource]time.Time),
|
||||
bus: eventBus,
|
||||
log: log,
|
||||
}
|
||||
}
|
||||
|
||||
// Record stores a snapshot. Always emits an event on exhaustion transition;
|
||||
// otherwise debounces to at most one update per resource per
|
||||
// rateUpdateDebounce window.
|
||||
func (r *RateTracker) Record(snap RateSnapshot) {
|
||||
if snap.Resource == "" {
|
||||
return
|
||||
}
|
||||
if snap.UpdatedAt.IsZero() {
|
||||
snap.UpdatedAt = time.Now().UTC()
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
prev := r.exhausted[snap.Resource]
|
||||
now := snap.Exhausted()
|
||||
r.snapshots[snap.Resource] = snap
|
||||
r.exhausted[snap.Resource] = now
|
||||
|
||||
transition := ""
|
||||
switch {
|
||||
case !prev && now:
|
||||
transition = "exhausted"
|
||||
case prev && !now:
|
||||
transition = "recovered"
|
||||
}
|
||||
|
||||
shouldEmit := transition != ""
|
||||
if !shouldEmit {
|
||||
last := r.lastEmitted[snap.Resource]
|
||||
if snap.UpdatedAt.Sub(last) >= rateUpdateDebounce {
|
||||
shouldEmit = true
|
||||
}
|
||||
}
|
||||
if shouldEmit {
|
||||
r.lastEmitted[snap.Resource] = snap.UpdatedAt
|
||||
}
|
||||
allSnap := r.allLocked()
|
||||
r.mu.Unlock()
|
||||
|
||||
if !shouldEmit {
|
||||
return
|
||||
}
|
||||
r.publish(allSnap, snap.Resource, transition)
|
||||
}
|
||||
|
||||
// Snapshot returns a copy of the current snapshot for resource.
|
||||
func (r *RateTracker) Snapshot(resource Resource) (RateSnapshot, bool) {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
snap, ok := r.snapshots[resource]
|
||||
return snap, ok
|
||||
}
|
||||
|
||||
// All returns a copy of every known snapshot.
|
||||
func (r *RateTracker) All() map[Resource]RateSnapshot {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.allLocked()
|
||||
}
|
||||
|
||||
func (r *RateTracker) allLocked() map[Resource]RateSnapshot {
|
||||
out := make(map[Resource]RateSnapshot, len(r.snapshots))
|
||||
for k, v := range r.snapshots {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// IsExhausted reports whether the resource bucket is currently exhausted.
|
||||
func (r *RateTracker) IsExhausted(resource Resource) bool {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.exhausted[resource]
|
||||
}
|
||||
|
||||
// WaitDuration returns how long callers should wait before retrying the
|
||||
// resource, or 0 if the bucket is not exhausted.
|
||||
func (r *RateTracker) WaitDuration(resource Resource) time.Duration {
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
if !r.exhausted[resource] {
|
||||
return 0
|
||||
}
|
||||
d := time.Until(r.snapshots[resource].ResetAt)
|
||||
if d < 0 {
|
||||
return 0
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func (r *RateTracker) publish(all map[Resource]RateSnapshot, trigger Resource, transition string) {
|
||||
if r.bus == nil {
|
||||
return
|
||||
}
|
||||
snaps := make([]RateSnapshot, 0, len(all))
|
||||
for _, s := range all {
|
||||
snaps = append(snaps, s)
|
||||
}
|
||||
evt := &RateLimitUpdatedEvent{
|
||||
Snapshots: snaps,
|
||||
Trigger: trigger,
|
||||
ExhaustionTransition: transition,
|
||||
}
|
||||
event := bus.NewEvent(events.GitHubRateLimitUpdated, "github", evt)
|
||||
if err := r.bus.Publish(context.Background(), events.GitHubRateLimitUpdated, event); err != nil && r.log != nil {
|
||||
r.log.Debug("publish rate-limit event failed", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// parseRateHeaders extracts a snapshot from a GitHub HTTP response. Returns
|
||||
// false when the response carries no rate-limit headers.
|
||||
func parseRateHeaders(resp *http.Response, defaultResource Resource) (RateSnapshot, bool) {
|
||||
if resp == nil {
|
||||
return RateSnapshot{}, false
|
||||
}
|
||||
limitStr := resp.Header.Get("X-RateLimit-Limit")
|
||||
remainingStr := resp.Header.Get("X-RateLimit-Remaining")
|
||||
resetStr := resp.Header.Get("X-RateLimit-Reset")
|
||||
if limitStr == "" && remainingStr == "" && resetStr == "" {
|
||||
return RateSnapshot{}, false
|
||||
}
|
||||
limit, _ := strconv.Atoi(limitStr)
|
||||
remaining, _ := strconv.Atoi(remainingStr)
|
||||
reset, _ := strconv.ParseInt(resetStr, 10, 64)
|
||||
|
||||
resource := defaultResource
|
||||
if r := resp.Header.Get("X-RateLimit-Resource"); r != "" {
|
||||
resource = Resource(strings.ToLower(r))
|
||||
}
|
||||
return RateSnapshot{
|
||||
Resource: resource,
|
||||
Limit: limit,
|
||||
Remaining: remaining,
|
||||
ResetAt: time.Unix(reset, 0).UTC(),
|
||||
UpdatedAt: time.Now().UTC(),
|
||||
}, true
|
||||
}
|
||||
|
||||
// markRateExhausted records an exhaustion snapshot for callers that detect
|
||||
// rate-limit conditions out-of-band (e.g. parsing `gh` stderr or 4xx bodies
|
||||
// without headers). Uses a conservative one-hour reset when no better value
|
||||
// is known.
|
||||
func (r *RateTracker) markRateExhausted(resource Resource, resetAt time.Time) {
|
||||
if resource == "" {
|
||||
return
|
||||
}
|
||||
if resetAt.IsZero() {
|
||||
resetAt = time.Now().Add(time.Hour).UTC()
|
||||
}
|
||||
prev, _ := r.Snapshot(resource)
|
||||
r.Record(RateSnapshot{
|
||||
Resource: resource,
|
||||
Limit: prev.Limit,
|
||||
Remaining: 0,
|
||||
ResetAt: resetAt,
|
||||
UpdatedAt: time.Now().UTC(),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
package github
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/kandev/kandev/internal/common/logger"
|
||||
"github.com/kandev/kandev/internal/events"
|
||||
"github.com/kandev/kandev/internal/events/bus"
|
||||
)
|
||||
|
||||
func newTestTrackerLogger(t *testing.T) *logger.Logger {
|
||||
t.Helper()
|
||||
log, err := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console"})
|
||||
if err != nil {
|
||||
t.Fatalf("logger: %v", err)
|
||||
}
|
||||
return log
|
||||
}
|
||||
|
||||
// captureBus collects events published to a single subject so tests can
|
||||
// assert on transitions without racing.
|
||||
type captureBus struct {
|
||||
mu sync.Mutex
|
||||
events []*bus.Event
|
||||
inner bus.EventBus
|
||||
}
|
||||
|
||||
func newCaptureBus(t *testing.T, subject string) *captureBus {
|
||||
t.Helper()
|
||||
log := newTestTrackerLogger(t)
|
||||
cb := &captureBus{inner: bus.NewMemoryEventBus(log)}
|
||||
if _, err := cb.inner.Subscribe(subject, func(_ context.Context, e *bus.Event) error {
|
||||
cb.mu.Lock()
|
||||
cb.events = append(cb.events, e)
|
||||
cb.mu.Unlock()
|
||||
return nil
|
||||
}); err != nil {
|
||||
t.Fatalf("subscribe: %v", err)
|
||||
}
|
||||
return cb
|
||||
}
|
||||
|
||||
func (c *captureBus) Publish(ctx context.Context, subject string, event *bus.Event) error {
|
||||
return c.inner.Publish(ctx, subject, event)
|
||||
}
|
||||
|
||||
func (c *captureBus) Subscribe(subject string, handler bus.EventHandler) (bus.Subscription, error) {
|
||||
return c.inner.Subscribe(subject, handler)
|
||||
}
|
||||
|
||||
func (c *captureBus) QueueSubscribe(subject, queue string, handler bus.EventHandler) (bus.Subscription, error) {
|
||||
return c.inner.QueueSubscribe(subject, queue, handler)
|
||||
}
|
||||
|
||||
func (c *captureBus) Request(ctx context.Context, subject string, event *bus.Event, timeout time.Duration) (*bus.Event, error) {
|
||||
return c.inner.Request(ctx, subject, event, timeout)
|
||||
}
|
||||
|
||||
func (c *captureBus) Close() { c.inner.Close() }
|
||||
func (c *captureBus) IsConnected() bool { return c.inner.IsConnected() }
|
||||
|
||||
// captured returns a snapshot of events captured so far. MemoryEventBus
|
||||
// dispatches to subscribers synchronously, so callers can read this directly
|
||||
// after Record() / Publish() returns without polling.
|
||||
func (c *captureBus) captured() []*bus.Event {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return append([]*bus.Event(nil), c.events...)
|
||||
}
|
||||
|
||||
func (c *captureBus) requireCount(t *testing.T, count int, msg string) []*bus.Event {
|
||||
t.Helper()
|
||||
evs := c.captured()
|
||||
if len(evs) != count {
|
||||
t.Fatalf("%s: expected %d events, got %d", msg, count, len(evs))
|
||||
}
|
||||
return evs
|
||||
}
|
||||
|
||||
func TestParseRateHeaders_DecodesAllFields(t *testing.T) {
|
||||
resp := &http.Response{Header: http.Header{}}
|
||||
resp.Header.Set("X-RateLimit-Limit", "5000")
|
||||
resp.Header.Set("X-RateLimit-Remaining", "4242")
|
||||
resp.Header.Set("X-RateLimit-Reset", "2000000000")
|
||||
resp.Header.Set("X-RateLimit-Resource", "graphql")
|
||||
|
||||
snap, ok := parseRateHeaders(resp, ResourceCore)
|
||||
if !ok {
|
||||
t.Fatalf("expected snapshot to be parsed")
|
||||
}
|
||||
if snap.Resource != ResourceGraphQL {
|
||||
t.Errorf("resource: got %q want graphql", snap.Resource)
|
||||
}
|
||||
if snap.Limit != 5000 || snap.Remaining != 4242 {
|
||||
t.Errorf("limit/remaining mismatch: %d / %d", snap.Limit, snap.Remaining)
|
||||
}
|
||||
if snap.ResetAt.Unix() != 2000000000 {
|
||||
t.Errorf("reset_at: got %d want 2000000000", snap.ResetAt.Unix())
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseRateHeaders_FallsBackToDefaultResource(t *testing.T) {
|
||||
resp := &http.Response{Header: http.Header{}}
|
||||
resp.Header.Set("X-RateLimit-Limit", "30")
|
||||
resp.Header.Set("X-RateLimit-Remaining", "29")
|
||||
resp.Header.Set("X-RateLimit-Reset", "2000000000")
|
||||
|
||||
snap, ok := parseRateHeaders(resp, ResourceSearch)
|
||||
if !ok || snap.Resource != ResourceSearch {
|
||||
t.Fatalf("expected default resource search, got %q (ok=%v)", snap.Resource, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseRateHeaders_NoHeadersReturnsFalse(t *testing.T) {
|
||||
resp := &http.Response{Header: http.Header{}}
|
||||
if _, ok := parseRateHeaders(resp, ResourceCore); ok {
|
||||
t.Fatalf("expected no snapshot when headers absent")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateTracker_RecordEmitsAndTracksExhaustion(t *testing.T) {
|
||||
log := newTestTrackerLogger(t)
|
||||
cb := newCaptureBus(t, events.GitHubRateLimitUpdated)
|
||||
tr := NewRateTracker(cb, log)
|
||||
|
||||
future := time.Now().Add(20 * time.Minute)
|
||||
|
||||
// First snapshot: healthy. Should publish once (initial update).
|
||||
tr.Record(RateSnapshot{
|
||||
Resource: ResourceGraphQL,
|
||||
Limit: 5000,
|
||||
Remaining: 4000,
|
||||
ResetAt: future,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
if tr.IsExhausted(ResourceGraphQL) {
|
||||
t.Fatalf("not exhausted yet")
|
||||
}
|
||||
|
||||
// Second snapshot: exhausted -> must publish (transition).
|
||||
tr.Record(RateSnapshot{
|
||||
Resource: ResourceGraphQL,
|
||||
Limit: 5000,
|
||||
Remaining: 0,
|
||||
ResetAt: future,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
if !tr.IsExhausted(ResourceGraphQL) {
|
||||
t.Fatalf("expected exhausted after Remaining=0 + future reset")
|
||||
}
|
||||
if d := tr.WaitDuration(ResourceGraphQL); d <= 0 || d > 21*time.Minute {
|
||||
t.Fatalf("WaitDuration out of range: %v", d)
|
||||
}
|
||||
|
||||
// Third snapshot: recovered -> publishes recovery transition.
|
||||
tr.Record(RateSnapshot{
|
||||
Resource: ResourceGraphQL,
|
||||
Limit: 5000,
|
||||
Remaining: 1500,
|
||||
ResetAt: future,
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
|
||||
evts := cb.requireCount(t, 3, "expected 3 events (initial, exhausted, recovered)")
|
||||
// Last two events should carry transition strings.
|
||||
gotTransitions := []string{}
|
||||
for _, e := range evts {
|
||||
payload, ok := e.Data.(*RateLimitUpdatedEvent)
|
||||
if !ok {
|
||||
t.Fatalf("unexpected event payload type %T", e.Data)
|
||||
}
|
||||
gotTransitions = append(gotTransitions, payload.ExhaustionTransition)
|
||||
}
|
||||
if gotTransitions[1] != "exhausted" || gotTransitions[2] != "recovered" {
|
||||
t.Errorf("transitions = %v, want [* exhausted recovered]", gotTransitions)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateTracker_DebouncesNonTransitionUpdates(t *testing.T) {
|
||||
log := newTestTrackerLogger(t)
|
||||
cb := newCaptureBus(t, events.GitHubRateLimitUpdated)
|
||||
tr := NewRateTracker(cb, log)
|
||||
now := time.Now()
|
||||
|
||||
// Two healthy snapshots within the debounce window -> only first emits.
|
||||
// MemoryEventBus.Publish dispatches synchronously, so reads after Record
|
||||
// see the final state without sleeping.
|
||||
tr.Record(RateSnapshot{Resource: ResourceCore, Limit: 5000, Remaining: 4999, ResetAt: now.Add(time.Hour), UpdatedAt: now})
|
||||
tr.Record(RateSnapshot{Resource: ResourceCore, Limit: 5000, Remaining: 4998, ResetAt: now.Add(time.Hour), UpdatedAt: now.Add(time.Second)})
|
||||
cb.requireCount(t, 1, "expected 1 debounced event")
|
||||
|
||||
// Outside the window -> emits again.
|
||||
tr.Record(RateSnapshot{Resource: ResourceCore, Limit: 5000, Remaining: 4997, ResetAt: now.Add(time.Hour), UpdatedAt: now.Add(2 * rateUpdateDebounce)})
|
||||
cb.requireCount(t, 2, "expected second event after debounce window")
|
||||
}
|
||||
|
||||
func TestRateTracker_MarkRateExhaustedSeedsUnknownLimit(t *testing.T) {
|
||||
log := newTestTrackerLogger(t)
|
||||
cb := newCaptureBus(t, events.GitHubRateLimitUpdated)
|
||||
tr := NewRateTracker(cb, log)
|
||||
|
||||
tr.markRateExhausted(ResourceGraphQL, time.Time{})
|
||||
|
||||
if !tr.IsExhausted(ResourceGraphQL) {
|
||||
t.Fatalf("expected exhausted after markRateExhausted")
|
||||
}
|
||||
snap, ok := tr.Snapshot(ResourceGraphQL)
|
||||
if !ok {
|
||||
t.Fatalf("snapshot not stored")
|
||||
}
|
||||
if snap.Remaining != 0 {
|
||||
t.Errorf("Remaining = %d, want 0", snap.Remaining)
|
||||
}
|
||||
// Default reset (~1h) means WaitDuration > 0.
|
||||
if d := tr.WaitDuration(ResourceGraphQL); d <= 0 {
|
||||
t.Errorf("WaitDuration should be positive, got %v", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateTracker_NilBusIsNoop(t *testing.T) {
|
||||
tr := NewRateTracker(nil, nil)
|
||||
tr.Record(RateSnapshot{Resource: ResourceCore, Limit: 5000, Remaining: 0, ResetAt: time.Now().Add(time.Minute)})
|
||||
if !tr.IsExhausted(ResourceCore) {
|
||||
t.Fatalf("tracker should still record when bus is nil")
|
||||
}
|
||||
}
|
||||
@@ -71,6 +71,7 @@ type Service struct {
|
||||
taskEventSubs []bus.Subscription
|
||||
searchCache *ttlCache
|
||||
prStatusCache *ttlCache
|
||||
rateTracker *RateTracker
|
||||
}
|
||||
|
||||
// NewService creates a new GitHub service.
|
||||
@@ -84,9 +85,80 @@ func NewService(client Client, authMethod string, secrets SecretProvider, store
|
||||
logger: log,
|
||||
searchCache: newTTLCache(),
|
||||
prStatusCache: newTTLCache(),
|
||||
rateTracker: NewRateTracker(eventBus, log),
|
||||
}
|
||||
}
|
||||
|
||||
// RateTracker exposes the service's rate-limit tracker so factory callers
|
||||
// can wire it into individual clients.
|
||||
func (s *Service) RateTracker() *RateTracker {
|
||||
return s.rateTracker
|
||||
}
|
||||
|
||||
// newPATClient builds a PATClient pre-wired with the service's shared rate
|
||||
// tracker. Centralizing this guards against forgetting the wiring on auth
|
||||
// flips (e.g. ConfigureToken), which would otherwise leave PAT calls
|
||||
// invisible to the rate-limit UI, health checks, and poller throttling.
|
||||
func (s *Service) newPATClient(token string) *PATClient {
|
||||
c := NewPATClient(token)
|
||||
attachRateTracker(c, s.rateTracker, s.logger)
|
||||
return c
|
||||
}
|
||||
|
||||
// ExhaustedRateLimit is a tiny DTO returned to the health package — the
|
||||
// health package can't import github (cycle), so it consumes a structural
|
||||
// shape via interface assertion.
|
||||
type ExhaustedRateLimit struct {
|
||||
Resource string
|
||||
ResetAt time.Time
|
||||
}
|
||||
|
||||
// ExhaustedRateLimits lists every resource bucket currently out of quota.
|
||||
// Returns an empty slice when none are exhausted, so callers can safely
|
||||
// `len()` the result.
|
||||
func (s *Service) ExhaustedRateLimits() []ExhaustedRateLimit {
|
||||
if s.rateTracker == nil {
|
||||
return nil
|
||||
}
|
||||
out := []ExhaustedRateLimit{}
|
||||
for resource, snap := range s.rateTracker.All() {
|
||||
if snap.Exhausted() {
|
||||
out = append(out, ExhaustedRateLimit{
|
||||
Resource: string(resource),
|
||||
ResetAt: snap.ResetAt,
|
||||
})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// rateLimitInfo materializes the tracker's snapshots into the DTO shape
|
||||
// returned by GetStatus and the rate-limit WS notification. Returns nil
|
||||
// when no buckets are known yet so the field stays omitted.
|
||||
func (s *Service) rateLimitInfo() *GitHubRateLimitInfo {
|
||||
if s.rateTracker == nil {
|
||||
return nil
|
||||
}
|
||||
all := s.rateTracker.All()
|
||||
if len(all) == 0 {
|
||||
return nil
|
||||
}
|
||||
info := &GitHubRateLimitInfo{}
|
||||
if snap, ok := all[ResourceCore]; ok {
|
||||
v := snap
|
||||
info.Core = &v
|
||||
}
|
||||
if snap, ok := all[ResourceGraphQL]; ok {
|
||||
v := snap
|
||||
info.GraphQL = &v
|
||||
}
|
||||
if snap, ok := all[ResourceSearch]; ok {
|
||||
v := snap
|
||||
info.Search = &v
|
||||
}
|
||||
return info
|
||||
}
|
||||
|
||||
// SetTaskDeleter sets the task deletion dependency for cleanup operations.
|
||||
func (s *Service) SetTaskDeleter(d TaskDeleter) { s.taskDeleter = d }
|
||||
|
||||
@@ -143,6 +215,7 @@ func (s *Service) GetStatus(ctx context.Context) (*GitHubStatus, error) {
|
||||
status := &GitHubStatus{
|
||||
AuthMethod: authMethod,
|
||||
RequiredScopes: RequiredGitHubScopes,
|
||||
RateLimit: s.rateLimitInfo(),
|
||||
}
|
||||
|
||||
// Check if a GITHUB_TOKEN secret is configured
|
||||
@@ -204,6 +277,7 @@ func (s *Service) retryClientCreation(ctx context.Context) {
|
||||
s.logger.Debug("GitHub client retry failed", zap.Error(err))
|
||||
return
|
||||
}
|
||||
attachRateTracker(client, s.rateTracker, s.logger)
|
||||
s.client = client
|
||||
s.authMethod = authMethod
|
||||
s.logger.Info("GitHub client recovered after retry",
|
||||
@@ -229,8 +303,10 @@ func (s *Service) ConfigureToken(ctx context.Context, token string) error {
|
||||
return fmt.Errorf("secret manager not configured")
|
||||
}
|
||||
|
||||
// Validate the token by testing authentication
|
||||
testClient := NewPATClient(token)
|
||||
// Validate the token by testing authentication. Wire the rate tracker
|
||||
// onto the test client up front so the validation request seeds the
|
||||
// shared quota and subsequent PAT-backed calls keep feeding it.
|
||||
testClient := s.newPATClient(token)
|
||||
user, err := testClient.GetAuthenticatedUser(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid token: %w", err)
|
||||
|
||||
@@ -262,3 +262,23 @@ func TestClearToken_NoSecretManager(t *testing.T) {
|
||||
t.Error("expected error when secretManager is nil")
|
||||
}
|
||||
}
|
||||
|
||||
// Regression: PAT clients minted via the service must share the rate tracker.
|
||||
// ConfigureToken previously created a bare NewPATClient and dropped quota
|
||||
// signals on the floor, so the UI / health / poller throttling all went stale
|
||||
// after a token reconfigure.
|
||||
func TestService_NewPATClient_AttachesRateTracker(t *testing.T) {
|
||||
log, _ := logger.NewLogger(logger.LoggingConfig{Level: "error", Format: "console"})
|
||||
svc := &Service{
|
||||
logger: log,
|
||||
rateTracker: NewRateTracker(nil, nil),
|
||||
}
|
||||
|
||||
c := svc.newPATClient("ghp_x")
|
||||
if c == nil {
|
||||
t.Fatalf("newPATClient returned nil")
|
||||
}
|
||||
if c.rateTracker != svc.rateTracker {
|
||||
t.Fatalf("expected rate tracker to be wired onto the new PAT client")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package health
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
@@ -12,9 +13,24 @@ type GitHubStatusProvider interface {
|
||||
AuthMethod() string
|
||||
}
|
||||
|
||||
// GitHubChecker checks GitHub CLI/auth availability.
|
||||
// GitHubRateLimitProvider exposes per-resource exhaustion state so the health
|
||||
// dialog can surface "GitHub API rate limit exhausted" alongside other setup
|
||||
// issues. Implemented by *github.Service.
|
||||
type GitHubRateLimitProvider interface {
|
||||
ExhaustedRateLimits() []GitHubRateLimitStatus
|
||||
}
|
||||
|
||||
// GitHubRateLimitStatus is a minimal DTO so the health package does not have
|
||||
// to import the github package (which would create a cycle).
|
||||
type GitHubRateLimitStatus struct {
|
||||
Resource string
|
||||
ResetAt time.Time
|
||||
}
|
||||
|
||||
// GitHubChecker checks GitHub CLI/auth availability and rate-limit health.
|
||||
type GitHubChecker struct {
|
||||
provider GitHubStatusProvider
|
||||
provider GitHubStatusProvider
|
||||
rateProvider GitHubRateLimitProvider
|
||||
}
|
||||
|
||||
// NewGitHubChecker creates a checker for GitHub integration status.
|
||||
@@ -23,6 +39,13 @@ func NewGitHubChecker(provider GitHubStatusProvider) *GitHubChecker {
|
||||
return &GitHubChecker{provider: provider}
|
||||
}
|
||||
|
||||
// WithRateLimitProvider attaches a provider so the checker also surfaces
|
||||
// rate-limit exhaustion. Optional — when nil the rate-limit issue is omitted.
|
||||
func (c *GitHubChecker) WithRateLimitProvider(p GitHubRateLimitProvider) *GitHubChecker {
|
||||
c.rateProvider = p
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *GitHubChecker) Check(_ context.Context) []Issue {
|
||||
if c.provider == nil {
|
||||
return []Issue{{
|
||||
@@ -46,7 +69,44 @@ func (c *GitHubChecker) Check(_ context.Context) []Issue {
|
||||
FixLabel: "Configure GitHub",
|
||||
}}
|
||||
}
|
||||
return nil
|
||||
return c.rateLimitIssues()
|
||||
}
|
||||
|
||||
// rateLimitIssues materializes one Issue per exhausted resource bucket.
|
||||
// Returns nil when the rate-limit provider is missing or every bucket has
|
||||
// quota — the issue auto-clears as soon as that happens.
|
||||
func (c *GitHubChecker) rateLimitIssues() []Issue {
|
||||
if c.rateProvider == nil {
|
||||
return nil
|
||||
}
|
||||
exhausted := c.rateProvider.ExhaustedRateLimits()
|
||||
if len(exhausted) == 0 {
|
||||
return nil
|
||||
}
|
||||
issues := make([]Issue, 0, len(exhausted))
|
||||
for _, status := range exhausted {
|
||||
issues = append(issues, Issue{
|
||||
ID: fmt.Sprintf("github_rate_limit_%s", status.Resource),
|
||||
Category: "github",
|
||||
Title: fmt.Sprintf("GitHub API rate limit exhausted (%s)", status.Resource),
|
||||
Message: rateLimitMessage(status),
|
||||
Severity: SeverityWarning,
|
||||
FixURL: "/settings/workspace/{workspaceId}/github",
|
||||
FixLabel: "View status",
|
||||
})
|
||||
}
|
||||
return issues
|
||||
}
|
||||
|
||||
func rateLimitMessage(status GitHubRateLimitStatus) string {
|
||||
if status.ResetAt.IsZero() {
|
||||
return "PR/issue checks are paused until the limit resets."
|
||||
}
|
||||
wait := time.Until(status.ResetAt).Round(time.Minute)
|
||||
if wait <= 0 {
|
||||
return "PR/issue checks are paused until the limit resets."
|
||||
}
|
||||
return fmt.Sprintf("PR/issue checks are paused; resets in %s.", wait)
|
||||
}
|
||||
|
||||
// AgentDiscoveryProvider abstracts the agent discovery check.
|
||||
|
||||
@@ -3,7 +3,9 @@ package health
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// --- GitHubChecker tests ---
|
||||
@@ -55,6 +57,41 @@ func TestGitHubChecker_Authenticated(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
type stubRateLimitProvider struct {
|
||||
exhausted []GitHubRateLimitStatus
|
||||
}
|
||||
|
||||
func (s *stubRateLimitProvider) ExhaustedRateLimits() []GitHubRateLimitStatus {
|
||||
return s.exhausted
|
||||
}
|
||||
|
||||
func TestGitHubChecker_AuthenticatedButRateLimited(t *testing.T) {
|
||||
rate := &stubRateLimitProvider{exhausted: []GitHubRateLimitStatus{
|
||||
{Resource: "graphql", ResetAt: time.Now().Add(15 * time.Minute)},
|
||||
}}
|
||||
checker := NewGitHubChecker(&mockGitHubProvider{authenticated: true, authMethod: "gh_cli"})
|
||||
checker.WithRateLimitProvider(rate)
|
||||
issues := checker.Check(context.Background())
|
||||
if len(issues) != 1 {
|
||||
t.Fatalf("expected 1 issue, got %d", len(issues))
|
||||
}
|
||||
if issues[0].ID != "github_rate_limit_graphql" {
|
||||
t.Errorf("ID = %q, want github_rate_limit_graphql", issues[0].ID)
|
||||
}
|
||||
if !strings.Contains(issues[0].Message, "resets in") {
|
||||
t.Errorf("message should include reset countdown, got %q", issues[0].Message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHubChecker_AuthenticatedRateProviderEmpty(t *testing.T) {
|
||||
rate := &stubRateLimitProvider{}
|
||||
checker := NewGitHubChecker(&mockGitHubProvider{authenticated: true})
|
||||
checker.WithRateLimitProvider(rate)
|
||||
if got := checker.Check(context.Background()); len(got) != 0 {
|
||||
t.Errorf("expected 0 issues when rate provider empty, got %d", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
// --- AgentChecker tests ---
|
||||
|
||||
type mockAgentProvider struct {
|
||||
|
||||
@@ -365,6 +365,7 @@ const (
|
||||
ActionGitHubPRFilesGet = "github.pr_files.get"
|
||||
ActionGitHubPRCommitsGet = "github.pr_commits.get"
|
||||
ActionGitHubTaskPRUpdated = "github.task_pr.updated" // Notification
|
||||
ActionGitHubRateLimitUpdated = "github.rate_limit.updated" // Notification
|
||||
ActionGitHubPRFeedbackNotify = "github.pr_feedback.notify" // Notification
|
||||
ActionGitHubNewReviewPRNotify = "github.new_review_pr.notify" // Notification
|
||||
ActionGitHubTaskPRSync = "github.task_pr.sync"
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
"use client";
|
||||
|
||||
import { useEffect, useState } from "react";
|
||||
import { formatDistanceToNow } from "date-fns";
|
||||
import { IconAlertTriangle } from "@tabler/icons-react";
|
||||
import { Alert, AlertDescription } from "@kandev/ui/alert";
|
||||
import type { GitHubRateLimitInfo, GitHubRateLimitSnapshot } from "@/lib/types/github";
|
||||
|
||||
const RESOURCE_LABELS: Record<string, string> = {
|
||||
core: "REST",
|
||||
graphql: "GraphQL",
|
||||
search: "Search",
|
||||
};
|
||||
|
||||
// useTickNow re-renders every intervalMs and exposes a stable now-value so the
|
||||
// render pass stays pure (Date.now() in the render body would be impure).
|
||||
function useTickNow(intervalMs: number): number {
|
||||
const [now, setNow] = useState(() => Date.now());
|
||||
useEffect(() => {
|
||||
const id = setInterval(() => setNow(Date.now()), intervalMs);
|
||||
return () => clearInterval(id);
|
||||
}, [intervalMs]);
|
||||
return now;
|
||||
}
|
||||
|
||||
function isExhausted(snap: GitHubRateLimitSnapshot, now: number): boolean {
|
||||
if (snap.remaining > 0) return false;
|
||||
const reset = new Date(snap.reset_at).getTime();
|
||||
return Number.isFinite(reset) && reset > now;
|
||||
}
|
||||
|
||||
function formatReset(snap: GitHubRateLimitSnapshot): string {
|
||||
const reset = new Date(snap.reset_at);
|
||||
if (!Number.isFinite(reset.getTime())) return "";
|
||||
return formatDistanceToNow(reset, { addSuffix: true });
|
||||
}
|
||||
|
||||
// latestReset returns the snapshot whose reset_at is furthest in the future.
|
||||
// When multiple buckets are exhausted, background checks remain paused until
|
||||
// the last one recovers, so the alert must anchor on that bucket — anchoring
|
||||
// on snapshots[0] would understate the pause window.
|
||||
function latestReset(snaps: GitHubRateLimitSnapshot[]): GitHubRateLimitSnapshot {
|
||||
return snaps.reduce((latest, s) =>
|
||||
new Date(s.reset_at).getTime() > new Date(latest.reset_at).getTime() ? s : latest,
|
||||
);
|
||||
}
|
||||
|
||||
function snapshotsFromInfo(info: GitHubRateLimitInfo): GitHubRateLimitSnapshot[] {
|
||||
const out: GitHubRateLimitSnapshot[] = [];
|
||||
if (info.core) out.push(info.core);
|
||||
if (info.graphql) out.push(info.graphql);
|
||||
if (info.search) out.push(info.search);
|
||||
return out;
|
||||
}
|
||||
|
||||
export function GitHubRateLimitDisplay({ info }: { info?: GitHubRateLimitInfo }) {
|
||||
const now = useTickNow(30_000);
|
||||
if (!info) return null;
|
||||
const snapshots = snapshotsFromInfo(info);
|
||||
if (snapshots.length === 0) return null;
|
||||
const exhausted = snapshots.filter((s) => isExhausted(s, now));
|
||||
|
||||
return (
|
||||
<div className="space-y-2" data-testid="github-rate-limit-display">
|
||||
{exhausted.length > 0 && (
|
||||
<Alert variant="destructive" data-testid="github-rate-limit-exhausted">
|
||||
<IconAlertTriangle className="h-4 w-4" />
|
||||
<AlertDescription className="text-sm">
|
||||
GitHub API rate limit exhausted on{" "}
|
||||
{exhausted.map((s) => RESOURCE_LABELS[s.resource] ?? s.resource).join(", ")}. Background
|
||||
PR/issue checks are paused until the limit resets {formatReset(latestReset(exhausted))}.
|
||||
</AlertDescription>
|
||||
</Alert>
|
||||
)}
|
||||
<div className="text-xs text-muted-foreground flex flex-wrap gap-x-4 gap-y-1">
|
||||
{snapshots.map((snap) => {
|
||||
const label = RESOURCE_LABELS[snap.resource] ?? snap.resource;
|
||||
const limit = snap.limit > 0 ? snap.limit : "?";
|
||||
const reset = formatReset(snap);
|
||||
return (
|
||||
<span key={snap.resource} data-testid={`github-rate-limit-${snap.resource}`}>
|
||||
{label}: <strong>{snap.remaining}</strong>/{limit}
|
||||
{reset ? ` · resets ${reset}` : ""}
|
||||
</span>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -18,6 +18,7 @@ import { useGitHubStatus } from "@/hooks/domains/github/use-github-status";
|
||||
import { useToast } from "@/components/toast-provider";
|
||||
import { configureGitHubToken, clearGitHubToken } from "@/lib/api/domains/github-api";
|
||||
import type { AuthDiagnostics } from "@/lib/types/github";
|
||||
import { GitHubRateLimitDisplay } from "./github-rate-limit";
|
||||
|
||||
function DiagnosticsOutput({ diagnostics }: { diagnostics: AuthDiagnostics }) {
|
||||
return (
|
||||
@@ -214,6 +215,7 @@ export function GitHubStatusCard() {
|
||||
Token stored in secrets. This token is used by remote agents for GitHub operations.
|
||||
</p>
|
||||
)}
|
||||
<GitHubRateLimitDisplay info={status.rate_limit} />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { create } from "zustand";
|
||||
import { immer } from "zustand/middleware/immer";
|
||||
import { createGitHubSlice } from "./github-slice";
|
||||
import type { GitHubSlice } from "./types";
|
||||
import type { GitHubStatus } from "@/lib/types/github";
|
||||
|
||||
function makeStore() {
|
||||
return create<GitHubSlice>()(immer((...a) => createGitHubSlice(...a)));
|
||||
}
|
||||
|
||||
const FUTURE_RESET = "2030-01-01T00:00:00Z";
|
||||
const NOW = "2026-05-04T12:00:00Z";
|
||||
|
||||
const baseStatus: GitHubStatus = {
|
||||
authenticated: true,
|
||||
username: "octocat",
|
||||
auth_method: "pat",
|
||||
token_configured: true,
|
||||
required_scopes: ["repo"],
|
||||
};
|
||||
|
||||
describe("applyGitHubRateLimitUpdate", () => {
|
||||
it("merges incoming snapshots into the existing status", () => {
|
||||
const store = makeStore();
|
||||
store.getState().setGitHubStatus({ ...baseStatus });
|
||||
|
||||
store.getState().applyGitHubRateLimitUpdate({
|
||||
trigger: "graphql",
|
||||
snapshots: [
|
||||
{
|
||||
resource: "graphql",
|
||||
remaining: 0,
|
||||
limit: 5000,
|
||||
reset_at: FUTURE_RESET,
|
||||
updated_at: NOW,
|
||||
},
|
||||
{
|
||||
resource: "core",
|
||||
remaining: 4500,
|
||||
limit: 5000,
|
||||
reset_at: FUTURE_RESET,
|
||||
updated_at: NOW,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const status = store.getState().githubStatus.status;
|
||||
expect(status?.rate_limit?.graphql?.remaining).toBe(0);
|
||||
expect(status?.rate_limit?.graphql?.limit).toBe(5000);
|
||||
expect(status?.rate_limit?.core?.remaining).toBe(4500);
|
||||
});
|
||||
|
||||
it("overwrites only the resources present in the update", () => {
|
||||
const store = makeStore();
|
||||
store.getState().setGitHubStatus({
|
||||
...baseStatus,
|
||||
rate_limit: {
|
||||
core: {
|
||||
resource: "core",
|
||||
remaining: 4500,
|
||||
limit: 5000,
|
||||
reset_at: FUTURE_RESET,
|
||||
updated_at: NOW,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
store.getState().applyGitHubRateLimitUpdate({
|
||||
trigger: "graphql",
|
||||
snapshots: [
|
||||
{
|
||||
resource: "graphql",
|
||||
remaining: 100,
|
||||
limit: 5000,
|
||||
reset_at: FUTURE_RESET,
|
||||
updated_at: NOW,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const rl = store.getState().githubStatus.status?.rate_limit;
|
||||
expect(rl?.core?.remaining).toBe(4500); // untouched
|
||||
expect(rl?.graphql?.remaining).toBe(100);
|
||||
});
|
||||
|
||||
it("is a no-op when status has not been hydrated yet", () => {
|
||||
const store = makeStore();
|
||||
expect(store.getState().githubStatus.status).toBeNull();
|
||||
|
||||
store.getState().applyGitHubRateLimitUpdate({
|
||||
trigger: "core",
|
||||
snapshots: [
|
||||
{
|
||||
resource: "core",
|
||||
remaining: 0,
|
||||
limit: 5000,
|
||||
reset_at: FUTURE_RESET,
|
||||
updated_at: NOW,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
expect(store.getState().githubStatus.status).toBeNull();
|
||||
});
|
||||
});
|
||||
@@ -157,6 +157,24 @@ function createActionPresetActions(
|
||||
};
|
||||
}
|
||||
|
||||
function createRateLimitActions(set: ImmerSet): Pick<GitHubSlice, "applyGitHubRateLimitUpdate"> {
|
||||
return {
|
||||
applyGitHubRateLimitUpdate: (update) =>
|
||||
set((draft) => {
|
||||
const existing = draft.githubStatus.status;
|
||||
if (!existing) {
|
||||
// Status not yet hydrated; defer until the SSR/HTTP fetch lands.
|
||||
return;
|
||||
}
|
||||
const rateLimit = { ...(existing.rate_limit ?? {}) };
|
||||
for (const snap of update.snapshots) {
|
||||
rateLimit[snap.resource] = snap;
|
||||
}
|
||||
draft.githubStatus.status = { ...existing, rate_limit: rateLimit };
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
export const createGitHubSlice: StateCreator<
|
||||
GitHubSlice,
|
||||
[["zustand/immer", never]],
|
||||
@@ -168,4 +186,5 @@ export const createGitHubSlice: StateCreator<
|
||||
...createTaskPRActions(set),
|
||||
...createWatchActions(set),
|
||||
...createActionPresetActions(set),
|
||||
...createRateLimitActions(set),
|
||||
});
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import type {
|
||||
GitHubStatus,
|
||||
GitHubRateLimitUpdate,
|
||||
TaskPR,
|
||||
PRWatch,
|
||||
ReviewWatch,
|
||||
@@ -70,6 +71,7 @@ export type GitHubSliceActions = {
|
||||
removeIssueWatch: (id: string) => void;
|
||||
setActionPresets: (workspaceId: string, presets: GitHubActionPresets) => void;
|
||||
setActionPresetsLoading: (workspaceId: string, loading: boolean) => void;
|
||||
applyGitHubRateLimitUpdate: (update: GitHubRateLimitUpdate) => void;
|
||||
};
|
||||
|
||||
export type GitHubSlice = GitHubSliceState & GitHubSliceActions;
|
||||
|
||||
@@ -11,6 +11,7 @@ import type {
|
||||
} from "@/lib/types/http";
|
||||
import type {
|
||||
GitHubStatus,
|
||||
GitHubRateLimitUpdate,
|
||||
TaskPR,
|
||||
PRWatch,
|
||||
ReviewWatch as GitHubReviewWatch,
|
||||
@@ -262,6 +263,7 @@ export type AppState = {
|
||||
removeIssueWatch: (id: string) => void;
|
||||
setActionPresets: (workspaceId: string, presets: GitHubActionPresets) => void;
|
||||
setActionPresetsLoading: (workspaceId: string, loading: boolean) => void;
|
||||
applyGitHubRateLimitUpdate: (update: GitHubRateLimitUpdate) => void;
|
||||
|
||||
// JIRA actions
|
||||
setJiraIssueWatches: (watches: JiraIssueWatch[]) => void;
|
||||
|
||||
@@ -63,7 +63,8 @@ export type BackendMessageType =
|
||||
| "secrets.updated"
|
||||
| "secrets.deleted"
|
||||
| "message.queue.status_changed"
|
||||
| "github.task_pr.updated";
|
||||
| "github.task_pr.updated"
|
||||
| "github.rate_limit.updated";
|
||||
|
||||
export type BackendMessage<T extends BackendMessageType, P> = {
|
||||
id?: string;
|
||||
@@ -82,7 +83,7 @@ import type {
|
||||
} from "@/lib/types/http";
|
||||
import type { SecretListItem } from "@/lib/types/http-secrets";
|
||||
import type { GitEventPayload } from "@/lib/types/git-events";
|
||||
import type { TaskPR } from "@/lib/types/github";
|
||||
import type { GitHubRateLimitUpdate, TaskPR } from "@/lib/types/github";
|
||||
import type { FileChangeNotificationPayload } from "./workspace-files";
|
||||
import type {
|
||||
AgentCapabilitiesPayload,
|
||||
@@ -605,6 +606,7 @@ export type BackendMessageMap = {
|
||||
QueueStatusChangedPayload
|
||||
>;
|
||||
"github.task_pr.updated": BackendMessage<"github.task_pr.updated", TaskPR>;
|
||||
"github.rate_limit.updated": BackendMessage<"github.rate_limit.updated", GitHubRateLimitUpdate>;
|
||||
};
|
||||
|
||||
// Workspace file types (extracted to reduce file size)
|
||||
|
||||
@@ -16,6 +16,29 @@ export type GitHubStatus = {
|
||||
token_secret_id?: string;
|
||||
required_scopes: string[];
|
||||
diagnostics?: AuthDiagnostics;
|
||||
rate_limit?: GitHubRateLimitInfo;
|
||||
};
|
||||
|
||||
export type GitHubRateLimitResource = "core" | "graphql" | "search";
|
||||
|
||||
export type GitHubRateLimitSnapshot = {
|
||||
resource: GitHubRateLimitResource;
|
||||
remaining: number;
|
||||
limit: number;
|
||||
reset_at: string;
|
||||
updated_at: string;
|
||||
};
|
||||
|
||||
export type GitHubRateLimitInfo = {
|
||||
core?: GitHubRateLimitSnapshot;
|
||||
graphql?: GitHubRateLimitSnapshot;
|
||||
search?: GitHubRateLimitSnapshot;
|
||||
};
|
||||
|
||||
export type GitHubRateLimitUpdate = {
|
||||
snapshots: GitHubRateLimitSnapshot[];
|
||||
trigger: GitHubRateLimitResource;
|
||||
exhaustion_transition?: "exhausted" | "recovered";
|
||||
};
|
||||
|
||||
export type GitHubPR = {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import type { StoreApi } from "zustand";
|
||||
import type { AppState } from "@/lib/state/store";
|
||||
import type { WsHandlers } from "@/lib/ws/handlers/types";
|
||||
import type { TaskPR } from "@/lib/types/github";
|
||||
import type { GitHubRateLimitUpdate, TaskPR } from "@/lib/types/github";
|
||||
|
||||
export function registerGitHubHandlers(store: StoreApi<AppState>): WsHandlers {
|
||||
return {
|
||||
@@ -11,5 +11,11 @@ export function registerGitHubHandlers(store: StoreApi<AppState>): WsHandlers {
|
||||
store.getState().setTaskPR(pr.task_id, pr);
|
||||
}
|
||||
},
|
||||
"github.rate_limit.updated": (message) => {
|
||||
const update = message.payload as GitHubRateLimitUpdate;
|
||||
if (update?.snapshots?.length) {
|
||||
store.getState().applyGitHubRateLimitUpdate(update);
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user