6krrt/tests/test_multi_provider_dispatch.py

"""Multi-provider dispatch edge cases.

Exercises the dispatcher where two configured providers diverge in telemetry
shape and in per-provider refusal state. Nothing here calls real providers.
"""
from __future__ import annotations

import json
import sqlite3
from pathlib import Path

import pytest
from starlette.testclient import TestClient

import dispatcher
import session_cache
from config import DispatchProvider

ROOT = Path(__file__).resolve().parent.parent
SCHEMA_SQL = (ROOT / "config" / "schema.sql").read_text()

CHEAP = "cheap-model"
DEAR = "dear-model"

# Stream lines carrying OpenRouter's keepalive comment *instead* of energy/cost.
OPENROUTER_STREAM_LINES = [
    'data: {"id":"chatcmpl-or-1","choices":[{"delta":{"content":"hel"}}]}',
    "",
    ': OPENROUTER PROCESSING',
    'data: {"id":"chatcmpl-or-1","choices":[{"delta":{"content":"lo"},'
    '"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":2}}',
    "",
    "data: [DONE]",
    "",
]


@pytest.fixture(autouse=True)
def _clean_state(monkeypatch):
    """Isolate the module-level refusal map between tests."""
    session_cache.clear()
    dispatcher._provider_refusal_since.clear()
    yield
    dispatcher._provider_refusal_since.clear()
    session_cache.clear()


@pytest.fixture
def router(tmp_path, monkeypatch):
    """A chat-completions fixture with two providers in the catalog."""
    db_path = tmp_path / "test.db"
    conn = sqlite3.connect(db_path)
    conn.executescript(SCHEMA_SQL)
    # Seed a cheap model under the default (neuralwatt) provider.
    conn.execute(
        """
        INSERT INTO models (
            model_id, provider, base_model_id, tier, context_window,
            effective_context_window, max_output_tokens,
            cost_per_1m_prompt, cost_per_1m_completion,
            supports_vision, supports_json_mode,
            latency_class, reasoning_mode, context_variant,
            access_level, availability, last_updated
        ) VALUES (?, ?, ?, 2, 262128, 192500, 16384, ?, ?,
                  1, 1, 'standard', 'default', 'full', 'public', 'active',
                  '2026-08-22T00:00:00+00:00')
        """,
        (CHEAP, "neuralwatt", CHEAP, 0.30 / 3, 0.30),
    )
    # Seed an OpenRouter-only frontier model.
    conn.execute(
        """
        INSERT INTO models (
            model_id, provider, base_model_id, tier, context_window,
            effective_context_window, max_output_tokens,
            cost_per_1m_prompt, cost_per_1m_completion,
            supports_vision, supports_json_mode,
            latency_class, reasoning_mode, context_variant,
            access_level, availability, last_updated
        ) VALUES (?, 'openrouter', ?, 3, 262128, 192500, 16384, 0.10, 20.0,
                  1, 1, 'standard', 'default', 'full', 'public', 'active',
                  '2026-08-22T00:00:00+00:00')
        """,
        ("openai/gpt-6-astra", "openai/gpt-6-astra"),
    )
    conn.commit()
    conn.close()

    monkeypatch.setattr(dispatcher.cfg.database, "path", str(db_path))
    monkeypatch.setattr(dispatcher.cfg.verification, "local_llm_enabled", False)
    monkeypatch.setattr(dispatcher.cfg.local_vision, "enabled", False)
    monkeypatch.setattr(dispatcher.cfg.session_cache, "enabled", False)
    monkeypatch.setattr(dispatcher.cfg.exploration, "epsilon", 0.0)
    monkeypatch.setenv("NEURALWATT_API_KEY", "test-key")
    monkeypatch.setenv("OPENROUTER_API_KEY", "or-key")

    calls = []

    def fake_post(url, headers=None, json=None, stream=False, timeout=None):
        calls.append({"url": url, "body": json, "stream": stream})
        if stream:
            return _FakeResponse(lines=OPENROUTER_STREAM_LINES)
        return _FakeResponse(_completion(json["model"]))

    monkeypatch.setattr(dispatcher.requests, "post", fake_post)
    yield TestClient(dispatcher.app), calls, db_path


class _FakeResponse:
    """Minimal requests.Response stand-in."""

    def __init__(self, payload=None, *, status_code=200, lines=None):
        self.status_code = status_code
        self._payload = payload or {}
        self._lines = lines or []
        self.text = json.dumps(self._payload)
        self.headers = {}
        self.request = None

    def json(self):
        return self._payload

    @property
    def encoding(self):
        """The charset requests would pick, by requests' own rule.

        Derived rather than asserted, so this fake tracks requests instead of
        restating a belief about it. The default is a charset-less
        ``text/event-stream`` on purpose: that is what OpenRouter actually
        sends, and `get_encoding_from_headers` answers ISO-8859-1 for any
        ``text/*`` without a charset. Decoding a UTF-8 stream with that and
        re-encoding it produced mojibake on every streamed non-ASCII
        character -- so the hostile case is the DEFAULT here, and a fake can
        no longer make a broken passthrough look correct.
        """
        from requests.structures import CaseInsensitiveDict
        from requests.utils import get_encoding_from_headers

        headers = CaseInsensitiveDict(
            getattr(self, "headers", None) or {"Content-Type": "text/event-stream"}
        )
        return get_encoding_from_headers(headers) or "utf-8"

    def iter_lines(self, decode_unicode=False):
        # Faithful to requests in both modes: BYTES unless decode_unicode is
        # set, and when it is set the charset comes from Content-Type.
        #
        # The `wire` line is the load-bearing one. Whatever a test wrote into
        # `lines`, what a provider actually puts on the wire is UTF-8 bytes,
        # so that is what gets decoded. An earlier version of this fake passed
        # str lines straight through under decode_unicode, which modelled a
        # stream that had ALREADY been decoded correctly -- and a fake that
        # hands back the right answer cannot reproduce a charset bug. It let
        # the whole streaming suite pass against a proxy that was mangling
        # every non-ASCII character.
        for line in self._lines:
            wire = line.encode("utf-8") if isinstance(line, str) else line
            yield wire.decode(self.encoding) if decode_unicode else wire

    def close(self):
        pass


def _completion(model):
    return {
        "id": "chatcmpl-test-1",
        "model": model,
        "choices": [
            {
                "message": {"role": "assistant", "content": "hello there"},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 31, "completion_tokens": 12},
    }


def _energy_row(db_path):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    row = conn.execute(
        "SELECT model_id, energy_kwh, cost_usd FROM energy_observations ORDER BY id DESC LIMIT 1"
    ).fetchone()
    conn.close()
    return row


def test_openrouter_streaming_ignores_keepalive_comments_and_logs_null_telemetry(router):
    """A provider with has_energy_telemetry=false can emit `: OPENROUTER ...`
    comments without crashing the stream or logging fake telemetry.
    """
    client, calls, db_path = router
    # Wire the openrouter provider config (already in config.yaml, but ensure
    # the test fixture has it for clarity).
    dispatcher.cfg.dispatch_providers["openrouter"] = DispatchProvider(
        base_url="https://openrouter.ai/api/v1",
        api_key_env="OPENROUTER_API_KEY",
        has_energy_telemetry=False,
        enabled=True,
    )

    resp = client.post(
        "/v1/chat/completions",
        json={
            "model": "openai/gpt-6-astra",
            "messages": [{"role": "user", "content": "hi"}],
            "stream": True,
        },
    )

    assert resp.status_code == 200
    assert resp.headers["content-type"].startswith("text/event-stream")
    # The keepalive comment must reach the client unchanged.
    assert ": OPENROUTER PROCESSING" in resp.text
    # No crash on the SSE framing.
    assert "data: [DONE]" in resp.text

    row = _energy_row(db_path)
    assert row is not None
    assert row["model_id"] == "openai/gpt-6-astra"
    assert row["energy_kwh"] is None
    assert row["cost_usd"] is None


def test_provider_refusal_is_isolated_between_providers(router, monkeypatch):
    """A refusal on neuralwatt must skip the cloud classifier only when the
    default provider is neuralwatt; a request whose default provider is
    openrouter must still be allowed to use the cloud classifier.
    """
    client, calls, db_path = router

    # Configure a cloud classifier so the cascade has something to skip.
    from types import SimpleNamespace

    cloud_called = {"neuralwatt": False, "openrouter": False}

    def fake_cloud_classify(*args, **kwargs):
        cloud_called[dispatcher.cfg.dispatch_settings.default_provider] = True
        return {
            "task_category": "reasoning_math",
            "task_tier": 3,
            "required_context_tokens": 100,
            "confidence": 0.8,
        }

    monkeypatch.setattr(
        dispatcher.cfg.classifier,
        "cloud_fallback",
        SimpleNamespace(
            base_url="https://cloud.example/v1",
            model="cloud-classifier",
            timeout_seconds=2,
            api_key_env=None,
            max_output_tokens=1024,
        ),
    )
    monkeypatch.setattr(dispatcher, "_cloud_classifier_client", lambda cf: object())
    monkeypatch.setattr(session_cache, "classify_one", fake_cloud_classify)
    # Ensure classify() degrades so the cascade reaches the cloud step.
    monkeypatch.setattr(dispatcher, "classify", lambda task, context: dispatcher.Classification(
        task_category="general_chat",
        task_tier=2,
        required_context_tokens=10,
        confidence=0.0,
        source="fallback",
    ))

    # Record a refusal for neuralwatt.
    dispatcher._record_provider_refusal("neuralwatt")

    # With default_provider=neuralwatt, the cloud step should be skipped.
    monkeypatch.setattr(dispatcher.cfg.dispatch_settings, "default_provider", "neuralwatt")
    assert dispatcher._classify_cascade("sess-nw", "sys", "user") is None
    assert cloud_called["neuralwatt"] is False

    # With default_provider=openrouter, the refusal on neuralwatt must not apply.
    monkeypatch.setattr(dispatcher.cfg.dispatch_settings, "default_provider", "openrouter")
    result = dispatcher._classify_cascade("sess-or", "sys", "user")
    assert result is not None
    assert result.source == "classifier_cloud"
    assert cloud_called["openrouter"] is True


def test_extract_telemetry_empty_payload_has_no_energy_or_cost():
    """An empty payload yields an all-None Telemetry rather than an error."""
    telemetry = dispatcher.extract_telemetry({})
    assert telemetry.energy_kwh is None
    assert telemetry.cost_usd is None


def test_sniff_telemetry_line_ignores_foreign_comments():
    """SSE comment lines that do not carry energy/cost (e.g. OpenRouter
    keepalives) are deliberately ignored."""
    assert dispatcher._sniff_telemetry_line(": OPENROUTER PROCESSING") is None


def test_sniff_telemetry_line_parses_energy_comment():
    """A `: energy {...}` SSE comment line is parsed into a truthy
    (word, payload) tuple."""
    parsed = dispatcher._sniff_telemetry_line(": energy {}")
    assert parsed is not None
    word, payload = parsed
    assert word == "energy"
    assert payload == {}


# --- WI-1e: cost capture ------------------------------------------------


def _stream_lines_with_usage_cost():
    """SSE stream whose final chunk carries `usage.cost` (OpenRouter shape:
    the provider reports its bill inside the usage block, not in a top-level
    `cost` comment, and has has_energy_telemetry=false so comment lines are
    stripped)."""
    return [
        'data: {"id":"chatcmpl-or-1","choices":[{"delta":{"content":"hel"}}]}',
        "",
        'data: {"id":"chatcmpl-or-1","choices":[{"delta":{"content":"lo"},'
        '"finish_reason":"stop"}],"usage":{"prompt_tokens":5,'
        '"completion_tokens":2,"cost":0.005}}',
        "",
        "data: [DONE]",
        "",
    ]


def _openrouter_provider():
    return DispatchProvider(
        base_url="https://openrouter.ai/api/v1",
        api_key_env="OPENROUTER_API_KEY",
        has_energy_telemetry=False,
        reports_cost_in_usage=True,
        enabled=True,
    )


def _neuralwatt_provider():
    return DispatchProvider(
        base_url="https://api.neuralwatt.com/v1",
        api_key_env="NEURALWATT_API_KEY",
        has_energy_telemetry=True,
        reports_cost_in_usage=False,
        enabled=True,
    )


def test_dispatch_neuralwatt_flag_off_omits_usage_include(router):
    """Case 1: reports_cost_in_usage=false -> upstream_body carries no `usage`
    include key on the wire."""
    client, calls, db_path = router
    dispatcher.cfg.dispatch_providers["neuralwatt"] = _neuralwatt_provider()

    resp = client.post(
        "/v1/chat/completions",
        json={
            "model": CHEAP,
            "messages": [{"role": "user", "content": "hi"}],
            "stream": True,
        },
    )

    assert resp.status_code == 200
    assert "data: [DONE]" in resp.text
    body = calls[0]["body"]
    assert "usage" not in body


def test_dispatch_openrouter_flag_on_injects_usage_include(router):
    """Case 2: reports_cost_in_usage=true -> upstream_body carries
    `usage.include=true` on the wire."""
    client, calls, db_path = router
    dispatcher.cfg.dispatch_providers["openrouter"] = _openrouter_provider()

    resp = client.post(
        "/v1/chat/completions",
        json={
            "model": "openai/gpt-6-astra",
            "messages": [{"role": "user", "content": "hi"}],
            "stream": True,
        },
    )

    assert resp.status_code == 200
    assert "data: [DONE]" in resp.text
    body = calls[0]["body"]
    assert body.get("usage") == {"include": True}


def test_extract_telemetry_cost_from_usage_buffered():
    """Case 3: flag on, buffered — cost_usd comes from usage.cost when the
    payload carries no cost.request_cost_usd."""
    telemetry = dispatcher.extract_telemetry({}, usage={"cost": 0.005})
    assert telemetry.cost_usd == 0.005


def test_extract_telemetry_cost_payload_wins_over_usage():
    """Case 5: cost.request_cost_usd takes precedence over usage.cost when
    both blocks are present."""
    telemetry = dispatcher.extract_telemetry(
        {"cost": {"request_cost_usd": 0.01}}, usage={"cost": 0.005}
    )
    assert telemetry.cost_usd == 0.01


def test_extract_telemetry_streamed_cost_from_usage(router, monkeypatch):
    """Case 4: the streamed path logs non-null cost_usd when usage.cost is
    present in the streamed finally block.

    This is the trap the whole fix exists for: comment lines carrying cost are
    stripped, so the cost must survive via the explicit `usage` parameter
    rather than re-derivation from the (stripped) payload.
    """
    client, calls, db_path = router
    dispatcher.cfg.dispatch_providers["openrouter"] = _openrouter_provider()

    def fake_post(url, headers=None, json=None, stream=False, timeout=None):
        calls.append({"url": url, "body": json, "stream": stream})
        if stream:
            return _FakeResponse(lines=_stream_lines_with_usage_cost())
        return _FakeResponse(_completion(json["model"]))

    monkeypatch.setattr(dispatcher.requests, "post", fake_post)

    resp = client.post(
        "/v1/chat/completions",
        json={
            "model": "openai/gpt-6-astra",
            "messages": [{"role": "user", "content": "hi"}],
            "stream": True,
        },
    )

    assert resp.status_code == 200
    assert "data: [DONE]" in resp.text
    row = _energy_row(db_path)
    assert row is not None
    assert row["cost_usd"] == 0.005


def test_extract_telemetry_empty_payload_cost_none():
    """Case 7: extract_telemetry({}) yields cost_usd=None."""
    telemetry = dispatcher.extract_telemetry({})
    assert telemetry.cost_usd is None


def test_extract_telemetry_cached_prompt_tokens_from_usage():
    """Case 8: cached_prompt_tokens read from usage.prompt_tokens_details
    even when the payload carries nothing else."""
    telemetry = dispatcher.extract_telemetry(
        {}, usage={"prompt_tokens_details": {"cached_tokens": 42}}
    )
    assert telemetry.cached_prompt_tokens == 42


# --- 1.1: the accounting opt-in reaches every dispatch path -----------------


def test_buffered_openrouter_carries_usage_include(router):
    """The NON-streaming path must carry the opt-in too.

    It builds its own `attempt_body` from the client body rather than from
    `upstream_body`, so the opt-in set alongside `stream_options` never
    reached it: every buffered OpenRouter dispatch asked for no accounting and
    logged NULL cost and NULL cached tokens.
    """
    client, calls, db_path = router
    dispatcher.cfg.dispatch_providers["openrouter"] = _openrouter_provider()

    resp = client.post(
        "/v1/chat/completions",
        json={
            "model": "openai/gpt-6-astra",
            "messages": [{"role": "user", "content": "hi"}],
        },
    )

    assert resp.status_code == 200
    body = calls[-1]["body"]
    assert calls[-1]["stream"] is False
    assert body.get("usage") == {"include": True}


def test_buffered_neuralwatt_omits_usage_include(router):
    """The other half of the same rule: a provider with the flag off must not
    have an opt-in invented for it on the buffered path either."""
    client, calls, db_path = router
    dispatcher.cfg.dispatch_providers["neuralwatt"] = _neuralwatt_provider()

    resp = client.post(
        "/v1/chat/completions",
        json={"model": CHEAP, "messages": [{"role": "user", "content": "hi"}]},
    )

    assert resp.status_code == 200
    assert "usage" not in calls[-1]["body"]


def test_client_supplied_usage_block_is_not_overwritten():
    """setdefault, not assignment: a client that sent its own `usage` block
    asked for something specific and keeps it."""
    dispatcher.cfg.dispatch_providers["openrouter"] = _openrouter_provider()
    body = dispatcher._apply_provider_usage_opts(
        {"usage": {"include": False}}, "openrouter"
    )
    assert body["usage"] == {"include": False}


def test_unknown_provider_name_adds_nothing():
    """A provider name with no config entry is not a crash and not an opt-in."""
    assert dispatcher._apply_provider_usage_opts({}, "no-such-provider") == {}


def test_streaming_failover_uses_the_candidates_flag_not_the_first_choice(
    router, monkeypatch
):
    """Failover crosses providers, and the opt-in belongs to whoever is being
    posted to.

    `upstream_body` was built once from the initially-selected provider, so a
    failover onto OpenRouter posted a body with no accounting opt-in -- the
    exact requests most likely to need one, since they are the retries.
    """
    client, calls, db_path = router
    dispatcher.cfg.dispatch_providers["neuralwatt"] = _neuralwatt_provider()
    dispatcher.cfg.dispatch_providers["openrouter"] = _openrouter_provider()
    monkeypatch.setattr(dispatcher.cfg.circuit_breaker, "enabled", False)

    stream_calls = []

    def fake_post(url, headers=None, json=None, stream=False, timeout=None):
        calls.append({"url": url, "body": json, "stream": stream})
        if not stream:
            return _FakeResponse(_completion(json["model"]))
        stream_calls.append(json)
        # First candidate fails, so the router fails over to the runner-up,
        # which the fixture places on the other provider.
        if len(stream_calls) == 1:
            return _FakeResponse({"error": "nope"}, status_code=500)
        return _FakeResponse(lines=_stream_lines_with_usage_cost())

    monkeypatch.setattr(dispatcher.requests, "post", fake_post)

    resp = client.post(
        "/v1/chat/completions",
        json={
            "model": "auto",
            "messages": [{"role": "user", "content": "refactor this"}],
            "stream": True,
        },
    )

    assert resp.status_code == 200
    assert len(stream_calls) == 2, "expected a failover onto the second provider"
    first, second = stream_calls
    # The two candidates sit on different providers, so exactly one body
    # carries the opt-in -- and which one is decided by the candidate, not by
    # the request's first choice.
    assert first["model"] != second["model"]
    by_provider = {
        CHEAP: "neuralwatt", "openai/gpt-6-astra": "openrouter",
    }
    for sent in (first, second):
        wants = by_provider[sent["model"]] == "openrouter"
        assert ("usage" in sent) is wants, (
            f"{sent['model']} on {by_provider[sent['model']]} "
            f"{'lost' if wants else 'gained'} the accounting opt-in"
        )


# --- 1.2: an absent cached_tokens is distinguishable from a reported zero ---
#
# Three provider shapes used to collapse into two stored outcomes, so "the
# provider said zero" could not be told from "the provider said nothing". The
# difference decides whether every cache rate computed from the rows that DO
# carry a count is conditioned on a hit having occurred.


def _cached_cell(db_path):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    row = conn.execute(
        "SELECT cached_prompt_tokens, cached_tokens_source "
        "FROM energy_observations ORDER BY id DESC LIMIT 1"
    ).fetchone()
    conn.close()
    return (row["cached_prompt_tokens"], row["cached_tokens_source"])


def test_cached_tokens_reported_zero_is_a_measurement():
    """(a) prompt_tokens_details present, cached_tokens: 0. An explicit zero
    is the provider measuring a full cache miss, and must survive as 0."""
    telemetry = dispatcher.extract_telemetry(
        {}, usage={"prompt_tokens_details": {"cached_tokens": 0}}
    )
    assert telemetry.cached_prompt_tokens == 0
    assert telemetry.cached_tokens_source == "reported"


def test_cached_tokens_details_without_the_key_is_not_a_zero():
    """(b) prompt_tokens_details present, cached_tokens absent. The provider
    described the prompt and said nothing about caching -- which is not a
    measurement of zero, so no value is invented."""
    telemetry = dispatcher.extract_telemetry(
        {}, usage={"prompt_tokens_details": {"audio_tokens": 0}}
    )
    assert telemetry.cached_prompt_tokens is None
    assert telemetry.cached_tokens_source == "details_no_count"


def test_cached_tokens_no_details_block_at_all():
    """(c) no prompt_tokens_details. Distinct from (b): a provider that never
    reports the breakdown is a different fact from one that reported a
    breakdown with no cached count in it."""
    telemetry = dispatcher.extract_telemetry(
        {}, usage={"prompt_tokens": 100, "completion_tokens": 5}
    )
    assert telemetry.cached_prompt_tokens is None
    assert telemetry.cached_tokens_source == "no_details"


def test_cached_tokens_json_null_groups_with_no_count():
    """A JSON null under an existing key is the provider producing the slot
    and declining to fill it. It groups with (b), never with a zero, which
    keeps the invariant that a non-NULL count means source='reported'."""
    telemetry = dispatcher.extract_telemetry(
        {}, usage={"prompt_tokens_details": {"cached_tokens": None}}
    )
    assert telemetry.cached_prompt_tokens is None
    assert telemetry.cached_tokens_source == "details_no_count"


def test_cached_tokens_details_not_a_dict_does_not_raise():
    """`prompt_tokens_details: null` reached `.get` on None and raised
    AttributeError inside the streamed finally block."""
    telemetry = dispatcher.extract_telemetry(
        {}, usage={"prompt_tokens_details": None}
    )
    assert telemetry.cached_prompt_tokens is None
    assert telemetry.cached_tokens_source == "no_details"


def test_cached_tokens_read_from_a_buffered_payloads_own_usage():
    """The buffered call sites pass a payload and no explicit `usage`, so the
    count was dropped for every one of them. Cost already resolved its source
    this way; cached tokens now do too."""
    telemetry = dispatcher.extract_telemetry(
        {"usage": {"prompt_tokens_details": {"cached_tokens": 7}}}
    )
    assert telemetry.cached_prompt_tokens == 7
    assert telemetry.cached_tokens_source == "reported"


def test_three_cached_shapes_store_three_distinct_outcomes(router):
    """The point of the whole item: (a), (b) and (c) must be distinguishable
    in the DB, not just in the parse."""
    client, calls, db_path = router
    shapes = {
        "reported_zero": {"prompt_tokens_details": {"cached_tokens": 0}},
        "details_no_count": {"prompt_tokens_details": {"audio_tokens": 0}},
        "no_details": {"prompt_tokens": 100},
    }
    stored = {}
    for name, usage in shapes.items():
        dispatcher.log_observation(
            CHEAP, "neuralwatt", "coding_general", f"req-{name}",
            prompt_tokens=100, completion_tokens=5,
            telemetry=dispatcher.extract_telemetry({}, usage=usage),
        )
        stored[name] = _cached_cell(db_path)

    assert stored["reported_zero"] == (0, "reported")
    assert stored["details_no_count"] == (None, "details_no_count")
    assert stored["no_details"] == (None, "no_details")
    assert len(set(stored.values())) == 3