Agent-System-Backend/internal/execenv/pip.go

package execenv

import (
	"bytes"
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

//go:embed base-requirements.txt
var baseRequirements string

// BaseRequirements 是基础层的包名清单（不带版本；版本由第一次解析锁定）。
func BaseRequirements() []string {
	var out []string
	for _, line := range strings.Split(baseRequirements, "\n") {
		line = strings.TrimSpace(line)
		if line != "" && !strings.HasPrefix(line, "#") {
			out = append(out, line)
		}
	}
	return out
}

// Pip 用宿主侧的 pip 下载 wheel。它不在沙箱里运行。
type Pip struct {
	Python string // 带 pip 的解释器
}

func (p Pip) python() string {
	if p.Python != "" {
		return p.Python
	}
	return "python3"
}

// Resolve 以基础层锁为 constraints 做 dry-run。冲突记入 Override，不改锁。
func (p Pip) Resolve(ctx context.Context, reqs []string, base Lock) (Result, error) {
	if len(reqs) == 0 {
		return Result{}, nil
	}
	dir, err := os.MkdirTemp("", "execenv-resolve-")
	if err != nil {
		return Result{}, err
	}
	defer os.RemoveAll(dir)
	cons := filepath.Join(dir, "constraints.txt")
	if err := os.WriteFile(cons, []byte(constraintText(base)), 0o644); err != nil {
		return Result{}, err
	}
	args := []string{"-m", "pip", "install", "--dry-run", "--report", "-", "--only-binary=:all:", "--constraint", cons}
	args = append(args, reqs...)
	cmd := exec.CommandContext(ctx, p.python(), args...)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err = cmd.Run()
	text := stderr.String() + stdout.String()
	if err != nil {
		if strings.Contains(text, "ResolutionImpossible") || strings.Contains(text, "conflicting") {
			return Result{Override: []string{firstLine(text)}}, nil
		}
		if strings.Contains(text, "only-binary") || strings.Contains(text, "No matching distribution") {
			return Result{SDist: reqs, Missing: []string{firstLine(text)}}, nil
		}
		return Result{}, fmt.Errorf("execenv: pip resolve: %v: %s", err, firstLine(text))
	}
	return parsePipReport(stdout.Bytes(), base)
}

func parsePipReport(raw []byte, base Lock) (Result, error) {
	var doc struct {
		Install []struct {
			Metadata struct {
				Name    string `json:"name"`
				Version string `json:"version"`
			} `json:"metadata"`
			DownloadInfo struct {
				ArchiveInfo struct {
					Hash string `json:"hash"`
				} `json:"archive_info"`
				URL string `json:"url"`
			} `json:"download_info"`
		} `json:"install"`
	}
	if err := json.Unmarshal(raw, &doc); err != nil {
		return Result{}, fmt.Errorf("execenv: pip report: %w", err)
	}
	pinned := map[string]LockedPkg{}
	for _, c := range base.Packages {
		pinned[normName(c.Name)] = c
	}
	var res Result
	for _, item := range doc.Install {
		name := item.Metadata.Name
		if pin, ok := pinned[normName(name)]; ok {
			if pin.Version != item.Metadata.Version {
				res.Override = append(res.Override, fmt.Sprintf("%s 需要 %s，基础层固定为 %s", name, item.Metadata.Version, pin.Version))
			}
			continue
		}
		hash := strings.TrimPrefix(item.DownloadInfo.ArchiveInfo.Hash, "sha256=")
		filename := name + "-" + item.Metadata.Version + ".whl"
		if item.DownloadInfo.URL != "" {
			filename = filepath.Base(strings.Split(item.DownloadInfo.URL, "?")[0])
		}
		res.Packages = append(res.Packages, LockedPkg{
			Name: name, Version: item.Metadata.Version, Filename: filename, SHA256: hash,
		})
	}
	return res, nil
}

// Fetch 把一个包下载进 wheelhouse，并返回落地路径。哈希不符则删除并报错。
func (p Pip) Fetch(ctx context.Context, wheelhouse string, pkg LockedPkg) (string, error) {
	if err := os.MkdirAll(wheelhouse, 0o755); err != nil {
		return "", err
	}
	dest := filepath.Join(wheelhouse, pkg.Filename)
	if pkg.SHA256 != "" {
		if sum, err := fileSHA256(dest); err == nil && strings.EqualFold(sum, pkg.SHA256) {
			return dest, nil
		}
	}
	cmd := exec.CommandContext(ctx, p.python(), "-m", "pip", "download",
		"--only-binary=:all:", "--no-deps", "-d", wheelhouse, pkg.Name+"=="+pkg.Version)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("execenv: pip download %s: %v: %s", pkg.Name, err, firstLine(string(out)))
	}
	found := dest
	if _, err := os.Stat(found); err != nil {
		matches, _ := filepath.Glob(filepath.Join(wheelhouse, normName(pkg.Name)+"*.whl"))
		if len(matches) == 0 {
			matches, _ = filepath.Glob(filepath.Join(wheelhouse, "*.whl"))
		}
		if len(matches) == 0 {
			return "", fmt.Errorf("execenv: download produced no wheel for %s", pkg.Name)
		}
		found = matches[len(matches)-1]
	}
	sum, err := fileSHA256(found)
	if err != nil {
		return "", err
	}
	if pkg.SHA256 != "" && !strings.EqualFold(sum, pkg.SHA256) {
		_ = os.Remove(found)
		return "", fmt.Errorf("execenv: hash mismatch for %s", pkg.Name)
	}
	stable := filepath.Join(wheelhouse, sum+".whl")
	if found != stable {
		if err := os.Rename(found, stable); err != nil {
			return found, nil
		}
	}
	return stable, nil
}

func constraintText(base Lock) string {
	var b strings.Builder
	for _, p := range base.Packages {
		if p.Name != "" && p.Version != "" {
			fmt.Fprintf(&b, "%s==%s\n", p.Name, p.Version)
		}
	}
	return b.String()
}

func firstLine(s string) string {
	s = strings.TrimSpace(s)
	if i := strings.IndexByte(s, '\n'); i >= 0 {
		s = s[:i]
	}
	if len(s) > 400 {
		s = s[:400]
	}
	return s
}