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