Add ssh agent-shim subcommand

experimental/ssh/cmd/agent_shim.go

@@ -0,0 +1,43 @@
+package ssh
+
+import (
+	"github.com/databricks/cli/cmd/root"
+	"github.com/databricks/cli/experimental/ssh/internal/client"
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/spf13/cobra"
+)
+
+func newAgentShimCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:    "agent-shim",
+		Short:  "Launch a coding agent configured against Unity AI Gateway (invoked on the remote)",
+		Hidden: true,
+	}
+	for _, agent := range client.SupportedAgentNames() {
+		cmd.AddCommand(newAgentShimAgentCommand(agent))
+	}
+	return cmd
+}
+
+func newAgentShimAgentCommand(agent string) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   agent,
+		Short: "Launch " + agent + " configured against Unity AI Gateway",
+		// Disable flag parsing: forward everything after the agent name to the agent verbatim.
+		DisableFlagParsing: true,
+	}
+
+	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
+		// Runs on the driver with injected env auth; no bundle, no prompt.
+		cmd.SetContext(root.SkipLoadBundle(cmd.Context()))
+		cmd.SetContext(root.SkipPrompt(cmd.Context()))
+		return root.MustWorkspaceClient(cmd, args)
+	}
+
+	cmd.RunE = func(cmd *cobra.Command, args []string) error {
+		ctx := cmd.Context()
+		return client.RunAgentShim(ctx, cmdctx.WorkspaceClient(ctx), agent, args)
+	}
+
+	return cmd
+}

experimental/ssh/cmd/ssh.go

@@ -22,6 +22,7 @@ Use ` + "`databricks ssh connect --help`" + ` to see all available flags.`,
 	cmd.AddCommand(newSetupCommand())
 	cmd.AddCommand(newConnectCommand())
 	cmd.AddCommand(newServerCommand())
+	cmd.AddCommand(newAgentShimCommand())
 
 	return cmd
 }