From e0b2c511445783f059a00a0a08c1d068356a4cc5 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Fri, 18 Sep 2026 23:41:20 +0000 Subject: [PATCH] fix(auto-router): validate JEV usage and clear stale context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../complexity_router/jev_classifier.py | 13 ++++---- .../complexity_router/test_jev_classifier.py | 31 +++++++++++++++++++ ...d_updated_complexity_router_config.test.ts | 22 +++++++++++++ .../edit_auto_router_modal.tsx | 4 +++ 4 files changed, 64 insertions(+), 6 deletions(-) diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index 11591b02461..de23824a5f6 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -55,8 +55,8 @@ class JevChoiceAnswer(BaseModel): class JevUsage(BaseModel): model_config = ConfigDict(frozen=True) - input_tokens: int = 0 - output_tokens: int = 0 + input_tokens: int = Field(default=0, ge=0, strict=True) + output_tokens: int = Field(default=0, ge=0, strict=True) class JevSystemOneResponse(BaseModel): @@ -111,6 +111,11 @@ class HttpJevClassifierClient: request_kwargs: Mapping[str, object] | None, start_time: datetime, ) -> None: + try: + body: Final = TypeAdapter(dict[str, object]).validate_json(response.content) + _ = TypeAdapter(JevUsage | None).validate_python(body.get("usage")) + except ValidationError: + return end_time: Final = datetime.now(timezone.utc) parent: Final = request_kwargs or MappingProxyType({}) parent_metadata: Final = { @@ -144,10 +149,6 @@ class HttpJevClassifierClient: optional_params={}, litellm_params=params, ) - try: - body: Final = TypeAdapter(dict[str, object]).validate_json(response.content) - except ValidationError: - return normalized: Final = TypeSafePassthroughLoggingHandler.typesafe_passthrough_handler( httpx_response=response, response_body=body, diff --git a/tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py b/tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py index d51690d8818..dae037ff47c 100644 --- a/tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/test_litellm/router_strategy/complexity_router/test_jev_classifier.py @@ -71,6 +71,37 @@ async def test_jev_http_errors_do_not_dispatch_successful_usage( assert recorder.calls == () +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["input_tokens", "output_tokens"]) +@pytest.mark.parametrize("tokens", [-1, True, 1.5, "3"]) +async def test_jev_invalid_usage_never_reaches_spend_callbacks( + monkeypatch: pytest.MonkeyPatch, field: str, tokens: object +) -> None: + recorder: Final = _UsageRecorder() + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + handler: Final = create_autospec(AsyncHTTPHandler, instance=True) + handler.post.return_value = httpx.Response( + 200, + request=httpx.Request("POST", "https://typesafe.test/v1/systemone"), + json={ + "model": "jev-accounting", + "usage": {"input_tokens": 3, "output_tokens": 2, field: tokens}, + "answers": {"tier": _answer().model_dump()}, + }, + ) + provider: Final = HttpJevClassifierClient("test", "https://typesafe.test", handler) + request: Final = build_jev_request( + "choose a tier", None, "jev-accounting", DEFAULT_JEV_INSTRUCTIONS, {"SIMPLE": "cheap"} + ) + + with pytest.raises(ValueError, match=field): + await provider.evaluate(request, timeout_s=3) + await GLOBAL_LOGGING_WORKER.flush() + + handler.post.assert_awaited_once() + assert recorder.calls == () + + @pytest.mark.asyncio @pytest.mark.parametrize("answer", ["SIMPLE", "UNAVAILABLE", "malformed"]) @pytest.mark.parametrize("private", [False, True]) diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 6a522b9ad4c..e5e2c61933c 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -291,6 +291,28 @@ describe("capability classifier configuration", () => { }); describe("buildUpdatedComplexityRouterConfig classifier context window", () => { + it.each(["llm", "jev"] as const)( + "drops the stored %s per-turn bound when switching to heuristic", + (classifier_type) => { + const stored = { ...STORED_LLM, classifier_type }; + const saved = buildUpdatedComplexityRouterConfig(stored, { + ...hydrateComplexityRouterConfig(stored, undefined), + classifier_type: "heuristic", + }); + + expect(saved).not.toHaveProperty("classifier_context_per_turn_chars"); + }, + ); + + it("does not resurrect an explicitly cleared per-turn bound", () => { + const saved = buildUpdatedComplexityRouterConfig(STORED_LLM, { + ...hydrateComplexityRouterConfig(STORED_LLM, undefined), + classifier_context_per_turn_chars: undefined, + }); + + expect(saved).not.toHaveProperty("classifier_context_per_turn_chars"); + }); + it.each(["llm", "jev"] as const)("saves the form's per-turn bound over the stored %s bound", (classifier_type) => { const formValue = { ...hydrateComplexityRouterConfig({ ...STORED_LLM, classifier_type }, undefined), diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index 56a851fba8c..10fa6fcb6be 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -1,4 +1,5 @@ import AutoRouterClassifierTabs from "../add_model/AutoRouterClassifierTabs"; +import { usesClassifierContext } from "../add_model/classifier_types"; import { defaultJevClassifierConfig, jevClassifierConfigSchema } from "../add_model/jev_classifier_config"; import type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; export type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; @@ -317,6 +318,9 @@ export const buildUpdatedComplexityRouterConfig = ( keywordMatching?: KeywordMatchingState, ): Record => { const isManaged = (key: string): boolean => { + if (key === "classifier_context_per_turn_chars") { + return !usesClassifierContext(effectiveClassifierType(value)) || Object.prototype.hasOwnProperty.call(value, key); + } if (MANAGED_COMPLEXITY_ROUTER_KEYS.has(key)) return true; if (key === "escalation_keywords" && isForecastClassifier(effectiveClassifierType(value))) return true; if (keywordMatching !== undefined && KEYWORD_MATCHING_KEYS.has(key)) return true;