fix: citations

private_gpt/components/chat/processors/events/citations/citations.py

@@ -144,7 +144,7 @@ async def process_citations(
                         delta = TextDelta.from_citations(delta_text, delta_citation)
 
                     yield RawContentBlockDeltaEvent(block_id=event.block_id, delta=delta)
-                    
+
                     send_text = cleaned_text
                     send_citations.extend(delta_citation)
 

private_gpt/components/engines/citations/parser.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from private_gpt.components.engines.citations.types import Document
+from private_gpt.components.text_processing import (
+    BacktickUnwrapRule,
+    DelimitedReferenceRule,
+    IncrementalTextProcessor,
+    ProcessingContext,
+)
+
+CitationFormatter = Callable[[int, Document, int], str]
+
+
+class CitationTextParser:
+    def __init__(
+        self,
+        documents: list[Document],
+        formatter: CitationFormatter,
+        *,
+        start_token: str,
+        end_token: str,
+        separator: str,
+        identifier_length: int,
+        citation_indices: dict[str, int] | None = None,
+    ) -> None:
+        self._documents_by_id = {
+            document.id.lower(): document for document in documents
+        }
+        self._start_token = start_token
+        self._end_token = end_token
+        self._formatter = formatter
+        self._initial_indices = dict(citation_indices or {})
+        self._context = ProcessingContext(
+            state={
+                "citation_indices": dict(self._initial_indices),
+                "citation_next_index": max(self._initial_indices.values(), default=-1)
+                + 1,
+                "citation_occurrence": 0,
+            }
+        )
+        reference_rule = DelimitedReferenceRule(
+            start_token=start_token,
+            end_token=end_token,
+            separator=separator,
+            resolve=self._resolve,
+            render=self._render,
+        )
+        self._processor = IncrementalTextProcessor(
+            [
+                BacktickUnwrapRule(reference_rule),
+                reference_rule,
+            ]
+        )
+
+    def parse(self, text: str, *, final: bool = False) -> tuple[str, dict[str, int]]:
+        normalized = text.replace("【", self._start_token).replace(
+            "】", self._end_token
+        )
+        result = self._processor.process(
+            normalized,
+            final=final,
+            context=self._context,
+        )
+        return result.text, dict(self._context.state["citation_indices"])
+
+    def _resolve(
+        self, identifiers: list[str], context: ProcessingContext
+    ) -> list[Document]:
+        return [
+            self._documents_by_id[identifier.lower()]
+            for identifier in identifiers
+            if identifier.lower() in self._documents_by_id
+        ]
+
+    def _render(
+        self, documents: list[Document], context: ProcessingContext
+    ) -> tuple[str, tuple[Document, ...]]:
+        indices: dict[str, int] = context.state["citation_indices"]
+        next_index: int = context.state["citation_next_index"]
+        occurrence: int = context.state["citation_occurrence"]
+        rendered = []
+
+        for document in documents:
+            if document.id_ in indices:
+                index = indices[document.id_]
+            else:
+                index = next_index
+                indices[document.id_] = index
+                next_index += 1
+            rendered.append(self._formatter(occurrence, document, index))
+            occurrence += 1
+
+        context.state["citation_next_index"] = next_index
+        context.state["citation_occurrence"] = occurrence
+        return ",".join(rendered), tuple(documents)

private_gpt/components/engines/citations/utils.py

@@ -12,6 +12,7 @@ from llama_index.core.schema import MetadataMode, NodeWithScore
 from private_gpt.components.chat.processors.chat_history.memory.utils.splitting import (
     get_user_blocks,
 )
+from private_gpt.components.engines.citations.parser import CitationTextParser
 from private_gpt.components.engines.citations.types import Citation, Document
 from private_gpt.components.ingest.metadata_helper import (
     MetadataFlags,
@@ -242,8 +243,6 @@ def _extract_citations_from_text(
     return cites
 
 
-
-
 def extract_citations_by_original_text(
     text: str,
     documents: list[Document],
@@ -254,187 +253,17 @@ def extract_citations_by_original_text(
     citation_indices: dict[str, int] | None = None,
     is_final: bool = False,
 ) -> tuple[str, list[Citation], dict[str, int]]:
-    # Initialize an empty string to store the cleaned text
-    citation_indices = citation_indices or {}
-
-    result = ""
-    start_len = len(start_token)
-    end_len = len(end_token)
-
-    # Model can generate brackets not normalized 【
-    text = text.replace("【", start_token).replace("】", end_token)
-
-    # Iterate through the text to remove malformed citations and save correct citations.
-    # Backticks directly wrapping a citation are formatting noise and are not emitted.
-    i = 0
-    docs = []
-    citation_placeholders: list[str] = []
-    code_delimiter: str | None = None
-    citation_wrapper_delimiter: str | None = None
-
-    while i < len(text):
-        if text[i] == "`":
-            delimiter_end = i + 1
-            while delimiter_end < len(text) and text[delimiter_end] == "`":
-                delimiter_end += 1
-            delimiter = text[i:delimiter_end]
-
-            if citation_wrapper_delimiter == delimiter:
-                citation_wrapper_delimiter = None
-                i = delimiter_end
-                continue
-
-            if code_delimiter == delimiter:
-                result += delimiter
-                code_delimiter = None
-                i = delimiter_end
-                continue
-
-            if delimiter_end == len(text):
-                if is_final:
-                    result += delimiter
-                break
-
-            if text[delimiter_end : delimiter_end + start_len] == start_token:
-                citation_start = delimiter_end
-                citation_end = citation_start + start_len
-                while (
-                    citation_end < len(text)
-                    and text[citation_end : citation_end + end_len] != end_token
-                ):
-                    citation_end += 1
-
-                if citation_end >= len(text):
-                    break
-
-                node_ids = [
-                    node_id.strip()
-                    for node_id in text[
-                        citation_start + start_len : citation_end
-                    ].split(split_token)
-                ]
-                valid_docs = [
-                    doc
-                    for node_id in node_ids
-                    if (
-                        doc := next(
-                            (
-                                document
-                                for document in documents
-                                if document.id.lower() == node_id.lower()
-                            ),
-                            None,
-                        )
-                    )
-                ]
-                if valid_docs:
-                    placeholders = []
-                    for doc in valid_docs:
-                        placeholder_index = "".join(
-                            f"n{digit}" for digit in str(len(docs))
-                        )
-                        placeholder = f"\ue000citation{placeholder_index}\ue001"
-                        docs.append(doc)
-                        citation_placeholders.append(placeholder)
-                        placeholders.append(placeholder)
-                    result += split_token.join(placeholders)
-                    i = citation_end + end_len
-                    if text[i : i + len(delimiter)] == delimiter:
-                        i += len(delimiter)
-                    else:
-                        citation_wrapper_delimiter = delimiter
-                    continue
-
-            result += delimiter
-            code_delimiter = delimiter
-            i = delimiter_end
-        elif text[i : i + start_len] == start_token:
-            # Check if we have a complete citation
-            j = i + start_len
-            while j < len(text) and text[j : j + end_len] != end_token:
-                j += 1
-
-            if j < len(text) and text[j : j + end_len] == end_token:
-                node_ids = [
-                    id.strip() for id in text[i + start_len : j].split(split_token)
-                ]
-                valid_docs = []
-                for node_id in node_ids:
-                    doc = next(
-                        (doc for doc in documents if doc.id.lower() == node_id.lower()),
-                        None,
-                    )
-                    if doc:
-                        valid_docs.append(doc)
-                if valid_docs:
-                    placeholders = []
-                    for doc in valid_docs:
-                        placeholder_index = "".join(
-                            f"n{digit}" for digit in str(len(docs))
-                        )
-                        placeholder = f"\ue000citation{placeholder_index}\ue001"
-                        docs.append(doc)
-                        citation_placeholders.append(placeholder)
-                        placeholders.append(placeholder)
-                    result += split_token.join(placeholders)
-                    i = j + end_len
-                else:
-                    # No valid docs in citation, treat as regular text
-                    result += text[i : j + end_len]
-                    i = j + end_len
-            else:
-                # Incomplete citation detected: drop the
-                # citation content and stop processing further.
-                # This change fixes the issue by not
-                # appending any incomplete citation text.
-                i = len(text)
-                continue
-        else:
-            result += text[i]
-            i += 1
-
-    # Process citations
-    pattern = re.compile(
-        rf"{re.escape(start_token)}?[A-Z0-9]{shorter_id_length}{re.escape(end_token)}?"
+    parser = CitationTextParser(
+        documents,
+        format_cite,
+        start_token=start_token,
+        end_token=end_token,
+        separator=split_token,
+        identifier_length=shorter_id_length,
+        citation_indices=citation_indices,
     )
-    for match in pattern.finditer(result):
-        # If citation is well-formed, skip
-        word = match.group(0) if match.groups() else ""
-        if not word or (word.startswith(start_token) and word.endswith(end_token)):
-            continue
-
-        # Try to find the related document, if no document found, skip
-        doc = next((doc for doc in documents if doc.id in word), None)
-        if not doc:
-            continue
-
-        # Remove doc reference
-        new_token = word.replace(doc.id, "", 1).lstrip().rstrip()
-        result = result[: match.start()] + new_token + result[match.end() :]
-
-    # Replace citations with sequential numbers, just first occurrence
-    max_index = max(citation_indices.values(), default=-1)
-    current_index = max_index + 1
-    processed_docs = []
-    for i, (doc, placeholder) in enumerate(
-        zip(docs, citation_placeholders, strict=True)
-    ):
-        if doc.id_ not in processed_docs:
-            processed_docs.append(doc.id_)
-
-        if doc.id_ in citation_indices:
-            # If we already have this document, use the existing index
-            index = citation_indices[doc.id_]
-        else:
-            # Otherwise, assign a new index
-            index = current_index
-            current_index += 1
-        citation_indices[doc.id_] = index
-
-        citation = format_cite(i, doc, index)
-        result = result.replace(placeholder, citation, 1)
-
-    return result, _extract_citations_from_text(result), citation_indices
+    result, updated_indices = parser.parse(text, final=is_final)
+    return result, _extract_citations_from_text(result), updated_indices
 
 
 async def deduplicate_documents_in_history(

private_gpt/components/text_processing/__init__.py

@@ -0,0 +1,29 @@
+from private_gpt.components.text_processing.engine import IncrementalTextProcessor
+from private_gpt.components.text_processing.models import (
+    Action,
+    ProbeResult,
+    ProbeStatus,
+    ProcessDelta,
+    ProcessingContext,
+    ProcessResult,
+)
+from private_gpt.components.text_processing.rules import (
+    BacktickUnwrapRule,
+    DelimitedReferenceRule,
+    LooseReferenceCleanupRule,
+    StreamRule,
+)
+
+__all__ = [
+    "Action",
+    "BacktickUnwrapRule",
+    "DelimitedReferenceRule",
+    "IncrementalTextProcessor",
+    "LooseReferenceCleanupRule",
+    "ProbeResult",
+    "ProbeStatus",
+    "ProcessDelta",
+    "ProcessResult",
+    "ProcessingContext",
+    "StreamRule",
+]

private_gpt/components/text_processing/engine.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import TYPE_CHECKING
+
+from private_gpt.components.text_processing.models import (
+    Action,
+    ProbeStatus,
+    ProcessDelta,
+    ProcessingContext,
+    ProcessResult,
+)
+
+if TYPE_CHECKING:
+    from private_gpt.components.text_processing.rules import StreamRule
+
+
+class IncrementalTextProcessor:
+    def __init__(
+        self,
+        rules: list[StreamRule],
+        initial_state: dict[str, object] | None = None,
+    ) -> None:
+        self._rules = sorted(rules, key=lambda rule: rule.priority, reverse=True)
+        self._initial_state = deepcopy(initial_state or {})
+        self._source = ""
+        self._emitted = ""
+        self._metadata_count = 0
+        self.context = ProcessingContext(state=deepcopy(self._initial_state))
+
+    def process(
+        self,
+        text: str,
+        *,
+        final: bool = False,
+        context: ProcessingContext | None = None,
+    ) -> ProcessResult:
+        active_context = context or ProcessingContext()
+        active_context.final = final
+        output: list[str] = []
+        metadata: list[object] = []
+        cursor = 0
+
+        while cursor < len(text):
+            probe = None
+            for rule in self._rules:
+                candidate = rule.probe(text, cursor, active_context)
+                if candidate.status != ProbeStatus.NO_MATCH:
+                    probe = candidate
+                    break
+
+            if probe is None:
+                output.append(text[cursor])
+                cursor += 1
+                continue
+
+            if probe.status == ProbeStatus.NEED_MORE:
+                break
+            if probe.consumed <= 0:
+                raise ValueError("A matching stream rule must consume source text")
+
+            source = text[cursor : cursor + probe.consumed]
+            if probe.action == Action.PASS:
+                output.append(
+                    probe.replacement if probe.replacement is not None else source
+                )
+            elif probe.action in (Action.REPLACE, Action.UNWRAP):
+                output.append(probe.replacement or "")
+            elif probe.action == Action.DROP:
+                pass
+            else:
+                raise ValueError(f"Unsupported matching action: {probe.action}")
+
+            for key in probe.state_deletes:
+                active_context.state.pop(key, None)
+            active_context.state.update(probe.state_updates)
+            metadata.extend(probe.metadata)
+            cursor += probe.consumed
+
+        return ProcessResult(
+            text="".join(output),
+            metadata=tuple(metadata),
+            pending=text[cursor:],
+            consumed=cursor,
+        )
+
+    def feed(self, chunk: str) -> ProcessDelta:
+        self._source += chunk
+        self.context = ProcessingContext(state=deepcopy(self._initial_state))
+        result = self.process(self._source, context=self.context)
+        if not result.text.startswith(self._emitted):
+            raise ValueError("A stream rule rewrote an already-emitted prefix")
+        delta = ProcessDelta(
+            text=result.text[len(self._emitted) :],
+            metadata=result.metadata[self._metadata_count :],
+            pending=result.pending,
+        )
+        self._emitted = result.text
+        self._metadata_count = len(result.metadata)
+        return delta
+
+    def finalize(self) -> ProcessDelta:
+        self.context = ProcessingContext(state=deepcopy(self._initial_state))
+        result = self.process(self._source, final=True, context=self.context)
+        if not result.text.startswith(self._emitted):
+            raise ValueError("Finalization rewrote an already-emitted prefix")
+        delta = ProcessDelta(
+            text=result.text[len(self._emitted) :],
+            metadata=result.metadata[self._metadata_count :],
+            pending=result.pending,
+        )
+        self._emitted = result.text
+        self._metadata_count = len(result.metadata)
+        return delta

private_gpt/components/text_processing/models.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum, auto
+from typing import Any
+
+
+class Action(Enum):
+    PASS = auto()
+    DROP = auto()
+    UNWRAP = auto()
+    REPLACE = auto()
+    HOLD = auto()
+
+
+class ProbeStatus(Enum):
+    NO_MATCH = auto()
+    MATCH = auto()
+    NEED_MORE = auto()
+
+
+@dataclass(frozen=True)
+class ProbeResult:
+    status: ProbeStatus
+    consumed: int = 0
+    action: Action = Action.PASS
+    replacement: str | None = None
+    metadata: tuple[Any, ...] = ()
+    state_updates: dict[str, Any] = field(default_factory=dict)
+    state_deletes: tuple[str, ...] = ()
+
+    @classmethod
+    def no_match(cls) -> ProbeResult:
+        return cls(status=ProbeStatus.NO_MATCH)
+
+    @classmethod
+    def need_more(cls) -> ProbeResult:
+        return cls(status=ProbeStatus.NEED_MORE, action=Action.HOLD)
+
+
+@dataclass
+class ProcessingContext:
+    final: bool = False
+    state: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class ProcessResult:
+    text: str
+    metadata: tuple[Any, ...]
+    pending: str
+    consumed: int
+
+
+@dataclass(frozen=True)
+class ProcessDelta:
+    text: str
+    metadata: tuple[Any, ...]
+    pending: str

private_gpt/components/text_processing/rules.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from private_gpt.components.text_processing.models import (
+    Action,
+    ProbeResult,
+    ProbeStatus,
+    ProcessingContext,
+)
+
+
+class StreamRule(Protocol):
+    name: str
+    priority: int
+
+    def probe(
+        self, text: str, position: int, context: ProcessingContext
+    ) -> ProbeResult: ...
+
+
+ResolveReferences = Callable[[list[str], ProcessingContext], list[Any]]
+RenderReferences = Callable[[list[Any], ProcessingContext], tuple[str, tuple[Any, ...]]]
+
+
+@dataclass
+class DelimitedReferenceRule:
+    start_token: str
+    end_token: str
+    separator: str
+    resolve: ResolveReferences
+    render: RenderReferences
+    name: str = "delimited_reference"
+    priority: int = 100
+
+    def probe(
+        self, text: str, position: int, context: ProcessingContext
+    ) -> ProbeResult:
+        if not text.startswith(self.start_token, position):
+            return ProbeResult.no_match()
+
+        end = text.find(self.end_token, position + len(self.start_token))
+        if end == -1:
+            return ProbeResult.need_more()
+
+        consumed = end + len(self.end_token) - position
+        identifiers = [
+            identifier.strip()
+            for identifier in text[position + len(self.start_token) : end].split(
+                self.separator
+            )
+        ]
+        references = self.resolve(identifiers, context)
+        if not references:
+            return ProbeResult(
+                status=ProbeStatus.MATCH,
+                consumed=consumed,
+                action=Action.PASS,
+            )
+
+        replacement, metadata = self.render(references, context)
+        return ProbeResult(
+            status=ProbeStatus.MATCH,
+            consumed=consumed,
+            action=Action.REPLACE,
+            replacement=replacement,
+            metadata=metadata,
+        )
+
+
+@dataclass
+class BacktickUnwrapRule:
+    inner: StreamRule
+    name: str = "backtick_unwrap"
+    priority: int = 200
+    code_state_key: str = "backtick_code_delimiter"
+    wrapper_state_key: str = "backtick_wrapper_delimiter"
+
+    def probe(
+        self, text: str, position: int, context: ProcessingContext
+    ) -> ProbeResult:
+        if text[position] != "`":
+            return ProbeResult.no_match()
+
+        delimiter_end = position + 1
+        while delimiter_end < len(text) and text[delimiter_end] == "`":
+            delimiter_end += 1
+        delimiter = text[position:delimiter_end]
+
+        if context.state.get(self.wrapper_state_key) == delimiter:
+            return ProbeResult(
+                status=ProbeStatus.MATCH,
+                consumed=len(delimiter),
+                action=Action.DROP,
+                state_deletes=(self.wrapper_state_key,),
+            )
+
+        if context.state.get(self.code_state_key) == delimiter:
+            return ProbeResult(
+                status=ProbeStatus.MATCH,
+                consumed=len(delimiter),
+                action=Action.PASS,
+                state_deletes=(self.code_state_key,),
+            )
+
+        if delimiter_end == len(text):
+            if not context.final:
+                return ProbeResult.need_more()
+            return ProbeResult(
+                status=ProbeStatus.MATCH,
+                consumed=len(delimiter),
+                action=Action.PASS,
+            )
+
+        inner_match = self.inner.probe(text, delimiter_end, context)
+        if inner_match.status == ProbeStatus.NEED_MORE:
+            return inner_match
+        if (
+            inner_match.status == ProbeStatus.MATCH
+            and inner_match.action == Action.REPLACE
+        ):
+            consumed = len(delimiter) + inner_match.consumed
+            updates = dict(inner_match.state_updates)
+            deletes = list(inner_match.state_deletes)
+            if text.startswith(delimiter, position + consumed):
+                consumed += len(delimiter)
+            else:
+                updates[self.wrapper_state_key] = delimiter
+            return ProbeResult(
+                status=ProbeStatus.MATCH,
+                consumed=consumed,
+                action=Action.UNWRAP,
+                replacement=inner_match.replacement,
+                metadata=inner_match.metadata,
+                state_updates=updates,
+                state_deletes=tuple(deletes),
+            )
+
+        return ProbeResult(
+            status=ProbeStatus.MATCH,
+            consumed=len(delimiter),
+            action=Action.PASS,
+            state_updates={self.code_state_key: delimiter},
+        )
+
+
+@dataclass
+class LooseReferenceCleanupRule:
+    start_token: str
+    end_token: str
+    identifier_length: int
+    identifiers: tuple[str, ...]
+    name: str = "loose_reference_cleanup"
+    priority: int = 50
+
+    def __post_init__(self) -> None:
+        self._pattern = re.compile(
+            rf"{re.escape(self.start_token)}?[A-Z0-9]"
+            rf"{{{self.identifier_length}}}{re.escape(self.end_token)}?"
+        )
+
+    def probe(
+        self, text: str, position: int, context: ProcessingContext
+    ) -> ProbeResult:
+        match = self._pattern.match(text, position)
+        if match is None:
+            return ProbeResult.no_match()
+        word = match.group(0)
+        if word.startswith(self.start_token) and word.endswith(self.end_token):
+            return ProbeResult.no_match()
+        identifier = next(
+            (identifier for identifier in self.identifiers if identifier in word),
+            None,
+        )
+        if identifier is None:
+            return ProbeResult.no_match()
+        replacement = word.replace(identifier, "", 1).strip()
+        return ProbeResult(
+            status=ProbeStatus.MATCH,
+            consumed=len(word),
+            action=Action.REPLACE,
+            replacement=replacement,
+        )

tests/components/text_processing/test_engine.py

@@ -0,0 +1,131 @@
+from dataclasses import dataclass
+
+from private_gpt.components.text_processing import (
+    Action,
+    IncrementalTextProcessor,
+    ProbeResult,
+    ProbeStatus,
+    ProcessingContext,
+)
+
+
+@dataclass
+class TokenRule:
+    token: str
+    action: Action
+    replacement: str | None = None
+    priority: int = 100
+    name: str = "token"
+
+    def probe(
+        self, text: str, position: int, context: ProcessingContext
+    ) -> ProbeResult:
+        remaining = text[position:]
+        if self.token.startswith(remaining) and len(remaining) < len(self.token):
+            return ProbeResult.need_more()
+        if not text.startswith(self.token, position):
+            return ProbeResult.no_match()
+        return ProbeResult(
+            status=ProbeStatus.MATCH,
+            consumed=len(self.token),
+            action=self.action,
+            replacement=self.replacement,
+        )
+
+
+def test_processor_passes_literal_text_without_rules() -> None:
+    result = IncrementalTextProcessor([]).process("plain text")
+
+    assert result.text == "plain text"
+    assert result.pending == ""
+
+
+def test_processor_supports_pass_drop_unwrap_and_replace() -> None:
+    processor = IncrementalTextProcessor(
+        [
+            TokenRule("<pass>", Action.PASS),
+            TokenRule("<drop>", Action.DROP),
+            TokenRule("<unwrap>", Action.UNWRAP, "body"),
+            TokenRule("<replace>", Action.REPLACE, "replacement"),
+        ]
+    )
+
+    result = processor.process("<pass>|<drop>|<unwrap>|<replace>")
+
+    assert result.text == "<pass>||body|replacement"
+
+
+def test_need_more_holds_only_the_unresolved_suffix() -> None:
+    processor = IncrementalTextProcessor(
+        [TokenRule("<replace>", Action.REPLACE, "done")]
+    )
+
+    result = processor.process("safe <repl")
+
+    assert result.text == "safe "
+    assert result.pending == "<repl"
+    assert result.consumed == len("safe ")
+
+
+def test_finalization_can_resolve_rule_specific_partial_behavior() -> None:
+    processor = IncrementalTextProcessor(
+        [TokenRule("<replace>", Action.REPLACE, "done")]
+    )
+
+    result = processor.process("safe <repl", final=True)
+
+    assert result.text == "safe "
+    assert result.pending == "<repl"
+
+
+def test_higher_priority_rule_wins_at_same_position() -> None:
+    processor = IncrementalTextProcessor(
+        [
+            TokenRule("token", Action.REPLACE, "low", priority=10),
+            TokenRule("token", Action.REPLACE, "high", priority=20),
+        ]
+    )
+
+    assert processor.process("token").text == "high"
+
+
+def test_feed_emits_only_new_safe_text() -> None:
+    processor = IncrementalTextProcessor(
+        [TokenRule("<replace>", Action.REPLACE, "done")]
+    )
+
+    first = processor.feed("before <rep")
+    second = processor.feed("lace> after")
+
+    assert first.text == "before "
+    assert first.pending == "<rep"
+    assert second.text == "done after"
+    assert second.pending == ""
+
+
+@dataclass
+class StatefulRule:
+    name: str = "stateful"
+    priority: int = 100
+
+    def probe(
+        self, text: str, position: int, context: ProcessingContext
+    ) -> ProbeResult:
+        if not text.startswith("#", position):
+            return ProbeResult.no_match()
+        count = int(context.state.get("count", 0)) + 1
+        return ProbeResult(
+            status=ProbeStatus.MATCH,
+            consumed=1,
+            action=Action.REPLACE,
+            replacement=str(count),
+            state_updates={"count": count},
+        )
+
+
+def test_feed_reparses_from_clean_initial_rule_state() -> None:
+    processor = IncrementalTextProcessor([StatefulRule()])
+
+    assert processor.feed("#").text == "1"
+    assert processor.feed("#").text == "2"
+    assert processor.context.state == {"count": 2}

tests/engines/test_citation_chunk_boundaries.py

@@ -7,7 +7,6 @@ from private_gpt.components.engines.citations.utils import (
     extract_citations_by_original_text,
 )
 
-
 TEXT = (
     "Format: `[XXXX]`. Correct: `[AB12]`. "
     "Invalid: `[[AB12]]`, `(AB12)`. "

tests/engines/test_citation_parser_robustness.py

@@ -118,10 +118,6 @@ def test_incomplete_citation_is_withheld_with_no_false_citation() -> None:
     assert citations == []
 
 
-@pytest.mark.xfail(
-    strict=True,
-    reason="Model output can currently collide with the internal citation placeholder.",
-)
 def test_placeholder_like_model_output_does_not_capture_real_citation() -> None:
     document = create_document("AB12")
     model_text = "Literal \ue000citationn0\ue001 then [AB12]."
@@ -132,8 +128,7 @@ def test_placeholder_like_model_output_does_not_capture_real_citation() -> None:
     )
 
     assert formatted == (
-        "Literal \ue000citationn0\ue001 then "
-        f"{format_cite(0, document, 0)}."
+        f"Literal \ue000citationn0\ue001 then {format_cite(0, document, 0)}."
     )
     assert len(citations) == 1
 

tests/engines/test_processor_citations.py

@@ -181,10 +181,10 @@ async def test_process_citations_trailing_backtick_with_stop_event() -> None:
 
     # start, delta (without backtick), delta (with backtick), stop
     assert len(result) == 4
-    
+
     combined = ""
     for r in result:
         if isinstance(r, RawContentBlockDeltaEvent) and isinstance(r.delta, TextDelta):
             combined += r.delta.text or ""
-            
+
     assert combined == "This is a test with a trailing backtick `"