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

internal/application/repository/message_memory_policy_test.go

@@ -0,0 +1,89 @@
+package repository
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMemoryAgentPolicyUsesPerTurnIdentityAndRejectsDeletedAgents(t *testing.T) {
+	db, sessions := newOwnerScopeDB(t, t.Name())
+	addMemoryWebOwner(t, db)
+	require.NoError(t, db.AutoMigrate(&types.CustomAgent{}))
+	require.NoError(t, db.Model(&types.Session{}).Where("id = ?", sessions["alice"]).Update("user_id", "alice").Error)
+	disabled := false
+	agents := []*types.CustomAgent{
+		{ID: "allowed", TenantID: 7, Enabled: true},
+		{ID: "disabled", TenantID: 7, Enabled: false},
+		{ID: "opted-out", TenantID: 7, Enabled: true, Config: types.CustomAgentConfig{MemoryEnabled: &disabled}},
+		{ID: "deleted", TenantID: 7, Enabled: true},
+		{ID: "foreign", TenantID: 8, Enabled: true},
+	}
+	for _, agent := range agents {
+		require.NoError(t, db.Create(agent).Error)
+	}
+	require.NoError(t, db.Model(agents[1]).Update("enabled", false).Error)
+	require.NoError(t, db.Delete(agents[3]).Error)
+	base := time.Now().Add(-time.Hour)
+	var inputs []*types.Message
+	for i, agentID := range []string{"allowed", "disabled", "opted-out", "deleted", "missing", "foreign", ""} {
+		input := &types.Message{ID: "user-" + agentID, SessionID: sessions["alice"], Role: "user", RequestID: "request-" + agentID, CreatedAt: base.Add(time.Duration(i) * time.Minute)}
+		answer := &types.Message{ID: "answer-" + agentID, SessionID: input.SessionID, Role: "assistant", RequestID: input.RequestID, AgentID: agentID, IsCompleted: true, CreatedAt: input.CreatedAt.Add(time.Second)}
+		require.NoError(t, db.Create(input).Error)
+		require.NoError(t, db.Create(answer).Error)
+		inputs = append(inputs, input)
+	}
+	repo := NewMessageRepository(db)
+	scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:alice"}
+	allowed, err := repo.MemoryAllowedMessageIDs(context.Background(), scope, sessions["alice"], inputs)
+	require.NoError(t, err)
+	require.Equal(t, map[string]bool{inputs[0].ID: true, inputs[6].ID: true}, allowed)
+	require.NoError(t, db.Model(agents[0]).Update("enabled", false).Error)
+	allowed, err = repo.MemoryAllowedMessageIDs(context.Background(), scope, sessions["alice"], inputs)
+	require.NoError(t, err)
+	require.Equal(t, map[string]bool{inputs[6].ID: true}, allowed)
+	allowed, err = repo.MemoryAllowedMessageIDs(context.Background(), interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:bob"}, sessions["alice"], inputs)
+	require.NoError(t, err)
+	require.Empty(t, allowed)
+}
+
+func TestMemorySharedAgentRechecksExactSourceAndCurrentMembership(t *testing.T) {
+	db, sessions := newOwnerScopeDB(t, t.Name())
+	addMemoryWebOwner(t, db)
+	require.NoError(t, db.AutoMigrate(&types.CustomAgent{}))
+	require.NoError(t, db.Model(&types.Session{}).Where("id = ?", sessions["alice"]).Update("user_id", "alice").Error)
+	for _, query := range []string{
+		`CREATE TABLE agent_shares (id TEXT PRIMARY KEY, agent_id TEXT, organization_id TEXT, source_tenant_id INTEGER, 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 ('org')`,
+		`INSERT INTO organization_tenant_members VALUES ('org', 7)`,
+		`INSERT INTO agent_shares(id, agent_id, organization_id, source_tenant_id) VALUES ('share', 'shared', 'org', 84)`,
+	} {
+		require.NoError(t, db.Exec(query).Error)
+	}
+	require.NoError(t, db.Create(&types.CustomAgent{ID: "shared", TenantID: 84, Enabled: true}).Error)
+	require.NoError(t, db.Create(&types.CustomAgent{ID: "shared", TenantID: 42, Enabled: true}).Error)
+	input := &types.Message{ID: "input", SessionID: sessions["alice"], Role: "user", RequestID: "request", CreatedAt: time.Now().Add(-time.Minute)}
+	answer := &types.Message{ID: "answer", SessionID: input.SessionID, Role: "assistant", RequestID: input.RequestID, AgentID: "shared", AgentTenantID: 84, IsCompleted: true, CreatedAt: time.Now()}
+	require.NoError(t, db.Create(input).Error)
+	require.NoError(t, db.Create(answer).Error)
+	repo := NewMessageRepository(db)
+	scope := interfaces.MemoryScope{TenantID: 7, SubjectID: "web_user:alice"}
+	check := func(want bool) {
+		t.Helper()
+		allowed, err := repo.MemoryAllowedMessageIDs(context.Background(), scope, input.SessionID, []*types.Message{input})
+		require.NoError(t, err)
+		require.Equal(t, want, allowed[input.ID])
+	}
+	check(true)
+	require.NoError(t, db.Model(answer).Update("agent_tenant_id", 42).Error)
+	check(false)
+	require.NoError(t, db.Model(answer).Update("agent_tenant_id", 84).Error)
+	require.NoError(t, db.Exec("DELETE FROM organization_tenant_members").Error)
+	check(false)
+}