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") } }