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

package execenv

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// Store 是申请与绑定的持久化。实现在 store 包，避免本包依赖数据库驱动。
type Store interface {
	InsertExecEnvRequest(Request) error
	GetExecEnvRequest(id string) (Request, error)
	ListExecEnvRequests(conversationID string, limit int) ([]Request, error)
	UpdateExecEnvRequest(id, status, layerID, detail, approver string) error
	GetExecEnvBinding(conversationID string) (Binding, error)
	PutExecEnvBinding(Binding) error
	ListExecEnvBindings() ([]Binding, error)
	InsertExecEnvAudit(owner, action, actor, target, outcome, entry string) error
}

// Request 是一次扩展包申请。
type Request struct {
	ID             string `json:"id"`
	ConversationID string `json:"conversationId"`
	Owner          string `json:"owner"`
	Spec           string `json:"spec"`
	Reason         string `json:"reason"`
	Status         string `json:"status"`
	LayerID        string `json:"layerId"`
	Detail         string `json:"detail"`
	Approver       string `json:"approver"`
	CreatedAt      string `json:"createdAt"`
	UpdatedAt      string `json:"updatedAt"`
}

// Binding 是会话当前挂载的层。Isolated 非空时只挂隔离层，不挂基础层。
type Binding struct {
	ConversationID string   `json:"conversationId"`
	BaseID         string   `json:"baseId"`
	Extensions     []string `json:"extensions"`
	Isolated       string   `json:"isolated"`
	UpdatedAt      string   `json:"updatedAt"`
}

// Outcome 是工具直接回给模型的结果。
type Outcome struct {
	Status    string `json:"status"`
	RequestID string `json:"requestId,omitempty"`
	LayerID   string `json:"layerId,omitempty"`
	Detail    string `json:"detail,omitempty"`
}

// Resolver 在宿主侧解析并下载 wheel。沙箱不实现这个接口。
type Resolver interface {
	Resolve(ctx context.Context, reqs []string, base Lock) (Result, error)
	Fetch(ctx context.Context, wheelhouse string, pkg LockedPkg) (string, error)
}

// Broker 在宿主侧解析、分级、构建。沙箱不调用它的构建路径。
type Broker struct {
	Root  *Root
	Store Store
	Pip   Resolver
	Now   func() time.Time
}

func NewBroker(root *Root, st Store, pip Resolver) *Broker {
	return &Broker{Root: root, Store: st, Pip: pip, Now: time.Now}
}

// StatusText 给模型看的紧凑状态。
func (b *Broker) StatusText(conv string) (string, error) {
	var buf strings.Builder
	baseID, err := b.Root.CurrentBase()
	if err != nil {
		return "", err
	}
	if baseID == "" {
		buf.WriteString("基础层：未建立（需要 agent-system env ensure）\n")
	} else {
		lock, err := b.Root.ReadLock(baseID)
		if err != nil {
			return "", err
		}
		buf.WriteString("基础层：" + baseID + "\n")
		for _, p := range lock.Packages {
			fmt.Fprintf(&buf, "- %s==%s\n", p.Name, p.Version)
		}
	}
	if conv != "" && b.Store != nil {
		bind, err := b.Store.GetExecEnvBinding(conv)
		if err != nil {
			return "", err
		}
		if bind.Isolated != "" {
			fmt.Fprintf(&buf, "本会话隔离层：%s\n", bind.Isolated)
		} else if len(bind.Extensions) > 0 {
			fmt.Fprintf(&buf, "本会话扩展层：%s\n", strings.Join(bind.Extensions, ", "))
		} else {
			buf.WriteString("本会话扩展层：无\n")
		}
		reqs, err := b.Store.ListExecEnvRequests(conv, 10)
		if err != nil {
			return "", err
		}
		for _, r := range reqs {
			if r.Status == "ready" || r.Status == "rejected" || r.Status == "failed" {
				continue
			}
			fmt.Fprintf(&buf, "申请 %s 状态 %s %s\n", r.ID, r.Status, r.Detail)
		}
	}
	return buf.String(), nil
}

// RequestPackages 解析并按分级处理。wheelPath 非空表示用户上传的 wheel，一律待审批。
func (b *Broker) RequestPackages(ctx context.Context, conv, owner, reason string, packages []string, wheelPath string) (Outcome, error) {
	if b.Store == nil {
		return Outcome{}, errors.New("execenv: store is not configured")
	}
	baseID, err := b.Root.CurrentBase()
	if err != nil {
		return Outcome{}, err
	}
	var base Lock
	if baseID != "" {
		base, err = b.Root.ReadLock(baseID)
		if err != nil {
			return Outcome{}, err
		}
	}
	cat, err := b.Root.LoadCatalog()
	if err != nil {
		return Outcome{}, err
	}
	var res Result
	uploaded := wheelPath != ""
	if uploaded {
		pkg, err := b.stageUpload(wheelPath)
		if err != nil {
			return Outcome{}, err
		}
		res.Packages = []LockedPkg{pkg}
	} else if len(packages) == 0 {
		return Outcome{}, errors.New("execenv: packages or wheel required")
	} else {
		res, err = b.Pip.Resolve(ctx, packages, base)
		if err != nil {
			return Outcome{}, err
		}
	}
	grade, detail := grade(cat, res, uploaded)
	specRaw, _ := json.Marshal(map[string]any{"packages": packages, "wheel": wheelPath, "resolved": res.Packages})
	rec := Request{
		ID: newID("env"), ConversationID: conv, Owner: owner,
		Spec: string(specRaw), Reason: reason, Status: grade, Detail: detail,
		CreatedAt: b.now(), UpdatedAt: b.now(),
	}
	if err := b.Store.InsertExecEnvRequest(rec); err != nil {
		return Outcome{}, err
	}
	_ = b.Store.InsertExecEnvAudit(owner, "execenv.request", owner, rec.ID, grade, detail)
	switch grade {
	case "rejected":
		return Outcome{Status: "rejected", RequestID: rec.ID, Detail: detail}, nil
	case "pending_approval":
		return Outcome{Status: "pending_approval", RequestID: rec.ID, Detail: detail}, nil
	default:
		id, err := b.materialize(ctx, conv, owner, base, res, false)
		if err != nil {
			_ = b.Store.UpdateExecEnvRequest(rec.ID, "failed", "", err.Error(), "")
			return Outcome{Status: "failed", RequestID: rec.ID, Detail: err.Error()}, nil
		}
		_ = b.Store.UpdateExecEnvRequest(rec.ID, "ready", id, "已绑定", "auto")
		return Outcome{Status: "ready", RequestID: rec.ID, LayerID: id, Detail: "已绑定，下一次 run_python 即可导入"}, nil
	}
}

func grade(cat Catalog, res Result, uploaded bool) (status, detail string) {
	if len(res.Denied) > 0 || hasDenied(cat, res) {
		return "rejected", "命中黑名单：" + strings.Join(deniedNames(cat, res), ", ")
	}
	if len(res.SDist) > 0 {
		return "rejected", "只有源码包，拒绝执行构建脚本：" + strings.Join(res.SDist, ", ")
	}
	if len(res.Missing) > 0 {
		return "rejected", "解析失败：" + strings.Join(res.Missing, "; ")
	}
	if len(res.Packages) == 0 && len(res.Override) == 0 {
		return "rejected", "没有需要安装的新包（基础层已满足，或规格为空）"
	}
	var why []string
	if uploaded {
		why = append(why, "用户上传的 wheel")
	}
	if len(res.Override) > 0 {
		why = append(why, strings.Join(res.Override, "; "))
	}
	for _, p := range res.Packages {
		if !cat.allowed(p.Name) {
			why = append(why, p.Name+" 不在白名单")
		}
		if cat.MaxWheelBytes > 0 && p.Size > cat.MaxWheelBytes {
			why = append(why, fmt.Sprintf("%s 体积 %d 超过上限", p.Name, p.Size))
		}
		if p.SHA256 == "" && !uploaded {
			why = append(why, p.Name+" 缺少哈希")
		}
	}
	if len(why) > 0 {
		return "pending_approval", strings.Join(why, "；")
	}
	return "approved", "白名单内自动批准"
}

func hasDenied(cat Catalog, res Result) bool {
	return len(deniedNames(cat, res)) > 0
}

func deniedNames(cat Catalog, res Result) []string {
	var out []string
	for _, p := range res.Packages {
		if cat.denied(p.Name) {
			out = append(out, p.Name)
		}
	}
	return out
}

// Decide 处理审批。isolate 仅在申请确实与基础层冲突时建立隔离层。
func (b *Broker) Decide(ctx context.Context, id, actor string, admin, allow, isolate bool) (Outcome, error) {
	rec, err := b.Store.GetExecEnvRequest(id)
	if err != nil {
		return Outcome{}, err
	}
	if rec.Status != "pending_approval" {
		return Outcome{}, fmt.Errorf("execenv: request %s is %s", id, rec.Status)
	}
	if !admin && actor != rec.Owner {
		return Outcome{}, errors.New("execenv: only the session owner or an admin can decide")
	}
	if !allow {
		if err := b.Store.UpdateExecEnvRequest(id, "rejected", "", "已拒绝", actor); err != nil {
			return Outcome{}, err
		}
		_ = b.Store.InsertExecEnvAudit(rec.Owner, "execenv.decide", actor, id, "rejected", "")
		return Outcome{Status: "rejected", RequestID: id, Detail: "已拒绝"}, nil
	}
	var spec struct {
		Packages []string    `json:"packages"`
		Resolved []LockedPkg `json:"resolved"`
	}
	_ = json.Unmarshal([]byte(rec.Spec), &spec)
	baseID, err := b.Root.CurrentBase()
	if err != nil {
		return Outcome{}, err
	}
	var base Lock
	if baseID != "" {
		base, err = b.Root.ReadLock(baseID)
		if err != nil {
			return Outcome{}, err
		}
	}
	var res Result
	if isolate {
		// 隔离层不继承基础层约束，整份闭包单独成层。
		res, err = b.Pip.Resolve(ctx, spec.Packages, Lock{})
		if err != nil {
			return Outcome{}, err
		}
	} else if len(spec.Resolved) > 0 {
		res.Packages = spec.Resolved
	} else if len(spec.Packages) > 0 {
		res, err = b.Pip.Resolve(ctx, spec.Packages, base)
		if err != nil {
			return Outcome{}, err
		}
	}
	if len(res.Override) > 0 && !isolate {
		msg := "与基础层版本冲突，批准时需要选择隔离层：" + strings.Join(res.Override, "; ")
		_ = b.Store.UpdateExecEnvRequest(id, "pending_approval", "", msg, "")
		return Outcome{Status: "pending_approval", RequestID: id, Detail: msg}, nil
	}
	layer, err := b.materialize(ctx, rec.ConversationID, rec.Owner, base, res, isolate)
	if err != nil {
		_ = b.Store.UpdateExecEnvRequest(id, "failed", "", err.Error(), actor)
		return Outcome{Status: "failed", RequestID: id, Detail: err.Error()}, nil
	}
	if err := b.Store.UpdateExecEnvRequest(id, "ready", layer, "已批准并绑定", actor); err != nil {
		return Outcome{}, err
	}
	_ = b.Store.InsertExecEnvAudit(rec.Owner, "execenv.decide", actor, id, "ready", layer)
	return Outcome{Status: "ready", RequestID: id, LayerID: layer, Detail: "已批准并绑定"}, nil
}

func (b *Broker) now() string { return b.Now().UTC().Format(time.RFC3339Nano) }

func newID(prefix string) string {
	var buf [8]byte
	_, _ = rand.Read(buf[:])
	return prefix + "_" + hex.EncodeToString(buf[:])
}

func (b *Broker) stageUpload(path string) (LockedPkg, error) {
	if filepath.Ext(path) != ".whl" {
		return LockedPkg{}, errors.New("execenv: upload must be a .whl")
	}
	f, err := os.Open(path)
	if err != nil {
		return LockedPkg{}, err
	}
	defer f.Close()
	sum := sha256New()
	n, err := io.Copy(sum, f)
	if err != nil {
		return LockedPkg{}, err
	}
	sha := sum.Sum()
	dest := filepath.Join(b.Root.Wheelhouse(), sha+".whl")
	if _, err := os.Stat(dest); err != nil {
		in, err := os.Open(path)
		if err != nil {
			return LockedPkg{}, err
		}
		defer in.Close()
		out, err := os.OpenFile(dest, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644)
		if err != nil && !errors.Is(err, os.ErrExist) {
			return LockedPkg{}, err
		}
		if err == nil {
			if _, err := io.Copy(out, in); err != nil {
				out.Close()
				return LockedPkg{}, err
			}
			out.Close()
		}
	}
	name := strings.Split(filepath.Base(path), "-")[0]
	return LockedPkg{Name: name, Version: "0", Filename: filepath.Base(path), SHA256: sha, Size: n}, nil
}