feat(auto-router): scope shadow eval jobs to multiple keys

A shadow eval job now covers a set of keys instead of exactly one, and each
key carries its own max_turns budget, so one key exhausting its budget leaves
its siblings sampling. Key, budget and stop state move to a new
LiteLLM_ShadowEvalJobKey child table, taking the one-active-per-key partial
unique index with them so the race-safe DB guarantee is unchanged. Attempts
now record api_key_id, which is what makes per-key turn counting and the new
by_key result slice possible

The dashboard picker becomes a multi-select and the detail view gains a per
key breakdown showing each key's own budget and status
This commit is contained in:
Tin Chi Lo 2026-08-13 19:09:49 -07:00
parent d7d10be063
commit 4360c9fc43
19 changed files with 1398 additions and 325 deletions

View file

@ -0,0 +1,59 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_ShadowEvalJobKey" (
"id" TEXT NOT NULL,
"job_id" TEXT NOT NULL,
"direction" TEXT NOT NULL,
"api_key_id" TEXT NOT NULL,
"max_turns" INTEGER NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"stopped_at" TIMESTAMP(3),
CONSTRAINT "LiteLLM_ShadowEvalJobKey_pkey" PRIMARY KEY ("id")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJobKey_job_id_idx" ON "LiteLLM_ShadowEvalJobKey"("job_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJobKey_api_key_id_idx" ON "LiteLLM_ShadowEvalJobKey"("api_key_id");
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_id_direction_key" ON "LiteLLM_ShadowEvalJob"("id", "direction");
DO $$ BEGIN
ALTER TABLE "LiteLLM_ShadowEvalJobKey" ADD CONSTRAINT "LiteLLM_ShadowEvalJobKey_job_id_direction_fkey"
FOREIGN KEY ("job_id", "direction") REFERENCES "LiteLLM_ShadowEvalJob"("id", "direction")
ON DELETE CASCADE ON UPDATE CASCADE;
EXCEPTION
WHEN duplicate_object THEN NULL;
END $$;
-- One active job per key per direction, enforced by the database rather than a read-then-create in
-- the start endpoint, which races against a concurrent start on another pod. Partial indexes are not
-- expressible in schema.prisma, so this lives here only. Active means not yet stopped; the start
-- endpoint stamps stopped_at on exhausted key rows before creating.
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJobKey_one_active_per_key_direction"
ON "LiteLLM_ShadowEvalJobKey"("api_key_id", "direction") WHERE "stopped_at" IS NULL;
INSERT INTO "LiteLLM_ShadowEvalJobKey" ("id", "job_id", "direction", "api_key_id", "max_turns", "created_at", "stopped_at")
SELECT 'jobkey_' || "id", "id", "direction", "api_key_id", "max_turns", "created_at", "stopped_at"
FROM "LiteLLM_ShadowEvalJob"
ON CONFLICT DO NOTHING;
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "api_key_id" TEXT;
UPDATE "LiteLLM_ShadowEvalAttempt" a SET "api_key_id" = j."api_key_id"
FROM "LiteLLM_ShadowEvalJob" j WHERE a."job_id" = j."id" AND a."api_key_id" IS NULL;
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ALTER COLUMN "api_key_id" SET NOT NULL;
CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalAttempt_job_id_api_key_id_idx" ON "LiteLLM_ShadowEvalAttempt"("job_id", "api_key_id");
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalAttempt_job_id_idx";
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_one_active_per_key_direction";
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_api_key_id_idx";
-- Key scope, per-key budget and stop state moved into LiteLLM_ShadowEvalJobKey above. Dropping them
-- in the same migration is safe because no tagged release ships shadow eval, so no pod is serving
-- the single key code and there is no mixed version window that needs these columns readable.
ALTER TABLE "LiteLLM_ShadowEvalJob" DROP COLUMN IF EXISTS "api_key_id",
DROP COLUMN IF EXISTS "max_turns",
DROP COLUMN IF EXISTS "stopped_at";

View file

@ -1450,36 +1450,57 @@ model LiteLLM_AutoRouterSession {
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
}
// Shadow eval: evaluation of an auto-router against a key's live traffic, in either
// direction. forward duplicates the requests the key did not route through the router
// through it, answering whether the key should adopt it; reverse duplicates the requests
// Shadow eval: evaluation of an auto-router against a set of keys' live traffic, in either
// direction. forward duplicates the requests the keys did not route through the router
// through it, answering whether they should adopt it; reverse duplicates the requests
// the router did serve against a fixed baseline model, answering whether a key already on
// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge
// compares real vs shadow responses blind. The job row is immutable config plus
// stopped_at; every count, status, and spend figure is derived from the append-only
// attempt rows, so nothing can disagree across pods or stop races.
// compares real vs shadow responses blind. The job row is immutable config; the keys it
// scopes, their budgets and their stop state live on LiteLLM_ShadowEvalJobKey, and every
// count, status, and spend figure is derived from the append-only attempt rows, so nothing
// can disagree across pods or stop races.
model LiteLLM_ShadowEvalJob {
id String @id @default(cuid())
api_key_id String // hashed virtual key whose traffic is shadowed
router_name String // the auto-router under evaluation, in either direction
direction String @default("forward") // forward | reverse
baseline_model String? // reverse only: the fixed model the router is judged against
judge_model String
shadow_percentage Float
max_turns Int // sample budget: judge at most this many turns
created_at DateTime @default(now())
created_by String?
ends_at DateTime
stopped_at DateTime?
keys LiteLLM_ShadowEvalJobKey[]
@@index([api_key_id])
@@unique([id, direction])
@@index([created_at])
}
// One row per key a job shadows. Budget and stop state are per key, so one key exhausting
// its turns never ends a sibling's sampling. "One active job per (key, direction)" is a
// partial unique index on (api_key_id, direction) WHERE stopped_at IS NULL, expressed only
// in the migration because schema.prisma cannot state partial indexes; it is what makes a
// concurrent start on another pod race-safe rather than read-then-create. direction is
// carried here rather than read off the job so that index can exist, and the composite
// foreign key makes it provably the job's own direction.
model LiteLLM_ShadowEvalJobKey {
id String @id @default(cuid())
job_id String
direction String
api_key_id String // hashed virtual key whose traffic is shadowed
max_turns Int // this key's sample budget: judge at most this many turns
created_at DateTime @default(now())
stopped_at DateTime?
job LiteLLM_ShadowEvalJob @relation(fields: [job_id, direction], references: [id, direction], onDelete: Cascade)
@@index([job_id])
@@index([api_key_id])
}
// One row per sampled pipeline: a blind verdict (real | shadow | tie) or an error.
model LiteLLM_ShadowEvalAttempt {
id String @id @default(cuid())
job_id String
api_key_id String // which of the job's scoped keys produced this turn
request_id String // the judged real request
outcome String // real | shadow | tie | error
tier String? // router's tier for the prompt, when classified
@ -1490,7 +1511,7 @@ model LiteLLM_ShadowEvalAttempt {
error String?
created_at DateTime @default(now())
@@index([job_id])
@@index([job_id, api_key_id])
}
// ---------------------------------------------------------------------------

View file

@ -4,7 +4,9 @@ each against the job's other arm in a detached task (the auto-router for a forwa
fixed baseline model for a reverse one), blind-judges real vs shadow, and appends one
``LiteLLM_ShadowEvalAttempt`` row (verdict or error) as the feature's only hot-path write.
Counts, status, and spend derive from those rows at read time, so nothing can disagree
across pods or stop races; the hook reads active jobs through a short-TTL cache."""
across pods or stop races; the hook reads active keys through a short-TTL cache. A job can
scope several keys, and each carries its own turn budget and stop state, so the unit this
module samples and budgets against is the key, not the job."""
import asyncio
import hashlib
@ -449,14 +451,16 @@ class _JudgeVerdict:
class ActiveShadowEvalJob(BaseModel):
"""One active job as the sampling path needs it, validated straight off the untyped
job row: immutable config plus the attempt count as of the cache fill (the turn
budget's staleness is bounded by the cache TTL). Every way a row can be unsamplable
is a validation error here, so a bad row is skipped rather than sampled wrongly."""
"""One key's active sampling as the path needs it: immutable job config plus that key's
own budget and attempt count as of the cache fill (the budget's staleness is bounded by
the cache TTL). A job over N keys yields N of these, one per key, with independent
budgets, so exhausting one leaves its siblings sampling. Every way a row can be
unsamplable is a validation error here, so a bad row is skipped rather than sampled wrongly."""
model_config = ConfigDict(frozen=True, from_attributes=True)
model_config = ConfigDict(frozen=True)
id: str
api_key_id: str
router_name: str
direction: ShadowEvalDirection = "forward"
baseline_model: str | None = None
@ -485,16 +489,15 @@ class ActiveShadowEvalJob(BaseModel):
return self.baseline_model or self.router_name
def _as_active_job(record: object, attempts: int) -> ActiveShadowEvalJob | None:
"""The sampling path's view of one job row, or None for a row it cannot sample: an
unknown direction, or a reverse job with no baseline model to duplicate against.
Failing closed here is what keeps the dispatch path total."""
def _as_active_job(payload: Mapping[str, object]) -> ActiveShadowEvalJob | None:
"""The sampling path's view of one key row and the job it belongs to, or None for a row
it cannot sample: an unknown direction, or a reverse job with no baseline model to
duplicate against. Failing closed here is what keeps the dispatch path total."""
try:
job: Final = ActiveShadowEvalJob.model_validate(record)
return ActiveShadowEvalJob.model_validate(payload)
except ValidationError as e:
verbose_logger.debug("shadow_eval: skipping unsamplable job row: %s", e)
verbose_logger.debug("shadow_eval: skipping unsamplable job key row: %s", e)
return None
return job.model_copy(update={"attempts": attempts})
_jobs_cache: Final = InMemoryCache(max_size_in_memory=4, default_ttl=_JOBS_CACHE_TTL_SECONDS)
@ -516,14 +519,16 @@ class ShadowEvalLogger(CustomLogger):
self._prisma_provider = prisma_provider or _default_prisma_provider
self._jobs_cache = jobs_cache or _jobs_cache
self._inflight_shadow_tasks: int = 0
# Starts per job since the last cache fill, never decremented within a
# generation; the refill absorbs written rows and resets.
self._job_starts: dict[str, int] = {} # mutable-ok: per-generation counter
# Starts per shadowed key row since the last cache fill, never decremented within a
# generation; the refill absorbs written rows and resets. Keyed by (job, key) rather
# than by either alone, since the turn budget it guards is one key's within one job.
self._job_starts: dict[tuple[str, str], int] = {} # mutable-ok: per-generation counter
async def _active_jobs(self) -> Mapping[str, tuple[ActiveShadowEvalJob, ...]]:
"""Active jobs by api_key_id, cache-first. A key holds at most one job per
direction, so the value is a collection. A DB fault returns empty without
caching, so sampling pauses for that request and the next one retries."""
"""Sampling key rows by api_key_id, cache-first. A key holds at most one unstopped
row per direction (partial unique index), so the value is a collection of at most
two. A DB fault returns empty without caching, so sampling pauses for that request
and the next one retries."""
cached: Final = await self._jobs_cache.async_get_cache(_JOBS_CACHE_KEY)
if cached is not None:
return cached # pyright: ignore[reportReturnType] # cache stores exactly this mapping shape
@ -531,32 +536,48 @@ class ShadowEvalLogger(CustomLogger):
if prisma is None:
return _EMPTY_JOBS
try:
records: Final = await prisma.db.litellm_shadowevaljob.find_many(
records: Final = await prisma.db.litellm_shadowevaljobkey.find_many(
where={ # mutable-ok: Prisma filter
"stopped_at": None,
"ends_at": {"gt": datetime.now(timezone.utc)}, # mutable-ok: Prisma filter
"job": {"is": {"ends_at": {"gt": datetime.now(timezone.utc)}}}, # mutable-ok: Prisma filter
},
include={"job": True}, # mutable-ok: Prisma payload
)
pairs: Final = [{"job_id": r.job_id, "api_key_id": r.api_key_id} for r in records] # mutable-ok: Prisma OR
grouped: Final = (
await prisma.db.litellm_shadowevalattempt.group_by(
by=["job_id"],
by=["job_id", "api_key_id"],
count=True,
where={"job_id": {"in": [str(record.id) for record in records]}}, # mutable-ok: Prisma filter
where={"OR": pairs}, # mutable-ok: Prisma filter
)
if records
else ()
)
attempt_counts: Final = {str(row["job_id"]): int(row["_count"]["_all"]) for row in grouped or []}
by_key: Final = tuple(
sorted(
(
(str(record.api_key_id), job)
for record in records or []
if (job := _as_active_job(record, attempt_counts.get(str(record.id), 0))) is not None
),
key=itemgetter(0),
)
attempt_counts: Final = MappingProxyType(
{(str(row["job_id"]), str(row["api_key_id"])): int(row["_count"]["_all"]) for row in grouped or ()}
)
candidates: Final = tuple(
(
str(record.api_key_id),
_as_active_job(
{
"id": record.job_id,
"api_key_id": record.api_key_id,
"router_name": job.router_name,
"direction": job.direction,
"baseline_model": job.baseline_model,
"shadow_percentage": job.shadow_percentage,
"judge_model": job.judge_model,
"max_turns": record.max_turns,
"ends_at": job.ends_at,
"attempts": attempt_counts.get((str(record.job_id), str(record.api_key_id)), 0),
}
),
)
for record in records or ()
if (job := record.job) is not None
)
by_key: Final = tuple(sorted(((key, job) for key, job in candidates if job is not None), key=itemgetter(0)))
jobs: Final = MappingProxyType(
{key: tuple(job for _, job in group) for key, group in groupby(by_key, key=itemgetter(0))}
)
@ -608,7 +629,7 @@ class ShadowEvalLogger(CustomLogger):
job
for job in (await self._active_jobs()).get(str(api_key_hash), ())
if datetime.now(timezone.utc) < job.ends_at
and job.attempts + self._job_starts.get(job.id, 0) < job.max_turns
and job.attempts + self._job_starts.get((job.id, job.api_key_id), 0) < job.max_turns
and _sample_hits(request_id, job.id, job.shadow_percentage)
and _request_was_routed_by(request_metadata, job.router_name) == (job.direction == "reverse")
)
@ -627,7 +648,7 @@ class ShadowEvalLogger(CustomLogger):
for job in eligible:
if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS:
return
self._job_starts[job.id] = self._job_starts.get(job.id, 0) + 1
self._job_starts[(job.id, job.api_key_id)] = self._job_starts.get((job.id, job.api_key_id), 0) + 1
self._inflight_shadow_tasks += 1
asyncio.create_task(
self._run_shadow_eval(
@ -731,6 +752,7 @@ class ShadowEvalLogger(CustomLogger):
await prisma.db.litellm_shadowevalattempt.create(
data={ # mutable-ok: Prisma payload
"job_id": job.id,
"api_key_id": job.api_key_id,
"request_id": request_id,
"outcome": outcome,
"tier": control_tier if job.direction == "reverse" else (shadow.tier if shadow else None),

View file

@ -528,12 +528,18 @@ GROUP BY 1
_ATTEMPT_AGG_BY_TIER_SQL: Final = "SELECT COALESCE(tier, 'UNCLASSIFIED') AS grp," + _ATTEMPT_AGG_SELECT
_ATTEMPT_AGG_BY_MODEL_SQL: Final = "SELECT COALESCE(real_model, 'unknown') AS grp," + _ATTEMPT_AGG_SELECT
_SWEEP_FINISHED_JOBS_SQL: Final = """
UPDATE "LiteLLM_ShadowEvalJob" j SET stopped_at = NOW()
WHERE j.api_key_id = $1 AND j.stopped_at IS NULL
_ATTEMPT_AGG_BY_KEY_SQL: Final = "SELECT api_key_id AS grp," + _ATTEMPT_AGG_SELECT
_SWEEP_FINISHED_KEYS_SQL: Final = """
UPDATE "LiteLLM_ShadowEvalJobKey" k SET stopped_at = NOW()
FROM "LiteLLM_ShadowEvalJob" j
WHERE k.job_id = j.id AND k.api_key_id = ANY($1::text[]) AND k.stopped_at IS NULL
AND (
j.ends_at <= NOW()
OR (SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a WHERE a.job_id = j.id) >= j.max_turns
OR (
SELECT COUNT(*) FROM "LiteLLM_ShadowEvalAttempt" a
WHERE a.job_id = k.job_id AND a.api_key_id = k.api_key_id
) >= k.max_turns
)
"""
@ -575,11 +581,12 @@ def _slices(rows: Sequence[_AttemptAggRow]) -> tuple[ShadowEvalSlice, ...]:
async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> ShadowEvalResult | None:
"""Both stratifications of one job's verdicts. Tier answers "where does the router do
well"; the model stratification groups by whichever model served the real arm, so it
answers "which of the models this key uses today would the router beat" forward, and
"for the turns the router sent to X, did X beat the baseline" in reverse. Reads are
bounded by the job's own attempts (<= max_turns) via the job_id index."""
"""All three stratifications of one job's verdicts. Tier answers "where does the router
do well"; the model stratification groups by whichever model served the real arm, so it
answers "which of the models these keys use today would the router beat" forward, and
"for the turns the router sent to X, did X beat the baseline" in reverse; key answers
"which key's traffic does the router suit". Reads are bounded by the job's own attempts
(<= the sum of its keys' max_turns) via the job_id index."""
by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_TIER_SQL, job_id) or ()
)
@ -588,10 +595,14 @@ async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> Sh
by_model: Final = _ATTEMPT_AGG_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_MODEL_SQL, job_id) or ()
)
by_key: Final = _ATTEMPT_AGG_ROWS.validate_python(
await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_KEY_SQL, job_id) or ()
)
total_turns: Final = sum(r.turn_count for r in by_tier)
return ShadowEvalResult(
by_tier=_slices(by_tier),
by_current_model=_slices(by_model),
by_key=_slices(by_key),
overall_shadow_win_rate_pct=_pct_of(sum(r.shadow_wins for r in by_tier), total_turns),
overall_tie_rate_pct=_pct_of(sum(r.ties for r in by_tier), total_turns),
)
@ -609,20 +620,21 @@ async def start_shadow_eval(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> ShadowEvalJobResponse:
"""
Start a shadow eval: duplicate a sampled slice of a key's live traffic against a second
arm, judge the two responses blind, and stratify win rates by tier and by the model that
served the real arm.
Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against
a second arm, judge the two responses blind, and stratify win rates by tier, by the model
that served the real arm, and by key.
A forward job answers whether the key should adopt router_name: it samples the requests
A forward job answers whether the keys should adopt router_name: it samples the requests
the router did not serve and duplicates them through it. A reverse job answers whether a
key already on the router still gains from it: it samples the requests the router did
serve and duplicates them against baseline_model. A key can hold one active job per
direction, so both questions can run at once.
Shadow responses are never served to users. The job samples until it has judged
max_turns turns, reaches the end of its window, or is stopped; sampling changes
propagate to pods within about 10 seconds. Shadow and judge calls bill to the
shadowed key but are excluded from request counts and auto-router adoption metrics.
Shadow responses are never served to users. Each key samples until it has judged
max_turns turns of its own traffic, the job's window ends, or the job is stopped, so a
busy key running out of budget does not end sampling for the others; sampling changes
propagate to pods within about 10 seconds. Shadow and judge calls bill to the shadowed
key but are excluded from request counts and auto-router adoption metrics.
"""
from litellm.proxy.proxy_server import llm_router, prisma_client
@ -634,48 +646,58 @@ async def start_shadow_eval(
_validate_plain_model(llm_router, data.judge_model, "judge_model")
if data.baseline_model is not None:
_validate_plain_model(llm_router, data.baseline_model, "baseline_model")
key_row: Final = await prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": data.api_key_id} # mutable-ok: Prisma filter
token_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": list(data.api_key_ids)}} # mutable-ok: Prisma filter
)
if key_row is None:
unknown: Final = tuple(sorted(frozenset(data.api_key_ids) - frozenset(row.token for row in token_rows)))
if unknown:
raise HTTPException(
status_code=400,
detail=(
f"api_key_id '{data.api_key_id}' is not a key on this proxy; pass the key's token hash, "
f"api_key_ids not on this proxy: {', '.join(unknown)}; pass each key's token hash, "
"the value the key list and key info endpoints report"
),
)
# A job that expired or exhausted its turn budget stopped sampling on its own, but
# still holds its slot in the per-key, per-direction partial unique index until
# stamped; free it so a new eval can start. Sweeping both directions is deliberate.
await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, data.api_key_id)
active: Final = await prisma_client.db.litellm_shadowevaljob.find_first(
# A key whose job expired or whose own turn budget ran out stopped sampling on its own, but
# still holds its slot in the per-key, per-direction partial unique index until stamped;
# free it so a new eval can start. Sweeping both directions is deliberate.
await prisma_client.db.execute_raw(_SWEEP_FINISHED_KEYS_SQL, list(data.api_key_ids)) # mutable-ok: query param
claimed: Final = await prisma_client.db.litellm_shadowevaljobkey.find_many(
where={ # mutable-ok: Prisma filter
"api_key_id": data.api_key_id,
"api_key_id": {"in": list(data.api_key_ids)},
"direction": data.direction,
"stopped_at": None,
},
)
if active is not None:
if claimed:
raise HTTPException(
status_code=409,
detail=f"Key already has an active {data.direction} shadow eval job ({active.id}). Stop it first.",
detail=(
f"Already in an active {data.direction} shadow eval job: "
+ ", ".join(sorted(f"{row.api_key_id} (job {row.job_id})" for row in claimed))
+ ". Stop it first."
),
)
now: Final = datetime.now(timezone.utc)
key_rows: Final = tuple(
{"api_key_id": api_key_id, "max_turns": data.max_turns} # mutable-ok: Prisma create payload
for api_key_id in data.api_key_ids
)
try:
job: Final = await prisma_client.db.litellm_shadowevaljob.create(
data={ # mutable-ok: Prisma payload
"api_key_id": data.api_key_id,
"router_name": data.router_name,
"direction": data.direction,
"baseline_model": data.baseline_model,
"judge_model": data.judge_model,
"shadow_percentage": data.shadow_percentage,
"max_turns": data.max_turns,
"created_by": user_api_key_dict.user_id,
"ends_at": now + timedelta(days=data.duration_days),
}
# direction is injected into each key row by Prisma from the composite FK
"keys": {"create": list(key_rows)}, # mutable-ok: Prisma nested create
},
include={"keys": True}, # mutable-ok: Prisma payload
)
except Exception as e:
if not _is_unique_violation(e):
@ -683,7 +705,7 @@ async def start_shadow_eval(
raise HTTPException(
status_code=409,
detail=(
f"Key already has an active {data.direction} shadow eval job (started concurrently). Stop it first."
f"A requested key was claimed by another {data.direction} shadow eval job concurrently. Stop it first."
),
) from e
return ShadowEvalJobResponse.model_validate(job, from_attributes=True)
@ -697,7 +719,9 @@ async def start_shadow_eval(
)
async def list_shadow_eval_jobs(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
api_key_id: Annotated[str | None, Query(description="Filter to jobs shadowing this key")] = None,
api_key_id: Annotated[
str | None, Query(description="Filter to jobs that shadow this key, alone or alongside others")
] = None,
limit: Annotated[int, Query(ge=1, le=200, description="Newest jobs to return")] = 50,
) -> tuple[ShadowEvalJobResponse, ...]:
"""List shadow eval jobs, newest first. Counts and results ride the detail endpoint only."""
@ -707,7 +731,8 @@ async def list_shadow_eval_jobs(
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
records: Final = await prisma_client.db.litellm_shadowevaljob.find_many(
where={"api_key_id": api_key_id} if api_key_id else {}, # mutable-ok: Prisma filter
where={"keys": {"some": {"api_key_id": api_key_id}}} if api_key_id else {}, # mutable-ok: Prisma filter
include={"keys": True}, # mutable-ok: Prisma payload
order={"created_at": "desc"}, # mutable-ok: Prisma order
take=limit,
)
@ -731,7 +756,8 @@ async def get_shadow_eval_job(
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
where={"id": job_id} # mutable-ok: Prisma filter
where={"id": job_id}, # mutable-ok: Prisma filter
include={"keys": True}, # mutable-ok: Prisma payload
)
if record is None:
raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
@ -763,22 +789,31 @@ async def stop_shadow_eval_job(
job_id: str,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> ShadowEvalJobResponse:
"""Stop an active shadow eval job. Attempts are kept; sampling halts within ~10s."""
"""Stop an active shadow eval job, every key it scopes at once. Attempts are kept;
sampling halts within ~10s. Keys that already stopped on their own budget keep the
stopped_at they earned, so one write records the outcome for the whole job."""
from litellm.proxy.proxy_server import prisma_client
_require_admin_writer(user_api_key_dict, "stop a shadow eval")
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
where={"id": job_id} # mutable-ok: Prisma filter
where={"id": job_id}, # mutable-ok: Prisma filter
include={"keys": True}, # mutable-ok: Prisma payload
)
if record is None:
raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
current: Final = ShadowEvalJobResponse.model_validate(record, from_attributes=True)
if current.status != "running":
raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}")
updated: Final = await prisma_client.db.litellm_shadowevaljob.update(
where={"id": job_id}, # mutable-ok: Prisma filter
await prisma_client.db.litellm_shadowevaljobkey.update_many(
where={"job_id": job_id, "stopped_at": None}, # mutable-ok: Prisma filter
data={"stopped_at": datetime.now(timezone.utc)}, # mutable-ok: Prisma payload
)
updated: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
where={"id": job_id}, # mutable-ok: Prisma filter
include={"keys": True}, # mutable-ok: Prisma payload
)
if updated is None:
raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
return ShadowEvalJobResponse.model_validate(updated, from_attributes=True)

View file

@ -1450,36 +1450,57 @@ model LiteLLM_AutoRouterSession {
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
}
// Shadow eval: evaluation of an auto-router against a key's live traffic, in either
// direction. forward duplicates the requests the key did not route through the router
// through it, answering whether the key should adopt it; reverse duplicates the requests
// Shadow eval: evaluation of an auto-router against a set of keys' live traffic, in either
// direction. forward duplicates the requests the keys did not route through the router
// through it, answering whether they should adopt it; reverse duplicates the requests
// the router did serve against a fixed baseline model, answering whether a key already on
// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge
// compares real vs shadow responses blind. The job row is immutable config plus
// stopped_at; every count, status, and spend figure is derived from the append-only
// attempt rows, so nothing can disagree across pods or stop races.
// compares real vs shadow responses blind. The job row is immutable config; the keys it
// scopes, their budgets and their stop state live on LiteLLM_ShadowEvalJobKey, and every
// count, status, and spend figure is derived from the append-only attempt rows, so nothing
// can disagree across pods or stop races.
model LiteLLM_ShadowEvalJob {
id String @id @default(cuid())
api_key_id String // hashed virtual key whose traffic is shadowed
router_name String // the auto-router under evaluation, in either direction
direction String @default("forward") // forward | reverse
baseline_model String? // reverse only: the fixed model the router is judged against
judge_model String
shadow_percentage Float
max_turns Int // sample budget: judge at most this many turns
created_at DateTime @default(now())
created_by String?
ends_at DateTime
stopped_at DateTime?
keys LiteLLM_ShadowEvalJobKey[]
@@index([api_key_id])
@@unique([id, direction])
@@index([created_at])
}
// One row per key a job shadows. Budget and stop state are per key, so one key exhausting
// its turns never ends a sibling's sampling. "One active job per (key, direction)" is a
// partial unique index on (api_key_id, direction) WHERE stopped_at IS NULL, expressed only
// in the migration because schema.prisma cannot state partial indexes; it is what makes a
// concurrent start on another pod race-safe rather than read-then-create. direction is
// carried here rather than read off the job so that index can exist, and the composite
// foreign key makes it provably the job's own direction.
model LiteLLM_ShadowEvalJobKey {
id String @id @default(cuid())
job_id String
direction String
api_key_id String // hashed virtual key whose traffic is shadowed
max_turns Int // this key's sample budget: judge at most this many turns
created_at DateTime @default(now())
stopped_at DateTime?
job LiteLLM_ShadowEvalJob @relation(fields: [job_id, direction], references: [id, direction], onDelete: Cascade)
@@index([job_id])
@@index([api_key_id])
}
// One row per sampled pipeline: a blind verdict (real | shadow | tie) or an error.
model LiteLLM_ShadowEvalAttempt {
id String @id @default(cuid())
job_id String
api_key_id String // which of the job's scoped keys produced this turn
request_id String // the judged real request
outcome String // real | shadow | tie | error
tier String? // router's tier for the prompt, when classified
@ -1490,7 +1511,7 @@ model LiteLLM_ShadowEvalAttempt {
error String?
created_at DateTime @default(now())
@@index([job_id])
@@index([job_id, api_key_id])
}
// ---------------------------------------------------------------------------

View file

@ -152,13 +152,15 @@ DEFAULT_SHADOW_EVAL_JUDGE_MODEL: Final[str] = "anthropic/claude-sonnet-5"
class StartShadowEvalRequest(BaseModel):
"""Start duplicating a key's traffic for blind comparison against an auto-router."""
"""Start duplicating one or more keys' traffic for blind comparison against an auto-router."""
api_key_id: str = Field(
api_key_ids: tuple[str, ...] = Field(
min_length=1,
description=(
"The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this "
"key's traffic; requests made with any other key are not sampled."
)
"The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these "
"keys' traffic; requests made with any other key are not sampled. Each key gets its own "
"max_turns budget, so one key exhausting its budget does not end sampling for the others."
),
)
router_name: str = Field(description="The auto-router under evaluation, in either direction")
direction: ShadowEvalDirection = Field(
@ -201,8 +203,9 @@ class StartShadowEvalRequest(BaseModel):
ge=1,
le=2000,
description=(
"Sample budget: the job judges at most this many turns, then completes. This is also the spend "
"bound; expected judge cost is roughly max_turns times one judge call"
"Per-key sample budget: the job judges at most this many turns of EACH scoped key's traffic, so a "
"job over N keys judges at most N times max_turns turns. This is also the spend bound; expected "
"judge cost is roughly that turn ceiling times one judge call"
),
)
@ -211,6 +214,13 @@ class StartShadowEvalRequest(BaseModel):
def _round_percentage(cls, value: float) -> float:
return round(value, 2)
@field_validator("api_key_ids")
@classmethod
def _dedupe_keys(cls, value: tuple[str, ...]) -> tuple[str, ...]:
"""Naming a key twice means the same thing as naming it once, and the second key row
would collide with the first on the one-active-per-(key, direction) index."""
return tuple(dict.fromkeys(value))
@model_validator(mode="after")
def _baseline_model_matches_direction(self) -> "StartShadowEvalRequest":
if self.direction == "reverse" and self.baseline_model is None:
@ -221,8 +231,8 @@ class StartShadowEvalRequest(BaseModel):
class ShadowEvalSlice(BaseModel):
"""Judge outcomes for one slice of a job's verdicts (a router tier, or one of the
models that served the real arm)."""
"""Judge outcomes for one slice of a job's verdicts (a router tier, one of the models
that served the real arm, or one of the keys the job is scoped to)."""
group: str
turn_count: int
@ -248,33 +258,56 @@ class ShadowEvalResult(BaseModel):
by_tier: tuple[ShadowEvalSlice, ...]
by_current_model: tuple[ShadowEvalSlice, ...] = Field(
description=(
"Sliced by the model that served the real arm: the key's incumbent models in forward mode, "
"Sliced by the model that served the real arm: the keys' incumbent models in forward mode, "
"and in reverse the models the router itself picked"
)
)
by_key: tuple[ShadowEvalSlice, ...] = Field(
default=(),
description=(
"One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job "
"scopes but has not yet judged a turn for are absent rather than reported as zero"
),
)
overall_shadow_win_rate_pct: float
overall_tie_rate_pct: float
class ShadowEvalJobKeyResponse(BaseModel):
"""One key a job shadows, with its own budget and stop state."""
api_key_id: str = Field(description="The hashed virtual key whose traffic this row scopes")
max_turns: int = Field(description="This key's own sample budget, independent of its siblings'")
stopped_at: datetime | None = Field(
default=None,
description=(
"When this key stopped sampling, whether its own budget ran out, the window closed, or an "
"operator stopped the job. The key reads completed whenever the job does, otherwise stopped "
"once this is set and running until then"
),
)
class ShadowEvalJobResponse(BaseModel):
"""A shadow-eval job. Validates directly from the prisma record (job_id reads the
row's id); status is derived from stopped_at and ends_at, never stored, so no writer
anywhere can produce an inconsistent one. Aggregate fields are populated by the
row's id); status is derived from the keys' stopped_at and ends_at, never stored, so no
writer anywhere can produce an inconsistent one. Aggregate fields are populated by the
detail endpoint only and stay None on list responses."""
model_config = ConfigDict(from_attributes=True, populate_by_name=True)
job_id: str = Field(validation_alias=AliasChoices("id", "job_id"))
api_key_id: str = Field(description="The hashed virtual key whose traffic this job evaluates, and only that key's")
keys: tuple[ShadowEvalJobKeyResponse, ...] = Field(
min_length=1,
description="The keys whose traffic this job evaluates, and only those keys', each with its own budget",
)
router_name: str
direction: ShadowEvalDirection = "forward"
baseline_model: str | None = None
judge_model: str
shadow_percentage: float
max_turns: int
created_at: datetime
ends_at: datetime
stopped_at: datetime | None = None
judged_count: int | None = Field(default=None, description="Verdicts recorded; detail endpoint only")
error_count: int | None = Field(default=None, description="Sampled attempts that errored; detail endpoint only")
@ -285,12 +318,13 @@ class ShadowEvalJobResponse(BaseModel):
@computed_field
@property
def status(self) -> ShadowEvalStatus:
"""A job whose window has passed reads completed even if a later sweep stamped
stopped_at; stopped means sampling ended before the window did."""
"""A job whose window has passed reads completed even if a sweep stamped its keys
first; stopped means every key ended sampling before the window did. One key
exhausting its own budget leaves the job running while any sibling still samples."""
if datetime.now(timezone.utc) >= (
self.ends_at if self.ends_at.tzinfo else self.ends_at.replace(tzinfo=timezone.utc)
):
return "completed"
if self.stopped_at is not None:
if all(key.stopped_at is not None for key in self.keys):
return "stopped"
return "running"

View file

@ -1450,36 +1450,57 @@ model LiteLLM_AutoRouterSession {
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
}
// Shadow eval: evaluation of an auto-router against a key's live traffic, in either
// direction. forward duplicates the requests the key did not route through the router
// through it, answering whether the key should adopt it; reverse duplicates the requests
// Shadow eval: evaluation of an auto-router against a set of keys' live traffic, in either
// direction. forward duplicates the requests the keys did not route through the router
// through it, answering whether they should adopt it; reverse duplicates the requests
// the router did serve against a fixed baseline model, answering whether a key already on
// it still benefits. Either way a sampled slice runs in a detached task and an LLM judge
// compares real vs shadow responses blind. The job row is immutable config plus
// stopped_at; every count, status, and spend figure is derived from the append-only
// attempt rows, so nothing can disagree across pods or stop races.
// compares real vs shadow responses blind. The job row is immutable config; the keys it
// scopes, their budgets and their stop state live on LiteLLM_ShadowEvalJobKey, and every
// count, status, and spend figure is derived from the append-only attempt rows, so nothing
// can disagree across pods or stop races.
model LiteLLM_ShadowEvalJob {
id String @id @default(cuid())
api_key_id String // hashed virtual key whose traffic is shadowed
router_name String // the auto-router under evaluation, in either direction
direction String @default("forward") // forward | reverse
baseline_model String? // reverse only: the fixed model the router is judged against
judge_model String
shadow_percentage Float
max_turns Int // sample budget: judge at most this many turns
created_at DateTime @default(now())
created_by String?
ends_at DateTime
stopped_at DateTime?
keys LiteLLM_ShadowEvalJobKey[]
@@index([api_key_id])
@@unique([id, direction])
@@index([created_at])
}
// One row per key a job shadows. Budget and stop state are per key, so one key exhausting
// its turns never ends a sibling's sampling. "One active job per (key, direction)" is a
// partial unique index on (api_key_id, direction) WHERE stopped_at IS NULL, expressed only
// in the migration because schema.prisma cannot state partial indexes; it is what makes a
// concurrent start on another pod race-safe rather than read-then-create. direction is
// carried here rather than read off the job so that index can exist, and the composite
// foreign key makes it provably the job's own direction.
model LiteLLM_ShadowEvalJobKey {
id String @id @default(cuid())
job_id String
direction String
api_key_id String // hashed virtual key whose traffic is shadowed
max_turns Int // this key's sample budget: judge at most this many turns
created_at DateTime @default(now())
stopped_at DateTime?
job LiteLLM_ShadowEvalJob @relation(fields: [job_id, direction], references: [id, direction], onDelete: Cascade)
@@index([job_id])
@@index([api_key_id])
}
// One row per sampled pipeline: a blind verdict (real | shadow | tie) or an error.
model LiteLLM_ShadowEvalAttempt {
id String @id @default(cuid())
job_id String
api_key_id String // which of the job's scoped keys produced this turn
request_id String // the judged real request
outcome String // real | shadow | tie | error
tier String? // router's tier for the prompt, when classified
@ -1490,7 +1511,7 @@ model LiteLLM_ShadowEvalAttempt {
error String?
created_at DateTime @default(now())
@@index([job_id])
@@index([job_id, api_key_id])
}
// ---------------------------------------------------------------------------

View file

@ -230,6 +230,7 @@ class TestMigrationSQLIdempotency:
"20250918083359_drop_spec_version_column_from_mcp_table",
"20260213170952_access_group_change_to_model_name",
"20260224203854_add_agent_object_permissions_table",
"20260814000000_shadow_eval_multi_key",
}
def test_no_drop_column_statements(self, all_migrations):

View file

@ -27,6 +27,7 @@ from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTE
def _job(**overrides) -> ActiveShadowEvalJob:
defaults = dict(
id="job-1",
api_key_id="key-hash",
router_name="my-router",
shadow_percentage=100.0,
judge_model="judge-model",
@ -37,30 +38,34 @@ def _job(**overrides) -> ActiveShadowEvalJob:
return ActiveShadowEvalJob(**{**defaults, **overrides})
def _prisma(jobs=(), attempt_counts=()) -> MagicMock:
def _prisma(keys=(), attempt_counts=()) -> MagicMock:
prisma = MagicMock()
prisma.db.litellm_shadowevaljob.find_many = AsyncMock(return_value=list(jobs))
prisma.db.litellm_shadowevaljobkey.find_many = AsyncMock(return_value=list(keys))
prisma.db.litellm_shadowevalattempt.group_by = AsyncMock(
return_value=[{"job_id": job_id, "_count": {"_all": count}} for job_id, count in attempt_counts]
return_value=[
{"job_id": job_id, "api_key_id": api_key_id, "_count": {"_all": count}}
for job_id, api_key_id, count in attempt_counts
]
)
prisma.db.litellm_shadowevalattempt.create = AsyncMock()
return prisma
def _job_record(job: ActiveShadowEvalJob, api_key_id="key-hash") -> MagicMock:
def _key_record(job: ActiveShadowEvalJob) -> MagicMock:
"""One LiteLLM_ShadowEvalJobKey row with its job relation loaded, as `include={"job": True}`
hands it back: per-key budget on the row, shared config on the relation."""
record = MagicMock()
for field, value in dict(job_id=job.id, api_key_id=job.api_key_id, max_turns=job.max_turns).items():
setattr(record, field, value)
for field, value in dict(
id=job.id,
api_key_id=api_key_id,
router_name=job.router_name,
direction=job.direction,
baseline_model=job.baseline_model,
shadow_percentage=job.shadow_percentage,
judge_model=job.judge_model,
max_turns=job.max_turns,
ends_at=job.ends_at,
).items():
setattr(record, field, value)
setattr(record.job, field, value)
return record
@ -97,7 +102,13 @@ def _logger(router=None, prisma=None, jobs=()) -> ShadowEvalLogger:
jobs_cache=cache,
)
if jobs:
cache.set_cache("shadow_eval:active_jobs", {"key-hash": tuple(jobs)})
cache.set_cache(
"shadow_eval:active_jobs",
{
key: tuple(job for job in jobs if job.api_key_id == key)
for key in dict.fromkeys(job.api_key_id for job in jobs)
},
)
return logger
@ -466,6 +477,7 @@ class TestSuccessHookSkipChain:
create.assert_awaited_once()
row = create.call_args.kwargs["data"]
assert row["job_id"] == "job-1"
assert row["api_key_id"] == "key-hash"
assert row["request_id"] == "req-1"
assert row["outcome"] in ("real", "shadow")
assert row["tier"] == "SIMPLE"
@ -474,7 +486,7 @@ class TestSuccessHookSkipChain:
assert row["confidence"] == 0.9
assert row["judge_cost"] == 0.005
assert row["error"] is None
assert prisma.db.litellm_shadowevaljob.find_many.await_count == 0
assert prisma.db.litellm_shadowevaljobkey.find_many.await_count == 0
@pytest.mark.parametrize(
"kwargs_mutation,job_mutation",
@ -503,13 +515,13 @@ class TestSuccessHookSkipChain:
starts = job_mutation.pop("_starts", 0)
prisma = _prisma()
logger = _logger(router=_router(), prisma=prisma, jobs=(_job(**job_mutation),))
logger._job_starts = {"job-1": starts}
logger._job_starts = {("job-1", "key-hash"): starts}
await logger.async_log_success_event(_success_kwargs(**kwargs_mutation), RESPONSE, None, None)
await _drain(logger)
prisma.db.litellm_shadowevalattempt.create.assert_not_called()
assert logger._job_starts.get("job-1", 0) == starts
assert logger._job_starts.get(("job-1", "key-hash"), 0) == starts
async def test_completed_pipelines_hold_turn_budget_within_a_cache_generation(self):
"""A finished pipeline frees its concurrency slot but not its slice of the turn
@ -524,6 +536,50 @@ class TestSuccessHookSkipChain:
assert prisma.db.litellm_shadowevalattempt.create.await_count == 1
async def test_one_keys_exhausted_budget_leaves_its_siblings_sampling(self):
"""The turn budget belongs to the key, not to the job, so a job scoping two keys
keeps shadowing the second after the first has spent its own max_turns."""
prisma = _prisma()
logger = _logger(
router=_router(),
prisma=prisma,
jobs=(
_job(api_key_id="spent-key", attempts=200, max_turns=200),
_job(api_key_id="fresh-key", attempts=0, max_turns=200),
),
)
await logger.async_log_success_event(_success_kwargs(api_key_hash="spent-key"), RESPONSE, None, None)
await _drain(logger)
await logger.async_log_success_event(_success_kwargs(api_key_hash="fresh-key"), RESPONSE, None, None)
await _drain(logger)
create = prisma.db.litellm_shadowevalattempt.create
create.assert_awaited_once()
assert create.call_args.kwargs["data"]["api_key_id"] == "fresh-key"
async def test_started_turns_are_held_against_the_starting_key_only(self):
"""The in-generation starts counter is keyed by key too. Spending the sibling's
last turn must not close the budget of a key that has one left."""
prisma = _prisma()
logger = _logger(
router=_router(),
prisma=prisma,
jobs=(
_job(api_key_id="key-a", attempts=199, max_turns=200),
_job(api_key_id="key-b", attempts=199, max_turns=200),
),
)
for api_key_hash in ("key-a", "key-a", "key-b"):
await logger.async_log_success_event(
_success_kwargs(request_id=f"req-{api_key_hash}", api_key_hash=api_key_hash), RESPONSE, None, None
)
await _drain(logger)
attributed = [call.kwargs["data"]["api_key_id"] for call in prisma.db.litellm_shadowevalattempt.create.call_args_list]
assert attributed == ["key-a", "key-b"]
async def test_v1_messages_surface_forwards_identity_from_litellm_metadata(self):
"""/v1/messages stores identity in litellm_params.litellm_metadata, so the hook
resolves the bucket through the shared helper; every surface forwards the same
@ -574,7 +630,7 @@ class TestSuccessHookSkipChain:
class TestActiveJobsCache:
async def test_cache_miss_reads_db_once_then_serves_from_cache(self):
job = _job()
prisma = _prisma(jobs=[_job_record(job)], attempt_counts=[("job-1", 7)])
prisma = _prisma(keys=[_key_record(job)], attempt_counts=[("job-1", "key-hash", 7)])
logger = ShadowEvalLogger(
router_provider=lambda: None,
prisma_provider=lambda: prisma,
@ -586,15 +642,58 @@ class TestActiveJobsCache:
assert [job.id for job in first["key-hash"]] == ["job-1"]
assert second["key-hash"][0].attempts == 7
assert prisma.db.litellm_shadowevaljob.find_many.await_count == 1
where = prisma.db.litellm_shadowevaljob.find_many.call_args.kwargs["where"]
assert prisma.db.litellm_shadowevaljobkey.find_many.await_count == 1
where = prisma.db.litellm_shadowevaljobkey.find_many.call_args.kwargs["where"]
assert where["stopped_at"] is None
assert "gt" in where["ends_at"]
assert "gt" in where["job"]["is"]["ends_at"]
count_where = prisma.db.litellm_shadowevalattempt.group_by.call_args.kwargs["where"]
assert count_where == {"job_id": {"in": ["job-1"]}}
assert count_where == {"OR": [{"job_id": "job-1", "api_key_id": "key-hash"}]}
async def test_the_attempt_count_filter_matches_exact_pairs_not_a_cross_product(self):
"""Two active keys under two jobs are two pairs. Filtering job_id IN (...) AND
api_key_id IN (...) also selects the pairs that never ran together, so every
other job's rows for the same key get scanned and grouped for nothing."""
alpha = _key_record(_job(id="job-a", api_key_id="key-a"))
beta = _key_record(_job(id="job-b", api_key_id="key-b"))
prisma = _prisma(keys=[alpha, beta])
logger = ShadowEvalLogger(
router_provider=lambda: None,
prisma_provider=lambda: prisma,
jobs_cache=InMemoryCache(max_size_in_memory=4, default_ttl=60),
)
await logger._active_jobs()
count_where = prisma.db.litellm_shadowevalattempt.group_by.call_args.kwargs["where"]
assert count_where == {
"OR": [{"job_id": "job-a", "api_key_id": "key-a"}, {"job_id": "job-b", "api_key_id": "key-b"}]
}
async def test_one_job_fills_one_entry_per_key_with_that_keys_own_count(self):
"""Two keys of one job carry independent budgets, so the fill must group attempts
by key as well as by job. Grouping by job alone hands both keys the same count."""
spent = _job(api_key_id="spent-key", max_turns=200)
fresh = _job(api_key_id="fresh-key", max_turns=50)
prisma = _prisma(
keys=[_key_record(spent), _key_record(fresh)],
attempt_counts=[("job-1", "spent-key", 200), ("job-1", "fresh-key", 3)],
)
logger = ShadowEvalLogger(
router_provider=lambda: None,
prisma_provider=lambda: prisma,
jobs_cache=InMemoryCache(max_size_in_memory=4, default_ttl=60),
)
jobs = await logger._active_jobs()
assert set(jobs) == {"spent-key", "fresh-key"}
assert jobs["spent-key"][0].attempts == 200 and jobs["spent-key"][0].max_turns == 200
assert jobs["fresh-key"][0].attempts == 3 and jobs["fresh-key"][0].max_turns == 50
assert jobs["fresh-key"][0].id == "job-1"
assert prisma.db.litellm_shadowevalattempt.group_by.call_args.kwargs["by"] == ["job_id", "api_key_id"]
async def test_no_active_jobs_is_cached_too(self):
prisma = _prisma(jobs=[])
prisma = _prisma(keys=[])
logger = ShadowEvalLogger(
router_provider=lambda: None,
prisma_provider=lambda: prisma,
@ -603,12 +702,12 @@ class TestActiveJobsCache:
assert await logger._active_jobs() == {}
assert await logger._active_jobs() == {}
assert prisma.db.litellm_shadowevaljob.find_many.await_count == 1
assert prisma.db.litellm_shadowevaljobkey.find_many.await_count == 1
prisma.db.litellm_shadowevalattempt.group_by.assert_not_called()
async def test_db_fault_returns_empty_without_caching_the_fault(self):
prisma = _prisma()
prisma.db.litellm_shadowevaljob.find_many = AsyncMock(side_effect=RuntimeError("db blip"))
prisma.db.litellm_shadowevaljobkey.find_many = AsyncMock(side_effect=RuntimeError("db blip"))
logger = ShadowEvalLogger(
router_provider=lambda: None,
prisma_provider=lambda: prisma,
@ -617,17 +716,17 @@ class TestActiveJobsCache:
assert await logger._active_jobs() == {}
assert await logger._active_jobs() == {}
assert prisma.db.litellm_shadowevaljob.find_many.await_count == 2
assert prisma.db.litellm_shadowevaljobkey.find_many.await_count == 2
async def test_cache_refill_resets_the_starts_counter(self):
job = _job()
prisma = _prisma(jobs=[_job_record(job)], attempt_counts=[("job-1", 7)])
prisma = _prisma(keys=[_key_record(job)], attempt_counts=[("job-1", "key-hash", 7)])
logger = ShadowEvalLogger(
router_provider=lambda: None,
prisma_provider=lambda: prisma,
jobs_cache=InMemoryCache(max_size_in_memory=4, default_ttl=60),
)
logger._job_starts = {"job-1": 5}
logger._job_starts = {("job-1", "key-hash"): 5}
await logger._active_jobs()
@ -872,7 +971,7 @@ class TestDirection:
rows = [call.kwargs["data"] for call in prisma.db.litellm_shadowevalattempt.create.call_args_list]
assert sorted(row["job_id"] for row in rows) == ["forward-job", "reverse-job"]
assert logger._job_starts == {"forward-job": 1, "reverse-job": 1}
assert logger._job_starts == {("forward-job", "key-hash"): 1, ("reverse-job", "key-hash"): 1}
@pytest.mark.asyncio
@ -880,10 +979,12 @@ class TestActiveJobsFailClosed:
async def test_a_row_the_sampler_cannot_read_is_dropped_not_guessed(self):
"""A reverse row with no baseline model has no second arm to call, so it is skipped
rather than silently dispatched at the router it is supposed to be judging."""
broken = _job_record(_job(id="job-broken"))
broken.direction = "reverse"
broken.baseline_model = None
prisma = _prisma(jobs=[broken, _job_record(_job(id="job-ok"))], attempt_counts=[("job-ok", 1)])
broken = _key_record(_job(id="job-broken"))
broken.job.direction = "reverse"
broken.job.baseline_model = None
prisma = _prisma(
keys=[broken, _key_record(_job(id="job-ok"))], attempt_counts=[("job-ok", "key-hash", 1)]
)
logger = ShadowEvalLogger(
router_provider=lambda: None,
prisma_provider=lambda: prisma,
@ -894,11 +995,11 @@ class TestActiveJobsFailClosed:
async def test_both_of_a_key_s_jobs_survive_the_lookup(self):
records = [
_job_record(_job(id="job-forward")),
_job_record(_reverse_job(id="job-reverse")),
_job_record(_job(id="job-other"), api_key_id="other-key"),
_key_record(_job(id="job-forward")),
_key_record(_reverse_job(id="job-reverse")),
_key_record(_job(id="job-other", api_key_id="other-key")),
]
prisma = _prisma(jobs=records, attempt_counts=[("job-reverse", 3)])
prisma = _prisma(keys=records, attempt_counts=[("job-reverse", "key-hash", 3)])
logger = ShadowEvalLogger(
router_provider=lambda: None,
prisma_provider=lambda: prisma,

View file

@ -318,9 +318,7 @@ class TestAutoRouterBenchmarks:
from litellm.proxy.management_endpoints.auto_router_endpoints import _benchmark_totals
totals = _benchmark_totals(self.ROW)
bucket_hits = (
totals.cache.same_model.hits + totals.cache.first_visit.hits + totals.cache.return_to_tier.hits
)
bucket_hits = totals.cache.same_model.hits + totals.cache.first_visit.hits + totals.cache.return_to_tier.hits
assert bucket_hits == 27
assert totals.cache.hit_rate_pct == pytest.approx(100.0 * 28 / 38, abs=0.1)
@ -500,43 +498,71 @@ def _shadow_router() -> MagicMock:
return router
def _job_record(**overrides: object) -> MagicMock:
def _spec_record(fields: dict) -> MagicMock:
"""Spec'd like a real prisma row: only the table's columns exist as attributes, so
from_attributes validation falls back to model defaults for everything else."""
defaults = {
"id": "job-1",
"api_key_id": "key-hash",
"router_name": "my-router",
"judge_model": "anthropic/claude-sonnet-5",
"shadow_percentage": 10.0,
"max_turns": 200,
"created_at": datetime(2026, 8, 11, tzinfo=timezone.utc),
"ends_at": datetime.now(timezone.utc) + timedelta(days=7),
"stopped_at": None,
}
fields = {**defaults, **overrides}
record = MagicMock(spec=list(fields))
for key, value in fields.items():
setattr(record, key, value)
return record
def _shadow_prisma(active_job=None, agg_rows=None) -> MagicMock:
def _key_row(api_key_id: str = "key-hash", max_turns: int = 200, stopped_at=None) -> MagicMock:
return _spec_record({"api_key_id": api_key_id, "max_turns": max_turns, "stopped_at": stopped_at})
def _claim(api_key_id: str, job_id: str = "job-9", direction: str = "forward", stopped_at=None) -> MagicMock:
"""A key row a prior job holds, still claiming the one-active-per-key-and-direction
index unless stopped."""
return _spec_record(
{"api_key_id": api_key_id, "job_id": job_id, "direction": direction, "stopped_at": stopped_at}
)
def _job_record(**overrides: object) -> MagicMock:
defaults = {
"id": "job-1",
"router_name": "my-router",
"direction": "forward",
"baseline_model": None,
"judge_model": "anthropic/claude-sonnet-5",
"shadow_percentage": 10.0,
"created_at": datetime(2026, 8, 11, tzinfo=timezone.utc),
"ends_at": datetime.now(timezone.utc) + timedelta(days=7),
"keys": (_key_row(),),
}
return _spec_record({**defaults, **overrides})
def _shadow_prisma(claimed=(), agg_rows=None, by_key_rows=None, known_keys=("key-hash", "key-hash-2")) -> MagicMock:
prisma = MagicMock()
prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=MagicMock())
prisma.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[MagicMock(token=token) for token in known_keys]
)
prisma.db.execute_raw = AsyncMock(return_value=0)
prisma.db.litellm_shadowevaljob.find_first = AsyncMock(return_value=active_job)
async def find_claimed_keys(where: dict, **_: object):
"""Honours the filter it is handed, so a claim read that forgets stopped_at reads the
stopped rows a real partial index would have released, and one that forgets direction
reads the opposite-direction rows a key is entitled to hold at the same time."""
scoped = [row for row in claimed if row.api_key_id in where["api_key_id"]["in"]]
by_direction = [row for row in scoped if row.direction == where.get("direction", row.direction)]
if "stopped_at" not in where:
return by_direction
return [row for row in by_direction if row.stopped_at is where["stopped_at"]]
prisma.db.litellm_shadowevaljobkey.find_many = AsyncMock(side_effect=find_claimed_keys)
prisma.db.litellm_shadowevaljobkey.update_many = AsyncMock(return_value=1)
prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=None)
prisma.db.litellm_shadowevaljob.find_many = AsyncMock(return_value=[])
prisma.db.litellm_shadowevaljob.create = AsyncMock(return_value=_job_record())
prisma.db.litellm_shadowevaljob.update = AsyncMock(
return_value=_job_record(stopped_at=datetime.now(timezone.utc))
)
prisma.db.litellm_shadowevalattempt.find_first = AsyncMock(return_value=None)
async def query_raw(sql: str, *params: object):
if "FILTER (WHERE outcome != 'error')::int AS judged_count" in sql:
return [{"judged_count": 10, "error_count": 2, "judge_spend": 0.031}]
if "api_key_id AS grp" in sql:
return by_key_rows if by_key_rows is not None else []
return agg_rows if agg_rows is not None else []
prisma.db.query_raw = AsyncMock(side_effect=query_raw)
@ -545,7 +571,7 @@ def _shadow_prisma(active_job=None, agg_rows=None) -> MagicMock:
def _start_request(**overrides: object) -> StartShadowEvalRequest:
payload = {
"api_key_id": "key-hash",
"api_key_ids": ("key-hash",),
"router_name": "my-router",
"shadow_percentage": 10.0,
"judge_model": "anthropic/claude-sonnet-5",
@ -559,27 +585,33 @@ def _start_request(**overrides: object) -> StartShadowEvalRequest:
@pytest.mark.asyncio
async def test_start_shadow_eval_creates_job_and_frees_expired_or_exhausted_ones(monkeypatch: pytest.MonkeyPatch):
"""Expiry and turn-budget exhaustion both end sampling on their own; either must
release the key's slot in the active-job index so a new eval can start."""
release the key's slot in the active-job index so a new eval can start. Every requested
key is swept and claimed on its own row, and each gets its own copy of the turn budget."""
import litellm.proxy.proxy_server as proxy_server
prisma = _shadow_prisma()
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
response = await start_shadow_eval(_start_request(), ADMIN)
response = await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN)
assert response.status == "running"
assert response.max_turns == 200
assert response.judged_count is None
sweep_sql, sweep_key = prisma.db.execute_raw.call_args.args
sweep_sql = prisma.db.execute_raw.call_args.args[0]
assert "stopped_at IS NULL" in sweep_sql
assert "ends_at <= NOW()" in sweep_sql
assert ">= j.max_turns" in sweep_sql
assert sweep_key == "key-hash"
assert "j.ends_at <= NOW()" in sweep_sql
assert ">= k.max_turns" in sweep_sql
assert "a.api_key_id = k.api_key_id" in sweep_sql
assert "k.api_key_id = ANY($1::text[])" in sweep_sql
assert prisma.db.litellm_shadowevaljob.create.call_args.kwargs["include"] == {"keys": True}
assert prisma.db.execute_raw.call_args.args[1] == ["key-hash", "key-hash-2"]
create_data = prisma.db.litellm_shadowevaljob.create.call_args.kwargs["data"]
assert create_data["api_key_id"] == "key-hash"
assert create_data["keys"]["create"] == [
{"api_key_id": "key-hash", "max_turns": 200},
{"api_key_id": "key-hash-2", "max_turns": 200},
]
assert create_data["created_by"] == "admin"
assert "status" not in create_data
assert "status" not in create_data and "api_key_id" not in create_data
@pytest.mark.asyncio
@ -591,7 +623,8 @@ async def test_start_shadow_eval_creates_job_and_frees_expired_or_exhausted_ones
(ADMIN, {"router_name": "not-a-router"}, None, 400),
(ADMIN, {"judge_model": "not/a real model!"}, None, 400),
(ADMIN, {"judge_model": "my-router"}, None, 400),
(ADMIN, {}, "active", 409),
(ADMIN, {}, ("key-hash",), 409),
(ADMIN, {"api_key_ids": ("key-hash", "key-hash-2")}, ("key-hash-2",), 409),
(ADMIN, {"direction": "reverse", "baseline_model": "my-router"}, None, 400),
(ADMIN, {"direction": "reverse", "baseline_model": "not/a real model!"}, None, 400),
(ADMIN, {"direction": "reverse", "baseline_model": "openai/gpt-4o", "router_name": "not-a-router"}, None, 400),
@ -603,6 +636,7 @@ async def test_start_shadow_eval_creates_job_and_frees_expired_or_exhausted_ones
"unresolvable-judge",
"router-as-judge",
"already-active",
"one-of-several-keys-already-active",
"router-as-baseline",
"unresolvable-baseline",
"reverse-still-needs-an-auto-router",
@ -613,7 +647,7 @@ async def test_start_shadow_eval_rejections(
):
import litellm.proxy.proxy_server as proxy_server
prisma = _shadow_prisma(active_job=_job_record() if active else None)
prisma = _shadow_prisma(claimed=[_claim(key) for key in active or ()])
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
@ -644,11 +678,8 @@ async def test_start_shadow_eval_reverse_records_its_arms_and_holds_its_own_slot
the slot must not block a reverse one. The second reverse start still 409s."""
import litellm.proxy.proxy_server as proxy_server
prisma = _shadow_prisma()
active = {"forward": _job_record()}
prisma.db.litellm_shadowevaljob.find_first = AsyncMock(
side_effect=lambda where, **_: active.get(str(where.get("direction")))
)
claimed = [_claim("key-hash", job_id="job-1", direction="forward")]
prisma = _shadow_prisma(claimed=claimed)
prisma.db.litellm_shadowevaljob.create = AsyncMock(
return_value=_job_record(direction="reverse", baseline_model="openai/gpt-4o")
)
@ -663,7 +694,7 @@ async def test_start_shadow_eval_reverse_records_its_arms_and_holds_its_own_slot
assert create_data["direction"] == "reverse"
assert create_data["baseline_model"] == "openai/gpt-4o"
active["reverse"] = _job_record(id="job-2", direction="reverse")
claimed.append(_claim("key-hash", job_id="job-2", direction="reverse"))
with pytest.raises(HTTPException) as exc:
await start_shadow_eval(reverse, ADMIN)
assert exc.value.status_code == 409
@ -685,19 +716,62 @@ async def test_start_shadow_eval_forward_leaves_the_baseline_column_empty(monkey
@pytest.mark.asyncio
async def test_start_shadow_eval_rejects_a_key_this_proxy_does_not_know(monkeypatch: pytest.MonkeyPatch):
"""A typo'd api_key_id would otherwise create a job no traffic can ever match."""
async def test_start_shadow_eval_rejects_a_key_already_active_in_another_job(monkeypatch: pytest.MonkeyPatch):
"""One active job per key holds across jobs, so a key busy elsewhere blocks the whole
start rather than being silently dropped from it, and the 409 names which key and where."""
import litellm.proxy.proxy_server as proxy_server
prisma = _shadow_prisma()
prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
prisma = _shadow_prisma(claimed=[_claim("key-hash-2", job_id="job-7")])
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
with pytest.raises(HTTPException) as exc:
await start_shadow_eval(_start_request(), ADMIN)
await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN)
assert exc.value.status_code == 409
assert "key-hash-2 (job job-7)" in exc.value.detail
prisma.db.litellm_shadowevaljob.create.assert_not_called()
@pytest.mark.asyncio
async def test_start_shadow_eval_reuses_a_key_whose_previous_job_already_stopped(monkeypatch: pytest.MonkeyPatch):
"""The claim is held by unstopped key rows only, matching the partial unique index. A read
that forgets that would strand every key that has ever finished a job."""
import litellm.proxy.proxy_server as proxy_server
prisma = _shadow_prisma(claimed=[_claim("key-hash", job_id="job-7", stopped_at=datetime.now(timezone.utc))])
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
job = await start_shadow_eval(_start_request(), ADMIN)
assert job.job_id == "job-1"
prisma.db.litellm_shadowevaljob.create.assert_awaited_once()
@pytest.mark.asyncio
async def test_start_shadow_eval_rejects_a_key_this_proxy_does_not_know(monkeypatch: pytest.MonkeyPatch):
"""A typo'd api_key_id would otherwise create a job no traffic can ever match. Every
unknown key is named at once, so a caller passing several fixes them in one round."""
import litellm.proxy.proxy_server as proxy_server
prisma = _shadow_prisma(known_keys=("key-hash",))
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "llm_router", _shadow_router())
with pytest.raises(HTTPException) as exc:
await start_shadow_eval(_start_request(api_key_ids=("key-hash", "typo-a", "typo-b")), ADMIN)
assert exc.value.status_code == 400
assert "not a key on this proxy" in exc.value.detail
assert "typo-a, typo-b" in exc.value.detail
assert "key-hash," not in exc.value.detail
prisma.db.litellm_shadowevaljob.create.assert_not_called()
def test_start_shadow_eval_request_dedupes_and_requires_a_key():
"""A key named twice would collide with itself on the one-active-per-key index, and a
job scoping no key samples nothing."""
assert _start_request(api_key_ids=("a", "b", "a")).api_key_ids == ("a", "b")
with pytest.raises(ValueError):
_start_request(api_key_ids=())
@pytest.mark.asyncio
@ -725,11 +799,15 @@ async def test_get_shadow_eval_job_derives_counts_spend_and_stratified_results(m
{"grp": "SIMPLE", "turn_count": 8, "real_wins": 2, "shadow_wins": 4, "ties": 2, "avg_confidence": 0.8},
{"grp": "REASONING", "turn_count": 2, "real_wins": 2, "shadow_wins": 0, "ties": 0, "avg_confidence": 0.9},
]
prisma = _shadow_prisma(agg_rows=tier_rows)
prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=_job_record())
prisma.db.litellm_shadowevalattempt.find_first = AsyncMock(
return_value=MagicMock(error="judge call failed: boom")
key_rows = [
{"grp": "key-hash", "turn_count": 6, "real_wins": 1, "shadow_wins": 4, "ties": 1, "avg_confidence": 0.7},
{"grp": "key-hash-2", "turn_count": 4, "real_wins": 3, "shadow_wins": 0, "ties": 1, "avg_confidence": 0.6},
]
prisma = _shadow_prisma(agg_rows=tier_rows, by_key_rows=key_rows)
prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(
return_value=_job_record(keys=(_key_row(), _key_row("key-hash-2", max_turns=50)))
)
prisma.db.litellm_shadowevalattempt.find_first = AsyncMock(return_value=MagicMock(error="judge call failed: boom"))
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
response = await get_shadow_eval_job("job-1", VIEWER)
@ -744,6 +822,10 @@ async def test_get_shadow_eval_job_derives_counts_spend_and_stratified_results(m
assert response.results.by_tier[0].shadow_win_rate_pct == 50.0
assert response.results.overall_shadow_win_rate_pct == 40.0
assert response.results.overall_tie_rate_pct == 20.0
assert [(key.api_key_id, key.max_turns) for key in response.keys] == [("key-hash", 200), ("key-hash-2", 50)]
assert [(s.group, s.turn_count) for s in response.results.by_key] == [("key-hash", 6), ("key-hash-2", 4)]
assert response.results.by_key[0].shadow_win_rate_pct == 66.7
assert prisma.db.litellm_shadowevaljob.find_unique.await_args.kwargs["include"] == {"keys": True}
@pytest.mark.asyncio
@ -765,44 +847,80 @@ async def test_get_shadow_eval_job_404s_and_gates_on_role(monkeypatch: pytest.Mo
async def test_list_shadow_eval_jobs_returns_derived_status_without_aggregates(monkeypatch: pytest.MonkeyPatch):
import litellm.proxy.proxy_server as proxy_server
stamp = datetime.now(timezone.utc)
prisma = _shadow_prisma()
prisma.db.litellm_shadowevaljob.find_many = AsyncMock(
return_value=[
_job_record(),
_job_record(id="job-2", ends_at=datetime.now(timezone.utc) - timedelta(days=1)),
_job_record(id="job-3", stopped_at=datetime.now(timezone.utc)),
_job_record(id="job-3", keys=(_key_row(stopped_at=stamp),)),
]
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
jobs = await list_shadow_eval_jobs(VIEWER, api_key_id=None, limit=50)
jobs = await list_shadow_eval_jobs(VIEWER, api_key_id="key-hash", limit=50)
assert [job.status for job in jobs] == ["running", "completed", "stopped"]
swept = ShadowEvalJobResponse.model_validate(
assert all(job.judged_count is None and job.results is None for job in jobs)
assert prisma.db.query_raw.await_count == 0
assert prisma.db.litellm_shadowevaljob.find_many.await_args.kwargs["where"] == {
"keys": {"some": {"api_key_id": "key-hash"}}
}
assert prisma.db.litellm_shadowevaljob.find_many.await_args.kwargs["include"] == {"keys": True}
@pytest.mark.parametrize(
("stopped_flags", "days_left", "expected"),
[
((False, False), 7, "running"),
((True, False), 7, "running"),
((True, True), 7, "stopped"),
((True, True), -1, "completed"),
((False, False), -1, "completed"),
],
)
def test_job_status_runs_until_every_key_stops_and_completed_outranks_stopped(
stopped_flags: tuple[bool, ...], days_left: int, expected: str
):
stamp = datetime.now(timezone.utc)
job = ShadowEvalJobResponse.model_validate(
_job_record(
id="job-4",
ends_at=datetime.now(timezone.utc) - timedelta(days=1),
stopped_at=datetime.now(timezone.utc),
ends_at=datetime.now(timezone.utc) + timedelta(days=days_left),
keys=tuple(
_key_row(f"key-{index}", stopped_at=stamp if stopped else None)
for index, stopped in enumerate(stopped_flags)
),
),
from_attributes=True,
)
assert swept.status == "completed"
assert all(job.judged_count is None and job.results is None for job in jobs)
assert prisma.db.query_raw.await_count == 0
assert job.status == expected
@pytest.mark.asyncio
async def test_stop_shadow_eval_sets_stopped_at_and_rejects_non_running(monkeypatch: pytest.MonkeyPatch):
async def test_stop_shadow_eval_stops_every_unstopped_key_row_and_rejects_non_running(
monkeypatch: pytest.MonkeyPatch,
):
import litellm.proxy.proxy_server as proxy_server
stamp = datetime.now(timezone.utc)
prisma = _shadow_prisma()
prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(return_value=_job_record())
prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(
side_effect=[
_job_record(keys=(_key_row(), _key_row("key-hash-2"))),
_job_record(keys=(_key_row(stopped_at=stamp), _key_row("key-hash-2", stopped_at=stamp))),
]
)
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
stopped = await stop_shadow_eval_job("job-1", ADMIN)
assert stopped.status == "stopped"
update = prisma.db.litellm_shadowevaljob.update.call_args.kwargs
assert [key.stopped_at for key in stopped.keys] == [stamp, stamp]
update = prisma.db.litellm_shadowevaljobkey.update_many.await_args.kwargs
assert update["where"] == {"job_id": "job-1", "stopped_at": None}
assert set(update["data"]) == {"stopped_at"}
assert prisma.db.litellm_shadowevaljob.find_unique.await_args.kwargs["include"] == {"keys": True}
assert not prisma.db.litellm_shadowevaljob.update.called
prisma.db.litellm_shadowevaljob.find_unique = AsyncMock(
return_value=_job_record(ends_at=datetime.now(timezone.utc) - timedelta(days=1))

View file

@ -1,4 +1,4 @@
import { fireEvent, render, screen } from "@testing-library/react";
import { fireEvent, render, screen, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import React from "react";
import { beforeEach, describe, expect, it, vi } from "vitest";
@ -72,12 +72,12 @@ import {
const job = (overrides: Partial<ShadowEvalJob> = {}): ShadowEvalJob => ({
job_id: "job-1",
status: "running",
router_name: "claude-auto",
direction: "forward",
router_name: "claude-auto",
baseline_model: null,
judge_model: "anthropic/claude-sonnet-5",
shadow_percentage: 10,
max_turns: 200,
keys: [{ api_key_id: "hashed-key-abc", max_turns: 200, stopped_at: null }],
judged_count: 42,
error_count: 1,
judge_spend: 3.21,
@ -110,13 +110,21 @@ const job = (overrides: Partial<ShadowEvalJob> = {}): ShadowEvalJob => ({
avg_judge_confidence: 0.8,
},
],
by_key: [
{
group: "hashed-key-abc",
turn_count: 42,
real_win_rate_pct: 30.0,
shadow_win_rate_pct: 45.0,
tie_rate_pct: 25.0,
avg_judge_confidence: 0.8,
},
],
overall_shadow_win_rate_pct: 48.0,
overall_tie_rate_pct: 22.0,
},
created_at: "2026-08-07T00:00:00Z",
ends_at: "2026-09-07T00:00:00Z",
stopped_at: null,
api_key_id: "hashed-key-abc",
last_error: null,
...overrides,
});
@ -197,8 +205,8 @@ describe("ShadowEvalSection", () => {
it("gives every active job its own card with a stop button, with the form still offered", () => {
mockHooks({
jobs: [
job({ job_id: "job-a", status: "running", api_key_id: "key-a" }),
job({ job_id: "job-b", status: "running", api_key_id: "key-b" }),
job({ job_id: "job-a", status: "running", keys: [{ api_key_id: "key-a", max_turns: 200, stopped_at: null }] }),
job({ job_id: "job-b", status: "running", keys: [{ api_key_id: "key-b", max_turns: 200, stopped_at: null }] }),
],
});
render(<ShadowEvalSection />);
@ -219,7 +227,7 @@ describe("ShadowEvalSection", () => {
render(<ShadowEvalSection />);
expect(screen.queryByText("Start a shadow eval")).not.toBeInTheDocument();
expect(screen.queryByRole("button", { name: "Stop" })).not.toBeInTheDocument();
expect(screen.getByText("running")).toBeInTheDocument();
expect(screen.getAllByText("running").filter((badge) => badge.closest("td") === null)).toHaveLength(1);
});
it("never labels a collapsed previous eval as empty from a countless list row", () => {
@ -334,21 +342,91 @@ describe("ShadowEvalSection", () => {
expect(screen.getByText("Start a shadow eval")).toBeInTheDocument();
});
it("breaks results down per key, so one key exhausting its own budget is visible while a sibling runs on", () => {
mockHooks({
jobs: [
job({
judged_count: 205,
keys: [
{ api_key_id: "hash-spent", max_turns: 200, stopped_at: "2026-08-08T00:00:00Z" },
{ api_key_id: "hash-hungry", max_turns: 500, stopped_at: null },
],
results: {
by_tier: [],
by_current_model: [],
by_key: [
{
group: "hash-spent",
turn_count: 200,
real_win_rate_pct: 20.0,
shadow_win_rate_pct: 60.0,
tie_rate_pct: 20.0,
avg_judge_confidence: 0.9,
},
],
overall_shadow_win_rate_pct: 60.0,
overall_tie_rate_pct: 20.0,
},
}),
],
});
render(<ShadowEvalSection />);
const spent = screen.getByText("hash-spent").closest("tr");
const hungry = screen.getByText("hash-hungry").closest("tr");
if (!spent || !hungry) throw new Error("expected a table row per scoped key");
expect(within(spent).getByText("stopped")).toBeInTheDocument();
expect(within(spent).getByText("200 / 200")).toBeInTheDocument();
expect(within(spent).getByText("60.0%")).toBeInTheDocument();
expect(within(hungry).getByText("running")).toBeInTheDocument();
expect(within(hungry).getByText("0 / 500")).toBeInTheDocument();
expect(within(hungry).getByText("No verdicts yet")).toBeInTheDocument();
expect(screen.getByText(/205 of 700 turns judged/)).toBeInTheDocument();
});
it("reads every key as completed once the job's window closes, whatever its own stop state", () => {
mockHooks({
jobs: [
job({
status: "completed",
keys: [
{ api_key_id: "hash-spent", max_turns: 200, stopped_at: "2026-08-08T00:00:00Z" },
{ api_key_id: "hash-hungry", max_turns: 500, stopped_at: null },
],
}),
],
});
render(<ShadowEvalSection />);
const hungry = screen.getByText("hash-hungry").closest("tr");
if (!hungry) throw new Error("expected a table row per scoped key");
expect(within(hungry).getByText("completed")).toBeInTheDocument();
expect(within(hungry).queryByText("running")).not.toBeInTheDocument();
});
it("renders nothing for non-admins when the proxy answers 403", () => {
mockHooks({ error: new ApiError("forbidden", 403, {}) });
const { container } = render(<ShadowEvalSection />);
expect(container).toBeEmptyDOMElement();
});
it("keeps the start button disabled until key, router, and judge model are picked, then submits them", async () => {
it("keeps the start button disabled until key, router, and judge model are picked, then submits every picked key", async () => {
const user = userEvent.setup();
const { start } = mockHooks({});
render(<ShadowEvalSection />);
expect(screen.getByText("Start shadow eval")).toBeDisabled();
await user.click(screen.getByPlaceholderText("Search keys by alias"));
await user.click(await screen.findByText("prod-alpha"));
const keyInput = screen.getByPlaceholderText("Search keys by alias");
await user.click(keyInput);
const keyList = await screen.findByTestId("paginated-multi-select-list");
await user.click(within(keyList).getByText("prod-alpha"));
await user.click(keyInput);
await user.click(within(keyList).getByText("staging-beta"));
await user.click(screen.getByPlaceholderText("Select an auto-router"));
await user.click(await screen.findByText("gpt-auto"));
@ -359,7 +437,7 @@ describe("ShadowEvalSection", () => {
await user.click(screen.getByText("Start shadow eval"));
const expectedBody = {
api_key_id: "hash-alpha",
api_key_ids: ["hash-alpha", "hash-beta"],
router_name: "gpt-auto",
direction: "forward",
shadow_percentage: 10,
@ -394,7 +472,7 @@ describe("ShadowEvalSection", () => {
await user.click(screen.getByText("Start shadow eval"));
const expectedBody = {
api_key_id: "hash-alpha",
api_key_ids: ["hash-alpha"],
router_name: "gpt-auto",
direction: "reverse",
baseline_model: "prod-claude",
@ -416,7 +494,7 @@ describe("ShadowEvalSection", () => {
expect(screen.getByText("52.0%")).toBeInTheDocument();
expect(screen.getByText(/Router won 30.0%/)).toBeInTheDocument();
expect(screen.getByText(/Baseline won 48.0%/)).toBeInTheDocument();
expect(screen.getAllByText("Baseline wins")).toHaveLength(2);
expect(screen.getAllByText("Baseline wins")).toHaveLength(3);
expect(screen.getByText("Router pick")).toBeInTheDocument();
expect(screen.queryByText(/Current model/)).not.toBeInTheDocument();
expect(screen.queryByText("Compared against")).not.toBeInTheDocument();
@ -444,5 +522,6 @@ describe("ShadowEvalSection", () => {
expect(await screen.findByText("SIMPLE")).toBeInTheDocument();
expect(screen.getByText("REASONING")).toBeInTheDocument();
expect(screen.getByText("42 / 200")).toBeInTheDocument();
});
});

View file

@ -6,7 +6,7 @@ import { useInfiniteKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap";
import { useAutoRouters, usePlainModelGroups } from "@/app/(dashboard)/hooks/models/useModels";
import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect";
import { PaginatedMultiSelect } from "@/components/shared/PaginatedMultiSelect";
import { SearchSelect, type SearchSelectOption } from "@/components/shared/SearchSelect";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
@ -24,7 +24,9 @@ import {
useStartShadowEval,
useStopShadowEval,
type ShadowEvalJob,
type ShadowEvalJobKey,
type ShadowEvalSlice,
type StartShadowEvalRequest,
} from "./useShadowEval";
const pct = (value: number): string => `${value.toFixed(1)}%`;
@ -164,6 +166,68 @@ const VerdictBar: React.FC<{ direction: ShadowEvalDirection; results: NonNullabl
);
};
const totalBudget = (job: ShadowEvalJob): number => job.keys.reduce((sum, key) => sum + key.max_turns, 0);
const keyStatus = (job: ShadowEvalJob, key: ShadowEvalJobKey): string => {
if (job.status === "completed") return "completed";
return key.stopped_at ? "stopped" : "running";
};
const KeyTable: React.FC<{ job: ShadowEvalJob }> = ({ job }) => {
const slices = new Map((job.results?.by_key ?? []).map((slice) => [slice.group, slice]));
return (
<Table>
<TableHeader>
<TableRow>
<TableHead>Key</TableHead>
<TableHead>Status</TableHead>
{["Judged of budget", "Router wins", `${otherArmLabel(job.direction)} wins`, "Ties", "Judge confidence"].map(
(label) => (
<TableHead key={label} className="text-right">
{label}
</TableHead>
),
)}
</TableRow>
</TableHeader>
<TableBody>
{job.keys.map((key) => {
const slice = slices.get(key.api_key_id);
return (
<TableRow key={key.api_key_id}>
<TableCell className="max-w-[16rem] truncate font-mono text-xs" title={key.api_key_id}>
{key.api_key_id}
</TableCell>
<TableCell>
<StatusBadge status={keyStatus(job, key)} />
</TableCell>
<TableCell className="text-right tabular-nums">
{(slice?.turn_count ?? 0).toLocaleString()} / {key.max_turns.toLocaleString()}
</TableCell>
{slice ? (
<>
<TableCell className="text-right font-medium tabular-nums text-foreground">
{pct(routerWinRate(job.direction, slice))}
</TableCell>
<TableCell className="text-right tabular-nums">
{pct(otherArmWinRate(job.direction, slice))}
</TableCell>
<TableCell className="text-right tabular-nums">{pct(slice.tie_rate_pct)}</TableCell>
<TableCell className="text-right tabular-nums">{slice.avg_judge_confidence.toFixed(2)}</TableCell>
</>
) : (
<TableCell colSpan={4} className="text-right text-xs text-muted-foreground">
No verdicts yet
</TableCell>
)}
</TableRow>
);
})}
</TableBody>
</Table>
);
};
const emptyResultsText = (job: ShadowEvalJob, resultsError: boolean): string => {
if (resultsError) return "Results could not be loaded. Retrying.";
if (isActive(job)) return "Collecting verdicts. Results appear as sampled requests are judged.";
@ -173,31 +237,41 @@ const emptyResultsText = (job: ShadowEvalJob, resultsError: boolean): string =>
const ResultsBody: React.FC<{ job: ShadowEvalJob; resultsError?: boolean }> = ({ job, resultsError = false }) => {
const results = job.results;
if (!results || (results.by_tier.length === 0 && results.by_current_model.length === 0)) {
return <p className="px-6 py-8 text-center text-sm text-muted-foreground">{emptyResultsText(job, resultsError)}</p>;
}
return (
<>
<div className="flex flex-col gap-1 border-b px-6 py-4">
<p className="text-[11px] uppercase tracking-wide text-muted-foreground">
Router matched or beat {job.direction === "reverse" ? "the baseline" : "your current model"}
</p>
<p className="text-3xl font-semibold text-foreground">{pct(routerMatchedOrBeatPct(job.direction, results))}</p>
<p className="text-xs text-muted-foreground">of {(job.judged_count ?? 0).toLocaleString()} judged responses</p>
{results && (results.by_tier.length > 0 || results.by_current_model.length > 0) ? (
<>
<div className="flex flex-col gap-1 border-b px-6 py-4">
<p className="text-[11px] uppercase tracking-wide text-muted-foreground">
Router matched or beat {job.direction === "reverse" ? "the baseline" : "your current model"}
</p>
<p className="text-3xl font-semibold text-foreground">
{pct(routerMatchedOrBeatPct(job.direction, results))}
</p>
<p className="text-xs text-muted-foreground">
of {(job.judged_count ?? 0).toLocaleString()} judged responses
</p>
</div>
<VerdictBar direction={job.direction} results={results} />
{results.by_current_model.length > 0 && (
<SliceTable
groupHeader={job.direction === "reverse" ? "Router pick" : "Compared against"}
direction={job.direction}
slices={results.by_current_model}
/>
)}
{results.by_tier.length > 0 && (
<div className={results.by_current_model.length > 0 ? "border-t" : ""}>
<SliceTable groupHeader="Prompt difficulty" direction={job.direction} slices={results.by_tier} />
</div>
)}
</>
) : (
<p className="px-6 py-8 text-center text-sm text-muted-foreground">{emptyResultsText(job, resultsError)}</p>
)}
<div className="border-t">
<KeyTable job={job} />
</div>
<VerdictBar direction={job.direction} results={results} />
{results.by_current_model.length > 0 && (
<SliceTable
groupHeader={job.direction === "reverse" ? "Router pick" : "Compared against"}
direction={job.direction}
slices={results.by_current_model}
/>
)}
{results.by_tier.length > 0 && (
<div className={results.by_current_model.length > 0 ? "border-t" : ""}>
<SliceTable groupHeader="Prompt difficulty" direction={job.direction} slices={results.by_tier} />
</div>
)}
</>
);
};
@ -219,7 +293,7 @@ const JobResults: React.FC<{
<div>
<p className="text-sm font-medium text-foreground">{jobHeadline(job)}</p>
<p className="text-xs text-muted-foreground">
{(job.judged_count ?? 0).toLocaleString()} of {job.max_turns.toLocaleString()} turns judged ·{" "}
{(job.judged_count ?? 0).toLocaleString()} of {totalBudget(job).toLocaleString()} turns judged ·{" "}
{(job.error_count ?? 0).toLocaleString()} errored · {usd(job.judge_spend ?? 0)} judge spend
{active && remaining ? ` · ${remaining}` : ""}
</p>
@ -294,9 +368,9 @@ const DIRECTION_OPTIONS: readonly { value: ShadowEvalDirection; label: string }[
const START_FORM_DESCRIPTION: Record<ShadowEvalDirection, string> = {
forward:
"Duplicates a sampled slice of the key's traffic through the auto-router and has an LLM judge compare both answers blind. The router's answers are never served to users; judge calls bill to the shadowed key.",
"Duplicates a sampled slice of each selected key's traffic through the auto-router and has an LLM judge compare both answers blind. The router's answers are never served to users; judge calls bill to the key whose traffic was sampled.",
reverse:
"Duplicates a sampled slice of the traffic the auto-router already serves against a fixed baseline model and has an LLM judge compare both answers blind. The baseline's answers are never served to users; judge calls bill to the shadowed key.",
"Duplicates a sampled slice of the traffic the auto-router already serves for each selected key against a fixed baseline model and has an LLM judge compare both answers blind. The baseline's answers are never served to users; judge calls bill to the key whose traffic was sampled.",
};
const DURATION_OPTIONS = [
@ -321,7 +395,7 @@ const Field: React.FC<{ label: string; htmlFor?: string; className?: string; chi
</div>
);
const KeySelect: React.FC<{ value: string; onChange: (token: string) => void }> = ({ value, onChange }) => {
const KeySelect: React.FC<{ value: string[]; onChange: (tokens: string[]) => void }> = ({ value, onChange }) => {
const [search, setSearch] = useState("");
const { data, isPending, isError, fetchNextPage, hasNextPage, isFetchingNextPage } = useInfiniteKeys(50, {
selectedKeyAlias: search || null,
@ -338,7 +412,7 @@ const KeySelect: React.FC<{ value: string; onChange: (token: string) => void }>
[data],
);
return (
<PaginatedSearchSelect
<PaginatedMultiSelect
inputId="shadow-eval-key"
options={options}
value={value}
@ -357,7 +431,7 @@ const KeySelect: React.FC<{ value: string; onChange: (token: string) => void }>
const StartForm: React.FC = () => {
const { accessToken } = useAuthorized();
const [apiKeyId, setApiKeyId] = useState("");
const [apiKeyIds, setApiKeyIds] = useState<string[]>([]);
const [routerName, setRouterName] = useState("");
const [direction, setDirection] = useState<ShadowEvalDirection>("forward");
const [baselineModel, setBaselineModel] = useState("");
@ -381,14 +455,13 @@ const StartForm: React.FC = () => {
const percentageValid = parsedPct >= 0.1 && parsedPct <= 100;
const parsedMaxTurns = Number.parseInt(maxTurns, 10);
const maxTurnsValid = parsedMaxTurns >= 1 && parsedMaxTurns <= 2000;
const filled =
[apiKeyId, routerName, judgeModel].every((field) => field !== "") &&
(direction === "forward" || baselineModel !== "");
const baselinePicked = direction === "forward" || baselineModel !== "";
const filled = apiKeyIds.length > 0 && [routerName, judgeModel].every((field) => field !== "") && baselinePicked;
const boundsValid = percentageValid && maxTurnsValid;
const valid = Boolean(accessToken) && filled && boundsValid;
const handleStart = () => {
const startBody = {
api_key_id: apiKeyId,
const startBody: StartShadowEvalRequest = {
api_key_ids: apiKeyIds,
router_name: routerName,
direction,
...(direction === "reverse" ? { baseline_model: baselineModel } : {}),
@ -425,8 +498,8 @@ const StartForm: React.FC = () => {
</SelectContent>
</Select>
</Field>
<Field label="Key to shadow" htmlFor="shadow-eval-key">
<KeySelect value={apiKeyId} onChange={setApiKeyId} />
<Field label="Keys to shadow" htmlFor="shadow-eval-key">
<KeySelect value={apiKeyIds} onChange={setApiKeyIds} />
</Field>
<Field label="Auto-router">
<SearchSelect
@ -481,7 +554,7 @@ const StartForm: React.FC = () => {
value={maxTurns}
onChange={(e) => setMaxTurns(e.target.value)}
/>
<span className="text-sm text-muted-foreground">turns judged, max</span>
<span className="text-sm text-muted-foreground">turns judged per key, max</span>
</div>
{maxTurns.trim() !== "" && !maxTurnsValid && (
<p className="text-xs text-destructive">Enter a value from 1 to 2000</p>

View file

@ -7,6 +7,7 @@ import { $api, fetchClient } from "@/lib/http/api";
import type { components } from "@/lib/http/schema";
export type ShadowEvalJob = components["schemas"]["ShadowEvalJobResponse"];
export type ShadowEvalJobKey = components["schemas"]["ShadowEvalJobKeyResponse"];
export type ShadowEvalSlice = components["schemas"]["ShadowEvalSlice"];
export type StartShadowEvalRequest = components["schemas"]["StartShadowEvalRequest"];

View file

@ -0,0 +1,193 @@
import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { useState } from "react";
import { describe, expect, it, vi } from "vitest";
import { PaginatedMultiSelect } from "./PaginatedMultiSelect";
import type { SearchSelectOption } from "./SearchSelect";
const OPTIONS: SearchSelectOption[] = [
{ label: "alias-alpha", value: "alias-alpha" },
{ label: "alias-beta", value: "alias-beta" },
{ label: "gamma-key", value: "gamma-key" },
];
function renderSelect(overrides: Partial<React.ComponentProps<typeof PaginatedMultiSelect>> = {}) {
const props: React.ComponentProps<typeof PaginatedMultiSelect> = {
options: OPTIONS,
onValueChange: vi.fn(),
onSearchChange: vi.fn(),
onLoadMore: vi.fn(),
...overrides,
};
render(<PaginatedMultiSelect {...props} />);
return props;
}
function setListMetrics(list: HTMLElement, metrics: { scrollTop: number; clientHeight: number; scrollHeight: number }) {
Object.defineProperty(list, "scrollTop", { value: metrics.scrollTop, configurable: true });
Object.defineProperty(list, "clientHeight", { value: metrics.clientHeight, configurable: true });
Object.defineProperty(list, "scrollHeight", { value: metrics.scrollHeight, configurable: true });
}
function chipRemoveButton(label: string): HTMLElement {
const chip = screen.getByText(label).closest('[data-slot="combobox-chip"]');
if (chip === null) throw new Error(`no chip found for ${label}`);
const button = chip.querySelector('[data-slot="combobox-chip-remove"]');
if (button === null) throw new Error(`no remove control found on chip for ${label}`);
return button as HTMLElement;
}
describe("PaginatedMultiSelect", () => {
it("reports the cleared query upstream after a selection, so the next open is not still filtered", async () => {
const user = userEvent.setup();
const onSearchChange = vi.fn();
function Controlled() {
const [value, setValue] = useState<string[]>([]);
return (
<PaginatedMultiSelect
options={OPTIONS}
value={value}
onValueChange={setValue}
onSearchChange={onSearchChange}
onLoadMore={vi.fn()}
/>
);
}
render(<Controlled />);
const input = screen.getByRole("combobox");
await user.click(input);
await user.type(input, "alias-a");
await waitFor(() => expect(onSearchChange).toHaveBeenCalledWith("alias-a"), { timeout: 2000 });
const list = await screen.findByTestId("paginated-multi-select-list");
await user.click(within(list).getByText("alias-alpha"));
expect(input).toHaveValue("");
await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith(""), { timeout: 2000 });
});
it("selects multiple values and reports them cumulatively", async () => {
const user = userEvent.setup();
const onValueChange = vi.fn();
function Controlled() {
const [value, setValue] = useState<string[]>([]);
return (
<PaginatedMultiSelect
options={OPTIONS}
value={value}
onValueChange={(next) => {
setValue(next);
onValueChange(next);
}}
onSearchChange={vi.fn()}
onLoadMore={vi.fn()}
/>
);
}
render(<Controlled />);
const input = screen.getByRole("combobox");
await user.click(input);
const list = await screen.findByTestId("paginated-multi-select-list");
await user.click(within(list).getByText("alias-alpha"));
await user.click(input);
await user.click(within(list).getByText("gamma-key"));
expect(onValueChange).toHaveBeenLastCalledWith(["alias-alpha", "gamma-key"]);
const chips = document.querySelector('[data-slot="combobox-chips"]') as HTMLElement;
expect(within(chips).getByText("alias-alpha")).toBeInTheDocument();
expect(within(chips).getByText("gamma-key")).toBeInTheDocument();
});
it("deselects one value via the chip remove control and keeps the rest", async () => {
const user = userEvent.setup();
const onValueChange = vi.fn();
function Controlled() {
const [value, setValue] = useState<string[]>(["alias-alpha", "alias-beta"]);
return (
<PaginatedMultiSelect
options={OPTIONS}
value={value}
onValueChange={(next) => {
setValue(next);
onValueChange(next);
}}
onSearchChange={vi.fn()}
onLoadMore={vi.fn()}
/>
);
}
render(<Controlled />);
await user.click(chipRemoveButton("alias-alpha"));
expect(onValueChange).toHaveBeenCalledWith(["alias-beta"]);
expect(screen.queryByText("alias-alpha")).not.toBeInTheDocument();
expect(screen.getByText("alias-beta")).toBeInTheDocument();
});
it("keeps a selected chip visible after the options page no longer contains it", () => {
const { rerender } = render(
<PaginatedMultiSelect
options={OPTIONS}
value={["ghost-key"]}
onValueChange={vi.fn()}
onSearchChange={vi.fn()}
onLoadMore={vi.fn()}
/>,
);
expect(screen.getByText("ghost-key")).toBeInTheDocument();
rerender(
<PaginatedMultiSelect
options={[{ label: "alias-beta", value: "alias-beta" }]}
value={["ghost-key"]}
onValueChange={vi.fn()}
onSearchChange={vi.fn()}
onLoadMore={vi.fn()}
/>,
);
expect(screen.getByText("ghost-key")).toBeInTheDocument();
});
it("does not request the next page on scroll when there is no next page", async () => {
const user = userEvent.setup();
const onLoadMore = vi.fn();
renderSelect({ onLoadMore, hasNextPage: false });
const input = screen.getByRole("combobox");
await user.click(input);
const list = await screen.findByTestId("paginated-multi-select-list");
setListMetrics(list, { scrollTop: 900, clientHeight: 100, scrollHeight: 1000 });
fireEvent.scroll(list);
expect(onLoadMore).not.toHaveBeenCalled();
});
it("requests the next page once scrolled past the threshold when a next page exists", async () => {
const user = userEvent.setup();
const onLoadMore = vi.fn();
renderSelect({ onLoadMore, hasNextPage: true });
const input = screen.getByRole("combobox");
await user.click(input);
const list = await screen.findByTestId("paginated-multi-select-list");
setListMetrics(list, { scrollTop: 0, clientHeight: 100, scrollHeight: 1000 });
fireEvent.scroll(list);
expect(onLoadMore).not.toHaveBeenCalled();
setListMetrics(list, { scrollTop: 850, clientHeight: 100, scrollHeight: 1000 });
fireEvent.scroll(list);
expect(onLoadMore).toHaveBeenCalledTimes(1);
});
});

View file

@ -0,0 +1,140 @@
"use client";
import { Loader2 } from "lucide-react";
import { useMemo, useState } from "react";
import {
Combobox,
ComboboxChip,
ComboboxChips,
ComboboxChipsInput,
ComboboxContent,
ComboboxEmpty,
ComboboxItem,
ComboboxList,
ComboboxValue,
} from "@/components/ui/combobox";
import type { SearchSelectOption } from "./SearchSelect";
import { usePaginatedCombobox } from "./usePaginatedCombobox";
interface PaginatedMultiSelectProps {
options: SearchSelectOption[];
value?: string[];
onValueChange: (value: string[]) => void;
onSearchChange: (query: string) => void;
onLoadMore: () => void;
hasNextPage?: boolean;
isLoading?: boolean;
isFetchingNextPage?: boolean;
placeholder?: string;
emptyText?: string;
errorText?: string;
loadingText?: string;
disabled?: boolean;
className?: string;
inputId?: string;
"aria-invalid"?: true | undefined;
"aria-describedby"?: string;
}
export function PaginatedMultiSelect({
options,
value = [],
onValueChange,
onSearchChange,
onLoadMore,
hasNextPage = false,
isLoading = false,
isFetchingNextPage = false,
placeholder = "Search…",
emptyText = "No results",
errorText,
loadingText = "Loading…",
disabled = false,
className,
inputId,
"aria-invalid": ariaInvalid,
"aria-describedby": ariaDescribedBy,
}: PaginatedMultiSelectProps) {
const [query, setQuery] = useState("");
const selected = useMemo<SearchSelectOption[]>(
() =>
value.map(
(selectedValue) =>
options.find((option) => option.value === selectedValue) ?? { label: selectedValue, value: selectedValue },
),
[options, value],
);
const items = useMemo<SearchSelectOption[]>(() => {
const missing = selected.filter((option) => !options.some((o) => o.value === option.value));
return missing.length === 0 ? options : [...missing, ...options];
}, [options, selected]);
const pagination = { onSearchChange, onLoadMore, hasNextPage, isFetchingNextPage };
const { handleSearchInput, handleScroll } = usePaginatedCombobox(pagination);
const handleInputValueChange = (next: string, reason: string) => {
setQuery(next);
handleSearchInput(next, reason);
};
return (
<Combobox
multiple
items={items}
value={selected}
onValueChange={(next: SearchSelectOption[]) => onValueChange(next.map((option) => option.value))}
inputValue={query}
onInputValueChange={(next, eventDetails) => handleInputValueChange(next, eventDetails.reason)}
isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value}
itemToStringLabel={(item: SearchSelectOption) => item.label}
filter={null}
disabled={disabled}
>
<ComboboxChips className={`min-h-8 py-1 text-sm ${className ?? ""}`}>
<ComboboxValue>
{(selectedItems: SearchSelectOption[]) =>
selectedItems.map((option) => (
<ComboboxChip key={option.value} aria-label={option.label}>
{option.label}
</ComboboxChip>
))
}
</ComboboxValue>
<ComboboxChipsInput
id={inputId}
aria-invalid={ariaInvalid}
aria-describedby={ariaDescribedBy}
placeholder={placeholder}
className="h-5 min-w-24 flex-1 border-0 bg-transparent py-0 text-sm"
aria-label={placeholder}
/>
</ComboboxChips>
<ComboboxContent>
<ComboboxEmpty className={errorText == null ? undefined : "text-destructive"}>
{errorText ?? (isLoading ? loadingText : emptyText)}
</ComboboxEmpty>
<ComboboxList onScroll={handleScroll} data-testid="paginated-multi-select-list">
{(item: SearchSelectOption) => (
<ComboboxItem key={item.value} value={item}>
<span className="flex min-w-0 flex-col">
<span className="truncate">{item.label}</span>
{item.sublabel != null && item.sublabel !== "" && (
<span className="truncate text-xs text-muted-foreground">{item.sublabel}</span>
)}
</span>
</ComboboxItem>
)}
</ComboboxList>
{isFetchingNextPage && (
<div className="flex justify-center py-2" data-testid="paginated-multi-select-loading-more">
<Loader2 className="size-4 animate-spin text-muted-foreground" />
</div>
)}
</ComboboxContent>
</Combobox>
);
}

View file

@ -1,8 +1,7 @@
"use client";
import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
import { Loader2 } from "lucide-react";
import { useMemo, type UIEvent } from "react";
import { useMemo } from "react";
import {
Combobox,
@ -12,13 +11,9 @@ import {
ComboboxItem,
ComboboxList,
} from "@/components/ui/combobox";
import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
import type { SearchSelectOption } from "./SearchSelect";
const SCROLL_THRESHOLD = 0.8;
const SEARCH_REASONS: ReadonlySet<string> = new Set(["input-change", "input-clear", "clear-press"]);
import { usePaginatedCombobox } from "./usePaginatedCombobox";
interface PaginatedSearchSelectProps {
options: SearchSelectOption[];
@ -70,28 +65,15 @@ export function PaginatedSearchSelect({
return [selected, ...options];
}, [options, selected]);
const debouncedSearch = useDebouncedCallback(onSearchChange, { wait: DEBOUNCE_WAIT_MS });
const handleInputValueChange = (next: string, reason: string) => {
if (!SEARCH_REASONS.has(reason)) return;
debouncedSearch(next);
};
const handleScroll = (event: UIEvent<HTMLDivElement>) => {
const target = event.currentTarget;
if (target.scrollHeight === 0) return;
const ratio = (target.scrollTop + target.clientHeight) / target.scrollHeight;
if (ratio >= SCROLL_THRESHOLD && hasNextPage && !isFetchingNextPage) {
onLoadMore();
}
};
const pagination = { onSearchChange, onLoadMore, hasNextPage, isFetchingNextPage };
const { handleSearchInput, handleScroll } = usePaginatedCombobox(pagination);
return (
<Combobox
items={items}
value={selected}
onValueChange={(item: SearchSelectOption | null) => onValueChange(item?.value ?? "")}
onInputValueChange={(next, eventDetails) => handleInputValueChange(next, eventDetails.reason)}
onInputValueChange={(next, eventDetails) => handleSearchInput(next, eventDetails.reason)}
isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value}
itemToStringLabel={(item: SearchSelectOption) => item.label}
filter={null}

View file

@ -0,0 +1,98 @@
import { act, renderHook } from "@testing-library/react";
import type { UIEvent } from "react";
import { describe, expect, it, vi } from "vitest";
import { usePaginatedCombobox } from "./usePaginatedCombobox";
const scrollEvent = (scrollTop: number, clientHeight: number, scrollHeight: number) =>
({ currentTarget: { scrollTop, clientHeight, scrollHeight } }) as UIEvent<HTMLDivElement>;
const setup = (overrides: Partial<Parameters<typeof usePaginatedCombobox>[0]> = {}) => {
const onLoadMore = vi.fn();
const onSearchChange = vi.fn();
const options = { onSearchChange, onLoadMore, hasNextPage: true, isFetchingNextPage: false, ...overrides };
const { result } = renderHook(() => usePaginatedCombobox(options));
return { result, onLoadMore, onSearchChange };
};
describe("usePaginatedCombobox", () => {
it("loads the next page only once the list is scrolled most of the way down", () => {
const { result, onLoadMore } = setup();
result.current.handleScroll(scrollEvent(0, 100, 1000));
expect(onLoadMore).not.toHaveBeenCalled();
result.current.handleScroll(scrollEvent(690, 100, 1000));
expect(onLoadMore).not.toHaveBeenCalled();
result.current.handleScroll(scrollEvent(700, 100, 1000));
expect(onLoadMore).toHaveBeenCalledTimes(1);
});
it("does not ask for a page that does not exist or is already in flight", () => {
const exhausted = setup({ hasNextPage: false });
exhausted.result.current.handleScroll(scrollEvent(900, 100, 1000));
expect(exhausted.onLoadMore).not.toHaveBeenCalled();
const inFlight = setup({ isFetchingNextPage: true });
inFlight.result.current.handleScroll(scrollEvent(900, 100, 1000));
expect(inFlight.onLoadMore).not.toHaveBeenCalled();
});
it("ignores a scroll on an unmeasured list, whose ratio divides by a zero scroll height", () => {
const { result, onLoadMore } = setup();
result.current.handleScroll(scrollEvent(0, 100, 0));
expect(onLoadMore).not.toHaveBeenCalled();
});
it("searches on typing and clearing, but not on reasons that leave the query unchanged", async () => {
vi.useFakeTimers();
try {
const { result, onSearchChange } = setup();
result.current.handleSearchInput("alpha", "item-press");
result.current.handleSearchInput("alpha", "list-navigation");
await act(async () => {
vi.advanceTimersByTime(1000);
});
expect(onSearchChange).not.toHaveBeenCalled();
result.current.handleSearchInput("alpha", "input-change");
await act(async () => {
vi.advanceTimersByTime(1000);
});
expect(onSearchChange).toHaveBeenCalledWith("alpha");
result.current.handleSearchInput("", "input-clear");
await act(async () => {
vi.advanceTimersByTime(1000);
});
expect(onSearchChange).toHaveBeenLastCalledWith("");
} finally {
vi.useRealTimers();
}
});
it("debounces typing so one search leaves for a burst of keystrokes", async () => {
vi.useFakeTimers();
try {
const { result, onSearchChange } = setup();
result.current.handleSearchInput("a", "input-change");
result.current.handleSearchInput("al", "input-change");
result.current.handleSearchInput("alp", "input-change");
await act(async () => {
vi.advanceTimersByTime(299);
});
expect(onSearchChange).not.toHaveBeenCalled();
await act(async () => {
vi.advanceTimersByTime(1);
});
expect(onSearchChange).toHaveBeenCalledTimes(1);
expect(onSearchChange).toHaveBeenCalledWith("alp");
} finally {
vi.useRealTimers();
}
});
});

View file

@ -0,0 +1,47 @@
"use client";
import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
import type { UIEvent } from "react";
import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
const SCROLL_THRESHOLD = 0.8;
const SEARCH_REASONS: ReadonlySet<string> = new Set(["input-change", "input-clear", "clear-press"]);
interface UsePaginatedComboboxOptions {
onSearchChange: (query: string) => void;
onLoadMore: () => void;
hasNextPage: boolean;
isFetchingNextPage: boolean;
}
interface PaginatedComboboxHandlers {
handleSearchInput: (next: string, reason: string) => void;
handleScroll: (event: UIEvent<HTMLDivElement>) => void;
}
export function usePaginatedCombobox({
onSearchChange,
onLoadMore,
hasNextPage,
isFetchingNextPage,
}: UsePaginatedComboboxOptions): PaginatedComboboxHandlers {
const debouncedSearch = useDebouncedCallback(onSearchChange, { wait: DEBOUNCE_WAIT_MS });
const handleSearchInput = (next: string, reason: string) => {
if (!SEARCH_REASONS.has(reason)) return;
debouncedSearch(next);
};
const handleScroll = (event: UIEvent<HTMLDivElement>) => {
const target = event.currentTarget;
if (target.scrollHeight === 0) return;
const ratio = (target.scrollTop + target.clientHeight) / target.scrollHeight;
if (ratio >= SCROLL_THRESHOLD && hasNextPage && !isFetchingNextPage) {
onLoadMore();
}
};
return { handleSearchInput, handleScroll };
}

View file

@ -838,20 +838,21 @@ export interface paths {
put?: never;
/**
* Start Shadow Eval
* @description Start a shadow eval: duplicate a sampled slice of a key's live traffic against a second
* arm, judge the two responses blind, and stratify win rates by tier and by the model that
* served the real arm.
* @description Start a shadow eval: duplicate a sampled slice of one or more keys' live traffic against
* a second arm, judge the two responses blind, and stratify win rates by tier, by the model
* that served the real arm, and by key.
*
* A forward job answers whether the key should adopt router_name: it samples the requests
* A forward job answers whether the keys should adopt router_name: it samples the requests
* the router did not serve and duplicates them through it. A reverse job answers whether a
* key already on the router still gains from it: it samples the requests the router did
* serve and duplicates them against baseline_model. A key can hold one active job per
* direction, so both questions can run at once.
*
* Shadow responses are never served to users. The job samples until it has judged
* max_turns turns, reaches the end of its window, or is stopped; sampling changes
* propagate to pods within about 10 seconds. Shadow and judge calls bill to the
* shadowed key but are excluded from request counts and auto-router adoption metrics.
* Shadow responses are never served to users. Each key samples until it has judged
* max_turns turns of its own traffic, the job's window ends, or the job is stopped, so a
* busy key running out of budget does not end sampling for the others; sampling changes
* propagate to pods within about 10 seconds. Shadow and judge calls bill to the shadowed
* key but are excluded from request counts and auto-router adoption metrics.
*/
post: operations["start_shadow_eval_auto_router_shadow_eval_start_post"];
delete?: never;
@ -891,7 +892,9 @@ export interface paths {
put?: never;
/**
* Stop Shadow Eval Job
* @description Stop an active shadow eval job. Attempts are kept; sampling halts within ~10s.
* @description Stop an active shadow eval job, every key it scopes at once. Attempts are kept;
* sampling halts within ~10s. Keys that already stopped on their own budget keep the
* stopped_at they earned, so one write records the outcome for the whole job.
*/
post: operations["stop_shadow_eval_job_auto_router_shadow_eval__job_id__stop_post"];
delete?: never;
@ -32744,19 +32747,35 @@ export interface components {
/** Timeout */
timeout?: number | null;
};
/**
* ShadowEvalJobKeyResponse
* @description One key a job shadows, with its own budget and stop state.
*/
ShadowEvalJobKeyResponse: {
/**
* Api Key Id
* @description The hashed virtual key whose traffic this row scopes
*/
api_key_id: string;
/**
* Max Turns
* @description This key's own sample budget, independent of its siblings'
*/
max_turns: number;
/**
* Stopped At
* @description When this key stopped sampling, whether its own budget ran out, the window closed, or an operator stopped the job. The key reads completed whenever the job does, otherwise stopped once this is set and running until then
*/
stopped_at?: string | null;
};
/**
* ShadowEvalJobResponse
* @description A shadow-eval job. Validates directly from the prisma record (job_id reads the
* row's id); status is derived from stopped_at and ends_at, never stored, so no writer
* anywhere can produce an inconsistent one. Aggregate fields are populated by the
* row's id); status is derived from the keys' stopped_at and ends_at, never stored, so no
* writer anywhere can produce an inconsistent one. Aggregate fields are populated by the
* detail endpoint only and stay None on list responses.
*/
ShadowEvalJobResponse: {
/**
* Api Key Id
* @description The hashed virtual key whose traffic this job evaluates, and only that key's
*/
api_key_id: string;
/** Baseline Model */
baseline_model?: string | null;
/**
@ -32794,13 +32813,16 @@ export interface components {
* @description Verdicts recorded; detail endpoint only
*/
judged_count?: number | null;
/**
* Keys
* @description The keys whose traffic this job evaluates, and only those keys', each with its own budget
*/
keys: components["schemas"]["ShadowEvalJobKeyResponse"][];
/**
* Last Error
* @description Most recent attempt error; detail endpoint only
*/
last_error?: string | null;
/** Max Turns */
max_turns: number;
/** @description Stratified verdicts; detail endpoint only */
results?: components["schemas"]["ShadowEvalResult"] | null;
/** Router Name */
@ -32809,13 +32831,12 @@ export interface components {
shadow_percentage: number;
/**
* Status
* @description A job whose window has passed reads completed even if a later sweep stamped
* stopped_at; stopped means sampling ended before the window did.
* @description A job whose window has passed reads completed even if a sweep stamped its keys
* first; stopped means every key ended sampling before the window did. One key
* exhausting its own budget leaves the job running while any sibling still samples.
* @enum {string}
*/
readonly status: "running" | "completed" | "stopped";
/** Stopped At */
stopped_at?: string | null;
};
/**
* ShadowEvalResult
@ -32824,9 +32845,15 @@ export interface components {
ShadowEvalResult: {
/**
* By Current Model
* @description Sliced by the model that served the real arm: the key's incumbent models in forward mode, and in reverse the models the router itself picked
* @description Sliced by the model that served the real arm: the keys' incumbent models in forward mode, and in reverse the models the router itself picked
*/
by_current_model: components["schemas"]["ShadowEvalSlice"][];
/**
* By Key
* @description One slice per scoped key that has judged verdicts, grouped on the raw key hash. Keys the job scopes but has not yet judged a turn for are absent rather than reported as zero
* @default []
*/
by_key: components["schemas"]["ShadowEvalSlice"][];
/** By Tier */
by_tier: components["schemas"]["ShadowEvalSlice"][];
/** Overall Shadow Win Rate Pct */
@ -32836,8 +32863,8 @@ export interface components {
};
/**
* ShadowEvalSlice
* @description Judge outcomes for one slice of a job's verdicts (a router tier, or one of the
* models that served the real arm).
* @description Judge outcomes for one slice of a job's verdicts (a router tier, one of the models
* that served the real arm, or one of the keys the job is scoped to).
*/
ShadowEvalSlice: {
/** Avg Judge Confidence */
@ -33034,14 +33061,14 @@ export interface components {
};
/**
* StartShadowEvalRequest
* @description Start duplicating a key's traffic for blind comparison against an auto-router.
* @description Start duplicating one or more keys' traffic for blind comparison against an auto-router.
*/
StartShadowEvalRequest: {
/**
* Api Key Id
* @description The hashed virtual key whose traffic will be shadowed. Shadow evaluation runs ONLY on this key's traffic; requests made with any other key are not sampled.
* Api Key Ids
* @description The hashed virtual keys whose traffic will be shadowed. Shadow evaluation runs ONLY on these keys' traffic; requests made with any other key are not sampled. Each key gets its own max_turns budget, so one key exhausting its budget does not end sampling for the others.
*/
api_key_id: string;
api_key_ids: string[];
/**
* Baseline Model
* @description Required when direction is reverse and rejected otherwise: the fixed model the router's own responses are judged against. Must be a plain model rather than another auto-router
@ -33068,7 +33095,7 @@ export interface components {
judge_model: string;
/**
* Max Turns
* @description Sample budget: the job judges at most this many turns, then completes. This is also the spend bound; expected judge cost is roughly max_turns times one judge call
* @description Per-key sample budget: the job judges at most this many turns of EACH scoped key's traffic, so a job over N keys judges at most N times max_turns turns. This is also the spend bound; expected judge cost is roughly that turn ceiling times one judge call
* @default 200
*/
max_turns: number;
@ -37370,7 +37397,7 @@ export interface operations {
list_shadow_eval_jobs_auto_router_shadow_eval_get: {
parameters: {
query?: {
/** @description Filter to jobs shadowing this key */
/** @description Filter to jobs that shadow this key, alone or alongside others */
api_key_id?: string | null;
/** @description Newest jobs to return */
limit?: number;