347 lines
12 KiB
Python
347 lines
12 KiB
Python
"""Task 3 acceptance: conversation identity wired into chat_completions.
|
|
|
|
These tests exercise the integration the wiring added end to end: a client
|
|
``X-Router-Conversation`` header must scope the recorded ``session_key`` (and
|
|
therefore the session incumbent, the cloud ``energy_observations`` ledger and
|
|
the local-dispatch energy ledger), while ``X-Router-Agent``/``X-Router-Parent``
|
|
must land in their columns. A headerless request must still behave exactly as
|
|
before — one fingerprint key, NULL agent/parent, byte-identical to the frozen
|
|
pre-wiring snapshot (``tests/fixtures/route_decision_no_header.json``).
|
|
|
|
The commit that added the snapshot fixture deliberately precedes this file, so
|
|
a clean checkout can read it without a gitignored ``.omo/evidence/`` path.
|
|
"""
|
|
|
|
import sqlite3
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
import dispatcher
|
|
|
|
from test_no_header_snapshot import (
|
|
assert_matches_snapshot,
|
|
make_router,
|
|
no_header_payload,
|
|
)
|
|
|
|
CHEAP = "cheap-model"
|
|
DEAR = "dear-model"
|
|
|
|
|
|
@pytest.fixture(autouse=True, scope="module")
|
|
def _clear_prefix_probe_state():
|
|
"""Clean up ``prefix_probe._store`` after our module so the snapshot
|
|
test (runs alphabetically after us) sees a clean prefix-probe cache.
|
|
The probe stores per-session fingerprints in process-memory and a
|
|
POST with the snapshot's payload here would otherwise make a later
|
|
test see a prior turn and report non-NULL prefix divergence."""
|
|
yield
|
|
import prefix_probe
|
|
prefix_probe._store.clear()
|
|
|
|
_SYS = "You work in /home/user/proj and follow the design docs."
|
|
|
|
|
|
def _messages(*roles):
|
|
return [
|
|
{"role": r, "content": c}
|
|
for r, c in zip(roles, ["system prompt", "user query one", "assistant reply", "user query two"])
|
|
][: len(roles)]
|
|
|
|
|
|
def _rows(db_path, table, where=None, order="id"):
|
|
conn = sqlite3.connect(db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
q = f"SELECT * FROM {table}"
|
|
if where:
|
|
q += f" WHERE {where}"
|
|
q += f" ORDER BY {order}"
|
|
rows = [dict(r) for r in conn.execute(q).fetchall()]
|
|
conn.close()
|
|
return rows
|
|
|
|
|
|
def _latest_decision(db_path):
|
|
rows = _rows(db_path, "route_decisions")
|
|
assert rows, "no route_decisions row"
|
|
return rows[-1]
|
|
|
|
|
|
def _post(client, *, headers=None, model="auto", msgs=None):
|
|
payload = {
|
|
"model": model,
|
|
"messages": msgs if msgs is not None else _messages("system", "user"),
|
|
}
|
|
return client.post(
|
|
"/v1/chat/completions", json=payload, headers=headers or {}
|
|
)
|
|
|
|
|
|
# --- Acceptance 1: same prompt, different conversation -> distinct keys -----
|
|
|
|
|
|
def test_same_system_prompt_distinct_conversations_get_distinct_session_keys(
|
|
tmp_path, monkeypatch
|
|
):
|
|
client, db_path = make_router(tmp_path, monkeypatch)
|
|
msgs = _messages("system", "user")
|
|
|
|
r1 = _post(client, headers={"X-Router-Conversation": "a"}, msgs=msgs)
|
|
r2 = _post(client, headers={"X-Router-Conversation": "b"}, msgs=msgs)
|
|
assert r1.status_code == 200 and r2.status_code == 200
|
|
|
|
keys = {r["session_key"] for r in _rows(db_path, "route_decisions")}
|
|
assert keys == {"c:a", "c:b"}, keys
|
|
|
|
|
|
# --- Acceptance 2: each conversation's incumbent is its OWN last model -----
|
|
|
|
|
|
def test_incumbent_is_conversation_scoped_not_global(
|
|
tmp_path, monkeypatch
|
|
):
|
|
client, db_path = make_router(tmp_path, monkeypatch)
|
|
conn = sqlite3.connect(db_path)
|
|
for sk, model in (
|
|
("c:a", CHEAP), # interleave so the GLOBAL last chat row is cheap
|
|
("c:b", DEAR),
|
|
("c:a", CHEAP),
|
|
):
|
|
conn.execute(
|
|
"""INSERT INTO route_decisions (
|
|
observed_at, kind, session_key, selected_provider,
|
|
selected_model, task_category, task_tier
|
|
) VALUES ('2026-09-13T00:00:00+00:00', 'chat', ?, 'neuralwatt',
|
|
?, 'coding_general', 2)""",
|
|
(sk, model),
|
|
)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
monkeypatch.setattr(
|
|
dispatcher.cfg.objective, "incumbent_cache_pricing", True
|
|
)
|
|
|
|
captured = {}
|
|
|
|
class Capturer:
|
|
def __call__(self, rows, **kwargs):
|
|
captured["incumbent"] = kwargs.get("incumbent")
|
|
from routing import rank_candidates as _real
|
|
return _real(rows, **kwargs)
|
|
|
|
monkeypatch.setattr(dispatcher, "rank_candidates", Capturer())
|
|
|
|
r = _post(client, headers={"X-Router-Conversation": "b"})
|
|
assert r.status_code == 200, r.text
|
|
|
|
# Conversation b's own last chat row is dear-model, even though the
|
|
# overall-last chat row belongs to conversation a (cheap-model).
|
|
assert captured["incumbent"] == ("neuralwatt", DEAR), captured["incumbent"]
|
|
|
|
|
|
# --- Acceptance 3: headerless requests land under one fingerprint, NULL
|
|
# agent/parent, byte-identical to the committed snapshot -----------------
|
|
|
|
|
|
def test_headerless_requests_reproduce_committed_snapshot(
|
|
tmp_path, monkeypatch
|
|
):
|
|
client, db_path = make_router(tmp_path, monkeypatch)
|
|
# The exact payload the committed snapshot was frozen from: the row must
|
|
# be byte-identical, so the messages (and therefore the fingerprint,
|
|
# pinch and prefix-probe values) must match the fixture's.
|
|
payload = no_header_payload()
|
|
|
|
# First POST: fresh session, no prior prefix-probe state → the row has
|
|
# NULL prefix columns and is byte-identical to the pre-wiring fixture.
|
|
r1 = _post(client, msgs=payload["messages"])
|
|
assert r1.status_code == 200, r1.text
|
|
assert_matches_snapshot(db_path)
|
|
|
|
# Second POST in the same session: the prefix probe now sees a prior
|
|
# turn and writes non-NULL prefix columns, so the byte-identical
|
|
# assertion cannot include this row. But both rows must share one
|
|
# fingerprint key and carry NULL agent/parent.
|
|
r2 = _post(client, msgs=payload["messages"])
|
|
assert r2.status_code == 200, r2.text
|
|
rows = _rows(db_path, "route_decisions")
|
|
assert len(rows) == 2
|
|
assert rows[0]["session_key"] == rows[1]["session_key"]
|
|
assert rows[0]["session_key"] is not None
|
|
for r in rows:
|
|
assert r["agent"] is None
|
|
assert r["parent_key"] is None
|
|
|
|
|
|
# --- Acceptance 4: agent and parent land in their columns -------------------
|
|
|
|
|
|
def test_agent_and_parent_land_in_columns(tmp_path, monkeypatch):
|
|
client, db_path = make_router(tmp_path, monkeypatch)
|
|
r = _post(
|
|
client,
|
|
headers={
|
|
"X-Router-Conversation": "conv1",
|
|
"X-Router-Agent": "my-agent",
|
|
"X-Router-Parent": "parent1",
|
|
},
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
row = _latest_decision(db_path)
|
|
assert row["session_key"] == "c:conv1"
|
|
assert row["agent"] == "my-agent"
|
|
assert row["parent_key"] == "c:parent1"
|
|
|
|
|
|
# --- Acceptance 5: invalid conversation header falls back to fingerprint ----
|
|
|
|
|
|
def test_invalid_conversation_header_falls_back_to_fingerprint(
|
|
tmp_path, monkeypatch
|
|
):
|
|
client, db_path = make_router(tmp_path, monkeypatch)
|
|
msgs = _messages("system", "user")
|
|
|
|
bad = _post(
|
|
client,
|
|
headers={"X-Router-Conversation": "has space and /slash"},
|
|
msgs=msgs,
|
|
)
|
|
assert bad.status_code == 200, bad.text
|
|
|
|
good = _post(client, msgs=msgs)
|
|
assert good.status_code == 200, good.text
|
|
|
|
rows = _rows(db_path, "route_decisions")
|
|
bad_key, fp_key = rows[0]["session_key"], rows[1]["session_key"]
|
|
# The invalid header is treated as absent: same fingerprint as no header.
|
|
assert bad_key == fp_key
|
|
assert not bad_key.startswith("c:")
|
|
|
|
|
|
# --- Acceptance 6: energy ledgers get the conversation key ------------------
|
|
|
|
def test_cloud_energy_session_key_starts_with_c(tmp_path, monkeypatch):
|
|
client, db_path = make_router(tmp_path, monkeypatch)
|
|
r = _post(
|
|
client,
|
|
headers={"X-Router-Conversation": "ener"},
|
|
model="auto",
|
|
msgs=_messages("system", "user"),
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
rows = _rows(db_path, "energy_observations")
|
|
assert rows, "expected an energy_observations row"
|
|
for row in rows:
|
|
assert (row["session_key"] or "").startswith("c:"), row
|
|
|
|
|
|
def _local_dispatch_harness(tmp_path, monkeypatch, session_key):
|
|
"""Drive one real _run_local_dispatch and return the produced row."""
|
|
client, db_path = make_router(tmp_path, monkeypatch)
|
|
model_id = "local-q"
|
|
|
|
monkeypatch.setattr(dispatcher.cfg.local_energy, "enabled", True)
|
|
# local_energy_dispatch_models is a read-only cached property on the
|
|
# Pydantic config; shadow it on the class so the metering gate sees our
|
|
# model as meterable.
|
|
monkeypatch.setattr(
|
|
type(dispatcher.cfg), "local_energy_dispatch_models",
|
|
property(lambda self: frozenset({model_id})),
|
|
)
|
|
monkeypatch.setattr(dispatcher.cfg.circuit_breaker, "enabled", False)
|
|
|
|
class FakeMeasurement:
|
|
avg_power_watts = 100.0
|
|
duration_seconds = 2.0
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
return False
|
|
|
|
monkeypatch.setattr(
|
|
dispatcher.local_energy, "measure",
|
|
lambda *a, **k: FakeMeasurement(),
|
|
)
|
|
|
|
def fake_post(url, headers=None, json=None, stream=False, timeout=None):
|
|
resp = SimpleNamespace()
|
|
resp.status_code = 200
|
|
resp.json = lambda: {
|
|
"id": "ollama-local-1",
|
|
"choices": [
|
|
{"message": {"role": "assistant", "content": "local answer"}}
|
|
],
|
|
"usage": {"prompt_tokens": 20, "completion_tokens": 8},
|
|
}
|
|
resp.text = "{}"
|
|
return resp
|
|
|
|
monkeypatch.setattr(dispatcher.requests, "post", fake_post)
|
|
|
|
entry = SimpleNamespace(
|
|
model_id=model_id,
|
|
base_url="http://localhost:11434/v1",
|
|
api_key_env=None,
|
|
max_output_tokens=2048,
|
|
timeout_seconds=60.0,
|
|
)
|
|
result = dispatcher._run_local_dispatch(
|
|
entry,
|
|
_messages("system", "user"),
|
|
category=None,
|
|
body={},
|
|
session_key=session_key,
|
|
)
|
|
assert result["payload"]["choices"][0]["message"]["content"] == "local answer"
|
|
assert result["request_id"].startswith("local-dispatch-")
|
|
rows = _rows(db_path, "local_energy_observations")
|
|
assert rows, "expected a local_energy_observations row"
|
|
return rows, client
|
|
|
|
|
|
def test_local_dispatch_answer_records_conversation_key(
|
|
tmp_path, monkeypatch
|
|
):
|
|
rows, _ = _local_dispatch_harness(tmp_path, monkeypatch, session_key="c:ener")
|
|
assert rows[-1]["session_key"] == "c:ener"
|
|
|
|
|
|
# --- Acceptance 7 [H5]: /dispatch (no header) keeps local energy NULL -------
|
|
|
|
|
|
def test_no_header_local_dispatch_keeps_energy_null_and_resolves(
|
|
tmp_path, monkeypatch
|
|
):
|
|
# /dispatch passes no session_key to _run_local_dispatch, so the keyword
|
|
# defaults to None and the local energy row must stay NULL.
|
|
rows, _ = _local_dispatch_harness(tmp_path, monkeypatch, session_key=None)
|
|
assert rows[-1]["session_key"] is None
|
|
|
|
|
|
# --- Acceptance 8 [H4]: cold start — sensible routing on a short history ----
|
|
|
|
|
|
def test_cold_start_conversation_routes_sensibly_on_short_history(
|
|
tmp_path, monkeypatch
|
|
):
|
|
client, db_path = make_router(tmp_path, monkeypatch)
|
|
# A short sub-agent history (3 messages) with a fresh conversation key
|
|
# that has NO prior route_decisions rows: must still resolve to a real
|
|
# model instead of failing or guessing an incumbent that cannot exist.
|
|
msgs = _messages("system", "user", "assistant")
|
|
r = _post(
|
|
client,
|
|
headers={"X-Router-Conversation": "cold"},
|
|
msgs=msgs,
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
row = _latest_decision(db_path)
|
|
assert row["kind"] == "chat"
|
|
assert row["session_key"] == "c:cold"
|
|
assert row["selected_model"] is not None
|