87 lines
3.8 KiB
Go
87 lines
3.8 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/accesspolicy"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
func agentExecutionHarness(t *testing.T) (*customAgentService, *lifecycleAgentRepo, *knowledgeBaseService, context.Context) {
|
||
|
|
t.Helper()
|
||
|
|
repo := &lifecycleAgentRepo{rows: map[string]*types.CustomAgent{
|
||
|
|
"root": {ID: "root", TenantID: 1, CreatedBy: "alice", Visibility: "private", Enabled: true, Config: types.CustomAgentConfig{KBSelectionMode: "selected", KnowledgeBases: []string{"owned", "unrelated"}, FusionAgentIDs: []string{"child"}}},
|
||
|
|
"child": {ID: "child", TenantID: 1, CreatedBy: "alice", Visibility: "private", Enabled: true, Config: types.CustomAgentConfig{KBSelectionMode: "selected", KnowledgeBases: []string{"nested"}}},
|
||
|
|
}}
|
||
|
|
kbs := newFakeKBRepo()
|
||
|
|
for _, id := range []string{"owned", "nested", "unbound"} {
|
||
|
|
kbs.rows[id] = &types.KnowledgeBase{ID: id, TenantID: 1, CreatorID: "alice"}
|
||
|
|
}
|
||
|
|
kbs.rows["unrelated"] = &types.KnowledgeBase{ID: "unrelated", TenantID: 1, CreatorID: "carol"}
|
||
|
|
kbService := newPR3KBService(kbs, &fakeRegistry{}, &fakeOwnership{})
|
||
|
|
ctx := accesspolicy.WithSnapshot(instructionContext(1, "bob", types.TenantRoleContributor), &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"agents.read": true, "knowledge.read": true}, Grants: map[string]bool{"agents:root": true}})
|
||
|
|
return &customAgentService{repo: repo, kbService: kbService}, repo, kbService, ctx
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestAgentExecutionDelegatesOnlyLiveOwnedBindings(t *testing.T) {
|
||
|
|
s, _, kbs, ctx := agentExecutionHarness(t)
|
||
|
|
run, err := s.PrepareAgentExecutionContext(ctx, "root")
|
||
|
|
require.NoError(t, err)
|
||
|
|
for _, id := range []string{"owned", "nested"} {
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(run, id)
|
||
|
|
require.NoError(t, err)
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(ctx, id)
|
||
|
|
require.Error(t, err, "execution must not mutate request grants")
|
||
|
|
}
|
||
|
|
for _, id := range []string{"unrelated", "unbound"} {
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(run, id)
|
||
|
|
require.Error(t, err)
|
||
|
|
}
|
||
|
|
_, err = s.GetAgentByID(run, "child")
|
||
|
|
require.NoError(t, err)
|
||
|
|
_, err = s.GetAgentByID(ctx, "child")
|
||
|
|
require.Error(t, err)
|
||
|
|
_, err = s.CopyAgent(run, "child")
|
||
|
|
require.Error(t, err, "runtime dependency must not grant copy or mutation")
|
||
|
|
listed, err := s.ListAgents(run)
|
||
|
|
require.NoError(t, err)
|
||
|
|
for _, a := range listed {
|
||
|
|
require.NotEqual(t, "child", a.ID)
|
||
|
|
}
|
||
|
|
list := filterReadableKnowledgeBases(run, []*types.KnowledgeBase{{ID: "owned", TenantID: 1, CreatorID: "alice"}})
|
||
|
|
require.Empty(t, list, "runtime dependency must not broaden resource list visibility")
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestAgentExecutionRechecksRevocationChangesAndCycles(t *testing.T) {
|
||
|
|
for _, test := range []string{"revoke-root", "disable-child", "different-owner", "cycle", "deny-agents"} {
|
||
|
|
t.Run(test, func(t *testing.T) {
|
||
|
|
s, repo, _, ctx := agentExecutionHarness(t)
|
||
|
|
switch test {
|
||
|
|
case "revoke-root":
|
||
|
|
accesspolicy.FromContext(ctx).Grants = nil
|
||
|
|
case "disable-child":
|
||
|
|
repo.rows["child"].Enabled = false
|
||
|
|
case "different-owner":
|
||
|
|
repo.rows["child"].CreatedBy = "carol"
|
||
|
|
case "cycle":
|
||
|
|
repo.rows["child"].Config.FusionAgentIDs = []string{"root"}
|
||
|
|
case "deny-agents":
|
||
|
|
accesspolicy.FromContext(ctx).Permissions["agents.read"] = false
|
||
|
|
}
|
||
|
|
_, err := s.PrepareAgentExecutionContext(ctx, "root")
|
||
|
|
require.Error(t, err)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
s, repo, kbs, ctx := agentExecutionHarness(t)
|
||
|
|
repo.rows["root"].Config.KnowledgeBases = nil
|
||
|
|
run, err := s.PrepareAgentExecutionContext(ctx, "root")
|
||
|
|
require.NoError(t, err)
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(run, "owned")
|
||
|
|
require.Error(t, err, "removed binding cannot be recreated from saved snapshots")
|
||
|
|
accesspolicy.FromContext(ctx).Permissions["knowledge.read"] = false
|
||
|
|
_, err = kbs.GetKnowledgeBaseByID(run, "nested")
|
||
|
|
require.Error(t, err, "module permission still gates runtime use")
|
||
|
|
}
|