fix(task_executor): fix Langfuse flush/shutdown deadlock that freezes document parsing (#16502)

agent/canvas.py

@@ -35,6 +35,7 @@ from api.db.services.file_service import FileService
 from api.db.services.llm_service import LLMBundle
 from api.db.services.task_service import has_canceled
 from common.constants import LLMType
+from common.llm_request_context import set_llm_request_context, reset_llm_request_context
 from common.exceptions import TaskCanceledException
 from common.misc_utils import get_uuid, hash_str2int
 from common.token_utils import token_usage_sink, langfuse_run_attrs
@@ -438,6 +439,14 @@ class Canvas(Graph):
             _lf_attrs["session_id"] = str(_session_id)[:200]
         sink_token = token_usage_sink.set(self._run_token_usage)
         attrs_token = langfuse_run_attrs.set(_lf_attrs)
+        # Forward the originating session/user to upstream LLM providers (as the
+        # OpenAI `user` field) for the duration of this run, and reset afterwards so
+        # the value never leaks to later calls in the same task. Reuse the same
+        # session/user already derived above so both integrations stay consistent.
+        _req_ctx_token = set_llm_request_context(
+            session_id=_session_id,
+            user_id=_user_id,
+        )
         try:
             async for ev in self._run_impl(**kwargs):
                 yield ev
@@ -454,6 +463,7 @@ class Canvas(Graph):
             except ValueError:
                 logging.debug("Failed to reset Langfuse run attributes ContextVar", exc_info=True)
                 langfuse_run_attrs.set(None)
+            reset_llm_request_context(_req_ctx_token)
 
     async def _run_impl(self, **kwargs):
         self.globals["sys.date"] = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
@@ -539,11 +549,13 @@ class Canvas(Graph):
                     if use_async:
                         await cpn_obj.invoke_async(**(call_kwargs or {}))
                         return
-                    # run_in_executor does not propagate context variables; copy the
-                    # current context so the token usage sink / Langfuse attributes set
-                    # by run() remain visible to LLMBundle calls inside sync components.
-                    ctx = contextvars.copy_context()
-                    await loop.run_in_executor(self._thread_pool, lambda: ctx.run(partial(sync_fn, **(call_kwargs or {}))))
+                    # run_in_executor does not carry context variables into the worker
+                    # thread; copy the current context so the LLM request context (the
+                    # `user` forwarding), token usage sink, and Langfuse attributes set
+                    # by run() remain visible to sync components.
+                    bound_call = partial(sync_fn, **(call_kwargs or {}))
+                    call_ctx = contextvars.copy_context()
+                    await loop.run_in_executor(self._thread_pool, partial(call_ctx.run, bound_call))
 
             i = f
             while i < t:

api/apps/restful_apis/document_api.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 #
 from io import BytesIO
+from datetime import datetime
 import logging
 import json
 import os
@@ -1473,6 +1474,11 @@ def _run_sync(user_id: str, req):
             has_unfinished_task = any((task.progress or 0) < 1 for task in tasks)
             if str(doc.run) in [TaskStatus.RUNNING.value, TaskStatus.CANCEL.value] or has_unfinished_task:
                 cancel_all_task_of(doc_id)
+                # Append a "stopped by user" marker so the history is preserved and
+                # the document no longer looks like it is still waiting in the queue.
+                cancel_doc_msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task stopped by user."
+                info["progress_msg"] = (doc.progress_msg or "") + cancel_doc_msg
+                logging.debug("Appended cancellation marker to progress_msg on cancel for doc %s", doc_id)
             else:
                 return RetCode.DATA_ERROR, "Cannot cancel a task that is not in RUNNING status"
         if all([rerun_with_delete, str(doc.run) == TaskStatus.DONE.value]):
@@ -1698,14 +1704,17 @@ async def stop_parse_documents(tenant_id, dataset_id):
                     continue
 
                 cancel_all_task_of(doc_id)
+                cancel_doc_msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task stopped by user."
                 DocumentService.update_by_id(
                     doc_id,
                     {
                         "run": str(TaskStatus.CANCEL.value),
                         "progress": 0,
                         "chunk_num": 0,
+                        "progress_msg": (doc.progress_msg or "") + cancel_doc_msg,
                     },
                 )
+                logging.debug("Appended cancellation marker to progress_msg on stop-parse for doc %s", doc_id)
                 index_name = search.index_name(tenant_id)
                 if settings.docStoreConn.index_exist(index_name, doc.kb_id):
                     settings.docStoreConn.delete({"doc_id": doc.id}, index_name, doc.kb_id)

api/apps/restful_apis/task_api.py

@@ -89,7 +89,12 @@ async def _cancel_task(task_id):
         if doc_id and doc_id not in (CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID):
             _, doc = DocumentService.get_by_id(doc_id)
             if doc and str(doc.run) in (TaskStatus.RUNNING.value, TaskStatus.SCHEDULE.value):
-                DocumentService.update_by_id(doc_id, {"run": TaskStatus.CANCEL.value, "progress": 0})
+                cancel_doc_msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task stopped by user."
+                DocumentService.update_by_id(
+                    doc_id,
+                    {"run": TaskStatus.CANCEL.value, "progress": 0, "progress_msg": (doc.progress_msg or "") + cancel_doc_msg},
+                )
+                logging.debug("Appended cancellation marker to progress_msg on task cancel: task_id=%s doc_id=%s", task_id, doc_id)
     except Exception as e:
         logging.warning("Failed to update document run status for task %s: %s", task_id, str(e))
 

api/db/services/canvas_service.py

@@ -337,7 +337,7 @@ async def completion(tenant_id, agent_id, session_id=None, **kwargs):
         "files": files,
         "user_id": user_id,
         "inputs": inputs,
-        # Used by Canvas.run to correlate RAGFlow's Langfuse generations by session.
+        # Forwarded to upstream LLM providers as the `user` field for session correlation.
         "session_id": session_id,
     }
     if chat_template_kwargs is not None:

api/db/services/document_service.py

@@ -16,6 +16,7 @@
 import logging
 import random
 from datetime import datetime
+from time import monotonic
 
 import xxhash
 from peewee import fn, Case, JOIN
@@ -954,6 +955,10 @@ class DocumentService(CommonService):
                 doc_progress = doc.progress if doc and doc.progress else 0.0
                 special_task_running = False
                 priority = 0
+                # Count this document's own not-yet-started tasks per priority so
+                # they can be excluded from the "tasks ahead in the queue" figure
+                # for the matching priority queue.
+                own_queued_by_priority = {}
                 for t in tsks:
                     task_type = (t.task_type or "").lower()
                     if task_type in PIPELINE_SPECIAL_PROGRESS_FREEZE_TASK_TYPES:
@@ -962,6 +967,8 @@ class DocumentService(CommonService):
                         finished = False
                     if t.progress == -1:
                         bad += 1
+                    if (t.progress or 0) == 0:
+                        own_queued_by_priority[t.priority] = own_queued_by_priority.get(t.priority, 0) + 1
                     prg += t.progress if t.progress >= 0 else 0
                     if (t.progress_msg or "").strip():
                         msg.append(t.progress_msg)
@@ -991,9 +998,14 @@ class DocumentService(CommonService):
                 if msg:
                     info["progress_msg"] = msg
                     if msg.endswith("created task graphrag") or msg.endswith("created task raptor") or msg.endswith("created task mindmap"):
-                        info["progress_msg"] += "\n%d tasks are ahead in the queue..." % get_queue_length(priority)
+                        # Exclude this document's own queued tasks in the same
+                        # priority queue: they are not "ahead" of itself, they
+                        # ARE the work being waited on.
+                        queue_ahead = max(0, get_queue_length(priority) - own_queued_by_priority.get(priority, 0))
+                        info["progress_msg"] += "\n%d tasks are ahead in the queue..." % queue_ahead
                 else:
-                    info["progress_msg"] = "%d tasks are ahead in the queue..." % get_queue_length(priority)
+                    queue_ahead = max(0, get_queue_length(priority) - own_queued_by_priority.get(priority, 0))
+                    info["progress_msg"] = "%d tasks are ahead in the queue..." % queue_ahead
                 info["update_time"] = current_timestamp()
                 info["update_date"] = get_format_time()
                 (cls.model.update(info).where((cls.model.id == d["id"]) & ((cls.model.run.is_null(True)) | (cls.model.run != TaskStatus.CANCEL.value))).execute())
@@ -1165,8 +1177,79 @@ def queue_per_doc_raptor_task(doc, priority):
     return task["id"]
 
 
+# Short-lived per-priority cache for the genuine queued-task backlog so the
+# per-document progress sync does not issue a COUNT query for every document
+# each cycle. Keyed by priority (None means "all priorities").
+_PENDING_TASK_COUNT_CACHE = {}
+_PENDING_TASK_COUNT_TTL_SECONDS = 3.0
+
+
+def get_pending_task_count(priority=None):
+    """Count tasks that are genuinely still waiting to be processed.
+
+    A task counts as "waiting" when it has not started yet (progress == 0) and
+    its document is neither cancelled nor failed. We deliberately do NOT require
+    the document to be RUNNING with progress in [0, 1): special tasks (graphrag/
+    raptor/mindmap) are queued via ``begin2parse(keep_progress=True)`` while the
+    document's own progress may already be 1, so requiring RUNNING/progress<1
+    would undercount them and wrongly drop the cap to 0 while Redis lag is still
+    non-zero. Only cancelled documents (run == CANCEL) and failed ones
+    (progress < 0) are excluded, plus soft-deleted (invalid) documents.
+
+    When ``priority`` is given, only tasks queued at that priority are counted,
+    so the figure stays consistent with the per-priority Redis queue it caps.
+
+    Returns None when the count cannot be determined, so callers can fall back
+    to the raw Redis stream lag.
+    """
+    now = monotonic()
+    cached = _PENDING_TASK_COUNT_CACHE.get(priority)
+    if cached and cached.get("expire_at", 0.0) > now:
+        return cached["value"]
+    try:
+        query = (
+            Task.select(fn.COUNT(Task.id))
+            .join(Document, on=(Task.doc_id == Document.id))
+            .where(
+                (Task.progress == 0)
+                & ((Document.run.is_null(True)) | (Document.run != TaskStatus.CANCEL.value))
+                & (Document.progress >= 0)
+                & (Document.status == StatusEnum.VALID.value)
+            )
+        )
+        if priority is not None:
+            query = query.where(Task.priority == priority)
+        count = int(query.scalar() or 0)
+    except Exception:
+        logging.exception("get_pending_task_count failed")
+        return None
+    _PENDING_TASK_COUNT_CACHE[priority] = {"value": count, "expire_at": now + _PENDING_TASK_COUNT_TTL_SECONDS}
+    return count
+
+
 def get_queue_length(priority, suffix="common"):
+    """Return how many tasks are ahead in the processing queue.
+
+    The Redis stream consumer-group ``lag`` counts every message that has not
+    yet been delivered to a task executor, including messages whose tasks were
+    already cancelled/stopped. Those messages only stop counting once an
+    executor happens to read them, so after a user stops parsing the lag can
+    stay inflated indefinitely and parsing appears to hang forever
+    ("N tasks are ahead in the queue...").
+
+    To keep the figure honest, the raw lag is capped by the number of tasks
+    that are genuinely still waiting in the database, which self-heals the
+    moment work is cancelled or completes.
+    """
     group_info = REDIS_CONN.queue_info(settings.get_svr_queue_name(priority, suffix), SVR_CONSUMER_GROUP_NAME)
-    if not group_info:
+    lag = int(group_info.get("lag", 0) or 0) if group_info else 0
+
+    # Nothing queued in Redis: the answer is 0 regardless of the DB backlog, so
+    # short-circuit to avoid a COUNT/JOIN on every progress-sync cycle.
+    if lag <= 0:
         return 0
-    return int(group_info.get("lag", 0) or 0)
+
+    pending = get_pending_task_count(priority)
+    if pending is None:
+        return lag
+    return min(lag, pending)

api/db/services/tenant_llm_service.py

@@ -534,27 +534,34 @@ class LLM4Tenant:
     def close(self):
         """Release resources held by this LLM4Tenant instance.
 
-        This method should be called when the instance is no longer needed
-        to properly release resources such as:
-        - Langfuse tracing client (flush and shutdown)
-        - Underlying model instance resources (HTTP sessions, etc.)
+        IMPORTANT: do NOT call ``langfuse.flush()`` or ``langfuse.shutdown()``
+        here. ``close()`` runs once per task, synchronously, on the asyncio
+        event-loop thread of the task executor. Two problems follow:
+
+        - ``flush()`` blocks on an unbounded ``queue.join()`` in the underlying
+          OpenTelemetry span processor. If the exporter cannot drain (slow or
+          unreachable Langfuse, or an already-shutdown processor) it never
+          returns.
+        - ``shutdown()`` permanently tears down the process-wide Langfuse /
+          OpenTelemetry tracer provider that every ``LLMBundle`` shares. After
+          the first task shuts it down, every subsequent ``flush()`` blocks
+          forever.
+
+        Because this runs on the event loop, a single stuck ``flush()`` freezes
+        the entire task executor: all in-flight parse tasks stop making
+        progress and no new tasks are ever picked up (observed as document
+        parsing being stuck with every executor thread parked on a lock).
+
+        Langfuse already exports spans from its own background processor and
+        flushes at process exit, so releasing the reference is sufficient here.
         """
-        # Flush and shutdown Langfuse client if it was initialized
-        if self.langfuse:
-            try:
-                self.langfuse.flush()
-                if hasattr(self.langfuse, "shutdown"):
-                    self.langfuse.shutdown()
-            except Exception:
-                # Ignore errors during cleanup
-                pass
-            finally:
-                self.langfuse = None
+        # Drop the Langfuse reference WITHOUT flushing/shutting down the shared
+        # client (see the docstring above for why this would deadlock).
+        self.langfuse = None
 
         # Release underlying model instance if it has a close method
-        if self.mdl and hasattr(self.mdl, "close") and callable(getattr(self.mdl, "close")):
+        if self.mdl and callable(getattr(self.mdl, "close", None)):
             try:
                 self.mdl.close()
             except Exception:
-                # Ignore errors during cleanup
-                pass
+                logging.warning("LLM4Tenant.close: error while closing model instance", exc_info=True)

common/llm_request_context.py

@@ -0,0 +1,71 @@
+#
+#  Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+"""Per-request identifiers forwarded to upstream LLM providers.
+
+An agent run (or any LLM-issuing flow) installs the originating ``session_id`` /
+``user_id`` here. The chat model layer reads it and forwards an end-user
+identifier as the OpenAI-standard ``user`` request field, which providers such as
+OpenAI and OpenRouter include in the request body. This lets upstream activity be
+correlated back to the session/user that produced it.
+
+The value is a small dict (or ``None`` when no request context is active), e.g.
+``{"session_id": "...", "user_id": "..."}``.
+"""
+
+import contextvars
+import logging
+
+llm_request_context: contextvars.ContextVar = contextvars.ContextVar("ragflow_llm_request_context", default=None)
+
+
+def set_llm_request_context(session_id: str | None = None, user_id: str | None = None):
+    """Install the current request identifiers and return the reset token.
+
+    Pass the returned token to ``reset_llm_request_context`` (typically in a
+    ``finally`` block) so the value does not leak to later calls in the same task.
+    """
+    ctx = {}
+    if session_id:
+        ctx["session_id"] = str(session_id)[:128]
+    if user_id:
+        ctx["user_id"] = str(user_id)[:128]
+    # Log only presence flags, never the raw identifiers.
+    logging.debug("Installing LLM request context (session=%s, user=%s)", bool(session_id), bool(user_id))
+    return llm_request_context.set(ctx or None)
+
+
+def reset_llm_request_context(token) -> None:
+    try:
+        llm_request_context.reset(token)
+    except (ValueError, RuntimeError):
+        # The context may be reset from a different context (e.g. an async generator
+        # closed on client disconnect -> ValueError) or with an already-consumed
+        # token (Python 3.13+ -> RuntimeError); fall back to clearing the value.
+        logging.debug("LLM request context reset failed; clearing active context", exc_info=True)
+        llm_request_context.set(None)
+
+
+def current_llm_user() -> str | None:
+    """Return the identifier to forward as the provider ``user`` field.
+
+    Prefers ``session_id`` (so upstream activity can be traced per chat session),
+    falling back to ``user_id``. Returns ``None`` when no context is active.
+    """
+    ctx = llm_request_context.get()
+    if not ctx:
+        return None
+    return ctx.get("session_id") or ctx.get("user_id") or None

rag/llm/chat_model.py

@@ -32,6 +32,7 @@ from openai import AsyncOpenAI, OpenAI
 from enum import StrEnum
 
 from common.misc_utils import thread_pool_exec
+from common.llm_request_context import current_llm_user
 from common.token_utils import num_tokens_from_string, total_token_count_from_response, usage_from_response
 from rag.llm import FACTORY_DEFAULT_BASE_URL, LITELLM_PROVIDER_PREFIX, SupportedLiteLLMProvider
 from rag.llm.key_utils import _normalize_replicate_key
@@ -2181,6 +2182,15 @@ class LiteLLMBase(ABC):
             "num_retries": self.max_retries,
             **kwargs,
         }
+        # Forward the originating session/user as the OpenAI-standard `user` field so
+        # providers (OpenAI, OpenRouter, ...) receive it in the request body and
+        # upstream activity can be correlated back to the session. An explicit
+        # caller-supplied `user` (including an empty string to suppress it) wins, so
+        # check key presence rather than truthiness.
+        if "user" not in completion_args:
+            request_user = current_llm_user()
+            if request_user:
+                completion_args["user"] = request_user
         if stream:
             completion_args.update(
                 {

test/unit_test/api/db/services/test_get_queue_length.py

@@ -0,0 +1,155 @@
+#
+#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+import logging
+import sys
+import types
+import warnings
+
+import pytest
+
+# xgboost imports pkg_resources and emits a deprecation warning that is promoted
+# to error in our pytest configuration; ignore it for this unit test module.
+warnings.filterwarnings(
+    "ignore",
+    message="pkg_resources is deprecated as an API.*",
+    category=UserWarning,
+)
+
+
+def _install_cv2_stub_if_unavailable():
+    try:
+        import cv2  # noqa: F401
+        return
+    except (ImportError, OSError) as exc:
+        # cv2 can fail to import with OSError too (e.g. missing shared libs),
+        # not just ImportError; fall back to the stub in both cases.
+        logging.debug("cv2 unavailable; installing test stub: %s", exc)
+
+    stub = types.ModuleType("cv2")
+
+    def _missing(*_args, **_kwargs):
+        raise RuntimeError("cv2 runtime call is unavailable in this test environment")
+
+    def _module_getattr(name):
+        if name.isupper():
+            return 0
+        return _missing
+
+    stub.__getattr__ = _module_getattr
+    sys.modules["cv2"] = stub
+
+
+_install_cv2_stub_if_unavailable()
+
+from api.db.services import document_service as ds  # noqa: E402
+
+
+@pytest.fixture(autouse=True)
+def _reset_pending_cache(monkeypatch):
+    """Disable the short-lived backlog cache so each test is deterministic."""
+    monkeypatch.setattr(ds, "_PENDING_TASK_COUNT_CACHE", {})
+    monkeypatch.setattr(ds, "_PENDING_TASK_COUNT_TTL_SECONDS", 0.0)
+    yield
+
+
+def _patch_lag(monkeypatch, lag):
+    """Make REDIS_CONN.queue_info report a given consumer-group lag."""
+    group_info = None if lag is None else {"lag": lag}
+    monkeypatch.setattr(ds.REDIS_CONN, "queue_info", lambda *_a, **_k: group_info)
+
+
+def _patch_pending(monkeypatch, pending):
+    monkeypatch.setattr(ds, "get_pending_task_count", lambda *_a, **_k: pending)
+
+
+@pytest.mark.p2
+class TestGetQueueLength:
+    def test_lag_capped_by_genuine_pending(self, monkeypatch):
+        # Redis still reports 34 undelivered messages, but only 5 tasks are
+        # genuinely waiting -> the user must not see "34 tasks ahead".
+        _patch_lag(monkeypatch, 34)
+        _patch_pending(monkeypatch, 5)
+        assert ds.get_queue_length(0) == 5
+
+    def test_self_heals_to_zero_after_stop(self, monkeypatch):
+        # Everything was cancelled: no genuine backlog -> queue length is 0
+        # even though stale messages are still sitting in the Redis stream.
+        _patch_lag(monkeypatch, 34)
+        _patch_pending(monkeypatch, 0)
+        assert ds.get_queue_length(0) == 0
+
+    def test_reports_lag_when_smaller_than_pending(self, monkeypatch):
+        # Some waiting tasks were already delivered (in flight), so lag is the
+        # tighter, truthful bound.
+        _patch_lag(monkeypatch, 2)
+        _patch_pending(monkeypatch, 9)
+        assert ds.get_queue_length(0) == 2
+
+    def test_falls_back_to_lag_when_db_unavailable(self, monkeypatch):
+        # If the backlog cannot be computed we keep the previous behaviour.
+        _patch_lag(monkeypatch, 7)
+        _patch_pending(monkeypatch, None)
+        assert ds.get_queue_length(0) == 7
+
+    def test_missing_group_info_is_zero(self, monkeypatch):
+        _patch_lag(monkeypatch, None)
+        _patch_pending(monkeypatch, 5)
+        assert ds.get_queue_length(0) == 0
+
+    def test_null_lag_value_is_treated_as_zero(self, monkeypatch):
+        monkeypatch.setattr(ds.REDIS_CONN, "queue_info", lambda *_a, **_k: {"lag": None})
+        _patch_pending(monkeypatch, 5)
+        assert ds.get_queue_length(0) == 0
+
+
+@pytest.mark.p2
+class TestGetPendingTaskCount:
+    def test_returns_none_on_db_error(self, monkeypatch):
+        def _boom(*_a, **_k):
+            raise RuntimeError("db down")
+
+        monkeypatch.setattr(ds.Task, "select", _boom)
+        assert ds.get_pending_task_count() is None
+
+    def test_uses_cache_within_ttl(self, monkeypatch):
+        monkeypatch.setattr(ds, "_PENDING_TASK_COUNT_TTL_SECONDS", 60.0)
+        # Cache is keyed by priority (None == "all priorities").
+        monkeypatch.setattr(
+            ds,
+            "_PENDING_TASK_COUNT_CACHE",
+            {None: {"value": 11, "expire_at": ds.monotonic() + 60.0}},
+        )
+
+        def _boom(*_a, **_k):
+            raise AssertionError("DB must not be queried while cache is valid")
+
+        monkeypatch.setattr(ds.Task, "select", _boom)
+        assert ds.get_pending_task_count() == 11
+
+    def test_cache_is_per_priority(self, monkeypatch):
+        monkeypatch.setattr(ds, "_PENDING_TASK_COUNT_TTL_SECONDS", 60.0)
+        # Priority 1 is cached; priority 0 is not -> only priority 0 hits the DB.
+        monkeypatch.setattr(
+            ds,
+            "_PENDING_TASK_COUNT_CACHE",
+            {1: {"value": 3, "expire_at": ds.monotonic() + 60.0}},
+        )
+
+        def _boom(*_a, **_k):
+            raise AssertionError("DB must not be queried for a cached priority")
+
+        monkeypatch.setattr(ds.Task, "select", _boom)
+        assert ds.get_pending_task_count(1) == 3

test/unit_test/common/test_llm_request_context.py

@@ -0,0 +1,138 @@
+#
+#  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+"""Unit tests for the LLM request-context user-forwarding precedence rules.
+
+Covers the behaviour flagged in review: (1) session_id is preferred over
+user_id, (2) an explicit caller-supplied ``user`` overrides the context value,
+and (3) a caller-supplied empty ``user`` suppresses forwarding entirely.
+"""
+import types
+
+import pytest
+
+from common.llm_request_context import (
+    current_llm_user,
+    llm_request_context,
+    reset_llm_request_context,
+    set_llm_request_context,
+)
+
+
+@pytest.fixture(autouse=True)
+def _isolate_context():
+    """Ensure every test starts and ends with no active request context."""
+    token = llm_request_context.set(None)
+    yield
+    try:
+        llm_request_context.reset(token)
+    except ValueError:
+        llm_request_context.set(None)
+
+
+@pytest.mark.p2
+class TestCurrentLlmUser:
+    def test_no_active_context_returns_none(self):
+        assert current_llm_user() is None
+
+    def test_session_id_preferred_over_user_id(self):
+        token = set_llm_request_context(session_id="sess-1", user_id="user-1")
+        try:
+            assert current_llm_user() == "sess-1"
+        finally:
+            reset_llm_request_context(token)
+
+    def test_falls_back_to_user_id_without_session(self):
+        token = set_llm_request_context(session_id=None, user_id="user-1")
+        try:
+            assert current_llm_user() == "user-1"
+        finally:
+            reset_llm_request_context(token)
+
+    def test_empty_identifiers_return_none(self):
+        token = set_llm_request_context(session_id=None, user_id=None)
+        try:
+            assert current_llm_user() is None
+        finally:
+            reset_llm_request_context(token)
+
+    def test_identifier_is_truncated_to_128_chars(self):
+        token = set_llm_request_context(session_id="s" * 200)
+        try:
+            assert current_llm_user() == "s" * 128
+        finally:
+            reset_llm_request_context(token)
+
+
+@pytest.mark.p2
+class TestSetResetLifecycle:
+    def test_reset_restores_previous_context(self):
+        outer = set_llm_request_context(session_id="outer")
+        try:
+            inner = set_llm_request_context(session_id="inner")
+            assert current_llm_user() == "inner"
+            reset_llm_request_context(inner)
+            assert current_llm_user() == "outer"
+        finally:
+            reset_llm_request_context(outer)
+
+    def test_reset_with_stale_token_does_not_raise(self):
+        token = set_llm_request_context(session_id="x")
+        reset_llm_request_context(token)
+        # Resetting again with the now-stale token must fall back, not raise
+        # (mirrors an async generator closed from a different context).
+        reset_llm_request_context(token)
+        assert current_llm_user() is None
+
+
+@pytest.mark.p2
+class TestCompletionArgsUserPrecedence:
+    """Exercise the provider chokepoint: LiteLLMBase._construct_completion_args."""
+
+    def _construct(self, **kwargs):
+        chat_model = pytest.importorskip("rag.llm.chat_model")
+        # provider="" keeps every provider-specific branch inert, so only the
+        # generic completion_args + user-forwarding path is exercised.
+        fake = types.SimpleNamespace(model_name="m", api_key="k", max_retries=0, provider="")
+        return chat_model.LiteLLMBase._construct_completion_args(fake, [], False, False, **kwargs)
+
+    def test_context_user_applied_when_caller_omits_it(self):
+        token = set_llm_request_context(session_id="sess-9", user_id="user-9")
+        try:
+            args = self._construct()
+            assert args["user"] == "sess-9"
+        finally:
+            reset_llm_request_context(token)
+
+    def test_caller_user_overrides_context(self):
+        token = set_llm_request_context(session_id="sess-9")
+        try:
+            args = self._construct(user="explicit-caller")
+            assert args["user"] == "explicit-caller"
+        finally:
+            reset_llm_request_context(token)
+
+    def test_caller_empty_user_suppresses_forwarding(self):
+        token = set_llm_request_context(session_id="sess-9")
+        try:
+            args = self._construct(user="")
+            # Key presence (not truthiness) is honoured: the empty string wins.
+            assert args["user"] == ""
+        finally:
+            reset_llm_request_context(token)
+
+    def test_no_context_leaves_user_unset(self):
+        args = self._construct()
+        assert "user" not in args