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

package runtime

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"agent-system-backend/internal/workflow"
	core "agent-system-core"
)

func scriptSettings() core.EffectiveSettings {
	return core.EffectiveSettings{"runtime": map[string]any{"workflowScripts": true}}
}

func TestScriptStepRunsEachCallback(t *testing.T) {
	body := `import agentwf as wf
print("tool", wf.tool("time_now", {}))
print("ask", wf.ask("继续？", [{"id":"yes","label":"好","impact":"继续"}]))
print("agent", wf.agent("research", "写一句"))
print("art", wf.artifact("note.md", "你好"))
print("card", wf.card("结果", "result"))
print("step", wf.step("s1", "一步", "done"))
print("file", wf.tool("make_file", {}))
`
	f := &wfFake{
		sandbox:  t.TempDir(),
		settings: scriptSettings(),
		hitl:     []core.HITLResponse{{ChoiceID: "yes"}},
		launch: workflow.Launch{
			Def: core.WorkflowDef{ID: "demo.script", Title: "脚本", Steps: []core.WorkflowStep{
				{ID: "run", Title: "跑", Kind: "script", Script: "workflows/scripts/demo.py"},
				{ID: "save", Title: "留下", Kind: "artifact", Prompt: "{{steps.run.output}}"},
			}},
			Scripts: map[string]string{"workflows/scripts/demo.py": body},
		},
	}
	if err := (WorkflowRunner{}).Execute(f); err != nil {
		t.Fatal(err)
	}
	saved := f.writes[len(f.writes)-1]
	for _, want := range []string{"tool 12:00", "ask yes", "agent 正文", "art art_aaaaaaaaaaaa", "card card_1", "file wrote"} {
		if !strings.Contains(saved, want) {
			t.Fatalf("missing %q in %s", want, saved)
		}
	}
	if len(f.tools) != 2 || f.tools[0].Name != "time_now" || f.tools[1].Name != "make_file" {
		t.Fatalf("tools %+v", f.tools)
	}
	if len(f.cards) != 1 || f.cards[0].Title != "结果" {
		t.Fatalf("cards %+v", f.cards)
	}
	if len(f.spawnIDs) != 1 || f.spawnIDs[0] != "agent.sub.research" {
		t.Fatalf("spawn %v", f.spawnIDs)
	}
	if f.writes[0] != "note.md=你好" {
		t.Fatalf("writes %v", f.writes)
	}
	sawStep := false
	var ids []string
	for _, plan := range f.plans {
		for _, step := range plan.Steps {
			if step.ID == "s1" && step.Status == "done" {
				sawStep = true
			}
		}
	}
	for _, part := range f.assistant {
		if part.Artifact != nil {
			ids = append(ids, string(part.Artifact.ArtifactID))
		}
	}
	if !sawStep {
		t.Fatalf("plans %+v", f.plans)
	}
	joined := strings.Join(ids, " ")
	if !strings.Contains(joined, "art_aaaaaaaaaaaa") || !strings.Contains(joined, "art_bbbbbbbbbbbb") {
		t.Fatalf("deliverables %v", ids)
	}
}

func TestScriptRejectsDeniedTool(t *testing.T) {
	dir := t.TempDir()
	body := `import agentwf as wf
try:
    wf.tool("run_python", {})
except Exception as exc:
    open("ERR", "w").write(str(exc))
`
	f := &wfFake{
		sandbox:  dir,
		settings: scriptSettings(),
		launch: workflow.Launch{
			Def: core.WorkflowDef{ID: "demo.script", Title: "脚本", Steps: []core.WorkflowStep{{
				ID: "run", Title: "跑", Kind: "script", Script: "a.py",
			}}},
			Scripts: map[string]string{"a.py": body},
		},
	}
	if err := (WorkflowRunner{}).Execute(f); err != nil {
		t.Fatal(err)
	}
	raw, err := os.ReadFile(filepath.Join(dir, "ERR"))
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(raw), "不能调用 run_python") || len(f.tools) != 0 {
		t.Fatalf("err %s tools %+v", raw, f.tools)
	}
}

func TestScriptStopsAfterCallCap(t *testing.T) {
	body := `import agentwf as wf
for i in range(60):
    wf.tool("time_now", {})
`
	f := &wfFake{
		sandbox:  t.TempDir(),
		settings: scriptSettings(),
		launch: workflow.Launch{
			Def: core.WorkflowDef{ID: "demo.script", Title: "脚本", Steps: []core.WorkflowStep{{
				ID: "run", Title: "跑", Kind: "script", Script: "a.py",
			}}},
			Scripts: map[string]string{"a.py": body},
		},
	}
	err := (WorkflowRunner{}).Execute(f)
	if err == nil || !strings.Contains(err.Error(), "超过 50") || len(f.tools) != 50 {
		t.Fatalf("err %v calls %d", err, len(f.tools))
	}
}

func TestScriptCancelKillsProcess(t *testing.T) {
	dir := t.TempDir()
	body := `import os, time
open("PID", "w").write(str(os.getpid()))
time.sleep(30)
open("AFTER", "w").write("ran")
`
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	f := &wfFake{
		ctx:      ctx,
		sandbox:  dir,
		settings: scriptSettings(),
		launch: workflow.Launch{
			Def: core.WorkflowDef{ID: "demo.script", Title: "脚本", Steps: []core.WorkflowStep{{
				ID: "run", Title: "跑", Kind: "script", Script: "a.py",
			}}},
			Scripts: map[string]string{"a.py": body},
		},
	}
	done := make(chan error, 1)
	go func() { done <- (WorkflowRunner{}).Execute(f) }()
	deadline := time.Now().Add(5 * time.Second)
	for {
		if _, err := os.Stat(filepath.Join(dir, "PID")); err == nil {
			break
		}
		if time.Now().After(deadline) {
			cancel()
			t.Fatal("script did not start")
		}
		time.Sleep(20 * time.Millisecond)
	}
	raw, err := os.ReadFile(filepath.Join(dir, "PID"))
	if err != nil {
		t.Fatal(err)
	}
	cancel()
	select {
	case err = <-done:
		if err == nil || !strings.Contains(err.Error(), "取消") {
			t.Fatalf("%v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("cancel did not stop the script")
	}
	if _, err := os.Stat("/proc/" + strings.TrimSpace(string(raw))); err == nil {
		t.Fatal("script process still alive")
	}
	if _, err := os.Stat(filepath.Join(dir, "AFTER")); err == nil {
		t.Fatal("script kept running after cancel")
	}
}

func TestScriptLineLimit(t *testing.T) {
	body := `import os
f = os.fdopen(3, "w", encoding="utf-8", buffering=1)
f.write("x" * (1024 * 1024 + 32) + "\n")
f.flush()
`
	f := &wfFake{
		sandbox:  t.TempDir(),
		settings: scriptSettings(),
		launch: workflow.Launch{
			Def: core.WorkflowDef{ID: "demo.script", Title: "脚本", Steps: []core.WorkflowStep{{
				ID: "run", Title: "跑", Kind: "script", Script: "a.py",
			}}},
			Scripts: map[string]string{"a.py": body},
		},
	}
	err := (WorkflowRunner{}).Execute(f)
	if err == nil || !strings.Contains(err.Error(), "1MB") {
		t.Fatal(err)
	}
}

func TestRegisteredWorkflowUsesItsOwnScriptCopy(t *testing.T) {
	f := &wfFake{sandbox: t.TempDir(), settings: scriptSettings()}
	res, err := (WorkflowRunner{}).run(f, workflow.Launch{
		Def: core.WorkflowDef{ID: "demo.script", Title: "脚本", Steps: []core.WorkflowStep{{
			ID: "run", Title: "跑", Kind: "script", Script: "workflows/scripts/demo.py",
		}}},
		Scripts: map[string]string{"workflows/scripts/demo.py": "print('from-launch')\n"},
	})
	if err != nil || !strings.Contains(res.Summary, "完成") {
		t.Fatalf("err %v summary %s", err, res.Summary)
	}
}

func TestModelScriptNeedsSwitchAndConfirm(t *testing.T) {
	f := &wfFake{sandbox: t.TempDir()}
	out := (&directRun{rc: f}).runRegisteredWorkflow(`{"script":"print('pong')"}`)
	if !strings.Contains(out, "未开启") || len(f.hitlQ) != 0 {
		t.Fatalf("%s questions %v", out, f.hitlQ)
	}

	f2 := &wfFake{sandbox: t.TempDir(), settings: scriptSettings(), hitl: []core.HITLResponse{{ChoiceID: "no"}}}
	out = (&directRun{rc: f2}).runRegisteredWorkflow(`{"script":"print('pong')"}`)
	if out != "未确认，未执行" || !strings.Contains(f2.hitlQ[0], "print('pong')") {
		t.Fatalf("%s questions %v", out, f2.hitlQ)
	}

	f3 := &wfFake{sandbox: t.TempDir(), settings: scriptSettings(), hitl: []core.HITLResponse{{ChoiceID: "yes"}}}
	out = (&directRun{rc: f3}).runRegisteredWorkflow(`{"script":"print('pong')"}`)
	if !strings.Contains(out, "pong") {
		t.Fatalf("%s", out)
	}
}