WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/service/digital_human_agent_access_test.go

@@ -0,0 +1,106 @@
+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)
+}