fix: isolate citation registries per pipeline

[ump45nose]

Summary

• isolate citation state behind a unique registry ID for each LightResearch pipeline run
• pass that ID through the generated tool metadata and release the registry when the pipeline completes
• add a focused regression test that interleaves two runs and verifies their citation counters remain independent

Root cause

init_citation_registry() reset a class-level dictionary shared by every request. Starting a second pipeline therefore erased the first pipeline's in-progress citation mappings.

Validation

• pytest -p no:cacheprovider tests/test_citation_registry.py -q
• ultrarag build examples/demos/LightResearch.yaml
• generated MCP metadata smoke for init, assign, and clear tool contracts
• ruff check --select I,F,E9 tests/test_citation_registry.py
• ruff check --select I servers/custom/src/custom.py

Fixes #394

[ump45nose]

@xhd0728 This focused LightResearch concurrency fix is ready for review. It scopes citation registries by pipeline-run ID, carries that ID through generated tool metadata, and releases the registry on completion; the regression interleaves two runs and verifies independent counters. The focused pytest, pipeline build, MCP metadata smoke, and targeted Ruff checks pass. Since you maintain and review servers/custom/src/custom.py and the related demo pipelines, could you take a look when convenient?

[xhd0728]

Thanks for the fix! The registry isolation works in our interleaved-request tests, but further testing found two issues:

1. The final answer is lost from the pipeline return value. Adding clear_citation_registry as the last step makes PipelineCall.final_result return None, even though the answer was generated successfully.
2. Cleanup is skipped on failure. If generation or retrieval raises before the last step, the registry remains in memory. Repeated failures in a reused server can accumulate these entries.

Could you move cleanup into a guaranteed lifecycle/finally path that preserves the generated result and handles failures and cancellation? Please also add regression tests for these cases before merging.

[MestreY0d4-Uninter]

Chiming in from #408 (the thread-local attempt this PR supersedes) — the two issues @xhd0728 found both trace to the same design choice: the registry state lives in a class-level global (CitationRegistry._instances) while only the small citation_registry_id handle flows through the pipeline context. That forces cleanup to be expressed as a pipeline step, which is exactly what breaks final_result (issue 1) and gets skipped on mid-pipeline failure (issue 2).

Two ways to remove the global-lifecycle problem entirely:

Option A — carry the state in the pipeline context (no global at all). The registry state is just {doc_hash: id, counter} per query index — JSON-serializable. assign_citation_ids_stateful could return the state as an output variable (citation_state) that the next call feeds back in via the YAML input mapping. Per-request isolation then comes for free from the pipeline's own per-run variable dict: nothing shared across requests, nothing to clean up on failure, and no trailing step so generation.generate stays last and final_result is preserved.

Option B — keep the registry-id design, replace the trailing step with allocation-time GC. Add a created_at timestamp per registry in create(), purge entries older than a TTL (and/or cap the number of live registries) at the top of create(), and drop custom.clear_citation_registry from the YAML. Cleanup then runs on a guaranteed path (every run allocates, so every run collects), memory stays bounded even when generation raises mid-pipeline, and the pipeline's return value is untouched.

Either way, the regression tests that would pin this down: (1) interleaved runs keep independent counters (already present), (2) a pipeline that raises after assign_citation_ids_stateful leaves no unbounded growth in a reused server, (3) final_result still carries the generation output on the happy path, (4) cancellation mid-run behaves like failure.

Happy to test either approach against the interleaved MCP-client scenario from the maintainer testing on this thread, or to send a patch if useful — we have the failing repro from #394 wired up locally.

[xhd0728]

Thanks for the suggestions. I traced this change through the current UI demo execution and citation-rendering paths. There is an additional UI regression beyond the pipeline return-value issue.

Citation data remains compatible

The registry isolation itself looks sound. citation_registry_id is an internal handle, while assign_citation_ids_stateful still returns passages in the existing format:

[1] Document A
[2] Document B

The pipeline executor extracts those IDs and emits sources events. The frontend stores the source documents in the message metadata and uses them to resolve citations. Clearing the server-side registry does not remove source documents that have already been sent to the frontend.

So the registry isolation does not inherently break citation rendering. The problem is the trailing cleanup step and its effect on answer delivery.

The cleanup step changes how the UI handles the final answer

This PR changes the end of LightResearch to:

- prompt.webnote_gen_answer
- generation.generate
- custom.clear_citation_registry

In src/ultrarag/client.py, is_final_step is determined by whether any non-memory-save steps remain. Since custom.clear_citation_registry follows generation, the final answer’s tokens are now emitted with is_final: false.

In ui/frontend/src/pages/ChatPage.tsx, tokens with is_final: false are appended to the thinking steps rather than the assistant’s answer text.

This means the final answer—including its citation markers—no longer streams into the normal answer body.

The final event does not reliably recover the answer

The cleanup tool returns {}, and its result replaces the generation result returned by the executor.

The UI backend then chooses the answer using:

answer = final_ans or mem_ans or "No answer generated"

There are two relevant cases:

• If final_result is None, the backend may recover the answer from the memory snapshot. That can restore the completed answer, but it does not restore normal streaming behavior.
• If final_result is {} or "{}", _extract_result() returns the non-empty string "{}". This prevents the memory fallback from being used.

The frontend subsequently replaces the answer body with the answer from the final event. In the latter case, users can end up with a source list but a body containing only {}, with no answer citations.

Focused reproduction

I reproduced the relevant execution path using the repository’s actual executor function and this PR’s citation functions, with external generation and MCP calls replaced by test doubles.

Behavior	Without trailing cleanup	With trailing cleanup
Emitted source IDs and contents	Correct	Identical
Final generation tokens	is_final: true	is_final: false
Extracted UI answer when cleanup returns a result with data={}	Generated answer with citations	"{}"

This is a focused code-path reproduction, not a full browser end-to-end test with live retrieval and generation services.

Suggested direction

Of the proposed alternatives, I would prefer Option A: carry citation state in the pipeline context.

The state can remain scoped to each execution, while the tool continues returning numbered passages in the existing format. This removes the shared registry lifecycle problem and allows generation.generate to remain the final pipeline step.

A few compatibility details would need to be preserved:

• Citation IDs must remain stable across retrieval iterations within one run.
• Interleaved runs must remain independent.
• State must flow correctly through the loop and branch input/output mappings.
• The numbered ret_psg output and sources event contract must remain compatible with the frontend. Source extraction currently also uses the tool name containing "citation" to identify citation tools.

A lifecycle cleanup solution could also work, provided cleanup is outside the normal answer-producing step sequence and runs on success, failure, and cancellation without replacing the answer.

I would be more cautious about Option B: TTL-based cleanup. Age-based eviction could remove a registry that a long-running pipeline is still using, and allocation-triggered cleanup only runs when another allocation occurs. It would need explicit handling of active registries and resource bounds rather than relying on TTL alone.

Validation before merging

Could we add coverage for:

1. Interleaved runs retaining independent citation counters.
2. Stable IDs for repeated documents across loop iterations.
3. Final generation tokens retaining is_final: true.
4. The final pipeline result preserving the generated answer.
5. The UI’s final event preserving that answer instead of replacing it with an empty cleanup result.
6. Citation markers resolving to the correct source documents after completion.
7. Failure and cancellation leaving no unbounded registry accumulation.

The isolation fix is valuable, but I would hold off on merging the current version. Preserving final_result alone is not sufficient: the final-token classification also needs to remain correct for UI demo streaming and citation display.