feat(ui): deployment affinity toggle for the auto-router

ui/litellm-dashboard/src/autorouter_presets.json

@@ -11,7 +11,8 @@
       },
       "classifier_type": "heuristic",
       "escalation_keywords": ["LITELLM ESCALATE"],
-      "session_affinity": false
+      "session_affinity": false,
+      "deployment_affinity": true
     }
   },
   "openai_family": {
@@ -26,7 +27,8 @@
       },
       "classifier_type": "heuristic",
       "escalation_keywords": ["LITELLM ESCALATE"],
-      "session_affinity": false
+      "session_affinity": false,
+      "deployment_affinity": true
     }
   }
 }

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

@@ -8,7 +8,7 @@
   { model_group: "gpt-3.5-turbo", mode: "chat" },
   { model_group: "claude-3-opus", mode: "chat" },
   { model_group: "text-embedding-3-small", mode: "embedding" },
 ] as any[];
 
 const defaultValue: ComplexityRouterConfigValue = {
   tiers: {
@@ -602,3 +602,32 @@
     expect(screen.getByTitle("Deep")).toBeInTheDocument();
   });
 });
+
+describe("ComplexityRouterConfig affinity panel", () => {
+  it("holds both affinity switches with their backend defaults", () => {
+    renderWithProviders(<ComplexityRouterConfig {...baseProps} />);
+    fireEvent.click(screen.getByText("Advanced: Affinity"));
+
+    expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).toBeChecked();
+    expect(screen.getByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked();
+  });
+
+  it("writes deployment_affinity through onChange without touching other keys", () => {
+    const onChange = vi.fn();
+    renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
+    fireEvent.click(screen.getByText("Advanced: Affinity"));
+
+    fireEvent.click(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" }));
+
+    expect(onChange).toHaveBeenCalledWith({ ...defaultValue, deployment_affinity: false });
+  });
+
+  it("renders a stored deployment_affinity=false as off", () => {
+    renderWithProviders(
+      <ComplexityRouterConfig {...baseProps} value={{ ...defaultValue, deployment_affinity: false }} />,
+    );
+    fireEvent.click(screen.getByText("Advanced: Affinity"));
+
+    expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).not.toBeChecked();
+  });
+});

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

@@ -15,6 +15,7 @@
 export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3;
 export const DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS = 200;
 export const DEFAULT_SESSION_AFFINITY = false;
+export const DEFAULT_DEPLOYMENT_AFFINITY = true;
 
 export interface ComplexityTiers {
   SIMPLE: string[];
@@ -56,6 +57,7 @@
   classifier_context_include_assistant_turns?: boolean;
   classifier_fallback?: ClassifierFallback;
   session_affinity?: boolean;
+  deployment_affinity?: boolean;
   adaptive?: boolean;
   adaptive_weights?: AdaptiveRouterWeights;
   tier_distance_penalty?: number;
@@ -132,11 +134,11 @@
   escalationKeywords = [],
   onEscalationKeywordsChange,
   showValidationErrors = false,
 }) => {
   // The deployment's default model is derived from the tiers on submit, mirroring the order
   // add_auto_router_tab uses, so the fallback option is offered exactly when one will exist.
   const hasDefaultModel = Boolean(
     value.tiers.MEDIUM[0] || value.tiers.SIMPLE[0] || value.tiers.COMPLEX[0] || value.tiers.REASONING[0],
   );
 
   // Embedding models can't serve a chat-completion role, so they're excluded here.
@@ -277,14 +279,26 @@
             children: <AdaptiveRoutingConfig value={value} onChange={onChange} />,
           },
           {
-            key: "session-affinity",
+            key: "affinity",
             label: (
               <Text strong style={{ color: "#374151" }}>
-                Advanced: Session Affinity
+                Advanced: Affinity
               </Text>
             ),
             children: (
               <>
+                <div className="flex items-center gap-2 mb-2">
+                  <Switch
+                    checked={value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY}
+                    onChange={(deploymentAffinity) => onChange({ ...value, deployment_affinity: deploymentAffinity })}
+                    aria-label="Pin a session to one deployment per model group"
+                  />
+                  <Text strong>Pin a session to one deployment per model group</Text>
+                </div>
+                <Text type="secondary" style={{ display: "block", fontSize: 12, marginBottom: 12 }}>
+                  Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off
+                  to load-balance every turn.
+                </Text>
                 <div className="flex items-center gap-2 mb-2">
                   <Switch
                     checked={value.session_affinity ?? DEFAULT_SESSION_AFFINITY}
@@ -294,10 +308,8 @@
                   <Text strong>Pin a session to its first model</Text>
                 </div>
                 <Text type="secondary" style={{ display: "block", fontSize: 12 }}>
-                  Off by default: every turn is classified on its own merits and routed to the cheapest adequate tier.
-                  Turn this on to reuse the model chosen on a session&apos;s first turn for every later turn, which
-                  preserves provider prompt caches and avoids cross-model conversation-history errors, at the cost of
-                  keeping the whole session on the first turn&apos;s tier.
+                  Keeps a session on its first turn&apos;s model instead of re-classifying each turn. Also pins the
+                  deployment.
                 </Text>
               </>
             ),

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

@@ -273,7 +273,7 @@ describe("AddAutoRouterTab", () => {
 
     await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router");
     expandDetailedConfiguration();
-    await user.click(screen.getByText("Advanced: Session Affinity"));
+    await user.click(screen.getByText("Advanced: Affinity"));
     expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked();
 
     await user.click(screen.getByRole("button", { name: /add auto router/i }));
@@ -292,7 +292,7 @@ describe("AddAutoRouterTab", () => {
 
     await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router");
     expandDetailedConfiguration();
-    await user.click(screen.getByText("Advanced: Session Affinity"));
+    await user.click(screen.getByText("Advanced: Affinity"));
     await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" }));
 
     await user.click(screen.getByRole("button", { name: /add auto router/i }));
@@ -303,6 +303,46 @@ describe("AddAutoRouterTab", () => {
     });
   });
 
+  it("defaults a new router to deployment affinity on, matching the backend field default", async () => {
+    const user = userEvent.setup();
+    vi.mocked(getMissingTiersError).mockReturnValue(null);
+
+    renderWithProviders(<Harness />);
+
+    await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router");
+    expandDetailedConfiguration();
+    await user.click(screen.getByText("Advanced: Affinity"));
+    expect(
+      await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }),
+    ).toBeChecked();
+
+    await user.click(screen.getByRole("button", { name: /add auto router/i }));
+
+    await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
+    expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
+      deployment_affinity: true,
+    });
+  });
+
+  it("carries deployment affinity turned off through to the create payload", async () => {
+    const user = userEvent.setup();
+    vi.mocked(getMissingTiersError).mockReturnValue(null);
+
+    renderWithProviders(<Harness />);
+
+    await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router");
+    expandDetailedConfiguration();
+    await user.click(screen.getByText("Advanced: Affinity"));
+    await user.click(await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }));
+
+    await user.click(screen.getByRole("button", { name: /add auto router/i }));
+
+    await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
+    expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
+      deployment_affinity: false,
+    });
+  });
+
   // Custom is the escape hatch, not the headline choice, so it's listed after every bundled preset
   // rather than first.
   it("lists Custom Configuration after the bundled presets", () => {

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

@@ -15,6 +15,7 @@
   ComplexityTiers,
   DEFAULT_ADAPTIVE_WEIGHTS,
   DEFAULT_SESSION_AFFINITY,
+  DEFAULT_DEPLOYMENT_AFFINITY,
   DEFAULT_TIER_DISTANCE_PENALTY,
 } from "./ComplexityRouterConfig";
 import { KeywordTierRule } from "./KeywordTierRules";
@@ -89,7 +90,7 @@
 const presets = getAllPresets();
 
 const resolveDefaultModel = (tiers: ComplexityTiers): string | undefined =>
   tiers.MEDIUM[0] || tiers.SIMPLE[0] || tiers.COMPLEX[0] || tiers.REASONING[0];
 
 // A one-line summary of what's configured, shown when the detailed section is collapsed so a
 // caller can see the shape of the config without opening it.
@@ -128,7 +129,7 @@
   userRole,
   userId,
   createScope = "unscoped-ok",
 }) => {
   const requiresTeamScope = createScope === "team-required";
   const [form] = Form.useForm();
   const [modelAccessGroups, setModelAccessGroups] = useState<string[]>([]);
@@ -162,7 +163,7 @@
   useEffect(() => {
     const fetchModelAccessGroups = async () => {
       const response = await modelAvailableCall(accessToken, "", "", false, null, true, true);
       setModelAccessGroups(response["data"].map((model: any) => model["id"]));
     };
     fetchModelAccessGroups();
   }, [accessToken]);
@@ -290,6 +291,7 @@
     classifierContextIncludeAssistantTurns: complexityRouterConfig.classifier_context_include_assistant_turns,
     classifierFallback: complexityRouterConfig.classifier_fallback,
     sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY,
+    deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
     customTechnicalKeywords,
     keywordTierRules,
     semanticMatchingEnabled,
@@ -353,7 +355,7 @@
 
     const defaultModel = resolveDefaultModel(tiers);
 
     form.setFieldsValue({
       custom_llm_provider: "auto_router",
       model: name,
       api_key: "not_required_for_auto_router",

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

@@ -26,6 +26,7 @@
   classifierContextIncludeAssistantTurns: undefined,
   classifierFallback: undefined,
   sessionAffinity: false,
+  deploymentAffinity: true,
   customTechnicalKeywords: [],
   keywordTierRules: [],
   semanticMatchingEnabled: false,
@@ -42,10 +43,11 @@
 describe("buildComplexityRouterConfig", () => {
   it("emits tiers, classifier_type, and escalation_keywords when nothing else is configured", () => {
     const config = buildComplexityRouterConfig(baseParams);
     expect(config).toEqual({
       tiers,
       classifier_type: "heuristic",
       session_affinity: false,
+      deployment_affinity: true,
       escalation_keywords: ["LITELLM ESCALATE"],
     });
   });
@@ -209,7 +211,7 @@
   });
 
   it("omits adaptive fields when adaptive is disabled even if weights linger in state", () => {
     const config = buildComplexityRouterConfig({
       ...baseParams,
       adaptive: false,
       adaptiveWeights: { quality: 0.9, cost: 0.1 },
@@ -243,7 +245,7 @@
   });
 
   it("includes tier_distance_penalty when adaptive is enabled with eligible='all'", () => {
     const config = buildComplexityRouterConfig({
       ...baseParams,
       adaptive: true,
       adaptiveWeights: { quality: 0.6, cost: 0.4 },

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

@@ -30,6 +30,7 @@ export interface BuildComplexityRouterConfigParams {
   classifierContextIncludeAssistantTurns: boolean | undefined;
   classifierFallback: ClassifierFallback | undefined;
   sessionAffinity: boolean;
+  deploymentAffinity: boolean;
   customTechnicalKeywords: string[];
   keywordTierRules: KeywordTierRule[];
   semanticMatchingEnabled: boolean;
@@ -53,6 +54,7 @@ export interface ComplexityRouterConfigPayload {
   classifier_context_include_assistant_turns?: boolean;
   classifier_fallback?: ClassifierFallback;
   session_affinity: boolean;
+  deployment_affinity: boolean;
   custom_technical_keywords?: string[];
   keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[];
   semantic_keyword_matching?: boolean;
@@ -137,6 +139,7 @@ export const buildComplexityRouterConfig = ({
   classifierContextIncludeAssistantTurns,
   classifierFallback,
   sessionAffinity,
+  deploymentAffinity,
   customTechnicalKeywords,
   keywordTierRules,
   semanticMatchingEnabled,
@@ -173,6 +176,7 @@ export const buildComplexityRouterConfig = ({
         classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns,
       }),
     session_affinity: sessionAffinity,
+    deployment_affinity: deploymentAffinity,
     ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }),
     ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }),
     escalation_keywords: cleanedEscalationKeywords,

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

@@ -228,6 +228,31 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => {
   });
 });
 
+describe("buildUpdatedComplexityRouterConfig deployment affinity", () => {
+  it("writes deployment_affinity=false when the toggle is off", () => {
+    const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, deployment_affinity: false });
+    expect(result.deployment_affinity).toBe(false);
+  });
+
+  it("writes deployment_affinity=true when the toggle is on", () => {
+    const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, deployment_affinity: true });
+    expect(result.deployment_affinity).toBe(true);
+  });
+
+  it("re-asserts the backend's on-by-default when the form value is absent, rather than dropping the key", () => {
+    const result = buildUpdatedComplexityRouterConfig({ ...STORED, deployment_affinity: false }, FORM_VALUE);
+    expect(result.deployment_affinity).toBe(true);
+  });
+
+  it("stops a stored deployment_affinity=false from surviving a save that turned the toggle back on", () => {
+    const result = buildUpdatedComplexityRouterConfig(
+      { ...STORED, deployment_affinity: false },
+      { ...FORM_VALUE, deployment_affinity: true },
+    );
+    expect(result.deployment_affinity).toBe(true);
+  });
+});
+
 describe("buildUpdatedComplexityRouterConfig tier labels", () => {
   const RENAMED = { ...STORED, tier_labels: { SIMPLE: "Cheap", REASONING: "Deep" } };
 

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

@@ -48,6 +48,7 @@ const expectedClassifiedTierConfig = {
   embedding_model: "voyage-4-large",
   match_threshold: 0.65,
   session_affinity: false,
+  deployment_affinity: true,
   adaptive: true,
   adaptive_weights: { quality: 0.4, cost: 0.6 },
   adaptive_eligible: "classified_tier",
@@ -68,6 +69,7 @@ const expectedAdaptiveDisabledConfig = {
   embedding_model: "voyage-4-large",
   match_threshold: 0.65,
   session_affinity: false,
+  deployment_affinity: true,
 };
 
 describe("buildUpdatedComplexityRouterConfig", () => {

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

@@ -342,7 +342,7 @@ describe("EditAutoRouterModal session affinity", () => {
     const user = userEvent.setup();
     renderWithStoredConfig(STORED_CONFIG);
 
-    await user.click(await screen.findByText("Advanced: Session Affinity"));
+    await user.click(await screen.findByText("Advanced: Affinity"));
     expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked();
 
     await user.click(screen.getByRole("button", { name: /save changes/i }));
@@ -355,7 +355,7 @@ describe("EditAutoRouterModal session affinity", () => {
     const user = userEvent.setup();
     renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true });
 
-    await user.click(await screen.findByText("Advanced: Session Affinity"));
+    await user.click(await screen.findByText("Advanced: Affinity"));
     expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).toBeChecked();
 
     await user.click(screen.getByRole("button", { name: /save changes/i }));
@@ -368,7 +368,7 @@ describe("EditAutoRouterModal session affinity", () => {
     const user = userEvent.setup();
     renderWithStoredConfig(STORED_CONFIG);
 
-    await user.click(await screen.findByText("Advanced: Session Affinity"));
+    await user.click(await screen.findByText("Advanced: Affinity"));
     await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" }));
 
     await user.click(screen.getByRole("button", { name: /save changes/i }));
@@ -381,7 +381,7 @@ describe("EditAutoRouterModal session affinity", () => {
     const user = userEvent.setup();
     renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true });
 
-    await user.click(await screen.findByText("Advanced: Session Affinity"));
+    await user.click(await screen.findByText("Advanced: Affinity"));
     await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" }));
 
     await user.click(screen.getByRole("button", { name: /save changes/i }));
@@ -391,6 +391,67 @@ describe("EditAutoRouterModal session affinity", () => {
   });
 });
 
+describe("EditAutoRouterModal deployment affinity", () => {
+  beforeEach(() => {
+    modelPatchUpdateCall.mockClear();
+  });
+
+  const renderWithStoredConfig = (complexity_router_config: Record<string, unknown>) =>
+    renderWithProviders(
+      <EditAutoRouterModal
+        isVisible
+        onCancel={vi.fn()}
+        onSuccess={vi.fn()}
+        modelData={{ ...MODEL_DATA, litellm_params: { ...MODEL_DATA.litellm_params, complexity_router_config } }}
+        accessToken="token"
+        userRole="Admin"
+      />,
+    );
+
+  it("shows a stored config with no deployment_affinity key as on, matching the backend default", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig(STORED_CONFIG);
+
+    await user.click(await screen.findByText("Advanced: Affinity"));
+    expect(
+      await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }),
+    ).toBeChecked();
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig().deployment_affinity).toBe(true);
+  });
+
+  it("shows a stored deployment_affinity=false as off and preserves it through an untouched save", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig({ ...STORED_CONFIG, deployment_affinity: false });
+
+    await user.click(await screen.findByText("Advanced: Affinity"));
+    expect(
+      await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }),
+    ).not.toBeChecked();
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig().deployment_affinity).toBe(false);
+  });
+
+  it("persists turning deployment affinity off", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig(STORED_CONFIG);
+
+    await user.click(await screen.findByText("Advanced: Affinity"));
+    await user.click(await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }));
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig().deployment_affinity).toBe(false);
+  });
+});
+
 describe("EditAutoRouterModal custom classifier prompt and fallback", () => {
   beforeEach(() => {
     modelPatchUpdateCall.mockClear();