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) }