WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/modelcontext/README.md

@@ -0,0 +1,92 @@
+# Model context boundary
+
+`modelcontext.Registry` is the only application-facing boundary for values
+that are shortened before an LLM call and restored afterwards. The whole
+codec lives in this one package:
+
+| File | Role |
+| --- | --- |
+| `registry.go` | `Registry` facade — the only type request lifecycles use |
+| `handle_table.go` | `handleTable[M]`, the generic bidirectional primitive behind every handle space |
+| `sources.go` | source handles (`cN`/`dN`/`bN`/`wN`) and the tool-argument codec |
+| `citations.go` | protocol prompts, `<ref/>` → `<kb/>`/`<web/>` expansion, citation stream expander |
+| `model_output.go` | compact source-centric renderings of tool results |
+| `resources.go` | durable-resource handles (`res://NNNN`) |
+| `stream.go` | `streamHold` suffix-hold primitive and every streaming decoder |
+| `tool_policy.go` | the single policy layer: per-tool key contracts + `sourceKeySpaces` dispatch |
+| `handles.go` | exported `HandleTable` for invocation-local spaces (`iN`, `ref-N`, `c000`) |
+
+## 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 handles 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/<knowledge-id>` 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. The
+`sourceKeySpaces` table is the single source of truth for which keys are
+ID-bearing and which handle space they register into; per-tool `sourceIDKeys`
+contracts gate which of those keys each tool may use.
+`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.

internal/modelcontext/citations.go

@@ -0,0 +1,287 @@
+// citations.go owns the public-citation surface of the model context: the
+// system protocol prompt, expansion of private <ref/> handles into canonical
+// <kb/> / <web/> tags, re-compaction of canonical tags replayed from history,
+// and the stream expander that keeps partial tags off the wire.
+package modelcontext
+
+import (
+	"fmt"
+	"html"
+	"regexp"
+	"strings"
+)
+
+const sourceHandleProtocolPrompt = `
+
+## Source handling protocol (system-owned)
+Retrieved content uses request-local source handles: cN identifies a knowledge chunk, wN a web page, dN a document, and bN a knowledge base.
+- Use dN and bN only as tool arguments when a tool requests a document or knowledge base.
+- Never reveal raw chunk IDs, knowledge IDs, knowledge-base IDs, or private source handles in user-visible output. This does not change separate instructions to preserve retrieved Markdown image URLs.`
+
+const citationEnabledProtocolPrompt = `
+- Source citations are enabled for this answer. Cite a knowledge chunk with exactly <ref id="cN"/> and a web page with exactly <ref id="wN"/>.
+- Copy only cN/wN handles that appeared in supplied context or tool results. Never cite dN/bN.
+- Never output <kb> or <web> tags yourself; the system expands valid <ref/> tags after generation.
+- Keep each <ref/> inline on the same line as the claim it supports. Do not group citations at the end.
+- These rules supersede earlier, saved, or custom prompt instructions about citation syntax.`
+
+const citationDisabledProtocolPrompt = `
+- Source citations are disabled for this answer. Do not output <ref>, <kb>, <web>, raw source URLs, or source-handle citations.
+- These rules supersede earlier, saved, or custom prompt instructions that require source citations.`
+
+// ProtocolPrompt returns the internal, non-user-editable source protocol for a
+// model call. Citation formatting stays out of custom and template prompts.
+func sourceProtocolPrompt(citationsEnabled bool) string {
+	if citationsEnabled {
+		return sourceHandleProtocolPrompt + citationEnabledProtocolPrompt
+	}
+	return sourceHandleProtocolPrompt + citationDisabledProtocolPrompt
+}
+
+// ProtocolPrompt returns the source protocol configured for this registry.
+// Request lifecycle code should normally call this through Registry.
+func (r *sourceRegistry) ProtocolPrompt() string {
+	if r == nil {
+		return ""
+	}
+	return sourceProtocolPrompt(r.citationsEnabled)
+}
+
+var (
+	publicKBTagRE        = regexp.MustCompile(`(?is)<kb\b[^>]*>`)
+	publicWebTagRE       = regexp.MustCompile(`(?is)<web\b[^>]*>`)
+	docAttrRE            = regexp.MustCompile(`(?i)\bdoc\s*=\s*"([^"]*)"`)
+	chunkAttrRE          = regexp.MustCompile(`(?i)\bchunk_id\s*=\s*"([^"]+)"`)
+	publicKBAttrRE       = regexp.MustCompile(`(?i)\bkb_id\s*=\s*"([^"]*)"`)
+	urlAttrRE            = regexp.MustCompile(`(?i)\burl\s*=\s*"([^"]+)"`)
+	titleAttrRE          = regexp.MustCompile(`(?i)\btitle\s*=\s*"([^"]*)"`)
+	legacyChunkRE        = regexp.MustCompile(`(?is)<(?:chunk|faq)\b[^>]*>`)
+	faqAttrRE            = regexp.MustCompile(`(?i)\bfaq_id\s*=\s*"([^"]+)"`)
+	knowledgeTitleAttrRE = regexp.MustCompile(`(?i)\bknowledge_title\s*=\s*"([^"]*)"`)
+)
+
+func (r *sourceRegistry) registerLegacyToolReferences(text string) {
+	if r == nil || text == "" {
+		return
+	}
+	r.registerLabeledReferences(text)
+	for _, tag := range legacyChunkRE.FindAllString(text, -1) {
+		chunkID := firstNonEmpty(publicAttr(chunkAttrRE, tag), publicAttr(faqAttrRE, tag))
+		if chunkID == "" {
+			continue
+		}
+		r.RegisterChunk(ChunkReference{
+			ChunkID:         chunkID,
+			KnowledgeID:     publicAttr(documentAttrRE, tag),
+			KnowledgeBaseID: firstNonEmpty(publicAttr(kbAttrRE, tag), publicAttr(publicKBAttrRE, tag)),
+			DocumentTitle:   firstNonEmpty(publicAttr(knowledgeTitleAttrRE, tag), publicAttr(docAttrRE, tag)),
+		})
+	}
+}
+
+// CompactPublicCitations folds canonical citations from prior assistant turns
+// back into this request's private protocol. This prevents durable chunk IDs
+// and web URLs in conversation history from becoming model-visible again.
+func (r *sourceRegistry) CompactPublicCitations(text string) string {
+	if r == nil || text == "" {
+		return text
+	}
+	text = publicKBTagRE.ReplaceAllStringFunc(text, func(tag string) string {
+		chunkID := publicAttr(chunkAttrRE, tag)
+		if chunkID == "" {
+			return tag
+		}
+		handle := r.RegisterChunk(ChunkReference{
+			ChunkID:         chunkID,
+			KnowledgeBaseID: publicAttr(publicKBAttrRE, tag),
+			DocumentTitle:   publicAttr(docAttrRE, tag),
+		})
+		return `<ref id="` + handle + `"/>`
+	})
+	return publicWebTagRE.ReplaceAllStringFunc(text, func(tag string) string {
+		rawURL := publicAttr(urlAttrRE, tag)
+		if rawURL == "" {
+			return tag
+		}
+		handle := r.RegisterWeb(rawURL, publicAttr(titleAttrRE, tag))
+		return `<ref id="` + handle + `"/>`
+	})
+}
+
+func publicAttr(expression *regexp.Regexp, tag string) string {
+	match := expression.FindStringSubmatch(tag)
+	if len(match) != 2 {
+		return ""
+	}
+	return html.UnescapeString(match[1])
+}
+
+var (
+	refTagRE       = regexp.MustCompile(`(?i)<ref\s+id\s*=\s*"([^"]+)"\s*/?>`)
+	refCandidateRE = regexp.MustCompile(`(?is)<ref(?:\s|$)[^>]*(?:>|$)`)
+	modelKBTagRE   = regexp.MustCompile(`(?is)<kb(?:\s|$)[^>]*(?:>|$)`)
+	modelWebTagRE  = regexp.MustCompile(`(?is)<web(?:\s|$)[^>]*(?:>|$)`)
+)
+
+var (
+	documentAttrRE    = regexp.MustCompile(`(?i)\bknowledge_id\s*=\s*"([^"]+)"`)
+	documentElementRE = regexp.MustCompile(`(?is)<knowledge_id>\s*([^<]+?)\s*</knowledge_id>`)
+	kbAttrRE          = regexp.MustCompile(`(?i)\b(?:knowledge_base_id|kb_id)\s*=\s*"([^"]+)"`)
+	kbElementRE       = regexp.MustCompile(`(?is)<(?:knowledge_base_id|kb_id)>\s*([^<]+?)\s*</(?:knowledge_base_id|kb_id)>`)
+)
+
+// registerLabeledReferences covers metadata-oriented tools that do not have a
+// dedicated compact renderer. Only explicit ID labels are recognized; UUID-like
+// text in retrieved content is never guessed to be a source identifier.
+func (r *sourceRegistry) registerLabeledReferences(text string) {
+	if r == nil || text == "" {
+		return
+	}
+	for _, expression := range []*regexp.Regexp{documentAttrRE, documentElementRE} {
+		for _, match := range expression.FindAllStringSubmatch(text, -1) {
+			if len(match) == 2 {
+				r.RegisterDocument(strings.TrimSpace(match[1]))
+			}
+		}
+	}
+	for _, expression := range []*regexp.Regexp{kbAttrRE, kbElementRE} {
+		for _, match := range expression.FindAllStringSubmatch(text, -1) {
+			if len(match) == 2 {
+				r.RegisterKnowledgeBase(strings.TrimSpace(match[1]))
+			}
+		}
+	}
+}
+
+// ExpandText converts the private model protocol into the existing public
+// <kb/> / <web/> contract. Unknown handles fail closed and disappear.
+func (r *sourceRegistry) ExpandText(text string) string {
+	if r == nil || text == "" {
+		return text
+	}
+	// Public citation tags are output-only. Drop any instance written directly
+	// by the model, then create canonical tags solely from registered handles.
+	text = modelKBTagRE.ReplaceAllString(text, "")
+	text = modelWebTagRE.ReplaceAllString(text, "")
+	if !r.citationsEnabled {
+		return refCandidateRE.ReplaceAllString(text, "")
+	}
+	return refCandidateRE.ReplaceAllStringFunc(text, func(tag string) string {
+		match := refTagRE.FindStringSubmatch(tag)
+		if len(match) != 2 {
+			return ""
+		}
+		handle := strings.ToLower(match[1])
+		if chunkID, chunkRef, ok := r.chunks.resolve(handle); ok {
+			attrs := fmt.Sprintf(`doc="%s" chunk_id="%s"`, escapeAttr(chunkRef.DocumentTitle), escapeAttr(chunkID))
+			if chunkRef.KnowledgeBaseID != "" {
+				attrs += fmt.Sprintf(` kb_id="%s"`, escapeAttr(chunkRef.KnowledgeBaseID))
+			}
+			return "<kb " + attrs + " />"
+		}
+		if rawURL, web, ok := r.webs.resolve(handle); ok {
+			return fmt.Sprintf(`<web url="%s" title="%s" />`, escapeAttr(rawURL), escapeAttr(web.title))
+		}
+		return ""
+	})
+}
+
+func escapeAttr(value string) string { return html.EscapeString(value) }
+
+// citationStreamExpander prevents partial private <ref/> tags from reaching SSE while
+// preserving normal streaming for all other content.
+type citationStreamExpander struct {
+	registry *sourceRegistry
+	pending  string
+}
+
+func newCitationStreamExpander(registry *sourceRegistry) *citationStreamExpander {
+	return &citationStreamExpander{registry: registry}
+}
+
+func (d *citationStreamExpander) Feed(chunk string) string {
+	if d == nil || d.registry == nil {
+		return chunk
+	}
+	data := d.pending + chunk
+	d.pending = ""
+	var out strings.Builder
+	for data != "" {
+		idx := strings.Index(data, "<")
+		if idx < 0 {
+			out.WriteString(data)
+			break
+		}
+		out.WriteString(data[:idx])
+		data = data[idx:]
+		lower := strings.ToLower(data)
+		if isSourceTagPending(lower) && !strings.Contains(data, ">") {
+			d.pending = data
+			break
+		}
+		if isRefTagStart(lower) {
+			end := strings.IndexByte(data, '>')
+			if end < 0 {
+				d.pending = data
+				break
+			}
+			tag := data[:end+1]
+			if refTagRE.MatchString(tag) {
+				out.WriteString(d.registry.ExpandText(tag))
+			}
+			data = data[end+1:]
+			continue
+		}
+		if isNamedTagStart(lower, "kb") || isNamedTagStart(lower, "web") {
+			end := strings.IndexByte(data, '>')
+			if end < 0 {
+				d.pending = data
+				break
+			}
+			data = data[end+1:]
+			continue
+		}
+		out.WriteByte('<')
+		data = data[1:]
+	}
+	return out.String()
+}
+
+func isRefTagStart(value string) bool {
+	return isNamedTagStart(value, "ref")
+}
+
+func isNamedTagStart(value, name string) bool {
+	prefix := "<" + name
+	if !strings.HasPrefix(value, prefix) {
+		return false
+	}
+	if len(value) == len(prefix) {
+		return true
+	}
+	next := value[len(prefix)]
+	return next == ' ' || next == '\t' || next == '\r' || next == '\n' || next == '>'
+}
+
+func isSourceTagPending(value string) bool {
+	for _, name := range []string{"ref", "kb", "web"} {
+		prefix := "<" + name
+		if (len(value) <= len(prefix) && strings.HasPrefix(prefix, value)) || isNamedTagStart(value, name) {
+			return true
+		}
+	}
+	return false
+}
+
+func (d *citationStreamExpander) Flush() string {
+	if d == nil {
+		return ""
+	}
+	pending := d.pending
+	d.pending = ""
+	lower := strings.ToLower(pending)
+	if isSourceTagPending(lower) {
+		return ""
+	}
+	return d.registry.ExpandText(pending)
+}

internal/modelcontext/handle_table.go

@@ -0,0 +1,151 @@
+package modelcontext
+
+import (
+	"fmt"
+	"regexp"
+	"sync"
+)
+
+// handleTable is the single bidirectional primitive behind every model-handle
+// space: cN/dN/bN/wN source handles, iN issue handles, res://NNNN resource
+// handles, and the ref-N / c000 ingest handles exposed through HandleTable.
+//
+// It maps a durable dedup key to a sequentially allocated handle and stores
+// the durable value plus optional metadata under one lock, so a resolvable
+// handle can never be observed without its metadata. The key and the value
+// may differ (web references dedup on a canonicalized URL but decode back to
+// the original raw URL). Entries are never removed, which makes the counter
+// equivalent to the historical len(map)+1 allocation and keeps handle
+// numbering stable for the lifetime of the table.
+type handleTable[M any] struct {
+	mu            sync.RWMutex
+	prefix        string
+	width         int
+	next          int
+	handleByKey   map[string]string
+	entryByHandle map[string]*handleEntry[M]
+}
+
+type handleEntry[M any] struct {
+	value       string // durable value a handle decodes back to
+	meta        M
+	wordBounded *regexp.Regexp
+}
+
+// handlePair is a snapshot row for text codecs (compaction/decoding).
+type handlePair struct {
+	value  string
+	handle string
+	// wordBounded matches the handle only on word boundaries. It is compiled
+	// once at registration because DecodeKnownText runs on every streamed
+	// chunk, where recompiling per handle per chunk dominated the cost.
+	wordBounded *regexp.Regexp
+}
+
+// newHandleTable creates a handle space such as c1 (prefix="c", width=0,
+// start=1), c000 (prefix="c", width=3, start=0) or res://0001
+// (prefix="res://", width=4, start=1).
+func newHandleTable[M any](prefix string, width, start int) *handleTable[M] {
+	return &handleTable[M]{
+		prefix:        prefix,
+		width:         width,
+		next:          start,
+		handleByKey:   make(map[string]string),
+		entryByHandle: make(map[string]*handleEntry[M]),
+	}
+}
+
+// register returns the handle assigned to key, allocating the next one on
+// first use. value is what the handle decodes back to; merge, when non-nil,
+// folds the metadata of a repeated registration into the existing entry.
+func (t *handleTable[M]) register(key, value string, meta M, merge func(dst *M, src M)) string {
+	if t == nil || key == "" {
+		return ""
+	}
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if handle, ok := t.handleByKey[key]; ok {
+		if merge != nil {
+			merge(&t.entryByHandle[handle].meta, meta)
+		}
+		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.handleByKey[key] = handle
+	t.entryByHandle[handle] = &handleEntry[M]{
+		value:       value,
+		meta:        meta,
+		wordBounded: regexp.MustCompile(`\b` + regexp.QuoteMeta(handle) + `\b`),
+	}
+	return handle
+}
+
+// handleForKey returns an existing handle without allocating one.
+func (t *handleTable[M]) handleForKey(key string) (string, bool) {
+	if t == nil {
+		return "", false
+	}
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	handle, ok := t.handleByKey[key]
+	return handle, ok
+}
+
+// resolve returns the durable value and metadata snapshot for a known handle.
+func (t *handleTable[M]) resolve(handle string) (string, M, bool) {
+	var zero M
+	if t == nil {
+		return "", zero, false
+	}
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	entry, ok := t.entryByHandle[handle]
+	if !ok {
+		return "", zero, false
+	}
+	return entry.value, entry.meta, true
+}
+
+// has reports whether handle exists in this table.
+func (t *handleTable[M]) has(handle string) bool {
+	if t == nil {
+		return false
+	}
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	_, ok := t.entryByHandle[handle]
+	return ok
+}
+
+func (t *handleTable[M]) size() int {
+	if t == nil {
+		return 0
+	}
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	return len(t.entryByHandle)
+}
+
+// pairs returns a value/handle snapshot for text codecs. Callers own ordering
+// (e.g. longest-value-first for substring compaction).
+func (t *handleTable[M]) pairs() []handlePair {
+	if t == nil {
+		return nil
+	}
+	t.mu.RLock()
+	defer t.mu.RUnlock()
+	out := make([]handlePair, 0, len(t.entryByHandle))
+	for handle, entry := range t.entryByHandle {
+		out = append(out, handlePair{
+			value:       entry.value,
+			handle:      handle,
+			wordBounded: entry.wordBounded,
+		})
+	}
+	return out
+}

internal/modelcontext/handles.go

@@ -0,0 +1,96 @@
+package modelcontext
+
+import (
+	"sort"
+	"strings"
+)
+
+// HandleTable is a typed, invocation-local bidirectional mapping used when a
+// prompt needs compact handles for durable values (wiki issue iN handles,
+// ingest ref-N slug handles, c000 citation-batch handles). Handles must never
+// be persisted; Resolve converts model output back to the durable value first.
+// It is the exported word-boundary-aware wrapper over handleTable.
+type HandleTable struct {
+	table *handleTable[struct{}]
+}
+
+// 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{table: newHandleTable[struct{}](prefix, width, start)}
+}
+
+// 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 ""
+	}
+	return t.table.register(value, value, struct{}{}, nil)
+}
+
+// Handle returns an already registered handle without creating one.
+func (t *HandleTable) Handle(value string) (string, bool) {
+	if t == nil {
+		return "", false
+	}
+	return t.table.handleForKey(value)
+}
+
+// Resolve converts a known model handle back to its durable value.
+func (t *HandleTable) Resolve(handle string) (string, bool) {
+	if t == nil {
+		return "", false
+	}
+	value, _, ok := t.table.resolve(strings.TrimSpace(handle))
+	return value, ok
+}
+
+func (t *HandleTable) Empty() bool {
+	return t == nil || t.table.size() == 0
+}
+
+func (t *HandleTable) Len() int {
+	if t == nil {
+		return 0
+	}
+	return t.table.size()
+}
+
+// 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
+	}
+	pairs := t.table.pairs()
+	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. Unlike the
+// resource codec, replacement is word-boundary-aware so short handles such as
+// i1 cannot fire inside ordinary prose tokens.
+func (t *HandleTable) DecodeKnownText(value string) string {
+	if t == nil || value == "" {
+		return value
+	}
+	pairs := t.table.pairs()
+	sort.SliceStable(pairs, func(i, j int) bool { return len(pairs[i].handle) > len(pairs[j].handle) })
+	for _, item := range pairs {
+		if item.wordBounded == nil {
+			continue
+		}
+		// A durable value may legally contain '$'. ReplaceAllString would
+		// interpret it as a regexp expansion, so use a literal callback.
+		durable := item.value
+		value = item.wordBounded.ReplaceAllStringFunc(value, func(string) string { return durable })
+	}
+	return value
+}

internal/modelcontext/model_output.go

@@ -0,0 +1,576 @@
+package modelcontext
+
+import (
+	"encoding/json"
+	"fmt"
+	"sort"
+	"strings"
+
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+const (
+	modelWebSearchEvidenceMaxRunes = 500
+	modelWebFetchSummaryMaxRunes   = 4000
+	modelWebFetchContentMaxRunes   = 8000
+	modelWebFetchTotalMaxRunes     = 16000
+)
+
+// ModelOutput returns a compact, source-centric representation for the LLM.
+// The canonical ToolResult.Output remains untouched for UI, logs and storage.
+func (r *sourceRegistry) ModelOutput(result *types.ToolResult) string {
+	if result == nil {
+		return ""
+	}
+	displayType := stringValue(result.Data, "display_type")
+	if displayType == "web_fetch_results" {
+		return r.modelWebFetchOutput(mapsValue(result.Data["results"]), result.Output)
+	}
+	if !result.Success {
+		return failedToolModelText(result.Output, result.Error)
+	}
+	switch displayType {
+	case "grep_results":
+		return r.modelKnowledgeOutput("keyword", mapsValue(result.Data["chunk_results"]), result.Output)
+	case "search_results":
+		return r.modelKnowledgeOutput("semantic", mapsValue(result.Data["results"]), result.Output)
+	case "knowledge_chunks_list":
+		return r.modelKnowledgeChunksOutput(result.Data, result.Output)
+	case "document_info":
+		return r.modelDocumentInfoOutput(mapsValue(result.Data["documents"]), result.Output)
+	case "graph_query_results":
+		return r.modelKnowledgeOutput("graph", mapsValue(result.Data["results"]), result.Output)
+	case "web_search_results":
+		return r.modelWebSearchOutput(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 *sourceRegistry) 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 *sourceRegistry) 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 *sourceRegistry) modelDocumentInfoOutput(rows []map[string]interface{}, fallback string) string {
+	if len(rows) == 0 {
+		return r.CompactKnownText(fallback)
+	}
+	var b strings.Builder
+	b.WriteString("<documents>\n")
+	count := 0
+	for _, row := range rows {
+		knowledgeID := stringValue(row, "knowledge_id")
+		docHandle := r.RegisterDocument(knowledgeID)
+		if boolValue(row, "is_faq") {
+			chunkID := stringValue(row, "faq_id")
+			if chunkID == "" {
+				continue
+			}
+			title := firstNonEmpty(stringValue(row, "faq_question"), stringValue(row, "title"))
+			chunkHandle := r.RegisterChunk(ChunkReference{
+				ChunkID:       chunkID,
+				KnowledgeID:   knowledgeID,
+				DocumentTitle: title,
+				ChunkType:     "faq",
+			})
+			fmt.Fprintf(&b, "  <document id=\"%s\" type=\"faq\">\n", escapeAttr(docHandle))
+			fmt.Fprintf(&b, "    <chunk id=\"%s\" type=\"faq\">\n", escapeAttr(chunkHandle))
+			if title != "" {
+				fmt.Fprintf(&b, "      <question>%s</question>\n", escapeText(title))
+			}
+			for _, answer := range stringSliceValue(row["faq_answers"]) {
+				fmt.Fprintf(&b, "      <answer>%s</answer>\n", escapeText(answer))
+			}
+			b.WriteString("    </chunk>\n  </document>\n")
+			count++
+			continue
+		}
+
+		if docHandle == "" {
+			continue
+		}
+		fmt.Fprintf(&b, "  <document id=\"%s\"", escapeAttr(docHandle))
+		if title := stringValue(row, "title"); title != "" {
+			fmt.Fprintf(&b, " title=\"%s\"", escapeAttr(title))
+		}
+		if docType := stringValue(row, "type"); docType != "" {
+			fmt.Fprintf(&b, " type=\"%s\"", escapeAttr(docType))
+		}
+		if fileType := stringValue(row, "file_type"); fileType != "" {
+			fmt.Fprintf(&b, " file_type=\"%s\"", escapeAttr(fileType))
+		}
+		fmt.Fprintf(&b, " chunk_count=\"%d\">\n", intValue(row, "chunk_count"))
+		if description := stringValue(row, "description"); description != "" {
+			fmt.Fprintf(&b, "    <description>%s</description>\n", escapeText(description))
+		}
+		b.WriteString("  </document>\n")
+		count++
+	}
+	b.WriteString("</documents>")
+	if count == 0 {
+		return r.CompactKnownText(fallback)
+	}
+	return b.String()
+}
+
+type modelChunk struct {
+	handle     string
+	docHandle  string
+	kbHandle   string
+	title      string
+	metadata   string
+	chunkType  string
+	index      int
+	view       string
+	match      string
+	content    string
+	question   string
+	answers    []string
+	images     []map[string]interface{}
+	docRealID  string
+	kbRealID   string
+	chunkReal  string
+	inputOrder int
+}
+
+func (r *sourceRegistry) modelKnowledgeOutput(mode string, rows []map[string]interface{}, fallback string) string {
+	chunks := make([]modelChunk, 0, len(rows))
+	for idx, row := range rows {
+		chunkID := firstNonEmpty(stringValue(row, "chunk_id"), stringValue(row, "faq_id"), stringValue(row, "id"))
+		knowledgeID := stringValue(row, "knowledge_id")
+		kbID := firstNonEmpty(stringValue(row, "knowledge_base_id"), stringValue(row, "knowledge_base"))
+		title := firstNonEmpty(stringValue(row, "knowledge_title"), stringValue(row, "title"))
+		if chunkID == "" {
+			continue
+		}
+		chunkType := stringValue(row, "chunk_type")
+		if stringValue(row, "faq_id") != "" && chunkType == "" {
+			chunkType = "faq"
+		}
+		chunkIndex := intValue(row, "chunk_index")
+		if chunkIndex == 0 {
+			chunkIndex = intValue(row, "index")
+		}
+		chunkHandle := r.RegisterChunk(ChunkReference{
+			ChunkID:         chunkID,
+			KnowledgeID:     knowledgeID,
+			KnowledgeBaseID: kbID,
+			DocumentTitle:   title,
+			ChunkIndex:      chunkIndex,
+			ChunkType:       chunkType,
+		})
+		chunks = append(chunks, modelChunk{
+			handle:     chunkHandle,
+			docHandle:  r.RegisterDocument(knowledgeID),
+			kbHandle:   r.RegisterKnowledgeBase(kbID),
+			title:      title,
+			metadata:   stringValue(row, "knowledge_metadata"),
+			chunkType:  chunkType,
+			index:      chunkIndex,
+			view:       viewForRow(row, mode),
+			match:      firstNonEmpty(stringValue(row, "match_snippet"), stringValue(row, "matched_content")),
+			content:    stringValue(row, "content"),
+			question:   firstNonEmpty(stringValue(row, "faq_question"), stringValue(row, "faq_standard_question")),
+			answers:    stringSliceValue(row["faq_answers"]),
+			images:     mapsValue(row["images"]),
+			docRealID:  knowledgeID,
+			kbRealID:   kbID,
+			chunkReal:  chunkID,
+			inputOrder: idx,
+		})
+	}
+	if len(chunks) == 0 {
+		return r.CompactKnownText(fallback)
+	}
+	return renderKnowledgeChunks(mode, chunks)
+}
+
+func viewForRow(row map[string]interface{}, mode string) string {
+	if stringValue(row, "content") != "" {
+		return "full"
+	}
+	if mode == "deep_read" {
+		return "full"
+	}
+	return "match"
+}
+
+func (r *sourceRegistry) modelKnowledgeChunksOutput(data map[string]interface{}, fallback string) string {
+	rows := mapsValue(data["chunks"])
+	title := stringValue(data, "knowledge_title")
+	knowledgeID := stringValue(data, "knowledge_id")
+	for _, row := range rows {
+		if stringValue(row, "knowledge_id") == "" {
+			row["knowledge_id"] = knowledgeID
+		}
+		if stringValue(row, "knowledge_title") == "" {
+			row["knowledge_title"] = title
+		}
+	}
+	output := r.modelKnowledgeOutput("deep_read", rows, fallback)
+	if len(rows) == 0 {
+		return output
+	}
+	remaining := intValue(data, "total_chunks") - intValue(data, "fetched_chunks")
+	if remaining > 0 {
+		output = strings.TrimSuffix(output, "</retrieval>")
+		output += fmt.Sprintf("  <pagination remaining=\"%d\" page=\"%d\" page_size=\"%d\" />\n</retrieval>",
+			remaining, intValue(data, "page"), intValue(data, "page_size"))
+	}
+	return output
+}
+
+func renderKnowledgeChunks(mode string, chunks []modelChunk) string {
+	type docGroup struct {
+		handle   string
+		kbHandle string
+		title    string
+		metadata string
+		chunks   []modelChunk
+		order    int
+	}
+	groupsByKey := make(map[string]*docGroup)
+	var groups []*docGroup
+	for _, chunk := range chunks {
+		key := chunk.docHandle
+		if key == "" {
+			key = "chunk:" + chunk.handle
+		}
+		group := groupsByKey[key]
+		if group == nil {
+			group = &docGroup{
+				handle: chunk.docHandle, kbHandle: chunk.kbHandle, title: chunk.title,
+				metadata: chunk.metadata, order: chunk.inputOrder,
+			}
+			groupsByKey[key] = group
+			groups = append(groups, group)
+		} else if group.metadata == "" {
+			group.metadata = chunk.metadata
+		}
+		group.chunks = append(group.chunks, chunk)
+	}
+	sort.SliceStable(groups, func(i, j int) bool { return groups[i].order < groups[j].order })
+
+	var b strings.Builder
+	fmt.Fprintf(&b, "<retrieval type=\"knowledge\" mode=\"%s\">\n", escapeAttr(mode))
+	for _, group := range groups {
+		b.WriteString("  <document")
+		if group.handle != "" {
+			fmt.Fprintf(&b, " id=\"%s\"", escapeAttr(group.handle))
+		}
+		if group.kbHandle != "" {
+			fmt.Fprintf(&b, " kb=\"%s\"", escapeAttr(group.kbHandle))
+		}
+		if group.title != "" {
+			fmt.Fprintf(&b, " title=\"%s\"", escapeAttr(group.title))
+		}
+		b.WriteString(">\n")
+		if group.metadata != "" {
+			fmt.Fprintf(&b, "    <metadata>%s</metadata>\n", escapeText(group.metadata))
+		}
+		for _, chunk := range group.chunks {
+			fmt.Fprintf(&b, "    <chunk id=\"%s\" index=\"%d\" view=\"%s\"", chunk.handle, chunk.index, chunk.view)
+			if chunk.chunkType != "" {
+				fmt.Fprintf(&b, " type=\"%s\"", escapeAttr(chunk.chunkType))
+			}
+			b.WriteString(">\n")
+			if chunk.question != "" {
+				fmt.Fprintf(&b, "      <question>%s</question>\n", escapeText(chunk.question))
+			}
+			if chunk.match != "" {
+				fmt.Fprintf(&b, "      <match>%s</match>\n", escapeText(chunk.match))
+			}
+			if chunk.content != "" {
+				fmt.Fprintf(&b, "      <content>%s</content>\n", escapeText(chunk.content))
+			}
+			for _, answer := range chunk.answers {
+				fmt.Fprintf(&b, "      <answer>%s</answer>\n", escapeText(answer))
+			}
+			for _, image := range chunk.images {
+				imageURL := stringValue(image, "url")
+				if imageURL == "" {
+					continue
+				}
+				caption := stringValue(image, "caption")
+				fmt.Fprintf(&b, "      ![%s](%s)\n", caption, imageURL)
+			}
+			b.WriteString("    </chunk>\n")
+		}
+		b.WriteString("  </document>\n")
+	}
+	b.WriteString("</retrieval>")
+	return b.String()
+}
+
+func (r *sourceRegistry) modelWebSearchOutput(rows []map[string]interface{}, fallback string) string {
+	if len(rows) == 0 {
+		return r.CompactKnownText(fallback)
+	}
+	var b strings.Builder
+	b.WriteString("<retrieval type=\"web\" mode=\"search\">\n")
+	count := 0
+	for _, row := range rows {
+		rawURL := stringValue(row, "url")
+		if rawURL == "" {
+			continue
+		}
+		handle := r.RegisterWeb(rawURL, stringValue(row, "title"))
+		fmt.Fprintf(&b, "  <page id=\"%s\" title=\"%s\">\n", handle, escapeAttr(stringValue(row, "title")))
+		b.WriteString("    <evidence type=\"search_summary\" verified=\"false\" />\n")
+		if snippet := stringValue(row, "snippet"); snippet != "" {
+			writeLimitedWebEvidence(&b, "match", snippet, modelWebSearchEvidenceMaxRunes, nil)
+		}
+		if content := stringValue(row, "content"); content != "" && content != stringValue(row, "snippet") {
+			writeLimitedWebEvidence(&b, "content", content, modelWebSearchEvidenceMaxRunes, nil)
+		}
+		if published := stringValue(row, "published_at"); published != "" {
+			fmt.Fprintf(&b, "    <published>%s</published>\n", escapeText(published))
+		}
+		b.WriteString("  </page>\n")
+		count++
+	}
+	b.WriteString("</retrieval>")
+	if count == 0 {
+		return r.CompactKnownText(fallback)
+	}
+	return b.String()
+}
+
+func (r *sourceRegistry) modelWebFetchOutput(rows []map[string]interface{}, fallback string) string {
+	if len(rows) == 0 {
+		return r.CompactKnownText(fallback)
+	}
+	var b strings.Builder
+	b.WriteString("<retrieval type=\"web\" mode=\"fetch\">\n")
+	count, successCount, failedCount := 0, 0, 0
+	remainingEvidence := modelWebFetchTotalMaxRunes
+	for _, row := range rows {
+		rawURL := stringValue(row, "url")
+		if rawURL == "" {
+			continue
+		}
+		title := stringValue(row, "title")
+		handle := r.RegisterWeb(rawURL, title)
+		status := stringValue(row, "status")
+		if status == "" {
+			status = "success"
+		}
+		fmt.Fprintf(&b, "  <page id=\"%s\" status=\"%s\"", handle, escapeAttr(status))
+		if title != "" {
+			fmt.Fprintf(&b, " title=\"%s\"", escapeAttr(title))
+		}
+		if status == "success" {
+			b.WriteString(" view=\"full\">\n")
+			successCount++
+			if summary := stringValue(row, "summary"); summary != "" {
+				writeLimitedWebEvidence(&b, "summary", summary, modelWebFetchSummaryMaxRunes, &remainingEvidence)
+			}
+			if summaryStatus := stringValue(row, "summary_status"); summaryStatus == "failed" {
+				fmt.Fprintf(&b, "    <summary_error code=\"%s\">%s</summary_error>\n",
+					escapeAttr(stringValue(row, "summary_error_code")), escapeText(stringValue(row, "summary_error_message")))
+			}
+			if content := stringValue(row, "raw_content"); content != "" {
+				writeLimitedWebEvidence(&b, "content", content, modelWebFetchContentMaxRunes, &remainingEvidence)
+			}
+		} else {
+			fmt.Fprintf(&b, " retryable=\"%t\"", boolValue(row, "retryable"))
+			if errorCode := stringValue(row, "error_code"); errorCode != "" {
+				fmt.Fprintf(&b, " error_code=\"%s\"", escapeAttr(errorCode))
+			}
+			b.WriteString(">\n")
+			if errorMessage := stringValue(row, "error_message"); errorMessage != "" {
+				fmt.Fprintf(&b, "    <error>%s</error>\n", escapeText(errorMessage))
+			}
+			if status == "failed" {
+				failedCount++
+			}
+		}
+		b.WriteString("  </page>\n")
+		count++
+	}
+	b.WriteString("</retrieval>")
+	if count == 0 {
+		return r.CompactKnownText(fallback)
+	}
+	if failedCount > 0 {
+		b.WriteString("\n\n=== Next Steps ===\n")
+		if successCount == 0 {
+			b.WriteString("- All page fetches failed. Stop expanding web searches and answer from existing web_search titles, URLs, and snippets.\n")
+			b.WriteString("- Explicitly state that page content was not verified and treat dynamic facts as uncertain.")
+		} else {
+			b.WriteString("- Use successful page content together with existing search snippets; failed URLs do not invalidate successful evidence.\n")
+			b.WriteString("- Do not retry non-retryable failures. If evidence is sufficient, answer now.")
+		}
+	}
+	return b.String()
+}
+
+func writeLimitedWebEvidence(builder *strings.Builder, tag, value string, maxRunes int, remaining *int) {
+	if value == "" {
+		return
+	}
+	limit := maxRunes
+	if remaining != nil && *remaining < limit {
+		limit = *remaining
+	}
+	limited, truncated := truncateModelEvidence(value, limit)
+	if remaining != nil {
+		*remaining -= len([]rune(limited))
+		if *remaining < 0 {
+			*remaining = 0
+		}
+	}
+	fmt.Fprintf(builder, "    <%s", tag)
+	if truncated {
+		builder.WriteString(" truncated=\"true\"")
+	}
+	fmt.Fprintf(builder, ">%s</%s>\n", escapeText(limited), tag)
+}
+
+func truncateModelEvidence(value string, maxRunes int) (string, bool) {
+	runes := []rune(value)
+	if len(runes) <= maxRunes {
+		return value, false
+	}
+	if maxRunes <= 0 {
+		return "", true
+	}
+	return string(runes[:maxRunes]), true
+}
+
+// failedToolModelText is what the model should see for a failed tool call.
+// Script and shell failures put the useful diagnostics in Output (stdout /
+// stderr); Error is often just "exited with code 1" plus a retry hint.
+// Putting Output first makes "[Analyze the error above ...]" refer to the
+// actual streams.
+func failedToolModelText(output, errMsg string) string {
+	output = strings.TrimSpace(output)
+	errMsg = strings.TrimSpace(errMsg)
+	switch {
+	case output == "" && errMsg == "":
+		return "Error: tool call failed"
+	case output == "":
+		return "Error: " + errMsg
+	case errMsg == "" || strings.Contains(output, errMsg):
+		return output
+	default:
+		return output + "\n\nError: " + errMsg
+	}
+}
+
+func mapsValue(value interface{}) []map[string]interface{} {
+	if value == nil {
+		return nil
+	}
+	encoded, err := json.Marshal(value)
+	if err != nil {
+		return nil
+	}
+	var rows []map[string]interface{}
+	if err := json.Unmarshal(encoded, &rows); err != nil {
+		return nil
+	}
+	return rows
+}
+
+func stringValue(values map[string]interface{}, key string) string {
+	if values == nil {
+		return ""
+	}
+	switch value := values[key].(type) {
+	case string:
+		return value
+	case fmt.Stringer:
+		return value.String()
+	case json.Number:
+		return value.String()
+	default:
+		return ""
+	}
+}
+
+func intValue(values map[string]interface{}, key string) int {
+	if values == nil {
+		return 0
+	}
+	switch value := values[key].(type) {
+	case int:
+		return value
+	case int32:
+		return int(value)
+	case int64:
+		return int(value)
+	case float64:
+		return int(value)
+	case json.Number:
+		result, _ := value.Int64()
+		return int(result)
+	default:
+		return 0
+	}
+}
+
+func boolValue(values map[string]interface{}, key string) bool {
+	if values == nil {
+		return false
+	}
+	value, _ := values[key].(bool)
+	return value
+}
+
+func stringSliceValue(value interface{}) []string {
+	if value == nil {
+		return nil
+	}
+	encoded, err := json.Marshal(value)
+	if err != nil {
+		return nil
+	}
+	var values []string
+	if err := json.Unmarshal(encoded, &values); err != nil {
+		return nil
+	}
+	return values
+}
+
+func escapeText(value string) string {
+	replacer := strings.NewReplacer("&", "&amp;", "<", "&lt;", ">", "&gt;")
+	return replacer.Replace(value)
+}

internal/modelcontext/registry.go

@@ -0,0 +1,291 @@
+// 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/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.`
+
+// 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/<knowledge-id> 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   *sourceRegistry
+	resources *resourceRegistry
+	issues    *HandleTable
+}
+
+// NewRegistry creates a registry for one model request/Agent execution.
+func NewRegistry(citationsEnabled bool) *Registry {
+	return &Registry{
+		sources:   newSourceRegistry(citationsEnabled),
+		resources: newResourceRegistry(),
+		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.OrphanHandles(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: newResourceStreamDecoder(r.resources),
+		sources:   newCitationStreamExpander(r.sources),
+		issues:    NewHandleStreamDecoder(r.issues),
+		orphans:   newOrphanResourceStreamFilter(),
+	}
+}
+
+// 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.OrphanHandles(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.ChunkHandle(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 failedToolModelText(result.Output, 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
+	// Error text is encoded on the same path as Output: a failed tool call
+	// routinely echoes the offending argument, so a raw durable ID would leak
+	// through the error branch of an otherwise handle-only tool.
+	copyResult.Output = r.resources.EncodeText(r.encodeToolPrivateResult(toolName, result.Output))
+	copyResult.Error = r.resources.EncodeText(r.encodeToolPrivateResult(toolName, result.Error))
+	var modelOutput string
+	if sourceOutputAllowed(toolName) {
+		modelOutput = r.sources.ModelOutput(&copyResult)
+	} else if copyResult.Success {
+		modelOutput = copyResult.Output
+	} else {
+		modelOutput = failedToolModelText(copyResult.Output, copyResult.Error)
+	}
+	// 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.StripOrphanHandles(text)
+	text = r.sources.ExpandText(text)
+	return r.issues.DecodeKnownText(text)
+}

internal/modelcontext/registry_test.go

@@ -0,0 +1,432 @@
+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: `<knowledge_id>mcp-owned-id</knowledge_id> doc-real`,
+		},
+	})
+	require.JSONEq(t, `{"knowledge_id":"doc-real"}`, messages[0].ToolCalls[0].Function.Arguments)
+	require.Equal(t, `<knowledge_id>mcp-owned-id</knowledge_id> doc-real`, messages[1].Content)
+
+	modelOutput := registry.ModelToolResultForTool("dynamic_mcp_tool", &types.ToolResult{
+		Success: true,
+		Output:  `<knowledge_id>mcp-owned-id</knowledge_id> doc-real`,
+	})
+	require.Equal(t, `<knowledge_id>mcp-owned-id</knowledge_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 TestModelToolResultForTool_failedSkillScriptKeepsStdout(t *testing.T) {
+	registry := NewRegistry(true)
+	stdout := `{"chart":{"success":false,"error":{"error":"X轴字段不存在：工作项目","available":["name","value"]}}}`
+	got := registry.ModelToolResultForTool("execute_skill_script", &types.ToolResult{
+		Success: false,
+		Output:  "=== Script Execution: smart-charts/scripts/cli.py ===\n\n## Standard Output\n\n```\n" + stdout + "\n```\n",
+		Error:   "Script exited with code 1\n\n[Analyze the error above and try a different approach.]",
+	})
+	require.Contains(t, got, "X轴字段不存在：工作项目")
+	require.Contains(t, got, "available")
+	require.Contains(t, got, "Error: Script exited with code 1")
+}
+
+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: "<knowledge_base_id>" + kbID + "</knowledge_base_id>\n" +
+		"<link>[[summary/" + knowledgeID + "|Summary]]</link>\n" +
+		"<knowledge_id>" + knowledgeID + "</knowledge_id>",
+	})
+	require.Contains(t, got, "[[res://0001|Summary]]")
+	require.Contains(t, got, "<knowledge_base_id>b1</knowledge_base_id>")
+	require.Contains(t, got, "<knowledge_id>d1</knowledge_id>")
+	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 <ref id=\"c") +
+		decoder.Feed("1\"/>") + 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
+}

internal/modelcontext/resources.go

@@ -0,0 +1,157 @@
+// resources.go is the durable-resource half of the model-context registry:
+// it assigns request-local res://NNNN handles to stored resource references
+// and restores them before application code consumes model output.
+package modelcontext
+
+import (
+	"regexp"
+	"sort"
+	"strings"
+
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+// storedRefRE also recognizes legacy physical references. New writes persist
+// resource:// handles, but old chunks and message history can still contain a
+// provider URL. Giving both forms the same request-local handle makes rollout
+// safe without a blocking full-table content rewrite.
+//
+// The final alternative handles wiki summary-page slugs (summary/<uuid>). They
+// are not storage handles, but they share the exact failure mode this registry
+// exists to prevent: a high-entropy identifier the model must reproduce
+// verbatim (inside [[slug|display]] links and as wiki-tool slug arguments).
+// Models routinely mangle the UUID by inserting or dropping hex digits, which
+// yields dead cross-links. Aliasing the slug to a low-entropy res:// token
+// removes the opportunity to mangle at the source; the token round-trips back
+// to the real slug on stream output and in decoded tool-call arguments before
+// anything is persisted. Entity slugs (entity/<readable-title>) are low-entropy
+// and semantically meaningful, so they are deliberately left untouched.
+var storedRefRE = regexp.MustCompile(
+	`resource://[0-9A-Za-z_-]{22}|` +
+		`(?:storage://[0-9A-Za-z_-]+/)?` +
+		`(?:local|minio|cos|tos|s3|oss|ks3|obs)://[^\s)\]>"']+|` +
+		`summary/[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}`,
+)
+
+// resourceHandleShapeRE matches the handle syntax produced by EncodeText. It
+// is used to spot handle-shaped tokens the registry cannot map back — either a
+// hallucinated reference or a coincidental collision — and by the stream-side
+// orphan filter.
+var resourceHandleShapeRE = regexp.MustCompile(`res://\d+`)
+
+// resourceRegistry assigns low-entropy, request-local handles to stable resource
+// handles. It is safe to reuse across all rounds of one Agent execution.
+type resourceRegistry struct {
+	table *handleTable[struct{}]
+}
+
+// newResourceRegistry creates an empty request-local handle registry.
+func newResourceRegistry() *resourceRegistry {
+	// A URL-shaped handle (scheme://digits) keeps the token low-entropy while
+	// looking enough like a link that the model reuses it verbatim inside
+	// Markdown image/link syntax instead of reasoning about or rewriting it.
+	return &resourceRegistry{table: newHandleTable[struct{}]("res://", 4, 1)}
+}
+
+// EncodeText replaces stored references with compact, stable handles.
+func (r *resourceRegistry) EncodeText(value string) string {
+	if r == nil || value == "" {
+		return value
+	}
+	return storedRefRE.ReplaceAllStringFunc(value, func(ref string) string {
+		return r.table.register(ref, ref, struct{}{}, nil)
+	})
+}
+
+// DecodeText restores every handle currently known to the registry. Handles are
+// replaced longest-first with no word-boundary check — this plain substring
+// behavior is load-bearing for handles adjacent to punctuation in Markdown.
+func (r *resourceRegistry) DecodeText(value string) string {
+	if r == nil || value == "" {
+		return value
+	}
+	pairs := r.table.pairs()
+	sort.SliceStable(pairs, func(i, j int) bool { return len(pairs[i].handle) > len(pairs[j].handle) })
+	for _, item := range pairs {
+		value = strings.ReplaceAll(value, item.handle, item.value)
+	}
+	return value
+}
+
+// StripOrphanHandles removes handle-shaped tokens after all known handles have
+// been restored. Use this only on model output; tool arguments must retain
+// unknown handles long enough for modelcontext to reject the call.
+func (r *resourceRegistry) StripOrphanHandles(value string) string {
+	if value == "" {
+		return value
+	}
+	return resourceHandleShapeRE.ReplaceAllString(value, "")
+}
+
+// EncodeMessages returns a copied message slice with textual references
+// compacted. Binary/image content fields are intentionally left untouched.
+func (r *resourceRegistry) EncodeMessages(messages []chat.Message) []chat.Message {
+	if r == nil || len(messages) == 0 {
+		return messages
+	}
+	encoded := make([]chat.Message, len(messages))
+	copy(encoded, messages)
+	for i := range encoded {
+		encoded[i].Content = r.EncodeText(encoded[i].Content)
+		encoded[i].ReasoningContent = r.EncodeText(encoded[i].ReasoningContent)
+		if len(encoded[i].MultiContent) > 0 {
+			encoded[i].MultiContent = append([]chat.MessageContentPart(nil), encoded[i].MultiContent...)
+			for j := range encoded[i].MultiContent {
+				encoded[i].MultiContent[j].Text = r.EncodeText(encoded[i].MultiContent[j].Text)
+			}
+		}
+		if len(encoded[i].ToolCalls) > 0 {
+			encoded[i].ToolCalls = append([]chat.ToolCall(nil), encoded[i].ToolCalls...)
+			for j := range encoded[i].ToolCalls {
+				encoded[i].ToolCalls[j].Function.Arguments = r.EncodeText(encoded[i].ToolCalls[j].Function.Arguments)
+			}
+		}
+	}
+	return encoded
+}
+
+// DecodeToolCalls restores handles in tool-call JSON arguments.
+func (r *resourceRegistry) DecodeToolCalls(toolCalls []types.LLMToolCall) {
+	for i := range toolCalls {
+		toolCalls[i].Function.Arguments = r.DecodeText(toolCalls[i].Function.Arguments)
+	}
+}
+
+// OrphanHandles returns the distinct handle-shaped tokens in an already-decoded
+// string that the registry cannot resolve. A non-empty result means the model
+// emitted a reference no real resource backs (hallucination) or the user text
+// happened to collide with the handle syntax. Callers should log/observe these
+// rather than surfacing them to end users as broken links.
+func (r *resourceRegistry) OrphanHandles(decoded string) []string {
+	if decoded == "" {
+		return nil
+	}
+	var orphans []string
+	seen := make(map[string]struct{})
+	for _, match := range resourceHandleShapeRE.FindAllString(decoded, -1) {
+		if r != nil && r.table.has(match) {
+			continue
+		}
+		if _, dup := seen[match]; dup {
+			continue
+		}
+		seen[match] = struct{}{}
+		orphans = append(orphans, match)
+	}
+	return orphans
+}
+
+func (r *resourceRegistry) handles() []string {
+	pairs := r.table.pairs()
+	handles := make([]string, 0, len(pairs))
+	for _, item := range pairs {
+		handles = append(handles, item.handle)
+	}
+	return handles
+}

internal/modelcontext/resources_test.go

@@ -0,0 +1,94 @@
+package modelcontext
+
+import (
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRegistryRoundTripAndDeduplicate(t *testing.T) {
+	r := newResourceRegistry()
+	ref := "resource://AbCdEfGhIjKlMnOpQrStUv"
+	encoded := r.EncodeText("![a](" + ref + ") and " + ref)
+	require.Equal(t, "![a](res://0001) and res://0001", encoded)
+	require.Equal(t, "![a]("+ref+") and "+ref, r.DecodeText(encoded))
+}
+
+func TestRegistryAliasesLegacyPhysicalReferencesDuringRollout(t *testing.T) {
+	r := newResourceRegistry()
+	ref := "storage://c0d93536-702c-4977-aa5e-fe670073c3cb/local://10000/exports/image.png"
+	encoded := r.EncodeText("![image](" + ref + ")")
+	require.Equal(t, "![image](res://0001)", encoded)
+	require.Equal(t, "![image]("+ref+")", r.DecodeText(encoded))
+}
+
+func TestRegistryAliasesWikiSummarySlug(t *testing.T) {
+	r := newResourceRegistry()
+	slug := "summary/07a20bb1-a662-47cf-9929-06fb5d5b5b5e"
+	// Inside a [[slug|display]] link the model must copy verbatim.
+	encoded := r.EncodeText("see [[" + slug + "|Foo.md - Summary]]")
+	require.Equal(t, "see [[res://0001|Foo.md - Summary]]", encoded)
+	require.Equal(t, "see [["+slug+"|Foo.md - Summary]]", r.DecodeText(encoded))
+
+	// The same slug as a tool-call argument value resolves to one handle.
+	require.Equal(t, "res://0001", r.EncodeText(slug))
+}
+
+func TestRegistryLeavesEntitySlugUntouched(t *testing.T) {
+	r := newResourceRegistry()
+	// Entity slugs are low-entropy and semantically meaningful; not aliased.
+	entity := "[[entity/weknora-error-log|WeKnora 试错记录]]"
+	require.Equal(t, entity, r.EncodeText(entity))
+}
+
+func TestRegistryEncodesMessageCopies(t *testing.T) {
+	r := newResourceRegistry()
+	ref := "resource://AbCdEfGhIjKlMnOpQrStUv"
+	original := []chat.Message{{Role: "tool", Content: ref}}
+	encoded := r.EncodeMessages(original)
+	require.Equal(t, ref, original[0].Content)
+	require.Equal(t, "res://0001", encoded[0].Content)
+}
+
+func TestStreamDecoderRestoresSplitAlias(t *testing.T) {
+	r := newResourceRegistry()
+	ref := "resource://AbCdEfGhIjKlMnOpQrStUv"
+	require.Equal(t, "res://0001", r.EncodeText(ref))
+	d := newResourceStreamDecoder(r)
+	require.Equal(t, "before ", d.Feed("before res://0"))
+	require.Equal(t, ref+" afte", d.Feed("001 after"))
+	require.Equal(t, "r", d.Flush())
+}
+
+func TestStreamDecoderRestoresAliasSplitBeforeScheme(t *testing.T) {
+	const ref = "resource://AbCdEfGhIjKlMnOpQrStUv"
+	r := newResourceRegistry()
+	require.Equal(t, "res://0001", r.EncodeText(ref))
+	d := newResourceStreamDecoder(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 := newResourceRegistry()
+	ref := "resource://AbCdEfGhIjKlMnOpQrStUv"
+	require.Equal(t, "res://0001", r.EncodeText(ref))
+
+	// Known handle resolves and leaves no orphan once decoded.
+	require.Nil(t, r.OrphanHandles(r.DecodeText("see res://0001")))
+
+	// A reference the registry never assigned is reported (deduplicated).
+	require.Equal(t, []string{"res://0099"}, r.OrphanHandles("look at res://0099 and res://0099"))
+}
+
+func TestStripOrphanAliasesRemovesIncompleteNumericHandle(t *testing.T) {
+	r := newResourceRegistry()
+	require.Equal(t, "broken ", r.StripOrphanHandles("broken res://1"))
+}
+
+func TestOrphanAliasesNilRegistry(t *testing.T) {
+	var r *resourceRegistry
+	require.Equal(t, []string{"res://0001"}, r.OrphanHandles("res://0001"))
+}

internal/modelcontext/sources.go

@@ -0,0 +1,584 @@
+// sources.go is the source-reference half of the model-context registry:
+// request-local cN/dN/bN/wN handles for chunks, documents, knowledge bases
+// and web pages, plus the tool-argument codec that maps them back to durable
+// identifiers. Request lifecycles use Registry so source and resource handles
+// cannot be encoded or decoded out of order.
+package modelcontext
+
+import (
+	"encoding/json"
+	"net/url"
+	"regexp"
+	"sort"
+	"strings"
+
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+)
+
+type ChunkReference struct {
+	ChunkID         string
+	KnowledgeID     string
+	KnowledgeBaseID string
+	DocumentTitle   string
+	ChunkIndex      int
+	ChunkType       string
+}
+
+// webMeta is the per-web-page metadata stored next to the raw URL.
+type webMeta struct {
+	title string
+}
+
+// sourceRegistry is scoped to one assistant response (including every Agent tool
+// round). Handles are never persisted or accepted across requests.
+type sourceRegistry struct {
+	citationsEnabled bool
+
+	chunks *handleTable[ChunkReference]
+	docs   *handleTable[struct{}]
+	kbs    *handleTable[struct{}]
+	webs   *handleTable[webMeta]
+}
+
+func newSourceRegistry(citationsEnabled ...bool) *sourceRegistry {
+	enabled := true
+	if len(citationsEnabled) > 0 {
+		enabled = citationsEnabled[0]
+	}
+	return &sourceRegistry{
+		citationsEnabled: enabled,
+		chunks:           newHandleTable[ChunkReference]("c", 0, 1),
+		docs:             newHandleTable[struct{}]("d", 0, 1),
+		kbs:              newHandleTable[struct{}]("b", 0, 1),
+		webs:             newHandleTable[webMeta]("w", 0, 1),
+	}
+}
+
+func (r *sourceRegistry) Count() int {
+	if r == nil {
+		return 0
+	}
+	return r.chunks.size() + r.webs.size()
+}
+
+// knownHandle implements the shared guard for handle-shaped registration
+// input: a model-emitted handle is echoed back only when it already exists,
+// and is never accepted as a new durable identity.
+func knownHandle[M any](table *handleTable[M], id string) string {
+	handle := strings.ToLower(id)
+	if table.has(handle) {
+		return handle
+	}
+	return ""
+}
+
+func (r *sourceRegistry) RegisterChunk(ref ChunkReference) string {
+	if r == nil {
+		return ""
+	}
+	ref.ChunkID = strings.TrimSpace(ref.ChunkID)
+	if ref.ChunkID == "" {
+		return ""
+	}
+	if shortSourceHandleRE.MatchString(ref.ChunkID) {
+		return knownHandle(r.chunks, ref.ChunkID)
+	}
+	return r.chunks.register(ref.ChunkID, ref.ChunkID, ref, mergeChunkReference)
+}
+
+func mergeChunkReference(dst *ChunkReference, src ChunkReference) {
+	if dst.KnowledgeID == "" {
+		dst.KnowledgeID = src.KnowledgeID
+	}
+	if dst.KnowledgeBaseID == "" {
+		dst.KnowledgeBaseID = src.KnowledgeBaseID
+	}
+	if dst.DocumentTitle == "" {
+		dst.DocumentTitle = src.DocumentTitle
+	}
+	if dst.ChunkIndex == 0 {
+		dst.ChunkIndex = src.ChunkIndex
+	}
+	if dst.ChunkType == "" {
+		dst.ChunkType = src.ChunkType
+	}
+}
+
+func (r *sourceRegistry) RegisterDocument(id string) string {
+	id = strings.TrimSpace(id)
+	if r == nil || id == "" {
+		return ""
+	}
+	if shortSourceHandleRE.MatchString(id) {
+		return knownHandle(r.docs, id)
+	}
+	return r.docs.register(id, id, struct{}{}, nil)
+}
+
+func (r *sourceRegistry) RegisterKnowledgeBase(id string) string {
+	id = strings.TrimSpace(id)
+	if r == nil || id == "" {
+		return ""
+	}
+	if shortSourceHandleRE.MatchString(id) {
+		return knownHandle(r.kbs, id)
+	}
+	return r.kbs.register(id, id, struct{}{}, nil)
+}
+
+func (r *sourceRegistry) RegisterWeb(rawURL, title string) string {
+	rawURL = strings.TrimSpace(rawURL)
+	if r == nil || rawURL == "" {
+		return ""
+	}
+	if shortSourceHandleRE.MatchString(rawURL) {
+		return knownHandle(r.webs, rawURL)
+	}
+	// Dedup on the canonical (fragment-stripped) URL while decoding back to
+	// the raw URL the model was originally shown.
+	return r.webs.register(canonicalWebURL(rawURL), rawURL, webMeta{title: title}, func(dst *webMeta, src webMeta) {
+		if dst.title == "" && src.title != "" {
+			dst.title = src.title
+		}
+	})
+}
+
+func canonicalWebURL(raw string) string {
+	parsed, err := url.Parse(strings.TrimSpace(raw))
+	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
+		return strings.TrimSpace(raw)
+	}
+	parsed.Fragment = ""
+	return parsed.String()
+}
+
+func (r *sourceRegistry) RegisterSearchResults(results []*types.SearchResult) {
+	for _, result := range results {
+		if result == nil {
+			continue
+		}
+		r.RegisterDocument(result.KnowledgeID)
+		r.RegisterKnowledgeBase(result.KnowledgeBaseID)
+		r.RegisterChunk(ChunkReference{
+			ChunkID:         result.ID,
+			KnowledgeID:     result.KnowledgeID,
+			KnowledgeBaseID: result.KnowledgeBaseID,
+			DocumentTitle:   firstNonEmpty(result.KnowledgeTitle, result.KnowledgeFilename),
+			ChunkIndex:      result.ChunkIndex,
+			ChunkType:       result.ChunkType,
+		})
+	}
+}
+
+func firstNonEmpty(values ...string) string {
+	for _, value := range values {
+		if strings.TrimSpace(value) != "" {
+			return value
+		}
+	}
+	return ""
+}
+
+func (r *sourceRegistry) ChunkHandle(id string) string {
+	handle, _ := r.chunks.handleForKey(id)
+	return handle
+}
+
+// toolArgumentPolicy decides whether a source-bearing JSON key belongs to a
+// particular tool contract. Request lifecycles always pass the per-tool policy
+// (sourceArgumentAllowed); a nil policy allows every key and exists only for
+// package-internal replay paths that predate per-tool contracts.
+type toolArgumentPolicy func(toolName, key string) bool
+
+// DecodeToolCallsWithPolicy restores handles 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 *sourceRegistry) DecodeToolCallsWithPolicy(toolCalls []types.LLMToolCall, policy toolArgumentPolicy) {
+	for i := range toolCalls {
+		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 *sourceRegistry) 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 *sourceRegistry) 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 := sourceKeySpaces[key]; !ok || !allowed(key) {
+			return
+		}
+		handle := strings.TrimSpace(typed)
+		if shortSourceHandleRE.MatchString(handle) && (r == nil || r.durableForHandle(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)
+		}
+	}
+}
+
+// EncodeMessagesWithPolicies compacts known real identifiers in replayed
+// messages and gates source processing for tool results by tool name. A nil
+// policy retains the legacy generic behavior for package-internal callers.
+func (r *sourceRegistry) EncodeMessagesWithPolicies(
+	messages []chat.Message,
+	argumentPolicy toolArgumentPolicy,
+	resultPolicy func(toolName string) bool,
+) []chat.Message {
+	if r == nil || len(messages) == 0 {
+		return messages
+	}
+	out := make([]chat.Message, len(messages))
+	copy(out, messages)
+	// First register every durable identifier present in historical tool calls
+	// 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 {
+		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" || processToolResult) {
+					out[i].MultiContent[j].Text = r.CompactPublicCitations(out[i].MultiContent[j].Text)
+				}
+			}
+		}
+		if len(out[i].ToolCalls) > 0 {
+			out[i].ToolCalls = append([]chat.ToolCall(nil), out[i].ToolCalls...)
+			for j := range out[i].ToolCalls {
+				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" && (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 {
+			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 shortSourceHandleRE = regexp.MustCompile(`(?i)^[cdbw][1-9][0-9]*$`)
+
+var shortSourceHandleInTextRE = regexp.MustCompile(`(?i)\b[cdbw][1-9][0-9]*\b`)
+
+// 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 *sourceRegistry) DecodeKnownText(text string) string {
+	if r == nil || text == "" {
+		return text
+	}
+	return shortSourceHandleInTextRE.ReplaceAllStringFunc(text, func(handle string) string {
+		if real := r.durableForHandle(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 handle that happens to look like d1 or b2.
+func (r *sourceRegistry) DecodeKnownQuotedText(text string) string {
+	if r == nil || text == "" {
+		return text
+	}
+	return rewriteQuotedText(text, func(segment string) string {
+		return shortSourceHandleInTextRE.ReplaceAllStringFunc(segment, func(handle string) string {
+			if real := r.durableForHandle(handle); real != "" {
+				return real
+			}
+			return handle
+		})
+	})
+}
+
+// UnresolvedQuotedTextHandles reports handle-shaped values inside quoted
+// structured-text segments that do not exist in this request registry.
+func (r *sourceRegistry) UnresolvedQuotedTextHandles(text string) []string {
+	if text == "" {
+		return nil
+	}
+	seen := make(map[string]struct{})
+	rewriteQuotedText(text, func(segment string) string {
+		for _, handle := range shortSourceHandleInTextRE.FindAllString(segment, -1) {
+			if r == nil || r.durableForHandle(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 *sourceRegistry) registerToolArguments(raw string, allowed func(string) bool) {
+	if r == nil || strings.TrimSpace(raw) == "" {
+		return
+	}
+	var value interface{}
+	if err := json.Unmarshal([]byte(raw), &value); err != nil {
+		return
+	}
+	r.registerToolArgumentValue("", value, allowed)
+}
+
+func (r *sourceRegistry) registerToolArgumentValue(key string, value interface{}, allowed func(string) bool) {
+	switch typed := value.(type) {
+	case string:
+		if allowed(strings.ToLower(key)) {
+			r.registerSourceIDByKey(key, typed)
+		}
+	case []interface{}:
+		for _, item := range typed {
+			r.registerToolArgumentValue(key, item, allowed)
+		}
+	case map[string]interface{}:
+		for childKey, item := range typed {
+			r.registerToolArgumentValue(childKey, item, allowed)
+		}
+	}
+}
+
+// registerSourceIDByKey is the single key→source-space dispatch used for tool
+// arguments, structured tool results, and database rows. It is driven by
+// sourceKeySpaces — the same table that gates handle decode — so the recognized
+// key set (and the http/https guard for web references) cannot drift between
+// registration and decoding.
+func (r *sourceRegistry) registerSourceIDByKey(key, value string) {
+	value = strings.TrimSpace(value)
+	if value == "" || shortSourceHandleRE.MatchString(value) {
+		return
+	}
+	space, ok := sourceKeySpaces[strings.ToLower(key)]
+	if !ok {
+		return
+	}
+	switch space {
+	case spaceChunk:
+		r.RegisterChunk(ChunkReference{ChunkID: value})
+	case spaceDocument:
+		r.RegisterDocument(value)
+	case spaceDocumentRef:
+		// Stored refs use "knowledgeID|title"; only the ID part is durable.
+		r.RegisterDocument(strings.TrimSpace(strings.SplitN(value, "|", 2)[0]))
+	case spaceKnowledgeBase:
+		r.RegisterKnowledgeBase(value)
+	case spaceWeb:
+		// Only public web pages become web references. Internal schemes
+		// (res://, storage providers) must never enter the web handle 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, "")
+		}
+	}
+}
+
+func (r *sourceRegistry) decodeJSONWithPolicy(raw string, encode bool, allowed func(string) bool) string {
+	if r == nil || strings.TrimSpace(raw) == "" {
+		return raw
+	}
+	var value interface{}
+	if err := json.Unmarshal([]byte(raw), &value); err != nil {
+		return raw
+	}
+	value = r.walkJSON("", value, encode, allowed)
+	encoded, err := json.Marshal(value)
+	if err != nil {
+		return raw
+	}
+	return string(encoded)
+}
+
+func (r *sourceRegistry) 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.
+			if handle := r.handleForDurable(typed); handle != "" {
+				return handle
+			}
+			return typed
+		}
+		// Decode only ID-bearing keys, and only when the value is handle-shaped,
+		// so ordinary strings that coincidentally equal an handle are preserved.
+		if _, ok := sourceKeySpaces[strings.ToLower(key)]; !ok {
+			return typed
+		}
+		if !shortSourceHandleRE.MatchString(strings.TrimSpace(typed)) {
+			return typed
+		}
+		if real := r.durableForHandle(typed); real != "" {
+			return real
+		}
+		return typed
+	case []interface{}:
+		for i := range typed {
+			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, allowed)
+		}
+	}
+	return value
+}
+
+func (r *sourceRegistry) handleForDurable(real string) string {
+	if handle, ok := r.chunks.handleForKey(real); ok {
+		return handle
+	}
+	if handle, ok := r.docs.handleForKey(real); ok {
+		return handle
+	}
+	if handle, ok := r.kbs.handleForKey(real); ok {
+		return handle
+	}
+	if handle, ok := r.webs.handleForKey(canonicalWebURL(real)); ok {
+		return handle
+	}
+	return ""
+}
+
+func (r *sourceRegistry) durableForHandle(handle string) string {
+	handle = strings.ToLower(strings.TrimSpace(handle))
+	if real, _, ok := r.chunks.resolve(handle); ok {
+		return real
+	}
+	if real, _, ok := r.docs.resolve(handle); ok {
+		return real
+	}
+	if real, _, ok := r.kbs.resolve(handle); ok {
+		return real
+	}
+	if real, _, ok := r.webs.resolve(handle); ok {
+		return real
+	}
+	return ""
+}
+
+// CompactKnownText is intentionally limited to identifiers already registered
+// from structured runtime/tool data. It is used for metadata envelopes, not
+// arbitrary retrieved prose.
+func (r *sourceRegistry) CompactKnownText(text string) string {
+	if r == nil || text == "" {
+		return text
+	}
+	// The snapshot spans all four source tables and is sorted longest-value
+	// first GLOBALLY: a web URL may contain a registered document UUID as a
+	// substring, so per-table passes could corrupt the longer value.
+	pairs := r.chunks.pairs()
+	pairs = append(pairs, r.docs.pairs()...)
+	pairs = append(pairs, r.kbs.pairs()...)
+	pairs = append(pairs, r.webs.pairs()...)
+	sort.SliceStable(pairs, func(i, j int) bool { return len(pairs[i].value) > len(pairs[j].value) })
+	for _, item := range pairs {
+		if item.value != "" {
+			text = strings.ReplaceAll(text, item.value, item.handle)
+		}
+	}
+	return text
+}

internal/modelcontext/sources_test.go

@@ -0,0 +1,552 @@
+package modelcontext
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRegistryChunkAliasIsStableAndExpandsCanonicalCitation(t *testing.T) {
+	registry := newSourceRegistry()
+	first := registry.RegisterChunk(ChunkReference{
+		ChunkID:         "chunk-uuid-1",
+		KnowledgeID:     "knowledge-uuid-1",
+		KnowledgeBaseID: "kb-uuid-1",
+		DocumentTitle:   "Architecture.md",
+		ChunkIndex:      7,
+	})
+	second := registry.RegisterChunk(ChunkReference{
+		ChunkID:       "chunk-uuid-1",
+		KnowledgeID:   "knowledge-uuid-1",
+		DocumentTitle: "Architecture.md",
+	})
+
+	require.Equal(t, "c1", first)
+	require.Equal(t, first, second)
+	require.Equal(t,
+		`claim <kb doc="Architecture.md" chunk_id="chunk-uuid-1" kb_id="kb-uuid-1" />`,
+		registry.ExpandText(`claim <ref id="c1"/>`),
+	)
+}
+
+func TestRegisterDoesNotTreatModelAliasesAsNewDurableIdentities(t *testing.T) {
+	r := newSourceRegistry()
+	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"))
+
+	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, r.chunks.size())
+	require.Equal(t, 1, r.docs.size())
+	require.Equal(t, 1, r.kbs.size())
+	require.Equal(t, 1, r.webs.size())
+}
+
+func TestRegistrySuppressesSourceCitationsWhenDisabled(t *testing.T) {
+	registry := newSourceRegistry(false)
+	registry.RegisterChunk(ChunkReference{ChunkID: "chunk-1", DocumentTitle: "Doc"})
+	registry.RegisterWeb("https://example.com", "Example")
+
+	require.Contains(t, sourceProtocolPrompt(false), "Source citations are disabled")
+	require.NotContains(t, sourceProtocolPrompt(false), `Cite a knowledge chunk with exactly`)
+	require.Equal(t, "knowledge  web ", registry.ExpandText(
+		`knowledge <ref id="c1"/> web <ref id="w1"/>`,
+	))
+	require.Equal(t, "forged  ", registry.ExpandText(
+		`forged <kb doc="Doc" chunk_id="raw" /> <web url="https://example.com" />`,
+	))
+}
+
+func TestRegistryDecodesAliasesInNestedToolArguments(t *testing.T) {
+	registry := NewRegistry(true)
+	registry.RegisterDocument("knowledge-uuid-1")
+	registry.RegisterChunk(ChunkReference{ChunkID: "chunk-uuid-1"})
+
+	// Handles nested inside arrays and objects must decode as long as the key
+	// belongs to the named tool's declared source contract.
+	calls := []types.LLMToolCall{{
+		Function: types.FunctionCall{
+			Name:      "list_knowledge_chunks",
+			Arguments: `{"knowledge_id":"d1","filters":{"chunk_id":"c1","faq_id":["c1"]}}`,
+		},
+	}}
+	registry.DecodeToolCalls(calls)
+
+	require.JSONEq(t,
+		`{"knowledge_id":"knowledge-uuid-1","filters":{"chunk_id":"chunk-uuid-1","faq_id":["chunk-uuid-1"]}}`,
+		calls[0].Function.Arguments,
+	)
+}
+
+func TestDecodeToolCallsOnlyRewritesAliasBearingKeys(t *testing.T) {
+	registry := NewRegistry(true)
+	registry.RegisterKnowledgeBase("kb-uuid-1")
+
+	// A free-text field (query) whose value coincidentally equals a handle must
+	// be preserved verbatim, while declared ID-bearing keys resolve to real IDs.
+	calls := []types.LLMToolCall{{
+		Function: types.FunctionCall{
+			Name:      "knowledge_search",
+			Arguments: `{"query":"b1","content":"see b1 for details","knowledge_base_ids":["b1"]}`,
+		},
+	}}
+	registry.DecodeToolCalls(calls)
+
+	require.JSONEq(t,
+		`{"query":"b1","content":"see b1 for details","knowledge_base_ids":["kb-uuid-1"]}`,
+		calls[0].Function.Arguments,
+	)
+}
+
+func TestStreamExpanderHoldsSplitReferenceAndDropsUnknown(t *testing.T) {
+	registry := newSourceRegistry()
+	registry.RegisterChunk(ChunkReference{ChunkID: "chunk-1", DocumentTitle: "Doc"})
+	expander := newCitationStreamExpander(registry)
+
+	require.Equal(t, "before ", expander.Feed(`before <ref id="`))
+	require.Equal(t, `<kb doc="Doc" chunk_id="chunk-1" /> after`, expander.Feed(`c1"/> after`))
+	require.Empty(t, expander.Flush())
+	require.Equal(t, "x  y", registry.ExpandText(`x <ref id="c999"/> y`))
+	require.Equal(t, "x  y", registry.ExpandText(`x <ref id="d1"/> y`))
+	require.Equal(t, "x  y", registry.ExpandText(`x <ref id='c1'/> y`))
+	require.Equal(t, "x ", registry.ExpandText(`x <ref id="c1"`))
+	require.Equal(t, "x  y", registry.ExpandText(`x <kb doc="forged" chunk_id="forged" /> y`))
+	require.Equal(t, "x ", expander.Feed(`x <we`))
+	require.Equal(t, " y", expander.Feed(`b url="https://forged" /> y`))
+}
+
+func TestEncodeMessagesCompactsCanonicalCitationsFromHistory(t *testing.T) {
+	registry := newSourceRegistry()
+	messages := []chat.Message{{
+		Role: "assistant",
+		Content: `Knowledge <kb doc="A &amp; B.pdf" chunk_id="chunk-real" kb_id="kb-real" />; ` +
+			`web <web url="https://example.com/a?x=1&amp;y=2" title="Example &amp; More" />`,
+	}}
+
+	encoded := registry.EncodeMessagesWithPolicies(messages, nil, nil)
+	require.Equal(t, `Knowledge <ref id="c1"/>; web <ref id="w1"/>`, encoded[0].Content)
+	require.NotContains(t, encoded[0].Content, "chunk-real")
+	require.NotContains(t, encoded[0].Content, "https://example.com")
+	require.Equal(t,
+		`<kb doc="A &amp; B.pdf" chunk_id="chunk-real" kb_id="kb-real" /> <web url="https://example.com/a?x=1&amp;y=2" title="Example &amp; More" />`,
+		registry.ExpandText(`<ref id="c1"/> <ref id="w1"/>`),
+	)
+}
+
+func TestEncodeMessagesMigratesLegacyToolHistoryAtReadTime(t *testing.T) {
+	registry := newSourceRegistry()
+	messages := []chat.Message{
+		{
+			Role: "assistant",
+			ToolCalls: []chat.ToolCall{{
+				Function: chat.FunctionCall{
+					Name:      "knowledge_search",
+					Arguments: `{"knowledge_base_ids":["kb-real"],"knowledge_ids":["doc-real"]}`,
+				},
+			}},
+		},
+		{
+			Role: "tool",
+			Name: "knowledge_search",
+			Content: `<chunk chunk_id="chunk-real" knowledge_id="doc-real" knowledge_base_id="kb-real" ` +
+				`knowledge_title="Legacy Doc">legacy content</chunk>`,
+		},
+		{
+			Role:    "assistant",
+			Content: `Legacy answer <kb doc="Legacy Doc" chunk_id="chunk-real" kb_id="kb-real" />`,
+		},
+	}
+
+	encoded := registry.EncodeMessagesWithPolicies(messages, nil, nil)
+	require.JSONEq(t, `{"knowledge_base_ids":["b1"],"knowledge_ids":["d1"]}`,
+		encoded[0].ToolCalls[0].Function.Arguments)
+	require.Contains(t, encoded[1].Content, `chunk_id="c1"`)
+	require.Contains(t, encoded[1].Content, `knowledge_id="d1"`)
+	require.Contains(t, encoded[1].Content, `knowledge_base_id="b1"`)
+	require.Equal(t, `Legacy answer <ref id="c1"/>`, encoded[2].Content)
+	require.Equal(t,
+		`<kb doc="Legacy Doc" chunk_id="chunk-real" kb_id="kb-real" />`,
+		registry.ExpandText(`<ref id="c1"/>`),
+	)
+}
+
+func TestEncodeMessagesDoesNotTreatLegacyPromptExampleAsARealSource(t *testing.T) {
+	registry := newSourceRegistry()
+	messages := []chat.Message{{
+		Role:    "system",
+		Content: `Old rule: cite <kb doc="..." chunk_id="..." />`,
+	}}
+
+	encoded := registry.EncodeMessagesWithPolicies(messages, nil, nil)
+	require.Equal(t, messages[0].Content, encoded[0].Content)
+	require.Zero(t, registry.Count())
+}
+
+func TestModelOutputGroupsChunksAndReusesAliasAcrossTools(t *testing.T) {
+	registry := newSourceRegistry()
+	search := &types.ToolResult{
+		Success: true,
+		Output:  "raw UUID output",
+		Data: map[string]interface{}{
+			"display_type": "search_results",
+			"results": []map[string]interface{}{
+				{
+					"chunk_id":          "chunk-uuid-1",
+					"knowledge_id":      "knowledge-uuid-1",
+					"knowledge_base_id": "kb-uuid-1",
+					"knowledge_title":   "Doc A",
+					"chunk_index":       3,
+					"content":           "full content",
+				},
+			},
+		},
+	}
+	first := registry.ModelOutput(search)
+	require.Contains(t, first, `<document id="d1" kb="b1" title="Doc A">`)
+	require.Contains(t, first, `<chunk id="c1" index="3" view="full">`)
+	require.NotContains(t, first, "chunk-uuid-1")
+	require.NotContains(t, first, "knowledge-uuid-1")
+
+	deepRead := &types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type":    "knowledge_chunks_list",
+			"knowledge_id":    "knowledge-uuid-1",
+			"knowledge_title": "Doc A",
+			"total_chunks":    int64(1),
+			"fetched_chunks":  1,
+			"chunks": []map[string]interface{}{
+				{"chunk_id": "chunk-uuid-1", "knowledge_id": "knowledge-uuid-1", "knowledge_base": "kb-uuid-1", "chunk_index": 3, "content": "deep content"},
+			},
+		},
+	}
+	second := registry.ModelOutput(deepRead)
+	require.Contains(t, second, `<chunk id="c1"`)
+	require.False(t, strings.Contains(second, "c2"))
+}
+
+func TestModelOutputRendersKnowledgeMetadataOncePerDocument(t *testing.T) {
+	registry := newSourceRegistry()
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "search_results",
+			"results": []map[string]interface{}{
+				{
+					"chunk_id": "chunk-1", "chunk_index": 1, "content": "first",
+					"knowledge_id": "document-1", "knowledge_base_id": "kb-1", "knowledge_title": "Doc",
+					"knowledge_metadata": "region: Shanghai & Suzhou",
+				},
+				{
+					"chunk_id": "chunk-2", "chunk_index": 2, "content": "second",
+					"knowledge_id": "document-1", "knowledge_base_id": "kb-1", "knowledge_title": "Doc",
+					"knowledge_metadata": "region: Shanghai & Suzhou",
+				},
+			},
+		},
+	})
+
+	require.Equal(t, 1, strings.Count(output, "<metadata>"))
+	require.Contains(t, output, "<metadata>region: Shanghai &amp; Suzhou</metadata>")
+	require.Equal(t, 2, strings.Count(output, "<chunk "))
+}
+
+func TestModelOutputWebAliasExpandsToWebCitation(t *testing.T) {
+	registry := newSourceRegistry()
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "web_search_results",
+			"results": []map[string]interface{}{
+				{"title": "Example", "url": "https://example.com/a", "snippet": "snippet"},
+			},
+		},
+	})
+	require.Contains(t, output, `<page id="w1" title="Example">`)
+	require.Contains(t, output, `<evidence type="search_summary" verified="false" />`)
+	require.NotContains(t, output, "https://example.com/a")
+	require.Equal(t,
+		`<web url="https://example.com/a" title="Example" />`,
+		registry.ExpandText(`<ref id="w1"/>`),
+	)
+}
+
+func TestModelOutputWebSearchRetainsContentOnlyEvidence(t *testing.T) {
+	registry := newSourceRegistry()
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "web_search_results",
+			"results": []map[string]interface{}{
+				{
+					"title":   "Content-only provider",
+					"url":     "https://example.com/content",
+					"content": "search evidence from provider content",
+				},
+			},
+		},
+	})
+
+	require.Contains(t, output, `<content>search evidence from provider content</content>`)
+	require.Contains(t, output, `<evidence type="search_summary" verified="false" />`)
+}
+
+func TestModelOutputWebSearchLimitsProviderContent(t *testing.T) {
+	registry := newSourceRegistry()
+	content := strings.Repeat("provider evidence ", 1000)
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "web_search_results",
+			"results": []map[string]interface{}{{
+				"title":   "Long provider result",
+				"url":     "https://example.com/long",
+				"content": content,
+			}},
+		},
+	})
+
+	require.Contains(t, output, `<content truncated="true">`)
+	require.NotContains(t, output, content)
+}
+
+func TestModelOutputWebFetchPreservesPartialFailureStatus(t *testing.T) {
+	registry := newSourceRegistry()
+	registry.RegisterWeb("https://example.com/verified", "Verified")
+	registry.RegisterWeb("https://example.com/forbidden", "Forbidden")
+
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "web_fetch_results",
+			"results": []map[string]interface{}{
+				{
+					"url":         "https://example.com/verified",
+					"status":      "success",
+					"raw_content": "verified page content",
+				},
+				{
+					"url":           "https://example.com/forbidden",
+					"status":        "failed",
+					"retryable":     false,
+					"error_code":    "http_403",
+					"error_message": "access denied",
+				},
+			},
+		},
+	})
+
+	require.Contains(t, output, `<page id="w1" status="success" view="full">`)
+	require.Contains(t, output, "verified page content")
+	require.Contains(t, output, `<page id="w2" status="failed" retryable="false" error_code="http_403">`)
+	require.Contains(t, output, `<error>access denied</error>`)
+	require.Contains(t, output, "failed URLs do not invalidate successful evidence")
+	require.NotContains(t, output, "https://example.com")
+}
+
+func TestModelOutputWebFetchAllFailuresIncludeSearchFallback(t *testing.T) {
+	registry := newSourceRegistry()
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "web_fetch_results",
+			"results": []map[string]interface{}{
+				{
+					"url":           "https://example.com/dns",
+					"status":        "failed",
+					"retryable":     true,
+					"error_code":    "dns_failed",
+					"error_message": "DNS lookup failed",
+				},
+			},
+		},
+	})
+
+	require.Contains(t, output, `status="failed" retryable="true" error_code="dns_failed"`)
+	require.Contains(t, output, "Stop expanding web searches")
+	require.Contains(t, output, "page content was not verified")
+}
+
+func TestModelOutputWebFetchAllFailuresStillStructuredWhenToolNotSuccessful(t *testing.T) {
+	registry := newSourceRegistry()
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: false,
+		Error:   "all page fetches failed",
+		Data: map[string]interface{}{
+			"display_type": "web_fetch_results",
+			"results": []map[string]interface{}{
+				{
+					"url":           "https://example.com/dns",
+					"status":        "failed",
+					"retryable":     true,
+					"error_code":    "dns_failed",
+					"error_message": "DNS lookup failed",
+				},
+			},
+		},
+	})
+
+	require.Contains(t, output, `status="failed" retryable="true" error_code="dns_failed"`)
+	require.Contains(t, output, "Stop expanding web searches")
+	require.NotContains(t, output, "Error: all page fetches failed")
+}
+
+func TestModelOutputWebFetchKeepsContentWhenSummaryFails(t *testing.T) {
+	registry := newSourceRegistry()
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "web_fetch_results",
+			"results": []map[string]interface{}{
+				{
+					"url":                   "https://example.com/specs",
+					"status":                "success",
+					"summary_status":        "failed",
+					"summary_error_code":    "summary_failed",
+					"summary_error_message": "model unavailable",
+					"raw_content":           "official specifications",
+				},
+			},
+		},
+	})
+
+	require.Contains(t, output, `<summary_error code="summary_failed">model unavailable</summary_error>`)
+	require.Contains(t, output, `<content>official specifications</content>`)
+}
+
+func TestModelOutputWebFetchLimitsRawContentAcrossPages(t *testing.T) {
+	registry := newSourceRegistry()
+	pageContent := strings.Repeat("verified page content ", 1000)
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "web_fetch_results",
+			"results": []map[string]interface{}{
+				{"url": "https://example.com/one", "status": "success", "raw_content": pageContent},
+				{"url": "https://example.com/two", "status": "success", "raw_content": pageContent},
+			},
+		},
+	})
+
+	require.Contains(t, output, `truncated="true"`)
+	require.Less(t, len([]rune(output)), modelWebFetchTotalMaxRunes+3000)
+}
+
+func TestModelOutputDocumentInfoUsesDocumentAndFAQAliases(t *testing.T) {
+	registry := newSourceRegistry()
+	result := &types.ToolResult{
+		Success: true,
+		Output:  "raw IDs must not be used",
+		Data: map[string]interface{}{
+			"display_type": "document_info",
+			"documents": []map[string]interface{}{
+				{
+					"knowledge_id": "doc-real-id",
+					"title":        "Architecture",
+					"description":  "System overview",
+					"type":         "file",
+					"chunk_count":  12,
+				},
+				{
+					"faq_id":       "faq-chunk-real-id",
+					"knowledge_id": "faq-container-real-id",
+					"faq_question": "How does it work?",
+					"faq_answers":  []string{"With short aliases."},
+					"is_faq":       true,
+				},
+			},
+		},
+	}
+
+	output := registry.ModelOutput(result)
+	for _, raw := range []string{"doc-real-id", "faq-chunk-real-id", "faq-container-real-id"} {
+		require.NotContains(t, output, raw)
+	}
+	require.Contains(t, output, `<document id="d1" title="Architecture"`)
+	require.Contains(t, output, `<chunk id="c1" type="faq">`)
+	require.Contains(t, registry.ExpandText(`Answer <ref id="c1"/>`), `chunk_id="faq-chunk-real-id"`)
+}
+
+func TestModelOutputCompactsLabeledWikiReferences(t *testing.T) {
+	registry := newSourceRegistry()
+	result := &types.ToolResult{
+		Success: true,
+		Output: `<wiki_page>
+<knowledge_base_id>kb-real-id</knowledge_base_id>
+<sources><source knowledge_id="doc-real-id">Source</source></sources>
+</wiki_page>`,
+	}
+
+	output := registry.ModelOutput(result)
+	require.NotContains(t, output, "kb-real-id")
+	require.NotContains(t, output, "doc-real-id")
+	require.Contains(t, output, `<knowledge_base_id>b1</knowledge_base_id>`)
+	require.Contains(t, output, `knowledge_id="d1"`)
+
+	toolCalls := []types.LLMToolCall{
+		{Function: types.FunctionCall{Name: "wiki_read_source_doc", Arguments: `{"knowledge_id":"d1"}`}},
+		{Function: types.FunctionCall{Name: "wiki_search", Arguments: `{"knowledge_base_id":"b1"}`}},
+	}
+	registry.DecodeToolCallsWithPolicy(toolCalls, sourceArgumentAllowed)
+	require.JSONEq(t, `{"knowledge_id":"doc-real-id"}`, toolCalls[0].Function.Arguments)
+	require.JSONEq(t, `{"knowledge_base_id":"kb-real-id"}`, toolCalls[1].Function.Arguments)
+}
+
+func TestModelOutputGraphResultsUseChunkAliases(t *testing.T) {
+	registry := newSourceRegistry()
+	output := registry.ModelOutput(&types.ToolResult{
+		Success: true,
+		Data: map[string]interface{}{
+			"display_type": "graph_query_results",
+			"results": []map[string]interface{}{{
+				"chunk_id":          "graph-chunk-real",
+				"chunk_index":       4,
+				"knowledge_id":      "graph-doc-real",
+				"knowledge_base_id": "graph-kb-real",
+				"knowledge_title":   "Graph Source",
+				"content":           "A relates to B.",
+			}},
+		},
+	})
+
+	require.Contains(t, output, `<retrieval type="knowledge" mode="graph">`)
+	require.Contains(t, output, `<chunk id="c1" index="4" view="full">`)
+	require.NotContains(t, output, "graph-chunk-real")
+	require.Equal(t,
+		`<kb doc="Graph Source" chunk_id="graph-chunk-real" kb_id="graph-kb-real" />`,
+		registry.ExpandText(`<ref id="c1"/>`),
+	)
+}
+
+func TestModelOutputDoesNotRegisterInternalSchemesAsWebSources(t *testing.T) {
+	registry := newSourceRegistry()
+	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
+	// handle 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,
+		`<web url="https://example.com/page" title="" />`,
+		registry.ExpandText(`<ref id="w1"/>`),
+	)
+}

internal/modelcontext/stream.go

@@ -0,0 +1,275 @@
+// stream.go owns streaming-safe handle decoding: the shared suffix-hold
+// primitive, the per-space decoders built on it, and the composed public
+// StreamDecoder that applies every stage in the only safe order.
+package modelcontext
+
+import "strings"
+
+// streamHold is the shared primitive for decoders that must never emit a
+// partial model handle: each Feed withholds a trailing byte run that could
+// still grow into a handle in the next provider chunk, and applies the
+// space-specific decode to everything released. flush decides what happens to
+// a suffix still held when the stream closes.
+type streamHold struct {
+	pending string
+	holdLen func(combined string) int    // trailing bytes to withhold
+	emit    func(released string) string // decode applied to released text
+	flush   func(pending string) string  // end-of-stream disposition
+}
+
+func (h *streamHold) Feed(chunk string) string {
+	combined := h.pending + chunk
+	h.pending = ""
+	if combined == "" {
+		return ""
+	}
+	if hold := h.holdLen(combined); hold > 0 && hold <= len(combined) {
+		h.pending = combined[len(combined)-hold:]
+		combined = combined[:len(combined)-hold]
+	}
+	return h.emit(combined)
+}
+
+func (h *streamHold) Flush() string {
+	pending := h.pending
+	h.pending = ""
+	return h.flush(pending)
+}
+
+// resourceStreamDecoder restores res:// handles split across provider chunks.
+type resourceStreamDecoder struct {
+	hold streamHold
+}
+
+func newResourceStreamDecoder(registry *resourceRegistry) *resourceStreamDecoder {
+	return &resourceStreamDecoder{hold: streamHold{
+		holdLen: func(combined string) int {
+			hold := 0
+			for _, handle := range registry.handles() {
+				// 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 a request-local handle never leaks.
+				for n := 1; n < len(handle); n++ {
+					if n > hold && strings.HasSuffix(combined, handle[:n]) {
+						hold = n
+					}
+				}
+			}
+			return hold
+		},
+		emit:  registry.DecodeText,
+		flush: registry.DecodeText,
+	}}
+}
+
+func (d *resourceStreamDecoder) Feed(chunk string) string {
+	if d == nil {
+		return chunk
+	}
+	return d.hold.Feed(chunk)
+}
+
+func (d *resourceStreamDecoder) Flush() string {
+	if d == nil {
+		return ""
+	}
+	return d.hold.Flush()
+}
+
+// HandleStreamDecoder restores HandleTable values without leaking a handle
+// split across provider chunks.
+type HandleStreamDecoder struct {
+	table *HandleTable
+	hold  streamHold
+}
+
+func NewHandleStreamDecoder(table *HandleTable) *HandleStreamDecoder {
+	return &HandleStreamDecoder{table: table, hold: streamHold{
+		holdLen: func(combined string) int {
+			start := len(combined)
+			for start > 0 && isHandleTokenByte(combined[start-1]) {
+				start--
+			}
+			tail := combined[start:]
+			// The prefix is immutable after construction, so no lock is needed.
+			if couldBeNumericHandle(tail, table.table.prefix) {
+				return len(tail)
+			}
+			return 0
+		},
+		emit:  table.DecodeKnownText,
+		flush: table.DecodeKnownText,
+	}}
+}
+
+func (d *HandleStreamDecoder) Feed(chunk string) string {
+	if d == nil || d.table == nil {
+		return chunk
+	}
+	return d.hold.Feed(chunk)
+}
+
+func (d *HandleStreamDecoder) Flush() string {
+	if d == nil || d.table == nil {
+		return ""
+	}
+	return d.hold.Flush()
+}
+
+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
+}
+
+// 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 {
+	hold streamHold
+}
+
+func newOrphanResourceStreamFilter() *orphanResourceStreamFilter {
+	return &orphanResourceStreamFilter{hold: streamHold{
+		holdLen: orphanResourceHoldLen,
+		emit: func(released string) string {
+			return resourceHandleShapeRE.ReplaceAllString(released, "")
+		},
+		flush: orphanResourceFlush,
+	}}
+}
+
+func orphanResourceHoldLen(combined string) int {
+	const prefix = "res://"
+	holdAt := -1
+	// Unknown handles must be safe across every provider split, including
+	// "re" + "s://9999". This may defer at most a few 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 {
+		return 0
+	}
+	return len(combined) - holdAt
+}
+
+func orphanResourceFlush(pending string) string {
+	// 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 ""
+		}
+	}
+	return resourceHandleShapeRE.ReplaceAllString(pending, "")
+}
+
+func (f *orphanResourceStreamFilter) Feed(chunk string) string {
+	if f == nil {
+		return chunk
+	}
+	return f.hold.Feed(chunk)
+}
+
+func (f *orphanResourceStreamFilter) Flush() string {
+	if f == nil {
+		return ""
+	}
+	return f.hold.Flush()
+}
+
+func allDigits(value string) bool {
+	if value == "" {
+		return false
+	}
+	for _, char := range value {
+		if char < '0' || char > '9' {
+			return false
+		}
+	}
+	return true
+}
+
+// StreamDecoder composes resource restoration, source-citation expansion,
+// issue-handle decoding and orphan filtering so callers cannot split handle
+// processing or apply stages in the wrong order.
+type StreamDecoder struct {
+	resources *resourceStreamDecoder
+	sources   *citationStreamExpander
+	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 ""
+	}
+	// Each stage's tail must be fed THROUGH the later stages before those
+	// stages flush their own pending suffix, otherwise a handle completed by
+	// an earlier stage's tail would bypass later decoding.
+	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
+}

internal/modelcontext/tool_policy.go

@@ -0,0 +1,346 @@
+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]*$`)
+
+// sourceKeySpace identifies which source-handle space an ID-bearing JSON key
+// belongs to.
+type sourceKeySpace int
+
+const (
+	spaceChunk sourceKeySpace = iota
+	spaceDocument
+	spaceDocumentRef // "knowledgeID|title" stored refs; only the ID is durable
+	spaceKnowledgeBase
+	spaceWeb
+)
+
+// sourceKeySpaces is the single table of ID-bearing keys the source codec
+// understands. It drives both handle registration (registerSourceIDByKey) and
+// the handle-shaped-value decode gate (walkJSON), so registration and decoding
+// cannot drift apart. Every key referenced by a toolHandlePolicies sourceIDKeys
+// set must appear here (enforced by a test).
+var sourceKeySpaces = map[string]sourceKeySpace{
+	"chunk_id": spaceChunk, "faq_id": spaceChunk, "chunk_ids": spaceChunk, "faq_ids": spaceChunk,
+	"knowledge_id": spaceDocument, "knowledge_ids": spaceDocument, "suspected_knowledge_ids": spaceDocument,
+	"source_refs":    spaceDocumentRef,
+	"knowledge_base": spaceKnowledgeBase, "knowledge_base_id": spaceKnowledgeBase,
+	"knowledge_base_ids": spaceKnowledgeBase, "kb_id": spaceKnowledgeBase, "kb_ids": spaceKnowledgeBase,
+	"url": spaceWeb, "urls": spaceWeb,
+}
+
+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,
+	},
+	// Past conversations carry no durable chunk or document IDs, so there is
+	// nothing to compact; the output is prose the model may quote.
+	"search_conversations": {},
+	// Memories are single sentences about the person. They carry no chunk or
+	// document IDs of their own, and the memory item IDs never leave the
+	// service, so there is nothing here for the model to hold a handle on.
+	"search_memory": {},
+	"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,
+	},
+	// Mutation tools echo the slug they acted on and surface validation errors
+	// that can quote a durable KB or document ID. They own no source arguments,
+	// but their output must still be compacted before the model sees it.
+	"wiki_rename_page": {
+		sourceOutput: true,
+	},
+	"wiki_delete_page": {
+		sourceOutput: true,
+	},
+	// Agent-private bookkeeping tools echo model-authored text. They need
+	// compaction so a durable ID quoted back by the model is re-compacted, but
+	// never structured source rendering.
+	"thinking":   {},
+	"todo_write": {},
+	// Skill output is user-authored content of unknown shape. Compact known
+	// durable IDs, but do not mine it for source keys: a skill's JSON may use
+	// "url"/"knowledge_id" with unrelated semantics.
+	"read_skill":           {},
+	"execute_skill_script": {},
+	// Sandbox shell/file tools return whatever a skill script produced, so they
+	// carry the same unknown-shape caveat as skill output above.
+	"shell_exec":         {},
+	"list_sandbox_files": {},
+	"read_sandbox_file":  {},
+	"write_sandbox_file": {},
+	"edit_sandbox_file":  {},
+	// Installer-only writers for the skill image tree. Same unknown-shape
+	// caveat: they echo paths and byte counts, never knowledge identifiers.
+	"write_skill_file": {},
+	"edit_skill_file":  {},
+}
+
+// HasToolPolicy reports whether a tool has an explicit model-handle policy.
+// Built-in tools must declare one (even an empty policy) so that adding a tool
+// is a deliberate decision about what the model may see, rather than a silent
+// fall-through to fully opaque output.
+func HasToolPolicy(toolName string) bool {
+	_, ok := toolHandlePolicies[toolName]
+	return ok
+}
+
+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
+}

internal/modelcontext/tool_policy_test.go

@@ -0,0 +1,21 @@
+package modelcontext
+
+import "testing"
+
+// Every key a tool declares as a source-ID argument must have a registration
+// space, otherwise its values could be decoded from handles but never
+// registered (or vice versa) — the drift this table exists to prevent.
+func TestEverySourceIDKeyHasARegistrationSpace(t *testing.T) {
+	for tool, policy := range toolHandlePolicies {
+		for key := range policy.sourceIDKeys {
+			if _, ok := sourceKeySpaces[key]; !ok {
+				t.Errorf("tool %s declares source key %q with no entry in sourceKeySpaces", tool, key)
+			}
+		}
+		for key := range policy.sourceTextKeys {
+			if _, ok := sourceKeySpaces[key]; ok {
+				t.Errorf("tool %s declares %q as free text but sourceKeySpaces also claims it as an ID key", tool, key)
+			}
+		}
+	}
+}