xinmi/internal/workbench/agent_duty_test.go

package workbench

import (
	"errors"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/types"
)

func TestAgentDutyEnforcedOnSave(t *testing.T) {
	s := sqliteWorkbench(t)
	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
	for _, role := range []string{"", types.AgentRoleGeneral, types.AgentRoleTrainingMap, types.AgentRoleProduction, types.AgentRoleQC} {
		s.agents = generationAgentStub{config: types.CustomAgentConfig{Role: role}}
		m := &Resource{Name: "map", Kind: "training-maps", Data: asJSON(MapData{Topic: "topic", AgentID: "agent"})}
		f := &Resource{Name: "flow", Kind: "flows", Data: asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "agent"}}})}
		for _, item := range []struct {
			resource *Resource
			duty     string
		}{{m, types.AgentRoleTrainingMap}, {f, types.AgentRoleProduction}} {
			err := s.validate(ctx, item.resource)
			want := (types.CustomAgentConfig{Role: role}).SupportsDuty(item.duty)
			if (err == nil) != want {
				t.Fatalf("role %q duty %q: %v", role, item.duty, err)
			}
		}
	}
}

func TestAgentContentTagsReachTaskDecomposition(t *testing.T) {
	s, ctx, job, model := generationHarness(t)
	s.agents = generationAgentStub{config: types.CustomAgentConfig{Role: types.AgentRoleProduction, ContentTags: []string{"课程标签"}}}
	f := seedResource(t, s, 1, "writer", "flows")
	f.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "agent"}}})
	if err := s.db.Save(f).Error; err != nil {
		t.Fatal(err)
	}
	job.Action = "decompose"
	model.output = `{"segments":[{"id":"s","title":"segment","node_id":"n","agent_id":"agent"}]}`
	r := &Resource{Kind: "content-tasks", Data: asJSON(TaskData{FlowID: f.ID, AgentID: "agent", SourceText: "framework"})}
	if err := s.executeTask(ctx, job, r, ActionRequest{}); err != nil {
		t.Fatal(err)
	}
	if len(model.messages) < 2 || !strings.Contains(model.messages[1].Content, "课程标签") {
		t.Fatal("decomposition omitted native content tags")
	}
}

func TestAgentDutyChangeBlocksSavedMapAndFlowExecution(t *testing.T) {
	s, ctx, job, model := generationHarness(t)
	s.agents = generationAgentStub{config: types.CustomAgentConfig{Role: types.AgentRoleQC}}
	job.Action = "generate"
	m := &Resource{Kind: "training-maps", Data: asJSON(MapData{Topic: "topic", AgentID: "agent"})}
	if err := s.executeMap(ctx, job, m, ActionRequest{}); !errors.Is(err, ErrInvalid) {
		t.Fatalf("map: %v", err)
	}
	if len(model.messages) != 0 {
		t.Fatal("incompatible agent invoked model")
	}
	f := seedResource(t, s, 1, "writer", "flows")
	f.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "n", AgentID: "agent"}}})
	if err := s.db.Save(f).Error; err != nil {
		t.Fatal(err)
	}
	snapshot := FlowData{Nodes: []FlowNode{{ID: "frozen", AgentID: "agent"}}}
	if _, err := s.taskFlow(ctx, TaskData{FlowID: f.ID, FlowSnapshot: &snapshot}); !errors.Is(err, ErrInvalid) {
		t.Fatalf("frozen flow: %v", err)
	}
}