From 27bec9f35515e125802bdbd1fa5e6c5bc0f9e551 Mon Sep 17 00:00:00 2001 From: Arnaud Meukam Date: Sat, 23 May 2026 14:09:34 +0200 Subject: [PATCH] Implements a Prow plugin triggered by /gemini-agent comments on GitHub issues and PRs. The plugin collects GitHub context (issue metadata, PR diffs, recent comments) and sends it to Gemini on Vertex AI for analysis. Assisted by Opus 4.6 Signed-off-by: Arnaud Meukam --- cmd/external-plugins/geminiagent/main.go | 127 +++++ .../geminiagent/plugin/gemini-agent.go | 506 ++++++++++++++++++ .../geminiagent/plugin/gemini-agent_test.go | 467 ++++++++++++++++ cmd/external-plugins/geminiagent/server.go | 99 ++++ .../geminiagent/server_test.go | 152 ++++++ go.mod | 1 + go.sum | 2 + pkg/plugins/config.go | 16 + pkg/plugins/plugin-config-documented.yaml | 11 + 9 files changed, 1381 insertions(+) create mode 100644 cmd/external-plugins/geminiagent/main.go create mode 100644 cmd/external-plugins/geminiagent/plugin/gemini-agent.go create mode 100644 cmd/external-plugins/geminiagent/plugin/gemini-agent_test.go create mode 100644 cmd/external-plugins/geminiagent/server.go create mode 100644 cmd/external-plugins/geminiagent/server_test.go diff --git a/cmd/external-plugins/geminiagent/main.go b/cmd/external-plugins/geminiagent/main.go new file mode 100644 index 0000000000..3f2431ebb9 --- /dev/null +++ b/cmd/external-plugins/geminiagent/main.go @@ -0,0 +1,127 @@ +/* +Copyright The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "flag" + "fmt" + "net/http" + "os" + "strconv" + "time" + + "github.com/sirupsen/logrus" + + "sigs.k8s.io/prow/cmd/external-plugins/geminiagent/plugin" + "sigs.k8s.io/prow/pkg/config/secret" + "sigs.k8s.io/prow/pkg/flagutil" + prowflagutil "sigs.k8s.io/prow/pkg/flagutil" + pluginsflagutil "sigs.k8s.io/prow/pkg/flagutil/plugins" + "sigs.k8s.io/prow/pkg/interrupts" + "sigs.k8s.io/prow/pkg/logrusutil" + "sigs.k8s.io/prow/pkg/pjutil" + "sigs.k8s.io/prow/pkg/pluginhelp/externalplugins" +) + +type options struct { + port int + + pluginsConfig pluginsflagutil.PluginOptions + dryRun bool + github prowflagutil.GitHubOptions + instrumentationOptions prowflagutil.InstrumentationOptions + logLevel string + + webhookSecretFile string +} + +const defaultHourlyTokens = 360 + +func (o *options) Validate() error { + for idx, group := range []flagutil.OptionGroup{&o.github} { + if err := group.Validate(o.dryRun); err != nil { + return fmt.Errorf("%d: %w", idx, err) + } + } + + return nil +} + +func gatherOptions() options { + o := options{} + fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError) + fs.IntVar(&o.port, "port", 8888, "Port to listen on.") + fs.BoolVar(&o.dryRun, "dry-run", true, "Dry run for testing. Uses API tokens but does not mutate.") + fs.StringVar(&o.webhookSecretFile, "hmac-secret-file", "/etc/webhook/hmac", "Path to the file containing the GitHub HMAC secret.") + fs.StringVar(&o.logLevel, "log-level", "debug", fmt.Sprintf("Log level is one of %v.", logrus.AllLevels)) + + o.github.AddCustomizedFlags(fs, prowflagutil.ThrottlerDefaults(defaultHourlyTokens, defaultHourlyTokens)) + + o.pluginsConfig.PluginConfigPathDefault = "/etc/plugins/plugins.yaml" + for _, group := range []flagutil.OptionGroup{&o.instrumentationOptions, &o.pluginsConfig} { + group.AddFlags(fs) + } + fs.Parse(os.Args[1:]) + return o +} + +func main() { + logrusutil.ComponentInit() + o := gatherOptions() + if err := o.Validate(); err != nil { + logrus.Fatalf("Invalid options: %v", err) + } + + logLevel, err := logrus.ParseLevel(o.logLevel) + if err != nil { + logrus.WithError(err).Fatal("Failed to parse loglevel") + } + logrus.SetLevel(logLevel) + log := logrus.StandardLogger().WithField("plugin", plugin.PluginName) + + if err := secret.Add(o.webhookSecretFile); err != nil { + logrus.WithError(err).Fatal("Error starting secrets agent.") + } + + pluginAgent, err := o.pluginsConfig.PluginAgent() + if err != nil { + log.WithError(err).Fatal("Error loading plugin config") + } + + githubClient, err := o.github.GitHubClient(o.dryRun) + if err != nil { + logrus.WithError(err).Fatal("Error getting GitHub client.") + } + + server := &Server{ + tokenGenerator: secret.GetTokenGenerator(o.webhookSecretFile), + ghc: githubClient, + log: log, + pluginAgent: pluginAgent, + handleGenericComment: plugin.HandleGenericComment, + } + + health := pjutil.NewHealthOnPort(o.instrumentationOptions.HealthPort) + health.ServeReady() + + mux := http.NewServeMux() + mux.Handle("/", server) + externalplugins.ServeExternalPluginHelp(mux, log, plugin.HelpProvider) + httpServer := &http.Server{Addr: ":" + strconv.Itoa(o.port), Handler: mux} + defer interrupts.WaitForGracefulShutdown() + interrupts.ListenAndServe(httpServer, 5*time.Second) +} diff --git a/cmd/external-plugins/geminiagent/plugin/gemini-agent.go b/cmd/external-plugins/geminiagent/plugin/gemini-agent.go new file mode 100644 index 0000000000..b67b4ac17e --- /dev/null +++ b/cmd/external-plugins/geminiagent/plugin/gemini-agent.go @@ -0,0 +1,506 @@ +/* +Copyright The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package plugin sends /gemini-agent requests to Gemini on Vertex AI. +package plugin + +import ( + "context" + "errors" + "fmt" + "math/rand/v2" + "net/http" + "os" + "regexp" + "slices" + "strings" + "time" + "unicode/utf8" + + "github.com/sirupsen/logrus" + "golang.org/x/oauth2/google" + "golang.org/x/time/rate" + "google.golang.org/genai" + + prowconfig "sigs.k8s.io/prow/pkg/config" + "sigs.k8s.io/prow/pkg/github" + "sigs.k8s.io/prow/pkg/markdown" + "sigs.k8s.io/prow/pkg/pluginhelp" + "sigs.k8s.io/prow/pkg/plugins" +) + +const ( + // PluginName is the configured external plugin name. + PluginName = "gemini-agent" + + defaultModel = "gemma-4-26b-a4b-it" + defaultLocation = "global" + cloudPlatformScope = "https://www.googleapis.com/auth/cloud-platform" + + maxContextComments = 20 + maxPRFiles = 25 + maxPatchBytes = 120_000 + maxResponseBytes = 60_000 + requestTimeout = 10 * time.Minute + + // Rate limiting: conservative for Tier 1 free accounts (typically 15 RPM + // for flash models). + defaultRPM = 10 + maxRetries = 3 + initialBackoff = 2 * time.Second + maxBackoff = 60 * time.Second +) + +var geminiAgentRe = regexp.MustCompile(`(?m)^/gemini-agent(?:[ \t]+(.+))?[ \t]*$`) + +// resolvedConfig is GeminiAgentConfig with defaults applied. +type resolvedConfig struct { + Model string + AllowedTeams []string +} + +func resolveConfig(cfg plugins.GeminiAgentConfig) resolvedConfig { + rc := resolvedConfig{ + Model: cfg.Model, + AllowedTeams: cfg.AllowedTeams, + } + if rc.Model == "" { + rc.Model = firstNonEmpty(os.Getenv("AI_AGENT_MODEL"), defaultModel) + } + return rc +} + +// geminiLimiter is a process-wide token bucket that proactively stays within +// the Gemini RPM quota. Requests that exceed the limit wait rather than +// hammering the API and eating 429s. +var geminiLimiter = rate.NewLimiter(rate.Limit(float64(defaultRPM)/60.0), defaultRPM) + +// GitHubClient is the GitHub API surface used by the Gemini agent. +type GitHubClient interface { + CreateComment(owner, repo string, number int, comment string) error + GetIssue(org, repo string, number int) (*github.Issue, error) + GetPullRequest(org, repo string, number int) (*github.PullRequest, error) + GetPullRequestChanges(org, repo string, number int) ([]github.PullRequestChange, error) + ListIssueComments(org, repo string, number int) ([]github.IssueComment, error) + TeamBySlugHasMember(org string, teamSlug string, memberLogin string) (bool, error) + IsMember(org, user string) (bool, error) + IsCollaborator(org, repo, user string) (bool, error) +} + +type geminiClient interface { + GenerateContent(ctx context.Context, prompt string) (string, error) +} + +type vertexGeminiClient struct { + client *genai.Client + model string +} + +// safetySettings enforces content filtering on all Gemini responses. +// Gemini 3.x defaults to OFF (no filtering), so we explicitly enable it. +// BLOCK_MEDIUM_AND_ABOVE is restrictive enough to stop harmful content while +// not tripping on code security discussions (vulnerability reports, CVEs, etc). +var safetySettings = []*genai.SafetySetting{ + {Category: genai.HarmCategoryHarassment, Threshold: genai.HarmBlockThresholdBlockMediumAndAbove}, + {Category: genai.HarmCategoryHateSpeech, Threshold: genai.HarmBlockThresholdBlockMediumAndAbove}, + {Category: genai.HarmCategorySexuallyExplicit, Threshold: genai.HarmBlockThresholdBlockMediumAndAbove}, + {Category: genai.HarmCategoryDangerousContent, Threshold: genai.HarmBlockThresholdBlockMediumAndAbove}, +} + +// isRateLimited checks whether an error is a 429 from the Gemini API. +func isRateLimited(err error) bool { + var apiErr genai.APIError + if errors.As(err, &apiErr) { + return apiErr.Code == http.StatusTooManyRequests + } + return false +} + +// HandleGenericComment handles a generalized GitHub comment event. +func HandleGenericComment(ghc GitHubClient, pluginConfig *plugins.Configuration, log *logrus.Entry, e github.GenericCommentEvent) error { + if e.Action != github.GenericCommentActionCreated { + return nil + } + task, ok := parseGeminiAgentTask(e.Body) + if !ok { + return nil + } + if task == "" { + return respond(ghc, e, "Usage: `/gemini-agent `") + } + + org := e.Repo.Owner.Login + repo := e.Repo.Name + + cfg := geminiAgentFor(pluginConfig, org, repo) + + trusted, err := isAllowed(ghc, cfg, org, repo, e.User.Login) + if err != nil { + return err + } + if !trusted { + return respond(ghc, e, "`/gemini-agent` can only be used by org members, repository collaborators, or members of configured teams.") + } + + ctx, cancel := context.WithTimeout(context.Background(), requestTimeout) + defer cancel() + + gemini, err := newVertexGeminiClient(ctx, cfg.Model) + if err != nil { + log.WithError(err).Error("Failed to initialize Gemini client.") + if commentErr := respond(ghc, e, "Failed to initialize the AI backend. Please contact the platform team if this persists."); commentErr != nil { + return errors.Join(err, commentErr) + } + return err + } + + return runAgent(ctx, ghc, newRateLimitedClient(gemini), log, e, task) +} + +func parseGeminiAgentTask(body string) (string, bool) { + matches := geminiAgentRe.FindStringSubmatch(markdown.DropCodeBlock(body)) + if matches == nil { + return "", false + } + if len(matches) < 2 { + return "", true + } + return strings.TrimSpace(matches[1]), true +} + +// geminiAgentFor finds the GeminiAgent config for a repo. Prioritizes +// repo-level ("org/repo") over org-level ("org") entries. +func geminiAgentFor(pc *plugins.Configuration, org, repo string) resolvedConfig { + fullName := fmt.Sprintf("%s/%s", org, repo) + for _, cfg := range pc.GeminiAgents { + if slices.Contains(cfg.Repos, fullName) { + return resolveConfig(cfg) + } + } + for _, cfg := range pc.GeminiAgents { + if slices.Contains(cfg.Repos, org) { + return resolveConfig(cfg) + } + } + return resolveConfig(plugins.GeminiAgentConfig{}) +} + +// isAllowed checks if a user can invoke /gemini-agent. A user is allowed if +// they are an org member, a repository collaborator, or a member of any +// configured AllowedTeams. +func isAllowed(ghc GitHubClient, cfg resolvedConfig, org, repo, user string) (bool, error) { + // Org members are always trusted. + if member, err := ghc.IsMember(org, user); err != nil { + return false, fmt.Errorf("check org membership: %w", err) + } else if member { + return true, nil + } + + // Repository collaborators are trusted. + if ok, err := ghc.IsCollaborator(org, repo, user); err != nil { + return false, fmt.Errorf("check collaborator: %w", err) + } else if ok { + return true, nil + } + + // Check configured team membership. + for _, team := range cfg.AllowedTeams { + isMember, err := ghc.TeamBySlugHasMember(org, team, user) + if err != nil { + return false, fmt.Errorf("check team %s membership: %w", team, err) + } + if isMember { + return true, nil + } + } + + return false, nil +} + +func newVertexGeminiClient(ctx context.Context, model string) (*vertexGeminiClient, error) { + projectID := firstNonEmpty(os.Getenv("AI_AGENT_PROJECT_ID"), os.Getenv("GOOGLE_CLOUD_PROJECT"), os.Getenv("GCLOUD_PROJECT")) + location := firstNonEmpty(os.Getenv("AI_AGENT_LOCATION"), os.Getenv("GOOGLE_CLOUD_LOCATION"), defaultLocation) + if projectID == "" { + creds, err := google.FindDefaultCredentials(ctx, cloudPlatformScope) + if err == nil { + projectID = creds.ProjectID + } + } + + client, err := genai.NewClient(ctx, &genai.ClientConfig{ + Backend: genai.BackendVertexAI, + Project: projectID, + Location: location, + }) + if err != nil { + return nil, err + } + return &vertexGeminiClient{client: client, model: model}, nil +} + +func (c *vertexGeminiClient) GenerateContent(ctx context.Context, prompt string) (string, error) { + resp, err := c.client.Models.GenerateContent(ctx, c.model, genai.Text(prompt), &genai.GenerateContentConfig{ + SafetySettings: safetySettings, + }) + if err != nil { + return "", err + } + text := strings.TrimSpace(resp.Text()) + if text == "" { + return "", errors.New("response was blocked by content safety filters") + } + return text, nil +} + +// rateLimitedClient wraps a geminiClient with proactive rate limiting (token +// bucket) and reactive retry with exponential backoff on 429 responses. +type rateLimitedClient struct { + inner geminiClient + limiter *rate.Limiter + maxRetries int + initialBackoff time.Duration + maxBackoff time.Duration +} + +func newRateLimitedClient(inner geminiClient) *rateLimitedClient { + return &rateLimitedClient{ + inner: inner, + limiter: geminiLimiter, + maxRetries: maxRetries, + initialBackoff: initialBackoff, + maxBackoff: maxBackoff, + } +} + +func (c *rateLimitedClient) GenerateContent(ctx context.Context, prompt string) (string, error) { + if err := c.limiter.Wait(ctx); err != nil { + return "", fmt.Errorf("rate limiter: %w", err) + } + + var lastErr error + for attempt := range c.maxRetries { + result, err := c.inner.GenerateContent(ctx, prompt) + if err == nil { + return result, nil + } + + if !isRateLimited(err) { + return "", err + } + + lastErr = err + backoff := c.retryBackoff(attempt) + select { + case <-time.After(backoff): + case <-ctx.Done(): + return "", fmt.Errorf("context cancelled while retrying rate limit (attempt %d): %w", attempt+1, ctx.Err()) + } + } + + return "", fmt.Errorf("rate limited after %d retries: %w", c.maxRetries, lastErr) +} + +func (c *rateLimitedClient) retryBackoff(attempt int) time.Duration { + backoff := min(c.initialBackoff*(1<= maxPRFiles { + fmt.Fprintf(b, "\n%d additional files omitted.\n\n", len(changes)-maxPRFiles) + return + } + fmt.Fprintf(b, "- %s (%s, +%d/-%d)\n", change.Filename, change.Status, change.Additions, change.Deletions) + if change.PreviousFilename != "" { + fmt.Fprintf(b, " previous filename: %s\n", change.PreviousFilename) + } + if change.Patch == "" { + continue + } + if remainingPatchBytes <= 0 { + fmt.Fprintln(b, " patch omitted: aggregate patch limit reached") + continue + } + patch := change.Patch + if len(patch) > remainingPatchBytes { + patch = truncate(patch, remainingPatchBytes) + } + remainingPatchBytes -= len(patch) + fmt.Fprintf(b, " patch:\n~~~diff\n%s\n~~~\n", patch) + } + fmt.Fprintln(b) +} + +func writeRecentComments(b *strings.Builder, comments []github.IssueComment) { + if len(comments) == 0 { + return + } + + start := max(len(comments)-maxContextComments, 0) + fmt.Fprintf(b, "Recent comments (%d total, showing %d):\n", len(comments), len(comments)-start) + for _, comment := range comments[start:] { + created := "" + if !comment.CreatedAt.IsZero() { + created = " at " + comment.CreatedAt.Format(time.RFC3339) + } + fmt.Fprintf(b, "- @%s%s:\n%s\n", comment.User.Login, created, truncate(comment.Body, maxPatchBytes/20)) + } +} + +func buildPrompt(task, githubContext string) string { + return fmt.Sprintf(`You are a senior Go developer and Kubernetes/Prow platform engineer. + +Answer the requested task using the GitHub issue or pull request context below. +Be concise, technical, and safe. If the request needs repository changes that you cannot make from this Vertex AI call, explain the exact next implementation steps instead of pretending the work was done. + +Requested task: +%s + +GitHub context: +%s + +Return a GitHub comment in Markdown.`, task, githubContext) +} + +func respond(ghc GitHubClient, e github.GenericCommentEvent, response string) error { + org := e.Repo.Owner.Login + repo := e.Repo.Name + return ghc.CreateComment(org, repo, e.Number, plugins.FormatResponseRaw(e.Body, e.HTMLURL, e.User.Login, response)) +} + +// HelpProvider constructs help for the external plugin help endpoint. +func HelpProvider(_ []prowconfig.OrgRepo) (*pluginhelp.PluginHelp, error) { + pluginHelp := &pluginhelp.PluginHelp{ + Description: "The gemini-agent plugin sends `/gemini-agent` requests and GitHub context to Gemini on Vertex AI.", + } + pluginHelp.AddCommand(pluginhelp.Command{ + Usage: "/gemini-agent ", + Description: "Ask Gemini on Vertex AI to analyze an issue or pull request with surrounding GitHub context.", + Featured: true, + WhoCanUse: "Org members, repository collaborators, or members of configured teams.", + Examples: []string{"/gemini-agent Explain why this PR is failing tests.", "/gemini-agent Summarize the risk in this change."}, + }) + return pluginHelp, nil +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + return "" +} + +func truncate(value string, limit int) string { + if limit <= 0 || len(value) <= limit { + return value + } + // Back up from the byte limit to avoid splitting a multi-byte UTF-8 char. + for limit > 0 && !utf8.RuneStart(value[limit]) { + limit-- + } + return value[:limit] + "\n...[truncated]" +} diff --git a/cmd/external-plugins/geminiagent/plugin/gemini-agent_test.go b/cmd/external-plugins/geminiagent/plugin/gemini-agent_test.go new file mode 100644 index 0000000000..7d12bf23e1 --- /dev/null +++ b/cmd/external-plugins/geminiagent/plugin/gemini-agent_test.go @@ -0,0 +1,467 @@ +/* +Copyright The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package plugin + +import ( + "context" + "errors" + "fmt" + "net/http" + "strings" + "testing" + "time" + + "github.com/sirupsen/logrus" + "golang.org/x/time/rate" + "google.golang.org/genai" + + "sigs.k8s.io/prow/pkg/github" + "sigs.k8s.io/prow/pkg/plugins" +) + +type fakeGitHubClient struct { + issue *github.Issue + pr *github.PullRequest + changes []github.PullRequestChange + comments []github.IssueComment + + createdComments []string + isMember bool + isCollaborator bool +} + +func (f *fakeGitHubClient) CreateComment(_, _ string, _ int, comment string) error { + f.createdComments = append(f.createdComments, comment) + return nil +} + +func (f *fakeGitHubClient) GetIssue(_, _ string, _ int) (*github.Issue, error) { + if f.issue == nil { + return nil, errors.New("unexpected GetIssue call") + } + return f.issue, nil +} + +func (f *fakeGitHubClient) GetPullRequest(_, _ string, _ int) (*github.PullRequest, error) { + if f.pr == nil { + return nil, errors.New("unexpected GetPullRequest call") + } + return f.pr, nil +} + +func (f *fakeGitHubClient) GetPullRequestChanges(_, _ string, _ int) ([]github.PullRequestChange, error) { + return f.changes, nil +} + +func (f *fakeGitHubClient) ListIssueComments(_, _ string, _ int) ([]github.IssueComment, error) { + return f.comments, nil +} + +func (f *fakeGitHubClient) TeamBySlugHasMember(_, _ string, _ string) (bool, error) { + return false, nil +} + +func (f *fakeGitHubClient) IsMember(_, _ string) (bool, error) { + return f.isMember, nil +} + +func (f *fakeGitHubClient) IsCollaborator(_, _, _ string) (bool, error) { + return f.isCollaborator, nil +} + +type fakeGeminiClient struct { + prompt string + response string +} + +func (f *fakeGeminiClient) GenerateContent(_ context.Context, prompt string) (string, error) { + f.prompt = prompt + return f.response, nil +} + +func TestParseGeminiAgentTask(t *testing.T) { + tests := []struct { + name string + body string + expected string + matched bool + }{ + { + name: "task on command line", + body: "/gemini-agent explain this failure", + expected: "explain this failure", + matched: true, + }, + { + name: "empty task", + body: "/gemini-agent", + matched: true, + }, + { + name: "ignores code block", + body: "```text\n/gemini-agent not a command\n```", + }, + { + name: "unrelated comment", + body: "/retest", + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + actual, matched := parseGeminiAgentTask(test.body) + if matched != test.matched { + t.Fatalf("matched %t, expected %t", matched, test.matched) + } + if actual != test.expected { + t.Fatalf("task %q, expected %q", actual, test.expected) + } + }) + } +} + +func TestRunAgentPullRequestContext(t *testing.T) { + ghc := &fakeGitHubClient{ + pr: &github.PullRequest{ + Title: "Fix the frobnicator", + Body: "This changes the frobnicator retry path.", + HTMLURL: "https://github.example/pr/7", + User: github.User{Login: "alice"}, + Base: github.PullRequestBranch{Ref: "main", SHA: "base-sha"}, + Head: github.PullRequestBranch{Ref: "feature", SHA: "head-sha"}, + }, + changes: []github.PullRequestChange{ + { + Filename: "pkg/frob/frob.go", + Status: string(github.PullRequestFileModified), + Additions: 3, + Deletions: 1, + Patch: "@@ -1 +1 @@\n-old\n+new", + }, + }, + comments: []github.IssueComment{ + { + Body: "previous reviewer context", + User: github.User{Login: "reviewer"}, + CreatedAt: time.Date(2026, 5, 19, 12, 0, 0, 0, time.UTC), + }, + }, + } + gemini := &fakeGeminiClient{response: "Gemini says the retry path is risky."} + event := github.GenericCommentEvent{ + Action: github.GenericCommentActionCreated, + Body: "/gemini-agent summarize risk", + HTMLURL: "https://github.example/comment/1", + Number: 7, + IsPR: true, + IssueState: "open", + IssueTitle: "stale title should be replaced by PR title", + IssueBody: "stale body should be replaced by PR body", + IssueHTMLURL: "https://github.example/issues/7", + IssueAuthor: github.User{Login: "issue-author"}, + User: github.User{Login: "bob"}, + Repo: github.Repo{ + Name: "repo", + Owner: github.User{Login: "org"}, + }, + } + + if err := runAgent(context.Background(), ghc, gemini, logrus.NewEntry(logrus.New()), event, "summarize risk"); err != nil { + t.Fatalf("runAgent returned error: %v", err) + } + + for _, expected := range []string{ + "Requested task:\nsummarize risk", + "Repository: org/repo", + "Pull request base: main@base-sha", + "pkg/frob/frob.go", + "previous reviewer context", + } { + if !strings.Contains(gemini.prompt, expected) { + t.Fatalf("Gemini prompt missing %q:\n%s", expected, gemini.prompt) + } + } + + if len(ghc.createdComments) != 1 { + t.Fatalf("created %d comments, expected 1", len(ghc.createdComments)) + } + if !strings.Contains(ghc.createdComments[0], "Gemini says the retry path is risky.") { + t.Fatalf("created comment missing Gemini response:\n%s", ghc.createdComments[0]) + } +} + +// fakeRateLimitedGeminiClient simulates 429 responses followed by success. +type fakeRateLimitedGeminiClient struct { + failCount int // number of 429s to return before succeeding + calls int + response string +} + +func (f *fakeRateLimitedGeminiClient) GenerateContent(_ context.Context, _ string) (string, error) { + f.calls++ + if f.calls <= f.failCount { + return "", genai.APIError{ + Code: http.StatusTooManyRequests, + Message: "Resource exhausted", + Status: "429 Too Many Requests", + } + } + return f.response, nil +} + +func TestIsRateLimited(t *testing.T) { + tests := []struct { + name string + err error + expected bool + }{ + { + name: "429 APIError", + err: genai.APIError{Code: http.StatusTooManyRequests, Message: "rate limited"}, + expected: true, + }, + { + name: "500 APIError", + err: genai.APIError{Code: http.StatusInternalServerError, Message: "server error"}, + expected: false, + }, + { + name: "wrapped 429 APIError", + err: fmt.Errorf("call failed: %w", genai.APIError{Code: http.StatusTooManyRequests}), + expected: true, + }, + { + name: "generic error", + err: errors.New("network timeout"), + expected: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + if got := isRateLimited(tc.err); got != tc.expected { + t.Errorf("isRateLimited() = %v, want %v", got, tc.expected) + } + }) + } +} + +func TestRetryBackoff(t *testing.T) { + rl := &rateLimitedClient{ + initialBackoff: 2 * time.Second, + maxBackoff: 60 * time.Second, + } + for attempt := range 4 { + d := rl.retryBackoff(attempt) + // Should be at least initialBackoff * 2^attempt and capped at maxBackoff. + minExpected := min(rl.initialBackoff*(1< maxExpected { + t.Errorf("attempt %d: backoff %v > expected maximum %v", attempt, d, maxExpected) + } + } +} + +// newTestRateLimitedClient creates a rateLimitedClient with minimal backoff for fast tests. +func newTestRateLimitedClient(inner geminiClient) *rateLimitedClient { + return &rateLimitedClient{ + inner: inner, + limiter: rate.NewLimiter(rate.Inf, 1), // no rate limit in tests + maxRetries: maxRetries, + initialBackoff: time.Millisecond, + maxBackoff: 5 * time.Millisecond, + } +} + +func TestRunAgentRetriesOnRateLimit(t *testing.T) { + gemini := &fakeRateLimitedGeminiClient{ + failCount: 2, + response: "Success after retries", + } + ghc := &fakeGitHubClient{ + issue: &github.Issue{ + Title: "Test issue", + Body: "Test body", + HTMLURL: "https://github.example/issues/1", + User: github.User{Login: "alice"}, + }, + comments: []github.IssueComment{}, + } + event := github.GenericCommentEvent{ + Action: github.GenericCommentActionCreated, + Body: "/gemini-agent do the thing", + HTMLURL: "https://github.example/comment/1", + Number: 1, + IssueState: "open", + IssueTitle: "Test issue", + IssueBody: "Test body", + User: github.User{Login: "bob"}, + Repo: github.Repo{ + Name: "repo", + Owner: github.User{Login: "org"}, + }, + } + + err := runAgent(context.Background(), ghc, newTestRateLimitedClient(gemini), logrus.NewEntry(logrus.New()), event, "do the thing") + if err != nil { + t.Fatalf("runAgent returned error: %v", err) + } + if gemini.calls != 3 { + t.Errorf("expected 3 calls (2 retries + 1 success), got %d", gemini.calls) + } + if len(ghc.createdComments) != 1 || !strings.Contains(ghc.createdComments[0], "Success after retries") { + t.Errorf("unexpected comment: %v", ghc.createdComments) + } +} + +func TestRunAgentExhaustsRetries(t *testing.T) { + gemini := &fakeRateLimitedGeminiClient{ + failCount: maxRetries + 1, // more failures than retries allowed + response: "never reached", + } + ghc := &fakeGitHubClient{ + issue: &github.Issue{ + Title: "Test issue", + Body: "Test body", + HTMLURL: "https://github.example/issues/1", + User: github.User{Login: "alice"}, + }, + comments: []github.IssueComment{}, + } + event := github.GenericCommentEvent{ + Action: github.GenericCommentActionCreated, + Body: "/gemini-agent do the thing", + HTMLURL: "https://github.example/comment/1", + Number: 1, + IssueState: "open", + IssueTitle: "Test issue", + IssueBody: "Test body", + User: github.User{Login: "bob"}, + Repo: github.Repo{ + Name: "repo", + Owner: github.User{Login: "org"}, + }, + } + + err := runAgent(context.Background(), ghc, newTestRateLimitedClient(gemini), logrus.NewEntry(logrus.New()), event, "do the thing") + // runAgent returns the error after posting the failure comment. + if err == nil { + t.Fatal("expected error from runAgent after exhausting retries") + } + if !strings.Contains(err.Error(), "rate limited after") { + t.Errorf("expected rate limit exhaustion error, got: %v", err) + } + // Should have posted an error comment. + if len(ghc.createdComments) != 1 { + t.Fatalf("expected 1 error comment, got %d", len(ghc.createdComments)) + } + if !strings.Contains(ghc.createdComments[0], "The AI request failed") { + t.Errorf("expected failure message in comment, got: %s", ghc.createdComments[0]) + } +} + +func TestGeminiAgentForLookup(t *testing.T) { + pc := &plugins.Configuration{ + GeminiAgents: []plugins.GeminiAgentConfig{ + { + Repos: []string{"org/specific-repo"}, + Model: "gemma4-27b", + AllowedTeams: []string{"ml-team"}, + }, + { + Repos: []string{"org"}, + Model: "gemini-3-flash", + AllowedTeams: []string{"platform-team"}, + }, + }, + } + + // Repo-level match wins over org-level. + cfg := geminiAgentFor(pc, "org", "specific-repo") + if cfg.Model != "gemma4-27b" { + t.Errorf("expected gemma4-27b for org/specific-repo, got %s", cfg.Model) + } + if len(cfg.AllowedTeams) != 1 || cfg.AllowedTeams[0] != "ml-team" { + t.Errorf("expected ml-team for org/specific-repo, got %v", cfg.AllowedTeams) + } + + // Org-level fallback. + cfg = geminiAgentFor(pc, "org", "other-repo") + if cfg.Model != "gemini-3-flash" { + t.Errorf("expected gemini-3-flash for org/other-repo, got %s", cfg.Model) + } + + // No config: defaults. + cfg = geminiAgentFor(pc, "unknown-org", "repo") + if cfg.Model != defaultModel { + t.Errorf("expected default model for unknown org, got %s", cfg.Model) + } +} + +func TestGeminiAgentForDefaultModel(t *testing.T) { + pc := &plugins.Configuration{} + cfg := geminiAgentFor(pc, "org", "repo") + if cfg.Model != defaultModel { + t.Errorf("expected %s, got %s", defaultModel, cfg.Model) + } +} + +func TestIsAllowed(t *testing.T) { + tests := []struct { + name string + isMember bool + isCollaborator bool + allowedTeams []string + expected bool + }{ + { + name: "org member is allowed", + isMember: true, + expected: true, + }, + { + name: "collaborator is allowed", + isCollaborator: true, + expected: true, + }, + { + name: "stranger is denied", + expected: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ghc := &fakeGitHubClient{ + isMember: tc.isMember, + isCollaborator: tc.isCollaborator, + } + cfg := resolvedConfig{AllowedTeams: tc.allowedTeams} + allowed, err := isAllowed(ghc, cfg, "org", "repo", "user") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if allowed != tc.expected { + t.Errorf("isAllowed() = %v, want %v", allowed, tc.expected) + } + }) + } +} diff --git a/cmd/external-plugins/geminiagent/server.go b/cmd/external-plugins/geminiagent/server.go new file mode 100644 index 0000000000..d5b923720c --- /dev/null +++ b/cmd/external-plugins/geminiagent/server.go @@ -0,0 +1,99 @@ +/* +Copyright The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "encoding/json" + "errors" + "fmt" + "net/http" + + "github.com/sirupsen/logrus" + + "sigs.k8s.io/prow/cmd/external-plugins/geminiagent/plugin" + "sigs.k8s.io/prow/pkg/github" + "sigs.k8s.io/prow/pkg/plugins" +) + +type genericCommentHandler func(plugin.GitHubClient, *plugins.Configuration, *logrus.Entry, github.GenericCommentEvent) error + +// Server implements http.Handler. It validates incoming GitHub webhooks and +// dispatches issue comments to the Gemini agent. +type Server struct { + tokenGenerator func() []byte + ghc plugin.GitHubClient + log *logrus.Entry + pluginAgent *plugins.ConfigAgent + + handleGenericComment genericCommentHandler +} + +// ServeHTTP validates an incoming webhook and dispatches it. +func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { + eventType, eventGUID, payload, ok, _ := github.ValidateWebhook(w, r, s.tokenGenerator) + if !ok { + return + } + fmt.Fprint(w, "Event received. Have a nice day.") + + if err := s.handleEvent(eventType, eventGUID, payload); err != nil { + s.log.WithError(err).Error("Error parsing event.") + } +} + +func (s *Server) handleEvent(eventType, eventGUID string, payload []byte) error { + l := s.log.WithFields(logrus.Fields{ + "event-type": eventType, + github.EventGUID: eventGUID, + }) + switch eventType { + case "issue_comment": + var ice github.IssueCommentEvent + if err := json.Unmarshal(payload, &ice); err != nil { + return err + } + ice.GUID = eventGUID + go func() { + if err := s.handleIssueComment(l, ice); err != nil { + l.WithError(err).Info("Gemini agent failed.") + } + }() + default: + l.Debugf("skipping event of type %q", eventType) + } + return nil +} + +func (s *Server) handleIssueComment(l *logrus.Entry, ice github.IssueCommentEvent) error { + if s.handleGenericComment == nil { + return errors.New("generic comment handler is nil") + } + + event, err := github.GeneralizeComment(ice) + if err != nil { + return err + } + + l = l.WithFields(logrus.Fields{ + github.OrgLogField: event.Repo.Owner.Login, + github.RepoLogField: event.Repo.Name, + github.PrLogField: event.Number, + "commenter": event.User.Login, + "url": event.HTMLURL, + }) + return s.handleGenericComment(s.ghc, s.pluginAgent.Config(), l, *event) +} diff --git a/cmd/external-plugins/geminiagent/server_test.go b/cmd/external-plugins/geminiagent/server_test.go new file mode 100644 index 0000000000..64b1b19dad --- /dev/null +++ b/cmd/external-plugins/geminiagent/server_test.go @@ -0,0 +1,152 @@ +/* +Copyright The Kubernetes Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "errors" + "testing" + + "github.com/sirupsen/logrus" + + "sigs.k8s.io/prow/cmd/external-plugins/geminiagent/plugin" + "sigs.k8s.io/prow/pkg/github" + "sigs.k8s.io/prow/pkg/plugins" +) + +type fakeGitHubClient struct{} + +func (f *fakeGitHubClient) CreateComment(_, _ string, _ int, _ string) error { + return errors.New("unexpected CreateComment call") +} + +func (f *fakeGitHubClient) GetIssue(_, _ string, _ int) (*github.Issue, error) { + return nil, errors.New("unexpected GetIssue call") +} + +func (f *fakeGitHubClient) GetPullRequest(_, _ string, _ int) (*github.PullRequest, error) { + return nil, errors.New("unexpected GetPullRequest call") +} + +func (f *fakeGitHubClient) GetPullRequestChanges(_, _ string, _ int) ([]github.PullRequestChange, error) { + return nil, errors.New("unexpected GetPullRequestChanges call") +} + +func (f *fakeGitHubClient) ListIssueComments(_, _ string, _ int) ([]github.IssueComment, error) { + return nil, errors.New("unexpected ListIssueComments call") +} + +func (f *fakeGitHubClient) TeamBySlugHasMember(_, _, _ string) (bool, error) { + return false, errors.New("unexpected TeamBySlugHasMember call") +} + +func (f *fakeGitHubClient) IsMember(_, _ string) (bool, error) { + return false, errors.New("unexpected IsMember call") +} + +func (f *fakeGitHubClient) IsCollaborator(_, _, _ string) (bool, error) { + return false, errors.New("unexpected IsCollaborator call") +} + +func TestHandleIssueCommentGeneralizesAndDispatches(t *testing.T) { + pluginAgent := plugins.NewFakeConfigAgent() + pluginAgent.Set(&plugins.Configuration{ + GeminiAgents: []plugins.GeminiAgentConfig{ + {Repos: []string{"org/repo"}, Model: "gemini-test-model"}, + }, + }) + + var got github.GenericCommentEvent + server := &Server{ + ghc: &fakeGitHubClient{}, + log: logrus.NewEntry(logrus.New()), + pluginAgent: &pluginAgent, + handleGenericComment: func(_ plugin.GitHubClient, cfg *plugins.Configuration, _ *logrus.Entry, event github.GenericCommentEvent) error { + got = event + if len(cfg.GeminiAgents) != 1 || cfg.GeminiAgents[0].Model != "gemini-test-model" { + t.Fatalf("unexpected plugin config: %#v", cfg.GeminiAgents) + } + return nil + }, + } + + err := server.handleIssueComment(logrus.NewEntry(logrus.New()), github.IssueCommentEvent{ + Action: github.IssueCommentActionCreated, + GUID: "event-guid", + Issue: github.Issue{ + ID: 42, + NodeID: "issue-node", + User: github.User{Login: "issue-author"}, + Number: 7, + Title: "Investigate the thing", + State: "open", + HTMLURL: "https://github.example/org/repo/issues/7", + Body: "Issue body", + PullRequest: &struct{}{}, + }, + Comment: github.IssueComment{ + ID: 99, + Body: "/gemini-agent explain this", + User: github.User{Login: "commenter"}, + HTMLURL: "https://github.example/org/repo/issues/7#issuecomment-99", + }, + Repo: github.Repo{ + Name: "repo", + Owner: github.User{Login: "org"}, + }, + }) + if err != nil { + t.Fatalf("handleIssueComment returned error: %v", err) + } + + if got.GUID != "event-guid" { + t.Fatalf("GUID %q, expected event-guid", got.GUID) + } + if !got.IsPR { + t.Fatal("expected generalized event to be marked as PR") + } + if got.Action != github.GenericCommentActionCreated { + t.Fatalf("action %q, expected %q", got.Action, github.GenericCommentActionCreated) + } + if got.Body != "/gemini-agent explain this" { + t.Fatalf("body %q, expected command body", got.Body) + } + if got.Number != 7 || got.Repo.Owner.Login != "org" || got.Repo.Name != "repo" { + t.Fatalf("unexpected repo coordinates: %s/%s#%d", got.Repo.Owner.Login, got.Repo.Name, got.Number) + } + if got.CommentID == nil || *got.CommentID != 99 { + t.Fatalf("comment ID %v, expected 99", got.CommentID) + } +} + +func TestHandleIssueCommentRequiresHandler(t *testing.T) { + pluginAgent := plugins.NewFakeConfigAgent() + server := &Server{ + ghc: &fakeGitHubClient{}, + log: logrus.NewEntry(logrus.New()), + pluginAgent: &pluginAgent, + } + + err := server.handleIssueComment(logrus.NewEntry(logrus.New()), github.IssueCommentEvent{ + Action: github.IssueCommentActionCreated, + Issue: github.Issue{PullRequest: &struct{}{}}, + }) + if err == nil { + t.Fatal("expected error for nil generic comment handler") + } +} + +var _ plugin.GitHubClient = (*fakeGitHubClient)(nil) diff --git a/go.mod b/go.mod index 69dd2218e8..2d51bca49f 100644 --- a/go.mod +++ b/go.mod @@ -67,6 +67,7 @@ require ( golang.org/x/time v0.12.0 gomodules.xyz/jsonpatch/v2 v2.5.0 google.golang.org/api v0.233.0 + google.golang.org/genai v1.57.0 google.golang.org/genproto v0.0.0-20250303144028-a0af3efb3deb google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 google.golang.org/grpc v1.79.3 diff --git a/go.sum b/go.sum index 89de382696..1aadc26c5b 100644 --- a/go.sum +++ b/go.sum @@ -1020,6 +1020,8 @@ google.golang.org/appengine v1.6.1/go.mod h1:i06prIuMbXzDqacNJfV5OdTW448YApPu5ww google.golang.org/appengine v1.6.5/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc= google.golang.org/appengine v1.6.6/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc= google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc= +google.golang.org/genai v1.57.0 h1:qTyG2ynz5dQy2jF4CvZdLHHVslhR0heMue+zM1a4GNM= +google.golang.org/genai v1.57.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk= google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= google.golang.org/genproto v0.0.0-20190307195333-5fe7a883aa19/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE= google.golang.org/genproto v0.0.0-20190418145605-e7d98fc518a7/go.mod h1:VzzqZJRnGkLBvHegQrXjBqPurQTc5/KpmUdxsrq26oE= diff --git a/pkg/plugins/config.go b/pkg/plugins/config.go index 4291a1c951..f3268aea5f 100644 --- a/pkg/plugins/config.go +++ b/pkg/plugins/config.go @@ -63,6 +63,9 @@ type Configuration struct { // external plugins. ExternalPlugins map[string][]ExternalPlugin `json:"external_plugins,omitempty"` + // GeminiAgents holds per-repo configuration for the gemini-agent external plugin. + GeminiAgents []GeminiAgentConfig `json:"gemini_agents,omitempty"` + // Owners contains configuration related to handling OWNERS files. Owners Owners `json:"owners,omitempty"` @@ -99,6 +102,19 @@ type Configuration struct { Help Help `json:"help,omitempty"` } +// GeminiAgentConfig holds per-repo settings for the gemini-agent external plugin. +type GeminiAgentConfig struct { + // Repos is either of the form org/repos or just org. + Repos []string `json:"repos,omitempty"` + + // Model is the Gemini model to use (e.g. "gemini-3-flash-preview"). + Model string `json:"model,omitempty"` + + // AllowedTeams is a list of GitHub team slugs whose members are allowed + // to invoke /gemini-agent, in addition to standard trigger trust. + AllowedTeams []string `json:"allowed_teams,omitempty"` +} + type Help struct { // HelpGuidelinesURL is the URL of the help page, which provides guidance on how and when to use the help wanted and good first issue labels. // The default value is "https://git.k8s.io/community/contributors/guide/help-wanted.md". diff --git a/pkg/plugins/plugin-config-documented.yaml b/pkg/plugins/plugin-config-documented.yaml index 62413b62d1..1b2ab071c6 100644 --- a/pkg/plugins/plugin-config-documented.yaml +++ b/pkg/plugins/plugin-config-documented.yaml @@ -413,6 +413,17 @@ dco: # external plugins. external_plugins: "": null +# GeminiAgents holds per-repo configuration for the gemini-agent external plugin. +gemini_agents: + - # AllowedTeams is a list of GitHub team slugs whose members are allowed + # to invoke /gemini-agent, in addition to standard trigger trust. + allowed_teams: + - "" + # Model is the Gemini model to use (e.g. "gemini-3-flash-preview"). + model: ' ' + # Repos is either of the form org/repos or just org. + repos: + - "" golint: # MinimumConfidence is the smallest permissible confidence # in (0,1] over which problems will be printed. Defaults to