From 28a41ed0701adec9222a9dc5851239af70950cd8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=96ndery?= Date: Tue, 7 Jul 2026 14:06:30 +0300 Subject: [PATCH] fix(task_executor): fix Langfuse flush/shutdown deadlock that freezes document parsing (#16502) --- agent/canvas.py | 22 ++- api/apps/restful_apis/document_api.py | 9 + api/apps/restful_apis/task_api.py | 7 +- api/db/services/canvas_service.py | 2 +- api/db/services/document_service.py | 91 +++++++++- api/db/services/tenant_llm_service.py | 43 +++-- common/llm_request_context.py | 71 ++++++++ rag/llm/chat_model.py | 10 ++ .../api/db/services/test_get_queue_length.py | 155 ++++++++++++++++++ .../common/test_llm_request_context.py | 138 ++++++++++++++++ 10 files changed, 519 insertions(+), 29 deletions(-) create mode 100644 common/llm_request_context.py create mode 100644 test/unit_test/api/db/services/test_get_queue_length.py create mode 100644 test/unit_test/common/test_llm_request_context.py diff --git a/agent/canvas.py b/agent/canvas.py index ed7fbe30ca..6874e50752 100644 --- a/agent/canvas.py +++ b/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: diff --git a/api/apps/restful_apis/document_api.py b/api/apps/restful_apis/document_api.py index 4bf149d9f9..f37d4e9cd7 100644 --- a/api/apps/restful_apis/document_api.py +++ b/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) diff --git a/api/apps/restful_apis/task_api.py b/api/apps/restful_apis/task_api.py index 8566c5583d..3688e71bc8 100644 --- a/api/apps/restful_apis/task_api.py +++ b/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)) diff --git a/api/db/services/canvas_service.py b/api/db/services/canvas_service.py index 308f495c13..9aa250c7a7 100644 --- a/api/db/services/canvas_service.py +++ b/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: diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py index 1fae046c95..b9642dffa0 100644 --- a/api/db/services/document_service.py +++ b/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) diff --git a/api/db/services/tenant_llm_service.py b/api/db/services/tenant_llm_service.py index d4b89cc6e6..f9caa2100d 100644 --- a/api/db/services/tenant_llm_service.py +++ b/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) diff --git a/common/llm_request_context.py b/common/llm_request_context.py new file mode 100644 index 0000000000..bfc95f2d43 --- /dev/null +++ b/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 diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py index c7ee06dea9..418c568c5c 100644 --- a/rag/llm/chat_model.py +++ b/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( { diff --git a/test/unit_test/api/db/services/test_get_queue_length.py b/test/unit_test/api/db/services/test_get_queue_length.py new file mode 100644 index 0000000000..a985fed519 --- /dev/null +++ b/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 diff --git a/test/unit_test/common/test_llm_request_context.py b/test/unit_test/common/test_llm_request_context.py new file mode 100644 index 0000000000..17af0ec80b --- /dev/null +++ b/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