xinmi/internal/application/service/session_agent_execution_test.go

package service

import (
	"context"
	"testing"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/stretchr/testify/require"
)

type executionShareStub struct {
	interfaces.AgentShareService
	agent *types.CustomAgent
	err   error
}

func (s *executionShareStub) GetSharedAgentForTenant(_ context.Context, tenant uint64, _ types.TenantRole, id string, source ...uint64) (*types.CustomAgent, error) {
	if tenant != 2 || id != "root" || len(source) != 1 || source[0] != 1 {
		return nil, ErrAgentSharePermission
	}
	return s.agent, s.err
}

func TestSessionAgentExecutionLoadsCurrentBindings(t *testing.T) {
	agents, repo, kbs, ctx := agentExecutionHarness(t)
	old := *repo.rows["root"]
	repo.rows["root"].Config.KnowledgeBases = []string{"nested"}
	s := &sessionService{customAgentService: agents}
	req := &types.QARequest{Session: &types.Session{TenantID: 1}, CustomAgent: &old}
	run, err := s.prepareQAAgentExecution(ctx, req)
	require.NoError(t, err)
	require.Equal(t, []string{"nested"}, req.CustomAgent.Config.KnowledgeBases)
	_, err = kbs.GetKnowledgeBaseByID(run, "nested")
	require.NoError(t, err)
	_, err = kbs.GetKnowledgeBaseByID(run, "owned")
	require.Error(t, err, "old request bindings must not survive a live edit")
}

func TestSessionAgentExecutionRechecksOrganizationShare(t *testing.T) {
	agents, repo, kbs, ctx := agentExecutionHarness(t)
	accesspolicy.FromContext(ctx).Grants = nil
	share := &executionShareStub{agent: repo.rows["root"]}
	s := &sessionService{customAgentService: agents, agentShareService: share}
	req := &types.QARequest{Session: &types.Session{TenantID: 2}, CustomAgent: repo.rows["root"], SharedAgentReadOnly: true}
	run, err := s.prepareQAAgentExecution(ctx, req)
	require.NoError(t, err)
	_, err = kbs.GetKnowledgeBaseByID(run, "owned")
	require.NoError(t, err)
	_, err = kbs.GetKnowledgeBaseByID(ctx, "owned")
	require.Error(t, err)
	share.err = ErrAgentSharePermission
	_, err = s.prepareQAAgentExecution(ctx, req)
	require.Error(t, err)
	share.err = nil
	repo.rows["root"].Enabled = false
	_, err = s.prepareQAAgentExecution(ctx, req)
	require.Error(t, err)
}

func TestSessionQABlocksRevokedAgentBeforeModelInvocation(t *testing.T) {
	for _, mode := range []string{"agent", "knowledge"} {
		t.Run(mode, func(t *testing.T) {
			agents, repo, _, ctx := agentExecutionHarness(t)
			accesspolicy.FromContext(ctx).Grants = nil
			s := &sessionService{customAgentService: agents}
			req := &types.QARequest{Session: &types.Session{ID: "test", TenantID: 1}, CustomAgent: repo.rows["root"]}
			var err error
			if mode == "agent" {
				err = s.AgentQA(ctx, req, nil)
			} else {
				err = s.KnowledgeQA(ctx, req, nil)
			}
			require.Error(t, err)
		})
	}
}