mirror of
https://github.com/langgenius/dify.git
synced 2026-09-05 08:48:10 +08:00
471 lines
14 KiB
Go
471 lines
14 KiB
Go
package agentcli
|
|
|
|
import (
|
|
"archive/zip"
|
|
"bytes"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestConfigPullRequestsURLThenDownloadsFromDataPlane(t *testing.T) {
|
|
skillArchive := zipFixture(t, map[string]string{"SKILL.md": "# Alpha\n", "reference.md": "guide"})
|
|
tests := []struct {
|
|
name string
|
|
kind string
|
|
assetName string
|
|
payload []byte
|
|
run func(*Environment, string) error
|
|
assertFiles func(*testing.T, string)
|
|
}{
|
|
{
|
|
name: "file",
|
|
kind: "file",
|
|
assetName: "guide.txt",
|
|
payload: []byte("guide"),
|
|
run: func(env *Environment, targetDir string) error {
|
|
return RunConfigFilesPull(env, []string{"guide.txt"}, targetDir, true)
|
|
},
|
|
assertFiles: func(t *testing.T, targetDir string) {
|
|
data, err := os.ReadFile(filepath.Join(targetDir, "guide.txt"))
|
|
if err != nil {
|
|
t.Fatalf("read pulled config file: %v", err)
|
|
}
|
|
if string(data) != "guide" {
|
|
t.Fatalf("pulled config file = %q", data)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "skill",
|
|
kind: "skill",
|
|
assetName: "alpha",
|
|
payload: skillArchive,
|
|
run: func(env *Environment, targetDir string) error {
|
|
return RunConfigSkillsPull(env, []string{"alpha"}, targetDir, true)
|
|
},
|
|
assertFiles: func(t *testing.T, targetDir string) {
|
|
data, err := os.ReadFile(filepath.Join(targetDir, "alpha", "SKILL.md"))
|
|
if err != nil {
|
|
t.Fatalf("read pulled skill: %v", err)
|
|
}
|
|
if string(data) != "# Alpha\n" {
|
|
t.Fatalf("pulled SKILL.md = %q", data)
|
|
}
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
var controlPayload map[string]any
|
|
dataPlaneCalls := 0
|
|
var server *httptest.Server
|
|
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/agent-stub/files/download-request":
|
|
if err := json.NewDecoder(r.Body).Decode(&controlPayload); err != nil {
|
|
http.Error(w, "bad request", http.StatusBadRequest)
|
|
return
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"filename": test.assetName,
|
|
"mime_type": "application/octet-stream",
|
|
"size": len(test.payload),
|
|
"download_url": server.URL + "/files/config-asset",
|
|
})
|
|
case "/files/config-asset":
|
|
dataPlaneCalls++
|
|
_, _ = w.Write(test.payload)
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
targetDir := t.TempDir()
|
|
err := test.run(&Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"}, targetDir)
|
|
if err != nil {
|
|
t.Fatalf("pull config %s: %v", test.kind, err)
|
|
}
|
|
|
|
config, ok := controlPayload["config"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("control payload config = %#v", controlPayload["config"])
|
|
}
|
|
if config["kind"] != test.kind || config["name"] != test.assetName {
|
|
t.Fatalf("control payload config = %#v", config)
|
|
}
|
|
if controlPayload["for_frontend"] != false {
|
|
t.Fatalf("for_frontend = %#v", controlPayload["for_frontend"])
|
|
}
|
|
if dataPlaneCalls != 1 {
|
|
t.Fatalf("data-plane calls = %d, want 1", dataPlaneCalls)
|
|
}
|
|
test.assertFiles(t, targetDir)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestConfigPullMultiItemFailuresIdentifyItemAndStage(t *testing.T) {
|
|
skillArchive := zipFixture(t, map[string]string{"SKILL.md": "# Alpha\n"})
|
|
tests := []struct {
|
|
name string
|
|
kind string
|
|
names []string
|
|
payload []byte
|
|
failureStage string
|
|
wantStage string
|
|
}{
|
|
{
|
|
name: "file control-plane URL acquisition",
|
|
kind: "file",
|
|
names: []string{"first.txt", "second.txt"},
|
|
payload: []byte("first file"),
|
|
failureStage: "control",
|
|
wantStage: `request config file "second.txt" download URL`,
|
|
},
|
|
{
|
|
name: "file signed-URL download",
|
|
kind: "file",
|
|
names: []string{"first.txt", "second.txt"},
|
|
payload: []byte("first file"),
|
|
failureStage: "data",
|
|
wantStage: `download config file "second.txt"`,
|
|
},
|
|
{
|
|
name: "skill control-plane URL acquisition",
|
|
kind: "skill",
|
|
names: []string{"alpha", "beta"},
|
|
payload: skillArchive,
|
|
failureStage: "control",
|
|
wantStage: `request config skill "beta" download URL`,
|
|
},
|
|
{
|
|
name: "skill signed-URL download",
|
|
kind: "skill",
|
|
names: []string{"alpha", "beta"},
|
|
payload: skillArchive,
|
|
failureStage: "data",
|
|
wantStage: `download config skill "beta"`,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
controlCalls := 0
|
|
dataPlaneCalls := 0
|
|
var server *httptest.Server
|
|
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/agent-stub/files/download-request":
|
|
controlCalls++
|
|
var request struct {
|
|
Config struct {
|
|
Kind string `json:"kind"`
|
|
Name string `json:"name"`
|
|
} `json:"config"`
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
|
|
t.Errorf("decode download request: %v", err)
|
|
http.Error(w, "bad request", http.StatusBadRequest)
|
|
return
|
|
}
|
|
if controlCalls > len(test.names) {
|
|
t.Errorf("unexpected control-plane call %d", controlCalls)
|
|
http.Error(w, "unexpected request", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
wantName := test.names[controlCalls-1]
|
|
if request.Config.Kind != test.kind || request.Config.Name != wantName {
|
|
t.Errorf("config request = (%q, %q), want (%q, %q)", request.Config.Kind, request.Config.Name, test.kind, wantName)
|
|
}
|
|
if test.failureStage == "control" && controlCalls == 2 {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
_, _ = w.Write([]byte(`{"detail":{"code":"agent_stub_authorization_expired","message":"expired"}}`))
|
|
return
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"filename": wantName,
|
|
"size": len(test.payload),
|
|
"download_url": server.URL + "/files/config-asset",
|
|
})
|
|
case "/files/config-asset":
|
|
dataPlaneCalls++
|
|
if test.failureStage == "data" && dataPlaneCalls == 2 {
|
|
http.Error(w, "data plane unavailable", http.StatusBadGateway)
|
|
return
|
|
}
|
|
_, _ = w.Write(test.payload)
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
targetDir := t.TempDir()
|
|
env := &Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"}
|
|
var err error
|
|
if test.kind == "file" {
|
|
err = RunConfigFilesPull(env, test.names, targetDir, true)
|
|
} else {
|
|
err = RunConfigSkillsPull(env, test.names, targetDir, true)
|
|
}
|
|
if err == nil {
|
|
t.Fatal("config pull succeeded, want second-item failure")
|
|
}
|
|
if !strings.Contains(err.Error(), test.wantStage) {
|
|
t.Errorf("error = %q, want stage %q", err, test.wantStage)
|
|
}
|
|
if controlCalls != 2 {
|
|
t.Errorf("control-plane calls = %d, want 2", controlCalls)
|
|
}
|
|
wantDataPlaneCalls := 2
|
|
if test.failureStage == "control" {
|
|
wantDataPlaneCalls = 1
|
|
for _, want := range []string{
|
|
"expired after 5 minutes",
|
|
"will not refresh automatically",
|
|
"start a new shell tool call",
|
|
"retry the command",
|
|
} {
|
|
if !strings.Contains(err.Error(), want) {
|
|
t.Errorf("error = %q, want expiration guidance %q", err, want)
|
|
}
|
|
}
|
|
}
|
|
if dataPlaneCalls != wantDataPlaneCalls {
|
|
t.Errorf("data-plane calls = %d, want %d", dataPlaneCalls, wantDataPlaneCalls)
|
|
}
|
|
|
|
if test.kind == "file" {
|
|
data, readErr := os.ReadFile(filepath.Join(targetDir, test.names[0]))
|
|
if readErr != nil || !bytes.Equal(data, test.payload) {
|
|
t.Errorf("first file was not completed: data=%q err=%v", data, readErr)
|
|
}
|
|
} else {
|
|
data, readErr := os.ReadFile(filepath.Join(targetDir, test.names[0], "SKILL.md"))
|
|
if readErr != nil || string(data) != "# Alpha\n" {
|
|
t.Errorf("first skill was not completed: data=%q err=%v", data, readErr)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestConfigPushMultiItemUploadFailuresIdentifyItemAndStage(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
failureStage string
|
|
makeSources func(*testing.T) []string
|
|
run func(*Environment, []string) error
|
|
wantItem string
|
|
}{
|
|
{
|
|
name: "file control-plane URL acquisition",
|
|
failureStage: "request upload URL",
|
|
makeSources: makeConfigFileSources,
|
|
run: RunConfigFilesPush,
|
|
wantItem: `config file "second.txt"`,
|
|
},
|
|
{
|
|
name: "file signed URL data transfer",
|
|
failureStage: "upload data",
|
|
makeSources: makeConfigFileSources,
|
|
run: RunConfigFilesPush,
|
|
wantItem: `config file "second.txt"`,
|
|
},
|
|
{
|
|
name: "skill control-plane URL acquisition",
|
|
failureStage: "request upload URL",
|
|
makeSources: makeConfigSkillSources,
|
|
run: RunConfigSkillsPush,
|
|
wantItem: `config skill "beta"`,
|
|
},
|
|
{
|
|
name: "skill signed URL data transfer",
|
|
failureStage: "upload data",
|
|
makeSources: makeConfigSkillSources,
|
|
run: RunConfigSkillsPush,
|
|
wantItem: `config skill "beta"`,
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
uploadRequestCalls := 0
|
|
dataPlaneCalls := 0
|
|
var server *httptest.Server
|
|
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/agent-stub/files/upload-request":
|
|
uploadRequestCalls++
|
|
if test.failureStage == "request upload URL" && uploadRequestCalls == 2 {
|
|
http.Error(w, "control plane unavailable", http.StatusServiceUnavailable)
|
|
return
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]string{"upload_url": server.URL + "/upload"})
|
|
case "/upload":
|
|
_, _ = io.Copy(io.Discard, r.Body)
|
|
dataPlaneCalls++
|
|
if test.failureStage == "upload data" && dataPlaneCalls == 2 {
|
|
http.Error(w, "data plane unavailable", http.StatusBadGateway)
|
|
return
|
|
}
|
|
_, _ = w.Write([]byte(`{"id":"tool-file-1"}`))
|
|
case "/agent-stub/config/push":
|
|
t.Error("final config push must not run after an item upload failure")
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
err := test.run(
|
|
&Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"},
|
|
test.makeSources(t),
|
|
)
|
|
if err == nil {
|
|
t.Fatal("config push succeeded, want second-item upload failure")
|
|
}
|
|
if !strings.Contains(err.Error(), test.wantItem) {
|
|
t.Errorf("error = %q, want current item %q", err, test.wantItem)
|
|
}
|
|
if !strings.Contains(err.Error(), test.failureStage) {
|
|
t.Errorf("error = %q, want stage %q", err, test.failureStage)
|
|
}
|
|
if uploadRequestCalls != 2 {
|
|
t.Errorf("upload request calls = %d, want 2", uploadRequestCalls)
|
|
}
|
|
wantDataPlaneCalls := 2
|
|
if test.failureStage == "request upload URL" {
|
|
wantDataPlaneCalls = 1
|
|
}
|
|
if dataPlaneCalls != wantDataPlaneCalls {
|
|
t.Errorf("data-plane calls = %d, want %d", dataPlaneCalls, wantDataPlaneCalls)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestConfigPushFinalFailureIdentifiesOperationAndExplainsExpiry(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
makeSource func(*testing.T) string
|
|
run func(*Environment, string) error
|
|
want string
|
|
}{
|
|
{
|
|
name: "file",
|
|
makeSource: func(t *testing.T) string {
|
|
path := filepath.Join(t.TempDir(), "guide.txt")
|
|
if err := os.WriteFile(path, []byte("guide"), 0o600); err != nil {
|
|
t.Fatalf("write config file: %v", err)
|
|
}
|
|
return path
|
|
},
|
|
run: func(env *Environment, path string) error { return RunConfigFilesPush(env, []string{path}) },
|
|
want: "push config files",
|
|
},
|
|
{
|
|
name: "skill",
|
|
makeSource: func(t *testing.T) string {
|
|
dir := filepath.Join(t.TempDir(), "alpha")
|
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
|
t.Fatalf("create config skill: %v", err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(dir, "SKILL.md"), []byte("# Alpha\n"), 0o600); err != nil {
|
|
t.Fatalf("write config skill: %v", err)
|
|
}
|
|
return dir
|
|
},
|
|
run: func(env *Environment, path string) error { return RunConfigSkillsPush(env, []string{path}) },
|
|
want: "push config skills",
|
|
},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
var server *httptest.Server
|
|
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/agent-stub/files/upload-request":
|
|
_ = json.NewEncoder(w).Encode(map[string]string{"upload_url": server.URL + "/upload"})
|
|
case "/upload":
|
|
_, _ = w.Write([]byte(`{"id":"tool-file-1"}`))
|
|
case "/agent-stub/config/push":
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
_, _ = w.Write([]byte(`{"detail":{"code":"agent_stub_authorization_expired","message":"expired"}}`))
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
err := test.run(
|
|
&Environment{URL: server.URL + "/agent-stub", AuthJWE: "token"},
|
|
test.makeSource(t),
|
|
)
|
|
if err == nil || !strings.Contains(err.Error(), test.want) {
|
|
t.Fatalf("error = %v, want operation %q", err, test.want)
|
|
}
|
|
if !strings.Contains(err.Error(), "start a new shell tool call") {
|
|
t.Fatalf("error = %v, want expiry recovery guidance", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func makeConfigFileSources(t *testing.T) []string {
|
|
t.Helper()
|
|
dir := t.TempDir()
|
|
paths := []string{filepath.Join(dir, "first.txt"), filepath.Join(dir, "second.txt")}
|
|
for _, path := range paths {
|
|
if err := os.WriteFile(path, []byte(filepath.Base(path)), 0o600); err != nil {
|
|
t.Fatalf("write config file: %v", err)
|
|
}
|
|
}
|
|
return paths
|
|
}
|
|
|
|
func makeConfigSkillSources(t *testing.T) []string {
|
|
t.Helper()
|
|
dir := t.TempDir()
|
|
paths := []string{filepath.Join(dir, "alpha"), filepath.Join(dir, "beta")}
|
|
for _, path := range paths {
|
|
if err := os.MkdirAll(path, 0o755); err != nil {
|
|
t.Fatalf("create config skill: %v", err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(path, "SKILL.md"), []byte("# "+filepath.Base(path)+"\n"), 0o600); err != nil {
|
|
t.Fatalf("write config skill: %v", err)
|
|
}
|
|
}
|
|
return paths
|
|
}
|
|
|
|
func zipFixture(t *testing.T, files map[string]string) []byte {
|
|
t.Helper()
|
|
var buffer bytes.Buffer
|
|
archive := zip.NewWriter(&buffer)
|
|
for name, content := range files {
|
|
writer, err := archive.Create(name)
|
|
if err != nil {
|
|
t.Fatalf("create zip member: %v", err)
|
|
}
|
|
if _, err := writer.Write([]byte(content)); err != nil {
|
|
t.Fatalf("write zip member: %v", err)
|
|
}
|
|
}
|
|
if err := archive.Close(); err != nil {
|
|
t.Fatalf("close zip fixture: %v", err)
|
|
}
|
|
return buffer.Bytes()
|
|
}
|