fix(complexity_router): escalate short turns in long conversations past heuristic_first

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/router_strategy/complexity_router/README.md

@@ -470,6 +470,7 @@ model_list:
       complexity_router_config:
         classifier_type: heuristic_first
         heuristic_first_max_tier: SIMPLE
+        heuristic_first_max_context_tokens: 8000
         classifier_llm_config:
           model: gpt-5-mini
           reasoning_effort: low
@@ -510,6 +511,11 @@ A request short-circuits, meaning it routes on the scorer's own tier with no cla
 two things hold: the scorer landed at or below `heuristic_first_max_tier`, and it produced at least
 one signal. Everything else goes to the classifier, which then decides as it normally would.
 
+Set `heuristic_first_max_context_tokens` to veto that shortcut when the estimated whole conversation
+exceeds the limit. The estimate counts all message text at approximately four characters per token,
+so a short newest nudge in a long agentic session still reaches the classifier. Leave it unset to
+keep the scorer's tier in control at any conversation size
+
 The signal requirement is what keeps this from quietly routing everything to your cheapest model.
 A prompt where no dimension fires scores exactly 0.0, which is below `simple_medium`, so the score
 to tier mapping calls it SIMPLE by default rather than by evidence. Around half of general traffic

litellm/router_strategy/complexity_router/complexity_router.py

@@ -460,6 +460,10 @@ def _message_text(content: object) -> str:
     return content if isinstance(content, str) else ""
 
 
+def _estimated_conversation_tokens(messages: Sequence[Mapping[str, object]] | None) -> int:
+    return sum(len(_message_text(message.get("content"))) // 4 for message in messages or ())
+
+
 def _reminder_block_spans(lowered: str, open_marker: str, close_marker: str) -> Iterator[tuple[int, int]]:
     """Span of each complete reminder block for one marker pair, left to right.
 
@@ -1925,6 +1929,9 @@ class ComplexityRouter(CustomLogger):
         A turn carrying images the classifier would see is never decided cheaply: the scorer reads
         text alone, so its confidence describes a request it has only partly seen, and a trivial
         caption beside a screenshot is exactly the misrouting vision classification exists to stop.
+
+        A configured conversation-size limit also vetoes the cheap decision because a short newest
+        turn can conceal a complex task in the preceding agentic context.
         """
         tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
         scored: Final = ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
@@ -1933,12 +1940,17 @@ class ComplexityRouter(CustomLogger):
             threshold is not None
             and bool(signals)
             and not self._classifier_image_parts(messages)
+            and not self._exceeds_heuristic_first_context(messages)
             and self._active_tier_severity(tier) <= self._active_tier_severity(threshold)
         )
         if decided_cheaply:
             return ClassificationOutcome(tier=tier, score=score, signals=signals, cause="heuristic_first_short_circuit")
         return await self._llm_classifier_outcome(prompt, system_prompt, request_kwargs, messages, scored=scored)
 
+    def _exceeds_heuristic_first_context(self, messages: Sequence[Mapping[str, object]] | None) -> bool:
+        limit: Final = self.config.heuristic_first_max_context_tokens
+        return limit is not None and _estimated_conversation_tokens(messages) > limit
+
     async def _classify_hybrid(
         self,
         prompt: str,
@@ -2691,7 +2703,7 @@ class ComplexityRouter(CustomLogger):
             else ()
         )
 
-        cumulative_tokens: Final = sum(len(_message_text(msg.get("content"))) // 4 for msg in messages or ())
+        cumulative_tokens: Final = _estimated_conversation_tokens(messages)
         trajectory_block: Final = (
             (f"\nConversation so far: ~{cumulative_tokens} tokens across the request",)
             if has_prior_conversation

litellm/router_strategy/complexity_router/config.py

@@ -1065,6 +1065,18 @@ class ComplexityRouterConfig(BaseModel):
             "may not name the highest one, since that would make the LLM classifier unreachable."
         ),
     )
+    heuristic_first_max_context_tokens: int | None = Field(
+        default=None,
+        gt=0,
+        description=(
+            "The estimated size of the whole conversation, counting all message text at approximately four "
+            "characters per token, above which the local scorer may not decide cheaply and the request goes to "
+            "the LLM classifier even when the newest turn scores at or below heuristic_first_max_tier. The "
+            "newest turn in a long agentic session is usually a short nudge such as 'run the tests' or 'why did "
+            "that fail?' whose token-count signal says nothing about the task living in the conversation. None "
+            "keeps the scorer's tier at any conversation size."
+        ),
+    )
     hybrid_boundary_margin: float | None = Field(
         default=None,
         ge=0,
@@ -1745,6 +1757,15 @@ class ComplexityRouterConfig(BaseModel):
             )
         return self
 
+    @model_validator(mode="after")
+    def _validate_heuristic_first_max_context_tokens(self) -> "ComplexityRouterConfig":
+        if self.classifier_type != "heuristic_first" and self.heuristic_first_max_context_tokens is not None:
+            raise ValueError(
+                f"heuristic_first_max_context_tokens is set but classifier_type is {self.classifier_type!r}; "
+                "set classifier_type 'heuristic_first' or remove heuristic_first_max_context_tokens"
+            )
+        return self
+
     @model_validator(mode="after")
     def _validate_hybrid_boundary_margin(self) -> "ComplexityRouterConfig":
         if self.classifier_type != "hybrid":

tests/e2e/router/test_heuristic_first_long_context_e2e.py

@@ -0,0 +1,198 @@
+"""Live e2e repros for heuristic-first classification of short turns with context."""
+
+from __future__ import annotations
+
+from typing import Final
+
+import pytest
+from complexity_router_client import ComplexityRouterClient
+from e2e_config import unique_marker
+from e2e_http import StreamingResponse, require_successful_call
+from lifecycle import ResourceManager
+from models import (
+    ChatAssistantTurn,
+    ChatBody,
+    ChatMessage,
+    ChatToolResultTurn,
+    KeyGenerateBody,
+    LiteLLMParamsBody,
+    ToolCall,
+    ToolCallFunction,
+)
+from proxy_client import ProxyClient
+
+pytestmark = pytest.mark.e2e
+
+ROUTER_BACKENDS: Final = ("gpt-5.5", "claude-haiku-4-5")
+SIMPLE_MODELS: Final = frozenset(("openai/gpt-5.5", "gpt-5.5"))
+
+
+@pytest.fixture(scope="module")
+def heuristic_first_router(proxy: ProxyClient, request: pytest.FixtureRequest) -> str:
+    router_name: Final = f"e2e-heuristic-first-router-{unique_marker()}"
+    model_id: Final = proxy.create_model(
+        router_name,
+        LiteLLMParamsBody(
+            model="auto_router/complexity_router",
+            complexity_router_config={
+                "classifier_type": "heuristic_first",
+                "heuristic_first_max_tier": "MEDIUM",
+                "heuristic_first_max_context_tokens": 8000,
+                "classifier_fallback": "heuristic",
+                "classifier_llm_config": {"model": "gpt-5.5"},
+                "tiers": {
+                    "SIMPLE": "gpt-5.5",
+                    "MEDIUM": "claude-haiku-4-5",
+                    "COMPLEX": "claude-haiku-4-5",
+                    "REASONING": "claude-haiku-4-5",
+                },
+            },
+        ),
+    )
+    request.addfinalizer(lambda: proxy.delete_model(model_id))
+    return router_name
+
+
+@pytest.fixture
+def heuristic_first_key(
+    resources: ResourceManager,
+    client: ComplexityRouterClient,
+    heuristic_first_router: str,
+) -> str:
+    key: Final = client.proxy.generate_key(
+        KeyGenerateBody(
+            models=[heuristic_first_router, *ROUTER_BACKENDS],
+            user_id=f"e2e-heuristic-first-{unique_marker()}",
+        )
+    )
+    resources.defer(lambda: client.proxy.delete_key(key))
+    return key
+
+
+def _agentic_messages(marker: str) -> tuple[ChatMessage | ChatAssistantTurn | ChatToolResultTurn, ...]:
+    system: Final = ChatMessage(
+        role="system",
+        content=(
+            f"You are a coding agent operating on a repository. Preserve the marker {marker}. "
+            "Use tools to inspect files, run tests, and diagnose failures. Keep track of prior "
+            "commands and their outputs before proposing a fix. Never discard relevant logs or "
+            "assume that a failed command was unrelated to the current change."
+        ),
+    )
+    rounds: Final = tuple(
+        turn
+        for round_index in range(5)
+        for turn in (
+            ChatMessage(
+                role="user",
+                content=(
+                    f"Inspect the repository state for debugging round {round_index} using marker {marker}. "
+                    "Run the relevant checks and report every warning, traceback, and changed file."
+                ),
+            ),
+            ChatAssistantTurn(
+                content=None,
+                tool_calls=[
+                    ToolCall(
+                        id=f"{marker}-call-{round_index}",
+                        type="function",
+                        function=ToolCallFunction(
+                            name="run_tests",
+                            arguments=f'{{"round": {round_index}, "marker": "{marker}"}}',
+                        ),
+                    )
+                ],
+            ),
+            ChatToolResultTurn(
+                tool_call_id=f"{marker}-call-{round_index}",
+                content="\n".join(
+                    f"{marker} round={round_index} line={line_index} "
+                    "synthetic test output records a failing assertion, a retry, a provider "
+                    "response, a stack frame, and the captured repository state for diagnosis"
+                    for line_index in range(120)
+                ),
+            ),
+        )
+    )
+    return (system, *rounds, ChatMessage(role="user", content=f"why did that fail? {marker}"))
+
+
+def _send(
+    client: ComplexityRouterClient,
+    key: str,
+    body: ChatBody,
+) -> StreamingResponse:
+    response: Final = client.proxy.transport.send(
+        "/chat/completions",
+        headers=client.proxy.transport.bearer(key),
+        json=body,
+        stream=False,
+    )
+    require_successful_call(response)
+    return response
+
+
+def _assert_classifier_consulted(response: StreamingResponse, context: str) -> None:
+    classifier_cost: Final = response.headers.get("x-litellm-classifier-cost")
+    assert classifier_cost is not None, (
+        f"{context}: classifier header missing; observed headers={response.headers!r}; body={response.body[:300]!r}"
+    )
+    try:
+        parsed_cost: Final = float(classifier_cost)
+    except ValueError as exc:
+        raise AssertionError(
+            f"{context}: classifier header was not parseable as float: {classifier_cost!r}; "
+            f"observed headers={response.headers!r}"
+        ) from exc
+    assert parsed_cost >= 0, f"{context}: classifier cost was negative: {parsed_cost}; headers={response.headers!r}"
+
+
+@pytest.mark.covers("reliability.routing.complexity_heuristic.scores_current_ask_only")
+class TestHeuristicFirstLongContext:
+    def test_short_turn_in_long_agentic_conversation_consults_classifier(
+        self,
+        client: ComplexityRouterClient,
+        heuristic_first_key: str,
+        heuristic_first_router: str,
+    ) -> None:
+        marker: Final = unique_marker()
+        response: Final = _send(
+            client,
+            heuristic_first_key,
+            ChatBody(
+                model=heuristic_first_router,
+                messages=_agentic_messages(marker),
+                max_tokens=16,
+            ),
+        )
+        _assert_classifier_consulted(
+            response,
+            f"long agentic context marker={marker}",
+        )
+
+    def test_short_single_turn_stays_on_heuristic_path(
+        self,
+        client: ComplexityRouterClient,
+        heuristic_first_key: str,
+        heuristic_first_router: str,
+    ) -> None:
+        marker: Final = unique_marker()
+        response: Final = _send(
+            client,
+            heuristic_first_key,
+            ChatBody(
+                model=heuristic_first_router,
+                messages=[ChatMessage(role="user", content=f"why did that fail? {marker}")],
+                max_tokens=16,
+            ),
+        )
+        assert "x-litellm-classifier-cost" not in response.headers, (
+            f"single-turn heuristic path unexpectedly consulted classifier; "
+            f"observed headers={response.headers!r}; body={response.body[:300]!r}"
+        )
+        rows: Final = client.proxy.poll_logs_for_key(heuristic_first_key, min_rows=1)
+        served: Final = tuple(row.model for row in rows if row.model is not None)
+        assert len(served) == 1 and served[0] in SIMPLE_MODELS, (
+            f"single-turn heuristic path should serve SIMPLE backend {sorted(SIMPLE_MODELS)!r}; "
+            f"observed spend-log models={served!r}; headers={response.headers!r}"
+        )

tests/test_litellm/router_strategy/test_complexity_router.py

@@ -49,6 +49,7 @@ from litellm.router_strategy.complexity_router.complexity_router import (
     KeywordOverride,
     _built_in_prompt,
     _ClassifierCircuitBreaker,
+    _estimated_conversation_tokens,
     _is_classifier_timeout,
     _matched_plan_mode_sentinel,
     classification_system_prompt,
@@ -13444,6 +13445,35 @@ class TestHeuristicFirstConfig:
         assert config.uses_llm_classifier is True
         assert ComplexityRouterConfig(tiers=dict(HEURISTIC_FIRST_TIERS)).uses_llm_classifier is False
 
+    def test_context_limit_is_accepted_on_heuristic_first(self):
+        config = ComplexityRouterConfig(
+            tiers=dict(HEURISTIC_FIRST_TIERS),
+            classifier_type="heuristic_first",
+            heuristic_first_max_tier="SIMPLE",
+            heuristic_first_max_context_tokens=8000,
+            classifier_llm_config={"model": "haiku-classifier"},
+        )
+        assert config.heuristic_first_max_context_tokens == 8000
+
+    def test_context_limit_is_rejected_on_llm(self):
+        with pytest.raises(ValidationError, match="heuristic_first_max_context_tokens is set but classifier_type"):
+            ComplexityRouterConfig(
+                tiers=dict(HEURISTIC_FIRST_TIERS),
+                classifier_type="llm",
+                heuristic_first_max_context_tokens=8000,
+                classifier_llm_config={"model": "haiku-classifier"},
+            )
+
+    def test_context_limit_rejects_zero(self):
+        with pytest.raises(ValidationError, match="greater than 0"):
+            ComplexityRouterConfig(
+                tiers=dict(HEURISTIC_FIRST_TIERS),
+                classifier_type="heuristic_first",
+                heuristic_first_max_tier="SIMPLE",
+                heuristic_first_max_context_tokens=0,
+                classifier_llm_config={"model": "haiku-classifier"},
+            )
+
 
 class TestHeuristicFirst:
     """Behavior of the heuristic-first chain: when the classifier call is skipped, and when it is not."""
@@ -13461,6 +13491,58 @@ class TestHeuristicFirst:
         assert outcome.signals
         assert outcome.classifier_cost is None
 
+    @pytest.mark.asyncio
+    async def test_long_context_vetoes_cheap_short_turn(self, mock_router_instance):
+        mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "MEDIUM"}'))
+        router = _heuristic_first_router(
+            mock_router_instance,
+            heuristic_first_max_tier="MEDIUM",
+            heuristic_first_max_context_tokens=10,
+        )
+        messages = [
+            {"role": "system", "content": "x" * 80},
+            {"role": "user", "content": "why did that fail?"},
+        ]
+
+        outcome = await router.aclassify("why did that fail?", messages=messages)
+
+        mock_router_instance.acompletion.assert_awaited_once()
+        assert outcome.cause != "heuristic_first_short_circuit"
+        assert outcome.cause == "llm_classifier"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("context_limit", [None, 100])
+    async def test_short_context_keeps_cheap_short_turn(self, mock_router_instance, context_limit):
+        mock_router_instance.acompletion = AsyncMock()
+        router = _heuristic_first_router(
+            mock_router_instance,
+            heuristic_first_max_tier="MEDIUM",
+            heuristic_first_max_context_tokens=context_limit,
+        )
+        messages = [{"role": "user", "content": "why did that fail?"}]
+
+        outcome = await router.aclassify("why did that fail?", messages=messages)
+
+        mock_router_instance.acompletion.assert_not_called()
+        assert outcome.cause == "heuristic_first_short_circuit"
+
+    @pytest.mark.parametrize(
+        "messages, expected",
+        [
+            (None, 0),
+            (
+                [
+                    {"role": "system", "content": "abcd"},
+                    {"role": "user", "content": [{"type": "text", "text": "efghij"}, {"type": "image_url"}]},
+                    {"role": "assistant", "content": "klmnopqr"},
+                ],
+                4,
+            ),
+        ],
+    )
+    def test_estimated_conversation_tokens_counts_text_parts(self, messages, expected):
+        assert _estimated_conversation_tokens(messages) == expected
+
     @pytest.mark.asyncio
     async def test_no_signal_prompt_escalates_even_though_it_scores_simple(self, mock_router_instance):
         """The core guard. This prompt scores 0.0 and the mapping calls it SIMPLE, which is at the

ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx

@@ -55,6 +55,7 @@ const CLASSIFIER_TIMEOUT_ID = "classifier-timeout-ms";
 const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size";
 const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars";
 const HYBRID_BOUNDARY_MARGIN_ID = "hybrid-boundary-margin";
+const HEURISTIC_FIRST_MAX_CONTEXT_TOKENS_ID = "heuristic-first-max-context-tokens";
 
 const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK =
   "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " +
@@ -268,6 +269,18 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
     onChange({ ...value, heuristic_first_max_tier: tier });
   };
 
+  const handleHeuristicFirstMaxContextTokensChange = (raw: string) => {
+    setDraft({ id: HEURISTIC_FIRST_MAX_CONTEXT_TOKENS_ID, raw });
+    if (raw.trim() === "") {
+      onChange({ ...value, heuristic_first_max_context_tokens: undefined });
+      return;
+    }
+    const parsed: number = Number(raw);
+    if (Number.isFinite(parsed)) {
+      onChange({ ...value, heuristic_first_max_context_tokens: Math.max(1, Math.round(parsed)) });
+    }
+  };
+
   const handleHybridBoundaryMarginChange = (raw: string) => {
     setDraft({ id: HYBRID_BOUNDARY_MARGIN_ID, raw });
     const parsed = Number(raw);
@@ -425,6 +438,21 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
               ))}
             </SelectContent>
           </Select>
+          <Label htmlFor={HEURISTIC_FIRST_MAX_CONTEXT_TOKENS_ID}>Max conversation tokens before classifier</Label>
+          <Input
+            id={HEURISTIC_FIRST_MAX_CONTEXT_TOKENS_ID}
+            type="text"
+            inputMode="numeric"
+            min={1}
+            value={
+              draft?.id === HEURISTIC_FIRST_MAX_CONTEXT_TOKENS_ID
+                ? draft.raw
+                : String(value.heuristic_first_max_context_tokens ?? "")
+            }
+            onChange={(event) => handleHeuristicFirstMaxContextTokensChange(event.target.value)}
+            onBlur={() => setDraft(null)}
+            className="w-full"
+          />
           <p className="text-sm text-muted-foreground">
             A request the scorer places at or below this tier routes there without a classifier call. Anything the
             scorer places higher, and anything it found no signal for at all, goes to the classifier instead

ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx

@@ -388,6 +388,8 @@ export interface ComplexityRouterConfigValue {
   classification_examples?: string;
   /** Highest tier the scorer may decide alone under heuristic_first. Required by that type, rejected by the others. */
   heuristic_first_max_tier?: string;
+  /** Conversation token estimate above which heuristic_first defers to the classifier. */
+  heuristic_first_max_context_tokens?: number;
   /** How near a tier boundary a score may land before hybrid defers to the classifier. Required by that type, rejected by the others. */
   hybrid_boundary_margin?: number;
   classification_mode?: ClassificationMode;

ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx

@@ -401,6 +401,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
     classificationPrompt: complexityRouterConfig.classification_prompt,
     classificationExamples: complexityRouterConfig.classification_examples,
     heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier,
+    heuristicFirstMaxContextTokens: complexityRouterConfig.heuristic_first_max_context_tokens,
     hybridBoundaryMargin: complexityRouterConfig.hybrid_boundary_margin,
     classificationMode: complexityRouterConfig.classification_mode,
     tierLabels: complexityRouterConfig.tier_labels,

ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts

@@ -970,6 +970,7 @@ describe("heuristic_first", () => {
     ...baseParams,
     classifierType: "heuristic_first",
     heuristicFirstMaxTier: "SIMPLE",
+    heuristicFirstMaxContextTokens: 8000,
     classifierLlmConfig: { model: "gpt-4o-mini", timeout_ms: 3000 },
     classifierContextWindowSize: 5,
     classifierContextBudgetChars: 4000,
@@ -980,6 +981,15 @@ describe("heuristic_first", () => {
     const config = buildComplexityRouterConfig(heuristicFirstParams);
     expect(config.classifier_type).toBe("heuristic_first");
     expect(config.heuristic_first_max_tier).toBe("SIMPLE");
+    expect(config.heuristic_first_max_context_tokens).toBe(8000);
+  });
+
+  it("omits heuristic_first_max_context_tokens when empty", () => {
+    const config = buildComplexityRouterConfig({
+      ...heuristicFirstParams,
+      heuristicFirstMaxContextTokens: undefined,
+    });
+    expect(config.heuristic_first_max_context_tokens).toBeUndefined();
   });
 
   it("keeps every classifier key the operator set, since heuristic_first still calls the classifier", () => {

ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts

@@ -141,6 +141,7 @@ export interface StoredComplexityRouterConfig {
   classification_prompt?: unknown;
   classification_examples?: unknown;
   heuristic_first_max_tier?: unknown;
+  heuristic_first_max_context_tokens?: unknown;
   hybrid_boundary_margin?: unknown;
   tier_labels?: unknown;
   classifier_type?: ClassifierType;
@@ -259,6 +260,7 @@ export interface ComplexityRouterConfigPayload {
   classification_prompt?: string;
   classification_examples?: string;
   heuristic_first_max_tier?: string;
+  heuristic_first_max_context_tokens?: number;
   hybrid_boundary_margin?: number;
   classification_mode: ClassificationMode;
   session_affinity: boolean;
@@ -506,6 +508,7 @@ const classifierWireFields = (
     classifierLlmConfig,
     classifierFallback,
     heuristicFirstMaxTier,
+    heuristicFirstMaxContextTokens,
     hybridBoundaryMargin,
     classifierContextWindowSize,
     classifierContextBudgetChars,
@@ -515,6 +518,7 @@ const classifierWireFields = (
     | "classifierLlmConfig"
     | "classifierFallback"
     | "heuristicFirstMaxTier"
+    | "heuristicFirstMaxContextTokens"
     | "hybridBoundaryMargin"
     | "classifierContextWindowSize"
     | "classifierContextBudgetChars"
@@ -532,6 +536,9 @@ const classifierWireFields = (
     ...(supportsFallback && classifierFallback !== undefined && { classifier_fallback: classifierFallback }),
     ...(effectiveType === "heuristic_first" &&
       heuristicFirstMaxTier?.trim() && { heuristic_first_max_tier: heuristicFirstMaxTier }),
+    ...(effectiveType === "heuristic_first" &&
+      heuristicFirstMaxContextTokens !== undefined &&
+      { heuristic_first_max_context_tokens: heuristicFirstMaxContextTokens }),
     ...(effectiveType === "hybrid" &&
       hybridBoundaryMargin !== undefined && { hybrid_boundary_margin: hybridBoundaryMargin }),
     ...(usesLlmClassifier(effectiveType) &&
@@ -567,6 +574,7 @@ export const buildComplexityRouterConfig = ({
   classificationPrompt,
   classificationExamples,
   heuristicFirstMaxTier,
+  heuristicFirstMaxContextTokens,
   hybridBoundaryMargin,
   classificationMode,
   sessionAffinity,
@@ -620,6 +628,7 @@ export const buildComplexityRouterConfig = ({
     classifierLlmConfig,
     classifierFallback,
     heuristicFirstMaxTier,
+    heuristicFirstMaxContextTokens,
     hybridBoundaryMargin,
     classifierContextWindowSize,
     classifierContextBudgetChars,

ui/litellm-dashboard/src/components/add_model/classifier_type_transition.test.ts

@@ -9,6 +9,7 @@ const standard: ComplexityRouterConfigValue = {
   classifier_context_budget_chars: 16000,
   classifier_context_include_assistant_turns: true,
   classifier_fallback: "default_model",
+  heuristic_first_max_context_tokens: 8000,
   tiers: { SIMPLE: ["efficient"], MEDIUM: ["middle"], COMPLEX: [], REASONING: ["capable"] },
 };
 
@@ -22,6 +23,7 @@ describe("transitionClassifierType", () => {
       classifier_context_budget_chars: 16000,
       classifier_context_include_assistant_turns: true,
       classifier_fallback: "default_model",
+      ...(target === "heuristic_first" && { heuristic_first_max_context_tokens: 8000 }),
     };
     expect(result).toMatchObject(expectedSettings);
   });
@@ -30,6 +32,7 @@ describe("transitionClassifierType", () => {
     const result = transitionClassifierType(standard, target);
     expect(result.classifier_llm_config).toEqual({ model: "judge", timeout_ms: 20000 });
     expect(result.classifier_fallback).toBeUndefined();
+    expect(result.heuristic_first_max_context_tokens).toBeUndefined();
     if (target === "capability") {
       expect(result.capability_classifier_config?.base_threshold).toBeNaN();
     } else {

ui/litellm-dashboard/src/components/add_model/classifier_type_transition.ts

@@ -42,6 +42,8 @@ export const transitionClassifierType = (
       classifierType === "heuristic_first"
         ? value.heuristic_first_max_tier ?? DEFAULT_HEURISTIC_FIRST_MAX_TIER
         : undefined,
+    heuristic_first_max_context_tokens:
+      classifierType === "heuristic_first" ? value.heuristic_first_max_context_tokens : undefined,
     hybrid_boundary_margin:
       classifierType === "hybrid" ? value.hybrid_boundary_margin ?? DEFAULT_HYBRID_BOUNDARY_MARGIN : undefined,
     ...nonReasoningTierFields(classifierType, value),

ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts

@@ -670,6 +670,7 @@ describe("managed keys survive an untouched open-and-save", () => {
     tier_labels: { SIMPLE: "Cheap" },
     classifier_type: "heuristic_first",
     heuristic_first_max_tier: "SIMPLE",
+    heuristic_first_max_context_tokens: 8000,
     classifier_llm_config: { model: "gpt-4o-mini", timeout_ms: 3000, reasoning_effort: "low" },
     classifier_context_window_size: 5,
     classifier_context_budget_chars: 4000,

ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx

@@ -158,6 +158,10 @@ export const hydrateComplexityRouterConfig = (
       typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== ""
         ? parsedConfig.heuristic_first_max_tier
         : undefined,
+    heuristic_first_max_context_tokens:
+      typeof parsedConfig.heuristic_first_max_context_tokens === "number"
+        ? parsedConfig.heuristic_first_max_context_tokens
+        : undefined,
     hybrid_boundary_margin:
       typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined,
     classification_mode:
@@ -226,6 +230,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
   "classification_prompt",
   "classification_examples",
   "heuristic_first_max_tier",
+  "heuristic_first_max_context_tokens",
   "hybrid_boundary_margin",
   "classification_mode",
   "session_affinity",
@@ -325,6 +330,7 @@ export const buildUpdatedComplexityRouterConfig = (
     classificationPrompt: value.classification_prompt,
     classificationExamples: value.classification_examples,
     heuristicFirstMaxTier: value.heuristic_first_max_tier,
+    heuristicFirstMaxContextTokens: value.heuristic_first_max_context_tokens,
     hybridBoundaryMargin: value.hybrid_boundary_margin,
     classificationMode: value.classification_mode,
     tierLabels: value.tier_labels,