From d24e830c85b98b57218fc49dd2e9b50a16f25c09 Mon Sep 17 00:00:00 2001 From: wizardchen Date: Mon, 29 Jun 2026 17:17:38 +0800 Subject: [PATCH] refactor(agent): fix grep scope loss and consolidate per-turn @mention scoping - Fix grep_chunks dropping the full-KB scope when a turn mixes @KB with @tag/@file: extract scopeClause and OR knowledge-ID and (kb,tenant) predicates together, matching knowledge_search's fan-out. - Consolidate the three scattered per-request @Skill/@MCP scope blocks in buildAgentConfig into applyPerRequestSkillScope / applyPerRequestMCPScope, each narrowing the whitelist and computing the must_use pin in one place. - Dedupe the two identical selection-mode normalizers in Input-field.vue. - Fix knowledge_create_test passing "" for the tagIDs []string parameter, which made the whole service test package fail to compile. - Add unit tests for scopeClause and the per-request scope helpers. --- frontend/src/components/Input-field.vue | 15 +-- internal/agent/tools/grep_chunks.go | 53 +++++--- .../agent/tools/grep_chunks_scope_test.go | 49 +++++++ .../service/knowledge_create_test.go | 8 +- .../application/service/session_agent_qa.go | 123 ++++++++++-------- .../service/session_agent_qa_scope_test.go | 38 ++++++ 6 files changed, 203 insertions(+), 83 deletions(-) create mode 100644 internal/agent/tools/grep_chunks_scope_test.go diff --git a/frontend/src/components/Input-field.vue b/frontend/src/components/Input-field.vue index 20da97cb..757fed98 100644 --- a/frontend/src/components/Input-field.vue +++ b/frontend/src/components/Input-field.vue @@ -313,13 +313,14 @@ const agentAllowedTools = computed(() => { return currentAgentConfig.value?.allowed_tools || []; }); -const normalizeMCPSelectionMode = (mode?: string): 'all' | 'selected' | 'none' => { +type SelectionMode = 'all' | 'selected' | 'none'; +const normalizeSelectionMode = (mode?: string): SelectionMode => { return mode === 'all' || mode === 'selected' || mode === 'none' ? mode : 'none'; }; -const agentMCPSelectionMode = computed<'all' | 'selected' | 'none'>(() => { +const agentMCPSelectionMode = computed(() => { if (!settingsStore.isAgentStreamMode || !hasAgentConfig.value) return 'none'; - return normalizeMCPSelectionMode(currentAgentConfig.value?.mcp_selection_mode); + return normalizeSelectionMode(currentAgentConfig.value?.mcp_selection_mode); }); const agentMCPServiceIds = computed(() => { @@ -335,13 +336,9 @@ const isMCPAllowedByAgent = (service: MCPService) => { return true; }; -const normalizeSkillsSelectionMode = (mode?: string): 'all' | 'selected' | 'none' => { - return mode === 'all' || mode === 'selected' || mode === 'none' ? mode : 'none'; -}; - -const agentSkillsSelectionMode = computed<'all' | 'selected' | 'none'>(() => { +const agentSkillsSelectionMode = computed(() => { if (!settingsStore.isAgentStreamMode || !hasAgentConfig.value) return 'none'; - return normalizeSkillsSelectionMode(currentAgentConfig.value?.skills_selection_mode); + return normalizeSelectionMode(currentAgentConfig.value?.skills_selection_mode); }); const agentSelectedSkills = computed(() => { diff --git a/internal/agent/tools/grep_chunks.go b/internal/agent/tools/grep_chunks.go index b60f4c57..ff492a8f 100644 --- a/internal/agent/tools/grep_chunks.go +++ b/internal/agent/tools/grep_chunks.go @@ -319,6 +319,31 @@ func (t *GrepChunksTool) listKnowledgeIDsByTags( return ids, err } +// scopeClause builds the SQL predicate restricting chunks to the active scope. +// Specific knowledge IDs and full-KB (kb, tenant) pairs are OR'd together so a +// turn that mentions both an entire KB and a tag/file in another KB searches +// every requested scope. Returns ("", nil) when no usable scope exists. +func scopeClause(kbIDs, knowledgeIDs []string, kbTenantMap map[string]uint64) (string, []interface{}) { + var clauses []string + var args []interface{} + if len(knowledgeIDs) > 0 { + clauses = append(clauses, "chunks.knowledge_id IN ?") + args = append(args, knowledgeIDs) + } + for _, kbID := range kbIDs { + tenantID := kbTenantMap[kbID] + if tenantID == 0 { + continue + } + clauses = append(clauses, "(chunks.knowledge_base_id = ? AND chunks.tenant_id = ?)") + args = append(args, kbID, tenantID) + } + if len(clauses) == 0 { + return "", nil + } + return "(" + strings.Join(clauses, " OR ") + ")", args +} + // searchChunks performs the database search using regex queries. func (t *GrepChunksTool) searchChunks( ctx context.Context, @@ -343,26 +368,16 @@ func (t *GrepChunksTool) searchChunks( Where("chunks.deleted_at IS NULL"). Where("knowledges.deleted_at IS NULL") - if len(knowledgeIDs) > 0 { - query = query.Where("chunks.knowledge_id IN ?", knowledgeIDs) - logger.Infof(ctx, "[Tool][GrepChunks] Filtering by %d specific knowledge IDs", len(knowledgeIDs)) - } else if len(kbIDs) > 0 { - var conditions []string - var args []interface{} - for _, kbID := range kbIDs { - tenantID := kbTenantMap[kbID] - if tenantID > 0 { - conditions = append(conditions, "(chunks.knowledge_base_id = ? AND chunks.tenant_id = ?)") - args = append(args, kbID, tenantID) - } - } - if len(conditions) > 0 { - query = query.Where("("+strings.Join(conditions, " OR ")+")", args...) - } else { - logger.Warnf(ctx, "[Tool][GrepChunks] No valid KB-tenant pairs found") - return nil, nil - } + // Combine specific knowledge IDs (incl. tag-resolved docs) and full-KB scopes + // with OR so that mixing @KB with @tag/@file searches BOTH, mirroring + // knowledge_search's concurrent fan-out instead of dropping one side. + scopeSQL, scopeArgs := scopeClause(kbIDs, knowledgeIDs, kbTenantMap) + if scopeSQL == "" { + logger.Warnf(ctx, "[Tool][GrepChunks] No valid scope (kb-tenant pairs / knowledge IDs)") + return nil, nil } + logger.Infof(ctx, "[Tool][GrepChunks] Scope: %d knowledge IDs, %d KBs", len(knowledgeIDs), len(kbIDs)) + query = query.Where(scopeSQL, scopeArgs...) // For MySQL/SQLite REGEXP case-insensitivity we rely on the column's default // collation (utf8mb4_general_ci etc.) OR the driver's REGEXP implementation, diff --git a/internal/agent/tools/grep_chunks_scope_test.go b/internal/agent/tools/grep_chunks_scope_test.go new file mode 100644 index 00000000..c2b04361 --- /dev/null +++ b/internal/agent/tools/grep_chunks_scope_test.go @@ -0,0 +1,49 @@ +package tools + +import ( + "strings" + "testing" +) + +func TestScopeClause_KnowledgeAndFullKBOredTogether(t *testing.T) { + sql, args := scopeClause( + []string{"kb-1"}, + []string{"doc-9"}, + map[string]uint64{"kb-1": 7}, + ) + + if !strings.Contains(sql, "chunks.knowledge_id IN ?") { + t.Fatalf("expected knowledge-id predicate, got %q", sql) + } + if !strings.Contains(sql, "chunks.knowledge_base_id = ? AND chunks.tenant_id = ?") { + t.Fatalf("expected kb-tenant predicate, got %q", sql) + } + if !strings.Contains(sql, " OR ") { + t.Fatalf("expected the two scopes to be OR'd, got %q", sql) + } + // args: knowledgeIDs slice + kbID + tenantID + if len(args) != 3 { + t.Fatalf("expected 3 args, got %d: %v", len(args), args) + } +} + +func TestScopeClause_SkipsKBWithoutTenant(t *testing.T) { + sql, args := scopeClause( + []string{"kb-1", "kb-2"}, + nil, + map[string]uint64{"kb-1": 7}, // kb-2 has no tenant mapping + ) + if strings.Contains(sql, "kb-2") { + t.Fatalf("kb-2 has no tenant and must be skipped, got %q", sql) + } + if len(args) != 2 { + t.Fatalf("expected only kb-1 (2 args), got %d: %v", len(args), args) + } +} + +func TestScopeClause_EmptyWhenNoUsableScope(t *testing.T) { + sql, args := scopeClause([]string{"kb-1"}, nil, map[string]uint64{}) + if sql != "" || args != nil { + t.Fatalf("expected empty scope, got sql=%q args=%v", sql, args) + } +} diff --git a/internal/application/service/knowledge_create_test.go b/internal/application/service/knowledge_create_test.go index 8cc9942b..f5ad1056 100644 --- a/internal/application/service/knowledge_create_test.go +++ b/internal/application/service/knowledge_create_test.go @@ -136,7 +136,7 @@ func TestCreateKnowledgeFromFileDoesNotPersistWhenStorageSaveFails(t *testing.T) nil, nil, "", - "", + nil, "", nil, ) @@ -167,7 +167,7 @@ func TestCreateKnowledgeFromFilePersistsStoredFilePathOnCreate(t *testing.T) { nil, nil, "", - "", + nil, "", nil, ) @@ -201,7 +201,7 @@ func TestCreateKnowledgeFromFileDeletesStoredFileWhenCreateFails(t *testing.T) { nil, nil, "", - "", + nil, "", nil, ) @@ -239,7 +239,7 @@ func TestCreateKnowledgeFromFile_PersistsProcessOverrides(t *testing.T) { map[string]string{"source": "test"}, nil, "", - "", + nil, "", overrides, ) diff --git a/internal/application/service/session_agent_qa.go b/internal/application/service/session_agent_qa.go index f26aaaa6..9ec05a92 100644 --- a/internal/application/service/session_agent_qa.go +++ b/internal/application/service/session_agent_qa.go @@ -249,46 +249,12 @@ func (s *sessionService) buildAgentConfig( } else { agentConfig.AllowedTools = tools.DefaultAllowedTools() } - if len(req.SkillNames) > 0 { - skillsMode := customAgent.Config.SkillsSelectionMode - if skillsMode == "none" || skillsMode == "" { - logger.Warnf(ctx, "Ignoring @skill mention: agent skills selection is disabled (mode=%s)", skillsMode) - } else if agentConfig.SkillsEnabled { - switch skillsMode { - case "selected": - agentConfig.AllowedSkills = intersectPreservingRequestOrder(req.SkillNames, agentConfig.AllowedSkills) - if len(agentConfig.AllowedSkills) == 0 { - agentConfig.SkillsEnabled = false - } - case "all": - agentConfig.AllowedSkills = dedupPreservingOrder(req.SkillNames) - } - logger.Infof(ctx, "Applied per-request @skill scope: requested=%v effective=%v", req.SkillNames, agentConfig.AllowedSkills) - } - } - if len(req.MCPServiceIDs) > 0 { - if agentConfig.MCPSelectionMode == "none" { - logger.Warnf(ctx, "Ignoring @MCP mention: agent MCP selection is disabled (mode=none)") - } else { - mentioned := dedupPreservingOrder(req.MCPServiceIDs) - isSharedAgent := req.Session != nil && req.Session.TenantID != customAgent.TenantID - effectiveMCP, mcpMode := resolvePerRequestMCPScope( - mentioned, - customAgent.Config.MCPServices, - agentConfig.MCPSelectionMode, - isSharedAgent, - ) - if len(effectiveMCP) > 0 { - agentConfig.MCPSelectionMode = mcpMode - agentConfig.MCPServices = effectiveMCP - } else if len(mentioned) > 0 { - logger.Warnf(ctx, "Ignoring @MCP scope outside agent preset: requested=%v agent=%v shared=%v", - req.MCPServiceIDs, customAgent.Config.MCPServices, isSharedAgent) - } - logger.Infof(ctx, "Applied per-request @MCP scope: requested=%v mode=%s effective=%v", - req.MCPServiceIDs, agentConfig.MCPSelectionMode, agentConfig.MCPServices) - } - } + // Apply per-turn @Skill / @MCP scope. Each helper narrows the agent's + // whitelist to the mentioned items and records the pinned set used for the + // hint, keeping all scope logic in one place per resource type. + isSharedAgent := req.Session != nil && req.Session.TenantID != customAgent.TenantID + applyPerRequestSkillScope(ctx, agentConfig, customAgent.Config.SkillsSelectionMode, req.SkillNames) + applyPerRequestMCPScope(ctx, agentConfig, customAgent.Config.MCPServices, isSharedAgent, req.MCPServiceIDs) // Use custom agent's system prompt if specified if customAgent.Config.SystemPrompt != "" { @@ -340,20 +306,75 @@ func (s *sessionService) buildAgentConfig( agentConfig.MaxContextTokens = types.DefaultMaxContextTokens } - if len(agentConfig.MCPServices) > 0 && len(req.MCPServiceIDs) > 0 { - agentConfig.PinnedMCPServiceIDs = intersectPreservingRequestOrder(req.MCPServiceIDs, agentConfig.MCPServices) - } - if len(req.SkillNames) > 0 && agentConfig.SkillsEnabled { - if len(agentConfig.AllowedSkills) > 0 { - agentConfig.PinnedSkillNames = intersectPreservingRequestOrder(req.SkillNames, agentConfig.AllowedSkills) - } else if customAgent.Config.SkillsSelectionMode == "all" { - agentConfig.PinnedSkillNames = dedupPreservingOrder(req.SkillNames) - } - } - return agentConfig, nil } +// applyPerRequestSkillScope narrows the agent's skill whitelist to the @Skill +// mentions for this turn and records the pinned set for the hint. +// It is a no-op when no skills were mentioned or skills are disabled. +func applyPerRequestSkillScope( + ctx context.Context, + agentConfig *types.AgentConfig, + skillsMode string, + requested []string, +) { + if len(requested) == 0 { + return + } + if skillsMode == "none" || skillsMode == "" { + logger.Warnf(ctx, "Ignoring @skill mention: agent skills selection is disabled (mode=%s)", skillsMode) + return + } + if !agentConfig.SkillsEnabled { + return + } + switch skillsMode { + case "selected": + agentConfig.AllowedSkills = intersectPreservingRequestOrder(requested, agentConfig.AllowedSkills) + if len(agentConfig.AllowedSkills) == 0 { + agentConfig.SkillsEnabled = false + } + case "all": + agentConfig.AllowedSkills = dedupPreservingOrder(requested) + } + if agentConfig.SkillsEnabled && len(agentConfig.AllowedSkills) > 0 { + agentConfig.PinnedSkillNames = intersectPreservingRequestOrder(requested, agentConfig.AllowedSkills) + } + logger.Infof(ctx, "Applied per-request @skill scope: requested=%v effective=%v pinned=%v", + requested, agentConfig.AllowedSkills, agentConfig.PinnedSkillNames) +} + +// applyPerRequestMCPScope narrows the agent's MCP services to the @MCP mentions +// for this turn and records the pinned set for the hint. It is a +// no-op when no services were mentioned or MCP selection is disabled. +func applyPerRequestMCPScope( + ctx context.Context, + agentConfig *types.AgentConfig, + agentPresetMCPs []string, + isSharedAgent bool, + requested []string, +) { + if len(requested) == 0 { + return + } + if agentConfig.MCPSelectionMode == "none" { + logger.Warnf(ctx, "Ignoring @MCP mention: agent MCP selection is disabled (mode=none)") + return + } + mentioned := dedupPreservingOrder(requested) + effective, mode := resolvePerRequestMCPScope(mentioned, agentPresetMCPs, agentConfig.MCPSelectionMode, isSharedAgent) + if len(effective) == 0 { + logger.Warnf(ctx, "Ignoring @MCP scope outside agent preset: requested=%v agent=%v shared=%v", + requested, agentPresetMCPs, isSharedAgent) + return + } + agentConfig.MCPSelectionMode = mode + agentConfig.MCPServices = effective + agentConfig.PinnedMCPServiceIDs = intersectPreservingRequestOrder(requested, agentConfig.MCPServices) + logger.Infof(ctx, "Applied per-request @MCP scope: requested=%v mode=%s effective=%v", + requested, agentConfig.MCPSelectionMode, agentConfig.MCPServices) +} + // resolvePerRequestMCPScope narrows MCP registration for a per-turn @mention. // selectionMode "none" rejects all mentions. Shared agents never register MCP // services outside the agent preset. diff --git a/internal/application/service/session_agent_qa_scope_test.go b/internal/application/service/session_agent_qa_scope_test.go index 58505877..5ebb320a 100644 --- a/internal/application/service/session_agent_qa_scope_test.go +++ b/internal/application/service/session_agent_qa_scope_test.go @@ -1,8 +1,10 @@ package service import ( + "context" "testing" + "github.com/Tencent/WeKnora/internal/types" "github.com/stretchr/testify/assert" ) @@ -60,3 +62,39 @@ func TestResolvePerRequestMCPScope_SharedAgentAllowsPreset(t *testing.T) { assert.Equal(t, "selected", mode) assert.Equal(t, []string{"mcp-a"}, effective) } + +func TestApplyPerRequestMCPScope_SelectedNarrowsAndPins(t *testing.T) { + cfg := &types.AgentConfig{MCPSelectionMode: "selected", MCPServices: []string{"mcp-a", "mcp-b"}} + applyPerRequestMCPScope(context.Background(), cfg, []string{"mcp-a", "mcp-b"}, false, []string{"mcp-b"}) + assert.Equal(t, "selected", cfg.MCPSelectionMode) + assert.Equal(t, []string{"mcp-b"}, cfg.MCPServices) + assert.Equal(t, []string{"mcp-b"}, cfg.PinnedMCPServiceIDs) +} + +func TestApplyPerRequestMCPScope_NoneIgnoresMentionAndDoesNotPin(t *testing.T) { + cfg := &types.AgentConfig{MCPSelectionMode: "none", MCPServices: []string{"mcp-a"}} + applyPerRequestMCPScope(context.Background(), cfg, []string{"mcp-a"}, false, []string{"mcp-a"}) + assert.Equal(t, "none", cfg.MCPSelectionMode) + assert.Empty(t, cfg.PinnedMCPServiceIDs) +} + +func TestApplyPerRequestSkillScope_SelectedEmptyIntersectionDisables(t *testing.T) { + cfg := &types.AgentConfig{SkillsEnabled: true, AllowedSkills: []string{"a", "b"}} + applyPerRequestSkillScope(context.Background(), cfg, "selected", []string{"c"}) + assert.False(t, cfg.SkillsEnabled) + assert.Empty(t, cfg.PinnedSkillNames) +} + +func TestApplyPerRequestSkillScope_AllPinsMentioned(t *testing.T) { + cfg := &types.AgentConfig{SkillsEnabled: true} + applyPerRequestSkillScope(context.Background(), cfg, "all", []string{"analysis", "analysis"}) + assert.True(t, cfg.SkillsEnabled) + assert.Equal(t, []string{"analysis"}, cfg.AllowedSkills) + assert.Equal(t, []string{"analysis"}, cfg.PinnedSkillNames) +} + +func TestApplyPerRequestSkillScope_NoneIgnores(t *testing.T) { + cfg := &types.AgentConfig{SkillsEnabled: true, AllowedSkills: []string{"a"}} + applyPerRequestSkillScope(context.Background(), cfg, "none", []string{"a"}) + assert.Empty(t, cfg.PinnedSkillNames) +}