chore: add CLI command to list aibridge interceptions (#19935)

Co-authored-by: Dean Sheather <dean@deansheather.com>

cli/exp.go

@@ -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
-}

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) {

docs/reference/cli/index.md

@@ -62,11 +62,11 @@ Coder — A tool for provisioning self-hosted development environments with Terr
 | [<code>whoami</code>](./whoami.md)                           | Fetch authenticated user info for Coder deployment                                                                           |
 | [<code>support</code>](./support.md)                         | Commands for troubleshooting issues with a Coder deployment.                                                                 |
 | [<code>server</code>](./server.md)                           | Start a Coder server                                                                                                         |
+| [<code>provisioner</code>](./provisioner.md)                 | View and manage provisioner daemons and jobs                                                                                 |
 | [<code>features</code>](./features.md)                       | List Enterprise features                                                                                                     |
 | [<code>licenses</code>](./licenses.md)                       | Add, delete, and list licenses                                                                                               |
 | [<code>groups</code>](./groups.md)                           | Manage groups                                                                                                                |
 | [<code>prebuilds</code>](./prebuilds.md)                     | Manage Coder prebuilds                                                                                                       |
-| [<code>provisioner</code>](./provisioner.md)                 | View and manage provisioner daemons and jobs                                                                                 |
 | [<code>external-workspaces</code>](./external-workspaces.md) | Create or manage external workspaces                                                                                         |
 
 ## Options

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
+		},
+	}
+}

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)
+	}
+}

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