mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
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 <krrishdholakia@gmail.com>
This commit is contained in:
Cursor Agent
parent
245a3d2b26
commit
878a333697
2 changed files with 182 additions and 34 deletions
|
|
@ -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: () => <div data-testid="mcp-server-selector" />,
|
||||
}));
|
||||
|
||||
vi.mock("../mcp_server_management/MCPToolPermissions", () => ({
|
||||
default: () => <div data-testid="mcp-tool-permissions" />,
|
||||
}));
|
||||
|
||||
vi.mock("../guardrails/GuardrailSelector", () => ({
|
||||
default: () => <div data-testid="guardrail-selector" />,
|
||||
}));
|
||||
|
||||
vi.mock("../shared/CreatedKeyDisplay", () => ({
|
||||
default: () => <div data-testid="created-key-display" />,
|
||||
}));
|
||||
|
||||
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(<AddAgentForm {...defaultProps} />);
|
||||
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(<AddAgentForm {...defaultProps} />);
|
||||
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();
|
||||
});
|
||||
});
|
||||
|
|
@ -43,7 +43,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
|
|||
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<AddAgentFormProps> = ({
|
|||
<div className="space-y-6">
|
||||
<div>
|
||||
<h4 className="text-sm font-medium text-gray-700 mb-3">Tracing</h4>
|
||||
<div className="space-y-4">
|
||||
<div className="flex items-center justify-between">
|
||||
<div>
|
||||
<span className="text-sm font-medium text-gray-700">
|
||||
Require x-litellm-trace-id on calls TO this agent
|
||||
</span>
|
||||
<p className="text-xs text-gray-500 mt-1">
|
||||
Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent).
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
checked={requireTraceIdInbound}
|
||||
onChange={setRequireTraceIdInbound}
|
||||
/>
|
||||
{!premiumUser ? (
|
||||
<div className="p-3 bg-yellow-50 border border-yellow-200 rounded-lg">
|
||||
<p className="text-sm text-yellow-800">
|
||||
Enforcing trace-id requirements on agents is a LiteLLM Enterprise feature. Get a trial key{" "}
|
||||
<a href="https://www.litellm.ai/#pricing" target="_blank" rel="noopener noreferrer" className="underline">
|
||||
here
|
||||
</a>
|
||||
.
|
||||
</p>
|
||||
</div>
|
||||
) : (
|
||||
<div className="space-y-4">
|
||||
<div className="flex items-center justify-between">
|
||||
<div>
|
||||
<span className="text-sm font-medium text-gray-700">
|
||||
Require x-litellm-trace-id on calls TO this agent
|
||||
</span>
|
||||
<p className="text-xs text-gray-500 mt-1">
|
||||
Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent).
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
checked={requireTraceIdInbound}
|
||||
onChange={setRequireTraceIdInbound}
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center justify-between">
|
||||
<div>
|
||||
<span className="text-sm font-medium text-gray-700">
|
||||
Require x-litellm-trace-id on calls BY this agent
|
||||
</span>
|
||||
<p className="text-xs text-gray-500 mt-1">
|
||||
Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking.
|
||||
</p>
|
||||
<div className="flex items-center justify-between">
|
||||
<div>
|
||||
<span className="text-sm font-medium text-gray-700">
|
||||
Require x-litellm-trace-id on calls BY this agent
|
||||
</span>
|
||||
<p className="text-xs text-gray-500 mt-1">
|
||||
Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking.
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
checked={requireTraceIdOutbound}
|
||||
onChange={(checked) => {
|
||||
setRequireTraceIdOutbound(checked);
|
||||
if (!checked) {
|
||||
setMaxIterations(null);
|
||||
setMaxBudgetPerSession(null);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
<Switch
|
||||
checked={requireTraceIdOutbound}
|
||||
onChange={(checked) => {
|
||||
setRequireTraceIdOutbound(checked);
|
||||
if (!checked) {
|
||||
setMaxIterations(null);
|
||||
setMaxBudgetPerSession(null);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<Divider className="my-0" />
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue