From 878a333697fcaf701be5011f733d782c1ffc56fa Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 17 Mar 2026 15:20:52 +0000 Subject: [PATCH] Gate agent tracing enforcement behind enterprise/premium license The tracing enforcement switches (require x-litellm-trace-id on calls TO/BY agent) in the Add Agent form are now only available to premium users. Non-premium users see a yellow notice directing them to upgrade. This follows the existing premium gating pattern used throughout the dashboard (key_edit_view, create_key_button, PremiumLoggingSettings, etc.). Co-authored-by: Krish Dholakia --- .../components/agents/add_agent_form.test.tsx | 136 ++++++++++++++++++ .../src/components/agents/add_agent_form.tsx | 80 ++++++----- 2 files changed, 182 insertions(+), 34 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/agents/add_agent_form.test.tsx diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.test.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.test.tsx new file mode 100644 index 00000000000..49159e0d78b --- /dev/null +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.test.tsx @@ -0,0 +1,136 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { screen, waitFor } from "@testing-library/react"; +import { renderWithProviders } from "../../../tests/test-utils"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(), +})); + +vi.mock("../networking", () => ({ + getAgentCreateMetadata: vi.fn().mockResolvedValue([]), + getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + keyListCall: vi.fn().mockResolvedValue({ keys: [] }), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), + createAgentCall: vi.fn().mockResolvedValue({ agent_id: "a1", agent_name: "test" }), + keyCreateForAgentCall: vi.fn().mockResolvedValue({ key: "sk-123" }), + keyUpdateCall: vi.fn().mockResolvedValue({}), +})); + +vi.mock("../mcp_server_management/MCPServerSelector", () => ({ + default: () =>
, +})); + +vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ + default: () =>
, +})); + +vi.mock("../guardrails/GuardrailSelector", () => ({ + default: () =>
, +})); + +vi.mock("../shared/CreatedKeyDisplay", () => ({ + default: () =>
, +})); + +import AddAgentForm from "./add_agent_form"; + +const baseAuthorized = { + token: "123", + accessToken: "123", + userId: "user-1", + userEmail: "user@example.com", + userRole: "Admin", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, +}; + +const defaultProps = { + visible: true, + onClose: vi.fn(), + accessToken: "test-token", + onSuccess: vi.fn(), + teams: [], +}; + +const navigateToGovernanceStep = async () => { + const nextButtons = screen.getAllByRole("button", { name: /next/i }); + const nextButton = nextButtons[nextButtons.length - 1]; + + // Step 0 -> 1: need agent_name to be filled + // The form validation may block, so we fill agent name first + const agentNameInput = screen.getByLabelText(/agent name/i); + const { fireEvent } = await import("@testing-library/react"); + const { act } = await import("react"); + await act(async () => { + fireEvent.change(agentNameInput, { target: { value: "test-agent" } }); + }); + + // Click Next to go to step 1 (Entitlements) + await act(async () => { + fireEvent.click(nextButton); + }); + + // Click Next to go to step 2 (Governance) + await waitFor(() => { + expect(screen.getByText("Entitlements")).toBeInTheDocument(); + }); + + const nextButtons2 = screen.getAllByRole("button", { name: /next/i }); + await act(async () => { + fireEvent.click(nextButtons2[nextButtons2.length - 1]); + }); + + await waitFor(() => { + expect(screen.getByText("Tracing")).toBeInTheDocument(); + }); +}; + +describe("AddAgentForm tracing enforcement premium gate", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should show enterprise upgrade notice when user is not premium", async () => { + vi.mocked(useAuthorized).mockReturnValue(baseAuthorized); + + renderWithProviders(); + await navigateToGovernanceStep(); + + expect( + screen.getByText(/enforcing trace-id requirements on agents is a litellm enterprise feature/i) + ).toBeInTheDocument(); + + expect( + screen.queryByText(/require x-litellm-trace-id on calls to this agent/i) + ).not.toBeInTheDocument(); + + expect( + screen.queryByText(/require x-litellm-trace-id on calls by this agent/i) + ).not.toBeInTheDocument(); + }); + + it("should show tracing switches when user is premium", async () => { + vi.mocked(useAuthorized).mockReturnValue({ + ...baseAuthorized, + premiumUser: true, + }); + + renderWithProviders(); + await navigateToGovernanceStep(); + + expect( + screen.queryByText(/enforcing trace-id requirements on agents is a litellm enterprise feature/i) + ).not.toBeInTheDocument(); + + expect( + screen.getByText(/require x-litellm-trace-id on calls to this agent/i) + ).toBeInTheDocument(); + + expect( + screen.getByText(/require x-litellm-trace-id on calls by this agent/i) + ).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx index c5518596b81..1fcd49ab6ff 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -43,7 +43,7 @@ const AddAgentForm: React.FC = ({ onSuccess, teams, }) => { - const { userId, userRole } = useAuthorized(); + const { userId, userRole, premiumUser } = useAuthorized(); const [form] = Form.useForm(); const [currentStep, setCurrentStep] = useState(0); const [isSubmitting, setIsSubmitting] = useState(false); @@ -437,43 +437,55 @@ const AddAgentForm: React.FC = ({

Tracing

-
-
-
- - Require x-litellm-trace-id on calls TO this agent - -

- Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent). -

-
- + {!premiumUser ? ( +
+

+ Enforcing trace-id requirements on agents is a LiteLLM Enterprise feature. Get a trial key{" "} + + here + + . +

+ ) : ( +
+
+
+ + Require x-litellm-trace-id on calls TO this agent + +

+ Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent). +

+
+ +
-
-
- - Require x-litellm-trace-id on calls BY this agent - -

- Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking. -

+
+
+ + Require x-litellm-trace-id on calls BY this agent + +

+ Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking. +

+
+ { + setRequireTraceIdOutbound(checked); + if (!checked) { + setMaxIterations(null); + setMaxBudgetPerSession(null); + } + }} + />
- { - setRequireTraceIdOutbound(checked); - if (!checked) { - setMaxIterations(null); - setMaxBudgetPerSession(null); - } - }} - />
-
+ )}