learningcircuit--local-deep-research/tests/test_token_counter_coverage.py

"""Comprehensive pytest tests for local_deep_research/metrics/token_counter.py.

Covers: TokenCountingCallback (init, on_llm_start, on_llm_end, on_llm_error,
        _get_context_overflow_fields, _save_to_db, get_counts) and
        TokenCounter (create_callback, _get_empty_metrics).
"""

import time
from unittest.mock import MagicMock, Mock, patch

import pytest

from local_deep_research.metrics.token_counter import (
    TokenCounter,
    TokenCountingCallback,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_callback(**kwargs):
    """Shorthand to build a callback without DB deps."""
    return TokenCountingCallback(**kwargs)


def _llm_result_with_token_usage(prompt=10, completion=20, total=None):
    """Return a mock LLMResult whose llm_output contains token_usage."""
    if total is None:
        total = prompt + completion
    response = Mock()
    response.llm_output = {
        "token_usage": {
            "prompt_tokens": prompt,
            "completion_tokens": completion,
            "total_tokens": total,
        }
    }
    response.generations = []
    return response


def _llm_result_with_usage_metadata(
    input_tokens=100, output_tokens=50, total_tokens=150
):
    """Return a mock LLMResult with usage_metadata on the message (Gemini/Google path)."""
    response = Mock()
    response.llm_output = None

    message = Mock()
    message.usage_metadata = {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": total_tokens,
    }
    message.response_metadata = {}

    generation = Mock()
    generation.message = message
    response.generations = [[generation]]
    return response


def _llm_result_with_ollama_response_metadata(
    prompt_eval_count=200, eval_count=80, total_duration=None
):
    """Return a mock LLMResult with Ollama-style response_metadata."""
    response = Mock()
    response.llm_output = None

    message = Mock()
    message.usage_metadata = None  # usage_metadata absent or None
    message.response_metadata = {
        "prompt_eval_count": prompt_eval_count,
        "eval_count": eval_count,
        "total_duration": total_duration or 5_000_000_000,
        "load_duration": 100_000_000,
        "prompt_eval_duration": 2_000_000_000,
        "eval_duration": 1_500_000_000,
    }

    generation = Mock()
    generation.message = message
    response.generations = [[generation]]
    return response


def _llm_result_empty():
    """Return a mock LLMResult with no token info at all."""
    response = Mock()
    response.llm_output = None
    response.generations = []
    return response


# ===========================================================================
# TokenCountingCallback — Initialization
# ===========================================================================


class TestTokenCountingCallbackInit:
    def test_defaults(self):
        cb = _make_callback()
        assert cb.research_id is None
        assert cb.research_context == {}
        assert cb.current_model is None
        assert cb.current_provider is None
        assert cb.preset_model is None
        assert cb.preset_provider is None
        assert cb.start_time is None
        assert cb.response_time_ms is None
        assert cb.success_status == "success"
        assert cb.error_type is None
        assert cb.calling_file is None
        assert cb.calling_function is None
        assert cb.call_stack is None
        assert cb.context_limit is None
        assert cb.context_truncated is False
        assert cb.tokens_truncated == 0
        assert cb.truncation_ratio == 0.0
        assert cb.original_prompt_estimate == 0
        assert cb.ollama_metrics == {}

    def test_counts_structure(self):
        cb = _make_callback()
        assert cb.counts == {
            "total_tokens": 0,
            "total_prompt_tokens": 0,
            "total_completion_tokens": 0,
            "by_model": {},
        }

    def test_research_id_stored(self):
        cb = _make_callback(research_id="abc-123")
        assert cb.research_id == "abc-123"

    def test_research_context_stored(self):
        ctx = {
            "research_query": "quantum computing",
            "research_mode": "detailed",
        }
        cb = _make_callback(research_context=ctx)
        assert cb.research_context is ctx

    def test_none_research_context_becomes_empty_dict(self):
        cb = _make_callback(research_context=None)
        assert cb.research_context == {}


# ===========================================================================
# on_llm_start — model/provider detection
# ===========================================================================


class TestOnLlmStart:
    def test_preset_model_takes_priority(self):
        cb = _make_callback()
        cb.preset_model = "my-custom-model"
        cb.preset_provider = "custom-provider"

        cb.on_llm_start(serialized={}, prompts=["hello"])

        assert cb.current_model == "my-custom-model"
        assert cb.current_provider == "custom-provider"

    def test_model_from_invocation_params(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={},
            prompts=["hello"],
            invocation_params={"model": "gpt-4"},
        )
        assert cb.current_model == "gpt-4"

    def test_model_from_invocation_params_model_name(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={},
            prompts=["hello"],
            invocation_params={"model_name": "gpt-3.5-turbo"},
        )
        assert cb.current_model == "gpt-3.5-turbo"

    def test_model_from_kwargs(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="claude-3")
        assert cb.current_model == "claude-3"

    def test_model_from_serialized_kwargs(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"kwargs": {"model": "llama-3.1"}},
            prompts=["hi"],
        )
        assert cb.current_model == "llama-3.1"

    def test_model_from_serialized_name(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"name": "ChatGPT"},
            prompts=["hi"],
        )
        assert cb.current_model == "ChatGPT"

    def test_model_from_ollama_type(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"_type": "ChatOllama", "kwargs": {"model": "mistral"}},
            prompts=["hi"],
        )
        assert cb.current_model == "mistral"

    def test_model_ollama_type_fallback(self):
        """When _type is ChatOllama but no model in kwargs, falls back to 'ollama'."""
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"_type": "ChatOllama"},
            prompts=["hi"],
        )
        assert cb.current_model == "ollama"

    def test_model_fallback_to_type(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"_type": "SomeCustomLLM"},
            prompts=["hi"],
        )
        assert cb.current_model == "SomeCustomLLM"

    def test_model_fallback_to_unknown(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"])
        assert cb.current_model == "unknown"

    # --- provider detection ---

    def test_provider_ollama(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"_type": "ChatOllama", "kwargs": {"model": "m"}},
            prompts=["hi"],
        )
        assert cb.current_provider == "ollama"

    def test_provider_openai(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}},
            prompts=["hi"],
        )
        assert cb.current_provider == "openai"

    def test_provider_anthropic(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"_type": "ChatAnthropic", "kwargs": {"model": "c3"}},
            prompts=["hi"],
        )
        assert cb.current_provider == "anthropic"

    def test_provider_from_kwargs(self):
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"_type": "SomethingElse"},
            prompts=["hi"],
            provider="azure",
        )
        assert cb.current_provider == "azure"

    def test_provider_unknown_fallback(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"])
        assert cb.current_provider == "unknown"

    # --- call count / model tracking ---

    def test_initializes_model_tracking(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="gpt-4")
        assert "gpt-4" in cb.counts["by_model"]
        assert cb.counts["by_model"]["gpt-4"]["calls"] == 1

    def test_increments_call_count(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="gpt-4")
        cb.on_llm_start(serialized={}, prompts=["hi"], model="gpt-4")
        assert cb.counts["by_model"]["gpt-4"]["calls"] == 2

    # --- prompt estimation ---

    def test_original_prompt_estimate(self):
        cb = _make_callback()
        # 400 chars -> ~100 estimated tokens
        cb.on_llm_start(serialized={}, prompts=["a" * 400])
        assert cb.original_prompt_estimate == 100

    def test_original_prompt_estimate_empty_prompts(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=[])
        assert cb.original_prompt_estimate == 0

    def test_original_prompt_estimate_multiple_prompts(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["a" * 100, "b" * 300])
        assert cb.original_prompt_estimate == 100  # 400 chars / 4

    # --- timing ---

    def test_start_time_is_set(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"])
        assert cb.start_time is not None
        assert cb.start_time <= time.time()

    # --- context_limit from research_context ---

    def test_context_limit_from_research_context(self):
        cb = _make_callback(research_context={"context_limit": 4096})
        cb.on_llm_start(serialized={}, prompts=["hi"])
        assert cb.context_limit == 4096


# ===========================================================================
# on_llm_end — token usage extraction
# ===========================================================================


class TestOnLlmEnd:
    def _start_and_end(self, cb, response):
        """Helper: call on_llm_start then on_llm_end."""
        cb.on_llm_start(serialized={}, prompts=["hi"], model="test-model")
        cb.on_llm_end(response)

    def test_token_usage_from_llm_output(self):
        cb = _make_callback()
        response = _llm_result_with_token_usage(prompt=10, completion=20)
        self._start_and_end(cb, response)

        assert cb.counts["total_prompt_tokens"] == 10
        assert cb.counts["total_completion_tokens"] == 20
        assert cb.counts["total_tokens"] == 30

    def test_token_usage_from_usage_key(self):
        """Token usage found under 'usage' key in llm_output."""
        response = Mock()
        response.llm_output = {
            "usage": {
                "prompt_tokens": 5,
                "completion_tokens": 15,
                "total_tokens": 20,
            }
        }
        response.generations = []

        cb = _make_callback()
        self._start_and_end(cb, response)
        assert cb.counts["total_tokens"] == 20

    def test_token_usage_from_usage_metadata(self):
        cb = _make_callback()
        response = _llm_result_with_usage_metadata(100, 50, 150)
        self._start_and_end(cb, response)

        assert cb.counts["total_prompt_tokens"] == 100
        assert cb.counts["total_completion_tokens"] == 50
        assert cb.counts["total_tokens"] == 150

    def test_token_usage_from_ollama_response_metadata(self):
        cb = _make_callback()
        response = _llm_result_with_ollama_response_metadata(200, 80)
        self._start_and_end(cb, response)

        assert cb.counts["total_prompt_tokens"] == 200
        assert cb.counts["total_completion_tokens"] == 80
        assert cb.counts["total_tokens"] == 280

    def test_ollama_metrics_captured(self):
        cb = _make_callback()
        response = _llm_result_with_ollama_response_metadata(
            200, 80, 5_000_000_000
        )
        self._start_and_end(cb, response)

        assert cb.ollama_metrics["prompt_eval_count"] == 200
        assert cb.ollama_metrics["eval_count"] == 80
        assert cb.ollama_metrics["total_duration"] == 5_000_000_000

    def test_no_token_usage_does_not_crash(self):
        cb = _make_callback()
        response = _llm_result_empty()
        self._start_and_end(cb, response)

        assert cb.counts["total_tokens"] == 0
        assert cb.counts["total_prompt_tokens"] == 0
        assert cb.counts["total_completion_tokens"] == 0

    def test_by_model_counts_updated(self):
        cb = _make_callback()
        response = _llm_result_with_token_usage(prompt=10, completion=20)
        self._start_and_end(cb, response)

        model_counts = cb.counts["by_model"]["test-model"]
        assert model_counts["prompt_tokens"] == 10
        assert model_counts["completion_tokens"] == 20
        assert model_counts["total_tokens"] == 30

    def test_accumulation_over_multiple_calls(self):
        cb = _make_callback()
        r1 = _llm_result_with_token_usage(prompt=10, completion=20)
        r2 = _llm_result_with_token_usage(prompt=5, completion=15)

        self._start_and_end(cb, r1)
        cb.on_llm_start(serialized={}, prompts=["hi"], model="test-model")
        cb.on_llm_end(r2)

        assert cb.counts["total_prompt_tokens"] == 15
        assert cb.counts["total_completion_tokens"] == 35
        assert cb.counts["total_tokens"] == 50

    def test_total_tokens_defaults_to_sum(self):
        """When total_tokens missing from dict, it's computed as prompt + completion."""
        response = Mock()
        response.llm_output = {
            "token_usage": {
                "prompt_tokens": 7,
                "completion_tokens": 3,
            }
        }
        response.generations = []

        cb = _make_callback()
        self._start_and_end(cb, response)
        assert cb.counts["total_tokens"] == 10

    def test_response_time_calculated(self):
        # audit: PUNCHLIST reviewed 2026-05 — issue resolved by prior PR (recommendation: keep but consider freezing).
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")
        # Simulate elapsed time
        cb.start_time = time.time() - 0.5  # 500ms ago
        cb.on_llm_end(_llm_result_with_token_usage())
        assert cb.response_time_ms is not None
        assert cb.response_time_ms >= 400  # at least ~400ms

    def test_save_to_db_called_when_research_id_present(self):
        cb = _make_callback(research_id="r-123")
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")

        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_end(_llm_result_with_token_usage(10, 20))
            mock_save.assert_called_once_with(10, 20)

    def test_save_to_db_not_called_without_research_id(self):
        cb = _make_callback()  # no research_id
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")

        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_end(_llm_result_with_token_usage(10, 20))
            mock_save.assert_not_called()


# ===========================================================================
# on_llm_end — context overflow detection (Ollama)
# ===========================================================================


class TestContextOverflowDetection:
    def test_context_truncated_when_near_limit(self):
        """When prompt_eval_count >= 80% of context_limit, flag truncation."""
        cb = _make_callback(research_context={"context_limit": 1000})
        cb.on_llm_start(serialized={}, prompts=["a" * 4000], model="m")
        # original_prompt_estimate = 4000/4 = 1000

        response = _llm_result_with_ollama_response_metadata(
            prompt_eval_count=960, eval_count=50
        )
        cb.on_llm_end(response)

        assert cb.context_truncated is True
        assert cb.tokens_truncated == 40  # 1000 - 960
        assert cb.truncation_ratio == pytest.approx(0.04, abs=0.001)  # 40/1000

    def test_context_not_truncated_when_below_threshold(self):
        cb = _make_callback(research_context={"context_limit": 1000})
        cb.on_llm_start(serialized={}, prompts=["a" * 400], model="m")
        # original_prompt_estimate = 100

        response = _llm_result_with_ollama_response_metadata(
            prompt_eval_count=100, eval_count=50
        )
        cb.on_llm_end(response)

        assert cb.context_truncated is False

    def test_context_no_limit_set(self):
        """No context_limit means no truncation detection."""
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["a" * 4000], model="m")

        response = _llm_result_with_ollama_response_metadata(
            prompt_eval_count=960, eval_count=50
        )
        cb.on_llm_end(response)

        assert cb.context_truncated is False


class TestContextOverflowViaUsageMetadata:
    """Verify overflow detection via usage_metadata branch (langchain-ollama v1.0.1+)."""

    def test_overflow_detected_via_usage_metadata_input_tokens(self):
        """input_tokens >= 80% of context_limit triggers truncation."""
        cb = _make_callback(research_context={"context_limit": 1000})
        cb.on_llm_start(serialized={}, prompts=["a" * 4000], model="m")

        # Build a response where usage_metadata is present (langchain-ollama v1.0.1)
        response = Mock()
        response.llm_output = None

        message = Mock()
        message.usage_metadata = {
            "input_tokens": 850,  # >= 1000 * 0.80
            "output_tokens": 50,
            "total_tokens": 900,
        }
        message.response_metadata = {}

        generation = Mock()
        generation.message = message
        response.generations = [[generation]]

        cb.on_llm_end(response)

        assert cb.context_truncated is True

    def test_no_overflow_below_threshold_via_usage_metadata(self):
        """input_tokens below 80% does not trigger truncation."""
        cb = _make_callback(research_context={"context_limit": 1000})
        cb.on_llm_start(serialized={}, prompts=["a" * 100], model="m")

        response = Mock()
        response.llm_output = None

        message = Mock()
        message.usage_metadata = {
            "input_tokens": 700,  # < 1000 * 0.80 = 800
            "output_tokens": 50,
            "total_tokens": 750,
        }
        message.response_metadata = {}

        generation = Mock()
        generation.message = message
        response.generations = [[generation]]

        cb.on_llm_end(response)

        assert cb.context_truncated is False

    def test_usage_metadata_takes_priority_over_response_metadata(self):
        """When both metadata sources exist, usage_metadata branch fires first."""
        cb = _make_callback(research_context={"context_limit": 1000})
        cb.on_llm_start(serialized={}, prompts=["a" * 4000], model="m")

        response = Mock()
        response.llm_output = None

        message = Mock()
        # usage_metadata present — this branch should handle detection
        message.usage_metadata = {
            "input_tokens": 900,
            "output_tokens": 50,
            "total_tokens": 950,
        }
        # response_metadata also present but should NOT be reached
        message.response_metadata = {
            "prompt_eval_count": 900,
            "eval_count": 50,
        }

        generation = Mock()
        generation.message = message
        response.generations = [[generation]]

        cb.on_llm_end(response)

        assert cb.context_truncated is True


# ===========================================================================
# on_llm_error
# ===========================================================================


class TestOnLlmError:
    def test_error_status_set(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")

        cb.on_llm_error(ValueError("bad input"))

        assert cb.success_status == "error"
        assert cb.error_type == "ValueError"

    def test_response_time_calculated_on_error(self):
        # audit: PUNCHLIST reviewed 2026-05 — KEEP (RACE_CONDITIONS).
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")
        cb.start_time = time.time() - 1.0  # 1 second ago

        cb.on_llm_error(RuntimeError("fail"))

        assert cb.response_time_ms is not None
        assert cb.response_time_ms >= 900

    def test_save_to_db_called_on_error_with_research_id(self):
        cb = _make_callback(research_id="r-err")
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")

        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_error(RuntimeError("fail"))
            mock_save.assert_called_once_with(0, 0)

    def test_save_to_db_not_called_on_error_without_research_id(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")

        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_error(RuntimeError("fail"))
            mock_save.assert_not_called()

    def test_error_type_captures_class_name(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")

        class CustomAPIError(Exception):
            pass

        cb.on_llm_error(CustomAPIError("rate limited"))
        assert cb.error_type == "CustomAPIError"


# ===========================================================================
# _get_context_overflow_fields
# ===========================================================================


class TestGetContextOverflowFields:
    def test_no_overflow(self):
        cb = _make_callback()
        fields = cb._get_context_overflow_fields()

        assert fields["context_limit"] is None
        assert fields["context_truncated"] is False
        assert fields["tokens_truncated"] is None
        assert fields["truncation_ratio"] is None

    def test_with_overflow(self):
        cb = _make_callback()
        cb.context_limit = 4096
        cb.context_truncated = True
        cb.tokens_truncated = 500
        cb.truncation_ratio = 0.12

        fields = cb._get_context_overflow_fields()

        assert fields["context_limit"] == 4096
        assert fields["context_truncated"] is True
        assert fields["tokens_truncated"] == 500
        assert fields["truncation_ratio"] == 0.12

    def test_ollama_metrics_in_fields(self):
        cb = _make_callback()
        cb.ollama_metrics = {
            "prompt_eval_count": 100,
            "eval_count": 50,
            "total_duration": 5_000_000_000,
            "load_duration": 200_000_000,
            "prompt_eval_duration": 1_000_000_000,
            "eval_duration": 800_000_000,
        }

        fields = cb._get_context_overflow_fields()
        assert fields["ollama_prompt_eval_count"] == 100
        assert fields["ollama_eval_count"] == 50
        assert fields["ollama_total_duration"] == 5_000_000_000

    def test_ollama_metrics_empty(self):
        cb = _make_callback()
        fields = cb._get_context_overflow_fields()
        assert fields["ollama_prompt_eval_count"] is None
        assert fields["ollama_eval_count"] is None


# ===========================================================================
# get_counts
# ===========================================================================


class TestGetCounts:
    def test_returns_counts_dict(self):
        cb = _make_callback()
        counts = cb.get_counts()
        assert counts is cb.counts

    def test_reflects_updates(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")
        cb.on_llm_end(_llm_result_with_token_usage(prompt=7, completion=3))

        counts = cb.get_counts()
        assert counts["total_tokens"] == 10
        assert counts["total_prompt_tokens"] == 7
        assert counts["total_completion_tokens"] == 3


# ===========================================================================
# TokenCounter — factory class
# ===========================================================================


class TestTokenCounter:
    def test_create_callback_returns_callback_instance(self):
        tc = TokenCounter()
        cb = tc.create_callback()
        assert isinstance(cb, TokenCountingCallback)

    def test_create_callback_passes_research_id(self):
        tc = TokenCounter()
        cb = tc.create_callback(research_id="r-1")
        assert cb.research_id == "r-1"

    def test_create_callback_passes_research_context(self):
        tc = TokenCounter()
        ctx = {"research_query": "test"}
        cb = tc.create_callback(research_context=ctx)
        assert cb.research_context is ctx

    def test_get_empty_metrics_structure(self):
        tc = TokenCounter()
        m = tc._get_empty_metrics()

        assert m["total_tokens"] == 0
        assert m["total_researches"] == 0
        assert m["by_model"] == []
        assert m["recent_researches"] == []
        assert "token_breakdown" in m


# ===========================================================================
# Edge cases
# ===========================================================================


class TestEdgeCases:
    def test_on_llm_end_without_on_llm_start(self):
        """on_llm_end should not crash if on_llm_start was never called."""
        cb = _make_callback()
        # current_model is None
        response = _llm_result_with_token_usage(prompt=5, completion=5)
        # Should not raise
        cb.on_llm_end(response)
        # Totals updated but no by_model entry
        assert cb.counts["total_tokens"] == 10

    def test_on_llm_error_without_start_time(self):
        """on_llm_error should not crash if start_time was never set."""
        cb = _make_callback()
        cb.on_llm_error(RuntimeError("oops"))
        assert cb.response_time_ms is None
        assert cb.success_status == "error"

    def test_llm_output_empty_dict(self):
        response = Mock()
        response.llm_output = {}
        response.generations = []

        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")
        cb.on_llm_end(response)
        assert cb.counts["total_tokens"] == 0

    def test_llm_output_none(self):
        response = Mock()
        response.llm_output = None
        response.generations = []

        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")
        cb.on_llm_end(response)
        assert cb.counts["total_tokens"] == 0

    def test_token_usage_with_zero_values(self):
        response = _llm_result_with_token_usage(prompt=0, completion=0, total=0)
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")
        cb.on_llm_end(response)
        assert cb.counts["total_tokens"] == 0

    def test_empty_string_prompt(self):
        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=[""])
        assert cb.original_prompt_estimate == 0

    def test_very_long_prompt_estimate(self):
        cb = _make_callback()
        long_text = "x" * 1_000_000  # 1M chars
        cb.on_llm_start(serialized={}, prompts=[long_text])
        assert cb.original_prompt_estimate == 250_000

    def test_multiple_models_tracked_separately(self):
        cb = _make_callback()

        cb.on_llm_start(serialized={}, prompts=["hi"], model="model-a")
        cb.on_llm_end(_llm_result_with_token_usage(prompt=10, completion=5))

        cb.on_llm_start(serialized={}, prompts=["hi"], model="model-b")
        cb.on_llm_end(_llm_result_with_token_usage(prompt=20, completion=10))

        assert cb.counts["by_model"]["model-a"]["total_tokens"] == 15
        assert cb.counts["by_model"]["model-b"]["total_tokens"] == 30
        assert cb.counts["total_tokens"] == 45

    def test_usage_metadata_with_none_value(self):
        """usage_metadata exists but is None — should fall through gracefully."""
        response = Mock()
        response.llm_output = None

        message = Mock()
        message.usage_metadata = None
        message.response_metadata = {}

        generation = Mock()
        generation.message = message
        response.generations = [[generation]]

        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")
        cb.on_llm_end(response)
        assert cb.counts["total_tokens"] == 0

    def test_generations_with_no_message_attr(self):
        """Generations without .message should not crash."""
        response = Mock()
        response.llm_output = None

        generation = Mock(spec=[])  # no attributes at all
        response.generations = [[generation]]

        cb = _make_callback()
        cb.on_llm_start(serialized={}, prompts=["hi"], model="m")
        cb.on_llm_end(response)
        assert cb.counts["total_tokens"] == 0

    def test_preset_model_and_provider(self):
        """preset_model/provider set before on_llm_start should be used."""
        cb = _make_callback()
        cb.preset_model = "preset-model"
        cb.preset_provider = "preset-provider"

        cb.on_llm_start(
            serialized={"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}},
            prompts=["hi"],
        )

        assert cb.current_model == "preset-model"
        assert cb.current_provider == "preset-provider"

    def test_serialized_kwargs_model_name(self):
        """model_name (not model) in serialized kwargs."""
        cb = _make_callback()
        cb.on_llm_start(
            serialized={"kwargs": {"model_name": "my-model"}},
            prompts=["hi"],
        )
        assert cb.current_model == "my-model"


# ===========================================================================
# _save_to_db — thread detection and error handling
# ===========================================================================


class TestSaveToDb:
    @patch("threading.current_thread")
    def test_background_thread_without_username_skips(self, mock_thread):
        """In a background thread without username, _save_to_db logs warning and returns."""
        mock_thread.return_value.name = "WorkerThread"

        cb = _make_callback(research_id="r-1", research_context={})
        cb.current_model = "m"
        cb.current_provider = "p"

        # Should not raise
        cb._save_to_db(10, 20)

    @patch("threading.current_thread")
    def test_background_thread_without_password_skips(self, mock_thread):
        """In a background thread with username but no password, skips save."""
        mock_thread.return_value.name = "WorkerThread"

        cb = _make_callback(
            research_id="r-1",
            research_context={"username": "alice"},  # no user_password
        )
        cb.current_model = "m"
        cb.current_provider = "p"

        # Should not raise
        cb._save_to_db(10, 20)

    @patch("threading.current_thread")
    def test_background_thread_with_credentials_writes_metrics(
        self, mock_thread
    ):
        """In a background thread with full credentials, calls metrics_writer."""
        mock_thread.return_value.name = "WorkerThread"

        mock_writer = MagicMock()
        cb = _make_callback(
            research_id="r-1",
            research_context={
                "username": "alice",
                "user_password": "secret",
            },
        )
        cb.current_model = "m"
        cb.current_provider = "p"

        with patch(
            "local_deep_research.metrics.token_counter.TokenCountingCallback._save_to_db",
            wraps=cb._save_to_db,
        ):
            with patch(
                "local_deep_research.database.thread_metrics.metrics_writer",
                mock_writer,
            ):
                cb._save_to_db(10, 20)

        mock_writer.set_user_password.assert_called_once_with("alice", "secret")
        mock_writer.write_token_metrics.assert_called_once()

    @patch("threading.current_thread")
    def test_main_thread_no_flask_session_skips(self, mock_thread):
        """In MainThread without flask session, save is skipped."""
        # audit: PUNCHLIST reviewed 2026-05 — issue resolved by prior PR (recommendation: delete or rewrite to actually invoke _save_to_db and assert no metrics write).
        mock_thread.return_value.name = "MainThread"

        cb = _make_callback(research_id="r-1")
        cb.current_model = "m"
        cb.current_provider = "p"

        with patch(
            "local_deep_research.metrics.token_counter.flask_session",
            create=True,
        ):
            # Patch at the import location used in the method
            with patch.dict(
                "sys.modules",
                {"flask": MagicMock()},
            ):
                # The method imports flask.session internally, so we patch it there
                mock_flask_mod = MagicMock()
                mock_flask_mod.session.get.return_value = None
                with patch(
                    "local_deep_research.metrics.token_counter.TokenCountingCallback._save_to_db",
                ) as _:
                    # Simply verify no exception is raised when there's no session
                    pass