Agent-System-Backend/internal/runtime/workflow_script.go

package runtime

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"agent-system-backend/internal/tools"
	"agent-system-backend/internal/workflow"
	core "agent-system-core"
	_ "embed"
)

//go:embed agentwf.py
var agentwfSource string

const (
	scriptDefaultTimeout = 600 * time.Second
	scriptMaxTimeout     = 1800 * time.Second
	scriptMaxCalls       = 50
	scriptMaxLine        = 1 << 20
)

var (
	agentwfOnce sync.Once
	agentwfDir  string
	agentwfErr  error
)

func agentwfHelperDir() (string, error) {
	agentwfOnce.Do(func() {
		dir, err := os.MkdirTemp("", "agentwf-*")
		if err != nil {
			agentwfErr = err
			return
		}
		if err := os.WriteFile(filepath.Join(dir, "agentwf.py"), []byte(agentwfSource), 0o644); err != nil {
			agentwfErr = err
			return
		}
		agentwfDir = dir
	})
	return agentwfDir, agentwfErr
}

func scriptsEnabled(rc RunCtx) bool {
	return rc.Settings().GetBool("runtime.workflowScripts", false)
}

func scriptSandbox(rc RunCtx) (string, error) {
	sandbox := strings.TrimSpace(conversationSandbox(rc))
	if sandbox == "" {
		return "", fmt.Errorf("workflow: 没有会话沙箱")
	}
	return sandbox, nil
}

func executeScriptStep(rc RunCtx, runner WorkflowRunner, step core.WorkflowStep, scripts map[string]string) (string, []*core.ArtifactRef, error) {
	body := scripts[step.Script]
	if strings.TrimSpace(body) == "" {
		var err error
		body, err = scriptBody(rc, step.Script)
		if err != nil {
			return "", nil, err
		}
	}
	sandbox, err := scriptSandbox(rc)
	if err != nil {
		return "", nil, err
	}
	path, err := writeScriptFile(sandbox, "step-"+step.ID, body)
	if err != nil {
		return "", nil, err
	}
	return runScriptFile(rc, runner, path)
}

func scriptBody(rc RunCtx, path string) (string, error) {
	src, ok := rc.(interface {
		Workflow() (workflow.Launch, bool)
	})
	if !ok {
		return "", fmt.Errorf("workflow: 脚本 %s 不在启动副本里", path)
	}
	launch, ok := src.Workflow()
	if !ok {
		return "", fmt.Errorf("workflow: 脚本 %s 不在启动副本里", path)
	}
	body, ok := launch.Scripts[path]
	if !ok || strings.TrimSpace(body) == "" {
		return "", fmt.Errorf("workflow: 脚本 %s 不在启动副本里", path)
	}
	return body, nil
}

func registeredWorkflowLaunch(rc RunCtx, def core.WorkflowDef, slots map[string]string) workflow.Launch {
	launch := workflow.Launch{Def: def, Slots: slots, Scripts: map[string]string{}}
	lookup, ok := rc.(interface {
		WorkflowScript(id, path string) (string, bool)
	})
	if !ok {
		return launch
	}
	for _, step := range def.Steps {
		if step.Kind != "script" || strings.TrimSpace(step.Script) == "" {
			continue
		}
		if body, found := lookup.WorkflowScript(def.ID, step.Script); found {
			launch.Scripts[step.Script] = body
		}
	}
	return launch
}

func runModelScript(rc RunCtx, runner WorkflowRunner, body string) (string, []*core.ArtifactRef, error) {
	if !scriptsEnabled(rc) {
		return "", nil, fmt.Errorf("workflow: 工作流脚本未开启")
	}
	preview := body
	if len(preview) > 4000 {
		preview = preview[:4000] + "\n…（脚本超长，已截断展示）"
	}
	resp, err := rc.AskHITL(core.HITLCard{
		Question: "运行这段脚本？\n" + preview, Kind: core.HITLKindConfirm, AllowFreeInput: false,
		Choices: []core.HITLChoice{
			{ID: "yes", Label: "运行", Description: "在会话沙箱里执行", Impact: "会调用工具并可能写入工件"},
			{ID: "no", Label: "取消", Description: "不执行", Impact: "不会改动任何东西"},
		},
	})
	if err != nil {
		return "", nil, err
	}
	if resp.ChoiceID != "yes" {
		return "未确认，未执行", nil, nil
	}
	sandbox, err := scriptSandbox(rc)
	if err != nil {
		return "", nil, err
	}
	path, err := writeScriptFile(sandbox, "model", body)
	if err != nil {
		return "", nil, err
	}
	return runScriptFile(rc, runner, path)
}

func writeScriptFile(sandbox, prefix, body string) (string, error) {
	dir := filepath.Join(sandbox, ".agent")
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return "", err
	}
	path := filepath.Join(dir, fmt.Sprintf("wf-%s-%d.py", prefix, time.Now().UnixNano()))
	if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
		return "", err
	}
	return path, nil
}

func runScriptFile(rc RunCtx, runner WorkflowRunner, path string) (string, []*core.ArtifactRef, error) {
	if !scriptsEnabled(rc) {
		return "", nil, fmt.Errorf("workflow: 工作流脚本未开启")
	}
	sandbox, err := scriptSandbox(rc)
	if err != nil {
		return "", nil, err
	}
	timeout := runner.ScriptTimeout
	if timeout <= 0 {
		timeout = scriptDefaultTimeout
	}
	if timeout > scriptMaxTimeout {
		timeout = scriptMaxTimeout
	}
	helper, err := agentwfHelperDir()
	if err != nil {
		return "", nil, err
	}
	parent := rc.Context()
	if parent == nil {
		parent = context.Background()
	}
	ctx, cancel := context.WithTimeout(parent, timeout)
	defer cancel()
	cmd, err := tools.ScriptCommand(ctx, sandbox, path)
	if err != nil {
		return "", nil, err
	}
	cmd.Env = prependPythonPath(cmd.Env, helper)
	reqR, reqW, err := os.Pipe()
	if err != nil {
		return "", nil, err
	}
	respR, respW, err := os.Pipe()
	if err != nil {
		reqR.Close()
		reqW.Close()
		return "", nil, err
	}
	cmd.ExtraFiles = []*os.File{reqW, respR}
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Start(); err != nil {
		reqR.Close()
		reqW.Close()
		respR.Close()
		respW.Close()
		return "", nil, err
	}
	reqW.Close()
	respR.Close()
	defer reqR.Close()
	defer respW.Close()
	arts, callErr := serveScript(ctx, rc, reqR, respW)
	if callErr != nil {
		cancel()
	}
	waitErr := cmd.Wait()
	text := clipScriptOutput(&stdout, &stderr)
	if callErr != nil {
		return text, arts, callErr
	}
	if ctx.Err() == context.DeadlineExceeded {
		return text, arts, fmt.Errorf("workflow: 脚本超时")
	}
	if ctx.Err() == context.Canceled {
		return text, arts, fmt.Errorf("workflow: 脚本已取消")
	}
	if waitErr != nil {
		return text, arts, fmt.Errorf("workflow: 脚本退出：%v\n%s", waitErr, strings.TrimSpace(stderr.String()))
	}
	return text, arts, nil
}

func clipScriptOutput(stdout, stderr *bytes.Buffer) string {
	text := strings.TrimSpace(stdout.String())
	errText := strings.TrimSpace(stderr.String())
	if errText != "" {
		if text != "" {
			text += "\n"
		}
		text += errText
	}
	if len(text) > 16000 {
		text = text[:16000] + "…（超长截断）"
	}
	return text
}

func prependPythonPath(env []string, dir string) []string {
	for i, kv := range env {
		if strings.HasPrefix(kv, "PYTHONPATH=") {
			env[i] = "PYTHONPATH=" + dir + string(os.PathListSeparator) + strings.TrimPrefix(kv, "PYTHONPATH=")
			return env
		}
	}
	return append(env, "PYTHONPATH="+dir)
}

type scriptRequest struct {
	ID       string            `json:"id"`
	Op       string            `json:"op"`
	Name     string            `json:"name"`
	Args     json.RawMessage   `json:"args"`
	Question string            `json:"question"`
	Choices  []core.HITLChoice `json:"choices"`
	Agent    string            `json:"agent"`
	Task     string            `json:"task"`
	Content  string            `json:"content"`
	Title    string            `json:"title"`
	Kind     string            `json:"kind"`
	Status   string            `json:"status"`
	StepID   string            `json:"stepId"`
}

type scriptResponse struct {
	ID     string `json:"id"`
	OK     bool   `json:"ok"`
	Output string `json:"output,omitempty"`
	Error  string `json:"error,omitempty"`
}

func serveScript(ctx context.Context, rc RunCtx, reqR, respW *os.File) ([]*core.ArtifactRef, error) {
	sc := bufio.NewScanner(reqR)
	sc.Buffer(make([]byte, 64*1024), scriptMaxLine)
	var arts []*core.ArtifactRef
	calls := 0
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" {
			continue
		}
		calls++
		var req scriptRequest
		if err := json.Unmarshal([]byte(line), &req); err != nil {
			_ = writeScriptResponse(respW, scriptResponse{OK: false, Error: "请求不是合法 JSON"})
			continue
		}
		if calls > scriptMaxCalls {
			_ = writeScriptResponse(respW, scriptResponse{ID: req.ID, OK: false, Error: "脚本回调超过 50 次"})
			return arts, fmt.Errorf("workflow: 脚本回调超过 50 次")
		}
		if ctx.Err() != nil {
			return arts, ctx.Err()
		}
		output, got, err := dispatchScript(rc, req)
		arts = append(arts, got...)
		if err != nil {
			_ = writeScriptResponse(respW, scriptResponse{ID: req.ID, OK: false, Error: err.Error()})
			if ctx.Err() != nil {
				return arts, ctx.Err()
			}
			continue
		}
		if err := writeScriptResponse(respW, scriptResponse{ID: req.ID, OK: true, Output: output}); err != nil {
			return arts, err
		}
	}
	if err := sc.Err(); err != nil {
		if err == bufio.ErrTooLong {
			return arts, fmt.Errorf("workflow: 单行请求超过 1MB")
		}
		if ctx.Err() != nil {
			return arts, ctx.Err()
		}
		return arts, err
	}
	return arts, nil
}

func writeScriptResponse(w io.Writer, resp scriptResponse) error {
	raw, err := json.Marshal(resp)
	if err != nil {
		return err
	}
	_, err = w.Write(append(raw, '\n'))
	return err
}

func dispatchScript(rc RunCtx, req scriptRequest) (string, []*core.ArtifactRef, error) {
	switch req.Op {
	case "tool.call":
		if _, denied := workflow.DeniedTools[req.Name]; denied {
			return "", nil, fmt.Errorf("不能调用 %s", req.Name)
		}
		res := rc.ExecuteTool(core.ToolCallRequest{Name: req.Name, Arguments: req.Args})
		if res.Status == "error" {
			return "", nil, fmt.Errorf("%s", res.Output)
		}
		var arts []*core.ArtifactRef
		if res.Artifact != nil {
			arts = append(arts, &core.ArtifactRef{ArtifactID: res.Artifact.ID, Name: res.Artifact.Name, Kind: res.Artifact.Kind})
		}
		arts = append(arts, artifactsIn(rc, res.Output)...)
		return res.Output, arts, nil
	case "hitl.ask":
		card := core.HITLCard{Question: req.Question, Kind: core.HITLKindInput, AllowFreeInput: true}
		if len(req.Choices) > 0 {
			card.Kind = core.HITLKindChoice
			card.AllowFreeInput = false
			card.Choices = req.Choices
			for i := range card.Choices {
				if card.Choices[i].Impact == "" {
					card.Choices[i].Impact = "会按这个选择继续"
				}
			}
		}
		resp, err := rc.AskHITL(card)
		if err != nil {
			return "", nil, err
		}
		if strings.TrimSpace(resp.FreeText) != "" {
			return strings.TrimSpace(resp.FreeText), nil, nil
		}
		return resp.ChoiceID, nil, nil
	case "agent.spawn":
		agent := req.Agent
		if strings.TrimSpace(agent) == "" {
			agent = "general"
		}
		var res core.AgentResultPayload
		var err error
		if sp, ok := rc.(interface {
			SpawnAgent(agentID, name, task, hint string) (core.AgentResultPayload, error)
		}); ok {
			res, err = sp.SpawnAgent(subagentProfileID(agent), req.Title, req.Task, "")
		} else {
			res, err = rc.SpawnSubagent(req.Title, req.Task, "")
		}
		if err != nil {
			return "", nil, err
		}
		text := res.Text
		if text == "" {
			text = res.Summary
		}
		return text, artifactsIn(rc, text), nil
	case "artifact.register":
		name := req.Name
		if name == "" {
			name = "script.md"
		}
		art, err := rc.WriteArtifact(name, "document", "text/markdown", []byte(req.Content))
		if err != nil {
			return "", nil, err
		}
		return string(art.ID), []*core.ArtifactRef{{ArtifactID: art.ID, Name: art.Name, Kind: art.Kind}}, nil
	case "card.create":
		kind := core.CardKind(req.Kind)
		if kind == "" {
			kind = core.CardKindResult
		}
		title := req.Title
		if title == "" {
			title = req.Name
		}
		created, err := rc.CreateCard(core.ActivityCard{Title: title, Kind: kind})
		if err != nil {
			return "", nil, err
		}
		return string(created.ID), nil, nil
	case "step":
		applyScriptPlan(rc, req.StepID, req.Title, req.Status)
		return req.Status, nil, nil
	default:
		return "", nil, fmt.Errorf("未知回调 %s", req.Op)
	}
}

func applyScriptPlan(rc RunCtx, id, title, status string) {
	switch status {
	case "pending", "in_progress", "done", "skipped":
	default:
		status = "done"
	}
	snap := core.PlanSnapshot{}
	if src, ok := rc.(interface {
		Plan() (core.PlanSnapshot, bool)
	}); ok {
		if cur, found := src.Plan(); found {
			snap = cur
			snap.Steps = append([]core.PlanStep(nil), cur.Steps...)
		}
	}
	found := false
	for i := range snap.Steps {
		if snap.Steps[i].ID == id {
			snap.Steps[i].Status = status
			if title != "" {
				snap.Steps[i].Text = title
			}
			found = true
			break
		}
	}
	if !found {
		snap.Steps = append(snap.Steps, core.PlanStep{ID: id, Text: title, Status: status})
	}
	_ = rc.UpdatePlan(snap)
}