From 0f1e09b5558160ad18db40387d63ebb90a1daf4a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:43:53 -0700 Subject: [PATCH 1/9] fix(cli): keep the refresh token in the OS keychain, not in token.json `lite login --pkce` mints a refresh token that buys a fresh key from the proxy on demand, so it is the credential just as much as the key is. Moving the key into the keychain left it behind in ~/.litellm/token.json, where any process running as the user can read it and renew the login for itself. It now travels with the key: `save_cli_token` writes both into the keychain entry, the token file keeps only metadata, and `lite logout` takes it out of the file whether or not the keychain answers. Upgrading finds one sign-in split across the two stores, the key already in the keychain and the refresh token still on disk. That case rejoins the two halves into a single entry before scrubbing the file, so the write never replaces a live key with nothing, and a machine that refuses the scrub keeps what it has rather than having the key rolled back out from under it. --- litellm/litellm_core_utils/cli_token_utils.py | 83 ++++++++--- litellm/proxy/client/README.md | 6 +- .../test_cli_token_utils.py | 130 +++++++++++++++++- 3 files changed, 194 insertions(+), 25 deletions(-) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index b45513c5ea3..ee506a69ef9 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -4,10 +4,10 @@ CLI Token Utilities SDK-level utilities for reading the credential minted by `lite login`. Non-secret metadata lives in ~/.litellm/token.json. The secret material (the -bearer key, plus a JWT when one is issued) lives in the OS keychain when the -machine has one, and in that same 0600 file otherwise. This module hides the -split from callers, and migrates a legacy plaintext file into the keychain the -first time it reads one. +bearer key, the refresh token that renews it, and a JWT when one is issued) +lives in the OS keychain when the machine has one, and in that same 0600 file +otherwise. This module hides the split from callers, and migrates a plaintext +file into the keychain the first time it reads one. This module has no dependencies on proxy code and can be safely imported at the SDK level. """ @@ -111,13 +111,20 @@ class CliTokenSecret(BaseModel): secret minted for one server is never handed to another, even if the metadata file is edited underneath us. `timestamp` is the sign-in this secret came from, which is what decides it against a secret still on disk. + + Every field a thief could sign in with belongs here, which is why the + refresh token is one of them: it buys a fresh key from the proxy on demand, + so leaving it on disk would leave the login readable there. `key` is + optional because the file can hold a refresh token without one, and moving + that into the keychain must not invent a key to go with it. """ model_config = ConfigDict(frozen=True) base_url: str - key: str + key: str | None = None jwt_token: str = "" + refresh_token: str | None = None timestamp: float = 0.0 @@ -155,7 +162,7 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN staged: Final = _stage_token_file(_without_secret(stamped)) if isinstance(staged, CredentialNotSaved): return staged - outcome: Final = SecretStored() if stamped.key is None else vault.write(_encode_secret(stamped, stamped.key)) + outcome: Final = vault.write(_encode_secret(stamped)) if _holds_a_secret(stamped) else SecretStored() if isinstance(outcome, SecretStored): return outcome if _commit_token_file(staged) else CredentialNotRecorded() discard_staged_json(staged) @@ -390,8 +397,9 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - secret is usually left on disk by a keychain that would not take it, which makes the file the fresher of the two. It is the older one when a login the keychain did take could not replace the file afterwards, and serving that one would put a superseded credential back in use. Equal - stamps are one login sitting in both stores, left by a migration whose scrub was refused, so - that branch retries the migration rather than trading one credential for another. + stamps are one login sitting in both stores, left by a migration whose scrub was refused or by + an upgrade that took the key into the keychain and left the refresh token behind, so that branch + rejoins the halves and retries the migration rather than trading one credential for another. A scrub the file refuses leaves that superseded secret where it lies, which is the state the login already named when it could not replace the file, and which `lite logout` reports rather @@ -400,21 +408,46 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - superseded one back out. """ secret: Final = _decode_secret(blob, record.base_url) - if secret is None or (record.key is not None and secret.timestamp <= record.timestamp): - return _migrate_file_secret(record, vault) + if secret is None or (_holds_a_secret(record) and secret.timestamp <= record.timestamp): + return _migrate_file_secret(_rejoined(record, secret), vault, replacing=secret) _scrub_file_secret(record) return record.model_copy( update=MappingProxyType( { "key": secret.key, "jwt_token": secret.jwt_token, + "refresh_token": secret.refresh_token, "timestamp": max(secret.timestamp, record.timestamp), } ) ) -def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord | None: +def _rejoined(record: CliTokenRecord, secret: CliTokenSecret | None) -> CliTokenRecord: + """Put one sign-in's secret material back together when each store holds part of it. + + Upgrading from the release that kept only the key in the keychain leaves the refresh token + behind in the file, so a single login sits across both stores. Filling in whatever the file is + missing before the migration writes its entry is what stops that write from replacing a live key + with nothing. Only a matching stamp is one login. Two stamps are two logins, and pairing one's + key with the other's refresh token would build a credential neither store ever held. + """ + if secret is None or secret.timestamp != record.timestamp: + return record + return record.model_copy( + update=MappingProxyType( + { + "key": record.key if record.key is not None else secret.key, + "jwt_token": record.jwt_token or secret.jwt_token, + "refresh_token": record.refresh_token if record.refresh_token is not None else secret.refresh_token, + } + ) + ) + + +def _migrate_file_secret( + record: CliTokenRecord, vault: SecretVault, *, replacing: CliTokenSecret | None = None +) -> CliTokenRecord | None: """Move a file-held secret into the vault, but only once the file's copy can be taken away. The scrubbed file is staged first so a directory that will not accept it stops the migration @@ -426,23 +459,28 @@ def _migrate_file_secret(record: CliTokenRecord, vault: SecretVault) -> CliToken asked to take the new entry back, so the migration finishes on a directory that would only ever have refused it. Rolling back is the last resort, and a rollback the keychain also refuses leaves the secret in both stores until the next read, which retries this same migration. + + Only an entry this migration put there is taken back. `replacing` names one that was already in + the keychain, whose material the new entry carries forward, so erasing it would take away the + half the file never had, and a machine that refuses the scrub is exactly the one with nowhere + else to keep it. The next read finds the same two halves and tries the move again. """ - if record.key is None: + if not _holds_a_secret(record): return None staged: Final = _stage_scrubbed_file(record) if staged is None: return record - if not isinstance(vault.write(_encode_secret(record, record.key)), SecretStored): + if not isinstance(vault.write(_encode_secret(record)), SecretStored): discard_staged_json(staged) return record - if not _commit_token_file(staged) and not _overwrite_file_secret(record): + if not _commit_token_file(staged) and not _overwrite_file_secret(record) and replacing is None: vault.erase() return record def _scrub_file_secret(record: CliTokenRecord) -> bool: """Leave no secret material in the token file once the vault holds it""" - if record.key is None and not record.jwt_token: + if not _holds_a_secret(record): return True staged: Final = _stage_scrubbed_file(record) if staged is not None and _commit_token_file(staged): @@ -487,13 +525,22 @@ def _commit_token_file(staged: str) -> bool: return True +def _holds_a_secret(record: CliTokenRecord) -> bool: + """Whether the record carries anything that would sign someone in as this user""" + return record.key is not None or bool(record.jwt_token) or record.refresh_token is not None + + def _without_secret(record: CliTokenRecord) -> CliTokenRecord: - return record.model_copy(update=MappingProxyType({"key": None, "jwt_token": ""})) + return record.model_copy(update=MappingProxyType({"key": None, "jwt_token": "", "refresh_token": None})) -def _encode_secret(record: CliTokenRecord, key: str) -> str: +def _encode_secret(record: CliTokenRecord) -> str: return CliTokenSecret( - base_url=record.base_url, key=key, jwt_token=record.jwt_token, timestamp=record.timestamp + base_url=record.base_url, + key=record.key, + jwt_token=record.jwt_token, + refresh_token=record.refresh_token, + timestamp=record.timestamp, ).model_dump_json() diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 3a39bec5ee6..1fff68677cc 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -331,7 +331,7 @@ sequenceDiagram CLI->>Proxy: Poll /sso/cli/poll/login_id with poll_secret header Proxy->>CLI: Return {"status": "ready", "key": "jwt"} - CLI->>CLI: Save key to the OS keychain (metadata to ~/.litellm/token.json) + CLI->>CLI: Save the secret to the OS keychain (metadata to ~/.litellm/token.json) ``` ### Authentication Commands @@ -365,7 +365,7 @@ The CLI provides these authentication commands: ### Token Storage -The key itself goes into the OS keychain (macOS Keychain, Windows Credential Manager, or the Linux Secret Service) under service `litellm-cli`, account `credential`. Only the non-secret session metadata is written to `~/.litellm/token.json`, in a `0700` directory with `0600` file permissions: +The key itself, together with the refresh token that renews a `--pkce` credential, goes into the OS keychain (macOS Keychain, Windows Credential Manager, or the Linux Secret Service) under service `litellm-cli`, account `credential`. Only the non-secret session metadata is written to `~/.litellm/token.json`, in a `0700` directory with `0600` file permissions: ```json { @@ -378,7 +378,7 @@ The key itself goes into the OS keychain (macOS Keychain, Windows Credential Man } ``` -Keychain storage needs the `keyring` package, which ships with `pip install 'litellm[cli]'`. Headless boxes and CI runners usually have no keychain either. In all of those cases the key stays in the same `0600` file alongside the metadata, exactly as it did before, and `lite login` names which one applies: the package is missing, the machine has no keychain, or you set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. +Keychain storage needs the `keyring` package, which ships with `pip install 'litellm[cli]'`. Headless boxes and CI runners usually have no keychain either. In all of those cases the key and the refresh token stay in the same `0600` file alongside the metadata, exactly as they did before, and `lite login` names which one applies: the package is missing, the machine has no keychain, or you set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. That includes a refresh token left behind by the release that moved only the key. `lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job. diff --git a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py index b8ac1987453..7e7eee5373f 100644 --- a/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_cli_token_utils.py @@ -80,8 +80,28 @@ def _write_metadata_only_file(home): return path -def _blob(base_url=SERVER, key="sk-vault", jwt_token="", timestamp=0.0): - return json.dumps({"base_url": base_url, "key": key, "jwt_token": jwt_token, "timestamp": timestamp}) +def _blob(base_url=SERVER, key="sk-vault", jwt_token="", timestamp=0.0, refresh_token=None): + return json.dumps( + { + "base_url": base_url, + "key": key, + "jwt_token": jwt_token, + "refresh_token": refresh_token, + "timestamp": timestamp, + } + ) + + +def _write_key_only_keychain_file(home, *, refresh_token="rt-live", timestamp=2000.0): + """What the release that kept only the key in the keychain left on disk: metadata, plus the + refresh token in the clear.""" + path = _token_file(home) + path.parent.mkdir(exist_ok=True) + path.write_text( + json.dumps({"base_url": SERVER, "user_id": "u-1", "refresh_token": refresh_token, "timestamp": timestamp}) + ) + path.chmod(0o600) + return path _REAL_MKSTEMP = tempfile.mkstemp @@ -161,6 +181,56 @@ class TestLoadCliToken: assert (record.key, record.jwt_token) == ("sk-a", "jwt-a") + def test_the_refresh_token_round_trips_through_the_vault(self, isolated_home, secret_vault_factory): + """A refresh token mints a fresh key from the proxy on demand, so it is the credential just + as much as the key is, and it has to come back out of the keychain to be usable.""" + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(blob=_blob(key="sk-a", refresh_token="rt-a")) + + record = load_cli_token(vault=vault) + + assert (record.key, record.refresh_token) == ("sk-a", "rt-a") + + def test_a_plaintext_refresh_token_is_moved_off_disk(self, isolated_home, secret_vault_factory): + path = _write_legacy_file(isolated_home, refresh_token="rt-legacy") + vault = secret_vault_factory() + + record = load_cli_token(vault=vault) + + assert record.refresh_token == "rt-legacy" + assert "rt-legacy" not in path.read_text() + assert json.loads(vault.blob)["refresh_token"] == "rt-legacy" + + def test_an_upgrade_that_left_the_refresh_token_on_disk_rejoins_it_with_the_key( + self, isolated_home, secret_vault_factory + ): + """The release before this one took the key into the keychain and left the refresh token + behind, so upgrading finds one sign-in split across both stores. The read has to end with + the whole credential in the keychain, not with whichever half it happened to prefer.""" + path = _write_key_only_keychain_file(isolated_home) + vault = secret_vault_factory(blob=_blob(key="sk-live", timestamp=2000.0)) + + record = load_cli_token(vault=vault) + + assert (record.key, record.refresh_token) == ("sk-live", "rt-live") + assert "rt-live" not in path.read_text() + assert json.loads(vault.blob)["key"] == "sk-live" + assert json.loads(vault.blob)["refresh_token"] == "rt-live" + + def test_a_superseded_refresh_token_on_disk_never_outlives_the_keychain( + self, isolated_home, secret_vault_factory + ): + """Two stores, two sign-ins, and the newer one is in the keychain. Handing back its key with + the older one's refresh token would build a credential neither store ever held, and would + renew the login the user already replaced.""" + path = _write_legacy_file(isolated_home, key="sk-old", refresh_token="rt-old", timestamp=1000.0) + vault = secret_vault_factory(blob=_blob(key="sk-new", refresh_token="rt-new", timestamp=2000.0)) + + record = load_cli_token(vault=vault) + + assert (record.key, record.refresh_token) == ("sk-new", "rt-new") + assert "rt-old" not in path.read_text() + def test_legacy_plaintext_file_still_authenticates_and_is_migrated(self, isolated_home, secret_vault_factory): """A token.json written by an older `lite` keeps working, and reading it moves the secret into the keychain and scrubs it from disk.""" @@ -361,6 +431,36 @@ class TestSaveCliToken: assert json.loads(vault.blob)["key"] == "sk-new" assert load_cli_token(vault=vault).key == "sk-new" + def test_the_refresh_token_goes_to_the_keychain_and_never_to_the_file( + self, isolated_home, secret_vault_factory + ): + vault = secret_vault_factory() + + stored = save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-new", refresh_token="rt-new", timestamp=time.time()), + vault=vault, + ) + + assert stored == SecretStored() + assert "rt-new" not in _token_file(isolated_home).read_text() + assert json.loads(vault.blob)["refresh_token"] == "rt-new" + assert load_cli_token(vault=vault).refresh_token == "rt-new" + + def test_the_refresh_token_falls_back_to_the_owner_only_file_with_the_key( + self, isolated_home, secret_vault_factory + ): + """A machine with no keychain keeps the whole credential in the 0600 file, refresh token + included, because a renewal that cannot be stored logs the user out on the next command.""" + vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) + + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-new", refresh_token="rt-new", timestamp=time.time()), + vault=vault, + ) + + assert json.loads(_token_file(isolated_home).read_text())["refresh_token"] == "rt-new" + assert load_cli_token(vault=vault).refresh_token == "rt-new" + def test_falls_back_to_the_owner_only_file_when_there_is_no_keychain(self, isolated_home, secret_vault_factory): stored = save_cli_token( CliTokenRecord(base_url=SERVER, key="sk-new", timestamp=time.time()), @@ -628,6 +728,26 @@ class TestScrubFailure: assert json.loads(path.read_text()).get("key") is None + @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") + def test_a_rejoin_the_file_refuses_never_takes_the_key_with_it( + self, isolated_home, secret_vault_factory, monkeypatch + ): + """Rolling the rejoined entry back would erase a key that was safely in the keychain before + this read began, and the file it would fall back to is the one that has just refused to be + rewritten. The duplicate refresh token stays until a later read can finish the move.""" + path = _write_key_only_keychain_file(isolated_home) + vault = secret_vault_factory(blob=_blob(key="sk-live", timestamp=2000.0)) + monkeypatch.setattr("litellm.litellm_core_utils.private_json.os.replace", _refuse_replace) + path.chmod(0o400) + + record = load_cli_token(vault=vault) + + assert (record.key, record.refresh_token) == ("sk-live", "rt-live") + assert json.loads(vault.blob)["key"] == "sk-live" + assert json.loads(vault.blob)["refresh_token"] == "rt-live" + assert vault.erases == 0 + + class TestClearCliToken: def test_removes_the_credential_from_both_stores(self, isolated_home, secret_vault_factory): vault = secret_vault_factory() @@ -717,12 +837,14 @@ class TestClearCliToken: ): """Keeping a record of the unreachable keychain must never mean keeping the cleartext copy the user just asked to be rid of.""" - _write_legacy_file(isolated_home) + _write_legacy_file(isolated_home, refresh_token="rt-legacy") vault = secret_vault_factory(available=False, failure=KeyringUnreachable()) clear_cli_token(vault=vault) - assert "sk-legacy" not in _token_file(isolated_home).read_text() + left_on_disk = _token_file(isolated_home).read_text() + assert "sk-legacy" not in left_on_disk + assert "rt-legacy" not in left_on_disk def test_a_repeat_logout_never_answers_its_own_warning_with_an_all_clear( self, isolated_home, secret_vault_factory From 79cac3656407d983326c364cac43b1a1ef9f87bc Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 20 Aug 2026 12:34:38 -0700 Subject: [PATCH 2/9] fix(ui): keep keyword tier rules that target operator-defined tiers when hydrating the edit modal (#37413) --- .../components/add_model/KeywordTierRules.tsx | 5 ++-- .../complexity_router_keywords.test.ts | 25 +++++++++++++++++++ .../add_model/complexity_router_keywords.ts | 14 +++++------ 3 files changed, 35 insertions(+), 9 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/complexity_router_keywords.test.ts diff --git a/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx b/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx index b5291a34cb6..0b47bc8500a 100644 --- a/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx +++ b/ui/litellm-dashboard/src/components/add_model/KeywordTierRules.tsx @@ -14,7 +14,8 @@ export type ComplexityTier = "SIMPLE" | "MEDIUM" | "COMPLEX" | "REASONING"; export interface KeywordTierRule { id: string; keywords: string[]; - tier: ComplexityTier; + /** A built-in tier name, or with a custom tier set, one of the defined tier names. */ + tier: string; } interface KeywordTierRulesProps { @@ -99,7 +100,7 @@ const KeywordTierRules: React.FC = ({ rules, onChange, ti { + setLogoUrlDarkInput(event.target.value); + setLogoUrlDark(event.target.value || null); + }} + /> +

+ Enter a URL for a logo suited to dark backgrounds, or leave empty to reuse the logo above +

+
diff --git a/ui/litellm-dashboard/src/components/ui/sonner.tsx b/ui/litellm-dashboard/src/components/ui/sonner.tsx index 034e093a118..5c55557b024 100644 --- a/ui/litellm-dashboard/src/components/ui/sonner.tsx +++ b/ui/litellm-dashboard/src/components/ui/sonner.tsx @@ -1,12 +1,15 @@ "use client"; import { CircleCheckIcon, InfoIcon, Loader2Icon, OctagonXIcon, TriangleAlertIcon } from "lucide-react"; +import { useTheme } from "next-themes"; import { Toaster as Sonner, type ToasterProps } from "sonner"; function Toaster({ ...props }: ToasterProps) { + const { resolvedTheme } = useTheme(); + return ( { - document.documentElement.classList.remove("dark"); -}); - -afterAll(() => { - document.documentElement.classList.remove("dark"); -}); - -describe("useIsDarkMode", () => { - it("reports the dark class already on the root element at mount", () => { - document.documentElement.classList.add("dark"); - - const { result } = renderHook(() => useIsDarkMode()); - - expect(result.current).toBe(true); - }); - - it("follows the root element's dark class as it is toggled", async () => { - const { result } = renderHook(() => useIsDarkMode()); - expect(result.current).toBe(false); - - document.documentElement.classList.add("dark"); - await waitFor(() => expect(result.current).toBe(true)); - - document.documentElement.classList.remove("dark"); - await waitFor(() => expect(result.current).toBe(false)); - }); - - it("stops observing the root element once unmounted", () => { - const disconnect = vi.spyOn(MutationObserver.prototype, "disconnect"); - - const { unmount } = renderHook(() => useIsDarkMode()); - unmount(); - - expect(disconnect).toHaveBeenCalled(); - disconnect.mockRestore(); - }); -}); diff --git a/ui/litellm-dashboard/src/hooks/useIsDarkMode.ts b/ui/litellm-dashboard/src/hooks/useIsDarkMode.ts deleted file mode 100644 index bccaa31cb3a..00000000000 --- a/ui/litellm-dashboard/src/hooks/useIsDarkMode.ts +++ /dev/null @@ -1,13 +0,0 @@ -import { useSyncExternalStore } from "react"; - -const subscribe = (onStoreChange: () => void): (() => void) => { - const observer = new MutationObserver(onStoreChange); - observer.observe(document.documentElement, { attributes: true, attributeFilter: ["class"] }); - return () => observer.disconnect(); -}; - -const getSnapshot = (): boolean => document.documentElement.classList.contains("dark"); - -const getServerSnapshot = (): boolean => false; - -export const useIsDarkMode = (): boolean => useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot); diff --git a/ui/litellm-dashboard/src/hooks/useSyntaxTheme.test.tsx b/ui/litellm-dashboard/src/hooks/useSyntaxTheme.test.tsx index 13b7d8136ef..53a152e6f44 100644 --- a/ui/litellm-dashboard/src/hooks/useSyntaxTheme.test.tsx +++ b/ui/litellm-dashboard/src/hooks/useSyntaxTheme.test.tsx @@ -1,47 +1,50 @@ import { act, renderHook } from "@testing-library/react"; +import { ThemeProvider, useTheme } from "next-themes"; +import type { ReactNode } from "react"; import { oneDark } from "react-syntax-highlighter/dist/esm/styles/prism"; import { afterAll, beforeEach, describe, expect, it } from "vitest"; import { useSyntaxTheme, type SyntaxTheme } from "./useSyntaxTheme"; const callerLightTheme: SyntaxTheme = { 'code[class*="language-"]': { color: "rebeccapurple" } }; -const setRootDark = async (enabled: boolean) => { - await act(async () => { - document.documentElement.classList.toggle("dark", enabled); - await Promise.resolve(); +const renderSyntaxTheme = (defaultTheme: string) => + renderHook(() => ({ syntax: useSyntaxTheme(callerLightTheme), setTheme: useTheme().setTheme }), { + wrapper: ({ children }: { children: ReactNode }) => ( + + {children} + + ), }); -}; beforeEach(() => { - document.documentElement.classList.remove("dark"); + localStorage.clear(); + document.documentElement.classList.remove("dark", "light"); }); afterAll(() => { - document.documentElement.classList.remove("dark"); + document.documentElement.classList.remove("dark", "light"); }); describe("useSyntaxTheme", () => { it("keeps the caller's own stylesheet in light mode", () => { - const { result } = renderHook(() => useSyntaxTheme(callerLightTheme)); + const { result } = renderSyntaxTheme("light"); - expect(result.current).toBe(callerLightTheme); + expect(result.current.syntax).toBe(callerLightTheme); }); - it("swaps to oneDark when the root element turns dark", async () => { - const { result } = renderHook(() => useSyntaxTheme(callerLightTheme)); + it("serves oneDark when the resolved theme is dark", () => { + const { result } = renderSyntaxTheme("dark"); - await setRootDark(true); - - expect(result.current).toBe(oneDark); + expect(result.current.syntax).toBe(oneDark); }); - it("restores the caller's stylesheet when dark mode is turned back off", async () => { - document.documentElement.classList.add("dark"); - const { result } = renderHook(() => useSyntaxTheme(callerLightTheme)); - expect(result.current).toBe(oneDark); + it("swaps stylesheets when the theme is changed at runtime", () => { + const { result } = renderSyntaxTheme("light"); - await setRootDark(false); + act(() => result.current.setTheme("dark")); + expect(result.current.syntax).toBe(oneDark); - expect(result.current).toBe(callerLightTheme); + act(() => result.current.setTheme("light")); + expect(result.current.syntax).toBe(callerLightTheme); }); }); diff --git a/ui/litellm-dashboard/src/hooks/useSyntaxTheme.ts b/ui/litellm-dashboard/src/hooks/useSyntaxTheme.ts index 80f5b514751..d107aa9ed18 100644 --- a/ui/litellm-dashboard/src/hooks/useSyntaxTheme.ts +++ b/ui/litellm-dashboard/src/hooks/useSyntaxTheme.ts @@ -1,8 +1,8 @@ import type { CSSProperties } from "react"; +import { useTheme } from "next-themes"; import { oneDark } from "react-syntax-highlighter/dist/esm/styles/prism"; -import { useIsDarkMode } from "./useIsDarkMode"; - export type SyntaxTheme = Record; -export const useSyntaxTheme = (light: SyntaxTheme): SyntaxTheme => (useIsDarkMode() ? oneDark : light); +export const useSyntaxTheme = (light: SyntaxTheme): SyntaxTheme => + useTheme().resolvedTheme === "dark" ? oneDark : light; From e07a7129c56c3ad54be965e63dc5b8688cc07a22 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 20 Aug 2026 13:12:55 -0700 Subject: [PATCH 7/9] feat(proxy): redact or drop individual batch records instead of rejecting the file (#37561) * feat(proxy): redact or drop individual batch records instead of rejecting the file A single record tripping a guardrail rejected the whole upload, which is unusable for a file holding thousands of rows. A record a guardrail rewrites is now submitted in its rewritten form, a record it blocks is left out, and the create response reports every changed record by both custom_id and line so a caller can reconcile against the file it sent. The same outcome is written to the proxy log and to request metadata, so it is not visible only to the caller. A rewritten record goes straight to a spool and only its offset is carried, so a masking guardrail touching most rows of a large upload does not build a second copy of the file on the heap, and the rewrite runs off the event loop the way the sibling full-file validation does. Both proxy-injected metadata keys are captured from the record and restored exactly, including an explicit null, so a masked row keeps the tags that decide how it is attributed. A record is dropped only when a guardrail judged its content. `GuardrailRaisedException` now carries `blocked_content` for that, because half its raise sites in the repo signal an unreachable or unparseable backend under a fail-closed policy, and treating those as blocks would turn "refuse this request" into "drop this record and submit the rest". The default is off, so a raise that does not say what it means aborts the upload instead of silently shrinking the file. * fix(proxy): only drop a batch record on a verdict the guardrail actually reached A guardrail that reports a technical failure as an HTTPException carrying a block status was read as a content block, so an unreachable backend under a fail-closed policy quietly shrank the file instead of failing the upload. Two in-tree integrations do exactly that, and one of them defaults to fail-closed, so the broken configuration was the default one. Such an exception is raised `from` the underlying error, which is a deliberate statement that something else caused it, and no content verdict in the repo is raised that way, so the chain now settles it. Implicit context is left alone, since a block raised inside an unrelated `except` would read as a failure. Two annotation errors in the same family: the one GuardrailRaisedException subclass in tree never opted into blocked_content, so a real block took the whole upload down with it, and straiker's block helper is reached both from its verdict and from its fail-closed handler, so it claimed a verdict for an outage. The helper now takes the flag from its caller. A record could also opt itself out of the chain. Guardrail selection reads a body-level `guardrails` key ahead of the proxy-injected list, and online that key can only add to the key and team selection, never replace it, so a batch record naming an empty list skipped every guardrail that was not default_on and was still reported as scanned. Every injected key is now stripped before dispatch and restored afterwards. A guardrail that reroutes a record to another model is honoured on the online path by rewriting the model, which the scan read as a rewrite and submitted in the same file, sending content to the provider the reroute existed to avoid. Every record of a batch file goes to one provider, so the upload is refused instead, naming the line. The scan spool is closed on the paths that never read it back. * fix(proxy): give the scan the metadata bag guardrails actually read, and close its spools The narrowed request metadata was installed under `litellm_metadata` only, but a record is scanned as the chat request it describes, and the guardrails that pick a policy from a request header read `metadata` instead. Noma choosing an application and Aim choosing a user both look there, so the header allowlist added for them did not reach either one and a batch record was still evaluated under the fallback policy. The scan metadata now goes into both bags, which are both stripped and restored, so neither survives into the record that ships. The scan spool was closed on the paths that abort, which are exactly the paths where it is empty, and left open on the one path where it holds the rewritten records. Nothing closed the rewrite output either, where before this feature the uploaded handle belonged to Starlette. The upload now owns both and closes them however it exits. * fix(proxy): register the scan spool before the rewrite can fail The scan spool was added to the request's cleanup list only after the rewrite returned, so a rewrite that raised, which for a spilled file can be as ordinary as the disk filling up, jumped to the handler with the list still empty and left the scan's own handle open. The rewrite also left its half-written output behind on that path, since nothing owns that handle until it is returned. Both now close. --- litellm/exceptions.py | 13 + litellm/integrations/custom_guardrail.py | 75 +- .../guardrail_hooks/deepkeep/deepkeep.py | 1 + .../generic_guardrail_api.py | 1 + .../guardrail_hooks/ovalix/ovalix.py | 1 + .../promptguard/promptguard.py | 1 + .../guardrail_hooks/singulr/singulr.py | 1 + .../guardrail_hooks/straiker/straiker.py | 4 + .../guardrail_hooks/tool_permission.py | 4 +- .../vigil_guard/vigil_guard.py | 2 + .../batch_guardrails.py | 332 +++++++-- .../openai_files_endpoints/files_endpoints.py | 85 ++- litellm/types/llms/openai.py | 40 ++ .../guardrail_hooks/test_straiker.py | 26 + .../test_batch_guardrails.py | 641 ++++++++++++++++-- .../test_files_endpoint.py | 160 ++++- 16 files changed, 1224 insertions(+), 163 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 2eb4232fef9..286f7528896 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1039,16 +1039,29 @@ class LiteLLMUnknownProvider(BadRequestError): class GuardrailRaisedException(Exception): + """ + Raised both when a guardrail judged content and when it could not judge it at all, since a + guardrail that fails closed refuses the request the same way a policy violation does. + + ``blocked_content`` separates the two. Set it only where the guardrail actually reached a + verdict on the payload; leave it alone for an unreachable backend, a timeout, or a response + the integration could not parse. Callers that treat a block as something other than a plain + failure, such as the batch path dropping one record and submitting the rest, must gate on it, + because dropping a record no guardrail ever inspected is a silent loss of enforcement. + """ + def __init__( self, guardrail_name: str | None = None, message: str = "", should_wrap_with_default_message: bool = True, status_code: int = 400, + blocked_content: bool = False, ): default_message: Final = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}" self.guardrail_name = guardrail_name self.status_code = status_code + self.blocked_content = blocked_content self.message = default_message if should_wrap_with_default_message else message super().__init__(self.message) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 0172c789d1e..f2e390625f5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -65,6 +65,41 @@ _guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.Cont ) +def is_guardrail_intervention(e: Exception) -> bool: + """ + Returns True if the exception represents an intentional guardrail block + (this was logged previously as an API failure - guardrail_failed_to_respond). + + Guardrails signal intentional blocks by raising: + - GuardrailRaisedException (generic guardrail API, tool permission) + - BlockedPiiEntityError (Presidio PII detection) + - SensitiveDataRouteException (sensitive-data reroute to on-premise model) + - HTTPException with a block-signalling status (400, 403, 422) + - ModifyResponseException (passthrough mode violation) + + Only the statuses guardrails use in-tree to signal a deliberate rejection + count as an intervention: 400 (content policy), 403 (e.g. akto) and 422 + (e.g. llm_as_a_judge). Other 4xx codes are commonly propagated from an + upstream guardrail provider response (401 bad key, 408 timeout, 429 rate + limit, or a raw upstream status), which are technical failures, not + blocks, so they stay guardrail_failed_to_respond. + """ + if isinstance(e, ModifyResponseException): + return True + if isinstance( + e, + ( + GuardrailRaisedException, + BlockedPiiEntityError, + SensitiveDataRouteException, + ), + ): + return True + if HTTPException is not None and isinstance(e, HTTPException) and e.status_code in _GUARDRAIL_BLOCK_STATUS_CODES: + return True + return False + + def _strict_guardrail_modes_enabled() -> bool: """Whether guardrail-mode validation raises (default) or logs a warning. @@ -429,11 +464,13 @@ class CustomGuardrail(CustomLogger): f"Sensitive data detected by {self.guardrail_name} (routing skipped: request has no session_id)" ), guardrail_name=self.guardrail_name, + blocked_content=True, ) else: raise GuardrailRaisedException( message=f"Sensitive data detected by {self.guardrail_name}", guardrail_name=self.guardrail_name, + blocked_content=True, ) @staticmethod @@ -1068,42 +1105,8 @@ class CustomGuardrail(CustomLogger): @staticmethod def _is_guardrail_intervention(e: Exception) -> bool: - """ - Returns True if the exception represents an intentional guardrail block - (this was logged previously as an API failure - guardrail_failed_to_respond). - - Guardrails signal intentional blocks by raising: - - GuardrailRaisedException (generic guardrail API, tool permission) - - BlockedPiiEntityError (Presidio PII detection) - - SensitiveDataRouteException (sensitive-data reroute to on-premise model) - - HTTPException with a block-signalling status (400, 403, 422) - - ModifyResponseException (passthrough mode violation) - - Only the statuses guardrails use in-tree to signal a deliberate rejection - count as an intervention: 400 (content policy), 403 (e.g. akto) and 422 - (e.g. llm_as_a_judge). Other 4xx codes are commonly propagated from an - upstream guardrail provider response (401 bad key, 408 timeout, 429 rate - limit, or a raw upstream status), which are technical failures, not - blocks, so they stay guardrail_failed_to_respond. - """ - if isinstance(e, ModifyResponseException): - return True - if isinstance( - e, - ( - GuardrailRaisedException, - BlockedPiiEntityError, - SensitiveDataRouteException, - ), - ): - return True - if ( - HTTPException is not None - and isinstance(e, HTTPException) - and e.status_code in _GUARDRAIL_BLOCK_STATUS_CODES - ): - return True - return False + """Retained spelling for existing callers; prefer ``is_guardrail_intervention``.""" + return is_guardrail_intervention(e) def _process_error( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index e1c0653ebd3..c0f72af7576 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py @@ -356,6 +356,7 @@ class DeepKeepGuardrail(CustomGuardrail): guardrail_name=GUARDRAIL_NAME, message=error_message, should_wrap_with_default_message=False, + blocked_content=True, ) return self._build_return_inputs( diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 16768a4b08f..e3cf645ceaf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -464,6 +464,7 @@ class GenericGuardrailAPI(CustomGuardrail): guardrail_name=GUARDRAIL_NAME, message=error_message, should_wrap_with_default_message=False, + blocked_content=True, ) return self._build_guardrail_return_inputs( diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 4a20adf0e82..6644a3d3902 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -54,6 +54,7 @@ class OvalixGuardrailBlockedException(GuardrailRaisedException): guardrail_name=guardrail_name, message=message, should_wrap_with_default_message=should_wrap_with_default_message, + blocked_content=True, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index 4775a8b3caa..c25f704567e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -169,6 +169,7 @@ class PromptGuardGuardrail(CustomGuardrail): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=(f"Blocked by PromptGuard: {threat_type} (confidence={confidence}, event_id={event_id})"), + blocked_content=True, ) if decision == "redact": diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index cd9da8a58b7..3865ba4ed0e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -211,6 +211,7 @@ class SingulrGuardrail(CustomGuardrail): raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=f"Blocked by Singulr: {result.blocking_due_to or 'unknown'}", + blocked_content=True, ) return inputs diff --git a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py index eee66f93b7a..7cca1ae2d63 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py +++ b/litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py @@ -536,12 +536,14 @@ class StraikerGuardrail(CustomGuardrail): request_data: dict, input_type: Literal["request", "response"], message: str, + blocked_content: bool = False, ) -> NoReturn: if input_type == "request": raise GuardrailRaisedException( guardrail_name=self.guardrail_name or GUARDRAIL_NAME, message=message, should_wrap_with_default_message=False, + blocked_content=blocked_content, ) raise ModifyResponseException( message=message, @@ -623,6 +625,7 @@ class StraikerGuardrail(CustomGuardrail): request_data=request_data, input_type=input_type, message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, + blocked_content=True, ) if parsed.action == "GUARDRAIL_INTERVENED": is_streamed_response: Final = input_type == "response" and _is_streamed_request(request_data) @@ -631,6 +634,7 @@ class StraikerGuardrail(CustomGuardrail): request_data=request_data, input_type=input_type, message=parsed.blocked_reason or DEFAULT_BLOCK_MESSAGE, + blocked_content=True, ) return self._intervened_inputs(inputs, parsed) return inputs diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 0514d2ab6f7..3c5625bc272 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -527,7 +527,9 @@ class ToolPermissionGuardrail(CustomGuardrail): if not is_allowed and message is not None: verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": - raise GuardrailRaisedException(guardrail_name=self.guardrail_name, message=message) + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, message=message, blocked_content=True + ) return tuple( ( diff --git a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py index ee1aade8ea6..6b8148645aa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py @@ -205,6 +205,7 @@ class VigilGuardGuardrail(CustomGuardrail): guardrail_name=self.guardrail_name, message=self._build_block_reason(analysis), should_wrap_with_default_message=False, + blocked_content=True, ) if decision == "SANITIZED": @@ -245,6 +246,7 @@ class VigilGuardGuardrail(CustomGuardrail): guardrail_name=self.guardrail_name, message=self._build_block_reason(analysis), should_wrap_with_default_message=False, + blocked_content=True, ) if decision == "SANITIZED": diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py index 349039a050a..5c886ca0e9b 100644 --- a/litellm/proxy/openai_files_endpoints/batch_guardrails.py +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -10,6 +10,8 @@ from __future__ import annotations import asyncio import copy import json +import re +import tempfile from collections.abc import Iterator, Mapping from dataclasses import dataclass from types import MappingProxyType @@ -19,8 +21,11 @@ from urllib.parse import urlsplit from fastapi import HTTPException from typing_extensions import assert_never +from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import is_guardrail_intervention from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.llms.openai import BatchGuardrailRecord, BatchGuardrailReport from litellm.types.utils import CallTypes, CallTypesLiteral if TYPE_CHECKING: @@ -30,12 +35,28 @@ EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) _SCAN_WINDOW: Final = 32 -_SCAN_METADATA_KEY: Final = "litellm_metadata" +# Past this the rewrite rolls to disk, keeping the router's per-deployment deepcopy of the handle +# as cheap as it is for the spooled upload this replaces. +_REWRITE_SPOOL_BYTES: Final = 1024 * 1024 -# `metadata` is dropped rather than diffed: guardrail dispatch writes its bookkeeping into it -# whenever the payload has one, and a record's own metadata is not scanned content on the -# online path either. -_INJECTED_KEYS: Final = frozenset({_SCAN_METADATA_KEY, "metadata"}) +# custom_id is caller-supplied and reaches a log line, so it is stripped of control characters +# and capped rather than rendered as given. +_CONTROL_CHARACTERS: Final = re.compile(r"[\x00-\x1f\x7f]") +_CUSTOM_ID_LOG_LIMIT: Final = 128 +_SUMMARY_LIMIT: Final = 50 + +_SCAN_METADATA_KEY: Final = "litellm_metadata" +_SCAN_METADATA_BAGS: Final = (_SCAN_METADATA_KEY, "metadata") + +# Set by pre_call_hook when a guardrail rerouted the request to a different model. +_ROUTE_APPLIED_KEY: Final = "sensitive_data_routing_applied" + +# Dropped before dispatch and restored afterwards rather than diffed. Guardrail dispatch writes +# its bookkeeping into `metadata`, and a record's own metadata is not scanned content on the +# online path either. `guardrails` is dropped because guardrail selection reads it ahead of the +# proxy-injected list, so leaving it would let a record's own body opt out of the chain its key +# and team selected; online that key can only add to the list, never replace it. +_INJECTED_KEYS: Final = frozenset({_SCAN_METADATA_KEY, "metadata", "guardrails"}) # Only what guardrail dispatch reads. The parent OTel span is deliberately left out: parenting one # guardrail span per record would put tens of thousands of spans on a single upload's trace. @@ -83,12 +104,80 @@ class UnscannableRecord: @dataclass(frozen=True, slots=True) -class RedactionRequired: +class UnroutableRecord: line_number: int custom_id: str | None + guardrail: str | None -BatchScanFailure: TypeAlias = UnparseableRecord | UnscannableRecord | RedactionRequired +BatchScanFailure: TypeAlias = UnparseableRecord | UnscannableRecord | UnroutableRecord + + +@dataclass(frozen=True, slots=True) +class _Redaction: + """A rewritten record on its way to the scan spool, held only for the window it was scanned in.""" + + line_number: int + custom_id: str | None + text: str + + +@dataclass(frozen=True, slots=True) +class RecordRedacted: + line_number: int + custom_id: str | None + offset: int + length: int + """Where the re-serialized record sits in the scan spool, so a large file's rewrites stay off the heap.""" + + +@dataclass(frozen=True, slots=True) +class RecordDropped: + line_number: int + custom_id: str | None + guardrail: str | None = None + + +_RecordChange: TypeAlias = RecordRedacted | RecordDropped +_ScanOutcome: TypeAlias = BatchScanFailure | _Redaction | RecordDropped + + +@dataclass(frozen=True, slots=True) +class BatchScanResult: + """What the scan decided, per record. Empty changes means the upload proceeds untouched.""" + + changes: tuple[_RecordChange, ...] + scanned_records: int + redactions: BinaryIO + """Spool holding every rewritten record, keyed by the offsets on each ``RecordRedacted``.""" + + @property + def submitted_records(self) -> int: + return self.scanned_records - sum(1 for change in self.changes if isinstance(change, RecordDropped)) + + def summary(self) -> str: + """Compact per-record outcome for the server-side log line, capped so one upload cannot flood it.""" + shown: Final = ", ".join( + f"line {change.line_number}{_describe(change.custom_id)} " + f"{'redacted' if isinstance(change, RecordRedacted) else 'dropped'}" + for change in self.changes[:_SUMMARY_LIMIT] + ) + remaining: Final = len(self.changes) - _SUMMARY_LIMIT + return shown if remaining <= 0 else f"{shown}, and {remaining} more" + + def report(self) -> BatchGuardrailReport: + return BatchGuardrailReport( + submitted_records=self.submitted_records, + modified_records=tuple( + BatchGuardrailRecord( + line=change.line_number, + custom_id=change.custom_id, + action="redacted" if isinstance(change, RecordRedacted) else "dropped", + guardrail=change.guardrail if isinstance(change, RecordDropped) else None, + ) + for change in self.changes + ), + ) @dataclass(frozen=True, slots=True) @@ -114,25 +203,75 @@ def raise_public(failure: BatchScanFailure) -> NoReturn: "and its body has no messages, prompt or input, so guardrails cannot read it. " "Give the record a chat, completion, embedding, responses or messages body" ) - case RedactionRequired(line_number=line_number, custom_id=custom_id): + case UnroutableRecord(line_number=line_number, custom_id=custom_id, guardrail=guardrail): raise _rejected( - f"A guardrail changed batch input line {line_number}{_describe(custom_id)}. " - "Per-record redaction is not enabled, so the file was rejected rather than modified" + f"Batch input line {line_number}{_describe(custom_id)} was routed to a different model by " + f"{guardrail or 'a guardrail'}, and every record of a batch file goes to one provider, so " + "the file cannot be submitted. Send that record outside the batch" ) case _: assert_never(failure) +def raise_nothing_to_submit() -> NoReturn: + """Every record was blocked, so there is no batch left to create.""" + raise _rejected( + "Every record in the batch input file was blocked by a guardrail, so there is nothing left to submit" + ) + + +def _is_content_block(exc: BaseException) -> bool: + """ + Whether the guardrail judged the record, as opposed to failing to judge it. + + Stricter than ``is_guardrail_intervention``, which answers a different question and counts + every ``GuardrailRaisedException`` as a block. Several integrations raise that same exception + for an unreachable backend or an unparseable response, and only when the operator configured + the guardrail to fail closed, so treating it as a block would turn "refuse this request" into + "drop this record and submit the rest", which is the silent loss of enforcement this whole + path exists to prevent. A guardrail that does not say it blocked content aborts the upload. + + Guardrails that report a technical failure as an ``HTTPException`` carrying a block status + are caught by ``__cause__``: raising ``from`` the underlying error is a deliberate statement + that something else caused this, which a verdict on content never is. Implicit context is + left alone, since a block raised inside an unrelated ``except`` would read as a failure. + """ + if isinstance(exc, GuardrailRaisedException): + return exc.blocked_content + if exc.__cause__ is not None: + return False + return is_guardrail_intervention(exc) + + +def _naming_guardrail(exc: BaseException) -> str | None: + """The guardrail that raised, from whichever place it recorded its own name.""" + named: Final = getattr(exc, "guardrail_name", None) + if isinstance(named, str): + return named + detail: Final = getattr(exc, "detail", None) + enriched: Final = detail.get("guardrail_name") if isinstance(detail, dict) else None + return enriched if isinstance(enriched, str) else None + + def _describe(custom_id: str | None) -> str: - return f" (custom_id {custom_id})" if custom_id else "" + if not custom_id: + return "" + safe: Final = _CONTROL_CHARACTERS.sub(" ", custom_id)[:_CUSTOM_ID_LOG_LIMIT] + return f" (custom_id {safe})" + + +def _iter_lines(source: BinaryIO) -> Iterator[tuple[int, str]]: + """Yield every non-blank line with its 1-based number, so both passes number records alike.""" + for line_number, raw_line in enumerate(source, start=1): + text = raw_line.decode("utf-8") + if text.strip(): + yield line_number, text def _iter_records(source: BinaryIO) -> Iterator[_ParsedRecord]: """Yield one record per line, relying on the upload validation that already ran.""" - for line_number, raw_line in enumerate(source, start=1): - text = raw_line.decode("utf-8") - if text.strip(): - yield _ParsedRecord(line_number=line_number, payload=json.loads(text)) + for line_number, text in _iter_lines(source): + yield _ParsedRecord(line_number=line_number, payload=json.loads(text)) def _call_type_from_url(url: str) -> CallTypesLiteral | None: @@ -204,7 +343,7 @@ async def _scan_record( scan_metadata: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, -) -> BatchScanFailure | None: +) -> _ScanOutcome | None: body: Final = record.payload.get("body") if not isinstance(body, dict): return UnparseableRecord(line_number=record.line_number) @@ -220,26 +359,50 @@ async def _scan_record( ) scan_input: Final[dict[str, object]] = copy.deepcopy(body) # mutable-ok: pre_call_hook mutates the dict it is given - scan_input.pop("metadata", None) - # Deep, and per record: `headers` and `tags` are nested containers shared with the upload - # request and with every other record in the window, and a guardrail that writes into one in - # place would otherwise leak across records and back into the request. The narrowing above - # already removed the values that cannot be copied. - scan_input[_SCAN_METADATA_KEY] = copy.deepcopy(dict(scan_metadata)) # mutable-ok: guardrails write here + own_injected: Final = MappingProxyType({key: body[key] for key in _INJECTED_KEYS if key in body}) + for injected in _INJECTED_KEYS: + scan_input.pop(injected, None) + # Both bags, because guardrails read whichever one their own route populates and a record + # scanned as chat reaches ones that only ever look at `metadata`; both are injected keys, so + # neither survives into the record that ships. Deep, and per bag per record, because `headers` + # and `tags` are nested containers otherwise shared with the upload request and with every + # other record in the window. The narrowing above already removed what cannot be copied. + for injected in _SCAN_METADATA_BAGS: + scan_input[injected] = copy.deepcopy(dict(scan_metadata)) # mutable-ok: guardrails write here - # The chain hands back the body it produced, which may be a replacement for the dict it was - # given rather than that same dict mutated, so this is what gets compared. - scanned: Final[dict] = await proxy_logging_obj.pre_call_hook( # mutable-ok: the guardrails' own dict - user_api_key_dict=user_api_key_dict, - data=scan_input, - call_type=call_type, - guardrails_only=True, - ) + try: + # The chain hands back the body it produced, which may be a replacement for the dict it was + # given rather than that same dict mutated, so this is what gets compared. + scanned: Final[dict] = await proxy_logging_obj.pre_call_hook( # mutable-ok: the guardrails' own dict + user_api_key_dict=user_api_key_dict, + data=scan_input, + call_type=call_type, + guardrails_only=True, + ) + except Exception as exc: + if _is_content_block(exc): + return RecordDropped(line_number=record.line_number, custom_id=custom_id, guardrail=_naming_guardrail(exc)) + raise + + rerouted: Final = scanned.get("metadata") + if isinstance(rerouted, dict) and rerouted.get(_ROUTE_APPLIED_KEY): + return UnroutableRecord( + line_number=record.line_number, + custom_id=custom_id, + guardrail=rerouted.get("sensitive_data_routing_guardrail"), + ) compared: Final = (frozenset(body) | frozenset(scanned)) - _INJECTED_KEYS - if _fingerprint(scanned, compared) != _fingerprint(body, compared): - return RedactionRequired(line_number=record.line_number, custom_id=custom_id) - return None + if _fingerprint(scanned, compared) == _fingerprint(body, compared): + return None + for injected in _INJECTED_KEYS: + scanned.pop(injected, None) + scanned.update(own_injected) + return _Redaction( + line_number=record.line_number, + custom_id=custom_id, + text=json.dumps({**record.payload, "body": scanned}), # mutable-ok: json.dumps needs a plain dict + ) async def _scan_window( @@ -247,7 +410,7 @@ async def _scan_window( scan_metadata: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, -) -> tuple[tuple[int, BatchScanFailure | BaseException], ...]: +) -> tuple[tuple[int, _ScanOutcome | BaseException], ...]: """``return_exceptions=True`` so one record raising never leaves its siblings unobserved.""" outcomes: Final = await asyncio.gather( *(_scan_record(record, scan_metadata, user_api_key_dict, proxy_logging_obj) for record in window), @@ -256,6 +419,20 @@ async def _scan_window( return tuple((record.line_number, outcome) for record, outcome in zip(window, outcomes) if outcome is not None) +def _spool(redactions: BinaryIO, redaction: _Redaction) -> RecordRedacted: + """Park the rewritten record on disk so only its location is carried for the rest of the scan.""" + encoded: Final = redaction.text.encode("utf-8") + redactions.seek(0, 2) + offset: Final = redactions.tell() + redactions.write(encoded) + return RecordRedacted( + line_number=redaction.line_number, + custom_id=redaction.custom_id, + offset=offset, + length=len(encoded), + ) + + def _worst(problems: tuple[tuple[int, BatchScanFailure | BaseException], ...]) -> BatchScanFailure | BaseException: """A guardrail that blocked outranks a record we merely refused; then earliest line wins.""" raised: Final = tuple(problem for problem in problems if isinstance(problem[1], BaseException)) @@ -268,20 +445,36 @@ async def scan_batch_input_file( request_metadata: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, -) -> BatchScanFailure | None: +) -> BatchScanFailure | BatchScanResult: """ Stream a batch input file and run the pre-call guardrail chain against every record. - Returns the record to reject, or None when every record passed. A guardrail that blocks raises - its own exception, which is re-raised untouched so its status code survives. + A record a guardrail rewrites is kept in its rewritten form and a record it blocks is dropped, + which is what the online path does per request. Both are returned for reporting. A guardrail + exception that is not a block is re-raised untouched so its status code survives, since dropping + a record that was never inspected is worse than refusing the file. """ scan_metadata: Final = build_scan_metadata(request_metadata) problems: Final[list[tuple[int, BatchScanFailure | BaseException]]] = [] # mutable-ok: spans windows + changes: Final[list[_RecordChange]] = [] # mutable-ok: accumulates across windows window: Final[list[_ParsedRecord]] = [] # mutable-ok: bounded read-ahead buffer + scanned: Final[list[int]] = [] # mutable-ok: counts records the scan actually reached + redactions: Final = tempfile.SpooledTemporaryFile( # noqa: SIM115 # the rewrite reads this back + max_size=_REWRITE_SPOOL_BYTES + ) async def drain() -> None: if window: - problems.extend(await _scan_window(tuple(window), scan_metadata, user_api_key_dict, proxy_logging_obj)) + scanned.append(len(window)) + for line_number, outcome in await _scan_window( + tuple(window), scan_metadata, user_api_key_dict, proxy_logging_obj + ): + if isinstance(outcome, _Redaction): + changes.append(_spool(redactions, outcome)) + elif isinstance(outcome, RecordDropped): + changes.append(outcome) + else: + problems.append((line_number, outcome)) window.clear() try: @@ -293,12 +486,63 @@ async def scan_batch_input_file( break if not problems: await drain() + except BaseException: + redactions.close() + raise finally: file_source.seek(0) - if not problems: - return None - worst: Final = _worst(tuple(problems)) - if isinstance(worst, BaseException): - raise worst - return worst + if problems: + redactions.close() + worst: Final = _worst(tuple(problems)) + if isinstance(worst, BaseException): + raise worst + return worst + if not changes: + redactions.close() + return BatchScanResult( + changes=tuple(sorted(changes, key=lambda change: change.line_number)), + scanned_records=sum(scanned), + redactions=redactions, + ) + + +def _read_spooled(redactions: BinaryIO, change: RecordRedacted) -> str: + redactions.seek(change.offset) + return redactions.read(change.length).decode("utf-8") + + +def rewrite_batch_input_file(file_source: BinaryIO, result: BatchScanResult) -> BinaryIO: + """ + Re-emit the file with redacted records rewritten and dropped records left out. + + Untouched records are copied through as written rather than re-serialized, so enabling the + feature does not reformat records no guardrail objected to. Blank lines between records are + not carried over, since they are not records. Rewritten records are read back from the scan's + spool rather than from memory, so a file whose records are mostly rewritten does not put a + second copy of itself on the heap. + """ + redacted: Final = MappingProxyType( + {change.line_number: change for change in result.changes if isinstance(change, RecordRedacted)} + ) # mutable-ok: MappingProxyType freezes the lookup table + dropped: Final = frozenset(change.line_number for change in result.changes if isinstance(change, RecordDropped)) + + output: Final = tempfile.SpooledTemporaryFile( # noqa: SIM115 # the caller uploads this handle + max_size=_REWRITE_SPOOL_BYTES + ) + wrote_any = False # rebind-ok: tracks whether a separator is needed + try: + for line_number, text in _iter_lines(file_source): + if line_number in dropped: + continue + change = redacted.get(line_number) + line = text.rstrip("\n") if change is None else _read_spooled(result.redactions, change) + output.write((("\n" if wrote_any else "") + line).encode("utf-8")) + wrote_any = True + except BaseException: + output.close() + raise + finally: + file_source.seek(0) + output.seek(0) + return output diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 5794d490bb3..37cfd9d073d 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -7,6 +7,7 @@ import asyncio import traceback +from collections.abc import Mapping from typing import Any, BinaryIO, Final, cast, get_args import httpx @@ -29,6 +30,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_cloud_storage_uri, ) +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -48,7 +50,10 @@ from litellm.proxy.openai_files_endpoints.batch_file_validation import ( ) from litellm.proxy.openai_files_endpoints.batch_guardrails import ( EMPTY_MAPPING, + BatchScanResult, + raise_nothing_to_submit, raise_public, + rewrite_batch_input_file, scan_batch_input_file, ) from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -111,6 +116,34 @@ def get_files_provider_config( return None +async def _scan_batch_upload( + *, + file_source: bytes | BinaryIO, + purpose: str, + request_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, +) -> BatchScanResult | None: + """Guardrail the records of a batch input file, or None when this upload has nothing to scan.""" + if ( + purpose != "batch" + or isinstance(file_source, bytes) + or not proxy_logging_obj.has_pre_call_guardrails(request_metadata) + ): + return None + outcome: Final = await scan_batch_input_file( + file_source=file_source, + request_metadata=request_metadata, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + if not isinstance(outcome, BatchScanResult): + raise_public(outcome) + if outcome.changes and outcome.submitted_records == 0: + raise_nothing_to_submit() + return outcome + + def get_first_json_object(file_source: bytes | BinaryIO) -> dict | None: try: if isinstance(file_source, (bytes, bytearray)): @@ -338,6 +371,10 @@ async def create_file( ) data: dict = {} + # Spools this request owns. Starlette owns the upload handle; anything the guardrail scan + # opens is ours, and a batch upload that fails after the scan would otherwise hold the + # descriptor and its disk blocks until the collector runs. + spools: Final[list[BinaryIO]] = [] # mutable-ok: filled as the scan opens handles try: # Batch uploads can be gigabytes. Starlette has already spooled the upload # to disk, so stream from that handle instead of reading it into memory. @@ -478,28 +515,42 @@ async def create_file( # /v1/files stores its proxy metadata under litellm_metadata, not metadata request_metadata: Final = data.get("metadata") or data.get("litellm_metadata") or EMPTY_MAPPING - if ( - purpose == "batch" - and not isinstance(file_source, bytes) - and proxy_logging_obj.has_pre_call_guardrails(request_metadata) - ): - scan_failure: Final = await scan_batch_input_file( - file_source=file_source, - request_metadata=request_metadata, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, + scan_result: Final = await _scan_batch_upload( + file_source=file_source, + purpose=purpose, + request_metadata=request_metadata, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + if scan_result is not None and scan_result.changes: + # The caller sees this in the response; a proxy admin needs it server side too, + # and it has to land before the post-call hook for logging callbacks to pick it up. + get_or_create_metadata_bucket(data)[1]["batch_guardrail"] = scan_result.report().model_dump() + verbose_proxy_logger.warning( + "batch guardrails changed %s of %s records in %s: %s", + len(scan_result.changes), + scan_result.scanned_records, + file.filename, + scan_result.summary(), ) - if scan_failure is not None: - raise_public(scan_failure) # Prepare the file data according to FileTypes - file_data: Final = (file.filename, file_source, file.content_type) + if scan_result is not None: + spools.append(scan_result.redactions) + upload_source: Final = ( + await asyncio.to_thread(rewrite_batch_input_file, file_source, scan_result) + if scan_result is not None and scan_result.changes + else file_source + ) + if upload_source is not file_source: + spools.append(upload_source) + file_data: Final = (file.filename, upload_source, file.content_type) ## check if model is a loadbalanced model router_model: str | None = None is_router_model = False if litellm.enable_loadbalancing_on_batch_endpoints is True: - json_obj: Final = get_first_json_object(file_source) + json_obj: Final = get_first_json_object(upload_source) if json_obj: router_model = get_model_from_json_obj(json_object=json_obj) is_router_model = is_known_model(model=router_model, llm_router=llm_router) @@ -567,6 +618,9 @@ async def create_file( if _response is not None and isinstance(_response, OpenAIFileObject): response = _response + if scan_result is not None and scan_result.changes: + response.litellm_batch_guardrail = scan_result.report() + ### RESPONSE HEADERS ### hidden_params: Final = getattr(response, "_hidden_params", {}) or {} model_id: Final = hidden_params.get("model_id", None) or "" @@ -606,6 +660,9 @@ async def create_file( param=getattr(e, "param", "None"), code=getattr(e, "status_code", 500), ) + finally: + for spool in spools: + spool.close() @router.get( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 6457b285cb5..1588c650177 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -279,6 +279,40 @@ OpenAIFilesPurpose = Literal[ ] +class BatchGuardrailRecord(BaseModel): + """One batch input record a guardrail acted on.""" + + line: int + """The 1-based line of the uploaded file the record started on.""" + + custom_id: str | None = None + """The record's own `custom_id`, when it carried one.""" + + action: Literal["redacted", "dropped"] + """`redacted` means the record was submitted with the guardrail's rewrite applied. + + `dropped` means the guardrail blocked it and it was left out of the submitted file. + """ + + guardrail: str | None = None + """Which guardrail dropped the record, when it named itself. + + Set for dropped records only. A guardrail refusing content and a guardrail that is + unreachable under a fail-closed setting raise the same way, so this names the guardrail + to check rather than claiming a reason it cannot distinguish. + """ + + +class BatchGuardrailReport(BaseModel): + """What guardrails did to a batch input file, per record.""" + + submitted_records: int + """How many records reached the provider.""" + + modified_records: tuple[BatchGuardrailRecord, ...] + """Every record that was redacted or dropped, in file order.""" + + class OpenAIFileObject(BaseModel): id: str """The file identifier, which can be referenced in the API endpoints.""" @@ -319,6 +353,12 @@ class OpenAIFileObject(BaseModel): `error` field on `fine_tuning.job`. """ + litellm_batch_guardrail: BatchGuardrailReport | None = None + """Set by the proxy when guardrails acted on a `purpose=batch` upload. + + Absent on every other upload, so OpenAI-shaped clients see an unchanged response. + """ + _hidden_params: dict = {"response_cost": 0.0} # no cost for writing a file def __contains__(self, key) -> bool: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py index 36a2e205ea7..a3d86034f70 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_straiker.py @@ -1067,3 +1067,29 @@ async def test_anthropic_non_streaming_response_reports_usage(): payload = _posted_payload(g) assert payload["usage"] == {"input_tokens": 10, "output_tokens": 5} assert payload["response"]["finish_reason"] == "end_turn" + + +def test_fail_closed_backend_failure_is_not_reported_as_a_content_verdict(): + """A drop-one-record consumer must be able to tell a verdict from an outage; _fail is not a verdict.""" + from litellm.exceptions import GuardrailRaisedException + + guardrail = _make_guardrail() + + with pytest.raises(GuardrailRaisedException) as unreachable: + guardrail._fail( + inputs={}, + request_data={"model": "m"}, + input_type="request", + error="connection refused", + is_unreachable=True, + ) + assert unreachable.value.blocked_content is False + + with pytest.raises(GuardrailRaisedException) as verdict: + guardrail._block( + request_data={"model": "m"}, + input_type="request", + message="blocked", + blocked_content=True, + ) + assert verdict.value.blocked_content is True diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py index c54debc2418..a05b8ae530c 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -4,9 +4,14 @@ import json import pytest from fastapi import HTTPException +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.batch_guardrails import ( - RedactionRequired, + BatchScanResult, + RecordDropped, + RecordRedacted, + rewrite_batch_input_file, UnparseableRecord, UnscannableRecord, raise_public, @@ -64,7 +69,7 @@ def _raise_on(needle, exc): return _hook -async def _scan(source, logging_obj, metadata=None): +async def _scan_full(source, logging_obj, metadata=None): return await scan_batch_input_file( file_source=source, request_metadata=metadata if metadata is not None else {}, @@ -73,6 +78,14 @@ async def _scan(source, logging_obj, metadata=None): ) +async def _scan(source, logging_obj, metadata=None): + """Collapses "the scan found nothing to do" to None so the reject-mode cases read plainly.""" + result = await _scan_full(source, logging_obj, metadata) + if isinstance(result, BatchScanResult): + return None if not result.changes else result + return result + + @pytest.mark.asyncio async def test_clean_file_passes_and_rewinds_the_handle(): source = _jsonl(_record("a"), _record("b"), _record("c")) @@ -98,7 +111,7 @@ async def test_redaction_is_reported_with_line_and_custom_id(): failure = await _scan(source, FakeProxyLogging(_redact_containing("secret"))) - assert failure == RedactionRequired(line_number=2, custom_id="dirty") + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(2, "dirty")] @pytest.mark.asyncio @@ -127,14 +140,28 @@ async def test_records_own_metadata_is_left_out_of_the_scan_and_the_diff(): seen = [] def _write_bookkeeping(data): - seen.append("metadata" in data) + seen.append(dict(data.get("metadata") or {})) data.setdefault("metadata", {})["applied_guardrails"] = ["g"] - assert await _scan(_jsonl(record), FakeProxyLogging(_write_bookkeeping)) is None - assert seen == [False], "the record's own metadata must not be handed to guardrail dispatch" + assert await _scan(_jsonl(record), FakeProxyLogging(_write_bookkeeping), metadata={"tags": ["t"]}) is None + assert seen == [{"tags": ["t"]}], "dispatch sees the proxy's metadata, never the record's own" assert record["body"]["metadata"] == {"team": "finance"} +@pytest.mark.asyncio +async def test_the_scan_metadata_reaches_guardrails_that_only_read_the_metadata_bag(): + """noma and aim read `metadata["headers"]`; a record scanned as chat must reach them too.""" + seen = [] + + await _scan( + _jsonl(_record("a")), + FakeProxyLogging(lambda d: seen.append((d.get("metadata") or {}).get("headers"))), + metadata={"guardrails": ["g"], "headers": {"x-noma-application-id": "app-1"}}, + ) + + assert seen == [{"x-noma-application-id": "app-1"}] + + @pytest.mark.asyncio async def test_request_metadata_is_narrowed_to_what_guardrails_read(): """An OTel-enabled proxy puts a lock-bearing span here; a per-record copy of it is a crash.""" @@ -198,7 +225,7 @@ async def test_guardrail_that_adds_a_key_is_detected(): failure = await _scan(_jsonl(_record("a")), FakeProxyLogging(_add_key)) - assert failure == RedactionRequired(line_number=1, custom_id="a") + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")] @pytest.mark.asyncio @@ -210,7 +237,7 @@ async def test_guardrail_that_adds_a_null_valued_key_is_detected(): failure = await _scan(_jsonl(_record("a")), FakeProxyLogging(_add_null_key)) - assert failure == RedactionRequired(line_number=1, custom_id="a") + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")] @pytest.mark.asyncio @@ -223,7 +250,7 @@ async def test_guardrail_that_drops_a_null_valued_key_is_detected(): failure = await _scan(_jsonl(record), FakeProxyLogging(_drop_null_key)) - assert failure == RedactionRequired(line_number=1, custom_id="a") + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")] @pytest.mark.asyncio @@ -262,26 +289,6 @@ async def test_empty_url_falls_back_to_its_body_shape(body, expected_call_type): assert logging_obj.seen[0][0] == expected_call_type -@pytest.mark.asyncio -async def test_blocking_guardrail_outranks_an_earlier_refused_record(): - """PR 2 turns RedactionRequired into a non-failure; a block must not be lost behind it.""" - blocked = HTTPException(status_code=403, detail={"error": "Violated guardrail policy"}) - - def _hook(data): - content = data["messages"][0]["content"] - if content == "raiser": - raise blocked - if content == "redact": - data["messages"][0]["content"] = "***" - - source = _jsonl(_record("a", content="redact"), _record("b", content="raiser")) - - with pytest.raises(HTTPException) as raised: - await _scan(source, FakeProxyLogging(_hook)) - - assert raised.value is blocked - - @pytest.mark.asyncio async def test_handle_is_rewound_even_when_a_record_is_refused(): source = _jsonl(_record("a", content="secret")) @@ -390,45 +397,14 @@ async def test_url_less_record_whose_body_shape_is_unknown_is_rejected(): @pytest.mark.asyncio async def test_blocking_guardrail_exception_propagates_unwrapped(): - blocked = HTTPException(status_code=403, detail={"error": "Violated guardrail policy"}) + blocked = HTTPException(status_code=503, detail={"error": "guardrail service unavailable"}) source = _jsonl(_record("a"), _record("b", content="tripwire")) with pytest.raises(HTTPException) as raised: await _scan(source, FakeProxyLogging(_raise_on("tripwire", blocked))) assert raised.value is blocked, "the guardrail's own exception must survive so its status code does" - assert raised.value.status_code == 403 - - -@pytest.mark.asyncio -async def test_earliest_refused_record_is_the_one_reported(): - source = _jsonl(_record("a"), _record("b", content="secret"), _record("c", content="secret")) - - failure = await _scan(source, FakeProxyLogging(_redact_containing("secret"))) - - assert failure == RedactionRequired(line_number=2, custom_id="b") - - -@pytest.mark.asyncio -async def test_earliest_failing_record_wins_when_the_raise_comes_first(): - blocked = HTTPException(status_code=400, detail="blocked") - - def _hook(data): - content = data["messages"][0]["content"] - if content == "raiser": - raise blocked - if content == "redact": - data["messages"][0]["content"] = "***" - - source = _jsonl( - _record("a", content="raiser"), - _record("b", content="redact"), - ) - - with pytest.raises(HTTPException) as raised: - await _scan(source, FakeProxyLogging(_hook)) - - assert raised.value is blocked + assert raised.value.status_code == 503 @pytest.mark.asyncio @@ -447,7 +423,6 @@ async def test_records_are_not_mutated_by_the_scan(): [ (UnparseableRecord(line_number=7), "line 7"), (UnscannableRecord(line_number=3, custom_id="x", url="/v1/audio/speech"), "custom_id x"), - (RedactionRequired(line_number=2, custom_id=None), "line 2"), ], ) def test_every_failure_maps_to_a_400_naming_the_record(failure, fragment): @@ -463,7 +438,8 @@ async def test_scan_does_not_mutate_the_parsed_record(): """The guardrail must redact a copy. Mutating the record would corrupt what PR 2 writes out.""" from litellm.proxy.openai_files_endpoints.batch_guardrails import _ParsedRecord, _scan_record - record = _ParsedRecord(line_number=1, payload=_record("a", content="my secret is here")) + payload = _record("a", content="my secret is here") + record = _ParsedRecord(line_number=1, payload=payload) failure = await _scan_record( record, @@ -472,7 +448,7 @@ async def test_scan_does_not_mutate_the_parsed_record(): FakeProxyLogging(_redact_containing("secret")), ) - assert failure == RedactionRequired(line_number=1, custom_id="a") + assert (failure.line_number, failure.custom_id) == (1, "a") assert record.payload["body"]["messages"][0]["content"] == "my secret is here", ( "the guardrail redacted the record itself instead of a copy" ) @@ -529,4 +505,537 @@ async def test_guardrail_that_returns_a_replacement_dict_is_detected(): failure = await _scan(_jsonl(_record("a", content="my secret is here")), ReplacingLogging()) - assert failure == RedactionRequired(line_number=1, custom_id="a") + assert [(c.line_number, c.custom_id) for c in failure.changes] == [(1, "a")] + + +def _blocking(needle, status_code=400, guardrail_name="block-guard"): + def _hook(data): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and needle in message["content"]: + raise HTTPException( + status_code=status_code, + detail={"error": "Violated guardrail policy", "guardrail_name": guardrail_name}, + ) + + return _hook + + +@pytest.mark.asyncio +async def test_redact_mode_keeps_a_masked_record_instead_of_rejecting(): + source = _jsonl(_record("a"), _record("b", content="my secret is here"), _record("c")) + + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + + assert [(c.line_number, c.custom_id) for c in result.changes] == [(2, "b")] + rewritten = json.loads(rewrite_batch_input_file(source, result).read().decode().splitlines()[1]) + assert rewritten["body"]["messages"][0]["content"] == "my *** is here" + assert "litellm_metadata" not in rewritten["body"], "proxy metadata must not reach the uploaded file" + assert result.submitted_records == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [400, 403, 422], ids=["content_policy", "akto", "llm_as_a_judge"]) +async def test_every_status_litellm_calls_a_block_drops_the_record(status_code): + """Follows CustomGuardrail._is_guardrail_intervention, so drop matches what litellm logs as a block.""" + source = _jsonl(_record("a"), _record("b", content="tripwire")) + + result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire", status_code))) + + assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail="block-guard"),) + + +@pytest.mark.asyncio +async def test_redact_mode_drops_a_blocked_record_and_submits_the_rest(): + source = _jsonl(_record("a"), _record("b", content="tripwire"), _record("c")) + + result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire"))) + + assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail="block-guard"),) + assert result.submitted_records == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status_code", [500, 502, 408, 429, 401]) +async def test_redact_mode_does_not_drop_a_record_on_an_infrastructure_failure(status_code): + """A guardrail service that is down must abort the upload, never silently cost the caller records.""" + source = _jsonl(_record("a"), _record("b", content="tripwire")) + + with pytest.raises(HTTPException) as raised: + await _scan_full(source, FakeProxyLogging(_blocking("tripwire", status_code))) + + assert raised.value.status_code == status_code + + +@pytest.mark.asyncio +async def test_every_record_blocked_leaves_nothing_to_submit(): + source = _jsonl(_record("a", content="tripwire"), _record("b", content="tripwire")) + + result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire"))) + + assert result.submitted_records == 0 + assert [change.line_number for change in result.changes] == [1, 2] + + +@pytest.mark.asyncio +async def test_rewrite_drops_blocked_records_and_masks_redacted_ones(): + records = [_record("a"), _record("b", content="my secret is here"), _record("c", content="tripwire"), _record("d")] + source = _jsonl(*records) + + def _hook(data): + _redact_containing("secret")(data) + _blocking("tripwire")(data) + + result = await _scan_full(source, FakeProxyLogging(_hook)) + rewritten = rewrite_batch_input_file(source, result) + + lines = [json.loads(line) for line in (rewritten.seek(0), rewritten.read().decode())[1].splitlines()] + assert [line["custom_id"] for line in lines] == ["a", "b", "d"] + assert lines[1]["body"]["messages"][0]["content"] == "my *** is here" + + +@pytest.mark.asyncio +async def test_rewrite_copies_untouched_records_byte_for_byte(): + """Enabling the feature must not reformat records no guardrail objected to.""" + untouched = '{"custom_id":"keep","url":"/v1/chat/completions","body":{"messages":[{"role":"user","content":"hi"}],"model":"m"}}' + dirty = json.dumps(_record("dirty", content="my secret is here")) + source = io.BytesIO((untouched + "\n" + dirty).encode()) + + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + rewritten = rewrite_batch_input_file(source, result) + + assert (rewritten.seek(0), rewritten.read().decode())[1].splitlines()[0] == untouched + + +@pytest.mark.asyncio +async def test_report_names_every_changed_record_in_file_order(): + records = [_record("a"), _record("b", content="tripwire"), _record("c", content="my secret is here")] + + def _hook(data): + _redact_containing("secret")(data) + _blocking("tripwire")(data) + + result = await _scan_full(_jsonl(*records), FakeProxyLogging(_hook)) + report = result.report() + + assert report.submitted_records == 2 + assert [(r.line, r.custom_id, r.action, r.guardrail) for r in report.modified_records] == [ + (2, "b", "dropped", "block-guard"), + (3, "c", "redacted", None), + ] + + +@pytest.mark.asyncio +async def test_clean_file_needs_no_rewrite(): + """A file nothing objected to keeps streaming off disk rather than being buffered in memory.""" + result = await _scan_full(_jsonl(_record("a"), _record("b")), FakeProxyLogging()) + + assert result.changes == () + assert result.submitted_records == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "exc", + [ + GuardrailRaisedException(guardrail_name="g", message="blocked", blocked_content=True), + BlockedPiiEntityError(entity_type="US_SSN", guardrail_name="presidio"), + ], + ids=["guardrail_raised", "blocked_pii_entity"], +) +async def test_litellm_native_block_exceptions_drop_the_record(exc): + """Presidio and friends raise these rather than an HTTPException; they are still policy blocks.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise exc + + source = _jsonl(_record("a"), _record("b", content="tripwire")) + + result = await _scan_full(source, FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail=exc.guardrail_name),) + assert result.submitted_records == 1 + + +@pytest.mark.asyncio +async def test_raising_a_native_block_exception_drops_whatever_status_it_carries(): + """Raising this type IS the block signal in litellm, so the drop set matches what it logs as a block.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise GuardrailRaisedException( + guardrail_name="g", message="refused", status_code=503, blocked_content=True + ) + + result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail="g"),) + + +@pytest.mark.asyncio +async def test_an_unreachable_guardrail_aborts_instead_of_quietly_dropping_the_record(): + """Several integrations raise this same exception when their backend is down and they fail closed.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise GuardrailRaisedException( + guardrail_name="g", message="Singulr API unreachable (block_on_error=True): timed out" + ) + + with pytest.raises(GuardrailRaisedException): + await _scan_full(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook)) + + +@pytest.mark.asyncio +async def test_a_guardrail_subclass_that_blocks_content_drops_only_that_record(): + """A subclass has to opt in too, or a real block takes the whole upload down with it.""" + from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix import OvalixGuardrailBlockedException + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise OvalixGuardrailBlockedException(guardrail_name="ovalix", message="blocked") + + result = await _scan_full(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=2, custom_id="b", guardrail="ovalix"),) + assert result.submitted_records == 1 + + +@pytest.mark.asyncio +async def test_a_record_a_guardrail_rerouted_aborts_rather_than_shipping_to_the_original_provider(): + """pre_call_hook honours a reroute by rewriting `model`; a batch file cannot follow it.""" + from litellm.proxy.openai_files_endpoints.batch_guardrails import UnroutableRecord + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + data["model"] = "on-prem-model" + data["metadata"] = { + "sensitive_data_routing_applied": True, + "sensitive_data_routing_guardrail": "router-guard", + } + + failure = await _scan(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert failure == UnroutableRecord(line_number=2, custom_id="b", guardrail="router-guard") + with pytest.raises(HTTPException) as caught: + raise_public(failure) + assert "routed to a different model" in str(caught.value.detail) + + +@pytest.mark.asyncio +async def test_the_scan_spool_is_closed_when_nothing_will_read_it(): + """The spool is opened for every scan, so a clean file must not leave a temp handle behind.""" + result = await _scan_full(_jsonl(_record("a")), FakeProxyLogging()) + + assert result.changes == () + assert result.redactions.closed + + +@pytest.mark.asyncio +async def test_the_scan_spool_is_closed_when_the_upload_is_refused(): + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise RuntimeError("infrastructure is down") + + source = _jsonl(_record("a"), _record("b", content="tripwire")) + spools = [] + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + bg.tempfile.SpooledTemporaryFile = _tracking + try: + with pytest.raises(RuntimeError): + await _scan_full(source, FakeProxyLogging(_hook)) + finally: + bg.tempfile.SpooledTemporaryFile = real + + assert spools and all(handle.closed for handle in spools) + + +@pytest.mark.asyncio +async def test_the_rewrite_closes_its_own_output_when_it_cannot_finish(): + """A half-written rewrite spool has no owner yet, so it has to clean up after itself.""" + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + + source = _jsonl(_record("a"), _record("b", content="my secret is here")) + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + + spools = [] + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + def _boom(*args, **kwargs): + raise OSError("no space left on device") + + bg.tempfile.SpooledTemporaryFile = _tracking + original_read = bg._read_spooled + bg._read_spooled = _boom + try: + with pytest.raises(OSError): + rewrite_batch_input_file(source, result) + finally: + bg.tempfile.SpooledTemporaryFile = real + bg._read_spooled = original_read + + assert spools and all(handle.closed for handle in spools) + + +@pytest.mark.asyncio +async def test_the_scan_spool_is_closed_when_a_record_escapes_the_iterator(): + """A raise from inside the read loop bypasses the per-record outcome path entirely.""" + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + + spools = [] + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + bg.tempfile.SpooledTemporaryFile = _tracking + try: + with pytest.raises(json.JSONDecodeError): + await _scan_full(io.BytesIO(b"{not json at all}\n"), FakeProxyLogging()) + finally: + bg.tempfile.SpooledTemporaryFile = real + + assert spools and all(handle.closed for handle in spools) + + +@pytest.mark.asyncio +async def test_a_technical_failure_dressed_as_a_block_status_still_aborts(): + """xecguard and purview report an unreachable backend as HTTPException(400) under fail-closed.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + try: + raise ConnectionError("backend unreachable") + except ConnectionError as exc: + raise HTTPException( + status_code=400, detail={"error": "XecGuard API unreachable (block_on_error=True)"} + ) from exc + + with pytest.raises(HTTPException): + await _scan_full(_jsonl(_record("a"), _record("b", content="tripwire")), FakeProxyLogging(_hook)) + + +@pytest.mark.asyncio +async def test_a_record_body_cannot_opt_itself_out_of_the_guardrail_chain(): + """Guardrail selection reads a body-level `guardrails` key first; online it can only add.""" + seen = [] + + await _scan_full( + _jsonl({**_record("a"), "body": {**_record("a")["body"], "guardrails": []}}), + FakeProxyLogging(lambda d: seen.append(sorted(d))), + metadata={"guardrails": ["team-guard"]}, + ) + + assert seen and "guardrails" not in seen[0] + + +@pytest.mark.asyncio +async def test_a_redacted_record_keeps_its_own_guardrails_key(): + """Stripping it for the scan must not rewrite what the caller asked the provider to run.""" + record = _record("m", content="my secret is here") + record["body"]["guardrails"] = ["extra-guard"] + + body = await _rewritten_body(record, _redact_containing("secret")) + + assert body["guardrails"] == ["extra-guard"] + + +@pytest.mark.asyncio +async def test_a_400_that_is_not_a_guardrail_decision_still_aborts(): + """A guardrail's own HTTP client can raise a 400 because OUR payload was rejected, not the content.""" + from litellm.exceptions import BadRequestError + + def _hook(data): + raise BadRequestError(message="guardrail service rejected the payload", model="m", llm_provider="p") + + with pytest.raises(BadRequestError): + await _scan_full(_jsonl(_record("a")), FakeProxyLogging(_hook)) + + +async def _rewritten_body(record, hook): + """Scan one record and hand back the body as it lands in the uploaded file.""" + source = _jsonl(record) + result = await _scan_full(source, FakeProxyLogging(hook)) + rewritten = rewrite_batch_input_file(source, result) + return json.loads(rewritten.read().decode())["body"] + + +@pytest.mark.asyncio +async def test_a_redacted_record_keeps_its_own_body_metadata(): + """`metadata` is a real chat-completions parameter; redaction must not silently drop it.""" + record = _record("m", content="my secret is here") + record["body"]["metadata"] = {"team": "finance"} + + body = await _rewritten_body(record, _redact_containing("secret")) + + assert body["metadata"] == {"team": "finance"} + assert body["messages"][0]["content"] == "my *** is here" + assert "litellm_metadata" not in body + + +@pytest.mark.asyncio +async def test_a_redacted_record_keeps_its_own_litellm_metadata(): + """Tags ride in litellm_metadata; a guardrail firing must not change how the record is attributed.""" + record = _record("m", content="my secret is here") + record["body"]["litellm_metadata"] = {"tags": ["cost-center-42"]} + + body = await _rewritten_body(record, _redact_containing("secret")) + + assert body["litellm_metadata"] == {"tags": ["cost-center-42"]} + + +@pytest.mark.asyncio +async def test_a_redacted_record_keeps_an_explicitly_null_metadata(): + """An absent key and a null one are different records, so redaction must not collapse them.""" + record = _record("m", content="my secret is here") + record["body"]["metadata"] = None + + body = await _rewritten_body(record, _redact_containing("secret")) + + assert "metadata" in body and body["metadata"] is None + + +@pytest.mark.asyncio +async def test_the_log_summary_cannot_be_used_to_forge_log_lines(): + """custom_id is caller-supplied and lands in a log line, so control characters must not survive.""" + forged = "a\nWARNING: proxy shutting down" + result = await _scan_full(_jsonl(_record(forged, content="tripwire")), FakeProxyLogging(_blocking("tripwire"))) + + summary = result.summary() + + assert "\n" not in summary + assert "a WARNING: proxy shutting down" in summary + + +@pytest.mark.asyncio +async def test_the_log_summary_is_capped_so_one_upload_cannot_flood_it(): + records = [_record(f"row-{index}", content="tripwire") for index in range(60)] + result = await _scan_full(_jsonl(*records), FakeProxyLogging(_blocking("tripwire"))) + + summary = result.summary() + + assert summary.endswith("and 10 more") + assert "row-49" in summary and "row-50" not in summary + + +@pytest.mark.asyncio +async def test_the_scan_keeps_rewritten_records_off_the_heap(): + """A file whose records are mostly rewritten must not build a second copy of itself in memory.""" + import dataclasses + + bulky = "my secret is here" + ("x" * 50_000) + result = await _scan_full( + _jsonl(*(_record(str(index), content=bulky) for index in range(4))), + FakeProxyLogging(_redact_containing("secret")), + ) + + retained = sum( + len(value) + for change in result.changes + for value in (getattr(change, field.name) for field in dataclasses.fields(change)) + if isinstance(value, str) + ) + assert len(result.changes) == 4 + assert retained < 100, f"{retained} bytes of record text retained per scan" + assert result.redactions.tell() > 200_000 + + +@pytest.mark.asyncio +async def test_the_uploaded_file_is_what_the_loadbalancing_model_sniff_reads(): + """If line 1 is dropped, the router must not pick its model from a record nobody submitted.""" + dropped_first = { + "custom_id": "gone", + "url": "/v1/chat/completions", + "body": {"model": "model-a", "messages": [{"role": "user", "content": "tripwire"}]}, + } + kept = { + "custom_id": "kept", + "url": "/v1/chat/completions", + "body": {"model": "model-b", "messages": [{"role": "user", "content": "fine"}]}, + } + source = _jsonl(dropped_first, kept) + + result = await _scan_full(source, FakeProxyLogging(_blocking("tripwire"))) + rewritten = rewrite_batch_input_file(source, result) + + first_line = json.loads((rewritten.seek(0), rewritten.read().decode())[1].splitlines()[0]) + assert first_line["custom_id"] == "kept" + assert first_line["body"]["model"] == "model-b" + + +@pytest.mark.asyncio +async def test_an_infrastructure_failure_outranks_a_redaction_and_aborts(): + """A record we could not inspect must abort the upload even when an earlier record was rewritten.""" + down = HTTPException(status_code=503, detail={"error": "guardrail service unavailable"}) + + def _hook(data): + content = data["messages"][0]["content"] + if content == "raiser": + raise down + if content == "redact": + data["messages"][0]["content"] = "***" + + source = _jsonl(_record("a", content="redact"), _record("b", content="raiser")) + + with pytest.raises(HTTPException) as raised: + await _scan_full(source, FakeProxyLogging(_hook)) + + assert raised.value is down + + +@pytest.mark.asyncio +async def test_the_earliest_unscannable_record_is_the_one_reported(): + source = _jsonl( + _record("a"), + {"custom_id": "bad-1", "url": "/v1/rerank", "body": {"model": "m"}}, + {"custom_id": "bad-2", "url": "/v1/rerank", "body": {"model": "m"}}, + ) + + failure = await _scan_full(source, FakeProxyLogging()) + + assert failure == UnscannableRecord(line_number=2, custom_id="bad-1", url="/v1/rerank") + + +@pytest.mark.asyncio +async def test_a_dropped_record_names_the_guardrail_from_an_enriched_http_detail(): + """litellm stamps guardrail_name into a block's detail dict; the report should carry it through.""" + blocked = HTTPException( + status_code=400, + detail={"error": "Violated guardrail policy", "guardrail_name": "zscaler"}, + ) + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise blocked + + result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail="zscaler"),) + + +@pytest.mark.asyncio +async def test_a_dropped_record_without_a_named_guardrail_reports_none(): + """An unnamed block still drops; the report just cannot say which guardrail did it.""" + + def _hook(data): + if "tripwire" in data["messages"][0]["content"]: + raise HTTPException(status_code=400, detail="blocked") + + result = await _scan_full(_jsonl(_record("b", content="tripwire")), FakeProxyLogging(_hook)) + + assert result.changes == (RecordDropped(line_number=1, custom_id="b", guardrail=None),) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 8f96fcd796a..bf9323cdc6a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3541,8 +3541,8 @@ def _batch_upload(client_, content: bytes, purpose: str = "batch"): b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",' b'"body":{"model":"gpt-3.5-turbo","messages":[{"role":"user","content":"leak me"}]}}\n', "batch", - 400, - "A guardrail changed batch input line 1", + 200, + None, ), ( b'{"custom_id":"r-0","method":"POST","url":"/v1/chat/completions",' @@ -3602,3 +3602,159 @@ def test_batch_upload_runs_guardrails_on_each_record( finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) ProxyLogging._callback_capabilities_cache.clear() + + +def test_batch_upload_redacts_per_record(monkeypatch, llm_router: Router): + """An offending record is submitted masked, matching what the online path does per request.""" + expected_custom_ids = ["keep-1", "dirty", "keep-2"] + import json as _json + + import litellm + import litellm.proxy.openai_files_endpoints.files_endpoints as fe + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + for message in data.get("messages") or []: + if isinstance(message.get("content"), str) and "leak" in message["content"]: + message["content"] = message["content"].replace("leak", "***") + return data + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + uploaded = {} + + async def fake_route_create_file(**kwargs): + handle = kwargs["_create_file_request"]["file"][1] + uploaded["body"] = handle.read() if hasattr(handle, "read") else handle + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + + def _row(custom_id, content): + return _json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": content}]}, + } + ) + + content = ("\n".join([_row("keep-1", "fine"), _row("dirty", "please leak this"), _row("keep-2", "fine")])).encode() + try: + resp = client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + rows = [_json.loads(line) for line in uploaded["body"].decode().splitlines()] + assert [row["custom_id"] for row in rows] == expected_custom_ids + assert rows[1]["body"]["messages"][0]["content"] == "please *** this" + report = resp.json()["litellm_batch_guardrail"] + assert report["submitted_records"] == 3 + assert report["modified_records"] == [ + {"line": 2, "custom_id": "dirty", "action": "redacted", "guardrail": None} + ] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + ProxyLogging._callback_capabilities_cache.clear() + + +def test_batch_upload_closes_the_spools_it_opened(monkeypatch, llm_router: Router): + """The scan and the rewrite each open a spool; the request owns both and must not leak them.""" + import json as _json + + import litellm + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + import litellm.proxy.openai_files_endpoints.files_endpoints as fe + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + for message in data.get("messages") or []: + if "leak" in (message.get("content") or ""): + message["content"] = message["content"].replace("leak", "***") + return data + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + spools = [] + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + monkeypatch.setattr(bg.tempfile, "SpooledTemporaryFile", _tracking) + + async def fake_route_create_file(**kwargs): + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + + def _row(custom_id, content): + return _json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": content}]}, + } + ) + + content = ("\n".join([_row("keep", "fine"), _row("dirty", "please leak this")])).encode() + try: + resp = client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + assert len(spools) == 2, f"expected a scan spool and a rewrite spool, saw {len(spools)}" + assert all(handle.closed for handle in spools), "the request must close every spool it opened" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + ProxyLogging._callback_capabilities_cache.clear() From 787edb123f85a5f4b7a86b6b3068d494b655013d Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 20 Aug 2026 13:29:43 -0700 Subject: [PATCH 8/9] refactor(ui): mark dark as beta in the theme menu, not the toolbar (#37680) The Experimental badge sat in the top bar next to the icon, which read as if the whole theme control were experimental and cost toolbar width for a caveat that only applies once. It moves into the menu as a Beta tag on the Dark entry, where it labels exactly the choice it is about and is visible before the choice is made rather than only after. --- .../ThemeToggle/ThemeToggle.test.tsx | 25 ++++--- .../components/ThemeToggle/ThemeToggle.tsx | 65 +++++++++---------- 2 files changed, 46 insertions(+), 44 deletions(-) diff --git a/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.test.tsx b/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.test.tsx index 3f7af94bfa0..41efedcd43a 100644 --- a/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.test.tsx +++ b/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.test.tsx @@ -16,7 +16,7 @@ const openMenu = async () => { await screen.findByRole("menu"); }; -const pick = async (label: string) => userEvent.click(screen.getByRole("menuitemradio", { name: label })); +const pick = async (label: string | RegExp) => userEvent.click(screen.getByRole("menuitemradio", { name: label })); beforeEach(() => { localStorage.clear(); @@ -33,7 +33,7 @@ describe("ThemeToggle", () => { await openMenu(); expect(screen.getByRole("menuitemradio", { name: "Light" })).toBeChecked(); - expect(screen.getByRole("menuitemradio", { name: "Dark" })).not.toBeChecked(); + expect(screen.getByRole("menuitemradio", { name: /^Dark/ })).not.toBeChecked(); expect(screen.getByRole("menuitemradio", { name: "System" })).not.toBeChecked(); }); @@ -41,7 +41,7 @@ describe("ThemeToggle", () => { renderToggle(); await openMenu(); - await pick("Dark"); + await pick(/^Dark/); expect(document.documentElement).toHaveClass("dark"); expect(localStorage.getItem("theme")).toBe("dark"); @@ -50,7 +50,7 @@ describe("ThemeToggle", () => { it("hands control back to the system preference when asked", async () => { renderToggle(); await openMenu(); - await pick("Dark"); + await pick(/^Dark/); await pick("System"); @@ -58,15 +58,20 @@ describe("ThemeToggle", () => { expect(document.documentElement).not.toHaveClass("dark"); }); - it("flags dark mode as experimental, and only while it is on", async () => { + it("marks dark as beta in the menu, and leaves the other choices unmarked", async () => { renderToggle(); await openMenu(); - expect(screen.queryByText("Experimental")).not.toBeInTheDocument(); - await pick("Dark"); - expect(screen.getByText("Experimental")).toBeInTheDocument(); + expect(screen.getByRole("menuitemradio", { name: /^Dark/ })).toHaveTextContent("Beta"); + expect(screen.getByRole("menuitemradio", { name: "Light" })).not.toHaveTextContent("Beta"); + expect(screen.getByRole("menuitemradio", { name: "System" })).not.toHaveTextContent("Beta"); + }); - await pick("Light"); - expect(screen.queryByText("Experimental")).not.toBeInTheDocument(); + it("keeps the beta marker inside the menu rather than in the toolbar", async () => { + renderToggle(); + await openMenu(); + await pick(/^Dark/); + + expect(screen.getByRole("button", { name: "Theme" })).not.toHaveTextContent("Beta"); }); }); diff --git a/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.tsx b/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.tsx index 6a006141523..3fbb3d2eb5b 100644 --- a/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.tsx +++ b/ui/litellm-dashboard/src/components/ThemeToggle/ThemeToggle.tsx @@ -15,46 +15,43 @@ import { } from "@/components/ui/dropdown-menu"; const THEMES = [ - { value: "system", label: "System", Icon: Monitor }, - { value: "light", label: "Light", Icon: Sun }, - { value: "dark", label: "Dark", Icon: Moon }, + { value: "system", label: "System", Icon: Monitor, beta: false }, + { value: "light", label: "Light", Icon: Sun, beta: false }, + { value: "dark", label: "Dark", Icon: Moon, beta: true }, ] as const; const ThemeToggle: React.FC = () => { const { theme, setTheme, resolvedTheme } = useTheme(); - const isDark = resolvedTheme === "dark"; return ( - - {isDark && ( - - Experimental - - )} - - - } - > - {isDark ? : } - - - - {THEMES.map(({ value, label, Icon }) => ( - - - {label} - - ))} - - - - + + + } + > + {resolvedTheme === "dark" ? : } + + + + {THEMES.map(({ value, label, Icon, beta }) => ( + + + {label} + {beta && ( + + Beta + + )} + + ))} + + + ); }; From 4af59d7c6e38f5b3fb54f70eb06f1a5726e45410 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 20 Aug 2026 13:30:34 -0700 Subject: [PATCH 9/9] ci: lint the test tree for undefined names and fix all 30 (#37671) ruff.toml excludes tests/* from `ruff check`, so nothing has ever checked the test tree for names that do not exist. That matters more in tests than in product code: a NameError inside a test whose body is wrapped in `except Exception: pass` is swallowed, and the test reports green forever. Adds ruff-tests.toml selecting F821 alone, wired into the lint workflow and `make lint-ruff`, and clears every existing violation: - 4 tests interpolated an unbound `e` into a `pytest.fail` message reached only on the failure path, so the NameError, not the assertion, is what ran. test_llm_guard_error_raising is the worst: it passes today with content safety disabled entirely. It now asserts the 400 and its detail body. - 5 sites construct BaseExceptionGroup, a 3.11 builtin, in a tree that still supports 3.10. Guarded behind the exceptiongroup backport that anyio already pulls in below 3.11. - 9 missing imports (json, openai, Any, Final, HTTPException), including one in a helper that catches HTTPException by a name it never imported, so the challenge path it exists to detect raises NameError instead. - 5 annotations naming types imported inside the function body, hoisted to module scope or TYPE_CHECKING. - 2 blocks of dead code: everything after a pytest.fail in test_claude_agent_sdk, and an unused helper in test_end_users calling a function defined in a different module. - 1 error-path f-string in the router-settings doc test that masked the real FileNotFoundError behind a NameError. Only F821 for now. Widening the select list means ratcheting thousands of pre-existing findings, so rules go in one at a time with their violations already fixed. --- .github/workflows/test-linting.yml | 5 ++ Makefile | 1 + ruff-tests.toml | 15 +++++ .../test_router_settings.py | 2 +- tests/local_testing/test_completion.py | 15 +++-- tests/local_testing/test_exceptions.py | 2 +- tests/local_testing/test_llm_guard.py | 10 ++-- .../test_claude_agent_sdk.py | 59 ------------------- .../test_custom_callback_input.py | 1 + tests/test_end_users.py | 41 ------------- .../llms/oci/test_oci_coverage_boost.py | 5 ++ .../mcp_server/faults/test_list_outcomes.py | 5 ++ .../mcp_server/faults/test_traversal.py | 5 ++ .../mcp_server/test_discoverable_endpoints.py | 6 ++ .../mcp_server/test_mcp_stale_session.py | 2 + .../mcp_server/test_rest_endpoints.py | 4 ++ .../mcp_server/test_semantic_tool_filter.py | 3 + .../proxy/auth/test_auth_checks.py | 9 +-- .../guardrails/test_pillar_guardrails.py | 2 +- .../proxy/test_common_request_processing.py | 2 +- 20 files changed, 72 insertions(+), 122 deletions(-) create mode 100644 ruff-tests.toml diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 5a180c13c53..f98077ea2f0 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -122,6 +122,11 @@ jobs: uv run --no-sync ruff check . cd .. + - name: Run Ruff linting (test tree) + if: steps.changes.outputs.decision != 'skip' + run: | + uv run --no-sync ruff check --config ruff-tests.toml tests + - name: Check strict-rule budget (delta vs base) if: steps.changes.outputs.decision != 'skip' run: | diff --git a/Makefile b/Makefile index 5ae2638fbaa..580d663ba53 100644 --- a/Makefile +++ b/Makefile @@ -160,6 +160,7 @@ lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) # Linting targets lint-ruff: $(LINT_DEP_INSTALL) cd litellm && $(UV_RUN) ruff check . && cd .. + $(UV_RUN) ruff check --config ruff-tests.toml tests # faster linter for developing ... # inspiration from: diff --git a/ruff-tests.toml b/ruff-tests.toml new file mode 100644 index 00000000000..c1bdcc755a7 --- /dev/null +++ b/ruff-tests.toml @@ -0,0 +1,15 @@ +# Lint config for the test tree, which ruff.toml excludes from `ruff check`. +# +# Deliberately one rule. F821 is the cheapest guard against a test that cannot fail: +# a name that does not exist raises NameError, and a test whose body is wrapped in +# `except Exception: pass` swallows that NameError and reports green. Widening this +# select list means ratcheting thousands of pre-existing findings, so new rules go in +# one at a time, each with its violations already fixed. +# +# No target-version here on purpose: it resolves from requires-python (>=3.10), so +# 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that +# still has to run on 3.10. + +line-length = 120 + +lint.select = ["F821"] diff --git a/tests/documentation_tests/test_router_settings.py b/tests/documentation_tests/test_router_settings.py index 290aa283af4..a1b6f1dac1d 100644 --- a/tests/documentation_tests/test_router_settings.py +++ b/tests/documentation_tests/test_router_settings.py @@ -61,7 +61,7 @@ try: documented_keys.update(doc_key_pattern.findall(table_content)) except Exception as e: raise Exception( - f"Error reading documentation: {e}, \n repo base - {os.listdir(repo_base)}" + f"Error reading documentation: {e}, \n repo base - {os.listdir(_repo_root)}" ) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index 6f58bb2eb35..01fd35cb42d 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -842,6 +842,8 @@ def test_completion_mistral_api_modified_input(): @pytest.mark.skip(reason="this test is flaky") def test_completion_gpt4_vision(): + import openai + try: litellm.set_verbose = True response = completion( @@ -1820,6 +1822,8 @@ def test_completion_openai_litellm_key(): @pytest.mark.skip(reason="Unresponsive endpoint.[TODO] Rehost this somewhere else") def test_completion_ollama_hosted(): + import openai + try: litellm.request_timeout = 20 # give ollama 20 seconds to response litellm.set_verbose = True @@ -2057,17 +2061,12 @@ def test_completion_openrouter_reasoning_effort(): def test_completion_hf_model_no_provider(): - try: - response = completion( + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + completion( model="WizardLM/WizardLM-70B-V1.0", messages=messages, max_tokens=5, ) - # Add any assertions here to check the response - print(response) - pytest.fail(f"Error occurred: {e}") - except Exception as e: - pass # test_completion_hf_model_no_provider() @@ -2546,7 +2545,7 @@ def test_completion_replicate_vicuna(): response_str = response["choices"][0]["message"]["content"] print("RESPONSE STRING\n", response_str) if type(response_str) != str: - pytest.fail(f"Error occurred: {e}") + pytest.fail(f"Expected a string response, got {type(response_str)}: {response_str}") except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index e02d9e21171..8c1df52e28e 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -573,7 +573,7 @@ def test_content_policy_violation_error_streaming(): num_finish_reason += 1 print("finish_reason", chunk["choices"][0].get("finish_reason")) - pytest.fail(f"Expected to return 400 error In streaming{e}") + pytest.fail("Expected a content-policy error in streaming, got a clean stream") except Exception as e: pass diff --git a/tests/local_testing/test_llm_guard.py b/tests/local_testing/test_llm_guard.py index 78bbd1c0af8..86fa80ee944 100644 --- a/tests/local_testing/test_llm_guard.py +++ b/tests/local_testing/test_llm_guard.py @@ -15,6 +15,8 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest +from fastapi import HTTPException + import litellm from litellm_enterprise.enterprise_callbacks.llm_guard import _ENTERPRISE_LLMGuard from litellm import Router, mock_completion @@ -128,7 +130,7 @@ async def test_llm_guard_error_raising(): user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) local_cache = DualCache() - try: + with pytest.raises(HTTPException) as exc_info: await llm_guard.async_moderation_hook( data={ "messages": [ @@ -141,9 +143,9 @@ async def test_llm_guard_error_raising(): user_api_key_dict=user_api_key_dict, call_type="completion", ) - pytest.fail(f"Should have failed - {str(e)}") - except Exception as e: - pass + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == {"error": "Violated content safety policy"} def test_llm_guard_key_specific_mode(): diff --git a/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py b/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py index 48eb7d85ec1..c1339ce6280 100644 --- a/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py +++ b/tests/proxy_e2e_anthropic_messages_tests/test_claude_agent_sdk.py @@ -148,65 +148,6 @@ async def test_claude_agent_sdk_streaming( f"Test failed for {model_name} ({model_description}) after {MAX_RETRIES} attempts: {last_error}" ) - # Test query - test_query = "Say 'Hello from LiteLLM!' and nothing else." - - # Track streaming - received_chunks = [] - full_response = "" - - try: - async with ClaudeSDKClient(options=options) as client: - await client.query(test_query) - - # Collect streaming response - async for msg in client.receive_response(): - # Handle different message types - if hasattr(msg, "type"): - if msg.type == "content_block_delta": - # Streaming text delta - if hasattr(msg, "delta") and hasattr(msg.delta, "text"): - chunk_text = msg.delta.text - received_chunks.append(chunk_text) - full_response += chunk_text - elif msg.type == "content_block_start": - # Start of content block - if hasattr(msg, "content_block") and hasattr( - msg.content_block, "text" - ): - chunk_text = msg.content_block.text - received_chunks.append(chunk_text) - full_response += chunk_text - - # Fallback to content handling - if hasattr(msg, "content"): - for content_block in msg.content: - if hasattr(content_block, "text"): - chunk_text = content_block.text - received_chunks.append(chunk_text) - full_response += chunk_text - - # Assertions - print(f"\nāœ… Received {len(received_chunks)} chunks") - print(f"šŸ“ Full response: {full_response[:100]}...") - - # Verify we got a response - assert len(full_response) > 0, f"No response received from {model_name}" - - # Verify streaming (should have multiple chunks for most responses) - # Note: Very short responses might come in 1 chunk, so we just verify we got content - assert len(received_chunks) > 0, f"No chunks received from {model_name}" - - # Verify response is non-empty (don't assert on specific LLM content ā€” it's non-deterministic) - assert ( - len(full_response.strip()) > 0 - ), f"Empty response received from {model_name}" - - print(f"āœ… Test passed for {model_name}") - - except Exception as e: - pytest.fail(f"Test failed for {model_name} ({model_description}): {str(e)}") - if __name__ == "__main__": # Run tests diff --git a/tests/proxy_unit_tests/test_custom_callback_input.py b/tests/proxy_unit_tests/test_custom_callback_input.py index 71a7e94b180..a032b8706bc 100644 --- a/tests/proxy_unit_tests/test_custom_callback_input.py +++ b/tests/proxy_unit_tests/test_custom_callback_input.py @@ -2,6 +2,7 @@ ## This test asserts the type of data passed into each method of the custom callback handler import asyncio import inspect +import json import os import sys import time diff --git a/tests/test_end_users.py b/tests/test_end_users.py index ff3cc4ec94b..bc1fcbb662d 100644 --- a/tests/test_end_users.py +++ b/tests/test_end_users.py @@ -14,47 +14,6 @@ from typing import Optional """ -async def chat_completion_with_headers(session, key, model="gpt-4"): - url = "http://0.0.0.0:4000/chat/completions" - headers = { - "Authorization": f"Bearer {key}", - "Content-Type": "application/json", - } - data = { - "model": model, - "messages": [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello!"}, - ], - } - - async with session.post(url, headers=headers, json=data) as response: - status = response.status - response_text = await response.text() - - print(response_text) - print() - - if status != 200: - raise Exception(f"Request did not return a 200 status code: {status}") - - response_header_check( - response - ) # calling the function to check response headers - - raw_headers = response.raw_headers - raw_headers_json = {} - - for ( - item - ) in ( - response.raw_headers - ): # ((b'date', b'Fri, 19 Apr 2024 21:17:29 GMT'), (), ) - raw_headers_json[item[0].decode("utf-8")] = item[1].decode("utf-8") - - return raw_headers_json - - async def generate_key( session, i, diff --git a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py b/tests/test_litellm/llms/oci/test_oci_coverage_boost.py index 0b7afa3775d..7c91ece70b5 100644 --- a/tests/test_litellm/llms/oci/test_oci_coverage_boost.py +++ b/tests/test_litellm/llms/oci/test_oci_coverage_boost.py @@ -11,11 +11,16 @@ All tests are self-contained and require no real OCI credentials or network acce """ import json +from typing import TYPE_CHECKING + import pytest from unittest.mock import patch, MagicMock, AsyncMock import httpx +if TYPE_CHECKING: + from litellm.llms.oci.chat.transformation import OCIStreamWrapper + from litellm import ModelResponse from litellm.llms.oci.chat.cohere import ( _extract_text_content, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py index 64afa52ab55..65e2faee1b2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_list_outcomes.py @@ -2,6 +2,11 @@ to exactly one category, wire values never carry upstream prose, and single-upstream HTTP statuses stay truthful to who failed.""" +import sys + +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + import httpx import pytest from mcp import McpError diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py index a12c02339e6..b8bf4da1dc4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/faults/test_traversal.py @@ -2,6 +2,11 @@ links win (the ``raise ... from`` cause subtree, then ExceptionGroup members in raise order, then the incidental ``__context__`` chain last), and adversarial shapes terminate.""" +import sys + +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + from litellm.proxy._experimental.mcp_server.faults import iter_exception_tree diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index bdaf1458fb0..34852850de6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4,6 +4,7 @@ import hashlib import json import time from base64 import urlsafe_b64encode +from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -11,6 +12,11 @@ from fastapi import HTTPException from litellm.types.mcp import MCPAuth +if TYPE_CHECKING: + import httpx + + from litellm.types.mcp_server.mcp_server_manager import MCPServer + # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index 7bdd3b36763..feab179570b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -1411,6 +1411,8 @@ async def _run_passthrough_connect( ): """Drive handle_streamable_http_mcp through the preemptive-401 gate and report whether it challenged (raised) or forwarded to the session manager. Returns (challenged, www_authenticate).""" + from fastapi import HTTPException + from litellm.proxy._experimental.mcp_server.server import ( handle_streamable_http_mcp, session_manager_stateless, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 7ba9f463197..d7fb121ef9b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -1,9 +1,13 @@ import asyncio import json +import sys from datetime import datetime from typing import Any, Dict, Optional from unittest.mock import AsyncMock, MagicMock +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + import httpx import pytest from fastapi import HTTPException diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index 83e4dcf5677..da7f43c7118 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -12,6 +12,9 @@ from unittest.mock import AsyncMock, Mock, patch import pytest +if sys.version_info < (3, 11): # BaseExceptionGroup is a builtin only from 3.11 + from exceptiongroup import BaseExceptionGroup + sys.path.insert(0, os.path.abspath("../..")) from mcp.types import Tool as MCPTool diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 1d1bd9ebf8a..7a3288dec37 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -13,7 +13,7 @@ from datetime import datetime, timedelta, timezone import httpx import pytest -from fastapi import status +from fastapi import Request, status import litellm from litellm.proxy._types import ( @@ -2760,11 +2760,9 @@ async def test_common_checks_metadata_route_keeps_key_tags_out_of_provider_metad assert "metadata" not in request_body -def _pass_through_request() -> "Request": +def _pass_through_request() -> Request: """A Request whose FastAPI-resolved endpoint carries the pass-through marker, i.e. the request was dispatched to a user-defined pass-through handler.""" - from fastapi import Request - from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) @@ -2776,10 +2774,9 @@ def _pass_through_request() -> "Request": return Request(scope={"type": "http", "headers": [], "endpoint": pass_through_endpoint}) -def _builtin_request() -> "Request": +def _builtin_request() -> Request: """A Request dispatched to a built-in (non-pass-through) handler, e.g. what a custom path colliding with a core route actually resolves to.""" - from fastapi import Request def chat_completions(): ... diff --git a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py index c977223bfab..48f6b3ba2b9 100644 --- a/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py @@ -9,7 +9,7 @@ and following LiteLLM testing patterns and best practices. import importlib import os import sys -from typing import Dict +from typing import Any, Dict from unittest.mock import Mock, patch # Add parent directory to path for imports diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 355c6d27eb2..510fb977a61 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3,7 +3,7 @@ import copy import datetime import json from types import SimpleNamespace -from typing import AsyncGenerator, Callable, Optional +from typing import AsyncGenerator, Callable, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx