Merge pull request #1619 from CREDO23/feat-run-citations

[Feat] Cite scraper runs as verifiable sources in chat

surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/system_prompt/prompts/citations/on.md

@@ -1,8 +1,9 @@
 <citations>
 Cite with one token: the bracket label `[n]`. Every citable result —
 prose from a `task` knowledge_base/research specialist (including the
-knowledge_base specialist's `[n]`-labelled workspace findings) — already
-carries `[n]` labels on a single shared count.
+knowledge_base specialist's `[n]`-labelled workspace findings) and
+scraper specialists' run-backed findings — already carries `[n]` labels
+on a single shared count.
 Those labels are the only citation you write; the server resolves each one
 back to its source after the turn.
 

surfsense_backend/app/agents/chat/multi_agent_chat/shared/citations/markers.py

@@ -1,10 +1,11 @@
 """Map a registered citation to the frontend ``[citation:<payload>]`` payload.
 
 The citation renderer understands a chunk id (``42``), a negative chunk id for
-anonymous uploads (``-3``), and a URL. This is the seam that turns a server-side
-source into one the renderer can resolve; it grows as more source kinds become
-renderable. Kinds with no renderable form yet return ``None`` so the marker is
-dropped rather than emitted broken.
+anonymous uploads (``-3``), a URL, and a scraper-run handle (``run_<uuid>``).
+This is the seam that turns a server-side source into one the renderer can
+resolve; it grows as more source kinds become renderable. Kinds with no
+renderable form yet return ``None`` so the marker is dropped rather than
+emitted broken.
 """
 
 from __future__ import annotations
@@ -22,6 +23,9 @@ def to_frontend_payload(entry: CitationEntry) -> str | None:
         case CitationSourceType.WEB_RESULT:
             url = locator.get("url")
             return url or None
+        case CitationSourceType.RUN:
+            run_id = locator.get("run_id")
+            return str(run_id) if run_id else None
         case _:
             # Connector items and chat turns have no client-side renderer yet
             # (the frontend resolves only chunk ids and URLs), so they stay

surfsense_backend/app/agents/chat/multi_agent_chat/shared/citations/models.py

@@ -17,6 +17,7 @@ class CitationSourceType(StrEnum):
     WEB_RESULT = "web_result"
     CHAT_TURN = "chat_turn"
     ANON_CHUNK = "anon_chunk"
+    RUN = "run"
 
 
 class CitationEntry(BaseModel):

surfsense_backend/app/agents/chat/multi_agent_chat/subagents/middleware_stack.py

@@ -15,6 +15,9 @@ from __future__ import annotations
 from typing import Any
 
 from app.agents.chat.multi_agent_chat.shared.feature_flags import AgentFeatureFlags
+from app.agents.chat.multi_agent_chat.shared.middleware.citation_state import (
+    build_citation_state_mw,
+)
 from app.agents.chat.multi_agent_chat.shared.middleware.resilience import (
     ResilienceMiddlewares,
 )
@@ -45,6 +48,8 @@ def build_subagent_middleware_stack(
     return {
         "todos": build_todos_mw(),
         "permission": permission,
+        # Declares the citation_registry channel so run [n]s merge up from tools.
+        "citation": build_citation_state_mw(),
         "retry": resilience.retry,
         "fallback": resilience.fallback,
         "model_call_limit": resilience.model_call_limit,

surfsense_backend/app/agents/chat/multi_agent_chat/subagents/shared/snippets/output_contract_base.md

@@ -5,3 +5,4 @@ Rules (universal):
 - `status=blocked` due to missing required inputs -> `missing_fields` must be non-null.
 - `assumptions`: any inferences you made about the user's intent; `null` when no inferences were needed.
 - The `evidence` object's fields are documented in your route-specific `<output_contract>` above; never invent fields the tool did not return.
+- When a finding is drawn from a scraper run, append that run's `[n]` (the tool result states `Cite this scraper run as [n]`) to the finding text so the citation survives into the final answer. Copy the label exactly; never invent one.

surfsense_backend/app/capabilities/core/access/agent.py

@@ -11,9 +11,13 @@ subagent can follow a truncation reference without extra wiring.
 
 from __future__ import annotations
 
+import json
 import time
 
+from langchain.tools import ToolRuntime
+from langchain_core.messages import ToolMessage
 from langchain_core.tools import BaseTool, StructuredTool
+from langgraph.types import Command
 
 from app.capabilities.core.billing import charge_capability, gate_capability
 from app.capabilities.core.progress import progress_scope
@@ -64,7 +68,7 @@ def _capability_tool(capability: Capability, workspace_id: int) -> BaseTool:
     executor = capability.executor
     name = capability.name
 
-    async def _run(**kwargs: object) -> dict | str:
+    async def _run(runtime: ToolRuntime, **kwargs: object) -> dict | str | Command:
         payload = input_model(**kwargs)
         input_dump = payload.model_dump(exclude_none=True)
         thread_id = _current_thread_id()
@@ -119,13 +123,42 @@ def _capability_tool(capability: Capability, workspace_id: int) -> BaseTool:
                     progress=reporter.coarse,
                 )
 
+        # No stored run to cite: keep the legacy return shape, no citation.
+        if run_id is None:
+            if serialized.char_count <= RUN_OUTPUT_CHAR_CAP:
+                return output.model_dump(exclude_none=True)
+            return _build_preview(serialized, run_id)
+
+        run_external_id = f"run_{run_id}"
         if serialized.char_count <= RUN_OUTPUT_CHAR_CAP:
             dump = output.model_dump(exclude_none=True)
-            if run_id is not None:
-                dump["run_id"] = f"run_{run_id}"
-            return dump
+            dump["run_id"] = run_external_id
+            content = json.dumps(dump, ensure_ascii=False, default=str)
+        else:
+            content = _build_preview(serialized, run_id)
 
-        return _build_preview(serialized, run_id)
+        # Deferred import: the citation spine imports from here; lazy avoids a cycle.
+        from app.agents.chat.multi_agent_chat.shared.citations import load_registry
+        from app.capabilities.core.access.run_citation import attach_run_citation
+
+        registry = load_registry(getattr(runtime, "state", None))
+        _, label = attach_run_citation(
+            registry, run_external_id=run_external_id, capability=name
+        )
+        return Command(
+            update={
+                "messages": [
+                    ToolMessage(
+                        content=content + label,
+                        tool_call_id=runtime.tool_call_id,
+                    )
+                ],
+                "citation_registry": registry,
+            }
+        )
+
+    # Un-stringify for StructuredTool's signature-based runtime injection.
+    _run.__annotations__["runtime"] = ToolRuntime
 
     return StructuredTool.from_function(
         coroutine=_run,

surfsense_backend/app/capabilities/core/access/run_citation.py

@@ -0,0 +1,23 @@
+"""Register a recorded scraper run as a citable ``[n]``."""
+
+from __future__ import annotations
+
+from app.agents.chat.multi_agent_chat.shared.citations import (
+    CitationRegistry,
+    CitationSourceType,
+)
+
+
+def attach_run_citation(
+    registry: CitationRegistry,
+    *,
+    run_external_id: str,
+    capability: str,
+) -> tuple[int, str]:
+    """Register the ``run_<uuid>`` handle; return its ``[n]`` and the label line."""
+    n = registry.register(
+        CitationSourceType.RUN,
+        {"run_id": run_external_id},
+        {"capability": capability},
+    )
+    return n, f"\n\nCite this scraper run as [{n}] after any claim drawn from its data."

surfsense_backend/tests/unit/agents/multi_agent_chat/shared/citations/test_markers.py

@@ -37,6 +37,18 @@ def test_web_result_maps_to_url() -> None:
     assert to_frontend_payload(entry) == "https://example.com/a"
 
 
+def test_run_maps_to_run_handle() -> None:
+    entry = _entry(CitationSourceType.RUN, {"run_id": "run_abc-123"})
+
+    assert to_frontend_payload(entry) == "run_abc-123"
+
+
+def test_run_without_handle_is_dropped() -> None:
+    entry = _entry(CitationSourceType.RUN, {})
+
+    assert to_frontend_payload(entry) is None
+
+
 def test_not_yet_renderable_kind_is_dropped() -> None:
     entry = _entry(CitationSourceType.CHAT_TURN, {"thread_id": 1, "turn": 2})
 

surfsense_backend/tests/unit/capabilities/access/test_agent_tools.py

@@ -69,6 +69,12 @@ def _verb_tool(tools, name: str):
     return next(t for t in tools if t.name == name)
 
 
+def _invoke(tool, text: str, *, state=None):
+    """Call the coroutine with a stand-in runtime (ToolNode injects it in prod)."""
+    runtime = SimpleNamespace(state=state or {}, tool_call_id="tc_1", context=None)
+    return tool.coroutine(runtime, text=text)
+
+
 async def test_registry_becomes_one_tool_per_verb_plus_readers(isolate):
     caps = [
         _capability(name="web.scrape", output=_EchoOutput(echoed="a")),
@@ -104,20 +110,51 @@ async def test_tool_runs_executor_and_returns_serialized_output(isolate):
     tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
     tool = _verb_tool(tools, "web_scrape")
 
-    result = await tool.ainvoke({"text": "ping"})
+    result = await _invoke(tool, "ping")
 
     # Fake session makes record_run fail -> no run_id key, plain serialized output.
     assert result == {"echoed": "hi there"}
     assert cap.executor.seen.text == "ping"
 
 
+async def test_tool_registers_run_citation_when_stored(isolate, monkeypatch):
+    from langgraph.types import Command
+
+    cap = _capability(name="web.scrape", output=_EchoOutput(echoed="hi"))
+    monkeypatch.setattr(isolate.module, "record_run", AsyncMock(return_value="abc-123"))
+    tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
+    tool = _verb_tool(tools, "web_scrape")
+
+    result = await _invoke(tool, "ping")
+
+    assert isinstance(result, Command)
+    registry = result.update["citation_registry"]
+    entry = registry.resolve(1)
+    assert entry is not None
+    assert entry.locator["run_id"] == "run_abc-123"
+    message = result.update["messages"][0]
+    assert "[1]" in message.content
+    assert "run_abc-123" in message.content
+
+
+async def test_runtime_survives_langchain_arg_parsing(isolate):
+    """runtime must survive langchain arg parsing (else ToolNode drops it)."""
+    cap = _capability(name="web.scrape", output=_EchoOutput(echoed="hi"))
+    tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
+    tool = _verb_tool(tools, "web_scrape")
+
+    parsed = tool._parse_input({"text": "x", "runtime": "RT"}, "tc_1")
+
+    assert parsed["runtime"] == "RT"
+
+
 async def test_tool_charges_owner(isolate):
     output = _EchoOutput(echoed="hi")
     cap = _capability(name="web.scrape", output=output)
     tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
     tool = _verb_tool(tools, "web_scrape")
 
-    await tool.ainvoke({"text": "ping"})
+    await _invoke(tool, "ping")
 
     isolate.charge.assert_awaited_once()
     (charged_output, unit, ctx), _ = isolate.charge.call_args
@@ -136,7 +173,7 @@ async def test_over_budget_returns_friendly_message(isolate):
     tools = isolate.module.build_capability_tools(workspace_id=7, capabilities=[cap])
     tool = _verb_tool(tools, "web_scrape")
 
-    result = await tool.ainvoke({"text": "ping"})
+    result = await _invoke(tool, "ping")
 
     assert isinstance(result, str)
     assert "credit" in result.lower()

surfsense_backend/tests/unit/capabilities/access/test_run_citation.py

@@ -0,0 +1,43 @@
+"""Unit tests for registering a scraper run as a citation."""
+
+from __future__ import annotations
+
+import pytest
+
+from app.agents.chat.multi_agent_chat.shared.citations import (
+    CitationRegistry,
+    CitationSourceType,
+    to_frontend_payload,
+)
+from app.capabilities.core.access.run_citation import attach_run_citation
+
+pytestmark = pytest.mark.unit
+
+
+def test_attaches_run_and_returns_label_with_ordinal() -> None:
+    registry = CitationRegistry()
+
+    n, label = attach_run_citation(
+        registry, run_external_id="run_abc-123", capability="walmart.scrape"
+    )
+
+    assert n == 1
+    assert f"[{n}]" in label
+    entry = registry.resolve(n)
+    assert entry is not None
+    assert entry.source_type is CitationSourceType.RUN
+    assert to_frontend_payload(entry) == "run_abc-123"
+
+
+def test_same_run_dedups_to_one_label() -> None:
+    registry = CitationRegistry()
+
+    first, _ = attach_run_citation(
+        registry, run_external_id="run_x", capability="walmart.scrape"
+    )
+    again, _ = attach_run_citation(
+        registry, run_external_id="run_x", capability="walmart.reviews"
+    )
+
+    assert first == again
+    assert len(registry.by_n) == 1

surfsense_web/atoms/citation/citation-panel.atom.ts

@@ -1,14 +1,19 @@
-import { atom } from "jotai";
+import { atom, type Getter, type Setter } from "jotai";
 import { rightPanelCollapsedAtom, rightPanelTabAtom } from "@/atoms/layout/right-panel.atom";
 
+/** The source the citation panel is showing: a KB chunk or a scraper run. */
+export type CitationTarget =
+	| { kind: "chunk"; chunkId: number }
+	| { kind: "run"; runId: string };
+
 interface CitationPanelState {
 	isOpen: boolean;
-	chunkId: number | null;
+	target: CitationTarget | null;
 }
 
 const initialState: CitationPanelState = {
 	isOpen: false,
-	chunkId: null,
+	target: null,
 };
 
 export const citationPanelAtom = atom<CitationPanelState>(initialState);
@@ -17,16 +22,21 @@ export const citationPanelOpenAtom = atom((get) => get(citationPanelAtom).isOpen
 
 const preCitationCollapsedAtom = atom<boolean | null>(null);
 
-export const openCitationPanelAtom = atom(null, (get, set, payload: { chunkId: number }) => {
+function openWithTarget(get: Getter, set: Setter, target: CitationTarget) {
 	if (!get(citationPanelAtom).isOpen) {
 		set(preCitationCollapsedAtom, get(rightPanelCollapsedAtom));
 	}
-	set(citationPanelAtom, {
-		isOpen: true,
-		chunkId: payload.chunkId,
-	});
+	set(citationPanelAtom, { isOpen: true, target });
 	set(rightPanelTabAtom, "citation");
 	set(rightPanelCollapsedAtom, false);
+}
+
+export const openCitationPanelAtom = atom(null, (get, set, payload: { chunkId: number }) => {
+	openWithTarget(get, set, { kind: "chunk", chunkId: payload.chunkId });
+});
+
+export const openRunCitationPanelAtom = atom(null, (get, set, payload: { runId: string }) => {
+	openWithTarget(get, set, { kind: "run", runId: payload.runId });
 });
 
 export const closeCitationPanelAtom = atom(null, (get, set) => {

surfsense_web/components/citations/citation-renderer.tsx

@@ -2,6 +2,7 @@
 
 import type { ReactNode } from "react";
 import { InlineCitation, UrlCitation } from "@/components/assistant-ui/inline-citation";
+import { RunCitation } from "@/components/citations/run-citation";
 import {
 	type CitationToken,
 	type CitationUrlMap,
@@ -21,6 +22,9 @@ export function renderCitationToken(token: CitationToken, ordinalKey: number): R
 	if (token.kind === "url") {
 		return <UrlCitation key={`citation-url-${ordinalKey}`} url={token.url} />;
 	}
+	if (token.kind === "run") {
+		return <RunCitation key={`citation-run-${token.runId}-${ordinalKey}`} runId={token.runId} />;
+	}
 	return (
 		<InlineCitation
 			key={`citation-${token.isDocsChunk ? "doc-" : ""}${token.chunkId}-${ordinalKey}`}

surfsense_web/components/citations/run-citation-panel.tsx

@@ -0,0 +1,47 @@
+"use client";
+
+import { XIcon } from "lucide-react";
+import { useParams } from "next/navigation";
+import type { FC } from "react";
+import { RunDetail } from "@/app/dashboard/[workspace_id]/playground/components/run-detail";
+import { Button } from "@/components/ui/button";
+
+/** Right-panel viewer for a cited scraper run. `runId` is the `run_<uuid>` handle. */
+export const RunCitationPanelContent: FC<{ runId: string; onClose?: () => void }> = ({
+	runId,
+	onClose,
+}) => {
+	const params = useParams();
+	const rawWorkspaceId = Array.isArray(params?.workspace_id)
+		? params.workspace_id[0]
+		: params?.workspace_id;
+	const workspaceId = Number(rawWorkspaceId);
+	const scraperRunId = runId.replace(/^run_/, "");
+
+	return (
+		<>
+			<div className="shrink-0 flex h-12 items-center justify-between px-3 border-b">
+				<h2 className="select-none text-lg font-semibold">Scraper run</h2>
+				{onClose && (
+					<Button
+						variant="ghost"
+						size="icon"
+						onClick={onClose}
+						className="h-8 w-8 rounded-full shrink-0 text-muted-foreground hover:text-accent-foreground"
+					>
+						<XIcon className="h-4 w-4" />
+						<span className="sr-only">Close run panel</span>
+					</Button>
+				)}
+			</div>
+
+			<div className="flex-1 overflow-y-auto">
+				{Number.isFinite(workspaceId) ? (
+					<RunDetail workspaceId={workspaceId} runId={scraperRunId} />
+				) : (
+					<p className="p-4 text-sm text-muted-foreground">Open a workspace to view this run.</p>
+				)}
+			</div>
+		</>
+	);
+};

surfsense_web/components/citations/run-citation.tsx

@@ -0,0 +1,31 @@
+"use client";
+
+import { useSetAtom } from "jotai";
+import { Database } from "lucide-react";
+import type { FC } from "react";
+import { openRunCitationPanelAtom } from "@/atoms/citation/citation-panel.atom";
+import { Button } from "@/components/ui/button";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
+
+/** Inline citation badge for a scraper run; opens the run in the citation panel. */
+export const RunCitation: FC<{ runId: string }> = ({ runId }) => {
+	const openRunPanel = useSetAtom(openRunCitationPanelAtom);
+
+	return (
+		<Tooltip>
+			<TooltipTrigger asChild>
+				<Button
+					type="button"
+					variant="ghost"
+					onClick={() => openRunPanel({ runId })}
+					className="ml-0.5 inline-flex h-5 min-w-5 items-center justify-center gap-0.5 rounded-md bg-popover px-1.5 text-[11px] font-medium text-popover-foreground/80 align-baseline"
+					aria-label="See where this came from"
+				>
+					<Database className="size-3" />
+					Source
+				</Button>
+			</TooltipTrigger>
+			<TooltipContent>See where this came from</TooltipContent>
+		</Tooltip>
+	);
+};

surfsense_web/components/layout/ui/right-panel/RightPanel.tsx

@@ -35,6 +35,14 @@ const CitationPanelContent = dynamic(
 	{ ssr: false, loading: () => null }
 );
 
+const RunCitationPanelContent = dynamic(
+	() =>
+		import("@/components/citations/run-citation-panel").then((m) => ({
+			default: m.RunCitationPanelContent,
+		})),
+	{ ssr: false, loading: () => null }
+);
+
 const HitlEditPanelContent = dynamic(
 	() =>
 		import("@/features/chat-messages/hitl").then((m) => ({
@@ -89,7 +97,7 @@ export function RightPanelToggleButton({
 				? !!editorState.memoryScope
 				: !!editorState.localFilePath);
 	const hitlEditOpen = hitlEditState.isOpen && !!hitlEditState.onSave;
-	const citationOpen = citationState.isOpen && citationState.chunkId != null;
+	const citationOpen = citationState.isOpen && citationState.target != null;
 	const hasContent = reportOpen || editorOpen || hitlEditOpen || citationOpen || artifactsOpen;
 	const label = collapsed ? "Expand panel" : "Collapse panel";
 
@@ -141,7 +149,7 @@ export function RightPanelExpandButton() {
 				? !!editorState.memoryScope
 				: !!editorState.localFilePath);
 	const hitlEditOpen = hitlEditState.isOpen && !!hitlEditState.onSave;
-	const citationOpen = citationState.isOpen && citationState.chunkId != null;
+	const citationOpen = citationState.isOpen && citationState.target != null;
 	const hasContent = reportOpen || editorOpen || hitlEditOpen || citationOpen || artifactsOpen;
 
 	if (!collapsed || !hasContent) return null;
@@ -212,7 +220,7 @@ export function RightPanel({ showTopBorder = false }: RightPanelProps) {
 				? !!editorState.memoryScope
 				: !!editorState.localFilePath);
 	const hitlEditOpen = hitlEditState.isOpen && !!hitlEditState.onSave;
-	const citationOpen = citationState.isOpen && citationState.chunkId != null;
+	const citationOpen = citationState.isOpen && citationState.target != null;
 
 	useEffect(() => {
 		if (!reportOpen && !editorOpen && !hitlEditOpen && !citationOpen && !artifactsOpen) return;
@@ -306,9 +314,19 @@ export function RightPanel({ showTopBorder = false }: RightPanelProps) {
 									/>
 								</div>
 							)}
-							{effectiveTab === "citation" && citationOpen && citationState.chunkId != null && (
+							{effectiveTab === "citation" && citationOpen && citationState.target && (
 								<div className="h-full flex flex-col">
-									<CitationPanelContent chunkId={citationState.chunkId} onClose={closeCitation} />
+									{citationState.target.kind === "run" ? (
+										<RunCitationPanelContent
+											runId={citationState.target.runId}
+											onClose={closeCitation}
+										/>
+									) : (
+										<CitationPanelContent
+											chunkId={citationState.target.chunkId}
+											onClose={closeCitation}
+										/>
+									)}
 								</div>
 							)}
 							{effectiveTab === "artifacts" && artifactsOpen && (

surfsense_web/lib/citations/citation-parser.test.ts

@@ -0,0 +1,35 @@
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import { parseTextWithCitations } from "./citation-parser";
+
+// Run with: pnpm exec tsx --test lib/citations/citation-parser.test.ts
+
+const NO_URLS = new Map<string, string>();
+
+test("parses a scraper-run handle into a run token", () => {
+	const segments = parseTextWithCitations(
+		"Price is $19 [citation:run_1b2c3d4e-5f60-7081-9abc-def012345678].",
+		NO_URLS
+	);
+
+	assert.deepEqual(segments, [
+		"Price is $19 ",
+		{ kind: "run", runId: "run_1b2c3d4e-5f60-7081-9abc-def012345678" },
+		".",
+	]);
+});
+
+test("run handles do not collide with numeric chunk citations", () => {
+	const segments = parseTextWithCitations(
+		"chunk [citation:42] vs run [citation:run_ab-12].",
+		NO_URLS
+	);
+
+	assert.deepEqual(segments, [
+		"chunk ",
+		{ kind: "chunk", chunkId: 42, isDocsChunk: false },
+		" vs run ",
+		{ kind: "run", runId: "run_ab-12" },
+		".",
+	]);
+});

surfsense_web/lib/citations/citation-parser.ts

@@ -11,18 +11,20 @@ import { FENCED_OR_INLINE_CODE } from "@/lib/markdown/code-regions";
 
 /**
  * Matches `[citation:...]` with numeric IDs (incl. negative, doc- prefix,
- * comma-separated), URL-based IDs from live web search, or `urlciteN`
- * placeholders produced by `preprocessCitationMarkdown`.
+ * comma-separated), URL-based IDs from live web search, `urlciteN`
+ * placeholders produced by `preprocessCitationMarkdown`, or a scraper-run
+ * handle (`run_<uuid>`).
  *
  * Also matches Chinese brackets 【】 and zero-width spaces that LLMs
  * sometimes emit.
  */
 export const CITATION_REGEX =
-	/[[【]\u200B?citation:\s*(https?:\/\/[^\]】\u200B]+|urlcite\d+|(?:doc-)?-?\d+(?:\s*,\s*(?:doc-)?-?\d+)*)\s*\u200B?[\]】]/g;
+	/[[【]\u200B?citation:\s*(https?:\/\/[^\]】\u200B]+|urlcite\d+|run_[0-9a-fA-F-]+|(?:doc-)?-?\d+(?:\s*,\s*(?:doc-)?-?\d+)*)\s*\u200B?[\]】]/g;
 
 /** A single parsed citation reference. */
 export type CitationToken =
 	| { kind: "url"; url: string }
+	| { kind: "run"; runId: string }
 	| { kind: "chunk"; chunkId: number; isDocsChunk: boolean };
 
 /** Output of `parseTextWithCitations` — interleaved text + citation tokens. */
@@ -102,6 +104,8 @@ export function parseTextWithCitations(text: string, urlMap: CitationUrlMap): Pa
 			if (url) {
 				segments.push({ kind: "url", url });
 			}
+		} else if (captured.startsWith("run_")) {
+			segments.push({ kind: "run", runId: captured.trim() });
 		} else {
 			const rawIds = captured.split(",").map((s) => s.trim());
 			for (const rawId of rawIds) {