xinmi/internal/workbench/execution_access_test.go

package workbench

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/types/interfaces"
	"github.com/gin-gonic/gin"
)

type boundedExecutionAgents struct {
	interfaces.CustomAgentService
	rows map[string]*types.CustomAgent
}

func (s boundedExecutionAgents) GetAgentByID(ctx context.Context, id string) (*types.CustomAgent, error) {
	a := s.rows[id]
	tenant, user, _, _ := identity(ctx)
	if a == nil || !a.Enabled || a.TenantID != tenant || !accesspolicy.Allowed(ctx, "agents.read") ||
		(!accesspolicy.Readable(ctx, "agents", id, a.CreatedBy == user, false) && !accesspolicy.CanExecuteDependency(ctx, a.TenantID, "agents", id, a.CreatedBy)) {
		return nil, ErrForbidden
	}
	return a, nil
}

func (s boundedExecutionAgents) PrepareAgentExecutionContext(ctx context.Context, id string) (context.Context, error) {
	a, err := s.GetAgentByID(ctx, id)
	if err != nil {
		return ctx, err
	}
	for _, kb := range a.Config.KnowledgeBases {
		ctx = accesspolicy.WithExecutionDependency(ctx, a.TenantID, "knowledge", kb, a.CreatedBy)
	}
	return ctx, nil
}

type boundedExecutionKnowledge struct {
	interfaces.KnowledgeBaseService
	base *types.KnowledgeBase
}

func TestMissingCanvasAgentReportsNodeAsBadRequest(t *testing.T) {
	s := sqliteWorkbench(t)
	flow := seedResource(t, s, 1, "alice", "flows")
	flow.Name = "慈溪"
	flow.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "opening", Label: "开场", AgentID: "deleted"}}})
	if err := s.db.Save(flow).Error; err != nil {
		t.Fatal(err)
	}
	s.agents = boundedExecutionAgents{rows: map[string]*types.CustomAgent{}}
	_, err := s.taskExecutionContext(workbenchContext(1, "alice", types.TenantRoleAdmin), TaskData{FlowID: flow.ID})
	if !errors.Is(err, ErrInvalid) || !strings.Contains(err.Error(), "慈溪") || !strings.Contains(err.Error(), "开场") {
		t.Fatalf("missing actionable canvas dependency error: %v", err)
	}
	response := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(response)
	c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
	respondError(c, err)
	if response.Code != http.StatusBadRequest || !strings.Contains(response.Body.String(), "重新选择") {
		t.Fatalf("missing Agent must return actionable HTTP 400: %d %s", response.Code, response.Body.String())
	}
}

func (s boundedExecutionKnowledge) GetKnowledgeBaseByID(ctx context.Context, id string) (*types.KnowledgeBase, error) {
	if id == s.base.ID && accesspolicy.CanExecuteDependency(ctx, s.base.TenantID, "knowledge", id, s.base.CreatorID) {
		return s.base, nil
	}
	return nil, ErrForbidden
}

func TestSharedCanvasExecutionUsesOnlyCurrentOwnedDependencies(t *testing.T) {
	s := sqliteWorkbench(t)
	flow := seedResource(t, s, 1, "owner", "flows")
	flow.Data = asJSON(FlowData{Nodes: []FlowNode{{ID: "node", AgentID: "private"}}})
	if err := s.db.Model(flow).Update("data", flow.Data).Error; err != nil {
		t.Fatal(err)
	}
	agents := boundedExecutionAgents{rows: map[string]*types.CustomAgent{
		"private": {ID: "private", TenantID: 1, CreatedBy: "owner", Enabled: true, Config: types.CustomAgentConfig{Role: types.AgentRoleProduction, KBSelectionMode: "selected", KnowledgeBases: []string{"private-kb"}}},
		"unbound": {ID: "unbound", TenantID: 1, CreatedBy: "owner", Enabled: true},
	}}
	s.agents = agents
	s.knowledgeBases = boundedExecutionKnowledge{base: &types.KnowledgeBase{ID: "private-kb", TenantID: 1, CreatorID: "owner"}}
	policy := &accesspolicy.Snapshot{Restricted: true, Permissions: map[string]bool{"flows.read": true, "agents.read": true, "knowledge.read": true}, Publications: map[string]bool{"flows:" + flow.ID: true}}
	ctx := accesspolicy.WithSnapshot(workbenchContext(1, "recipient", types.TenantRoleContributor), policy)
	d := TaskData{FlowID: flow.ID}
	execution, err := s.taskExecutionContext(ctx, d)
	if err != nil {
		t.Fatal(err)
	}
	resolved, err := s.taskFlow(execution, d)
	if err != nil {
		t.Fatal(err)
	}
	if err := s.inheritFlow(execution, &d, resolved); err != nil || len(d.KBIDs) != 1 || d.KBIDs[0] != "private-kb" {
		t.Fatalf("shared canvas cannot execute private Agent knowledge: %+v %v", d, err)
	}
	if _, err := agents.GetAgentByID(ctx, "private"); err == nil {
		t.Fatal("private Agent exposed outside execution")
	}
	if _, err := s.knowledgeBases.GetKnowledgeBaseByID(ctx, "private-kb"); err == nil {
		t.Fatal("private knowledge exposed outside execution")
	}
	for _, id := range []string{"unbound"} {
		if _, err := agents.GetAgentByID(execution, id); err == nil {
			t.Fatal("unbound Agent exposed")
		}
	}
	d.FlowSnapshot = &FlowData{Nodes: []FlowNode{{ID: "forged", AgentID: "unbound"}}}
	if _, err := s.taskFlow(execution, d); err == nil {
		t.Fatal("snapshot added an unauthorized private Agent")
	}
	d.FlowSnapshot = nil
	agents.rows["private"].Enabled = false
	if _, err := s.taskExecutionContext(ctx, d); err == nil {
		t.Fatal("disabled Agent accepted")
	}
	agents.rows["private"].Enabled = true
	agents.rows["private"].CreatedBy = "third-party"
	if _, err := s.taskExecutionContext(ctx, d); err == nil {
		t.Fatal("canvas delegated another owner's private Agent")
	}
	agents.rows["private"].CreatedBy = "owner"
	policy.Publications["flows:"+flow.ID] = false
	if _, err := s.taskExecutionContext(ctx, d); err == nil {
		t.Fatal("revoked canvas remained executable")
	}
	policy.Publications["flows:"+flow.ID] = true
	foreign := context.WithValue(ctx, types.TenantIDContextKey, uint64(2))
	if _, err := s.taskExecutionContext(foreign, d); err == nil {
		t.Fatal("canvas executable from another tenant")
	}
	delete(policy.Permissions, "agents.read")
	if _, err := s.taskExecutionContext(ctx, d); err == nil {
		t.Fatal("canvas bypassed Agent module permission")
	}
}