fix: support keenable component (#16884)

### Summary

As title


<img width="3763" height="2042" alt="image"
src="https://github.com/user-attachments/assets/0bcf2d41-9475-4d22-b8f0-64a527e22aba"
/>

internal/agent/canvas/runner.go

@@ -107,15 +107,15 @@ type NodeFinishedData struct {
 
 // MessageEvent is the JSON payload for Type=="message" frames.
 type MessageEvent struct {
-	Content   string        `json:"content"`
-	Reference []interface{} `json:"reference,omitempty"`
+	Content   string      `json:"content"`
+	Reference interface{} `json:"reference,omitempty"`
 }
 
 // MessageEndEvent is the JSON payload for Type=="message_end" frames.
 type MessageEndEvent struct {
 	Status     *string       `json:"status,omitempty"`
 	Attachment []interface{} `json:"attachment,omitempty"`
-	Reference  []interface{} `json:"reference,omitempty"`
+	Reference  interface{}   `json:"reference,omitempty"`
 }
 
 // WaitingForUserEvent is the JSON payload for Type=="waiting_for_user"

internal/agent/component/fixture_stubs.go

@@ -521,6 +521,7 @@ func init() {
 	Register("GitHub", newGitHubComponent)
 	Register("Wikipedia", newWikipediaComponent)
 	Register("DuckDuckGo", newDuckDuckGoComponent)
+	Register("KeenableSearch", newKeenableSearchComponent)
 	Register("Google", newGoogleComponent)
 	Register("GoogleScholar", newGoogleScholarComponent)
 	Register("ArXiv", newArxivComponent)

internal/agent/component/keenable_component_test.go

@@ -0,0 +1,169 @@
+//
+//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+package component
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"testing"
+
+	einotool "github.com/cloudwego/eino/components/tool"
+
+	"ragflow/internal/agent/runtime"
+)
+
+type fakeKeenableInvoker struct {
+	args map[string]any
+}
+
+func (f *fakeKeenableInvoker) InvokableRun(_ context.Context, argsJSON string, _ ...einotool.Option) (string, error) {
+	if err := json.Unmarshal([]byte(argsJSON), &f.args); err != nil {
+		return "", err
+	}
+	return `{"results":[{"title":"Keenable result","url":"https://example.com/item","description":"Fresh search result"}]}`, nil
+}
+
+func TestKeenableSearch_RegisteredFactory(t *testing.T) {
+	t.Parallel()
+
+	c, err := New("KeenableSearch", map[string]any{"api_key": "key-1", "mode": "realtime", "top_n": float64(3)})
+	if err != nil {
+		t.Fatalf("New(KeenableSearch) errored: %v", err)
+	}
+	kc, ok := c.(*keenableSearchComponent)
+	if !ok {
+		t.Fatalf("New(KeenableSearch) returned %T, want *keenableSearchComponent", c)
+	}
+	if kc.apiKey != "key-1" {
+		t.Fatalf("apiKey = %q, want key-1", kc.apiKey)
+	}
+	if kc.mode != "realtime" {
+		t.Fatalf("mode = %q, want realtime", kc.mode)
+	}
+	if kc.topN != 3 {
+		t.Fatalf("topN = %d, want 3", kc.topN)
+	}
+
+	formGetter, ok := c.(interface{ GetInputForm() map[string]any })
+	if !ok {
+		t.Fatal("KeenableSearch component does not expose GetInputForm")
+	}
+	form := formGetter.GetInputForm()
+	query, ok := form["query"].(map[string]any)
+	if !ok {
+		t.Fatalf("GetInputForm()[query] has type %T, want map", form["query"])
+	}
+	if query["type"] != "line" {
+		t.Fatalf("GetInputForm()[query][type] = %v, want line", query["type"])
+	}
+	site, ok := form["site"].(map[string]any)
+	if !ok {
+		t.Fatalf("GetInputForm()[site] has type %T, want map", form["site"])
+	}
+	if site["type"] != "line" {
+		t.Fatalf("GetInputForm()[site][type] = %v, want line", site["type"])
+	}
+	if _, ok := c.Outputs()["formalized_content"]; !ok {
+		t.Fatal("Outputs() missing formalized_content")
+	}
+	if _, ok := c.Outputs()["json"]; !ok {
+		t.Fatal("Outputs() missing json")
+	}
+}
+
+func TestKeenableSearch_InvokeAdaptsCanvasInputsAndOutputs(t *testing.T) {
+	t.Parallel()
+
+	fake := &fakeKeenableInvoker{}
+	c := newKeenableSearchComponentWithInvoker(fake, map[string]any{
+		"mode":  "pro",
+		"top_n": 5,
+		"site":  "example.com",
+	})
+
+	state := runtime.NewCanvasState("run-keenable", "task-keenable")
+	ctx := runtime.WithState(context.Background(), state)
+
+	out, err := c.Invoke(ctx, map[string]any{
+		"query": "  agent search  ",
+		"mode":  "realtime",
+		"top_n": float64(2),
+	})
+	if err != nil {
+		t.Fatalf("Invoke errored: %v", err)
+	}
+
+	if got := fake.args["query"]; got != "agent search" {
+		t.Errorf("query arg = %v, want trimmed query", got)
+	}
+	if got := fake.args["mode"]; got != "realtime" {
+		t.Errorf("mode arg = %v, want realtime", got)
+	}
+	if got := fake.args["top_n"]; got != float64(2) && got != 2 {
+		t.Errorf("top_n arg = %v, want 2", got)
+	}
+	if got := fake.args["site"]; got != "example.com" {
+		t.Errorf("site arg = %v, want default site", got)
+	}
+
+	formalized, _ := out["formalized_content"].(string)
+	for _, want := range []string{"Keenable result", "https://example.com/item", "Fresh search result"} {
+		if !strings.Contains(formalized, want) {
+			t.Errorf("formalized_content missing %q: %s", want, formalized)
+		}
+	}
+
+	results, ok := out["json"].([]any)
+	if !ok {
+		t.Fatalf("json output has type %T, want []any", out["json"])
+	}
+	if len(results) != 1 {
+		t.Fatalf("json output length = %d, want 1", len(results))
+	}
+
+	reference := state.GetRetrievalReference()
+	chunks, _ := reference["chunks"].([]any)
+	if len(chunks) != 1 {
+		t.Fatalf("reference chunks length = %d, want 1", len(chunks))
+	}
+	chunk, _ := chunks[0].(map[string]any)
+	if chunk["document_name"] != "Keenable result" || chunk["url"] != "https://example.com/item" {
+		t.Fatalf("reference chunk metadata = %#v", chunk)
+	}
+	docAggs, _ := reference["doc_aggs"].([]any)
+	if len(docAggs) != 1 {
+		t.Fatalf("reference doc_aggs length = %d, want 1", len(docAggs))
+	}
+}
+
+func TestKeenableSearch_InvokeEmptyQueryReturnsEmptyPayload(t *testing.T) {
+	t.Parallel()
+
+	c := newKeenableSearchComponentWithInvoker(&fakeKeenableInvoker{}, nil)
+	out, err := c.Invoke(context.Background(), map[string]any{"query": "   "})
+	if err != nil {
+		t.Fatalf("Invoke errored: %v", err)
+	}
+	if got := out["formalized_content"]; got != "" {
+		t.Errorf("formalized_content = %v, want empty string", got)
+	}
+	results, ok := out["json"].([]any)
+	if !ok || len(results) != 0 {
+		t.Fatalf("json output = %#v, want empty []any", out["json"])
+	}
+}

internal/agent/component/message.go

@@ -175,7 +175,7 @@ func (m *MessageComponent) Invoke(ctx context.Context, inputs map[string]any) (m
 		return nil, fmt.Errorf("Message: nil canvas state")
 	}
 
-	text, _ := inputs["text"].(string)
+	text := extractMessageText(inputs)
 	if text == "" {
 		text = m.text
 	}

internal/agent/component/message_test.go

@@ -98,6 +98,20 @@ func TestMessage_NoTemplate(t *testing.T) {
 	}
 }
 
+func TestMessage_RuntimeContentInput(t *testing.T) {
+	c, _ := NewMessageComponent(nil)
+	state := canvas.NewCanvasState("run-4", "task-4")
+	ctx := withStateForTest(context.Background(), state)
+
+	out, err := c.Invoke(ctx, map[string]any{"content": "from upstream", "stream": false})
+	if err != nil {
+		t.Fatalf("Invoke: %v", err)
+	}
+	if got, _ := out["content"].(string); got != "from upstream" {
+		t.Errorf("content: got %q, want %q", got, "from upstream")
+	}
+}
+
 // withStateForTest is a thin alias for canvas.WithState kept for
 // readability at the test call sites (also defined in begin_test.go).
 // Same package → same symbol; defining it twice is a compile error,

internal/agent/component/universe_a_wrappers.go

@@ -317,6 +317,10 @@ type duckDuckGoInvoker interface {
 	InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error)
 }
 
+type keenableInvoker interface {
+	InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error)
+}
+
 // bgptComponent delegates to internal/agent/tool/BGPTTool and adapts
 // the tool envelope to the BGPT canvas output contract.
 type bgptComponent struct {
@@ -574,6 +578,199 @@ func (c *duckDuckGoComponent) Stream(_ context.Context, _ map[string]any) (<-cha
 	return nil, nil
 }
 
+// keenableSearchComponent delegates to internal/agent/tool/KeenableTool.
+type keenableSearchComponent struct {
+	inner  keenableInvoker
+	apiKey string
+	mode   string
+	topN   int
+	site   string
+}
+
+func newKeenableSearchComponent(params map[string]any) (Component, error) {
+	apiKey := stringParam(params["api_key"])
+	return newKeenableSearchComponentWithInvoker(agenttool.NewKeenableToolWithAPIKey(nil, apiKey), params), nil
+}
+
+func newKeenableSearchComponentWithInvoker(inner keenableInvoker, params map[string]any) Component {
+	if inner == nil {
+		inner = agenttool.NewKeenableTool()
+	}
+	mode := strings.TrimSpace(stringParam(params["mode"]))
+	if mode == "" {
+		mode = "pro"
+	}
+	topN := toIntParam(params["top_n"])
+	if topN <= 0 {
+		topN = 10
+	}
+	return &keenableSearchComponent{
+		inner:  inner,
+		apiKey: stringParam(params["api_key"]),
+		mode:   mode,
+		topN:   topN,
+		site:   stringParam(params["site"]),
+	}
+}
+
+func (c *keenableSearchComponent) Name() string { return "KeenableSearch" }
+
+func (c *keenableSearchComponent) Inputs() map[string]string {
+	return map[string]string{
+		"query":   "Search query.",
+		"site":    "Optional single-domain filter.",
+		"api_key": "Optional Keenable API key.",
+		"mode":    "Search mode: pro or realtime.",
+		"top_n":   "Maximum number of results.",
+	}
+}
+
+func (c *keenableSearchComponent) GetInputForm() map[string]any {
+	return map[string]any{
+		"query": map[string]any{
+			"name": "Query",
+			"type": "line",
+		},
+		"site": map[string]any{
+			"name": "Site",
+			"type": "line",
+		},
+	}
+}
+
+func (c *keenableSearchComponent) Outputs() map[string]string {
+	return map[string]string{
+		"formalized_content": "Rendered search results for downstream LLM prompts.",
+		"json":               "Raw Keenable result list.",
+	}
+}
+
+func (c *keenableSearchComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) {
+	query := strings.TrimSpace(stringParam(inputs["query"]))
+	if query == "" {
+		return map[string]any{"formalized_content": "", "json": []any{}}, nil
+	}
+
+	args := map[string]any{
+		"query": query,
+		"mode":  c.mode,
+		"top_n": c.topN,
+	}
+	if mode := strings.TrimSpace(stringParam(inputs["mode"])); mode != "" {
+		args["mode"] = mode
+	}
+	if topN := toIntParam(inputs["top_n"]); topN > 0 {
+		args["top_n"] = topN
+	}
+	site := strings.TrimSpace(stringParam(inputs["site"]))
+	if site == "" {
+		site = strings.TrimSpace(c.site)
+	}
+	if site != "" {
+		args["site"] = site
+	}
+
+	invoker := c.inner
+	if apiKey := strings.TrimSpace(stringParam(inputs["api_key"])); apiKey != "" && apiKey != strings.TrimSpace(c.apiKey) {
+		invoker = agenttool.NewKeenableToolWithAPIKey(nil, apiKey)
+	}
+
+	argsJSON, _ := json.Marshal(args)
+	out, err := invoker.InvokableRun(ctx, string(argsJSON))
+	decoded := parseToolEnvelope(out)
+	results := anySlice(decoded["results"])
+	if existing, _ := decoded["_ERROR"].(string); strings.TrimSpace(existing) != "" {
+		return map[string]any{
+			"formalized_content": "",
+			"json":               results,
+			"_ERROR":             existing,
+		}, nil
+	}
+	if err != nil {
+		if len(decoded) > 0 {
+			return map[string]any{
+				"formalized_content": "",
+				"json":               results,
+				"_ERROR":             decoded["_ERROR"],
+			}, nil
+		}
+		return nil, fmt.Errorf("canvas: KeenableSearch: %w", err)
+	}
+	chunks, docAggs := buildKeenableReferences(results)
+	if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil {
+		state.SetRetrievalReferences(chunks, docAggs)
+	}
+	return map[string]any{
+		"formalized_content": renderKeenableReferences(chunks),
+		"json":               results,
+	}, nil
+}
+
+func (c *keenableSearchComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) {
+	return nil, nil
+}
+
+func buildKeenableReferences(results []any) ([]map[string]any, []map[string]any) {
+	chunks := make([]map[string]any, 0, len(results))
+	docAggs := make([]map[string]any, 0, len(results))
+	for _, item := range results {
+		m, ok := item.(map[string]any)
+		if !ok {
+			continue
+		}
+		content := truncateRunes(strings.TrimSpace(keenableValueString(m["description"])), 10000)
+		if content == "" || content == "None" {
+			continue
+		}
+		documentID := strconv.FormatInt(githubHashInt(content, 100000000), 10)
+		title := keenableValueString(m["title"])
+		url := keenableValueString(m["url"])
+		displayID := strconv.FormatInt(githubHashInt(documentID, 500), 10)
+		chunks = append(chunks, map[string]any{
+			"id":            displayID,
+			"chunk_id":      documentID,
+			"content":       content,
+			"doc_id":        documentID,
+			"document_id":   documentID,
+			"docnm_kwd":     title,
+			"document_name": title,
+			"similarity":    1,
+			"score":         1,
+			"url":           url,
+		})
+		docAggs = append(docAggs, map[string]any{
+			"doc_name": title,
+			"doc_id":   documentID,
+			"count":    1,
+			"url":      url,
+		})
+	}
+	return chunks, docAggs
+}
+
+func renderKeenableReferences(chunks []map[string]any) string {
+	if len(chunks) == 0 {
+		return ""
+	}
+	blocks := make([]string, 0, len(chunks))
+	for _, chunk := range chunks {
+		blocks = append(blocks, strings.Join([]string{
+			"\nID: " + keenableValueString(chunk["id"]),
+			"├── Title: " + keenableValueString(chunk["docnm_kwd"]),
+			"├── URL: " + keenableValueString(chunk["url"]),
+			"└── Content:\n" + keenableValueString(chunk["content"]),
+		}, "\n"))
+	}
+	return strings.Join(blocks, "\n")
+}
+
+func keenableValueString(value any) string {
+	if value == nil {
+		return "None"
+	}
+	return strings.TrimSpace(fmt.Sprint(value))
+}
+
 func renderDuckDuckGoResults(results []any) string {
 	if len(results) == 0 {
 		return ""

internal/agent/runtime/state.go

@@ -33,6 +33,7 @@ package runtime
 import (
 	"encoding/json"
 	"fmt"
+	"sort"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -368,6 +369,76 @@ func (s *CanvasState) GetRetrievalChunks() []map[string]any {
 	return out
 }
 
+// GetRetrievalReference returns the run-level reference payload consumed by
+// the agent chat stream. It mirrors Python canvas.py's message_end.reference
+// shape while keeping doc_aggs as a list for the current Go frontend path.
+func (s *CanvasState) GetRetrievalReference() map[string]any {
+	if s == nil {
+		return nil
+	}
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	if len(s.Retrieval) == 0 {
+		return nil
+	}
+
+	chunks := copyRetrievalList(s.Retrieval["chunks"])
+	docAggs := copyRetrievalDocAggs(s.Retrieval["doc_aggs"])
+	if len(chunks) == 0 && len(docAggs) == 0 {
+		return nil
+	}
+	return map[string]any{
+		"chunks":   chunks,
+		"doc_aggs": docAggs,
+		"total":    len(chunks),
+	}
+}
+
+func copyRetrievalList(value any) []any {
+	switch list := value.(type) {
+	case []any:
+		out := make([]any, len(list))
+		copy(out, list)
+		return out
+	case []map[string]any:
+		out := make([]any, 0, len(list))
+		for _, item := range list {
+			out = append(out, item)
+		}
+		return out
+	default:
+		return nil
+	}
+}
+
+func copyRetrievalDocAggs(value any) []any {
+	switch aggs := value.(type) {
+	case []any:
+		out := make([]any, len(aggs))
+		copy(out, aggs)
+		return out
+	case []map[string]any:
+		out := make([]any, 0, len(aggs))
+		for _, item := range aggs {
+			out = append(out, item)
+		}
+		return out
+	case map[string]any:
+		keys := make([]string, 0, len(aggs))
+		for key := range aggs {
+			keys = append(keys, key)
+		}
+		sort.Strings(keys)
+		out := make([]any, 0, len(keys))
+		for _, key := range keys {
+			out = append(out, aggs[key])
+		}
+		return out
+	default:
+		return nil
+	}
+}
+
 // SetRetrievalChunks records the supplied chunks into
 // state.Retrieval["chunks"]. Existing entries are replaced
 // (last-writer-wins) so a multi-tool canvas reflects the most

internal/agent/runtime/state_test.go

@@ -133,6 +133,29 @@ func TestCanvasState_SetRetrievalReferencesMergesCalls(t *testing.T) {
 	}
 }
 
+func TestCanvasState_GetRetrievalReferenceReturnsFrontendPayload(t *testing.T) {
+	t.Parallel()
+
+	state := NewCanvasState("run-1", "task-1")
+	state.SetRetrievalReferences(
+		[]map[string]any{{"id": "chunk-1", "document_name": "Doc 1"}},
+		[]map[string]any{{"doc_name": "Doc 1", "doc_id": "doc-1", "count": 1}},
+	)
+
+	reference := state.GetRetrievalReference()
+	chunks, _ := reference["chunks"].([]any)
+	if len(chunks) != 1 {
+		t.Fatalf("chunks length = %d, want 1", len(chunks))
+	}
+	docAggs, _ := reference["doc_aggs"].([]any)
+	if len(docAggs) != 1 {
+		t.Fatalf("doc_aggs length = %d, want 1", len(docAggs))
+	}
+	if reference["total"] != 1 {
+		t.Fatalf("total = %v, want 1", reference["total"])
+	}
+}
+
 func contains(s, substr string) bool {
 	for i := 0; i+len(substr) <= len(s); i++ {
 		if s[i:i+len(substr)] == substr {

internal/agent/tool/keenable.go

@@ -209,6 +209,19 @@ func (k *KeenableTool) Info(_ context.Context) (*schema.ToolInfo, error) {
 	}, nil
 }
 
+func (k *KeenableTool) GetInputForm() map[string]any {
+	return map[string]any{
+		"query": map[string]any{
+			"name": "Query",
+			"type": "line",
+		},
+		"site": map[string]any{
+			"name": "Site",
+			"type": "line",
+		},
+	}
+}
+
 // InvokableRun performs the Keenable search.
 func (k *KeenableTool) InvokableRun(ctx context.Context, argsJSON string, _ ...tool.Option) (string, error) {
 	var p keenableParams

internal/agent/tool/keenable_test.go

@@ -389,4 +389,14 @@ func TestKeenable_Info(t *testing.T) {
 	if strings.Contains(string(paramsJSON), "api_key") {
 		t.Fatalf("Info ParamsOneOf unexpectedly exposes api_key: %s", string(paramsJSON))
 	}
+	form := tool.GetInputForm()
+	for _, key := range []string{"query", "site"} {
+		field, ok := form[key].(map[string]any)
+		if !ok {
+			t.Fatalf("GetInputForm()[%s] = %#v, want field map", key, form[key])
+		}
+		if field["type"] != "line" {
+			t.Fatalf("GetInputForm()[%s][type] = %v, want line", key, field["type"])
+		}
+	}
 }

internal/service/agent.go

@@ -1115,7 +1115,7 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
 		// node_finished events are already emitted per-node by the
 		// statePost wrappers in scheduler.go.
 		var answer string
-		var reference []interface{}
+		var legacyReference []interface{}
 		var downloads any
 		now := float64(time.Now().UnixNano()) / 1e9
 		for _, bucket := range state.Snapshot() {
@@ -1131,12 +1131,13 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
 				answer = v
 			}
 			if v, ok := bucket["reference"].([]interface{}); ok {
-				reference = append(reference, v...)
+				legacyReference = append(legacyReference, v...)
 			}
 			if v, ok := bucket["downloads"]; ok && !emptyDownloadValue(v) {
 				downloads = v
 			}
 		}
+		referencePayload := agentRunReferencePayload(state, legacyReference)
 
 		if err != nil {
 			common.Debug("RunAgent invoke err",
@@ -1149,30 +1150,30 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
 			if canvas.IsInterruptError(err) {
 				s.markRunFailed(ctx2, runID, "interrupt: "+err.Error())
 				if answer != "" {
-					s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference)
+					s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload)
 					msgData, _ := json.Marshal(canvas.MessageEvent{
 						Content:   answer,
-						Reference: reference,
+						Reference: referencePayload,
 					})
 					emit("message", string(msgData))
 
 					meData, _ := json.Marshal(canvas.MessageEndEvent{
-						Reference: reference,
+						Reference: referencePayload,
 					})
 					emit("message_end", string(meData))
 				}
 				return state, err
 			}
 			if shouldTreatAsCompletedLoopRun(err, answer) {
-				s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference)
+				s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload)
 				msgData, _ := json.Marshal(canvas.MessageEvent{
 					Content:   answer,
-					Reference: reference,
+					Reference: referencePayload,
 				})
 				emit("message", string(msgData))
 
 				meData, _ := json.Marshal(canvas.MessageEndEvent{
-					Reference: reference,
+					Reference: referencePayload,
 				})
 				emit("message_end", string(meData))
 
@@ -1196,15 +1197,15 @@ func (s *AgentService) buildRunFunc(canvasID string, versionRow *entity.UserCanv
 		}
 
 		// Emit message + message_end (mirrors Python's ans dict).
-		s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, reference)
+		s.persistAgentRunSession(canvasID, sessionID, messageID, userInput, answer, referencePayload)
 		msgData, _ := json.Marshal(canvas.MessageEvent{
 			Content:   answer,
-			Reference: reference,
+			Reference: referencePayload,
 		})
 		emit("message", string(msgData))
 
 		meData, _ := json.Marshal(canvas.MessageEndEvent{
-			Reference: reference,
+			Reference: referencePayload,
 		})
 		emit("message_end", string(meData))
 
@@ -1261,7 +1262,7 @@ func emptyDownloadValue(value any) bool {
 	}
 }
 
-func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID string, userInput any, answer string, reference []interface{}) {
+func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID string, userInput any, answer string, reference map[string]interface{}) {
 	if sessionID == "" || s == nil || s.api4ConversationDAO == nil || dao.DB == nil {
 		return
 	}
@@ -1283,7 +1284,7 @@ func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID stri
 		session.Message = raw
 	}
 	references := parseAgentSessionReferences(session.Reference)
-	references = append(references, normalizeAgentReferenceEntry(map[string]interface{}{"chunks": reference}))
+	references = append(references, normalizeAgentReferenceEntry(reference))
 	if raw, err := json.Marshal(references); err == nil {
 		session.Reference = raw
 	}
@@ -1292,6 +1293,22 @@ func (s *AgentService) persistAgentRunSession(agentID, sessionID, messageID stri
 	}
 }
 
+func agentRunReferencePayload(state *canvas.CanvasState, legacyChunks []interface{}) map[string]interface{} {
+	if state != nil {
+		if reference := state.GetRetrievalReference(); len(reference) > 0 {
+			return reference
+		}
+	}
+	if len(legacyChunks) == 0 {
+		return nil
+	}
+	return map[string]interface{}{
+		"chunks":   legacyChunks,
+		"doc_aggs": []interface{}{},
+		"total":    len(legacyChunks),
+	}
+}
+
 func stringifyAgentUserInput(userInput any) string {
 	switch v := userInput.(type) {
 	case nil: