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

package execenv

import (
	"context"
	"errors"
	"fmt"
	"os"
	"time"
)

func (b *Broker) materialize(ctx context.Context, conv, owner string, base Lock, res Result, isolate bool) (string, error) {
	pkgs := append([]LockedPkg(nil), res.Packages...)
	for i := range pkgs {
		if pkgs[i].SHA256 == "" || !wheelExists(b.Root, pkgs[i]) {
			path, err := b.Pip.Fetch(ctx, b.Root.Wheelhouse(), pkgs[i])
			if err != nil {
				return "", err
			}
			sum, err := fileSHA256(path)
			if err != nil {
				return "", err
			}
			if pkgs[i].SHA256 != "" && sum != pkgs[i].SHA256 {
				return "", fmt.Errorf("execenv: hash mismatch for %s", pkgs[i].Name)
			}
			pkgs[i].SHA256 = sum
			pkgs[i].Filename = pkgs[i].SHA256 + ".whl"
			stable := b.Root.Wheelhouse() + "/" + pkgs[i].Filename
			if path != stable {
				if err := os.Rename(path, stable); err != nil && !errors.Is(err, os.ErrExist) {
					// Fetch 可能已经改名。忽略已就位的文件。
					if _, statErr := os.Stat(stable); statErr != nil {
						return "", err
					}
				}
			}
		}
	}
	kind := "extension"
	baseID := ""
	if id, _ := b.Root.CurrentBase(); id != "" {
		baseID = id
	}
	lockPkgs := pkgs
	if isolate {
		kind = "isolated"
	} else if baseID == "" {
		kind = "base"
	}
	if !isolate && baseID != "" {
		lockPkgs = res.Packages
	}
	lock := Lock{
		Schema: 1, Kind: kind, PythonTag: base.PythonTag, Platform: base.Platform,
		BaseID: baseID, Packages: lockPkgs,
	}
	if lock.PythonTag == "" {
		lock.PythonTag = "cp314"
	}
	if lock.Platform == "" {
		lock.Platform = "any"
	}
	id, err := b.Root.BuildLayer(lock)
	if err != nil {
		return "", err
	}
	if kind == "base" {
		if err := b.Root.SwitchCurrent(id); err != nil {
			return "", err
		}
		return id, nil
	}
	bind, err := b.Store.GetExecEnvBinding(conv)
	if err != nil {
		return "", err
	}
	bind.ConversationID = conv
	bind.BaseID = baseID
	bind.UpdatedAt = b.now()
	if isolate {
		bind.Isolated = id
		bind.Extensions = nil
	} else {
		bind.Isolated = ""
		bind.Extensions = appendUnique(bind.Extensions, id)
	}
	if err := b.Store.PutExecEnvBinding(bind); err != nil {
		return "", err
	}
	_ = owner
	return id, nil
}

func wheelExists(r *Root, pkg LockedPkg) bool {
	if pkg.SHA256 == "" {
		return false
	}
	_, err := os.Stat(r.Wheelhouse() + "/" + pkg.SHA256 + ".whl")
	return err == nil
}

func appendUnique(ss []string, id string) []string {
	for _, s := range ss {
		if s == id {
			return ss
		}
	}
	return append(ss, id)
}

// LayersFor 返回本次执行要挂载的 site-packages。isolated 非空时调用方不要再挂基础层。
func (b *Broker) LayersFor(conv string) (extra []string, isolated string) {
	if b.Store == nil || conv == "" {
		return nil, ""
	}
	bind, err := b.Store.GetExecEnvBinding(conv)
	if err != nil {
		return nil, ""
	}
	if bind.Isolated != "" {
		return nil, b.Root.LayerDir(bind.Isolated)
	}
	for _, id := range bind.Extensions {
		extra = append(extra, b.Root.LayerDir(id))
	}
	return extra, ""
}

// Wait 轮询申请直到终态或超时。
func (b *Broker) Wait(ctx context.Context, id string, timeout time.Duration) (Outcome, error) {
	deadline := time.Now().Add(timeout)
	if timeout <= 0 {
		deadline = time.Now().Add(2 * time.Minute)
	}
	for {
		rec, err := b.Store.GetExecEnvRequest(id)
		if err != nil {
			return Outcome{}, err
		}
		switch rec.Status {
		case "ready", "rejected", "failed":
			return Outcome{Status: rec.Status, RequestID: id, LayerID: rec.LayerID, Detail: rec.Detail}, nil
		}
		if time.Now().After(deadline) || ctx.Err() != nil {
			return Outcome{Status: rec.Status, RequestID: id, Detail: "仍在等待审批或构建"}, nil
		}
		select {
		case <-ctx.Done():
			return Outcome{Status: rec.Status, RequestID: id, Detail: "等待被取消"}, nil
		case <-time.After(200 * time.Millisecond):
		}
	}
}

// EnsureBase 解析基础层需求并建立 current。已有 current 时直接返回其 ID。
func (b *Broker) EnsureBase(ctx context.Context) (string, error) {
	if id, err := b.Root.CurrentBase(); err != nil {
		return "", err
	} else if id != "" {
		return id, nil
	}
	res, err := b.Pip.Resolve(ctx, BaseRequirements(), Lock{})
	if err != nil {
		return "", err
	}
	if !res.OK() {
		return "", fmt.Errorf("execenv: base resolve: override=%v sdist=%v missing=%v", res.Override, res.SDist, res.Missing)
	}
	return b.materialize(ctx, "", "", Lock{Kind: "base"}, res, false)
}

// Rebase 强制按当前需求重建基础层并移动 current。已有会话绑定保持旧 ID。
func (b *Broker) Rebase(ctx context.Context) (string, error) {
	res, err := b.Pip.Resolve(ctx, BaseRequirements(), Lock{})
	if err != nil {
		return "", err
	}
	if !res.OK() {
		return "", fmt.Errorf("execenv: rebase: override=%v sdist=%v missing=%v", res.Override, res.SDist, res.Missing)
	}
	lock := Lock{Schema: 1, Kind: "base", PythonTag: "cp314", Platform: "any", Packages: res.Packages}
	for i := range lock.Packages {
		path, err := b.Pip.Fetch(ctx, b.Root.Wheelhouse(), lock.Packages[i])
		if err != nil {
			return "", err
		}
		sum, err := fileSHA256(path)
		if err != nil {
			return "", err
		}
		lock.Packages[i].SHA256 = sum
		lock.Packages[i].Filename = sum + ".whl"
		stable := b.Root.Wheelhouse() + "/" + sum + ".whl"
		if path != stable {
			_ = os.Rename(path, stable)
		}
	}
	id, err := b.Root.BuildLayer(lock)
	if err != nil {
		return "", err
	}
	if err := b.Root.SwitchCurrent(id); err != nil {
		return "", err
	}
	return id, nil
}