"""Opt-in model timing remains independent of busy tools and slow log sinks.""" import asyncio import logging import threading from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from uuid import uuid4 import pytest from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import AIMessageChunk from langchain_core.outputs import ChatGenerationChunk from deerflow.config import runtime_performance_config from deerflow.models import timing_callback from deerflow.models.timing_callback import ModelTimingCallback, _TimingLogDispatcher @pytest.fixture def timing_logs(monkeypatch, caplog): monkeypatch.setattr(runtime_performance_config, "get_runtime_performance_config", lambda: SimpleNamespace(model_timing_nonblocking_enabled=True)) dispatcher = _TimingLogDispatcher() monkeypatch.setattr(timing_callback, "_timing_logs", dispatcher) caplog.set_level(logging.INFO, logger=timing_callback.logger.name) yield dispatcher dispatcher.close(timeout=2) class _NativeAsyncModel(BaseChatModel): @property def _llm_type(self): return "timing-callback-test" def _generate(self, messages, stop=None, run_manager=None, **kwargs): raise AssertionError("The model must use native async streaming") async def _astream(self, messages, stop=None, run_manager=None, **kwargs): yield ChatGenerationChunk(message=AIMessageChunk(content="first")) yield ChatGenerationChunk(message=AIMessageChunk(content="second")) class _SlowHandler(logging.Handler): def __init__(self): super().__init__() self.entered = threading.Event() self.release_event = threading.Event() self.worker_threads = [] def emit(self, record): self.worker_threads.append(threading.current_thread().name) self.entered.set() self.release_event.wait(timeout=5) async def _wait_until_set(event): while not event.is_set(): await asyncio.sleep(0.001) def test_default_off_keeps_synchronous_logging_and_executor_dispatch(monkeypatch, caplog): config = runtime_performance_config.RuntimePerformanceConfig() monkeypatch.setattr(runtime_performance_config, "get_runtime_performance_config", lambda: config) caplog.set_level(logging.INFO, logger=timing_callback.logger.name) callback = ModelTimingCallback(model_name="test", base_url=None) assert callback.run_inline is False def reject_queue(*args): raise AssertionError("Default-off callbacks must not use the background log queue") monkeypatch.setattr(timing_callback._timing_logs, "emit", reject_queue) run_id = uuid4() callback.on_chat_model_start({}, [[]], run_id=run_id) assert "LLM_CONNECT_START" in caplog.text callback.on_llm_new_token("token", run_id=run_id) assert "LLM_CONNECT_OK" in caplog.text callback.on_llm_end(None, run_id=run_id) assert "LLM_REQUEST_DONE" in caplog.text def test_stream_progresses_with_saturated_tool_pool_and_blocked_log_handler(timing_logs, caplog): """Exercise LangChain dispatch, not just direct callback method calls.""" sink = _SlowHandler() timing_callback.logger.addHandler(sink) tool_started = threading.Event() release_tool = threading.Event() def long_tool(): tool_started.set() release_tool.wait(timeout=5) async def scenario(): loop = asyncio.get_running_loop() # A fresh asyncio.run loop owns this executor and cleans it up. loop.set_default_executor(ThreadPoolExecutor(max_workers=1, thread_name_prefix="busy-tool")) blocked_tool = loop.run_in_executor(None, long_tool) try: await asyncio.wait_for(_wait_until_set(tool_started), timeout=1) callback = ModelTimingCallback(model_name="test", base_url="https://models.invalid/v1") model = _NativeAsyncModel(callbacks=[callback]) async def consume(): return [chunk.content async for chunk in model.astream("private prompt")] chunks = await asyncio.wait_for(consume(), timeout=1) await asyncio.wait_for(_wait_until_set(sink.entered), timeout=1) assert "".join(chunks) == "firstsecond" assert not blocked_tool.done() assert not sink.release_event.is_set() assert callback._calls == {} finally: release_tool.set() sink.release_event.set() await blocked_tool try: asyncio.run(scenario()) timing_logs.close(timeout=2) finally: release_tool.set() sink.release_event.set() timing_callback.logger.removeHandler(sink) assert sink.worker_threads assert set(sink.worker_threads) == {"deerflow-model-timing-log"} assert "LLM_CONNECT_START" in caplog.text assert "LLM_CONNECT_OK" in caplog.text assert "LLM_REQUEST_DONE" in caplog.text assert "private prompt" not in caplog.text def test_sync_callback_preserves_sanitized_lifecycle_and_only_logs_first_token(timing_logs, caplog): callback = ModelTimingCallback(model_name="test", base_url="https://user:secret@models.invalid:443/private?key=secret") run_id = uuid4() callback.on_chat_model_start({}, [["private input"]], run_id=run_id) callback.on_llm_new_token("", run_id=run_id) for _ in range(100): callback.on_llm_new_token("private output", run_id=run_id) callback.on_llm_end(None, run_id=run_id) timing_logs.close(timeout=2) records = [record for record in caplog.records if record.name == timing_callback.logger.name] assert len(records) == 3 assert "base_url=https://models.invalid:443" in caplog.text assert "first_chunk_ms=" in caplog.text assert f"run_id={run_id}" in caplog.text for secret in ("private", "secret", "user:"): assert secret not in caplog.text assert callback._calls == {} @pytest.mark.parametrize( ("error", "expected", "level"), [ (RuntimeError("private upstream payload"), "LLM_REQUEST_FAILED", logging.ERROR), (asyncio.CancelledError("private cancellation"), "LLM_REQUEST_CANCELLED", logging.INFO), (GeneratorExit("private closure"), "LLM_REQUEST_CANCELLED", logging.INFO), ], ) def test_failure_and_cancellation_are_content_free(timing_logs, caplog, error, expected, level): callback = ModelTimingCallback(model_name="test", base_url=None) run_id = uuid4() callback.on_chat_model_start({}, [[]], run_id=run_id) callback.on_llm_error(error, run_id=run_id) timing_logs.close(timeout=2) terminal = [record for record in caplog.records if expected in record.getMessage()] assert len(terminal) == 1 assert terminal[0].levelno == level assert "private" not in caplog.text assert callback._calls == {} def test_full_log_queue_drops_without_blocking_and_reports_recovery(monkeypatch, caplog): dispatcher = _TimingLogDispatcher(capacity=1) monkeypatch.setattr(timing_callback, "_timing_logs", dispatcher) caplog.set_level(logging.INFO, logger=timing_callback.logger.name) sink = _SlowHandler() timing_callback.logger.addHandler(sink) try: dispatcher.emit(logging.INFO, "first") assert sink.entered.wait(timeout=1) dispatcher.emit(logging.INFO, "queued") dispatcher.emit(logging.INFO, "dropped") assert dispatcher._queue.qsize() == 1 assert dispatcher._dropped == 1 sink.release_event.set() dispatcher.close(timeout=2) finally: sink.release_event.set() dispatcher.close(timeout=2) timing_callback.logger.removeHandler(sink) assert "LLM_TIMING_LOG_DROPPED count=1 reason=queue_full" in caplog.text assert "queued" in caplog.text assert not any(record.getMessage() == "dropped" for record in caplog.records)