#!/usr/bin/env bats
# Seam: the frank-medium proxy's routing decisions.
# The proxy is the only thing between the fleet's clients and the two
# single-GPU replicas, and its routing rules carry the performance work: spread
# single-turn requests (batch members, one-shot questions) so both cards prefill
# in parallel, and pin continuations so an agent loop keeps hitting the replica
# whose prefix cache holds its context. This seam drives the routing policy
# directly; the first version of these rules shipped a bug that only appeared on
# a continuation request, which is exactly the class of mistake a test this
# cheap should catch.

load helpers

ROUTING="$REPO_ROOT/dankserver/frank-medium/proxy/routing.py"
PROXY="$REPO_ROOT/dankserver/frank-medium/proxy/affinity-proxy.py"

setup() {
    DRIVER="$BATS_TEST_TMPDIR/route.py"
    touch "$DRIVER"
}

@test "the proxy and its routing policy exist and compile" {
    [ -f "$ROUTING" ]
    [ -f "$PROXY" ]
    python3 -m py_compile "$ROUTING"
    python3 -m py_compile "$PROXY"
    # The policy must not drag the HTTP stack in: it is imported by tests and
    # reasoned about without aiohttp, a network or a GPU.
    ! grep -q 'import aiohttp' "$ROUTING"
}

@test "a batch of single-turn requests spreads across both cards" {
    cat > "$DRIVER" <<PY
import importlib.util, json, os, sys
os.environ["ROUTING"] = "adaptive"
spec = importlib.util.spec_from_file_location("routing", "$ROUTING")
mod = importlib.util.module_from_spec(spec)
sys.modules["routing"] = mod
spec.loader.exec_module(mod)
class R:
    remote = "127.0.0.1"
shared = "shared document " + "y" * 8000
picks = []
for i in range(6):
    text = shared if i % 2 else f"distinct prompt {i} " + "x" * 8000
    body = json.dumps({"messages": [{"role": "user", "content": text}]}).encode()
    picks.append(mod.pick(*mod.affinity_key(body, R.remote), [True, True]))
print(" ".join(str(p) for p in picks))
PY
    run python3 "$DRIVER"
    [ "$status" -eq 0 ]
    read -r -a picks <<< "$output"
    used=$(printf '%s\n' "${picks[@]}" | sort -u | wc -l)
    # Both cards must work on a batch, or one card prefills an entire benchmark
    # while the other idles.
    [ "$used" -eq 2 ]
    # Consecutive members alternate (the counter is advanced only by requests
    # that are not continuations), including the two that share a document.
    for ((i = 0; i < 5; i++)); do
        [ "${picks[i]}" != "${picks[i + 1]}" ]
    done
}

@test "a continuation is pinned to the replica that served its prefix" {
    cat > "$DRIVER" <<PY
import importlib.util, json, os, sys
os.environ["ROUTING"] = "adaptive"
spec = importlib.util.spec_from_file_location("routing", "$ROUTING")
mod = importlib.util.module_from_spec(spec)
sys.modules["routing"] = mod
spec.loader.exec_module(mod)
class R:
    remote = "127.0.0.1"
prefix = "system prompt " + "z" * 8000
first = json.dumps({"messages": [{"role": "user", "content": prefix + " turn one"}]}).encode()
start = mod.pick(*mod.affinity_key(first, R.remote), [True, True])
# Same leading bytes, longer history and a different tail each turn: an agent
# loop resending its context.
turns = []
history = [{"role": "assistant", "content": "ack"}]
for tail in ("turn two", "turn three", "turn four"):
    history.append({"role": "user", "content": "question " + tail})
    body = json.dumps({"messages": [{"role": "user", "content": prefix + " turn one"}] + history}).encode()
    turns.append(mod.pick(*mod.affinity_key(body, R.remote), [True, True]))
    history.append({"role": "assistant", "content": "ack again"})
print(start, " ".join(str(t) for t in turns))
PY
    run python3 "$DRIVER"
    [ "$status" -eq 0 ]
    read -r -a out <<< "$output"
    # Every continuation agrees with the first turn's replica: that is where
    # the prefix cache entry lives.
    [ "${out[1]}" = "${out[0]}" ]
    [ "${out[2]}" = "${out[0]}" ]
    [ "${out[3]}" = "${out[0]}" ]
}

@test "no routing mode ever picks a replica that is down" {
    for mode in adaptive affinity roundrobin; do
        cat > "$DRIVER" <<PY
import importlib.util, json, os, sys
os.environ["ROUTING"] = "$mode"
spec = importlib.util.spec_from_file_location("routing", "$ROUTING")
mod = importlib.util.module_from_spec(spec)
sys.modules["routing"] = mod
spec.loader.exec_module(mod)
class R:
    remote = "127.0.0.1"
body = json.dumps({"messages": [{"role": "user", "content": "x" * 400}]}).encode()
one_down = mod.pick(*mod.affinity_key(body, R.remote), [False, True])
both_down = mod.pick(*mod.affinity_key(body, R.remote), [False, False])
print(one_down, both_down)
PY
        run python3 "$DRIVER"
        [ "$status" -eq 0 ]
        # Replica 0 down -> replica 1; both down -> -1, and the proxy answers
        # 503 rather than forwarding into the void.
        [ "$output" = "1 -1" ]
    done
}

@test "affinity is deterministic per prefix, roundrobin is not" {
    cat > "$DRIVER" <<PY
import importlib.util, json, os, sys
os.environ["ROUTING"] = "affinity"
spec = importlib.util.spec_from_file_location("routing", "$ROUTING")
mod = importlib.util.module_from_spec(spec)
sys.modules["routing"] = mod
spec.loader.exec_module(mod)
class R:
    remote = "127.0.0.1"
def pick(text):
    body = json.dumps({"messages": [{"role": "user", "content": text}]}).encode()
    return mod.pick(*mod.affinity_key(body, R.remote), [True, True])
doc = "d" * 8000
# Same leading 4 KiB, different tails: affinity keys on the prefix, so both go
# to one replica. Different documents may hash anywhere — that is the point of
# hashing — so only the prefix property is asserted.
print(pick(doc + " a"), pick(doc + " b"))
PY
    run python3 "$DRIVER"
    [ "$status" -eq 0 ]
    read -r -a aff <<< "$output"
    [ "${aff[0]}" = "${aff[1]}" ]
}