xinmi/internal/application/service/instruction_skill_test.go

package service

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

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

type instructionAgentStub struct {
	interfaces.CustomAgentService
	db *gorm.DB
}

func (s instructionAgentStub) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
	tenant, _ := types.TenantIDFromContext(ctx)
	var a types.CustomAgent
	err := s.db.Where("tenant_id = ? AND id = ?", tenant, id).First(&a).Error
	if err == gorm.ErrRecordNotFound && types.IsBuiltinAgentID(id) {
		return types.GetBuiltinAgent(id, tenant), nil
	}
	return &a, err
}

func instructionContext(tenant uint64, user string, role types.TenantRole) context.Context {
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, tenant)
	ctx = context.WithValue(ctx, types.UserIDContextKey, user)
	return context.WithValue(ctx, types.TenantRoleContextKey, role)
}

func instructionHarness(t *testing.T) (*InstructionSkillService, context.Context) {
	t.Helper()
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = sqlDB.Close() })
	require.NoError(t, db.AutoMigrate(&types.CustomAgent{}, &types.TenantSkillCatalogEntity{}))
	require.NoError(t, db.Exec("CREATE UNIQUE INDEX uq_test_instruction_name ON tenant_skill_catalog(tenant_id,created_by,name,version_source_id) WHERE deleted_at IS NULL AND instruction_only = 1").Error)
	return NewInstructionSkillService(db, instructionAgentStub{db: db}), instructionContext(1, "alice", types.TenantRoleContributor)
}

func TestInstructionVisibilityReflectsManagedPublication(t *testing.T) {
	row := types.TenantSkillCatalogEntity{ID: "skill", Visibility: "private", Enabled: true}
	ctx := accesspolicy.WithSnapshot(context.Background(), &accesspolicy.Snapshot{Publications: map[string]bool{"skills:skill": true}})
	require.Equal(t, "global", instructionSkillView(ctx, row, false).Visibility)
	row.Visibility = "global"
	ctx = accesspolicy.WithSnapshot(context.Background(), &accesspolicy.Snapshot{Publications: map[string]bool{"skills:skill": false}})
	require.Equal(t, "private", instructionSkillView(ctx, row, true).Visibility)
	require.Equal(t, "global", instructionSkillView(context.Background(), row, true).Visibility)
}

func TestInstructionMarkdownAcceptsPlainAndSkillFrontmatter(t *testing.T) {
	plain, err := parseInstructionMarkdown("写作.md", "", "中文说明", []byte("# 写作规范\n保持内容连贯"))
	require.NoError(t, err)
	require.Equal(t, "写作", plain.Name)
	require.Equal(t, "中文说明", plain.Description)
	require.Equal(t, "# 写作规范\n保持内容连贯", plain.Instructions)
	withHeader, err := parseInstructionMarkdown("SKILL.md", "", "", []byte("---\nname: writing\ndescription: Keep continuity\nversion: 1.2\n---\nBody rules"))
	require.NoError(t, err)
	require.Equal(t, "writing", withHeader.Name)
	require.Equal(t, "Body rules", withHeader.Instructions)
	require.True(t, withHeader.InstructionOnly)
	for _, tc := range []struct {
		filename string
		body     []byte
	}{
		{"x.zip", []byte("not a markdown file")},
		{"x.md", []byte{0, 1, 2}},
		{"x.md", []byte{0xff}},
		{"x.md", []byte("   ")},
		{"x.md", []byte(strings.Repeat("x", MaxInstructionSkillBytes+1))},
		{"x.md", []byte("---\nname: missing-end")},
	} {
		_, err := parseInstructionMarkdown(tc.filename, "", "", tc.body)
		require.Error(t, err)
	}
}

func TestInstructionImportPreservesOwnershipAndSandboxDefinitions(t *testing.T) {
	s, ctx := instructionHarness(t)
	first, err := s.Import(ctx, "writing.md", "", "desc", []byte("original rules"))
	require.NoError(t, err)
	updated, err := s.Import(ctx, "writing.md", "", "updated", []byte("updated rules"))
	require.NoError(t, err)
	require.Equal(t, first.ID, updated.ID)
	require.Equal(t, "alice", updated.CreatedBy)
	bob, err := s.Import(instructionContext(1, "bob", types.TenantRoleContributor), "writing.md", "", "", []byte("bob independent rules"))
	require.NoError(t, err)
	require.NotEqual(t, first.ID, bob.ID)
	_, err = s.Get(ctx, bob.ID)
	require.Error(t, err)
	listed, err := s.List(ctx)
	require.NoError(t, err)
	require.Len(t, listed, 1)
	row, err := s.Get(ctx, first.ID)
	require.NoError(t, err)
	require.Equal(t, "updated rules", row.Instructions)
	_, err = s.Import(instructionContext(2, "bob", types.TenantRoleContributor), "writing.md", "", "", []byte("other workspace"))
	require.NoError(t, err)
	_, err = s.Get(instructionContext(2, "bob", types.TenantRoleViewer), first.ID)
	require.Error(t, err)
	require.NoError(t, s.db.Create(&types.TenantSkillCatalogEntity{ID: "sandbox", TenantID: 1, Name: "executable"}).Error)
	textSkill, err := s.Import(ctx, "executable.md", "", "", []byte("independent text skill"))
	require.NoError(t, err)
	nativeSkill, err := repository.NewTenantSkillRepository(s.db).GetCatalogByName(ctx, 1, "executable")
	require.NoError(t, err)
	require.Equal(t, "sandbox", nativeSkill.ID)
	require.NotEqual(t, nativeSkill.ID, textSkill.ID)
	_, err = s.Import(instructionContext(1, "viewer", types.TenantRoleViewer), "new.md", "", "", []byte("rules"))
	require.Error(t, err)
}

func TestInstructionImportDoesNotOverwriteSameNamePersonalVersion(t *testing.T) {
	s, ctx := instructionHarness(t)
	copy := types.TenantSkillCatalogEntity{ID: "personal", TenantID: 1, CreatedBy: "alice", Name: "writing", InstructionOnly: true, VersionSourceID: "global", Instructions: "personal rules"}
	require.NoError(t, s.db.Create(&copy).Error)
	first, err := s.Import(ctx, "writing.md", "", "", []byte("original rules"))
	require.NoError(t, err)
	require.NotEqual(t, copy.ID, first.ID)
	updated, err := s.Import(ctx, "writing.md", "", "", []byte("updated original rules"))
	require.NoError(t, err)
	require.Equal(t, first.ID, updated.ID)
	personal, err := s.Get(ctx, copy.ID)
	require.NoError(t, err)
	require.Equal(t, "personal rules", personal.Instructions)
	require.Equal(t, "global", personal.VersionSourceID)
}

func TestInstructionMountIsAtomicAndPreservesSandboxConfiguration(t *testing.T) {
	s, ctx := instructionHarness(t)
	skill, err := s.Import(ctx, "writing.md", "", "", []byte("Use consistent terminology."))
	require.NoError(t, err)
	for _, a := range []types.CustomAgent{
		{ID: "a", TenantID: 1, CreatedBy: "alice", Name: "A", Config: types.CustomAgentConfig{SystemPrompt: "base", SkillsSelectionMode: "selected", SelectedSkills: []string{"sandbox-skill"}, SandboxConfigID: "sandbox-config"}},
		{ID: "b", TenantID: 1, CreatedBy: "alice", Name: "B"},
		{ID: "c", TenantID: 1, CreatedBy: "bob", Name: "C"},
		{ID: "d", TenantID: 2, CreatedBy: "alice", Name: "D"},
	} {
		require.NoError(t, s.db.Create(&a).Error)
	}
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{"a", "c"}, SkillIDs: []string{skill.ID}})
	require.Error(t, err)
	var untouched types.CustomAgent
	require.NoError(t, s.db.First(&untouched, "id = ?", "a").Error)
	require.Empty(t, untouched.Config.InstructionSkillIDs)
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{"a", "d"}, SkillIDs: []string{skill.ID}})
	require.Error(t, err)
	rows, err := s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{"a", "b", "a"}, SkillIDs: []string{skill.ID, skill.ID}})
	require.NoError(t, err)
	require.Len(t, rows, 2)
	var mounted types.CustomAgent
	require.NoError(t, s.db.First(&mounted, "id = ?", "a").Error)
	require.Equal(t, []string{skill.ID}, mounted.Config.InstructionSkillIDs)
	require.Equal(t, []string{"sandbox-skill"}, mounted.Config.SelectedSkills)
	require.Equal(t, "sandbox-config", mounted.Config.SandboxConfigID)
	require.Equal(t, "base", mounted.Config.SystemPrompt)
	prompt, err := s.ResolveAgentInstructions(ctx, &mounted)
	require.NoError(t, err)
	require.Contains(t, prompt, "Use consistent terminology.")
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{"a"}, SkillIDs: []string{skill.ID}, Mode: "remove"})
	require.NoError(t, err)
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{"b"}, Mode: "replace"})
	require.NoError(t, err)
	for _, id := range []string{"a", "b"} {
		var a types.CustomAgent
		require.NoError(t, s.db.First(&a, "id = ?", id).Error)
		require.Empty(t, a.Config.InstructionSkillIDs)
	}
}

func TestInstructionResolverRejectsForeignMissingAndExecutableSkills(t *testing.T) {
	s, ctx := instructionHarness(t)
	for _, row := range []types.TenantSkillCatalogEntity{
		{ID: "foreign", TenantID: 2, Name: "foreign", InstructionOnly: true, Instructions: "private"},
		{ID: "script", TenantID: 1, Name: "script", Instructions: "execute.sh"},
	} {
		require.NoError(t, s.db.Create(&row).Error)
	}
	for _, id := range []string{"missing", "foreign", "script"} {
		_, err := ResolveAgentInstructionPrompt(ctx, repository.NewTenantSkillRepository(s.db), &types.CustomAgent{TenantID: 1, CreatedBy: "alice", Config: types.CustomAgentConfig{InstructionSkillIDs: []string{id}}})
		require.Error(t, err)
	}
	_, err := s.ResolveAgentInstructions(ctx, &types.CustomAgent{TenantID: 2, Config: types.CustomAgentConfig{InstructionSkillIDs: []string{"foreign"}}})
	require.ErrorContains(t, err, "跨工作空间")
}

func TestInstructionUnmountRepairsInvalidBindingsWithoutSandbox(t *testing.T) {
	s, ctx := instructionHarness(t)
	skill, err := s.Import(ctx, "writing.md", "", "", []byte("Use clear terminology."))
	require.NoError(t, err)
	agent := types.CustomAgent{ID: "repair", TenantID: 1, CreatedBy: "alice", Name: "Repair",
		Config: types.CustomAgentConfig{InstructionSkillIDs: []string{skill.ID, "removed-skill"}}}
	require.NoError(t, s.db.Create(&agent).Error)
	_, err = s.ResolveAgentInstructions(ctx, &agent)
	require.Error(t, err)

	rows, err := s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{agent.ID}, SkillIDs: []string{skill.ID}, Mode: "remove"})
	require.NoError(t, err)
	require.Equal(t, []string{"removed-skill"}, rows[0].SkillIDs)
	var saved types.CustomAgent
	require.NoError(t, s.db.First(&saved, "id = ?", agent.ID).Error)
	require.Empty(t, saved.Config.SandboxConfigID)
	_, err = s.ResolveAgentInstructions(ctx, &saved)
	require.Error(t, err, "partial repair must not allow executing an invalid dependency")
	_, err = s.Mount(instructionContext(1, "bob", types.TenantRoleContributor), InstructionSkillMountRequest{AgentIDs: []string{agent.ID}, SkillIDs: []string{"removed-skill"}, Mode: "remove"})
	require.Error(t, err, "repair still requires Agent ownership")

	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{agent.ID}, SkillIDs: []string{"removed-skill"}, Mode: "remove"})
	require.NoError(t, err)
	require.NoError(t, s.db.First(&saved, "id = ?", agent.ID).Error)
	prompt, err := s.ResolveAgentInstructions(ctx, &saved)
	require.NoError(t, err)
	require.Empty(t, prompt)
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{agent.ID}, SkillIDs: []string{skill.ID}})
	require.NoError(t, err)
	require.NoError(t, s.db.First(&saved, "id = ?", agent.ID).Error)
	prompt, err = s.ResolveAgentInstructions(ctx, &saved)
	require.NoError(t, err)
	require.Contains(t, prompt, "Use clear terminology.")
	require.Empty(t, saved.Config.SandboxConfigID)
}

func TestInstructionMountRejectsCombinedContextOverflow(t *testing.T) {
	s, ctx := instructionHarness(t)
	first, err := s.Import(ctx, "one.md", "", "", []byte(strings.Repeat("规则。", 3500)))
	require.NoError(t, err)
	second, err := s.Import(ctx, "two.md", "", "", []byte(strings.Repeat("规则。", 3500)))
	require.NoError(t, err)
	require.NoError(t, s.db.Create(&types.CustomAgent{ID: "a", TenantID: 1, CreatedBy: "alice", Name: "A"}).Error)
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{"a"}, SkillIDs: []string{first.ID, second.ID}})
	require.ErrorContains(t, err, "合计正文过长")
	var a types.CustomAgent
	require.NoError(t, s.db.First(&a, "id = ?", "a").Error)
	require.Empty(t, a.Config.InstructionSkillIDs)
}

func TestPrivateInstructionSkillsCannotBeSharedThroughAnotherOwnersAgent(t *testing.T) {
	s, ctx := instructionHarness(t)
	skill, err := s.Import(ctx, "private.md", "", "", []byte("alice private rules"))
	require.NoError(t, err)
	for _, a := range []types.CustomAgent{{ID: "alice-agent", TenantID: 1, CreatedBy: "alice", Name: "A", Visibility: "private", Enabled: true}, {ID: "bob-agent", TenantID: 1, CreatedBy: "bob", Name: "B"}} {
		require.NoError(t, s.db.Create(&a).Error)
	}
	admin := instructionContext(1, "admin", types.TenantRoleAdmin)
	_, err = s.Get(admin, skill.ID)
	require.NoError(t, err, "admin may inspect")
	_, err = s.Mount(admin, InstructionSkillMountRequest{AgentIDs: []string{"bob-agent"}, SkillIDs: []string{skill.ID}})
	require.ErrorContains(t, err, "同一创建者")
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{"alice-agent"}, SkillIDs: []string{skill.ID}})
	require.NoError(t, err)
	var agent types.CustomAgent
	require.NoError(t, s.db.First(&agent, "id = ?", "alice-agent").Error)
	_, err = s.ResolveAgentInstructions(instructionContext(1, "bob", types.TenantRoleContributor), &agent)
	require.Error(t, err)
	_, err = s.ResolveAgentInstructions(admin, &agent)
	require.NoError(t, err)
	_, err = s.Mount(instructionContext(1, "bob", types.TenantRoleContributor), InstructionSkillMountRequest{AgentIDs: []string{"bob-agent"}, SkillIDs: []string{skill.ID}})
	require.Error(t, err)
}

func TestSharedAgentSkillExecutionDoesNotBroadenResourceAccess(t *testing.T) {
	s, alice := instructionHarness(t)
	skill, err := s.Import(alice, "owned.md", "", "", []byte("Only run in the shared Agent."))
	require.NoError(t, err)
	agent := &types.CustomAgent{ID: "shared-agent", TenantID: 1, CreatedBy: "alice", Visibility: "private", Enabled: true, Config: types.CustomAgentConfig{InstructionSkillIDs: []string{skill.ID}}}
	bob := accesspolicy.WithSnapshot(instructionContext(1, "bob", types.TenantRoleContributor), &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"agents.read": true}, Grants: map[string]bool{"agents:shared-agent": true}})
	prompt, err := s.ResolveAgentInstructions(bob, agent)
	require.NoError(t, err)
	require.Contains(t, prompt, "Only run in the shared Agent.")
	_, err = s.Get(bob, skill.ID)
	require.Error(t, err)
	require.False(t, accesspolicy.Readable(bob, "skills", skill.ID, false, false))
	denied := accesspolicy.WithSnapshot(bob, &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"agents.read": true}})
	_, err = s.ResolveAgentInstructions(denied, agent)
	require.Error(t, err, "revoking the Agent must revoke its dependency execution")
	foreign, err := s.Import(instructionContext(1, "carol", types.TenantRoleContributor), "foreign.md", "", "", []byte("Unrelated private rules"))
	require.NoError(t, err)
	agent.Config.InstructionSkillIDs = []string{foreign.ID}
	_, err = s.ResolveAgentInstructions(bob, agent)
	require.Error(t, err, "an Agent must not smuggle another creator's private rules")
}

func TestGlobalInstructionSkillsMountToMultipleOwnersAndVirtualBuiltins(t *testing.T) {
	require.NoError(t, types.LoadBuiltinAgentsConfig(filepath.Join("..", "..", "..", "config")))
	s, ctx := instructionHarness(t)
	skill, err := s.Import(ctx, "shared.md", "", "", []byte("Use consistent terms."))
	require.NoError(t, err)
	require.NoError(t, s.db.Model(&types.TenantSkillCatalogEntity{}).Where("id = ?", skill.ID).Update("visibility", "global").Error)
	for _, agent := range []types.CustomAgent{
		{ID: "alice-agent", TenantID: 1, CreatedBy: "alice", Name: "A"},
		{ID: "bob-agent", TenantID: 1, CreatedBy: "bob", Name: "B"},
	} {
		require.NoError(t, s.db.Create(&agent).Error)
	}
	admin := instructionContext(1, "admin", types.TenantRoleAdmin)
	ids := []string{"alice-agent", "bob-agent", types.BuiltinQuickAnswerID, types.BuiltinSmartReasoningID}
	rows, err := s.Mount(admin, InstructionSkillMountRequest{AgentIDs: ids, SkillIDs: []string{skill.ID}})
	require.NoError(t, err)
	require.Len(t, rows, len(ids))
	for _, id := range ids {
		var agent types.CustomAgent
		require.NoError(t, s.db.Where("tenant_id = ? AND id = ?", 1, id).First(&agent).Error)
		require.Equal(t, []string{skill.ID}, agent.Config.InstructionSkillIDs)
		_, err := s.ResolveAgentInstructions(admin, &agent)
		require.NoError(t, err)
	}
	// A second workspace gets its own built-in override, never the first one's.
	other := instructionContext(2, "admin", types.TenantRoleAdmin)
	otherSkill, err := s.Import(other, "shared.md", "", "", []byte("Other workspace rules."))
	require.NoError(t, err)
	require.NoError(t, s.db.Model(&types.TenantSkillCatalogEntity{}).Where("id = ?", otherSkill.ID).Update("visibility", "global").Error)
	_, err = s.Mount(other, InstructionSkillMountRequest{AgentIDs: []string{types.BuiltinQuickAnswerID}, SkillIDs: []string{otherSkill.ID}})
	require.NoError(t, err)
	var original types.CustomAgent
	require.NoError(t, s.db.Where("tenant_id = ? AND id = ?", 1, types.BuiltinQuickAnswerID).First(&original).Error)
	require.Equal(t, []string{skill.ID}, original.Config.InstructionSkillIDs)
	// Visibility grants use of the skill, not edit permission on another Agent.
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{"bob-agent"}, SkillIDs: []string{skill.ID}})
	require.Error(t, err)
	_, err = s.Mount(ctx, InstructionSkillMountRequest{AgentIDs: []string{types.BuiltinQuickAnswerID}, SkillIDs: []string{skill.ID}})
	require.Error(t, err)
}