feat: add agentfake scaletest subcommand (#25072)

This PR builds on top of https://github.com/coder/coder/pull/25070 to add a way of running the larger "fake agent" manager via the existing CLI, pulling in the URL/credentials already set. With this, we can run a pod per scaletest region to act as all the workspaces in that region. This is in a new subcommand `scaletest agentfake` currently. --------- Signed-off-by: Callum Styan <callumstyan@gmail.com>
2026-06-02 20:48:20 +00:00 · 2026-05-15 14:36:54 -07:00
parent d9976768db
commit 191dd230ae
7 changed files with 112 additions and 12 deletions
@@ -0,0 +1,100 @@
+//go:build !slim
+
+package cli
+
+import (
+	"os/signal"
+
+	"golang.org/x/xerrors"
+
+	agplcli "github.com/coder/coder/v2/cli"
+	"github.com/coder/coder/v2/enterprise/scaletest/agentfake"
+	"github.com/coder/serpent"
+)
+
+// AGPLExperimental shadows the embedded RootCmd.AGPLExperimental to inject the
+// enterprise-only agentfake scaletest subcommand into the scaletest subtree.
+func (r *RootCmd) AGPLExperimental() []*serpent.Command {
+	cmds := r.RootCmd.AGPLExperimental()
+	for _, cmd := range cmds {
+		if cmd.Use == "scaletest" {
+			cmd.Children = append(cmd.Children, r.scaletestAgentFake())
+		}
+	}
+	return cmds
+}
+
+func (r *RootCmd) scaletestAgentFake() *serpent.Command {
+	var (
+		template string
+		owner    string
+	)
+
+	cmd := &serpent.Command{
+		Use:   "agentfake",
+		Short: "Run fake external agents against workspaces of the given template.",
+		Long: agplcli.FormatExamples(
+			agplcli.Example{
+				Description: "Connect a fake agent for every external-agent workspace built from the template named " +
+					"\"agentfake-runner\".",
+				Command: "coder exp scaletest agentfake --template agentfake-runner",
+			},
+		) + "\n\n" +
+			"Enumerates external-agent workspaces matching --template (optionally filtered by --owner), " +
+			"fetches each workspace agent's external-agent credentials, and supervises one in-process fake " +
+			"agent per token until the command is interrupted.\n\n" +
+			"Requires a session token whose user is template-admin (or higher) on a deployment licensed " +
+			"for the workspace external-agent feature; both the workspace builds and the credentials " +
+			"endpoint are gated server-side. Pair with `coder exp scaletest create-workspaces " +
+			"--no-wait-for-agents` to seed the workspaces this command will pick up. Workspaces created " +
+			"after this command starts are NOT picked up; rerun the command after seeding more.",
+		Handler: func(inv *serpent.Invocation) error {
+			ctx := inv.Context()
+			client, err := r.InitClient(inv)
+			if err != nil {
+				return err
+			}
+
+			notifyCtx, stop := signal.NotifyContext(ctx, agplcli.StopSignals...)
+			defer stop()
+			ctx = notifyCtx
+
+			if _, err := agplcli.RequireAdmin(ctx, client); err != nil {
+				return err
+			}
+
+			if template == "" {
+				return xerrors.New("--template is required")
+			}
+
+			logger := inv.Logger
+			mgr := agentfake.NewManager(client, logger, agentfake.ManagerOptions{
+				Template: template,
+				Owner:    owner,
+			})
+			defer mgr.Close()
+
+			if err := mgr.Run(ctx); err != nil {
+				return xerrors.Errorf("run agentfake manager: %w", err)
+			}
+			return nil
+		},
+	}
+
+	cmd.Options = serpent.OptionSet{
+		{
+			Flag:        "template",
+			Env:         "CODER_SCALETEST_AGENTFAKE_TEMPLATE",
+			Description: "Name of the template whose external-agent workspaces should be supervised. Required.",
+			Value:       serpent.StringOf(&template),
+		},
+		{
+			Flag:        "owner",
+			Env:         "CODER_SCALETEST_AGENTFAKE_OWNER",
+			Description: "Optional workspace-owner filter (username). When empty, all owners' workspaces of the template are included.",
+			Value:       serpent.StringOf(&owner),
+		},
+	}
+
+	return cmd
+}