xinmi/internal/application/service/digital_human_agent_access_test.go

package service

import (
	"context"
	"testing"

	"github.com/Tencent/WeKnora/internal/application/repository"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/google/uuid"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

func TestDigitalHumanNativeAgentOwnerIsolation(t *testing.T) {
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	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.DigitalHuman{}, &types.DigitalHumanVersion{}, &types.DigitalHumanBinding{}))
	require.NoError(t, db.Exec("CREATE TABLE im_channels (id TEXT PRIMARY KEY, agent_id TEXT, tenant_id INTEGER, enabled BOOLEAN, updated_at DATETIME, deleted_at DATETIME)").Error)
	dhRepo := repository.NewDigitalHumanRepository(db)
	human := &types.DigitalHuman{ID: uuid.NewString(), TenantID: 7, OwnerID: "owner", Name: "Private Persona", PrimaryAgentID: uuid.NewString(), Status: "draft", LatestVersion: 1, Revision: 1}
	require.NoError(t, dhRepo.Create(context.Background(), human, map[string]string{"tone": "private expression"}))
	scope := types.DigitalHumanScope{TenantID: 7, OwnerID: "owner"}
	require.NoError(t, dhRepo.Review(context.Background(), scope, human.ID, 1, true, ""))
	extra := &types.CustomAgent{ID: uuid.NewString(), TenantID: 7, CreatedBy: "owner", Name: "Additional Agent", Config: types.CustomAgentConfig{SystemPrompt: "base"}}
	require.NoError(t, db.Create(extra).Error)
	require.NoError(t, dhRepo.Bind(context.Background(), scope, human.ID, extra.ID))
	ordinary := &types.CustomAgent{ID: uuid.NewString(), TenantID: 7, CreatedBy: "owner", Name: "Workspace Agent"}
	require.NoError(t, db.Create(ordinary).Error)
	repo := repository.NewCustomAgentRepository(db)
	svc := &customAgentService{repo: repo}
	ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7))
	ctx = context.WithValue(ctx, types.UserIDContextKey, "other")
	ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleContributor)
	for _, id := range []string{human.PrimaryAgentID, extra.ID} {
		_, err := svc.GetAgentByID(ctx, id)
		require.ErrorIs(t, err, ErrAgentNotFound)
		_, err = svc.GetAgentByIDAndTenant(ctx, id, 7)
		require.ErrorIs(t, err, ErrAgentNotFound)
		_, err = svc.CopyAgent(ctx, id)
		require.ErrorIs(t, err, ErrAgentNotFound)
	}
	listed, err := svc.ListAgents(ctx)
	require.NoError(t, err)
	for _, agent := range listed {
		require.NotEqual(t, human.PrimaryAgentID, agent.ID)
		require.NotEqual(t, extra.ID, agent.ID)
	}
	_, err = svc.GetAgentByID(ctx, ordinary.ID)
	require.NoError(t, err)
	ownerCtx := context.WithValue(ctx, types.UserIDContextKey, "owner")
	agent, err := svc.GetAgentByID(ownerCtx, extra.ID)
	require.NoError(t, err)
	require.Contains(t, agent.Config.SystemPrompt, "private expression")
	copy, err := svc.CopyAgent(ownerCtx, extra.ID)
	require.NoError(t, err)
	require.Contains(t, copy.Config.SystemPrompt, "private expression")
	_, err = svc.GetAgentByID(ctx, copy.ID)
	require.ErrorIs(t, err, ErrAgentNotFound)
	_, err = svc.GetAgentByID(ownerCtx, copy.ID)
	require.NoError(t, err)
	require.ErrorIs(t, svc.DeleteAgent(ownerCtx, human.PrimaryAgentID), types.ErrDigitalHumanPrimaryAgent)
	adminCtx := context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleAdmin)
	agent, err = svc.GetAgentByID(adminCtx, human.PrimaryAgentID)
	require.NoError(t, err)
	require.Contains(t, agent.Config.SystemPrompt, "private expression")
	// Verified organization shares use the repository after validating the share.
	shared, err := repo.GetAgentByID(ctx, extra.ID, 7)
	require.NoError(t, err)
	require.Contains(t, shared.Config.SystemPrompt, "private expression")

	shareSvc := &agentShareService{agentRepo: repo}
	_, err = shareSvc.ShareAgent(ctx, extra.ID, "organization", "other", 7, types.OrgRoleViewer)
	require.ErrorIs(t, err, ErrNotAgentOwner)
	for _, statement := range []string{
		`CREATE TABLE agent_shares (id TEXT PRIMARY KEY, agent_id TEXT, organization_id TEXT, shared_by_user_id TEXT, source_tenant_id INTEGER, permission TEXT, created_at DATETIME, updated_at DATETIME, deleted_at DATETIME)`,
		`CREATE TABLE organization_tenant_members (organization_id TEXT, tenant_id INTEGER)`,
		`CREATE TABLE organizations (id TEXT PRIMARY KEY, deleted_at DATETIME)`,
		`INSERT INTO organizations(id) VALUES ('organization')`,
		`INSERT INTO organization_tenant_members(organization_id, tenant_id) VALUES ('organization', 8)`,
	} {
		require.NoError(t, db.Exec(statement).Error)
	}
	require.NoError(t, db.Exec(`INSERT INTO agent_shares(id, agent_id, organization_id, source_tenant_id) VALUES (?, ?, 'organization', 7)`, uuid.NewString(), extra.ID).Error)
	shareSvc.shareRepo = repository.NewAgentShareRepository(db)
	receiverCtx := context.WithValue(ctx, types.TenantIDContextKey, uint64(8))
	shared, err = shareSvc.GetSharedAgentForTenant(receiverCtx, 8, types.TenantRoleContributor, extra.ID, 7)
	require.NoError(t, err)
	require.Contains(t, shared.Config.SystemPrompt, "private expression")
	_, err = shareSvc.GetSharedAgentForTenant(receiverCtx, 9, types.TenantRoleContributor, extra.ID, 7)
	require.ErrorIs(t, err, ErrAgentSharePermission)

	require.NoError(t, db.Exec("INSERT INTO im_channels (id, agent_id, tenant_id, enabled) VALUES ('extra-channel', ?, 7, true)", extra.ID).Error)
	deletedChannels, err := svc.DeleteAgentWithIMChannels(ownerCtx, extra.ID)
	require.NoError(t, err)
	require.Equal(t, []string{"extra-channel"}, deletedChannels)
	var bindings int64
	require.NoError(t, db.Model(&types.DigitalHumanBinding{}).Where("agent_id = ?", extra.ID).Count(&bindings).Error)
	require.Zero(t, bindings)
	_, err = repo.GetAgentByID(ownerCtx, extra.ID, 7)
	require.ErrorIs(t, err, repository.ErrCustomAgentNotFound)
}