ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: surface Anthropic document citations inline + in Sources panel (#5718)

Browse source 
* Studio: surface Anthropic document citations inline + in Sources panel

Anthropic's Messages API streams ``citations_delta`` events on
``content_block_delta`` when the request enables
``citations: {enabled: true}`` on document blocks. Each event carries
one citation pointing at the source document; previously they were
silently dropped, so reader-visible references never reached the chat
UI even when the model was citing properly.

The proxy now:
- dedupes by the type-specific anchor (char_location / page_location /
  content_block_location / search_result_location) so re-cites of the
  same span collapse onto a single footnote;
- injects ``[N]`` inline right after the matching text run;
- forwards the full list as a synthetic ``document_citations``
  tool_event at ``message_stop`` so the Sources panel can render
  per-document footnotes next to web_search / web_fetch citations.

Streams that never emit ``citations_delta`` stay byte-identical.

References:
- https://platform.claude.com/docs/en/build-with-claude/citations
- https://platform.claude.com/docs/en/build-with-claude/search-results

Tests (5 in test_anthropic_citations.py): passthrough, single
char_location, dedup of repeat citations, distinct sources get
distinct numbers, search_result_location supported.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: surface Anthropic document_citations in the Sources panel

The PR added a backend _toolEvent.type='document_citations' on
message_stop and an inline [N] marker in the assistant text, but the
chat-adapter only handles container_*/tool_*/sources from
web_search and web_fetch tool calls. Reviewers flagged that the
inline [N] markers had no matching footnote entries in the Sources
panel.

Capture the new event into a documentCitationParts buffer, convert
each citation dict into a Sources-panel source entry (using
document_title or search-result source URL plus cited_text as the
snippet), dedupe by id, and append to the final yield alongside
the existing web_search/web_fetch sourceParts.

* Studio: dedupe search_result_location citations by search_result_index

Anthropic's documented search_result_location citation shape carries
search_result_index, source, title, and start/end_block_index --
NOT document_index/document_title. The previous key keyed on
document_index + document_title + source + start_block_index, so
two distinct search results from the same source collapsed onto the
same footnote and the second [N] marker was lost.

Switch the search_result_location branch to key on the documented
fields, and pin the behaviour with a regression test asserting that
two citations sharing source/title but with different
search_result_index get distinct [1] [2] markers.

* Studio: keep each citation distinct across the end-anchor

Codex follow-ups on the citations PR:

  * Backend _anthropic_citation_key now includes the end anchor for
    every variant (end_char_index, end_page_number,
    end_block_index). Anthropic ranges are start-AND-end pairs, so
    a same-start / different-end pair is two distinct citations
    that previously collapsed onto one footnote.

  * Frontend documentCitationToSource ids include the position
    fields (search_result_index, start/end char/page/block) instead
    of being keyed on URL alone. Two citations from the same
    document or two search_result_locations with the same source
    now produce distinct Sources-panel entries, matching the
    inline [N] numbering.

* Studio: key Sources list by per-citation id instead of url

Codex flagged that the Sources renderer keys badges on source.url,
so two Anthropic document citations sharing the same source URL
collide as React keys and one badge gets dropped (or duplicated).

The chat-adapter already mints a per-citation id that folds the
position fields (search_result_index, start/end char/page/block)
into the URL, so the two citations have distinct ids even when
their URL matches. Plumb that id through SourceData and use it as
the React key for both the measurement badges and the visible
SourceBadge list. Falls back to the URL when no id is supplied
(web_search and web_fetch source parts).

* Studio: enable Anthropic doc citations on input_document blocks

Plumb citations: {enabled: true} onto the translated Anthropic document
block (both base64 and URL source branches) so the upstream actually
emits citations_delta events. Without this opt-in the inline [N] +
Sources panel plumbing added in this PR is a no-op for real user
PDF / doc uploads.

Refs https://platform.claude.com/docs/en/build-with-claude/citations

Also add edge-case coverage for the citations_delta path:
malformed citations, mixed types per document, reversed indices,
missing document_index, non-int block indices, unknown citation
type, internal _key never leaking, footnote numbering across
content blocks, and the input_document wire-through itself.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Reject unsafe citation sources, bound cited_text payload

Three follow-ups on top of #5718 surfaced by a deeper review pass:

1) javascript: / data: / vbscript: in citation source is XSS-able.
   ``documentCitationToSource`` was assigning ``cit.source`` straight
   into ``Source.url`` and rendering it as an <a href>. A hostile
   model emitting ``cit.source = "javascript:alert(document.domain)"``
   would execute on click (openLink only intercepts URLs that contain
   "://" or start with "mailto:", which both miss the javascript:
   scheme). Restrict the navigable path to http(s):// only; anything
   else falls back to the existing #anthropic-doc anchor and the
   source title still renders the raw identifier for context. Also
   reject CR/LF inside the URL string.

2) Frontend sources collapse distinct backend footnotes when the
   citation type differs but positions match. char_location(0,5) and
   page_location(0,5) over the same source previously deduped into
   one entry because the id only carried position. Fold citation
   type into the id anchor so the 1:1 mapping with inline [N]
   markers is preserved across every citation shape.

3) ``cited_text`` was forwarded unbounded inside the synthetic
   document_citations tool_event. The Sources panel trims to 240
   chars for display anyway; for large RAG / search_result spans
   (~10kB cited_text is plausible) this inflates SSE bytes 40x
   for no UI benefit. Truncate server-side at 512 chars with an
   ellipsis so the description-trim downstream still has room to
   work and the wire stays bounded.

Tests grow from 21 to 22; existing 7 + edge 15 still green. Frontend
typecheck clean.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: apply http(s) URL guard to all Sources-panel link sources

The previous round only filtered ``cit.source`` inside
``documentCitationToSource``. Two parallel code paths still copied
provider/tool-controlled ``URL:`` text directly into clickable
``<a href>`` Sources-panel links:

  * ``parseSourcesFromResult`` in chat-adapter.ts (legacy web_search /
    web_fetch tool result parser)
  * ``parseSearchResults`` in tool-ui-web-search.tsx (inline tool card)

A hostile tool response like ``URL: javascript:alert(1)`` or
``URL: data:text/html,...`` was therefore still rendered as a
navigable badge in the Sources panel.

Centralise the safe-URL test (``isSafeNavigableSourceUrl``,
``isSafeHttpUrl``) using ``new URL()`` + protocol allowlist + CR/LF
rejection, and apply it to both parsers. Unsafe blocks are dropped
rather than rewritten to a hash anchor because the web_search /
web_fetch parsers have no document-index fallback.

Citation conversion now uses the same helper so the in-place
http(s) regex and CR/LF check stay in one place.

* Shorten citation comments for PR #5718

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-05-25 23:39:02 -07:00 • committed by GitHub
commit 4854d4579f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 1326 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

119
studio/backend/core/inference/external_provider.py
View file

@ -325,6 +325,65 @@ def _anthropic_supports_fast_mode(model: str) -> bool:
)
# Cap on ``cited_text`` forwarded in document_citations tool_events;
# keeps SSE bytes bounded on multi-KB cited spans (frontend trims to
# 240 chars anyway).
_CITED_TEXT_MAX_LEN = 512
def _anthropic_citation_key(citation: dict[str, Any]) -> tuple:
"""Stable dedup key for an Anthropic ``citations_delta.citation``.
Anchor fields vary per type (char_location, page_location,
content_block_location, search_result_location); both start AND
exclusive end indices are part of the key so same-start /
different-end pairs stay distinct. search_result_location keys on
``search_result_index`` + ``source`` instead of document_index so
distinct results with the same source don't collapse. Unknown
shapes fall back to a stringified copy (more entries, never
collisions). See
https://platform.claude.com/docs/en/build-with-claude/citations
and https://platform.claude.com/docs/en/build-with-claude/search-results.
"""
ctype = citation.get("type")
doc = citation.get("document_index")
title = citation.get("document_title") or ""
if ctype == "char_location":
return (
ctype,
doc,
title,
citation.get("start_char_index"),
citation.get("end_char_index"),
)
if ctype == "page_location":
return (
ctype,
doc,
title,
citation.get("start_page_number"),
citation.get("end_page_number"),
)
if ctype == "content_block_location":
return (
ctype,
doc,
title,
citation.get("start_block_index"),
citation.get("end_block_index"),
)
if ctype == "search_result_location":
return (
ctype,
citation.get("search_result_index"),
citation.get("source"),
citation.get("title") or "",
citation.get("start_block_index"),
citation.get("end_block_index"),
)
return (ctype, _json.dumps(citation, sort_keys = True))
class _MistralThinkingSpec(NamedTuple):
models: tuple[str, ...]
style: Literal["prompt_mode", "reasoning_effort", "disabled"]
@ -1460,6 +1519,11 @@ class ExternalProviderClient:
"media_type": media_type,
"data": b64data,
},
# Opt into Anthropic's natural-citation
# pipeline; without this no citations_delta
# events fire. See
# https://platform.claude.com/docs/en/build-with-claude/citations
"citations": {"enabled": True},
}
if title:
doc_block["title"] = title
@ -1471,6 +1535,7 @@ class ExternalProviderClient:
"type": "url",
"url": url,
},
"citations": {"enabled": True},
}
if title:
doc_block["title"] = title
@ -1925,6 +1990,12 @@ class ExternalProviderClient:
# the next turn.
current_compaction: Optional[dict[str, Any]] = None
compaction_blocks_seen = 0
# Document citations from ``citations_delta`` events.
# Deduped by type-specific anchor key; inline [N] is
# injected after each cited run, and the full list is
# forwarded as a synthetic document_citations tool_event
# on message_stop for the Sources panel.
document_citations: list[dict[str, Any]] = []
# Counts surfaced in the final log line so reports of
# "Code execution did nothing" can be triaged at a
# glance. generated_files_count is interesting for the
@ -2276,10 +2347,27 @@ class ExternalProviderClient:
thinking_open = False
if text:
yield _content_chunk(text)
# Citations on text deltas are attached
# per-call by Anthropic via the
# `web_search_tool_result` block; we don't
# need to scrape them off the text events.
# web_search citations: web_search_tool_result.
# User-doc citations: citations_delta below.
elif delta_type == "citations_delta":
# One citation per event; collapse onto a
# numbered footnote list and inject [N]
# inline. See
# https://platform.claude.com/docs/en/build-with-claude/citations
cit = delta.get("citation")
if isinstance(cit, dict):
key = _anthropic_citation_key(cit)
idx_for_marker: Optional[int] = None
for idx, existing in enumerate(
document_citations, start = 1
):
if existing.get("_key") == key:
idx_for_marker = idx
break
if idx_for_marker is None:
document_citations.append({**cit, "_key": key})
idx_for_marker = len(document_citations)
yield _content_chunk(f"[{idx_for_marker}]")
elif delta_type == "input_json_delta":
# Streamed partial_json carrying tool inputs
# — the search query for web_search, or the
@ -2609,6 +2697,29 @@ class ExternalProviderClient:
if thinking_open:
yield _content_chunk("</think>")
thinking_open = False
# Forward document_citations so the Sources
# panel can render the inline [N] footnotes.
# ``cited_text`` is truncated server-side to
# keep SSE bytes bounded on long spans.
if document_citations:
clean_cits = []
for c in document_citations:
entry = {k: v for k, v in c.items() if k != "_key"}
cited = entry.get("cited_text")
if (
isinstance(cited, str)
and len(cited) > _CITED_TEXT_MAX_LEN
):
entry["cited_text"] = (
cited[:_CITED_TEXT_MAX_LEN] + "…"
)
clean_cits.append(entry)
yield _emit_tool_event(
{
"type": "document_citations",
"citations": clean_cits,
}
)
# Final include_usage-style chunk so callers can
# see cache_creation / cache_read without
# scraping the server log.

353
studio/backend/tests/test_anthropic_citations.py Normal file
View file

@ -0,0 +1,353 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for Anthropic ``citations_delta`` handling in the streaming proxy.
Verifies the proxy injects inline ``[N]`` markers after cited text,
dedupes by type-specific anchor (char_location, page_location,
content_block_location, search_result_location), forwards a synthetic
``document_citations`` tool_event at message_stop, and stays inert when
no citations_delta events fire. See
https://platform.claude.com/docs/en/build-with-claude/citations
"""
import asyncio
import json
import httpx
from core.inference import external_provider as ep_mod
from core.inference.external_provider import ExternalProviderClient
def _drive(coro):
return asyncio.new_event_loop().run_until_complete(coro)
def _make_client() -> ExternalProviderClient:
return ExternalProviderClient(
provider_type = "anthropic",
base_url = "https://api.anthropic.com/v1",
api_key = "sk-ant-test",
)
def _sse(events: list[dict]) -> bytes:
out = []
for e in events:
ev = e.get("type", "message")
out.append(f"event: {ev}\ndata: {json.dumps(e)}\n\n")
return "".join(out).encode("utf-8")
def _capture(monkeypatch, events: list[dict]) -> list[str]:
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
content = _sse(events),
headers = {"content-type": "text/event-stream"},
)
monkeypatch.setattr(
ep_mod,
"_http_client",
httpx.AsyncClient(transport = httpx.MockTransport(handler)),
)
lines: list[str] = []
async def run():
client = _make_client()
try:
async for line in client.stream_chat_completion(
messages = [{"role": "user", "content": "what color is grass?"}],
model = "claude-opus-4-7",
max_tokens = 64,
):
lines.append(line)
finally:
await client.close()
_drive(run())
return lines
def _message_start() -> dict:
return {
"type": "message_start",
"message": {
"id": "m1",
"content": [],
"model": "claude-opus-4-7",
"role": "assistant",
"stop_reason": None,
"usage": {"input_tokens": 5, "output_tokens": 2},
},
}
def _content_block_start_text() -> dict:
return {
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
}
def _text_delta(text: str, index: int = 0) -> dict:
return {
"type": "content_block_delta",
"index": index,
"delta": {"type": "text_delta", "text": text},
}
def _citations_delta(citation: dict, index: int = 0) -> dict:
return {
"type": "content_block_delta",
"index": index,
"delta": {"type": "citations_delta", "citation": citation},
}
def _content_block_stop(index: int = 0) -> dict:
return {"type": "content_block_stop", "index": index}
def _message_delta_end() -> dict:
return {"type": "message_delta", "delta": {"stop_reason": "end_turn"}}
def _message_stop() -> dict:
return {"type": "message_stop"}
def _joined(lines: list[str]) -> str:
return "\n".join(lines)
def test_no_citations_stream_unchanged(monkeypatch):
"""Plain text streams pass through with no inline markers and no
document_citations tool_event."""
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Grass is green."),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "Grass is green." in body
assert "document_citations" not in body
assert "[1]" not in body
def test_single_char_location_emits_inline_marker(monkeypatch):
cit = {
"type": "char_location",
"cited_text": "The grass is green.",
"document_index": 0,
"document_title": "Example",
"start_char_index": 0,
"end_char_index": 20,
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Grass is green."),
_citations_delta(cit),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "Grass is green." in body
assert "[1]" in body, body
assert "document_citations" in body, body
assert '"document_index": 0' in body, body
assert "_key" not in body, body
def test_duplicate_citation_dedupes_to_same_number(monkeypatch):
cit = {
"type": "char_location",
"document_index": 0,
"document_title": "Example",
"start_char_index": 0,
"end_char_index": 20,
"cited_text": "The grass is green.",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Grass."),
_citations_delta(cit),
_text_delta(" Still green."),
_citations_delta(cit),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert body.count("[1]") == 2, body
citation_blob = body[body.index("document_citations") :]
assert citation_blob.count('"start_char_index"') == 1, citation_blob
def test_distinct_sources_get_distinct_numbers(monkeypatch):
cit1 = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc A",
"start_char_index": 0,
"end_char_index": 5,
}
cit2 = {
"type": "page_location",
"document_index": 1,
"document_title": "Doc B",
"start_page_number": 3,
"end_page_number": 4,
}
cit3 = {
"type": "content_block_location",
"document_index": 2,
"document_title": "Doc C",
"start_block_index": 0,
"end_block_index": 1,
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("First"),
_citations_delta(cit1),
_text_delta(" Second"),
_citations_delta(cit2),
_text_delta(" Third"),
_citations_delta(cit3),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body and "[2]" in body and "[3]" in body, body
assert body.index("[1]") < body.index("[2]") < body.index("[3]")
def test_search_result_location_supported(monkeypatch):
cit = {
"type": "search_result_location",
"document_index": 0,
"document_title": "Anthropic Search Results",
"source": "https://example.com/doc.html",
"start_block_index": 0,
"end_block_index": 1,
"cited_text": "blah",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Some sourced fact."),
_citations_delta(cit),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body
assert "search_result_location" in body
assert "example.com/doc.html" in body
def test_same_start_different_end_offsets_get_distinct_numbers(monkeypatch):
"""Same start_char_index + different end_char_index = distinct spans,
so they must get distinct footnote numbers (ranges use exclusive end)."""
cit_a = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc",
"start_char_index": 100,
"end_char_index": 150,
"cited_text": "first half",
}
cit_b = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc",
"start_char_index": 100,
"end_char_index": 250,
"cited_text": "wider span",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("A "),
_citations_delta(cit_a),
_text_delta(" and B "),
_citations_delta(cit_b),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body, body
assert "[2]" in body, body
def test_search_result_location_different_indices_get_distinct_numbers(monkeypatch):
"""Same source + different search_result_index = distinct footnotes
(matches the Anthropic search-result citation contract)."""
cit_a = {
"type": "search_result_location",
"search_result_index": 0,
"source": "https://example.com/result.html",
"title": "Result",
"start_block_index": 0,
"end_block_index": 1,
"cited_text": "first",
}
cit_b = {
"type": "search_result_location",
"search_result_index": 1,
"source": "https://example.com/result.html",
"title": "Result",
"start_block_index": 0,
"end_block_index": 1,
"cited_text": "second",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("A "),
_citations_delta(cit_a),
_text_delta(" and B "),
_citations_delta(cit_b),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body, body
assert "[2]" in body, body

690
studio/backend/tests/test_anthropic_citations_edge.py Normal file
View file

@ -0,0 +1,690 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Edge-case tests for Anthropic ``citations_delta`` handling.
Complements ``test_anthropic_citations.py``. Covers malformed payloads,
unusual orderings, mixed citation types, and the ``citations:
{enabled: true}`` opt-in attached to translated ``input_document``
blocks. See
https://platform.claude.com/docs/en/build-with-claude/citations and
https://platform.claude.com/docs/en/build-with-claude/search-results.
"""
import asyncio
import json
import httpx
from core.inference import external_provider as ep_mod
from core.inference.external_provider import ExternalProviderClient
# ── shared SSE harness ───────────────────────────────────────
def _drive(coro):
return asyncio.new_event_loop().run_until_complete(coro)
def _make_client() -> ExternalProviderClient:
return ExternalProviderClient(
provider_type = "anthropic",
base_url = "https://api.anthropic.com/v1",
api_key = "sk-ant-test",
)
def _sse(events: list[dict]) -> bytes:
out = []
for e in events:
ev = e.get("type", "message")
out.append(f"event: {ev}\ndata: {json.dumps(e)}\n\n")
return "".join(out).encode("utf-8")
def _capture(
monkeypatch,
events: list[dict],
*,
messages: list[dict] | None = None,
captured_body: dict | None = None,
) -> list[str]:
"""Drive ``stream_chat_completion`` against a mocked Anthropic
response and return the SSE lines. Pass ``captured_body`` to also
capture the outgoing request body for assertions on the translated
Anthropic shape.
"""
def handler(request: httpx.Request) -> httpx.Response:
if captured_body is not None:
try:
captured_body.update(json.loads(request.content.decode("utf-8")))
except Exception: # pragma: no cover -- diagnostic only
pass
return httpx.Response(
200,
content = _sse(events),
headers = {"content-type": "text/event-stream"},
)
monkeypatch.setattr(
ep_mod,
"_http_client",
httpx.AsyncClient(transport = httpx.MockTransport(handler)),
)
lines: list[str] = []
async def run():
client = _make_client()
try:
async for line in client.stream_chat_completion(
messages = messages
or [{"role": "user", "content": "what color is grass?"}],
model = "claude-opus-4-7",
max_tokens = 64,
):
lines.append(line)
finally:
await client.close()
_drive(run())
return lines
def _message_start() -> dict:
return {
"type": "message_start",
"message": {
"id": "m1",
"content": [],
"model": "claude-opus-4-7",
"role": "assistant",
"stop_reason": None,
"usage": {"input_tokens": 5, "output_tokens": 2},
},
}
def _content_block_start_text() -> dict:
return {
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
}
def _text_delta(text: str, index: int = 0) -> dict:
return {
"type": "content_block_delta",
"index": index,
"delta": {"type": "text_delta", "text": text},
}
def _citations_delta(citation: dict, index: int = 0) -> dict:
return {
"type": "content_block_delta",
"index": index,
"delta": {"type": "citations_delta", "citation": citation},
}
def _content_block_stop(index: int = 0) -> dict:
return {"type": "content_block_stop", "index": index}
def _message_delta_end() -> dict:
return {"type": "message_delta", "delta": {"stop_reason": "end_turn"}}
def _message_stop() -> dict:
return {"type": "message_stop"}
def _joined(lines: list[str]) -> str:
return "\n".join(lines)
def _citation_payload(body: str) -> dict:
"""Pull the ``document_citations`` synthetic tool_event from the
SSE body and return its payload. Raises if absent."""
assert "document_citations" in body, body
for line in body.splitlines():
if not line.startswith("data: "):
continue
try:
payload = json.loads(line[len("data: ") :])
except json.JSONDecodeError:
continue
tool_event = payload.get("_toolEvent") if isinstance(payload, dict) else None
if (
isinstance(tool_event, dict)
and tool_event.get("type") == "document_citations"
):
return tool_event
raise AssertionError("document_citations event not parsed out of SSE body")
# ── edge cases ───────────────────────────────────────────────
def test_citation_with_no_preceding_text_still_emits_marker(monkeypatch):
"""citations_delta before any text_delta must not crash; marker
lands at the start of the block."""
cit = {
"type": "char_location",
"document_index": 0,
"document_title": "X",
"start_char_index": 0,
"end_char_index": 5,
"cited_text": "x",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_citations_delta(cit),
_text_delta("hello"),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body, body
assert "document_citations" in body, body
def test_citations_delta_with_non_dict_citation_is_ignored(monkeypatch):
"""Non-dict ``delta.citation`` must not crash, emit a marker, or
poison the document_citations list."""
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Hello."),
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "citations_delta", "citation": "not-a-dict"},
},
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "Hello." in body
assert "[1]" not in body
assert "document_citations" not in body
def test_citations_delta_with_missing_citation_field_is_ignored(monkeypatch):
"""Missing ``citation`` field is treated like a non-dict citation:
skip without crashing."""
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Hello."),
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "citations_delta"},
},
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "Hello." in body
assert "[1]" not in body
assert "document_citations" not in body
def test_char_location_with_reversed_indices_does_not_crash(monkeypatch):
"""Malformed char_location with reversed indices must not crash;
the dedup key accepts any int pair and still surfaces a footnote."""
cit = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc",
"start_char_index": 300,
"end_char_index": 50,
"cited_text": "?",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Weird."),
_citations_delta(cit),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body, body
payload = _citation_payload(body)
assert payload["citations"][0]["start_char_index"] == 300
assert payload["citations"][0]["end_char_index"] == 50
def test_page_location_missing_document_index_does_not_crash(monkeypatch):
"""page_location missing ``document_index`` still produces a
footnote; dedup key falls back to ``None`` for the missing field."""
cit = {
"type": "page_location",
"document_title": "Untitled PDF",
"start_page_number": 1,
"end_page_number": 2,
"cited_text": "p1",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("From the PDF:"),
_citations_delta(cit),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body, body
payload = _citation_payload(body)
assert payload["citations"][0].get("document_index") is None
def test_content_block_location_with_non_int_block_index_does_not_crash(monkeypatch):
"""content_block_location with string block indices must not crash;
dedup key tolerates non-int values."""
cit = {
"type": "content_block_location",
"document_index": 0,
"document_title": "Custom",
"start_block_index": "0",
"end_block_index": "1",
"cited_text": "anything",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Cite."),
_citations_delta(cit),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body, body
payload = _citation_payload(body)
assert payload["citations"][0]["start_block_index"] == "0"
def test_unknown_citation_type_falls_back_to_stringified_key(monkeypatch):
"""Unknown citation ``type`` (forward-compat) still dedupes:
identical ones collapse, differing ones get distinct numbers."""
cit_a = {
"type": "future_shape_location",
"anchor": "abc",
"cited_text": "blah",
}
cit_b = {
"type": "future_shape_location",
"anchor": "xyz",
"cited_text": "blah",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("A"),
_citations_delta(cit_a),
_text_delta(" again"),
_citations_delta(cit_a),
_text_delta(" B"),
_citations_delta(cit_b),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
# cit_a dedupes onto [1], cit_b gets [2].
assert body.count("[1]") == 2, body
assert body.count("[2]") == 1, body
payload = _citation_payload(body)
assert len(payload["citations"]) == 2
def test_mixed_citation_types_same_document_get_distinct_keys(monkeypatch):
"""char_location and page_location on the same document_index are
distinct shapes; dedup key uses citation type as its first slot."""
cit_char = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc",
"start_char_index": 0,
"end_char_index": 10,
}
cit_page = {
"type": "page_location",
"document_index": 0,
"document_title": "Doc",
"start_page_number": 1,
"end_page_number": 2,
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("char-cite"),
_citations_delta(cit_char),
_text_delta(" page-cite"),
_citations_delta(cit_page),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body and "[2]" in body, body
payload = _citation_payload(body)
assert len(payload["citations"]) == 2
def test_cited_text_is_preserved_in_synthetic_event(monkeypatch):
"""``cited_text`` must survive into the synthetic event so the
Sources panel can render it as a tooltip. Anthropic does not bill
cited_text against output tokens, so preserving it is free."""
cit = {
"type": "char_location",
"document_index": 0,
"document_title": "Trustworthy Doc",
"start_char_index": 0,
"end_char_index": 20,
"cited_text": "The grass is green.",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Grass is green."),
_citations_delta(cit),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
payload = _citation_payload(body)
assert payload["citations"][0]["cited_text"] == "The grass is green."
def test_internal_key_field_never_leaks_to_client(monkeypatch):
"""The internal ``_key`` dedup sentinel must be stripped before
the synthetic event is forwarded; it is not an Anthropic field."""
cit = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc",
"start_char_index": 0,
"end_char_index": 5,
"cited_text": "..",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("hi"),
_citations_delta(cit),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
payload = _citation_payload(body)
assert payload["citations"], payload
for c in payload["citations"]:
assert "_key" not in c, c
def test_citation_across_multiple_content_blocks_numbers_continue(monkeypatch):
"""Footnote numbering is per-message, not per-content-block:
citations across separate blocks emit [1] then [2]."""
cit_a = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc",
"start_char_index": 0,
"end_char_index": 5,
}
cit_b = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc",
"start_char_index": 100,
"end_char_index": 105,
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("first"),
_citations_delta(cit_a, index = 0),
_content_block_stop(0),
{
"type": "content_block_start",
"index": 1,
"content_block": {"type": "text", "text": ""},
},
_text_delta(" second", index = 1),
_citations_delta(cit_b, index = 1),
_content_block_stop(1),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "[1]" in body and "[2]" in body, body
assert body.index("[1]") < body.index("[2]")
payload = _citation_payload(body)
assert len(payload["citations"]) == 2
def test_inline_marker_lands_after_text_run(monkeypatch):
"""Inline ``[N]`` must land AFTER the cited text run: Anthropic
streams text then citation, so the proxy emits ``"...green.[1]"``
not ``"[1]green"``."""
cit = {
"type": "char_location",
"document_index": 0,
"document_title": "Doc",
"start_char_index": 0,
"end_char_index": 20,
"cited_text": "grass",
}
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Grass is green."),
_citations_delta(cit),
_text_delta(" Sky is blue."),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
grass = body.index("Grass is green.")
marker = body.index("[1]")
sky = body.index("Sky is blue.")
assert grass < marker < sky, body
def test_no_synthetic_event_when_only_text_deltas(monkeypatch):
"""No citations_delta means no synthetic ``document_citations``
event; Sources panel relies on absence to suppress the section."""
lines = _capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("Just some prose. "),
_text_delta("More prose."),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
)
body = _joined(lines)
assert "document_citations" not in body
assert "[1]" not in body
def test_input_document_translation_enables_citations(monkeypatch):
"""``input_document`` must translate to an Anthropic ``document``
block carrying ``citations: {enabled: true}`` (both base64 and url
source branches) so upstream emits citations_delta."""
captured_b64: dict = {}
_capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("ok"),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
messages = [
{
"role": "user",
"content": [
{
"type": "input_document",
"file_data": "data:application/pdf;base64,QUJD",
"filename": "spec.pdf",
},
{"type": "text", "text": "summarise"},
],
}
],
captured_body = captured_b64,
)
user_msg = captured_b64["messages"][0]
doc_block = next(p for p in user_msg["content"] if p.get("type") == "document")
assert doc_block["source"]["type"] == "base64", doc_block
assert doc_block.get("citations") == {"enabled": True}, doc_block
captured_url: dict = {}
_capture(
monkeypatch,
[
_message_start(),
_content_block_start_text(),
_text_delta("ok"),
_content_block_stop(),
_message_delta_end(),
_message_stop(),
],
messages = [
{
"role": "user",
"content": [
{
"type": "input_document",
"file_url": "https://example.com/doc.pdf",
"filename": "doc.pdf",
},
{"type": "text", "text": "summarise"},
],
}
],
captured_body = captured_url,
)
user_msg = captured_url["messages"][0]
doc_block = next(p for p in user_msg["content"] if p.get("type") == "document")
assert doc_block["source"]["type"] == "url", doc_block
assert doc_block.get("citations") == {"enabled": True}, doc_block
# ── cited_text truncation + safe-url citation conversion ────────
def test_cited_text_truncated_in_synthetic_event(monkeypatch):
"""``cited_text`` is capped server-side so multi-KB spans do not
balloon the SSE payload."""
from core.inference.external_provider import _CITED_TEXT_MAX_LEN
long_quote = "x" * (_CITED_TEXT_MAX_LEN + 4000)
events = [
{
"type": "message_start",
"message": {
"id": "msg_1",
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "claim "},
},
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "citations_delta",
"citation": {
"type": "char_location",
"document_index": 0,
"document_title": "doc",
"start_char_index": 0,
"end_char_index": 5,
"cited_text": long_quote,
},
},
},
{"type": "content_block_stop", "index": 0},
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 1},
},
{"type": "message_stop"},
]
chunks = _capture(monkeypatch, events)
tool_events = [c for c in chunks if "_toolEvent" in c and "document_citations" in c]
assert tool_events, "no document_citations tool event"
payload = json.loads(tool_events[0].split("data: ", 1)[1])
cited = payload["_toolEvent"]["citations"][0]["cited_text"]
assert len(cited) <= _CITED_TEXT_MAX_LEN + 1, len(cited)
assert cited.endswith("…")

6
studio/backend/tests/test_multimodal_document.py
View file

@ -117,6 +117,8 @@ def test_anthropic_base64_pdf_becomes_document_block(monkeypatch):
types = [p.get("type") for p in parts]
assert "document" in types, parts
doc = _strip_cache(next(p for p in parts if p.get("type") == "document"))
# citations: {enabled: true} opts into Anthropic's natural-citation
# pipeline; without it the citations_delta handler is a no-op.
assert doc == {
"type": "document",
"source": {
@ -124,6 +126,7 @@ def test_anthropic_base64_pdf_becomes_document_block(monkeypatch):
"media_type": "application/pdf",
"data": _TINY_PDF_B64,
},
"citations": {"enabled": True},
"title": "paper.pdf",
}
@ -151,6 +154,7 @@ def test_anthropic_url_pdf_becomes_document_block(monkeypatch):
assert doc == {
"type": "document",
"source": {"type": "url", "url": "https://example.com/doc.pdf"},
"citations": {"enabled": True},
}
@ -255,6 +259,7 @@ def test_anthropic_empty_data_uri_falls_back_to_file_url(monkeypatch):
assert doc == {
"type": "document",
"source": {"type": "url", "url": "https://example.com/doc.pdf"},
"citations": {"enabled": True},
"title": "doc.pdf",
}
@ -283,6 +288,7 @@ def test_anthropic_whitespace_only_data_uri_falls_back_to_file_url(monkeypatch):
assert doc == {
"type": "document",
"source": {"type": "url", "url": "https://example.com/doc.pdf"},
"citations": {"enabled": True},
}

18
studio/frontend/src/components/assistant-ui/sources.tsx
View file

@ -127,6 +127,12 @@ function Source({
// ── Source badge with hover card ─────────────────────────────
interface SourceData {
/**
* Stable per-citation key. Two Anthropic document citations into
* different spans of the same source share a ``url``, so React keys
* on ``id`` to keep each footnote distinct.
*/
id: string;
url: string;
title: string;
description?: string;
@ -190,8 +196,14 @@ const SourcesGroup: FC = () => {
"url" in part &&
part.url
) {
const url = part.url as string;
const partId =
typeof (part as { id?: unknown }).id === "string"
? ((part as { id: string }).id)
: url;
sources.push({
url: part.url as string,
id: partId,
url,
title: (part as { title?: string }).title || "",
description: (part as { metadata?: { description?: string } })
.metadata?.description,
@ -258,7 +270,7 @@ const SourcesGroup: FC = () => {
className="flex w-full flex-wrap gap-1 invisible absolute pointer-events-none"
>
{sources.map((source) => (
<span key={source.url} className="inline-block">
<span key={source.id} className="inline-block">
<Source href={source.url}>
<SourceIcon url={source.url} />
<SourceTitle>{source.title || extractDomain(source.url)}</SourceTitle>
@ -270,7 +282,7 @@ const SourcesGroup: FC = () => {
{/* Visible container */}
<div className="flex flex-wrap gap-1">
{displayedSources.map((source) => (
<SourceBadge key={source.url} source={source} />
<SourceBadge key={source.id} source={source} />
))}
{shouldCollapse && !expanded && (
<button

31
studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx
View file

@ -24,6 +24,22 @@ const RE_TITLE = /Title:\s*(.+)/;
const RE_URL = /URL:\s*(.+)/;
const RE_SNIPPET = /Snippet:\s*(.+)/s;
/**
* Reject anything that is not a real http(s) URL. Web-search / web-fetch
* output is provider-controlled, so hostile ``javascript:`` / ``data:``
* lines must not reach the Source badge's <a href>.
*/
function isSafeHttpUrl(raw: string): boolean {
const value = raw.trim();
if (!value || /[\r\n]/.test(value)) return false;
try {
const parsed = new URL(value);
return parsed.protocol === "http:" || parsed.protocol === "https:";
} catch {
return false;
}
}
/** Parse the backend's "Title: ...\nURL: ...\nSnippet: ...\n---" format into structured sources. */
function parseSearchResults(raw: string): ParsedSource[] {
if (!raw) {
@ -35,13 +51,14 @@ function parseSearchResults(raw: string): ParsedSource[] {
const titleMatch = block.match(RE_TITLE);
const urlMatch = block.match(RE_URL);
const snippetMatch = block.match(RE_SNIPPET);
if (titleMatch && urlMatch) {
sources.push({
title: titleMatch[1].trim(),
url: urlMatch[1].trim(),
snippet: snippetMatch?.[1]?.trim() ?? "",
});
}
if (!titleMatch || !urlMatch) continue;
const url = urlMatch[1].trim();
if (!isSafeHttpUrl(url)) continue;
sources.push({
title: titleMatch[1].trim(),
url,
snippet: snippetMatch?.[1]?.trim() ?? "",
});
}
return sources;
}

124
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -151,6 +151,91 @@ async function updateStoredChatThreadEventually(
}
}
/**
* Return ``raw`` when it is a safe-to-navigate http(s) URL, or "" otherwise.
* Rejects non-string input, CR/LF (header injection), and non-http(s)
* schemes (``javascript:`` / ``data:`` / ``vbscript:``) so provider /
* tool-controlled strings cannot land in an <a href>.
*/
function isSafeNavigableSourceUrl(raw: unknown): string {
if (typeof raw !== "string") return "";
const value = raw.trim();
if (!value || /[\r\n]/.test(value)) return "";
try {
const parsed = new URL(value);
if (parsed.protocol === "http:" || parsed.protocol === "https:") {
return value;
}
} catch {
// Fall through.
}
return "";
}
/** Convert an Anthropic document citation dict into a Sources-panel source. */
function documentCitationToSource(
cit: Record<string, unknown>,
fallbackIdx: number,
): {
type: "source";
sourceType: "url";
id: string;
url: string;
title: string;
metadata?: { description: string };
} | null {
const source =
typeof cit.source === "string" && cit.source ? cit.source : "";
const docTitle =
(typeof cit.document_title === "string" && cit.document_title) ||
(typeof cit.title === "string" && cit.title) ||
"";
const docIndex =
typeof cit.document_index === "number" ? cit.document_index : undefined;
// Only treat ``source`` as a navigable URL when it is real http(s);
// search_result_location can carry a free-form id (e.g. ``kb-doc-42``)
// or a hostile ``javascript:`` / ``data:`` / ``vbscript:`` string.
// Fall back to a stable doc anchor otherwise.
const url =
isSafeNavigableSourceUrl(source) || `#anthropic-doc-${docIndex ?? fallbackIdx}`;
const title = docTitle || source || `Document ${fallbackIdx + 1}`;
const cited =
typeof cit.cited_text === "string" ? cit.cited_text.trim() : "";
// Trim the cited snippet so the Sources panel stays scannable.
const description =
cited.length > 240 ? `${cited.slice(0, 240)}...` : cited;
// Anthropic numbers inline [N] per citation, not per source URL.
// Fold citation type + position-bearing fields into the id so two
// distinct citations on the same source (or two search_result_locations
// with different search_result_index) keep separate Sources entries.
const citationType =
typeof cit.type === "string" ? String(cit.type) : "";
const positionParts = [
cit.search_result_index,
cit.start_char_index,
cit.end_char_index,
cit.start_page_number,
cit.end_page_number,
cit.start_block_index,
cit.end_block_index,
]
.filter((v) => typeof v === "number")
.map((v) => String(v))
.join(":");
const idAnchor = positionParts
? `${citationType}:${positionParts}`
: `${citationType}:${fallbackIdx}`;
const id = `${url}#${idAnchor}`;
return {
type: "source" as const,
sourceType: "url" as const,
id,
url,
title,
...(description ? { metadata: { description } } : {}),
};
}
/** Parse "Title: ...\nURL: ...\nSnippet: ..." blocks into source content parts. */
function parseSourcesFromResult(raw: string): {
type: "source";
@ -175,7 +260,11 @@ function parseSourcesFromResult(raw: string): {
const urlMatch = block.match(/URL:\s*(.+)/);
const snippetMatch = block.match(/Snippet:\s*(.+)/);
if (titleMatch && urlMatch) {
const url = urlMatch[1].trim();
// Drop blocks whose ``URL:`` is not safe http(s); provider/tool
// output is attacker-controllable so a hostile ``javascript:`` /
// ``data:`` line must not reach the Sources panel <a href>.
const url = isSafeNavigableSourceUrl(urlMatch[1]);
if (!url) continue;
const snippet = snippetMatch?.[1]?.trim();
sources.push({
type: "source" as const,
@ -1192,6 +1281,17 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// Tool call content parts — accumulated and yielded cumulatively.
// result is set directly on the tool-call part when tool_end arrives.
const toolCallParts: ToolCallMessagePart[] = [];
// Anthropic document_citations tool_event payload, converted to
// Sources-panel source parts at end-of-stream so the inline [N]
// markers have matching entries.
const documentCitationParts: Array<{
type: "source";
sourceType: "url";
id: string;
url: string;
title: string;
metadata?: { description: string };
}> = [];
// Latched on the `anthropic_refusal` tool event; stamped onto the
// final assistant metadata as `custom.anthropicRefusal` to drive
// the history-prune above.
@ -1648,6 +1748,27 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
}
continue;
}
if (toolEvent.type === "document_citations") {
// Convert Anthropic citations_delta footnotes into
// Sources-panel entries matching the inline [N] markers.
const cits = toolEvent.citations;
if (Array.isArray(cits)) {
cits.forEach((entry, idx) => {
if (!entry || typeof entry !== "object") return;
const part = documentCitationToSource(
entry as Record<string, unknown>,
idx,
);
if (
part &&
!documentCitationParts.some((p) => p.id === part.id)
) {
documentCitationParts.push(part);
}
});
}
continue;
}
if (toolEvent.type === "container_invalidated") {
if (resolvedThreadId) {
const field =
@ -1995,6 +2116,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
...toolCallParts,
...parseAssistantContent(cumulativeText),
...sourceParts,
...documentCitationParts,
],
metadata: {
timing: finalTiming,