Merge pull request #704 from chaitin/feature/extension-package-import

feat: 支持团队扩展包导入

Author: yokowu

backend/biz/host/usecase/host.go

@@ -194,9 +194,9 @@ func (h *HostUsecase) GetInstallCommand(ctx context.Context, user *domain.User)
 
 // InstallScript implements domain.HostUsecase.
 func (h *HostUsecase) InstallScript(ctx context.Context, token *domain.InstallReq) (string, error) {
-	key := fmt.Sprintf("host:token:%s", token.Token)
-	if _, err := h.redis.Get(ctx, key).Result(); err != nil {
-		return "", errcode.ErrInvalidInstallToken
+	teamID, err := h.teamIDFromInstallToken(ctx, token.Token)
+	if err != nil {
+		return "", err
 	}
 
 	tplName := "install"
@@ -212,19 +212,36 @@ func (h *HostUsecase) InstallScript(ctx context.Context, token *domain.InstallRe
 	}
 	buf := bytes.NewBuffer([]byte(""))
 	param := map[string]any{
-		"token":              token.Token,
-		"grpc_url":           h.cfg.TaskFlow.GrpcURL,
-		"base_url":           h.cfg.Server.BaseURL,
-		"installer_url":      h.installerURL(),
-		"docker_bundle_path": h.installerBundlePath("docker.tgz"),
-		"host_bundle_path":   h.hostBundlePath(),
+		"token":                          token.Token,
+		"grpc_url":                       h.cfg.TaskFlow.GrpcURL,
+		"base_url":                       h.cfg.Server.BaseURL,
+		"installer_url":                  h.installerURL(),
+		"docker_bundle_path":             h.installerBundlePath("docker.tgz"),
+		"host_bundle_path":               h.hostBundlePath(),
+		"extension_images_manifest_path": h.extensionImagesManifestPath(teamID),
 	}
 	if err := tmp.Execute(buf, param); err != nil {
 		return "", fmt.Errorf("failed to execute template %s", err)
 	}
 	return buf.String(), nil
 }
 
+func (h *HostUsecase) teamIDFromInstallToken(ctx context.Context, token string) (uuid.UUID, error) {
+	key := fmt.Sprintf("host:token:%s", token)
+	raw, err := h.redis.Get(ctx, key).Result()
+	if err != nil {
+		return uuid.Nil, errcode.ErrInvalidInstallToken
+	}
+	var user domain.User
+	if err := json.Unmarshal([]byte(raw), &user); err != nil {
+		return uuid.Nil, nil
+	}
+	if user.Team == nil {
+		return uuid.Nil, nil
+	}
+	return user.Team.ID, nil
+}
+
 func (h *HostUsecase) installerURL() string {
 	if h.cfg.Server.BaseURL == "" {
 		return ""
@@ -249,6 +266,13 @@ func (h *HostUsecase) installerBundlePath(name string) string {
 	return "/" + strings.Trim(h.cfg.StaticFiles.RoutePrefix, "/") + "/installer/{{.arch}}/" + name
 }
 
+func (h *HostUsecase) extensionImagesManifestPath(teamID uuid.UUID) string {
+	if teamID == uuid.Nil {
+		return ""
+	}
+	return "/" + strings.Trim(h.cfg.StaticFiles.RoutePrefix, "/") + "/extensions/teams/" + teamID.String() + "/images/{{.arch}}/manifest.json"
+}
+
 // List implements domain.HostUsecase.
 func (h *HostUsecase) List(ctx context.Context, uid uuid.UUID) (*domain.HostListResp, error) {
 	user, err := h.userRepo.Get(ctx, uid)

backend/biz/host/usecase/host_test.go

@@ -2,6 +2,7 @@ package usecase
 
 import (
 	"context"
+	"encoding/json"
 	"io"
 	"log/slog"
 	"net/url"
@@ -148,6 +149,61 @@ func TestInstallScriptUsesOfflineBundle(t *testing.T) {
 	assertInstallScriptChecksAVX(t, script)
 }
 
+func TestInstallScriptIncludesExtensionImagesManifestPath(t *testing.T) {
+	t.Parallel()
+
+	mr := miniredis.RunT(t)
+	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
+	t.Cleanup(func() { _ = rdb.Close() })
+
+	token := "install-token"
+	teamID := uuid.New()
+	rawUser, err := json.Marshal(&domain.User{
+		ID: uuid.New(),
+		Team: &domain.Team{
+			ID: teamID,
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := rdb.Set(context.Background(), "host:token:"+token, string(rawUser), time.Minute).Err(); err != nil {
+		t.Fatal(err)
+	}
+	u := &HostUsecase{
+		cfg: &config.Config{
+			Server: struct {
+				Addr    string `mapstructure:"addr"`
+				BaseURL string `mapstructure:"base_url"`
+			}{BaseURL: "http://monkeycode.local"},
+			TaskFlow: config.TaskFlow{GrpcURL: "127.0.0.1:50443"},
+			StaticFiles: config.StaticFilesConfig{
+				RoutePrefix: "/static",
+			},
+			HostInstaller: config.HostInstaller{
+				Mode:       "offline",
+				BundlePath: "installer/{{.arch}}/host.tgz",
+			},
+		},
+		redis: rdb,
+	}
+
+	script, err := u.InstallScript(context.Background(), &domain.InstallReq{Token: token})
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := "/static/extensions/teams/" + teamID.String() + "/images/{{.arch}}/manifest.json"
+	if !strings.Contains(script, "EXTENSION_IMAGES_MANIFEST_PATH=\""+want+"\"") {
+		t.Fatalf("script missing extension manifest path %q:\n%s", want, script)
+	}
+	if !strings.Contains(script, "EXTENSION_IMAGES_MANIFEST_PATH=${EXTENSION_IMAGES_MANIFEST_PATH//\\{\\{.arch\\}\\}/$ARCH}") {
+		t.Fatalf("script missing extension manifest arch replacement:\n%s", script)
+	}
+	if !strings.Contains(script, "MCAI_EXTENSION_IMAGES_MANIFEST_PATH=\"$EXTENSION_IMAGES_MANIFEST_PATH\"") {
+		t.Fatalf("script missing installer extension manifest env:\n%s", script)
+	}
+}
+
 func assertInstallScriptChecksAVX(t *testing.T, script string) {
 	t.Helper()
 

backend/biz/team/handler/http/v1/extension_package.go

@@ -0,0 +1,96 @@
+package v1
+
+import (
+	"fmt"
+	"io"
+	"mime/multipart"
+
+	"github.com/GoYoko/web"
+	"github.com/samber/do"
+
+	"github.com/chaitin/MonkeyCode/backend/config"
+	"github.com/chaitin/MonkeyCode/backend/domain"
+	"github.com/chaitin/MonkeyCode/backend/errcode"
+	"github.com/chaitin/MonkeyCode/backend/middleware"
+)
+
+type TeamExtensionPackageHandler struct {
+	usecase domain.TeamExtensionPackageUsecase
+	cfg     *config.Config
+}
+
+func NewTeamExtensionPackageHandler(i *do.Injector) (*TeamExtensionPackageHandler, error) {
+	w := do.MustInvoke[*web.Web](i)
+	auth := do.MustInvoke[*middleware.AuthMiddleware](i)
+	audit := do.MustInvoke[*middleware.AuditMiddleware](i)
+
+	h := &TeamExtensionPackageHandler{
+		usecase: do.MustInvoke[domain.TeamExtensionPackageUsecase](i),
+		cfg:     do.MustInvoke[*config.Config](i),
+	}
+
+	g := w.Group("/api/v1/teams/extension-packages")
+	g.Use(auth.TeamAuth())
+	g.POST("", web.BindHandler(h.Import), audit.Audit("import_team_extension_package"))
+
+	return h, nil
+}
+
+type importTeamExtensionPackageFormReq struct {
+	File *multipart.FileHeader `form:"file" validate:"required"`
+}
+
+// Import 上传团队扩展包
+//
+//	@Summary		上传团队扩展包
+//	@Description	上传团队扩展包并导入 Skills 和团队镜像记录
+//	@Tags			【Team 管理员】扩展包管理
+//	@Accept			multipart/form-data
+//	@Produce		json
+//	@Security		MonkeyCodeAITeamAuth
+//	@Param			file	formData	file													true	"扩展包 zip"
+//	@Success		200		{object}	web.Resp{data=domain.ImportTeamExtensionPackageResp}	"成功"
+//	@Failure		401		{object}	web.Resp												"未授权"
+//	@Failure		500		{object}	web.Resp												"服务器内部错误"
+//	@Router			/api/v1/teams/extension-packages [post]
+func (h *TeamExtensionPackageHandler) Import(c *web.Context, req importTeamExtensionPackageFormReq) error {
+	teamUser := middleware.GetTeamUser(c)
+	data, err := h.readPackageFile(req.File)
+	if err != nil {
+		return err
+	}
+	resp, err := h.usecase.Import(c.Request().Context(), teamUser, &domain.ImportTeamExtensionPackageReq{
+		Filename: req.File.Filename,
+		Data:     data,
+	})
+	if err != nil {
+		return err
+	}
+	return c.Success(resp)
+}
+
+func (h *TeamExtensionPackageHandler) readPackageFile(fileHeader *multipart.FileHeader) ([]byte, error) {
+	if fileHeader == nil {
+		return nil, errcode.ErrBadRequest.Wrap(fmt.Errorf("file is required"))
+	}
+	maxSize := h.cfg.ObjectStorage.MaxSize
+	if maxSize <= 0 {
+		maxSize = 50 << 20
+	}
+	if fileHeader.Size > maxSize {
+		return nil, errcode.ErrBadRequest.Wrap(fmt.Errorf("file exceeds limit"))
+	}
+	file, err := fileHeader.Open()
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+	data, err := io.ReadAll(io.LimitReader(file, maxSize+1))
+	if err != nil {
+		return nil, err
+	}
+	if int64(len(data)) > maxSize {
+		return nil, errcode.ErrBadRequest.Wrap(fmt.Errorf("file exceeds limit"))
+	}
+	return data, nil
+}

backend/biz/team/handler/http/v1/extension_package_test.go

@@ -0,0 +1,45 @@
+package v1
+
+import (
+	"bytes"
+	"mime/multipart"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/chaitin/MonkeyCode/backend/config"
+)
+
+func TestTeamExtensionPackageHandlerReadPackageFileUsesConfiguredLimit(t *testing.T) {
+	h := &TeamExtensionPackageHandler{
+		cfg: &config.Config{},
+	}
+	h.cfg.ObjectStorage.MaxSize = 3
+	file := makeMultipartFileHeader(t, "extension.zip", "abcd")
+
+	if _, err := h.readPackageFile(file); err == nil {
+		t.Fatal("expected file size limit error")
+	}
+}
+
+func makeMultipartFileHeader(t *testing.T, filename, content string) *multipart.FileHeader {
+	t.Helper()
+	var body bytes.Buffer
+	writer := multipart.NewWriter(&body)
+	part, err := writer.CreateFormFile("file", filename)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := part.Write([]byte(content)); err != nil {
+		t.Fatal(err)
+	}
+	if err := writer.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	req := httptest.NewRequest("POST", "/upload", &body)
+	req.Header.Set("Content-Type", writer.FormDataContentType())
+	if err := req.ParseMultipartForm(1 << 20); err != nil {
+		t.Fatal(err)
+	}
+	return req.MultipartForm.File["file"][0]
+}

backend/biz/team/register.go

@@ -27,6 +27,9 @@ func ProvideTeam(i *do.Injector) {
 	do.Provide(i, repo.NewTeamSkillRepo)
 	do.Provide(i, usecase.NewTeamSkillUsecase)
 	do.Provide(i, v1.NewTeamSkillHandler)
+	do.Provide(i, repo.NewTeamExtensionPackageRepo)
+	do.Provide(i, usecase.NewTeamExtensionPackageUsecase)
+	do.Provide(i, v1.NewTeamExtensionPackageHandler)
 	do.Provide(i, repo.NewTeamHostRepo)
 	do.Provide(i, usecase.NewTeamHostUsecase)
 	do.Provide(i, v1.NewTeamHostHandler)
@@ -51,6 +54,7 @@ func InvokeTeam(i *do.Injector) {
 	do.MustInvoke[*v1.TeamModelHandler](i)
 	do.MustInvoke[*v1.TeamImageHandler](i)
 	do.MustInvoke[*v1.TeamSkillHandler](i)
+	do.MustInvoke[*v1.TeamExtensionPackageHandler](i)
 	do.MustInvoke[*v1.TeamHostHandler](i)
 	do.MustInvoke[*v1.TeamPolicyHandler](i)
 	do.MustInvoke[*v1.TeamOIDCHandler](i)