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

Studio: rewrite OpenAI Responses citation markers to markdown links (#5713)

Browse source 
* Studio: rewrite OpenAI Responses citation markers to markdown links

OpenAI's /v1/responses stream interleaves text deltas with inline
citation markers built from private-use codepoints (U+E200 / U+E201 /
U+E202) shaped like `citeSOURCE_ID`. The codepoints render
as garbled "E202" glyphs or empty boxes in most fonts, and the
markdown layer further strips them, leaving run-on text like
"citeturn1view0turn1view1turn3view0...". The url list still arrived in
the Sources panel via url_citation annotations, but the inline cite
hand-off into the prose was unreadable.

Rewrite each marker into `[N](URL)` when the matching url_citation
has already been recorded on this stream, and drop the marker
silently otherwise. The lookup uses a new `source_id` field captured
on `_record_url_citation` (accepts source_id / id / locator across
Responses API revisions). Annotations are now applied BEFORE the
delta text is rewritten so that markers and their resolving
annotation arriving in the same SSE event still resolve.

Reference: https://developers.openai.com/api/docs/guides/citation-formatting

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

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

* Preserve every source_id alias for a deduplicated url_citation

OpenAI's Responses stream cites the same URL under multiple
source_id markers when the model references different spans of the
same page. The previous dedup-by-URL kept only the first alias and
dropped the rest, so subsequent markers for the same URL never
resolved and got stripped from the prose. Switch the citation
record to a ``source_ids`` list and append new aliases on every
duplicate. The rewriter resolves any alias back to the same
citation number so the inline markers all collapse onto one footnote
rather than fanning out into bogus repeats.

Also collapse the two passes over ``all_url_citations`` in
``_record_url_citation`` into a single loop for clarity. Adds two
regression tests covering the alias-collision and mixed-shape cases.

* ci: re-trigger after flake in Studio GGUF Tool calling (rebased on main #5741 already)

* ci: re-run after transient CodeQL Python checkout auth flake

* Fix split-marker buffer + multi-source ids for PR #5713

The original rewriter only handles markers that arrive whole inside a
single response.output_text.delta event. OpenAI's stream chunks text
on byte-buffer boundaries with no awareness of the marker grammar,
so a marker can straddle two deltas (delta-1 ends with
"citetu", delta-2 starts with "rn0view0"). Each delta
was rewritten in isolation, so the half-marker leaked as garbled
"E200/E202" glyphs in the rendered prose.

Buffer the unterminated tail across deltas and concatenate it onto
the front of the next one so the rewriter sees a complete marker.
Flush the held-over tail on response.completed / response.incomplete /
[DONE], stripping any leftover private-use bytes so a never-closed
marker (truncated stream, missing annotation) never leaks.

Also handle the multi-source marker shape from the OpenAI docs --
citeid1id2 should expand to one bracket
link per resolvable id. The previous regex captured only the first
source id and silently dropped id2/id3.

Reference: https://developers.openai.com/api/docs/guides/citation-formatting

Tests: 21 new cases covering multi-source, locator suffix, marker
split across two and three deltas, unterminated marker on truncation,
late annotation resolving a buffered marker, idempotency, and the
head/tail split helper directly.

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

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

* Defer citation segments until url_citation annotation arrives

The split-marker buffer already concatenates a marker that straddles
two response.output_text.delta events. But when the annotation event
for a url_citation arrives AFTER the delta that contains its inline
marker (the typical OpenAI Responses ordering), the rewriter still
saw an empty lookup table at delta time and silently stripped the
marker. The URL kept showing up in the sources panel but the inline
link reference was permanently gone.

Add _rewrite_citation_markers_partial which leaves an unresolved
marker verbatim and reports has_unresolved=True. The streaming loop
buffers any closed segment that contains an unresolved marker into a
pending_citation_segments FIFO and drains the queue on every later
annotation event, on response.completed, on response.incomplete, and
on the [DONE] sentinel. Drain order is preserved so later clean text
does not leapfrog an earlier deferred segment. End-of-stream forces a
strip so no codepoint leaks if the annotation never arrived.

Add six regression tests covering single-pass resolution, the late-
annotation two-pass case, multi-source markers with partial
resolution, mixed known and pending markers in one segment, and
idempotency on marker-free input.

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

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

* Drop unterminated citation tail to prevent cite-prefix plain-text leak

`_flush_pending_marker_tail` stripped the three private-use citation
codepoints from the held-over buffer, but left the literal ``cite``
keyword plus the source id behind as plain text. A stream ending
mid-marker therefore emitted user-visible garbage like
``Some text citeturn0view0`` instead of the intended clean prose.

``pending_marker_tail`` is by construction the suffix that starts at
an unclosed ``\\ue200`` opener -- the split helper guarantees there is
no closing ``\\ue201`` byte. Without that close the marker is
meaningless: the source id cannot be resolved to a URL and the user
prose before the opener was already emitted as ``head`` on the
originating delta. Bail out before the strip step and return the
empty string. As a belt-and-braces measure also drop any orphan
``cite<sid>`` literal at the head of the buffer in case a future
caller passes a partially-terminated tail.

Update the matching ``_simulate_delta_stream`` harness in the edge
tests so it mirrors the new flush logic, and add four regression
tests covering unterminated marker with surrounding prose, marker-
only inputs, prefix-only outputs, and the split-then-close path that
still must resolve to a link.

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

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

* Defer multi-source markers until all ids resolve for PR #5713

`_rewrite_citation_markers_partial` previously treated a marker as
resolved when even one token in a multi-source marker resolved,
dropping any still-pending source ids. In streamed Responses events
the annotations for a multi-source marker can arrive across separate
`annotation.added` chunks, so the caller no longer buffered that
segment for retry and the late source id was lost from the inline
citation entirely.

Flag the marker unresolved whenever any token misses the lookup so
the streamer keeps the segment pending. End-of-stream force flush
still drops unresolved tokens through `_replace_openai_citation_markers`
so locator-style suffixes (which look like unresolved ids at the token
level but only appear at end-of-stream) render cleanly.

Updated the multi-source test to assert the new pending-then-flush
behavior; locator output now lands at force-flush rather than mid
stream.

* Shorten citation marker comments for PR #5713

---------

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:37:16 -07:00 • committed by GitHub
commit 063e1e497b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 994 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -68,6 +68,136 @@ _ANTHROPIC_4_7_SAMPLING_REMOVED = re.compile(
)
_OPENAI_REASONING_SUMMARY_UNSUPPORTED = re.compile(r"^o3(?:[-.]|$)")
# OpenAI Responses inline citation markers: `citeSOURCE_ID[id2...][LOCATOR]`
# using private-use codepoints (see
# https://developers.openai.com/api/docs/guides/citation-formatting).
# Group 1 holds the delim-separated tokens; each resolvable token expands
# to `[[N]](URL)`, unresolved tokens (locators, unknown ids) drop silently
# so no garbled glyph reaches the renderer.
_OPENAI_CITE_OPEN = "cite"
_OPENAI_CITE_STOP = ""
_OPENAI_CITE_DELIM = ""
_OPENAI_CITATION_MARKER = re.compile(
f"{_OPENAI_CITE_OPEN}([^{_OPENAI_CITE_STOP}]+){_OPENAI_CITE_STOP}"
)
def _build_citation_lookup(
url_citations: list[dict[str, Any]],
) -> dict[str, tuple[int, str]]:
"""Map every known ``source_id`` alias to ``(citation_index, url)``.
Accepts singular ``source_id`` and plural ``source_ids``. First-seen
wins on alias collision so an earlier citation keeps its number.
"""
by_source: dict[str, tuple[int, str]] = {}
for idx, cit in enumerate(url_citations, start = 1):
url = cit.get("url")
if not isinstance(url, str) or not url:
continue
aliases: list[str] = []
sid = cit.get("source_id")
if isinstance(sid, str) and sid:
aliases.append(sid)
sids = cit.get("source_ids")
if isinstance(sids, list):
aliases.extend(s for s in sids if isinstance(s, str) and s)
for alias in aliases:
by_source.setdefault(alias, (idx, url))
return by_source
def _replace_openai_citation_markers(
text: str,
url_citations: list[dict[str, Any]],
) -> str:
"""Rewrite `\\ue200cite\\ue202SOURCE_ID[\\ue202LOCATOR]\\ue201` markers into
`[[N]](URL)` per resolvable id. Multi-source markers expand to one link
per id; unresolved tokens drop silently. Idempotent on text without
private-use codepoints.
"""
if not text or _OPENAI_CITE_STOP not in text:
return text
by_source = _build_citation_lookup(url_citations)
def _sub(match: re.Match[str]) -> str:
# Try every delim-split token; unresolved tokens drop silently.
# Handles multi-source (all resolve) and source+locator (only the
# id resolves, locator drops). Empty result strips the marker.
rendered: list[str] = []
for tok in match.group(1).split(_OPENAI_CITE_DELIM):
if not tok:
continue
hit = by_source.get(tok)
if hit is None:
continue
idx, url = hit
rendered.append(f"[[{idx}]]({url})")
return "".join(rendered)
return _OPENAI_CITATION_MARKER.sub(_sub, text)
def _rewrite_citation_markers_partial(
text: str,
url_citations: list[dict[str, Any]],
) -> tuple[str, bool]:
"""Like ``_replace_openai_citation_markers`` but also reports whether
any marker referenced a source_id not yet in ``url_citations``.
The ``annotation.added`` event for a url_citation typically arrives
AFTER the delta carrying the marker referencing it. Callers buffer the
segment until a later event records the annotation; unresolved markers
are left verbatim so a follow-up pass still parses cleanly.
"""
if not text or _OPENAI_CITE_STOP not in text:
return text, False
by_source = _build_citation_lookup(url_citations)
has_unresolved = False
def _sub(match: re.Match[str]) -> str:
nonlocal has_unresolved
tokens = [t for t in match.group(1).split(_OPENAI_CITE_DELIM) if t]
rendered: list[str] = []
any_unresolved = False
for tok in tokens:
hit = by_source.get(tok)
if hit is None:
any_unresolved = True
continue
idx, url = hit
rendered.append(f"[[{idx}]]({url})")
# Leave the whole marker verbatim if any token is unresolved so the
# caller can re-run once the late annotation lands; partial emission
# would lose the unresolved ids once the source text is dropped.
if any_unresolved:
has_unresolved = True
return match.group(0)
return "".join(rendered)
return _OPENAI_CITATION_MARKER.sub(_sub, text), has_unresolved
def _split_pending_citation_tail(text: str) -> tuple[str, str]:
"""Split ``text`` into ``(head, pending_tail)`` for streamed deltas.
A citation marker can straddle two SSE deltas (e.g. delta-1 ends with
``\\ue200citetu`` and delta-2 starts with ``rn0view0\\ue201``); the
unterminated tail is buffered and prepended onto the next delta so the
rewriter sees a complete marker. ``pending_tail`` is the longest suffix
starting with ``\\ue200`` and lacking ``\\ue201``; ``head`` is safe to
emit. Empty tail when ``text`` has no open marker or a fully closed one.
"""
if not text:
return text, ""
last_open = text.rfind("")
if last_open == -1:
return text, ""
# Stop byte after the last open byte means the marker closed in this delta.
if _OPENAI_CITE_STOP in text[last_open:]:
return text, ""
return text[:last_open], text[last_open:]
class _AnthropicThinkingSpec(NamedTuple):
prefixes: tuple[str, ...]
@ -3000,6 +3130,65 @@ class ExternalProviderClient:
# see.
latched_container_id: Optional[str] = None
container_id_emitted = False
# Buffer for a citation marker straddling two delta events;
# prepended onto the next delta. See _split_pending_citation_tail.
pending_marker_tail: str = ""
# Segments deferred while their markers reference unseen
# source_ids; held in arrival order so output never
# leapfrogs an earlier deferred segment. Flushed on
# annotation events and force-flushed at end-of-stream
# with leftover private-use codepoints stripped.
pending_citation_segments: list[str] = []
def _drain_pending_segments(force: bool) -> str:
"""Re-attempt resolution on buffered segments in order.
Stops at the first still-unresolved segment unless
``force`` (end-of-stream), where lingering markers are stripped."""
out: list[str] = []
while pending_citation_segments:
seg = pending_citation_segments[0]
rewritten, unresolved = _rewrite_citation_markers_partial(
seg,
all_url_citations,
)
if unresolved and not force:
pending_citation_segments[0] = rewritten
break
if unresolved and force:
rewritten = _replace_openai_citation_markers(
rewritten,
all_url_citations,
)
pending_citation_segments.pop(0)
if rewritten:
out.append(rewritten)
return "".join(out)
def _flush_pending_marker_tail(tail: str) -> str:
"""Render any leftover citation tail at end-of-stream.
Unterminated tails drop (no annotation to bind to). If the
close byte arrived concatenated, rewrite then scrub any
residual private-use bytes and any orphan ``cite<sid>``
literal so the renderer never sees raw markup. url_citations
are aggregated separately and applied to web_search tool_end.
"""
if not tail:
return ""
if _OPENAI_CITE_STOP not in tail:
# Unterminated: drop the whole tail, otherwise the
# residual ``cite<sid>`` would leak as plain text.
return ""
rendered = _replace_openai_citation_markers(
tail, all_url_citations
)
# Scrub residual private-use bytes (e.g. a partial opener).
for ch in ("", "", ""):
rendered = rendered.replace(ch, "")
# Drop any orphan ``cite<sid>`` literal -- meaningless
# without its closing byte and matching url_citation.
rendered = re.sub(r"^cite\S*", "", rendered)
return rendered
def _emit_tool_event(payload: dict[str, Any]) -> str:
chunk = {
@ -3057,16 +3246,35 @@ class ExternalProviderClient:
def _record_url_citation(payload: dict[str, Any]) -> None:
"""Append a url_citation onto the shared all_url_citations
list. Dedup by URL — the same source can be cited multiple
times across deltas. We do NOT try to attribute citations
to individual web_search_call invocations because OpenAI's
annotation events don't carry that linkage."""
list. Dedup by URL — the same URL can be cited many
times under different ``source_id`` aliases (one per
span/locator), so collect every alias we see onto
the matching entry's ``source_ids`` list. The
delta-text rewriter resolves any of those aliases
back to this entry's URL. The id may live under
``source_id``, ``id``, or ``locator`` across the
Responses API revisions."""
if payload.get("type") != "url_citation":
return
url = payload.get("url", "")
if not url:
return
if any(c["url"] == url for c in all_url_citations):
source_id = (
payload.get("source_id")
or payload.get("id")
or payload.get("locator")
or ""
)
# Single pass: either backfill aliases onto an
# existing URL entry (and return) or fall through
# to append a fresh one.
for c in all_url_citations:
if c["url"] != url:
continue
if source_id:
aliases = c.setdefault("source_ids", [])
if source_id not in aliases:
aliases.append(source_id)
return
title = payload.get("title") or url
snippet = payload.get("snippet") or payload.get("quote") or ""
@ -3075,6 +3283,7 @@ class ExternalProviderClient:
"url": url,
"title": title,
"snippet": snippet,
"source_ids": [source_id] if source_id else [],
}
)
@ -3131,6 +3340,28 @@ class ExternalProviderClient:
if not data_str:
continue
if data_str == "[DONE]":
# Flush any held-over partial marker; strip
# private-use bytes so garbled glyphs don't leak.
if pending_marker_tail:
flushed = _flush_pending_marker_tail(
pending_marker_tail
)
pending_marker_tail = ""
if flushed:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
yield _chunk_with_text(flushed)
# Force-drain any segment still awaiting an
# annotation; lingering codepoints are stripped.
tail_flushed = _drain_pending_segments(
force = True,
)
if tail_flushed:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
yield _chunk_with_text(tail_flushed)
if not done_emitted:
yield "data: [DONE]"
done_emitted = True
@ -3145,22 +3376,57 @@ class ExternalProviderClient:
if event_type == "response.output_text.delta":
delta_text = event.get("delta", "")
if delta_text:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
yield _chunk_with_text(delta_text)
# Some API versions inline url citations on the
# delta event itself rather than as a separate
# response.output_text.annotation.added event.
# Process inline annotations first so source_ids
# referenced by same-delta markers are in the lookup
# before the rewriter runs. Some API versions inline
# url citations on the delta event itself.
for ann in event.get("annotations") or []:
if isinstance(ann, dict):
_record_url_citation(ann)
if delta_text or pending_marker_tail:
# Prepend any held-over tail so a marker
# straddling two SSE events resolves cleanly.
combined = pending_marker_tail + delta_text
head, pending_marker_tail = (
_split_pending_citation_tail(combined)
)
if head:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
# Re-attempt earlier deferred segments first
# so output stays in order; the needed
# annotation may have arrived inline above.
flushed = _drain_pending_segments(
force = False,
)
if flushed:
yield _chunk_with_text(flushed)
head_rewritten, has_unresolved = (
_rewrite_citation_markers_partial(
head,
all_url_citations,
)
)
if has_unresolved or pending_citation_segments:
pending_citation_segments.append(
head_rewritten
)
elif head_rewritten:
yield _chunk_with_text(head_rewritten)
elif event_type == "response.output_text.annotation.added":
ann = event.get("annotation")
if isinstance(ann, dict):
_record_url_citation(ann)
flushed = _drain_pending_segments(
force = False,
)
if flushed:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
yield _chunk_with_text(flushed)
elif event_type == "response.output_item.added":
# Track the call early but do NOT emit tool_start
@ -3399,6 +3665,34 @@ class ExternalProviderClient:
)
if isinstance(completed_usage, dict):
last_usage = completed_usage
# Flush any unterminated citation tail
# held over from the last delta. By
# the time we get here every annotation
# has been recorded so a late-arriving
# source_id may resolve cleanly; if it
# still doesn't, the helper strips the
# private-use bytes so no garbled
# glyph reaches the user.
if pending_marker_tail:
flushed = _flush_pending_marker_tail(
pending_marker_tail
)
pending_marker_tail = ""
if flushed:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
yield _chunk_with_text(flushed)
# Force-drain any segment still awaiting an
# annotation; lingering codepoints are stripped.
tail_flushed = _drain_pending_segments(
force = True,
)
if tail_flushed:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
yield _chunk_with_text(tail_flushed)
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
@ -3491,6 +3785,29 @@ class ExternalProviderClient:
)
if isinstance(incomplete_usage, dict):
last_usage = incomplete_usage
# Same flush as response.completed --
# truncated streams can leave a half-
# marker in the buffer.
if pending_marker_tail:
flushed = _flush_pending_marker_tail(
pending_marker_tail
)
pending_marker_tail = ""
if flushed:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
yield _chunk_with_text(flushed)
# Force-drain any segment still awaiting an
# annotation; lingering codepoints are stripped.
tail_flushed = _drain_pending_segments(
force = True,
)
if tail_flushed:
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False
yield _chunk_with_text(tail_flushed)
if reasoning_open:
yield _chunk_with_text("</think>")
reasoning_open = False

251
studio/backend/tests/test_openai_citation_markers.py Normal file
View file

@ -0,0 +1,251 @@
# 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 the OpenAI Responses-API citation marker rewriter.
The stream interleaves text deltas with ``\\ue200cite\\ue202SOURCE_ID\\ue201``
markers. The rewriter resolves each to `[N](URL)` when the annotation has
arrived and drops it otherwise; the URL list still flows to Sources via
`_record_url_citation`.
Reference: https://developers.openai.com/api/docs/guides/citation-formatting
"""
import pytest
from core.inference.external_provider import (
_replace_openai_citation_markers,
_rewrite_citation_markers_partial,
)
# Citation marker control codepoints (private-use area):
CITE_START = ""
CITE_STOP = ""
CITE_DELIM = ""
def _marker(source_id: str, locator: str | None = None) -> str:
payload = f"{CITE_START}cite{CITE_DELIM}{source_id}"
if locator:
payload = f"{payload}{CITE_DELIM}{locator}"
return f"{payload}{CITE_STOP}"
def _has_marker_codepoints(text: str) -> bool:
return any(c in text for c in (CITE_START, CITE_STOP, CITE_DELIM))
def test_passthrough_when_no_marker_present():
text = "Plain text with no citation markers."
assert _replace_openai_citation_markers(text, []) == text
def test_marker_rewritten_to_link_when_annotation_known():
text = f"The capital is Paris {_marker('turn0view0')}."
citations = [
{
"source_id": "turn0view0",
"url": "https://example.com/paris",
"title": "Paris",
},
]
out = _replace_openai_citation_markers(text, citations)
assert not _has_marker_codepoints(out)
assert "[[1]](https://example.com/paris)" in out
def test_unknown_source_marker_dropped_silently():
text = f"Foo {_marker('turn9view9')} bar."
out = _replace_openai_citation_markers(text, [])
# Marker stripped, no garbled "E202" glyph leaks through, and the
# surrounding text stays intact.
assert not _has_marker_codepoints(out)
assert "E202" not in out
assert "turn9view9" not in out
assert "Foo" in out and "bar" in out
def test_multiple_concatenated_markers_resolved_in_order():
"""Real-world wire shape: a string of markers butted up against each other
after a sentence, as in the user-reported bug."""
markers = "".join(_marker(f"turn{i}view{j}") for i, j in [(1, 0), (1, 1), (3, 0)])
text = f"All animals ranked. {markers}"
citations = [
{"source_id": "turn1view0", "url": "https://a.example/dog", "title": "Dog"},
{"source_id": "turn1view1", "url": "https://a.example/cat", "title": "Cat"},
{"source_id": "turn3view0", "url": "https://a.example/tiger", "title": "Tiger"},
]
out = _replace_openai_citation_markers(text, citations)
assert "[[1]](https://a.example/dog)" in out
assert "[[2]](https://a.example/cat)" in out
assert "[[3]](https://a.example/tiger)" in out
assert not _has_marker_codepoints(out)
def test_marker_with_locator_resolves():
text = f"See {_marker('turn2file0', 'L8-L13')}."
citations = [
{"source_id": "turn2file0", "url": "https://example.com/doc.txt"},
]
out = _replace_openai_citation_markers(text, citations)
assert "[[1]](https://example.com/doc.txt)" in out
assert "L8-L13" not in out # locator detail dropped; we just link.
assert not _has_marker_codepoints(out)
def test_mixed_known_and_unknown_markers():
known = _marker("turn0view0")
unknown = _marker("turn0view99")
text = f"Known {known} and unknown {unknown}."
citations = [
{"source_id": "turn0view0", "url": "https://example.com/known"},
]
out = _replace_openai_citation_markers(text, citations)
assert "[[1]](https://example.com/known)" in out
# Unknown markers leave no trace, but surrounding prose stays.
assert "Known" in out and "unknown" in out
assert not _has_marker_codepoints(out)
assert "E202" not in out
def test_empty_text_returns_verbatim():
assert _replace_openai_citation_markers("", []) == ""
def test_idempotent_on_pre_stripped_text():
"""Pre-stripped text (no private-use codepoints) returns verbatim."""
text = "citeturn1view0 plain"
assert _replace_openai_citation_markers(text, []) == text
@pytest.mark.parametrize(
"citation",
[
{"url": "https://example.com/a"}, # no source_id at all
{"source_id": None, "url": "https://example.com/b"},
{"source_id": "", "url": "https://example.com/c"},
],
)
def test_citation_without_source_id_does_not_crash(citation):
text = f"X {_marker('turnXviewY')} Y"
out = _replace_openai_citation_markers(text, [citation])
# No mapping, marker stripped. Crash-free is the contract.
assert not _has_marker_codepoints(out)
assert "turnXviewY" not in out
def test_multiple_source_id_aliases_resolve_to_same_url():
"""Every alias for the same URL must resolve, not just the first.
Regression for the Codex P1 on the original PR."""
a = _marker("turn0view0")
b = _marker("turn0view0_span_1")
c = _marker("turn0view0_span_2")
text = f"Triple {a}{b}{c} cite."
citations = [
{
"source_ids": ["turn0view0", "turn0view0_span_1", "turn0view0_span_2"],
"url": "https://example.com/paris",
"title": "Paris",
},
]
out = _replace_openai_citation_markers(text, citations)
# All three aliases collapse onto citation [1] -- the URL is the
# same so it would be misleading to show three different numbers.
assert out.count("[[1]](https://example.com/paris)") == 3
assert not _has_marker_codepoints(out)
def test_source_ids_list_and_legacy_source_id_both_resolve():
"""Mixed-shape citation: legacy ``source_id`` plus newer
``source_ids`` aliases both resolve."""
legacy = _marker("legacy_id")
alias = _marker("alias_id")
text = f"Both {legacy} and {alias} work."
citations = [
{
"source_id": "legacy_id",
"source_ids": ["alias_id"],
"url": "https://example.com/doc",
},
]
out = _replace_openai_citation_markers(text, citations)
assert out.count("[[1]](https://example.com/doc)") == 2
assert not _has_marker_codepoints(out)
# ---------------------------------------------------------------------------
# _rewrite_citation_markers_partial: deferred-annotation tests. OpenAI emits
# url_citation annotations on a subsequent SSE event; this helper reports
# `has_unresolved` so the stream loop defers emission. See PR #5713 audit.
# ---------------------------------------------------------------------------
def test_partial_known_marker_resolves_and_clears_unresolved():
text = f"Foo {_marker('s1')} bar."
out, unresolved = _rewrite_citation_markers_partial(
text,
[{"source_id": "s1", "url": "https://example.com/a"}],
)
assert "[[1]](https://example.com/a)" in out
assert unresolved is False
assert not _has_marker_codepoints(out)
def test_partial_unknown_marker_preserves_verbatim_and_flags():
text = f"Foo {_marker('s1')} bar."
out, unresolved = _rewrite_citation_markers_partial(text, [])
assert unresolved is True
# Codepoints must remain so a follow-up pass can re-parse.
assert _has_marker_codepoints(out)
assert "Foo" in out and "bar." in out
def test_partial_resolves_after_late_annotation():
"""Two-pass: first call sees no citations, second resolves after annotation."""
text = f"See {_marker('s1')} for details."
out1, unresolved1 = _rewrite_citation_markers_partial(text, [])
assert unresolved1 is True
citations = [{"source_id": "s1", "url": "https://example.com/x"}]
out2, unresolved2 = _rewrite_citation_markers_partial(out1, citations)
assert unresolved2 is False
assert "[[1]](https://example.com/x)" in out2
assert not _has_marker_codepoints(out2)
def test_partial_multi_source_partial_resolution_keeps_marker_pending():
"""Any unresolved token in a multi-source marker leaves the whole marker
verbatim with ``unresolved`` True; defer until every id resolves or
end-of-stream forces a flush (dropping unresolved tokens then)."""
cite = f"{CITE_START}cite{CITE_DELIM}known{CITE_DELIM}locator{CITE_STOP}"
text = f"Pre {cite} post."
citations = [{"source_id": "known", "url": "https://example.com/y"}]
out, unresolved = _rewrite_citation_markers_partial(text, citations)
assert unresolved is True
assert cite in out
# End-of-stream force flush: drop the unresolved token, keep the
# resolved link. The streamer routes pending segments through
# `_replace_openai_citation_markers` at force=True for this.
forced = _replace_openai_citation_markers(out, citations)
assert "[[1]](https://example.com/y)" in forced
assert "locator" not in forced
assert not _has_marker_codepoints(forced)
def test_partial_idempotent_on_marker_free_text():
text = "Plain text."
out, unresolved = _rewrite_citation_markers_partial(text, [])
assert out == text
assert unresolved is False
def test_partial_mixed_known_and_pending_markers_flags_unresolved():
known = _marker("known")
pending = _marker("pending")
text = f"{known} {pending}"
citations = [{"source_id": "known", "url": "https://example.com/k"}]
out, unresolved = _rewrite_citation_markers_partial(text, citations)
assert unresolved is True # the pending marker drives the flag
assert "[[1]](https://example.com/k)" in out
# The pending marker stays verbatim for the next pass.
assert CITE_START in out and "pending" in out

413
studio/backend/tests/test_openai_citation_markers_edge.py Normal file
View file

@ -0,0 +1,413 @@
# 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 the OpenAI Responses citation marker rewriter.
Covers multi-source markers, source+locator, marker SPLIT across SSE deltas,
unterminated tails at end-of-stream, multiple markers per delta, late
annotation ordering, and idempotency.
Reference: https://developers.openai.com/api/docs/guides/citation-formatting
"""
import importlib
# Streaming integration is exercised by ``_simulate_delta_stream`` further
# down, mirroring the head/buffer/flush dance from ``_stream_openai_responses``.
_module = importlib.import_module("core.inference.external_provider")
_replace_openai_citation_markers = _module._replace_openai_citation_markers
_split_pending_citation_tail = _module._split_pending_citation_tail
CITE_START = ""
CITE_STOP = ""
CITE_DELIM = ""
def _marker(*source_ids: str, locator: str | None = None) -> str:
"""Build a ``\\ue200cite\\ue202<sid>[\\ue202<sid>...][\\ue202<loc>]\\ue201``
marker. Accepts one or many ``source_ids`` plus an optional ``locator``."""
payload = f"{CITE_START}cite{CITE_DELIM}" + CITE_DELIM.join(source_ids)
if locator:
payload = f"{payload}{CITE_DELIM}{locator}"
return f"{payload}{CITE_STOP}"
def _no_private_use(text: str) -> bool:
return all(c not in text for c in (CITE_START, CITE_STOP, CITE_DELIM))
# Harness mirroring the head/pending-tail/flush dance in
# `_stream_openai_responses`, so streaming tests skip the httpx mock.
def _simulate_delta_stream(
deltas: list[str],
citations: list[dict],
*,
flush: bool = True,
) -> str:
pending = ""
emitted: list[str] = []
for delta in deltas:
combined = pending + delta
head, pending = _split_pending_citation_tail(combined)
if head:
head = _replace_openai_citation_markers(head, citations)
if head:
emitted.append(head)
if flush and pending:
# Mirror `_flush_pending_marker_tail`: drop the tail entirely if no
# closing stop byte arrived; the literal ``cite<sid>`` would leak otherwise.
if CITE_STOP not in pending:
rendered = ""
else:
rendered = _replace_openai_citation_markers(pending, citations)
for ch in (CITE_START, CITE_STOP, CITE_DELIM):
rendered = rendered.replace(ch, "")
import re as _re
rendered = _re.sub(r"^cite\S*", "", rendered)
if rendered:
emitted.append(rendered)
return "".join(emitted)
# ---------------------------------------------------------------------------
# 1. Multi-source markers per the OpenAI docs.
# ---------------------------------------------------------------------------
def test_multi_source_marker_all_resolve():
"""\\ue200cite\\ue202id1\\ue202id2\\ue202id3\\ue201 expands to three links
when every id is known. Earlier regex captured only id1 and dropped id2/id3."""
text = f"All three: {_marker('id1', 'id2', 'id3')}"
citations = [
{"source_id": "id1", "url": "https://example.com/1"},
{"source_id": "id2", "url": "https://example.com/2"},
{"source_id": "id3", "url": "https://example.com/3"},
]
out = _replace_openai_citation_markers(text, citations)
assert "[[1]](https://example.com/1)" in out
assert "[[2]](https://example.com/2)" in out
assert "[[3]](https://example.com/3)" in out
assert _no_private_use(out)
def test_multi_source_marker_partial_resolution():
"""Known ids render, unknown ids drop silently, no glyph leaks."""
text = f"Mixed: {_marker('known', 'unknown', 'also_known')}"
citations = [
{"source_id": "known", "url": "https://k.example"},
{"source_id": "also_known", "url": "https://ak.example"},
]
out = _replace_openai_citation_markers(text, citations)
assert "[[1]](https://k.example)" in out
assert "[[2]](https://ak.example)" in out
assert "unknown" not in out
assert _no_private_use(out)
# ---------------------------------------------------------------------------
# 2. Source + locator: locator is dropped, link still resolves.
# ---------------------------------------------------------------------------
def test_marker_with_numeric_locator():
text = f"See {_marker('tu0', locator = '42')}."
citations = [{"source_id": "tu0", "url": "https://example.com/doc"}]
out = _replace_openai_citation_markers(text, citations)
assert "[[1]](https://example.com/doc)" in out
assert "42" not in out
assert _no_private_use(out)
def test_marker_with_range_locator():
text = f"See {_marker('tu0', locator = 'L8-L13')}."
citations = [{"source_id": "tu0", "url": "https://example.com/code"}]
out = _replace_openai_citation_markers(text, citations)
assert "[[1]](https://example.com/code)" in out
assert "L8-L13" not in out
assert _no_private_use(out)
# ---------------------------------------------------------------------------
# 3. Marker SPLIT across two SSE deltas -- the codex-flagged P1.
# ---------------------------------------------------------------------------
def test_marker_split_in_source_id():
"""Delta-1 ends mid-source-id (``\\ue200cite\\ue202tu``), delta-2 starts
with the rest (``rn0view0\\ue201``). The buffer stitches the halves
back together so they resolve to one link instead of leaking."""
full = f"See {_marker('turn0view0')} now."
# Cut right after the second delim + "tu" inside the source id.
cut = full.index("tu", full.index(CITE_START)) + len("tu")
d1, d2 = full[:cut], full[cut:]
# Sanity check: delta-1 actually contains a partial marker.
assert CITE_START in d1 and CITE_STOP not in d1
assert CITE_STOP in d2
citations = [{"source_id": "turn0view0", "url": "https://x"}]
out = _simulate_delta_stream([d1, d2], citations)
assert out == "See [[1]](https://x) now."
assert _no_private_use(out)
def test_marker_split_at_start_byte():
"""Split exactly after the opening ``\\ue200`` byte; the buffer must
hold the lone open byte until the rest arrives."""
full = f"Text {_marker('sid')} done"
cut = full.index(CITE_START) + 1 # right AFTER the open byte
d1, d2 = full[:cut], full[cut:]
citations = [{"source_id": "sid", "url": "https://y"}]
out = _simulate_delta_stream([d1, d2], citations)
assert out == "Text [[1]](https://y) done"
assert _no_private_use(out)
def test_marker_split_across_three_deltas():
"""Worst case: marker chopped into three pieces across three deltas."""
full = f"A {_marker('threesplit')} B"
# cut at two points inside the marker
open_pos = full.index(CITE_START)
stop_pos = full.index(CITE_STOP)
cut1 = open_pos + 4
cut2 = stop_pos - 2
parts = [full[:cut1], full[cut1:cut2], full[cut2:]]
citations = [{"source_id": "threesplit", "url": "https://z"}]
out = _simulate_delta_stream(parts, citations)
assert out == "A [[1]](https://z) B"
assert _no_private_use(out)
def test_marker_split_with_trailing_text_after_close():
"""Delta-2 closes the marker AND carries trailing prose; both emit cleanly."""
full = f"X {_marker('sid')} after"
cut = full.index("cite") + len("ci")
d1, d2 = full[:cut], full[cut:]
citations = [{"source_id": "sid", "url": "https://a"}]
out = _simulate_delta_stream([d1, d2], citations)
assert out == "X [[1]](https://a) after"
assert _no_private_use(out)
def test_split_marker_unknown_source_is_dropped_cleanly():
"""Split marker for an unknown source drops silently on flush."""
full = f"Pre {_marker('never_seen')} post"
cut = full.index(CITE_START) + 3
d1, d2 = full[:cut], full[cut:]
out = _simulate_delta_stream([d1, d2], [])
assert out == "Pre post"
assert _no_private_use(out)
# ---------------------------------------------------------------------------
# 4. Unterminated marker at end-of-stream -- truncation safety.
# ---------------------------------------------------------------------------
def test_unterminated_marker_at_stream_end_dropped_on_flush():
"""Stream ends mid-marker (e.g. response.incomplete); the tail is
flushed with private-use bytes stripped, no `E202` text leaks."""
deltas = ["Some text ", f"{CITE_START}citetu", "rn0view0"] # no STOP ever
out = _simulate_delta_stream(deltas, [], flush = True)
assert _no_private_use(out)
assert "E200" not in out and "E202" not in out
# Surrounding prose stays; we don't assert exact marker remainder.
assert "Some text " in out
def test_flush_resolves_marker_when_late_annotation_arrives():
"""Marker in a delta, matching annotation arrives later (on
response.output_text.annotation.added after the final delta). The
rewriter reads ``all_url_citations`` LIVE at flush, so the buffered
marker still resolves."""
deltas = ["Look ", f"{CITE_START}cite{CITE_DELIM}late_sid"]
pending = ""
citations: list[dict] = []
emitted: list[str] = []
for d in deltas:
combined = pending + d
head, pending = _split_pending_citation_tail(combined)
if head:
emitted.append(_replace_openai_citation_markers(head, citations))
# Annotation arrives AFTER all deltas but BEFORE flush.
citations.append({"source_id": "late_sid", "url": "https://late.example"})
# Append the STOP byte that closed the marker in a later delta.
pending = pending + CITE_STOP
flushed = _replace_openai_citation_markers(pending, citations)
for ch in (CITE_START, CITE_STOP, CITE_DELIM):
flushed = flushed.replace(ch, "")
emitted.append(flushed)
out = "".join(emitted)
assert "[[1]](https://late.example)" in out
assert _no_private_use(out)
# ---------------------------------------------------------------------------
# 5. Multiple unrelated markers in a single delta.
# ---------------------------------------------------------------------------
def test_three_markers_in_one_delta_resolve_independently():
text = f"alpha {_marker('a')} beta {_marker('b')} gamma {_marker('c')} end"
citations = [
{"source_id": "a", "url": "https://example.com/a"},
{"source_id": "b", "url": "https://example.com/b"},
{"source_id": "c", "url": "https://example.com/c"},
]
out = _replace_openai_citation_markers(text, citations)
assert out == (
"alpha [[1]](https://example.com/a) beta "
"[[2]](https://example.com/b) gamma "
"[[3]](https://example.com/c) end"
)
# ---------------------------------------------------------------------------
# 6. Idempotency.
# ---------------------------------------------------------------------------
def test_rewriter_idempotent_on_already_rewritten_text():
"""Running the rewriter twice does not double-link or corrupt brackets."""
text = f"alpha {_marker('a')} omega"
citations = [{"source_id": "a", "url": "https://example.com/a"}]
once = _replace_openai_citation_markers(text, citations)
twice = _replace_openai_citation_markers(once, citations)
assert once == twice
assert _no_private_use(once)
def test_rewriter_idempotent_on_marker_free_text():
"""No-op when there is nothing to rewrite."""
text = "Plain prose with no citations and no private-use bytes."
out = _replace_openai_citation_markers(text, [])
assert out is text or out == text
# ---------------------------------------------------------------------------
# 7. Edge / robustness.
# ---------------------------------------------------------------------------
def test_only_marker_no_surrounding_text():
"""A delta that is JUST a marker (no prose) still renders correctly;
used to leak without the empty-string short-circuit in the split helper."""
text = _marker("solo")
citations = [{"source_id": "solo", "url": "https://solo.example"}]
out = _replace_openai_citation_markers(text, citations)
assert out == "[[1]](https://solo.example)"
def test_back_to_back_markers_with_no_separator():
"""Adjacent markers resolve to concatenated links, no joining whitespace."""
text = f"{_marker('x')}{_marker('y')}"
citations = [
{"source_id": "x", "url": "https://x.example"},
{"source_id": "y", "url": "https://y.example"},
]
out = _replace_openai_citation_markers(text, citations)
assert out == "[[1]](https://x.example)[[2]](https://y.example)"
def test_split_helper_buffers_only_after_last_open_byte():
"""A complete marker followed by an unterminated one: head includes
the complete marker, buffer holds only the trailing partial."""
complete = _marker("done")
partial = f"{CITE_START}cite{CITE_DELIM}half" # no STOP
text = f"pre {complete} mid {partial}"
head, tail = _split_pending_citation_tail(text)
assert head == f"pre {complete} mid "
assert tail == partial
# And the head, once rewritten, drops every private-use byte.
rewritten = _replace_openai_citation_markers(
head, [{"source_id": "done", "url": "https://d"}]
)
assert rewritten == "pre [[1]](https://d) mid "
def test_split_helper_empty_input():
head, tail = _split_pending_citation_tail("")
assert head == "" and tail == ""
def test_split_helper_no_open_byte():
head, tail = _split_pending_citation_tail("nothing to see here")
assert head == "nothing to see here" and tail == ""
def test_split_helper_complete_marker_only():
"""A delta ending with a closed marker leaves the buffer empty."""
text = f"alpha {_marker('a')}"
head, tail = _split_pending_citation_tail(text)
assert head == text and tail == ""
# ---------------------------------------------------------------------------
# 8. Sources-panel: marker drop must not affect citation aggregation.
# Indices come from the url_citations list, not the marker stream.
# ---------------------------------------------------------------------------
def test_unknown_marker_does_not_perturb_citation_indexing():
"""Unknown source_id markers drop without consuming an index slot."""
text = f"A {_marker('unknown')} B {_marker('real_a')} C {_marker('real_b')}"
citations = [
{"source_id": "real_a", "url": "https://example.com/a"},
{"source_id": "real_b", "url": "https://example.com/b"},
]
out = _replace_openai_citation_markers(text, citations)
# real_a is index 1; unknown does not take a slot.
assert "[[1]](https://example.com/a)" in out
assert "[[2]](https://example.com/b)" in out
assert _no_private_use(out)
# ---------------------------------------------------------------------------
# Regression: unterminated marker tail must NOT leak the residual
# ``cite``-prefixed source id as plain text. PR #5713 audit P1.
# ---------------------------------------------------------------------------
def test_unterminated_marker_does_not_leak_cite_residue():
"""Stream ends mid-marker: drop the whole tail rather than strip
codepoints and leave ``cite<sid>`` behind."""
half = f"Hi there {CITE_START}cite{CITE_DELIM}turn0view0"
out = _simulate_delta_stream([half], [], flush = True)
# Prose before the marker stays; no private-use bytes or cite residue.
assert "Hi there" in out
assert _no_private_use(out)
assert "citeturn0view0" not in out
assert "cite" not in out.split("Hi there", 1)[1]
def test_unterminated_marker_only_no_prefix_drops_entirely():
"""A delta that is purely an unterminated marker flushes to ""."""
half = f"{CITE_START}cite{CITE_DELIM}turn0view0"
out = _simulate_delta_stream([half], [], flush = True)
assert out == ""
def test_unterminated_marker_with_prefix_emits_only_prefix():
"""Prose then unterminated marker: prose emits, marker remnant drops."""
half = f"prefix prose {CITE_START}cite{CITE_DELIM}abc"
out = _simulate_delta_stream([half], [], flush = True)
assert out == "prefix prose "
def test_closing_byte_arrives_after_pending_buffered_split():
"""Closing byte arrives in a later delta after opener + source id were
buffered; link resolves with no residue."""
cuts = [
f"a {CITE_START}cite{CITE_DELIM}",
f"sid{CITE_STOP} b",
]
out = _simulate_delta_stream(
cuts,
[{"source_id": "sid", "url": "https://example.com/x"}],
flush = True,
)
assert "[[1]](https://example.com/x)" in out
assert "a " in out and "b" in out
assert _no_private_use(out)
assert "citesid" not in out