diff --git a/internal/agent/act.go b/internal/agent/act.go index 92222bc82..f8f25e690 100644 --- a/internal/agent/act.go +++ b/internal/agent/act.go @@ -13,6 +13,7 @@ import ( "github.com/Tencent/WeKnora/internal/common" "github.com/Tencent/WeKnora/internal/event" "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/modelcontext" "github.com/Tencent/WeKnora/internal/tracing/langfuse" "github.com/Tencent/WeKnora/internal/types" "golang.org/x/sync/errgroup" @@ -51,6 +52,52 @@ func argKeys(args map[string]any) []string { return keys } +// traceArgumentValue keeps valid model-emitted JSON structured in Langfuse +// while preserving malformed payloads verbatim for diagnosis. +func traceArgumentValue(raw string) interface{} { + var value interface{} + if err := json.Unmarshal([]byte(raw), &value); err != nil { + return raw + } + return value +} + +// buildToolSpanInput exposes both sides of the model-context boundary: +// model_arguments is exactly what the model emitted (including temporary +// handles), while resolved_arguments is the durable payload actually executed. +// Langfuse never performs mapping itself; it only observes the registry's audit. +func buildToolSpanInput(tc types.LLMToolCall, resolvedArgs map[string]any, sensitive bool) map[string]interface{} { + modelArguments := tc.ModelArguments + if modelArguments == "" { + modelArguments = tc.Function.Arguments + } + resolution := tc.ArgumentResolution + if resolution == "" { + resolution = modelcontext.ArgumentResolutionUnchanged + } + if sensitive { + modelArgKeys := []string(nil) + if parsed, ok := traceArgumentValue(modelArguments).(map[string]interface{}); ok { + modelArgKeys = argKeys(parsed) + } + return map[string]interface{}{ + "tool_call_id": tc.ID, + "model_arg_keys": modelArgKeys, + "resolved_arg_keys": argKeys(resolvedArgs), + "argument_resolution": resolution, + "unresolved_handle_count": len(tc.UnresolvedHandles), + "args_redacted": true, + } + } + return map[string]interface{}{ + "tool_call_id": tc.ID, + "model_arguments": traceArgumentValue(modelArguments), + "resolved_arguments": resolvedArgs, + "argument_resolution": resolution, + "unresolved_handles": tc.UnresolvedHandles, + } +} + // finishToolSpan serialises a completed tool call into a Langfuse span // update. Extracted from runToolCall so the tool-call pipeline keeps // a single assignment per line and the observability-specific logic @@ -330,7 +377,29 @@ func (e *AgentEngine) runToolCall( } } logger.Warnf(ctx, "%s Repaired malformed JSON arguments", toolTag) - tc.Function.Arguments = repaired + // The initial model-context pass could not inspect malformed JSON. + // Decode the repaired payload before execution, while preserving the + // exact provider payload already stored in tc.ModelArguments. + decoded := tc + decoded.ModelArguments = "" + decoded.Function.Arguments = repaired + decodedCalls := []types.LLMToolCall{decoded} + e.modelContext.DecodeToolCalls(decodedCalls) + tc.Function.Arguments = decodedCalls[0].Function.Arguments + tc.ArgumentResolution = decodedCalls[0].ArgumentResolution + tc.UnresolvedHandles = decodedCalls[0].UnresolvedHandles + if err := json.Unmarshal([]byte(tc.Function.Arguments), &args); err != nil { + return types.ToolCall{ + ID: tc.ID, + Name: tc.Function.Name, + Args: map[string]any{"_raw": tc.Function.Arguments}, + ProviderMetadata: tc.ProviderMetadata, + Result: &types.ToolResult{ + Success: false, + Error: fmt.Sprintf("Failed to parse repaired tool arguments: %v", err), + }, + } + } } logger.Debugf(ctx, "%s Args: %s", toolTag, tc.Function.Arguments) @@ -365,30 +434,23 @@ func (e *AgentEngine) runToolCall( // any nested generations (embedding/rerank/VLM) that the tool itself // triggers. No-op when Langfuse is disabled. mgr := langfuse.GetManager() - toolSpanInput := map[string]interface{}{ - "arguments": args, - "tool_call_id": tc.ID, - } // database_query's SQL is treated as sensitive by the UI hint layer // (toolHintSensitiveArgs) because it exposes implementation details. // Mirror that policy for Langfuse: redact raw arguments to avoid // leaking raw SQL into the observability backend. - if toolHintSensitiveArgs[tc.Function.Name] { - toolSpanInput = map[string]interface{}{ - "tool_call_id": tc.ID, - "arg_keys": argKeys(args), - "args_redacted": true, - } - } + toolSpanInput := buildToolSpanInput(tc, args, toolHintSensitiveArgs[tc.Function.Name]) + argumentResolution, _ := toolSpanInput["argument_resolution"].(string) toolCtx, toolSpan := mgr.StartSpan(ctx, langfuse.SpanOptions{ Name: "agent.tool." + tc.Function.Name, Input: toolSpanInput, Metadata: map[string]interface{}{ - "iteration": iteration, - "round": round, - "tool_index": i + 1, - "tool_call_id": tc.ID, - "session_id": sessionID, + "iteration": iteration, + "round": round, + "tool_index": i + 1, + "tool_call_id": tc.ID, + "session_id": sessionID, + "argument_resolution": argumentResolution, + "unresolved_handle_count": len(tc.UnresolvedHandles), }, }) @@ -405,12 +467,21 @@ func (e *AgentEngine) runToolCall( ExecTimeout: defaultToolExecTimeout, }) - execCtx, toolCancel := context.WithTimeout(toolExecCtx, defaultToolExecTimeout) - result, err := e.toolRegistry.ExecuteTool( - execCtx, tc.Function.Name, - json.RawMessage(tc.Function.Arguments), - ) - toolCancel() + var result *types.ToolResult + var err error + if len(tc.UnresolvedHandles) > 0 { + // A temporary handle is not an application identity. Fail before tool + // execution so a hallucinated/stale cN/dN/bN/wN/iN/res:// token can + // never reach persistence, an external service, or a routing decision. + err = fmt.Errorf("tool arguments contain unresolved model handles: %v", tc.UnresolvedHandles) + } else { + execCtx, toolCancel := context.WithTimeout(toolExecCtx, defaultToolExecTimeout) + result, err = e.toolRegistry.ExecuteTool( + execCtx, tc.Function.Name, + json.RawMessage(tc.Function.Arguments), + ) + toolCancel() + } duration := time.Since(toolCallStartTime).Milliseconds() toolCall := types.ToolCall{ diff --git a/internal/agent/engine.go b/internal/agent/engine.go index 9877023f2..56188daa1 100644 --- a/internal/agent/engine.go +++ b/internal/agent/engine.go @@ -14,9 +14,8 @@ import ( "github.com/Tencent/WeKnora/internal/common" appconfig "github.com/Tencent/WeKnora/internal/config" "github.com/Tencent/WeKnora/internal/event" - "github.com/Tencent/WeKnora/internal/llmreference" - "github.com/Tencent/WeKnora/internal/llmresource" "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/modelcontext" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/tracing/langfuse" "github.com/Tencent/WeKnora/internal/types" @@ -51,8 +50,7 @@ type AgentEngine struct { memoryConsolidator *agentmemory.Consolidator // Memory consolidator for LLM-powered summarization (optional) lastUsage types.TokenUsage // Token usage from the most recent LLM call lastSentMsgCount int // Number of messages sent in the most recent LLM call - resourceRefs *llmresource.Registry // request-local aliases for durable resource references - sourceRefs *llmreference.Registry // request-local chunk/document/web aliases and citations + modelContext *modelcontext.Registry // single request-local boundary for every model handle } // ImageDescriberFunc generates a text description of an image. @@ -87,8 +85,7 @@ func NewAgentEngine( sessionID: sessionID, systemPromptTemplate: systemPromptTemplate, tokenEstimator: tokenEst, - resourceRefs: llmresource.NewRegistry(), - sourceRefs: llmreference.NewRegistry(config.CitationsEnabled()), + modelContext: modelcontext.NewRegistry(config.CitationsEnabled()), } // Initialize memory consolidator if context window management is configured @@ -125,7 +122,7 @@ func (e *AgentEngine) buildSystemPrompt(ctx context.Context) string { e.systemPromptOptions(ctx), e.systemPromptTemplate, ) - return strings.TrimRight(prompt, " \t\r\n") + llmreference.ProtocolPrompt(e.config.CitationsEnabled()) + return strings.TrimRight(prompt, " \t\r\n") + e.modelContext.ProtocolPrompt() } // NewAgentEngineWithSkills creates a new agent engine with skills support diff --git a/internal/agent/engine_test.go b/internal/agent/engine_test.go index 6bb86958c..20869d034 100644 --- a/internal/agent/engine_test.go +++ b/internal/agent/engine_test.go @@ -2,18 +2,34 @@ package agent import ( "context" + "encoding/json" "fmt" "sync" "testing" + agenttools "github.com/Tencent/WeKnora/internal/agent/tools" "github.com/Tencent/WeKnora/internal/event" - "github.com/Tencent/WeKnora/internal/llmreference" + "github.com/Tencent/WeKnora/internal/modelcontext" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +type countingTool struct { + agenttools.BaseTool + calls int +} + +func newCountingTool(name string) *countingTool { + return &countingTool{BaseTool: agenttools.NewBaseTool(name, "test", json.RawMessage(`{"type":"object"}`))} +} + +func (t *countingTool) Execute(context.Context, json.RawMessage) (*types.ToolResult, error) { + t.calls++ + return &types.ToolResult{Success: true, Output: "executed"}, nil +} + // --------------------------------------------------------------------------- // Mock: chat.Chat // --------------------------------------------------------------------------- @@ -72,11 +88,8 @@ func TestStreamLLMResourceAliasesRoundTrip(t *testing.T) { // TestStreamLLMSummarySlugSurvivesDocumentCompaction is the regression guard for // the mangled `summary/` → `summary/d1` bug. A wiki summary-page slug -// embeds a document's UUID, and sourceRefs.EncodeMessages compacts that same -// UUID into a citation alias (d1) via a blind substring replace. resourceRefs -// must alias the slug to a res:// token BEFORE citation compaction runs; -// otherwise the model is shown a mangled slug it then copies into a wiki tool -// call that 404s. See the EncodeMessages ordering in streamLLMToEventBus. +// embeds a document's UUID. The unified model-context registry owns the +// resource-before-source encoding order so the slug cannot become summary/d1. func TestStreamLLMSummarySlugSurvivesDocumentCompaction(t *testing.T) { const knowledgeID = "07a20bb1-a662-47cf-9929-06fb5d5b5b5e" const summarySlug = "summary/" + knowledgeID @@ -101,7 +114,7 @@ func TestStreamLLMSummarySlugSurvivesDocumentCompaction(t *testing.T) { engine := newTestEngine(t, model) // The document UUID is registered as citation alias d1, exactly as the RAG // context (…) would have registered it upstream. - require.Equal(t, "d1", engine.sourceRefs.RegisterDocument(knowledgeID)) + require.Equal(t, "d1", engine.modelContext.RegisterDocument(knowledgeID)) toolMsg := chat.Message{ Role: "tool", @@ -131,7 +144,7 @@ func TestStreamLLMChunkReferenceExpandsBeforeEmission(t *testing.T) { {ResponseType: types.ResponseTypeAnswer, Content: `c1"/>`, Done: true}, }}}} engine := newTestEngine(t, model) - engine.sourceRefs.RegisterChunk(llmreference.ChunkReference{ + engine.modelContext.RegisterChunk(modelcontext.ChunkReference{ ChunkID: "chunk-1", KnowledgeBaseID: "kb-1", DocumentTitle: "Doc", @@ -141,6 +154,71 @@ func TestStreamLLMChunkReferenceExpandsBeforeEmission(t *testing.T) { require.Equal(t, `answer `, result.Content) } +func TestRunToolCallRejectsUnresolvedHandlesBeforeExecution(t *testing.T) { + engine := newTestEngine(t, &mockChat{}) + engine.toolRegistry = agenttools.NewToolRegistry() + tool := newCountingTool("test_unresolved") + engine.toolRegistry.RegisterTool(tool) + + result := engine.runToolCall( + context.Background(), + types.LLMToolCall{ + ID: "call-1", + Function: types.FunctionCall{ + Name: tool.Name(), + Arguments: `{"knowledge_id":"d99"}`, + }, + ModelArguments: `{"knowledge_id":"d99"}`, + ArgumentResolution: modelcontext.ArgumentResolutionUnresolved, + UnresolvedHandles: []string{"d99"}, + }, + 0, 0, 1, "session", "message", + ) + require.Zero(t, tool.calls) + require.NotNil(t, result.Result) + require.False(t, result.Result.Success) + require.Contains(t, result.Result.Error, "unresolved model handles") +} + +func TestRunToolCallDecodesHandlesAfterJSONRepair(t *testing.T) { + newEngine := func() (*AgentEngine, *countingTool) { + engine := newTestEngine(t, &mockChat{}) + engine.toolRegistry = agenttools.NewToolRegistry() + tool := newCountingTool(agenttools.ToolListKnowledgeChunks) + engine.toolRegistry.RegisterTool(tool) + return engine, tool + } + + unknownEngine, unknownTool := newEngine() + unknown := unknownEngine.runToolCall( + context.Background(), + types.LLMToolCall{ + ID: "call-unknown", + Function: types.FunctionCall{Name: unknownTool.Name(), Arguments: `{"knowledge_id":"d99",}`}, + ModelArguments: `{"knowledge_id":"d99",}`, + }, + 0, 0, 1, "session", "message", + ) + require.Zero(t, unknownTool.calls) + require.False(t, unknown.Result.Success) + require.Contains(t, unknown.Result.Error, "unresolved model handles") + + knownEngine, knownTool := newEngine() + knownEngine.modelContext.RegisterDocument("doc-real") + known := knownEngine.runToolCall( + context.Background(), + types.LLMToolCall{ + ID: "call-known", + Function: types.FunctionCall{Name: knownTool.Name(), Arguments: `{"knowledge_id":"d1",}`}, + ModelArguments: `{"knowledge_id":"d1",}`, + }, + 0, 0, 1, "session", "message", + ) + require.Equal(t, 1, knownTool.calls) + require.True(t, known.Result.Success) + require.Equal(t, "doc-real", known.Args["knowledge_id"]) +} + func (m *mockChat) Chat(_ context.Context, _ []chat.Message, _ *chat.ChatOptions) (*types.ChatResponse, error) { return nil, fmt.Errorf("not implemented") } @@ -501,9 +579,9 @@ func TestStreamFinalAnswerToEventBus_EmitsDoneWhenProviderEndsWithEmptyChunk(t * require.NoError(t, err) require.Len(t, finalAnswerEvents, 2) - assert.Equal(t, "final answer", finalAnswerEvents[0].Content) assert.False(t, finalAnswerEvents[0].Done) - assert.Empty(t, finalAnswerEvents[1].Content) assert.True(t, finalAnswerEvents[1].Done) + assert.Equal(t, "final answer", finalAnswerEvents[0].Content+finalAnswerEvents[1].Content, + "a decoder may hold a short suffix until Done to rule out a split model handle") assert.Equal(t, "final answer", state.FinalAnswer) } diff --git a/internal/agent/finalize.go b/internal/agent/finalize.go index 003f442aa..c5c4b2fa6 100644 --- a/internal/agent/finalize.go +++ b/internal/agent/finalize.go @@ -58,7 +58,7 @@ func (e *AgentEngine) streamFinalAnswerToEventBus( if searchutil.MarkdownImageRegex.MatchString(toolCall.Result.Output) { hasRetrievedImage = true } - modelOutput := e.sourceRefs.ModelOutput(toolCall.Result) + modelOutput := e.modelContext.ModelToolResultForTool(toolCall.Name, toolCall.Result) messages = append(messages, chat.Message{ Role: "user", Content: fmt.Sprintf("Tool %s returned: %s", toolCall.Name, modelOutput), diff --git a/internal/agent/langfuse_test.go b/internal/agent/langfuse_test.go index 793817d05..6e648b2b7 100644 --- a/internal/agent/langfuse_test.go +++ b/internal/agent/langfuse_test.go @@ -8,6 +8,61 @@ import ( "github.com/Tencent/WeKnora/internal/types" ) +func TestBuildToolSpanInputKeepsModelAndResolvedArguments(t *testing.T) { + tc := types.LLMToolCall{ + ID: "call-1", + Function: types.FunctionCall{ + Name: "wiki_read_page", + Arguments: `{"knowledge_base_id":"kb-real","slugs":["summary/uuid"]}`, + }, + ModelArguments: `{"knowledge_base_id":"b1","slugs":["res://0001"]}`, + ArgumentResolution: "resolved", + } + resolved := map[string]any{ + "knowledge_base_id": "kb-real", + "slugs": []interface{}{"summary/uuid"}, + } + got := buildToolSpanInput(tc, resolved, false) + + if !reflect.DeepEqual(got["model_arguments"], map[string]interface{}{ + "knowledge_base_id": "b1", + "slugs": []interface{}{"res://0001"}, + }) { + t.Fatalf("model_arguments = %#v", got["model_arguments"]) + } + if !reflect.DeepEqual(got["resolved_arguments"], resolved) { + t.Fatalf("resolved_arguments = %#v", got["resolved_arguments"]) + } + if got["argument_resolution"] != "resolved" { + t.Fatalf("argument_resolution = %#v", got["argument_resolution"]) + } +} + +func TestBuildToolSpanInputRedactsBothArgumentViews(t *testing.T) { + tc := types.LLMToolCall{ + ID: "call-sensitive", + Function: types.FunctionCall{Arguments: `{"sql":"select secret","knowledge_base_id":"kb-real"}`}, + ModelArguments: `{"sql":"select secret","knowledge_base_id":"b1"}`, + ArgumentResolution: "resolved", + UnresolvedHandles: []string{"d99"}, + } + got := buildToolSpanInput(tc, map[string]any{"sql": "select secret", "knowledge_base_id": "kb-real"}, true) + + if _, ok := got["model_arguments"]; ok { + t.Fatal("sensitive model arguments leaked") + } + if _, ok := got["resolved_arguments"]; ok { + t.Fatal("sensitive resolved arguments leaked") + } + wantKeys := []string{"knowledge_base_id", "sql"} + if !reflect.DeepEqual(got["model_arg_keys"], wantKeys) || !reflect.DeepEqual(got["resolved_arg_keys"], wantKeys) { + t.Fatalf("redacted keys = model:%v resolved:%v", got["model_arg_keys"], got["resolved_arg_keys"]) + } + if got["unresolved_handle_count"] != 1 { + t.Fatalf("unresolved_handle_count = %#v", got["unresolved_handle_count"]) + } +} + // TestTruncateForLangfuse verifies rune-aware truncation. // We specifically check CJK and ASCII inputs to ensure the "…" marker is // appended exactly once without splitting multi-byte characters. diff --git a/internal/agent/observe.go b/internal/agent/observe.go index 695e9defc..45e599d88 100644 --- a/internal/agent/observe.go +++ b/internal/agent/observe.go @@ -11,8 +11,8 @@ import ( agenttools "github.com/Tencent/WeKnora/internal/agent/tools" "github.com/Tencent/WeKnora/internal/common" "github.com/Tencent/WeKnora/internal/event" - "github.com/Tencent/WeKnora/internal/llmreference" "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/modelcontext" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" ) @@ -377,7 +377,7 @@ func commonStringPrefix(a, b string) string { func (e *AgentEngine) RenderUserTurnContent(sessionID, query string) string { e.registerRuntimeReferences() runtimeCtx := buildRuntimeContextBlock(sessionID, e.knowledgeBasesInfo, e.selectedDocs) - runtimeCtx = e.sourceRefs.CompactKnownText(runtimeCtx) + runtimeCtx = e.modelContext.CompactKnownText(runtimeCtx) mustUse := buildMustUseBlock(e.pinnedMCPServices, e.pinnedSkills) return composeUserTurnContent(runtimeCtx, mustUse, query) } @@ -385,22 +385,22 @@ func (e *AgentEngine) RenderUserTurnContent(sessionID, query string) string { // registerRuntimeReferences makes bound KBs, pinned documents and recent // chunks addressable without exposing their durable IDs to the model. func (e *AgentEngine) registerRuntimeReferences() { - if e == nil || e.sourceRefs == nil { + if e == nil || e.modelContext == nil { return } for _, kb := range e.knowledgeBasesInfo { if kb == nil { continue } - e.sourceRefs.RegisterKnowledgeBase(kb.ID) + e.modelContext.RegisterKnowledgeBase(kb.ID) for _, doc := range kb.RecentDocs { - e.sourceRefs.RegisterDocument(doc.KnowledgeID) + e.modelContext.RegisterDocument(doc.KnowledgeID) if doc.ChunkID != "" { title := doc.Title if title == "" { title = doc.FileName } - e.sourceRefs.RegisterChunk(llmreference.ChunkReference{ + e.modelContext.RegisterChunk(modelcontext.ChunkReference{ ChunkID: doc.ChunkID, KnowledgeID: doc.KnowledgeID, KnowledgeBaseID: firstNonEmptyAgent(doc.KnowledgeBaseID, kb.ID), @@ -414,8 +414,8 @@ func (e *AgentEngine) registerRuntimeReferences() { if doc == nil { continue } - e.sourceRefs.RegisterDocument(doc.KnowledgeID) - e.sourceRefs.RegisterKnowledgeBase(doc.KnowledgeBaseID) + e.modelContext.RegisterDocument(doc.KnowledgeID) + e.modelContext.RegisterKnowledgeBase(doc.KnowledgeBaseID) } } @@ -512,7 +512,7 @@ func (e *AgentEngine) appendToolResults( // Add tool result messages (role: "tool", following OpenAI format) for _, toolCall := range step.ToolCalls { - resultContent := e.sourceRefs.ModelOutput(toolCall.Result) + resultContent := e.modelContext.ModelToolResultForTool(toolCall.Name, toolCall.Result) toolMsg := chat.Message{ Role: "tool", diff --git a/internal/agent/observe_test.go b/internal/agent/observe_test.go index 64e57cc29..2ceb2cab0 100644 --- a/internal/agent/observe_test.go +++ b/internal/agent/observe_test.go @@ -8,7 +8,7 @@ import ( "time" agenttools "github.com/Tencent/WeKnora/internal/agent/tools" - "github.com/Tencent/WeKnora/internal/llmreference" + "github.com/Tencent/WeKnora/internal/modelcontext" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" "github.com/stretchr/testify/assert" @@ -274,7 +274,7 @@ func TestRenderUserTurnContent_IncludesScopeBlocks(t *testing.T) { func TestBuildMessagesWithLLMContextRegistersBoundScopeBeforeFirstModelCall(t *testing.T) { engine := &AgentEngine{ - sourceRefs: llmreference.NewRegistry(), + modelContext: modelcontext.NewRegistry(true), knowledgeBasesInfo: []*KnowledgeBaseInfo{{ ID: "kb-real-id", Name: "Docs", @@ -298,7 +298,7 @@ func TestBuildMessagesWithLLMContextRegistersBoundScopeBeforeFirstModelCall(t *t assert.Contains(t, userContent, `knowledge_base id="b1"`) assert.Contains(t, userContent, `knowledge_id="d1"`) assert.Contains(t, userContent, `knowledge_id="d2"`) - assert.Equal(t, "c1", engine.sourceRefs.ChunkAlias("chunk-real-id")) + assert.Equal(t, "c1", engine.modelContext.ChunkHandle("chunk-real-id")) assert.NotContains(t, userContent, "kb-real-id") assert.NotContains(t, userContent, "chunk-real-id") assert.NotContains(t, userContent, "doc-real-id") diff --git a/internal/agent/prompts_wiki.go b/internal/agent/prompts_wiki.go index a57d6b79f..3827689ab 100644 --- a/internal/agent/prompts_wiki.go +++ b/internal/agent/prompts_wiki.go @@ -321,7 +321,7 @@ Now apply the instructions above to the chunks and output ONLY the JSON.` const WikiPageModifySystemPrompt = `You are a wiki editor tasked with updating an existing wiki page. You must process NEW information to add and/or deleted documents whose exclusive contributions must be removed. ### SOURCE GROUNDING & MERGE RULES (CRITICAL): -1. **No Inline Chunk IDs:** Chunk aliases such as [c003] are internal processing metadata. NEVER output them in the page body or summary, and remove any legacy inline chunk aliases from existing content while editing. Source associations are stored separately by the system. +1. **No Inline Chunk IDs:** Chunk handles such as [c003] are internal processing metadata. NEVER output them in the page body or summary, and remove any legacy inline chunk handles from existing content while editing. Source associations are stored separately by the system. 2. **Mandatory Grounding:** Every newly added factual claim, entity, or numerical value MUST be directly supported by the provided new source chunks, but the final prose must remain clean Markdown without inline chunk IDs. 3. **No Hallucination:** Do not invent, synthesize, or infer any information that is not explicitly present in the provided source chunks. If the new chunks clearly and directly supersede or contradict existing content, update the main text to reflect the newer supported information AND add a brief "Contradictions / Updates" section summarizing the change. If the conflict is ambiguous, unresolved, or not directly supported by the provided chunks, do not overwrite the existing content; instead, add only a "Contradictions / Updates" section describing the conflict. 4. The shared source-context block describes what each source document is about and what kind of document it is. Use it only to calibrate scope, attribution, and tone. Never copy source-context wording into the page as factual evidence. diff --git a/internal/agent/prompts_wiki_test.go b/internal/agent/prompts_wiki_test.go index c85c03a89..5214492e1 100644 --- a/internal/agent/prompts_wiki_test.go +++ b/internal/agent/prompts_wiki_test.go @@ -124,7 +124,7 @@ func TestWikiChunkCitationPrompt_PreservesPlaceholders(t *testing.T) { } } -func TestWikiPageModifyUserPrompt_HidesInternalChunkAliases(t *testing.T) { +func TestWikiPageModifyUserPrompt_HidesInternalChunkHandles(t *testing.T) { combined := WikiPageModifySystemPrompt + "\n" + WikiPageModifyUserPrompt for _, guidance := range []string{ "NEVER output them in the page body or summary", @@ -132,7 +132,7 @@ func TestWikiPageModifyUserPrompt_HidesInternalChunkAliases(t *testing.T) { "clean Markdown without inline chunk IDs", } { if !strings.Contains(combined, guidance) { - t.Errorf("WikiPageModifyUserPrompt missing chunk-alias guidance %q", guidance) + t.Errorf("WikiPageModifyUserPrompt missing chunk-handle guidance %q", guidance) } } diff --git a/internal/agent/think.go b/internal/agent/think.go index ccac69dc3..d22d7f3db 100644 --- a/internal/agent/think.go +++ b/internal/agent/think.go @@ -9,8 +9,6 @@ import ( agenttools "github.com/Tencent/WeKnora/internal/agent/tools" "github.com/Tencent/WeKnora/internal/common" "github.com/Tencent/WeKnora/internal/event" - "github.com/Tencent/WeKnora/internal/llmreference" - "github.com/Tencent/WeKnora/internal/llmresource" "github.com/Tencent/WeKnora/internal/logger" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" @@ -39,17 +37,8 @@ func (e *AgentEngine) streamLLMToEventBus( llmCtx, llmCancel := context.WithTimeout(ctx, e.getLLMCallTimeout()) defer llmCancel() - // Order matters. resourceRefs must alias durable resource handles FIRST, - // because wiki summary-page slugs (summary/) embed a document's - // UUID and sourceRefs.EncodeMessages compacts that same UUID into a document - // citation alias (d1, d2, …) via a blind substring replace. Running - // sourceRefs first mangled `[[summary/|Title]]` into `[[summary/d1]]` - // — an unrecoverable dead link the model then copied into wiki tool calls. - // Encoding resourceRefs first turns the slug into a res:// token so the UUID - // is gone before citation compaction runs; the two alias spaces (res://NNNN - // vs d/c/b/w-N) are disjoint, so decode order is unaffected. - messages = e.resourceRefs.EncodeMessages(messages) - messages = e.sourceRefs.EncodeMessages(messages) + // Model-context encoding owns codec ordering and temporary-handle lifecycle. + messages = e.modelContext.EncodeMessages(messages) prefixFingerprint := chat.PromptPrefixFingerprint(messages, opts) llmCtx = types.WithLLMCallMetadata(llmCtx, "agent_round", prefixFingerprint) stream, err := e.chatModel.ChatStream(llmCtx, messages, opts) @@ -62,10 +51,8 @@ func (e *AgentEngine) streamLLMToEventBus( chunkCount := 0 responseTypeCounts := make(map[string]int) firstChunkTime := time.Time{} - answerDecoder := llmresource.NewStreamDecoder(e.resourceRefs) - thinkingDecoder := llmresource.NewStreamDecoder(e.resourceRefs) - answerRefExpander := llmreference.NewStreamExpander(e.sourceRefs) - thinkingRefExpander := llmreference.NewStreamExpander(e.sourceRefs) + answerDecoder := e.modelContext.StreamDecoder() + thinkingDecoder := e.modelContext.StreamDecoder() for chunk := range stream { chunkCount++ @@ -82,12 +69,17 @@ func (e *AgentEngine) streamLLMToEventBus( continue } if chunk.ResponseType == types.ResponseTypeThinking { - chunk.Content = thinkingRefExpander.Feed(thinkingDecoder.Feed(chunk.Content)) + chunk.Content = thinkingDecoder.Feed(chunk.Content) + if chunk.Done { + chunk.Content += thinkingDecoder.Flush() + } } else { - chunk.Content = answerRefExpander.Feed(answerDecoder.Feed(chunk.Content)) + chunk.Content = answerDecoder.Feed(chunk.Content) + if chunk.Done { + chunk.Content += answerDecoder.Flush() + } } - e.resourceRefs.DecodeToolCalls(chunk.ToolCalls) - e.sourceRefs.DecodeToolCalls(chunk.ToolCalls) + e.modelContext.DecodeToolCalls(chunk.ToolCalls) if chunk.Content != "" { isExtracted := chunk.Data != nil && chunk.Data["source"] != nil @@ -116,17 +108,33 @@ func (e *AgentEngine) streamLLMToEventBus( emitFunc(&chunk, result.Content) } } - result.Content += answerRefExpander.Feed(answerDecoder.Flush()) - result.Content += answerRefExpander.Flush() - result.ReasoningContent += thinkingRefExpander.Feed(thinkingDecoder.Flush()) - result.ReasoningContent += thinkingRefExpander.Flush() + answerTail := answerDecoder.Flush() + thinkingTail := thinkingDecoder.Flush() + result.Content += answerTail + result.ReasoningContent += thinkingTail + if emitFunc != nil { + if thinkingTail != "" { + emitFunc(&types.StreamResponse{ResponseType: types.ResponseTypeThinking, Content: thinkingTail}, result.Content) + } + if answerTail != "" { + emitFunc(&types.StreamResponse{ResponseType: types.ResponseTypeAnswer, Content: answerTail}, result.Content) + } + } // Some providers stream tool-call argument fragments but expose the // accumulated call in the final chunk. Decode once more after assembly so - // aliases split across provider chunks cannot leak into tool execution. - e.resourceRefs.DecodeToolCalls(result.ToolCalls) - e.sourceRefs.DecodeToolCalls(result.ToolCalls) - if orphans := e.resourceRefs.OrphanAliases(result.Content); len(orphans) > 0 { - logger.Warnf(ctx, "[Agent][Stream] Model emitted %d unresolvable resource alias(es): %v", + // handles split across provider chunks cannot leak into tool execution. + e.modelContext.DecodeToolCalls(result.ToolCalls) + for _, toolCall := range result.ToolCalls { + if len(toolCall.UnresolvedHandles) == 0 { + continue + } + logger.Warnf(ctx, + "[Agent][Stream] Tool %s (%s) contains unresolvable model handle(s): %v", + toolCall.Function.Name, toolCall.ID, toolCall.UnresolvedHandles, + ) + } + if orphans := e.modelContext.OrphanResourceHandles(result.Content); len(orphans) > 0 { + logger.Warnf(ctx, "[Agent][Stream] Model emitted %d unresolvable resource handle(s): %v", len(orphans), orphans) } diff --git a/internal/agent/tools/data_analysis.go b/internal/agent/tools/data_analysis.go index cda648838..237f500d7 100644 --- a/internal/agent/tools/data_analysis.go +++ b/internal/agent/tools/data_analysis.go @@ -124,6 +124,14 @@ type DataAnalysisTool struct { // after startup, causing GetFile to look in the wrong directory (#1040). localBaseDir string storageResolver interfaces.StorageBackendResolver + searchTargets types.SearchTargets +} + +// WithSearchTargets enables the Agent-only authorization boundary. Other +// internal data-analysis callers retain their existing service-owned scope. +func (t *DataAnalysisTool) WithSearchTargets(searchTargets types.SearchTargets) *DataAnalysisTool { + t.searchTargets = searchTargets + return t } func NewDataAnalysisTool( @@ -201,6 +209,11 @@ func (t *DataAnalysisTool) Execute(ctx context.Context, args json.RawMessage) (* Error: fmt.Sprintf("Failed to parse input args: %v", err), }, err } + if len(t.searchTargets) > 0 { + if _, err := authorizeKnowledgeInSearchTargets(ctx, t.searchTargets, input.KnowledgeID, t.knowledgeService); err != nil { + return &types.ToolResult{Success: false, Error: err.Error()}, err + } + } schema, err := t.LoadFromKnowledgeID(ctx, input.KnowledgeID) if err != nil { @@ -659,7 +672,10 @@ func (t *DataAnalysisTool) materializeKnowledgeFile(ctx context.Context, knowled func (t *DataAnalysisTool) LoadFromKnowledgeID(ctx context.Context, knowledgeID string) (*TableSchema, error) { // Use GetKnowledgeByIDOnly to support cross-tenant shared KB knowledge, err := t.knowledgeService.GetKnowledgeByIDOnly(ctx, knowledgeID) - if err != nil { + if err != nil || knowledge == nil { + if err == nil { + err = fmt.Errorf("knowledge service returned an empty result") + } logger.Errorf(ctx, "[Tool][DataAnalysis] Failed to get knowledge by ID '%s': %v", knowledgeID, err) return nil, fmt.Errorf("failed to get knowledge by ID: %w", err) } diff --git a/internal/agent/tools/data_schema.go b/internal/agent/tools/data_schema.go index ef2f89ebf..ffb219006 100644 --- a/internal/agent/tools/data_schema.go +++ b/internal/agent/tools/data_schema.go @@ -25,6 +25,13 @@ type DataSchemaTool struct { knowledgeService interfaces.KnowledgeService chunkRepo interfaces.ChunkRepository targetChunkTypes []types.ChunkType + searchTargets types.SearchTargets +} + +// WithSearchTargets enables Agent request-scope authorization. +func (t *DataSchemaTool) WithSearchTargets(searchTargets types.SearchTargets) *DataSchemaTool { + t.searchTargets = searchTargets + return t } func NewDataSchemaTool(knowledgeService interfaces.KnowledgeService, chunkRepo interfaces.ChunkRepository, targetChunkTypes ...types.ChunkType) *DataSchemaTool { @@ -50,8 +57,17 @@ func (t *DataSchemaTool) Execute(ctx context.Context, args json.RawMessage) (*ty } // Get knowledge to get TenantID (use IDOnly to support cross-tenant shared KB) - knowledge, err := t.knowledgeService.GetKnowledgeByIDOnly(ctx, input.KnowledgeID) - if err != nil { + var knowledge *types.Knowledge + var err error + if len(t.searchTargets) > 0 { + knowledge, err = authorizeKnowledgeInSearchTargets(ctx, t.searchTargets, input.KnowledgeID, t.knowledgeService) + } else { + knowledge, err = t.knowledgeService.GetKnowledgeByIDOnly(ctx, input.KnowledgeID) + } + if err != nil || knowledge == nil { + if err == nil { + err = fmt.Errorf("knowledge service returned an empty result") + } return &types.ToolResult{ Success: false, Error: fmt.Sprintf("Failed to get knowledge '%s': %v", input.KnowledgeID, err), @@ -73,10 +89,10 @@ func (t *DataSchemaTool) Execute(ctx context.Context, args json.RawMessage) (*ty page, chunkTypes, nil, // tagIDs - "", // keyword - "", // searchField - "", // sortOrder - "", // knowledgeType + "", // keyword + "", // searchField + "", // sortOrder + "", // knowledgeType ) if err != nil { return &types.ToolResult{ diff --git a/internal/agent/tools/database_query.go b/internal/agent/tools/database_query.go index b3aaa972e..f364daf7b 100644 --- a/internal/agent/tools/database_query.go +++ b/internal/agent/tools/database_query.go @@ -72,11 +72,6 @@ Count documents by status: "sql": "SELECT parse_status, COUNT(*) as count FROM knowledges GROUP BY parse_status" } -Find recent sessions: -{ - "sql": "SELECT id, title, created_at FROM sessions ORDER BY created_at DESC LIMIT 5" -} - Get storage usage: { "sql": "SELECT SUM(storage_size) as total_storage FROM knowledges" @@ -255,13 +250,17 @@ func (t *DatabaseQueryTool) Execute(ctx context.Context, args json.RawMessage) ( // validateAndSecureSQL validates the SQL query and injects tenant_id conditions func (t *DatabaseQueryTool) validateAndSecureSQL(sqlQuery string, tenantID uint64) (string, error) { + searchScopes := searchScopesFromTargets(t.searchTargets) + if len(searchScopes) == 0 { + return "", fmt.Errorf("no effective Agent knowledge scope is available") + } securedSQL, validationResult, err := utils.ValidateAndSecureSQL( sqlQuery, utils.WithSecurityDefaults(tenantID), utils.WithSoftDeleteFilter("knowledge_bases", "knowledges", "chunks"), utils.WithHiddenKBFilter(), utils.WithInjectionRiskCheck(), - utils.WithSearchScopes(searchScopesFromTargets(t.searchTargets)), + utils.WithSearchScopes(searchScopes), ) if err != nil { return "", err @@ -284,10 +283,16 @@ func searchScopesFromTargets(searchTargets types.SearchTargets) []utils.SearchSc if target == nil || target.KnowledgeBaseID == "" { continue } + tagIDs := effectiveSearchTargetTagIDs(target) + wholeKB := target.Type == types.SearchTargetTypeKnowledgeBase && + len(target.KnowledgeIDs) == 0 && len(tagIDs) == 0 + if !wholeKB && len(target.KnowledgeIDs) == 0 && len(tagIDs) == 0 { + continue + } scopes = append(scopes, utils.SearchScope{ KnowledgeBaseID: target.KnowledgeBaseID, KnowledgeIDs: append([]string(nil), target.KnowledgeIDs...), - TagIDs: append([]string(nil), target.TagIDs...), + TagIDs: tagIDs, }) } return scopes diff --git a/internal/agent/tools/get_document_info.go b/internal/agent/tools/get_document_info.go index dfaa4a058..1a7974dfd 100644 --- a/internal/agent/tools/get_document_info.go +++ b/internal/agent/tools/get_document_info.go @@ -131,27 +131,12 @@ func (t *GetDocumentInfoTool) Execute(ctx context.Context, args json.RawMessage) wg.Add(1) go func(id string) { defer wg.Done() - chunk, err := t.chunkService.GetChunkByIDOnly(ctx, id) - if err != nil || chunk == nil { + chunk, err := authorizeChunkInSearchTargets( + ctx, t.searchTargets, id, t.chunkService, t.knowledgeService, + ) + if err != nil { mu.Lock() - results["faq:"+id] = &docInfo{err: fmt.Errorf("faq entry not found: %v", err)} - mu.Unlock() - return - } - if !t.searchTargets.ContainsKB(chunk.KnowledgeBaseID) { - mu.Lock() - results["faq:"+id] = &docInfo{err: fmt.Errorf("knowledge base %s is not accessible", chunk.KnowledgeBaseID)} - mu.Unlock() - return - } - allowed, scopeErr := searchTargetsAllowKnowledgeID(ctx, t.searchTargets, chunk.KnowledgeID, chunk.KnowledgeBaseID, t.knowledgeService) - if scopeErr != nil || !allowed { - mu.Lock() - if scopeErr != nil { - results["faq:"+id] = &docInfo{err: fmt.Errorf("failed to validate FAQ scope: %v", scopeErr)} - } else { - results["faq:"+id] = &docInfo{err: fmt.Errorf("FAQ entry %s is not within the current @mention scope", id)} - } + results["faq:"+id] = &docInfo{err: fmt.Errorf("FAQ entry is not accessible: %v", err)} mu.Unlock() return } @@ -170,8 +155,9 @@ func (t *GetDocumentInfoTool) Execute(ctx context.Context, args json.RawMessage) go func(id string) { defer wg.Done() - // Get knowledge metadata without tenant filter to support shared KB - knowledge, err := t.knowledgeService.GetKnowledgeByIDOnly(ctx, id) + knowledge, err := authorizeKnowledgeInSearchTargets( + ctx, t.searchTargets, id, t.knowledgeService, + ) if err != nil { mu.Lock() results[id] = &docInfo{ @@ -181,27 +167,6 @@ func (t *GetDocumentInfoTool) Execute(ctx context.Context, args json.RawMessage) return } - // Verify the knowledge's KB is in searchTargets (permission check) - if !t.searchTargets.ContainsKB(knowledge.KnowledgeBaseID) { - mu.Lock() - results[id] = &docInfo{ - err: fmt.Errorf("knowledge base %s is not accessible", knowledge.KnowledgeBaseID), - } - mu.Unlock() - return - } - allowed, scopeErr := searchTargetsAllowKnowledgeID(ctx, t.searchTargets, knowledge.ID, knowledge.KnowledgeBaseID, t.knowledgeService) - if scopeErr != nil || !allowed { - mu.Lock() - if scopeErr != nil { - results[id] = &docInfo{err: fmt.Errorf("failed to validate document scope: %v", scopeErr)} - } else { - results[id] = &docInfo{err: fmt.Errorf("document %s is not within the current @mention scope", knowledge.ID)} - } - mu.Unlock() - return - } - // Use knowledge's actual tenant_id for chunk query (supports cross-tenant shared KB). // Keep chunk-type filter aligned with list_knowledge_chunks so the // "chunk_count" reported here matches what that tool can page over. diff --git a/internal/agent/tools/grep_chunks.go b/internal/agent/tools/grep_chunks.go index c53371dfe..1449bdd8c 100644 --- a/internal/agent/tools/grep_chunks.go +++ b/internal/agent/tools/grep_chunks.go @@ -264,19 +264,12 @@ func (t *GrepChunksTool) resolveGrepScope() (fullKBIDs, knowledgeIDs []string, t continue } switch { - case len(target.KnowledgeIDs) > 0: - for _, kid := range target.KnowledgeIDs { - if kid != "" && !seenKnowledge[kid] { - seenKnowledge[kid] = true - knowledgeIDs = append(knowledgeIDs, kid) - } - } - case len(target.TagIDs) > 0: + case len(effectiveSearchTargetTagIDs(target)) > 0: tenantID := target.TenantID if tenantID == 0 { tenantID = t.searchTargets.GetTenantIDForKB(target.KnowledgeBaseID) } - tagIDs := dedupNonEmptyStrings(target.TagIDs) + tagIDs := effectiveSearchTargetTagIDs(target) if len(tagIDs) == 0 || tenantID == 0 { continue } @@ -291,6 +284,13 @@ func (t *GrepChunksTool) resolveGrepScope() (fullKBIDs, knowledgeIDs []string, t TenantID: tenantID, TagIDs: tagIDs, }) + case len(target.KnowledgeIDs) > 0: + for _, kid := range target.KnowledgeIDs { + if kid != "" && !seenKnowledge[kid] { + seenKnowledge[kid] = true + knowledgeIDs = append(knowledgeIDs, kid) + } + } default: if !seenKB[target.KnowledgeBaseID] { seenKB[target.KnowledgeBaseID] = true diff --git a/internal/agent/tools/knowledge_search.go b/internal/agent/tools/knowledge_search.go index 76deeacd9..8423f7b26 100644 --- a/internal/agent/tools/knowledge_search.go +++ b/internal/agent/tools/knowledge_search.go @@ -180,6 +180,9 @@ func (t *KnowledgeSearchTool) Execute(ctx context.Context, args json.RawMessage) var userSpecifiedKBs []string if len(input.KnowledgeBaseIDs) > 0 { userSpecifiedKBs = input.KnowledgeBaseIDs + if err := validateKnowledgeBaseIDsInSearchTargets(t.searchTargets, userSpecifiedKBs); err != nil { + return &types.ToolResult{Success: false, Error: err.Error()}, err + } logger.Infof(ctx, "[Tool][KnowledgeSearch] User specified %d knowledge bases: %v", len(userSpecifiedKBs), userSpecifiedKBs) } @@ -193,6 +196,9 @@ func (t *KnowledgeSearchTool) Execute(ctx context.Context, args json.RawMessage) } var filteredTargets types.SearchTargets for _, target := range t.searchTargets { + if target == nil { + continue + } if userKBSet[target.KnowledgeBaseID] { filteredTargets = append(filteredTargets, target) } @@ -460,6 +466,9 @@ func (t *KnowledgeSearchTool) concurrentSearchByTargets( } filteredTargets := make(types.SearchTargets, 0, len(searchTargets)) for _, st := range searchTargets { + if st == nil || st.KnowledgeBaseID == "" { + continue + } if searchableKBs[st.KnowledgeBaseID] { filteredTargets = append(filteredTargets, st) continue @@ -482,6 +491,9 @@ func (t *KnowledgeSearchTool) concurrentSearchByTargets( groups := make(map[string][]*types.SearchTarget) for _, st := range searchTargets { + if st == nil || st.KnowledgeBaseID == "" { + continue + } key := modelKeyMap[st.KnowledgeBaseID] groups[key] = append(groups[key], st) } @@ -1124,7 +1136,7 @@ func (t *KnowledgeSearchTool) formatOutput( // Count results by KB kbCounts := make(map[string]int) for _, r := range results { - kbCounts[r.KnowledgeID]++ + kbCounts[r.KnowledgeBaseID]++ } // Format individual results as XML. Tag names are kept in sync with diff --git a/internal/agent/tools/list_knowledge_chunks.go b/internal/agent/tools/list_knowledge_chunks.go index 5ec6bea11..ca2fd0c98 100644 --- a/internal/agent/tools/list_knowledge_chunks.go +++ b/internal/agent/tools/list_knowledge_chunks.go @@ -117,36 +117,14 @@ func (t *ListKnowledgeChunksTool) Execute(ctx context.Context, args json.RawMess }, fmt.Errorf("missing id parameter") } - // Get knowledge info without tenant filter to support shared KB - knowledge, err := t.knowledgeService.GetKnowledgeByIDOnly(ctx, knowledgeID) + knowledge, err := authorizeKnowledgeInSearchTargets(ctx, t.searchTargets, knowledgeID, t.knowledgeService) if err != nil { return &types.ToolResult{ Success: false, - Error: fmt.Sprintf("Knowledge not found: %v", err), + Error: fmt.Sprintf("Knowledge is not accessible: %v", err), }, err } - // Verify the knowledge's KB is in searchTargets (permission check) - if !t.searchTargets.ContainsKB(knowledge.KnowledgeBaseID) { - return &types.ToolResult{ - Success: false, - Error: fmt.Sprintf("Knowledge base %s is not accessible", knowledge.KnowledgeBaseID), - }, fmt.Errorf("knowledge base not in search targets") - } - allowed, err := searchTargetsAllowKnowledgeID(ctx, t.searchTargets, knowledge.ID, knowledge.KnowledgeBaseID, t.knowledgeService) - if err != nil { - return &types.ToolResult{ - Success: false, - Error: fmt.Sprintf("failed to validate knowledge scope: %v", err), - }, err - } - if !allowed { - return &types.ToolResult{ - Success: false, - Error: fmt.Sprintf("Knowledge %s is not within the current @mention scope", knowledge.ID), - }, fmt.Errorf("knowledge not in search target scope") - } - // Use the knowledge's actual tenant_id for chunk query (supports cross-tenant shared KB) effectiveTenantID := knowledge.TenantID @@ -296,32 +274,15 @@ func (t *ListKnowledgeChunksTool) Execute(ctx context.Context, args json.RawMess // executeByChunkID loads one chunk by faq_id / chunk_id (FAQ entry or any chunk). func (t *ListKnowledgeChunksTool) executeByChunkID(ctx context.Context, chunkID string) (*types.ToolResult, error) { - chunk, err := t.chunkService.GetChunkByIDOnly(ctx, chunkID) - if err != nil || chunk == nil { + chunk, err := authorizeChunkInSearchTargets( + ctx, t.searchTargets, chunkID, t.chunkService, t.knowledgeService, + ) + if err != nil { return &types.ToolResult{ Success: false, - Error: fmt.Sprintf("chunk not found: %v", err), + Error: fmt.Sprintf("chunk is not accessible: %v", err), }, err } - if !t.searchTargets.ContainsKB(chunk.KnowledgeBaseID) { - return &types.ToolResult{ - Success: false, - Error: fmt.Sprintf("knowledge base %s is not accessible", chunk.KnowledgeBaseID), - }, fmt.Errorf("knowledge base not in search targets") - } - allowed, scopeErr := searchTargetsAllowKnowledgeID(ctx, t.searchTargets, chunk.KnowledgeID, chunk.KnowledgeBaseID, t.knowledgeService) - if scopeErr != nil { - return &types.ToolResult{ - Success: false, - Error: fmt.Sprintf("failed to validate chunk scope: %v", scopeErr), - }, scopeErr - } - if !allowed { - return &types.ToolResult{ - Success: false, - Error: fmt.Sprintf("chunk %s is not within the current @mention scope", chunk.ID), - }, fmt.Errorf("chunk not in search target scope") - } chunks := []*types.Chunk{chunk} if chunk.ImageInfo == "" { diff --git a/internal/agent/tools/query_knowledge_graph.go b/internal/agent/tools/query_knowledge_graph.go index 8751d3793..327e6b180 100644 --- a/internal/agent/tools/query_knowledge_graph.go +++ b/internal/agent/tools/query_knowledge_graph.go @@ -68,15 +68,34 @@ type QueryKnowledgeGraphInput struct { // QueryKnowledgeGraphTool queries the knowledge graph for entities and relationships type QueryKnowledgeGraphTool struct { BaseTool - knowledgeService interfaces.KnowledgeBaseService + knowledgeService interfaces.KnowledgeBaseService + scopeKnowledgeService interfaces.KnowledgeService + searchTargets types.SearchTargets +} + +// WithKnowledgeScope enables document/tag-level result filtering for Agent +// calls. The graph backend queries by KB, so the tool must enforce narrower +// SearchTargets before returning any result to the model. +func (t *QueryKnowledgeGraphTool) WithKnowledgeScope( + knowledgeService interfaces.KnowledgeService, +) *QueryKnowledgeGraphTool { + t.scopeKnowledgeService = knowledgeService + return t } // NewQueryKnowledgeGraphTool creates a new query knowledge graph tool -func NewQueryKnowledgeGraphTool(knowledgeService interfaces.KnowledgeBaseService) *QueryKnowledgeGraphTool { - return &QueryKnowledgeGraphTool{ +func NewQueryKnowledgeGraphTool( + knowledgeService interfaces.KnowledgeBaseService, + searchTargets ...types.SearchTargets, +) *QueryKnowledgeGraphTool { + tool := &QueryKnowledgeGraphTool{ BaseTool: queryKnowledgeGraphTool, knowledgeService: knowledgeService, } + if len(searchTargets) > 0 { + tool.searchTargets = searchTargets[0] + } + return tool } // Execute performs the knowledge graph query with concurrent KB processing @@ -105,6 +124,11 @@ func (t *QueryKnowledgeGraphTool) Execute(ctx context.Context, args json.RawMess Error: "knowledge_base_ids must contain at most 10 KB IDs", }, fmt.Errorf("too many KB IDs") } + if len(t.searchTargets) > 0 { + if err := validateKnowledgeBaseIDsInSearchTargets(t.searchTargets, input.KnowledgeBaseIDs); err != nil { + return &types.ToolResult{Success: false, Error: err.Error()}, err + } + } query := input.Query if query == "" { @@ -137,7 +161,7 @@ func (t *QueryKnowledgeGraphTool) Execute(ctx context.Context, args json.RawMess defer wg.Done() // Get knowledge base to check graph configuration - kb, err := t.knowledgeService.GetKnowledgeBaseByID(ctx, id) + kb, err := t.knowledgeService.GetKnowledgeBaseByIDOnly(ctx, id) if err != nil { mu.Lock() kbResults[id] = &graphQueryResult{kbID: id, err: fmt.Errorf("failed to get knowledge base: %v", err)} @@ -161,6 +185,17 @@ func (t *QueryKnowledgeGraphTool) Execute(ctx context.Context, args json.RawMess mu.Unlock() return } + if len(t.searchTargets) > 0 { + results, err = filterSearchResultsInSearchTargets( + ctx, t.searchTargets, id, results, t.scopeKnowledgeService, + ) + if err != nil { + mu.Lock() + kbResults[id] = &graphQueryResult{kbID: id, kb: kb, err: err} + mu.Unlock() + return + } + } mu.Lock() kbResults[id] = &graphQueryResult{kbID: id, kb: kb, results: results} diff --git a/internal/agent/tools/scope_authorization.go b/internal/agent/tools/scope_authorization.go index 8937c84f4..8c0fa223d 100644 --- a/internal/agent/tools/scope_authorization.go +++ b/internal/agent/tools/scope_authorization.go @@ -2,11 +2,149 @@ package tools import ( "context" + "fmt" + "strings" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" ) +func effectiveSearchTargetTagIDs(target *types.SearchTarget) []string { + if target == nil { + return nil + } + return dedupNonEmptyStrings(append( + append([]string(nil), target.TagIDs...), target.ScopeTagIDs..., + )) +} + +// authorizeKnowledgeInSearchTargets is the shared authorization boundary for +// every Agent tool that accepts a model-visible dN/knowledge_id. Handle +// decoding is necessary but never sufficient: the durable document must also +// belong to the server-owned search scope for this Agent execution. +func authorizeKnowledgeInSearchTargets( + ctx context.Context, + searchTargets types.SearchTargets, + knowledgeID string, + knowledgeService interfaces.KnowledgeService, +) (*types.Knowledge, error) { + knowledgeID = strings.TrimSpace(knowledgeID) + if knowledgeID == "" { + return nil, fmt.Errorf("knowledge_id is required") + } + if knowledgeService == nil { + return nil, fmt.Errorf("knowledge service is unavailable") + } + knowledge, err := knowledgeService.GetKnowledgeByIDOnly(ctx, knowledgeID) + if err != nil || knowledge == nil { + if err == nil { + err = fmt.Errorf("empty result") + } + return nil, fmt.Errorf("document %s not found: %w", knowledgeID, err) + } + if !searchTargets.ContainsKB(knowledge.KnowledgeBaseID) { + return nil, fmt.Errorf("knowledge base %s is not within the current Agent scope", knowledge.KnowledgeBaseID) + } + allowed, err := searchTargetsAllowKnowledgeID( + ctx, searchTargets, knowledge.ID, knowledge.KnowledgeBaseID, knowledgeService, + ) + if err != nil { + return nil, fmt.Errorf("failed to validate document scope: %w", err) + } + if !allowed { + return nil, fmt.Errorf("document %s is not within the current @mention scope", knowledge.ID) + } + return knowledge, nil +} + +// authorizeChunkInSearchTargets is the chunk/FAQ counterpart of +// authorizeKnowledgeInSearchTargets. A chunk ID is accepted only after the +// server resolves its owning document and validates that document against the +// full KB/document/tag scope. +func authorizeChunkInSearchTargets( + ctx context.Context, + searchTargets types.SearchTargets, + chunkID string, + chunkService interfaces.ChunkService, + knowledgeService interfaces.KnowledgeService, +) (*types.Chunk, error) { + chunkID = strings.TrimSpace(chunkID) + if chunkID == "" { + return nil, fmt.Errorf("chunk_id is required") + } + if chunkService == nil { + return nil, fmt.Errorf("chunk service is unavailable") + } + chunk, err := chunkService.GetChunkByIDOnly(ctx, chunkID) + if err != nil || chunk == nil { + if err == nil { + err = fmt.Errorf("empty result") + } + return nil, fmt.Errorf("chunk %s not found: %w", chunkID, err) + } + if !searchTargets.ContainsKB(chunk.KnowledgeBaseID) { + return nil, fmt.Errorf("knowledge base %s is not within the current Agent scope", chunk.KnowledgeBaseID) + } + allowed, err := searchTargetsAllowKnowledgeID( + ctx, searchTargets, chunk.KnowledgeID, chunk.KnowledgeBaseID, knowledgeService, + ) + if err != nil { + return nil, fmt.Errorf("failed to validate chunk scope: %w", err) + } + if !allowed { + return nil, fmt.Errorf("chunk %s is not within the current @mention scope", chunk.ID) + } + return chunk, nil +} + +// validateKnowledgeBaseIDsInSearchTargets rejects hallucinated, stale, or +// out-of-scope bN values after the model-context registry resolves them. +func validateKnowledgeBaseIDsInSearchTargets(searchTargets types.SearchTargets, kbIDs []string) error { + for _, kbID := range dedupNonEmptyStrings(kbIDs) { + if !searchTargets.ContainsKB(kbID) { + return fmt.Errorf("knowledge base %s is not within the current Agent scope", kbID) + } + } + return nil +} + +// resolveAuthorizedSourceRefs validates Wiki source_refs/suspected IDs against +// the same Agent scope and rebuilds the stored "uuid|title" representation +// from server data instead of trusting a model-supplied title suffix. +func resolveAuthorizedSourceRefs( + ctx context.Context, + searchTargets types.SearchTargets, + refs []string, + knowledgeService interfaces.KnowledgeService, +) ([]string, error) { + resolved := make([]string, 0, len(refs)) + seen := make(map[string]struct{}, len(refs)) + for _, ref := range refs { + knowledgeID := strings.TrimSpace(strings.SplitN(ref, "|", 2)[0]) + if knowledgeID == "" { + continue + } + knowledge, err := authorizeKnowledgeInSearchTargets(ctx, searchTargets, knowledgeID, knowledgeService) + if err != nil { + return nil, err + } + if _, exists := seen[knowledge.ID]; exists { + continue + } + seen[knowledge.ID] = struct{}{} + title := strings.TrimSpace(knowledge.Title) + if title == "" { + title = strings.TrimSpace(knowledge.FileName) + } + if title == "" { + resolved = append(resolved, knowledge.ID) + } else { + resolved = append(resolved, knowledge.ID+"|"+title) + } + } + return resolved, nil +} + type knowledgeTagsFetcher func(context.Context, []string) (map[string][]*types.KnowledgeTag, error) func searchTargetsAllowKnowledgeID( @@ -27,7 +165,9 @@ func searchTargetsAllowKnowledgeID( continue } matchedKB = true - if target.Type == types.SearchTargetTypeKnowledgeBase && len(target.TagIDs) == 0 { + targetTagIDs := effectiveSearchTargetTagIDs(target) + if target.Type == types.SearchTargetTypeKnowledgeBase && + len(target.KnowledgeIDs) == 0 && len(targetTagIDs) == 0 { return true, nil } for _, allowedID := range target.KnowledgeIDs { @@ -35,7 +175,7 @@ func searchTargetsAllowKnowledgeID( return true, nil } } - tagIDs = append(tagIDs, target.TagIDs...) + tagIDs = append(tagIDs, targetTagIDs...) } if !matchedKB || len(tagIDs) == 0 || knowledgeService == nil { return false, nil @@ -48,6 +188,83 @@ func searchTargetsAllowKnowledgeID( return matches[knowledgeID], nil } +// filterSearchResultsInSearchTargets applies the same whole-KB/document/tag +// union semantics to tools whose backend can only query by KB (notably the +// knowledge graph). It batches tag lookup and rejects results without enough +// provenance instead of turning a narrow mention into whole-KB access. +func filterSearchResultsInSearchTargets( + ctx context.Context, + searchTargets types.SearchTargets, + kbID string, + results []*types.SearchResult, + knowledgeService interfaces.KnowledgeService, +) ([]*types.SearchResult, error) { + var explicitIDs []string + var tagIDs []string + matchedKB := false + for _, target := range searchTargets { + if target == nil || target.KnowledgeBaseID != kbID { + continue + } + matchedKB = true + targetTagIDs := effectiveSearchTargetTagIDs(target) + if target.Type == types.SearchTargetTypeKnowledgeBase && + len(target.KnowledgeIDs) == 0 && len(targetTagIDs) == 0 { + return results, nil + } + explicitIDs = append(explicitIDs, target.KnowledgeIDs...) + tagIDs = append(tagIDs, targetTagIDs...) + } + if !matchedKB { + return nil, fmt.Errorf("knowledge base %s is not within the current Agent scope", kbID) + } + + explicitSet := make(map[string]struct{}, len(explicitIDs)) + for _, id := range dedupNonEmptyStrings(explicitIDs) { + explicitSet[id] = struct{}{} + } + remainingIDs := make([]string, 0, len(results)) + for _, result := range results { + if result == nil || result.KnowledgeID == "" { + continue + } + if result.KnowledgeBaseID != "" && result.KnowledgeBaseID != kbID { + return nil, fmt.Errorf( + "graph result document %s belongs to knowledge base %s, expected %s", + result.KnowledgeID, result.KnowledgeBaseID, kbID, + ) + } + if _, ok := explicitSet[result.KnowledgeID]; !ok { + remainingIDs = append(remainingIDs, result.KnowledgeID) + } + } + var tagMatches map[string]bool + if len(tagIDs) > 0 { + if knowledgeService == nil { + return nil, fmt.Errorf("knowledge service is unavailable for tag-scoped graph filtering") + } + var err error + tagMatches, err = knowledgeIDsMatchingAnyTag( + ctx, remainingIDs, tagIDs, knowledgeService.GetKnowledgeTags, + ) + if err != nil { + return nil, fmt.Errorf("failed to validate graph result scope: %w", err) + } + } + + filtered := make([]*types.SearchResult, 0, len(results)) + for _, result := range results { + if result == nil || result.KnowledgeID == "" { + continue + } + _, explicit := explicitSet[result.KnowledgeID] + if explicit || tagMatches[result.KnowledgeID] { + filtered = append(filtered, result) + } + } + return filtered, nil +} + func knowledgeIDsMatchingAnyTag( ctx context.Context, knowledgeIDs []string, diff --git a/internal/agent/tools/scope_authorization_test.go b/internal/agent/tools/scope_authorization_test.go index 8e8a93259..966d21553 100644 --- a/internal/agent/tools/scope_authorization_test.go +++ b/internal/agent/tools/scope_authorization_test.go @@ -5,8 +5,18 @@ import ( "testing" "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" ) +type scopeKnowledgeService struct { + interfaces.KnowledgeService + knowledge *types.Knowledge +} + +func (s *scopeKnowledgeService) GetKnowledgeByIDOnly(context.Context, string) (*types.Knowledge, error) { + return s.knowledge, nil +} + func testKnowledgeTag(id string) *types.KnowledgeTag { return &types.KnowledgeTag{ID: id} } @@ -34,6 +44,40 @@ func TestKnowledgeIDsMatchingAnyTag(t *testing.T) { } } +func TestAuthorizeKnowledgeInSearchTargetsRejectsSiblingKB(t *testing.T) { + service := &scopeKnowledgeService{knowledge: &types.Knowledge{ID: "doc-2", KnowledgeBaseID: "kb-2"}} + targets := types.SearchTargets{{ + Type: types.SearchTargetTypeKnowledgeBase, + KnowledgeBaseID: "kb-1", + }} + if _, err := authorizeKnowledgeInSearchTargets(context.Background(), targets, "doc-2", service); err == nil { + t.Fatal("document from a sibling KB must be rejected") + } +} + +func TestAuthorizeKnowledgeInSearchTargetsAllowsBoundDocument(t *testing.T) { + service := &scopeKnowledgeService{knowledge: &types.Knowledge{ID: "doc-1", KnowledgeBaseID: "kb-1"}} + targets := types.SearchTargets{{ + Type: types.SearchTargetTypeKnowledge, + KnowledgeBaseID: "kb-1", + KnowledgeIDs: []string{"doc-1"}, + }} + got, err := authorizeKnowledgeInSearchTargets(context.Background(), targets, "doc-1", service) + if err != nil || got == nil || got.ID != "doc-1" { + t.Fatalf("bound document authorization failed: got=%+v err=%v", got, err) + } +} + +func TestValidateKnowledgeBaseIDsInSearchTargetsRejectsMixedScope(t *testing.T) { + targets := types.SearchTargets{{ + Type: types.SearchTargetTypeKnowledgeBase, + KnowledgeBaseID: "kb-1", + }} + if err := validateKnowledgeBaseIDsInSearchTargets(targets, []string{"kb-1", "kb-2"}); err == nil { + t.Fatal("mixed valid and out-of-scope KB IDs must be rejected") + } +} + func TestPagePassesWikiScope_TagScope(t *testing.T) { page := &types.WikiPage{ Slug: "entity/acme", @@ -76,3 +120,89 @@ func TestPagePassesWikiScope_TagScope(t *testing.T) { t.Fatalf("page should be filtered when no source document has the mentioned tag") } } + +func TestNewWikiScopesFromSearchTargetsMergesAlternativesPerKB(t *testing.T) { + targets := types.SearchTargets{ + {Type: types.SearchTargetTypeKnowledge, KnowledgeBaseID: "kb-1", KnowledgeIDs: []string{"doc-1"}}, + {Type: types.SearchTargetTypeKnowledgeBase, KnowledgeBaseID: "kb-1", TagIDs: []string{"tag-1"}}, + {Type: types.SearchTargetTypeKnowledgeBase, KnowledgeBaseID: "kb-2"}, + {Type: types.SearchTargetTypeKnowledge, KnowledgeBaseID: "kb-outside", KnowledgeIDs: []string{"doc-x"}}, + } + scopes := NewWikiScopesFromSearchTargets(targets, []string{"kb-1", "kb-2"}) + if len(scopes) != 2 { + t.Fatalf("scope count = %d, want 2: %+v", len(scopes), scopes) + } + if len(scopes[0].KnowledgeIDs) != 1 || scopes[0].KnowledgeIDs[0] != "doc-1" || + len(scopes[0].TagIDs) != 1 || scopes[0].TagIDs[0] != "tag-1" { + t.Fatalf("kb-1 alternatives were not merged: %+v", scopes[0]) + } + if len(scopes[1].KnowledgeIDs) != 0 || len(scopes[1].TagIDs) != 0 { + t.Fatalf("whole-KB target must remain unrestricted: %+v", scopes[1]) + } +} + +func TestPagePassesWikiScopeUsesDocumentOrTagUnion(t *testing.T) { + page := &types.WikiPage{SourceRefs: []string{"doc-by-tag"}} + passes, err := pagePassesWikiScope( + context.Background(), + page, + WikiScope{ + KnowledgeBaseID: "kb-1", + KnowledgeIDs: []string{"different-explicit-doc"}, + TagIDs: []string{"tag-1"}, + }, + func(context.Context, []string) (map[string][]*types.KnowledgeTag, error) { + return map[string][]*types.KnowledgeTag{ + "doc-by-tag": {testKnowledgeTag("tag-1")}, + }, nil + }, + ) + if err != nil || !passes { + t.Fatalf("tag-authorized page should pass the union scope: passes=%v err=%v", passes, err) + } +} + +func TestPagePassesWikiScopeFailsClosedWithoutProvenance(t *testing.T) { + scope := WikiScope{KnowledgeBaseID: "kb-1", KnowledgeIDs: []string{"doc-1"}} + for _, page := range []*types.WikiPage{ + {Slug: "concept/uncited", PageType: types.WikiPageTypeConcept}, + {Slug: "index", PageType: types.WikiPageTypeIndex}, + } { + passes, err := pagePassesWikiScope(context.Background(), page, scope, nil) + if err != nil { + t.Fatalf("pagePassesWikiScope() error = %v", err) + } + if passes { + t.Fatalf("restricted Wiki scope must reject page without attributable provenance: %+v", page) + } + } +} + +func TestNewWikiScopesUsesLogicalTagsAndSkipsEmptyDocumentTargets(t *testing.T) { + targets := types.SearchTargets{ + {Type: types.SearchTargetTypeKnowledge, KnowledgeBaseID: "kb-tag", ScopeTagIDs: []string{"tag-logical"}}, + {Type: types.SearchTargetTypeKnowledge, KnowledgeBaseID: "kb-empty"}, + } + scopes := NewWikiScopesFromSearchTargets(targets, []string{"kb-tag", "kb-empty"}) + if len(scopes) != 1 || scopes[0].KnowledgeBaseID != "kb-tag" || + len(scopes[0].TagIDs) != 1 || scopes[0].TagIDs[0] != "tag-logical" { + t.Fatalf("unexpected Wiki scopes: %+v", scopes) + } +} + +func TestFilterGraphResultsHonorsDocumentScope(t *testing.T) { + targets := types.SearchTargets{{ + Type: types.SearchTargetTypeKnowledge, KnowledgeBaseID: "kb-1", KnowledgeIDs: []string{"doc-allowed"}, + }} + results := []*types.SearchResult{ + {ID: "chunk-1", KnowledgeID: "doc-allowed", KnowledgeBaseID: "kb-1"}, + {ID: "chunk-2", KnowledgeID: "doc-outside", KnowledgeBaseID: "kb-1"}, + } + filtered, err := filterSearchResultsInSearchTargets(context.Background(), targets, "kb-1", results, nil) + if err != nil { + t.Fatalf("filterSearchResultsInSearchTargets() error = %v", err) + } + if len(filtered) != 1 || filtered[0].KnowledgeID != "doc-allowed" { + t.Fatalf("graph result scope leaked: %+v", filtered) + } +} diff --git a/internal/agent/tools/wiki_delete_page.go b/internal/agent/tools/wiki_delete_page.go index a208c32eb..55c6a4659 100644 --- a/internal/agent/tools/wiki_delete_page.go +++ b/internal/agent/tools/wiki_delete_page.go @@ -15,10 +15,15 @@ type wikiDeletePageTool struct { BaseTool wikiPageService interfaces.WikiPageService kbIDs []string + routes *WikiRouteResolver } // NewWikiDeletePageTool creates a new wiki_delete_page tool -func NewWikiDeletePageTool(wikiPageService interfaces.WikiPageService, kbIDs []string) types.Tool { +func NewWikiDeletePageTool( + wikiPageService interfaces.WikiPageService, + kbIDs []string, + routes ...*WikiRouteResolver, +) types.Tool { return &wikiDeletePageTool{ BaseTool: NewBaseTool( ToolWikiDeletePage, @@ -36,6 +41,7 @@ func NewWikiDeletePageTool(wikiPageService interfaces.WikiPageService, kbIDs []s ), wikiPageService: wikiPageService, kbIDs: kbIDs, + routes: firstWikiRoute(routes), } } @@ -51,60 +57,54 @@ func (t *wikiDeletePageTool) Execute(ctx context.Context, args json.RawMessage) if len(t.kbIDs) == 0 { return &types.ToolResult{Success: false, Error: "No knowledge bases available for editing"}, nil } - kbID := t.kbIDs[0] - if params.Slug == "" { return &types.ToolResult{Success: false, Error: "slug is required"}, nil } + normalizedSlug, slugErr := normalizeAndValidateWikiSlug(params.Slug) + if slugErr != nil { + return &types.ToolResult{Success: false, Error: slugErr.Error()}, nil + } + params.Slug = normalizedSlug // Fetch page to get its incoming links before deleting - existingPage, err := t.wikiPageService.GetPageBySlug(ctx, kbID, params.Slug) + existingPage, kbID, err := resolveUniqueWikiPage(ctx, t.wikiPageService, params.Slug, t.kbIDs, t.routes) if err != nil { return &types.ToolResult{Success: false, Error: "Failed to fetch page to delete: " + err.Error()}, nil } inLinks := make([]string, len(existingPage.InLinks)) copy(inLinks, existingPage.InLinks) + parts := strings.Split(params.Slug, "/") + readableName := strings.ReplaceAll(parts[len(parts)-1], "-", " ") + pipeLinkRE := regexp.MustCompile(`\[\[` + regexp.QuoteMeta(params.Slug) + `\|([^\]]+)\]\]`) + changes, updatedSlugs, rewriteErr := applyIncomingWikiContentRewrite( + ctx, t.wikiPageService, kbID, inLinks, + func(content string) (string, bool) { + updated := strings.ReplaceAll(content, "[["+params.Slug+"]]", readableName) + updated = pipeLinkRE.ReplaceAllString(updated, "$1") + return updated, updated != content + }, + ) + if rewriteErr != nil { + rollbackErr := rollbackWikiContentChanges(ctx, t.wikiPageService, changes) + return &types.ToolResult{ + Success: false, + Error: "Delete aborted while cleaning incoming links: " + + joinWikiMutationErrors(rewriteErr, rollbackErr), + }, nil + } + err = t.wikiPageService.DeletePage(ctx, kbID, params.Slug) if err != nil { - return &types.ToolResult{Success: false, Error: "Failed to delete page: " + err.Error()}, nil - } - - // Clean up incoming links to prevent dead links - updatedCount := 0 - var updatedSlugs []string - for _, sourceSlug := range inLinks { - sourcePage, err := t.wikiPageService.GetPageBySlug(ctx, kbID, sourceSlug) - if err == nil { - changed := false - - // Replace [[deleted-slug]] with readable name - parts := strings.Split(params.Slug, "/") - readableName := parts[len(parts)-1] - readableName = strings.ReplaceAll(readableName, "-", " ") - - link1 := "[[" + params.Slug + "]]" - if strings.Contains(sourcePage.Content, link1) { - sourcePage.Content = strings.ReplaceAll(sourcePage.Content, link1, readableName) - changed = true - } - - // Replace [[deleted-slug|Text]] with just Text - re := regexp.MustCompile(`\[\[` + regexp.QuoteMeta(params.Slug) + `\|([^\]]+)\]\]`) - if re.MatchString(sourcePage.Content) { - sourcePage.Content = re.ReplaceAllString(sourcePage.Content, "$1") - changed = true - } - - if changed { - _, updateErr := t.wikiPageService.UpdatePage(ctx, sourcePage) - if updateErr == nil { - updatedCount++ - updatedSlugs = append(updatedSlugs, sourceSlug) - } - } - } + rollbackErr := rollbackWikiContentChanges(ctx, t.wikiPageService, changes) + return &types.ToolResult{ + Success: false, + Error: "Delete aborted because the page could not be removed: " + + joinWikiMutationErrors(err, rollbackErr), + }, nil } + t.routes.forget(params.Slug, kbID) + updatedCount := len(updatedSlugs) outputMsg := fmt.Sprintf("Successfully deleted page [[%s]] and cleaned up %d incoming links.", params.Slug, updatedCount) if updatedCount > 0 { @@ -115,11 +115,11 @@ func (t *wikiDeletePageTool) Execute(ctx context.Context, args json.RawMessage) Success: true, Output: outputMsg, Data: map[string]interface{}{ - "display_type": "wiki_delete_page", - "slug": params.Slug, - "title": existingPage.Title, - "updated_count": updatedCount, - "affected_pages": updatedSlugs, + "display_type": "wiki_delete_page", + "slug": params.Slug, + "title": existingPage.Title, + "updated_count": updatedCount, + "affected_pages": updatedSlugs, }, }, nil } diff --git a/internal/agent/tools/wiki_flag_issue.go b/internal/agent/tools/wiki_flag_issue.go index d3a6833a3..76959192c 100644 --- a/internal/agent/tools/wiki_flag_issue.go +++ b/internal/agent/tools/wiki_flag_issue.go @@ -12,11 +12,18 @@ import ( type wikiFlagIssueTool struct { BaseTool - wikiService interfaces.WikiPageService - kbIDs []string + wikiService interfaces.WikiPageService + kbIDs []string + routes *WikiRouteResolver + knowledgeService interfaces.KnowledgeService + searchTargets types.SearchTargets } -func NewWikiFlagIssueTool(wikiService interfaces.WikiPageService, kbIDs []string) types.Tool { +func NewWikiFlagIssueTool( + wikiService interfaces.WikiPageService, + kbIDs []string, + routes ...*WikiRouteResolver, +) *wikiFlagIssueTool { return &wikiFlagIssueTool{ BaseTool: NewBaseTool( ToolWikiFlagIssue, @@ -50,9 +57,19 @@ This will log an issue for human review or automated maintenance.`, ), wikiService: wikiService, kbIDs: kbIDs, + routes: firstWikiRoute(routes), } } +func (t *wikiFlagIssueTool) WithKnowledgeScope( + knowledgeService interfaces.KnowledgeService, + searchTargets types.SearchTargets, +) *wikiFlagIssueTool { + t.knowledgeService = knowledgeService + t.searchTargets = searchTargets + return t +} + func (t *wikiFlagIssueTool) Execute(ctx context.Context, args json.RawMessage) (*types.ToolResult, error) { var params struct { Slug string `json:"slug"` @@ -65,21 +82,32 @@ func (t *wikiFlagIssueTool) Execute(ctx context.Context, args json.RawMessage) ( } slug := strings.TrimSpace(params.Slug) - if slug == "" { - return &types.ToolResult{Success: false, Error: "slug is required"}, nil + normalizedSlug, slugErr := normalizeAndValidateWikiSlug(slug) + if slugErr != nil { + return &types.ToolResult{Success: false, Error: slugErr.Error()}, nil } + slug = normalizedSlug if len(t.kbIDs) == 0 { return &types.ToolResult{Success: false, Error: "No knowledge bases available for issue tracking"}, nil } - // Default to first KB ID if multiple (normally there's only one in this context) - kbID := t.kbIDs[0] - - // Verify the page exists - page, err := t.wikiService.GetPageBySlug(ctx, kbID, slug) - if err != nil || page == nil { - return &types.ToolResult{Success: false, Error: fmt.Sprintf("Wiki page with slug '%s' not found", slug)}, nil + page, kbID, err := resolveUniqueWikiPage(ctx, t.wikiService, slug, t.kbIDs, t.routes) + if err != nil { + return &types.ToolResult{Success: false, Error: err.Error()}, nil + } + suspectedKnowledgeIDs := params.SuspectedKnowledgeIDs + if len(t.searchTargets) > 0 && len(suspectedKnowledgeIDs) > 0 { + resolved, scopeErr := resolveAuthorizedSourceRefs( + ctx, t.searchTargets, suspectedKnowledgeIDs, t.knowledgeService, + ) + if scopeErr != nil { + return &types.ToolResult{Success: false, Error: "Invalid suspected_knowledge_ids: " + scopeErr.Error()}, nil + } + suspectedKnowledgeIDs = make([]string, 0, len(resolved)) + for _, ref := range resolved { + suspectedKnowledgeIDs = append(suspectedKnowledgeIDs, strings.SplitN(ref, "|", 2)[0]) + } } issue := &types.WikiPageIssue{ @@ -88,7 +116,7 @@ func (t *wikiFlagIssueTool) Execute(ctx context.Context, args json.RawMessage) ( Slug: slug, IssueType: params.IssueType, Description: params.Description, - SuspectedKnowledgeIDs: params.SuspectedKnowledgeIDs, + SuspectedKnowledgeIDs: suspectedKnowledgeIDs, ReportedBy: "wiki-researcher-agent", Status: "pending", } diff --git a/internal/agent/tools/wiki_link_mutation.go b/internal/agent/tools/wiki_link_mutation.go new file mode 100644 index 000000000..8135aec8b --- /dev/null +++ b/internal/agent/tools/wiki_link_mutation.go @@ -0,0 +1,82 @@ +package tools + +import ( + "context" + "fmt" + "strings" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +type wikiContentRewrite func(string) (string, bool) + +type appliedWikiContentChange struct { + page *types.WikiPage + originalContent string +} + +// applyIncomingWikiContentRewrite updates machine-maintained links without +// incrementing user-visible page versions. It stops on the first failure and +// returns the already-applied changes so the caller can compensate. +func applyIncomingWikiContentRewrite( + ctx context.Context, + service interfaces.WikiPageService, + kbID string, + inLinks []string, + rewrite wikiContentRewrite, +) ([]appliedWikiContentChange, []string, error) { + var changes []appliedWikiContentChange + var updatedSlugs []string + for _, sourceSlug := range dedupNonEmptyStrings(inLinks) { + page, err := service.GetPageBySlug(ctx, kbID, sourceSlug) + if err != nil || page == nil { + if err == nil { + err = fmt.Errorf("empty result") + } + return changes, updatedSlugs, fmt.Errorf("load incoming page %s: %w", sourceSlug, err) + } + updatedContent, changed := rewrite(page.Content) + if !changed { + continue + } + original := page.Content + page.Content = updatedContent + if err := service.UpdateAutoLinkedContent(ctx, page); err != nil { + page.Content = original + return changes, updatedSlugs, fmt.Errorf("update incoming page %s: %w", sourceSlug, err) + } + changes = append(changes, appliedWikiContentChange{page: page, originalContent: original}) + updatedSlugs = append(updatedSlugs, sourceSlug) + } + return changes, updatedSlugs, nil +} + +func rollbackWikiContentChanges( + ctx context.Context, + service interfaces.WikiPageService, + changes []appliedWikiContentChange, +) error { + var failures []string + for i := len(changes) - 1; i >= 0; i-- { + change := changes[i] + change.page.Content = change.originalContent + if err := service.UpdateAutoLinkedContent(ctx, change.page); err != nil { + failures = append(failures, fmt.Sprintf("%s: %v", change.page.Slug, err)) + } + } + if len(failures) > 0 { + return fmt.Errorf("failed to roll back incoming pages: %s", strings.Join(failures, "; ")) + } + return nil +} + +func joinWikiMutationErrors(primary error, extras ...error) string { + parts := []string{primary.Error()} + for _, err := range extras { + if err != nil { + parts = append(parts, err.Error()) + } + } + return strings.Join(parts, "; ") +} diff --git a/internal/agent/tools/wiki_link_mutation_test.go b/internal/agent/tools/wiki_link_mutation_test.go new file mode 100644 index 000000000..7e26e4f7b --- /dev/null +++ b/internal/agent/tools/wiki_link_mutation_test.go @@ -0,0 +1,71 @@ +package tools + +import ( + "context" + "errors" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +type linkMutationWikiService struct { + interfaces.WikiPageService + pages map[string]*types.WikiPage + failUpdates map[string]bool + updateCalls int +} + +func (s *linkMutationWikiService) GetPageBySlug(_ context.Context, _, slug string) (*types.WikiPage, error) { + page := s.pages[slug] + if page == nil { + return nil, errors.New("not found") + } + copyPage := *page + return ©Page, nil +} + +func (s *linkMutationWikiService) UpdateAutoLinkedContent(_ context.Context, page *types.WikiPage) error { + s.updateCalls++ + if s.failUpdates[page.Slug] { + return errors.New("write failed") + } + copyPage := *page + s.pages[page.Slug] = ©Page + return nil +} + +func TestIncomingWikiRewriteCanBeCompensated(t *testing.T) { + service := &linkMutationWikiService{ + pages: map[string]*types.WikiPage{ + "source/a": {Slug: "source/a", Content: "see [[concept/old]]"}, + "source/b": {Slug: "source/b", Content: "also [[concept/old]]"}, + }, + failUpdates: map[string]bool{"source/b": true}, + } + changes, updated, err := applyIncomingWikiContentRewrite( + context.Background(), service, "kb-1", []string{"source/a", "source/b"}, + func(content string) (string, bool) { + if content == "see [[concept/old]]" { + return "see [[concept/new]]", true + } + if content == "also [[concept/old]]" { + return "also [[concept/new]]", true + } + return content, false + }, + ) + if err == nil || len(changes) != 1 || len(updated) != 1 { + t.Fatalf("expected one applied change before failure: changes=%d updated=%v err=%v", len(changes), updated, err) + } + delete(service.failUpdates, "source/b") + if rollbackErr := rollbackWikiContentChanges(context.Background(), service, changes); rollbackErr != nil { + t.Fatalf("rollback failed: %v", rollbackErr) + } + if got := service.pages["source/a"].Content; got != "see [[concept/old]]" { + t.Fatalf("rollback content = %q", got) + } + if service.updateCalls != 3 { + t.Fatalf("UpdateAutoLinkedContent calls = %d, want apply + failed apply + rollback", service.updateCalls) + } +} diff --git a/internal/agent/tools/wiki_read_issue.go b/internal/agent/tools/wiki_read_issue.go index 8148c86ec..31c829835 100644 --- a/internal/agent/tools/wiki_read_issue.go +++ b/internal/agent/tools/wiki_read_issue.go @@ -25,7 +25,7 @@ func NewWikiReadIssueTool(wikiService interfaces.WikiPageService, kbIDs []string "properties": { "issue_id": { "type": "string", - "description": "Optional: The ID of a specific issue to read." + "description": "Optional: The short iN ID of a specific issue from an earlier wiki_read_issue result." }, "slug": { "type": "string", @@ -60,27 +60,31 @@ func (t *wikiReadIssueTool) Execute(ctx context.Context, args json.RawMessage) ( return &types.ToolResult{Success: false, Error: "No knowledge bases available"}, nil } - kbID := t.kbIDs[0] - if issueID != "" { - // Just reuse ListIssues since there's no GetIssueByID yet - issues, err := t.wikiService.ListIssues(ctx, kbID, "", "") + issue, err := resolveWikiIssue(ctx, t.wikiService, issueID, t.kbIDs) + if err != nil { + return &types.ToolResult{Success: false, Error: err.Error()}, nil + } + out, _ := json.MarshalIndent(issue, "", " ") + return &types.ToolResult{Success: true, Output: string(out)}, nil + } + + var issues []*types.WikiPageIssue + for _, kbID := range dedupNonEmptyStrings(t.kbIDs) { + kbIssues, err := t.wikiService.ListIssues(ctx, kbID, slug, "pending") if err != nil { return &types.ToolResult{Success: false, Error: "Failed to list issues: " + err.Error()}, nil } - - for _, issue := range issues { - if issue.ID == issueID { - out, _ := json.MarshalIndent(issue, "", " ") - return &types.ToolResult{Success: true, Output: string(out)}, nil + for _, issue := range kbIssues { + if issue != nil && issue.KnowledgeBaseID != "" && issue.KnowledgeBaseID != kbID { + return &types.ToolResult{ + Success: false, + Error: "Issue result returned knowledge base " + issue.KnowledgeBaseID + + " while resolving allowed scope " + kbID, + }, nil } } - return &types.ToolResult{Success: false, Error: "Issue not found"}, nil - } - - issues, err := t.wikiService.ListIssues(ctx, kbID, slug, "pending") - if err != nil { - return &types.ToolResult{Success: false, Error: "Failed to list issues: " + err.Error()}, nil + issues = append(issues, kbIssues...) } if len(issues) == 0 { diff --git a/internal/agent/tools/wiki_read_source_doc.go b/internal/agent/tools/wiki_read_source_doc.go index b77518492..963c12ec9 100644 --- a/internal/agent/tools/wiki_read_source_doc.go +++ b/internal/agent/tools/wiki_read_source_doc.go @@ -16,10 +16,15 @@ type wikiReadSourceDocTool struct { BaseTool knowledgeService interfaces.KnowledgeService chunkService interfaces.ChunkService + searchTargets types.SearchTargets } -func NewWikiReadSourceDocTool(knowledgeService interfaces.KnowledgeService, chunkService interfaces.ChunkService) types.Tool { - return &wikiReadSourceDocTool{ +func NewWikiReadSourceDocTool( + knowledgeService interfaces.KnowledgeService, + chunkService interfaces.ChunkService, + searchTargets ...types.SearchTargets, +) types.Tool { + tool := &wikiReadSourceDocTool{ BaseTool: NewBaseTool( ToolWikiReadSourceDoc, `Read or search within a specific source document to drill down for details omitted from the wiki. @@ -52,6 +57,10 @@ If neither query nor range is provided, it returns the beginning of the document knowledgeService: knowledgeService, chunkService: chunkService, } + if len(searchTargets) > 0 { + tool.searchTargets = searchTargets[0] + } + return tool } // enrichChunkImageInfo populates chunk.ImageInfo for a batch of parent text @@ -123,8 +132,17 @@ func (t *wikiReadSourceDocTool) Execute(ctx context.Context, args json.RawMessag return &types.ToolResult{Success: false, Error: "knowledge_id is required"}, nil } - knowledge, err := t.knowledgeService.GetKnowledgeByIDOnly(ctx, knowledgeID) - if err != nil { + var knowledge *types.Knowledge + var err error + if len(t.searchTargets) > 0 { + knowledge, err = authorizeKnowledgeInSearchTargets(ctx, t.searchTargets, knowledgeID, t.knowledgeService) + } else { + knowledge, err = t.knowledgeService.GetKnowledgeByIDOnly(ctx, knowledgeID) + } + if err != nil || knowledge == nil { + if err == nil { + err = fmt.Errorf("knowledge service returned an empty result") + } return &types.ToolResult{Success: false, Error: fmt.Sprintf("Document not found: %v", err)}, nil } diff --git a/internal/agent/tools/wiki_rename_page.go b/internal/agent/tools/wiki_rename_page.go index b4c64a25c..efa1b220b 100644 --- a/internal/agent/tools/wiki_rename_page.go +++ b/internal/agent/tools/wiki_rename_page.go @@ -14,10 +14,15 @@ type wikiRenamePageTool struct { BaseTool wikiPageService interfaces.WikiPageService kbIDs []string + routes *WikiRouteResolver } // NewWikiRenamePageTool creates a new wiki_rename_page tool -func NewWikiRenamePageTool(wikiPageService interfaces.WikiPageService, kbIDs []string) types.Tool { +func NewWikiRenamePageTool( + wikiPageService interfaces.WikiPageService, + kbIDs []string, + routes ...*WikiRouteResolver, +) types.Tool { return &wikiRenamePageTool{ BaseTool: NewBaseTool( ToolWikiRenamePage, @@ -39,6 +44,7 @@ func NewWikiRenamePageTool(wikiPageService interfaces.WikiPageService, kbIDs []s ), wikiPageService: wikiPageService, kbIDs: kbIDs, + routes: firstWikiRoute(routes), } } @@ -55,19 +61,27 @@ func (t *wikiRenamePageTool) Execute(ctx context.Context, args json.RawMessage) if len(t.kbIDs) == 0 { return &types.ToolResult{Success: false, Error: "No knowledge bases available for editing"}, nil } - kbID := t.kbIDs[0] - if params.NewSlug == "" { return &types.ToolResult{Success: false, Error: "new_slug is required"}, nil } + normalizedSlug, slugErr := normalizeAndValidateWikiSlug(params.Slug) + if slugErr != nil { + return &types.ToolResult{Success: false, Error: slugErr.Error()}, nil + } + params.Slug = normalizedSlug + normalizedNewSlug, slugErr := normalizeAndValidateWikiSlug(params.NewSlug) + if slugErr != nil { + return &types.ToolResult{Success: false, Error: slugErr.Error()}, nil + } + params.NewSlug = normalizedNewSlug if params.NewSlug == params.Slug { return &types.ToolResult{Success: false, Error: "new_slug must be different from old slug"}, nil } // Get existing page - existingPage, err := t.wikiPageService.GetPageBySlug(ctx, kbID, params.Slug) + existingPage, kbID, err := resolveUniqueWikiPage(ctx, t.wikiPageService, params.Slug, t.kbIDs, t.routes) if err != nil { - return &types.ToolResult{Success: false, Error: fmt.Sprintf("Page %s not found. Cannot rename a non-existent page.", params.Slug)}, nil + return &types.ToolResult{Success: false, Error: "Failed to resolve page to rename: " + err.Error()}, nil } inLinks := make([]string, len(existingPage.InLinks)) @@ -75,58 +89,64 @@ func (t *wikiRenamePageTool) Execute(ctx context.Context, args json.RawMessage) // Create new page with new slug but same content newPage := &types.WikiPage{ + TenantID: existingPage.TenantID, KnowledgeBaseID: kbID, Slug: params.NewSlug, Title: existingPage.Title, Summary: existingPage.Summary, Content: existingPage.Content, PageType: existingPage.PageType, - Aliases: existingPage.Aliases, + Status: existingPage.Status, + Aliases: append(types.StringArray(nil), existingPage.Aliases...), + ParentSlug: existingPage.ParentSlug, + FolderID: existingPage.FolderID, + SortOrder: existingPage.SortOrder, + SourceRefs: append(types.StringArray(nil), existingPage.SourceRefs...), + ChunkRefs: append(types.StringArray(nil), existingPage.ChunkRefs...), + InLinks: append(types.StringArray(nil), existingPage.InLinks...), + PageMetadata: append(types.JSON(nil), existingPage.PageMetadata...), } _, err = t.wikiPageService.CreatePage(ctx, newPage) if err != nil { return &types.ToolResult{Success: false, Error: "Failed to create renamed page: " + err.Error()}, nil } - // Update incoming links in other pages - updatedCount := 0 - var updatedSlugs []string - for _, sourceSlug := range inLinks { - sourcePage, err := t.wikiPageService.GetPageBySlug(ctx, kbID, sourceSlug) - if err == nil { - changed := false - - // Replace [[old-slug]] with [[new-slug]] - link1 := "[[" + params.Slug + "]]" - newLink1 := "[[" + params.NewSlug + "]]" - if strings.Contains(sourcePage.Content, link1) { - sourcePage.Content = strings.ReplaceAll(sourcePage.Content, link1, newLink1) - changed = true - } - - // Replace [[old-slug|text]] with [[new-slug|text]] - link2 := "[[" + params.Slug + "|" - newLink2 := "[[" + params.NewSlug + "|" - if strings.Contains(sourcePage.Content, link2) { - sourcePage.Content = strings.ReplaceAll(sourcePage.Content, link2, newLink2) - changed = true - } - - if changed { - _, updateErr := t.wikiPageService.UpdatePage(ctx, sourcePage) - if updateErr == nil { - updatedCount++ - updatedSlugs = append(updatedSlugs, sourceSlug) - } - } - } + changes, updatedSlugs, rewriteErr := applyIncomingWikiContentRewrite( + ctx, t.wikiPageService, kbID, inLinks, + func(content string) (string, bool) { + updated := strings.ReplaceAll( + content, "[["+params.Slug+"]]", "[["+params.NewSlug+"]]", + ) + updated = strings.ReplaceAll( + updated, "[["+params.Slug+"|", "[["+params.NewSlug+"|", + ) + return updated, updated != content + }, + ) + if rewriteErr != nil { + rollbackErr := rollbackWikiContentChanges(ctx, t.wikiPageService, changes) + cleanupErr := t.wikiPageService.DeletePage(ctx, kbID, params.NewSlug) + return &types.ToolResult{ + Success: false, + Error: "Rename aborted while updating incoming links: " + + joinWikiMutationErrors(rewriteErr, rollbackErr, cleanupErr), + }, nil } + updatedCount := len(updatedSlugs) // Delete old page err = t.wikiPageService.DeletePage(ctx, kbID, params.Slug) if err != nil { - return &types.ToolResult{Success: false, Error: "Successfully created new page and updated links, but failed to delete old page: " + err.Error()}, nil + rollbackErr := rollbackWikiContentChanges(ctx, t.wikiPageService, changes) + cleanupErr := t.wikiPageService.DeletePage(ctx, kbID, params.NewSlug) + return &types.ToolResult{ + Success: false, + Error: "Rename aborted because the old page could not be deleted: " + + joinWikiMutationErrors(err, rollbackErr, cleanupErr), + }, nil } + t.routes.forget(params.Slug, kbID) + t.routes.remember(params.NewSlug, kbID) // Inject cross-links so other pages know about this new slug t.wikiPageService.InjectCrossLinks(ctx, kbID, []string{params.NewSlug}) @@ -143,12 +163,12 @@ func (t *wikiRenamePageTool) Execute(ctx context.Context, args json.RawMessage) Success: true, Output: outputMsg, Data: map[string]interface{}{ - "display_type": "wiki_rename_page", - "old_slug": params.Slug, - "new_slug": params.NewSlug, - "title": existingPage.Title, - "updated_count": updatedCount, - "affected_pages": updatedSlugs, + "display_type": "wiki_rename_page", + "old_slug": params.Slug, + "new_slug": params.NewSlug, + "title": existingPage.Title, + "updated_count": updatedCount, + "affected_pages": updatedSlugs, }, }, nil } diff --git a/internal/agent/tools/wiki_replace_text.go b/internal/agent/tools/wiki_replace_text.go index 63480f8d4..9d19024f6 100644 --- a/internal/agent/tools/wiki_replace_text.go +++ b/internal/agent/tools/wiki_replace_text.go @@ -15,10 +15,17 @@ type wikiReplaceTextTool struct { wikiPageService interfaces.WikiPageService knowledgeService interfaces.KnowledgeService kbIDs []string + routes *WikiRouteResolver + searchTargets types.SearchTargets } // NewWikiReplaceTextTool creates a new wiki_replace_text tool -func NewWikiReplaceTextTool(wikiPageService interfaces.WikiPageService, kbIDs []string, knowledgeService interfaces.KnowledgeService) types.Tool { +func NewWikiReplaceTextTool( + wikiPageService interfaces.WikiPageService, + kbIDs []string, + knowledgeService interfaces.KnowledgeService, + routes ...*WikiRouteResolver, +) *wikiReplaceTextTool { return &wikiReplaceTextTool{ BaseTool: NewBaseTool( ToolWikiReplaceText, @@ -50,15 +57,21 @@ func NewWikiReplaceTextTool(wikiPageService interfaces.WikiPageService, kbIDs [] wikiPageService: wikiPageService, knowledgeService: knowledgeService, kbIDs: kbIDs, + routes: firstWikiRoute(routes), } } +func (t *wikiReplaceTextTool) WithSearchTargets(searchTargets types.SearchTargets) *wikiReplaceTextTool { + t.searchTargets = searchTargets + return t +} + func (t *wikiReplaceTextTool) Execute(ctx context.Context, args json.RawMessage) (*types.ToolResult, error) { var params struct { - Slug string `json:"slug"` - OldText string `json:"old_text"` - NewText string `json:"new_text"` - SourceRefs []string `json:"source_refs"` + Slug string `json:"slug"` + OldText string `json:"old_text"` + NewText string `json:"new_text"` + SourceRefs *[]string `json:"source_refs"` } if err := json.Unmarshal(args, ¶ms); err != nil { @@ -68,14 +81,17 @@ func (t *wikiReplaceTextTool) Execute(ctx context.Context, args json.RawMessage) if len(t.kbIDs) == 0 { return &types.ToolResult{Success: false, Error: "No knowledge bases available for editing"}, nil } - kbID := t.kbIDs[0] - if params.OldText == "" { return &types.ToolResult{Success: false, Error: "old_text is required"}, nil } + normalizedSlug, slugErr := normalizeAndValidateWikiSlug(params.Slug) + if slugErr != nil { + return &types.ToolResult{Success: false, Error: slugErr.Error()}, nil + } + params.Slug = normalizedSlug // Get the existing page - existingPage, err := t.wikiPageService.GetPageBySlug(ctx, kbID, params.Slug) + existingPage, _, err := resolveUniqueWikiPage(ctx, t.wikiPageService, params.Slug, t.kbIDs, t.routes) if err != nil { return &types.ToolResult{Success: false, Error: fmt.Sprintf("Failed to fetch page %s: %v", params.Slug, err)}, nil } @@ -86,8 +102,16 @@ func (t *wikiReplaceTextTool) Execute(ctx context.Context, args json.RawMessage) existingPage.Content = strings.Replace(existingPage.Content, params.OldText, params.NewText, 1) - if len(params.SourceRefs) > 0 { - existingPage.SourceRefs = resolveSourceRefs(ctx, t.knowledgeService, params.SourceRefs) + if params.SourceRefs != nil { + if len(t.searchTargets) > 0 { + resolvedRefs, scopeErr := resolveAuthorizedSourceRefs(ctx, t.searchTargets, *params.SourceRefs, t.knowledgeService) + if scopeErr != nil { + return &types.ToolResult{Success: false, Error: "Invalid source_refs: " + scopeErr.Error()}, nil + } + existingPage.SourceRefs = resolvedRefs + } else { + existingPage.SourceRefs = resolveSourceRefs(ctx, t.knowledgeService, *params.SourceRefs) + } } _, err = t.wikiPageService.UpdatePage(ctx, existingPage) diff --git a/internal/agent/tools/wiki_route_resolver.go b/internal/agent/tools/wiki_route_resolver.go new file mode 100644 index 000000000..28ffdcc06 --- /dev/null +++ b/internal/agent/tools/wiki_route_resolver.go @@ -0,0 +1,293 @@ +package tools + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + + "github.com/Tencent/WeKnora/internal/application/repository" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +var ( + errWikiPageNotFoundInScope = errors.New("wiki page not found in current scope") + errWikiPageAmbiguous = errors.New("wiki page exists in multiple knowledge bases") +) + +// WikiRouteResolver stores request-local provenance for wiki slugs. It is +// server-side routing state, not a model-handle registry: only KBs already in +// the current WikiScope can ever be returned. +type WikiRouteResolver struct { + mu sync.RWMutex + bySlug map[string]map[string]struct{} +} + +func NewWikiRouteResolver() *WikiRouteResolver { + return &WikiRouteResolver{bySlug: make(map[string]map[string]struct{})} +} + +func (r *WikiRouteResolver) remember(slug, kbID string) { + if r == nil || slug == "" || kbID == "" { + return + } + r.mu.Lock() + defer r.mu.Unlock() + owners := r.bySlug[slug] + if owners == nil { + owners = make(map[string]struct{}) + r.bySlug[slug] = owners + } + owners[kbID] = struct{}{} +} + +func (r *WikiRouteResolver) forget(slug, kbID string) { + if r == nil || slug == "" || kbID == "" { + return + } + r.mu.Lock() + defer r.mu.Unlock() + owners := r.bySlug[slug] + delete(owners, kbID) + if len(owners) == 0 { + delete(r.bySlug, slug) + } +} + +func (r *WikiRouteResolver) rememberPage(page *types.WikiPage, kbID string) { + if r == nil || page == nil || kbID == "" { + return + } + r.remember(page.Slug, kbID) + for _, slug := range page.OutLinks { + r.remember(slug, kbID) + } + for _, slug := range page.InLinks { + r.remember(slug, kbID) + } +} + +// scopesForSlug returns cached owners that are still present in the current +// server-owned scope. An empty result means the caller should search all +// scopes. Scope order and per-KB document/tag filters are preserved. +func (r *WikiRouteResolver) scopesForSlug(slug string, scopes []WikiScope) []WikiScope { + if r == nil || slug == "" || len(scopes) == 0 { + return nil + } + r.mu.RLock() + defer r.mu.RUnlock() + owners := r.bySlug[slug] + if len(owners) == 0 { + return nil + } + matched := make([]WikiScope, 0, len(owners)) + for _, scope := range scopes { + if _, ok := owners[scope.KnowledgeBaseID]; ok { + matched = append(matched, scope) + } + } + return matched +} + +func scopesOutsideKBs(scopes []WikiScope, excluded []WikiScope) []WikiScope { + if len(excluded) == 0 { + return scopes + } + excludedKBs := make(map[string]struct{}, len(excluded)) + for _, scope := range excluded { + excludedKBs[scope.KnowledgeBaseID] = struct{}{} + } + remaining := make([]WikiScope, 0, len(scopes)) + for _, scope := range scopes { + if _, ok := excludedKBs[scope.KnowledgeBaseID]; !ok { + remaining = append(remaining, scope) + } + } + return remaining +} + +// resolveUniqueWikiPage is the shared mutation/issue routing boundary. It +// checks every allowed KB (cached provenance only affects order) and refuses +// ambiguous slugs instead of silently mutating the first KB. +func resolveUniqueWikiPage( + ctx context.Context, + service interfaces.WikiPageService, + slug string, + kbIDs []string, + routes *WikiRouteResolver, +) (*types.WikiPage, string, error) { + slug = strings.TrimSpace(slug) + if slug == "" { + return nil, "", fmt.Errorf("slug is required") + } + scopes := NewWikiScopesFromKBIDs(kbIDs) + preferred := routes.scopesForSlug(slug, scopes) + ordered := append(append([]WikiScope(nil), preferred...), scopesOutsideKBs(scopes, preferred)...) + type hit struct { + page *types.WikiPage + kbID string + } + var hits []hit + for _, scope := range ordered { + page, err := service.GetPageBySlug(ctx, scope.KnowledgeBaseID, slug) + if err != nil { + if errors.Is(err, repository.ErrWikiPageNotFound) { + continue + } + return nil, "", fmt.Errorf( + "failed to resolve wiki page %s in knowledge base %s: %w", + slug, scope.KnowledgeBaseID, err, + ) + } + if page == nil { + continue + } + if page.KnowledgeBaseID != "" && page.KnowledgeBaseID != scope.KnowledgeBaseID { + return nil, "", fmt.Errorf( + "wiki page %s returned knowledge base %s while resolving allowed scope %s", + slug, page.KnowledgeBaseID, scope.KnowledgeBaseID, + ) + } + kbID := scope.KnowledgeBaseID + hits = append(hits, hit{page: page, kbID: kbID}) + routes.rememberPage(page, kbID) + } + switch len(hits) { + case 0: + return nil, "", fmt.Errorf("%w: %s", errWikiPageNotFoundInScope, slug) + case 1: + return hits[0].page, hits[0].kbID, nil + default: + owners := make([]string, 0, len(hits)) + for _, item := range hits { + owners = append(owners, item.kbID) + } + return nil, "", fmt.Errorf("%w: slug %s belongs to %s", errWikiPageAmbiguous, slug, strings.Join(owners, ", ")) + } +} + +// resolveWikiCreateKB selects a creation target only when server-side context +// is unambiguous: one cached provenance owner or one Wiki KB in scope. +func resolveWikiCreateKB( + slug string, + kbIDs []string, + routes *WikiRouteResolver, + serverHints ...string, +) (string, error) { + scopes := NewWikiScopesFromKBIDs(kbIDs) + preferred := routes.scopesForSlug(strings.TrimSpace(slug), scopes) + allowed := make(map[string]struct{}, len(scopes)) + for _, scope := range scopes { + allowed[scope.KnowledgeBaseID] = struct{}{} + } + candidates := make([]string, 0, len(preferred)+len(serverHints)) + for _, scope := range preferred { + candidates = append(candidates, scope.KnowledgeBaseID) + } + for _, kbID := range serverHints { + if _, ok := allowed[kbID]; ok { + candidates = append(candidates, kbID) + } + } + candidates = dedupNonEmptyStrings(candidates) + if len(candidates) == 1 { + return candidates[0], nil + } + if len(candidates) > 1 { + return "", fmt.Errorf( + "cannot choose a knowledge base for new wiki page %s: server provenance conflicts across %s", + slug, strings.Join(candidates, ", "), + ) + } + if len(scopes) == 1 { + return scopes[0].KnowledgeBaseID, nil + } + return "", fmt.Errorf("cannot choose a knowledge base for new wiki page %s from %d allowed scopes", slug, len(scopes)) +} + +func wikiKnowledgeBasesForSourceRefs( + ctx context.Context, + refs []string, + knowledgeService interfaces.KnowledgeService, + allowedKBIDs []string, +) ([]string, error) { + if len(refs) == 0 { + return nil, nil + } + if knowledgeService == nil { + return nil, fmt.Errorf("knowledge service is unavailable") + } + allowed := make(map[string]struct{}, len(allowedKBIDs)) + for _, kbID := range dedupNonEmptyStrings(allowedKBIDs) { + allowed[kbID] = struct{}{} + } + var kbIDs []string + for _, ref := range refs { + knowledgeID := strings.TrimSpace(strings.SplitN(ref, "|", 2)[0]) + if knowledgeID == "" { + continue + } + knowledge, err := knowledgeService.GetKnowledgeByIDOnly(ctx, knowledgeID) + if err != nil || knowledge == nil { + if err == nil { + err = fmt.Errorf("empty result") + } + return nil, fmt.Errorf("failed to resolve source document %s: %w", knowledgeID, err) + } + if _, ok := allowed[knowledge.KnowledgeBaseID]; !ok { + return nil, fmt.Errorf( + "source document %s belongs to non-Wiki or unauthorized knowledge base %s", + knowledgeID, knowledge.KnowledgeBaseID, + ) + } + kbIDs = append(kbIDs, knowledge.KnowledgeBaseID) + } + return dedupNonEmptyStrings(kbIDs), nil +} + +func firstWikiRoute(routes []*WikiRouteResolver) *WikiRouteResolver { + if len(routes) > 0 && routes[0] != nil { + return routes[0] + } + return NewWikiRouteResolver() +} + +func resolveWikiIssue( + ctx context.Context, + service interfaces.WikiPageService, + issueID string, + kbIDs []string, +) (*types.WikiPageIssue, error) { + issueID = strings.TrimSpace(issueID) + if issueID == "" { + return nil, fmt.Errorf("issue_id is required") + } + var match *types.WikiPageIssue + for _, kbID := range dedupNonEmptyStrings(kbIDs) { + issues, err := service.ListIssues(ctx, kbID, "", "") + if err != nil { + return nil, fmt.Errorf("failed to list issues in knowledge base %s: %w", kbID, err) + } + for _, issue := range issues { + if issue == nil || issue.ID != issueID { + continue + } + if issue.KnowledgeBaseID != "" && issue.KnowledgeBaseID != kbID { + return nil, fmt.Errorf( + "issue_id %s returned knowledge base %s while resolving allowed scope %s", + issueID, issue.KnowledgeBaseID, kbID, + ) + } + if match != nil { + return nil, fmt.Errorf("issue_id %s is ambiguous across current Wiki scopes", issueID) + } + match = issue + } + } + if match == nil { + return nil, fmt.Errorf("issue_id %s is not within the current Wiki scope", issueID) + } + return match, nil +} diff --git a/internal/agent/tools/wiki_tools.go b/internal/agent/tools/wiki_tools.go index b389de6a9..661a5c8f7 100644 --- a/internal/agent/tools/wiki_tools.go +++ b/internal/agent/tools/wiki_tools.go @@ -3,11 +3,13 @@ package tools import ( "context" "encoding/json" + "errors" "fmt" "regexp" "strings" "sync" + "github.com/Tencent/WeKnora/internal/application/repository" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" ) @@ -110,12 +112,72 @@ type WikiScope struct { // NewWikiScopesFromKBIDs is a convenience constructor for callers that only // carry plain KB IDs and don't need per-document filtering (e.g. legacy tests). func NewWikiScopesFromKBIDs(kbIDs []string) []WikiScope { - scopes := make([]WikiScope, 0, len(kbIDs)) - for _, id := range kbIDs { - if id == "" { + // Deduplicate defensively: resolveUniqueWikiPage counts one hit per scope + // as a distinct owner, so a duplicated KB ID would misreport a unique slug + // as ambiguous. + deduped := dedupNonEmptyStrings(kbIDs) + scopes := make([]WikiScope, 0, len(deduped)) + for _, id := range deduped { + scopes = append(scopes, WikiScope{KnowledgeBaseID: id}) + } + return scopes +} + +// NewWikiScopesFromSearchTargets merges all targets for a Wiki KB into one +// scope. SearchTargets are alternatives selected by the user, so document and +// tag constraints are combined as a union; an unrestricted whole-KB target +// supersedes narrower targets for that KB. +func NewWikiScopesFromSearchTargets(searchTargets types.SearchTargets, wikiKBIDs []string) []WikiScope { + allowed := make(map[string]struct{}, len(wikiKBIDs)) + for _, kbID := range dedupNonEmptyStrings(wikiKBIDs) { + allowed[kbID] = struct{}{} + } + type accumulatedScope struct { + WikiScope + unrestricted bool + } + byKB := make(map[string]*accumulatedScope, len(allowed)) + for _, target := range searchTargets { + if target == nil { continue } - scopes = append(scopes, WikiScope{KnowledgeBaseID: id}) + if _, ok := allowed[target.KnowledgeBaseID]; !ok { + continue + } + targetTagIDs := effectiveSearchTargetTagIDs(target) + wholeKB := target.Type == types.SearchTargetTypeKnowledgeBase && + len(target.KnowledgeIDs) == 0 && len(targetTagIDs) == 0 + if !wholeKB && len(target.KnowledgeIDs) == 0 && len(targetTagIDs) == 0 { + // A malformed empty document target must not silently become + // whole-KB authorization. + continue + } + scope := byKB[target.KnowledgeBaseID] + if scope == nil { + scope = &accumulatedScope{WikiScope: WikiScope{KnowledgeBaseID: target.KnowledgeBaseID}} + byKB[target.KnowledgeBaseID] = scope + } + if wholeKB { + scope.unrestricted = true + continue + } + scope.KnowledgeIDs = append(scope.KnowledgeIDs, target.KnowledgeIDs...) + scope.TagIDs = append(scope.TagIDs, targetTagIDs...) + } + + scopes := make([]WikiScope, 0, len(allowed)) + for _, kbID := range dedupNonEmptyStrings(wikiKBIDs) { + scope := byKB[kbID] + if scope == nil { + continue + } + if scope.unrestricted { + scopes = append(scopes, WikiScope{KnowledgeBaseID: kbID}) + continue + } + scope.KnowledgeIDs = dedupNonEmptyStrings(scope.KnowledgeIDs) + scope.TagIDs = dedupNonEmptyStrings(scope.TagIDs) + scopes = append(scopes, scope.WikiScope) } return scopes } @@ -241,23 +303,28 @@ func pagePassesWikiScope( scope WikiScope, fetchTags knowledgeTagsFetcher, ) (bool, error) { - if allowed, has := scopeKnowledgeFilter(scope); has { - if !pageIntersectsKnowledgeIDs(page, allowed) { - return false, nil - } - } - + allowed, hasKnowledgeFilter := scopeKnowledgeFilter(scope) tagIDs := dedupNonEmptyStrings(scope.TagIDs) - if len(tagIDs) == 0 { + if !hasKnowledgeFilter && len(tagIDs) == 0 { return true, nil } + // Under a document/tag-constrained scope, every surfaced page must prove + // provenance. Structural pages and uncited pages describe the whole Wiki + // and cannot be safely attributed to the selected subset; wiki_search is + // still available for scoped discovery. if isStructuralPage(page) { - return true, nil + return false, nil } sourceKnowledgeIDs := extractSourceKnowledgeIDs(page) if len(sourceKnowledgeIDs) == 0 { + return false, nil + } + if hasKnowledgeFilter && pageIntersectsKnowledgeIDs(page, allowed) { return true, nil } + if len(tagIDs) == 0 { + return false, nil + } matches, err := knowledgeIDsMatchingAnyTag(ctx, sourceKnowledgeIDs, tagIDs, fetchTags) if err != nil { return false, err @@ -270,6 +337,7 @@ type wikiReadPageTool struct { wikiService interfaces.WikiPageService knowledgeService interfaces.KnowledgeService scopes []WikiScope + routes *WikiRouteResolver seenLinks map[string]bool mu sync.Mutex } @@ -278,13 +346,17 @@ func NewWikiReadPageTool( wikiService interfaces.WikiPageService, knowledgeService interfaces.KnowledgeService, scopes []WikiScope, + routes *WikiRouteResolver, ) types.Tool { + if routes == nil { + routes = NewWikiRouteResolver() + } return &wikiReadPageTool{ BaseTool: NewBaseTool( ToolWikiReadPage, `Read one or more wiki pages by their slugs. Returns the full markdown content, metadata, and links. Use this to read specific wiki pages when you know their slug (e.g. "entity/acme-corp", "concept/rag"). -When the same slug exists in multiple knowledge bases, all matching pages are returned (each tagged with its short bN knowledge_base_id). Pass that bN value in "knowledge_base_id" to limit to a specific KB.`, +Knowledge-base routing is automatic. Known link/search provenance is preferred; otherwise every wiki knowledge base in scope is checked and ambiguous matches are returned.`, json.RawMessage(`{ "type": "object", "properties": { @@ -292,10 +364,6 @@ When the same slug exists in multiple knowledge bases, all matching pages are re "type": "array", "items": { "type": "string" }, "description": "List of wiki page slugs to read (e.g. ['entity/acme-corp', 'index'])" - }, - "knowledge_base_id": { - "type": "string", - "description": "Optional: specific short bN knowledge base ID. If omitted, reads the slug from every wiki KB in scope (all matches returned)." } }, "required": ["slugs"] @@ -304,6 +372,7 @@ When the same slug exists in multiple knowledge bases, all matching pages are re wikiService: wikiService, knowledgeService: knowledgeService, scopes: scopes, + routes: routes, seenLinks: make(map[string]bool), } } @@ -316,9 +385,8 @@ func seenLinkKey(kbID, slug string) string { func (t *wikiReadPageTool) Execute(ctx context.Context, args json.RawMessage) (*types.ToolResult, error) { var params struct { - Slug any `json:"slug"` - Slugs any `json:"slugs"` - KnowledgeBaseID string `json:"knowledge_base_id"` + Slug any `json:"slug"` + Slugs any `json:"slugs"` } if err := json.Unmarshal(args, ¶ms); err != nil { return &types.ToolResult{Success: false, Error: "Invalid parameters: " + err.Error()}, nil @@ -332,27 +400,6 @@ func (t *wikiReadPageTool) Execute(ctx context.Context, args json.RawMessage) (* return &types.ToolResult{Success: false, Error: "Missing 'slugs' parameter"}, nil } - // Build effective scope list. If the caller pinned a knowledge_base_id, - // limit to that KB while preserving any scope-level knowledge_ids filter - // (the filter itself is never exposed as a tool argument — it comes from - // the server-side scope so the model doesn't have to reason about it). - effectiveScopes := t.scopes - if params.KnowledgeBaseID != "" { - filtered := make([]WikiScope, 0, 1) - for _, sc := range t.scopes { - if sc.KnowledgeBaseID == params.KnowledgeBaseID { - filtered = append(filtered, sc) - break - } - } - if len(filtered) == 0 { - // Not in the agent's scope list — still allow direct addressing - // but without any pin (scopes were the source of the pin). - filtered = append(filtered, WikiScope{KnowledgeBaseID: params.KnowledgeBaseID}) - } - effectiveScopes = filtered - } - var outputs []string var errs []string // Per-slug list of KB IDs where the slug was found. A slug may exist in @@ -420,6 +467,11 @@ func (t *wikiReadPageTool) Execute(ctx context.Context, args json.RawMessage) (* if page.PageType == types.WikiPageTypeIndex { if overview, err := t.wikiService.GetIndexView(ctx, kbID, nil, wikiIndexAgentTopK, ""); err == nil && overview != nil { contentBody = renderIndexOverviewForAgent(overview) + for _, group := range overview.Groups { + for _, item := range group.Items { + t.routes.remember(item.Slug, kbID) + } + } } } @@ -458,6 +510,7 @@ func (t *wikiReadPageTool) Execute(ctx context.Context, args json.RawMessage) (* // Track slugs that were found in the raw lookup but filtered out by a // knowledge_ids whitelist, so we can surface a clearer error. filteredOut := make(map[string][]string) // slug -> list of KB IDs where filtered + lookupFailed := make(map[string]bool) var fetchTags knowledgeTagsFetcher if t.knowledgeService != nil { fetchTags = t.knowledgeService.GetKnowledgeTags @@ -467,61 +520,86 @@ func (t *wikiReadPageTool) Execute(ctx context.Context, args json.RawMessage) (* page *types.WikiPage kbID string } - for _, sc := range effectiveScopes { - kbID := sc.KnowledgeBaseID - if kbID == "" { - continue - } - page, err := t.wikiService.GetPageBySlug(ctx, kbID, slug) - if err != nil || page == nil { - continue - } - actualKBID := kbID - if page.KnowledgeBaseID != "" { - actualKBID = page.KnowledgeBaseID - } + cachedScopes := t.routes.scopesForSlug(slug, t.scopes) + // Provenance is an ordering hint, never an authority boundary. Always + // inspect every legal scope so a cached hit cannot hide a duplicate slug + // in another Wiki KB and make the model act on an apparently unique page. + effectiveScopes := append( + append([]WikiScope(nil), cachedScopes...), + scopesOutsideKBs(t.scopes, cachedScopes)..., + ) - // Apply server-enforced source-document / tag scope (silent; never - // exposed to the model as a tool argument). - passesScope, scopeErr := pagePassesWikiScope(ctx, page, sc, fetchTags) - if scopeErr != nil { - errs = append(errs, fmt.Sprintf("Failed to validate wiki scope for '%s' in KB %s: %v", slug, actualKBID, scopeErr)) - continue - } - if !passesScope { - filteredOut[slug] = append(filteredOut[slug], actualKBID) - continue - } + lookupScopes := func(scopes []WikiScope) { + for _, sc := range scopes { + kbID := sc.KnowledgeBaseID + if kbID == "" { + continue + } + page, err := t.wikiService.GetPageBySlug(ctx, kbID, slug) + if err != nil { + if !errors.Is(err, repository.ErrWikiPageNotFound) { + lookupFailed[slug] = true + errs = append(errs, fmt.Sprintf("Failed to read Wiki page '%s' in KB %s: %v", slug, kbID, err)) + } + continue + } + if page == nil { + continue + } + if page.KnowledgeBaseID != "" && page.KnowledgeBaseID != kbID { + lookupFailed[slug] = true + errs = append(errs, fmt.Sprintf( + "Wiki page '%s' returned KB %s while resolving allowed KB %s", + slug, page.KnowledgeBaseID, kbID, + )) + continue + } + actualKBID := kbID - hits = append(hits, struct { - page *types.WikiPage - kbID string - }{page, actualKBID}) - foundKBs[slug] = append(foundKBs[slug], actualKBID) - // Also register the page's neighbours so that when the model - // echoes a link like `[[summary/xyz]]` from this page's body, - // the frontend can resolve it to the same KB without guessing. - registerLinkedSlugs(foundKBs, page, actualKBID) - t.mu.Lock() - t.seenLinks[seenLinkKey(actualKBID, slug)] = true - t.mu.Unlock() + // Apply server-enforced source-document / tag scope (silent; never + // exposed to the model as a tool argument). + passesScope, scopeErr := pagePassesWikiScope(ctx, page, sc, fetchTags) + if scopeErr != nil { + errs = append(errs, fmt.Sprintf("Failed to validate wiki scope for '%s' in KB %s: %v", slug, actualKBID, scopeErr)) + continue + } + if !passesScope { + filteredOut[slug] = append(filteredOut[slug], actualKBID) + continue + } + + hits = append(hits, struct { + page *types.WikiPage + kbID string + }{page, actualKBID}) + t.routes.rememberPage(page, actualKBID) + foundKBs[slug] = append(foundKBs[slug], actualKBID) + // Also register the page's neighbours so that when the model + // echoes a link like `[[summary/xyz]]` from this page's body, + // the frontend can resolve it to the same KB without guessing. + registerLinkedSlugs(foundKBs, page, actualKBID) + t.mu.Lock() + t.seenLinks[seenLinkKey(actualKBID, slug)] = true + t.mu.Unlock() + } } + lookupScopes(effectiveScopes) + if len(hits) == 0 { if kbs := filteredOut[slug]; len(kbs) > 0 { errs = append(errs, fmt.Sprintf( "Wiki page '%s' exists in %v but none of its source documents are within the scope pinned by the user", slug, kbs, )) - } else { + } else if !lookupFailed[slug] { errs = append(errs, fmt.Sprintf("Wiki page '%s' not found", slug)) } continue } - // When the same slug exists in multiple KBs (and the caller did not - // specify a knowledge_base_id), emit all pages so the model can pick - // the right one or compare them explicitly. + // When routing remains ambiguous, emit every matched page so the model + // can use or compare them explicitly. for _, h := range hits { outputs = append(outputs, renderPage(h.page, h.kbID)) } @@ -562,6 +640,7 @@ type wikiSearchTool struct { wikiService interfaces.WikiPageService knowledgeService interfaces.KnowledgeService scopes []WikiScope + routes *WikiRouteResolver seenSlugs map[string]bool mu sync.Mutex } @@ -570,7 +649,11 @@ func NewWikiSearchTool( wikiService interfaces.WikiPageService, knowledgeService interfaces.KnowledgeService, scopes []WikiScope, + routes *WikiRouteResolver, ) types.Tool { + if routes == nil { + routes = NewWikiRouteResolver() + } return &wikiSearchTool{ BaseTool: NewBaseTool( ToolWikiSearch, @@ -607,6 +690,7 @@ Use this to find relevant wiki pages when you don't know the exact slug.`, wikiService: wikiService, knowledgeService: knowledgeService, scopes: scopes, + routes: routes, seenSlugs: make(map[string]bool), } } @@ -645,12 +729,17 @@ func (t *wikiSearchTool) Execute(ctx context.Context, args json.RawMessage) (*ty } } if len(filtered) == 0 { - filtered = append(filtered, WikiScope{KnowledgeBaseID: params.KnowledgeBaseID}) + return &types.ToolResult{ + Success: false, + Error: "knowledge_base_id is not within the current wiki scope", + }, nil } effectiveScopes = filtered } var allOutputs []string + var searchErrors []string + successfulSearchCalls := 0 // Per-slug list of KB IDs that produced a match. Multiple KBs may share a // slug when the agent has several wiki KBs in scope, so we keep the full list. foundKBs := make(map[string][]string) @@ -674,22 +763,35 @@ func (t *wikiSearchTool) Execute(ctx context.Context, args json.RawMessage) (*ty } pages, err := t.wikiService.SearchPages(ctx, kbID, query, params.Limit) if err != nil { + searchErrors = append(searchErrors, fmt.Sprintf("Wiki search %q failed in KB %s: %v", query, kbID, err)) continue } + successfulSearchCalls++ for _, p := range pages { if p == nil { continue } passesScope, scopeErr := pagePassesWikiScope(ctx, p, sc, fetchTags) - if scopeErr != nil || !passesScope { + if scopeErr != nil { + searchErrors = append(searchErrors, fmt.Sprintf( + "Failed to validate Wiki search result %q in KB %s: %v", p.Slug, kbID, scopeErr, + )) + continue + } + if !passesScope { filteredCount++ continue } - actualKBID := kbID - if p.KnowledgeBaseID != "" { - actualKBID = p.KnowledgeBaseID + if p.KnowledgeBaseID != "" && p.KnowledgeBaseID != kbID { + searchErrors = append(searchErrors, fmt.Sprintf( + "Wiki search result %q returned KB %s while resolving allowed KB %s", + p.Slug, p.KnowledgeBaseID, kbID, + )) + continue } + actualKBID := kbID allHits = append(allHits, searchHit{page: p, kbID: actualKBID}) + t.routes.rememberPage(p, actualKBID) foundKBs[p.Slug] = append(foundKBs[p.Slug], actualKBID) // Register neighbour slugs so links surfaced from this // page's body can be routed to the same KB by the frontend. @@ -737,9 +839,16 @@ func (t *wikiSearchTool) Execute(ctx context.Context, args json.RawMessage) (*ty allOutputs = append(allOutputs, sb.String()) } + if successfulSearchCalls == 0 && len(searchErrors) > 0 { + return &types.ToolResult{Success: false, Error: strings.Join(searchErrors, "; ")}, nil + } + output := strings.Join(allOutputs, "\n\n") + if len(searchErrors) > 0 { + output += "\n\n\n" + strings.Join(searchErrors, "\n") + "\n" + } return &types.ToolResult{ Success: true, - Output: strings.Join(allOutputs, "\n\n"), + Output: output, Data: map[string]interface{}{ "found_kbs": foundKBs, }, diff --git a/internal/agent/tools/wiki_tools_test.go b/internal/agent/tools/wiki_tools_test.go index fc7e17def..998a3817a 100644 --- a/internal/agent/tools/wiki_tools_test.go +++ b/internal/agent/tools/wiki_tools_test.go @@ -1,9 +1,72 @@ package tools import ( + "context" + "encoding/json" + "errors" + "reflect" + "strings" + "sync" "testing" + + "github.com/Tencent/WeKnora/internal/application/repository" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" ) +type fakeWikiPageService struct { + interfaces.WikiPageService + + mu sync.Mutex + pages map[string]*types.WikiPage + searchResults map[string][]*types.WikiPage + getErrors map[string]error + searchErrors map[string]error + getCalls []string + searchCalls []string +} + +func wikiPageKey(kbID, slug string) string { + return kbID + "\x00" + slug +} + +func (f *fakeWikiPageService) GetPageBySlug(_ context.Context, kbID, slug string) (*types.WikiPage, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.getCalls = append(f.getCalls, wikiPageKey(kbID, slug)) + if err := f.getErrors[wikiPageKey(kbID, slug)]; err != nil { + return nil, err + } + return f.pages[wikiPageKey(kbID, slug)], nil +} + +func (f *fakeWikiPageService) SearchPages(_ context.Context, kbID, _ string, _ int) ([]*types.WikiPage, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.searchCalls = append(f.searchCalls, kbID) + if err := f.searchErrors[kbID]; err != nil { + return nil, err + } + return f.searchResults[kbID], nil +} + +func (f *fakeWikiPageService) resetCalls() { + f.mu.Lock() + defer f.mu.Unlock() + f.getCalls = nil + f.searchCalls = nil +} + +func newTestWikiPage(kbID, slug string) *types.WikiPage { + return &types.WikiPage{ + KnowledgeBaseID: kbID, + Slug: slug, + Title: slug, + PageType: types.WikiPageTypeConcept, + Content: "content for " + slug, + } +} + func TestTruncateForSummary(t *testing.T) { tests := []struct { name string @@ -95,3 +158,294 @@ func TestWikiToolsInAvailableDefinitions(t *testing.T) { } } } + +func TestWikiReadPageSchemaUsesAutomaticKBRouting(t *testing.T) { + tool := NewWikiReadPageTool( + &fakeWikiPageService{}, nil, + NewWikiScopesFromKBIDs([]string{"kb-1"}), NewWikiRouteResolver(), + ) + if strings.Contains(string(tool.Parameters()), "knowledge_base_id") { + t.Fatalf("wiki_read_page schema must not ask the model to choose a knowledge base: %s", tool.Parameters()) + } +} + +func TestWikiReadPageRoutesEachSlugIndependently(t *testing.T) { + pageA := newTestWikiPage("kb-1", "concept/a") + pageB := newTestWikiPage("kb-2", "concept/b") + service := &fakeWikiPageService{pages: map[string]*types.WikiPage{ + wikiPageKey("kb-1", pageA.Slug): pageA, + wikiPageKey("kb-2", pageB.Slug): pageB, + }} + routes := NewWikiRouteResolver() + routes.remember(pageA.Slug, "kb-1") + routes.remember(pageB.Slug, "kb-2") + tool := NewWikiReadPageTool( + service, nil, NewWikiScopesFromKBIDs([]string{"kb-1", "kb-2"}), routes, + ) + + result, err := tool.Execute(context.Background(), json.RawMessage(`{"slugs":["concept/a","concept/b"]}`)) + if err != nil || result == nil || !result.Success { + t.Fatalf("wiki_read_page failed: result=%+v err=%v", result, err) + } + wantCalls := []string{ + wikiPageKey("kb-1", "concept/a"), + wikiPageKey("kb-2", "concept/a"), + wikiPageKey("kb-2", "concept/b"), + wikiPageKey("kb-1", "concept/b"), + } + if !reflect.DeepEqual(service.getCalls, wantCalls) { + t.Fatalf("GetPageBySlug calls = %v, want %v", service.getCalls, wantCalls) + } +} + +func TestWikiReadPageFallsBackWhenCachedRouteIsStale(t *testing.T) { + page := newTestWikiPage("kb-2", "concept/target") + service := &fakeWikiPageService{pages: map[string]*types.WikiPage{ + wikiPageKey("kb-2", page.Slug): page, + }} + routes := NewWikiRouteResolver() + routes.remember(page.Slug, "kb-1") + tool := NewWikiReadPageTool( + service, nil, NewWikiScopesFromKBIDs([]string{"kb-1", "kb-2"}), routes, + ) + + result, err := tool.Execute(context.Background(), json.RawMessage(`{"slugs":["concept/target"]}`)) + if err != nil || result == nil || !result.Success { + t.Fatalf("wiki_read_page should recover from a stale route: result=%+v err=%v", result, err) + } + wantCalls := []string{ + wikiPageKey("kb-1", page.Slug), + wikiPageKey("kb-2", page.Slug), + } + if !reflect.DeepEqual(service.getCalls, wantCalls) { + t.Fatalf("GetPageBySlug calls = %v, want stale route then fallback %v", service.getCalls, wantCalls) + } +} + +func TestWikiReadPageDoesNotLetLegacyKBArgumentPinAllSlugs(t *testing.T) { + page := newTestWikiPage("kb-2", "concept/target") + service := &fakeWikiPageService{pages: map[string]*types.WikiPage{ + wikiPageKey("kb-2", page.Slug): page, + }} + tool := NewWikiReadPageTool( + service, nil, + NewWikiScopesFromKBIDs([]string{"kb-1", "kb-2"}), NewWikiRouteResolver(), + ) + + result, err := tool.Execute(context.Background(), json.RawMessage( + `{"slugs":["concept/target"],"knowledge_base_id":"b2"}`, + )) + if err != nil || result == nil || !result.Success { + t.Fatalf("legacy KB argument must not pin wiki_read_page: result=%+v err=%v", result, err) + } + wantCalls := []string{ + wikiPageKey("kb-1", page.Slug), + wikiPageKey("kb-2", page.Slug), + } + if !reflect.DeepEqual(service.getCalls, wantCalls) { + t.Fatalf("GetPageBySlug calls = %v, want automatic scope scan %v", service.getCalls, wantCalls) + } +} + +func TestWikiSearchSharesRoutesWithWikiReadPage(t *testing.T) { + page := newTestWikiPage("kb-2", "concept/target") + service := &fakeWikiPageService{ + pages: map[string]*types.WikiPage{ + wikiPageKey("kb-2", page.Slug): page, + }, + searchResults: map[string][]*types.WikiPage{ + "kb-2": {page}, + }, + } + scopes := NewWikiScopesFromKBIDs([]string{"kb-1", "kb-2"}) + routes := NewWikiRouteResolver() + searchTool := NewWikiSearchTool(service, nil, scopes, routes) + readTool := NewWikiReadPageTool(service, nil, scopes, routes) + + searchResult, err := searchTool.Execute(context.Background(), json.RawMessage(`{"queries":["target"]}`)) + if err != nil || searchResult == nil || !searchResult.Success { + t.Fatalf("wiki_search failed: result=%+v err=%v", searchResult, err) + } + service.resetCalls() + + readResult, err := readTool.Execute(context.Background(), json.RawMessage(`{"slugs":["concept/target"]}`)) + if err != nil || readResult == nil || !readResult.Success { + t.Fatalf("wiki_read_page failed: result=%+v err=%v", readResult, err) + } + wantCalls := []string{ + wikiPageKey("kb-2", page.Slug), + wikiPageKey("kb-1", page.Slug), + } + if !reflect.DeepEqual(service.getCalls, wantCalls) { + t.Fatalf("GetPageBySlug calls = %v, want cached owner first and full ambiguity scan %v", service.getCalls, wantCalls) + } +} + +func TestWikiReadPageCachedRouteCannotHideDuplicateSlug(t *testing.T) { + slug := "concept/shared" + service := &fakeWikiPageService{pages: map[string]*types.WikiPage{ + wikiPageKey("kb-1", slug): newTestWikiPage("kb-1", slug), + wikiPageKey("kb-2", slug): newTestWikiPage("kb-2", slug), + }} + routes := NewWikiRouteResolver() + routes.remember(slug, "kb-1") + tool := NewWikiReadPageTool( + service, nil, NewWikiScopesFromKBIDs([]string{"kb-1", "kb-2"}), routes, + ) + + result, err := tool.Execute(context.Background(), json.RawMessage(`{"slugs":["concept/shared"]}`)) + if err != nil || result == nil || !result.Success { + t.Fatalf("wiki_read_page failed: result=%+v err=%v", result, err) + } + if !strings.Contains(result.Output, "kb-1") || + !strings.Contains(result.Output, "kb-2") { + t.Fatalf("cached provenance hid an ambiguous page: %s", result.Output) + } +} + +func TestWikiReadPageReturnsSameSlugFromAllScopesOnCacheMiss(t *testing.T) { + slug := "concept/shared" + pageA := newTestWikiPage("kb-1", slug) + pageB := newTestWikiPage("kb-2", slug) + service := &fakeWikiPageService{pages: map[string]*types.WikiPage{ + wikiPageKey("kb-1", slug): pageA, + wikiPageKey("kb-2", slug): pageB, + }} + tool := NewWikiReadPageTool( + service, nil, NewWikiScopesFromKBIDs([]string{"kb-1", "kb-2"}), NewWikiRouteResolver(), + ) + + result, err := tool.Execute(context.Background(), json.RawMessage(`{"slugs":["concept/shared"]}`)) + if err != nil || result == nil || !result.Success { + t.Fatalf("wiki_read_page failed: result=%+v err=%v", result, err) + } + if !strings.Contains(result.Output, "kb-1") || + !strings.Contains(result.Output, "kb-2") { + t.Fatalf("expected pages from both scopes, got: %s", result.Output) + } +} + +func TestWikiSearchRejectsKnowledgeBaseOutsideScope(t *testing.T) { + service := &fakeWikiPageService{} + tool := NewWikiSearchTool( + service, nil, NewWikiScopesFromKBIDs([]string{"kb-1"}), NewWikiRouteResolver(), + ) + + result, err := tool.Execute(context.Background(), json.RawMessage(`{"queries":["target"],"knowledge_base_id":"kb-outside"}`)) + if err != nil { + t.Fatalf("wiki_search returned unexpected error: %v", err) + } + if result == nil || result.Success || !strings.Contains(result.Error, "not within") { + t.Fatalf("expected scope rejection, got: %+v", result) + } + if len(service.searchCalls) != 0 { + t.Fatalf("out-of-scope search must not reach the service, calls=%v", service.searchCalls) + } +} + +func TestWikiReadPageSurfacesBackendFailureInsteadOfNotFound(t *testing.T) { + service := &fakeWikiPageService{getErrors: map[string]error{ + wikiPageKey("kb-1", "concept/a"): errors.New("database unavailable"), + }} + tool := NewWikiReadPageTool( + service, nil, NewWikiScopesFromKBIDs([]string{"kb-1"}), NewWikiRouteResolver(), + ) + result, err := tool.Execute(context.Background(), json.RawMessage(`{"slugs":["concept/a"]}`)) + if err != nil || result == nil || result.Success || !strings.Contains(result.Error, "database unavailable") { + t.Fatalf("backend failure must be preserved: result=%+v err=%v", result, err) + } + if strings.Contains(result.Error, "not found") { + t.Fatalf("backend failure must not be rewritten as not found: %s", result.Error) + } +} + +func TestWikiSearchFailsWhenEveryAllowedBackendLookupFails(t *testing.T) { + service := &fakeWikiPageService{searchErrors: map[string]error{ + "kb-1": errors.New("search backend unavailable"), + }} + tool := NewWikiSearchTool( + service, nil, NewWikiScopesFromKBIDs([]string{"kb-1"}), NewWikiRouteResolver(), + ) + result, err := tool.Execute(context.Background(), json.RawMessage(`{"queries":["target"]}`)) + if err != nil || result == nil || result.Success || !strings.Contains(result.Error, "search backend unavailable") { + t.Fatalf("search backend failure must be preserved: result=%+v err=%v", result, err) + } +} + +func TestResolveUniqueWikiPageRejectsAmbiguousMutationTarget(t *testing.T) { + slug := "concept/shared" + service := &fakeWikiPageService{pages: map[string]*types.WikiPage{ + wikiPageKey("kb-1", slug): newTestWikiPage("kb-1", slug), + wikiPageKey("kb-2", slug): newTestWikiPage("kb-2", slug), + }} + _, _, err := resolveUniqueWikiPage( + context.Background(), service, slug, []string{"kb-1", "kb-2"}, NewWikiRouteResolver(), + ) + if err == nil || !strings.Contains(err.Error(), "multiple knowledge bases") { + t.Fatalf("expected ambiguous target rejection, got %v", err) + } +} + +func TestResolveUniqueWikiPageUsesActualOwnerWithoutModelKBArgument(t *testing.T) { + slug := "concept/target" + service := &fakeWikiPageService{pages: map[string]*types.WikiPage{ + wikiPageKey("kb-2", slug): newTestWikiPage("kb-2", slug), + }} + page, kbID, err := resolveUniqueWikiPage( + context.Background(), service, slug, []string{"kb-1", "kb-2"}, NewWikiRouteResolver(), + ) + if err != nil || page == nil || kbID != "kb-2" { + t.Fatalf("target resolution failed: page=%+v kb=%s err=%v", page, kbID, err) + } +} + +func TestResolveUniqueWikiPageSkipsOnlyNotFoundErrors(t *testing.T) { + slug := "concept/target" + service := &fakeWikiPageService{ + pages: map[string]*types.WikiPage{ + wikiPageKey("kb-2", slug): newTestWikiPage("kb-2", slug), + }, + getErrors: map[string]error{ + wikiPageKey("kb-1", slug): repository.ErrWikiPageNotFound, + }, + } + page, kbID, err := resolveUniqueWikiPage( + context.Background(), service, slug, []string{"kb-1", "kb-2"}, NewWikiRouteResolver(), + ) + if err != nil || page == nil || kbID != "kb-2" { + t.Fatalf("not-found should continue to other allowed scopes: page=%+v kb=%s err=%v", page, kbID, err) + } + + service.getErrors[wikiPageKey("kb-1", slug)] = errors.New("database unavailable") + _, _, err = resolveUniqueWikiPage( + context.Background(), service, slug, []string{"kb-1", "kb-2"}, NewWikiRouteResolver(), + ) + if err == nil || !strings.Contains(err.Error(), "database unavailable") { + t.Fatalf("non-not-found lookup errors must fail closed, got %v", err) + } +} + +func TestResolveUniqueWikiPageRejectsMismatchedServiceOwner(t *testing.T) { + slug := "concept/target" + service := &fakeWikiPageService{pages: map[string]*types.WikiPage{ + wikiPageKey("kb-1", slug): newTestWikiPage("kb-outside", slug), + }} + _, _, err := resolveUniqueWikiPage( + context.Background(), service, slug, []string{"kb-1"}, NewWikiRouteResolver(), + ) + if err == nil || !strings.Contains(err.Error(), "returned knowledge base kb-outside") { + t.Fatalf("mismatched service owner must be rejected, got %v", err) + } +} + +func TestResolveWikiCreateKBRequiresUnambiguousServerContext(t *testing.T) { + routes := NewWikiRouteResolver() + if _, err := resolveWikiCreateKB("concept/new", []string{"kb-1", "kb-2"}, routes); err == nil { + t.Fatal("new page creation across multiple unhinted KBs must be rejected") + } + routes.remember("concept/new", "kb-2") + kbID, err := resolveWikiCreateKB("concept/new", []string{"kb-1", "kb-2"}, routes) + if err != nil || kbID != "kb-2" { + t.Fatalf("cached provenance should choose kb-2: kb=%s err=%v", kbID, err) + } +} diff --git a/internal/agent/tools/wiki_update_issue.go b/internal/agent/tools/wiki_update_issue.go index 75a20de56..55b2ce913 100644 --- a/internal/agent/tools/wiki_update_issue.go +++ b/internal/agent/tools/wiki_update_issue.go @@ -25,7 +25,7 @@ func NewWikiUpdateIssueTool(wikiService interfaces.WikiPageService, kbIDs []stri "properties": { "issue_id": { "type": "string", - "description": "The ID of the issue to update." + "description": "The short iN issue ID from wiki_read_issue." }, "status": { "type": "string", @@ -60,8 +60,11 @@ func (t *wikiUpdateIssueTool) Execute(ctx context.Context, args json.RawMessage) if len(t.kbIDs) == 0 { return &types.ToolResult{Success: false, Error: "No knowledge bases available"}, nil } + if _, err := resolveWikiIssue(ctx, t.wikiService, params.IssueID, t.kbIDs); err != nil { + return &types.ToolResult{Success: false, Error: err.Error()}, nil + } - // Update issue status + // Update only after the issue has been proven to belong to an allowed KB. err := t.wikiService.UpdateIssueStatus(ctx, params.IssueID, params.Status) if err != nil { return &types.ToolResult{Success: false, Error: "Failed to update issue status: " + err.Error()}, nil diff --git a/internal/agent/tools/wiki_write_page.go b/internal/agent/tools/wiki_write_page.go index d603f5c60..6933c0821 100644 --- a/internal/agent/tools/wiki_write_page.go +++ b/internal/agent/tools/wiki_write_page.go @@ -7,7 +7,6 @@ import ( "fmt" "strings" - "github.com/Tencent/WeKnora/internal/application/repository" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" ) @@ -17,10 +16,17 @@ type wikiWritePageTool struct { wikiPageService interfaces.WikiPageService knowledgeService interfaces.KnowledgeService kbIDs []string + routes *WikiRouteResolver + searchTargets types.SearchTargets } // NewWikiWritePageTool creates a new wiki_write_page tool -func NewWikiWritePageTool(wikiPageService interfaces.WikiPageService, kbIDs []string, knowledgeService interfaces.KnowledgeService) types.Tool { +func NewWikiWritePageTool( + wikiPageService interfaces.WikiPageService, + kbIDs []string, + knowledgeService interfaces.KnowledgeService, + routes ...*WikiRouteResolver, +) *wikiWritePageTool { return &wikiWritePageTool{ BaseTool: NewBaseTool( ToolWikiWritePage, @@ -51,7 +57,7 @@ func NewWikiWritePageTool(wikiPageService interfaces.WikiPageService, kbIDs []st "aliases": { "type": "array", "items": {"type": "string"}, - "description": "A list of aliases for the page (optional)" + "description": "A list of aliases for the page (optional). If provided, these will COMPLETELY REPLACE the existing aliases of the page." }, "source_refs": { "type": "array", @@ -65,18 +71,24 @@ func NewWikiWritePageTool(wikiPageService interfaces.WikiPageService, kbIDs []st wikiPageService: wikiPageService, knowledgeService: knowledgeService, kbIDs: kbIDs, + routes: firstWikiRoute(routes), } } +func (t *wikiWritePageTool) WithSearchTargets(searchTargets types.SearchTargets) *wikiWritePageTool { + t.searchTargets = searchTargets + return t +} + func (t *wikiWritePageTool) Execute(ctx context.Context, args json.RawMessage) (*types.ToolResult, error) { var params struct { - Slug string `json:"slug"` - Title string `json:"title"` - Summary string `json:"summary"` - Content string `json:"content"` - PageType string `json:"page_type"` - Aliases []string `json:"aliases"` - SourceRefs []string `json:"source_refs"` + Slug string `json:"slug"` + Title string `json:"title"` + Summary string `json:"summary"` + Content string `json:"content"` + PageType string `json:"page_type"` + Aliases *[]string `json:"aliases"` + SourceRefs *[]string `json:"source_refs"` } if err := json.Unmarshal(args, ¶ms); err != nil { @@ -86,8 +98,6 @@ func (t *wikiWritePageTool) Execute(ctx context.Context, args json.RawMessage) ( if len(t.kbIDs) == 0 { return &types.ToolResult{Success: false, Error: "No knowledge bases available for editing"}, nil } - kbID := t.kbIDs[0] - if params.Title == "" || params.PageType == "" || params.Content == "" || params.Summary == "" { return &types.ToolResult{Success: false, Error: "title, summary, content, and page_type are required for write action"}, nil } @@ -101,17 +111,43 @@ func (t *wikiWritePageTool) Execute(ctx context.Context, args json.RawMessage) ( } params.Slug = normalizedSlug - // Try to get the existing page - existingPage, err := t.wikiPageService.GetPageBySlug(ctx, kbID, params.Slug) - if err != nil && !errors.Is(err, repository.ErrWikiPageNotFound) { - return &types.ToolResult{Success: false, Error: "Failed to check existing page: " + err.Error()}, nil + // Resolve and authorize provenance before choosing a creation target. In a + // multi-Wiki Agent, source documents provide a server-owned KB hint and + // remove the need for a model-visible knowledge_base_id argument. + var resolvedRefs []string + var err error + if params.SourceRefs != nil { + if len(t.searchTargets) > 0 { + resolvedRefs, err = resolveAuthorizedSourceRefs(ctx, t.searchTargets, *params.SourceRefs, t.knowledgeService) + if err != nil { + return &types.ToolResult{Success: false, Error: "Invalid source_refs: " + err.Error()}, nil + } + } else { + resolvedRefs = resolveSourceRefs(ctx, t.knowledgeService, *params.SourceRefs) + } + } + // Resolve existing pages across every legal Wiki KB. Cached provenance only + // influences lookup order; ambiguous slugs are never silently written to + // the first KB. New pages require one unambiguous creation target. + existingPage, kbID, err := resolveUniqueWikiPage(ctx, t.wikiPageService, params.Slug, t.kbIDs, t.routes) + if errors.Is(err, errWikiPageNotFoundInScope) { + existingPage = nil + sourceKBHints, sourceErr := wikiKnowledgeBasesForSourceRefs( + ctx, resolvedRefs, t.knowledgeService, t.kbIDs, + ) + if sourceErr != nil { + return &types.ToolResult{Success: false, Error: "Failed to resolve source_refs routing: " + sourceErr.Error()}, nil + } + kbID, err = resolveWikiCreateKB(params.Slug, t.kbIDs, t.routes, sourceKBHints...) + } + if err != nil { + return &types.ToolResult{Success: false, Error: "Failed to resolve wiki target: " + err.Error()}, nil } // Summary pages are system-owned: they are generated deterministically // from a source document and keyed by its knowledge ID // (summary/). Letting the agent CREATE one with an // arbitrary/hand-typed slug is precisely how mangled-UUID ghost summary - // pages (e.g. summary/…b171c…) get into the KB. Allow updating an // existing summary page, but never fabricating a new one. if existingPage == nil && (isSummaryNamespace(params.Slug) || strings.EqualFold(params.PageType, types.WikiPageTypeSummary)) { @@ -131,8 +167,6 @@ func (t *wikiWritePageTool) Execute(ctx context.Context, args json.RawMessage) ( params.Content = repaired } - resolvedRefs := resolveSourceRefs(ctx, t.knowledgeService, params.SourceRefs) - var action string if existingPage != nil { // Update @@ -140,9 +174,11 @@ func (t *wikiWritePageTool) Execute(ctx context.Context, args json.RawMessage) ( existingPage.Summary = params.Summary existingPage.Content = params.Content existingPage.PageType = params.PageType - existingPage.Aliases = params.Aliases + if params.Aliases != nil { + existingPage.Aliases = *params.Aliases + } - if len(resolvedRefs) > 0 { + if params.SourceRefs != nil { existingPage.SourceRefs = resolvedRefs } @@ -160,15 +196,18 @@ func (t *wikiWritePageTool) Execute(ctx context.Context, args json.RawMessage) ( Summary: params.Summary, Content: params.Content, PageType: params.PageType, - Aliases: params.Aliases, SourceRefs: resolvedRefs, } + if params.Aliases != nil { + newPage.Aliases = *params.Aliases + } _, err = t.wikiPageService.CreatePage(ctx, newPage) if err != nil { return &types.ToolResult{Success: false, Error: "Failed to create page: " + err.Error()}, nil } action = "created" } + t.routes.remember(params.Slug, kbID) // Inject cross-links so other pages know about this new/updated entity t.wikiPageService.InjectCrossLinks(ctx, kbID, []string{params.Slug}) @@ -177,10 +216,10 @@ func (t *wikiWritePageTool) Execute(ctx context.Context, args json.RawMessage) ( _ = t.wikiPageService.RebuildIndexPage(ctx, kbID) output := fmt.Sprintf("Successfully %s page [[%s]].\n- Title: %s\n- Type: %s\n- Summary: %s\n- Content length: %d chars", action, params.Slug, params.Title, params.PageType, params.Summary, len(params.Content)) - if len(params.Aliases) > 0 { - output += fmt.Sprintf("\n- Aliases: %s", strings.Join(params.Aliases, ", ")) + if params.Aliases != nil && len(*params.Aliases) > 0 { + output += fmt.Sprintf("\n- Aliases: %s", strings.Join(*params.Aliases, ", ")) } - if len(resolvedRefs) > 0 { + if params.SourceRefs != nil { output += fmt.Sprintf("\n- Source refs: %d document(s)", len(resolvedRefs)) } diff --git a/internal/agent/tools/wiki_write_page_test.go b/internal/agent/tools/wiki_write_page_test.go index f32b1b646..275459096 100644 --- a/internal/agent/tools/wiki_write_page_test.go +++ b/internal/agent/tools/wiki_write_page_test.go @@ -1,6 +1,41 @@ package tools -import "testing" +import ( + "context" + "encoding/json" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" +) + +type sourceRefWikiService struct { + interfaces.WikiPageService + page *types.WikiPage + createdKB string +} + +func (s *sourceRefWikiService) GetPageBySlug(context.Context, string, string) (*types.WikiPage, error) { + return s.page, nil +} + +func (s *sourceRefWikiService) RepairContentLinks(_ context.Context, _, _, content string) (string, bool, error) { + return content, false, nil +} + +func (s *sourceRefWikiService) UpdatePage(_ context.Context, page *types.WikiPage) (*types.WikiPage, error) { + s.page = page + return page, nil +} + +func (s *sourceRefWikiService) CreatePage(_ context.Context, page *types.WikiPage) (*types.WikiPage, error) { + s.page = page + s.createdKB = page.KnowledgeBaseID + return page, nil +} + +func (s *sourceRefWikiService) InjectCrossLinks(context.Context, string, []string) {} +func (s *sourceRefWikiService) RebuildIndexPage(context.Context, string) error { return nil } func TestNormalizeAndValidateWikiSlug(t *testing.T) { tests := []struct { @@ -51,3 +86,106 @@ func TestIsSummaryNamespace(t *testing.T) { t.Fatal("entity/summary-of-x must not count as the summary namespace") } } + +func TestWikiWritePageDistinguishesOmittedAndExplicitEmptySourceRefs(t *testing.T) { + for _, test := range []struct { + name string + args string + want int + }{ + { + name: "omitted preserves provenance", + args: `{"slug":"concept/a","title":"A","summary":"S","content":"C","page_type":"concept"}`, + want: 1, + }, + { + name: "empty clears provenance", + args: `{"slug":"concept/a","title":"A","summary":"S","content":"C","page_type":"concept","source_refs":[]}`, + want: 0, + }, + } { + t.Run(test.name, func(t *testing.T) { + service := &sourceRefWikiService{page: &types.WikiPage{ + KnowledgeBaseID: "kb-1", + Slug: "concept/a", + SourceRefs: types.StringArray{"doc-real|Document"}, + }} + tool := NewWikiWritePageTool(service, []string{"kb-1"}, nil, NewWikiRouteResolver()) + result, err := tool.Execute(context.Background(), json.RawMessage(test.args)) + if err != nil || result == nil || !result.Success { + t.Fatalf("write failed: result=%+v err=%v", result, err) + } + if got := len(service.page.SourceRefs); got != test.want { + t.Fatalf("source_refs length = %d, want %d", got, test.want) + } + }) + } +} + +func TestWikiWritePageDistinguishesOmittedAndExplicitEmptyAliases(t *testing.T) { + for _, test := range []struct { + name string + args string + want []string + }{ + { + name: "omitted preserves stored aliases", + args: `{"slug":"concept/a","title":"A","summary":"S","content":"C","page_type":"concept"}`, + want: []string{"kept"}, + }, + { + name: "empty clears stored aliases", + args: `{"slug":"concept/a","title":"A","summary":"S","content":"C","page_type":"concept","aliases":[]}`, + want: nil, + }, + { + name: "explicit list replaces stored aliases", + args: `{"slug":"concept/a","title":"A","summary":"S","content":"C","page_type":"concept","aliases":["fresh"]}`, + want: []string{"fresh"}, + }, + } { + t.Run(test.name, func(t *testing.T) { + service := &sourceRefWikiService{page: &types.WikiPage{ + KnowledgeBaseID: "kb-1", + Slug: "concept/a", + Aliases: types.StringArray{"kept"}, + }} + tool := NewWikiWritePageTool(service, []string{"kb-1"}, nil, NewWikiRouteResolver()) + result, err := tool.Execute(context.Background(), json.RawMessage(test.args)) + if err != nil || result == nil || !result.Success { + t.Fatalf("write failed: result=%+v err=%v", result, err) + } + if len(service.page.Aliases) != len(test.want) { + t.Fatalf("aliases = %v, want %v", service.page.Aliases, test.want) + } + for i := range test.want { + if service.page.Aliases[i] != test.want[i] { + t.Fatalf("aliases = %v, want %v", service.page.Aliases, test.want) + } + } + }) + } +} + +func TestWikiWritePageRoutesNewPageFromAuthorizedSourceRefs(t *testing.T) { + service := &sourceRefWikiService{} + knowledgeService := &scopeKnowledgeService{knowledge: &types.Knowledge{ + ID: "doc-2", KnowledgeBaseID: "kb-2", Title: "Document 2", + }} + searchTargets := types.SearchTargets{{ + Type: types.SearchTargetTypeKnowledge, KnowledgeBaseID: "kb-2", KnowledgeIDs: []string{"doc-2"}, + }} + tool := NewWikiWritePageTool( + service, []string{"kb-1", "kb-2"}, knowledgeService, NewWikiRouteResolver(), + ).WithSearchTargets(searchTargets) + + result, err := tool.Execute(context.Background(), json.RawMessage( + `{"slug":"concept/new","title":"New","summary":"Summary","content":"Content","page_type":"concept","source_refs":["doc-2"]}`, + )) + if err != nil || result == nil || !result.Success { + t.Fatalf("write failed: result=%+v err=%v", result, err) + } + if service.createdKB != "kb-2" { + t.Fatalf("new page routed to %q, want source-owned kb-2", service.createdKB) + } +} diff --git a/internal/application/repository/wiki_page.go b/internal/application/repository/wiki_page.go index 9e85f3084..2763d8d20 100644 --- a/internal/application/repository/wiki_page.go +++ b/internal/application/repository/wiki_page.go @@ -112,6 +112,7 @@ func (r *wikiPageRepository) UpdateMeta(ctx context.Context, page *types.WikiPag Updates(map[string]interface{}{ "in_links": page.InLinks, "out_links": page.OutLinks, + "aliases": page.Aliases, "status": page.Status, "source_refs": page.SourceRefs, "chunk_refs": page.ChunkRefs, diff --git a/internal/application/service/agent_service.go b/internal/application/service/agent_service.go index 3e79d0813..6d5cfc782 100644 --- a/internal/application/service/agent_service.go +++ b/internal/application/service/agent_service.go @@ -440,10 +440,13 @@ func (s *agentService) registerTools( } // ---- Capability detection from SearchTargets ---- - var hasVectorKB, hasWikiKB bool + var hasVectorKB bool var wikiKBIDs []string - var wikiScopes []tools.WikiScope + wikiRoutes := tools.NewWikiRouteResolver() for _, target := range config.SearchTargets { + if target == nil || target.KnowledgeBaseID == "" { + continue + } kb, err := s.knowledgeBaseService.GetKnowledgeBaseByIDOnly(ctx, target.KnowledgeBaseID) if err != nil { continue @@ -452,21 +455,16 @@ func (s *agentService) registerTools( hasVectorKB = true } if kb.IsWikiEnabled() { - hasWikiKB = true wikiKBIDs = append(wikiKBIDs, kb.ID) - // When the user @mentioned specific documents, carry the document - // whitelist into the wiki scope so wiki_search / wiki_read_page - // only surface pages whose SourceRefs intersect the pinned docs. - scope := tools.WikiScope{KnowledgeBaseID: kb.ID} - if target.Type == types.SearchTargetTypeKnowledge && len(target.KnowledgeIDs) > 0 { - scope.KnowledgeIDs = append([]string(nil), target.KnowledgeIDs...) - } - if len(target.TagIDs) > 0 { - scope.TagIDs = append([]string(nil), target.TagIDs...) - } - wikiScopes = append(wikiScopes, scope) } } + wikiKBIDs = dedupStrings(wikiKBIDs) + wikiScopes := tools.NewWikiScopesFromSearchTargets(config.SearchTargets, wikiKBIDs) + wikiKBIDs = wikiKBIDs[:0] + for _, scope := range wikiScopes { + wikiKBIDs = append(wikiKBIDs, scope.KnowledgeBaseID) + } + hasWikiKB := len(wikiKBIDs) > 0 // Filter out knowledge base tools if no knowledge scope is configured for this turn. hasKnowledge := agentHasKnowledgeScope(config) @@ -604,7 +602,8 @@ func (s *agentService) registerTools( case tools.ToolListKnowledgeChunks: toolToRegister = tools.NewListKnowledgeChunksTool(s.knowledgeService, s.chunkService, config.SearchTargets) case tools.ToolQueryKnowledgeGraph: - toolToRegister = tools.NewQueryKnowledgeGraphTool(s.knowledgeBaseService) + toolToRegister = tools.NewQueryKnowledgeGraphTool(s.knowledgeBaseService, config.SearchTargets). + WithKnowledgeScope(s.knowledgeService) case tools.ToolGetDocumentInfo: toolToRegister = tools.NewGetDocumentInfoTool(s.knowledgeService, s.chunkService, config.SearchTargets) case tools.ToolDatabaseQuery: @@ -626,34 +625,39 @@ func (s *agentService) registerTools( logger.Infof(ctx, "Registered web_fetch tool for session: %s", sessionID) case tools.ToolDataAnalysis: - toolToRegister = tools.NewDataAnalysisTool(s.knowledgeBaseService, s.knowledgeService, s.tenantService, s.fileService, s.duckdb, sessionID, s.storageResolver) + toolToRegister = tools.NewDataAnalysisTool(s.knowledgeBaseService, s.knowledgeService, s.tenantService, s.fileService, s.duckdb, sessionID, s.storageResolver). + WithSearchTargets(config.SearchTargets) logger.Infof(ctx, "Registered data_analysis tool for session: %s", sessionID) case tools.ToolDataSchema: - toolToRegister = tools.NewDataSchemaTool(s.knowledgeService, s.chunkService.GetRepository()) + toolToRegister = tools.NewDataSchemaTool(s.knowledgeService, s.chunkService.GetRepository()). + WithSearchTargets(config.SearchTargets) logger.Infof(ctx, "Registered data_schema tool") // Wiki tools — only registered when wiki KBs are detected case tools.ToolWikiReadPage: - toolToRegister = tools.NewWikiReadPageTool(s.wikiPageService, s.knowledgeService, wikiScopes) + toolToRegister = tools.NewWikiReadPageTool(s.wikiPageService, s.knowledgeService, wikiScopes, wikiRoutes) case tools.ToolWikiSearch: - toolToRegister = tools.NewWikiSearchTool(s.wikiPageService, s.knowledgeService, wikiScopes) + toolToRegister = tools.NewWikiSearchTool(s.wikiPageService, s.knowledgeService, wikiScopes, wikiRoutes) case tools.ToolWikiReadSourceDoc: - toolToRegister = tools.NewWikiReadSourceDocTool(s.knowledgeService, s.chunkService) + toolToRegister = tools.NewWikiReadSourceDocTool(s.knowledgeService, s.chunkService, config.SearchTargets) case tools.ToolWikiFlagIssue: - toolToRegister = tools.NewWikiFlagIssueTool(s.wikiPageService, wikiKBIDs) + toolToRegister = tools.NewWikiFlagIssueTool(s.wikiPageService, wikiKBIDs, wikiRoutes). + WithKnowledgeScope(s.knowledgeService, config.SearchTargets) case tools.ToolWikiReadIssue: toolToRegister = tools.NewWikiReadIssueTool(s.wikiPageService, wikiKBIDs) case tools.ToolWikiUpdateIssue: toolToRegister = tools.NewWikiUpdateIssueTool(s.wikiPageService, wikiKBIDs) case tools.ToolWikiWritePage: - toolToRegister = tools.NewWikiWritePageTool(s.wikiPageService, wikiKBIDs, s.knowledgeService) + toolToRegister = tools.NewWikiWritePageTool(s.wikiPageService, wikiKBIDs, s.knowledgeService, wikiRoutes). + WithSearchTargets(config.SearchTargets) case tools.ToolWikiReplaceText: - toolToRegister = tools.NewWikiReplaceTextTool(s.wikiPageService, wikiKBIDs, s.knowledgeService) + toolToRegister = tools.NewWikiReplaceTextTool(s.wikiPageService, wikiKBIDs, s.knowledgeService, wikiRoutes). + WithSearchTargets(config.SearchTargets) case tools.ToolWikiRenamePage: - toolToRegister = tools.NewWikiRenamePageTool(s.wikiPageService, wikiKBIDs) + toolToRegister = tools.NewWikiRenamePageTool(s.wikiPageService, wikiKBIDs, wikiRoutes) case tools.ToolWikiDeletePage: - toolToRegister = tools.NewWikiDeletePageTool(s.wikiPageService, wikiKBIDs) + toolToRegister = tools.NewWikiDeletePageTool(s.wikiPageService, wikiKBIDs, wikiRoutes) default: logger.Warnf(ctx, "Unknown tool: %s", toolName) diff --git a/internal/application/service/chat_pipeline/chat_completion.go b/internal/application/service/chat_pipeline/chat_completion.go index f88f452ff..5cbb39d15 100644 --- a/internal/application/service/chat_pipeline/chat_completion.go +++ b/internal/application/service/chat_pipeline/chat_completion.go @@ -3,7 +3,6 @@ package chatpipeline import ( "context" - "github.com/Tencent/WeKnora/internal/llmresource" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" ) @@ -51,10 +50,8 @@ func (p *PluginChatCompletion) OnEvent( pipelineInfo(ctx, "Completion", "messages_ready", map[string]interface{}{ "message_count": len(chatManage.History) + 2, }) - chatMessages, sourceRefs := prepareMessagesWithReferences(ctx, chatManage) - resourceRefs := llmresource.NewRegistry() - chatMessages = sourceRefs.EncodeMessages(chatMessages) - chatMessages = resourceRefs.EncodeMessages(chatMessages) + chatMessages, modelContext := prepareMessagesWithModelContext(ctx, chatManage) + chatMessages = modelContext.EncodeMessages(chatMessages) ctx = withPromptCacheMetadata(ctx, chatModel, chatMessages, opt, "knowledge_qa") // Call the chat model to generate response @@ -69,12 +66,11 @@ func (p *PluginChatCompletion) OnEvent( }) return ErrModelCall.WithError(err) } - resourceRefs.DecodeResponse(chatResponse) - sourceRefs.ExpandResponse(chatResponse) - if orphans := resourceRefs.OrphanAliases(chatResponse.Content); len(orphans) > 0 { - pipelineWarn(ctx, "Completion", "orphan_resource_aliases", map[string]interface{}{ + modelContext.DecodeResponse(chatResponse) + if orphans := modelContext.OrphanResourceHandles(chatResponse.Content); len(orphans) > 0 { + pipelineWarn(ctx, "Completion", "orphan_resource_handles", map[string]interface{}{ "session_id": chatManage.SessionID, - "aliases": orphans, + "handles": orphans, }) } diff --git a/internal/application/service/chat_pipeline/chat_completion_stream.go b/internal/application/service/chat_pipeline/chat_completion_stream.go index edd549c11..767695513 100644 --- a/internal/application/service/chat_pipeline/chat_completion_stream.go +++ b/internal/application/service/chat_pipeline/chat_completion_stream.go @@ -6,8 +6,6 @@ import ( "fmt" "github.com/Tencent/WeKnora/internal/event" - "github.com/Tencent/WeKnora/internal/llmreference" - "github.com/Tencent/WeKnora/internal/llmresource" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/google/uuid" @@ -56,10 +54,8 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context, // Prepare base messages without history - chatMessages, sourceRefs := prepareMessagesWithReferences(ctx, chatManage) - resourceRefs := llmresource.NewRegistry() - chatMessages = sourceRefs.EncodeMessages(chatMessages) - chatMessages = resourceRefs.EncodeMessages(chatMessages) + chatMessages, modelContext := prepareMessagesWithModelContext(ctx, chatManage) + chatMessages = modelContext.EncodeMessages(chatMessages) ctx = withPromptCacheMetadata(ctx, chatModel, chatMessages, opt, "knowledge_qa") pipelineInfo(ctx, "Stream", "messages_ready", map[string]interface{}{ "message_count": len(chatMessages), @@ -111,10 +107,8 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context, // The goroutine monitors ctx.Done() to avoid leaking when the context is cancelled // and the upstream channel is not closed promptly. go func() { - answerDecoder := llmresource.NewStreamDecoder(resourceRefs) - thinkingDecoder := llmresource.NewStreamDecoder(resourceRefs) - answerRefExpander := llmreference.NewStreamExpander(sourceRefs) - thinkingRefExpander := llmreference.NewStreamExpander(sourceRefs) + answerDecoder := modelContext.StreamDecoder() + thinkingDecoder := modelContext.StreamDecoder() thinkingID := fmt.Sprintf("%s-thinking", uuid.New().String()[:8]) answerID := fmt.Sprintf("%s-answer", uuid.New().String()[:8]) thinkingOpen := false @@ -134,13 +128,13 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context, thinkingOpen = false } - // flushDecoders drains any alias suffix the stream decoders held back to + // flushDecoders drains any handle suffix the stream decoders held back to // bridge references split across provider chunks. Both the normal close // and the cancellation path must call this, otherwise a resource // reference in flight at teardown is silently dropped (and never // persisted, since the assistant message is saved from these events). flushDecoders := func() { - thinkingTail := thinkingRefExpander.Feed(thinkingDecoder.Flush()) + thinkingRefExpander.Flush() + thinkingTail := thinkingDecoder.Flush() if thinkingTail != "" { _ = eventBus.Emit(ctx, types.Event{ ID: thinkingID, @@ -149,7 +143,7 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context, Data: event.AgentThoughtData{Content: thinkingTail}, }) } - answerTail := answerRefExpander.Feed(answerDecoder.Flush()) + answerRefExpander.Flush() + answerTail := answerDecoder.Flush() if answerTail != "" { _ = eventBus.Emit(ctx, types.Event{ ID: answerID, @@ -199,7 +193,10 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context, } if response.ResponseType == types.ResponseTypeThinking { - response.Content = thinkingRefExpander.Feed(thinkingDecoder.Feed(response.Content)) + response.Content = thinkingDecoder.Feed(response.Content) + if response.Done { + response.Content += thinkingDecoder.Flush() + } if response.Content != "" { thinkingOpen = true eventBus.Emit(ctx, types.Event{ @@ -219,7 +216,10 @@ func (p *PluginChatCompletionStream) OnEvent(ctx context.Context, } if response.ResponseType == types.ResponseTypeAnswer { - response.Content = answerRefExpander.Feed(answerDecoder.Feed(response.Content)) + response.Content = answerDecoder.Feed(response.Content) + if response.Done { + response.Content += answerDecoder.Flush() + } closeThinking() eventBus.Emit(ctx, types.Event{ ID: answerID, diff --git a/internal/application/service/chat_pipeline/chat_completion_stream_test.go b/internal/application/service/chat_pipeline/chat_completion_stream_test.go index ba985385b..49e095253 100644 --- a/internal/application/service/chat_pipeline/chat_completion_stream_test.go +++ b/internal/application/service/chat_pipeline/chat_completion_stream_test.go @@ -78,10 +78,9 @@ func (s *stubModelService) GetChatModel(context.Context, string) (chat.Chat, err return s.model, nil } -// TestStreamFlushesHeldAliasOnCancel verifies that when the request is cancelled -// mid-stream, the decoder's held-back alias suffix is flushed (emitted) rather -// than silently dropped. Without the ctx.Done() flush, "res://0" would be lost. -func TestStreamFlushesHeldAliasOnCancel(t *testing.T) { +// TestStreamDropsIncompleteHandleOnCancel verifies that cancellation cannot +// leak a model-context protocol fragment to the client. +func TestStreamDropsIncompleteHandleOnCancel(t *testing.T) { const ref = "resource://AbCdEfGhIjKlMnOpQrStUv" bus := &syncEventBus{} model := &openStreamChat{chunks: []types.StreamResponse{ @@ -111,13 +110,9 @@ func TestStreamFlushesHeldAliasOnCancel(t *testing.T) { cancel() - // After cancel, the held "res://0" suffix must be flushed as a final-answer chunk. - require.Eventually(t, func() bool { - for _, c := range bus.finalAnswerContents() { - if c == "res://0" { - return true - } - } - return false - }, 2*time.Second, 5*time.Millisecond) + // Give the cancellation path time to flush, then assert that only the + // meaningful prefix was emitted; the incomplete private handle is dropped. + require.Eventually(t, func() bool { return len(bus.finalAnswerContents()) >= 1 }, 2*time.Second, 5*time.Millisecond) + time.Sleep(20 * time.Millisecond) + require.Equal(t, []string{"hello "}, bus.finalAnswerContents()) } diff --git a/internal/application/service/chat_pipeline/query_expansion.go b/internal/application/service/chat_pipeline/query_expansion.go index 9e38f8aac..529604654 100644 --- a/internal/application/service/chat_pipeline/query_expansion.go +++ b/internal/application/service/chat_pipeline/query_expansion.go @@ -47,6 +47,9 @@ func (p *PluginSearch) runQueryExpansion(ctx context.Context, chatManage *types. }) for _, q := range expansions { for _, target := range chatManage.SearchTargets { + if target == nil || target.KnowledgeBaseID == "" { + continue + } wgExp.Add(1) go func(q string, t *types.SearchTarget) { defer wgExp.Done() diff --git a/internal/application/service/chat_pipeline/references.go b/internal/application/service/chat_pipeline/references.go index 182d640eb..92af016e8 100644 --- a/internal/application/service/chat_pipeline/references.go +++ b/internal/application/service/chat_pipeline/references.go @@ -4,30 +4,29 @@ import ( "context" "strings" - "github.com/Tencent/WeKnora/internal/llmreference" + "github.com/Tencent/WeKnora/internal/modelcontext" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" ) -// prepareMessagesWithReferences replaces the pipeline's positional context IDs -// with request-local chunk references in the copy sent to the model. Persisted -// rendered_content remains unchanged; public citations are expanded only when -// the request's agent setting enables them. -func prepareMessagesWithReferences( +// prepareMessagesWithModelContext replaces positional retrieval IDs with +// request-local model handles. Persisted rendered_content remains unchanged; +// public citations are expanded only when the request setting enables them. +func prepareMessagesWithModelContext( ctx context.Context, chatManage *types.ChatManage, -) ([]chat.Message, *llmreference.Registry) { +) ([]chat.Message, *modelcontext.Registry) { citationsEnabled := chatManage == nil || chatManage.CitationsEnabled() - refs := llmreference.NewRegistry(citationsEnabled) + registry := modelcontext.NewRegistry(citationsEnabled) if chatManage == nil { - return nil, refs + return nil, registry } messages := prepareMessagesWithHistory(chatManage) if len(messages) > 0 { - messages[0].Content = strings.TrimRight(messages[0].Content, " \t\r\n") + llmreference.ProtocolPrompt(citationsEnabled) + messages[0].Content = strings.TrimRight(messages[0].Content, " \t\r\n") + registry.ProtocolPrompt() } if len(chatManage.MergeResult) == 0 || len(messages) == 0 { - return messages, refs + return messages, registry } ordered := orderedPipelineReferences(chatManage) @@ -58,10 +57,10 @@ func prepareMessagesWithReferences( "content": getEnrichedPassageForChat(ctx, result), }) } - refs.RegisterSearchResults(knowledgeResults) + registry.RegisterSearchResults(knowledgeResults) var contextParts []string if len(knowledgeRows) > 0 { - contextParts = append(contextParts, refs.ModelOutput(&types.ToolResult{ + contextParts = append(contextParts, registry.ModelToolResult(&types.ToolResult{ Success: true, Data: map[string]interface{}{ "display_type": "search_results", @@ -70,7 +69,7 @@ func prepareMessagesWithReferences( })) } if len(webRows) > 0 { - contextParts = append(contextParts, refs.ModelOutput(&types.ToolResult{ + contextParts = append(contextParts, registry.ModelToolResult(&types.ToolResult{ Success: true, Data: map[string]interface{}{ "display_type": "web_search_results", @@ -80,7 +79,7 @@ func prepareMessagesWithReferences( } modelContexts := strings.Join(contextParts, "\n") if strings.TrimSpace(modelContexts) == "" { - return messages, refs + return messages, registry } last := len(messages) - 1 @@ -94,7 +93,7 @@ func prepareMessagesWithReferences( if !replaced { messages[last].Content = modelContexts + "\n\n" + messages[last].Content } - return messages, refs + return messages, registry } func isPipelineWebReference(result *types.SearchResult) bool { diff --git a/internal/application/service/chat_pipeline/references_test.go b/internal/application/service/chat_pipeline/references_test.go index e7e18435b..78140e36e 100644 --- a/internal/application/service/chat_pipeline/references_test.go +++ b/internal/application/service/chat_pipeline/references_test.go @@ -9,7 +9,7 @@ import ( "github.com/stretchr/testify/require" ) -func TestPrepareMessagesWithReferencesUsesChunkCentricContext(t *testing.T) { +func TestPrepareMessagesWithModelContextUsesChunkCentricContext(t *testing.T) { rendered := `first contentsecond content` manage := &types.ChatManage{ PipelineRequest: types.PipelineRequest{ @@ -28,7 +28,7 @@ func TestPrepareMessagesWithReferencesUsesChunkCentricContext(t *testing.T) { }, } - messages, refs := prepareMessagesWithReferences(context.Background(), manage) + messages, refs := prepareMessagesWithModelContext(context.Background(), manage) require.Len(t, messages, 2) require.Contains(t, messages[0].Content, "Source handling protocol") require.Contains(t, messages[1].Content, ``) @@ -37,11 +37,11 @@ func TestPrepareMessagesWithReferencesUsesChunkCentricContext(t *testing.T) { require.False(t, strings.Contains(messages[1].Content, "chunk-1")) require.Equal(t, ``, - refs.ExpandText(``), + refs.DecodeOutputText(``), ) } -func TestPrepareMessagesWithReferencesReplacesSystemPromptContextAndHistoryCitations(t *testing.T) { +func TestPrepareMessagesWithModelContextReplacesSystemPromptContextAndHistoryCitations(t *testing.T) { rendered := `first content` manage := &types.ChatManage{ PipelineRequest: types.PipelineRequest{ @@ -63,7 +63,7 @@ func TestPrepareMessagesWithReferencesReplacesSystemPromptContextAndHistoryCitat }, } - messages, refs := prepareMessagesWithReferences(context.Background(), manage) + messages, refs := prepareMessagesWithModelContext(context.Background(), manage) messages = refs.EncodeMessages(messages) require.NotContains(t, messages[0].Content, rendered) require.Contains(t, messages[0].Content, ``) require.Contains(t, messages[1].Content, ``) require.NotContains(t, messages[1].Content, ``, - refs.ExpandText(``), + refs.DecodeOutputText(``), ) } -func TestPrepareMessagesWithReferencesCompactsHistoryWithoutCurrentRetrieval(t *testing.T) { +func TestPrepareMessagesWithModelContextCompactsHistoryWithoutCurrentRetrieval(t *testing.T) { manage := &types.ChatManage{ PipelineRequest: types.PipelineRequest{Query: "follow-up", SummaryConfig: types.SummaryConfig{Prompt: "system"}}, PipelineState: types.PipelineState{ @@ -108,14 +108,14 @@ func TestPrepareMessagesWithReferencesCompactsHistoryWithoutCurrentRetrieval(t * }, } - messages, refs := prepareMessagesWithReferences(context.Background(), manage) + messages, refs := prepareMessagesWithModelContext(context.Background(), manage) messages = refs.EncodeMessages(messages) require.Contains(t, messages[0].Content, "Source handling protocol") require.Contains(t, messages[2].Content, ``) require.NotContains(t, messages[2].Content, "https://example.com/old") } -func TestPrepareMessagesWithReferencesSuppressesCitationsWhenDisabled(t *testing.T) { +func TestPrepareMessagesWithModelContextSuppressesCitationsWhenDisabled(t *testing.T) { disabled := false manage := &types.ChatManage{ PipelineRequest: types.PipelineRequest{ @@ -131,9 +131,9 @@ func TestPrepareMessagesWithReferencesSuppressesCitationsWhenDisabled(t *testing }, } - messages, refs := prepareMessagesWithReferences(context.Background(), manage) + messages, refs := prepareMessagesWithModelContext(context.Background(), manage) require.Contains(t, messages[0].Content, "Source citations are disabled") require.Contains(t, messages[1].Content, ``)) + require.Equal(t, "answer ", refs.DecodeOutputText(`answer `)) } diff --git a/internal/application/service/chat_pipeline/wiki_boost.go b/internal/application/service/chat_pipeline/wiki_boost.go index edd6c93d3..fd0dbcc51 100644 --- a/internal/application/service/chat_pipeline/wiki_boost.go +++ b/internal/application/service/chat_pipeline/wiki_boost.go @@ -66,6 +66,9 @@ func (p *PluginWikiBoost) OnEvent( // Confirm at least one search target is actually a wiki KB. hasWikiKB := false for _, target := range chatManage.SearchTargets { + if target == nil || target.KnowledgeBaseID == "" { + continue + } kb, err := p.kbService.GetKnowledgeBaseByIDOnly(ctx, target.KnowledgeBaseID) if err == nil && kb != nil && kb.IsWikiEnabled() { hasWikiKB = true diff --git a/internal/application/service/fallback_messages_test.go b/internal/application/service/fallback_messages_test.go index 239283a40..bbe685dd7 100644 --- a/internal/application/service/fallback_messages_test.go +++ b/internal/application/service/fallback_messages_test.go @@ -93,7 +93,7 @@ func TestPrepareFallbackMessagesMigratesHistoricalCitations(t *testing.T) { require.Equal(t, `Previous `, messages[2].Content) require.Equal(t, ``, - refs.ExpandText(``), + refs.DecodeOutputText(``), ) } @@ -109,5 +109,5 @@ func TestPrepareFallbackMessagesSuppressesCitationsWhenDisabled(t *testing.T) { messages, refs := prepareFallbackMessages(cm, "legacy fallback prompt") require.Contains(t, messages[0].Content, "Source citations are disabled") require.Equal(t, `Previous `, messages[2].Content) - require.Equal(t, "answer ", refs.ExpandText(`answer `)) + require.Equal(t, "answer ", refs.DecodeOutputText(`answer `)) } diff --git a/internal/application/service/session_knowledge_qa.go b/internal/application/service/session_knowledge_qa.go index 5f8d62f16..938f1048f 100644 --- a/internal/application/service/session_knowledge_qa.go +++ b/internal/application/service/session_knowledge_qa.go @@ -10,8 +10,8 @@ import ( chatpipeline "github.com/Tencent/WeKnora/internal/application/service/chat_pipeline" "github.com/Tencent/WeKnora/internal/common" "github.com/Tencent/WeKnora/internal/event" - "github.com/Tencent/WeKnora/internal/llmreference" "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/modelcontext" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/tracing/langfuse" "github.com/Tencent/WeKnora/internal/types" @@ -972,7 +972,7 @@ func (s *sessionService) handleModelFallback(ctx context.Context, chatManage *ty } // Start streaming response - fallbackMessages, sourceRefs := prepareFallbackMessages(chatManage, promptContent) + fallbackMessages, modelContext := prepareFallbackMessages(chatManage, promptContent) responseChan, err := chatModel.ChatStream(ctx, fallbackMessages, opt) if err != nil { logger.Errorf(ctx, "Failed to start streaming fallback response: %v, falling back to fixed response", err) @@ -987,22 +987,22 @@ func (s *sessionService) handleModelFallback(ctx context.Context, chatManage *ty } // Start goroutine to consume stream and emit events - go s.consumeFallbackStream(ctx, chatManage, responseChan, sourceRefs) + go s.consumeFallbackStream(ctx, chatManage, responseChan, modelContext) } func prepareFallbackMessages( chatManage *types.ChatManage, promptContent string, -) ([]chat.Message, *llmreference.Registry) { +) ([]chat.Message, *modelcontext.Registry) { messages := buildFallbackMessages(chatManage, promptContent) citationsEnabled := chatManage == nil || chatManage.CitationsEnabled() - refs := llmreference.NewRegistry(citationsEnabled) + registry := modelcontext.NewRegistry(citationsEnabled) if len(messages) > 0 && messages[0].Role == "system" { - messages[0].Content = strings.TrimRight(messages[0].Content, " \t\r\n") + llmreference.ProtocolPrompt(citationsEnabled) + messages[0].Content = strings.TrimRight(messages[0].Content, " \t\r\n") + registry.ProtocolPrompt() } else { - messages = append([]chat.Message{{Role: "system", Content: strings.TrimSpace(llmreference.ProtocolPrompt(citationsEnabled))}}, messages...) + messages = append([]chat.Message{{Role: "system", Content: strings.TrimSpace(registry.ProtocolPrompt())}}, messages...) } - return refs.EncodeMessages(messages), refs + return registry.EncodeMessages(messages), registry } func buildFallbackMessages(chatManage *types.ChatManage, promptContent string) []chat.Message { @@ -1135,20 +1135,20 @@ func (s *sessionService) consumeFallbackStream( ctx context.Context, chatManage *types.ChatManage, responseChan <-chan types.StreamResponse, - sourceRefs *llmreference.Registry, + modelContext *modelcontext.Registry, ) { fallbackID := generateEventID("fallback") eventBus := chatManage.EventBus var finalContent string streamCompleted := false - refExpander := llmreference.NewStreamExpander(sourceRefs) + decoder := modelContext.StreamDecoder() for response := range responseChan { // Emit event for each answer chunk if response.ResponseType == types.ResponseTypeAnswer { - response.Content = refExpander.Feed(response.Content) + response.Content = decoder.Feed(response.Content) if response.Done { - response.Content += refExpander.Flush() + response.Content += decoder.Flush() } finalContent += response.Content if err := eventBus.Emit(ctx, types.Event{ @@ -1209,7 +1209,7 @@ func (s *sessionService) emitFallbackAnswer(ctx context.Context, chatManage *typ return } if !chatManage.CitationsEnabled() { - content = llmreference.NewRegistry(false).ExpandText(content) + content = modelcontext.NewRegistry(false).DecodeOutputText(content) } fallbackID := generateEventID("fallback") diff --git a/internal/application/service/wiki_ingest_batch.go b/internal/application/service/wiki_ingest_batch.go index e458b365c..e1f9bd6d8 100644 --- a/internal/application/service/wiki_ingest_batch.go +++ b/internal/application/service/wiki_ingest_batch.go @@ -2004,26 +2004,26 @@ func (s *wikiIngestService) reduceSlugUpdates( if len(additions) > 0 || len(retracts) > 0 { titles := batchCtx.SlugTitleMany(ctx, []string(page.OutLinks)) - // aliaser escapes high-entropy slugs behind short reference tokens + // slugHandles escape high-entropy slugs behind short reference handles // (ref-1, ref-2, …) for the editor LLM, then translates them back to - // real slugs after generation (see dealiasContent below). - aliaser := newSlugAliaser() + // real slugs after generation (see decodeContent below). + slugHandles := newWikiSlugHandles() - // Escape every out-link slug behind a reference token so the editor + // Escape every out-link slug behind a request-local handle so the editor // never has to retype a real slug — this is where UUID-based summary // slugs (summary/) got mangled into 404-ing links. Both the // listing AND the [[...]] refs already present in - // are aliased with the SAME token map, so the + // use the SAME handle table, so the // model sees one consistent, copy-safe identifier space; we translate - // the tokens back to real slugs after generation. + // the handles back to real slugs after generation. known := make(map[string]struct{}, len(page.OutLinks)) for _, outSlug := range page.OutLinks { known[outSlug] = struct{}{} - aliaser.token(outSlug) // pre-assign for stable ordering + slugHandles.handle(outSlug) // pre-assign for stable ordering } for _, outSlug := range page.OutLinks { if title := titles[outSlug]; title != "" { - fmt.Fprintf(&relatedSlugs, "- %s (%s)\n", aliaser.token(outSlug), title) + fmt.Fprintf(&relatedSlugs, "- %s (%s)\n", slugHandles.handle(outSlug), title) } } @@ -2031,7 +2031,7 @@ func (s *wikiIngestService) reduceSlugUpdates( // [c003]. They are internal ingest metadata; keep the editor context // clean so a subsequent update cannot copy them into rewritten prose. existingContent := stripWikiInlineChunkCitations(page.Content) - existingContent = aliaser.aliasContent(existingContent, known) + existingContent = slugHandles.encodeContent(existingContent, known) if !exists || existingContent == "" { existingContent = "(New page)" } @@ -2079,10 +2079,10 @@ func (s *wikiIngestService) reduceSlugUpdates( }) if err == nil && updatedContent != "" { - // Translate reference tokens (ref-N) the model copied from the - // aliased context back to their real slugs BEFORE the content is + // Translate request-local handles (ref-N) the model copied from the + // encoded context back to their real slugs BEFORE the content is // parsed/stored, so out_links reflect real pages again. - updatedContent = aliaser.dealiasContent(updatedContent) + updatedContent = slugHandles.decodeContent(updatedContent) updatedSummary, updatedBody := splitSummaryLine(updatedContent) if updatedBody != "" { page.Content = updatedBody diff --git a/internal/application/service/wiki_ingest_cite.go b/internal/application/service/wiki_ingest_cite.go index ea112ef42..cf2f5fee5 100644 --- a/internal/application/service/wiki_ingest_cite.go +++ b/internal/application/service/wiki_ingest_cite.go @@ -10,6 +10,7 @@ import ( "github.com/Tencent/WeKnora/internal/agent" "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/modelcontext" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" "golang.org/x/sync/errgroup" @@ -135,7 +136,7 @@ func (s *wikiIngestService) extractCandidateSlugs( // chunkBatch groups chunks that will be sent in a single WikiChunkCitationPrompt call. type chunkBatch struct { chunks []*types.Chunk - aliasToID map[string]string + handles *modelcontext.HandleTable totalRuneLen int } @@ -143,9 +144,9 @@ type chunkBatch struct { // rune count stays under maxRunesPerCitationBatch. Chunk order (by ChunkIndex) // is preserved, and a chunk that by itself exceeds the budget occupies its // own batch so we never silently drop content. Each chunk is assigned a short -// alias ("c000", "c001", ...) that the prompt uses in place of the raw UUID; -// the alias → real ID map is kept on the batch so the caller can translate -// LLM responses back to stable chunk IDs. +// handle ("c000", "c001", ...) that the prompt uses in place of the raw UUID. +// The invocation-local table translates model output back to stable chunk IDs +// before any result reaches application state. func splitChunksIntoCitationBatches(chunks []*types.Chunk) []chunkBatch { // Only cite text chunks — image/ocr chunks are already merged into the // text content via reconstructEnrichedContent and the LLM doesn't see @@ -173,15 +174,17 @@ func splitChunksIntoCitationBatches(chunks []*types.Chunk) []chunkBatch { }) var batches []chunkBatch - current := chunkBatch{aliasToID: make(map[string]string)} - aliasCounter := 0 + newBatch := func() chunkBatch { + return chunkBatch{handles: modelcontext.NewHandleTable("c", 3, 0)} + } + current := newBatch() flush := func() { if len(current.chunks) == 0 { return } batches = append(batches, current) - current = chunkBatch{aliasToID: make(map[string]string)} + current = newBatch() } for _, c := range filtered { @@ -194,26 +197,12 @@ func splitChunksIntoCitationBatches(chunks []*types.Chunk) []chunkBatch { flush() } - alias := fmt.Sprintf("c%03d", aliasCounter) - aliasCounter++ - current.aliasToID[alias] = c.ID + current.handles.Register(c.ID) current.chunks = append(current.chunks, c) current.totalRuneLen += runeLen } flush() - // Restart alias numbering per batch so the LLM sees a small, local ID - // space (helps it follow exact-match citation rules). We do this by - // rebuilding the alias map in a second pass. - for bi := range batches { - newAliasToID := make(map[string]string, len(batches[bi].chunks)) - for idx, c := range batches[bi].chunks { - alias := fmt.Sprintf("c%03d", idx) - newAliasToID[alias] = c.ID - } - batches[bi].aliasToID = newAliasToID - } - return batches } @@ -245,23 +234,12 @@ func renderCandidateSlugsXML(entities, concepts []extractedItem) string { } // renderChunksXML formats one batch's chunks into the block, using -// per-batch aliases (c000, c001, ...) instead of raw UUIDs. +// per-batch handles (c000, c001, ...) instead of raw UUIDs. func renderChunksXML(batch chunkBatch) string { - aliases := make([]string, 0, len(batch.aliasToID)) - for a := range batch.aliasToID { - aliases = append(aliases, a) - } - sort.Strings(aliases) - - idToAlias := make(map[string]string, len(batch.aliasToID)) - for a, id := range batch.aliasToID { - idToAlias[id] = a - } - var sb strings.Builder for _, c := range batch.chunks { - alias := idToAlias[c.ID] - fmt.Fprintf(&sb, "\n%s\n\n", alias, c.ChunkIndex, c.Content) + handle, _ := batch.handles.Handle(c.ID) + fmt.Fprintf(&sb, "\n%s\n\n", handle, c.ChunkIndex, c.Content) } return sb.String() } @@ -273,7 +251,7 @@ func renderChunksXML(batch chunkBatch) string { // "new_slugs" that Pass 0 missed are collected separately. // // Returns (citations, newSlugs, batchCount). citations is keyed by slug and -// contains real chunk UUIDs (already translated from batch aliases). newSlugs +// contains real chunk UUIDs (already translated from batch handles). newSlugs // likewise carry real chunk UUIDs in SourceChunks. func (s *wikiIngestService) classifyChunkCitations( ctx context.Context, @@ -319,11 +297,11 @@ func (s *wikiIngestService) classifyChunkCitations( return nil } - // Translate aliases → real chunk UUIDs; drop unknown aliases. + // Translate handles → real chunk UUIDs; drop unknown handles. mu.Lock() defer mu.Unlock() - for slug, aliasList := range parsed.Citations { + for slug, handleList := range parsed.Citations { if slug == "" { continue } @@ -332,10 +310,10 @@ func (s *wikiIngestService) classifyChunkCitations( set = make(map[string]bool) citationSet[slug] = set } - for _, alias := range aliasList { - realID, known := batch.aliasToID[alias] + for _, handle := range handleList { + realID, known := batch.handles.Resolve(handle) if !known { - logger.Warnf(ectx, "wiki ingest: citation batch %d referenced unknown chunk alias %q for slug %s", batchIdx, alias, slug) + logger.Warnf(ectx, "wiki ingest: citation batch %d referenced unknown chunk handle %q for slug %s", batchIdx, handle, slug) continue } set[realID] = true @@ -347,8 +325,8 @@ func (s *wikiIngestService) classifyChunkCitations( continue } real := make([]string, 0, len(ns.SourceChunks)) - for _, alias := range ns.SourceChunks { - if id, ok := batch.aliasToID[alias]; ok { + for _, handle := range ns.SourceChunks { + if id, ok := batch.handles.Resolve(handle); ok { real = append(real, id) } } diff --git a/internal/application/service/wiki_ingest_cite_test.go b/internal/application/service/wiki_ingest_cite_test.go index 453b3a441..78076ff21 100644 --- a/internal/application/service/wiki_ingest_cite_test.go +++ b/internal/application/service/wiki_ingest_cite_test.go @@ -156,10 +156,10 @@ func TestSplitChunksIntoCitationBatches_RespectsBudgetAndOrder(t *testing.T) { t.Errorf("batch order = %v, want %v", seen, wantOrder) } - // Verify alias → id map is populated per batch with unique aliases. + // Verify the typed handle table is populated per batch. for bi, b := range batches { - if len(b.aliasToID) != len(b.chunks) { - t.Errorf("batch %d alias count %d != chunk count %d", bi, len(b.aliasToID), len(b.chunks)) + if b.handles.Len() != len(b.chunks) { + t.Errorf("batch %d handle count %d != chunk count %d", bi, b.handles.Len(), len(b.chunks)) } } } diff --git a/internal/application/service/wiki_page.go b/internal/application/service/wiki_page.go index 3657defde..39aa592b8 100644 --- a/internal/application/service/wiki_page.go +++ b/internal/application/service/wiki_page.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "regexp" + "slices" "sort" "strconv" "strings" @@ -21,9 +22,9 @@ import ( // wikiLinkRegex matches [[wiki-link]] syntax in markdown content var wikiLinkRegex = regexp.MustCompile(`\[\[([^\]]+)\]\]`) -// wikiInlineChunkCitationRegex matches internal short aliases emitted by the +// wikiInlineChunkCitationRegex matches internal short handles emitted by the // wiki ingest prompt while it classifies supporting chunks. The stable source -// relationship lives in WikiPage.ChunkRefs, so these aliases have no meaning +// relationship lives in WikiPage.ChunkRefs, so these handles have no meaning // to readers and must not leak into generated Markdown. var wikiInlineChunkCitationRegex = regexp.MustCompile(`[ \t]*\[c\d{3,}(?:\s*[,;]\s*c\d{3,})*\]`) @@ -130,12 +131,14 @@ func (s *wikiPageService) UpdatePage(ctx context.Context, page *types.WikiPage) existing.Content != page.Content || existing.Summary != page.Summary || existing.PageType != page.PageType || - existing.Status != page.Status + existing.Status != page.Status || + !slices.Equal(existing.Aliases, page.Aliases) existing.Title = page.Title existing.Content = page.Content existing.Summary = page.Summary existing.PageType = page.PageType + existing.Aliases = append(types.StringArray(nil), page.Aliases...) existing.SourceRefs = page.SourceRefs existing.ChunkRefs = page.ChunkRefs existing.PageMetadata = page.PageMetadata diff --git a/internal/application/service/wiki_page_test.go b/internal/application/service/wiki_page_test.go index 948dcaa8c..5c96d9cfd 100644 --- a/internal/application/service/wiki_page_test.go +++ b/internal/application/service/wiki_page_test.go @@ -70,6 +70,38 @@ func TestStripWikiInlineChunkCitationsPreservesOrdinaryMarkdown(t *testing.T) { } } +func TestUpdateWikiPagePersistsAndClearsAliases(t *testing.T) { + db, err := gorm.Open(sqlite.Open(fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name())), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&types.WikiFolder{}, &types.WikiPage{})) + + ctx := context.Background() + repo := repository.NewWikiPageRepository(db) + svc := NewWikiPageService(repo, nil, nil, nil, nil) + page, err := svc.CreatePage(ctx, &types.WikiPage{ + TenantID: 1, KnowledgeBaseID: "kb-alias", Slug: "concept/alias", + Title: "Alias", Summary: "summary", Content: "content", + PageType: types.WikiPageTypeConcept, Aliases: types.StringArray{"old"}, + }) + require.NoError(t, err) + + page.Aliases = types.StringArray{"new", "alternate"} + updated, err := svc.UpdatePage(ctx, page) + require.NoError(t, err) + require.Equal(t, types.StringArray{"new", "alternate"}, updated.Aliases) + require.Equal(t, 2, updated.Version) + + updated.Aliases = types.StringArray{} + cleared, err := svc.UpdatePage(ctx, updated) + require.NoError(t, err) + require.Empty(t, cleared.Aliases) + require.Equal(t, 3, cleared.Version) + + stored, err := svc.GetPageBySlug(ctx, "kb-alias", "concept/alias") + require.NoError(t, err) + require.Empty(t, stored.Aliases) +} + func TestParseOutLinks(t *testing.T) { svc := &wikiPageService{} diff --git a/internal/application/service/wiki_slug_alias.go b/internal/application/service/wiki_slug_alias.go deleted file mode 100644 index 28d35223f..000000000 --- a/internal/application/service/wiki_slug_alias.go +++ /dev/null @@ -1,88 +0,0 @@ -package service - -import ( - "fmt" - "strings" -) - -// slugAliaser maps real wiki slugs to short, low-entropy reference tokens -// (ref-1, ref-2, …) so the ingest editor LLM never has to reproduce a -// high-entropy slug verbatim — most importantly the UUID-based summary slugs -// (summary/) that models routinely mangle by inserting or -// dropping hex digits. The model copies the tiny token; we translate tokens -// back to real slugs on output. -// -// This is the generation-stage counterpart to wiki_write_page's slug -// validation: rather than repairing a mangled slug after the fact, it removes -// the opportunity to mangle at the source. It mirrors the chunk-alias -// (c000/c001) indirection already used by the chunk-citation pass in -// wiki_ingest_cite.go. -// -// A slugAliaser is single-use and NOT safe for concurrent use; construct one -// per LLM call. -type slugAliaser struct { - toToken map[string]string // real slug -> token - toSlug map[string]string // token -> real slug -} - -func newSlugAliaser() *slugAliaser { - return &slugAliaser{ - toToken: make(map[string]string), - toSlug: make(map[string]string), - } -} - -// token returns the stable token for a real slug, assigning a fresh one on -// first use. An empty slug returns "". Tokens have no '/' so they can never -// collide with a namespaced real slug (entity/…, summary/…). -func (a *slugAliaser) token(realSlug string) string { - realSlug = strings.TrimSpace(realSlug) - if realSlug == "" { - return "" - } - if t, ok := a.toToken[realSlug]; ok { - return t - } - t := fmt.Sprintf("ref-%d", len(a.toToken)+1) - a.toToken[realSlug] = t - a.toSlug[t] = realSlug - return t -} - -// empty reports whether no tokens have been assigned yet (nothing to alias). -func (a *slugAliaser) empty() bool { return len(a.toToken) == 0 } - -// aliasContent rewrites [[realSlug|disp]] / [[realSlug]] occurrences to their -// token form, but only for slugs present in `known`. Unknown links are left -// untouched. Assigns tokens on demand so a slug that appears only in the body -// (not in the listing) still gets a consistent token. -func (a *slugAliaser) aliasContent(content string, known map[string]struct{}) string { - if content == "" || len(known) == 0 { - return content - } - out, _ := rewriteDeadWikiLinks(content, func(norm, _ string) (string, bool) { - if _, ok := known[norm]; !ok { - return "", false - } - return a.token(norm), true - }) - return out -} - -// dealiasContent rewrites [[token|disp]] / [[token]] occurrences back to their -// real slug. Tokens with no mapping are left untouched — they fall through to -// the ordinary parse/dead-link-cleanup path, which is the correct behavior for -// anything the model invented that isn't a known reference. -func (a *slugAliaser) dealiasContent(content string) string { - if content == "" || a.empty() { - return content - } - out, _ := rewriteDeadWikiLinks(content, func(norm, _ string) (string, bool) { - real, ok := a.toSlug[norm] - if !ok { - return "", false - } - return real, true - }) - return out -} diff --git a/internal/application/service/wiki_slug_handles.go b/internal/application/service/wiki_slug_handles.go new file mode 100644 index 000000000..40743b2d6 --- /dev/null +++ b/internal/application/service/wiki_slug_handles.go @@ -0,0 +1,74 @@ +package service + +import "github.com/Tencent/WeKnora/internal/modelcontext" + +// wikiSlugHandles maps real wiki slugs to short, low-entropy model handles +// (ref-1, ref-2, …) so the ingest editor LLM never has to reproduce a +// high-entropy slug verbatim — most importantly the UUID-based summary slugs +// (summary/) that models routinely mangle by inserting or +// dropping hex digits. The model copies the tiny handle; we translate handles +// back to real slugs on output. +// +// This is the generation-stage counterpart to wiki_write_page's slug +// validation: rather than repairing a mangled slug after the fact, it removes +// the opportunity to mangle at the source. It mirrors the chunk-handle +// (c000/c001) indirection already used by the chunk-citation pass in +// wiki_ingest_cite.go. +// +// Handle allocation and reverse lookup are owned by modelcontext.HandleTable; +// this wrapper contains only the Wiki-link rewrite semantics. +type wikiSlugHandles struct { + table *modelcontext.HandleTable +} + +func newWikiSlugHandles() *wikiSlugHandles { + return &wikiSlugHandles{table: modelcontext.NewHandleTable("ref-", 0, 1)} +} + +// handle returns the stable handle for a real slug, assigning a fresh one on +// first use. An empty slug returns "". Handles have no '/' so they can never +// collide with a namespaced real slug (entity/…, summary/…). +func (h *wikiSlugHandles) handle(realSlug string) string { + if h == nil { + return "" + } + return h.table.Register(realSlug) +} + +// empty reports whether no handles have been assigned yet. +func (h *wikiSlugHandles) empty() bool { return h == nil || h.table.Empty() } + +// encodeContent rewrites [[realSlug|disp]] / [[realSlug]] occurrences to their +// handle form, but only for slugs present in `known`. Unknown links are left +// untouched. Assigns handles on demand so a slug that appears only in the body +// (not in the listing) still gets a consistent handle. +func (h *wikiSlugHandles) encodeContent(content string, known map[string]struct{}) string { + if content == "" || len(known) == 0 { + return content + } + out, _ := rewriteDeadWikiLinks(content, func(norm, _ string) (string, bool) { + if _, ok := known[norm]; !ok { + return "", false + } + return h.handle(norm), true + }) + return out +} + +// decodeContent rewrites [[handle|disp]] / [[handle]] occurrences back to their +// real slug. Handles with no mapping are left untouched — they fall through to +// the ordinary parse/dead-link-cleanup path, which is the correct behavior for +// anything the model invented that isn't a known reference. +func (h *wikiSlugHandles) decodeContent(content string) string { + if content == "" || h.empty() { + return content + } + out, _ := rewriteDeadWikiLinks(content, func(norm, _ string) (string, bool) { + real, ok := h.table.Resolve(norm) + if !ok { + return "", false + } + return real, true + }) + return out +} diff --git a/internal/application/service/wiki_slug_alias_test.go b/internal/application/service/wiki_slug_handles_test.go similarity index 69% rename from internal/application/service/wiki_slug_alias_test.go rename to internal/application/service/wiki_slug_handles_test.go index a3476d11f..774b1999c 100644 --- a/internal/application/service/wiki_slug_alias_test.go +++ b/internal/application/service/wiki_slug_handles_test.go @@ -5,15 +5,15 @@ import ( "testing" ) -func TestSlugAliaser_RoundTripEscapesUUIDSlugs(t *testing.T) { +func TestWikiSlugHandles_RoundTripEscapesUUIDSlugs(t *testing.T) { summarySlug := "summary/07a20bb1-a662-47cf-9929-06fb5d5b5b5e" entitySlug := "entity/mongodb" known := map[string]struct{}{summarySlug: {}, entitySlug: {}} - a := newSlugAliaser() + h := newWikiSlugHandles() // Pre-assign tokens the way the ingest listing loop does. - tSummary := a.token(summarySlug) - tEntity := a.token(entitySlug) + tSummary := h.handle(summarySlug) + tEntity := h.handle(entitySlug) if tSummary == summarySlug || tEntity == entitySlug { t.Fatalf("tokens must differ from real slugs: %q %q", tSummary, tEntity) } @@ -22,7 +22,7 @@ func TestSlugAliaser_RoundTripEscapesUUIDSlugs(t *testing.T) { } body := "See [[" + summarySlug + "|Weknora 试错记录.md - Summary]] and [[" + entitySlug + "]]." - aliased := a.aliasContent(body, known) + aliased := h.encodeContent(body, known) // The high-entropy UUID must no longer be present in what the model sees. if strings.Contains(aliased, "06fb5d5b5b5e") { @@ -36,37 +36,37 @@ func TestSlugAliaser_RoundTripEscapesUUIDSlugs(t *testing.T) { } // Round-trip must restore the original body exactly. - if got := a.dealiasContent(aliased); got != body { + if got := h.decodeContent(aliased); got != body { t.Fatalf("dealias round-trip mismatch:\n got %q\n want %q", got, body) } } -func TestSlugAliaser_UnknownTokensAndSlugsUntouched(t *testing.T) { - a := newSlugAliaser() - a.token("summary/real-uuid") +func TestWikiSlugHandles_UnknownHandlesAndSlugsUntouched(t *testing.T) { + h := newWikiSlugHandles() + h.handle("summary/real-uuid") // A slug not in `known` must not be aliased. known := map[string]struct{}{"summary/real-uuid": {}} body := "keep [[entity/invented-by-model]] as-is" - if got := a.aliasContent(body, known); got != body { + if got := h.encodeContent(body, known); got != body { t.Fatalf("unknown slug should be left untouched, got %q", got) } // An unmapped token in output must be left untouched (falls through to the // normal dead-link cleanup path). out := "stray [[ref-999|X]] token" - if got := a.dealiasContent(out); got != out { + if got := h.decodeContent(out); got != out { t.Fatalf("unmapped token should be left untouched, got %q", got) } } -func TestSlugAliaser_EmptyIsNoop(t *testing.T) { - a := newSlugAliaser() - if !a.empty() { - t.Fatal("fresh aliaser must be empty") +func TestWikiSlugHandles_EmptyIsNoop(t *testing.T) { + h := newWikiSlugHandles() + if !h.empty() { + t.Fatal("fresh handle table must be empty") } body := "[[summary/x]] unchanged" - if got := a.dealiasContent(body); got != body { - t.Fatalf("empty aliaser dealias must be a no-op, got %q", got) + if got := h.decodeContent(body); got != body { + t.Fatalf("empty handle table decode must be a no-op, got %q", got) } } diff --git a/internal/llmreference/model_output.go b/internal/llmreference/model_output.go index bd3845e26..20cc6050d 100644 --- a/internal/llmreference/model_output.go +++ b/internal/llmreference/model_output.go @@ -37,12 +37,52 @@ func (r *Registry) ModelOutput(result *types.ToolResult) string { return r.modelWebSearchOutput(mapsValue(result.Data["results"]), result.Output) case "web_fetch_results": return r.modelWebFetchOutput(mapsValue(result.Data["results"]), result.Output) + case "database_query": + return r.modelDatabaseQueryOutput(mapsValue(result.Data["rows"]), result.Output) default: r.registerLabeledReferences(result.Output) + r.registerStructuredReferences(result.Output) return r.CompactKnownText(result.Output) } } +// registerStructuredReferences registers durable IDs found under explicitly +// labeled keys of a JSON tool result, using the same key dispatch as tool +// arguments. Non-JSON output is a no-op. +func (r *Registry) registerStructuredReferences(raw string) { + var value interface{} + if json.Unmarshal([]byte(raw), &value) != nil { + return + } + var walk func(string, interface{}) + walk = func(key string, value interface{}) { + switch typed := value.(type) { + case string: + r.registerSourceIDByKey(key, typed) + case []interface{}: + for _, item := range typed { + walk(key, item) + } + case map[string]interface{}: + for childKey, item := range typed { + walk(childKey, item) + } + } + } + walk("", value) +} + +func (r *Registry) modelDatabaseQueryOutput(rows []map[string]interface{}, fallback string) string { + for _, row := range rows { + for key, raw := range row { + if value, ok := raw.(string); ok { + r.registerSourceIDByKey(key, value) + } + } + } + return r.CompactKnownText(fallback) +} + func (r *Registry) modelDocumentInfoOutput(rows []map[string]interface{}, fallback string) string { if len(rows) == 0 { return r.CompactKnownText(fallback) diff --git a/internal/llmreference/registry.go b/internal/llmreference/registry.go index e0051f981..387332e2a 100644 --- a/internal/llmreference/registry.go +++ b/internal/llmreference/registry.go @@ -1,7 +1,6 @@ -// Package llmreference keeps internal source identifiers out of LLM-facing -// context. It assigns request-local aliases to chunks, documents, knowledge -// bases and web pages, decodes aliases used in tool arguments, and expands the -// model's compact citations before application code consumes them. +// Package llmreference is the low-level source codec used by modelcontext. +// Application request lifecycles should depend on modelcontext.Registry so +// source and resource handles cannot be encoded or decoded out of order. package llmreference import ( @@ -96,6 +95,15 @@ func NewRegistry(citationsEnabled ...bool) *Registry { } } +// ProtocolPrompt returns the source protocol configured for this registry. +// Request lifecycle code should normally call this through modelcontext.Registry. +func (r *Registry) ProtocolPrompt() string { + if r == nil { + return "" + } + return ProtocolPrompt(r.citationsEnabled) +} + func (r *Registry) Count() int { if r == nil { return 0 @@ -106,7 +114,21 @@ func (r *Registry) Count() int { } func (r *Registry) RegisterChunk(ref ChunkReference) string { - if r == nil || strings.TrimSpace(ref.ChunkID) == "" { + if r == nil { + return "" + } + ref.ChunkID = strings.TrimSpace(ref.ChunkID) + if ref.ChunkID == "" { + return "" + } + if shortSourceAliasRE.MatchString(ref.ChunkID) { + alias := strings.ToLower(ref.ChunkID) + r.mu.RLock() + _, known := r.chunkByAlias[alias] + r.mu.RUnlock() + if known { + return alias + } return "" } r.mu.Lock() @@ -141,7 +163,18 @@ func mergeChunkReference(dst *ChunkReference, src ChunkReference) { } func (r *Registry) RegisterDocument(id string) string { - if r == nil || strings.TrimSpace(id) == "" { + id = strings.TrimSpace(id) + if r == nil || id == "" { + return "" + } + if shortSourceAliasRE.MatchString(id) { + alias := strings.ToLower(id) + r.mu.RLock() + _, known := r.aliasToDoc[alias] + r.mu.RUnlock() + if known { + return alias + } return "" } r.mu.Lock() @@ -156,7 +189,18 @@ func (r *Registry) RegisterDocument(id string) string { } func (r *Registry) RegisterKnowledgeBase(id string) string { - if r == nil || strings.TrimSpace(id) == "" { + id = strings.TrimSpace(id) + if r == nil || id == "" { + return "" + } + if shortSourceAliasRE.MatchString(id) { + alias := strings.ToLower(id) + r.mu.RLock() + _, known := r.aliasToKB[alias] + r.mu.RUnlock() + if known { + return alias + } return "" } r.mu.Lock() @@ -171,7 +215,18 @@ func (r *Registry) RegisterKnowledgeBase(id string) string { } func (r *Registry) RegisterWeb(rawURL, title string) string { - if r == nil || strings.TrimSpace(rawURL) == "" { + rawURL = strings.TrimSpace(rawURL) + if r == nil || rawURL == "" { + return "" + } + if shortSourceAliasRE.MatchString(rawURL) { + alias := strings.ToLower(rawURL) + r.mu.RLock() + _, known := r.webByAlias[alias] + r.mu.RUnlock() + if known { + return alias + } return "" } key := canonicalWebURL(rawURL) @@ -253,13 +308,94 @@ func (r *Registry) KnowledgeBaseAlias(id string) string { // DecodeToolCalls restores exact alias-valued JSON strings before tools parse // or validate their arguments. It never performs substring replacement. func (r *Registry) DecodeToolCalls(toolCalls []types.LLMToolCall) { + r.DecodeToolCallsWithPolicy(toolCalls, nil) +} + +// ToolArgumentPolicy decides whether a source-bearing JSON key belongs to a +// particular tool contract. A nil policy retains the codec's legacy generic +// behavior; application request lifecycles should provide an explicit policy +// through modelcontext.Registry. +type ToolArgumentPolicy func(toolName, key string) bool + +// DecodeToolCallsWithPolicy restores aliases only for fields explicitly owned +// by the named tool. This prevents dynamic tools with coincidentally named +// fields from inheriting built-in source semantics. +func (r *Registry) DecodeToolCallsWithPolicy(toolCalls []types.LLMToolCall, policy ToolArgumentPolicy) { for i := range toolCalls { - toolCalls[i].Function.Arguments = r.decodeJSON(toolCalls[i].Function.Arguments, false) + toolName := toolCalls[i].Function.Name + toolCalls[i].Function.Arguments = r.decodeJSONWithPolicy( + toolCalls[i].Function.Arguments, + false, + func(key string) bool { return policy == nil || policy(toolName, key) }, + ) + } +} + +// UnresolvedToolHandlesWithPolicy reports unknown handles only in fields that +// belong to the named tool's declared source contract. +func (r *Registry) UnresolvedToolHandlesWithPolicy( + toolName, raw string, + policy ToolArgumentPolicy, +) []string { + if strings.TrimSpace(raw) == "" { + return nil + } + var value interface{} + if err := json.Unmarshal([]byte(raw), &value); err != nil { + return nil + } + seen := make(map[string]struct{}) + r.collectUnresolvedToolHandles( + "", value, seen, + func(key string) bool { return policy == nil || policy(toolName, key) }, + ) + result := make([]string, 0, len(seen)) + for handle := range seen { + result = append(result, handle) + } + sort.Strings(result) + return result +} + +func (r *Registry) collectUnresolvedToolHandles( + key string, + value interface{}, + seen map[string]struct{}, + allowed func(string) bool, +) { + switch typed := value.(type) { + case string: + key = strings.ToLower(key) + if _, ok := decodableKeys[key]; !ok || !allowed(key) { + return + } + handle := strings.TrimSpace(typed) + if shortSourceAliasRE.MatchString(handle) && (r == nil || r.realForAlias(handle) == "") { + seen[handle] = struct{}{} + } + case []interface{}: + for _, item := range typed { + r.collectUnresolvedToolHandles(key, item, seen, allowed) + } + case map[string]interface{}: + for childKey, item := range typed { + r.collectUnresolvedToolHandles(childKey, item, seen, allowed) + } } } // EncodeMessages compacts known real identifiers in assistant tool-call replay. func (r *Registry) EncodeMessages(messages []chat.Message) []chat.Message { + return r.EncodeMessagesWithPolicies(messages, nil, nil) +} + +// EncodeMessagesWithPolicies additionally gates source processing for tool +// result messages. A nil result policy retains the legacy generic behavior. +func (r *Registry) EncodeMessagesWithPolicies( + messages []chat.Message, + argumentPolicy ToolArgumentPolicy, + resultPolicy func(toolName string) bool, +) []chat.Message { if r == nil || len(messages) == 0 { return messages } @@ -269,14 +405,15 @@ func (r *Registry) EncodeMessages(messages []chat.Message) []chat.Message { // and canonical assistant citations. This two-pass shape lets an early tool // message reuse metadata that appears only in the turn's final answer. for i := range out { - if out[i].Role == "assistant" || out[i].Role == "tool" { + processToolResult := out[i].Role == "tool" && (resultPolicy == nil || resultPolicy(out[i].Name)) + if out[i].Role == "assistant" || processToolResult { out[i].Content = r.CompactPublicCitations(out[i].Content) out[i].ReasoningContent = r.CompactPublicCitations(out[i].ReasoningContent) } if len(out[i].MultiContent) > 0 { out[i].MultiContent = append([]chat.MessageContentPart(nil), out[i].MultiContent...) for j := range out[i].MultiContent { - if out[i].MultiContent[j].Type == "text" && (out[i].Role == "assistant" || out[i].Role == "tool") { + if out[i].MultiContent[j].Type == "text" && (out[i].Role == "assistant" || processToolResult) { out[i].MultiContent[j].Text = r.CompactPublicCitations(out[i].MultiContent[j].Text) } } @@ -284,25 +421,126 @@ func (r *Registry) EncodeMessages(messages []chat.Message) []chat.Message { if len(out[i].ToolCalls) > 0 { out[i].ToolCalls = append([]chat.ToolCall(nil), out[i].ToolCalls...) for j := range out[i].ToolCalls { - r.registerToolArguments(out[i].ToolCalls[j].Function.Arguments) + toolName := out[i].ToolCalls[j].Function.Name + r.registerToolArguments( + out[i].ToolCalls[j].Function.Arguments, + func(key string) bool { return argumentPolicy == nil || argumentPolicy(toolName, key) }, + ) } } } for i := range out { - if out[i].Role == "tool" { + if out[i].Role == "tool" && (resultPolicy == nil || resultPolicy(out[i].Name)) { r.registerLegacyToolReferences(out[i].Content) out[i].Content = r.CompactKnownText(out[i].Content) } for j := range out[i].ToolCalls { - out[i].ToolCalls[j].Function.Arguments = r.decodeJSON(out[i].ToolCalls[j].Function.Arguments, true) + toolName := out[i].ToolCalls[j].Function.Name + out[i].ToolCalls[j].Function.Arguments = r.decodeJSONWithPolicy( + out[i].ToolCalls[j].Function.Arguments, + true, + func(key string) bool { return argumentPolicy == nil || argumentPolicy(toolName, key) }, + ) } } return out } var shortSourceAliasRE = regexp.MustCompile(`(?i)^[cdbw][1-9][0-9]*$`) +var shortSourceAliasInTextRE = regexp.MustCompile(`(?i)\b[cdbw][1-9][0-9]*\b`) -func (r *Registry) registerToolArguments(raw string) { +// DecodeKnownText restores registered source handles embedded in a structured +// expression such as a built-in SQL tool argument. It must not be used for +// arbitrary prose; modelcontext owns the small tool/key policy that calls it. +func (r *Registry) DecodeKnownText(text string) string { + if r == nil || text == "" { + return text + } + return shortSourceAliasInTextRE.ReplaceAllStringFunc(text, func(handle string) string { + if real := r.realForAlias(handle); real != "" { + return real + } + return handle + }) +} + +// DecodeKnownQuotedText restores source handles only inside single-quoted, +// double-quoted, or backtick-quoted segments. It is intended for structured +// expressions such as SQL, where replacing an unquoted token could corrupt a +// legitimate table/column alias that happens to look like d1 or b2. +func (r *Registry) DecodeKnownQuotedText(text string) string { + if r == nil || text == "" { + return text + } + return rewriteQuotedText(text, func(segment string) string { + return shortSourceAliasInTextRE.ReplaceAllStringFunc(segment, func(handle string) string { + if real := r.realForAlias(handle); real != "" { + return real + } + return handle + }) + }) +} + +// UnresolvedQuotedTextHandles reports alias-shaped values inside quoted +// structured-text segments that do not exist in this request registry. +func (r *Registry) UnresolvedQuotedTextHandles(text string) []string { + if text == "" { + return nil + } + seen := make(map[string]struct{}) + rewriteQuotedText(text, func(segment string) string { + for _, handle := range shortSourceAliasInTextRE.FindAllString(segment, -1) { + if r == nil || r.realForAlias(handle) == "" { + seen[handle] = struct{}{} + } + } + return segment + }) + result := make([]string, 0, len(seen)) + for handle := range seen { + result = append(result, handle) + } + sort.Strings(result) + return result +} + +func rewriteQuotedText(text string, rewrite func(string) string) string { + var out strings.Builder + out.Grow(len(text)) + for i := 0; i < len(text); { + quote := text[i] + if quote != '\'' && quote != '"' && quote != '`' { + out.WriteByte(text[i]) + i++ + continue + } + start := i + i++ + for i < len(text) { + if text[i] == '\\' && i+1 < len(text) { + i += 2 + continue + } + if text[i] != quote { + i++ + continue + } + // SQL escapes a quote by doubling it (''). Keep scanning the + // same literal instead of treating the first quote as its end. + if i+1 < len(text) && text[i+1] == quote { + i += 2 + continue + } + i++ + break + } + out.WriteString(rewrite(text[start:i])) + } + return out.String() +} + +func (r *Registry) registerToolArguments(raw string, allowed func(string) bool) { if r == nil || strings.TrimSpace(raw) == "" { return } @@ -310,35 +548,51 @@ func (r *Registry) registerToolArguments(raw string) { if err := json.Unmarshal([]byte(raw), &value); err != nil { return } - r.registerToolArgumentValue("", value) + r.registerToolArgumentValue("", value, allowed) } -func (r *Registry) registerToolArgumentValue(key string, value interface{}) { +func (r *Registry) registerToolArgumentValue(key string, value interface{}, allowed func(string) bool) { switch typed := value.(type) { case string: - value := strings.TrimSpace(typed) - if value == "" || shortSourceAliasRE.MatchString(value) { - return - } - switch strings.ToLower(key) { - case "chunk_id", "faq_id", "chunk_ids", "faq_ids": - r.RegisterChunk(ChunkReference{ChunkID: value}) - case "knowledge_id", "knowledge_ids", "suspected_knowledge_ids", "source_refs": - r.RegisterDocument(value) - case "knowledge_base", "knowledge_base_id", "knowledge_base_ids", "kb_id", "kb_ids": - r.RegisterKnowledgeBase(value) - case "url", "urls": - if parsed, err := url.Parse(value); err == nil && (parsed.Scheme == "http" || parsed.Scheme == "https") { - r.RegisterWeb(value, "") - } + if allowed(strings.ToLower(key)) { + r.registerSourceIDByKey(key, typed) } case []interface{}: for _, item := range typed { - r.registerToolArgumentValue(key, item) + r.registerToolArgumentValue(key, item, allowed) } case map[string]interface{}: for childKey, item := range typed { - r.registerToolArgumentValue(childKey, item) + r.registerToolArgumentValue(childKey, item, allowed) + } + } +} + +// registerSourceIDByKey is the single key→source-space dispatch used for tool +// arguments, structured tool results, and database rows. Keeping one dispatch +// stops the recognized key set (and the http/https guard for web references) +// from drifting between call sites; it must stay in sync with decodableKeys. +func (r *Registry) registerSourceIDByKey(key, value string) { + value = strings.TrimSpace(value) + if value == "" || shortSourceAliasRE.MatchString(value) { + return + } + switch strings.ToLower(key) { + case "chunk_id", "faq_id", "chunk_ids", "faq_ids": + r.RegisterChunk(ChunkReference{ChunkID: value}) + case "knowledge_id", "knowledge_ids", "suspected_knowledge_ids": + r.RegisterDocument(value) + case "source_refs": + // Stored refs use "knowledgeID|title"; only the ID part is durable. + r.RegisterDocument(strings.TrimSpace(strings.SplitN(value, "|", 2)[0])) + case "knowledge_base", "knowledge_base_id", "knowledge_base_ids", "kb_id", "kb_ids": + r.RegisterKnowledgeBase(value) + case "url", "urls": + // Only public web pages become web references. Internal schemes + // (res://, storage providers) must never enter the web alias space, + // where CompactKnownText would rewrite them a second time. + if parsed, err := url.Parse(value); err == nil && (parsed.Scheme == "http" || parsed.Scheme == "https") { + r.RegisterWeb(value, "") } } } @@ -412,7 +666,7 @@ func publicAttr(expression *regexp.Regexp, tag string) string { return html.UnescapeString(match[1]) } -func (r *Registry) decodeJSON(raw string, encode bool) string { +func (r *Registry) decodeJSONWithPolicy(raw string, encode bool, allowed func(string) bool) string { if r == nil || strings.TrimSpace(raw) == "" { return raw } @@ -420,7 +674,7 @@ func (r *Registry) decodeJSON(raw string, encode bool) string { if err := json.Unmarshal([]byte(raw), &value); err != nil { return raw } - value = r.walkJSON("", value, encode) + value = r.walkJSON("", value, encode, allowed) encoded, err := json.Marshal(value) if err != nil { return raw @@ -439,9 +693,12 @@ var decodableKeys = map[string]struct{}{ "url": {}, "urls": {}, } -func (r *Registry) walkJSON(key string, value interface{}, encode bool) interface{} { +func (r *Registry) walkJSON(key string, value interface{}, encode bool, allowed func(string) bool) interface{} { switch typed := value.(type) { case string: + if !allowed(strings.ToLower(key)) { + return typed + } if encode { // Encode matches on exact real identifiers (UUIDs/URLs), which do // not collide with prose, so it stays key-agnostic. @@ -464,11 +721,11 @@ func (r *Registry) walkJSON(key string, value interface{}, encode bool) interfac return typed case []interface{}: for i := range typed { - typed[i] = r.walkJSON(key, typed[i], encode) + typed[i] = r.walkJSON(key, typed[i], encode, allowed) } case map[string]interface{}: for childKey, item := range typed { - typed[childKey] = r.walkJSON(childKey, item, encode) + typed[childKey] = r.walkJSON(childKey, item, encode, allowed) } } return value @@ -493,6 +750,7 @@ func (r *Registry) aliasForRealValue(real string) string { } func (r *Registry) realForAlias(alias string) string { + alias = strings.ToLower(strings.TrimSpace(alias)) r.mu.RLock() defer r.mu.RUnlock() if ref := r.chunkByAlias[alias]; ref != nil { @@ -536,47 +794,12 @@ func (r *Registry) CompactKnownText(text string) string { sort.SliceStable(pairs, func(i, j int) bool { return len(pairs[i].real) > len(pairs[j].real) }) for _, item := range pairs { if item.real != "" { - text = replaceIDExceptWikiSummarySlug(text, item.real, item.alias) + text = strings.ReplaceAll(text, item.real, item.alias) } } return text } -// wikiSummarySlugPrefix guards wiki summary-page slugs (summary/) -// from citation compaction. A summary slug embeds a document's knowledge ID -// verbatim, so a blind ReplaceAll of that ID into its citation alias (d1) would -// mangle the slug into summary/d1 — an unrecoverable dead link the model then -// copies into wiki tool calls. Only the slug-embedded occurrence is preserved; -// the same document ID elsewhere in the text still compacts normally. -const wikiSummarySlugPrefix = "summary/" - -// replaceIDExceptWikiSummarySlug replaces every occurrence of real with alias, -// except one that is immediately preceded by "summary/" (i.e. it is the -// identifier segment of a wiki summary-page slug). Go's regexp package has no -// lookbehind, so the scan is done manually. -func replaceIDExceptWikiSummarySlug(text, real, alias string) string { - if real == "" || !strings.Contains(text, real) { - return text - } - var b strings.Builder - b.Grow(len(text)) - for { - i := strings.Index(text, real) - if i < 0 { - b.WriteString(text) - break - } - b.WriteString(text[:i]) - if strings.HasSuffix(text[:i], wikiSummarySlugPrefix) { - b.WriteString(real) // part of a summary slug — keep intact - } else { - b.WriteString(alias) - } - text = text[i+len(real):] - } - return b.String() -} - var ( refTagRE = regexp.MustCompile(`(?i)`) refCandidateRE = regexp.MustCompile(`(?is)]*(?:>|$)`) diff --git a/internal/llmreference/registry_test.go b/internal/llmreference/registry_test.go index 6b74f0cdc..ac4573bce 100644 --- a/internal/llmreference/registry_test.go +++ b/internal/llmreference/registry_test.go @@ -32,28 +32,23 @@ func TestRegistryChunkAliasIsStableAndExpandsCanonicalCitation(t *testing.T) { ) } -func TestCompactKnownTextPreservesWikiSummarySlug(t *testing.T) { - const knowledgeID = "07a20bb1-a662-47cf-9929-06fb5d5b5b5e" - const kbID = "250368ff-f5a2-4e9e-868a-07bc9b857c44" +func TestRegisterDoesNotTreatModelAliasesAsNewDurableIdentities(t *testing.T) { + r := NewRegistry() + require.Equal(t, "c1", r.RegisterChunk(ChunkReference{ChunkID: "chunk-real"})) + require.Equal(t, "d1", r.RegisterDocument("doc-real")) + require.Equal(t, "b1", r.RegisterKnowledgeBase("kb-real")) + require.Equal(t, "w1", r.RegisterWeb("https://example.com", "Example")) - registry := NewRegistry() - require.Equal(t, "d1", registry.RegisterDocument(knowledgeID)) - require.Equal(t, "b1", registry.RegisterKnowledgeBase(kbID)) - - // A wiki_search result: the summary slug embeds the document UUID, and the - // same UUID also appears as a standalone label. - in := "" + kbID + "\n" + - "[[summary/" + knowledgeID + "|Doc - Summary]]\n" + - "" + knowledgeID + "" - - got := registry.CompactKnownText(in) - - // The slug-embedded UUID must survive verbatim — no summary/d1 mangling. - require.Contains(t, got, "[[summary/"+knowledgeID+"|Doc - Summary]]") - require.NotContains(t, got, "summary/d1") - // The standalone label and the KB id still compact normally. - require.Contains(t, got, "d1") - require.Contains(t, got, "b1") + require.Equal(t, "c1", r.RegisterChunk(ChunkReference{ChunkID: "c1"})) + require.Equal(t, "d1", r.RegisterDocument("d1")) + require.Equal(t, "b1", r.RegisterKnowledgeBase("b1")) + require.Equal(t, "w1", r.RegisterWeb("w1", "")) + require.Empty(t, r.RegisterDocument("d99")) + require.Empty(t, r.RegisterDocument("c1"), "a chunk handle must not be accepted as a document identity") + require.Equal(t, 1, len(r.chunkByAlias)) + require.Equal(t, 1, len(r.aliasToDoc)) + require.Equal(t, 1, len(r.aliasToKB)) + require.Equal(t, 1, len(r.webByAlias)) } func TestRegistrySuppressesSourceCitationsWhenDisabled(t *testing.T) { @@ -338,3 +333,27 @@ func TestModelOutputGraphResultsUseChunkAliases(t *testing.T) { registry.ExpandText(``), ) } + +func TestModelOutputDoesNotRegisterInternalSchemesAsWebSources(t *testing.T) { + registry := NewRegistry() + output := registry.ModelOutput(&types.ToolResult{ + Success: true, + Output: `{"url":"res://0001","knowledge_id":"doc-real-id"}`, + }) + + // An internal handle in a url-labeled field must never enter the web + // alias space, where CompactKnownText would rewrite it a second time. + require.Contains(t, output, "res://0001") + require.NotContains(t, output, "w1") + require.Contains(t, output, "d1") + require.NotContains(t, output, "doc-real-id") + + registry.ModelOutput(&types.ToolResult{ + Success: true, + Output: `{"url":"https://example.com/page"}`, + }) + require.Equal(t, + ``, + registry.ExpandText(``), + ) +} diff --git a/internal/llmresource/registry.go b/internal/llmresource/registry.go index 705553acb..8e2daf459 100644 --- a/internal/llmresource/registry.go +++ b/internal/llmresource/registry.go @@ -1,5 +1,6 @@ -// Package llmresource shortens stored-resource references while they are in an -// LLM context and restores them before application code consumes the output. +// Package llmresource is the low-level durable-resource codec used by +// modelcontext. Application request lifecycles should not coordinate it +// independently from source-reference handling. package llmresource import ( @@ -38,7 +39,7 @@ var storedRefRE = regexp.MustCompile( // aliasShapeRE matches the alias syntax produced by EncodeText. It is used only // to spot alias-shaped tokens the model emitted that the registry cannot map // back — either a hallucinated reference or a coincidental collision. -var aliasShapeRE = regexp.MustCompile(`res://\d{4,}`) +var aliasShapeRE = regexp.MustCompile(`res://\d+`) // Registry assigns low-entropy, request-local aliases to stable resource // handles. It is safe to reuse across all rounds of one Agent execution. @@ -96,6 +97,16 @@ func (r *Registry) DecodeText(value string) string { return decoded } +// StripOrphanAliases removes alias-shaped tokens after all known aliases have +// been restored. Use this only on model output; tool arguments must retain +// unknown aliases long enough for modelcontext to reject the call. +func (r *Registry) StripOrphanAliases(value string) string { + if value == "" { + return value + } + return aliasShapeRE.ReplaceAllString(value, "") +} + // EncodeMessages returns a copied message slice with textual references // compacted. Binary/image content fields are intentionally left untouched. func (r *Registry) EncodeMessages(messages []chat.Message) []chat.Message { @@ -201,6 +212,9 @@ func (d *StreamDecoder) Feed(chunk string) string { d.pending = "" hold := 0 for _, alias := range d.registry.aliases() { + // A provider may split the token at any byte boundary, including + // "re" + "s://0001". Holding at most the short matching suffix is the + // only way to guarantee that a request-local handle never leaks. for n := 1; n < len(alias); n++ { if n > hold && strings.HasSuffix(combined, alias[:n]) { hold = n diff --git a/internal/llmresource/registry_test.go b/internal/llmresource/registry_test.go index 2483abb69..2187651f3 100644 --- a/internal/llmresource/registry_test.go +++ b/internal/llmresource/registry_test.go @@ -61,6 +61,16 @@ func TestStreamDecoderRestoresSplitAlias(t *testing.T) { require.Equal(t, "r", d.Flush()) } +func TestStreamDecoderRestoresAliasSplitBeforeScheme(t *testing.T) { + const ref = "resource://AbCdEfGhIjKlMnOpQrStUv" + r := NewRegistry() + require.Equal(t, "res://0001", r.EncodeText(ref)) + d := NewStreamDecoder(r) + require.Equal(t, "see ", d.Feed("see re")) + require.Equal(t, ref, d.Feed("s://0001")) + require.Empty(t, d.Flush()) +} + func TestOrphanAliasesReportsUnresolvableTokens(t *testing.T) { r := NewRegistry() ref := "resource://AbCdEfGhIjKlMnOpQrStUv" @@ -73,6 +83,11 @@ func TestOrphanAliasesReportsUnresolvableTokens(t *testing.T) { require.Equal(t, []string{"res://0099"}, r.OrphanAliases("look at res://0099 and res://0099")) } +func TestStripOrphanAliasesRemovesIncompleteNumericHandle(t *testing.T) { + r := NewRegistry() + require.Equal(t, "broken ", r.StripOrphanAliases("broken res://1")) +} + func TestOrphanAliasesNilRegistry(t *testing.T) { var r *Registry require.Equal(t, []string{"res://0001"}, r.OrphanAliases("res://0001")) diff --git a/internal/modelcontext/README.md b/internal/modelcontext/README.md new file mode 100644 index 000000000..6df2c74af --- /dev/null +++ b/internal/modelcontext/README.md @@ -0,0 +1,76 @@ +# Model context boundary + +`modelcontext.Registry` is the only application-facing boundary for values +that are shortened before an LLM call and restored afterwards. + +## Identity rules + +- Durable application identities: UUIDs, Wiki slugs, URLs, and + `resource://...` references. +- Request-local model handles: `cN`, `dN`, `bN`, `wN`, and `res://NNNN`. +- Tool-private request handles: `iN` for Wiki issues. +- Ingest-call handles: `cNNN` and `ref-N`, allocated with `HandleTable`. +- Temporary handles are decoded before tool execution, persistence, or UI + delivery. They are never accepted as durable authorization or routing data. + +## Request lifecycle + +1. Create one `Registry` for the complete request or Agent execution. +2. Register structured source identities as they enter the context. +3. Call `EncodeMessages` immediately before the model call. +4. Decode tool calls through `DecodeToolCalls` before parsing or execution. +5. Decode complete responses with `DecodeResponse`, or streaming text with one + `StreamDecoder` per response channel. + +Unknown handle-shaped values in declared tool fields are rejected before tool +execution. The same per-tool policy is applied when historical assistant tool +calls and tool results are replayed, so aliases neither disappear nor drift to +new numbers across Agent rounds. + +The registry owns codec ordering. Resource references are encoded before +source IDs so a Wiki slug such as `summary/` cannot be corrupted +into `summary/d1`. + +## Observability + +Langfuse generation observations contain the exact encoded payload sent to and +returned by the model. Agent tool spans contain both `model_arguments` (the +model-emitted handles) and `resolved_arguments` (the durable values actually +executed), plus `argument_resolution` and any unresolved handles. Sensitive +tools report only argument keys and resolution counts. + +Langfuse is an observer of this boundary; it must not implement another handle +mapping. + +## Built-in tool policies + +Canonical ID-bearing arguments for knowledge, graph, data, web, and Wiki tools +are decoded centrally through a tool-name plus JSON-field allowlist. +`database_query.sql` and `data_analysis.sql` have an explicit policy because +source handles can be embedded inside quoted SQL values; unquoted SQL aliases +and arbitrary prose are never rewritten. Wiki issue IDs use the same lifecycle +via an `iN` handle space. + +Dynamic MCP tools are intentionally opaque in both arguments and results. +Their schemas and identity semantics are controlled by the MCP server, so the +application must not guess that an arbitrary `id`, `knowledge_id`, or `url` +field is durable or rewrite it. Durable resource handles inside MCP arguments +still use the normal `res://NNNN` codec. A future MCP ID mapping must be an +explicit server/tool annotation and should plug into this registry rather than +create a parallel mapper. + +## Wiki routing + +`WikiRouteResolver` is intentionally separate. It stores server-side +slug-to-knowledge-base provenance gathered during Wiki search/read operations, +and can only return knowledge bases already present in the Agent's authorized +Wiki scope. It is routing state, not a model handle table. Reads scan every +legal Wiki KB so cached provenance cannot hide a duplicate slug; mutations +require one unique owner. New pages use link provenance, a single Wiki scope, +or a single source-document owner and reject ambiguous routing instead of +asking the model for a durable KB ID. + +Document/tag-constrained Wiki and graph requests are filtered by the same +server-owned `SearchTargets` provenance. Uncited/global Wiki pages and graph +results outside that subset fail closed; only an explicit whole-KB target +authorizes whole-KB content. diff --git a/internal/modelcontext/handles.go b/internal/modelcontext/handles.go new file mode 100644 index 000000000..cf0dc6f55 --- /dev/null +++ b/internal/modelcontext/handles.go @@ -0,0 +1,208 @@ +package modelcontext + +import ( + "fmt" + "regexp" + "sort" + "strings" + "sync" +) + +// HandleTable is a typed, invocation-local bidirectional mapping used when a +// prompt needs compact handles for durable values. Handles must never be +// persisted; Resolve converts model output back to the durable value first. +type HandleTable struct { + mu sync.RWMutex + prefix string + width int + next int + toHandle map[string]string + toValue map[string]string +} + +// NewHandleTable creates a handle space such as c000 (prefix=c, width=3, +// start=0) or ref-1 (prefix=ref-, width=0, start=1). +func NewHandleTable(prefix string, width, start int) *HandleTable { + return &HandleTable{ + prefix: prefix, + width: width, + next: start, + toHandle: make(map[string]string), + toValue: make(map[string]string), + } +} + +// Register returns the stable handle assigned to value in this table. +func (t *HandleTable) Register(value string) string { + if t == nil { + return "" + } + value = strings.TrimSpace(value) + if value == "" { + return "" + } + t.mu.Lock() + defer t.mu.Unlock() + if handle := t.toHandle[value]; handle != "" { + return handle + } + number := fmt.Sprintf("%d", t.next) + if t.width > 0 { + number = fmt.Sprintf("%0*d", t.width, t.next) + } + t.next++ + handle := t.prefix + number + t.toHandle[value] = handle + t.toValue[handle] = value + return handle +} + +// Handle returns an already registered handle without creating one. +func (t *HandleTable) Handle(value string) (string, bool) { + if t == nil { + return "", false + } + t.mu.RLock() + defer t.mu.RUnlock() + handle, ok := t.toHandle[value] + return handle, ok +} + +// Resolve converts a known model handle back to its durable value. +func (t *HandleTable) Resolve(handle string) (string, bool) { + if t == nil { + return "", false + } + t.mu.RLock() + defer t.mu.RUnlock() + value, ok := t.toValue[strings.TrimSpace(handle)] + return value, ok +} + +func (t *HandleTable) Empty() bool { + if t == nil { + return true + } + t.mu.RLock() + defer t.mu.RUnlock() + return len(t.toHandle) == 0 +} + +func (t *HandleTable) Len() int { + if t == nil { + return 0 + } + t.mu.RLock() + defer t.mu.RUnlock() + return len(t.toHandle) +} + +// EncodeKnownText replaces already-registered durable values with their model +// handles. Longer values are processed first to avoid substring shadowing. +func (t *HandleTable) EncodeKnownText(value string) string { + if t == nil || value == "" { + return value + } + t.mu.RLock() + type pair struct{ value, handle string } + pairs := make([]pair, 0, len(t.toHandle)) + for durable, handle := range t.toHandle { + pairs = append(pairs, pair{durable, handle}) + } + t.mu.RUnlock() + sort.SliceStable(pairs, func(i, j int) bool { return len(pairs[i].value) > len(pairs[j].value) }) + for _, item := range pairs { + value = strings.ReplaceAll(value, item.value, item.handle) + } + return value +} + +// DecodeKnownText restores registered handles in complete text. +func (t *HandleTable) DecodeKnownText(value string) string { + if t == nil || value == "" { + return value + } + t.mu.RLock() + handles := make([]string, 0, len(t.toValue)) + for handle := range t.toValue { + handles = append(handles, handle) + } + t.mu.RUnlock() + sort.SliceStable(handles, func(i, j int) bool { return len(handles[i]) > len(handles[j]) }) + for _, handle := range handles { + pattern := regexp.MustCompile(`\b` + regexp.QuoteMeta(handle) + `\b`) + if durable, ok := t.Resolve(handle); ok { + // A durable value may legally contain '$'. ReplaceAllString would + // interpret it as a regexp expansion, so use a literal callback. + value = pattern.ReplaceAllStringFunc(value, func(string) string { return durable }) + } + } + return value +} + +// HandleStreamDecoder restores HandleTable values without leaking a handle +// split across provider chunks. +type HandleStreamDecoder struct { + table *HandleTable + pending string +} + +func NewHandleStreamDecoder(table *HandleTable) *HandleStreamDecoder { + return &HandleStreamDecoder{table: table} +} + +func (d *HandleStreamDecoder) Feed(chunk string) string { + if d == nil || d.table == nil { + return chunk + } + combined := d.pending + chunk + d.pending = "" + if combined == "" { + return "" + } + start := len(combined) + for start > 0 && isHandleTokenByte(combined[start-1]) { + start-- + } + tail := combined[start:] + d.table.mu.RLock() + prefix := d.table.prefix + d.table.mu.RUnlock() + if couldBeNumericHandle(tail, prefix) { + d.pending = tail + combined = combined[:start] + } + return d.table.DecodeKnownText(combined) +} + +func (d *HandleStreamDecoder) Flush() string { + if d == nil || d.table == nil { + return "" + } + tail := d.table.DecodeKnownText(d.pending) + d.pending = "" + return tail +} + +func isHandleTokenByte(value byte) bool { + return value >= 'a' && value <= 'z' || value >= 'A' && value <= 'Z' || + value >= '0' && value <= '9' || value == '-' || value == '_' +} + +func couldBeNumericHandle(value, prefix string) bool { + if value == "" || prefix == "" { + return false + } + if strings.HasPrefix(prefix, value) { + return true + } + if !strings.HasPrefix(value, prefix) || len(value) == len(prefix) { + return false + } + for _, char := range value[len(prefix):] { + if char < '0' || char > '9' { + return false + } + } + return true +} diff --git a/internal/modelcontext/orphan_stream.go b/internal/modelcontext/orphan_stream.go new file mode 100644 index 000000000..424df30b8 --- /dev/null +++ b/internal/modelcontext/orphan_stream.go @@ -0,0 +1,83 @@ +package modelcontext + +import ( + "regexp" + "strings" +) + +var orphanResourceHandleRE = regexp.MustCompile(`res://\d+`) + +// orphanResourceStreamFilter removes unresolved resource handles only after +// the registered resource decoder has had a chance to restore known ones. It +// buffers a possible handle suffix so provider chunk boundaries cannot leak a +// partial internal token to the UI. +type orphanResourceStreamFilter struct { + pending string +} + +func (f *orphanResourceStreamFilter) Feed(chunk string) string { + if f == nil { + return chunk + } + combined := f.pending + chunk + f.pending = "" + if combined == "" { + return "" + } + + const prefix = "res://" + holdAt := -1 + // Unknown handles must be safe across every provider split, including + // "re" + "s://9999". This may defer at most three ordinary characters + // until the next chunk; Flush preserves them when they are normal prose. + for n := 1; n < len(prefix); n++ { + if strings.HasSuffix(combined, prefix[:n]) { + holdAt = len(combined) - n + } + } + if idx := strings.LastIndex(combined, prefix); idx >= 0 { + suffix := combined[idx+len(prefix):] + if suffix == "" || allDigits(suffix) { + holdAt = idx + } + } + if holdAt >= 0 { + f.pending = combined[holdAt:] + combined = combined[:holdAt] + } + return orphanResourceHandleRE.ReplaceAllString(combined, "") +} + +func (f *orphanResourceStreamFilter) Flush() string { + if f == nil { + return "" + } + pending := f.pending + f.pending = "" + // A stream that ends mid-token must not surface the model-context + // protocol fragment. Preserve ordinary r/re/res prose, but discard any + // suffix that has already crossed into the reserved URL-like syntax. + if strings.HasPrefix("res://", pending) && len(pending) >= len("res:") { + return "" + } + if strings.HasPrefix(pending, "res://") { + digits := strings.TrimPrefix(pending, "res://") + if digits == "" || allDigits(digits) { + return "" + } + } + pending = orphanResourceHandleRE.ReplaceAllString(pending, "") + return pending +} + +func allDigits(value string) bool { + if value == "" { + return false + } + for _, char := range value { + if char < '0' || char > '9' { + return false + } + } + return true +} diff --git a/internal/modelcontext/registry.go b/internal/modelcontext/registry.go new file mode 100644 index 000000000..4f9e123bc --- /dev/null +++ b/internal/modelcontext/registry.go @@ -0,0 +1,345 @@ +// Package modelcontext owns every request-local handle exposed to a language +// model. Application code should use Registry instead of coordinating source +// references and durable-resource codecs independently. +// +// Its boundaries are deliberate: +// - UUIDs, wiki slugs, URLs, and resource:// handles are durable identities. +// - cN/dN/bN/wN/iN/res://NNNN/ref-N values are temporary model handles. +// - temporary handles are never persisted or accepted outside their registry. +// - every model response is decoded before tools, storage, or UI consume it. +package modelcontext + +import ( + "encoding/json" + "reflect" + "sort" + + "github.com/Tencent/WeKnora/internal/llmreference" + "github.com/Tencent/WeKnora/internal/llmresource" + "github.com/Tencent/WeKnora/internal/models/chat" + "github.com/Tencent/WeKnora/internal/types" +) + +const ( + ArgumentResolutionUnchanged = "unchanged" + ArgumentResolutionResolved = "resolved" + ArgumentResolutionPartiallyResolved = "partially_resolved" + ArgumentResolutionUnresolved = "unresolved" +) + +const resourceHandleProtocolPrompt = ` + +## Resource handle protocol (system-owned) +Some durable resources and high-entropy Wiki slugs are represented by request-local res://NNNN handles. Wiki issues may use iN handles. +- Copy only handles that appeared in supplied context or tool results, preserving them exactly in links, images, and tool arguments. +- Never invent, edit, or expand any handle. The system restores it after generation.` + +// ChunkReference describes a retrieved chunk registered in model context. +// It aliases the codec type so callers do not need to import llmreference. +type ChunkReference = llmreference.ChunkReference + +// Registry is the single request-scoped boundary between durable application +// identities and temporary model handles. +// +// Durable resource references are encoded before source identifiers. This +// ordering is intentionally private: summary/ wiki slugs must be +// protected as one resource-like handle before the embedded document ID can be +// compacted to dN. Callers therefore cannot accidentally reverse the codecs. +type Registry struct { + sources *llmreference.Registry + resources *llmresource.Registry + issues *HandleTable +} + +// NewRegistry creates a registry for one model request/Agent execution. +func NewRegistry(citationsEnabled bool) *Registry { + return &Registry{ + sources: llmreference.NewRegistry(citationsEnabled), + resources: llmresource.NewRegistry(), + issues: NewHandleTable("i", 0, 1), + } +} + +// ProtocolPrompt returns the system-owned model handle and citation protocol. +func (r *Registry) ProtocolPrompt() string { + if r == nil || r.sources == nil { + return "" + } + return r.sources.ProtocolPrompt() + resourceHandleProtocolPrompt +} + +// EncodeMessages returns a model-facing copy with every temporary handle +// encoded in the only safe order. +func (r *Registry) EncodeMessages(messages []chat.Message) []chat.Message { + if r == nil { + return messages + } + messages = r.resources.EncodeMessages(messages) + // Register tool-private IDs from all replayed results before encoding any + // assistant call. The scan is intentionally order-independent, matching the + // source codec's two-pass replay behavior. + for i := range messages { + if messages[i].Role == "tool" { + messages[i].Content = r.encodeToolPrivateResult(messages[i].Name, messages[i].Content) + } + } + messages = r.sources.EncodeMessagesWithPolicies(messages, sourceArgumentAllowed, sourceOutputAllowed) + for i := range messages { + if len(messages[i].ToolCalls) == 0 { + continue + } + messages[i].ToolCalls = append([]chat.ToolCall(nil), messages[i].ToolCalls...) + for j := range messages[i].ToolCalls { + r.encodeReplayedToolPolicies(&messages[i].ToolCalls[j]) + } + } + return messages +} + +// DecodeToolCalls restores all temporary handles in tool-call arguments. +func (r *Registry) DecodeToolCalls(toolCalls []types.LLMToolCall) { + if r == nil { + return + } + for i := range toolCalls { + if toolCalls[i].ModelArguments == "" { + toolCalls[i].ModelArguments = toolCalls[i].Function.Arguments + } + } + r.resources.DecodeToolCalls(toolCalls) + r.sources.DecodeToolCallsWithPolicy(toolCalls, sourceArgumentAllowed) + for i := range toolCalls { + r.decodeToolPolicies(&toolCalls[i]) + resolved := toolCalls[i].Function.Arguments + unresolved := append( + r.resources.OrphanAliases(resolved), + r.sources.UnresolvedToolHandlesWithPolicy( + toolCalls[i].Function.Name, resolved, sourceArgumentAllowed, + )..., + ) + unresolved = append(unresolved, r.unresolvedPrivateToolHandles(toolCalls[i].Function.Name, resolved)...) + toolCalls[i].UnresolvedHandles = uniqueSorted(unresolved) + changed := !jsonEquivalent(toolCalls[i].ModelArguments, resolved) + switch { + case changed && len(toolCalls[i].UnresolvedHandles) > 0: + toolCalls[i].ArgumentResolution = ArgumentResolutionPartiallyResolved + case len(toolCalls[i].UnresolvedHandles) > 0: + toolCalls[i].ArgumentResolution = ArgumentResolutionUnresolved + case changed: + toolCalls[i].ArgumentResolution = ArgumentResolutionResolved + default: + toolCalls[i].ArgumentResolution = ArgumentResolutionUnchanged + } + } +} + +func jsonEquivalent(left, right string) bool { + var leftValue interface{} + var rightValue interface{} + if json.Unmarshal([]byte(left), &leftValue) != nil || json.Unmarshal([]byte(right), &rightValue) != nil { + return left == right + } + return reflect.DeepEqual(leftValue, rightValue) +} + +func uniqueSorted(values []string) []string { + seen := make(map[string]struct{}, len(values)) + for _, value := range values { + if value != "" { + seen[value] = struct{}{} + } + } + result := make([]string, 0, len(seen)) + for value := range seen { + result = append(result, value) + } + sort.Strings(result) + return result +} + +// DecodeResponse restores resources, expands citations, and decodes tool +// arguments in a non-streaming response. +func (r *Registry) DecodeResponse(response *types.ChatResponse) { + if r == nil || response == nil { + return + } + response.Content = r.DecodeOutputText(response.Content) + response.ReasoningContent = r.DecodeOutputText(response.ReasoningContent) + r.DecodeToolCalls(response.ToolCalls) +} + +// StreamDecoder creates one ordered decoder for a response text channel. +func (r *Registry) StreamDecoder() *StreamDecoder { + if r == nil { + return &StreamDecoder{} + } + return &StreamDecoder{ + resources: llmresource.NewStreamDecoder(r.resources), + sources: llmreference.NewStreamExpander(r.sources), + issues: NewHandleStreamDecoder(r.issues), + orphans: &orphanResourceStreamFilter{}, + } +} + +// OrphanResourceHandles reports model-generated resource handles with no +// backing durable reference. +func (r *Registry) OrphanResourceHandles(decoded string) []string { + if r == nil { + return nil + } + return r.resources.OrphanAliases(decoded) +} + +func (r *Registry) RegisterChunk(ref ChunkReference) string { + if r == nil || r.sources == nil { + return "" + } + return r.sources.RegisterChunk(ref) +} + +func (r *Registry) RegisterDocument(id string) string { + if r == nil || r.sources == nil { + return "" + } + return r.sources.RegisterDocument(id) +} + +func (r *Registry) RegisterKnowledgeBase(id string) string { + if r == nil || r.sources == nil { + return "" + } + return r.sources.RegisterKnowledgeBase(id) +} + +func (r *Registry) RegisterWeb(rawURL, title string) string { + if r == nil || r.sources == nil { + return "" + } + return r.sources.RegisterWeb(rawURL, title) +} + +func (r *Registry) RegisterSearchResults(results []*types.SearchResult) { + if r == nil || r.sources == nil { + return + } + r.sources.RegisterSearchResults(results) +} + +func (r *Registry) ChunkHandle(id string) string { + if r == nil || r.sources == nil { + return "" + } + return r.sources.ChunkAlias(id) +} + +// CompactKnownText replaces only previously registered durable source IDs. +func (r *Registry) CompactKnownText(text string) string { + if r == nil || r.sources == nil { + return text + } + text = r.resources.EncodeText(text) + return r.sources.CompactKnownText(text) +} + +// ModelToolResult renders a tool result using registered model handles. +func (r *Registry) ModelToolResult(result *types.ToolResult) string { + return r.ModelToolResultForTool("", result) +} + +// ModelToolResultForTool renders a result and applies any explicit private-ID +// policy owned by that built-in tool family. +func (r *Registry) ModelToolResultForTool(toolName string, result *types.ToolResult) string { + if result == nil { + return "" + } + if r == nil || r.sources == nil { + if result.Success { + return result.Output + } + return result.Error + } + // Protect durable resources and UUID-bearing summary slugs before the + // source codec sees any embedded document IDs. Encode once more afterwards + // for resource references rendered from structured ToolResult.Data. + copyResult := *result + copyResult.Output = r.encodeToolPrivateResult(toolName, result.Output) + copyResult.Output = r.resources.EncodeText(copyResult.Output) + copyResult.Error = r.resources.EncodeText(result.Error) + var modelOutput string + if sourceOutputAllowed(toolName) { + modelOutput = r.sources.ModelOutput(©Result) + } else if copyResult.Success { + modelOutput = copyResult.Output + } else if copyResult.Error != "" { + modelOutput = "Error: " + copyResult.Error + } else { + modelOutput = "Error: tool call failed" + } + // Even tools without structured source results can surface a known durable + // ID in validation errors or status text. Compact only explicitly declared + // built-ins; dynamic MCP output remains fully opaque. + if sourceCompactionAllowed(toolName) { + modelOutput = r.sources.CompactKnownText(modelOutput) + } + return r.resources.EncodeText(modelOutput) +} + +// DecodeOutputText applies the public citation policy to complete text. It is +// primarily used by non-streaming cleanup/fallback paths. +func (r *Registry) DecodeOutputText(text string) string { + if r == nil { + return text + } + text = r.resources.DecodeText(text) + text = r.resources.StripOrphanAliases(text) + text = r.sources.ExpandText(text) + return r.issues.DecodeKnownText(text) +} + +// StreamDecoder composes resource restoration and source-citation expansion +// so callers cannot split handle processing or apply stages in the wrong order. +type StreamDecoder struct { + resources *llmresource.StreamDecoder + sources *llmreference.StreamExpander + issues *HandleStreamDecoder + orphans *orphanResourceStreamFilter +} + +func (d *StreamDecoder) Feed(chunk string) string { + if d == nil { + return chunk + } + if d.resources != nil { + chunk = d.resources.Feed(chunk) + } + if d.sources != nil { + chunk = d.sources.Feed(chunk) + } + if d.issues != nil { + chunk = d.issues.Feed(chunk) + } + if d.orphans != nil { + chunk = d.orphans.Feed(chunk) + } + return chunk +} + +func (d *StreamDecoder) Flush() string { + if d == nil { + return "" + } + var tail string + if d.resources != nil { + tail = d.resources.Flush() + } + if d.sources != nil { + tail = d.sources.Feed(tail) + d.sources.Flush() + } + if d.issues != nil { + tail = d.issues.Feed(tail) + d.issues.Flush() + } + if d.orphans != nil { + tail = d.orphans.Feed(tail) + d.orphans.Flush() + } + return tail +} diff --git a/internal/modelcontext/registry_test.go b/internal/modelcontext/registry_test.go new file mode 100644 index 000000000..061dc4b5e --- /dev/null +++ b/internal/modelcontext/registry_test.go @@ -0,0 +1,419 @@ +package modelcontext + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/Tencent/WeKnora/internal/models/chat" + "github.com/Tencent/WeKnora/internal/types" +) + +func TestRegistryOwnsEncodingOrderForSummarySlugs(t *testing.T) { + const knowledgeID = "36283d3d-a495-40d2-a13f-e925fc322996" + registry := NewRegistry(true) + require.Equal(t, "d1", registry.RegisterDocument(knowledgeID)) + + messages := registry.EncodeMessages([]chat.Message{{ + Role: "tool", + Content: "[[summary/" + knowledgeID + "|Summary]] document=" + knowledgeID, + }}) + require.Contains(t, messages[0].Content, "[[res://0001|Summary]]") + require.Contains(t, messages[0].Content, "document=d1") + require.NotContains(t, messages[0].Content, "summary/d1") + + calls := []types.LLMToolCall{{Function: types.FunctionCall{ + Name: "wiki_read_source_doc", + Arguments: `{"slugs":["res://0001"],"knowledge_id":"d1"}`, + }}} + registry.DecodeToolCalls(calls) + require.JSONEq(t, `{"slugs":["summary/`+knowledgeID+`"],"knowledge_id":"`+knowledgeID+`"}`, calls[0].Function.Arguments) + require.JSONEq(t, `{"slugs":["res://0001"],"knowledge_id":"d1"}`, calls[0].ModelArguments) + require.Equal(t, ArgumentResolutionResolved, calls[0].ArgumentResolution) + require.Empty(t, calls[0].UnresolvedHandles) + persisted, err := json.Marshal(calls[0]) + require.NoError(t, err) + require.NotContains(t, string(persisted), `"res://0001"`) + require.NotContains(t, string(persisted), `"d1"`) + require.NotContains(t, string(persisted), ArgumentResolutionResolved) +} + +func TestRegistryProtocolOwnsResourceHandleRules(t *testing.T) { + prompt := NewRegistry(true).ProtocolPrompt() + require.Contains(t, prompt, "Source handling protocol") + require.Contains(t, prompt, "Resource handle protocol") + require.Contains(t, prompt, "res://NNNN") +} + +func TestRegistryAuditsUnresolvedAndPartiallyResolvedToolHandles(t *testing.T) { + registry := NewRegistry(true) + registry.RegisterKnowledgeBase("kb-real") + + calls := []types.LLMToolCall{{Function: types.FunctionCall{ + Name: "knowledge_search", + Arguments: `{"knowledge_base_ids":["b1","b99"],"query":"b98","asset":"res://0099"}`, + }}} + registry.DecodeToolCalls(calls) + + require.JSONEq(t, + `{"knowledge_base_ids":["kb-real","b99"],"query":"b98","asset":"res://0099"}`, + calls[0].Function.Arguments, + ) + require.Equal(t, ArgumentResolutionPartiallyResolved, calls[0].ArgumentResolution) + require.Equal(t, []string{"b99", "res://0099"}, calls[0].UnresolvedHandles) + // Alias-shaped free text is not an ID and must not be reported as a source + // resolution failure. + require.NotContains(t, calls[0].UnresolvedHandles, "b98") +} + +func TestRegistryAuditsUnchangedJSONFormattingAsUnchanged(t *testing.T) { + registry := NewRegistry(true) + calls := []types.LLMToolCall{{Function: types.FunctionCall{Arguments: `{ "query": "hello" }`}}} + registry.DecodeToolCalls(calls) + require.Equal(t, ArgumentResolutionUnchanged, calls[0].ArgumentResolution) +} + +func TestRegistryDecodesKnownHandlesEmbeddedInBuiltInSQL(t *testing.T) { + registry := NewRegistry(true) + registry.RegisterDocument("doc-real") + registry.RegisterKnowledgeBase("kb-real") + calls := []types.LLMToolCall{{ + Function: types.FunctionCall{ + Name: "database_query", + Arguments: `{"sql":"SELECT * FROM knowledges WHERE id = 'd1' AND knowledge_base_id = 'b1'"}`, + }, + }} + registry.DecodeToolCalls(calls) + require.JSONEq(t, + `{"sql":"SELECT * FROM knowledges WHERE id = 'doc-real' AND knowledge_base_id = 'kb-real'"}`, + calls[0].Function.Arguments, + ) + require.Equal(t, ArgumentResolutionResolved, calls[0].ArgumentResolution) +} + +func TestRegistryOnlyDecodesQuotedSQLHandlesAndRejectsUnknownOnes(t *testing.T) { + registry := NewRegistry(true) + registry.RegisterDocument("doc-real") + + calls := []types.LLMToolCall{{ + Function: types.FunctionCall{ + Name: "database_query", + Arguments: `{"sql":"SELECT d1.id FROM documents d1 WHERE d1.id = 'd1' AND d1.parent_id = 'd99'"}`, + }, + }} + registry.DecodeToolCalls(calls) + require.JSONEq(t, + `{"sql":"SELECT d1.id FROM documents d1 WHERE d1.id = 'doc-real' AND d1.parent_id = 'd99'"}`, + calls[0].Function.Arguments, + ) + require.Equal(t, []string{"d99"}, calls[0].UnresolvedHandles) + require.Equal(t, ArgumentResolutionPartiallyResolved, calls[0].ArgumentResolution) +} + +func TestRegistryReplaysStructuredAndPrivateHandlesWithoutAliasDrift(t *testing.T) { + registry := NewRegistry(true) + registry.RegisterDocument("doc-real") + registry.RegisterKnowledgeBase("kb-real") + registry.ModelToolResultForTool("wiki_read_issue", &types.ToolResult{ + Success: true, + Output: `{"id":"issue-real","knowledge_base_id":"kb-real","suspected_knowledge_ids":["doc-real"]}`, + }) + + messages := []chat.Message{ + { + Role: "assistant", + ToolCalls: []chat.ToolCall{{Function: chat.FunctionCall{ + Name: "database_query", + Arguments: `{"sql":"SELECT * FROM knowledges WHERE id = 'doc-real' AND knowledge_base_id = 'kb-real'"}`, + }}}, + }, + { + Role: "tool", + Name: "wiki_read_issue", + Content: `{"id":"issue-real","knowledge_base_id":"kb-real","suspected_knowledge_ids":["doc-real"]}`, + }, + { + Role: "assistant", + ToolCalls: []chat.ToolCall{{Function: chat.FunctionCall{ + Name: "wiki_update_issue", + Arguments: `{"issue_id":"issue-real","status":"resolved"}`, + }}}, + }, + } + + first := registry.EncodeMessages(messages) + require.JSONEq(t, + `{"sql":"SELECT * FROM knowledges WHERE id = 'd1' AND knowledge_base_id = 'b1'"}`, + first[0].ToolCalls[0].Function.Arguments, + ) + require.JSONEq(t, + `{"id":"i1","knowledge_base_id":"b1","suspected_knowledge_ids":["d1"]}`, + first[1].Content, + ) + require.JSONEq(t, `{"issue_id":"i1","status":"resolved"}`, first[2].ToolCalls[0].Function.Arguments) + + second := registry.EncodeMessages(first) + require.JSONEq(t, first[0].ToolCalls[0].Function.Arguments, second[0].ToolCalls[0].Function.Arguments) + require.JSONEq(t, first[1].Content, second[1].Content) + require.JSONEq(t, first[2].ToolCalls[0].Function.Arguments, second[2].ToolCalls[0].Function.Arguments) +} + +func TestRegistryOwnsWikiIssueHandles(t *testing.T) { + registry := NewRegistry(true) + modelResult := registry.ModelToolResultForTool("wiki_read_issue", &types.ToolResult{ + Success: true, + Output: `[{"id":"issue-uuid","knowledge_base_id":"kb-real","suspected_knowledge_ids":["doc-real"],"slug":"concept/a","status":"pending"}]`, + }) + require.JSONEq(t, `[{"id":"i1","knowledge_base_id":"b1","suspected_knowledge_ids":["d1"],"slug":"concept/a","status":"pending"}]`, modelResult) + + calls := []types.LLMToolCall{{Function: types.FunctionCall{ + Name: "wiki_update_issue", + Arguments: `{"issue_id":"i1","status":"resolved"}`, + }}} + registry.DecodeToolCalls(calls) + require.JSONEq(t, `{"issue_id":"issue-uuid","status":"resolved"}`, calls[0].Function.Arguments) + require.Equal(t, ArgumentResolutionResolved, calls[0].ArgumentResolution) + + decoder := registry.StreamDecoder() + got := decoder.Feed("updated i") + decoder.Feed("1 successfully") + decoder.Flush() + require.Equal(t, "updated issue-uuid successfully", got) +} + +func TestRegistryDecodesWikiReadIssueHandleArgument(t *testing.T) { + registry := NewRegistry(true) + registry.ModelToolResultForTool("wiki_read_issue", &types.ToolResult{ + Success: true, + Output: `[{"id":"issue-uuid","slug":"concept/a","status":"pending"}]`, + }) + + // The wiki_read_issue schema instructs the model to pass the short iN + // handle back as issue_id; it must resolve to the durable issue ID before + // the tool executes. + calls := []types.LLMToolCall{{Function: types.FunctionCall{ + Name: "wiki_read_issue", + Arguments: `{"issue_id":"i1"}`, + }}} + registry.DecodeToolCalls(calls) + require.JSONEq(t, `{"issue_id":"issue-uuid"}`, calls[0].Function.Arguments) + require.Equal(t, ArgumentResolutionResolved, calls[0].ArgumentResolution) + require.Empty(t, calls[0].UnresolvedHandles) + + unknown := []types.LLMToolCall{{Function: types.FunctionCall{ + Name: "wiki_read_issue", + Arguments: `{"issue_id":"i9"}`, + }}} + registry.DecodeToolCalls(unknown) + require.Equal(t, ArgumentResolutionUnresolved, unknown[0].ArgumentResolution) + require.Equal(t, []string{"i9"}, unknown[0].UnresolvedHandles) +} + +func TestRegistryReportsUnknownWikiIssueHandle(t *testing.T) { + registry := NewRegistry(true) + calls := []types.LLMToolCall{{Function: types.FunctionCall{ + Name: "wiki_update_issue", + Arguments: `{"issue_id":"i9","status":"resolved"}`, + }}} + registry.DecodeToolCalls(calls) + require.Equal(t, ArgumentResolutionUnresolved, calls[0].ArgumentResolution) + require.Equal(t, []string{"i9"}, calls[0].UnresolvedHandles) +} + +func TestRegistryDoesNotApplyBuiltInFieldPoliciesToDynamicTools(t *testing.T) { + registry := NewRegistry(true) + registry.RegisterDocument("doc-real") + registry.RegisterWeb("https://example.com", "Example") + registry.ModelToolResultForTool("wiki_read_issue", &types.ToolResult{ + Success: true, + Output: `{"id":"issue-real"}`, + }) + calls := []types.LLMToolCall{{Function: types.FunctionCall{ + Name: "dynamic_mcp_tool", + Arguments: `{"issue_id":"i1","knowledge_id":"d1","url":"w1","sql":"SELECT 'd1'"}`, + }}} + registry.DecodeToolCalls(calls) + require.JSONEq(t, `{"issue_id":"i1","knowledge_id":"d1","url":"w1","sql":"SELECT 'd1'"}`, calls[0].Function.Arguments) + require.Equal(t, ArgumentResolutionUnchanged, calls[0].ArgumentResolution) + require.Empty(t, calls[0].UnresolvedHandles) + + messages := registry.EncodeMessages([]chat.Message{ + { + Role: "assistant", + ToolCalls: []chat.ToolCall{{Function: chat.FunctionCall{ + Name: "dynamic_mcp_tool", + Arguments: `{"knowledge_id":"doc-real"}`, + }}}, + }, + { + Role: "tool", + Name: "dynamic_mcp_tool", + Content: `mcp-owned-id doc-real`, + }, + }) + require.JSONEq(t, `{"knowledge_id":"doc-real"}`, messages[0].ToolCalls[0].Function.Arguments) + require.Equal(t, `mcp-owned-id doc-real`, messages[1].Content) + + modelOutput := registry.ModelToolResultForTool("dynamic_mcp_tool", &types.ToolResult{ + Success: true, + Output: `mcp-owned-id doc-real`, + }) + require.Equal(t, `mcp-owned-id doc-real`, modelOutput) +} + +func TestRegistryCompactsKnownIDsInBuiltInValidationErrors(t *testing.T) { + registry := NewRegistry(true) + registry.RegisterDocument("doc-real") + registry.RegisterKnowledgeBase("kb-real") + + got := registry.ModelToolResultForTool("wiki_write_page", &types.ToolResult{ + Success: false, + Error: "document doc-real belongs to knowledge base kb-real", + }) + require.Equal(t, "Error: document d1 belongs to knowledge base b1", got) +} + +func TestRegistryCompactsDatabaseQueryIDColumnsForBuiltInFollowUps(t *testing.T) { + registry := NewRegistry(true) + modelOutput := registry.ModelToolResultForTool("database_query", &types.ToolResult{ + Success: true, + Output: "knowledge_id | knowledge_base_id\ndoc-real | kb-real", + Data: map[string]interface{}{ + "display_type": "database_query", + "rows": []map[string]interface{}{{ + "knowledge_id": "doc-real", + "knowledge_base_id": "kb-real", + }}, + }, + }) + require.Contains(t, modelOutput, "d1 | b1") + + calls := []types.LLMToolCall{{Function: types.FunctionCall{ + Name: "wiki_read_source_doc", + Arguments: `{"knowledge_id":"d1"}`, + }}} + registry.DecodeToolCalls(calls) + require.JSONEq(t, `{"knowledge_id":"doc-real"}`, calls[0].Function.Arguments) +} + +func TestRegistryDecodesCanonicalArgumentsForEveryBuiltInReferenceTool(t *testing.T) { + registry := NewRegistry(true) + registry.RegisterDocument("doc-real") + registry.RegisterKnowledgeBase("kb-real") + registry.RegisterChunk(ChunkReference{ChunkID: "chunk-real", KnowledgeID: "doc-real", KnowledgeBaseID: "kb-real"}) + registry.RegisterWeb("https://example.com/page", "Example") + registry.EncodeMessages([]chat.Message{{Role: "user", Content: "summary/00000000-0000-0000-0000-000000000001"}}) + registry.ModelToolResultForTool("wiki_read_issue", &types.ToolResult{Success: true, Output: `{"id":"issue-real"}`}) + + tests := []struct { + name string + tool string + raw string + want string + }{ + {"knowledge search KB", "knowledge_search", `{"queries":["d1"],"knowledge_base_ids":["b1"]}`, `{"queries":["d1"],"knowledge_base_ids":["kb-real"]}`}, + {"list document", "list_knowledge_chunks", `{"knowledge_id":"d1"}`, `{"knowledge_id":"doc-real"}`}, + {"list chunk", "list_knowledge_chunks", `{"chunk_id":"c1"}`, `{"chunk_id":"chunk-real"}`}, + {"document info", "get_document_info", `{"knowledge_ids":["d1"],"faq_ids":["c1"]}`, `{"knowledge_ids":["doc-real"],"faq_ids":["chunk-real"]}`}, + {"knowledge graph", "query_knowledge_graph", `{"knowledge_base_ids":["b1"],"query":"topic"}`, `{"knowledge_base_ids":["kb-real"],"query":"topic"}`}, + {"data analysis SQL", "data_analysis", `{"knowledge_id":"d1","sql":"SELECT * FROM 'd1'"}`, `{"knowledge_id":"doc-real","sql":"SELECT * FROM 'doc-real'"}`}, + {"data schema", "data_schema", `{"knowledge_id":"d1"}`, `{"knowledge_id":"doc-real"}`}, + {"database SQL", "database_query", `{"sql":"SELECT * FROM chunks WHERE knowledge_base_id='b1'"}`, `{"sql":"SELECT * FROM chunks WHERE knowledge_base_id='kb-real'"}`}, + {"web fetch", "web_fetch", `{"items":[{"url":"w1","prompt":"read"}]}`, `{"items":[{"url":"https://example.com/page","prompt":"read"}]}`}, + {"wiki source", "wiki_read_source_doc", `{"knowledge_id":"d1"}`, `{"knowledge_id":"doc-real"}`}, + {"wiki source refs", "wiki_write_page", `{"slug":"res://0001","source_refs":["d1"]}`, `{"slug":"summary/00000000-0000-0000-0000-000000000001","source_refs":["doc-real"]}`}, + {"wiki suspected refs", "wiki_flag_issue", `{"slug":"concept/a","suspected_knowledge_ids":["d1"]}`, `{"slug":"concept/a","suspected_knowledge_ids":["doc-real"]}`}, + {"wiki search KB", "wiki_search", `{"queries":["topic"],"knowledge_base_id":"b1"}`, `{"queries":["topic"],"knowledge_base_id":"kb-real"}`}, + {"wiki issue", "wiki_update_issue", `{"issue_id":"i1","status":"resolved"}`, `{"issue_id":"issue-real","status":"resolved"}`}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + calls := []types.LLMToolCall{{Function: types.FunctionCall{Name: test.tool, Arguments: test.raw}}} + registry.DecodeToolCalls(calls) + require.JSONEq(t, test.want, calls[0].Function.Arguments) + require.NotEqual(t, ArgumentResolutionUnresolved, calls[0].ArgumentResolution) + }) + } +} + +func TestModelToolResultProtectsSummarySlugBeforeSourceCompaction(t *testing.T) { + const knowledgeID = "07a20bb1-a662-47cf-9929-06fb5d5b5b5e" + const kbID = "250368ff-f5a2-4e9e-868a-07bc9b857c44" + registry := NewRegistry(true) + registry.RegisterDocument(knowledgeID) + registry.RegisterKnowledgeBase(kbID) + + got := registry.ModelToolResult(&types.ToolResult{Success: true, Output: "" + kbID + "\n" + + "[[summary/" + knowledgeID + "|Summary]]\n" + + "" + knowledgeID + "", + }) + require.Contains(t, got, "[[res://0001|Summary]]") + require.Contains(t, got, "b1") + require.Contains(t, got, "d1") + require.NotContains(t, got, "summary/d1") +} + +func TestRegistryStreamDecoderRestoresSplitResourceAndCitationHandles(t *testing.T) { + registry := NewRegistry(true) + registry.EncodeMessages([]chat.Message{{Role: "user", Content: "resource://AbCdEfGhIjKlMnOpQrStUv"}}) + registry.RegisterChunk(ChunkReference{ + ChunkID: "chunk-real", + KnowledgeID: "doc-real", + KnowledgeBaseID: "kb-real", + DocumentTitle: "Doc", + }) + + decoder := registry.StreamDecoder() + got := decoder.Feed("image res://0") + + decoder.Feed("001 claim ") + decoder.Flush() + require.Contains(t, got, "resource://AbCdEfGhIjKlMnOpQrStUv") + require.Contains(t, got, `chunk_id="chunk-real"`) +} + +func TestRegistryDropsUnknownResourceHandlesFromCompleteAndStreamOutput(t *testing.T) { + registry := NewRegistry(true) + require.Equal(t, "broken ", registry.DecodeOutputText("broken res://9999")) + + decoder := registry.StreamDecoder() + got := decoder.Feed("broken res:/") + decoder.Feed("/99") + decoder.Feed("99 end") + decoder.Flush() + require.Equal(t, "broken end", got) +} + +func TestHandleTableRoundTripAndIsolation(t *testing.T) { + chunks := NewHandleTable("c", 3, 0) + require.Equal(t, "c000", chunks.Register("chunk-a")) + require.Equal(t, "c000", chunks.Register("chunk-a")) + require.Equal(t, "c001", chunks.Register("chunk-b")) + require.Equal(t, "chunk-b", mustResolve(t, chunks, "c001")) + + slugs := NewHandleTable("ref-", 0, 1) + require.Equal(t, "ref-1", slugs.Register("summary/uuid")) + require.Equal(t, "summary/uuid", mustResolve(t, slugs, "ref-1")) + require.Equal(t, "", mustNotResolve(t, slugs, "c000")) +} + +func TestHandleTableKnownTextRoundTrip(t *testing.T) { + table := NewHandleTable("i", 0, 1) + require.Equal(t, "i1", table.Register("issue-uuid")) + require.Equal(t, "read i1", table.EncodeKnownText("read issue-uuid")) + require.Equal(t, "read issue-uuid", table.DecodeKnownText("read i1")) + require.Equal(t, "read issue-$1", func() string { + literal := NewHandleTable("i", 0, 1) + literal.Register("issue-$1") + return literal.DecodeKnownText("read i1") + }()) +} + +func mustResolve(t *testing.T, table *HandleTable, handle string) string { + t.Helper() + value, ok := table.Resolve(handle) + require.True(t, ok) + return value +} + +func mustNotResolve(t *testing.T, table *HandleTable, handle string) string { + t.Helper() + value, ok := table.Resolve(handle) + require.False(t, ok) + return value +} diff --git a/internal/modelcontext/tool_policy.go b/internal/modelcontext/tool_policy.go new file mode 100644 index 000000000..fac2f40bb --- /dev/null +++ b/internal/modelcontext/tool_policy.go @@ -0,0 +1,274 @@ +package modelcontext + +import ( + "encoding/json" + "regexp" + "strings" + + "github.com/Tencent/WeKnora/internal/models/chat" + "github.com/Tencent/WeKnora/internal/types" +) + +const ( + toolDatabaseQuery = "database_query" + toolDataAnalysis = "data_analysis" + toolWikiReadIssue = "wiki_read_issue" + toolWikiUpdateIssue = "wiki_update_issue" +) + +var issueHandleShapeRE = regexp.MustCompile(`^i[1-9][0-9]*$`) + +type toolHandlePolicy struct { + sourceIDKeys map[string]struct{} + sourceTextKeys map[string]struct{} + sourceOutput bool + decodedIssueIDKeys map[string]struct{} + encodedIssueIDKeys map[string]struct{} + encodeKnownIssueIDs bool +} + +// toolHandlePolicies is the complete allowlist for built-in tool fields whose +// identifiers need more than the generic resource/source codecs. Field names +// alone are deliberately insufficient: a dynamic MCP tool may use the same +// name with unrelated semantics and must remain opaque. +var toolHandlePolicies = map[string]toolHandlePolicy{ + "knowledge_search": { + sourceIDKeys: map[string]struct{}{"knowledge_base_ids": {}}, + sourceOutput: true, + }, + "grep_chunks": { + sourceOutput: true, + }, + "list_knowledge_chunks": { + sourceIDKeys: map[string]struct{}{"knowledge_id": {}, "faq_id": {}, "chunk_id": {}}, + sourceOutput: true, + }, + "get_document_info": { + sourceIDKeys: map[string]struct{}{"knowledge_ids": {}, "faq_ids": {}}, + sourceOutput: true, + }, + "query_knowledge_graph": { + sourceIDKeys: map[string]struct{}{"knowledge_base_ids": {}}, + sourceOutput: true, + }, + toolDatabaseQuery: { + sourceTextKeys: map[string]struct{}{"sql": {}}, + sourceOutput: true, + }, + toolDataAnalysis: { + sourceIDKeys: map[string]struct{}{"knowledge_id": {}}, + sourceTextKeys: map[string]struct{}{"sql": {}}, + }, + "data_schema": { + sourceIDKeys: map[string]struct{}{"knowledge_id": {}}, + }, + "web_fetch": { + sourceIDKeys: map[string]struct{}{"url": {}, "urls": {}}, + sourceOutput: true, + }, + "web_search": { + sourceOutput: true, + }, + "wiki_read_page": { + sourceOutput: true, + }, + "wiki_read_source_doc": { + sourceIDKeys: map[string]struct{}{"knowledge_id": {}}, + sourceOutput: true, + }, + "wiki_write_page": { + sourceIDKeys: map[string]struct{}{"source_refs": {}}, + }, + "wiki_replace_text": { + sourceIDKeys: map[string]struct{}{"source_refs": {}}, + }, + "wiki_flag_issue": { + sourceIDKeys: map[string]struct{}{"suspected_knowledge_ids": {}}, + }, + "wiki_search": { + sourceIDKeys: map[string]struct{}{"knowledge_base_id": {}}, + sourceOutput: true, + }, + toolWikiReadIssue: { + decodedIssueIDKeys: map[string]struct{}{"issue_id": {}}, + encodedIssueIDKeys: map[string]struct{}{"id": {}}, + encodeKnownIssueIDs: true, + sourceOutput: true, + }, + toolWikiUpdateIssue: { + decodedIssueIDKeys: map[string]struct{}{"issue_id": {}}, + encodeKnownIssueIDs: true, + }, +} + +func sourceArgumentAllowed(toolName, key string) bool { + policy, ok := toolHandlePolicies[toolName] + if !ok { + return false + } + _, ok = policy.sourceIDKeys[strings.ToLower(key)] + return ok +} + +func sourceOutputAllowed(toolName string) bool { + if toolName == "" { + // Preserve the generic low-level facade for non-Agent callers. Agent + // lifecycles always provide the concrete tool name. + return true + } + policy, ok := toolHandlePolicies[toolName] + return ok && policy.sourceOutput +} + +func sourceCompactionAllowed(toolName string) bool { + if toolName == "" { + return true + } + _, ok := toolHandlePolicies[toolName] + return ok +} + +// decodeToolPolicies handles the deliberately small set of arguments whose +// model handles are embedded in structured text or belong to a tool-private +// identity space. Generic free text is never rewritten. +func (r *Registry) decodeToolPolicies(call *types.LLMToolCall) { + if r == nil || call == nil { + return + } + policy, ok := toolHandlePolicies[call.Function.Name] + if !ok { + return + } + call.Function.Arguments = rewriteJSONStringValues( + call.Function.Arguments, + func(key, value string) string { + if _, ok := policy.sourceTextKeys[key]; ok { + return r.sources.DecodeKnownQuotedText(value) + } + if _, ok := policy.decodedIssueIDKeys[key]; ok && issueHandleShapeRE.MatchString(strings.TrimSpace(value)) { + if durable, ok := r.issues.Resolve(value); ok { + return durable + } + } + return value + }, + ) +} + +// encodeReplayedToolPolicies applies the exact inverse of decodeToolPolicies +// to assistant tool calls replayed into a later model round. Only durable +// values already proven and registered in this request are compacted. +func (r *Registry) encodeReplayedToolPolicies(call *chat.ToolCall) { + if r == nil || call == nil { + return + } + policy, ok := toolHandlePolicies[call.Function.Name] + if !ok { + return + } + call.Function.Arguments = rewriteJSONStringValues( + call.Function.Arguments, + func(key, value string) string { + if _, ok := policy.sourceTextKeys[key]; ok { + return r.sources.CompactKnownText(value) + } + if _, ok := policy.decodedIssueIDKeys[key]; ok { + return r.issues.EncodeKnownText(value) + } + return value + }, + ) +} + +func (r *Registry) unresolvedPrivateToolHandles(toolName, raw string) []string { + if r == nil || raw == "" { + return nil + } + policy, ok := toolHandlePolicies[toolName] + if !ok || (len(policy.decodedIssueIDKeys) == 0 && len(policy.sourceTextKeys) == 0) { + return nil + } + var unresolved []string + _ = walkJSONStringValues(raw, func(key, value string) string { + value = strings.TrimSpace(value) + if _, ok := policy.sourceTextKeys[key]; ok { + unresolved = append(unresolved, r.sources.UnresolvedQuotedTextHandles(value)...) + } + if _, ok := policy.decodedIssueIDKeys[key]; ok && issueHandleShapeRE.MatchString(value) { + if _, ok := r.issues.Resolve(value); !ok { + unresolved = append(unresolved, value) + } + } + return value + }) + return unresolved +} + +// encodeToolPrivateResult registers and compacts identifiers that are local to +// one built-in tool family. Wiki issue IDs are the only such identity today; +// MCP tools remain opaque unless they add an explicit policy here. +func (r *Registry) encodeToolPrivateResult(toolName, output string) string { + if r == nil || output == "" { + return output + } + policy, ok := toolHandlePolicies[toolName] + if !ok { + return output + } + if len(policy.encodedIssueIDKeys) > 0 { + output = rewriteJSONStringValues(output, func(key, value string) string { + if _, ok := policy.encodedIssueIDKeys[key]; ok && strings.TrimSpace(value) != "" { + if issueHandleShapeRE.MatchString(strings.TrimSpace(value)) { + // Model-facing tool messages can be replayed across several + // rounds. Never treat an existing temporary handle as a new + // durable issue identity and allocate i2/i3 drift. + return value + } + return r.issues.Register(value) + } + return value + }) + } + if policy.encodeKnownIssueIDs { + output = r.issues.EncodeKnownText(output) + } + return output +} + +func rewriteJSONStringValues(raw string, rewrite func(key, value string) string) string { + var value interface{} + if json.Unmarshal([]byte(raw), &value) != nil { + return raw + } + value = walkJSONValue("", value, rewrite) + encoded, err := json.Marshal(value) + if err != nil { + return raw + } + return string(encoded) +} + +func walkJSONStringValues(raw string, rewrite func(key, value string) string) error { + var value interface{} + if err := json.Unmarshal([]byte(raw), &value); err != nil { + return err + } + walkJSONValue("", value, rewrite) + return nil +} + +func walkJSONValue(key string, value interface{}, rewrite func(key, value string) string) interface{} { + switch typed := value.(type) { + case string: + return rewrite(strings.ToLower(key), typed) + case []interface{}: + for i := range typed { + typed[i] = walkJSONValue(key, typed[i], rewrite) + } + case map[string]interface{}: + for childKey, item := range typed { + typed[childKey] = walkJSONValue(childKey, item, rewrite) + } + } + return value +} diff --git a/internal/models/chat/langfuse_wrapper.go b/internal/models/chat/langfuse_wrapper.go index 64990676d..d1cf1abff 100644 --- a/internal/models/chat/langfuse_wrapper.go +++ b/internal/models/chat/langfuse_wrapper.go @@ -114,7 +114,11 @@ func (l *langfuseChat) ChatStream(ctx context.Context, messages []Message, opts usage = resp.Usage } if len(resp.ToolCalls) > 0 { - toolCalls = resp.ToolCalls + // The downstream model-context registry decodes arguments in + // place before tool execution. Snapshot the provider payload so + // the generation observation remains the exact model output and + // cannot be changed through the shared slice backing array. + toolCalls = snapshotLangfuseToolCalls(resp.ToolCalls) } if resp.FinishReason != "" { finishReason = resp.FinishReason @@ -130,6 +134,10 @@ func (l *langfuseChat) ChatStream(ctx context.Context, messages []Message, opts return wrapped, nil } +func snapshotLangfuseToolCalls(toolCalls []types.LLMToolCall) []types.LLMToolCall { + return append([]types.LLMToolCall(nil), toolCalls...) +} + func buildLangfuseMessages(messages []Message) []map[string]interface{} { out := make([]map[string]interface{}, 0, len(messages)) for _, m := range messages { diff --git a/internal/models/chat/langfuse_wrapper_test.go b/internal/models/chat/langfuse_wrapper_test.go index 841ff8a6e..c911108ec 100644 --- a/internal/models/chat/langfuse_wrapper_test.go +++ b/internal/models/chat/langfuse_wrapper_test.go @@ -32,6 +32,19 @@ func TestBuildLangfuseGenerationOutput(t *testing.T) { } } +func TestSnapshotLangfuseToolCallsKeepsModelArguments(t *testing.T) { + providerCalls := []types.LLMToolCall{{ + ID: "call_1", + Function: types.FunctionCall{Name: "wiki_read_page", Arguments: `{"slugs":["res://0001"]}`}, + }} + snapshot := snapshotLangfuseToolCalls(providerCalls) + providerCalls[0].Function.Arguments = `{"slugs":["summary/uuid"]}` + + if got := snapshot[0].Function.Arguments; got != `{"slugs":["res://0001"]}` { + t.Fatalf("Langfuse snapshot was mutated to %s", got) + } +} + func TestBuildLangfuseMessagesReasoningContent(t *testing.T) { msgs := buildLangfuseMessages([]Message{ {Role: "assistant", ReasoningContent: "chain of thought", ToolCalls: []ToolCall{{ID: "tc1"}}}, diff --git a/internal/types/chat.go b/internal/types/chat.go index a7fcc9b26..a05c7c186 100644 --- a/internal/types/chat.go +++ b/internal/types/chat.go @@ -79,6 +79,15 @@ type LLMToolCall struct { Type string `json:"type"` // "function" Function FunctionCall `json:"function"` ProviderMetadata ToolCallMetadata `json:"provider_metadata,omitempty"` + + // ModelArguments and the resolution fields are request-local observability + // state. ModelArguments preserves the exact JSON emitted by the model while + // Function.Arguments is decoded to durable application identifiers before + // tool execution. These fields must never be sent back to a provider or + // persisted in chat history. + ModelArguments string `json:"-"` + ArgumentResolution string `json:"-"` + UnresolvedHandles []string `json:"-"` } // ToolCallMetadata carries provider-specific tool-call state that must round-trip diff --git a/internal/types/search.go b/internal/types/search.go index 7544b1684..de031b98a 100644 --- a/internal/types/search.go +++ b/internal/types/search.go @@ -77,7 +77,26 @@ func HasKnowledgeRetrievalScope( knowledgeBaseIDs []string, knowledgeIDs []string, ) bool { - return len(searchTargets) > 0 || len(knowledgeBaseIDs) > 0 || len(knowledgeIDs) > 0 + for _, id := range knowledgeBaseIDs { + if id != "" { + return true + } + } + for _, id := range knowledgeIDs { + if id != "" { + return true + } + } + for _, target := range searchTargets { + if target == nil || target.KnowledgeBaseID == "" { + continue + } + if target.Type == SearchTargetTypeKnowledgeBase || len(target.KnowledgeIDs) > 0 || + len(target.TagIDs) > 0 || len(target.ScopeTagIDs) > 0 { + return true + } + } + return false } // GetAllKnowledgeBaseIDs returns all unique knowledge base IDs from the search targets @@ -85,6 +104,9 @@ func (st SearchTargets) GetAllKnowledgeBaseIDs() []string { seen := make(map[string]bool) var result []string for _, t := range st { + if t == nil || t.KnowledgeBaseID == "" { + continue + } if !seen[t.KnowledgeBaseID] { seen[t.KnowledgeBaseID] = true result = append(result, t.KnowledgeBaseID) @@ -97,7 +119,7 @@ func (st SearchTargets) GetAllKnowledgeBaseIDs() []string { func (st SearchTargets) GetKBTenantMap() map[string]uint64 { result := make(map[string]uint64) for _, t := range st { - if t.KnowledgeBaseID != "" { + if t != nil && t.KnowledgeBaseID != "" { result[t.KnowledgeBaseID] = t.TenantID } } @@ -108,7 +130,7 @@ func (st SearchTargets) GetKBTenantMap() map[string]uint64 { // Returns 0 if not found func (st SearchTargets) GetTenantIDForKB(kbID string) uint64 { for _, t := range st { - if t.KnowledgeBaseID == kbID { + if t != nil && t.KnowledgeBaseID == kbID { return t.TenantID } } @@ -118,7 +140,7 @@ func (st SearchTargets) GetTenantIDForKB(kbID string) uint64 { // ContainsKB checks if the search targets contain a given knowledge base ID func (st SearchTargets) ContainsKB(kbID string) bool { for _, t := range st { - if t.KnowledgeBaseID == kbID { + if t != nil && t.KnowledgeBaseID == kbID { return true } } diff --git a/internal/types/search_scope_test.go b/internal/types/search_scope_test.go index 6e4966633..fbe9a18b4 100644 --- a/internal/types/search_scope_test.go +++ b/internal/types/search_scope_test.go @@ -11,6 +11,8 @@ func TestHasKnowledgeRetrievalScope(t *testing.T) { want bool }{ {name: "empty", want: false}, + {name: "nil target", searchTargets: SearchTargets{nil}, want: false}, + {name: "empty document target", searchTargets: SearchTargets{{Type: SearchTargetTypeKnowledge, KnowledgeBaseID: "kb-1"}}, want: false}, {name: "knowledge base IDs", knowledgeBaseIDs: []string{"kb-1"}, want: true}, {name: "knowledge IDs", knowledgeIDs: []string{"doc-1"}, want: true}, { diff --git a/internal/utils/inject.go b/internal/utils/inject.go index 5d96a06b4..782de8a67 100644 --- a/internal/utils/inject.go +++ b/internal/utils/inject.go @@ -1118,16 +1118,16 @@ func buildKnowledgeScopeCondition(alias string, scopes []SearchScope) string { } kbID := quoteString(scope.KnowledgeBaseID) switch { - case len(scope.KnowledgeIDs) > 0: - clauses = append(clauses, fmt.Sprintf( - "(%s.knowledge_base_id = %s AND %s.id IN (%s))", - alias, kbID, alias, strings.Join(quoteStringSlice(scope.KnowledgeIDs), ", "), - )) case len(scope.TagIDs) > 0: clauses = append(clauses, fmt.Sprintf( "(%s.knowledge_base_id = %s AND EXISTS (SELECT 1 FROM knowledge_tag_relations ktr WHERE ktr.knowledge_id = %s.id AND ktr.tag_id IN (%s)))", alias, kbID, alias, strings.Join(quoteStringSlice(scope.TagIDs), ", "), )) + case len(scope.KnowledgeIDs) > 0: + clauses = append(clauses, fmt.Sprintf( + "(%s.knowledge_base_id = %s AND %s.id IN (%s))", + alias, kbID, alias, strings.Join(quoteStringSlice(scope.KnowledgeIDs), ", "), + )) default: clauses = append(clauses, fmt.Sprintf("%s.knowledge_base_id = %s", alias, kbID)) } @@ -1143,16 +1143,16 @@ func buildChunkScopeCondition(alias string, scopes []SearchScope) string { } kbID := quoteString(scope.KnowledgeBaseID) switch { - case len(scope.KnowledgeIDs) > 0: - clauses = append(clauses, fmt.Sprintf( - "(%s.knowledge_base_id = %s AND %s.knowledge_id IN (%s))", - alias, kbID, alias, strings.Join(quoteStringSlice(scope.KnowledgeIDs), ", "), - )) case len(scope.TagIDs) > 0: clauses = append(clauses, fmt.Sprintf( "(%s.knowledge_base_id = %s AND EXISTS (SELECT 1 FROM knowledge_tag_relations ktr WHERE ktr.knowledge_id = %s.knowledge_id AND ktr.tag_id IN (%s)))", alias, kbID, alias, strings.Join(quoteStringSlice(scope.TagIDs), ", "), )) + case len(scope.KnowledgeIDs) > 0: + clauses = append(clauses, fmt.Sprintf( + "(%s.knowledge_base_id = %s AND %s.knowledge_id IN (%s))", + alias, kbID, alias, strings.Join(quoteStringSlice(scope.KnowledgeIDs), ", "), + )) default: clauses = append(clauses, fmt.Sprintf("%s.knowledge_base_id = %s", alias, kbID)) }