From 65f2895c0d3f70361c58f41d3603d8230c2f984c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Fri, 26 Sep 2025 16:58:12 +0200 Subject: [PATCH] chore: add CLI command to list aibridge interceptions (#19935) Co-authored-by: Dean Sheather --- cli/exp.go | 23 --- cli/root.go | 31 +++- docs/reference/cli/index.md | 2 +- enterprise/cli/exp_aibridge.go | 166 ++++++++++++++++++++++ enterprise/cli/exp_aibridge_test.go | 211 ++++++++++++++++++++++++++++ enterprise/cli/root.go | 17 ++- 6 files changed, 422 insertions(+), 28 deletions(-) delete mode 100644 cli/exp.go create mode 100644 enterprise/cli/exp_aibridge.go create mode 100644 enterprise/cli/exp_aibridge_test.go diff --git a/cli/exp.go b/cli/exp.go deleted file mode 100644 index e20d1e28d5..0000000000 --- a/cli/exp.go +++ /dev/null @@ -1,23 +0,0 @@ -package cli - -import "github.com/coder/serpent" - -func (r *RootCmd) expCmd() *serpent.Command { - cmd := &serpent.Command{ - Use: "exp", - Short: "Internal commands for testing and experimentation. These are prone to breaking changes with no notice.", - Handler: func(i *serpent.Invocation) error { - return i.Command.HelpHandler(i) - }, - Hidden: true, - Children: []*serpent.Command{ - r.scaletestCmd(), - r.errorExample(), - r.mcpCommand(), - r.promptExample(), - r.rptyCommand(), - r.tasksCommand(), - }, - } - return cmd -} diff --git a/cli/root.go b/cli/root.go index 828dc9944f..64c3da672f 100644 --- a/cli/root.go +++ b/cli/root.go @@ -128,7 +128,6 @@ func (r *RootCmd) CoreSubcommands() []*serpent.Command { // Hidden r.connectCmd(), - r.expCmd(), gitssh(), r.support(), r.vpnDaemon(), @@ -137,15 +136,45 @@ func (r *RootCmd) CoreSubcommands() []*serpent.Command { } } +// AGPLExperimental returns all AGPL experimental subcommands. +func (r *RootCmd) AGPLExperimental() []*serpent.Command { + return []*serpent.Command{ + r.scaletestCmd(), + r.errorExample(), + r.mcpCommand(), + r.promptExample(), + r.rptyCommand(), + r.tasksCommand(), + } +} + +// AGPL returns all AGPL commands including any non-core commands that are +// duplicated in the Enterprise CLI. func (r *RootCmd) AGPL() []*serpent.Command { all := append( r.CoreSubcommands(), r.Server( /* Do not import coderd here. */ nil), r.Provisioners(), + ExperimentalCommand(r.AGPLExperimental()), ) return all } +// ExperimentalCommand creates an experimental command that is hidden and has +// the given subcommands. +func ExperimentalCommand(subcommands []*serpent.Command) *serpent.Command { + cmd := &serpent.Command{ + Use: "exp", + Short: "Internal commands for testing and experimentation. These are prone to breaking changes with no notice.", + Handler: func(i *serpent.Invocation) error { + return i.Command.HelpHandler(i) + }, + Hidden: true, + Children: subcommands, + } + return cmd +} + // RunWithSubcommands runs the root command with the given subcommands. // It is abstracted to enable the Enterprise code to add commands. func (r *RootCmd) RunWithSubcommands(subcommands []*serpent.Command) { diff --git a/docs/reference/cli/index.md b/docs/reference/cli/index.md index 101186eeea..c298f8bcb6 100644 --- a/docs/reference/cli/index.md +++ b/docs/reference/cli/index.md @@ -62,11 +62,11 @@ Coder — A tool for provisioning self-hosted development environments with Terr | [whoami](./whoami.md) | Fetch authenticated user info for Coder deployment | | [support](./support.md) | Commands for troubleshooting issues with a Coder deployment. | | [server](./server.md) | Start a Coder server | +| [provisioner](./provisioner.md) | View and manage provisioner daemons and jobs | | [features](./features.md) | List Enterprise features | | [licenses](./licenses.md) | Add, delete, and list licenses | | [groups](./groups.md) | Manage groups | | [prebuilds](./prebuilds.md) | Manage Coder prebuilds | -| [provisioner](./provisioner.md) | View and manage provisioner daemons and jobs | | [external-workspaces](./external-workspaces.md) | Create or manage external workspaces | ## Options diff --git a/enterprise/cli/exp_aibridge.go b/enterprise/cli/exp_aibridge.go new file mode 100644 index 0000000000..722f7bf239 --- /dev/null +++ b/enterprise/cli/exp_aibridge.go @@ -0,0 +1,166 @@ +package cli + +import ( + "encoding/json" + "fmt" + "time" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/codersdk" + "github.com/coder/serpent" +) + +const maxInterceptionsLimit = 1000 + +func (r *RootCmd) aibridge() *serpent.Command { + cmd := &serpent.Command{ + Use: "aibridge", + Short: "Manage AIBridge.", + Handler: func(inv *serpent.Invocation) error { + return inv.Command.HelpHandler(inv) + }, + Children: []*serpent.Command{ + r.aibridgeInterceptions(), + }, + } + return cmd +} + +func (r *RootCmd) aibridgeInterceptions() *serpent.Command { + cmd := &serpent.Command{ + Use: "interceptions", + Short: "Manage AIBridge interceptions.", + Handler: func(inv *serpent.Invocation) error { + return inv.Command.HelpHandler(inv) + }, + Children: []*serpent.Command{ + r.aibridgeInterceptionsList(), + }, + } + return cmd +} + +func (r *RootCmd) aibridgeInterceptionsList() *serpent.Command { + var ( + initiator string + startedBeforeRaw string + startedAfterRaw string + provider string + model string + afterIDRaw string + limit int64 + ) + + return &serpent.Command{ + Use: "list", + Short: "List AIBridge interceptions as JSON.", + Options: serpent.OptionSet{ + { + Flag: "initiator", + Description: `Only return interceptions initiated by this user. Accepts a user ID, username, or "me".`, + Default: "", + Value: serpent.StringOf(&initiator), + }, + { + Flag: "started-before", + Description: fmt.Sprintf("Only return interceptions started before this time. Must be after 'started-after' if set. Accepts a time in the RFC 3339 format, e.g. %q.", time.RFC3339), + Default: "", + Value: serpent.StringOf(&startedBeforeRaw), + }, + { + Flag: "started-after", + Description: fmt.Sprintf("Only return interceptions started after this time. Must be before 'started-before' if set. Accepts a time in the RFC 3339 format, e.g. %q.", time.RFC3339), + Default: "", + Value: serpent.StringOf(&startedAfterRaw), + }, + { + Flag: "provider", + Description: `Only return interceptions from this provider.`, + Default: "", + Value: serpent.StringOf(&provider), + }, + { + Flag: "model", + Description: `Only return interceptions from this model.`, + Default: "", + Value: serpent.StringOf(&model), + }, + { + Flag: "after-id", + Description: "The ID of the last result on the previous page to use as a pagination cursor.", + Default: "", + Value: serpent.StringOf(&afterIDRaw), + }, + { + Flag: "limit", + Description: fmt.Sprintf(`The limit of results to return. Must be between 1 and %d.`, maxInterceptionsLimit), + Default: "100", + Value: serpent.Int64Of(&limit), + }, + }, + Handler: func(inv *serpent.Invocation) error { + client, err := r.InitClient(inv) + if err != nil { + return err + } + + startedBefore := time.Time{} + if startedBeforeRaw != "" { + startedBefore, err = time.Parse(time.RFC3339, startedBeforeRaw) + if err != nil { + return xerrors.Errorf("parse started before filter value %q: %w", startedBeforeRaw, err) + } + } + + startedAfter := time.Time{} + if startedAfterRaw != "" { + startedAfter, err = time.Parse(time.RFC3339, startedAfterRaw) + if err != nil { + return xerrors.Errorf("parse started after filter value %q: %w", startedAfterRaw, err) + } + } + + afterID := uuid.Nil + if afterIDRaw != "" { + afterID, err = uuid.Parse(afterIDRaw) + if err != nil { + return xerrors.Errorf("parse after_id filter value %q: %w", afterIDRaw, err) + } + } + + if limit < 1 || limit > maxInterceptionsLimit { + return xerrors.Errorf("limit value must be between 1 and %d", maxInterceptionsLimit) + } + + expCli := codersdk.NewExperimentalClient(client) + resp, err := expCli.AIBridgeListInterceptions(inv.Context(), codersdk.AIBridgeListInterceptionsFilter{ + Pagination: codersdk.Pagination{ + AfterID: afterID, + // #nosec G115 - Checked above. + Limit: int(limit), + }, + Initiator: initiator, + StartedBefore: startedBefore, + StartedAfter: startedAfter, + Provider: provider, + Model: model, + }) + if err != nil { + return xerrors.Errorf("list interceptions: %w", err) + } + + // We currently only support JSON output, so we don't use a + // formatter. + enc := json.NewEncoder(inv.Stdout) + enc.SetIndent("", " ") + err = enc.Encode(resp.Results) + if err != nil { + return err + } + + return err + }, + } +} diff --git a/enterprise/cli/exp_aibridge_test.go b/enterprise/cli/exp_aibridge_test.go new file mode 100644 index 0000000000..454e81d070 --- /dev/null +++ b/enterprise/cli/exp_aibridge_test.go @@ -0,0 +1,211 @@ +package cli_test + +import ( + "bytes" + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/cli/clitest" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" + "github.com/coder/coder/v2/testutil" +) + +func TestAIBridgeListInterceptions(t *testing.T) { + t.Parallel() + + t.Run("OK", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + client, db, owner := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + memberClient, member := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + now := dbtime.Now() + interception1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: member.ID, + StartedAt: now.Add(-time.Hour), + }) + interception2 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: member.ID, + StartedAt: now, + }) + // Should not be returned because the user can't see it. + _ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: owner.UserID, + StartedAt: now.Add(-2 * time.Hour), + }) + + args := []string{ + "exp", + "aibridge", + "interceptions", + "list", + } + inv, root := newCLI(t, args...) + clitest.SetupConfig(t, memberClient, root) + + ctx := testutil.Context(t, testutil.WaitLong) + + out := bytes.NewBuffer(nil) + inv.Stdout = out + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + + // Reverse order because the order is `started_at ASC`. + requireHasInterceptions(t, out.Bytes(), []uuid.UUID{interception2.ID, interception1.ID}) + }) + + t.Run("Filter", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + client, db, owner := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + memberClient, member := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + + now := dbtime.Now() + + // This interception should be returned since it matches all filters. + goodInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: member.ID, + Provider: "real-provider", + Model: "real-model", + StartedAt: now, + }) + + // These interceptions should not be returned since they don't match the + // filters. + _ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: owner.UserID, + Provider: goodInterception.Provider, + Model: goodInterception.Model, + StartedAt: goodInterception.StartedAt, + }) + _ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: goodInterception.InitiatorID, + Provider: "bad-provider", + Model: goodInterception.Model, + StartedAt: goodInterception.StartedAt, + }) + _ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: goodInterception.InitiatorID, + Provider: goodInterception.Provider, + Model: "bad-model", + StartedAt: goodInterception.StartedAt, + }) + _ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: goodInterception.InitiatorID, + Provider: goodInterception.Provider, + Model: goodInterception.Model, + // Violates the started after filter. + StartedAt: now.Add(-2 * time.Hour), + }) + _ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: goodInterception.InitiatorID, + Provider: goodInterception.Provider, + Model: goodInterception.Model, + // Violates the started before filter. + StartedAt: now.Add(2 * time.Hour), + }) + + args := []string{ + "exp", + "aibridge", + "interceptions", + "list", + "--started-after", now.Add(-time.Hour).Format(time.RFC3339), + "--started-before", now.Add(time.Hour).Format(time.RFC3339), + "--initiator", codersdk.Me, + "--provider", goodInterception.Provider, + "--model", goodInterception.Model, + } + inv, root := newCLI(t, args...) + clitest.SetupConfig(t, memberClient, root) + + ctx := testutil.Context(t, testutil.WaitLong) + + out := bytes.NewBuffer(nil) + inv.Stdout = out + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + + requireHasInterceptions(t, out.Bytes(), []uuid.UUID{goodInterception.ID}) + }) + + t.Run("Pagination", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + client, db, owner := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + memberClient, member := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID) + + now := dbtime.Now() + firstInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: member.ID, + StartedAt: now, + }) + returnedInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: member.ID, + StartedAt: now.Add(-time.Hour), + }) + _ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: member.ID, + StartedAt: now.Add(-2 * time.Hour), + }) + + args := []string{ + "exp", + "aibridge", + "interceptions", + "list", + "--limit", "1", + "--after-id", firstInterception.ID.String(), + } + inv, root := newCLI(t, args...) + clitest.SetupConfig(t, memberClient, root) + + ctx := testutil.Context(t, testutil.WaitLong) + + out := bytes.NewBuffer(nil) + inv.Stdout = out + err := inv.WithContext(ctx).Run() + require.NoError(t, err) + + // Only contains the second interception because after_id is the first + // interception, and we set a limit of 1. + requireHasInterceptions(t, out.Bytes(), []uuid.UUID{returnedInterception.ID}) + }) +} + +func requireHasInterceptions(t *testing.T, out []byte, ids []uuid.UUID) { + t.Helper() + + var results []codersdk.AIBridgeInterception + require.NoError(t, json.Unmarshal(out, &results)) + require.Len(t, results, len(ids)) + for i, id := range ids { + require.Equal(t, id, results[i].ID) + } +} diff --git a/enterprise/cli/root.go b/enterprise/cli/root.go index ed54a76f90..3cec119703 100644 --- a/enterprise/cli/root.go +++ b/enterprise/cli/root.go @@ -1,28 +1,39 @@ package cli import ( - "github.com/coder/coder/v2/cli" + agplcli "github.com/coder/coder/v2/cli" "github.com/coder/serpent" ) type RootCmd struct { - cli.RootCmd + agplcli.RootCmd } func (r *RootCmd) enterpriseOnly() []*serpent.Command { return []*serpent.Command{ + // These commands exist in AGPL, but we use a different implementation + // in enterprise: r.Server(nil), + r.provisionerDaemons(), + agplcli.ExperimentalCommand(append(r.AGPLExperimental(), r.enterpriseExperimental()...)), + + // New commands that don't exist in AGPL: r.workspaceProxy(), r.features(), r.licenses(), r.groups(), r.prebuilds(), - r.provisionerDaemons(), r.provisionerd(), r.externalWorkspaces(), } } +func (r *RootCmd) enterpriseExperimental() []*serpent.Command { + return []*serpent.Command{ + r.aibridge(), + } +} + func (r *RootCmd) EnterpriseSubcommands() []*serpent.Command { all := append(r.CoreSubcommands(), r.enterpriseOnly()...) return all