diff --git a/.circleci/config.yml b/.circleci/config.yml index 32983c92e27..38fdaf3609d 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -158,6 +158,8 @@ jobs: CHOCOLATEY_CONFIRM_ALL: "true" - run: name: Install Dependencies + environment: + UV_HTTP_TIMEOUT: "300" command: | $installer = Join-Path $env:TEMP "uv-install.ps1" Invoke-WebRequest -Uri https://astral.sh/uv/0.10.9/install.ps1 -OutFile $installer @@ -2475,10 +2477,15 @@ jobs: DISABLE_SCHEMA_UPDATE: "true" SERVER_ROOT_PATH: "" PROXY_LOGOUT_URL: "" + # LITELLM_LICENSE is forwarded from the project env so premium-gated + # UI flows can be exercised. license.spec.ts asserts the resulting + # JWT carries premium_user=true; if it ever stops being passed, that + # test fails loudly rather than silently regressing premium coverage. command: | - uv run --no-sync python -m litellm.proxy.proxy_cli \ - --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ - --port 4000 + LITELLM_LICENSE="$LITELLM_LICENSE" \ + uv run --no-sync python -m litellm.proxy.proxy_cli \ + --config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \ + --port 4000 background: true - run: name: Wait for proxy to be ready @@ -2495,9 +2502,12 @@ jobs: exit 1 - run: name: Run Playwright E2E tests + # Forward LITELLM_LICENSE so license.spec.ts can detect that the + # proxy was launched with a license and assert premium_user=true. command: | cd ui/litellm-dashboard - npx playwright test --config e2e_tests/playwright.config.ts + LITELLM_LICENSE="$LITELLM_LICENSE" \ + npx playwright test --config e2e_tests/playwright.config.ts no_output_timeout: 10m - store_artifacts: path: ui/litellm-dashboard/test-results diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 49a36aa23f0..2d4e85630dc 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -218,6 +218,7 @@ jobs: tests/proxy_unit_tests/test_gemini_agents_endpoints.py tests/proxy_unit_tests/test_get_favicon.py tests/proxy_unit_tests/test_get_image.py + tests/proxy_unit_tests/test_reducto_ocr_route.py tests/proxy_unit_tests/test_ui_path_detection.py tests/proxy_unit_tests/test_prompt_test_endpoint.py tests/proxy_unit_tests/test_check_batch_cost.py diff --git a/.github/workflows/test-unit-proxy-mgmt-behavior.yml b/.github/workflows/test-unit-proxy-mgmt-behavior.yml new file mode 100644 index 00000000000..e73997323a4 --- /dev/null +++ b/.github/workflows/test-unit-proxy-mgmt-behavior.yml @@ -0,0 +1,34 @@ +name: "Unit Tests: Proxy Management-Endpoint Behavior Pinning" + +on: + pull_request: + branches: + - main + - litellm_internal_staging + - litellm_oss_branch + - "litellm_**" + +permissions: + contents: read + id-token: write + pull-requests: write + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + proxy-mgmt-behavior: + uses: ./.github/workflows/_test-unit-services-base.yml + with: + test-path: tests/proxy_behavior + # workers=0 (no xdist): the world seed is a single shared Postgres + # state — two xdist workers both call seed_world() and race on the + # ``behavior-pin-budget`` row, producing UniqueViolation + cascading + # missing-membership FK failures. The whole suite is ~7s sequentially, + # so the cost of disabling parallelism here is negligible. + workers: 0 + reruns: 0 + enable-postgres: true + artifact-name: proxy-mgmt-behavior + timeout-minutes: 15 diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 4de4a55981d..2729babb6d6 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -24,7 +24,8 @@ RUN for i in 1 2 3; do \ curl \ openssl \ libsndfile \ - nodejs && break || sleep 5; \ + nodejs \ + npm && break || sleep 5; \ done ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ diff --git a/license_cache.json b/license_cache.json index 803db8fdd1b..dc061b48f4f 100644 --- a/license_cache.json +++ b/license_cache.json @@ -50,6 +50,7 @@ "h11:0.16.0": "MIT", "requests-toolbelt:1.0.0": "Apache 2.0", "tornado:6.5.4": "Apache-2.0", + "granian:2.5.7": "BSD-3-Clause", "mlflow:3.11.1": "Copyright 2018 Databricks, Inc. All rights reserved.\n \n \t\t\t\tApache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n \n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n \n 1. Definitions.\n \n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n \n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n \n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n \n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n \n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n \n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n \n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n \n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n \n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n \n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n \n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n \n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n \n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n \n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n \n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n \n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n \n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n \n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n \n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n \n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n \n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n \n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n \n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n \n END OF TERMS AND CONDITIONS\n APPENDIX: How to apply the Apache License to your work.\n \n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n \n Copyright [yyyy] [name of copyright owner]\n \n Licensed under the Apache License, Version 2.0 (the \"License\");\n you may not use this file except in compliance with the License.\n You may obtain a copy of the License at\n \n http://www.apache.org/licenses/LICENSE-2.0\n \n Unless required by applicable law or agreed to in writing, software\n distributed under the License is distributed on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n See the License for the specific language governing permissions and\n limitations under the License.\n ", "nvidia-riva-client:2.15.0": "MIT", "numpy:1.26.0": "Copyright (c) 2005-2023, NumPy Developers. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. * Neither the name of the NumPy Developers nor the names of any contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ---- This binary distribution of NumPy also bundles the following software: Name: GCC runtime library Files: .dylibs/* Description: dynamically linked to files compiled with gcc Availability: https://gcc.gnu.org/viewcvs/gcc/ License: GPLv3 + runtime exception Copyright (C) 2002-2017 Free Software Foundation, Inc. Libgfortran is free software; you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation; either version 3, or (at your option) any later version. Libgfortran is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. Under Section 7 of GPL version 3, you are granted additional permissions described in the GCC Runtime Library Exception, version 3.1, as published by the Free Software Foundation. You should have received a copy of the GNU General Public License and a copy of the GCC Runtime Library Exception along with this program; see the files COPYING3 and COPYING.RUNTIME respectively. If not, see . ---- Full text of license texts referred to above follows (that they are listed below does not necessarily imply the conditions apply to the present binary release): ---- GCC RUNTIME LIBRARY EXCEPTION Version 3.1, 31 March 2009 Copyright (C) 2009 Free Software Foundation, Inc. Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed. This GCC Runtime Library Exception (\"Exception\") is an additional permission under section 7 of the GNU General Public License, version 3 (\"GPLv3\"). It applies to a given file (the \"Runtime Library\") that bears a notice placed by the copyright holder of the file stating that the file is governed by GPLv3 along with this Exception. When you use GCC to compile a program, GCC may combine portions of certain GCC header files and runtime libraries with the compiled program. The purpose of this Exception is to allow compilation of non-GPL (including proprietary) programs to use, in this way, the header files and runtime libraries covered by this Exception. 0. Definitions. A file is an \"Independent Module\" if it either requires the Runtime Library for execution after a Compilation Process, or makes use of an interface provided by the Runtime Library, but is not otherwise based on the Runtime Library. \"GCC\" means a version of the GNU Compiler Collection, with or without modifications, governed by version 3 (or a specified later version) of the GNU General Public License (GPL) with the option of using any subsequent versions published by the FSF. \"GPL-compatible Software\" is software whose conditions of propagation, modification and use would permit combination with GCC in accord with the license of GCC. \"Target Code\" refers to output from any compiler for a real or virtual target processor architecture, in executable form or suitable for input to an assembler, loader, linker and/or execution phase. Notwithstanding that, Target Code does not include data in any format that is used as a compiler intermediate representation, or used for producing a compiler intermediate representation. The \"Compilation Process\" transforms code entirely represented in non-intermediate languages designed for human-written code, and/or in Java Virtual Machine byte code, into Target Code. Thus, for example, use of source code generators and preprocessors need not be considered part of the Compilation Process, since the Compilation Process can be understood as starting with the output of the generators or preprocessors. A Compilation Process is \"Eligible\" if it is done using GCC, alone or with other GPL-compatible software, or if it is done without using any work based on GCC. For example, using non-GPL-compatible Software to optimize any GCC intermediate representations would not qualify as an Eligible Compilation Process. 1. Grant of Additional Permission. You have permission to propagate a work of Target Code formed by combining the Runtime Library with Independent Modules, even if such propagation would otherwise violate the terms of GPLv3, provided that all Target Code was generated by Eligible Compilation Processes. You may then convey such a combination under terms of your choice, consistent with the licensing of the Independent Modules. 2. No Weakening of GCC Copyleft. The availability of this Exception does not imply any general presumption that third-party software is unaffected by the copyleft requirements of the license of GCC. ---- GNU GENERAL PUBLIC LICENSE Version 3, 29 June 2007 Copyright (C) 2007 Free Software Foundation, Inc. Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed. Preamble The GNU General Public License is a free, copyleft license for software and other kinds of works. The licenses for most software and other practical works are designed to take away your freedom to share and change the works. By contrast, the GNU General Public License is intended to guarantee your freedom to share and change all versions of a program--to make sure it remains free software for all its users. We, the Free Software Foundation, use the GNU General Public License for most of our software; it applies also to any other work released this way by its authors. You can apply it to your programs, too. When we speak of free software, we are referring to freedom, not price. Our General Public Licenses are designed to make sure that you have the freedom to distribute copies of free software (and charge for them if you wish), that you receive source code or can get it if you want it, that you can change the software or use pieces of it in new free programs, and that you know you can do these things. To protect your rights, we need to prevent others from denying you these rights or asking you to surrender the rights. Therefore, you have certain responsibilities if you distribute copies of the software, or if you modify it: responsibilities to respect the freedom of others. For example, if you distribute copies of such a program, whether gratis or for a fee, you must pass on to the recipients the same freedoms that you received. You must make sure that they, too, receive or can get the source code. And you must show them these terms so they know their rights. Developers that use the GNU GPL protect your rights with two steps: (1) assert copyright on the software, and (2) offer you this License giving you legal permission to copy, distribute and/or modify it. For the developers' and authors' protection, the GPL clearly explains that there is no warranty for this free software. For both users' and authors' sake, the GPL requires that modified versions be marked as changed, so that their problems will not be attributed erroneously to authors of previous versions. Some devices are designed to deny users access to install or run modified versions of the software inside them, although the manufacturer can do so. This is fundamentally incompatible with the aim of protecting users' freedom to change the software. The systematic pattern of such abuse occurs in the area of products for individuals to use, which is precisely where it is most unacceptable. Therefore, we have designed this version of the GPL to prohibit the practice for those products. If such problems arise substantially in other domains, we stand ready to extend this provision to those domains in future versions of the GPL, as needed to protect the freedom of users. Finally, every program is threatened constantly by software patents. States should not allow patents to restrict development and use of software on general-purpose computers, but in those that do, we wish to avoid the special danger that patents applied to a free program could make it effectively proprietary. To prevent this, the GPL assures that patents cannot be used to render the program non-free. The precise terms and conditions for copying, distribution and modification follow. TERMS AND CONDITIONS 0. Definitions. \"This License\" refers to version 3 of the GNU General Public License. \"Copyright\" also means copyright-like laws that apply to other kinds of works, such as semiconductor masks. \"The Program\" refers to any copyrightable work licensed under this License. Each licensee is addressed as \"you\". \"Licensees\" and \"recipients\" may be individuals or organizations. To \"modify\" a work means to copy from or adapt all or part of the work in a fashion requiring copyright permission, other than the making of an exact copy. The resulting work is called a \"modified version\" of the earlier work or a work \"based on\" the earlier work. A \"covered work\" means either the unmodified Program or a work based on the Program. To \"propagate\" a work means to do anything with it that, without permission, would make you directly or secondarily liable for infringement under applicable copyright law, except executing it on a computer or modifying a private copy. Propagation includes copying, distribution (with or without modification), making available to the public, and in some countries other activities as well. To \"convey\" a work means any kind of propagation that enables other parties to make or receive copies. Mere interaction with a user through a computer network, with no transfer of a copy, is not conveying. An interactive user interface displays \"Appropriate Legal Notices\" to the extent that it includes a convenient and prominently visible feature that (1) displays an appropriate copyright notice, and (2) tells the user that there is no warranty for the work (except to the extent that warranties are provided), that licensees may convey the work under this License, and how to view a copy of this License. If the interface presents a list of user commands or options, such as a menu, a prominent item in the list meets this criterion. 1. Source Code. The \"source code\" for a work means the preferred form of the work for making modifications to it. \"Object code\" means any non-source form of a work. A \"Standard Interface\" means an interface that either is an official standard defined by a recognized standards body, or, in the case of interfaces specified for a particular programming language, one that is widely used among developers working in that language. The \"System Libraries\" of an executable work include anything, other than the work as a whole, that (a) is included in the normal form of packaging a Major Component, but which is not part of that Major Component, and (b) serves only to enable use of the work with that Major Component, or to implement a Standard Interface for which an implementation is available to the public in source code form. A \"Major Component\", in this context, means a major essential component (kernel, window system, and so on) of the specific operating system (if any) on which the executable work runs, or a compiler used to produce the work, or an object code interpreter used to run it. The \"Corresponding Source\" for a work in object code form means all the source code needed to generate, install, and (for an executable work) run the object code and to modify the work, including scripts to control those activities. However, it does not include the work's System Libraries, or general-purpose tools or generally available free programs which are used unmodified in performing those activities but which are not part of the work. For example, Corresponding Source includes interface definition files associated with source files for the work, and the source code for shared libraries and dynamically linked subprograms that the work is specifically designed to require, such as by intimate data communication or control flow between those subprograms and other parts of the work. The Corresponding Source need not include anything that users can regenerate automatically from other parts of the Corresponding Source. The Corresponding Source for a work in source code form is that same work. 2. Basic Permissions. All rights granted under this License are granted for the term of copyright on the Program, and are irrevocable provided the stated conditions are met. This License explicitly affirms your unlimited permission to run the unmodified Program. The output from running a covered work is covered by this License only if the output, given its content, constitutes a covered work. This License acknowledges your rights of fair use or other equivalent, as provided by copyright law. You may make, run and propagate covered works that you do not convey, without conditions so long as your license otherwise remains in force. You may convey covered works to others for the sole purpose of having them make modifications exclusively for you, or provide you with facilities for running those works, provided that you comply with the terms of this License in conveying all material for which you do not control copyright. Those thus making or running the covered works for you must do so exclusively on your behalf, under your direction and control, on terms that prohibit them from making any copies of your copyrighted material outside their relationship with you. Conveying under any other circumstances is permitted solely under the conditions stated below. Sublicensing is not allowed; section 10 makes it unnecessary. 3. Protecting Users' Legal Rights From Anti-Circumvention Law. No covered work shall be deemed part of an effective technological measure under any applicable law fulfilling obligations under article 11 of the WIPO copyright treaty adopted on 20 December 1996, or similar laws prohibiting or restricting circumvention of such measures. When you convey a covered work, you waive any legal power to forbid circumvention of technological measures to the extent such circumvention is effected by exercising rights under this License with respect to the covered work, and you disclaim any intention to limit operation or modification of the work as a means of enforcing, against the work's users, your or third parties' legal rights to forbid circumvention of technological measures. 4. Conveying Verbatim Copies. You may convey verbatim copies of the Program's source code as you receive it, in any medium, provided that you conspicuously and appropriately publish on each copy an appropriate copyright notice; keep intact all notices stating that this License and any non-permissive terms added in accord with section 7 apply to the code; keep intact all notices of the absence of any warranty; and give all recipients a copy of this License along with the Program. You may charge any price or no price for each copy that you convey, and you may offer support or warranty protection for a fee. 5. Conveying Modified Source Versions. You may convey a work based on the Program, or the modifications to produce it from the Program, in the form of source code under the terms of section 4, provided that you also meet all of these conditions: a) The work must carry prominent notices stating that you modified it, and giving a relevant date. b) The work must carry prominent notices stating that it is released under this License and any conditions added under section 7. This requirement modifies the requirement in section 4 to \"keep intact all notices\". c) You must license the entire work, as a whole, under this License to anyone who comes into possession of a copy. This License will therefore apply, along with any applicable section 7 additional terms, to the whole of the work, and all its parts, regardless of how they are packaged. This License gives no permission to license the work in any other way, but it does not invalidate such permission if you have separately received it. d) If the work has interactive user interfaces, each must display Appropriate Legal Notices; however, if the Program has interactive interfaces that do not display Appropriate Legal Notices, your work need not make them do so. A compilation of a covered work with other separate and independent works, which are not by their nature extensions of the covered work, and which are not combined with it such as to form a larger program, in or on a volume of a storage or distribution medium, is called an \"aggregate\" if the compilation and its resulting copyright are not used to limit the access or legal rights of the compilation's users beyond what the individual works permit. Inclusion of a covered work in an aggregate does not cause this License to apply to the other parts of the aggregate. 6. Conveying Non-Source Forms. You may convey a covered work in object code form under the terms of sections 4 and 5, provided that you also convey the machine-readable Corresponding Source under the terms of this License, in one of these ways: a) Convey the object code in, or embodied in, a physical product (including a physical distribution medium), accompanied by the Corresponding Source fixed on a durable physical medium customarily used for software interchange. b) Convey the object code in, or embodied in, a physical product (including a physical distribution medium), accompanied by a written offer, valid for at least three years and valid for as long as you offer spare parts or customer support for that product model, to give anyone who possesses the object code either (1) a copy of the Corresponding Source for all the software in the product that is covered by this License, on a durable physical medium customarily used for software interchange, for a price no more than your reasonable cost of physically performing this conveying of source, or (2) access to copy the Corresponding Source from a network server at no charge. c) Convey individual copies of the object code with a copy of the written offer to provide the Corresponding Source. This alternative is allowed only occasionally and noncommercially, and only if you received the object code with such an offer, in accord with subsection 6b. d) Convey the object code by offering access from a designated place (gratis or for a charge), and offer equivalent access to the Corresponding Source in the same way through the same place at no further charge. You need not require recipients to copy the Corresponding Source along with the object code. If the place to copy the object code is a network server, the Corresponding Source may be on a different server (operated by you or a third party) that supports equivalent copying facilities, provided you maintain clear directions next to the object code saying where to find the Corresponding Source. Regardless of what server hosts the Corresponding Source, you remain obligated to ensure that it is available for as long as needed to satisfy these requirements. e) Convey the object code using peer-to-peer transmission, provided you inform other peers where the object code and Corresponding Source of the work are being offered to the general public at no charge under subsection 6d. A separable portion of the object code, whose source code is excluded from the Corresponding Source as a System Library, need not be included in conveying the object code work. A \"User Product\" is either (1) a \"consumer product\", which means any tangible personal property which is normally used for personal, family, or household purposes, or (2) anything designed or sold for incorporation into a dwelling. In determining whether a product is a consumer product, doubtful cases shall be resolved in favor of coverage. For a particular product received by a particular user, \"normally used\" refers to a typical or common use of that class of product, regardless of the status of the particular user or of the way in which the particular user actually uses, or expects or is expected to use, the product. A product is a consumer product regardless of whether the product has substantial commercial, industrial or non-consumer uses, unless such uses represent the only significant mode of use of the product. \"Installation Information\" for a User Product means any methods, procedures, authorization keys, or other information required to install and execute modified versions of a covered work in that User Product from a modified version of its Corresponding Source. The information must suffice to ensure that the continued functioning of the modified object code is in no case prevented or interfered with solely because modification has been made. If you convey an object code work under this section in, or with, or specifically for use in, a User Product, and the conveying occurs as part of a transaction in which the right of possession and use of the User Product is transferred to the recipient in perpetuity or for a fixed term (regardless of how the transaction is characterized), the Corresponding Source conveyed under this section must be accompanied by the Installation Information. But this requirement does not apply if neither you nor any third party retains the ability to install modified object code on the User Product (for example, the work has been installed in ROM). The requirement to provide Installation Information does not include a requirement to continue to provide support service, warranty, or updates for a work that has been modified or installed by the recipient, or for the User Product in which it has been modified or installed. Access to a network may be denied when the modification itself materially and adversely affects the operation of the network or violates the rules and protocols for communication across the network. Corresponding Source conveyed, and Installation Information provided, in accord with this section must be in a format that is publicly documented (and with an implementation available to the public in source code form), and must require no special password or key for unpacking, reading or copying. 7. Additional Terms. \"Additional permissions\" are terms that supplement the terms of this License by making exceptions from one or more of its conditions. Additional permissions that are applicable to the entire Program shall be treated as though they were included in this License, to the extent that they are valid under applicable law. If additional permissions apply only to part of the Program, that part may be used separately under those permissions, but the entire Program remains governed by this License without regard to the additional permissions. When you convey a copy of a covered work, you may at your option remove any additional permissions from that copy, or from any part of it. (Additional permissions may be written to require their own removal in certain cases when you modify the work.) You may place additional permissions on material, added by you to a covered work, for which you have or can give appropriate copyright permission. Notwithstanding any other provision of this License, for material you add to a covered work, you may (if authorized by the copyright holders of that material) supplement the terms of this License with terms: a) Disclaiming warranty or limiting liability differently from the terms of sections 15 and 16 of this License; or b) Requiring preservation of specified reasonable legal notices or author attributions in that material or in the Appropriate Legal Notices displayed by works containing it; or c) Prohibiting misrepresentation of the origin of that material, or requiring that modified versions of such material be marked in reasonable ways as different from the original version; or d) Limiting the use for publicity purposes of names of licensors or authors of the material; or e) Declining to grant rights under trademark law for use of some trade names, trademarks, or service marks; or f) Requiring indemnification of licensors and authors of that material by anyone who conveys the material (or modified versions of it) with contractual assumptions of liability to the recipient, for any liability that these contractual assumptions directly impose on those licensors and authors. All other non-permissive additional terms are considered \"further restrictions\" within the meaning of section 10. If the Program as you received it, or any part of it, contains a notice stating that it is governed by this License along with a term that is a further restriction, you may remove that term. If a license document contains a further restriction but permits relicensing or conveying under this License, you may add to a covered work material governed by the terms of that license document, provided that the further restriction does not survive such relicensing or conveying. If you add terms to a covered work in accord with this section, you must place, in the relevant source files, a statement of the additional terms that apply to those files, or a notice indicating where to find the applicable terms. Additional terms, permissive or non-permissive, may be stated in the form of a separately written license, or stated as exceptions; the above requirements apply either way. 8. Termination. You may not propagate or modify a covered work except as expressly provided under this License. Any attempt otherwise to propagate or modify it is void, and will automatically terminate your rights under this License (including any patent licenses granted under the third paragraph of section 11). However, if you cease all violation of this License, then your license from a particular copyright holder is reinstated (a) provisionally, unless and until the copyright holder explicitly and finally terminates your license, and (b) permanently, if the copyright holder fails to notify you of the violation by some reasonable means prior to 60 days after the cessation. Moreover, your license from a particular copyright holder is reinstated permanently if the copyright holder notifies you of the violation by some reasonable means, this is the first time you have received notice of violation of this License (for any work) from that copyright holder, and you cure the violation prior to 30 days after your receipt of the notice. Termination of your rights under this section does not terminate the licenses of parties who have received copies or rights from you under this License. If your rights have been terminated and not permanently reinstated, you do not qualify to receive new licenses for the same material under section 10. 9. Acceptance Not Required for Having Copies. You are not required to accept this License in order to receive or run a copy of the Program. Ancillary propagation of a covered work occurring solely as a consequence of using peer-to-peer transmission to receive a copy likewise does not require acceptance. However, nothing other than this License grants you permission to propagate or modify any covered work. These actions infringe copyright if you do not accept this License. Therefore, by modifying or propagating a covered work, you indicate your acceptance of this License to do so. 10. Automatic Licensing of Downstream Recipients. Each time you convey a covered work, the recipient automatically receives a license from the original licensors, to run, modify and propagate that work, subject to this License. You are not responsible for enforcing compliance by third parties with this License. An \"entity transaction\" is a transaction transferring control of an organization, or substantially all assets of one, or subdividing an organization, or merging organizations. If propagation of a covered work results from an entity transaction, each party to that transaction who receives a copy of the work also receives whatever licenses to the work the party's predecessor in interest had or could give under the previous paragraph, plus a right to possession of the Corresponding Source of the work from the predecessor in interest, if the predecessor has it or can get it with reasonable efforts. You may not impose any further restrictions on the exercise of the rights granted or affirmed under this License. For example, you may not impose a license fee, royalty, or other charge for exercise of rights granted under this License, and you may not initiate litigation (including a cross-claim or counterclaim in a lawsuit) alleging that any patent claim is infringed by making, using, selling, offering for sale, or importing the Program or any portion of it. 11. Patents. A \"contributor\" is a copyright holder who authorizes use under this License of the Program or a work on which the Program is based. The work thus licensed is called the contributor's \"contributor version\". A contributor's \"essential patent claims\" are all patent claims owned or controlled by the contributor, whether already acquired or hereafter acquired, that would be infringed by some manner, permitted by this License, of making, using, or selling its contributor version, but do not include claims that would be infringed only as a consequence of further modification of the contributor version. For purposes of this definition, \"control\" includes the right to grant patent sublicenses in a manner consistent with the requirements of this License. Each contributor grants you a non-exclusive, worldwide, royalty-free patent license under the contributor's essential patent claims, to make, use, sell, offer for sale, import and otherwise run, modify and propagate the contents of its contributor version. In the following three paragraphs, a \"patent license\" is any express agreement or commitment, however denominated, not to enforce a patent (such as an express permission to practice a patent or covenant not to sue for patent infringement). To \"grant\" such a patent license to a party means to make such an agreement or commitment not to enforce a patent against the party. If you convey a covered work, knowingly relying on a patent license, and the Corresponding Source of the work is not available for anyone to copy, free of charge and under the terms of this License, through a publicly available network server or other readily accessible means, then you must either (1) cause the Corresponding Source to be so available, or (2) arrange to deprive yourself of the benefit of the patent license for this particular work, or (3) arrange, in a manner consistent with the requirements of this License, to extend the patent license to downstream recipients. \"Knowingly relying\" means you have actual knowledge that, but for the patent license, your conveying the covered work in a country, or your recipient's use of the covered work in a country, would infringe one or more identifiable patents in that country that you have reason to believe are valid. If, pursuant to or in connection with a single transaction or arrangement, you convey, or propagate by procuring conveyance of, a covered work, and grant a patent license to some of the parties receiving the covered work authorizing them to use, propagate, modify or convey a specific copy of the covered work, then the patent license you grant is automatically extended to all recipients of the covered work and works based on it. A patent license is \"discriminatory\" if it does not include within the scope of its coverage, prohibits the exercise of, or is conditioned on the non-exercise of one or more of the rights that are specifically granted under this License. You may not convey a covered work if you are a party to an arrangement with a third party that is in the business of distributing software, under which you make payment to the third party based on the extent of your activity of conveying the work, and under which the third party grants, to any of the parties who would receive the covered work from you, a discriminatory patent license (a) in connection with copies of the covered work conveyed by you (or copies made from those copies), or (b) primarily for and in connection with specific products or compilations that contain the covered work, unless you entered into that arrangement, or that patent license was granted, prior to 28 March 2007. Nothing in this License shall be construed as excluding or limiting any implied license or other defenses to infringement that may otherwise be available to you under applicable patent law. 12. No Surrender of Others' Freedom. If conditions are imposed on you (whether by court order, agreement or otherwise) that contradict the conditions of this License, they do not excuse you from the conditions of this License. If you cannot convey a covered work so as to satisfy simultaneously your obligations under this License and any other pertinent obligations, then as a consequence you may not convey it at all. For example, if you agree to terms that obligate you to collect a royalty for further conveying from those to whom you convey the Program, the only way you could satisfy both those terms and this License would be to refrain entirely from conveying the Program. 13. Use with the GNU Affero General Public License. Notwithstanding any other provision of this License, you have permission to link or combine any covered work with a work licensed under version 3 of the GNU Affero General Public License into a single combined work, and to convey the resulting work. The terms of this License will continue to apply to the part which is the covered work, but the special requirements of the GNU Affero General Public License, section 13, concerning interaction through a network will apply to the combination as such. 14. Revised Versions of this License. The Free Software Foundation may publish revised and/or new versions of the GNU General Public License from time to time. Such new versions will be similar in spirit to the present version, but may differ in detail to address new problems or concerns. Each version is given a distinguishing version number. If the Program specifies that a certain numbered version of the GNU General Public License \"or any later version\" applies to it, you have the option of following the terms and conditions either of that numbered version or of any later version published by the Free Software Foundation. If the Program does not specify a version number of the GNU General Public License, you may choose any version ever published by the Free Software Foundation. If the Program specifies that a proxy can decide which future versions of the GNU General Public License can be used, that proxy's public statement of acceptance of a version permanently authorizes you to choose that version for the Program. Later license versions may give you additional or different permissions. However, no additional obligations are imposed on any author or copyright holder as a result of your choosing to follow a later version. 15. Disclaimer of Warranty. THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM \"AS IS\" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING, REPAIR OR CORRECTION. 16. Limitation of Liability. IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. 17. Interpretation of Sections 15 and 16. If the disclaimer of warranty and limitation of liability provided above cannot be given local legal effect according to their terms, reviewing courts shall apply local law that most closely approximates an absolute waiver of all civil liability in connection with the Program, unless a warranty or assumption of liability accompanies a copy of the Program in return for a fee. END OF TERMS AND CONDITIONS How to Apply These Terms to Your New Programs If you develop a new program, and you want it to be of the greatest possible use to the public, the best way to achieve this is to make it free software which everyone can redistribute and change under these terms. To do so, attach the following notices to the program. It is safest to attach them to the start of each source file to most effectively state the exclusion of warranty; and each file should have at least the \"copyright\" line and a pointer to where the full notice is found. Copyright (C) This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . Also add information on how to contact you by electronic and paper mail. If the program does terminal interaction, make it output a short notice like this when it starts in an interactive mode: Copyright (C) This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'. This is free software, and you are welcome to redistribute it under certain conditions; type `show c' for details. The hypothetical commands `show w' and `show c' should show the appropriate parts of the General Public License. Of course, your program's commands might be different; for a GUI interface, you would use an \"about box\". You should also get your employer (if you work as a programmer) or school, if any, to sign a \"copyright disclaimer\" for the program, if necessary. For more information on this, and how to apply and follow the GNU GPL, see . The GNU General Public License does not permit incorporating your program into proprietary programs. If your program is a subroutine library, you may consider it more useful to permit linking proprietary applications with the library. If this is what you want to do, use the GNU Lesser General Public License instead of this License. But first, please read .", diff --git a/litellm/__init__.py b/litellm/__init__.py index 1e8f8613fba..3365abe3256 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -225,6 +225,10 @@ use_chat_completions_url_for_anthropic_messages: bool = bool( route_all_chat_openai_to_responses: bool = ( os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true" ) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge +use_legacy_interactions_schema: bool = ( + os.getenv("LITELLM_USE_LEGACY_INTERACTIONS_SCHEMA", "false").lower() == "true" +) # When True, sends Api-Revision: 2026-05-07 to Google so responses use the legacy `outputs` +# schema instead of the new `steps` schema. Remove this flag after June 8, 2026. retry = True ### AUTH ### api_key: Optional[str] = None @@ -409,6 +413,12 @@ internal_user_budget_duration: Optional[str] = None tag_budget_config: Optional[Dict[str, "BudgetConfig"]] = None max_end_user_budget: Optional[float] = None max_end_user_budget_id: Optional[str] = None +# When True, end-user IDs extracted from requests are validated against +# LiteLLM_EndUserTable / LiteLLM_UserTable. Values that do not resolve to a +# known row are dropped before reaching spend logs. Defaults to False for +# backwards compatibility — arbitrary client-supplied identifiers still +# pass through unchanged. +validate_end_user_id_in_db: bool = False disable_end_user_cost_tracking: Optional[bool] = None disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None @@ -632,6 +642,7 @@ minimax_models: Set = set() aws_polly_models: Set = set() gigachat_models: Set = set() llamagate_models: Set = set() +reducto_models: Set = set() bedrock_mantle_models: Set = set() @@ -899,6 +910,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None): gigachat_models.add(key) elif value.get("litellm_provider") == "llamagate": llamagate_models.add(key) + elif value.get("litellm_provider") == "reducto": + reducto_models.add(key) elif value.get("litellm_provider") == "bedrock_mantle": bedrock_mantle_models.add(key) @@ -1010,6 +1023,7 @@ model_list = list( | ovhcloud_models | lemonade_models | docker_model_runner_models + | reducto_models | bedrock_mantle_models | set(clarifai_models) ) @@ -1116,6 +1130,7 @@ models_by_provider: dict = { "aws_polly": aws_polly_models, "gigachat": gigachat_models, "llamagate": llamagate_models, + "reducto": reducto_models, "bedrock_mantle": bedrock_mantle_models, } diff --git a/litellm/_uuid.py b/litellm/_uuid.py index 52acf647dd8..2b7c3b82d35 100644 --- a/litellm/_uuid.py +++ b/litellm/_uuid.py @@ -6,7 +6,6 @@ Always uses fastuuid for performance. import fastuuid as _uuid # type: ignore - # Expose a module-like alias so callers can use: uuid.uuid4() uuid = _uuid diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index 28020e763f4..4548185bbdc 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -9,7 +9,6 @@ from typing import Dict, Optional from .exceptions import AnthropicErrorResponse, AnthropicErrorType - # HTTP status code -> Anthropic error type # Source: https://docs.anthropic.com/en/api/errors ANTHROPIC_ERROR_TYPE_MAP: Dict[int, AnthropicErrorType] = { diff --git a/litellm/anthropic_interface/exceptions/exceptions.py b/litellm/anthropic_interface/exceptions/exceptions.py index 984390fa702..b289e493e6b 100644 --- a/litellm/anthropic_interface/exceptions/exceptions.py +++ b/litellm/anthropic_interface/exceptions/exceptions.py @@ -2,7 +2,6 @@ from typing_extensions import Literal, Required, TypedDict - # Known Anthropic error types # Source: https://docs.anthropic.com/en/api/errors AnthropicErrorType = Literal[ diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index e3cbf422e5d..51abbbf729b 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -30,6 +30,11 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( CompletionTransformationBridge, ) +from litellm.responses.sse_output_recovery import ( + parse_sse_json_chunk, + record_output_item_chunk, + record_output_text_chunk, +) from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionReasoningItem, @@ -97,7 +102,7 @@ def _build_reasoning_item( def _reasoning_item_to_response_input( - r_item: Union[ChatCompletionReasoningItem, Dict[str, Any]] + r_item: Union[ChatCompletionReasoningItem, Dict[str, Any]], ) -> Dict[str, Any]: """Convert a stored ChatCompletionReasoningItem back to a Responses API input item.""" r_input: Dict[str, Any] = { @@ -601,6 +606,79 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return choices + @classmethod + def _extract_output_from_completed_event( + cls, parsed_chunk: Dict[str, Any] + ) -> Optional[List[Dict[str, Any]]]: + response_payload = parsed_chunk.get("response") + if not isinstance(response_payload, dict): + return None + response_output = response_payload.get("output") + if not isinstance(response_output, list) or len(response_output) == 0: + return None + return cast(List[Dict[str, Any]], response_output) + + @classmethod + def _recover_output_items_from_raw_sse( + cls, raw_sse: Optional[str] + ) -> List[Dict[str, Any]]: + if not raw_sse or not isinstance(raw_sse, str): + return [] + + recovered_output_items: Dict[int, Dict[str, Any]] = {} + recovered_text_only_items: Dict[int, Dict[str, Any]] = {} + + for chunk in raw_sse.splitlines(): + parsed_chunk = parse_sse_json_chunk(chunk) + if parsed_chunk is None: + continue + + event_type = parsed_chunk.get("type") + + if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + recovered_output = cls._extract_output_from_completed_event( + parsed_chunk + ) + if recovered_output is not None: + return recovered_output + continue + + if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: + record_output_item_chunk( + parsed_chunk=parsed_chunk, + output_items=recovered_output_items, + ) + continue + + if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE: + record_output_text_chunk( + parsed_chunk=parsed_chunk, + output_items=recovered_output_items, + text_only_items=recovered_text_only_items, + ) + continue + + # Merge text-only items into the recovered output items. Real + # OUTPUT_ITEM_DONE events take precedence at any given output_index, + # but text-only items at indices without a matching OUTPUT_ITEM_DONE + # must still be preserved (e.g. multi-output responses where some + # indices only emitted OUTPUT_TEXT_DONE). + merged_items: Dict[int, Dict[str, Any]] = {**recovered_text_only_items} + merged_items.update(recovered_output_items) + + if merged_items: + return [item for _, item in sorted(merged_items.items())] + + return [] + + @classmethod + def _recover_output_items_from_logging( + cls, logging_obj: "LiteLLMLoggingObj" + ) -> List[Dict[str, Any]]: + model_call_details = getattr(logging_obj, "model_call_details", {}) or {} + original_response = model_call_details.get("original_response") + return cls._recover_output_items_from_raw_sse(original_response) + def transform_response( # noqa: PLR0915 self, model: str, @@ -625,9 +703,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if raw_response.error is not None: raise ValueError(f"Error in response: {raw_response.error}") + output_items = raw_response.output + if len(output_items) == 0: + recovered_output_items = self._recover_output_items_from_logging( + logging_obj + ) + if recovered_output_items: + output_items = cast(Any, recovered_output_items) + raw_response.output = cast(Any, recovered_output_items) + verbose_logger.warning( + "Recovered empty Responses API output from raw SSE for model=%s", + model, + ) + # Convert response output to choices using the static helper choices = self._convert_response_output_to_choices( - output_items=raw_response.output, + output_items=output_items, handle_raw_dict_callback=self._handle_raw_dict_response_item, ) @@ -641,7 +732,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) else: raise ValueError( - f"Unknown items in responses API response: {raw_response.output}" + f"Unknown items in responses API response: {output_items}" ) setattr(model_response, "choices", choices) @@ -1237,7 +1328,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): raise ValueError( f"Chat provider: Invalid function argument delta {parsed_chunk}" ) - elif event_type == "response.output_item.done": + elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: # New output item added output_item = parsed_chunk.get("item", {}) if output_item.get("type") == "function_call": diff --git a/litellm/compression/content_detection.py b/litellm/compression/content_detection.py index 0655a42daf5..975117eb608 100644 --- a/litellm/compression/content_detection.py +++ b/litellm/compression/content_detection.py @@ -5,7 +5,6 @@ Auto-detect content type per message: code, JSON, or text. import json import re - _CODE_KEYWORDS = re.compile( r"\b(?:def |function |class |import |from |require\(|#include|fn |func |const |let |var |public |private |static )\b" ) diff --git a/litellm/constants.py b/litellm/constants.py index e36746326cc..fb765c0226c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1443,6 +1443,12 @@ CLI_JWT_EXPIRATION_HOURS = int( or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS") or 24 ) +# Comma-separated allowlisted OIDC claim map for CLI SSO polling, e.g. +# "employment_type->acme_employment_type,org_info.department->department" +CLI_SSO_CLAIM_MAP = ( + os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or "" +) +CLI_SSO_CLAIM_MAX_SCALAR_LENGTH = 1024 ########################### UI SESSION DURATION ########################### # Duration for UI login session (username/password, SSO, invitation links). Format: "30s", "30m", "24h", "7d" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 2257861aff6..98e00cf5788 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1879,10 +1879,6 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") - pages_processed = response.usage_info.pages_processed - if pages_processed is None: - raise ValueError("OCR response pages_processed is None") - try: model_info: Optional[ModelInfo] = litellm.get_model_info( model=model, custom_llm_provider=custom_llm_provider @@ -1890,9 +1886,49 @@ def ocr_cost( except Exception: model_info = None - ocr_cost_per_page: float = 0.0 + credits = getattr(response.usage_info, "credits", None) + cost_per_credit = None if model_info is not None: - ocr_cost_per_page = model_info.get("ocr_cost_per_page") or 0.0 + cost_per_credit = model_info.get("ocr_cost_per_credit") + if credits is not None and cost_per_credit is not None: + return cost_per_credit * credits, 0.0 + + ocr_cost_per_page: Optional[float] = None + if model_info is not None: + ocr_cost_per_page = model_info.get("ocr_cost_per_page") + + pages_processed = response.usage_info.pages_processed + if pages_processed is None: + if cost_per_credit is not None or ocr_cost_per_page is None: + # Surface missing usage data instead of silently under-reporting + # cost. The previous behavior raised ValueError; we now return 0.0 + # for credit-priced or unpriced models, so log a warning to keep + # the regression visible to operators. + verbose_logger.warning( + "OCR cost: model=%s custom_llm_provider=%s response.usage_info." + "pages_processed is None and credits=%s; returning 0.0 cost.", + model, + custom_llm_provider, + credits, + ) + return 0.0, 0.0 + raise ValueError("OCR response pages_processed is None") + + if ocr_cost_per_page is None: + # No per-page pricing configured. Either the model is on credit-based + # pricing (and credits weren't returned, so the credit branch above did + # not match) or the model has no OCR pricing entry at all. Surface a + # warning so that missing pricing entries are visible rather than + # silently producing zero cost for billable usage. + verbose_logger.warning( + "OCR cost: model=%s custom_llm_provider=%s reported " + "pages_processed=%s but no ocr_cost_per_page is configured; " + "returning 0.0 cost.", + model, + custom_llm_provider, + pages_processed, + ) + return 0.0, 0.0 total_ocr_processing_cost: float = ocr_cost_per_page * pages_processed return total_ocr_processing_cost, 0.0 diff --git a/litellm/files/types.py b/litellm/files/types.py index 688bc86f0cf..ba42a39f666 100644 --- a/litellm/files/types.py +++ b/litellm/files/types.py @@ -1,6 +1,5 @@ from typing import AsyncIterator, Dict, Iterator, Literal, NamedTuple, Union - FileContentProvider = Literal[ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus" ] diff --git a/litellm/google_genai/adapters/__init__.py b/litellm/google_genai/adapters/__init__.py index bfa9e712678..6fbe7d95a55 100644 --- a/litellm/google_genai/adapters/__init__.py +++ b/litellm/google_genai/adapters/__init__.py @@ -1,10 +1,10 @@ """ Google GenAI Adapters for LiteLLM -This module provides adapters for transforming Google GenAI generate_content requests +This module provides adapters for transforming Google GenAI generate_content requests to/from LiteLLM completion format with full support for: - Text content transformation -- Tool calling (function declarations, function calls, function responses) +- Tool calling (function declarations, function calls, function responses) - Streaming (both regular and tool calling) - Mixed content (text + tool calls) """ diff --git a/litellm/integrations/SlackAlerting/batching_handler.py b/litellm/integrations/SlackAlerting/batching_handler.py index fdce2e04793..828f3eb4175 100644 --- a/litellm/integrations/SlackAlerting/batching_handler.py +++ b/litellm/integrations/SlackAlerting/batching_handler.py @@ -1,9 +1,9 @@ """ -Handles Batching + sending Httpx Post requests to slack +Handles Batching + sending Httpx Post requests to slack -Slack alerts are sent every 10s or when events are greater than X events +Slack alerts are sent every 10s or when events are greater than X events -see custom_batch_logger.py for more details / defaults +see custom_batch_logger.py for more details / defaults """ from typing import TYPE_CHECKING, Any diff --git a/litellm/integrations/SlackAlerting/utils.py b/litellm/integrations/SlackAlerting/utils.py index e695266c88b..e2580768178 100644 --- a/litellm/integrations/SlackAlerting/utils.py +++ b/litellm/integrations/SlackAlerting/utils.py @@ -18,7 +18,7 @@ else: def process_slack_alerting_variables( - alert_to_webhook_url: Optional[Dict[AlertType, Union[List[str], str]]] + alert_to_webhook_url: Optional[Dict[AlertType, Union[List[str], str]]], ) -> Optional[Dict[AlertType, Union[List[str], str]]]: """ process alert_to_webhook_url diff --git a/litellm/integrations/additional_logging_utils.py b/litellm/integrations/additional_logging_utils.py index 795afd81d41..59319140a18 100644 --- a/litellm/integrations/additional_logging_utils.py +++ b/litellm/integrations/additional_logging_utils.py @@ -1,5 +1,5 @@ """ -Base class for Additional Logging Utils for CustomLoggers +Base class for Additional Logging Utils for CustomLoggers - Health Check for the logging util - Get Request / Response Payload for the logging util diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index f9d4496c21f..8f4844501c3 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -1,5 +1,5 @@ """ -Custom Logger that handles batching logic +Custom Logger that handles batching logic Use this if you want your logs to be stored in memory and flushed periodically. """ @@ -14,22 +14,38 @@ from litellm.integrations.custom_logger import CustomLogger class CustomBatchLogger(CustomLogger): + preserve_events_added_during_flush = False + + # Default cap on the in-memory log queue. Prevents unbounded memory growth + # if ``async_send_batch`` consistently fails (e.g. the destination is + # unreachable) and events are preserved across flush attempts. Subclasses + # may override by passing ``max_queue_size`` or by setting the attribute + # directly (see ``RubrikLogger`` for an example). + DEFAULT_MAX_QUEUE_SIZE = 50_000 + def __init__( self, flush_lock: Optional[asyncio.Lock] = None, batch_size: Optional[int] = None, flush_interval: Optional[int] = None, + max_queue_size: Optional[int] = None, **kwargs, ) -> None: """ Args: flush_lock (Optional[asyncio.Lock], optional): Lock to use when flushing the queue. Defaults to None. Only used for custom loggers that do batching + max_queue_size (Optional[int], optional): Maximum number of events to retain in ``log_queue``. When the limit is exceeded (e.g. because the send destination is unreachable and events are preserved for retry), the oldest events are dropped. Defaults to ``DEFAULT_MAX_QUEUE_SIZE``. """ self.log_queue: List = [] self.flush_interval = flush_interval or litellm.DEFAULT_FLUSH_INTERVAL_SECONDS self.batch_size: int = batch_size or litellm.DEFAULT_BATCH_SIZE self.last_flush_time = time.time() self.flush_lock = flush_lock + self.max_queue_size: int = ( + max_queue_size + if max_queue_size is not None + else self.DEFAULT_MAX_QUEUE_SIZE + ) super().__init__(**kwargs) @@ -47,11 +63,40 @@ class CustomBatchLogger(CustomLogger): async with self.flush_lock: if self.log_queue: + log_queue_length = len(self.log_queue) verbose_logger.debug( "CustomLogger: Flushing batch of %s events", len(self.log_queue) ) - await self.async_send_batch() - self.log_queue.clear() + try: + await self.async_send_batch() + except Exception: + # If the underlying batch send raised, do NOT drop the + # in-flight events. They will be retried on the next flush. + # Most existing async_send_batch implementations swallow + # their own errors, so this only affects loggers that opt + # in to surfacing failures (e.g. Rubrik). + verbose_logger.exception( + "CustomLogger: async_send_batch raised; preserving " + "%s events in queue for retry", + log_queue_length, + ) + # Guard against unbounded queue growth if the destination + # is persistently unreachable. Drop the oldest events + # beyond ``max_queue_size``. + overflow = len(self.log_queue) - self.max_queue_size + if overflow > 0: + del self.log_queue[:overflow] + verbose_logger.warning( + "CustomLogger: log queue exceeded max_queue_size=%s; " + "dropped %s oldest events.", + self.max_queue_size, + overflow, + ) + return + if self.preserve_events_added_during_flush: + del self.log_queue[:log_queue_length] + else: + self.log_queue.clear() self.last_flush_time = time.time() async def async_send_batch(self, *args, **kwargs): diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index b7d28e3dbb9..6f4433b4a05 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -9,7 +9,6 @@ import polars as pl from .schema import FOCUS_NORMALIZED_SCHEMA - _TAG_KEYS = ( "team_id", "team_alias", diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index a70574952b8..6c8510380a8 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -726,9 +726,57 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): exception_logging_span.set_status(Status(StatusCode.ERROR)) exception_logging_span.end(end_time=self._to_ns(datetime.now())) + # Emit guardrail spans for any guardrail invocations that + # ran during this request. _handle_failure typically does this, + # but for pre-call guardrail blocks the standard_logging_object + # may not carry guardrail_information by the time _handle_failure + # fires (the data lives only in request_data["metadata"]). Pull + # directly from request_data so the span is recorded either way; + # _emit_once dedupes if _handle_failure already emitted it. + self._emit_guardrail_spans_from_request_data( + request_data=request_data, + parent_span=parent_otel_span, + ) + # End Parent OTEL Sspan parent_otel_span.end(end_time=self._to_ns(datetime.now())) + def _emit_guardrail_spans_from_request_data( + self, + request_data: dict, + parent_span: Optional[Any], + ) -> None: + """Emit ``guardrail`` spans from ``request_data["metadata"] + ["standard_logging_guardrail_information"]``. + + Routed through ``_create_guardrail_span`` so the dedupe state in + ``_otel_internal`` is honoured — if ``_handle_failure`` already + emitted these spans for the same kwargs, this is a no-op. + """ + from opentelemetry import trace as _trace + + metadata = (request_data or {}).get("metadata") or {} + guardrail_information = metadata.get("standard_logging_guardrail_information") + if not guardrail_information: + return + + # _create_guardrail_span reads guardrail_information from + # kwargs["standard_logging_object"] and shares its dedupe state via + # kwargs["litellm_params"]["metadata"]["_otel_internal"]. Pass the + # SAME metadata dict the proxy populated so _handle_failure and + # this hook see the same dedupe markers. + kwargs: Dict[str, Any] = { + "litellm_params": {"metadata": metadata}, + "standard_logging_object": { + "guardrail_information": guardrail_information, + "metadata": metadata, + }, + } + context = ( + _trace.set_span_in_context(parent_span) if parent_span is not None else None + ) + self._create_guardrail_span(kwargs=kwargs, context=context) + async def async_post_call_success_hook( self, data: dict, @@ -1611,12 +1659,43 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): "masked_entity_count", safe_dumps(masked_entity_count) ) + guardrail_response = guardrail_information.get("guardrail_response") + if guardrail_response is not None: + guardrail_span.set_attribute( + "guardrail_response", safe_dumps(guardrail_response) + ) + + # Surface guardrail_status (success / guardrail_intervened / + # guardrail_failed_to_respond / not_run) as a top-level span + # attribute so trace backends can filter on it without parsing + # guardrail_response. self.safe_set_attribute( span=guardrail_span, - key="guardrail_response", - value=guardrail_information.get("guardrail_response"), + key="guardrail_status", + value=guardrail_information.get("guardrail_status"), ) + # Provider's raw top-level action (e.g. Bedrock's + # ``GUARDRAIL_INTERVENED`` / ``NONE``). Populated by the provider + # hook onto StandardLoggingGuardrailInformation so this integration + # stays provider-agnostic — we only read a normalised string. + guardrail_action = guardrail_information.get("guardrail_action") + if guardrail_action: + guardrail_span.set_attribute("guardrail_action", guardrail_action) + + # The provider hook (e.g. Bedrock) extracts violation_categories + # from the raw response BEFORE redaction and stamps them onto + # StandardLoggingGuardrailInformation. Surfacing them here as a + # queryable attribute lets dashboards group by violation category + # without parsing the redacted guardrail_response blob. + violation_categories = guardrail_information.get("violation_categories") + if violation_categories: + # OTel sequence attributes must be homogeneous primitives; + # serialise to JSON once so set_attribute never coerces. + guardrail_span.set_attribute( + "guardrail_violation_categories", safe_dumps(violation_categories) + ) + self._set_team_attributes_from_kwargs(guardrail_span, kwargs) guardrail_span.end(end_time=self._to_ns(end_time_datetime)) diff --git a/litellm/integrations/opik/utils.py b/litellm/integrations/opik/utils.py index b0ab5991c91..43577505c11 100644 --- a/litellm/integrations/opik/utils.py +++ b/litellm/integrations/opik/utils.py @@ -105,7 +105,7 @@ def _remove_nulls(x: Dict[str, Any]) -> Dict[str, Any]: def get_traces_and_spans_from_payload( - payload: List[Dict[str, Any]] + payload: List[Dict[str, Any]], ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: """ Separate traces and spans from payload. diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py new file mode 100644 index 00000000000..af396ecdc73 --- /dev/null +++ b/litellm/integrations/rubrik.py @@ -0,0 +1,605 @@ +"""Rubrik LiteLLM Plugin for tool blocking and batch logging.""" + +import asyncio +import os +import random +import time +import urllib.parse +import uuid +from collections import Counter +from typing import TYPE_CHECKING, Any, Literal, Optional + +import httpx +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) +from litellm.litellm_core_utils.core_helpers import safe_deep_copy +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Function, + GenericGuardrailAPIInputs, + StandardLoggingPayload, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + +_ENDPOINT_ANTHROPIC_MESSAGES = "/v1/messages" +_WEBHOOK_PATH_TOOL_BLOCKING = "/v1/after_completion/openai/v1" +_WEBHOOK_PATH_LOGGING_BATCH = "/v1/litellm/batch" +_MAX_QUEUE_SIZE = 10_000 +_DROP_WARNING_INTERVAL_SECONDS = 60.0 + + +class _MalformedToolBlockingResponseError(Exception): + """Raised when the tool blocking service returns a structurally invalid + response (e.g. empty ``choices``). + + Distinct from transient network/HTTP errors so callers can surface a + louder, misconfiguration-style log instead of treating it as a routine + fail-open. + """ + + +class RubrikLogger(CustomGuardrail, CustomBatchLogger): + def __init__( + self, + api_key: str | None = None, + api_base: str | None = None, + **kwargs, + ): + self.flush_lock = asyncio.Lock() + kwargs.setdefault("guardrail_name", "rubrik") + # `initialize_guardrail` always passes these kwargs explicitly, with + # value `None` when the user omits `mode` / `default_on` from the + # guardrail config. Coerce None (omitted) to the desired default + # while preserving any explicit value the caller did set -- + # in particular `default_on=False` if the user wants the guardrail + # off by default. + kwargs["event_hook"] = kwargs.get("event_hook") or GuardrailEventHooks.post_call + if kwargs.get("default_on") is None: + kwargs["default_on"] = True + super().__init__( + flush_lock=self.flush_lock, + **kwargs, + ) + + verbose_logger.debug("initializing rubrik logger") + + self.sampling_rate = 1.0 + rbrk_sampling_rate = os.getenv("RUBRIK_SAMPLING_RATE") + if rbrk_sampling_rate is not None: + try: + parsed_rate = float(rbrk_sampling_rate.strip()) + self.sampling_rate = max(0.0, min(1.0, parsed_rate)) + if parsed_rate != self.sampling_rate: + verbose_logger.warning( + f"RUBRIK_SAMPLING_RATE={parsed_rate} clamped to " + f"{self.sampling_rate}" + ) + except ValueError: + verbose_logger.warning( + f"Invalid RUBRIK_SAMPLING_RATE: {rbrk_sampling_rate!r}, using 1.0" + ) + + self.key = api_key or os.getenv("RUBRIK_API_KEY") + if not self.key: + verbose_logger.warning( + "Rubrik: No API key configured. Requests will be unauthenticated." + ) + _batch_size = os.getenv("RUBRIK_BATCH_SIZE") + + if _batch_size: + try: + self.batch_size = int(_batch_size) + except ValueError: + verbose_logger.warning( + f"Invalid RUBRIK_BATCH_SIZE: {_batch_size!r}, using default" + ) + + # Cap the in-memory retry queue so a Rubrik webhook outage cannot let + # authenticated traffic accumulate prompt/response payloads until the + # proxy runs out of memory. Once the cap is reached, oldest events are + # dropped to make room for fresh ones (drop-oldest backpressure). + self.max_queue_size = _MAX_QUEUE_SIZE + self._dropped_since_warning = 0 + self._last_drop_warning_time = 0.0 + + _webhook_url = api_base or os.getenv("RUBRIK_WEBHOOK_URL") + + if _webhook_url is None: + raise ValueError( + "Rubrik webhook URL not configured. " + "Set RUBRIK_WEBHOOK_URL or pass api_base." + ) + + _webhook_url = _webhook_url.rstrip("/").removesuffix("/v1") + self.tool_blocking_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_TOOL_BLOCKING}" + self.logging_endpoint = f"{_webhook_url}{_WEBHOOK_PATH_LOGGING_BATCH}" + + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + + self.tool_blocking_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback, + params={"timeout": httpx.Timeout(5.0, connect=2.0)}, + ) + + self._headers: dict[str, str] = {"Content-Type": "application/json"} + if self.key: + self._headers["Authorization"] = f"Bearer {self.key}" + + # Periodic flush is started lazily on the first log event so that + # low-traffic deployments still get their batches drained even when the + # logger is instantiated outside a running event loop (sync init). + self._flush_task: Optional[asyncio.Task[Any]] = ( + self._start_periodic_flush_task() + ) + + def _start_periodic_flush_task(self) -> Optional[asyncio.Task[Any]]: + """Start the periodic flush task only when an event loop is already running.""" + try: + loop = asyncio.get_running_loop() + except RuntimeError: + verbose_logger.debug( + "Rubrik logger init: no running event loop, " + "periodic flush will start on first log event." + ) + return None + return loop.create_task(self.periodic_flush()) + + def _ensure_periodic_flush_task(self) -> None: + # Synchronous helper: in asyncio's cooperative model there is no await + # between the check and assignment, so two callers cannot race here. + if self._flush_task is None or self._flush_task.done(): + self._flush_task = self._start_periodic_flush_task() + + async def aclose(self): + """Close the dedicated HTTP clients used by this logger.""" + # Cancel the periodic flush task before closing the HTTP clients so + # the loop doesn't wake up and try to POST via a closed client. + if self._flush_task is not None and not self._flush_task.done(): + self._flush_task.cancel() + try: + await self._flush_task + except (asyncio.CancelledError, Exception): + pass + self._flush_task = None + await self.tool_blocking_client.close() + await self.async_httpx_client.close() + + # -- Guardrail hook -------------------------------------------------------- + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """Validate tool calls against the blocking service (fail-open).""" + if input_type != "response": + return inputs + + tool_calls = inputs.get("tool_calls") + if not tool_calls: + return inputs + + try: + return await self._check_tool_calls( + inputs, tool_calls, request_data, logging_obj + ) + except ModifyResponseException: + raise + except _MalformedToolBlockingResponseError as e: + # Distinct from transient errors: the service responded but the + # payload was structurally invalid, which usually indicates a + # misconfigured webhook or a breaking change in its response + # format. Log loudly so operators notice their tool-blocking + # policy is not actually being enforced. + verbose_logger.critical( + "Tool blocking service returned a malformed response: %s. " + "Tool calls are NOT being checked -- verify the webhook " + "configuration. Returning original response unchanged.", + e, + exc_info=True, + ) + return inputs + except Exception as e: + verbose_logger.error( + f"Tool blocking hook failed: {e}. " + "Returning original response unchanged.", + exc_info=True, + ) + return inputs + + async def _check_tool_calls( + self, + inputs: GenericGuardrailAPIInputs, + tool_calls: Any, + request_data: dict, + logging_obj: Optional["LiteLLMLoggingObj"], + ) -> GenericGuardrailAPIInputs: + """Send tool calls to blocking service, raise if any are blocked.""" + message_tool_calls = self._normalize_tool_calls(tool_calls) + + call_details = ( + getattr(logging_obj, "model_call_details", {}) if logging_obj else {} + ) + response = request_data.get("response") + request_id = getattr(response, "id", None) if response else None + if logging_obj and not call_details: + verbose_logger.warning( + "Rubrik: logging_obj present but model_call_details is empty " + "-- request context will be missing" + ) + + response_data = self._build_tool_call_payload(message_tool_calls, request_id) + req_data = self._extract_request_data(call_details) + + service_response = await self._post_to_tool_blocking_service( + response_data, req_data + ) + blocked_explanation = self._extract_blocked_tools( + service_response, message_tool_calls + ) + + if blocked_explanation is not None: + model = self._resolve_model(request_data, call_details) + raise ModifyResponseException( + message=blocked_explanation, + model=model, + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + + return inputs + + @staticmethod + def _normalize_tool_calls(tool_calls: Any) -> list[ChatCompletionMessageToolCall]: + """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" + result = [] + for tc in tool_calls: + if isinstance(tc, ChatCompletionMessageToolCall): + result.append(tc) + elif isinstance(tc, dict): + func = tc.get("function", {}) + result.append( + ChatCompletionMessageToolCall( + id=tc.get("id", ""), + type=tc.get("type", "function"), + function=Function( + name=func.get("name", ""), + arguments=func.get("arguments", ""), + ), + ) + ) + elif hasattr(tc, "id") and hasattr(tc, "function"): + result.append( + ChatCompletionMessageToolCall( + id=tc.id or "", + type=getattr(tc, "type", None) or "function", + function=tc.function, + ) + ) + else: + raise TypeError( + f"Cannot normalize tool_call of type {type(tc).__name__}" + ) + return result + + @staticmethod + def _build_tool_call_payload( + tool_calls: list[ChatCompletionMessageToolCall], + request_id: str | None, + ) -> dict[str, Any]: + """Build a full OpenAI ChatCompletion-format dict for the blocking service.""" + return { + "id": request_id or f"chatcmpl-{uuid.uuid4()}", + "object": "chat.completion", + "created": int(time.time()), + "model": "", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + tc.model_dump(exclude_none=True) for tc in tool_calls + ], + }, + "finish_reason": "tool_calls", + } + ], + } + + @staticmethod + def _extract_request_data(call_details: dict[str, Any]) -> dict[str, Any]: + """Extract original request data from model_call_details.""" + if not call_details: + return {} + litellm_params = call_details.get("litellm_params", {}) or {} + return { + "messages": call_details.get("messages"), + "model": call_details.get("model"), + "proxy_server_request": RubrikLogger._sanitize_proxy_server_request( + litellm_params.get("proxy_server_request") + ), + } + + @staticmethod + def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: + """Allowlist only routing fields (``url``, ``method``) when forwarding + ``proxy_server_request`` to the external Rubrik webhook, dropping + inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw + request ``body`` so proxy credentials are not exfiltrated.""" + if not isinstance(proxy_server_request, dict): + return proxy_server_request + return { + key: proxy_server_request[key] + for key in ("url", "method") + if key in proxy_server_request + } + + @staticmethod + def _resolve_model( + request_data: dict[str, Any], call_details: dict[str, Any] + ) -> str: + """Get the model name for the ModifyResponseException.""" + response = request_data.get("response") + if response and hasattr(response, "model"): + return response.model or "unknown" + return call_details.get("model", "unknown") + + # -- Logging hooks --------------------------------------------------------- + + async def _prepare_log_payload( + self, kwargs: dict, event_type: str + ) -> StandardLoggingPayload | None: + """Shared logic for success and failure logging.""" + if random.random() > self.sampling_rate: + verbose_logger.debug( + f"Skipping Rubrik {event_type} logging " + f"(sampling_rate={self.sampling_rate})" + ) + return None + + # Deep-copy so mutations don't affect other callbacks sharing this object + standard_logging_payload: StandardLoggingPayload = safe_deep_copy( + kwargs["standard_logging_object"] + ) + + # For Anthropic /v1/messages requests, LiteLLM creates a separate + # ModelResponse (with a generated chatcmpl-* id) for logging, which + # differs from the original Anthropic msg-* id on the response dict. + # Normalize to litellm_call_id so that the logging and tool-blocking + # endpoints see the same request identifier. + litellm_params = kwargs.get("litellm_params", {}) or {} + proxy_request = litellm_params.get("proxy_server_request", {}) or {} + url_path = urllib.parse.urlparse(proxy_request.get("url", "")).path + if url_path.endswith(_ENDPOINT_ANTHROPIC_MESSAGES): + _litellm_call_id = kwargs.get("litellm_call_id") + if _litellm_call_id: + standard_logging_payload["id"] = _litellm_call_id # type: ignore[literal-required] + + if "system" in kwargs: + system_prompt_msg_list = kwargs["system"] + try: + if system_prompt_msg_list: + system_scaffold = { + "role": "system", + "content": system_prompt_msg_list, + } + if isinstance(standard_logging_payload["messages"], list): + standard_logging_payload["messages"].insert(0, system_scaffold) + elif isinstance(standard_logging_payload["messages"], (dict, str)): + standard_logging_payload["messages"] = [ + system_scaffold, + standard_logging_payload["messages"], + ] + except Exception as e: + verbose_logger.warning( + f"Rubrik: failed to prepend system prompt: {e}", + exc_info=True, + ) + + return standard_logging_payload + + async def _enqueue_log_event(self, kwargs: dict, event_type: str): + try: + self._ensure_periodic_flush_task() + payload = await self._prepare_log_payload(kwargs, event_type) + if payload is None: + return + + self.log_queue.append(payload) + self._enforce_max_queue_size() + + if len(self.log_queue) >= self.batch_size: + await self.flush_queue() + except Exception as e: + verbose_logger.error( + f"Rubrik {event_type} logging hook failed: {e}. " + "Skipping logging for this event.", + exc_info=True, + ) + + def _enforce_max_queue_size(self) -> None: + overflow = len(self.log_queue) - self.max_queue_size + if overflow <= 0: + return + del self.log_queue[:overflow] + self._dropped_since_warning += overflow + now = time.time() + if now - self._last_drop_warning_time >= _DROP_WARNING_INTERVAL_SECONDS: + verbose_logger.warning( + "Rubrik: log queue exceeded max_queue_size=%s; dropped %s " + "oldest events since the last warning. The Rubrik webhook may " + "be unhealthy or undersized for current traffic.", + self.max_queue_size, + self._dropped_since_warning, + ) + self._dropped_since_warning = 0 + self._last_drop_warning_time = now + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await self._enqueue_log_event(kwargs, "success") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await self._enqueue_log_event(kwargs, "failure") + + # -- Batch logging --------------------------------------------------------- + + async def _log_batch_to_rubrik(self, data): + # NOTE: this method intentionally re-raises on failure so the parent + # CustomBatchLogger.flush_queue keeps the unsent events in the queue + # for the next flush attempt instead of silently dropping them. + try: + response = await self.async_httpx_client.post( + url=self.logging_endpoint, + json=data, + headers=self._headers, + ) + response.raise_for_status() + except httpx.HTTPStatusError as e: + verbose_logger.exception( + f"Rubrik HTTP Error: {e.response.status_code} - {e.response.text}" + ) + raise + except Exception: + verbose_logger.exception("Rubrik Layer Error") + raise + + async def async_send_batch(self): + """Handles sending batches of responses to Rubrik. + + Note: the canonical flush path is :meth:`flush_queue`, which takes a + single snapshot used for both sending and queue draining. This method + is kept for direct callers / tests; it intentionally does NOT remove + events from the queue. + """ + if not self.log_queue: + return + + log_queue_snapshot = list(self.log_queue) + verbose_logger.debug( + "Rubrik: Flushing batch of %s events", len(log_queue_snapshot) + ) + await self._log_batch_to_rubrik( + data=log_queue_snapshot, + ) + + async def flush_queue(self): + """Snapshot, send, and drain in one consistent step. + + Overrides the base implementation so the same snapshot drives both + the HTTP send and the queue truncation. This avoids the subtle + coupling where the base class captures `len(self.log_queue)` + separately from the snapshot taken inside `async_send_batch`, + which could otherwise drift in a future refactor and cause + duplicate deliveries to Rubrik. + """ + if self.flush_lock is None: + return + + async with self.flush_lock: + if not self.log_queue: + return + snapshot = list(self.log_queue) + verbose_logger.debug("Rubrik: Flushing batch of %s events", len(snapshot)) + try: + await self._log_batch_to_rubrik(data=snapshot) + except Exception: + # Already logged with traceback inside _log_batch_to_rubrik. + # Preserve the in-flight events for retry on the next flush. + return + del self.log_queue[: len(snapshot)] + self.last_flush_time = time.time() + + # -- Tool blocking service ------------------------------------------------- + + async def _post_to_tool_blocking_service( + self, + response_data: dict[str, Any], + request_data: dict[str, Any], + ) -> dict[str, Any]: + """Post a payload to the tool blocking service and return the response. + + Args: + response_data: The OpenAI-formatted response payload to send. + request_data: Original LLM request data to include alongside + the response for additional context. Empty dict if unavailable. + + Raises: + Exception: If the service is unavailable or returns an error. + """ + envelope = { + "request": request_data, + "response": response_data, + } + verbose_logger.debug( + f"Sending request to tool blocking service: " + f"{self.tool_blocking_endpoint}" + ) + http_response = await self.tool_blocking_client.post( + self.tool_blocking_endpoint, + json=envelope, + headers=self._headers, + ) + http_response.raise_for_status() + result: dict[str, Any] = http_response.json() + return result + + @staticmethod + def _extract_blocked_tools( + service_response: dict[str, Any], + all_tool_calls: list[ChatCompletionMessageToolCall], + ) -> Optional[str]: + """Return the blocking explanation if any tool calls were blocked. + + Compares the service response (which contains only allowed tools) against + the full set of tool calls. Returns ``None`` if all tools are allowed, or + the explanation string (prefixed with newlines) otherwise. + + Expects service_response in OpenAI chat completion format: + {"choices": [{"message": {"tool_calls": [...], "content": "..."}}]} + """ + choices = service_response.get("choices", []) + if not choices: + raise _MalformedToolBlockingResponseError( + "Tool blocking service returned empty response" + ) + + message = choices[0].get("message", {}) + returned_tool_calls = message.get("tool_calls") or [] + blocking_explanation = message.get("content", "") + + allowed_id_counts: Counter = Counter( + tc["id"] + for tc in returned_tool_calls + if isinstance(tc, dict) and tc.get("id") + ) + required_id_counts: Counter = Counter(tc.id for tc in all_tool_calls if tc.id) + + all_allowed = len(returned_tool_calls) >= len(all_tool_calls) and all( + allowed_id_counts.get(tc_id, 0) >= count + for tc_id, count in required_id_counts.items() + ) + + if all_allowed: + return None + + explanation = blocking_explanation or "Tool call blocked by policy." + return f"\n\n{explanation}" diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 332e84dd07d..4ed8a809a13 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -1,8 +1,8 @@ """ s3 Bucket Logging Integration -async_log_success_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3 -async_log_failure_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3 +async_log_success_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3 +async_log_failure_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3 NOTE 1: S3 does not provide a BATCH PUT API endpoint, so we create tasks to upload each element individually """ diff --git a/litellm/interactions/agents/main.py b/litellm/interactions/agents/main.py index 7375fd6273f..f56c6f3ed5e 100644 --- a/litellm/interactions/agents/main.py +++ b/litellm/interactions/agents/main.py @@ -49,7 +49,6 @@ from litellm.types.interactions import InteractionEnvironment from litellm.types.router import GenericLiteLLMParams from litellm.utils import client - # ------------------------------------------------------------------ # # Shared helpers # # ------------------------------------------------------------------ # diff --git a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py index 567e9b523e8..4a3eb63084e 100644 --- a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py +++ b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py @@ -2,7 +2,17 @@ Streaming iterator for transforming Responses API stream to Interactions API stream. """ -from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, cast +from collections import deque +from typing import ( + Any, + AsyncIterator, + Deque, + Dict, + Iterator, + List, + Optional, + cast, +) from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, @@ -15,7 +25,6 @@ from litellm.types.interactions import ( InteractionsAPIStreamingResponse, ) from litellm.types.llms.openai import ( - ContentPartAddedEvent, OutputTextDeltaEvent, ResponseCompletedEvent, ResponseCreatedEvent, @@ -30,7 +39,13 @@ class LiteLLMResponsesInteractionsStreamingIterator: This class handles both sync and async iteration, transforming Responses API streaming events (output.text.delta, response.completed, etc.) to Interactions - API streaming events (content.delta, interaction.complete, etc.). + API streaming events. + + Schema selection: + - New schema (default, use_legacy_interactions_schema=False): + interaction.created -> step.start -> step.delta ... -> step.stop -> interaction.completed + - Legacy schema (use_legacy_interactions_schema=True, remove after June 8 2026): + interaction.start -> content.start -> content.delta ... -> content.stop -> interaction.complete """ def __init__( @@ -42,6 +57,8 @@ class LiteLLMResponsesInteractionsStreamingIterator: custom_llm_provider: Optional[str] = None, litellm_metadata: Optional[Dict[str, Any]] = None, ): + import litellm + self.model = model self.responses_stream_iterator = litellm_custom_stream_wrapper self.request_input = request_input @@ -52,100 +69,156 @@ class LiteLLMResponsesInteractionsStreamingIterator: self.collected_text = "" self.sent_interaction_start = False self.sent_content_start = False - self._pending_events: List[InteractionsAPIStreamingResponse] = [] + # Capture the schema flag once at construction time so all events + # emitted by this stream use a consistent schema, even if the global + # flag is mutated mid-stream (e.g. by a config reload). + self._use_legacy: bool = litellm.use_legacy_interactions_schema + # Buffer of events that have been derived from upstream chunks but not + # yet returned to the caller. A single Responses API chunk may expand + # into multiple Interactions API events (e.g. the first text delta + # produces interaction.created + step.start + step.delta), and the + # terminal sequence on stream end may also span multiple events + # (step.stop + interaction.completed). + self._pending_events: Deque[InteractionsAPIStreamingResponse] = deque() + # Tracks whether we've already emitted a terminal completion event so + # the StopIteration fallback path doesn't double-emit. + self._sent_completion_event = False + # ID resolved from the first upstream chunk (item_id on a text delta or + # response.id on response.created). Persisted so the EOF terminal + # events stay correlated with the start events delivered earlier. + self._interaction_id: Optional[str] = None - def _transform_responses_chunk_to_interactions_chunk( - self, - responses_chunk: ResponsesAPIStreamingResponse, - ) -> Optional[InteractionsAPIStreamingResponse]: + # ------------------------------------------------------------------ + # Event builders + # ------------------------------------------------------------------ + + def _build_interaction_start_event( + self, interaction_id: str + ) -> InteractionsAPIStreamingResponse: + event_type = "interaction.start" if self._use_legacy else "interaction.created" + return InteractionsAPIStreamingResponse( + event_type=event_type, + id=interaction_id, + object="interaction", + status="in_progress", + model=self.model, + ) + + def _build_content_start_event( + self, interaction_id: str + ) -> InteractionsAPIStreamingResponse: + if self._use_legacy: + return InteractionsAPIStreamingResponse( + event_type="content.start", + id=interaction_id, + object="content", + delta={"type": "text", "text": ""}, + ) + return InteractionsAPIStreamingResponse( + event_type="step.start", + index=0, + step={"type": "model_output", "content": []}, + ) + + def _build_text_delta_event( + self, interaction_id: str, delta_text: str + ) -> InteractionsAPIStreamingResponse: + if self._use_legacy: + return InteractionsAPIStreamingResponse( + event_type="content.delta", + id=interaction_id, + object="content", + delta={"type": "text", "text": delta_text}, + ) + return InteractionsAPIStreamingResponse( + event_type="step.delta", + index=0, + delta={"type": "text", "text": delta_text}, + ) + + def _build_content_stop_event( + self, interaction_id: Optional[str] + ) -> InteractionsAPIStreamingResponse: + if self._use_legacy: + return InteractionsAPIStreamingResponse( + event_type="content.stop", + id=interaction_id, + object="content", + delta={"type": "text", "text": self.collected_text}, + ) + return InteractionsAPIStreamingResponse( + event_type="step.stop", + index=0, + ) + + def _build_completion_event( + self, response_id: str + ) -> InteractionsAPIStreamingResponse: + if self._use_legacy: + return InteractionsAPIStreamingResponse( + event_type="interaction.complete", + id=response_id, + object="interaction", + status="completed", + model=self.model, + outputs=[{"type": "text", "text": self.collected_text}], + ) + return InteractionsAPIStreamingResponse( + event_type="interaction.completed", + id=response_id, + object="interaction", + status="completed", + model=self.model, + steps=[ + { + "type": "model_output", + "content": [{"type": "text", "text": self.collected_text}], + } + ], + ) + + # ------------------------------------------------------------------ + # Per-chunk transform (returns a list of events to enqueue) + # ------------------------------------------------------------------ + + def _events_for_chunk( + self, responses_chunk: ResponsesAPIStreamingResponse + ) -> List[InteractionsAPIStreamingResponse]: """ - Transform a Responses API streaming chunk to an Interactions API streaming chunk. + Translate a single upstream Responses API chunk into the list of + Interactions API events it should produce. - Responses API events: - - output.text.delta -> content.delta - - response.completed -> interaction.complete - - Interactions API events: - - interaction.start - - content.start - - content.delta - - content.stop - - interaction.complete + Returning a list (rather than a single event) lets a chunk emit any + synthetic start events that haven't been sent yet *together with* the + actual delta event, so we never silently drop the chunk's payload. """ if not responses_chunk: - return None + return [] - # Handle OutputTextDeltaEvent -> content.delta + # Text delta: emit any missing start events, then the delta itself. if isinstance(responses_chunk, OutputTextDeltaEvent): delta_text = ( responses_chunk.delta if isinstance(responses_chunk.delta, str) else "" ) self.collected_text += delta_text - - # Fallback: emit interaction.start, and queue content.start carrying this - # delta so the first token is preserved in the stream. - if not self.sent_interaction_start: - self.sent_interaction_start = True - self.sent_content_start = True - self._pending_events.append( - InteractionsAPIStreamingResponse( - event_type="content.start", - id=getattr(responses_chunk, "item_id", None), - object="content", - delta={"type": "text", "text": delta_text}, - ) - ) - return InteractionsAPIStreamingResponse( - event_type="interaction.start", - id=getattr(responses_chunk, "item_id", None) - or f"interaction_{id(self)}", - object="interaction", - status="in_progress", - model=self.model, - ) - - # Fallback: emit content.start if ContentPartAddedEvent never arrived - if not self.sent_content_start: - self.sent_content_start = True - return InteractionsAPIStreamingResponse( - event_type="content.start", - id=getattr(responses_chunk, "item_id", None), - object="content", - delta={"type": "text", "text": delta_text}, - ) - - # Normal path: emit content.delta with type field - return InteractionsAPIStreamingResponse( - event_type="content.delta", - id=getattr(responses_chunk, "item_id", None), - object="content", - delta={"type": "text", "text": delta_text}, + interaction_id = ( + getattr(responses_chunk, "item_id", None) or f"interaction_{id(self)}" ) + if self._interaction_id is None: + self._interaction_id = interaction_id - # Handle ContentPartAddedEvent -> content.start (arrives before text deltas) - if isinstance(responses_chunk, ContentPartAddedEvent): - # Fallback: emit interaction.start if ResponseCreatedEvent never arrived + events: List[InteractionsAPIStreamingResponse] = [] if not self.sent_interaction_start: self.sent_interaction_start = True - return InteractionsAPIStreamingResponse( - event_type="interaction.start", - id=getattr(responses_chunk, "item_id", None) - or f"interaction_{id(self)}", - object="interaction", - status="in_progress", - model=self.model, - ) + events.append(self._build_interaction_start_event(interaction_id)) if not self.sent_content_start: self.sent_content_start = True - return InteractionsAPIStreamingResponse( - event_type="content.start", - id=getattr(responses_chunk, "item_id", None), - object="content", - delta={"type": "text", "text": ""}, - ) - return None + events.append(self._build_content_start_event(interaction_id)) + events.append(self._build_text_delta_event(interaction_id, delta_text)) + return events - # Handle ResponseCreatedEvent or ResponseInProgressEvent -> interaction.start + # Response created / in-progress: synthesize interaction start if we + # haven't already sent one. if isinstance(responses_chunk, (ResponseCreatedEvent, ResponseInProgressEvent)): if not self.sent_interaction_start: self.sent_interaction_start = True @@ -153,177 +226,136 @@ class LiteLLMResponsesInteractionsStreamingIterator: getattr(responses_chunk.response, "id", None) if hasattr(responses_chunk, "response") else None - ) - return InteractionsAPIStreamingResponse( - event_type="interaction.start", - id=response_id or f"interaction_{id(self)}", - object="interaction", - status="in_progress", - model=self.model, - ) + ) or f"interaction_{id(self)}" + if self._interaction_id is None: + self._interaction_id = response_id + return [self._build_interaction_start_event(response_id)] + return [] - # Handle ResponseCompletedEvent -> interaction.complete + # Response completed: emit step.stop (if content was started) followed + # by the terminal completion event. Prefer the interaction id already + # established by earlier events so consumers can correlate the start + # and completion events by id (response.id may differ from the item_id + # used to derive the initial id when the stream starts directly with a + # text delta). if isinstance(responses_chunk, ResponseCompletedEvent): self.finished = True response = responses_chunk.response - - # Send content.stop first if content was started - if self.sent_content_start: - # Note: We'll send this in the iterator, not here - pass - - # Send interaction.complete - return InteractionsAPIStreamingResponse( - event_type="interaction.complete", - id=getattr(response, "id", None) or f"interaction_{id(self)}", - object="interaction", - status="completed", - model=self.model, - outputs=[ - { - "type": "text", - "text": self.collected_text, - } - ], + response_id = ( + self._interaction_id + or getattr(response, "id", None) + or f"interaction_{id(self)}" ) - # For other event types, return None (skip) - return None + terminal: List[InteractionsAPIStreamingResponse] = [] + if self.sent_content_start: + terminal.append(self._build_content_stop_event(response_id)) + terminal.append(self._build_completion_event(response_id)) + self._sent_completion_event = True + return terminal + + return [] + + def _build_terminal_events_on_eof( + self, + ) -> List[InteractionsAPIStreamingResponse]: + """ + Build the events to flush when the upstream stream ends without a + ResponseCompletedEvent. Ensures consumers always observe a terminal + interaction.completed/interaction.complete carrying the full text. + """ + if self._sent_completion_event: + return [] + + fallback_id = self._interaction_id or f"interaction_{id(self)}" + terminal: List[InteractionsAPIStreamingResponse] = [] + if self.sent_content_start: + terminal.append(self._build_content_stop_event(fallback_id)) + if self.sent_interaction_start or self.collected_text: + terminal.append(self._build_completion_event(fallback_id)) + self._sent_completion_event = True + return terminal + + # ------------------------------------------------------------------ + # Iteration + # ------------------------------------------------------------------ def __iter__(self) -> Iterator[InteractionsAPIStreamingResponse]: - """Sync iterator implementation.""" return self def __next__(self) -> InteractionsAPIStreamingResponse: - """Get next chunk in sync mode.""" + if self._pending_events: + return self._pending_events.popleft() + if self.finished: raise StopIteration - # Check if we have a pending interaction.complete to send - if hasattr(self, "_pending_interaction_complete"): - pending: InteractionsAPIStreamingResponse = getattr( - self, "_pending_interaction_complete" - ) - delattr(self, "_pending_interaction_complete") - return pending - - # Drain events queued from a prior chunk (e.g. content.start emitted alongside - # the interaction.start fallback for the first OutputTextDeltaEvent). - if self._pending_events: - return self._pending_events.pop(0) - # Use a loop instead of recursion to avoid stack overflow sync_iterator = cast( SyncResponsesAPIStreamingIterator, self.responses_stream_iterator ) while True: try: - # Get next chunk from responses API stream chunk = next(sync_iterator) - - # Transform chunk (chunk is already a ResponsesAPIStreamingResponse) - transformed = self._transform_responses_chunk_to_interactions_chunk( - chunk - ) - - if transformed: - # If we finished and content was started, send content.stop before interaction.complete - if ( - self.finished - and self.sent_content_start - and transformed.event_type == "interaction.complete" - ): - # Send content.stop first - content_stop = InteractionsAPIStreamingResponse( - event_type="content.stop", - id=transformed.id, - object="content", - delta={"type": "text", "text": self.collected_text}, - ) - # Store the interaction.complete to send next - self._pending_interaction_complete = transformed - return content_stop - return transformed - - # If no transformation, continue to next chunk (loop continues) - except StopIteration: self.finished = True + self._pending_events.extend(self._build_terminal_events_on_eof()) + if self._pending_events: + return self._pending_events.popleft() + raise - # Send final events if needed - if self.sent_content_start: - return InteractionsAPIStreamingResponse( - event_type="content.stop", - object="content", - delta={"type": "text", "text": self.collected_text}, - ) - - raise StopIteration + events = self._events_for_chunk(chunk) + if events: + self._pending_events.extend(events) + return self._pending_events.popleft() def __aiter__(self) -> AsyncIterator[InteractionsAPIStreamingResponse]: - """Async iterator implementation.""" return self async def __anext__(self) -> InteractionsAPIStreamingResponse: - """Get next chunk in async mode.""" + if self._pending_events: + return self._pending_events.popleft() + if self.finished: raise StopAsyncIteration - # Check if we have a pending interaction.complete to send - if hasattr(self, "_pending_interaction_complete"): - pending: InteractionsAPIStreamingResponse = getattr( - self, "_pending_interaction_complete" - ) - delattr(self, "_pending_interaction_complete") - return pending - - # Drain events queued from a prior chunk (e.g. content.start emitted alongside - # the interaction.start fallback for the first OutputTextDeltaEvent). - if self._pending_events: - return self._pending_events.pop(0) - # Use a loop instead of recursion to avoid stack overflow async_iterator = cast( ResponsesAPIStreamingIterator, self.responses_stream_iterator ) while True: try: - # Get next chunk from responses API stream chunk = await async_iterator.__anext__() - - # Transform chunk (chunk is already a ResponsesAPIStreamingResponse) - transformed = self._transform_responses_chunk_to_interactions_chunk( - chunk - ) - - if transformed: - # If we finished and content was started, send content.stop before interaction.complete - if ( - self.finished - and self.sent_content_start - and transformed.event_type == "interaction.complete" - ): - # Send content.stop first - content_stop = InteractionsAPIStreamingResponse( - event_type="content.stop", - id=transformed.id, - object="content", - delta={"type": "text", "text": self.collected_text}, - ) - # Store the interaction.complete to send next - self._pending_interaction_complete = transformed - return content_stop - return transformed - - # If no transformation, continue to next chunk (loop continues) - except StopAsyncIteration: self.finished = True + self._pending_events.extend(self._build_terminal_events_on_eof()) + if self._pending_events: + return self._pending_events.popleft() + raise - # Send final events if needed - if self.sent_content_start: - return InteractionsAPIStreamingResponse( - event_type="content.stop", - object="content", - delta={"type": "text", "text": self.collected_text}, - ) + events = self._events_for_chunk(chunk) + if events: + self._pending_events.extend(events) + return self._pending_events.popleft() - raise StopAsyncIteration + # ------------------------------------------------------------------ + # Backwards-compatible single-chunk transform (used by tests and any + # external callers that drove the iterator chunk-by-chunk pre-fix). + # ------------------------------------------------------------------ + + def _transform_responses_chunk_to_interactions_chunk( + self, + responses_chunk: ResponsesAPIStreamingResponse, + ) -> Optional[InteractionsAPIStreamingResponse]: + """ + Compatibility shim: returns the *first* event produced for this chunk + and queues any remaining events on ``self._pending_events`` so they + are surfaced on subsequent calls/iterations. + + Prefer ``_events_for_chunk`` in new code. + """ + events = self._events_for_chunk(responses_chunk) + if not events: + return None + first = events[0] + if len(events) > 1: + self._pending_events.extend(events[1:]) + return first diff --git a/litellm/interactions/litellm_responses_transformation/transformation.py b/litellm/interactions/litellm_responses_transformation/transformation.py index 100300af7b5..173d4ca8764 100644 --- a/litellm/interactions/litellm_responses_transformation/transformation.py +++ b/litellm/interactions/litellm_responses_transformation/transformation.py @@ -226,29 +226,37 @@ class LiteLLMResponsesInteractionsConfig: - Map status - Extract usage """ - # Extract text from outputs - outputs = [] + # Extract text from outputs and build both `outputs` (legacy) and `steps` (new schema). + outputs: List[Dict[str, Any]] = [] + steps: List[Dict[str, Any]] = [] if hasattr(responses_response, "output") and responses_response.output: for output_item in responses_response.output: # Use getattr with None default to safely access content content = getattr(output_item, "content", None) if content is not None: content_items = content if isinstance(content, list) else [content] + model_output_contents: List[Dict[str, Any]] = [] for content_item in content_items: # Check if content_item has text attribute text = getattr(content_item, "text", None) if text is not None: - outputs.append( - { - "type": "text", - "text": text, - } - ) + # Use independent dict instances so mutations to one + # of `outputs` / `steps` don't leak into the other. + outputs.append({"type": "text", "text": text}) + model_output_contents.append({"type": "text", "text": text}) elif ( isinstance(content_item, dict) and content_item.get("type") == "text" ): - outputs.append(content_item) + outputs.append({**content_item}) + model_output_contents.append({**content_item}) + if model_output_contents: + steps.append( + { + "type": "model_output", + "content": model_output_contents, + } + ) # Convert created_at to ISO string created_at = getattr(responses_response, "created_at", None) @@ -270,12 +278,14 @@ class LiteLLMResponsesInteractionsConfig: else: interactions_status = status - # Build interactions response + # Build interactions response — populate both `outputs` (legacy schema) and + # `steps` (new schema) so callers work regardless of which schema they expect. interactions_response_dict: Dict[str, Any] = { "id": getattr(responses_response, "id", ""), "object": "interaction", "status": interactions_status, "outputs": outputs, + "steps": steps, "model": model or getattr(responses_response, "model", ""), "created": created, } diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index c6eca410fa7..d99cc3d11c7 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -8,25 +8,25 @@ Per OpenAPI spec (https://ai.google.dev/static/api/interactions.openapi.json): Usage: import litellm - + # Create an interaction with a model response = litellm.interactions.create( model="gemini-2.5-flash", input="Hello, how are you?" ) - + # Create an interaction with an agent response = litellm.interactions.create( agent="deep-research-pro-preview-12-2025", input="Research the current state of cancer research" ) - + # Async version response = await litellm.interactions.acreate(...) - + # Get an interaction response = litellm.interactions.get(interaction_id="...") - + # Delete an interaction result = litellm.interactions.delete(interaction_id="...") """ diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index a5a7f9e06e5..561686a3e1b 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -101,10 +101,14 @@ class BaseInteractionsAPIStreamingIterator: ) ) - # Store the completed response (check for status=completed) - if ( - streaming_response - and getattr(streaming_response, "status", None) == "completed" + # Store the completed response. + # Legacy schema signals completion via status="completed". + # New schema (Api-Revision: 2026-05-20) uses event_type="interaction.completed". + # Remove the legacy check after June 8, 2026. + if streaming_response and ( + getattr(streaming_response, "status", None) == "completed" + or getattr(streaming_response, "event_type", None) + == "interaction.completed" ): self.completed_response = streaming_response self._handle_logging_completed_response() diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 9d8bd7523db..b8cdc8210fc 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -11,6 +11,7 @@ def get_supported_openai_params( # noqa: PLR0915 request_type: Literal[ "chat_completion", "embeddings", "transcription" ] = "chat_completion", + base_model: Optional[str] = None, ) -> Optional[list]: """ Returns the supported openai params for a given model + provider @@ -20,6 +21,11 @@ def get_supported_openai_params( # noqa: PLR0915 get_supported_openai_params(model="anthropic.claude-3", custom_llm_provider="bedrock") ``` + Args: + base_model: For Azure, the true underlying model (e.g. ``"azure/gpt-5.2"``) + when the deployment name differs. Used for model-type detection so that + non-standard deployment names route to the correct config. + Returns: - List if custom_llm_provider is mapped - None if unmapped @@ -32,17 +38,21 @@ def get_supported_openai_params( # noqa: PLR0915 if custom_llm_provider in LlmProvidersSet: provider_config = litellm.ProviderConfigManager.get_provider_chat_config( - model=model, provider=LlmProviders(custom_llm_provider) + model=model, + provider=LlmProviders(custom_llm_provider), + base_model=base_model, ) elif custom_llm_provider.split("/")[0] in LlmProvidersSet: provider_config = litellm.ProviderConfigManager.get_provider_chat_config( - model=model, provider=LlmProviders(custom_llm_provider.split("/")[0]) + model=model, + provider=LlmProviders(custom_llm_provider.split("/")[0]), + base_model=base_model, ) else: provider_config = None if provider_config and request_type == "chat_completion": - return provider_config.get_supported_openai_params(model=model) + return provider_config.get_supported_openai_params(model=base_model or model) if custom_llm_provider == "bedrock": return litellm.AmazonConverseConfig().get_supported_openai_params(model=model) @@ -130,16 +140,23 @@ def get_supported_openai_params( # noqa: PLR0915 model=model ) elif custom_llm_provider == "azure": - if litellm.AzureOpenAIO1Config().is_o_series_model(model=model): + _azure_detection_model = base_model or model + if litellm.AzureOpenAIO1Config().is_o_series_model( + model=_azure_detection_model + ): return litellm.AzureOpenAIO1Config().get_supported_openai_params( - model=model + model=_azure_detection_model ) - elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=model): + elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model( + model=_azure_detection_model + ): return litellm.AzureOpenAIGPT5Config().get_supported_openai_params( - model=model + model=_azure_detection_model ) else: - return litellm.AzureOpenAIConfig().get_supported_openai_params(model=model) + return litellm.AzureOpenAIConfig().get_supported_openai_params( + model=_azure_detection_model + ) elif custom_llm_provider == "openrouter": return litellm.OpenrouterConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "vercel_ai_gateway": diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 876f1b167db..2ab037afb0d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -994,10 +994,8 @@ class Logging(LiteLLMLoggingBaseClass): try: # [Non-blocking Extra Debug Information in metadata] if turn_off_message_logging is True: - _metadata["raw_request"] = ( - "redacted by litellm. \ + _metadata["raw_request"] = "redacted by litellm. \ 'litellm.turn_off_message_logging=True'" - ) else: curl_command = self._get_request_curl_command( api_base=additional_args.get("api_base", ""), @@ -1031,12 +1029,8 @@ class Logging(LiteLLMLoggingBaseClass): error=str(e), ) ) - _metadata["raw_request"] = ( - "Unable to Log \ - raw request: {}".format( - str(e) - ) - ) + _metadata["raw_request"] = "Unable to Log \ + raw request: {}".format(str(e)) if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: self.logger_fn( @@ -1769,9 +1763,12 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = 0.0 elif "response_cost" in hidden_params: self.model_call_details["response_cost"] = hidden_params["response_cost"] - elif self.model_call_details.get("response_cost") is not None: + elif ( + existing_cost := self.model_call_details.get("response_cost") + ) is not None and existing_cost != 0: # Preserve response_cost if already calculated (e.g., by pass-through - # handlers like Gemini/Vertex which call completion_cost directly) + # handlers like Gemini/Vertex which call completion_cost directly). + # Do not preserve 0 from failure_handler on intermediate router retries. pass else: self.model_call_details["response_cost"] = self._response_cost_calculator( diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index d6c011990c1..f169f86079a 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1344,6 +1344,7 @@ def _get_dummy_thought_signature() -> str: def convert_to_gemini_tool_call_invoke( message: ChatCompletionAssistantMessage, model: Optional[str] = None, + custom_llm_provider: Optional[str] = None, ) -> List[VertexPartType]: """ OpenAI tool invokes: @@ -1394,7 +1395,10 @@ def convert_to_gemini_tool_call_invoke( ) forward_tool_call_id = bool( - model and VertexGeminiConfig._is_gemini_3_or_newer(model) + model + and VertexGeminiConfig._forward_gemini_function_call_id( + model, custom_llm_provider + ) ) if tool_calls is not None: @@ -1475,6 +1479,7 @@ def convert_to_gemini_tool_call_result( # noqa: PLR0915 message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage], last_message_with_tool_calls: Optional[dict], model: Optional[str] = None, + custom_llm_provider: Optional[str] = None, ) -> Union[VertexPartType, List[VertexPartType]]: """ OpenAI message with a tool result looks like: @@ -1616,14 +1621,16 @@ def convert_to_gemini_tool_call_result( # noqa: PLR0915 name = tool.get("function", {}).get("name", "") # Echo the OpenAI tool_call_id on functionResponse (strip thought-signature suffix). - # Only Gemini 3+ accepts (and returns) an `id` on function_response parts; - # older Gemini models reject the field with a 400. + # Only Google AI Studio Gemini 3+ accepts `id` on function_response parts. + # Vertex AI and older Gemini models reject the field with HTTP 400. from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) gemini_call_id: Optional[str] = None - if model and VertexGeminiConfig._is_gemini_3_or_newer(model): + if model and VertexGeminiConfig._forward_gemini_function_call_id( + model, custom_llm_provider + ): raw_tool_call_id = message.get("tool_call_id") if raw_tool_call_id and isinstance(raw_tool_call_id, str): stripped_id = raw_tool_call_id.split(THOUGHT_SIGNATURE_SEPARATOR, 1)[0] @@ -5583,9 +5590,7 @@ def default_response_schema_prompt(response_schema: dict) -> str: prompt_str = """Use this JSON schema: ```json {} - ```""".format( - response_schema - ) + ```""".format(response_schema) return prompt_str diff --git a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py index 13341f27a61..0a6a4e82c72 100644 --- a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py +++ b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py @@ -1,9 +1,9 @@ """ This is a cache for LangfuseLoggers. -Langfuse Python SDK initializes a thread for each client. +Langfuse Python SDK initializes a thread for each client. -This ensures we do +This ensures we do 1. Proper cleanup of Langfuse initialized clients. 2. Re-use created langfuse clients. """ diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 1ce80207552..0b56eb86d9c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1506,9 +1506,21 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params["metadata"] = {"user_id": value} elif param == "thinking": optional_params["thinking"] = value - elif param == "reasoning_effort" and isinstance(value, str): + elif param == "reasoning_effort": + # Accept both string ("low") and dict ({"effort": "low", + # "summary": "concise"}). The Responses->Chat parser keeps the + # full dict when `summary` is set (see #25359), so a dict here + # is the standard shape Otto/OpenAI-Responses-Bridge callers + # send. Coerce to the effort string before mapping — same + # shape-tolerance the GPT-5 path already implements in + # `_normalize_reasoning_effort_for_chat_completion`. + effort_value = value + if isinstance(effort_value, dict): + effort_value = effort_value.get("effort") + if not isinstance(effort_value, str): + continue mapped_thinking = AnthropicConfig._map_reasoning_effort( - reasoning_effort=value, + reasoning_effort=effort_value, model=model, llm_provider=self.custom_llm_provider or "anthropic", ) @@ -1519,12 +1531,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params["thinking"] = mapped_thinking if AnthropicConfig._is_adaptive_thinking_model(model): mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get( - value + effort_value ) if mapped_effort is None: AnthropicConfig._raise_invalid_reasoning_effort( model=model, - value=value, + value=effort_value, llm_provider=self.custom_llm_provider or "anthropic", ) optional_params["output_config"] = {"effort": mapped_effort} diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 0e198daf089..51a1e739a0f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1476,7 +1476,7 @@ class LiteLLMAnthropicMessagesAdapter: for choice in choices: if choice.delta.content is not None and len(choice.delta.content) > 0: text += choice.delta.content - if choice.delta.tool_calls is not None: + if choice.delta.tool_calls: partial_json = "" for tool in choice.delta.tool_calls: if ( diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index d0780c82d06..d693d50b8e5 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -13,7 +13,6 @@ from typing import Any, AsyncIterator, Dict, List, Optional, cast from litellm._logging import verbose_logger - # --------------------------------------------------------------------------- # SSE parsing helpers (module-level to keep the class lean) # --------------------------------------------------------------------------- diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 9291269d153..734b8ecef16 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -239,7 +239,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) data = {"model": None, "messages": messages, **optional_params} - elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=model): + elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model( + model=litellm_params.get("base_model") or model + ): data = litellm.AzureOpenAIGPT5Config().transform_request( model=model, messages=messages, diff --git a/litellm/llms/azure/chat/o_series_transformation.py b/litellm/llms/azure/chat/o_series_transformation.py index cae7513245c..0a73597a4e4 100644 --- a/litellm/llms/azure/chat/o_series_transformation.py +++ b/litellm/llms/azure/chat/o_series_transformation.py @@ -4,10 +4,10 @@ Support for o1 and o3 model families https://platform.openai.com/docs/guides/reasoning Translations handled by LiteLLM: -- modalities: image => drop param (if user opts in to dropping param) -- role: system ==> translate to role 'user' -- streaming => faked by LiteLLM -- Tools, response_format => drop param (if user opts in to dropping param) +- modalities: image => drop param (if user opts in to dropping param) +- role: system ==> translate to role 'user' +- streaming => faked by LiteLLM +- Tools, response_format => drop param (if user opts in to dropping param) - Logprobs => drop param (if user opts in to dropping param) - Temperature => drop param (if user opts in to dropping param) """ diff --git a/litellm/llms/azure/containers/transformation.py b/litellm/llms/azure/containers/transformation.py index 586b2e379a0..cd897511585 100644 --- a/litellm/llms/azure/containers/transformation.py +++ b/litellm/llms/azure/containers/transformation.py @@ -1,9 +1,16 @@ from typing import Optional +from urllib.parse import parse_qs, urlparse, urlunparse from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.containers.transformation import OpenAIContainerConfig from litellm.types.router import GenericLiteLLMParams +# Endpoint-specific path suffixes that may appear in a deployment's api_base +# (e.g. the responses endpoint URL is stored as api_base for Azure models). +# Strip these before building the containers URL so we always start from the +# resource root (https://resource.cognitiveservices.azure.com). +_AZURE_ENDPOINT_PATHS = ("/openai/responses",) + class AzureContainerConfig(OpenAIContainerConfig): """ @@ -27,6 +34,27 @@ class AzureContainerConfig(OpenAIContainerConfig): litellm_params=GenericLiteLLMParams(api_key=api_key), ) + @staticmethod + def _normalize_api_base(api_base: Optional[str]) -> Optional[str]: + """Strip endpoint-specific path suffixes from api_base to get the resource root.""" + if not api_base: + return api_base + parsed = urlparse(api_base) + path = parsed.path.rstrip("/") + for ep in _AZURE_ENDPOINT_PATHS: + if path.endswith(ep): + return urlunparse( + (parsed.scheme, parsed.netloc, path[: -len(ep)], "", "", "") + ) + return api_base + + @staticmethod + def _extract_api_version(api_base: Optional[str]) -> Optional[str]: + """Return the api-version query param from api_base if present.""" + if not api_base: + return None + return parse_qs(urlparse(api_base).query).get("api-version", [None])[0] + def get_complete_url( self, api_base: Optional[str], @@ -39,10 +67,19 @@ class AzureContainerConfig(OpenAIContainerConfig): {endpoint}/openai/v1/containers when api_version is 'v1', 'latest', or 'preview'; otherwise: {endpoint}/openai/containers + + The deployment's api_base may be the responses endpoint URL + (e.g. .../openai/responses?api-version=2025-04-01-preview). We + prefer the api-version embedded there over the deployment's + api_version field, which may point to an older chat API version. """ + effective_params = dict(litellm_params) + api_version_from_base = self._extract_api_version(api_base) + if api_version_from_base: + effective_params["api_version"] = api_version_from_base return BaseAzureLLM._get_base_azure_url( - api_base=api_base, - litellm_params=litellm_params, + api_base=self._normalize_api_base(api_base), + litellm_params=effective_params, route="/openai/containers", default_api_version="v1", ) diff --git a/litellm/llms/azure_ai/embed/cohere_transformation.py b/litellm/llms/azure_ai/embed/cohere_transformation.py index 64433c21b61..bbbfb60fbde 100644 --- a/litellm/llms/azure_ai/embed/cohere_transformation.py +++ b/litellm/llms/azure_ai/embed/cohere_transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic from OpenAI /v1/embeddings format to Azure AI Cohere's /v1/embed. +Transformation logic from OpenAI /v1/embeddings format to Azure AI Cohere's /v1/embed. Why separate file? Make it easy to see how transformation works diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index b5993040ea0..f64133afa8b 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -1,5 +1,5 @@ """ -Translate between Cohere's `/rerank` format and Azure AI's `/rerank` format. +Translate between Cohere's `/rerank` format and Azure AI's `/rerank` format. """ from typing import Optional diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index b7f4d8e3b2d..263e0c094ce 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -54,6 +54,7 @@ class OCRUsageInfo(LiteLLMPydanticObjectBase): """Usage information from OCR response.""" pages_processed: Optional[int] = None + credits: Optional[float] = None doc_size_bytes: Optional[int] = None model_config = {"extra": "allow"} diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 9dd2b055a12..8b316a587b4 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -450,6 +450,24 @@ class BaseAWSLLM: model_id = BaseAWSLLM.encode_model_id(model_id=model_id) else: model_id = model + # Strip LiteLLM routing prefixes (e.g. "bedrock/", "invoke/", + # "bedrock/invoke/", "bedrock/converse/") that are not part of the + # actual Bedrock model ID. The converse path already does this; the + # invoke path must do the same so that ARN models such as + # bedrock/arn:aws:bedrock:…:inference-profile/global.anthropic.… + # are not forwarded verbatim to the Bedrock API, which would produce + # a malformed URL and cause botocore's EventStreamBuffer to receive + # a JSON error body instead of a binary event-stream — surfaced as a + # misleading ChecksumMismatch (0x223a7b22 == ':{"'). + # Use strip_bedrock_routing_prefix (no break) so compound prefixes + # like "bedrock/invoke/arn:..." are fully stripped in one call. + from litellm.llms.bedrock.common_utils import strip_bedrock_routing_prefix + + model_id = strip_bedrock_routing_prefix(model_id) + # URL-encode ARNs so colons and slashes are safe in the URL path. + if model_id.startswith("arn:"): + model_id = BaseAWSLLM.encode_model_id(model_id=model_id) + return model_id model_id = model_id.replace("invoke/", "", 1) if provider == "llama" and "llama/" in model_id: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index c883ab68dff..d9599b8b9c4 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, List, Optional import httpx from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers +from litellm.litellm_core_utils.litellm_logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) @@ -22,6 +23,7 @@ from litellm.llms.bedrock.common_utils import ( from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse +from litellm.utils import _supports_factory if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -169,6 +171,24 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("model", None) anthropic_request.pop("stream", None) anthropic_request.pop("output_format", None) + if not ( + _supports_factory( + model=model, + custom_llm_provider="bedrock", + key="supports_output_config", + ) + or AnthropicConfig._model_supports_effort_param(model) + ): + if anthropic_request.pop("output_config", None) is not None: + verbose_logger.warning( + "Bedrock Invoke: stripping unsupported `output_config` for " + "model=%s — neither `supports_output_config` nor any " + "`supports_*_reasoning_effort` flag is set in " + "model_prices_and_context_window.json. Add the capability " + "flag to the model JSON entry if this model accepts " + "`output_config`.", + model, + ) if "anthropic_version" not in anthropic_request: anthropic_request["anthropic_version"] = self.anthropic_version diff --git a/litellm/llms/bedrock/claude_platform/common_utils.py b/litellm/llms/bedrock/claude_platform/common_utils.py index 121221518c8..3abb8710de7 100644 --- a/litellm/llms/bedrock/claude_platform/common_utils.py +++ b/litellm/llms/bedrock/claude_platform/common_utils.py @@ -4,7 +4,6 @@ import litellm from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.secret_managers.main import get_secret_str - CLAUDE_PLATFORM_SERVICE_NAME: Literal["aws-external-anthropic"] = ( "aws-external-anthropic" ) diff --git a/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py b/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py index 2747551af81..64a79b73273 100644 --- a/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_titan_g1_transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic from OpenAI /v1/embeddings format to Bedrock Amazon Titan G1 /invoke format. +Transformation logic from OpenAI /v1/embeddings format to Bedrock Amazon Titan G1 /invoke format. Why separate file? Make it easy to see how transformation works diff --git a/litellm/llms/bedrock/embed/cohere_transformation.py b/litellm/llms/bedrock/embed/cohere_transformation.py index 2c0dc834144..9570ff1a14c 100644 --- a/litellm/llms/bedrock/embed/cohere_transformation.py +++ b/litellm/llms/bedrock/embed/cohere_transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic from OpenAI /v1/embeddings format to Bedrock Cohere /invoke format. +Transformation logic from OpenAI /v1/embeddings format to Bedrock Cohere /invoke format. Why separate file? Make it easy to see how transformation works """ diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 151e0e404a0..69b61298d33 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -45,6 +45,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import GenericStreamingChunk from litellm.types.utils import GenericStreamingChunk as GChunk from litellm.types.utils import ModelResponseStream +from litellm.utils import _supports_factory if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -557,7 +558,29 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_request=anthropic_messages_request, ) - # 5a. Remove `custom` field from tools (Bedrock doesn't support it) + # 5a. Bedrock Invoke supports output_config (effort) for Claude 4.6+ models, + # but older models do not — strip it to avoid request rejection. + # Ref: https://github.com/BerriAI/litellm/issues/22797 + if not ( + _supports_factory( + model=model, + custom_llm_provider="bedrock", + key="supports_output_config", + ) + or AnthropicConfig._model_supports_effort_param(model) + ): + if anthropic_messages_request.pop("output_config", None) is not None: + verbose_logger.warning( + "Bedrock Invoke: stripping unsupported `output_config` for " + "model=%s — neither `supports_output_config` nor any " + "`supports_*_reasoning_effort` flag is set in " + "model_prices_and_context_window.json. Add the capability " + "flag to the model JSON entry if this model accepts " + "`output_config`.", + model, + ) + + # 5b. Remove `custom` field from tools (Bedrock doesn't support it) # Claude Code sends `custom: {defer_loading: true}` on tool definitions, # which causes Bedrock to reject the request with "Extra inputs are not permitted" # Ref: https://github.com/BerriAI/litellm/issues/22847 diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index e413bb22b2d..81a56030a5c 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -16,7 +16,6 @@ from litellm.secret_managers.main import get_secret_str from ...openai_like.chat.transformation import OpenAILikeChatConfig - BEDROCK_MANTLE_DEFAULT_REGION = "us-east-1" diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 66acd933416..56b61b66c84 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -1,7 +1,5 @@ -import json -from typing import Any, Optional +from typing import Any, Dict, Optional -from litellm.constants import STREAM_SSE_DONE_STRING from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -9,13 +7,17 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.responses.sse_output_recovery import ( + parse_sse_json_chunk, + record_output_item_chunk, + record_output_text_chunk, +) from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamEvents, ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders -from litellm.utils import CustomStreamWrapper from ..authenticator import Authenticator from ..common_utils import ( @@ -111,86 +113,139 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): raw_response: Any, logging_obj: Any, ): - content_type = (raw_response.headers or {}).get("content-type", "") body_text = raw_response.text or "" - if "text/event-stream" not in content_type.lower(): - trimmed_body = body_text.lstrip() - if not ( - trimmed_body.startswith("event:") - or trimmed_body.startswith("data:") - or "\nevent:" in body_text - or "\ndata:" in body_text - ): - return super().transform_response_api_response( - model=model, - raw_response=raw_response, - logging_obj=logging_obj, - ) + if not self._should_parse_as_sse( + raw_response=raw_response, body_text=body_text + ): + return super().transform_response_api_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) logging_obj.post_call( original_response=raw_response.text, additional_args={"complete_input_dict": {}}, ) - completed_response = None - error_message = None - for chunk in body_text.splitlines(): - stripped_chunk = CustomStreamWrapper._strip_sse_data_from_chunk(chunk) - if not stripped_chunk: - continue - stripped_chunk = stripped_chunk.strip() - if not stripped_chunk: - continue - if stripped_chunk == STREAM_SSE_DONE_STRING: - break - try: - parsed_chunk = json.loads(stripped_chunk) - except json.JSONDecodeError: - continue - if not isinstance(parsed_chunk, dict): - continue - event_type = parsed_chunk.get("type") - if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: - response_payload = parsed_chunk.get("response") - if isinstance(response_payload, dict): - response_payload = dict(response_payload) - if "created_at" in response_payload: - response_payload["created_at"] = _safe_convert_created_field( - response_payload["created_at"] - ) - try: - completed_response = ResponsesAPIResponse(**response_payload) - except Exception: - completed_response = ResponsesAPIResponse.model_construct( - **response_payload - ) - break - if event_type in ( - ResponsesAPIStreamEvents.RESPONSE_FAILED, - ResponsesAPIStreamEvents.ERROR, - ): - error_obj = parsed_chunk.get("error") or ( - parsed_chunk.get("response") or {} - ).get("error") - if error_obj is not None: - if isinstance(error_obj, dict): - error_message = error_obj.get("message") or str(error_obj) - else: - error_message = str(error_obj) - + completed_response, error_message = self._extract_completed_response_from_sse( + body_text=body_text + ) if completed_response is None: raise OpenAIError( message=error_message or raw_response.text, status_code=raw_response.status_code, ) + self._attach_response_headers( + completed_response=completed_response, raw_response=raw_response + ) + return completed_response + + def _should_parse_as_sse(self, raw_response: Any, body_text: str) -> bool: + content_type = (raw_response.headers or {}).get("content-type", "") + if "text/event-stream" in content_type.lower(): + return True + trimmed_body = body_text.lstrip() + return bool( + trimmed_body.startswith("event:") + or trimmed_body.startswith("data:") + or "\nevent:" in body_text + or "\ndata:" in body_text + ) + + def _extract_completed_response_from_sse( + self, body_text: str + ) -> tuple[Optional[ResponsesAPIResponse], Optional[str]]: + completed_response = None + error_message = None + streamed_output_items: Dict[int, dict] = {} + text_only_output_items: Dict[int, dict] = {} + for chunk in body_text.splitlines(): + parsed_chunk = parse_sse_json_chunk(chunk) + if parsed_chunk is None: + continue + + event_type = parsed_chunk.get("type") + if event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: + record_output_item_chunk( + parsed_chunk=parsed_chunk, + output_items=streamed_output_items, + ) + continue + + if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE: + record_output_text_chunk( + parsed_chunk=parsed_chunk, + output_items=streamed_output_items, + text_only_items=text_only_output_items, + ) + continue + + if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED: + # Real OUTPUT_ITEM_DONE events take precedence at any given + # output_index, but text-only items at indices without a + # matching OUTPUT_ITEM_DONE must still be preserved (e.g. + # providers that emit only OUTPUT_TEXT_DONE for some indices). + merged_items: Dict[int, dict] = {**text_only_output_items} + merged_items.update(streamed_output_items) + completed_response = self._build_completed_response_from_chunk( + parsed_chunk=parsed_chunk, + streamed_output_items=merged_items, + ) + break + + if event_type in ( + ResponsesAPIStreamEvents.RESPONSE_FAILED, + ResponsesAPIStreamEvents.ERROR, + ): + extracted_error = self._extract_error_message(parsed_chunk) + if extracted_error is not None: + error_message = extracted_error + + return completed_response, error_message + + def _build_completed_response_from_chunk( + self, parsed_chunk: Dict[str, Any], streamed_output_items: Dict[int, dict] + ) -> Optional[ResponsesAPIResponse]: + response_payload = parsed_chunk.get("response") + if not isinstance(response_payload, dict): + return None + response_payload = dict(response_payload) + if not response_payload.get("output") and streamed_output_items: + response_payload["output"] = [ + item for _, item in sorted(streamed_output_items.items()) + ] + if "created_at" in response_payload: + response_payload["created_at"] = _safe_convert_created_field( + response_payload["created_at"] + ) + try: + return ResponsesAPIResponse(**response_payload) + except Exception: + return ResponsesAPIResponse.model_construct(**response_payload) + + def _extract_error_message(self, parsed_chunk: Dict[str, Any]) -> Optional[str]: + error_obj = parsed_chunk.get("error") or ( + parsed_chunk.get("response") or {} + ).get("error") + if error_obj is None: + return None + if isinstance(error_obj, dict): + return error_obj.get("message") or str(error_obj) + return str(error_obj) + + def _attach_response_headers( + self, + completed_response: ResponsesAPIResponse, + raw_response: Any, + ) -> None: raw_headers = dict(raw_response.headers) processed_headers = process_response_headers(raw_headers) if not hasattr(completed_response, "_hidden_params"): setattr(completed_response, "_hidden_params", {}) completed_response._hidden_params["additional_headers"] = processed_headers completed_response._hidden_params["headers"] = raw_headers - return completed_response def get_complete_url( self, diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 3ab8baf7ba8..81b6a1c7aec 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/litellm/llms/cohere/embed/handler.py @@ -1,5 +1,5 @@ """ -Legacy /v1/embedding handler for Bedrock Cohere. +Legacy /v1/embedding handler for Bedrock Cohere. """ import json diff --git a/litellm/llms/cohere/embed/v1_transformation.py b/litellm/llms/cohere/embed/v1_transformation.py index feca9cb5b88..82c901e7eca 100644 --- a/litellm/llms/cohere/embed/v1_transformation.py +++ b/litellm/llms/cohere/embed/v1_transformation.py @@ -110,15 +110,35 @@ class CohereEmbeddingConfig: additional_args={"complete_input_dict": data}, original_response=response_json, ) + return self._populate_embedding_response( + response_json=response_json, + model_response=model_response, + model=model, + encoding=encoding, + input=input, + ) + + def _populate_embedding_response( + self, + response_json: dict, + model_response: EmbeddingResponse, + model: str, + encoding: Any, + input: list, + ) -> EmbeddingResponse: """ - response + Parse a Cohere embed response body into an OpenAI-style EmbeddingResponse. + + Split out from `_transform_response` so callers that already log + `post_call` themselves (e.g. SageMaker's embedding handler) can reuse + the parsing without triggering a second `post_call`. + + Response shape: { 'object': "list", - 'data': [ - - ] - 'model', - 'usage' + 'data': [...], + 'model', + 'usage', } """ embeddings = response_json["embeddings"] @@ -149,9 +169,6 @@ class CohereEmbeddingConfig: model_response.object = "list" model_response.data = output_data model_response.model = model - input_tokens = 0 - for text in input: - input_tokens += len(encoding.encode(text)) setattr( model_response, diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index 599cd705ebf..501390d840b 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -257,14 +257,19 @@ class GenericContainerHandler: returns_binary = endpoint_config.get("returns_binary", False) is_multipart = endpoint_config.get("is_multipart", False) + # An empty dict passed as `params` to httpx strips any existing query + # string from the URL (e.g. ?api-version=...). Use None instead so + # httpx leaves the URL's own query string intact. + effective_params = query_params or None + try: if method == "GET": response = http_client.get( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) elif method == "DELETE": response = http_client.delete( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) elif method == "POST": if is_multipart and "file" in kwargs: @@ -272,11 +277,11 @@ class GenericContainerHandler: kwargs["file"], headers ) response = http_client.post( - url=url, headers=headers, params=query_params, files=files + url=url, headers=headers, params=effective_params, files=files ) else: response = http_client.post( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) else: raise ValueError(f"Unsupported HTTP method: {method}") @@ -376,14 +381,19 @@ class GenericContainerHandler: returns_binary = endpoint_config.get("returns_binary", False) is_multipart = endpoint_config.get("is_multipart", False) + # An empty dict passed as `params` to httpx strips any existing query + # string from the URL (e.g. ?api-version=...). Use None instead so + # httpx leaves the URL's own query string intact. + effective_params = query_params or None + try: if method == "GET": response = await http_client.get( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) elif method == "DELETE": response = await http_client.delete( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) elif method == "POST": if is_multipart and "file" in kwargs: @@ -391,11 +401,11 @@ class GenericContainerHandler: kwargs["file"], headers ) response = await http_client.post( - url=url, headers=headers, params=query_params, files=files + url=url, headers=headers, params=effective_params, files=files ) else: response = await http_client.post( - url=url, headers=headers, params=query_params + url=url, headers=headers, params=effective_params ) else: raise ValueError(f"Unsupported HTTP method: {method}") diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2ff63cc2d7f..96fdf4494f9 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1409,6 +1409,8 @@ class BaseLLMHTTPHandler: document=document, optional_params=optional_params, headers=headers, + api_key=api_key, + api_base=api_base, ) # All providers return OCRRequestData @@ -1477,6 +1479,8 @@ class BaseLLMHTTPHandler: document=document, optional_params=optional_params, headers=headers, + api_key=api_key, + api_base=api_base, ) # All providers return OCRRequestData @@ -7834,7 +7838,7 @@ class BaseLLMHTTPHandler: response = sync_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_list_response( @@ -7911,7 +7915,7 @@ class BaseLLMHTTPHandler: response = await async_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_list_response( @@ -8001,7 +8005,7 @@ class BaseLLMHTTPHandler: response = sync_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_retrieve_response( @@ -8078,7 +8082,7 @@ class BaseLLMHTTPHandler: response = await async_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_retrieve_response( @@ -8168,7 +8172,7 @@ class BaseLLMHTTPHandler: response = sync_httpx_client.delete( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_delete_response( @@ -8245,7 +8249,7 @@ class BaseLLMHTTPHandler: response = await async_httpx_client.delete( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_delete_response( @@ -8341,7 +8345,7 @@ class BaseLLMHTTPHandler: response = sync_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_file_list_response( @@ -8420,7 +8424,7 @@ class BaseLLMHTTPHandler: response = await async_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_file_list_response( @@ -8508,7 +8512,7 @@ class BaseLLMHTTPHandler: response = sync_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_file_content_response( @@ -8584,7 +8588,7 @@ class BaseLLMHTTPHandler: response = await async_httpx_client.get( url=url, headers=headers, - params=params, + params=params or None, ) return container_provider_config.transform_container_file_content_response( diff --git a/litellm/llms/custom_httpx/mock_transport.py b/litellm/llms/custom_httpx/mock_transport.py index c9844753e0e..ad93cc134ee 100644 --- a/litellm/llms/custom_httpx/mock_transport.py +++ b/litellm/llms/custom_httpx/mock_transport.py @@ -13,7 +13,6 @@ from typing import Tuple import httpx - # --------------------------------------------------------------------------- # Pre-built response templates # --------------------------------------------------------------------------- diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 9b3e3851162..8bb7f605b82 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -1,5 +1,5 @@ """ -Cost calculator for Dashscope Chat models. +Cost calculator for Dashscope Chat models. Handles tiered pricing and prompt caching scenarios. """ diff --git a/litellm/llms/datarobot/chat/transformation.py b/litellm/llms/datarobot/chat/transformation.py index 23ce63c25b2..f81e2420930 100644 --- a/litellm/llms/datarobot/chat/transformation.py +++ b/litellm/llms/datarobot/chat/transformation.py @@ -1,5 +1,5 @@ """ -Support for OpenAI's `/v1/chat/completions` endpoint. +Support for OpenAI's `/v1/chat/completions` endpoint. Calls done in OpenAI/openai.py as DataRobot is openai-compatible. """ diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 276735f4758..e4bfbcb2513 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -1,5 +1,5 @@ """ -Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format. +Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format. """ from typing import Any, Dict, List, Optional, Union diff --git a/litellm/llms/deepseek/cost_calculator.py b/litellm/llms/deepseek/cost_calculator.py index 0f4490cb3df..e652ebeac54 100644 --- a/litellm/llms/deepseek/cost_calculator.py +++ b/litellm/llms/deepseek/cost_calculator.py @@ -1,5 +1,5 @@ """ -Cost calculator for DeepSeek Chat models. +Cost calculator for DeepSeek Chat models. Handles prompt caching scenario. """ diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index 6a59911701b..612fc687ef9 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -22,7 +22,6 @@ from litellm.types.utils import all_litellm_params from ..common_utils import ElevenLabsException - if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.llms.openai import HttpxBinaryResponseContent diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index eaf01c5fe18..d39adf0b6f4 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -4,6 +4,7 @@ from typing import Any, List, Literal, Optional, Tuple, Union, cast import httpx import litellm +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -26,6 +27,7 @@ from litellm.types.utils import ( ProviderSpecificModelInfo, ) from litellm.utils import ( + get_model_cost_mutation_generation, supports_function_calling, supports_reasoning, supports_tool_choice, @@ -112,6 +114,19 @@ class FireworksAIConfig(OpenAIGPTConfig): # Only add tools for models that support function calling if supports_function_calling(model=model, custom_llm_provider="fireworks_ai"): supported_params.append("tools") + supported_params.append("parallel_tool_calls") + else: + # Historically every Fireworks model advertised tool support, so a + # JSON entry that flips `supports_function_calling` to false will + # silently drop `tools` from requests. Surface this so users can + # tell why their tool calls suddenly stop working. + verbose_logger.debug( + "fireworks_ai model %r is marked as not supporting " + "function calling in model_prices_and_context_window.json; " + "`tools` and `parallel_tool_calls` will be dropped from the " + "request.", + model, + ) # Only add tool_choice for models that explicitly support it if supports_tool_choice(model=model, custom_llm_provider="fireworks_ai"): @@ -251,34 +266,100 @@ class FireworksAIConfig(OpenAIGPTConfig): return messages - def get_provider_info(self, model: str) -> ProviderSpecificModelInfo: - # Models that support reasoning_effort - reasoning_supported_models = [ - "qwen3-8b", - "qwen3-32b", - "qwen3-coder-480b-a35b-instruct", - "deepseek-v3p1", - "deepseek-v3p2", - "glm-4p5", - "glm-4p5-air", - "glm-4p6", - "gpt-oss-120b", - "gpt-oss-20b", + # Cached index of fireworks_ai/* entries from litellm.model_cost. Building + # this index requires a full scan of model_cost (tens of thousands of + # entries), so we memoize it. The cache key is (id(model_cost), + # mutation_generation): the generation counter is bumped on every + # register_model / reload path, so add+remove or in-place value + # replacement (which can leave id and len unchanged) still invalidates. + _fireworks_index_cache: Optional[Tuple[int, int, List[Tuple[str, dict]]]] = None + + @classmethod + def _get_fireworks_index(cls) -> List[Tuple[str, dict]]: + model_cost = litellm.model_cost + signature = (id(model_cost), get_model_cost_mutation_generation()) + cached = cls._fireworks_index_cache + if ( + cached is not None + and cached[0] == signature[0] + and cached[1] == signature[1] + ): + return cached[2] + + index: List[Tuple[str, dict]] = [] + for key, model_info in model_cost.items(): + if not key.startswith("fireworks_ai/"): + continue + if not isinstance(model_info, dict): + continue + key_short = key[len("fireworks_ai/") :] + if key_short.startswith("accounts/fireworks/models/"): + key_short = key_short[len("accounts/fireworks/models/") :] + if not key_short: + continue + index.append((key_short, model_info)) + + cls._fireworks_index_cache = (signature[0], signature[1], index) + return index + + @staticmethod + def _matches_on_hyphen_boundary(short_name: str, key_short: str) -> bool: + """Return True if `key_short` appears in `short_name` aligned to + hyphen-separated word boundaries (or end-of-string). This avoids + spurious substring matches like `"some-model"` matching + `"awesome-model"`.""" + if short_name == key_short: + return True + if short_name.startswith(key_short + "-"): + return True + if short_name.endswith("-" + key_short): + return True + return ("-" + key_short + "-") in short_name + + def _get_model_cost_capability(self, model: str, capability: str) -> Optional[bool]: + short_name = model + if short_name.startswith("fireworks_ai/"): + short_name = short_name[len("fireworks_ai/") :] + if short_name.startswith("accounts/fireworks/models/"): + short_name = short_name[len("accounts/fireworks/models/") :] + + candidate_keys = [ + model, + f"fireworks_ai/{short_name}", + f"fireworks_ai/accounts/fireworks/models/{short_name}", ] - # Normalize model name - remove prefix if present - normalized_model = model - if model.startswith("fireworks_ai/"): - normalized_model = model.replace("fireworks_ai/", "") - if normalized_model.startswith("accounts/fireworks/models/"): - normalized_model = normalized_model.replace( - "accounts/fireworks/models/", "" - ) + for candidate_key in candidate_keys: + model_info = litellm.model_cost.get(candidate_key) + if model_info is not None and model_info.get(capability) is not None: + return cast(Optional[bool], model_info.get(capability)) - # Check if model supports reasoning - supports_reasoning_value = any( - reasoning_model in normalized_model - for reasoning_model in reasoning_supported_models + # Fallback: preserve historical substring matching for model name + # variants (e.g. fine-tuned or regionally-suffixed versions of a + # known model). Pick the *longest* matching entry so a more specific + # known model (e.g. "qwen3-8b-instruct") wins over a less specific + # one (e.g. "qwen3-8b") when the query model is more specific still. + # Use hyphen-aligned matching to avoid false positives where a short + # known model name is an unrelated substring of a longer one. + best_match_short: Optional[str] = None + best_match_value: Optional[bool] = None + for key_short, model_info in self._get_fireworks_index(): + if model_info.get(capability) is None: + continue + if not self._matches_on_hyphen_boundary(short_name, key_short): + continue + if best_match_short is None or len(key_short) > len(best_match_short): + best_match_short = key_short + best_match_value = cast(Optional[bool], model_info.get(capability)) + + return best_match_value + + def get_provider_info(self, model: str) -> ProviderSpecificModelInfo: + supports_function_calling_value = self._get_model_cost_capability( + model=model, capability="supports_function_calling" + ) + supports_reasoning_value = self._get_model_cost_capability( + model=model, capability="supports_reasoning" ) provider_specific_model_info: ProviderSpecificModelInfo = { @@ -288,9 +369,16 @@ class FireworksAIConfig(OpenAIGPTConfig): "supports_vision": True, # via document inlining } + if supports_function_calling_value is not None: + provider_specific_model_info["supports_function_calling"] = ( + supports_function_calling_value + ) + # Only include supports_reasoning if True if supports_reasoning_value: - provider_specific_model_info["supports_reasoning"] = True + provider_specific_model_info["supports_reasoning"] = ( + supports_reasoning_value + ) return provider_specific_model_info diff --git a/litellm/llms/gemini/agents/transformation.py b/litellm/llms/gemini/agents/transformation.py index 150918c4737..f6e0b95cf28 100644 --- a/litellm/llms/gemini/agents/transformation.py +++ b/litellm/llms/gemini/agents/transformation.py @@ -23,7 +23,6 @@ from litellm.types.agents import ( AgentVersionsResponse, ) - # Keys inside litellm_params that should be forwarded to the Gemini # create-agent body verbatim. _GEMINI_AGENT_BODY_KEYS = ("base_agent", "instructions", "base_environment") diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index 16e17dcc876..b69b7e1913e 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -164,5 +164,8 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig): # If conversion fails, leave as is and let the API handle it pass return _gemini_convert_messages_with_history( - messages=messages, model=model, litellm_params=litellm_params + messages=messages, + model=model, + litellm_params=litellm_params, + custom_llm_provider="gemini", ) diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index 73435c8db6a..b18b6a28ce4 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -6,13 +6,18 @@ Per OpenAPI spec (https://ai.google.dev/static/api/interactions.openapi.json): - Get: GET https://generativelanguage.googleapis.com/{api_version}/interactions/{interaction_id} - Delete: DELETE https://generativelanguage.googleapis.com/{api_version}/interactions/{interaction_id} -This is a thin wrapper - no transformation needed since we follow the spec directly. +Schema versioning: +- Default (Api-Revision: 2026-05-20): new `steps` schema. +- Legacy (Api-Revision: 2026-05-07): old `outputs` schema, controlled via + litellm.use_legacy_interactions_schema = True. Remove flag after June 8, 2026. """ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import httpx +import litellm + from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -84,6 +89,15 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): api_key = GeminiModelInfo.get_api_key(litellm_params.get("api_key")) if api_key: headers["x-goog-api-key"] = api_key + + # Inject the Api-Revision header to select the response schema. + # Default to the new `steps` schema unless the operator has opted out. + # Remove this conditional after June 8, 2026 and always use 2026-05-20. + if litellm.use_legacy_interactions_schema: + headers["Api-Revision"] = "2026-05-07" + else: + headers["Api-Revision"] = "2026-05-20" + return headers def get_complete_url( @@ -119,8 +133,19 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): headers: dict, ) -> Dict: """ - Build request body per OpenAPI spec - minimal transformation. + Build request body per OpenAPI spec. + + When on the new schema (use_legacy_interactions_schema=False, the default): + - ``response_mime_type`` is folded into ``response_format`` and stripped from + the body (the field was removed in Api-Revision 2026-05-20). + - ``generation_config.image_config`` is moved to a ``response_format`` entry + with ``"type": "image"`` (also removed from generation_config in 2026-05-20). + + When on the legacy schema (use_legacy_interactions_schema=True): + - All fields are forwarded as-is. """ + use_legacy: bool = litellm.use_legacy_interactions_schema + request_body: Dict[str, Any] = {} # Model or Agent (one required) @@ -135,24 +160,81 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): if input is not None: request_body["input"] = input - # Pass through optional params directly (they match the spec) + # Pass through optional params — legacy schema keeps all fields as-is. optional_keys = [ "tools", "system_instruction", - "generation_config", "stream", "store", "background", "environment", "response_modalities", - "response_format", - "response_mime_type", "previous_interaction_id", ] for key in optional_keys: if optional_params.get(key) is not None: request_body[key] = optional_params[key] + if use_legacy: + # Legacy schema: forward response_mime_type and response_format as-is. + for key in ("response_format", "response_mime_type", "generation_config"): + if optional_params.get(key) is not None: + request_body[key] = optional_params[key] + else: + # New schema (Api-Revision: 2026-05-20): + # response_mime_type is removed — fold it into response_format. + response_format = optional_params.get("response_format") + response_mime_type = optional_params.get("response_mime_type") + + if ( + response_mime_type + and not isinstance(response_format, list) + and ( + not isinstance(response_format, dict) + or "mime_type" not in response_format + ) + ): + # Wrap the legacy schema into the new polymorphic format. + new_rf: Dict[str, Any] = { + "type": "text", + "mime_type": response_mime_type, + } + if response_format is not None: + new_rf["schema"] = response_format + response_format = new_rf + + if response_format is not None: + request_body["response_format"] = response_format + + # image_config moves out of generation_config into response_format. + generation_config: Optional[Dict[str, Any]] = optional_params.get( + "generation_config" + ) + if generation_config is not None: + image_config = None + if isinstance(generation_config, dict): + generation_config = dict( + generation_config + ) # avoid mutating the caller's dict + image_config = generation_config.pop("image_config", None) + if not generation_config: + generation_config = None + + if generation_config is not None: + request_body["generation_config"] = generation_config + + if image_config is not None: + # Move image_config to response_format with type=image. + image_rf: Dict[str, Any] = {"type": "image", **image_config} + existing_rf = request_body.get("response_format") + if existing_rf is None: + request_body["response_format"] = image_rf + elif isinstance(existing_rf, list): + request_body["response_format"] = [*existing_rf, image_rf] + else: + # Convert single entry to array for multimodal output. + request_body["response_format"] = [existing_rf, image_rf] + return request_body def transform_response( diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index c7116940b22..9714c8a3923 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -55,7 +55,7 @@ def _convert_image_to_gemini_format(image_file) -> Dict[str, str]: def _usage_video_resolution_from_parameters( - parameters: Dict[str, Any] + parameters: Dict[str, Any], ) -> Optional[str]: """Normalize Veo ``parameters.resolution`` for usage and cost tracking.""" res = parameters.get("resolution") diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index 314bf2f8a36..b9804605454 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank` format. +Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank` format. Why separate file? Make it easy to see how transformation works """ diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index ad4416925a6..56be754fc34 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic from Cohere's /v1/rerank format to Jina AI's `/v1/rerank` format. +Transformation logic from Cohere's /v1/rerank format to Jina AI's `/v1/rerank` format. Why separate file? Make it easy to see how transformation works diff --git a/litellm/llms/lm_studio/embed/transformation.py b/litellm/llms/lm_studio/embed/transformation.py index 1285550c30f..87f4f6e73d5 100644 --- a/litellm/llms/lm_studio/embed/transformation.py +++ b/litellm/llms/lm_studio/embed/transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic from OpenAI /v1/embeddings format to LM Studio's `/v1/embeddings` format. +Transformation logic from OpenAI /v1/embeddings format to LM Studio's `/v1/embeddings` format. Why separate file? Make it easy to see how transformation works diff --git a/litellm/llms/novita/chat/transformation.py b/litellm/llms/novita/chat/transformation.py index c05d2d7b2c5..5a64a124ade 100644 --- a/litellm/llms/novita/chat/transformation.py +++ b/litellm/llms/novita/chat/transformation.py @@ -1,5 +1,5 @@ """ -Support for OpenAI's `/v1/chat/completions` endpoint. +Support for OpenAI's `/v1/chat/completions` endpoint. Calls done in OpenAI/openai.py as Novita AI is openai-compatible. diff --git a/litellm/llms/nvidia_nim/chat/transformation.py b/litellm/llms/nvidia_nim/chat/transformation.py index b8f8b04eb53..2ef92a90626 100644 --- a/litellm/llms/nvidia_nim/chat/transformation.py +++ b/litellm/llms/nvidia_nim/chat/transformation.py @@ -1,7 +1,7 @@ """ -Nvidia NIM endpoint: https://docs.api.nvidia.com/nim/reference/databricks-dbrx-instruct-infer +Nvidia NIM endpoint: https://docs.api.nvidia.com/nim/reference/databricks-dbrx-instruct-infer -This is OpenAI compatible +This is OpenAI compatible This file only contains param mapping logic diff --git a/litellm/llms/nvidia_nim/embed.py b/litellm/llms/nvidia_nim/embed.py index 24c6cc34e4d..61c8e8244e4 100644 --- a/litellm/llms/nvidia_nim/embed.py +++ b/litellm/llms/nvidia_nim/embed.py @@ -1,7 +1,7 @@ """ Nvidia NIM embeddings endpoint: https://docs.api.nvidia.com/nim/reference/nvidia-nv-embedqa-e5-v5-infer -This is OpenAI compatible +This is OpenAI compatible This file only contains param mapping logic diff --git a/litellm/llms/openai/chat/o_series_transformation.py b/litellm/llms/openai/chat/o_series_transformation.py index 02ae2cc9750..8db7ecf7b3a 100644 --- a/litellm/llms/openai/chat/o_series_transformation.py +++ b/litellm/llms/openai/chat/o_series_transformation.py @@ -1,14 +1,14 @@ """ -Support for o1/o3 model family +Support for o1/o3 model family https://platform.openai.com/docs/guides/reasoning Translations handled by LiteLLM: -- modalities: image => drop param (if user opts in to dropping param) -- role: system ==> translate to role 'user' -- streaming => faked by LiteLLM -- Tools, response_format => drop param (if user opts in to dropping param) -- Logprobs => drop param (if user opts in to dropping param) +- modalities: image => drop param (if user opts in to dropping param) +- role: system ==> translate to role 'user' +- streaming => faked by LiteLLM +- Tools, response_format => drop param (if user opts in to dropping param) +- Logprobs => drop param (if user opts in to dropping param) """ from typing import Any, Coroutine, List, Literal, Optional, Union, cast, overload diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index c13a976c1b9..381f215a13f 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -201,7 +201,7 @@ class BaseOpenAILLM: @staticmethod def get_openai_client_initialization_param_fields( - client_type: Literal["openai", "azure"] + client_type: Literal["openai", "azure"], ) -> Tuple[str, ...]: """Returns a tuple of fields that are used to initialize the OpenAI client""" if client_type == "openai": diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index b7d5340d8d4..5043d25ee37 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -126,9 +126,21 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Dict: - """No transform applied since inputs are in OpenAI spec already""" + """Strip Anthropic-only `cache_control` markers before sending to OpenAI. + + OpenAI's Responses API rejects unknown fields on input content blocks + with HTTP 400 ("Unknown parameter: 'input[0].content[0].cache_control'"). + Chat Completions strips these in + `remove_cache_control_flag_from_messages_and_tools`; mirror that here. + """ input = self._validate_input_param(input) + tools = response_api_optional_request_params.get("tools") + input, tools = self.remove_cache_control_flag_from_input_and_tools( + model=model, input=input, tools=tools + ) + if tools is not None: + response_api_optional_request_params["tools"] = tools final_request_params = dict( ResponsesAPIRequestParams( model=model, input=input, **response_api_optional_request_params @@ -137,6 +149,38 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return final_request_params + def remove_cache_control_flag_from_input_and_tools( + self, + model: str, # allows overrides to selectively run this + input: Union[str, ResponseInputParam], + tools: Optional[List[ALL_RESPONSES_API_TOOL_PARAMS]] = None, + ) -> Tuple[ + Union[str, ResponseInputParam], + Optional[List[ALL_RESPONSES_API_TOOL_PARAMS]], + ]: + """Sibling of `remove_cache_control_flag_from_messages_and_tools` on + the chat path. Strips Anthropic-only `cache_control` markers from + Responses API input content blocks and tools. + + `filter_value_from_dict` mutates each dict in place, so the same + objects are returned. + """ + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + filter_value_from_dict, + ) + + if isinstance(input, list): + for item in input: + if isinstance(item, dict): + filter_value_from_dict(cast(dict, item), "cache_control") + + if tools is not None: + for tool in tools: + if isinstance(tool, dict): + filter_value_from_dict(cast(dict, tool), "cache_control") + + return input, tools + def _validate_input_param( self, input: Union[str, ResponseInputParam] ) -> Union[str, ResponseInputParam]: @@ -604,6 +648,12 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): url = str(parsed_url.copy_with(path=compact_path)) input = self._validate_input_param(input) + tools = response_api_optional_request_params.get("tools") + input, tools = self.remove_cache_control_flag_from_input_and_tools( + model=model, input=input, tools=tools + ) + if tools is not None: + response_api_optional_request_params["tools"] = tools data = dict( ResponsesAPIRequestParams( model=model, input=input, **response_api_optional_request_params diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index a55716a5e50..9c2293eb3f1 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -49,7 +49,6 @@ from litellm.types.utils import ( ) from litellm.llms.openrouter.common_utils import OpenRouterException - if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj else: diff --git a/litellm/llms/reducto/__init__.py b/litellm/llms/reducto/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/litellm/llms/reducto/__init__.py @@ -0,0 +1 @@ + diff --git a/litellm/llms/reducto/common.py b/litellm/llms/reducto/common.py new file mode 100644 index 00000000000..4e7d96dbe87 --- /dev/null +++ b/litellm/llms/reducto/common.py @@ -0,0 +1,159 @@ +import base64 +import binascii +from collections import defaultdict +from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Optional, Tuple + +from litellm.constants import request_timeout + +REDUCTO_API_BASE = "https://platform.reducto.ai" +REDUCTO_ID_PREFIX = "reducto://" + +if TYPE_CHECKING: + from litellm.llms.base_llm.ocr.transformation import OCRPage + + +def _normalize_api_base(api_base: Optional[str]) -> str: + return (api_base or REDUCTO_API_BASE).rstrip("/") + + +def _raise_bad_request(message: str, model: str) -> NoReturn: + import litellm + + raise litellm.BadRequestError( + message=message, + model=model, + llm_provider="reducto", + ) + + +def extract_file_id_or_bytes( + source_url: str, + model: str, +) -> Tuple[Optional[str], Optional[bytes], Optional[str]]: + if source_url.startswith(REDUCTO_ID_PREFIX): + return source_url, None, None + + if source_url.startswith("http://") or source_url.startswith("https://"): + _raise_bad_request( + "Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first.", + model=model, + ) + + if not source_url.startswith("data:"): + _raise_bad_request( + "Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing.", + model=model, + ) + + try: + header, encoded = source_url.split(",", 1) + except ValueError: + _raise_bad_request("Invalid Reducto data URI provided.", model=model) + + if ";base64" not in header: + _raise_bad_request( + "Reducto only supports base64-encoded data URIs.", model=model + ) + + mime = header.removeprefix("data:").split(";")[0] or "application/octet-stream" + try: + raw_bytes = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError): + _raise_bad_request("Invalid Reducto base64 payload provided.", model=model) + + return None, raw_bytes, mime + + +def _extract_file_id_from_upload_response(response: Any) -> str: + try: + payload = response.json() + except ValueError as exc: + raise ValueError( + "Reducto /upload returned a non-JSON 200 response: {}".format(response.text) + ) from exc + file_id = (payload or {}).get("file_id") if isinstance(payload, dict) else None + if not isinstance(file_id, str) or not file_id: + raise ValueError( + "Reducto /upload returned 200 without a file_id; got payload={}".format( + payload + ) + ) + return file_id + + +def upload_bytes_sync( + raw_bytes: bytes, + mime: Optional[str], + api_key: str, + api_base: Optional[str], +) -> str: + import litellm + + response = litellm.module_level_client.post( + url="{}{}".format(_normalize_api_base(api_base), "/upload"), + headers={"Authorization": f"Bearer {api_key}"}, + files={"file": ("document", raw_bytes, mime or "application/octet-stream")}, + timeout=request_timeout, + ) + response.raise_for_status() + return _extract_file_id_from_upload_response(response) + + +async def upload_bytes_async( + raw_bytes: bytes, + mime: Optional[str], + api_key: str, + api_base: Optional[str], +) -> str: + import litellm + + response = await litellm.module_level_aclient.post( + url="{}{}".format(_normalize_api_base(api_base), "/upload"), + headers={"Authorization": f"Bearer {api_key}"}, + files={"file": ("document", raw_bytes, mime or "application/octet-stream")}, + timeout=request_timeout, + ) + response.raise_for_status() + return _extract_file_id_from_upload_response(response) + + +def build_pages_from_reducto(result: Dict[str, Any]) -> List["OCRPage"]: + from litellm.llms.base_llm.ocr.transformation import OCRPage + + chunks = result.get("chunks", []) or [] + blocks_by_page: Dict[int, List[Dict[str, Any]]] = defaultdict(list) + + for chunk in chunks: + for block in chunk.get("blocks", []) or []: + page_no = (block.get("bbox") or {}).get("page") + if page_no is None: + continue + try: + normalized_page = int(page_no) + except (TypeError, ValueError): + continue + blocks_by_page[normalized_page].append(block) + + if not blocks_by_page: + fallback_markdown = "\n\n".join( + chunk.get("content", "") for chunk in chunks if chunk.get("content") + ) + if fallback_markdown == "": + return [] + return [OCRPage(index=0, markdown=fallback_markdown)] + + pages: List["OCRPage"] = [] + for page_no, blocks in sorted(blocks_by_page.items()): + markdown = "\n\n".join( + block.get("content", "") for block in blocks if block.get("content") + ) + page_index = max(page_no - 1, 0) + page = OCRPage( + index=page_index, + markdown=markdown, + ) + # OCRPage accepts extra keys at runtime; assign blocks after construction + # so static typing does not reject provider-specific metadata. + setattr(page, "blocks", blocks) + pages.append(page) + return pages diff --git a/litellm/llms/reducto/ocr/__init__.py b/litellm/llms/reducto/ocr/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/litellm/llms/reducto/ocr/__init__.py @@ -0,0 +1 @@ + diff --git a/litellm/llms/reducto/ocr/transformation.py b/litellm/llms/reducto/ocr/transformation.py new file mode 100644 index 00000000000..cc338ecc484 --- /dev/null +++ b/litellm/llms/reducto/ocr/transformation.py @@ -0,0 +1,241 @@ +from typing import Any, Dict, Optional, Tuple + +import httpx + +from litellm.llms.base_llm.ocr.transformation import ( + BaseOCRConfig, + DocumentType, + OCRRequestData, + OCRResponse, + OCRUsageInfo, +) +from litellm.llms.reducto.common import ( + REDUCTO_API_BASE, + build_pages_from_reducto, + extract_file_id_or_bytes, + upload_bytes_async, + upload_bytes_sync, +) + + +class _BaseReductoOCRConfig(BaseOCRConfig): + def map_ocr_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + ) -> dict: + mapped_params = dict(optional_params) + supported_params = self.get_supported_ocr_params(model=model) + for param, value in non_default_params.items(): + if param in supported_params: + mapped_params[param] = value + return mapped_params + + def validate_environment( + self, + headers: Dict, + model: str, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + litellm_params: Optional[dict] = None, + **kwargs, + ) -> Dict: + from litellm.secret_managers.main import get_secret_str + + resolved_key = api_key or get_secret_str("REDUCTO_API_KEY") + if resolved_key is None: + raise ValueError( + "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" + ) + + return { + "Authorization": f"Bearer {resolved_key}", + "Content-Type": "application/json", + **headers, + } + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: dict, + litellm_params: Optional[dict] = None, + **kwargs, + ) -> str: + return "{}/parse".format((api_base or REDUCTO_API_BASE).rstrip("/")) + + def _get_source_url(self, document: DocumentType, model: str) -> str: + source_url = document.get("document_url") or document.get("image_url") + if source_url is None: + raise ValueError( + "Reducto expected OCR preprocessing to produce document_url or image_url for model={}".format( + model + ) + ) + return source_url + + @staticmethod + def _resolve_credentials( + api_key: Optional[str], api_base: Optional[str] + ) -> Tuple[str, str]: + from litellm.secret_managers.main import get_secret_str + + resolved_key = api_key or get_secret_str("REDUCTO_API_KEY") + if resolved_key is None: + raise ValueError( + "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" + ) + resolved_base = (api_base or REDUCTO_API_BASE).rstrip("/") + return resolved_key, resolved_base + + def _ensure_file_id_sync( + self, + model: str, + document: DocumentType, + api_key: Optional[str], + api_base: Optional[str], + ) -> str: + source_url = self._get_source_url(document=document, model=model) + file_id, raw_bytes, mime = extract_file_id_or_bytes(source_url, model=model) + if file_id is not None: + return file_id + resolved_key, resolved_base = self._resolve_credentials(api_key, api_base) + return upload_bytes_sync( + raw_bytes=raw_bytes or b"", + mime=mime, + api_key=resolved_key, + api_base=resolved_base, + ) + + async def _ensure_file_id_async( + self, + model: str, + document: DocumentType, + api_key: Optional[str], + api_base: Optional[str], + ) -> str: + source_url = self._get_source_url(document=document, model=model) + file_id, raw_bytes, mime = extract_file_id_or_bytes(source_url, model=model) + if file_id is not None: + return file_id + resolved_key, resolved_base = self._resolve_credentials(api_key, api_base) + return await upload_bytes_async( + raw_bytes=raw_bytes or b"", + mime=mime, + api_key=resolved_key, + api_base=resolved_base, + ) + + def transform_ocr_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: Any, + **kwargs, + ) -> OCRResponse: + response_json = raw_response.json() + result = response_json.get("result", response_json) or {} + usage = response_json.get("usage", {}) or {} + response = OCRResponse( + pages=build_pages_from_reducto(result), + model=model, + usage_info=OCRUsageInfo( + pages_processed=usage.get("num_pages"), + credits=usage.get("credits"), + ), + object="ocr", + ) + response._hidden_params["reducto_raw"] = response_json + return response + + +class ReductoParseV3Config(_BaseReductoOCRConfig): + def get_supported_ocr_params(self, model: str) -> list: + return ["formatting", "retrieval", "settings"] + + def transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + file_id = self._ensure_file_id_sync( + model=model, + document=document, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + ) + return OCRRequestData(data={"input": file_id, **optional_params}, files=None) + + async def async_transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + file_id = await self._ensure_file_id_async( + model=model, + document=document, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + ) + return OCRRequestData(data={"input": file_id, **optional_params}, files=None) + + +class ReductoParseLegacyConfig(_BaseReductoOCRConfig): + def get_supported_ocr_params(self, model: str) -> list: + return ["enhance"] + + def _build_legacy_body(self, file_id: str, optional_params: dict) -> Dict[str, Any]: + body: Dict[str, Any] = {"document_url": file_id} + enhance = optional_params.get("enhance") + if enhance is not None: + body["options"] = {"enhance": enhance} + return body + + def transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + file_id = self._ensure_file_id_sync( + model=model, + document=document, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + ) + return OCRRequestData( + data=self._build_legacy_body( + file_id=file_id, optional_params=optional_params + ), + files=None, + ) + + async def async_transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + file_id = await self._ensure_file_id_async( + model=model, + document=document, + api_key=kwargs.get("api_key"), + api_base=kwargs.get("api_base"), + ) + return OCRRequestData( + data=self._build_legacy_body( + file_id=file_id, optional_params=optional_params + ), + files=None, + ) diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index efbb218f575..de7be18e8ba 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -578,7 +578,7 @@ class SagemakerLLM(BaseAWSLLM): logger_fn=None, ): """ - Supports both Huggingface Jumpstart embeddings and Voyage models + Supports Hugging Face (TGI), Voyage, and Cohere embedding endpoints """ ### BOTO3 INIT import boto3 diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index 3e4e2460cdb..8fd32bc4460 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -1,7 +1,7 @@ """ Translate from OpenAI's `/v1/chat/completions` to Sagemaker's `/invoke` -In the Huggingface TGI format. +In the Huggingface TGI format. """ import json diff --git a/litellm/llms/sagemaker/embedding/cohere_transformation.py b/litellm/llms/sagemaker/embedding/cohere_transformation.py new file mode 100644 index 00000000000..fdb67202ebb --- /dev/null +++ b/litellm/llms/sagemaker/embedding/cohere_transformation.py @@ -0,0 +1,141 @@ +""" +Translate from OpenAI's `/v1/embeddings` to Sagemaker's `/invoke` + +In the native Cohere embed format for self-hosted Cohere endpoints +(AWS Marketplace / JumpStart). Cohere containers expect +`{"texts": [...], "input_type": "..."}` and reject the HuggingFace TGI shape +`{"inputs": [...]}` with `422 EmbedReqV2.inputs is of type string but should +be of type Object`. + +Reference: https://docs.cohere.com/v2/reference/embed +""" + +from typing import TYPE_CHECKING, Any, List, Optional, Union, cast + +if TYPE_CHECKING: + from litellm.types.llms.openai import AllEmbeddingInputValues + +from httpx._models import Headers, Response + +import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.llms.bedrock.embed.cohere_transformation import ( + BedrockCohereEmbeddingConfig, +) +from litellm.llms.cohere.embed.v1_transformation import CohereEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +from ..common_utils import SagemakerError + + +class SagemakerCohereEmbeddingConfig(BaseEmbeddingConfig): + """ + SageMaker invoke payload for self-hosted Cohere embed models. + """ + + def __init__(self) -> None: + pass + + def get_supported_openai_params(self, model: str) -> List[str]: + return ["encoding_format", "dimensions", "input_type"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + optional_params = BedrockCohereEmbeddingConfig().map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + ) + if "input_type" in non_default_params: + optional_params["input_type"] = non_default_params["input_type"] + return optional_params + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, Headers] + ) -> BaseLLMException: + return SagemakerError( + message=error_message, status_code=status_code, headers=headers + ) + + def transform_embedding_request( + self, + model: str, + input: "AllEmbeddingInputValues", + optional_params: dict, + headers: dict, + ) -> dict: + """ + Transform embedding request for Cohere models on SageMaker + """ + if isinstance(input, str): + input_list: List[str] = [input] + elif isinstance(input, list): + if input and (isinstance(input[0], list) or isinstance(input[0], int)): + raise ValueError("Input must be a list of strings") + input_list = cast(List[str], input) + else: + input_list = [str(input)] + + return dict( + BedrockCohereEmbeddingConfig()._transform_request( + model=model, + input=input_list, + inference_params=optional_params, + ) + ) + + def transform_embedding_response( + self, + model: str, + raw_response: Response, + model_response: "EmbeddingResponse", + logging_obj: Any, + api_key: Optional[str] = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> "EmbeddingResponse": + """ + Transform embedding response for Cohere models on SageMaker. + + Uses `CohereEmbeddingConfig._populate_embedding_response` (not + `_transform_response`) so we do not log `post_call` a second time + — the SageMaker embedding handler already logs `post_call` before + invoking this transform. + """ + input_value = ( + logging_obj.model_call_details.get("input") + or request_data.get("texts") + or request_data.get("images") + or [] + ) + if isinstance(input_value, str): + input_value = [input_value] + + return CohereEmbeddingConfig()._populate_embedding_response( + response_json=raw_response.json(), + model_response=model_response, + model=model, + encoding=litellm.encoding, + input=input_value, + ) + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[Any], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate environment for SageMaker Cohere embeddings + """ + return {"Content-Type": "application/json"} diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 04430171187..5e2aa99534f 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -1,7 +1,7 @@ """ Translate from OpenAI's `/v1/embeddings` to Sagemaker's `/invoke` -In the Huggingface TGI format. +In the Huggingface TGI format. """ from typing import TYPE_CHECKING, Any, List, Optional, Union @@ -11,12 +11,13 @@ if TYPE_CHECKING: from httpx._models import Headers, Response -from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.types.utils import Usage, EmbeddingResponse +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig +from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import SagemakerError +from .cohere_transformation import SagemakerCohereEmbeddingConfig class SagemakerEmbeddingConfig(BaseEmbeddingConfig): @@ -38,17 +39,20 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): Returns: Appropriate embedding config instance """ - if "voyage" in model.lower(): + model_lower = model.lower() + if "voyage" in model_lower: return VoyageEmbeddingConfig() - else: - return cls() + if "cohere" in model_lower: + return SagemakerCohereEmbeddingConfig() + return cls() def get_supported_openai_params(self, model: str) -> List[str]: - # Check if this is an embedding model - if "voyage" in model.lower(): + model_lower = model.lower() + if "voyage" in model_lower: return VoyageEmbeddingConfig().get_supported_openai_params(model) - else: - return [] + if "cohere" in model_lower: + return SagemakerCohereEmbeddingConfig().get_supported_openai_params(model) + return [] def map_openai_params( self, diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index 0ae351783e8..dd307ddf496 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -207,7 +207,7 @@ def resolve_resource_group(sources: List[Source]) -> Optional[str]: def _parse_service_key_once( - service_key: Optional[Union[str, dict]] + service_key: Optional[Union[str, dict]], ) -> Optional[Dict[str, Any]]: """ Pre-parse service_key if it's a string to avoid repeated JSON parsing. diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 3e590680a75..23bb6f44757 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -14,7 +14,6 @@ from ...openai_like.chat.transformation import OpenAIGPTConfig from ..utils import SnowflakeBaseConfig - if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj diff --git a/litellm/llms/together_ai/chat.py b/litellm/llms/together_ai/chat.py index 7efb12fc1b2..238849cc1ec 100644 --- a/litellm/llms/together_ai/chat.py +++ b/litellm/llms/together_ai/chat.py @@ -1,5 +1,5 @@ """ -Support for OpenAI's `/v1/chat/completions` endpoint. +Support for OpenAI's `/v1/chat/completions` endpoint. Calls done in OpenAI/openai.py as TogetherAI is openai-compatible. diff --git a/litellm/llms/together_ai/embed.py b/litellm/llms/together_ai/embed.py index 577df0256cc..6a39b94acfc 100644 --- a/litellm/llms/together_ai/embed.py +++ b/litellm/llms/together_ai/embed.py @@ -1,5 +1,5 @@ """ -Support for OpenAI's `/v1/embeddings` endpoint. +Support for OpenAI's `/v1/embeddings` endpoint. Calls done in OpenAI/openai.py as TogetherAI is openai-compatible. diff --git a/litellm/llms/together_ai/rerank/transformation.py b/litellm/llms/together_ai/rerank/transformation.py index 63b593dfe42..f4d642bd25a 100644 --- a/litellm/llms/together_ai/rerank/transformation.py +++ b/litellm/llms/together_ai/rerank/transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic from Cohere's /v1/rerank format to Together AI's `/v1/rerank` format. +Transformation logic from Cohere's /v1/rerank format to Together AI's `/v1/rerank` format. Why separate file? Make it easy to see how transformation works """ diff --git a/litellm/llms/vertex_ai/context_caching/transformation.py b/litellm/llms/vertex_ai/context_caching/transformation.py index 950edbeb478..f73eb220cc6 100644 --- a/litellm/llms/vertex_ai/context_caching/transformation.py +++ b/litellm/llms/vertex_ai/context_caching/transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic for context caching. +Transformation logic for context caching. Why separate file? Make it easy to see how transformation works """ @@ -19,7 +19,7 @@ from ..gemini.transformation import ( def get_first_continuous_block_idx( - filtered_messages: List[Tuple[int, AllMessageValues]] # (idx, message) + filtered_messages: List[Tuple[int, AllMessageValues]], # (idx, message) ) -> int: """ Find the array index that ends the first continuous sequence of message blocks. @@ -174,7 +174,9 @@ def transform_openai_messages_to_gemini_context_caching( ) transformed_messages = _gemini_convert_messages_with_history( - messages=new_messages, model=model + messages=new_messages, + model=model, + custom_llm_provider=custom_llm_provider, ) model_name = "models/{}".format(model) diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index ac0f07b8e0b..3f945adca0d 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -41,7 +41,7 @@ class ContextCachingEndpoints(VertexBase): """ def __init__(self) -> None: - pass + super().__init__() def _get_token_and_url_context_caching( self, diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index ea57339e35f..4f5846cc5b6 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -682,6 +682,7 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 messages: List[AllMessageValues], model: Optional[str] = None, litellm_params: Optional[dict] = None, + custom_llm_provider: Optional[str] = None, ) -> List[ContentType]: """ Converts given messages from OpenAI format to Gemini format @@ -983,7 +984,9 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 or assistant_msg.get("function_call") is not None ): # support assistant tool invoke conversion gemini_tool_call_parts = convert_to_gemini_tool_call_invoke( - assistant_msg, model=model + assistant_msg, + model=model, + custom_llm_provider=custom_llm_provider, ) ## check if gemini_tool_call already exists in assistant_content for gemini_tool_call_part in gemini_tool_call_parts: @@ -1045,6 +1048,7 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 messages[msg_i], # type: ignore last_message_with_tool_calls, # type: ignore model=model, + custom_llm_provider=custom_llm_provider, ) msg_i += 1 # Handle both single part and list of parts (for Computer Use with images) @@ -1069,16 +1073,14 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 contents.append(ContentType(role="user", parts=tool_call_responses)) if len(contents) == 0: - verbose_logger.warning( - """ + verbose_logger.warning(""" No contents in messages. Contents are required. See https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.publishers.models/generateContent#request-body. If the original request did not comply to OpenAI API requirements it should have failed by now, but LiteLLM does not check for missing messages. Setting an empty content to prevent an 400 error. Relevant Issue - https://github.com/BerriAI/litellm/issues/9733 - """ - ) + """) contents.append(ContentType(role="user", parts=[PartType(text=" ")])) return contents except Exception as e: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index e99f69fcd3e..189ac7a7f6a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -289,6 +289,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return True return False + @staticmethod + def _forward_gemini_function_call_id( + model: str, custom_llm_provider: Optional[str] = None + ) -> bool: + """ + Whether to include `id` on function_call / function_response parts. + + Gemini 3+ on Google AI Studio accepts (and returns) `id` for strict + tool-call matching. Vertex AI rejects the field with HTTP 400. + """ + if custom_llm_provider != "gemini": + return False + return VertexGeminiConfig._is_gemini_3_or_newer(model) + def _supports_penalty_parameters(self, model: str) -> bool: # Gemini 3 models do not support penalty parameters if VertexGeminiConfig._is_gemini_3_or_newer(model): @@ -2649,7 +2663,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): litellm_params: Optional[dict] = None, ) -> List[ContentType]: return _gemini_convert_messages_with_history( - messages=messages, model=model, litellm_params=litellm_params + messages=messages, + model=model, + litellm_params=litellm_params, + custom_llm_provider="vertex_ai", ) def get_error_class( diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index e1b365c9f42..ba6e6f0c056 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -1,5 +1,5 @@ """ -Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batchEmbedContents format. +Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batchEmbedContents format. Why separate file? Make it easy to see how transformation works """ diff --git a/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py b/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py index 9d9015c2b91..b835ad7d8fa 100644 --- a/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py +++ b/litellm/llms/vertex_ai/text_to_speech/text_to_speech_handler.py @@ -139,7 +139,7 @@ class VertexTextToSpeechAPI(VertexLLM): ########## End of logging ############ ####### Send the request ################### if _is_async is True: - return self.async_audio_speech( # type:ignore + return self.async_audio_speech( # type: ignore logging_obj=logging_obj, url=url, headers=headers, request=request ) sync_handler = _get_httpx_client() diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index eb67e3aa828..13aa2a5350e 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -45,7 +45,7 @@ class PartnerModelPrefixes(str, Enum): class VertexAIPartnerModels(VertexBase): def __init__(self) -> None: - pass + super().__init__() @staticmethod def is_vertex_partner_model(model: str): @@ -116,9 +116,6 @@ class VertexAIPartnerModels(VertexBase): CodestralTextCompletion, ) from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexLLM, - ) except Exception as e: raise VertexAIError( status_code=400, @@ -133,9 +130,7 @@ class VertexAIPartnerModels(VertexBase): message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""", ) try: - vertex_httpx_logic = VertexLLM() - - access_token, project_id = vertex_httpx_logic._ensure_access_token( + access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/main.py b/litellm/llms/vertex_ai/vertex_gemma_models/main.py index 82cfe6de984..b6bf2f73b72 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/main.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/main.py @@ -31,7 +31,7 @@ from ..vertex_llm_base import VertexBase class VertexAIGemmaModels(VertexBase): def __init__(self) -> None: - pass + super().__init__() def completion( self, @@ -62,9 +62,6 @@ class VertexAIGemmaModels(VertexBase): try: import vertexai - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexLLM, - ) from litellm.llms.vertex_ai.vertex_gemma_models.transformation import ( VertexGemmaConfig, ) @@ -83,9 +80,8 @@ class VertexAIGemmaModels(VertexBase): ) try: model = get_vertex_base_model_name(model=model) - vertex_httpx_logic = VertexLLM() - access_token, project_id = vertex_httpx_logic._ensure_access_token( + access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 6c6446958bc..35cd54d65f6 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -91,6 +91,10 @@ class VertexGemmaConfig(OpenAIGPTConfig): "stream", None ) # Streaming not supported, will be faked client-side openai_request.pop("stream_options", None) # Stream options not supported + # Vertex Gemma's chatCompletions wrapper does not understand + # `context_management` (an Anthropic/Responses API concept). Strip it + # so the upstream endpoint does not 400 on the unknown field. + openai_request.pop("context_management", None) # Wrap in Vertex Gemma format return { diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 6f687dae7e8..990063bb9fb 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -4,8 +4,10 @@ Base Vertex, Google AI Studio LLM Class Handles Authentication and generating request urls for Vertex AI and Google AI Studio """ +import asyncio import json import os +import threading from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple import litellm @@ -30,6 +32,7 @@ GOOGLE_IMPORT_ERROR_MESSAGE = ( if TYPE_CHECKING: from google.auth.credentials import Credentials as GoogleCredentialsObject + from google.auth.credentials import TokenState else: GoogleCredentialsObject = Any @@ -42,10 +45,28 @@ class VertexBase: self._credentials: Optional[GoogleCredentialsObject] = None self._credentials_project_mapping: Dict[ Tuple[Optional[VERTEX_CREDENTIALS_TYPES], Optional[str]], - Tuple[GoogleCredentialsObject, str], + Tuple[GoogleCredentialsObject, Optional[str]], ] = {} self.project_id: Optional[str] = None self.async_handler: Optional[AsyncHTTPHandler] = None + # Per-credential-key asyncio.Lock for single-flight async refresh. + # Prevents thundering herd when token expires under high concurrency. + # Uses a regular dict (not WeakValueDictionary) so the lock identity is + # stable across concurrent callers — a weak reference can be GC'd + # between two coroutines arriving at the lock, breaking single-flight. + # An explicit refcount tracks the number of coroutines currently using + # each lock; the entry is pruned when the count reaches zero, so the + # dict stays bounded even in long-running high-cardinality deployments + # without depending on any private asyncio internals. + self._async_refresh_locks: Dict[tuple, asyncio.Lock] = {} + self._async_refresh_lock_refcounts: Dict[tuple, int] = {} + # Tracks in-flight background refresh tasks to avoid duplicate refreshes. + self._background_refresh_tasks: Dict[tuple, asyncio.Task] = {} + # Protects the sync get_access_token refresh path. + # Use RLock so that the reauthentication retry path (which calls + # back into get_access_token while still holding the lock) can + # re-acquire it without deadlocking the current thread. + self._sync_refresh_lock = threading.RLock() def get_vertex_region(self, vertex_region: Optional[str], model: str) -> str: import litellm @@ -77,7 +98,9 @@ class VertexBase: return vertex_region or "us-central1" def load_auth( - self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], project_id: Optional[str] + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], ) -> Tuple[Any, str]: if credentials is not None: if isinstance(credentials, str): @@ -343,7 +366,241 @@ class VertexBase: except ImportError: raise ImportError(GOOGLE_IMPORT_ERROR_MESSAGE) - credentials.refresh(Request()) + # Serialize all refreshes on this VertexBase across threads. + # ``credentials.refresh()`` is not safe to call concurrently on the + # same credentials object, and this method is invoked from three + # places that can run on different threads: + # - sync ``get_access_token`` (already holds ``_sync_refresh_lock``) + # - the async slow path (via ``asyncify`` in a worker thread) + # - the background proactive refresh task (via ``asyncify``) + # ``_sync_refresh_lock`` is an ``RLock`` so reentrant acquisition + # from the sync path is safe. + with self._sync_refresh_lock: + credentials.refresh(Request()) + + def _acquire_async_refresh_lock(self, credential_cache_key: tuple) -> asyncio.Lock: + """Increment the refcount and return the lock for ``credential_cache_key``. + + Every call must be paired with ``_release_async_refresh_lock`` once the + caller is done with the lock so the entry can be pruned when no other + coroutine is holding or waiting on it. + """ + lock = self._async_refresh_locks.setdefault( + credential_cache_key, asyncio.Lock() + ) + self._async_refresh_lock_refcounts[credential_cache_key] = ( + self._async_refresh_lock_refcounts.get(credential_cache_key, 0) + 1 + ) + return lock + + def _release_async_refresh_lock( + self, credential_cache_key: tuple, lock: asyncio.Lock + ) -> None: + """Decrement the refcount and drop the lock entry when it reaches zero. + + Must be called only after the caller has released ``lock`` (i.e. once + the surrounding ``async with`` has exited). asyncio is cooperative, so + the decrement-then-pop sequence below runs atomically with respect to + other coroutines. + """ + remaining = self._async_refresh_lock_refcounts.get(credential_cache_key, 0) - 1 + if remaining > 0: + self._async_refresh_lock_refcounts[credential_cache_key] = remaining + return + self._async_refresh_lock_refcounts.pop(credential_cache_key, None) + if self._async_refresh_locks.get(credential_cache_key) is lock: + self._async_refresh_locks.pop(credential_cache_key, None) + + def _try_get_cached_token( + self, + credential_cache_key: tuple, + project_id: Optional[str], + ) -> Optional[Tuple[str, str]]: + """ + Look up cached credentials and return (token, project_id) if the token + is FRESH. Returns None if not cached or not fresh. + """ + from google.auth.credentials import TokenState + + creds, cached_project_id = self._unpack_cached_credentials(credential_cache_key) + if ( + creds is not None + and self._get_token_state(creds) == TokenState.FRESH + and creds.token is not None + and isinstance(creds.token, str) + ): + resolved_project = project_id or cached_project_id + if resolved_project: + return creds.token, resolved_project + return None + + def _try_get_usable_cached_token( + self, + credential_cache_key: tuple, + project_id: Optional[str], + ) -> Optional[Tuple[str, str, "TokenState", Any, Optional[str]]]: + """ + Look up cached credentials and return usable token info for FRESH or + STALE tokens (both are still valid for outbound requests). STALE + tokens are returned along with their state and the underlying + credentials object so the caller can schedule a background refresh + without holding the per-key async lock. + """ + from google.auth.credentials import TokenState + + creds, cached_project_id = self._unpack_cached_credentials(credential_cache_key) + if creds is None: + return None + token_state = self._get_token_state(creds) + if token_state not in (TokenState.FRESH, TokenState.STALE): + return None + if creds.token is None or not isinstance(creds.token, str): + return None + resolved_project = project_id or cached_project_id + if not resolved_project: + return None + return creds.token, resolved_project, token_state, creds, cached_project_id + + def _unpack_cached_credentials( + self, credential_cache_key: tuple + ) -> Tuple[Any, Optional[str]]: + """ + Return (credentials, project_id) from the cache, or (None, None) if + not cached. Handles both tuple and legacy cache formats. + """ + if credential_cache_key not in self._credentials_project_mapping: + return None, None + cached_entry = self._credentials_project_mapping[credential_cache_key] + if isinstance(cached_entry, tuple): + return cached_entry + return cached_entry, cached_entry.quota_project_id or getattr( + cached_entry, "project_id", None + ) + + def _get_token_state(self, credentials: Any) -> "TokenState": + """ + Return the token state using google-auth's TokenState enum. + + Falls back to expired/valid checks if token_state is unavailable + (e.g. older google-auth versions or mock objects in tests). + """ + from google.auth.credentials import TokenState as _TokenState + + token_state = getattr(credentials, "token_state", None) + if isinstance(token_state, _TokenState): + return token_state + # Fallback for credentials without a real token_state (e.g. mocks) + if getattr(credentials, "expired", True): + return _TokenState.INVALID + if getattr(credentials, "valid", False): + return _TokenState.FRESH + return _TokenState.INVALID + + async def _load_and_cache_credentials( + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], + credential_cache_key: tuple, + ) -> Tuple[Any, Optional[str]]: + """Load credentials via load_auth (in thread) and cache the result.""" + try: + _credentials, credential_project_id = await asyncify(self.load_auth)( + credentials=credentials, + project_id=project_id, + ) + except Exception as e: + verbose_logger.exception("Failed to load vertex credentials: %s", str(e)) + raise + if _credentials is None: + raise ValueError("Could not resolve credentials") + self._credentials_project_mapping[credential_cache_key] = ( + _credentials, + credential_project_id, + ) + return _credentials, credential_project_id + + async def _background_refresh_credentials( + self, + credentials: Any, + credential_cache_key: tuple, + credential_project_id: Optional[str], + ) -> None: + """ + Refresh credentials in the background without blocking the calling request. + + Called when the token is still valid but nearing expiry (proactive refresh). + Errors are logged but not raised — the current token is still usable. + """ + try: + verbose_logger.debug("Background proactive credential refresh") + await asyncify(self.refresh_auth)(credentials) + # Only update the cache if it still points at the credentials + # object we just refreshed. The per-key async lock is not held + # here, so a concurrent INVALID path may have already replaced + # this entry (e.g. via _handle_reauthentication_async, which + # creates a fresh credentials object). In that case our write + # would clobber the newer entry with a stale reference. + cached_creds, _ = self._unpack_cached_credentials(credential_cache_key) + if cached_creds is credentials: + self._credentials_project_mapping[credential_cache_key] = ( + credentials, + credential_project_id, + ) + except Exception: + verbose_logger.debug( + "Background credential refresh failed, will retry on next request", + exc_info=True, + ) + + async def _await_in_flight_background_refresh( + self, credential_cache_key: tuple + ) -> None: + """Wait for an in-flight background refresh to finish, if any. + + google-auth's ``Credentials.refresh()`` is not safe to invoke + concurrently on the same credentials object. Coroutines that need a + blocking refresh must first drain any background refresh that was + scheduled while a previous STALE token was being served. + """ + existing_task = self._background_refresh_tasks.get(credential_cache_key) + if existing_task is None or existing_task.done(): + return + try: + await existing_task + except Exception: + # Background refresh failures are already logged inside + # _background_refresh_credentials; the caller will fall through + # to its own blocking refresh. + pass + + def _schedule_background_refresh( + self, + credentials: Any, + credential_cache_key: tuple, + credential_project_id: Optional[str], + ) -> None: + """Kick off a single background refresh for ``credential_cache_key``. + + Skips scheduling if a refresh is already in flight. The done-callback + guards against removing a newer task that has replaced this one in the + tracking dict (done_callbacks are scheduled via ``call_soon``). + """ + existing = self._background_refresh_tasks.get(credential_cache_key) + if existing is not None and not existing.done(): + return + self._background_refresh_tasks.pop(credential_cache_key, None) + task = asyncio.create_task( + self._background_refresh_credentials( + credentials, credential_cache_key, credential_project_id + ) + ) + + def _drop_background_refresh_task(_fut: asyncio.Future[Any]) -> None: + if self._background_refresh_tasks.get(credential_cache_key) is _fut: + self._background_refresh_tasks.pop(credential_cache_key, None) + + task.add_done_callback(_drop_background_refresh_task) + self._background_refresh_tasks[credential_cache_key] = task def _ensure_access_token( self, @@ -563,6 +820,65 @@ class VertexBase: # Re-raise the original error for better context raise error + async def _handle_reauthentication_async( + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], + credential_cache_key: Tuple, + error: Exception, + ) -> Tuple[str, str]: + """ + Async reauthentication retry that stays within the per-key async lock. + """ + verbose_logger.debug( + f"Handling async reauthentication for project_id: {project_id}. " + f"Clearing cache and retrying once." + ) + + self._credentials_project_mapping.pop(credential_cache_key, None) + + try: + _credentials, credential_project_id = ( + await self._load_and_cache_credentials( + credentials=credentials, + project_id=project_id, + credential_cache_key=credential_cache_key, + ) + ) + if project_id is None and isinstance(credential_project_id, str): + project_id = credential_project_id + cache_credentials = ( + json.dumps(credentials) + if isinstance(credentials, dict) + else credentials + ) + resolved_cache_key = (cache_credentials, project_id) + # Always overwrite — any pre-existing entry at the resolved key + # references the OLD credentials object we just replaced, and + # leaving it would force the next request to do a redundant + # refresh/reauth before realizing the cached creds are stale. + self._credentials_project_mapping[resolved_cache_key] = ( + _credentials, + credential_project_id, + ) + + if _credentials.token is None or not isinstance(_credentials.token, str): + raise ValueError( + "Could not resolve credentials token. Got None or non-string token (type={})".format( + type(_credentials.token).__name__ + ) + ) + if project_id is None: + raise ValueError("Could not resolve project_id") + + return _credentials.token, project_id + except Exception as retry_error: + verbose_logger.error( + f"Async reauthentication retry failed for project_id: {project_id}. " + f"Original error: {str(error)}. Retry error: {str(retry_error)}" + ) + raise error + def get_access_token( self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], @@ -646,7 +962,7 @@ class VertexBase: ) ## VALIDATE CREDENTIALS - verbose_logger.debug(f"Validating credentials for project_id: {project_id}") + verbose_logger.debug("Validating credentials") if ( project_id is None and credential_project_id is not None @@ -666,26 +982,27 @@ class VertexBase: raise ValueError("Credentials are None after loading") if _credentials.expired: - try: - verbose_logger.debug( - f"Credentials expired, refreshing for project_id: {project_id}" - ) - self.refresh_auth(_credentials) - self._credentials_project_mapping[credential_cache_key] = ( - _credentials, - credential_project_id, - ) - except Exception as e: - # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login` - # in this case, we should try to reload the credentials by clearing the cache and retrying - if "Reauthentication is needed" in str(e) and not _retry_reauth: - return self._handle_reauthentication( - credentials=credentials, - project_id=project_id, - credential_cache_key=credential_cache_key, - error=e, - ) - raise e + with self._sync_refresh_lock: + # Double-check after acquiring lock + if _credentials.expired: + try: + verbose_logger.debug("Credentials expired, refreshing") + self.refresh_auth(_credentials) + self._credentials_project_mapping[credential_cache_key] = ( + _credentials, + credential_project_id, + ) + except Exception as e: + # if refresh fails, it's possible the user has re-authenticated via `gcloud auth application-default login` + # in this case, we should try to reload the credentials by clearing the cache and retrying + if "Reauthentication is needed" in str(e) and not _retry_reauth: + return self._handle_reauthentication( + credentials=credentials, + project_id=project_id, + credential_cache_key=credential_cache_key, + error=e, + ) + raise e ## VALIDATION STEP if _credentials.token is None or not isinstance(_credentials.token, str): @@ -700,6 +1017,149 @@ class VertexBase: return _credentials.token, project_id + async def get_access_token_async( + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], + ) -> Tuple[str, str]: + """ + Async version of get_access_token with single-flight refresh coordination. + + Prevents thundering herd: when credentials expire under high concurrency, + only one coroutine refreshes while others wait on the lock. Uses native + async refresh for service_account and authorized_user credentials. + """ + from google.auth.credentials import TokenState + + cache_credentials = ( + json.dumps(credentials) if isinstance(credentials, dict) else credentials + ) + credential_cache_key = (cache_credentials, project_id) + + # === FAST PATH (no lock) === + # If credentials are FRESH or STALE, return immediately without + # touching the per-key async lock. STALE tokens are still usable; + # we kick off a deduplicated background refresh so subsequent + # requests get a fresh token, but we must not serialize concurrent + # callers on the lock just to schedule that refresh. + usable = self._try_get_usable_cached_token(credential_cache_key, project_id) + if usable is not None: + cached_token, resolved_project, token_state, creds, cached_project_id = ( + usable + ) + if token_state == TokenState.STALE: + self._schedule_background_refresh( + creds, credential_cache_key, cached_project_id + ) + return cached_token, resolved_project + + # === SLOW PATH (per-key lock) === + lock = self._acquire_async_refresh_lock(credential_cache_key) + try: + async with lock: + # Double-check after acquiring lock — another coroutine may have refreshed. + cached = self._try_get_cached_token(credential_cache_key, project_id) + if cached is not None: + return cached + + _credentials, credential_project_id = self._unpack_cached_credentials( + credential_cache_key + ) + + # Load credentials if not cached + if _credentials is None: + _credentials, credential_project_id = ( + await self._load_and_cache_credentials( + credentials, project_id, credential_cache_key + ) + ) + + # Resolve project_id from credentials if not provided + if project_id is None and isinstance(credential_project_id, str): + project_id = credential_project_id + resolved_cache_key = (cache_credentials, project_id) + # Always overwrite — a pre-existing entry at the resolved + # key may reference stale credentials (e.g. from before a + # reauth that only repopulated the unresolved key), which + # would force the next request through an unnecessary + # refresh/reauth cycle. + self._credentials_project_mapping[resolved_cache_key] = ( + _credentials, + credential_project_id, + ) + + # Use google-auth's token_state to decide refresh strategy: + # - STALE: token is usable but within REFRESH_THRESHOLD (3:45) of + # expiry — return it immediately and refresh in the background. + # - INVALID: token is expired or missing — must block on refresh. + token_state = self._get_token_state(_credentials) + + if token_state == TokenState.STALE: + if project_id is None: + raise ValueError("Could not resolve project_id") + current_token = _credentials.token + if current_token is None or not isinstance(current_token, str): + # Token is malformed despite STALE state — block on a full + # refresh using the same path as INVALID credentials. + token_state = TokenState.INVALID + else: + self._schedule_background_refresh( + _credentials, + credential_cache_key, + credential_project_id, + ) + return current_token, project_id + + if token_state == TokenState.INVALID: + # Drain any in-flight background refresh before invoking + # refresh_auth ourselves; google-auth's + # Credentials.refresh() is not safe to call concurrently + # on the same credentials object, and the background task + # runs outside this lock. + await self._await_in_flight_background_refresh(credential_cache_key) + cached = self._try_get_cached_token( + credential_cache_key, project_id + ) + if cached is not None: + return cached + + # Token is expired or missing — must block until refresh completes. + try: + verbose_logger.debug("Credentials expired, refreshing") + await asyncify(self.refresh_auth)(_credentials) + self._credentials_project_mapping[credential_cache_key] = ( + _credentials, + credential_project_id, + ) + except Exception as e: + if "Reauthentication is needed" in str(e): + verbose_logger.debug( + "Reauthentication needed, clearing cache and retrying" + ) + return await self._handle_reauthentication_async( + credentials=credentials, + project_id=project_id, + credential_cache_key=credential_cache_key, + error=e, + ) + raise + + # Final validation + if _credentials.token is None or not isinstance( + _credentials.token, str + ): + raise ValueError( + "Could not resolve credentials token. Got None or non-string token (type={})".format( + type(_credentials.token).__name__ + ) + ) + if project_id is None: + raise ValueError("Could not resolve project_id") + + return _credentials.token, project_id + finally: + self._release_async_refresh_lock(credential_cache_key, lock) + async def _ensure_access_token_async( self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], @@ -714,13 +1174,10 @@ class VertexBase: if custom_llm_provider == "gemini": return "", "" else: - try: - return await asyncify(self.get_access_token)( - credentials=credentials, - project_id=project_id, - ) - except Exception as e: - raise e + return await self.get_access_token_async( + credentials=credentials, + project_id=project_id, + ) def set_headers( self, auth_header: Optional[str], extra_headers: Optional[dict] diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index 7240d9dce57..732d5f90dc2 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -57,7 +57,7 @@ def create_vertex_url( class VertexAIModelGardenModels(VertexBase): def __init__(self) -> None: - pass + super().__init__() def completion( self, @@ -89,9 +89,6 @@ class VertexAIModelGardenModels(VertexBase): import vertexai from litellm.llms.openai_like.chat.handler import OpenAILikeChatHandler - from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( - VertexLLM, - ) except Exception as e: raise VertexAIError( status_code=400, @@ -107,9 +104,8 @@ class VertexAIModelGardenModels(VertexBase): ) try: model = get_vertex_base_model_name(model=model) - vertex_httpx_logic = VertexLLM() - access_token, project_id = vertex_httpx_logic._ensure_access_token( + access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, custom_llm_provider="vertex_ai", diff --git a/litellm/llms/vllm/completion/transformation.py b/litellm/llms/vllm/completion/transformation.py index ec4c07e95d8..e03b07f9897 100644 --- a/litellm/llms/vllm/completion/transformation.py +++ b/litellm/llms/vllm/completion/transformation.py @@ -1,5 +1,5 @@ """ -Translates from OpenAI's `/v1/chat/completions` to the VLLM sdk `llm.generate`. +Translates from OpenAI's `/v1/chat/completions` to the VLLM sdk `llm.generate`. NOT RECOMMENDED FOR PRODUCTION USE. Use `hosted_vllm/` instead. """ diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 40328062e09..1f5ca99f47d 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -1,6 +1,6 @@ """ -This module is used to transform the request and response for the Voyage contextualized embeddings API. -This would be used for all the contextualized embeddings models in Voyage. +This module is used to transform the request and response for the Voyage contextualized embeddings API. +This would be used for all the contextualized embeddings models in Voyage. """ from typing import List, Optional, Union diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 6300868a641..7325c0596a6 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple, Union +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple, Union import httpx @@ -26,6 +26,7 @@ from ...openai.chat.gpt_transformation import ( class XAIChatConfig(OpenAIGPTConfig): + @property def custom_llm_provider(self) -> Optional[str]: return "xai" @@ -225,21 +226,57 @@ class XAIChatConfig(OpenAIGPTConfig): verbose_logger.debug(f"Error extracting X.AI web search usage: {e}") self._fold_reasoning_tokens_into_completion(response) + self._normalize_openai_compatible_usage_totals(getattr(response, "usage", None)) return response @staticmethod - def _fold_reasoning_tokens_into_completion(model_response: ModelResponse) -> None: + def _fold_reasoning_tokens_into_completion( + target: Union[ModelResponse, Usage, Dict[str, Any], None], + ) -> None: """Reconcile xAI Usage to the OpenAI invariant. xAI accounts ``reasoning_tokens`` separately from ``completion_tokens`` while still summing them into ``total_tokens``. OpenAI's contract (o1/o3) folds reasoning into ``completion_tokens``, so fold here to keep ``total = prompt + completion``. Idempotent. + + Accepts a ``ModelResponse`` (non-streaming), a ``Usage`` object, or a + raw usage ``dict`` (streaming chunk) so streaming and non-streaming + paths stay in sync. """ - usage = getattr(model_response, "usage", None) + if target is None: + return + + if isinstance(target, ModelResponse): + usage: Union[Usage, Dict[str, Any], None] = getattr(target, "usage", None) + else: + usage = target if usage is None: return + if isinstance(usage, dict): + details = usage.get("completion_tokens_details") or {} + if isinstance(details, dict): + reasoning_tokens = int(details.get("reasoning_tokens") or 0) + else: + reasoning_tokens = int(getattr(details, "reasoning_tokens", 0) or 0) + if reasoning_tokens <= 0: + return + + prompt_tokens = int(usage.get("prompt_tokens") or 0) + completion_tokens = int(usage.get("completion_tokens") or 0) + total_tokens = int(usage.get("total_tokens") or 0) + + if total_tokens == prompt_tokens + completion_tokens: + return + + # Guard against double-counting if xAI changes accounting. + if total_tokens != prompt_tokens + completion_tokens + reasoning_tokens: + return + + usage["completion_tokens"] = completion_tokens + reasoning_tokens + return + details = getattr(usage, "completion_tokens_details", None) reasoning_tokens = ( int(getattr(details, "reasoning_tokens", 0) or 0) if details else 0 @@ -284,6 +321,25 @@ class XAIChatConfig(OpenAIGPTConfig): setattr(usage, "num_sources_used", int(num_sources_used)) verbose_logger.debug(f"X.AI web search sources used: {num_sources_used}") + @staticmethod + def _normalize_openai_compatible_usage_totals( + usage: Union[Usage, Dict[str, Any], None], + ) -> None: + if usage is None: + return + if isinstance(usage, dict): + prompt_tokens = int(usage.get("prompt_tokens") or 0) + completion_tokens = int(usage.get("completion_tokens") or 0) + expected_total = prompt_tokens + completion_tokens + if int(usage.get("total_tokens") or 0) < expected_total: + usage["total_tokens"] = expected_total + return + prompt_tokens = int(usage.prompt_tokens or 0) + completion_tokens = int(usage.completion_tokens or 0) + expected_total = prompt_tokens + completion_tokens + if int(usage.total_tokens or 0) < expected_total: + usage.total_tokens = expected_total + class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): def chunk_parser(self, chunk: dict) -> ModelResponseStream: @@ -304,4 +360,8 @@ class XAIChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): # Add a dummy choice with empty delta to ensure proper processing chunk["choices"] = [{"index": 0, "delta": {}, "finish_reason": None}] + if "usage" in chunk and chunk["usage"] is not None: + XAIChatConfig._fold_reasoning_tokens_into_completion(chunk["usage"]) + XAIChatConfig._normalize_openai_compatible_usage_totals(chunk["usage"]) + return super().chunk_parser(chunk) diff --git a/litellm/main.py b/litellm/main.py index b5364f8ba17..e17a5ad9a48 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1491,7 +1491,9 @@ def completion( # type: ignore # noqa: PLR0915 provider.value for provider in LlmProviders ]: provider_config = ProviderConfigManager.get_provider_chat_config( - model=model, provider=LlmProviders(custom_llm_provider) + model=model, + provider=LlmProviders(custom_llm_provider), + base_model=base_model, ) if provider_config is not None: @@ -1550,6 +1552,7 @@ def completion( # type: ignore # noqa: PLR0915 "safety_identifier": safety_identifier, "service_tier": service_tier, "allowed_openai_params": kwargs.get("allowed_openai_params"), + "base_model": base_model, } optional_params = get_optional_params( **optional_param_args, **non_default_params @@ -1670,6 +1673,10 @@ def completion( # type: ignore # noqa: PLR0915 reasoning_summary=_reasoning_summary_for_bridge, ) + # Use base_model (the true underlying model) for Azure model-type + # detection when the deployment name differs from the model name. + _azure_detection_model = base_model or model + if responses_api_model_info.get("mode") == "responses": from litellm.completion_extras import responses_api_bridge @@ -1713,7 +1720,9 @@ def completion( # type: ignore # noqa: PLR0915 and OpenAIGPT5Config.is_model_gpt_5_model(model) ) or ( custom_llm_provider == "azure" - and litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model) + and litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model( + _azure_detection_model + ) ): optional_params, _ = strip_reasoning_summary_aliases_from_optional_params( optional_params @@ -1766,7 +1775,9 @@ def completion( # type: ignore # noqa: PLR0915 if max_retries is not None: optional_params["max_retries"] = max_retries - if litellm.AzureOpenAIO1Config().is_o_series_model(model=model): + if litellm.AzureOpenAIO1Config().is_o_series_model( + model=_azure_detection_model + ): ## LOAD CONFIG - if set config = litellm.AzureOpenAIO1Config.get_config() for k, v in config.items(): diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a1d656aa2db..d9aa97585dd 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1448,6 +1448,35 @@ "supports_native_structured_output": true, "supports_minimal_reasoning_effort": true }, + "jp.anthropic.claude-sonnet-4-6": { + "cache_creation_input_token_cost": 4.125e-06, + "cache_read_input_token_cost": 3.3e-07, + "input_cost_per_token": 3.3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.65e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 346, + "supports_native_structured_output": true, + "supports_minimal_reasoning_effort": true + }, "anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -9606,6 +9635,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -9799,6 +9829,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -9832,6 +9863,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -9865,6 +9897,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -9899,6 +9932,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_adaptive_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -13952,6 +13986,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/accounts/fireworks/models/glm-5p1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 202800, + "max_output_tokens": 202800, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://fireworks.ai/models/fireworks/glm-5p1", + "supports_function_calling": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false + }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, "litellm_provider": "fireworks_ai", @@ -14218,6 +14267,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/glm-5p1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 202800, + "max_output_tokens": 202800, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://fireworks.ai/models/fireworks/glm-5p1", + "supports_function_calling": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false + }, "fireworks_ai/kimi-k2p5": { "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 6e-07, @@ -14887,7 +14951,65 @@ "mode": "chat", "output_cost_per_reasoning_token": 1.5e-06, "output_cost_per_token": 1.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/models", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, + "gemini-3.1-flash-lite": { + "cache_read_input_token_cost": 4.5e-08, + "cache_read_input_token_cost_per_audio_token": 9e-08, + "input_cost_per_audio_token": 9e-07, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 2.7e-06, + "output_cost_per_token": 2.7e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -16991,6 +17113,66 @@ "web_search_billing_unit": "per_query", "supports_service_tier": true }, + "gemini/gemini-3.1-flash-lite": { + "cache_read_input_token_cost": 4.5e-08, + "cache_read_input_token_cost_per_audio_token": 9e-08, + "input_cost_per_audio_token": 9e-07, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 2.7e-06, + "output_cost_per_token": 2.7e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 250000, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "gemini/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, "input_cost_per_audio_token": 1e-06, @@ -24293,6 +24475,21 @@ "supports_tool_choice": true, "supports_vision": true }, + "mistral/ministral-8b-2512": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-tiny": { "input_cost_per_token": 2.5e-07, "litellm_provider": "mistral", @@ -27137,6 +27334,58 @@ "supports_web_search": true, "tpm": 800000 }, + "openrouter/google/gemini-3.1-flash-lite": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "openrouter", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "tpm": 800000 + }, "openrouter/google/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -28911,6 +29160,24 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "reducto/parse-legacy": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.015, + "source": "https://reducto.ai/pricing", + "supported_endpoints": [ + "/v1/ocr" + ] + }, + "reducto/parse-v3": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.015, + "source": "https://reducto.ai/pricing", + "supported_endpoints": [ + "/v1/ocr" + ] + }, "recraft/recraftv2": { "litellm_provider": "recraft", "mode": "image_generation", @@ -33613,6 +33880,64 @@ }, "web_search_billing_unit": "per_query" }, + "vertex_ai/gemini-3.1-flash-lite": { + "cache_read_input_token_cost": 4.5e-08, + "cache_read_input_token_cost_per_audio_token": 9e-08, + "input_cost_per_audio_token": 9e-07, + "input_cost_per_token": 4.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 2.7e-06, + "output_cost_per_token": 2.7e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index d1b49039e8e..a72e8e34a49 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -1212,11 +1212,17 @@ class MCPServerManager: return [] # Get server-specific auth header if available - server_auth_header = None - if mcp_server_auth_headers and server.alias: - server_auth_header = mcp_server_auth_headers.get(server.alias) - elif mcp_server_auth_headers and server.server_name: - server_auth_header = mcp_server_auth_headers.get(server.server_name) + server_auth_header: Optional[Union[str, Dict[str, str]]] = None + if mcp_server_auth_headers: + from litellm.proxy._experimental.mcp_server.utils import ( + lookup_mcp_server_auth_in_headers, + ) + + server_auth_header = lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers, + alias=server.alias, + server_name=server.server_name, + ) # Fall back to deprecated mcp_auth_header if no server-specific header found if server_auth_header is None: @@ -1226,6 +1232,7 @@ class MCPServerManager: tools = await self._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, + user_api_key_auth=user_api_key_auth, ) return tools except Exception as e: @@ -1406,6 +1413,7 @@ class MCPServerManager: extra_headers: Optional[Dict[str, str]] = None, add_prefix: bool = True, raw_headers: Optional[Dict[str, str]] = None, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[MCPTool]: """ Helper method to get tools from a single MCP server with prefixed names. @@ -1432,6 +1440,46 @@ class MCPServerManager: extra_headers = {} extra_headers.update(server.static_headers) + # MCPJWTSigner: inject signed JWT for tools/list (list path skips pre_call_hook). + # Skip entirely when the signer is not configured (avoid an unnecessary + # dict copy on every list call), when the server has its own static + # Authorization header, when a per-user mcp_auth_header has already + # been resolved, or when the caller already supplied an Authorization + # entry in extra_headers (e.g. a per-user OAuth token resolved + # upstream) — admin-configured static auth and per-user OAuth must + # take precedence so the signer doesn't silently overwrite e.g. an + # upstream API key or a user's OAuth token (MCPClient._get_auth_headers + # applies extra_headers after writing Authorization from auth_value, so + # an injected JWT would otherwise clobber the per-user token). + if user_api_key_auth is not None and not server.spec_path: + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + get_mcp_jwt_signer, + inject_mcp_jwt_headers_for_upstream, + ) + + static_headers = server.static_headers or {} + has_static_authorization = any( + isinstance(k, str) and k.lower() == "authorization" + for k in static_headers.keys() + ) + has_extra_authorization = bool(extra_headers) and any( + isinstance(k, str) and k.lower() == "authorization" + for k in (extra_headers or {}).keys() + ) + + if ( + get_mcp_jwt_signer() is not None + and not has_static_authorization + and not mcp_auth_header + and not has_extra_authorization + ): + extra_headers = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_api_key_auth, + extra_headers=extra_headers, + raw_headers=raw_headers, + for_list_tools=True, + ) + stdio_env = self._build_stdio_env(server, raw_headers) client = await self._create_mcp_client( @@ -2665,16 +2713,15 @@ class MCPServerManager: server_auth_header: Optional[Union[Dict[str, str], str]] = None if mcp_server_auth_headers: # Normalize keys for case-insensitive lookup - normalized_headers = { - k.lower(): v for k, v in mcp_server_auth_headers.items() - } + from litellm.proxy._experimental.mcp_server.utils import ( + lookup_mcp_server_auth_in_headers, + ) - if mcp_server.alias: - server_auth_header = normalized_headers.get(mcp_server.alias.lower()) - if server_auth_header is None and mcp_server.server_name: - server_auth_header = normalized_headers.get( - mcp_server.server_name.lower() - ) + server_auth_header = lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers, + alias=mcp_server.alias, + server_name=mcp_server.server_name, + ) # Fall back to deprecated mcp_auth_header if no server-specific header found if server_auth_header is None: @@ -2791,6 +2838,112 @@ class MCPServerManager: return cast(CallToolResult, result) + def _resolve_mcp_server_for_tool_call( + self, + server_name: str, + name: str, + ) -> MCPServer: + """Resolve MCP server for call_tool (prefixed name, registry, fallback).""" + prefixed_tool_name = add_server_prefix_to_name(name, server_name) + mcp_server = self._get_mcp_server_from_tool_name(prefixed_tool_name) + resolved_by_server_name_only = False + normalized_server_name = normalize_server_name(server_name) + + def _candidate_matches_server_name(candidate: MCPServer) -> bool: + for identifier in ( + candidate.alias, + candidate.server_name, + candidate.name, + ): + if identifier and normalize_server_name(identifier) == ( + normalized_server_name + ): + return True + return False + + if mcp_server is None: + for candidate in self.get_registry().values(): + if _candidate_matches_server_name(candidate): + mcp_server = candidate + resolved_by_server_name_only = True + break + if mcp_server is None: + fallback = self._get_mcp_server_from_tool_name(name) + if fallback is not None and ( + not server_name or _candidate_matches_server_name(fallback) + ): + mcp_server = fallback + if mcp_server is None: + raise ValueError(f"Tool {name} not found") + + if resolved_by_server_name_only: + tool_known = ( + name in self.tool_name_to_mcp_server_name_mapping + or prefixed_tool_name in self.tool_name_to_mcp_server_name_mapping + ) + if not tool_known: + raise ValueError(f"Tool {name} not found") + + return mcp_server + + async def _resolve_oauth2_headers_for_tool_call( + self, + mcp_server: MCPServer, + oauth2_headers: Optional[Dict[str, str]], + user_api_key_auth: Optional[UserAPIKeyAuth], + ) -> Optional[Dict[str, str]]: + """Look up per-user OAuth headers when the client did not supply a token.""" + if ( + not mcp_server.needs_user_oauth_token + or oauth2_headers + or user_api_key_auth is None + ): + return oauth2_headers + + user_id = getattr(user_api_key_auth, "user_id", None) + if not user_id: + return oauth2_headers + + try: + from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 + _get_user_oauth_extra_headers_from_db, + ) + + stored_headers = await _get_user_oauth_extra_headers_from_db( + server=mcp_server, + user_api_key_auth=user_api_key_auth, + ) + if stored_headers: + return stored_headers + except Exception as _lookup_exc: + verbose_logger.debug( + "call_tool: per-user token lookup failed for " "user=%s server=%s: %s", + user_id, + mcp_server.server_id, + _lookup_exc, + ) + return oauth2_headers + + async def _gather_openapi_tool_tasks( + self, + tasks: List[Any], + proxy_logging_obj: Optional[ProxyLogging], + ) -> CallToolResult: + """Await OpenAPI tool tasks and return the tool call result.""" + try: + mcp_responses = await asyncio.gather(*tasks) + result_index = 1 if proxy_logging_obj else 0 + return cast(CallToolResult, mcp_responses[result_index]) + except ( + BlockedPiiEntityError, + GuardrailRaisedException, + HTTPException, + ) as e: + verbose_logger.error( + f"Guardrail blocked MCP tool call during result check: {str(e)}" + ) + raise e + async def call_tool( self, server_name: str, @@ -2821,12 +2974,7 @@ class MCPServerManager: CallToolResult from the MCP server """ start_time = datetime.datetime.now() - - # Get the MCP server - prefixed_tool_name = add_server_prefix_to_name(name, server_name) - mcp_server = self._get_mcp_server_from_tool_name(prefixed_tool_name) - if mcp_server is None: - raise ValueError(f"Tool {name} not found") + mcp_server = self._resolve_mcp_server_for_tool_call(server_name, name) ######################################################### # Pre MCP Tool Call Hook @@ -2860,36 +3008,9 @@ class MCPServerManager: ) tasks.append(during_hook_task) - # For per-user OAuth servers: if the client didn't supply a token in - # oauth2_headers, look up the stored token from Redis / DB. This is the - # call_tool equivalent of _get_user_oauth_extra_headers_from_db used in - # list_tools. - if ( - mcp_server.needs_user_oauth_token - and not oauth2_headers - and user_api_key_auth is not None - ): - user_id = getattr(user_api_key_auth, "user_id", None) - if user_id: - try: - from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 - _get_user_oauth_extra_headers_from_db, - ) - - stored_headers = await _get_user_oauth_extra_headers_from_db( - server=mcp_server, - user_api_key_auth=user_api_key_auth, - ) - if stored_headers: - oauth2_headers = stored_headers - except Exception as _lookup_exc: - verbose_logger.debug( - "call_tool: per-user token lookup failed for " - "user=%s server=%s: %s", - user_id, - mcp_server.server_id, - _lookup_exc, - ) + oauth2_headers = await self._resolve_oauth2_headers_for_tool_call( + mcp_server, oauth2_headers, user_api_key_auth + ) # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: @@ -2925,26 +3046,7 @@ class MCPServerManager: hook_extra_headers=hook_result.get("extra_headers"), ) - # For OpenAPI tools, await outside the client context - try: - mcp_responses = await asyncio.gather(*tasks) - - # If proxy_logging_obj is None, the tool call result is at index 0 - # If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task) - result_index = 1 if proxy_logging_obj else 0 - result = mcp_responses[result_index] - - return cast(CallToolResult, result) - except ( - BlockedPiiEntityError, - GuardrailRaisedException, - HTTPException, - ) as e: - # Re-raise guardrail exceptions to properly fail the MCP call - verbose_logger.error( - f"Guardrail blocked MCP tool call during result check: {str(e)}" - ) - raise e + return await self._gather_openapi_tool_tasks(tasks, proxy_logging_obj) ######################################################### # End of Methods that call the upstream MCP servers diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 8541a691e88..e8b591c39cf 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -3,8 +3,8 @@ import os from ipaddress import ip_address -from typing import List, Optional -from urllib.parse import urlparse, urlunparse +from typing import Any, Dict, List, NoReturn, Optional +from urllib.parse import ParseResult, urlparse, urlunparse from fastapi import HTTPException, Request @@ -29,10 +29,47 @@ _DEFAULT_PORTS = {"http": 80, "https": 443} # subdomain. HTTPS only. _TRUSTED_REDIRECT_ORIGINS_ENV = "MCP_TRUSTED_REDIRECT_ORIGINS" +# Comma-separated private-use URI allowlist for native MCP clients. +# A trailing ``*`` is a prefix match; end the prefix with ``/`` (e.g. +# ``myapp://host/oauth/*``) so ``.../oauth/callback*`` does not also +# match ``.../oauth/callback-2``. +_TRUSTED_NATIVE_REDIRECT_URIS_ENV = "MCP_TRUSTED_NATIVE_REDIRECT_URIS" + +# Default allowlist for trusted native redirect URIs. +_DEFAULT_NATIVE_REDIRECT_URIS: List[str] = [ + "cursor://anysphere.cursor-mcp/oauth/callback", +] _warned_invalid_proxy_base_url: Optional[str] = None +def _oauth_invalid_request( + error_description: str, + *, + hint: Optional[str] = None, + **extra: Any, +) -> NoReturn: + """Raise ``invalid_request`` (RFC 6749) with a debuggable description. + + FastAPI serializes ``detail`` as JSON. Callers still see ``error``: + ``invalid_request``; ``error_description`` and ``hint`` explain what + failed and how to fix it (e.g. reverse-proxy / PROXY_BASE_URL issues). + """ + detail: Dict[str, Any] = { + "error": "invalid_request", + "error_description": error_description, + } + if hint: + detail["hint"] = hint + detail.update(extra) + raise HTTPException(status_code=400, detail=detail) + + +def _origin_label(scheme: str, netloc: str) -> str: + """Human-readable origin for error messages (scheme + host[:port]).""" + return f"{scheme}://{netloc}" if netloc else f"{scheme}://" + + def _resolve_proxy_base_url_env() -> Optional[str]: global _warned_invalid_proxy_base_url configured = os.environ.get("PROXY_BASE_URL", "").strip() @@ -108,17 +145,15 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None: ``"127.0.0.1"`` alone would miss ``127.0.0.2`` and the full-form IPv6 loopback ``0:0:0:0:0:0:0:1``. """ - try: - parsed = urlparse(redirect_uri) - except ValueError: - raise HTTPException(status_code=400, detail="invalid_request") + parsed = _parse_redirect_uri_for_validation(redirect_uri) if parsed.scheme not in ("http", "https"): - raise HTTPException(status_code=400, detail="invalid_request") - # Fragments are not allowed in OAuth redirect URIs (RFC 6749 §3.1.2) - # — rejecting them prevents a ``http://127.0.0.1/cb#frag?code=...`` - # from silently eating the authorization code. + _oauth_invalid_request( + f"redirect_uri scheme {parsed.scheme!r} is not allowed; use http or https.", + ) if parsed.fragment: - raise HTTPException(status_code=400, detail="invalid_request") + _oauth_invalid_request( + "redirect_uri must not contain a URL fragment (#...).", + ) host = (parsed.hostname or "").lower() if host == "localhost": return @@ -129,7 +164,10 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None: # Unparseable host (malformed IPv6, etc.) — treat as invalid, # don't let it bubble up as a 500. pass - raise HTTPException(status_code=400, detail="invalid_request") + _oauth_invalid_request( + "redirect_uri must use a loopback host (localhost or 127.0.0.0/8).", + hint="Native MCP clients should register a callback on http://127.0.0.1:/...", + ) def _strip_default_port(scheme: str, netloc: str) -> str: @@ -212,10 +250,256 @@ def _matches_trusted_origin_entry(netloc: str, entry: str) -> bool: return netloc == entry +def _normalize_native_redirect_uri( + parsed, +) -> str: + """Lowercase scheme, netloc, and path for allowlist comparison.""" + return urlunparse( + ( + (parsed.scheme or "").lower(), + (parsed.netloc or "").lower(), + (parsed.path or "").lower(), + "", + "", + "", + ) + ) + + +def _parse_trusted_native_redirect_uris() -> List[str]: + """Built-in native MCP callbacks plus ``MCP_TRUSTED_NATIVE_REDIRECT_URIS``.""" + entries: List[str] = [uri.lower() for uri in _DEFAULT_NATIVE_REDIRECT_URIS] + raw = os.environ.get(_TRUSTED_NATIVE_REDIRECT_URIS_ENV, "").strip() + if not raw: + return entries + for token in raw.split(","): + entry = token.strip().lower() + if entry and entry not in entries: + entries.append(entry) + return entries + + +def _native_wildcard_prefix_matches(normalized: str, prefix: str) -> bool: + """Prefix match for ``entry*`` allowlist rows. + + When the prefix does not end with ``/``, only exact matches or + deeper path segments (``prefix/...``) are accepted — not siblings + like ``prefix-2``. + """ + if not normalized.startswith(prefix): + return False + suffix = normalized[len(prefix) :] + if not suffix: + return True + if prefix.endswith("/"): + return True + return suffix[0] == "/" + + +def _matches_trusted_native_redirect_uri(parsed) -> bool: + """Allowlisted private-use / custom-scheme OAuth callbacks for native MCP clients.""" + if parsed.fragment: + return False + # Query strings are not part of registered redirect_uris (RFC 6749 §3.1.2). + # Rejecting them prevents allowlist bypass via ``.../callback?injected=...``. + if parsed.query: + return False + if not parsed.netloc: + return False + if parsed.username is not None or parsed.password is not None: + return False + if "\\" in parsed.netloc: + return False + + normalized = _normalize_native_redirect_uri(parsed) + for entry in _parse_trusted_native_redirect_uris(): + if entry.endswith("*"): + if _native_wildcard_prefix_matches(normalized, entry[:-1]): + return True + elif normalized == entry: + return True + return False + + +def _parse_redirect_uri_for_validation(redirect_uri: str) -> ParseResult: + try: + return urlparse(redirect_uri) + except ValueError: + _oauth_invalid_request( + "redirect_uri is not a valid URL.", + hint="Use a full absolute URL for redirect_uri (e.g. https://your-host/ui/mcp/oauth/callback).", + ) + + +def _validate_trusted_http_redirect_shape(parsed: ParseResult) -> bool: + """Return True when ``parsed`` is an allowlisted native callback (caller may return).""" + if parsed.scheme not in ("http", "https"): + if _matches_trusted_native_redirect_uri(parsed): + return True + _oauth_invalid_request( + f"redirect_uri scheme {parsed.scheme!r} is not allowed; use http/https " + "or a registered native callback (e.g. cursor://).", + hint="Add the full URI to MCP_TRUSTED_NATIVE_REDIRECT_URIS for custom native clients.", + ) + if parsed.fragment: + _oauth_invalid_request( + "redirect_uri must not contain a URL fragment (#...).", + ) + if not parsed.netloc: + _oauth_invalid_request( + "redirect_uri must include a host (e.g. https://your-host/path).", + ) + if parsed.username is not None or parsed.password is not None: + _oauth_invalid_request( + "redirect_uri must not contain userinfo (user:pass@host).", + ) + if "\\" in parsed.netloc: + _oauth_invalid_request( + "redirect_uri host must not contain backslashes.", + ) + return False + + +def _resolve_proxy_base_for_redirect(request: Request) -> Optional[str]: + try: + return get_request_base_url(request) + except Exception as exc: + verbose_logger.warning( + "validate_trusted_redirect_uri: could not determine proxy origin, " + "falling back to loopback + allowlist. error=%s", + exc, + ) + return None + + +def _trusted_redirect_uri_is_allowed( + parsed: ParseResult, + redirect_netloc: str, + proxy_base: Optional[str], +) -> bool: + if proxy_base: + proxy_parsed = urlparse(proxy_base) + if ( + parsed.scheme == proxy_parsed.scheme + and redirect_netloc + == _strip_default_port(proxy_parsed.scheme, proxy_parsed.netloc) + ): + return True + + host = (parsed.hostname or "").lower() + if host == "localhost": + return True + try: + if ip_address(host).is_loopback: + return True + except ValueError: + pass + + if parsed.scheme == "https": + for entry in _parse_trusted_redirect_origins(): + if _matches_trusted_origin_entry(redirect_netloc, entry): + return True + return False + + +def _build_trusted_redirect_rejection_message( + redirect_uri: str, + parsed: ParseResult, + redirect_netloc: str, + proxy_base: Optional[str], +) -> str: + """Build a client-facing rejection message. + + Intentionally omits the proxy's resolved scheme / host / port to avoid + leaking internal network topology (e.g. ``http://litellm-internal:4000``) + through an unauthenticated endpoint. Full diagnostic detail — including + the computed proxy base — is logged server-side by the caller. + """ + redirect_origin = _origin_label(parsed.scheme, redirect_netloc) + proxy_parsed = urlparse(proxy_base) if proxy_base else None + proxy_netloc_norm = ( + _strip_default_port(proxy_parsed.scheme, proxy_parsed.netloc) + if proxy_parsed and proxy_parsed.netloc + else "" + ) + + mismatch_parts: List[str] = [] + if proxy_parsed and proxy_parsed.netloc: + if parsed.scheme != proxy_parsed.scheme: + mismatch_parts.append( + f"scheme: redirect_uri uses {parsed.scheme!r}, but the proxy " + "resolved a different scheme " + "(TLS often terminates at ingress — set PROXY_BASE_URL to https://… " + "or trust X-Forwarded-Proto from your ingress)" + ) + if redirect_netloc != proxy_netloc_norm: + mismatch_parts.append( + f"host/port: redirect_uri {redirect_netloc!r} does not match " + "the proxy origin" + ) + + if mismatch_parts: + return ( + f"redirect_uri origin ({redirect_origin}) does not match the proxy " + "origin. " + "; ".join(mismatch_parts) + ) + return ( + f"redirect_uri ({redirect_uri!r}) is not allowed: not same-origin with " + f"the proxy origin, not loopback, and not listed in " + f"{_TRUSTED_REDIRECT_ORIGINS_ENV}." + ) + + +def _raise_trusted_redirect_uri_rejected( + request: Request, + redirect_uri: str, + parsed: ParseResult, + redirect_netloc: str, + proxy_base: Optional[str], +) -> NoReturn: + description = _build_trusted_redirect_rejection_message( + redirect_uri, parsed, redirect_netloc, proxy_base + ) + + hint = ( + "Align the proxy public URL with the browser URL. Set PROXY_BASE_URL to your " + "HTTPS origin (e.g. https://litellm.example.com), or enable " + "general_settings.use_x_forwarded_for with mcp_trusted_proxy_ranges for your " + "ingress. Verify: curl https:///.well-known/oauth-authorization-server " + "| jq .issuer — issuer must match window.location.origin in the UI." + ) + + verbose_logger.warning( + "MCP OAuth: rejecting redirect_uri %r. %s " + "Computed proxy base=%r (PROXY_BASE_URL=%r). " + "Inbound headers: X-Forwarded-Proto=%r X-Forwarded-Host=%r " + "X-Forwarded-Port=%r Host=%r. " + "Trusted-redirect-origins env=%r. " + "Trusted-native-redirect-uris env=%r.", + redirect_uri, + description, + proxy_base, + os.environ.get("PROXY_BASE_URL"), + request.headers.get("X-Forwarded-Proto"), + request.headers.get("X-Forwarded-Host"), + request.headers.get("X-Forwarded-Port"), + request.headers.get("Host"), + os.environ.get(_TRUSTED_REDIRECT_ORIGINS_ENV), + os.environ.get(_TRUSTED_NATIVE_REDIRECT_URIS_ENV), + ) + + _oauth_invalid_request( + description, + hint=hint, + redirect_uri=redirect_uri, + ) + + def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: """Accept ``redirect_uri`` when it is (a) same-origin with the - proxy's own request origin, (b) loopback, or (c) listed in the - ``MCP_TRUSTED_REDIRECT_ORIGINS`` ops allowlist. + proxy's own request origin, (b) loopback, (c) listed in the + ``MCP_TRUSTED_REDIRECT_ORIGINS`` ops allowlist, or (d) a built-in / + env-configured native MCP client callback (e.g. ``cursor://``). Same-origin is VERIA-57's threat-model-safe equivalent of loopback: an attacker who can host content on the proxy's own HTTPS origin @@ -234,92 +518,13 @@ def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: BYOK endpoints, which only serve native MCP clients, retain :func:`validate_loopback_redirect_uri`. """ - try: - parsed = urlparse(redirect_uri) - except ValueError: - raise HTTPException(status_code=400, detail="invalid_request") - if parsed.scheme not in ("http", "https"): - raise HTTPException(status_code=400, detail="invalid_request") - if parsed.fragment: - raise HTTPException(status_code=400, detail="invalid_request") - if not parsed.netloc or parsed.username is not None or parsed.password is not None: - raise HTTPException(status_code=400, detail="invalid_request") - # Reject userinfo (``user:pass@host``) outright: OAuth redirect_uris - # have no legitimate reason to carry credentials, and allowing them - # opens a host-confusion attack where the netloc *looks* allowlisted - # (``app.example.com:443@attacker.example``) but the browser navigates - # to the post-``@`` host and hands the authorization code to the - # attacker. We compare against ``hostname`` after this, but defense in - # depth keeps malformed netloc strings from reaching the wildcard - # splitter. - if parsed.username is not None or parsed.password is not None: - raise HTTPException(status_code=400, detail="invalid_request") - # Reject backslash in netloc: urlparse keeps ``\`` as part of netloc, - # but browsers normalize ``\`` to ``/`` for http(s) URLs and treat it - # as the start of the path. An attacker can exploit that split by - # crafting ``https://attacker.net\app.example.com/cb`` — urlparse sees - # ``attacker.net\app.example.com`` (matches ``*.example.com``) while - # the browser navigates to ``attacker.net`` with the auth code. - if "\\" in parsed.netloc: - raise HTTPException(status_code=400, detail="invalid_request") - - redirect_netloc = _strip_default_port(parsed.scheme, parsed.netloc) - - # (a) Same-origin. Swallow ``get_request_base_url`` failures so the - # loopback + allowlist paths remain reachable when the origin can't - # be determined (e.g. request came from an untrusted proxy and - # ``get_request_base_url`` raised). - proxy_base: Optional[str] = None - try: - proxy_base = get_request_base_url(request) - except Exception as exc: - verbose_logger.warning( - "validate_trusted_redirect_uri: could not determine proxy origin, " - "falling back to loopback + allowlist. error=%s", - exc, - ) - proxy_base = None - if proxy_base: - proxy_parsed = urlparse(proxy_base) - if ( - parsed.scheme == proxy_parsed.scheme - and redirect_netloc - == _strip_default_port(proxy_parsed.scheme, proxy_parsed.netloc) - ): - return - - # (b) Loopback — same rule as validate_loopback_redirect_uri. - host = (parsed.hostname or "").lower() - if host == "localhost": + parsed = _parse_redirect_uri_for_validation(redirect_uri) + if _validate_trusted_http_redirect_shape(parsed): return - try: - if ip_address(host).is_loopback: - return - except ValueError: - pass - - # (c) Ops allowlist. https only. - if parsed.scheme == "https": - for entry in _parse_trusted_redirect_origins(): - if _matches_trusted_origin_entry(redirect_netloc, entry): - return - - verbose_logger.warning( - "MCP OAuth: rejecting redirect_uri %r as invalid_request. " - "Computed proxy base=%r (PROXY_BASE_URL=%r). " - "Inbound headers: X-Forwarded-Proto=%r X-Forwarded-Host=%r " - "X-Forwarded-Port=%r Host=%r. " - "Trusted-redirect-origins env=%r. " - "If this should be accepted, either align ingress X-Forwarded-* " - "with the browser URL, set PROXY_BASE_URL to your public origin, " - "or add the redirect_uri host to MCP_TRUSTED_REDIRECT_ORIGINS.", - redirect_uri, - proxy_base, - os.environ.get("PROXY_BASE_URL"), - request.headers.get("X-Forwarded-Proto"), - request.headers.get("X-Forwarded-Host"), - request.headers.get("X-Forwarded-Port"), - request.headers.get("Host"), - os.environ.get(_TRUSTED_REDIRECT_ORIGINS_ENV), + redirect_netloc = _strip_default_port(parsed.scheme, parsed.netloc) + proxy_base = _resolve_proxy_base_for_redirect(request) + if _trusted_redirect_uri_is_allowed(parsed, redirect_netloc, proxy_base): + return + _raise_trusted_redirect_uri_rejected( + request, redirect_uri, parsed, redirect_netloc, proxy_base ) - raise HTTPException(status_code=400, detail="invalid_request") diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 271517bb1e6..de70fe1331e 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -305,7 +305,7 @@ def build_input_schema(operation: Dict[str, Any]) -> Dict[str, Any]: def _merge_openapi_tool_request_headers( - static_headers: Dict[str, str] + static_headers: Dict[str, str], ) -> Dict[str, str]: """Merge static closure headers with per-request ContextVar overrides. diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 829863d2dbb..cec5224e183 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,6 +1,17 @@ import importlib from datetime import datetime -from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Set, Union +from typing import ( + Any, + Awaitable, + Callable, + Dict, + List, + Literal, + Optional, + Set, + Tuple, + Union, +) from fastapi import APIRouter, Depends, HTTPException, Query, Request, status @@ -51,20 +62,16 @@ if MCP_AVAILABLE: mcp_auth_header: Optional[str], ) -> Optional[Union[Dict[str, str], str]]: """Helper function to get server-specific auth header with case-insensitive matching.""" - if mcp_server_auth_headers and server.alias: - normalized_server_alias = server.alias.lower() - normalized_headers = { - k.lower(): v for k, v in mcp_server_auth_headers.items() - } - server_auth = normalized_headers.get(normalized_server_alias) - if server_auth is not None: - return server_auth - elif mcp_server_auth_headers and server.server_name: - normalized_server_name = server.server_name.lower() - normalized_headers = { - k.lower(): v for k, v in mcp_server_auth_headers.items() - } - server_auth = normalized_headers.get(normalized_server_name) + from litellm.proxy._experimental.mcp_server.utils import ( + lookup_mcp_server_auth_in_headers, + ) + + if mcp_server_auth_headers: + server_auth = lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers, + alias=getattr(server, "alias", None), + server_name=getattr(server, "server_name", None), + ) if server_auth is not None: return server_auth return mcp_auth_header @@ -231,11 +238,32 @@ if MCP_AVAILABLE: ) return mcp_auth_header, mcp_server_auth_headers, raw_headers + def _resolve_mcp_server_id_for_rest( + server_id: str, + allowed_server_ids: Union[Set[str], List[str]], + client_ip: Optional[str] = None, + ) -> str: + """ + Map REST ``server_id`` (UUID, server_name, or alias) to canonical server_id. + + tools/list already did this; tools/call must match so clients can pass + server names like ``order_status_mcp`` instead of only UUIDs. + """ + allowed = set(allowed_server_ids) + if server_id in allowed: + return server_id + by_name = global_mcp_server_manager.get_mcp_server_by_name( + server_id, client_ip=client_ip + ) + if by_name is not None and by_name.server_id in allowed: + return by_name.server_id + return server_id + async def _resolve_allowed_mcp_servers_with_ip_filter( request: Request, user_api_key_dict: UserAPIKeyAuth, server_id: str, - ) -> List[MCPServer]: + ) -> Tuple[List[MCPServer], str]: """ Resolve allowed MCP servers for a tool call with IP filtering. @@ -245,10 +273,10 @@ if MCP_AVAILABLE: server_id: The server ID to validate access for Returns: - List of allowed MCPServer objects + Tuple of (allowed MCPServer objects, canonical server_id) Raises: - HTTPException: If the server_id is not allowed + HTTPException: If the server_id is not allowed or not found """ # Get all auth contexts auth_contexts = await build_effective_auth_contexts(user_api_key_dict) @@ -268,8 +296,41 @@ if MCP_AVAILABLE: ) ) - # Check if the specified server_id is allowed - if server_id not in allowed_server_ids_set: + canonical_server_id = _resolve_mcp_server_id_for_rest( + server_id, allowed_server_ids_set, _rest_client_ip + ) + + if canonical_server_id not in allowed_server_ids_set: + _server = global_mcp_server_manager.get_mcp_server_by_id( + server_id + ) or global_mcp_server_manager.get_mcp_server_by_name(server_id) + if ( + _server is not None + and _rest_client_ip is not None + and not global_mcp_server_manager._is_server_accessible_from_ip( + _server, _rest_client_ip + ) + ): + raise HTTPException( + status_code=403, + detail={ + "error": "ip_filtering", + "message": ( + f"MCP server '{server_id}' is not accessible from your IP address " + f"({_rest_client_ip}). This server is restricted to internal " + "networks only. To make it externally accessible, set " + "'available_on_public_internet: true' in the server configuration." + ), + }, + ) + if _server is None: + raise HTTPException( + status_code=404, + detail={ + "error": "server_not_found", + "message": f"MCP server '{server_id}' was not found", + }, + ) raise HTTPException( status_code=403, detail={ @@ -285,7 +346,7 @@ if MCP_AVAILABLE: if server is not None: allowed_mcp_servers.append(server) - return allowed_mcp_servers + return allowed_mcp_servers, canonical_server_id async def _get_tools_for_single_server( server, @@ -301,6 +362,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=False, raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) # Filter tools based on allowed_tools configuration @@ -753,7 +815,7 @@ if MCP_AVAILABLE: }, ) - tool_arguments = data.get("arguments") + tool_arguments = data.get("arguments") or {} proxy_base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -786,14 +848,18 @@ if MCP_AVAILABLE: data["user_api_key_auth"] = data["metadata"]["user_api_key_auth"] # Resolve allowed MCP servers with IP filtering - allowed_mcp_servers = await _resolve_allowed_mcp_servers_with_ip_filter( + ( + allowed_mcp_servers, + canonical_server_id, + ) = await _resolve_allowed_mcp_servers_with_ip_filter( request, user_api_key_dict, server_id ) # Look up per-user OAuth headers for this server (mirrors list_tool_rest_api). user_oauth_extra_headers: Optional[Dict[str, str]] = None target_server = next( - (s for s in allowed_mcp_servers if s.server_id == server_id), None + (s for s in allowed_mcp_servers if s.server_id == canonical_server_id), + None, ) if target_server is not None: user_oauth_extra_headers = await _get_user_oauth_extra_headers( @@ -812,6 +878,7 @@ if MCP_AVAILABLE: oauth2_headers=user_oauth_extra_headers or data.get("oauth2_headers"), raw_headers=data.get("raw_headers"), litellm_logging_obj=data.get("litellm_logging_obj"), + requested_server_id=canonical_server_id, ) return result except BlockedPiiEntityError as e: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 0a74a92f9ce..5205426edf3 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1114,10 +1114,16 @@ if MCP_AVAILABLE: ) -> Tuple[Optional[Union[Dict[str, str], str]], Optional[Dict[str, str]]]: """Build auth and extra headers for a server.""" server_auth_header: Optional[Union[Dict[str, str], str]] = None - if mcp_server_auth_headers and server.alias is not None: - server_auth_header = mcp_server_auth_headers.get(server.alias) - elif mcp_server_auth_headers and server.server_name is not None: - server_auth_header = mcp_server_auth_headers.get(server.server_name) + if mcp_server_auth_headers: + from litellm.proxy._experimental.mcp_server.utils import ( + lookup_mcp_server_auth_in_headers, + ) + + server_auth_header = lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers, + alias=server.alias, + server_name=server.server_name, + ) extra_headers: Optional[Dict[str, str]] = None if server.auth_type == MCPAuth.oauth2: @@ -1368,6 +1374,7 @@ if MCP_AVAILABLE: extra_headers=extra_headers, add_prefix=True, # Always add server prefix raw_headers=raw_headers, + user_api_key_auth=user_api_key_auth, ) filtered_tools = filter_tools_by_allowed_tools(tools, server) @@ -2074,6 +2081,7 @@ if MCP_AVAILABLE: """ # Track resolved MCP server for both permission checks and dispatch mcp_server: Optional[MCPServer] = None + requested_server_id: Optional[str] = kwargs.get("requested_server_id") # If the client called with a display-name override (e.g. "Get Pet"), # translate it back to the original prefixed name before any routing. @@ -2082,14 +2090,55 @@ if MCP_AVAILABLE: # Remove prefix from tool name for logging and processing original_tool_name, server_name = split_server_prefix_from_name(name) + requested_server: Optional[MCPServer] = None + if requested_server_id: + requested_server = next( + (s for s in allowed_mcp_servers if s.server_id == requested_server_id), + None, + ) + # Resolve the actual MCP server up-front so the permission check uses # the canonical server.name even when the tool name is prefixed with a # short ID (LITELLM_USE_SHORT_MCP_TOOL_PREFIX) that doesn't match the # server's display name directly. mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + if mcp_server is None and requested_server is not None: + # REST callers may pass the raw tool name (no prefix) plus a + # ``requested_server_id``. The mapping might only contain the + # prefixed form, so retry the lookup with every known prefix of + # the requested server before treating the tool as unresolved — + # otherwise the tool_server_mismatch guard below is silently + # bypassed. + for known_prefix in iter_known_server_prefixes(requested_server): + candidate = global_mcp_server_manager._get_mcp_server_from_tool_name( + add_server_prefix_to_name(name, known_prefix) + ) + if candidate is not None: + mcp_server = candidate + break if mcp_server is not None: server_name = mcp_server.name + # REST /mcp-rest/tools/call passes server_id — tool must belong to that server + if requested_server is not None: + if ( + mcp_server is not None + and mcp_server.server_id != requested_server.server_id + ): + raise HTTPException( + status_code=403, + detail={ + "error": "tool_server_mismatch", + "message": ( + f"Tool '{name}' belongs to MCP server '{mcp_server.name}' " + f"but request specified server_id for '{requested_server.name}'." + ), + }, + ) + if mcp_server is None: + mcp_server = requested_server + server_name = requested_server.name + # Only enforce server-level permissions when we can resolve a server if server_name: if not MCPRequestHandler.is_tool_allowed( diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index df5705c3425..b8b9207555e 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -2,7 +2,8 @@ MCP Server Utilities """ -from typing import Any, Dict, Iterator, Mapping, Optional, Tuple +import re +from typing import Any, Dict, Iterator, Mapping, Optional, Tuple, Union import hashlib import importlib @@ -117,6 +118,50 @@ def normalize_server_name(server_name: str) -> str: return server_name.replace(" ", "_") +_MCP_ALIAS_HEADER_INVALID_RE = re.compile(r"[^a-z0-9_]") + + +def sanitize_mcp_alias_for_header(alias: str) -> str: + """ + Sanitize an MCP server alias for x-mcp-{alias}-{header} HTTP headers. + + Must stay in sync with ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts. + """ + sanitized = _MCP_ALIAS_HEADER_INVALID_RE.sub("_", alias.lower().strip()) + sanitized = re.sub(r"_+", "_", sanitized) + return sanitized.strip("_") + + +def lookup_mcp_server_auth_in_headers( + mcp_server_auth_headers: Mapping[str, Union[str, Dict[str, str]]], + *, + alias: Optional[str] = None, + server_name: Optional[str] = None, +) -> Optional[Union[str, Dict[str, str]]]: + """ + Resolve server-specific auth headers with case-insensitive matching. + + Tries the raw alias/server_name (lowercased) and the header-safe sanitized + alias so dashboard clients using sanitize_mcp_alias_for_header() still match. + """ + if not mcp_server_auth_headers: + return None + + normalized_headers = {k.lower(): v for k, v in mcp_server_auth_headers.items()} + + for identifier in (alias, server_name): + if not identifier: + continue + keys_to_try = [identifier.lower()] + sanitized = sanitize_mcp_alias_for_header(identifier) + if sanitized and sanitized not in keys_to_try: + keys_to_try.append(sanitized) + for key in keys_to_try: + if key in normalized_headers: + return normalized_headers[key] + return None + + def validate_and_normalize_mcp_server_payload(payload: Any) -> None: """ Validate and normalize MCP server payload fields (server_name and alias). diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/_not-found.html b/litellm/proxy/_experimental/out/_not-found/index.html similarity index 100% rename from litellm/proxy/_experimental/out/_not-found.html rename to litellm/proxy/_experimental/out/_not-found/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/chat.html b/litellm/proxy/_experimental/out/chat/index.html similarity index 100% rename from litellm/proxy/_experimental/out/chat.html rename to litellm/proxy/_experimental/out/chat/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/skills.html b/litellm/proxy/_experimental/out/skills/index.html similarity index 100% rename from litellm/proxy/_experimental/out/skills.html rename to litellm/proxy/_experimental/out/skills/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index eea6974193f..27cdc483d4a 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3171,7 +3171,7 @@ ] }, "post": { - "description": "Create a new agent\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/agents\" \\\n -H \"Authorization: Bearer \" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent\": {\n \"agent_name\": \"my-custom-agent\",\n \"agent_card_params\": {\n \"protocolVersion\": \"1.0\",\n \"name\": \"Hello World Agent\",\n \"description\": \"Just a hello world agent\",\n \"url\": \"http://localhost:9999/\",\n \"version\": \"1.0.0\",\n \"defaultInputModes\": [\"text\"],\n \"defaultOutputModes\": [\"text\"],\n \"capabilities\": {\n \"streaming\": true\n },\n \"skills\": [\n {\n \"id\": \"hello_world\",\n \"name\": \"Returns hello world\",\n \"description\": \"just returns hello world\",\n \"tags\": [\"hello world\"],\n \"examples\": [\"hi\", \"hello world\"]\n }\n ]\n },\n \"litellm_params\": {\n \"make_public\": true\n }\n }\n }'\n```", + "description": "Create a new agent\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/v1/agents\" \\\n -H \"Authorization: Bearer \" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent_name\": \"my-custom-agent\",\n \"agent_card_params\": {\n \"protocolVersion\": \"1.0\",\n \"name\": \"Hello World Agent\",\n \"description\": \"Just a hello world agent\",\n \"url\": \"http://localhost:9999/\",\n \"version\": \"1.0.0\",\n \"defaultInputModes\": [\"text\"],\n \"defaultOutputModes\": [\"text\"],\n \"capabilities\": {\n \"streaming\": true\n },\n \"skills\": [\n {\n \"id\": \"hello_world\",\n \"name\": \"Returns hello world\",\n \"description\": \"just returns hello world\",\n \"tags\": [\"hello world\"],\n \"examples\": [\"hi\", \"hello world\"]\n }\n ]\n },\n \"litellm_params\": {\n \"make_public\": true\n }\n }'\n```", "operationId": "create_agent_v1_agents_post", "requestBody": { "content": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9337aa7c8ea..004f33e630a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2361,6 +2361,30 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): database_connection_timeout: Optional[float] = Field( 60, description="default timeout for a connection to the database" ) + database_connect_timeout: Optional[float] = Field( + None, + description=( + "Prisma `connect_timeout` URL param (seconds). Bounds how long the " + "engine waits to establish a new connection before failing. Defaults " + "to Prisma's built-in value when unset." + ), + ) + database_socket_timeout: Optional[float] = Field( + None, + description=( + "Prisma `socket_timeout` URL param (seconds). When set, an idle/slow " + "connection that has not produced data within this window is closed. " + "This is the main knob for capping idle DB connections from LiteLLM." + ), + ) + database_extra_connection_params: Optional[Dict[str, Any]] = Field( + None, + description=( + "Escape hatch: extra key/value pairs appended verbatim to the Prisma " + "DATABASE_URL / DIRECT_URL query string (e.g. `sslmode`, `pgbouncer`, " + "`statement_cache_size`). Keys here override any default LiteLLM sets." + ), + ) database_type: Optional[Literal["dynamo_db"]] = Field( None, description="to use dynamodb instead of postgres db" ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 13381c7a6c9..14f198e0f12 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1187,6 +1187,127 @@ async def get_end_user_object( return None +_END_USER_VALIDATION_NEGATIVE_TTL = 60 +_END_USER_VALIDATION_POSITIVE_TTL = 300 + + +async def resolve_and_validate_end_user_id( + raw_end_user_id: Optional[str], + prisma_client: Optional[PrismaClient], + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Optional[Span] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, + route: str = "", +) -> Optional[str]: + """Optionally drop end-user ids that don't resolve to a known DB row. + + Default: pass-through. LiteLLM's documented pattern is that the `user` + field is an arbitrary caller-supplied identifier, so validation is + opt-in behind ``litellm.validate_end_user_id_in_db`` to preserve + backwards compatibility. + + When the flag is set: accept the id when it matches any of + - LiteLLM_EndUserTable.user_id + - LiteLLM_UserTable.user_id + - LiteLLM_UserTable.user_email (case-insensitive) + + If the id doesn't match but ``litellm.max_end_user_budget_id`` is set, + we still preserve the id so the default end-user budget is applied + downstream; otherwise we return None. + + DB lookups reuse ``get_end_user_object`` / ``get_user_object`` so they + share the same cache as the rest of the auth path instead of adding new + raw Prisma queries. + """ + if raw_end_user_id is None: + return None + if not litellm.validate_end_user_id_in_db: + return raw_end_user_id + if prisma_client is None: + return raw_end_user_id + + cache_key = f"end_user_validation:{raw_end_user_id}" + cached = await user_api_key_cache.async_get_cache(key=cache_key) + if cached == "valid": + return raw_end_user_id + if cached == "invalid": + return raw_end_user_id if litellm.max_end_user_budget_id else None + + is_valid = await _end_user_id_exists_in_db( + end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + + await user_api_key_cache.async_set_cache( + key=cache_key, + value="valid" if is_valid else "invalid", + ttl=( + _END_USER_VALIDATION_POSITIVE_TTL + if is_valid + else _END_USER_VALIDATION_NEGATIVE_TTL + ), + ) + + if is_valid: + return raw_end_user_id + # Preserve id so the caller can still apply litellm.max_end_user_budget_id. + if litellm.max_end_user_budget_id: + return raw_end_user_id + return None + + +async def _end_user_id_exists_in_db( + end_user_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Optional[Span] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, + route: str = "", +) -> bool: + """True when the id matches an EndUser, User, or user_email row.""" + try: + end_user_obj = await get_end_user_object( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if end_user_obj is not None: + return True + except litellm.BudgetExceededError: + raise + except Exception as e: + verbose_proxy_logger.debug( + f"end_user validation: get_end_user_object lookup failed: {e}" + ) + + try: + user_obj = await get_user_object( + user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + check_db_only=False, + user_email=end_user_id if "@" in end_user_id else None, + ) + if user_obj is not None: + return True + except Exception as e: + verbose_proxy_logger.debug( + f"end_user validation: get_user_object lookup failed: {e}" + ) + + return False + + @log_db_metrics async def get_tag_objects_batch( tag_names: List[str], @@ -2232,7 +2353,9 @@ class ExperimentalUIJWTToken: @staticmethod def get_cli_jwt_auth_token( - user_info: LiteLLM_UserTable, team_id: Optional[str] = None + user_info: LiteLLM_UserTable, + team_id: Optional[str] = None, + team_alias: Optional[str] = None, ) -> str: """ Generate a JWT token for CLI authentication with configurable expiration. @@ -2243,6 +2366,7 @@ class ExperimentalUIJWTToken: Args: user_info: User information from the database team_id: Team ID for the user (optional, uses user's team if available) + team_alias: Team alias for the selected team, if available Returns: Encrypted JWT token string @@ -2276,6 +2400,7 @@ class ExperimentalUIJWTToken: expires=expires, user_id=user_info.user_id, team_id=_team_id, + team_alias=team_alias, models=user_info.models, max_parallel_requests=None, user_role=LitellmUserRoles(user_info.user_role), diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 637a4a070c4..c4dcca764b2 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -10,6 +10,7 @@ import litellm from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS +from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS @@ -1008,12 +1009,47 @@ def _get_customer_id_from_standard_headers( for standard_header in STANDARD_CUSTOMER_ID_HEADERS: for header_name, header_value in request_headers.items(): if header_name.lower() == standard_header.lower(): - user_id_str = str(header_value) if header_value is not None else "" - if user_id_str.strip(): + user_id_str = _coerce_user_id_to_str(header_value) + if user_id_str: return user_id_str return None +def _coerce_user_id_to_str(value: Any) -> Optional[str]: + """Return a usable end-user identifier string, or None if the value isn't one. + + Always drops non-string structured values (dict/list/tuple/set) because + stringifying them produces garbage spend-log rows like + ``"{'device_id': ...}"``. Strings that *decode* to a structured payload + are only rejected when ``litellm.validate_end_user_id_in_db`` is enabled + — operators who currently pass JSON-encoded identifiers keep their + existing behavior until they opt in. See + auth_utils.py:get_end_user_id_from_request_body for the extraction chain. + """ + if value is None: + return None + if isinstance(value, bool): + # bool is an int subclass; handle explicitly to avoid "True"/"False". + return None + if isinstance(value, (int, float)): + return str(value) + if isinstance(value, str): + stripped = value.strip() + if not stripped: + return None + # Reject strings that decode to a structured payload (JSON object/array) + # only when the operator has opted into end-user validation. Gating + # behind the flag preserves backwards compatibility for deployments + # that intentionally pass JSON-encoded user identifiers. + if litellm.validate_end_user_id_in_db and stripped[:1] in ("{", "["): + parsed = safe_json_loads(stripped) + if isinstance(parsed, (dict, list)): + return None + return stripped + # dict, list, tuple, set, arbitrary objects -> drop. + return None + + def get_end_user_id_from_request_body( request_body: dict, request_headers: Optional[dict] = None ) -> Optional[str]: @@ -1052,23 +1088,22 @@ def get_end_user_id_from_request_body( if isinstance(custom_header_name_to_check, list): headers_lower = {k.lower(): v for k, v in request_headers.items()} for expected_header in custom_header_name_to_check: - header_value = headers_lower.get(expected_header) - if header_value is not None: - user_id_str = str(header_value) - if user_id_str.strip(): - return user_id_str + user_id_str = _coerce_user_id_to_str(headers_lower.get(expected_header)) + if user_id_str: + return user_id_str elif isinstance(custom_header_name_to_check, str): for header_name, header_value in request_headers.items(): if header_name.lower() == custom_header_name_to_check.lower(): - user_id_str = str(header_value) if header_value is not None else "" - if user_id_str.strip(): + user_id_str = _coerce_user_id_to_str(header_value) + if user_id_str: return user_id_str # Check 3: 'user' field in request_body (commonly OpenAI) - if "user" in request_body and request_body["user"] is not None: - user_from_body_user_field = request_body["user"] - return str(user_from_body_user_field) + if "user" in request_body: + user_id_str = _coerce_user_id_to_str(request_body["user"]) + if user_id_str: + return user_id_str def _as_dict(value: Any) -> dict: # metadata / litellm_metadata can arrive as JSON strings from @@ -1077,32 +1112,30 @@ def get_end_user_id_from_request_body( if isinstance(value, dict): return value if isinstance(value, str): - from litellm.litellm_core_utils.safe_json_loads import safe_json_loads - parsed = safe_json_loads(value) return parsed if isinstance(parsed, dict) else {} return {} # Check 4: 'litellm_metadata.user' in request_body (commonly Anthropic) litellm_metadata = _as_dict(request_body.get("litellm_metadata")) - user_from_litellm_metadata = litellm_metadata.get("user") - if user_from_litellm_metadata is not None: - return str(user_from_litellm_metadata) + user_id_str = _coerce_user_id_to_str(litellm_metadata.get("user")) + if user_id_str: + return user_id_str # Check 5: 'metadata.user_id' in request_body (another common pattern) metadata_dict = _as_dict(request_body.get("metadata")) - user_id_from_metadata_field = metadata_dict.get("user_id") - if user_id_from_metadata_field is not None: - return str(user_id_from_metadata_field) + user_id_str = _coerce_user_id_to_str(metadata_dict.get("user_id")) + if user_id_str: + return user_id_str # Check 6: 'safety_identifier' in request body (OpenAI Responses API parameter) # SECURITY NOTE: safety_identifier can be set by any caller in the request body. # Only use this for end-user identification in trusted environments where you control # the calling application. For untrusted callers, prefer using headers or server-side # middleware to set the end_user_id to prevent impersonation. - if request_body.get("safety_identifier") is not None: - user_from_body_user_field = request_body["safety_identifier"] - return str(user_from_body_user_field) + user_id_str = _coerce_user_id_to_str(request_body.get("safety_identifier")) + if user_id_str: + return user_id_str return None diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index bf76f99db69..d364b52c676 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -4,12 +4,15 @@ from typing import Dict, List, Optional, Set import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth from litellm.router import Router from litellm.router_utils.fallback_event_handlers import get_fallback_model_group -from litellm.types.router import LiteLLM_Params +from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params from litellm.utils import get_valid_models +_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) + def _check_wildcard_routing(model: str) -> bool: """ @@ -178,6 +181,7 @@ def get_complete_model_list( model_access_groups: Dict[str, List[str]] = {}, include_model_access_groups: Optional[bool] = False, only_model_access_groups: Optional[bool] = False, + team_id: Optional[str] = None, ) -> List[str]: """Logic for returning complete model list for a given key + team pair""" @@ -222,6 +226,7 @@ def get_complete_model_list( unique_models=unique_models, return_wildcard_routes=return_wildcard_routes, llm_router=llm_router, + team_id=team_id, ) complete_model_list = unique_models + all_wildcard_models @@ -229,6 +234,29 @@ def get_complete_model_list( return complete_model_list +def _hydrate_litellm_credential_name( + litellm_params: Optional[LiteLLM_Params], +) -> Optional[LiteLLM_Params]: + if litellm_params is None or litellm_params.litellm_credential_name is None: + return litellm_params + + credential_values = CredentialAccessor.get_credential_values( + litellm_params.litellm_credential_name + ) + if not credential_values: + return litellm_params + + litellm_params = litellm_params.model_copy() + for key, value in credential_values.items(): + if ( + key in _CREDENTIAL_LITELLM_PARAM_FIELDS + and getattr(litellm_params, key, None) is None + ): + setattr(litellm_params, key, value) + litellm_params.litellm_credential_name = None + return litellm_params + + def get_known_models_from_wildcard( wildcard_model: str, litellm_params: Optional[LiteLLM_Params] = None ) -> List[str]: @@ -247,7 +275,7 @@ def get_known_models_from_wildcard( else: provider = wildcard_provider_prefix - # get all known provider models + litellm_params = _hydrate_litellm_credential_name(litellm_params) wildcard_models = get_provider_models( provider=provider, litellm_params=litellm_params @@ -285,6 +313,7 @@ def _get_wildcard_models( unique_models: List[str], return_wildcard_routes: Optional[bool] = False, llm_router: Optional[Router] = None, + team_id: Optional[str] = None, ) -> List[str]: models_to_remove = set() all_wildcard_models = [] @@ -297,7 +326,9 @@ def _get_wildcard_models( ## get litellm params from model if llm_router is not None: - model_list = llm_router.get_model_list(model_name=model) + model_list = llm_router.get_model_list( + model_name=model, team_id=team_id + ) if model_list: for router_model in model_list: wildcard_models = get_known_models_from_wildcard( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 30b5d36e14a..6974860a22a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -12,7 +12,7 @@ import fnmatch import re import secrets from datetime import datetime, timezone -from typing import Any, Iterator, List, Optional, Tuple, Union, cast +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -44,6 +44,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_object, get_user_object, is_valid_fallback_model, + resolve_and_validate_end_user_id, ) from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_utils import ( @@ -333,8 +334,22 @@ def _apply_budget_limits_to_end_user_params( async def user_api_key_auth_websocket(websocket: WebSocket): # Accept the WebSocket connection - scope_headers = list(websocket.scope.get("headers") or []) - request = Request(scope={"type": "http", "headers": scope_headers}) + ws_scope = websocket.scope or {} + scope_headers = list(ws_scope.get("headers") or []) + # ``get_request_route`` falls back to ``request.url.path`` when + # ``scope["path"]`` is absent. On WebSockets that fallback reads + # ``websocket.url``, which Starlette reconstructs from the (poisonable) + # Host header. Carry the ASGI scope's path / root_path so the lookup + # never reaches the fallback. + synthetic_scope: Dict[str, Any] = { + "type": "http", + "headers": scope_headers, + "path": ws_scope.get("path", ""), + } + for key in ("root_path", "app_root_path"): + if key in ws_scope: + synthetic_scope[key] = ws_scope[key] + request = Request(scope=synthetic_scope) request._url = websocket.url @@ -1057,9 +1072,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 _end_user_object = None end_user_params = {} - end_user_id = get_end_user_id_from_request_body( + raw_end_user_id = get_end_user_id_from_request_body( request_data, _safe_get_request_headers(request) ) + end_user_id = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) if end_user_id: try: end_user_params["end_user_id"] = end_user_id @@ -1745,7 +1768,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached @tracer.wrap() -async def _run_centralized_common_checks( +async def _run_centralized_common_checks( # noqa: PLR0915 user_api_key_auth_obj: UserAPIKeyAuth, request: Request, request_data: dict, @@ -1823,9 +1846,23 @@ async def _run_centralized_common_checks( return parent_otel_span = user_api_key_auth_obj.parent_otel_span - end_user_id = get_end_user_id_from_request_body( - request_data, _safe_get_request_headers(request) - ) + # In the integrated auth flow ``_user_api_key_auth_builder`` has already + # resolved the end-user id and attached it here. Reuse that to avoid a + # second extraction pass; fall back to extracting locally when the + # function is invoked in isolation (e.g. in direct unit tests). + end_user_id = user_api_key_auth_obj.end_user_id + if end_user_id is None: + raw_end_user_id = get_end_user_id_from_request_body( + request_data, _safe_get_request_headers(request) + ) + end_user_id = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) fetch_coros = [] if user_api_key_auth_obj.team_id is not None: @@ -2156,11 +2193,33 @@ async def user_api_key_auth( api_key=api_key, ) - end_user_id = get_end_user_id_from_request_body( - request_data, _safe_get_request_headers(request) - ) - if end_user_id is not None: - user_api_key_auth_obj.end_user_id = end_user_id + # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return + # paths (no master key, /user/auth route, JWT short-circuits) that bypass + # the end-user resolution block. If those paths produced an auth obj + # without an ``end_user_id`` set, fall back to extracting from the request + # body so spend logs are still attributed correctly. Validation honours + # ``litellm.validate_end_user_id_in_db``. + if user_api_key_auth_obj.end_user_id is None: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + raw_end_user_id = get_end_user_id_from_request_body( + request_data, _safe_get_request_headers(request) + ) + if raw_end_user_id is not None: + resolved_end_user_id = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if resolved_end_user_id is not None: + user_api_key_auth_obj.end_user_id = resolved_end_user_id user_api_key_auth_obj.request_route = normalize_request_route(route) return user_api_key_auth_obj diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 935b96a0e39..166ef7a66d0 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -523,6 +523,10 @@ async def retrieve_batch( # noqa: PLR0915 custom_llm_provider=custom_llm_provider, **data # type: ignore ) + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + # FIX: Update the database with the latest state from provider await update_batch_in_database( batch_id=batch_id, @@ -533,19 +537,9 @@ async def retrieve_batch( # noqa: PLR0915 verbose_proxy_logger=verbose_proxy_logger, db_batch_object=db_batch_object, operation="retrieve", + user_api_key_dict=user_api_key_dict, ) - ### CALL HOOKS ### - modify outgoing data - response = await proxy_logging_obj.post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response - ) - - # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id - # Resolve raw provider file IDs (input, output, error) to unified IDs. - if unified_batch_id: - await resolve_input_file_id_to_unified(response, prisma_client) - await resolve_output_file_ids_to_unified(response, prisma_client) - ### ALERTING ### asyncio.create_task( proxy_logging_obj.update_request_status( @@ -917,10 +911,14 @@ async def cancel_batch( **_cancel_batch_data, ) - # FIX: Update the database with the new cancelled state managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") from litellm.proxy.proxy_server import prisma_client + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + # FIX: Update the database with the new cancelled state await update_batch_in_database( batch_id=batch_id, unified_batch_id=unified_batch_id, @@ -929,11 +927,7 @@ async def cancel_batch( prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, operation="cancel", - ) - - ### CALL HOOKS ### - modify outgoing data - response = await proxy_logging_obj.post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response + user_api_key_dict=user_api_key_dict, ) ### ALERTING ### diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index adf562d69c5..9fbc6f2197d 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -350,7 +350,7 @@ The CLI provides three authentication commands: 4. **User Authentication**: User completes SSO authentication in browser 5. **Callback Processing**: SSO provider redirects back to proxy with state parameter 6. **User Code Verification**: Browser confirms the verification code shown in the CLI -7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready +7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready. When `CLI_SSO_CLAIM_MAP` is configured on the proxy, the poll response may include `attribution_metadata` (allowlisted scalar OIDC claims for client attribution). 8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json` ### Benefits of This Approach diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index a93749c3952..fa3cb02195b 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -324,7 +324,7 @@ class CustomOpenAPISpec: @staticmethod def add_chat_completion_request_schema( - openapi_schema: Dict[str, Any] + openapi_schema: Dict[str, Any], ) -> Dict[str, Any]: """ Add ProxyChatCompletionRequest schema to chat completion endpoints for documentation. @@ -380,7 +380,7 @@ class CustomOpenAPISpec: @staticmethod def add_responses_api_request_schema( - openapi_schema: Dict[str, Any] + openapi_schema: Dict[str, Any], ) -> Dict[str, Any]: """ Add ResponsesAPIRequestParams schema to responses API endpoints for documentation. @@ -410,7 +410,7 @@ class CustomOpenAPISpec: @staticmethod def add_llm_api_request_schema_body( - openapi_schema: Dict[str, Any] + openapi_schema: Dict[str, Any], ) -> Dict[str, Any]: """ Add LLM API request schema bodies to OpenAPI specification for documentation. diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 71abdfa5e9e..2ce3fda6297 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -12,6 +12,33 @@ from litellm.proxy.common_utils.callback_utils import ( ) from litellm.types.router import Deployment +_FORM_CONTENT_TYPES: frozenset[str] = frozenset( + {"application/x-www-form-urlencoded", "multipart/form-data"} +) + + +def _normalize_media_type(content_type: str) -> str: + """Return the bare media type per RFC 7231: strip params, trim, lowercase.""" + if not content_type: + return "" + return content_type.split(";", 1)[0].strip().lower() + + +def _is_form_content_type(content_type: str) -> bool: + """ + True iff Starlette's ``request.form()`` will actually parse this body. + + Substring matching ``"form"`` is unsafe: ``request.form()`` returns empty + ``FormData`` for non-canonical types without consuming the body, leaving + the auth-time pre-read and the handler's read seeing different payloads. + """ + return _normalize_media_type(content_type) in _FORM_CONTENT_TYPES + + +def _is_json_content_type(content_type: str) -> bool: + """True iff the body should be parsed as JSON.""" + return _normalize_media_type(content_type) == "application/json" + async def _read_request_body(request: Optional[Request]) -> Dict: """ @@ -37,8 +64,24 @@ async def _read_request_body(request: Optional[Request]) -> Dict: _request_headers: dict = _safe_get_request_headers(request=request) content_type = _request_headers.get("content-type", "") - if "form" in content_type: - parsed_body = dict(await request.form()) + if _is_form_content_type(content_type): + try: + form_data = await request.form() + except Exception as e: + # ``request.form()`` raises on malformed multipart (missing + # boundary, malformed chunk encoding, …). Surface as 400 so + # the auth-time pre-read does not silently cache ``{}`` while + # a later raw-body re-read sees the original payload — + # banned-param checks must see the same body the handler + # acts on. + verbose_proxy_logger.error(f"Invalid form payload: {e}") + raise ProxyException( + message=f"Invalid form payload: {e}", + type="invalid_request_error", + param="request_body", + code=status.HTTP_400_BAD_REQUEST, + ) + parsed_body = dict(form_data) if "metadata" in parsed_body and isinstance(parsed_body["metadata"], str): parsed_body["metadata"] = json.loads(parsed_body["metadata"]) else: @@ -257,7 +300,7 @@ async def get_form_data(request: Request) -> Dict[str, Any]: async def convert_upload_files_to_file_data( - form_data: Dict[str, Any] + form_data: Dict[str, Any], ) -> Dict[str, Any]: """ Convert FastAPI UploadFile objects to file data tuples for litellm. @@ -306,18 +349,13 @@ async def get_request_body(request: Request) -> Dict[str, Any]: Read the request body and parse it as JSON. """ if request.method == "POST": - if request.headers.get("content-type", "") == "application/json": + content_type = request.headers.get("content-type", "") + if _is_json_content_type(content_type): return await _read_request_body(request) - elif "multipart/form-data" in request.headers.get( - "content-type", "" - ) or "application/x-www-form-urlencoded" in request.headers.get( - "content-type", "" - ): + elif _is_form_content_type(content_type): return await get_form_data(request) else: - raise ValueError( - f"Unsupported content type: {request.headers.get('content-type')}" - ) + raise ValueError(f"Unsupported content type: {content_type}") return {} diff --git a/litellm/proxy/common_utils/openai_endpoint_utils.py b/litellm/proxy/common_utils/openai_endpoint_utils.py index c4bfe11aec1..905967fa465 100644 --- a/litellm/proxy/common_utils/openai_endpoint_utils.py +++ b/litellm/proxy/common_utils/openai_endpoint_utils.py @@ -1,5 +1,5 @@ """ -Contains utils used by OpenAI compatible endpoints +Contains utils used by OpenAI compatible endpoints """ from typing import Optional, Set diff --git a/litellm/proxy/config_management_endpoints/pass_through_endpoints.py b/litellm/proxy/config_management_endpoints/pass_through_endpoints.py index 5ff02b8bce0..4ebd989dc53 100644 --- a/litellm/proxy/config_management_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/config_management_endpoints/pass_through_endpoints.py @@ -1,5 +1,5 @@ """ -What is this? +What is this? CRUD endpoints for managing pass-through endpoints """ diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 9650604bf81..fc1f77bb684 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -328,7 +328,7 @@ async def retrieve_container( custom_llm_provider=custom_llm_provider, ) data.update( - get_container_forwarding_params( + await get_container_forwarding_params( container_id, original_container_id, custom_llm_provider, @@ -433,7 +433,7 @@ async def delete_container( custom_llm_provider=custom_llm_provider, ) data.update( - get_container_forwarding_params( + await get_container_forwarding_params( container_id, original_container_id, custom_llm_provider, diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 4284cdd5d4a..7eeb11fc372 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -196,10 +196,12 @@ async def _process_binary_request( ) data: Dict[str, Any] = { "file_id": file_id, - **get_container_forwarding_params( - container_id=container_id, - original_container_id=original_container_id, - custom_llm_provider=resolved_provider, + **( + await get_container_forwarding_params( + container_id=container_id, + original_container_id=original_container_id, + custom_llm_provider=resolved_provider, + ) ), } processor = ProxyBaseLLMRequestProcessing(data=data) @@ -316,7 +318,7 @@ async def _process_multipart_upload_request( ) data.update( - get_container_forwarding_params( + await get_container_forwarding_params( container_id=container_id, original_container_id=original_container_id, custom_llm_provider=resolved_provider, @@ -396,7 +398,7 @@ async def _process_request( ) ) data.update( - get_container_forwarding_params( + await get_container_forwarding_params( container_id=path_params["container_id"], original_container_id=original_container_id, custom_llm_provider=resolved_provider, diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 568eca523ae..57de6c4a63d 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -23,6 +23,13 @@ CONTAINER_OBJECT_PURPOSE = "container" _NEGATIVE_OWNER_SENTINEL = "__litellm_container_no_owner__" _CONTAINER_OWNER_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) +# Caches the stored ``unified_object_id`` (the encoded container ID +# captured at create time) so ``get_container_forwarding_params`` can +# recover the deployment ``model_id`` for native upstream IDs without +# re-hitting Prisma on every retrieve/delete. +_NEGATIVE_STORED_ID_SENTINEL = "__litellm_container_no_stored_id__" +_CONTAINER_STORED_ID_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) + # Per-caller-scope cache for ``GET /v1/containers`` list filtering. Without # this, every list call issues a fresh ``find_many`` against # ``litellm_managedobjecttable``. The cache key is the sorted owner-scope @@ -56,7 +63,7 @@ def decode_container_id_for_ownership( return original_container_id, custom_llm_provider -def get_container_forwarding_params( +async def get_container_forwarding_params( container_id: str, original_container_id: str, custom_llm_provider: str ) -> Dict[str, str]: params = { @@ -65,6 +72,20 @@ def get_container_forwarding_params( } decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) model_id = decoded.get("model_id") + if not (isinstance(model_id, str) and model_id): + # Native upstream IDs (e.g. Azure ``cntr_``) carry no LiteLLM + # routing payload, so decoding the user-supplied id yields no + # ``model_id``. Recover it from the encoded ``unified_object_id`` + # captured on the ownership row at create time — when the router + # selected a specific deployment that ID embeds the model_id. + stored_id = await _get_stored_container_id( + original_container_id, custom_llm_provider + ) + if stored_id and stored_id != container_id: + stored_decoded = ResponsesAPIRequestUtils._decode_container_id(stored_id) + stored_model_id = stored_decoded.get("model_id") + if isinstance(stored_model_id, str) and stored_model_id: + model_id = stored_model_id if isinstance(model_id, str) and model_id: params["model_id"] = model_id return params @@ -168,6 +189,7 @@ async def record_container_owner( ) _CONTAINER_OWNER_CACHE.set_cache(model_object_id, owner) + _CONTAINER_STORED_ID_CACHE.set_cache(model_object_id, container_id) # Drop the caller's own list-cache entry so the just-created container # shows up on their next ``GET /v1/containers``. Other callers with # disjoint scope tuples have their own entries; intersecting-scope @@ -207,9 +229,60 @@ async def _get_container_owner( _CONTAINER_OWNER_CACHE.set_cache( model_object_id, owner if owner is not None else _NEGATIVE_OWNER_SENTINEL ) + stored_id = getattr(row, "unified_object_id", None) if row is not None else None + _CONTAINER_STORED_ID_CACHE.set_cache( + model_object_id, + ( + stored_id + if isinstance(stored_id, str) and stored_id + else _NEGATIVE_STORED_ID_SENTINEL + ), + ) return owner +async def _get_stored_container_id( + original_container_id: str, custom_llm_provider: str +) -> Optional[str]: + """Return the ``unified_object_id`` stored at create time, if any. + + Used by :func:`get_container_forwarding_params` to recover the + deployment ``model_id`` for native upstream container IDs: the stored + value is the encoded form produced by ``encode_container_id_in_response`` + when the router selected a specific deployment. + """ + model_object_id = _container_model_object_id( + original_container_id, custom_llm_provider + ) + + cached = _CONTAINER_STORED_ID_CACHE.get_cache(model_object_id) + if cached == _NEGATIVE_STORED_ID_SENTINEL: + return None + if isinstance(cached, str) and cached: + return cached + + prisma_client = await _get_prisma_client() + if prisma_client is None: + return None + + row = await prisma_client.db.litellm_managedobjecttable.find_first( + where={ + "model_object_id": model_object_id, + "file_purpose": CONTAINER_OBJECT_PURPOSE, + } + ) + stored_id = getattr(row, "unified_object_id", None) if row is not None else None + _CONTAINER_STORED_ID_CACHE.set_cache( + model_object_id, + ( + stored_id + if isinstance(stored_id, str) and stored_id + else _NEGATIVE_STORED_ID_SENTINEL + ), + ) + return stored_id if isinstance(stored_id, str) and stored_id else None + + async def assert_user_can_access_container( container_id: str, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index d84cebcf05a..97525a528d0 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -34,8 +34,7 @@ async def create_missing_views(db: _db): # noqa: PLR0915 if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS): raise # If an error occurs, the view does not exist, so create it - await db.execute_raw( - """ + await db.execute_raw(""" CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -47,8 +46,7 @@ async def create_missing_views(db: _db): # noqa: PLR0915 FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id = p.project_id; - """ - ) + """) verbose_logger.debug("LiteLLM_VerificationTokenView Created!") diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 19ec6699390..e7c5fa3f72c 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -178,15 +178,28 @@ class SpendCounterReseed: if db_spend is None: return None # Warm even when 0 so subsequent reads hit cache, not DB. + # + # Seed via SET NX (cross-pod safe): only one pod initializes the + # Redis key with db_spend; concurrent seeders read the winner's + # value. INCRBYFLOAT-of-db_spend from N pods would multiply the + # counter (N x db_spend) and trigger spurious budget alerts. + current_value: float = float(db_spend) try: if spend_counter_cache.redis_cache is not None: - current_value = ( - await spend_counter_cache.redis_cache.async_increment( - key=counter_key, - value=db_spend, - refresh_ttl=True, - ) + seeded = await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, + value=db_spend, + nx=True, ) + if seeded: + current_value = float(db_spend) + else: + cached = await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key + ) + current_value = ( + float(cached) if cached is not None else float(db_spend) + ) spend_counter_cache.in_memory_cache.set_cache( key=counter_key, value=current_value, @@ -202,7 +215,7 @@ class SpendCounterReseed: ) if require_cache_warm: raise - return db_spend + return current_value @staticmethod async def window_from_spend_logs( diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 7cad1352a79..766ef0cf9f6 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -10,7 +10,6 @@ every text fragment. from typing import Any, Callable, Dict, FrozenSet, Iterator, List - # Call types whose body carries free-form chat / prompt text that # text-content guardrails (banned keywords, content moderation, secret # detection, …) should inspect. The proxy ingress passes ``route_type`` diff --git a/litellm/proxy/guardrails/guardrail_hooks/akto/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/akto/__init__.py index c4aaea709ba..1e3dd906b9f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/akto/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/akto/__init__.py @@ -4,7 +4,6 @@ from litellm.types.guardrails import SupportedGuardrailIntegrations from .akto import AktoGuardrail - if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index bb1db3d62d2..765c419479e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -63,6 +63,7 @@ from litellm.types.utils import ( CallTypesLiteral, Choices, GuardrailStatus, + GuardrailTracingDetail, Message, ModelResponse, ModelResponseStream, @@ -509,6 +510,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Add guardrail information to request trace ######################################################### _json_response = httpx_response.json() + tracing_detail = self._build_tracing_detail(_json_response) + # Raw Bedrock JSON is passed here; match/regex redaction runs once inside # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. self.add_standard_logging_guardrail_information_to_request_data( @@ -522,6 +525,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), event_type=event_type, + tracing_detail=tracing_detail or None, ) ######################################################### if httpx_response.status_code == 200: @@ -640,6 +644,55 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return (status_code, err) return (status_code, message) + def _build_tracing_detail( + self, response: BedrockGuardrailResponse + ) -> GuardrailTracingDetail: + """ + Build the tracing detail from the raw Bedrock response, before + redaction, so downstream loggers (OTEL, Langfuse, ...) get the + actual category names rather than the "[REDACTED]" sentinel that + replaces customWords.match later. Bedrock's top-level ``action`` + field ("GUARDRAIL_INTERVENED" or "NONE") is also surfaced so the + OTEL integration can expose it as a queryable span attribute + without re-parsing the redacted guardrail_response blob. + """ + tracing_detail: GuardrailTracingDetail = {} + violation_categories = self._extract_violation_category_names(response) + if violation_categories: + tracing_detail["violation_categories"] = violation_categories + bedrock_action = response.get("action") + if isinstance(bedrock_action, str): + tracing_detail["guardrail_action"] = bedrock_action + return tracing_detail + + def _extract_violation_category_names( + self, response: BedrockGuardrailResponse + ) -> List[str]: + """ + Flatten the BLOCKED assessments into a list of human-readable category + names suitable for queryable OTEL / standard-logging attributes. + + SECURITY: only emits the non-sensitive policy *label* (topic name, + content-filter type, PII entity type, named-regex name). The raw + ``match`` field is intentionally NOT used — it carries the user's + original input that triggered the rule (e.g. a credit-card number + that hit a regex, or the literal custom word). Surfacing it to + telemetry would re-introduce the sensitive content the guardrail + was supposed to keep out. Entries that only have a ``match`` (bare + customWords, unnamed regexes) are therefore skipped — operators + can still see the count in ``_extract_blocked_assessments`` which + feeds the HTTP error detail. + """ + names: List[str] = [] + for block in self._extract_blocked_assessments(response): + for match in block.get("matches", []) or []: + # Allow-list non-sensitive labels only. Never fall back to + # `match.get("match")` — that's user-submitted content. + label = match.get("name") or match.get("type") + if isinstance(label, str) and label: + names.append(label) + return names + def _extract_blocked_assessments( self, response: BedrockGuardrailResponse ) -> List[dict]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 5502076829f..0f299f4c5f7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -92,6 +92,8 @@ from litellm.types.utils import CallTypesLiteral # Module-level singleton for the JWKS discovery endpoint to access. _mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None +_MCP_JWT_CALL_TYPES = frozenset({"call_mcp_tool", "list_mcp_tools"}) + # Simple in-memory JWKS cache: keyed by JWKS URI → (keys_list, fetched_at). _jwks_cache: Dict[str, tuple] = {} _JWKS_CACHE_TTL = 3600 # 1 hour @@ -603,17 +605,23 @@ class MCPJWTSigner(CustomGuardrail): # FR-10: Scope building # ------------------------------------------------------------------ - def _build_scope(self, raw_tool_name: str) -> str: + def _build_scope( + self, + raw_tool_name: str, + call_type: Optional[CallTypesLiteral] = None, + ) -> str: """ Build the JWT scope string. When allowed_scopes is configured: join them verbatim. Otherwise auto-generate minimal, least-privilege scopes: - Tool call → mcp:tools/call mcp:tools/:call - - No tool → mcp:tools/call mcp:tools/list + - No tool → mcp:tools/list NOTE: tools/list is intentionally NOT granted on tool-call JWTs to prevent callers from enumerating tools they didn't ask to use. + Conversely, tools/call is NOT granted on tools/list-only JWTs so an + intercepted list token cannot be replayed to invoke tools. """ if self.allowed_scopes is not None: return " ".join(self.allowed_scopes) @@ -623,8 +631,14 @@ class MCPJWTSigner(CustomGuardrail): ) if tool_name: scopes = ["mcp:tools/call", f"mcp:tools/{tool_name}:call"] + elif call_type == "call_mcp_tool": + # Tool-call request reached the signer without a tool name (e.g. + # missing mcp_tool_name in hook data). Fall back to a generic + # tools/call scope so the upstream server still accepts the + # invocation rather than rejecting it as a tools/list-only token. + scopes = ["mcp:tools/call"] else: - scopes = ["mcp:tools/call", "mcp:tools/list"] + scopes = ["mcp:tools/list"] return " ".join(scopes) # ------------------------------------------------------------------ @@ -673,6 +687,7 @@ class MCPJWTSigner(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, data: dict, jwt_claims: Optional[Dict[str, Any]] = None, + call_type: Optional[CallTypesLiteral] = None, ) -> Dict[str, Any]: """ Build JWT claims for the outbound MCP access token. @@ -713,7 +728,7 @@ class MCPJWTSigner(CustomGuardrail): # scope (FR-10) raw_tool_name: str = data.get("mcp_tool_name", "") - claims["scope"] = self._build_scope(raw_tool_name) + claims["scope"] = self._build_scope(raw_tool_name, call_type=call_type) # optional_claims passthrough (FR-15) claims = self._passthrough_optional_claims(claims, jwt_claims) @@ -779,16 +794,20 @@ class MCPJWTSigner(CustomGuardrail): Verifies the incoming token (when configured), validates required claims, then signs an outbound JWT and injects it as the Authorization header. - All non-MCP call types pass through unchanged. + Signs outbound MCP tool calls and tools/list requests. """ - if call_type != "call_mcp_tool": + if call_type not in _MCP_JWT_CALL_TYPES: return data + hook_data = dict(data) + if call_type == "list_mcp_tools": + hook_data["mcp_tool_name"] = "" + # ------------------------------------------------------------------ # FR-5: Verify incoming token before re-signing # ------------------------------------------------------------------ jwt_claims: Optional[Dict[str, Any]] = None - raw_token: Optional[str] = data.get("incoming_bearer_token") + raw_token: Optional[str] = hook_data.get("incoming_bearer_token") if self.access_token_discovery_uri and raw_token: # Three-dot pattern → JWT; otherwise opaque. @@ -837,7 +856,9 @@ class MCPJWTSigner(CustomGuardrail): # ------------------------------------------------------------------ # Build outbound access token # ------------------------------------------------------------------ - claims = self._build_claims(user_api_key_dict, data, jwt_claims) + claims = self._build_claims( + user_api_key_dict, hook_data, jwt_claims, call_type=call_type + ) signed_token = jwt.encode( claims, @@ -848,7 +869,7 @@ class MCPJWTSigner(CustomGuardrail): # Merge into existing extra_headers — a prior guardrail in the chain may # have already injected tracing headers or correlation IDs. - existing_headers: Dict[str, str] = data.get("extra_headers") or {} + existing_headers: Dict[str, str] = hook_data.get("extra_headers") or {} new_headers: Dict[str, str] = { **existing_headers, "Authorization": f"Bearer {signed_token}", @@ -875,17 +896,74 @@ class MCPJWTSigner(CustomGuardrail): claims, self._kid ) - data["extra_headers"] = new_headers + hook_data["extra_headers"] = new_headers verbose_proxy_logger.debug( "MCPJWTSigner: signed JWT sub=%s act=%s tool=%s exp=%d " - "verified=%s channel=%s", + "verified=%s channel=%s call_type=%s", claims.get("sub"), claims.get("act", {}).get("sub"), - data.get("mcp_tool_name"), + hook_data.get("mcp_tool_name"), claims["exp"], jwt_claims is not None, bool(self.channel_token_audience), + call_type, ) - return data + return hook_data + + +async def inject_mcp_jwt_headers_for_upstream( + user_api_key_dict: Optional[UserAPIKeyAuth], + extra_headers: Optional[Dict[str, str]] = None, + raw_headers: Optional[Dict[str, str]] = None, + *, + for_list_tools: bool = False, + mcp_tool_name: str = "", +) -> Dict[str, str]: + """ + Sign outbound MCP headers when MCPJWTSigner is configured. + + Used by tools/list paths that do not go through proxy pre_call_hook. + """ + merged = dict(extra_headers or {}) + signer = get_mcp_jwt_signer() + if signer is None or user_api_key_dict is None: + return merged + + normalized_raw = {k.lower(): v for k, v in (raw_headers or {}).items()} + incoming_bearer_token: Optional[str] = None + auth_hdr = normalized_raw.get("authorization", "") + if auth_hdr.lower().startswith("bearer "): + incoming_bearer_token = auth_hdr[len("bearer ") :] + + hook_data: Dict[str, Any] = { + "mcp_tool_name": "" if for_list_tools else mcp_tool_name, + "incoming_bearer_token": incoming_bearer_token, + "extra_headers": merged, + } + call_type: CallTypesLiteral = ( + "list_mcp_tools" if for_list_tools else "call_mcp_tool" + ) + try: + from litellm.proxy.proxy_server import ( # noqa: PLC0415 + proxy_logging_obj as _proxy_logging, + ) + + shared_cache = ( + _proxy_logging.internal_usage_cache.dual_cache + if _proxy_logging is not None + else DualCache() + ) + except Exception: + shared_cache = DualCache() + + result = await signer.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=shared_cache, + data=hook_data, + call_type=call_type, + ) + if isinstance(result, dict) and result.get("extra_headers"): + merged.update(result["extra_headers"]) + return merged diff --git a/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py new file mode 100644 index 00000000000..ab347130a30 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/rubrik/__init__.py @@ -0,0 +1,35 @@ +"""Rubrik guardrail integration for LiteLLM.""" + +from typing import TYPE_CHECKING + +from litellm.integrations.rubrik import RubrikLogger +from litellm.types.guardrails import SupportedGuardrailIntegrations + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail( + litellm_params: "LitellmParams", guardrail: "Guardrail" +) -> RubrikLogger: + import litellm + + rubrik_callback = RubrikLogger( + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + + litellm.logging_callback_manager.add_litellm_callback(rubrik_callback) + return rubrik_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.RUBRIK.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.RUBRIK.value: RubrikLogger, +} diff --git a/litellm/proxy/hooks/litellm_skills/__init__.py b/litellm/proxy/hooks/litellm_skills/__init__.py index 057cf3d8b38..1507b652ab4 100644 --- a/litellm/proxy/hooks/litellm_skills/__init__.py +++ b/litellm/proxy/hooks/litellm_skills/__init__.py @@ -6,7 +6,7 @@ The actual skill logic is in litellm/llms/litellm_proxy/skills/. Usage: from litellm.proxy.hooks.litellm_skills import SkillsInjectionHook - + # Register hook in proxy litellm.callbacks.append(SkillsInjectionHook()) """ diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 60dc7827a6f..2eda1b30c5d 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -1,9 +1,9 @@ """ BUDGET MANAGEMENT -All /budget management endpoints +All /budget management endpoints -/budget/new +/budget/new /budget/info /budget/update /budget/delete diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 4889f0b7f80..1fd8320db20 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -1,9 +1,9 @@ """ CUSTOMER MANAGEMENT -All /customer management endpoints +All /customer management endpoints -/customer/new +/customer/new /customer/info /customer/update /customer/delete diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 472306eb818..f2d8ec8fb55 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -546,7 +546,7 @@ async def _update_existing_team_model_assignment( """ def _get_team_public_model_name( - model_info: Optional[Union[dict, str]] + model_info: Optional[Union[dict, str]], ) -> Optional[str]: if isinstance(model_info, dict): value = model_info.get("team_public_model_name") diff --git a/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py b/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py index 191212d6f0b..04e44c623d1 100644 --- a/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py +++ b/litellm/proxy/management_endpoints/sso/custom_microsoft_sso.py @@ -7,7 +7,7 @@ variables. Environment Variables: - MICROSOFT_AUTHORIZATION_ENDPOINT: Custom authorization endpoint URL -- MICROSOFT_TOKEN_ENDPOINT: Custom token endpoint URL +- MICROSOFT_TOKEN_ENDPOINT: Custom token endpoint URL - MICROSOFT_USERINFO_ENDPOINT: Custom userinfo endpoint URL If these are not set, the default Microsoft endpoints are used. diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 86c4d6dcd9a..0b2f93d817a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4381,9 +4381,7 @@ async def list_team( except Exception as e: team_exception = """Invalid team object for team_id: {}. team_object={}. Error: {} - """.format( - team.team_id, team.model_dump(), str(e) - ) + """.format(team.team_id, team.model_dump(), str(e)) verbose_proxy_logger.exception(team_exception) continue # Sort the responses by team_alias diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 6e2e2bedac1..d6082899c02 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -43,6 +43,8 @@ from litellm.caching.dual_cache import DualCache from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( + CLI_SSO_CLAIM_MAP, + CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, CLI_SSO_SESSION_CACHE_KEY_PREFIX, CLI_SSO_SESSION_TTL_SECONDS, LITELLM_CLI_SOURCE_IDENTIFIER, @@ -140,6 +142,20 @@ _CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS = 60 _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30 _CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" _CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$") +_CLI_SSO_SCALAR_TYPES = (str, int, float, bool) +_CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$") +_CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset( + { + "access_token", + "api_key", + "client_secret", + "id_token", + "password", + "private_key", + "refresh_token", + "secret", + } +) def _hash_cli_sso_secret(secret: str) -> str: @@ -225,6 +241,239 @@ def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool: return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash) +def _parse_cli_sso_claim_map() -> List[Tuple[str, str]]: + """ + Parse CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP. + + Format: comma-separated ``source_claim->metadata_key`` pairs, e.g. + ``employment_type->acme_employment_type,org_info.department->department``. + Destination keys may use an optional ``metadata.`` prefix; values are stored + on the LiteLLM user's ``metadata`` JSON column. + """ + claim_map_raw = CLI_SSO_CLAIM_MAP.strip() + if not claim_map_raw: + return [] + + parsed: List[Tuple[str, str]] = [] + for entry in claim_map_raw.split(","): + entry = entry.strip() + if not entry or "->" not in entry: + continue + source_claim, dest_key = entry.split("->", 1) + source_claim = source_claim.strip() + dest_key = dest_key.strip() + if dest_key.startswith("metadata."): + dest_key = dest_key[len("metadata.") :] + if source_claim and dest_key: + parsed.append((source_claim, dest_key)) + return parsed + + +def _is_safe_cli_sso_metadata_dest_key(dest_key: str) -> bool: + if not dest_key or not _CLI_SSO_DEST_KEY_RE.fullmatch(dest_key): + return False + lowered = dest_key.lower() + return not any(fragment in lowered for fragment in _CLI_SSO_SECRET_KEY_FRAGMENTS) + + +def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool: + if not isinstance(value, _CLI_SSO_SCALAR_TYPES): + return False + if isinstance(value, str): + if len(value) > CLI_SSO_CLAIM_MAX_SCALAR_LENGTH: + return False + if value.startswith("eyJ") and value.count(".") >= 2: + return False + return True + + +def _sso_result_to_dict(result: Union[CustomOpenID, OpenID, dict]) -> Dict[str, Any]: + if isinstance(result, dict): + return result + if hasattr(result, "model_dump"): + dumped = result.model_dump() + if isinstance(dumped, dict): + return cast(Dict[str, Any], dumped) + return {} + + +def _get_nested_claim_value(data: Dict[str, Any], claim_path: str) -> Any: + """Resolve a dot-notation claim path against an SSO result dict. + + Unlike ``get_nested_value``, this does not strip a leading ``metadata.`` + prefix, since OIDC claims may legitimately use ``metadata`` as a top-level + key. + """ + if not claim_path: + return None + if claim_path in data: + return data[claim_path] + placeholder = "\x00" + parts = claim_path.replace("\\.", placeholder).split(".") + parts = [p.replace(placeholder, ".") for p in parts] + current: Any = data + for part in parts: + if isinstance(current, dict) and part in current: + current = current[part] + else: + return None + return current + + +def _extract_sso_claim_value( + result: Union[CustomOpenID, OpenID, dict], claim_path: str +) -> Any: + extra_fields = getattr(result, "extra_fields", None) + if isinstance(extra_fields, dict): + if claim_path in extra_fields: + return extra_fields[claim_path] + nested = _get_nested_claim_value(extra_fields, claim_path) + if nested is not None: + return nested + + if isinstance(result, dict): + return _get_nested_claim_value(result, claim_path) + + result_dict = _sso_result_to_dict(result) + return _get_nested_claim_value(result_dict, claim_path) + + +def _set_nested_metadata_value( + metadata: Dict[str, Any], key_path: str, value: Any +) -> None: + placeholder = "\x00" + parts = key_path.replace("\\.", placeholder).split(".") + parts = [p.replace(placeholder, ".") for p in parts] + current: Any = metadata + for part in parts[:-1]: + existing = current.get(part) + if not isinstance(existing, dict): + existing = {} + current[part] = existing + current = existing + current[parts[-1]] = value + + +def _flatten_cli_sso_metadata_for_poll( + metadata: Dict[str, Any], +) -> Dict[str, Union[str, int, float, bool]]: + """Expose scalar attribution metadata as a flat dict for CLI poll responses.""" + flattened: Dict[str, Union[str, int, float, bool]] = {} + stack: List[Tuple[str, Any]] = [("", metadata)] + while stack: + prefix, value = stack.pop() + if isinstance(value, dict): + for key, nested in value.items(): + nested_prefix = f"{prefix}.{key}" if prefix else key + stack.append((nested_prefix, nested)) + elif _is_safe_cli_sso_scalar_claim_value(value): + flattened[prefix] = value + return flattened + + +def build_cli_sso_attribution_metadata( + result: Union[CustomOpenID, OpenID, dict], +) -> Dict[str, Any]: + """ + Build allowlisted, non-secret scalar attribution metadata from an SSO result. + + Sources are configured via CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP and + may include claims captured by GENERIC_USER_EXTRA_ATTRIBUTES on CustomOpenID. + """ + claim_map = _parse_cli_sso_claim_map() + if not claim_map: + return {} + + metadata: Dict[str, Any] = {} + for source_claim, dest_key in claim_map: + if not _is_safe_cli_sso_metadata_dest_key(dest_key): + verbose_proxy_logger.debug( + f"Skipping unsafe CLI SSO metadata destination key: {dest_key}" + ) + continue + + raw_value = _extract_sso_claim_value(result=result, claim_path=source_claim) + if not _is_safe_cli_sso_scalar_claim_value(raw_value): + continue + + _set_nested_metadata_value( + metadata=metadata, key_path=dest_key, value=raw_value + ) + + return metadata + + +def _merge_cli_sso_attribution_metadata( + existing_metadata: Dict[str, Any], attribution_metadata: Dict[str, Any] +) -> Dict[str, Any]: + """Merge attribution metadata into existing user metadata in-place. + + Preserves original value types (in particular, string claim values that + happen to look numeric are NOT coerced to ``int``/``float``). Nested dicts + are merged iteratively so attribution claims do not clobber unrelated keys + under the same parent. + """ + pending: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [ + (existing_metadata, attribution_metadata) + ] + while pending: + target, source = pending.pop() + for key, value in source.items(): + if value is None: + continue + existing_value = target.get(key) + if isinstance(value, dict) and isinstance(existing_value, dict): + pending.append((existing_value, value)) + else: + target[key] = value + return existing_metadata + + +async def _persist_cli_sso_user_metadata( + prisma_client: PrismaClient, + user_id: str, + attribution_metadata: Dict[str, Any], +) -> None: + if not attribution_metadata: + return + + try: + user_row = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_id} + ) + existing_metadata: Dict[str, Any] = {} + if user_row is not None: + row_metadata = user_row.metadata + if isinstance(row_metadata, dict): + existing_metadata = deepcopy(row_metadata) + + merged_metadata = _merge_cli_sso_attribution_metadata( + existing_metadata=existing_metadata, + attribution_metadata=attribution_metadata, + ) + await prisma_client.db.litellm_usertable.update_many( + where={"user_id": user_id}, + data={"metadata": merged_metadata}, + ) + verbose_proxy_logger.info( + f"Persisted CLI SSO attribution metadata for user {user_id}: " + f"{list(_flatten_cli_sso_metadata_for_poll(attribution_metadata).keys())}" + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}" + ) + + +def _cli_poll_attribution_metadata_from_session( + session_data: Dict[str, Any], +) -> Dict[str, Union[str, int, float, bool]]: + stored = session_data.get("attribution_metadata") + if isinstance(stored, dict): + return _flatten_cli_sso_metadata_for_poll(stored) + return {} + + def _render_cli_sso_verification_page( verify_url: str, browser_complete_token: str ) -> str: @@ -1674,7 +1923,12 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: key_id = state_parts[1] if len(state_parts) > 1 else None verbose_proxy_logger.info("CLI SSO callback detected") - return await cli_sso_callback(request=request, key=key_id, result=result) + return await cli_sso_callback( + request=request, + key=key_id, + result=result, + received_response=received_response, + ) # Control-plane cross-origin: read return_to from cookie. # Starlette's cookie_parser already handles RFC 2109 unquoting. @@ -1692,15 +1946,144 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: ) +async def _build_cli_sso_user_defined_values( + result: Union[OpenID, dict], + parsed_openid_result: ParsedOpenIDResult, +) -> Optional[SSOUserDefinedValues]: + from litellm.proxy.proxy_server import user_custom_sso + + user_id = parsed_openid_result.get("user_id") + if user_custom_sso is not None: + if inspect.iscoroutinefunction(user_custom_sso): + return await user_custom_sso(result) # type: ignore + raise ValueError("user_custom_sso must be a coroutine function") + if user_id is None: + return None + return SSOUserDefinedValues( + models=[], + user_id=user_id, + user_email=parsed_openid_result.get("user_email"), + max_budget=litellm.max_internal_user_budget, + user_role=parsed_openid_result.get("user_role"), + budget_duration=litellm.internal_user_budget_duration, + ) + + +async def _fetch_cli_sso_team_details( + prisma_client: PrismaClient, + teams: List[str], +) -> List[Dict[str, Any]]: + team_details: List[Dict[str, Any]] = [] + try: + if teams: + prisma_teams = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": teams}} + ) + for team_row in prisma_teams: + team_dict = team_row.model_dump() + team_details.append( + { + "team_id": team_dict.get("team_id"), + "team_alias": team_dict.get("team_alias"), + } + ) + except Exception as e: + verbose_proxy_logger.error( + f"Error fetching team details for CLI SSO session: {e}" + ) + return team_details + + +async def _complete_cli_sso_callback_session( + *, + request: Request, + key: str, + flow: dict, + result: Union[OpenID, dict], + parsed_openid_result: ParsedOpenIDResult, + user_defined_values: Optional[SSOUserDefinedValues], + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +): + from fastapi.responses import HTMLResponse + + user_id = parsed_openid_result.get("user_id") + user_email = parsed_openid_result.get("user_email") + user_info = await get_user_info_from_db( + result=result, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + user_email=user_email, + user_defined_values=user_defined_values, + alternate_user_id=user_id, + ) + if user_info is None: + raise HTTPException( + status_code=500, detail="Failed to retrieve user information from SSO" + ) + if not user_info.user_id: + raise HTTPException( + status_code=500, detail="Failed to retrieve user information from SSO" + ) + + teams: List[str] = [] + if hasattr(user_info, "teams") and user_info.teams: + teams = user_info.teams if isinstance(user_info.teams, list) else [] + + team_details = await _fetch_cli_sso_team_details( + prisma_client=prisma_client, teams=teams + ) + attribution_metadata = build_cli_sso_attribution_metadata(result=result) + if attribution_metadata: + await _persist_cli_sso_user_metadata( + prisma_client=prisma_client, + user_id=cast(str, user_info.user_id), + attribution_metadata=attribution_metadata, + ) + + flow["session_data"] = { + "user_id": cast(str, user_info.user_id), + "user_role": user_info.user_role, + "models": user_info.models if hasattr(user_info, "models") else [], + "user_email": user_email, + "teams": teams, + "team_details": team_details, + "attribution_metadata": attribution_metadata, + } + flow["sso_complete"] = True + browser_complete_token = secrets.token_urlsafe(32) + flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token) + _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow) + + verbose_proxy_logger.info( + f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}" + ) + verify_url = get_custom_url( + request_base_url=str(request.base_url), + route=f"sso/cli/complete/{key}", + ) + return HTMLResponse( + content=_render_cli_sso_verification_page( + verify_url=verify_url, + browser_complete_token=browser_complete_token, + ), + status_code=200, + ) + + async def cli_sso_callback( request: Request, key: Optional[str] = None, result: Optional[Union[OpenID, dict]] = None, + received_response: Optional[dict] = None, ): """CLI SSO callback - stores session info for JWT generation on polling""" verbose_proxy_logger.info("CLI SSO callback") from litellm.proxy.proxy_server import ( + general_settings, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -1722,89 +2105,40 @@ async def cli_sso_callback( # After None check, cast to non-None type for type checker result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result) - parsed_openid_result = SSOAuthenticationHandler._get_user_email_and_id_from_result( - result=result_non_none - ) - verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}") - try: - # Get full user info from DB - user_info = await get_user_info_from_db( + parsed_openid_result = ( + SSOAuthenticationHandler._get_user_email_and_id_from_result( + result=result_non_none, + generic_client_id=os.getenv("GENERIC_CLIENT_ID", None), + ) + ) + verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}") + user_defined_values = await _build_cli_sso_user_defined_values( result=result_non_none, + parsed_openid_result=parsed_openid_result, + ) + + SSOAuthenticationHandler.verify_user_in_restricted_sso_group( + general_settings=general_settings, + result=result_non_none, + received_response=received_response, + ) + + return await _complete_cli_sso_callback_session( + request=request, + key=cast(str, key), + flow=flow, + result=result_non_none, + parsed_openid_result=parsed_openid_result, + user_defined_values=user_defined_values, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, - user_email=parsed_openid_result.get("user_email"), - user_defined_values=None, - alternate_user_id=parsed_openid_result.get("user_id"), ) - - if user_info is None: - raise HTTPException( - status_code=500, detail="Failed to retrieve user information from SSO" - ) - - # Get all teams from user_info - CLI will let user select which one - teams: List[str] = [] - if hasattr(user_info, "teams") and user_info.teams: - teams = user_info.teams if isinstance(user_info.teams, list) else [] - - # Also fetch team aliases for a better CLI UX. We keep the original - # "teams" list of IDs for backwards compatibility and add an - # optional "team_details" field containing objects with both - # team_id and team_alias. - team_details: List[Dict[str, Any]] = [] - try: - if teams: - prisma_teams = await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": teams}} - ) - for team_row in prisma_teams: - team_dict = team_row.model_dump() - team_details.append( - { - "team_id": team_dict.get("team_id"), - "team_alias": team_dict.get("team_alias"), - } - ) - except Exception as e: - # If anything goes wrong here, fall back gracefully without - # impacting the SSO flow. - verbose_proxy_logger.error( - f"Error fetching team details for CLI SSO session: {e}" - ) - - session_data = { - "user_id": user_info.user_id, - "user_role": user_info.user_role, - "models": user_info.models if hasattr(user_info, "models") else [], - "user_email": parsed_openid_result.get("user_email"), - "teams": teams, - # Optional rich metadata for clients that want nicer display - "team_details": team_details, - } - - flow["session_data"] = session_data - flow["sso_complete"] = True - browser_complete_token = secrets.token_urlsafe(32) - flow["browser_complete_token_hash"] = _hash_cli_sso_secret( - browser_complete_token - ) - _set_cli_sso_flow(login_id=cast(str, key), cache=user_api_key_cache, flow=flow) - - verbose_proxy_logger.info( - f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}" - ) - - from fastapi.responses import HTMLResponse - - verify_url = str(request.url_for("cli_sso_complete", login_id=key)) - html_content = _render_cli_sso_verification_page( - verify_url=verify_url, - browser_complete_token=browser_complete_token, - ) - return HTMLResponse(content=html_content, status_code=200) - + except ProxyException: + raise + except HTTPException: + raise except Exception as e: verbose_proxy_logger.error(f"Error with CLI SSO callback: {e}") raise HTTPException( @@ -1870,13 +2204,19 @@ async def cli_poll_key( team_details_response = [ {"team_id": t, "team_alias": None} for t in user_teams ] - return { + poll_response: Dict[str, Any] = { "status": "ready", "user_id": user_id, "teams": user_teams, "team_details": team_details_response, "requires_team_selection": True, } + attribution_metadata = _cli_poll_attribution_metadata_from_session( + session_data + ) + if attribution_metadata: + poll_response["attribution_metadata"] = attribution_metadata + return poll_response # Validate team_id if provided if team_id is not None: @@ -1889,6 +2229,17 @@ async def cli_poll_key( # If no team_id provided and user has 0 or 1 team, use first team (or None) team_id = user_teams[0] if len(user_teams) > 0 else None + team_alias = None + if team_id and isinstance(user_team_details, list): + team_alias = next( + ( + team.get("team_alias") + for team in user_team_details + if team.get("team_id") == team_id + ), + None, + ) + # Create user object for JWT generation user_info = LiteLLM_UserTable( user_id=user_id, @@ -1900,7 +2251,7 @@ async def cli_poll_key( # Generate CLI JWT on-demand (expiration configurable via LITELLM_CLI_JWT_EXPIRATION_HOURS) # Pass selected team_id to ensure JWT has correct team jwt_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info=user_info, team_id=team_id + user_info=user_info, team_id=team_id, team_alias=team_alias ) # Delete cache entry (single-use) @@ -1909,7 +2260,7 @@ async def cli_poll_key( verbose_proxy_logger.info( f"CLI JWT generated for user: {user_id}, team: {team_id}" ) - return { + poll_response = { "status": "ready", "key": jwt_token, "user_id": user_id, @@ -1919,6 +2270,12 @@ async def cli_poll_key( # present nicer information if needed. "team_details": user_team_details, } + attribution_metadata = _cli_poll_attribution_metadata_from_session( + session_data + ) + if attribution_metadata: + poll_response["attribution_metadata"] = attribution_metadata + return poll_response else: return {"status": "pending"} diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index 872b6fa2250..ebd276fbee5 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -3,7 +3,7 @@ User Agent Analytics Endpoints This module provides optimized endpoints for tracking user agent activity metrics including: - Daily Active Users (DAU) by tags for configurable number of days -- Weekly Active Users (WAU) by tags for configurable number of weeks +- Weekly Active Users (WAU) by tags for configurable number of weeks - Monthly Active Users (MAU) by tags for configurable number of months - Summary analytics by tags diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index 4f31c762df1..e32fee6afc5 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -178,6 +178,24 @@ async def _parse_ocr_request(request: Request) -> Dict[str, Any]: "For JSON requests, use 'document_url' or 'image_url' document types." ) + # Security: reject provider-native file IDs (e.g. reducto://) received via + # JSON. These IDs are not scoped to the LiteLLM proxy user/key, so an + # authenticated user who obtains another user's file ID could submit it + # here and receive the OCR result using the proxy's shared provider + # credentials. Force callers to upload fresh content per request via + # multipart/form-data or an inline base64 data URI, both of which produce + # a server-mediated upload bound to the current request. + if isinstance(doc, dict): + for url_field in ("document_url", "image_url"): + url_value = doc.get(url_field) + if isinstance(url_value, str) and url_value.startswith("reducto://"): + raise ValueError( + "reducto:// file IDs are not accepted through the proxy " + "OCR API; upload the file in the same request via " + "multipart/form-data with a 'file' field, or pass an " + "inline base64 data URI as the document URL." + ) + return data diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 30c78ed5ba7..0415bb456ec 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -727,6 +727,76 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None: pass +async def ensure_batch_response_managed_file_ids( + response, + managed_files_obj, + prisma_client, + verbose_proxy_logger, + user_api_key_dict=None, + db_batch_object=None, +) -> None: + """Normalize batch file IDs to managed unified IDs before DB persistence.""" + await resolve_input_file_id_to_unified(response, prisma_client) + await resolve_output_file_ids_to_unified(response, prisma_client) + + if managed_files_obj is None: + return + + hidden_params = getattr(response, "_hidden_params", None) or {} + model_id = hidden_params.get("model_id") + if not model_id: + return + + model_name = hidden_params.get("model_name") + unified_file_id = hidden_params.get("unified_file_id") + if not model_name and isinstance(unified_file_id, str): + decoded_unified_file_id = ( + _is_base64_encoded_unified_file_id(unified_file_id) or unified_file_id + ) + target_model_names = get_models_from_unified_file_id(decoded_unified_file_id) + if target_model_names: + model_name = ",".join(target_model_names) + + if user_api_key_dict is None and db_batch_object is not None: + from litellm.proxy._types import UserAPIKeyAuth + + user_api_key_dict = UserAPIKeyAuth( + user_id=getattr(db_batch_object, "created_by", None) or "default-user-id", + team_id=getattr(db_batch_object, "team_id", None), + ) + if user_api_key_dict is None: + return + + for file_attr in ("output_file_id", "error_file_id"): + raw_file_id = getattr(response, file_attr, None) + if not raw_file_id or _is_base64_encoded_unified_file_id(raw_file_id): + continue + try: + new_unified_file_id = managed_files_obj.get_unified_output_file_id( + output_file_id=raw_file_id, + model_id=model_id, + model_name=model_name, + ) + await managed_files_obj.store_unified_file_id( + file_id=new_unified_file_id, + file_object=None, + litellm_parent_otel_span=getattr( + user_api_key_dict, "parent_otel_span", None + ), + model_mappings={model_id: raw_file_id}, + user_api_key_dict=user_api_key_dict, + ) + setattr(response, file_attr, new_unified_file_id) + verbose_proxy_logger.debug( + f"Converted batch {file_attr} {raw_file_id!r} to managed ID before DB write" + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to convert batch {file_attr}={raw_file_id!r} to managed ID " + f"before DB write: {e}" + ) + + async def get_batch_from_database( batch_id: str, unified_batch_id: Union[str, Literal[False]], @@ -800,6 +870,7 @@ async def update_batch_in_database( verbose_proxy_logger, db_batch_object=None, operation: str = "update", + user_api_key_dict=None, ): """ Update batch status and object in ManagedObjectTable. @@ -813,6 +884,7 @@ async def update_batch_in_database( verbose_proxy_logger: Logger instance db_batch_object: Optional existing database object (for comparison) operation: Description of operation ("update", "cancel", etc.) + user_api_key_dict: Optional auth context for creating managed file IDs """ import litellm.utils @@ -823,6 +895,18 @@ async def update_batch_in_database( if not prisma_client: return + # Always normalize the response's file IDs to unified managed IDs + # (mutates in place) so the caller returns unified IDs to the user + # even when we skip the DB update below for an unchanged status. + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=managed_files_obj, + prisma_client=prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=db_batch_object, + ) + # Only update if status has changed (when db_batch_object is provided) if db_batch_object and response.status == db_batch_object.status: return diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py index a104f962630..e7696e5a18a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cursor_passthrough_logging_handler.py @@ -18,7 +18,6 @@ from litellm.litellm_core_utils.litellm_logging import ( from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.types.utils import StandardPassThroughResponseObject - CURSOR_AGENT_ENDPOINTS: Dict[str, str] = { "POST /v0/agents": "cursor:agent:create", "GET /v0/agents": "cursor:agent:list", diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 5fc8c44b2d8..c0246f234a8 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -6,6 +6,7 @@ import random import subprocess import sys import urllib.parse as urlparse +from pathlib import Path from typing import TYPE_CHECKING, Any, Optional, Union import click @@ -38,6 +39,35 @@ class LiteLLMDatabaseConnectionPool(Enum): database_connection_pool_timeout = 60 +def _build_db_connection_url_params( + connection_limit: int, + pool_timeout: Optional[Union[int, float]], + connect_timeout: Optional[Union[int, float]] = None, + socket_timeout: Optional[Union[int, float]] = None, + extra_params: Optional[dict] = None, +) -> dict: + """Build the Prisma DATABASE_URL query params controlling connection pool behavior. + + `connect_timeout` / `socket_timeout` map to the Prisma URL params of the same + name (https://www.prisma.io/docs/orm/overview/databases/postgresql) and are + omitted when None so Prisma's defaults apply. `extra_params` is an + untyped passthrough — keys it provides win over the named arguments above, + so it can be used to override any default we set here. + """ + params: dict = { + "connection_limit": connection_limit, + } + if pool_timeout is not None: + params["pool_timeout"] = pool_timeout + if connect_timeout is not None: + params["connect_timeout"] = connect_timeout + if socket_timeout is not None: + params["socket_timeout"] = socket_timeout + if extra_params: + params.update(extra_params) + return params + + def append_query_params(url: Optional[str], params: dict) -> str: from litellm._logging import verbose_proxy_logger @@ -264,6 +294,62 @@ class ProxyInitializationHelpers: # hypercorn serve raises a type warning when passing a fast api app - even though fast API is a valid type asyncio.run(serve(app, config)) # type: ignore + @staticmethod + def _init_granian_server( + host: str, + port: int, + num_workers: int, + ssl_certfile_path: Optional[str], + ssl_keyfile_path: Optional[str], + max_requests_before_restart: Optional[int], + ciphers: Optional[str], + granian_runtime_threads: Optional[int] = None, + ) -> None: + """ + Run the proxy with Granian (Rust-backed ASGI server, HTTP/1 + HTTP/2). + + Uses a string import path so workers load ``litellm.proxy.proxy_server:app`` + the same way as uvicorn's ``app=`` string target. + """ + from granian import Granian + from granian.constants import Interfaces + + print( # noqa + f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} using Granian\033[0m\n" + ) + if max_requests_before_restart is not None: + print( # noqa + "\033[1;33mLiteLLM: --max_requests_before_restart is not supported by Granian " + "(Granian uses workers_lifetime in seconds, not a per-request limit).\033[0m\n" + ) + if ciphers is not None: + print( # noqa + "\033[1;33mLiteLLM: --ciphers is not applied when using --run_granian.\033[0m\n" + ) + + kwargs: dict[str, Any] = { + "target": "litellm.proxy.proxy_server:app", + "address": host, + "port": port, + "workers": max(1, num_workers), + "interface": Interfaces.ASGI, + "websockets": True, + } + if granian_runtime_threads is not None: + kwargs["runtime_threads"] = granian_runtime_threads + if ssl_certfile_path is not None and ssl_keyfile_path is not None: + print( # noqa + f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" + ) + kwargs["ssl_cert"] = Path(ssl_certfile_path) + kwargs["ssl_key"] = Path(ssl_keyfile_path) + elif ssl_certfile_path is not None or ssl_keyfile_path is not None: + raise click.ClickException( + "Both --ssl_certfile_path and --ssl_keyfile_path are required for SSL." + ) + + Granian(**kwargs).serve() + @staticmethod def _run_gunicorn_server( host: str, @@ -292,9 +378,7 @@ class ProxyInitializationHelpers: _endpoint_str = ( f"curl --location 'http://0.0.0.0:{port}/chat/completions' \\" ) - curl_command = ( - _endpoint_str - + """ + curl_command = _endpoint_str + """ --header 'Content-Type: application/json' \\ --data ' { "model": "gpt-3.5-turbo", @@ -307,7 +391,6 @@ class ProxyInitializationHelpers: }' \n """ - ) print() # noqa print( # noqa '\033[1;34mLiteLLM: Test your local proxy with: "litellm --test" This runs an openai.ChatCompletion request to your proxy [In a new terminal tab]\033[0m\n' @@ -383,11 +466,9 @@ class ProxyInitializationHelpers: with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - print( # noqa - f""" + print(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` - """ - ) # noqa + """) # noqa # noqa @staticmethod def _is_port_in_use(port): @@ -459,9 +540,23 @@ class ProxyInitializationHelpers: @click.option( "--num_workers", default=DEFAULT_NUM_WORKERS_LITELLM_PROXY, - help="Number of uvicorn / gunicorn workers to spin up. Default is 1 (from DEFAULT_NUM_WORKERS_LITELLM_PROXY)", + help=( + "Number of worker processes for uvicorn / gunicorn, or Granian worker processes " + "(--workers). Default is 1 (from DEFAULT_NUM_WORKERS_LITELLM_PROXY). " + "With --run_granian, use --granian_threads for runtime threads per worker." + ), envvar="NUM_WORKERS", ) +@click.option( + "--granian_threads", + default=None, + type=click.IntRange(min=1), + help=( + "Only with --run_granian: runtime threads per worker process " + "(Granian --runtime-threads / GRANIAN_RUNTIME_THREADS). Omit to use Granian's default (1)." + ), + envvar="GRANIAN_RUNTIME_THREADS", +) @click.option("--api_base", default=None, help="API base URL.") @click.option( "--api_version", @@ -600,6 +695,15 @@ class ProxyInitializationHelpers: is_flag=True, help="Starts proxy via hypercorn, instead of uvicorn (supports HTTP/2)", ) +@click.option( + "--run_granian", + default=False, + is_flag=True, + help=( + "Starts proxy via Granian (Rust ASGI server) instead of uvicorn. " + "Requires Python 3.10+ and the `granian` package." + ), +) @click.option( "--ssl_keyfile_path", default=None, @@ -704,6 +808,7 @@ def run_server( # noqa: PLR0915 test, local, num_workers, + granian_threads, test_async, iam_token_db_auth, num_requests, @@ -713,6 +818,7 @@ def run_server( # noqa: PLR0915 version, run_gunicorn, run_hypercorn, + run_granian, ssl_keyfile_path, ssl_certfile_path, ciphers, @@ -797,16 +903,29 @@ def run_server( # noqa: PLR0915 config=config, use_queue=use_queue, ) - try: - import uvicorn - except Exception: - raise ImportError( - "uvicorn, gunicorn needs to be imported. Run - `pip install 'litellm[proxy]'`" - ) + if run_granian: + try: + import granian # noqa: F401 + except ImportError as e: + raise ImportError( + "granian must be installed to use --run_granian. " + "Run `pip install granian` or `pip install 'litellm[proxy]'` " + "(Granian requires Python 3.10+)." + ) from e + else: + try: + import uvicorn + except Exception: + raise ImportError( + "uvicorn, gunicorn needs to be imported. Run - `pip install 'litellm[proxy]'`" + ) db_connection_pool_limit = 100 # Starts optional due to config fallback checks; guaranteed non-None before use. db_connection_timeout: Optional[Union[int, float]] = 60 + db_connect_timeout: Optional[Union[int, float]] = None + db_socket_timeout: Optional[Union[int, float]] = None + db_extra_connection_params: Optional[dict] = None general_settings = {} ### GET DB TOKEN FOR IAM AUTH ### @@ -924,6 +1043,11 @@ def run_server( # noqa: PLR0915 db_connection_timeout = ( LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value ) + db_connect_timeout = general_settings.get("database_connect_timeout") + db_socket_timeout = general_settings.get("database_socket_timeout") + db_extra_connection_params = general_settings.get( + "database_extra_connection_params" + ) if database_url and database_url.startswith("os.environ/"): original_dir = os.getcwd() # set the working directory to where this script is @@ -963,27 +1087,26 @@ def run_server( # noqa: PLR0915 try: from litellm.secret_managers.main import get_secret + connection_url_params = _build_db_connection_url_params( + connection_limit=db_connection_pool_limit, + pool_timeout=db_connection_timeout, + connect_timeout=db_connect_timeout, + socket_timeout=db_socket_timeout, + extra_params=db_extra_connection_params, + ) if os.getenv("DATABASE_URL", None) is not None: - ### add connection pool + pool timeout args - params = { - "connection_limit": db_connection_pool_limit, - "pool_timeout": db_connection_timeout, - } database_url = get_secret("DATABASE_URL", default_value=None) modified_url = append_query_params( - str(database_url) if database_url else None, params + str(database_url) if database_url else None, + connection_url_params, ) os.environ["DATABASE_URL"] = modified_url if os.getenv("DIRECT_URL", None) is not None: - ### add connection pool + pool timeout args - params = { - "connection_limit": db_connection_pool_limit, - "pool_timeout": db_connection_timeout, - } database_url = os.getenv("DIRECT_URL") - modified_url = append_query_params(database_url, params) + modified_url = append_query_params( + database_url, connection_url_params + ) os.environ["DIRECT_URL"] = modified_url - ### subprocess.run(["prisma"], capture_output=True) is_prisma_runnable = True except FileNotFoundError: @@ -1081,7 +1204,7 @@ def run_server( # noqa: PLR0915 # Optional: recycle uvicorn workers after N requests if max_requests_before_restart is not None: uvicorn_args["limit_max_requests"] = max_requests_before_restart - if run_gunicorn is False and run_hypercorn is False: + if run_gunicorn is False and run_hypercorn is False and run_granian is False: if ssl_certfile_path is not None and ssl_keyfile_path is not None: print( # noqa f"\033[1;32mLiteLLM Proxy: Using SSL with certfile: {ssl_certfile_path} and keyfile: {ssl_keyfile_path}\033[0m\n" # noqa @@ -1123,6 +1246,17 @@ def run_server( # noqa: PLR0915 ssl_keyfile_path=ssl_keyfile_path, ciphers=ciphers, ) + elif run_granian is True: + ProxyInitializationHelpers._init_granian_server( + host=host, + port=port, + num_workers=num_workers, + ssl_certfile_path=ssl_certfile_path, + ssl_keyfile_path=ssl_keyfile_path, + max_requests_before_restart=max_requests_before_restart, + ciphers=ciphers, + granian_runtime_threads=granian_threads, + ) if __name__ == "__main__": diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 879914e5ac6..759534a32a1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2710,11 +2710,9 @@ def run_ollama_serve(): with open(os.devnull, "w") as devnull: subprocess.Popen(command, stdout=devnull, stderr=devnull) except Exception as e: - verbose_proxy_logger.debug( - f""" + verbose_proxy_logger.debug(f""" LiteLLM Warning: proxy started with `ollama` model\n`ollama serve` failed with Exception{e}. \nEnsure you run `ollama serve` - """ - ) + """) def _get_process_rss_mb() -> Optional[float]: @@ -4328,6 +4326,19 @@ class ProxyConfig: "health_check_concurrency", None ) health_check_details = general_settings.get("health_check_details", True) + ### INTERACTIONS API SCHEMA ### + _use_legacy_interactions_schema = general_settings.get( + "use_legacy_interactions_schema" + ) + if _use_legacy_interactions_schema is not None: + if isinstance(_use_legacy_interactions_schema, str): + litellm.use_legacy_interactions_schema = ( + _use_legacy_interactions_schema.lower() == "true" + ) + else: + litellm.use_legacy_interactions_schema = bool( + _use_legacy_interactions_schema + ) # Health-check-driven routing (opt-in, passes through to Router later) _enable_hc_routing = general_settings.get( "enable_health_check_routing", False diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index d030fabe8b5..e3019801aae 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3184,16 +3184,14 @@ async def provider_budgets() -> ProviderBudgetResponse: async def get_spend_by_tags( prisma_client: PrismaClient, start_date=None, end_date=None ): - response = await prisma_client.db.query_raw( - """ + response = await prisma_client.db.query_raw(""" SELECT jsonb_array_elements_text(request_tags) AS individual_request_tag, COUNT(*) AS log_count, SUM(spend) AS total_spend FROM "LiteLLM_SpendLogs" GROUP BY individual_request_tag; - """ - ) + """) return response diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 32c887f17b2..032ab6c63b2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2979,8 +2979,7 @@ class PrismaClient: required_view = "LiteLLM_VerificationTokenView" expected_views_str = ", ".join(f"'{view}'" for view in expected_views) pg_schema = os.getenv("DATABASE_SCHEMA", "public") - ret = await self.db.query_raw( - f""" + ret = await self.db.query_raw(f""" WITH existing_views AS ( SELECT viewname FROM pg_views @@ -2992,8 +2991,7 @@ class PrismaClient: (SELECT COUNT(*) FROM existing_views) AS view_count, ARRAY_AGG(viewname) AS view_names FROM existing_views - """ - ) + """) expected_total_views = len(expected_views) if ret[0]["view_count"] == expected_total_views: verbose_proxy_logger.info("All necessary views exist!") @@ -3002,8 +3000,7 @@ class PrismaClient: ## check if required view exists ## if ret[0]["view_names"] and required_view not in ret[0]["view_names"]: await self.health_check() # make sure we can connect to db - await self.db.execute_raw( - """ + await self.db.execute_raw(""" CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -3013,8 +3010,7 @@ class PrismaClient: t.rpm_limit AS team_rpm_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; - """ - ) + """) verbose_proxy_logger.info( "LiteLLM_VerificationTokenView Created in DB!" @@ -6068,6 +6064,8 @@ async def get_available_models_for_user( include_model_access_groups=include_model_access_groups, ) + effective_team_id = team_id or user_api_key_dict.team_id + # Get complete model list all_models = get_complete_model_list( key_models=key_models, @@ -6080,6 +6078,7 @@ async def get_available_models_for_user( model_access_groups=model_access_groups, include_model_access_groups=include_model_access_groups, only_model_access_groups=only_model_access_groups, + team_id=effective_team_id, ) return all_models diff --git a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py index 8ce1bedcf90..b47f6a747db 100644 --- a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py @@ -1,5 +1,5 @@ """ -What is this? +What is this? Logging Pass-Through Endpoints """ diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index f730a089624..03a2f339bea 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -65,8 +65,7 @@ class LiteLLMCompletionTransformationHandler: litellm_completion_response: Union[ ModelResponse, litellm.CustomStreamWrapper ] = litellm.completion( - **litellm_completion_request, - **kwargs, + **completion_args, ) if isinstance(litellm_completion_response, ModelResponse): diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 4ee9235af7d..35680889d86 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1115,6 +1115,7 @@ def responses( stream=stream, extra_headers=extra_headers, extra_body=extra_body, + timeout=timeout if timeout is not None else request_timeout, **kwargs, ) diff --git a/litellm/responses/sse_output_recovery.py b/litellm/responses/sse_output_recovery.py new file mode 100644 index 00000000000..5c18770a611 --- /dev/null +++ b/litellm/responses/sse_output_recovery.py @@ -0,0 +1,136 @@ +""" +Shared helpers for recovering Responses API output items from raw SSE chunks. + +The same recovery logic is needed in multiple places (e.g. the ChatGPT +Responses transformation and the LiteLLM Responses-to-Chat-Completions +bridge). Keep the implementation in a single module so a fix in one +caller automatically applies to all of them. +""" + +import json +from typing import Any, Dict, Optional + +from litellm.constants import STREAM_SSE_DONE_STRING + +_MAX_CONTENT_INDEX = 1024 + + +def parse_sse_json_chunk(chunk: str) -> Optional[Dict[str, Any]]: + """Parse a single raw SSE line into a JSON object dict. + + Returns ``None`` for empty lines, ``event:`` lines, ``[DONE]`` markers, + invalid JSON, or non-dict payloads. Centralizes the parsing step that + feeds into the recovery helpers in this module so behavior stays + consistent across all callers. + """ + # Import locally to avoid a circular import with the streaming handler. + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + stripped_chunk = ( + CustomStreamWrapper._strip_sse_data_from_chunk(chunk.strip()) or "" + ).strip() + if ( + not stripped_chunk + or stripped_chunk == STREAM_SSE_DONE_STRING + or stripped_chunk.startswith("event:") + ): + return None + try: + parsed_chunk = json.loads(stripped_chunk) + except json.JSONDecodeError: + return None + if not isinstance(parsed_chunk, dict): + return None + return parsed_chunk + + +def record_output_item_chunk( + parsed_chunk: Dict[str, Any], + output_items: Dict[int, Dict[str, Any]], +) -> None: + """Record an OUTPUT_ITEM_DONE chunk into ``output_items`` keyed by + ``output_index`` (falling back to the next free slot when missing). + """ + item = parsed_chunk.get("item") + if not isinstance(item, dict): + return + try: + output_index_raw = parsed_chunk.get("output_index") + if output_index_raw is None: + raise ValueError("missing output_index") + output_index = int(output_index_raw) + except (TypeError, ValueError): + output_index = len(output_items) + output_items[output_index] = item + + +def record_output_text_chunk( + parsed_chunk: Dict[str, Any], + output_items: Dict[int, Dict[str, Any]], + text_only_items: Dict[int, Dict[str, Any]], +) -> None: + """Record an OUTPUT_TEXT_DONE chunk as a synthetic message item in + ``text_only_items``. Real OUTPUT_ITEM_DONE events already captured in + ``output_items`` take precedence at the same ``output_index``. + """ + text = parsed_chunk.get("text") + if not isinstance(text, str): + return + + try: + output_index_raw = parsed_chunk.get("output_index") + if output_index_raw is None: + raise ValueError("missing output_index") + output_index = int(output_index_raw) + except (TypeError, ValueError): + output_index = len(text_only_items) + + if output_index in output_items: + return + + item = text_only_items.get(output_index) + if item is None: + item = { + "type": "message", + "id": parsed_chunk.get("item_id") or f"msg_{output_index}", + "role": "assistant", + "status": "completed", + "content": [], + } + text_only_items[output_index] = item + + content = item.setdefault("content", []) + if not isinstance(content, list): + return + + try: + content_index_raw = parsed_chunk.get("content_index") + if content_index_raw is None: + raise ValueError("missing content_index") + content_index = int(content_index_raw) + except (TypeError, ValueError): + content_index = len(content) + + if content_index < 0 or content_index > _MAX_CONTENT_INDEX: + return + + while len(content) <= content_index: + content.append( + { + "type": "output_text", + "text": "", + "annotations": [], + } + ) + + content_item = content[content_index] + if not isinstance(content_item, dict): + content_item = {} + content[content_index] = content_item + + content_item["type"] = "output_text" + content_item["text"] = text + if parsed_chunk.get("annotations") is not None: + content_item["annotations"] = parsed_chunk["annotations"] + else: + content_item.setdefault("annotations", []) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index da8da1b486f..c4e72cb7dc5 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -9,6 +9,7 @@ from functools import lru_cache from typing import Any, Dict, List, Literal, Optional import httpx +from openai._streaming import SSEDecoder import litellm from litellm.constants import ( @@ -27,7 +28,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import CallTypes -from litellm.utils import CustomStreamWrapper, async_post_call_success_deployment_hook +from litellm.utils import async_post_call_success_deployment_hook @lru_cache(maxsize=1) @@ -120,10 +121,10 @@ class BaseResponsesAPIStreamingIterator: if not chunk: return None - # Handle SSE format (data: {...}) - chunk = CustomStreamWrapper._strip_sse_data_from_chunk(chunk) - if chunk is None: - return None + # NOTE: ``SSEDecoder`` already strips the SSE ``data:`` field prefix, so + # the value passed in here is the raw field content. Do not re-run + # ``_strip_sse_data_from_chunk`` on it — doing so would incorrectly mangle + # payloads whose actual JSON value happens to start with ``data:``. # Handle "[DONE]" marker if chunk == STREAM_SSE_DONE_STRING: @@ -634,7 +635,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): request_data, call_type, ) - self.stream_iterator = response.aiter_lines() + self.stream_iterator = SSEDecoder().aiter_bytes(response.aiter_bytes()) def __aiter__(self): return self @@ -645,13 +646,13 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): while True: # Get the next chunk from the stream try: - chunk = await self.stream_iterator.__anext__() + sse = await self.stream_iterator.__anext__() except StopAsyncIteration: self.finished = True raise StopAsyncIteration self._check_max_streaming_duration() - result = self._process_chunk(chunk) + result = self._process_chunk(sse.data) if self.finished: raise StopAsyncIteration @@ -708,7 +709,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): request_data, call_type, ) - self.stream_iterator = response.iter_lines() + self.stream_iterator = SSEDecoder().iter_bytes(response.iter_bytes()) def __iter__(self): return self @@ -719,13 +720,13 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): while True: # Get the next chunk from the stream try: - chunk = next(self.stream_iterator) + sse = next(self.stream_iterator) except StopIteration: self.finished = True raise StopIteration self._check_max_streaming_duration() - result = self._process_chunk(chunk) + result = self._process_chunk(sse.data) if self.finished: raise StopIteration diff --git a/litellm/router.py b/litellm/router.py index 019f565e5c4..debccb0e83f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -208,6 +208,15 @@ if TYPE_CHECKING: from litellm.router_strategy.quality_router.quality_router import ( QualityRouter, ) + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseInputParam, + ResponsesAPIResponse, + ) Span = Union[_Span, Any] else: @@ -839,7 +848,7 @@ class Router: @staticmethod def _normalize_strategy( - strategy: Union[RoutingStrategy, str, None] + strategy: Union[RoutingStrategy, str, None], ) -> Optional[str]: if strategy is None: return None @@ -2246,6 +2255,388 @@ class Router: return FallbackStreamWrapper(stream_with_fallbacks()) + @staticmethod + def _extract_partial_responses_usage( + source_iterator: "BaseResponsesAPIStreamingIterator", + ) -> Optional["ResponseAPIUsage"]: + """ + Best-effort: pull partial token usage from a Responses-API streaming + iterator that errored mid-stream, normalized to ResponseAPIUsage so + the caller can combine without crossing token-naming conventions. + + Two sources, in priority order: + 1. The bridge path (LiteLLMCompletionStreamingIterator) accumulates + chat-completion chunks while streaming — feed them through + stream_chunk_builder to recover chat Usage, then translate + (prompt_tokens → input_tokens, completion_tokens → output_tokens). + 2. The native path (ResponsesAPIStreamingIterator) only has a + completed_response object if the stream reached + RESPONSE_COMPLETED before erroring — uncommon mid-stream but + worth checking. Already ResponseAPIUsage-shaped. + + Returns None when no partial usage is recoverable. + """ + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ) + + # Bridge subclass is the only iterator that accumulates chat-completion + # chunks. isinstance narrows the type so we can read the attribute + # directly instead of getattr-ing on the base class. + if isinstance(source_iterator, LiteLLMCompletionStreamingIterator): + chunks = source_iterator.collected_chat_completion_chunks + if chunks: + try: + from litellm.main import stream_chunk_builder + + built = stream_chunk_builder(chunks=chunks) + # stream_chunk_builder returns ModelResponse | + # TextCompletionResponse | None. ModelResponse sets .usage + # in __init__ rather than declaring it as a class field, so + # static narrowing doesn't expose it. Mirror the sync path + # (_completion_streaming_iterator) and pull via getattr. + chat = getattr(built, "usage", None) if built is not None else None + if chat is not None: + # getattr-with-default because the test path may + # substitute a SimpleNamespace lacking some fields; + # real Usage instances always have them. + prompt = int(getattr(chat, "prompt_tokens", 0) or 0) + completion = int(getattr(chat, "completion_tokens", 0) or 0) + total = int( + getattr(chat, "total_tokens", prompt + completion) + or (prompt + completion) + ) + return ResponseAPIUsage( + input_tokens=prompt, + output_tokens=completion, + total_tokens=total, + ) + except Exception: + # Builder is best-effort — fall through to native path. + pass + + # Native path: completed_response is set only if RESPONSE_COMPLETED + # arrived before the error (uncommon mid-stream but worth checking). + # Already ResponseAPIUsage-shaped — return as-is. + completed = source_iterator.completed_response + if isinstance( + completed, + (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent), + ): + return completed.response.usage + return None + + @staticmethod + def _combine_responses_fallback_usage( + fallback_item: "BaseLiteLLMOpenAIResponseObject", + partial_usage: "ResponseAPIUsage", + ) -> None: + """ + Merge partial-stream usage with fallback-stream usage on a + Responses-API streaming event. + + Only mutates events that carry a `response` with a `usage` field + (response.completed / response.failed / response.incomplete). Other + events pass through unchanged. + + Both inputs are ResponseAPIUsage-shaped (see + _extract_partial_responses_usage which normalizes the bridge path), + so we can sum input_tokens / output_tokens / total_tokens directly + and produce a clean ResponseAPIUsage — no token-naming split, no + setattr bypass. + """ + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ) + + if not isinstance( + fallback_item, + (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent), + ): + return + response = fallback_item.response + if response.usage is None: + return + + fb = response.usage + response.usage = ResponseAPIUsage( + input_tokens=(partial_usage.input_tokens or 0) + (fb.input_tokens or 0), + output_tokens=(partial_usage.output_tokens or 0) + (fb.output_tokens or 0), + total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), + ) + + @staticmethod + def _build_responses_continuation_input( + input_val: Optional[Union[str, "ResponseInputParam"]], + generated_content: str, + ) -> "ResponseInputParam": + """ + Convert Responses-API input + partial assistant output into a + continuation input that asks the fallback model to pick up where the + prior assistant message stopped. + + Best effort across providers. The chat-completions path uses + Anthropic's `prefix: True` prefill trick on the assistant message; + the Responses-API input schema has no direct equivalent, so we + append an instruction (developer role) plus a prior assistant + message containing the partial output. Providers without prefill + semantics (OpenAI, Vertex) treat this as conversational context + and may regenerate — same trade-off as the chat-completions path + for non-Anthropic fallbacks. + """ + # base/continuation are List[Any] because ResponseInputParam items + # are a wide Union of TypedDicts (EasyInputMessageParam, Message, + # ResponseOutputMessageParam, ...) — annotating as List[Dict[str, Any]] + # rejects the list() spread of input_val. We cast the combined list to + # ResponseInputParam at the return. + base: List[Any] + if isinstance(input_val, str): + base = [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": input_val}], + } + ] + elif isinstance(input_val, list): + base = list(input_val) + else: + base = [] + continuation: List[Any] = [ + { + "type": "message", + "role": "developer", + "content": [ + { + "type": "input_text", + "text": ( + "The previous assistant response was interrupted " + "mid-stream. Continue exactly where it stopped — " + "do not repeat any of its content. Your response " + "must read as a seamless continuation." + ), + } + ], + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": generated_content}], + }, + ] + return cast("ResponseInputParam", base + continuation) + + async def _aresponses_streaming_iterator( + self, + response: "BaseResponsesAPIStreamingIterator", + initial_kwargs: Dict[str, Any], + ) -> "BaseResponsesAPIStreamingIterator": + """ + Wrap a Responses-API streaming iterator so MidStreamFallbackError + triggers the Router's fallback chain (parity with + _acompletion_streaming_iterator for the chat-completions path). + + The Responses-API streaming path goes through + _ageneric_api_call_with_fallbacks rather than _acompletion, so the + returned iterator is never wrapped by the chat completions + fallback handler. Without this wrapper, MidStreamFallbackError + raised mid-stream from the underlying CustomStreamWrapper (used by + LiteLLMCompletionStreamingIterator when the Responses API is + served via the completion bridge) propagates unhandled and the + configured cross-provider fallback never fires. + + Full parity with the chat-completions path: + - Pre-first-chunk: retry with the original input unchanged. + - Partial content: inject a developer instruction + prior + assistant message carrying the generated text so the fallback + model continues rather than restarts. + - Usage combining: merge partial-stream usage onto the fallback's + response.completed event so accounting reflects both attempts. + - Stream cleanup: shielded aclose() on both source and fallback + iterators on terminate. + """ + from litellm.exceptions import MidStreamFallbackError + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + source_iterator = response + + class FallbackResponsesStreamWrapper(BaseResponsesAPIStreamingIterator): + """ + Subclasses BaseResponsesAPIStreamingIterator only for isinstance + compatibility (proxy + interactions code paths check the type). + Bypasses the parent constructor and delegates iteration to an + async generator. + """ + + def __init__(self, async_generator: AsyncGenerator): + import time + from datetime import datetime + + self._async_generator = async_generator + # Mirror every attribute BaseResponsesAPIStreamingIterator.__init__ + # would have set. The wrapper bypasses super().__init__ (it has no + # httpx.Response of its own and no provider config to drive), so + # we copy from source_iterator where applicable and use safe + # defaults elsewhere. This keeps inherited methods (e.g. + # _check_max_streaming_duration, _handle_failure) safe to call. + # + # The bridge path (LiteLLMCompletionStreamingIterator used by + # Anthropic/Bedrock/Vertex) does not call super().__init__ and + # is missing many of these attributes — use getattr fallbacks + # so wrapper construction never raises AttributeError. The + # bridge stores the logging object as `litellm_logging_obj`. + self.response = getattr(source_iterator, "response", None) + self.model = getattr(source_iterator, "model", None) + self.logging_obj = getattr( + source_iterator, + "logging_obj", + getattr(source_iterator, "litellm_logging_obj", None), + ) + self.finished = False + self.responses_api_provider_config = getattr( + source_iterator, "responses_api_provider_config", None + ) + self.completed_response = None + self.start_time = getattr(source_iterator, "start_time", datetime.now()) + self._failure_handled = False + self._completed_response_cached = False + self._completed_response_logged = False + self._completed_response_cache_hit = None + self._persist_completed_response_before_logging = True + self._stream_created_time = time.time() + self.litellm_metadata = getattr( + source_iterator, "litellm_metadata", None + ) + self.custom_llm_provider = getattr( + source_iterator, "custom_llm_provider", None + ) + self.request_data = getattr(source_iterator, "request_data", {}) or {} + self.call_type = getattr(source_iterator, "call_type", None) + # Preserve hidden params so response headers (model_id, + # api_base, additional_headers) keep flowing. + self._hidden_params = dict( + getattr(source_iterator, "_hidden_params", None) or {} + ) + + def __aiter__(self): + return self + + async def __anext__(self): + return await self._async_generator.__anext__() + + async def aclose(self): + # async generators always expose aclose — no defensive check needed. + await self._async_generator.aclose() + + async def stream_with_fallbacks(): + fallback_response = None + try: + async for item in source_iterator: + yield item + except MidStreamFallbackError as e: + partial_usage = Router._extract_partial_responses_usage(source_iterator) + try: + model_group = cast(str, initial_kwargs.get("model")) + fallbacks: Optional[List] = initial_kwargs.get( + "fallbacks", self.fallbacks + ) + context_window_fallbacks: Optional[List] = initial_kwargs.get( + "context_window_fallbacks", self.context_window_fallbacks + ) + content_policy_fallbacks: Optional[List] = initial_kwargs.get( + "content_policy_fallbacks", self.content_policy_fallbacks + ) + # Re-enter via the per-attempt helper so the fallback chain + # picks deployments through + # _ageneric_api_call_with_fallbacks_helper. + # original_generic_function is preserved by the caller so + # the helper knows what underlying API to invoke per attempt. + initial_kwargs["original_function"] = ( + self._ageneric_api_call_with_fallbacks_helper + ) + if e.is_pre_first_chunk or not e.generated_content: + # No content generated before the error — retry with the + # original input. Adding a continuation prompt would + # waste tokens and confuse the model. + pass + else: + initial_kwargs["input"] = ( + Router._build_responses_continuation_input( + initial_kwargs.get("input"), + e.generated_content, + ) + ) + # The Responses-API path stores observability metadata + # under "litellm_metadata" (not the default "metadata") — + # see _ageneric_api_call_with_fallbacks. Mirroring that + # here ensures model_group, model_group_alias, and trace + # ids land in the same key litellm.aresponses reads from. + self._update_kwargs_before_fallbacks( + model=model_group, + kwargs=initial_kwargs, + metadata_variable_name="litellm_metadata", + ) + fallback_response = ( + await self.async_function_with_fallbacks_common_utils( + e=e, + disable_fallbacks=False, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + model_group=model_group, + args=(), + kwargs=initial_kwargs, + ) + ) + + if hasattr(fallback_response, "__aiter__"): + async for fallback_item in fallback_response: # type: ignore + if partial_usage is not None: + Router._combine_responses_fallback_usage( + fallback_item, partial_usage + ) + yield fallback_item + else: + yield fallback_response + except Exception as fallback_error: + verbose_router_logger.error( + f"Responses streaming fallback also failed: {fallback_error}" + ) + raise fallback_error + finally: + with anyio.CancelScope(shield=True): + if hasattr(source_iterator, "aclose"): + try: + await source_iterator.aclose() # type: ignore[func-returns-value] + except BaseException as exc: + verbose_router_logger.debug( + "stream_with_fallbacks(aresponses): error closing source: %s", + exc, + ) + if fallback_response is not None and hasattr( + fallback_response, "aclose" + ): + try: + await fallback_response.aclose() + except BaseException as exc: + verbose_router_logger.debug( + "stream_with_fallbacks(aresponses): error closing fallback: %s", + exc, + ) + + return FallbackResponsesStreamWrapper(stream_with_fallbacks()) + def _completion_streaming_iterator( # noqa: PLR0915 self, model_response: CustomStreamWrapper, @@ -4292,6 +4683,61 @@ class Router: self.fail_calls[model] += 1 raise e + async def _aresponses_with_streaming_fallbacks( + self, original_function: Callable, **kwargs: Any + ) -> Union["ResponsesAPIResponse", "BaseResponsesAPIStreamingIterator"]: + """ + _ageneric_api_call_with_fallbacks for the Responses API, with the + addition of mid-stream fallback handling. + + When stream=True and the underlying call returns a + BaseResponsesAPIStreamingIterator, wrap it with + _aresponses_streaming_iterator so MidStreamFallbackError raised + during iteration triggers the Router's cross-provider fallback chain. + """ + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + from litellm.litellm_core_utils.core_helpers import safe_deep_copy + + # Snapshot the request kwargs before _ageneric_api_call_with_fallbacks + # mutates them. A shallow copy alone is not enough: the primary + # attempt mutates nested dicts in place — notably `litellm_metadata`, + # which `_update_kwargs_with_deployment` populates with + # deployment-specific fields (`deployment`, `model_info`, `api_base`, + # tags, etc.). Without an explicit copy of that dict, the shallow + # copy would still share its reference, leaking primary-deployment + # metadata into the mid-stream fallback request. + # + # We avoid deep-copying the full kwargs because it can contain + # non-deepcopyable objects (logging handles, async clients, etc.); + # `safe_deep_copy` deep-copies the metadata dicts key-by-key with a + # fallback to the original reference for any non-picklable value. + # The original_generic_function is preserved so the per-attempt + # helper knows which underlying API to call on fallback. + fallback_kwargs: Dict[str, Any] = kwargs.copy() + if isinstance(fallback_kwargs.get("litellm_metadata"), dict): + fallback_kwargs["litellm_metadata"] = safe_deep_copy( + fallback_kwargs["litellm_metadata"] + ) + if isinstance(fallback_kwargs.get("metadata"), dict): + fallback_kwargs["metadata"] = safe_deep_copy(fallback_kwargs["metadata"]) + fallback_kwargs["original_generic_function"] = original_function + + response = await self._ageneric_api_call_with_fallbacks( + original_function=original_function, **kwargs + ) + + if kwargs.get("stream") and isinstance( + response, BaseResponsesAPIStreamingIterator + ): + return await self._aresponses_streaming_iterator( + response=response, + initial_kwargs=fallback_kwargs, + ) + return response + def _generic_api_call_with_fallbacks( self, model: str, original_function: Callable, **kwargs ): @@ -5511,9 +5957,13 @@ class Router: custom_llm_provider=custom_llm_provider, **kwargs, ) + elif call_type == "aresponses": + return await self._aresponses_with_streaming_fallbacks( + original_function=original_function, + **kwargs, + ) elif call_type in ( "anthropic_messages", - "aresponses", "_arealtime", "_aresponses_websocket", "acreate_fine_tuning_job", @@ -5670,6 +6120,7 @@ class Router: from litellm.responses.utils import ResponsesAPIRequestUtils container_id = kwargs.get("container_id") + _forwarded_model_id = kwargs.get("model_id") if isinstance(container_id, str): decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) original_id = decoded.get("response_id", container_id) @@ -5678,7 +6129,14 @@ class Router: decoded_provider = decoded.get("custom_llm_provider") if decoded_provider and kwargs.get("custom_llm_provider") == "openai": kwargs["custom_llm_provider"] = decoded_provider - model_id = decoded.get("model_id") + # Fall back to the model_id forwarded by the proxy when the container_id + # is a native upstream ID (e.g. Azure hex cntr_) that carries no LiteLLM + # routing payload, so deployment credentials (api_base, api_key) are applied. + model_id = decoded.get("model_id") or ( + _forwarded_model_id.strip() + if isinstance(_forwarded_model_id, str) and _forwarded_model_id.strip() + else None + ) if model_id: kwargs["model"] = model_id return await self._ageneric_api_call_with_fallbacks( @@ -7320,6 +7778,38 @@ class Router: _shared_model_info = { k: v for k, v in _model_info.items() if k not in _custom_pricing_fields } + _existing_shared_mode = ( + cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {} + ).get("mode") + _deployment_mode = _shared_model_info.get("mode") + # Keep the built-in bridge mode stable for shared backend keys. + # Multiple aliases can point at the same provider/model backend, + # but their deployment-level overrides should not downgrade the + # backend from responses -> chat via last-write-wins registration. + # Only preserve in that specific direction so legitimate upgrades + # (e.g. chat -> responses) and unrelated mode changes still apply, + # and so a missing deployment mode does not silently clear the + # existing shared backend mode. + _is_responses_to_chat_downgrade = ( + _existing_shared_mode == "responses" and _deployment_mode == "chat" + ) + _would_clear_existing_mode = ( + _existing_shared_mode is not None and _deployment_mode is None + ) + if _is_responses_to_chat_downgrade or _would_clear_existing_mode: + if _deployment_mode is not None: + verbose_router_logger.warning( + "Router: preserving existing mode=%s for shared backend " + "key %s instead of the deployment-specified mode=%s " + "(prevents alias registration from downgrading the " + "shared backend mode).", + _existing_shared_mode, + _model_name, + _deployment_mode, + ) + _shared_model_info["mode"] = _existing_shared_mode + + # Always register the (possibly mode-preserved) shared backend info. _backend_alias_cost = {_model_name: _shared_model_info} if "responses/" in _model_name: _stripped_model_name = _model_name.replace("responses/", "") diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 9e346006ac1..99fe5e26f7f 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -103,7 +103,7 @@ def _last_user_content(messages: Optional[List[Dict[str, Any]]]) -> Optional[str def _recent_tool_results( - messages: Optional[List[Dict[str, Any]]] + messages: Optional[List[Dict[str, Any]]], ) -> List[Dict[str, Any]]: """Extract the current turn's tool result payloads from the request messages. diff --git a/litellm/router_strategy/adaptive_router/signals.py b/litellm/router_strategy/adaptive_router/signals.py index a48bdea1eb6..5e33a64d27f 100644 --- a/litellm/router_strategy/adaptive_router/signals.py +++ b/litellm/router_strategy/adaptive_router/signals.py @@ -24,7 +24,6 @@ from litellm.router_strategy.adaptive_router.config import ( TOOL_CALL_HISTORY_MAX, ) - # ---- Public types --------------------------------------------------------- diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index be27b852478..da41577e99a 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -10,11 +10,11 @@ This means you can use this with weighted-pick, lowest-latency, simple-shuffle, Example: ``` openai: - budget_limit: 0.000000000001 - time_period: 1d + budget_limit: 0.000000000001 + time_period: 1d anthropic: - budget_limit: 100 - time_period: 7d + budget_limit: 100 + time_period: 7d ``` """ diff --git a/litellm/router_utils/get_retry_from_policy.py b/litellm/router_utils/get_retry_from_policy.py index ec326ebb50d..162d6428f85 100644 --- a/litellm/router_utils/get_retry_from_policy.py +++ b/litellm/router_utils/get_retry_from_policy.py @@ -1,5 +1,5 @@ """ -Get num retries for an exception. +Get num retries for an exception. - Account for retry policy by exception type. """ diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 17b453d6031..48f85a83411 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -34,7 +34,7 @@ class PatternUtils: @staticmethod def sorted_patterns( - patterns: Dict[str, List[Dict]] + patterns: Dict[str, List[Dict]], ) -> List[Tuple[str, List[Dict]]]: """ Cached property for patterns sorted by specificity. diff --git a/litellm/router_utils/router_callbacks/track_deployment_metrics.py b/litellm/router_utils/router_callbacks/track_deployment_metrics.py index 1f226879d03..9039b0df8e6 100644 --- a/litellm/router_utils/router_callbacks/track_deployment_metrics.py +++ b/litellm/router_utils/router_callbacks/track_deployment_metrics.py @@ -1,5 +1,5 @@ """ -Helper functions to get/set num success and num failures per deployment +Helper functions to get/set num success and num failures per deployment set_deployment_failures_for_current_minute diff --git a/litellm/secret_managers/aws_secret_manager.py b/litellm/secret_managers/aws_secret_manager.py index fbe951e6492..60d0a713eff 100644 --- a/litellm/secret_managers/aws_secret_manager.py +++ b/litellm/secret_managers/aws_secret_manager.py @@ -4,7 +4,7 @@ This is a file for the AWS Secret Manager Integration Relevant issue: https://github.com/BerriAI/litellm/issues/1883 Requires: -* `os.environ["AWS_REGION_NAME"], +* `os.environ["AWS_REGION_NAME"], * `pip install boto3>=1.28.57` """ diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index c1b4d019dcf..4461e34396e 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -10,7 +10,7 @@ Handles Async Operations for: Relevant issue: https://github.com/BerriAI/litellm/issues/1883 Requires: -* `os.environ["AWS_REGION_NAME"], +* `os.environ["AWS_REGION_NAME"], * `pip install boto3>=1.28.57` """ diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 751113400d3..0a51ce3d456 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -100,6 +100,7 @@ class SupportedGuardrailIntegrations(Enum): MCP_JWT_SIGNER = "mcp_jwt_signer" LLM_AS_A_JUDGE = "llm_as_a_judge" QOSTODIAN_NEXUS = "qostodian_nexus" + RUBRIK = "rubrik" class Role(Enum): diff --git a/litellm/types/interactions/__init__.py b/litellm/types/interactions/__init__.py index 0f934fa0152..78d0b04ef3b 100644 --- a/litellm/types/interactions/__init__.py +++ b/litellm/types/interactions/__init__.py @@ -36,9 +36,13 @@ from litellm.types.interactions.generated import ( GoogleSearchResultContent, ImageContent, Interaction, + InteractionCompleted, + InteractionCreated, InteractionEvent, InteractionEnvironment, + InteractionInProgress, InteractionInput, + InteractionRequiresAction, InteractionsAPIOptionalRequestParams, InteractionsAPIResponse, InteractionsAPIStreamingResponse, @@ -50,6 +54,9 @@ from litellm.types.interactions.generated import ( McpServerToolResultContent, ModelOption, ResponseModality, + StepDelta, + StepStart, + StepStop, ) from litellm.types.interactions.generated import ( Status3 as InteractionStatus, # Main request/response types; Content types; Turn for multi-turn conversations; Tool types; Config types; Usage; Status enum; Events for streaming; Agent configs; Model/Agent options; Response modality; Annotation; LiteLLM types; Backwards compat aliases @@ -115,6 +122,14 @@ __all__ = [ "AgentOption", "ResponseModality", "Annotation", + # New schema SSE event types (Api-Revision: 2026-05-20) + "StepStart", + "StepDelta", + "StepStop", + "InteractionCreated", + "InteractionInProgress", + "InteractionCompleted", + "InteractionRequiresAction", # LiteLLM types "InteractionEnvironment", "InteractionInput", diff --git a/litellm/types/interactions/generated.py b/litellm/types/interactions/generated.py index 2ce6331b448..d546e897891 100644 --- a/litellm/types/interactions/generated.py +++ b/litellm/types/interactions/generated.py @@ -1151,9 +1151,114 @@ class InteractionEvent(BaseModel): ) +# --------------------------------------------------------------- +# New schema SSE event types (Api-Revision: 2026-05-20) +# These replace the legacy content.* / interaction.start|complete +# events and will become the only events after June 8, 2026. +# --------------------------------------------------------------- + + +class StepStart(BaseModel): + """Emitted when a new step begins (replaces content.start).""" + + event_type: Literal["step.start"] = "step.start" + index: Optional[int] = None + step: Optional[Dict[str, Any]] = Field( + None, + description="The initial step data (type, content, signature, etc.).", + ) + event_id: Optional[str] = Field( + None, + description="The event_id token to be used to resume the interaction stream.", + ) + + +class StepDelta(BaseModel): + """Emitted for incremental step content (replaces content.delta).""" + + event_type: Literal["step.delta"] = "step.delta" + index: Optional[int] = None + delta: Optional[Dict[str, Any]] = Field( + None, + description="Incremental content delta (e.g. text, arguments_delta for function calls).", + ) + event_id: Optional[str] = Field( + None, + description="The event_id token to be used to resume the interaction stream.", + ) + + +class StepStop(BaseModel): + """Emitted when a step finishes (replaces content.stop).""" + + event_type: Literal["step.stop"] = "step.stop" + index: Optional[int] = None + status: Optional[str] = Field( + None, + description="Step completion status (e.g. 'done').", + ) + event_id: Optional[str] = Field( + None, + description="The event_id token to be used to resume the interaction stream.", + ) + + +class InteractionCreated(BaseModel): + """Emitted when the interaction is first created (replaces interaction.start).""" + + event_type: Literal["interaction.created"] = "interaction.created" + interaction: Optional[Dict[str, Any]] = None + event_id: Optional[str] = Field( + None, + description="The event_id token to be used to resume the interaction stream.", + ) + + +class InteractionInProgress(BaseModel): + """Emitted while the interaction is running.""" + + event_type: Literal["interaction.in_progress"] = "interaction.in_progress" + interaction_id: Optional[str] = None + event_id: Optional[str] = Field( + None, + description="The event_id token to be used to resume the interaction stream.", + ) + + +class InteractionCompleted(BaseModel): + """Emitted when the interaction finishes (replaces interaction.complete).""" + + event_type: Literal["interaction.completed"] = "interaction.completed" + interaction: Optional[Dict[str, Any]] = None + event_id: Optional[str] = Field( + None, + description="The event_id token to be used to resume the interaction stream.", + ) + + +class InteractionRequiresAction(BaseModel): + """Emitted when the interaction is paused waiting for a tool result.""" + + event_type: Literal["interaction.requires_action"] = "interaction.requires_action" + interaction_id: Optional[str] = None + event_id: Optional[str] = Field( + None, + description="The event_id token to be used to resume the interaction stream.", + ) + + class InteractionSseEvent( RootModel[ Union[ + # New schema events (Api-Revision: 2026-05-20) + StepStart, + StepDelta, + StepStop, + InteractionCreated, + InteractionInProgress, + InteractionCompleted, + InteractionRequiresAction, + # Legacy schema events (Api-Revision: 2026-05-07, removed June 8 2026) InteractionEvent, InteractionStatusUpdate, ContentStart, @@ -1164,6 +1269,15 @@ class InteractionSseEvent( ] ): root: Union[ + # New schema events (Api-Revision: 2026-05-20) + StepStart, + StepDelta, + StepStop, + InteractionCreated, + InteractionInProgress, + InteractionCompleted, + InteractionRequiresAction, + # Legacy schema events (Api-Revision: 2026-05-07, removed June 8 2026) InteractionEvent, InteractionStatusUpdate, ContentStart, @@ -1193,6 +1307,11 @@ class InteractionsAPIResponse(BaseLiteLLMOpenAIResponseObject): Response from the Interactions API. Wraps the API response with LiteLLM-specific hidden params. + + Schema notes: + - New schema (Api-Revision: 2026-05-20, default): response contains ``steps``. + - Legacy schema (Api-Revision: 2026-05-07, removed June 8 2026): response contains ``outputs``. + Both fields are kept here so callers work with either schema. """ id: Optional[str] = None @@ -1203,7 +1322,10 @@ class InteractionsAPIResponse(BaseLiteLLMOpenAIResponseObject): created: Optional[str] = None updated: Optional[str] = None role: Optional[str] = None + # Legacy schema field (Api-Revision: 2026-05-07). Remove after June 8, 2026. outputs: Optional[List[Dict[str, Any]]] = None + # New schema field (Api-Revision: 2026-05-20). + steps: Optional[List[Dict[str, Any]]] = None usage: Optional[Dict[str, Any]] = None _hidden_params: dict = PrivateAttr(default_factory=dict) @@ -1213,7 +1335,12 @@ class InteractionsAPIStreamingResponse(BaseLiteLLMOpenAIResponseObject): """ Streaming response chunk from the Interactions API. - Event types per OpenAPI spec: + New schema event types (Api-Revision: 2026-05-20): + - interaction.created, interaction.in_progress, interaction.completed, + interaction.requires_action + - step.start, step.delta, step.stop + + Legacy event types (Api-Revision: 2026-05-07, removed June 8 2026): - interaction.start, interaction.status_update, interaction.complete - content.start, content.delta, content.stop - error @@ -1228,9 +1355,17 @@ class InteractionsAPIStreamingResponse(BaseLiteLLMOpenAIResponseObject): created: Optional[str] = None updated: Optional[str] = None role: Optional[str] = None + # Legacy schema field (Api-Revision: 2026-05-07). Remove after June 8, 2026. outputs: Optional[List[Dict[str, Any]]] = None + # New schema field (Api-Revision: 2026-05-20). + steps: Optional[List[Dict[str, Any]]] = None usage: Optional[Dict[str, Any]] = None delta: Optional[Dict[str, Any]] = None + # New schema streaming fields + index: Optional[int] = None + step: Optional[Dict[str, Any]] = None + interaction_id: Optional[str] = None + interaction: Optional[Dict[str, Any]] = None _hidden_params: dict = PrivateAttr(default_factory=dict) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index b357b64156d..a1d53978761 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -16,15 +16,15 @@ GeminiEmbeddingInput = Union[EmbeddingInput, List[List[str]]] class FunctionResponse(TypedDict, total=False): # `id` correlates this response with the originating `functionCall` part. - # Required by Gemini 3.5+ for strict function-calling response matching. + # Supported on Google AI Studio Gemini 3.5+; Vertex AI rejects this field. id: str name: Required[str] response: Optional[dict] class FunctionCall(TypedDict, total=False): - # `id` is returned by Gemini 3.5+ to correlate the corresponding - # `functionResponse`. Older Gemini models omit this field. + # `id` correlates the corresponding `functionResponse` on Google AI Studio + # Gemini 3.5+. Vertex AI and older Gemini models omit/reject this field. id: str name: Required[str] args: Optional[dict] @@ -52,8 +52,8 @@ class PartType(TypedDict, total=False): class HttpxFunctionCall(TypedDict, total=False): - # `id` is returned by Gemini 3.5+ to correlate the corresponding - # `functionResponse`. Older Gemini models omit this field. + # `id` correlates the corresponding `functionResponse` on Google AI Studio + # Gemini 3.5+. Vertex AI and older Gemini models omit/reject this field. id: str name: Required[str] args: dict diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6084f14e2df..282baff07fe 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -147,6 +147,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_low_reasoning_effort: Optional[bool] supports_xhigh_reasoning_effort: Optional[bool] supports_max_reasoning_effort: Optional[bool] + supports_output_config: Optional[bool] class SearchContextCostPerQuery(TypedDict, total=False): @@ -243,6 +244,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): float ] # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) ocr_cost_per_page: Optional[float] # for OCR models + ocr_cost_per_credit: Optional[float] # for OCR models priced by credit annotation_cost_per_page: Optional[float] # for OCR models search_context_cost_per_query: Optional[ SearchContextCostPerQuery @@ -260,6 +262,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): "chat", "audio_transcription", "responses", + "ocr", ] ] tpm: Optional[int] @@ -2765,6 +2768,20 @@ class StandardLoggingGuardrailInformation(TypedDict, total=False): risk_score: Optional[float] """Risk score 0-10 indicating how risky the request was (higher = riskier). Computed by the guardrail provider.""" + violation_categories: Optional[List[str]] + """Names of the policy items that intervened on this request (e.g. Bedrock + topic-policy topic names, content-policy filter types, PII entity types). + Populated by the provider hook before redaction so downstream loggers + (OTEL, Langfuse, ...) can filter by violation category without parsing + the raw guardrail_response blob. Empty/absent when the guardrail allowed + the request through.""" + + guardrail_action: Optional[str] + """Provider's raw top-level action string (e.g. Bedrock's ``GUARDRAIL_INTERVENED`` + or ``NONE``). Populated by the provider hook so the OTEL integration can + surface it as a queryable span attribute without parsing the raw + guardrail_response blob.""" + class EvalVerdict(TypedDict, total=False): criterion_name: str @@ -2806,6 +2823,8 @@ class GuardrailTracingDetail(TypedDict, total=False): patterns_checked: Optional[int] alert_recipients: Optional[List[str]] risk_score: Optional[float] + violation_categories: Optional[List[str]] + guardrail_action: Optional[str] StandardLoggingPayloadStatus = Literal["success", "failure"] @@ -3219,6 +3238,7 @@ class LlmProviders(str, Enum): ANTHROPIC_TEXT = "anthropic_text" BYTEZ = "bytez" REPLICATE = "replicate" + REDUCTO = "reducto" RUNWAYML = "runwayml" AWS_POLLY = "aws_polly" HUGGINGFACE = "huggingface" diff --git a/litellm/utils.py b/litellm/utils.py index 001c89fee4c..c28a88e0f1c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2933,6 +2933,13 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 except Exception: existing_model = {} model_cost_key = key + # ``get_model_info`` returns ``litellm_provider: None`` when the + # provider is unknown (e.g. custom deployments registered via + # ``Router.add_deployment``). Persisting that None into + # ``litellm.model_cost`` causes ``_check_provider_match`` to drop + # custom pricing on subsequent cost lookups. + if existing_model.get("litellm_provider") is None: + existing_model.pop("litellm_provider", None) ## override / add new keys to the existing model cost dictionary updated_dictionary = _update_dictionary(existing_model, value) litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary) @@ -3343,6 +3350,21 @@ def get_optional_params_embeddings( # noqa: PLR0915 model=model, drop_params=drop_params if drop_params is not None else False, ) + # Provider-only params (e.g. Cohere input_type) are not in + # OPENAI_EMBEDDING_PARAMS, so embedding_pre_process drops them from + # non_default_params before map_openai_params. Restore only those extras + # from passed_params — skip OPENAI_EMBEDDING_PARAMS to avoid duplicating + # values already mapped (e.g. dimensions -> output_dimension). + if supported_params: + for param in supported_params: + if param in OPENAI_EMBEDDING_PARAMS: + continue + if ( + param in passed_params + and passed_params[param] is not None + and param not in optional_params + ): + optional_params[param] = passed_params[param] ## raise exception if non-default value passed for non-openai/azure embedding calls elif custom_llm_provider == "openai": # 'dimensions` is only supported in `text-embedding-3` and later models @@ -4019,16 +4041,23 @@ def get_optional_params( # noqa: PLR0915 thinking: Optional[AnthropicThinkingParam] = None, web_search_options: Optional[OpenAIWebSearchOptions] = None, safety_identifier: Optional[str] = None, + base_model: Optional[str] = None, **kwargs, ): passed_params = locals().copy() special_params = passed_params.pop("kwargs") + # Remove base_model from passed_params so it doesn't interfere with + # non_default_params / _check_valid_arg — it's a routing hint, not an + # OpenAI param. + passed_params.pop("base_model", None) provider_config: Optional[BaseConfig] = None if custom_llm_provider is not None and custom_llm_provider in [ provider.value for provider in LlmProviders ]: provider_config = ProviderConfigManager.get_provider_chat_config( - model=model, provider=LlmProviders(custom_llm_provider) + model=model, + provider=LlmProviders(custom_llm_provider), + base_model=base_model, ) non_default_params = pre_process_non_default_params( passed_params=passed_params, @@ -4091,7 +4120,7 @@ def get_optional_params( # noqa: PLR0915 sys.modules[__name__], "get_supported_openai_params" ) supported_params = get_supported_openai_params( - model=model, custom_llm_provider=custom_llm_provider + model=model, custom_llm_provider=custom_llm_provider, base_model=base_model ) if supported_params is None: supported_params = get_supported_openai_params( @@ -4702,22 +4731,27 @@ def get_optional_params( # noqa: PLR0915 ), ) elif custom_llm_provider == "azure": - if litellm.AzureOpenAIO1Config().is_o_series_model(model=model): + _azure_detection_model = base_model or model + if litellm.AzureOpenAIO1Config().is_o_series_model( + model=_azure_detection_model + ): optional_params = litellm.AzureOpenAIO1Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, - model=model, + model=_azure_detection_model, drop_params=( drop_params if drop_params is not None and isinstance(drop_params, bool) else False ), ) - elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=model): + elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model( + model=_azure_detection_model + ): optional_params = litellm.AzureOpenAIGPT5Config().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, - model=model, + model=_azure_detection_model, drop_params=( drop_params if drop_params is not None and isinstance(drop_params, bool) @@ -4739,7 +4773,7 @@ def get_optional_params( # noqa: PLR0915 optional_params = litellm.AzureOpenAIConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, - model=model, + model=_azure_detection_model, api_version=api_version, # type: ignore drop_params=( drop_params @@ -5387,6 +5421,16 @@ def _strip_model_name(model: str, custom_llm_provider: Optional[str]) -> str: # Global case-insensitive lookup map for model_cost (built eagerly at module import) _model_cost_lowercase_map: Optional[Dict[str, str]] = None +# Monotonic counter bumped on every model_cost mutation. Consumers that +# memoize derived state (e.g. provider-specific indices) can include this +# value in their cache key so they invalidate even when key add+remove or +# in-place value replacement leaves len/id unchanged. +_model_cost_mutation_generation: int = 0 + + +def get_model_cost_mutation_generation() -> int: + return _model_cost_mutation_generation + def _invalidate_model_cost_lowercase_map() -> None: """Invalidate the case-insensitive lookup map for model_cost. @@ -5394,8 +5438,9 @@ def _invalidate_model_cost_lowercase_map() -> None: Call this whenever litellm.model_cost is modified to ensure the map is rebuilt. Also clears related LRU caches that depend on model_cost data. """ - global _model_cost_lowercase_map + global _model_cost_lowercase_map, _model_cost_mutation_generation _model_cost_lowercase_map = None + _model_cost_mutation_generation += 1 # Clear LRU caches that depend on model_cost data get_model_info.cache_clear() @@ -5499,9 +5544,15 @@ def _get_model_info_from_model_cost(key: str) -> dict: def _check_provider_match(model_info: dict, custom_llm_provider: Optional[str]) -> bool: """ Check if the model info provider matches the custom provider. + + A missing ``litellm_provider`` key and a ``litellm_provider`` set to + ``None`` both mean "no specific provider constraint" and are treated + as a wildcard match. ``register_model`` may persist ``None`` here via + ``get_model_info`` when a deployment is registered without a provider, + so normalising the two cases keeps custom pricing applied consistently. """ if custom_llm_provider and ( - "litellm_provider" in model_info + model_info.get("litellm_provider") is not None and model_info["litellm_provider"] != custom_llm_provider ): if custom_llm_provider == "vertex_ai" and model_info[ @@ -5986,6 +6037,7 @@ def _get_model_info_helper( # noqa: PLR0915 tpm=_model_info.get("tpm", None), rpm=_model_info.get("rpm", None), ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None), + ocr_cost_per_credit=_model_info.get("ocr_cost_per_credit", None), annotation_cost_per_page=_model_info.get( "annotation_cost_per_page", None ), @@ -8112,10 +8164,8 @@ class ProviderConfigManager: # Format: (factory_function, needs_model_parameter: bool) LlmProviders.OPENAI: (lambda: litellm.OpenAIGPTConfig(), False), LlmProviders.ANTHROPIC: (lambda: litellm.AnthropicConfig(), False), - LlmProviders.AZURE: ( - lambda model: ProviderConfigManager._get_azure_config(model), - True, - ), + # AZURE is handled as a special case in get_provider_chat_config() + # so that base_model can be threaded through for model-type detection. LlmProviders.AZURE_AI: ( lambda model: ProviderConfigManager._get_azure_ai_config(model), True, @@ -8255,11 +8305,19 @@ class ProviderConfigManager: } @staticmethod - def _get_azure_config(model: str) -> BaseConfig: - """Get Azure config based on model type.""" - if litellm.AzureOpenAIO1Config().is_o_series_model(model=model): + def _get_azure_config(model: str, base_model: Optional[str] = None) -> BaseConfig: + """Get Azure config based on model type. + + When *base_model* is provided (e.g. ``"azure/gpt-5.2"``), it is used + for model-type detection instead of *model* (the deployment name). + This allows non-standard deployment names like ``"azure/foo"`` to be + routed through the correct config when the user specifies the true + underlying model via ``base_model``. + """ + detection_model = base_model or model + if litellm.AzureOpenAIO1Config().is_o_series_model(model=detection_model): return litellm.AzureOpenAIO1Config() - if litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=model): + if litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=detection_model): return litellm.AzureOpenAIGPT5Config() return litellm.AzureOpenAIConfig() @@ -8317,13 +8375,18 @@ class ProviderConfigManager: @staticmethod def get_provider_chat_config( # noqa: PLR0915 - model: str, provider: LlmProviders + model: str, + provider: LlmProviders, + base_model: Optional[str] = None, ) -> Optional[BaseConfig]: """ Returns the provider config for a given provider. Uses O(1) dictionary lookup for fast provider resolution. Python classes take priority over JSON (they have custom overrides). + + For Azure, *base_model* (when set) drives model-type detection so that + non-standard deployment names still route to the correct config. """ # Handle OpenAI special cases (O-series and GPT-5 models) if provider == LlmProviders.OPENAI: @@ -8332,6 +8395,12 @@ class ProviderConfigManager: if litellm.OpenAIGPT5Config.is_model_gpt_5_model(model=model): return litellm.OpenAIGPT5Config() + # Handle Azure before the generic map so base_model can be threaded through + if provider == LlmProviders.AZURE: + return ProviderConfigManager._get_azure_config( + model=model, base_model=base_model + ) + # Initialize provider config map lazily (avoids circular imports) if ProviderConfigManager._PROVIDER_CONFIG_MAP is None: ProviderConfigManager._PROVIDER_CONFIG_MAP = ( @@ -9241,6 +9310,18 @@ class ProviderConfigManager: return get_vertex_ai_ocr_config(model=model) + if provider == litellm.LlmProviders.REDUCTO: + from litellm.llms.reducto.ocr.transformation import ( + ReductoParseLegacyConfig, + ReductoParseV3Config, + ) + + if model == "parse-v3": + return ReductoParseV3Config() + if model == "parse-legacy": + return ReductoParseLegacyConfig() + return None + MistralOCRConfig = getattr(sys.modules[__name__], "MistralOCRConfig") PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, diff --git a/litellm/vector_store_files/utils.py b/litellm/vector_store_files/utils.py index ffe73516bda..1ee5b47e306 100644 --- a/litellm/vector_store_files/utils.py +++ b/litellm/vector_store_files/utils.py @@ -21,7 +21,7 @@ class VectorStoreFileRequestUtils: @staticmethod def get_create_request_params( - params: Dict[str, Any] + params: Dict[str, Any], ) -> VectorStoreFileCreateRequest: filtered = VectorStoreFileRequestUtils._filter_params( params=params, model=VectorStoreFileCreateRequest @@ -37,7 +37,7 @@ class VectorStoreFileRequestUtils: @staticmethod def get_update_request_params( - params: Dict[str, Any] + params: Dict[str, Any], ) -> VectorStoreFileUpdateRequest: filtered = VectorStoreFileRequestUtils._filter_params( params=params, model=VectorStoreFileUpdateRequest diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8582072f2d4..8aef1cbbda1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1011,6 +1011,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1041,6 +1042,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1071,6 +1073,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1100,6 +1103,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1129,6 +1133,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -1328,6 +1333,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "global.anthropic.claude-sonnet-4-6": { @@ -1358,6 +1364,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "us.anthropic.claude-sonnet-4-6": { @@ -1388,6 +1395,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "eu.anthropic.claude-sonnet-4-6": { @@ -1417,6 +1425,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "au.anthropic.claude-sonnet-4-6": { @@ -1446,6 +1455,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "jp.anthropic.claude-sonnet-4-6": { @@ -1475,6 +1485,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "anthropic.claude-sonnet-4-20250514-v1:0": { @@ -1996,6 +2007,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 159, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -2093,6 +2105,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 346, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "azure/computer-use-preview": { @@ -9647,6 +9660,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 346, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "claude-sonnet-4-5-20250929-v1:0": { @@ -9844,6 +9858,7 @@ "us": 1.1, "fast": 6.0 }, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -9879,7 +9894,8 @@ "fast": 6.0 }, "supports_max_reasoning_effort": true, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "supports_output_config": true }, "claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, @@ -9914,7 +9930,8 @@ "us": 1.1, "fast": 6.0 }, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "supports_output_config": true }, "claude-opus-4-7-20260416": { "cache_creation_input_token_cost": 6.25e-06, @@ -9949,7 +9966,8 @@ "us": 1.1, "fast": 6.0 }, - "supports_minimal_reasoning_effort": true + "supports_minimal_reasoning_effort": true, + "supports_output_config": true }, "claude-sonnet-4-20250514": { "deprecation_date": "2026-05-14", @@ -13986,6 +14004,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/accounts/fireworks/models/glm-5p1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 202800, + "max_output_tokens": 202800, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://fireworks.ai/models/fireworks/glm-5p1", + "supports_function_calling": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false + }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, "litellm_provider": "fireworks_ai", @@ -14252,6 +14285,21 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/glm-5p1": { + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 202800, + "max_output_tokens": 202800, + "max_tokens": 202800, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://fireworks.ai/models/fireworks/glm-5p1", + "supports_function_calling": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": false + }, "fireworks_ai/kimi-k2p5": { "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 6e-07, @@ -14961,6 +15009,73 @@ "web_search_billing_unit": "per_query", "supports_service_tier": true }, + "gemini-3.1-flash-lite": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_batches": 1.25e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "cache_read_input_token_cost_priority": 4.5e-08, + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_priority": 2.7e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -17025,6 +17140,75 @@ "web_search_billing_unit": "per_query", "supports_service_tier": true }, + "gemini/gemini-3.1-flash-lite": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_batches": 1.25e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "cache_read_input_token_cost_priority": 4.5e-08, + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_priority": 2.7e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "tpm": 250000, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "gemini/gemini-3-flash-preview": { "cache_read_input_token_cost": 5e-08, "input_cost_per_audio_token": 1e-06, @@ -24327,6 +24511,21 @@ "supports_tool_choice": true, "supports_vision": true }, + "mistral/ministral-8b-2512": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://mistral.ai/pricing", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "mistral/mistral-tiny": { "input_cost_per_token": 2.5e-07, "litellm_provider": "mistral", @@ -27171,6 +27370,58 @@ "supports_web_search": true, "tpm": 800000 }, + "openrouter/google/gemini-3.1-flash-lite": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "openrouter", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-lite", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "tpm": 800000 + }, "openrouter/google/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -27980,10 +28231,10 @@ "supports_tool_choice": true }, "openrouter/xiaomi/mimo-v2-flash": { - "input_cost_per_token": 9e-08, - "output_cost_per_token": 2.9e-07, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, "cache_creation_input_token_cost": 0.0, - "cache_read_input_token_cost": 0.0, + "cache_read_input_token_cost": 1e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -27993,7 +28244,43 @@ "supports_tool_choice": true, "supports_reasoning": true, "supports_vision": false, - "supports_prompt_caching": false + "supports_prompt_caching": true + }, + "openrouter/xiaomi/mimo-v2.5-pro": { + "input_cost_per_token": 1e-06, + "output_cost_per_token": 3e-06, + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 2e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_vision": false, + "supports_response_schema": true, + "supports_prompt_caching": true + }, + "openrouter/xiaomi/mimo-v2.5": { + "input_cost_per_token": 4e-07, + "output_cost_per_token": 2e-06, + "cache_creation_input_token_cost": 0.0, + "cache_read_input_token_cost": 8e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true, + "supports_response_schema": true, + "supports_prompt_caching": true }, "openrouter/z-ai/glm-4.7": { "input_cost_per_token": 4e-07, @@ -28724,14 +29011,16 @@ "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "supports_output_config": true }, "perplexity/anthropic/claude-opus-4-7": { "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, "supports_reasoning": false, - "supports_function_calling": true + "supports_function_calling": true, + "supports_output_config": true }, "perplexity/anthropic/claude-opus-4-5": { "litellm_provider": "perplexity", @@ -28945,6 +29234,24 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "reducto/parse-legacy": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.015, + "source": "https://reducto.ai/pricing", + "supported_endpoints": [ + "/v1/ocr" + ] + }, + "reducto/parse-v3": { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.015, + "source": "https://reducto.ai/pricing", + "supported_endpoints": [ + "/v1/ocr" + ] + }, "recraft/recraftv2": { "litellm_provider": "recraft", "mode": "image_generation", @@ -33124,6 +33431,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 346, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -33152,6 +33460,7 @@ "supports_tool_choice": true, "supports_vision": true, "tool_use_system_prompt_tokens": 346, + "supports_output_config": true, "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, @@ -33265,6 +33574,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "vertex_ai/claude-sonnet-4-5@20250929": { @@ -33647,6 +33957,73 @@ }, "web_search_billing_unit": "per_query" }, + "vertex_ai/gemini-3.1-flash-lite": { + "cache_read_input_token_cost": 2.5e-08, + "cache_read_input_token_cost_batches": 1.25e-08, + "cache_read_input_token_cost_flex": 1.25e-08, + "cache_read_input_token_cost_per_audio_token": 5e-08, + "cache_read_input_token_cost_priority": 4.5e-08, + "input_cost_per_audio_token": 5e-07, + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "input_cost_per_token_flex": 1.25e-07, + "input_cost_per_token_priority": 4.5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_pdf_size_mb": 30, + "max_tokens": 65536, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_reasoning_token": 1.5e-06, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "output_cost_per_token_flex": 7.5e-07, + "output_cost_per_token_priority": 2.7e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_audio_output": false, + "supports_code_execution": true, + "supports_file_search": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "supports_native_streaming": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "vertex_ai/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -40319,6 +40696,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, + "supports_output_config": true, "supports_minimal_reasoning_effort": true }, "duckduckgo/search": { diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 1d577213a1b..388752b032e 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1904,6 +1904,23 @@ "rerank": false } }, + "reducto": { + "display_name": "Reducto (`reducto`)", + "url": "https://docs.litellm.ai/docs/providers/reducto", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "ocr": true + } + }, "replicate": { "display_name": "Replicate (`replicate`)", "url": "https://docs.litellm.ai/docs/providers/replicate", diff --git a/pyproject.toml b/pyproject.toml index 70681c4ed6c..8dedca241ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,12 +33,14 @@ Homepage = "https://litellm.ai" Repository = "https://github.com/BerriAI/litellm" Documentation = "https://docs.litellm.ai" -# Dependencies pinned from the published `litellm[proxy]==1.83.0` resolution. -# Docker and CI should prefer `uv.lock` rather than maintaining parallel installers. +# Optional extras retain exact pins because they are consumed by Docker images +# where exact reproducibility matters. The core SDK uses ranges so downstream +# consumers can coexist with other packages without forced downgrades. [project.optional-dependencies] proxy = [ "gunicorn==23.0.0", "uvicorn==0.33.0", + "granian==2.5.7", "uvloop==0.21.0; sys_platform != 'win32'", "fastapi==0.124.4", "backoff==2.2.1", @@ -131,7 +133,7 @@ litellm-proxy = "litellm.proxy.client.cli:cli" dev = [ "diff-cover==9.7.2", "flake8==7.3.0", - "black==24.10.0", + "black==26.3.1", "mypy==1.19.0", "pytest==9.0.3", "pytest-mock==3.15.1", @@ -287,6 +289,7 @@ paths_to_mutate = [ ] tests_dir = [ "tests/test_litellm/proxy/management_endpoints/", + "tests/proxy_behavior/management/", ] also_copy = [ "litellm/", @@ -312,3 +315,4 @@ pytest_add_cli_args = [ [tool.coverage.run] source = ["litellm"] relative_files = true + diff --git a/tests/code_coverage_tests/check_licenses.py b/tests/code_coverage_tests/check_licenses.py index 668aefa8024..5fb2b495c24 100644 --- a/tests/code_coverage_tests/check_licenses.py +++ b/tests/code_coverage_tests/check_licenses.py @@ -31,6 +31,13 @@ DEFAULT_TRANSITIVE_PIN_PACKAGES = ( "wheel", ) +# SPDX license expressions (PEP 639 "License-Expression") join identifiers with +# the uppercase operators OR / AND / WITH. The split is case-sensitive: the +# lowercase "-or-later" inside an identifier such as "GPL-2.0-or-later" is part +# of the identifier, not an operator. +_SPDX_OPERATOR_SPLIT = re.compile(r"\s+(?:OR|AND)\s+") +_SPDX_WITH_SUFFIX = re.compile(r"\s+WITH\s+.*", re.DOTALL) + @dataclass class PackageLicense: @@ -109,21 +116,86 @@ class LicenseChecker: def get_package_license_from_pypi( self, package_name: str, version: str ) -> Optional[str]: - """Fetch license information for a package from PyPI.""" + """Fetch license information for a package from PyPI. + + Prefers the PEP 639 SPDX expression (``info.license_expression``), + falls back to the legacy free-text ``info.license`` field, and as a + last resort derives the license from the ``License :: OSI Approved :: + ...`` trove classifiers. + """ try: url = f"https://pypi.org/pypi/{package_name}/{version}/json" response = requests.get(url, timeout=10) response.raise_for_status() - data = response.json() - return data.get("info", {}).get("license") + info = response.json().get("info", {}) or {} + return ( + info.get("license_expression") + or info.get("license") + or self._license_from_classifiers(info.get("classifiers") or []) + ) except Exception as e: print( f"Warning: Failed to fetch license for {package_name} {version}: {str(e)}" ) return None - def is_license_acceptable(self, license_str: str) -> Tuple[bool, str]: - """Check if a license is acceptable based on configured lists.""" + @staticmethod + def _license_from_classifiers(classifiers: List[str]) -> Optional[str]: + """Derive a license name from the ``License :: OSI Approved :: ...`` trove classifiers.""" + prefix = "License :: OSI Approved :: " + for classifier in classifiers: + if classifier.startswith(prefix): + license_name = classifier[len(prefix) :].strip() + if license_name: + return license_name + return None + + @staticmethod + def _split_spdx_expression(license_str: str) -> Optional[List[str]]: + """Split an SPDX license expression into its component identifiers. + + Returns ``None`` when the string is not a recognizable SPDX expression + (for example a free-text license blob), so callers fall back to + whole-string matching. + """ + if "OR" not in license_str and "AND" not in license_str: + return None + + components: List[str] = [] + normalized = license_str.replace("(", " ").replace(")", " ") + for part in _SPDX_OPERATOR_SPLIT.split(normalized): + # Drop any "WITH " suffix: the exception qualifies the + # preceding license, it is not itself a license to authorize. + identifier = _SPDX_WITH_SUFFIX.sub("", part).strip() + if not identifier: + continue + # SPDX short-form identifiers are single whitespace-free tokens; a + # component with internal whitespace means this is free text. + if any(char.isspace() for char in identifier): + return None + components.append(identifier) + + return components if len(components) > 1 else None + + def is_license_acceptable(self, license_str: Optional[str]) -> Tuple[bool, str]: + """Check if a license (or compound SPDX expression) is acceptable.""" + if not license_str: + return False, "Unknown license" + + components = self._split_spdx_expression(license_str) + if components is None: + return self._is_single_license_acceptable(license_str) + + # Compound SPDX expression: conservatively require every component to + # be acceptable on its own (the safe direction for a CI gate). + for component in components: + is_acceptable, reason = self._is_single_license_acceptable(component) + if not is_acceptable: + return False, f"{reason} (in SPDX expression '{license_str}')" + return True, f"All SPDX components authorized: {', '.join(components)}" + + def _is_single_license_acceptable(self, license_str: str) -> Tuple[bool, str]: + """Check if a single license identifier is acceptable based on configured lists.""" if not license_str: return False, "Unknown license" diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index 3d53ecede7b..5a09403c570 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -90,7 +90,6 @@ jinja2: >=3.1.4 # BSD 3-Clause License litellm-proxy-extras: >=0.1.1 # MIT License litellm-enterprise: >=0.1.1 # LiteLLM Enterprise License a2a-sdk: >=0.3.22 # Apache 2.0 license -pydantic-settings: >=2.14.1 # MIT License manually verified (uses PEP 639 License-Expression: MIT, not the legacy License field, so liccheck reports it as unknown) anyio: >=4.5.0 # Unknown license httpx-aiohttp: >=0.1.4 # Unknown license backoff: >=2.2.1 # Unknown license diff --git a/tests/llm_responses_api_testing/test_anthropic_responses_api.py b/tests/llm_responses_api_testing/test_anthropic_responses_api.py index 6537f67acb9..68ff22e8938 100644 --- a/tests/llm_responses_api_testing/test_anthropic_responses_api.py +++ b/tests/llm_responses_api_testing/test_anthropic_responses_api.py @@ -3,7 +3,7 @@ import sys import pytest import asyncio from typing import Optional -from unittest.mock import patch, AsyncMock +from unittest.mock import patch, AsyncMock, MagicMock from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) @@ -130,6 +130,26 @@ def test_multiturn_tool_calls(): print("follow_up_response=", follow_up_response) +def test_response_api_handler_merges_metadata_and_service_tier_without_error(): + """Sync path must merge kwargs like async; double-splat raises TypeError.""" + handler = LiteLLMCompletionTransformationHandler() + + with patch("litellm.completion", new_callable=MagicMock) as mock_completion: + mock_completion.return_value = ModelResponse( + id="id", created=0, model="test", object="chat.completion", choices=[] + ) + handler.response_api_handler( + model="test", + input="hi", + responses_api_request={}, + metadata={"trace": "abc"}, + service_tier="auto", + ) + assert mock_completion.call_count == 1 + assert mock_completion.call_args.kwargs["metadata"] == {"trace": "abc"} + assert mock_completion.call_args.kwargs["service_tier"] == "auto" + + @pytest.mark.asyncio async def test_async_response_api_handler_merges_trace_id_without_error(): handler = LiteLLMCompletionTransformationHandler() @@ -158,3 +178,39 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): assert ( mock_acompletion.call_args.kwargs["litellm_trace_id"] == "session-trace" ) + + +@pytest.mark.asyncio +async def test_aresponses_forwards_timeout_to_acompletion(): + """Regression test: timeout passed to aresponses() must reach acompletion() + on the completion transformation path (Anthropic, Bedrock, Vertex etc.). + + Previously, `timeout` was a named param of `responses()` but was NOT + forwarded to `litellm_completion_transformation_handler.response_api_handler`, + so it was silently dropped — `Router(timeout=N)` was a no-op for Anthropic + and similar providers, with calls falling back to the provider SDK default + (~600s for Anthropic). + """ + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: + mock_acompletion.return_value = ModelResponse( + id="id", + created=0, + model="anthropic/claude-sonnet-4-5", + object="chat.completion", + choices=[], + ) + + await litellm.aresponses( + model="anthropic/claude-sonnet-4-5", + input="hello", + timeout=42, + api_key="sk-ant-fake", + ) + + assert mock_acompletion.call_count == 1 + forwarded_timeout = mock_acompletion.call_args.kwargs.get("timeout") + assert forwarded_timeout == 42, ( + f"timeout was not forwarded to acompletion (got {forwarded_timeout!r}); " + "this means Router(timeout=N) silently fails for providers on the " + "completion transformation path." + ) diff --git a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py index e2c50810cc2..37fcc602d37 100644 --- a/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py +++ b/tests/llm_responses_api_testing/test_base_responses_api_streaming_iterator.py @@ -41,6 +41,62 @@ from litellm.types.llms.openai import ( class TestBaseResponsesAPIStreamingIterator: """Test cases for BaseResponsesAPIStreamingIterator""" + @pytest.mark.asyncio + async def test_responses_streaming_iterator_parses_u2028_in_sse_json(self): + """ + U+2028 inside JSON must not split the SSE event. httpx aiter_lines uses + str.splitlines() and drops response.completed; OpenAI SSEDecoder does not. + """ + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + u2028 = "\u2028" + payload = json.dumps( + { + "type": "response.completed", + "response": {"instructions": f"eligible{u2028}promo"}, + } + ) + sse_bytes = f"data: {payload}\n\n".encode("utf-8") + + async def mock_aiter_bytes(): + yield sse_bytes + + mock_response = Mock() + mock_response.headers = {} + mock_response.aiter_bytes = mock_aiter_bytes + + mock_logging_obj = Mock(spec=LiteLLMLoggingObj) + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_config = Mock(spec=BaseResponsesAPIConfig) + + mock_responses_api_response = Mock(spec=ResponsesAPIResponse) + mock_responses_api_response.id = "resp_u2028" + mock_completed_event = Mock(spec=ResponseCompletedEvent) + mock_completed_event.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED + mock_completed_event.response = mock_responses_api_response + mock_config.transform_streaming_response.return_value = mock_completed_event + + iterator = ResponsesAPIStreamingIterator( + response=mock_response, + model="gpt-5.5", + responses_api_provider_config=mock_config, + logging_obj=mock_logging_obj, + litellm_metadata={"model_info": {"id": "model_123"}}, + custom_llm_provider="openai", + ) + + chunks = [] + with ( + patch("asyncio.create_task"), + patch("litellm.responses.streaming_iterator.executor"), + ): + async for chunk in iterator: + chunks.append(chunk) + + assert len(chunks) == 1 + assert chunks[0].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert iterator.completed_response is not None + def test_process_chunk_with_response_completed_event(self): """ Test that _process_chunk correctly processes a ResponseCompletedEvent @@ -270,7 +326,7 @@ class TestBaseResponsesAPIStreamingIterator: # Mock dependencies mock_response = Mock() mock_response.headers = {} - mock_response.aiter_lines = Mock() + mock_response.aiter_bytes = Mock() mock_logging_obj = Mock(spec=LiteLLMLoggingObj) mock_logging_obj.model_call_details = {"litellm_params": {}} mock_logging_obj.async_success_handler = Mock() @@ -334,12 +390,10 @@ class TestBaseResponsesAPIStreamingIterator: mock_response = Mock() mock_response.headers = {} - # Create an async iterator that raises StopAsyncIteration after yielding one chunk - async def mock_aiter_lines(): - yield 'data: {"type": "response.output_text.delta", "delta": "test"}' - # Normal end of stream - raise StopAsyncIteration + async def mock_aiter_bytes(): + yield b'data: {"type": "response.output_text.delta", "delta": "test"}\n\n' - mock_response.aiter_lines = mock_aiter_lines + mock_response.aiter_bytes = mock_aiter_bytes mock_logging_obj = Mock(spec=LiteLLMLoggingObj) mock_logging_obj.model_call_details = {"litellm_params": {}} @@ -396,12 +450,10 @@ class TestBaseResponsesAPIStreamingIterator: mock_response = Mock() mock_response.headers = {} - # Create a sync iterator that raises StopIteration after yielding one chunk - def mock_iter_lines(): - yield 'data: {"type": "response.output_text.delta", "delta": "test"}' - # Normal end of stream - raise StopIteration + def mock_iter_bytes(): + yield b'data: {"type": "response.output_text.delta", "delta": "test"}\n\n' - mock_response.iter_lines = mock_iter_lines + mock_response.iter_bytes = mock_iter_bytes mock_logging_obj = Mock(spec=LiteLLMLoggingObj) mock_logging_obj.model_call_details = {"litellm_params": {}} @@ -450,7 +502,7 @@ class TestBaseResponsesAPIStreamingIterator: mock_response = Mock() mock_response.headers = {} - mock_response.aiter_lines = Mock() + mock_response.aiter_bytes = Mock() mock_logging_obj = Mock(spec=LiteLLMLoggingObj) mock_logging_obj.model_call_details = {"litellm_params": {}} mock_logging_obj.async_failure_handler = Mock() @@ -532,7 +584,7 @@ class TestBaseResponsesAPIStreamingIterator: mock_response = Mock() mock_response.headers = {} - mock_response.aiter_lines = Mock() + mock_response.aiter_bytes = Mock() mock_logging_obj = Mock(spec=LiteLLMLoggingObj) mock_logging_obj.model_call_details = {"litellm_params": {}} mock_logging_obj.async_failure_handler = Mock() diff --git a/tests/llm_translation/realtime/base_realtime_tests.py b/tests/llm_translation/realtime/base_realtime_tests.py index 1d55f13b00d..1a2c6ff6a9c 100644 --- a/tests/llm_translation/realtime/base_realtime_tests.py +++ b/tests/llm_translation/realtime/base_realtime_tests.py @@ -10,7 +10,7 @@ import json import os import sys from abc import ABC, abstractmethod -from typing import Optional +from typing import Optional, Tuple, Union import pytest import websockets @@ -79,7 +79,7 @@ class RealTimeWebSocketClient: def _is_initial_event(self, msg_type: str) -> bool: """Check if message type is an initial connection event""" - # OpenAI sends "session.created", xAI sends "conversation.created" + # OpenAI and xAI send "session.created"; some providers send "conversation.created" return msg_type in ["session.created", "conversation.created"] async def receive_text(self): @@ -153,8 +153,14 @@ class BaseRealtimeTest(ABC): pass @abstractmethod - def get_initial_event_type(self) -> str: - """Return the expected initial event type (e.g., 'session.created' or 'conversation.created')""" + def get_initial_event_type(self) -> Union[str, Tuple[str, ...]]: + """Return the expected initial event type(s). + + May return a single event type (e.g. ``'session.created'``) or a tuple + of acceptable types when the upstream provider can legitimately emit + more than one initial event (e.g. xAI's Grok Voice Agent has shipped + both ``conversation.created`` and ``session.created``). + """ pass def get_skip_reason(self) -> str: @@ -229,9 +235,14 @@ class BaseRealtimeTest(ABC): # Verify initial event initial_event = websocket_client.messages_received[0] + expected_event_type = self.get_initial_event_type() + if isinstance(expected_event_type, str): + allowed_event_types: Tuple[str, ...] = (expected_event_type,) + else: + allowed_event_types = tuple(expected_event_type) assert ( - initial_event["type"] == self.get_initial_event_type() - ), f"Expected {self.get_initial_event_type()}, got {initial_event.get('type')}" + initial_event["type"] in allowed_event_types + ), f"Expected one of {allowed_event_types}, got {initial_event.get('type')}" @pytest.mark.asyncio async def test_realtime_with_query_params(self): diff --git a/tests/llm_translation/realtime/test_xai_realtime.py b/tests/llm_translation/realtime/test_xai_realtime.py index 0bb7a59bb1a..8ffcb3db30d 100644 --- a/tests/llm_translation/realtime/test_xai_realtime.py +++ b/tests/llm_translation/realtime/test_xai_realtime.py @@ -7,6 +7,7 @@ Uses the base test class to ensure consistent behavior across providers. import os import sys +from typing import Tuple import pytest @@ -19,10 +20,12 @@ class TestXAIRealtime(BaseRealtimeTest): """ E2E tests for xAI Realtime API. - xAI's Grok Voice Agent API is OpenAI-compatible but uses: - - Different initial event: "conversation.created" instead of "session.created" - - Different endpoint: wss://api.x.ai/v1/realtime + xAI's Grok Voice Agent API is OpenAI-compatible: + - Endpoint: wss://api.x.ai/v1/realtime - Model: grok-4-1-fast-non-reasoning + - Initial event: historically "conversation.created"; xAI has since shipped + "session.created" (matching OpenAI). Accept either to avoid spurious + failures whenever xAI flips the wire format. """ def get_model(self) -> str: @@ -31,5 +34,5 @@ class TestXAIRealtime(BaseRealtimeTest): def get_api_key_env_var(self) -> str: return "XAI_API_KEY" - def get_initial_event_type(self) -> str: - return "conversation.created" + def get_initial_event_type(self) -> Tuple[str, ...]: + return ("conversation.created", "session.created") diff --git a/tests/local_testing/conftest.py b/tests/local_testing/conftest.py index 6a746041f15..acb79a7577d 100644 --- a/tests/local_testing/conftest.py +++ b/tests/local_testing/conftest.py @@ -22,6 +22,18 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm +# ``litellm.model_cost`` is loaded at import time from the URL pinned to +# ``main`` (``LITELLM_MODEL_COST_MAP_URL``). The in-tree backup ships with +# this branch and can include pricing entries that main has not yet picked +# up (e.g. an upstream provider rotates a model id and the test cassette +# records the new name). Backfill any entries that are missing from the +# remote-fetched map so cost-calculator lookups in tests succeed against +# the cassette state the branch is being tested with. +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + +for _k, _v in GetModelCostMap.load_local_model_cost_map().items(): + litellm.model_cost.setdefault(_k, _v) + from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, _pin_multipart_boundary, diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 9782bf3c2af..2382b8a5197 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -4223,7 +4223,13 @@ def test_gemini_google_maps_tool_simple(): ) print(f"Response: {response.model_dump_json(indent=4)}") assert response.choices[0].message.content is not None - except litellm.RateLimitError: + except (litellm.RateLimitError, litellm.InternalServerError): + # Transient Vertex-side failures (rate limiting, 500 INTERNAL from the + # Google Maps grounding backend) are not LiteLLM bugs — don't fail CI. pass + except litellm.InternalServerError: + pytest.skip( + "Google Maps Platform returned a transient 500 (upstream flake); skipping." + ) except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 409f4fad99a..c20fb09eeba 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -382,6 +382,11 @@ async def test_mcp_http_transport_tool_not_found(): } ) + # Mapping populated for this server but not for the requested tool + test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = ( + "test_http_server" + ) + # Try to call a tool that doesn't exist in mapping with pytest.raises(ValueError, match="Tool nonexistent_tool not found"): await test_manager.call_tool( @@ -881,6 +886,7 @@ async def test_get_tools_from_mcp_servers(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): if server.server_id == "server1_id": return [mock_tool_1] @@ -1764,6 +1770,26 @@ def test_get_server_auth_header_fallback_to_default(): assert result == "Bearer default_token" +def test_get_server_auth_header_hyphenated_alias_sanitized_header_key(): + """Header keys use sanitized alias; lookup must match legacy hyphenated aliases.""" + from litellm.proxy._experimental.mcp_server.rest_endpoints import ( + _get_server_auth_header, + ) + + mock_server = MagicMock() + mock_server.alias = "GitHub-MCP" + mock_server.server_name = "github_mcp_server" + + mcp_server_auth_headers = { + "github_mcp": {"Authorization": "Bearer github-mcp-token"}, + } + + result = _get_server_auth_header( + mock_server, mcp_server_auth_headers, "Bearer default_token" + ) + assert result == {"Authorization": "Bearer github-mcp-token"} + + def test_get_server_auth_header_no_auth_headers(): """Test _get_server_auth_header function with no auth headers.""" from litellm.proxy._experimental.mcp_server.rest_endpoints import ( @@ -1856,6 +1882,7 @@ async def test_get_tools_for_single_server(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ) # Verify the result diff --git a/tests/pass_through_tests/package-lock.json b/tests/pass_through_tests/package-lock.json index 8aa33340b16..2f8e7fe21b2 100644 --- a/tests/pass_through_tests/package-lock.json +++ b/tests/pass_through_tests/package-lock.json @@ -951,13 +951,12 @@ "license": "BSD-3-Clause" }, "node_modules/@protobufjs/fetch": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/@protobufjs/fetch/-/fetch-1.1.0.tgz", - "integrity": "sha512-lljVXpqXebpsijW71PZaCYeIcE5on1w5DlQy5WH6GLbFryLUrBD4932W/E2BSpfRJWseIL4v/KPgBFxDOIdKpQ==", + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@protobufjs/fetch/-/fetch-1.1.1.tgz", + "integrity": "sha512-GpptLrs57adMSuHi3VNj0mAF8dwh36LMaYF6XyJ6JMWlVsc+t42tm1HSEDmOs3A8fC9yyeisgLhsTVQokOZ0zw==", "license": "BSD-3-Clause", "dependencies": { - "@protobufjs/aspromise": "^1.1.1", - "@protobufjs/inquire": "^1.1.0" + "@protobufjs/aspromise": "^1.1.1" } }, "node_modules/@protobufjs/float": { @@ -967,9 +966,9 @@ "license": "BSD-3-Clause" }, "node_modules/@protobufjs/inquire": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/@protobufjs/inquire/-/inquire-1.1.1.tgz", - "integrity": "sha512-mnzgDV26ueAvk7rsbt9L7bE0SuAoqyuys/sMMrmVcN5x9VsxpcG3rqAUSgDyLp0UZlmNfIbQ4fHfCtreVBk8Ew==", + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/@protobufjs/inquire/-/inquire-1.1.2.tgz", + "integrity": "sha512-pa0vFRuws4wkvaXKK1uXZMAwAX4/t8ANaJo45iw/oQHNQ9q5xUzwgFmVJGXiga2BeN+zpX7Vf9vmsiIa2J+MUw==", "license": "BSD-3-Clause" }, "node_modules/@protobufjs/path": { @@ -3510,9 +3509,9 @@ } }, "node_modules/protobufjs": { - "version": "7.5.6", - "resolved": "https://registry.npmjs.org/protobufjs/-/protobufjs-7.5.6.tgz", - "integrity": "sha512-M71sTMB146U3u0di3yup8iM+zv8yPRNQVr1KK4tyBitl3qFvEGucq/rGDRShD2rsJhtN02RJaJ7j5X5hmy8SJg==", + "version": "7.6.0", + "resolved": "https://registry.npmjs.org/protobufjs/-/protobufjs-7.6.0.tgz", + "integrity": "sha512-LtESOsMPTZgyYtwxhvdgdjGL0HmXEaRA/hVD6sol4zA60hVXXXP/SGmxnqDbgGE8gy7pYex7cym+5vYPcmaXBQ==", "hasInstallScript": true, "license": "BSD-3-Clause", "dependencies": { @@ -3520,14 +3519,14 @@ "@protobufjs/base64": "^1.1.2", "@protobufjs/codegen": "^2.0.5", "@protobufjs/eventemitter": "^1.1.0", - "@protobufjs/fetch": "^1.1.0", + "@protobufjs/fetch": "^1.1.1", "@protobufjs/float": "^1.0.2", - "@protobufjs/inquire": "^1.1.1", + "@protobufjs/inquire": "^1.1.2", "@protobufjs/path": "^1.1.2", "@protobufjs/pool": "^1.1.0", "@protobufjs/utf8": "^1.1.1", "@types/node": ">=13.7.0", - "long": "^5.0.0" + "long": "^5.3.2" }, "engines": { "node": ">=12.0.0" @@ -4035,9 +4034,9 @@ } }, "node_modules/ws": { - "version": "8.20.0", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.20.0.tgz", - "integrity": "sha512-sAt8BhgNbzCtgGbt2OxmpuryO63ZoDk/sqaB/znQm94T4fCEsy/yV+7CdC1kJhOU9lboAEU7R3kquuycDoibVA==", + "version": "8.20.1", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.20.1.tgz", + "integrity": "sha512-It4dO0K5v//JtTXuPkfEOaI3uUN87iYPnqo/ZzqCoG3g8uhA66QUMs/SrM0YK7/NAu+r4LMh/9dq2A7k+rHs+w==", "license": "MIT", "engines": { "node": ">=10.0.0" diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/login_to_ui.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/login_to_ui.spec.ts deleted file mode 100644 index e5a397a6a66..00000000000 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/login_to_ui.spec.ts +++ /dev/null @@ -1,51 +0,0 @@ -/* - -Login to Admin UI -Basic UI Test - -Click on all the tabs ensure nothing is broken -*/ - -import { test, expect } from "@playwright/test"; - -test("admin login test", async ({ page }) => { - // Go to the specified URL - await page.goto("http://localhost:4000/ui"); - await page.waitForLoadState("networkidle"); - - await page.screenshot({ path: "test-results/login_before.png" }); - - // Enter "admin" in the username input field - await page.fill('input[placeholder="Enter your username"]', "admin"); - - // Enter "gm" in the password input field - await page.fill('input[placeholder="Enter your password"]', "gm"); - - page.screenshot({ path: "test-results/login_after_inputs.png" }); - - // Optionally, you can add an assertion to verify the login button is enabled - const loginButton = page.getByRole("button", { name: "Login" }); - await expect(loginButton).toBeEnabled(); - - // Optionally, you can click the login button to submit the form - await loginButton.click(); - const tabs = [ - "Virtual Keys", - "Playground", - "Models", - "Usage", - "Teams", - "Internal User", - "Settings", - "Experimental", - "API Reference", - "AI Hub", - ]; - - for (const tab of tabs) { - const tabElement = page.locator("span.ant-menu-title-content", { - hasText: tab, - }); - await tabElement.click(); - } -}); diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/redirect-fail-screenshot.png b/tests/proxy_admin_ui_tests/e2e_ui_tests/redirect-fail-screenshot.png deleted file mode 100644 index b2e33251260..00000000000 Binary files a/tests/proxy_admin_ui_tests/e2e_ui_tests/redirect-fail-screenshot.png and /dev/null differ diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/require_auth_for_dashboard.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/require_auth_for_dashboard.spec.ts deleted file mode 100644 index 4e4bd2fcd93..00000000000 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/require_auth_for_dashboard.spec.ts +++ /dev/null @@ -1,37 +0,0 @@ -// tests/auth.spec.ts -import { test, expect } from "@playwright/test"; - -test.describe("Authentication Checks", () => { - test("should redirect unauthenticated user from a protected page", async ({ - page, - }) => { - test.setTimeout(30000); - - page.on("console", (msg) => console.log("PAGE LOG:", msg.text())); - - const protectedPageUrl = "http://localhost:4000/ui?page=llm-playground"; - const expectedRedirectUrl = "http://localhost:4000/ui/login/"; - - console.log( - `Attempting to navigate to protected page: ${protectedPageUrl}` - ); - - await page.goto(protectedPageUrl); - - console.log(`Navigation initiated. Current URL: ${page.url()}`); - - try { - await page.waitForURL(expectedRedirectUrl, { timeout: 10000 }); - console.log(`Waited for URL. Current URL is now: ${page.url()}`); - } catch (error) { - console.error( - `Timeout waiting for URL: ${expectedRedirectUrl}. Current URL: ${page.url()}` - ); - await page.screenshot({ path: "redirect-fail-screenshot.png" }); - throw error; - } - - await expect(page).toHaveURL(expectedRedirectUrl); - console.log(`Assertion passed: Page URL is ${expectedRedirectUrl}`); - }); -}); diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/search_users.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/search_users.spec.ts deleted file mode 100644 index d72c44ab8cc..00000000000 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/search_users.spec.ts +++ /dev/null @@ -1,222 +0,0 @@ -/* -Search Users in Admin UI -E2E Test for user search functionality - -Tests: -1. Navigate to Internal Users tab -2. Verify search input exists -3. Test search functionality -4. Verify results update -5. Test filtering by email, user ID, and SSO user ID -*/ - -import { test, expect } from "@playwright/test"; - -test("user search test", async ({ page }) => { - // Set a longer timeout for the entire test - test.setTimeout(60000); - - // Enable console logging - page.on("console", (msg) => console.log("PAGE LOG:", msg.text())); - - // Login first - await page.goto("http://localhost:4000/ui"); - await page.waitForLoadState("networkidle"); - console.log("Navigated to login page"); - - page.screenshot({ path: "test-results/search_users_before_login.png" }); - - // Wait for login form to be visible - await page.waitForSelector('input[placeholder="Enter your username"]', { - timeout: 10000, - }); - console.log("Login form is visible"); - - await page.fill('input[placeholder="Enter your username"]', "admin"); - await page.fill('input[placeholder="Enter your password"]', "gm"); - console.log("Filled login credentials"); - - const loginButton = page.getByRole("button", { name: "Login" }); - await expect(loginButton).toBeEnabled(); - await loginButton.click(); - console.log("Clicked login button"); - - // Wait for navigation to complete and dashboard to load - await page.waitForLoadState("networkidle"); - console.log("Page loaded after login"); - - // Take a screenshot for debugging - await page.screenshot({ path: "after-login.png" }); - console.log("Took screenshot after login"); - - // Try to find the Internal User tab with more debugging - console.log("Looking for Internal User tab..."); - const internalUserTab = page.locator("span.ant-menu-title-content", { - hasText: "Internal User", - }); - - // Wait for the tab to be visible - await internalUserTab.waitFor({ state: "visible", timeout: 10000 }); - console.log("Internal User tab is visible"); - - // Take another screenshot before clicking - await page.screenshot({ path: "before-tab-click.png" }); - console.log("Took screenshot before tab click"); - - await internalUserTab.click(); - console.log("Clicked Internal User tab"); - - // Wait for the page to load and table to be visible - await page.waitForSelector("tbody tr", { timeout: 30000 }); - await page.waitForTimeout(2000); // Additional wait for table to stabilize - console.log("Table is visible"); - - // Take a final screenshot - await page.screenshot({ path: "after-tab-click.png" }); - console.log("Took screenshot after tab click"); - - // Verify search input exists - const searchInput = page.locator('input[placeholder="Search by email..."]'); - await expect(searchInput).toBeVisible(); - console.log("Search input is visible"); - - // Test search functionality - const initialUserCount = await page.locator("tbody tr").count(); - console.log(`Initial user count: ${initialUserCount}`); - - // Perform a search - const testEmail = "test@"; - await searchInput.fill(testEmail); - console.log("Filled search input"); - - // Wait for the debounced search to complete - await page.waitForTimeout(500); - console.log("Waited for debounce"); - - // Wait for the results count to update - await page.waitForFunction((initialCount) => { - const currentCount = document.querySelectorAll("tbody tr").length; - return currentCount !== initialCount; - }, initialUserCount); - console.log("Results updated"); - - const filteredUserCount = await page.locator("tbody tr").count(); - console.log(`Filtered user count: ${filteredUserCount}`); - - expect(filteredUserCount).toBeDefined(); - - // Clear the search - await searchInput.clear(); - console.log("Cleared search"); - - await page.waitForTimeout(500); - console.log("Waited for debounce after clear"); - - await page.waitForFunction((initialCount) => { - const currentCount = document.querySelectorAll("tbody tr").length; - return currentCount === initialCount; - }, initialUserCount); - console.log("Results reset"); - - const resetUserCount = await page.locator("tbody tr").count(); - console.log(`Reset user count: ${resetUserCount}`); - - expect(resetUserCount).toBe(initialUserCount); -}); - -test("user filter test", async ({ page }) => { - // Set a longer timeout for the entire test - test.setTimeout(60000); - - // Enable console logging - page.on("console", (msg) => console.log("PAGE LOG:", msg.text())); - - // Login first - await page.goto("http://localhost:4000/ui"); - await page.waitForLoadState("networkidle"); - console.log("Navigated to login page"); - - // Wait for login form to be visible - await page.waitForSelector('input[placeholder="Enter your username"]', { - timeout: 10000, - }); - console.log("Login form is visible"); - - await page.fill('input[placeholder="Enter your username"]', "admin"); - await page.fill('input[placeholder="Enter your password"]', "gm"); - console.log("Filled login credentials"); - - const loginButton = page.getByRole("button", { name: "Login" }); - await expect(loginButton).toBeEnabled(); - await loginButton.click(); - console.log("Clicked login button"); - - // Wait for navigation to complete and dashboard to load - await page.waitForLoadState("networkidle"); - console.log("Page loaded after login"); - - // Navigate to Internal Users tab - const internalUserTab = page.locator("span.ant-menu-title-content", { - hasText: "Internal User", - }); - await internalUserTab.waitFor({ state: "visible", timeout: 10000 }); - await internalUserTab.click(); - console.log("Clicked Internal User tab"); - - // Wait for the page to load and table to be visible - await page.waitForSelector("tbody tr", { timeout: 30000 }); - await page.waitForTimeout(2000); // Additional wait for table to stabilize - console.log("Table is visible"); - - // Get initial user count - const initialUserCount = await page.locator("tbody tr").count(); - console.log(`Initial user count: ${initialUserCount}`); - - // Click the filter button to show additional filters - const filterButton = page.getByRole("button", { - name: "Filters", - exact: true, - }); - await filterButton.click(); - console.log("Clicked filter button"); - await page.waitForTimeout(500); // Wait for filters to appear - - // Test user ID filter - const userIdInput = page.locator('input[placeholder="Filter by User ID"]'); - await expect(userIdInput).toBeVisible(); - console.log("User ID filter is visible"); - - await userIdInput.fill("user"); - console.log("Filled user ID filter"); - await page.waitForTimeout(1000); - const userIdFilteredCount = await page.locator("tbody tr").count(); - console.log(`User ID filtered count: ${userIdFilteredCount}`); - expect(userIdFilteredCount).toBeLessThan(initialUserCount); - - // Clear user ID filter - await userIdInput.clear(); - await page.waitForTimeout(1000); - console.log("Cleared user ID filter"); - - // Test SSO user ID filter - const ssoUserIdInput = page.locator('input[placeholder="Filter by SSO ID"]'); - await expect(ssoUserIdInput).toBeVisible(); - console.log("SSO user ID filter is visible"); - - await ssoUserIdInput.fill("sso"); - console.log("Filled SSO user ID filter"); - await page.waitForTimeout(1000); - const ssoUserIdFilteredCount = await page.locator("tbody tr").count(); - console.log(`SSO user ID filtered count: ${ssoUserIdFilteredCount}`); - expect(ssoUserIdFilteredCount).toBeLessThan(initialUserCount); - - // Clear SSO user ID filter - await ssoUserIdInput.clear(); - await page.waitForTimeout(5000); - console.log("Cleared SSO user ID filter"); - - // Verify count returns to initial after clearing all filters - const finalUserCount = await page.locator("tbody tr").count(); - console.log(`Final user count: ${finalUserCount}`); - expect(finalUserCount).toBe(initialUserCount); -}); diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/team_admin.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/team_admin.spec.ts deleted file mode 100644 index a753c724b37..00000000000 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/team_admin.spec.ts +++ /dev/null @@ -1,250 +0,0 @@ -import { test, expect } from "@playwright/test"; -import { loginToUI } from "../utils/login"; - -// test.describe("Invite User, Set Password, and Login", () => { -// let testEmail: string; -// const testPassword = "Password123!"; // Define a password -// const teamName1 = `team-invite-test-1-${Date.now()}`; -// const teamName2 = `team-invite-test-2-${Date.now()}`; -// const keyName1 = `key-${teamName1}`; -// const keyName2 = `key-${teamName2}`; - -// test.beforeEach(async ({ page }) => { -// await loginToUI(page); // Login as admin first -// await page.goto("http://localhost:4000/ui?page=teams"); - -// // --- Create Team 1 --- -// await page.getByRole("button", { name: "+ Create New Team" }).click(); -// await page -// .getByLabel("Team Name") -// .waitFor({ state: "visible", timeout: 5000 }); // Wait for label -// await page.getByLabel("Team Name").click(); -// await page.getByLabel("Team Name").fill(teamName1); -// await page.getByRole("button", { name: "Create Team" }).click(); -// // Wait for the modal to close or for a success message if applicable -// await expect( -// page.locator(".ant-modal-wrap").filter({ hasText: "Create New Team" }) -// ).not.toBeVisible({ timeout: 10000 }); -// console.log(`Created Team 1: ${teamName1}`); - -// // --- Create Team 2 --- -// await page.getByRole("button", { name: "+ Create New Team" }).click(); -// await page -// .getByLabel("Team Name") -// .waitFor({ state: "visible", timeout: 5000 }); // Wait for label -// await page.getByLabel("Team Name").click(); -// await page.getByLabel("Team Name").fill(teamName2); -// await page.getByRole("button", { name: "Create Team" }).click(); -// // Wait for the modal to close or for a success message if applicable -// await expect( -// page.locator(".ant-modal-wrap").filter({ hasText: "Create New Team" }) -// ).not.toBeVisible({ timeout: 10000 }); -// console.log(`Created Team 2: ${teamName2}`); - -// // // Verify both teams are listed -// // await page.goto("http://localhost:4000/ui?page=teams"); // Refresh or ensure on teams page -// // await page.waitForTimeout(3000); -// await expect(page.getByText(teamName1)).toBeVisible({ timeout: 10000 }); -// await expect(page.getByText(teamName2)).toBeVisible({ timeout: 10000 }); - -// // --- Navigate to Keys Page --- -// await page.goto("http://localhost:4000/ui?page=api-keys"); -// await page.waitForTimeout(3000); -// await expect( -// page.getByRole("button", { name: "+ Create New Key" }) -// ).toBeVisible(); // Wait for page load - -// // --- Create Key for Team 1 --- -// await page.getByRole("button", { name: "+ Create New Key" }).click(); -// const createKeyModal1 = page -// .locator(".ant-modal-wrap") -// .filter({ hasText: "Key Ownership" }); -// await expect(createKeyModal1).toBeVisible(); - -// // Select Team 1 -// await createKeyModal1 -// .locator(".ant-select-selector >> input") -// .first() -// .click(); // Click to open team dropdown -// await createKeyModal1 -// .locator(".ant-select-selector >> input") -// .first() -// .fill(teamName1); - -// await page -// .locator(".ant-select-item-option") -// .filter({ hasText: teamName1 }) -// .first() -// .click(); // Click specific team name - -// // Enter Key Name 1 -// await page.fill('input[id="key_alias"]', keyName1); - -// // Click on models dropdown -// await page.locator("input#models").click(); -// await page.waitForSelector( -// '.ant-select-item-option[title="All Team Models"]' -// ); -// await page -// .locator('.ant-select-item-option[title="All Team Models"]') -// .click(); - -// // Click Create Key -// await createKeyModal1.getByRole("button", { name: "Create Key" }).click(); - -// // Close the Key Generated modal (which appears after successful creation) -// const keyGeneratedModal1 = page -// .locator(".ant-modal-wrap") -// .filter({ hasText: "Save your Key" }); -// await expect(keyGeneratedModal1).toBeVisible({ timeout: 10000 }); -// await keyGeneratedModal1.locator('button[aria-label="Close"]').click(); -// await expect(keyGeneratedModal1).not.toBeVisible(); // Wait for close -// console.log(`Created Key 1: ${keyName1} for Team: ${teamName1}`); - -// // --- Create Key for Team 2 --- -// await page.getByRole("button", { name: "+ Create New Key" }).click(); -// const createKeyModal2 = page -// .locator(".ant-modal-wrap") -// .filter({ hasText: "Key Ownership" }); -// await expect(createKeyModal2).toBeVisible(); - -// // Select Team 2 -// await createKeyModal2 -// .locator(".ant-select-selector >> input") -// .first() -// .click(); // Click to open team dropdown -// await page -// .locator(".ant-select-item-option") -// .filter({ hasText: teamName2 }) -// .click(); // Click specific team name - -// // Enter Key Name 2 -// await page.fill('input[id="key_alias"]', keyName2); - -// // Click on models dropdown -// await page.locator("input#models").click(); -// await page.waitForSelector( -// '.ant-select-item-option[title="All Team Models"]' -// ); -// await page -// .locator('.ant-select-item-option[title="All Team Models"]') -// .click(); - -// // Click Create Key -// await createKeyModal2.getByRole("button", { name: "Create Key" }).click(); - -// // Close the Key Generated modal -// const keyGeneratedModal2 = page -// .locator(".ant-modal-wrap") -// .filter({ hasText: "Save your Key" }); -// await expect(keyGeneratedModal2).toBeVisible({ timeout: 10000 }); -// await keyGeneratedModal2.locator('button[aria-label="Close"]').click(); -// await expect(keyGeneratedModal2).not.toBeVisible(); // Wait for close -// console.log(`Created Key 2: ${keyName2} for Team: ${teamName2}`); -// }); - -// test("Invite user, set password via link, and login", async ({ page }) => { -// // Navigate to Users page -// await page.goto("http://localhost:4000/ui?page=users"); - -// // Go to Internal User tab -// const internalUserTab = page.locator("span.ant-menu-title-content", { -// hasText: "Internal User", -// }); -// await internalUserTab.waitFor({ state: "visible", timeout: 10000 }); -// await internalUserTab.click(); - -// // --- Invite User Flow --- -// await page.getByRole("button", { name: "+ Invite User" }).click(); - -// // Wait for the invite user modal to be visible -// const inviteModal = page -// .locator(".ant-modal-wrap") -// .filter({ hasText: "Invite User" }); -// await expect(inviteModal).toBeVisible(); - -// testEmail = `test-${Date.now()}@litellm.ai`; // Use a unique email -// // Assuming the email input is the first one with 'base-input' test id inside the modal -// await inviteModal.getByTestId("base-input").first().fill(testEmail); - -// // Select Global Admin Role (or another appropriate role) -// const globalRoleLabel = inviteModal.getByLabel("Global Proxy Role"); -// await globalRoleLabel.click(); -// // Wait for the dropdown option to be visible before clicking -// const adminRoleOption = page.getByTitle("Admin (All Permissions)", { -// exact: true, -// }); -// await adminRoleOption.waitFor({ state: "visible", timeout: 5000 }); -// await adminRoleOption.click(); - -// // Select Team - Add explicit wait before clicking -// const teamIdLabel = inviteModal.getByLabel("Team ID"); -// // Wait for the label associated with the Team ID select to be visible -// await teamIdLabel.waitFor({ state: "visible", timeout: 10000 }); // Increased timeout for safety -// await teamIdLabel.click(); - -// // Wait for the team name option to be visible in the dropdown -// const teamNameOption = page.getByText(teamName1, { exact: true }); -// await teamNameOption.waitFor({ state: "visible", timeout: 5000 }); -// await teamNameOption.click(); - -// // Create User -// await inviteModal.getByRole("button", { name: "Create User" }).click(); - -// // --- Capture Invitation Link --- -// const invitationModal = page -// .locator(".ant-modal-wrap") -// .filter({ hasText: "Invitation Link" }); -// await expect(invitationModal).toBeVisible({ timeout: 15000 }); // Wait longer for modal - -// // Locate the text element containing the URL more reliably -// const invitationUrl = await page -// .locator("div.flex.justify-between.pt-5.pb-2") // find the correct div -// .filter({ hasText: "Invitation Link" }) // find the div that has text "Invitation Link" -// .locator("p") // find all

inside that div -// .nth(1) // pick the second

(index 1) -// .innerText(); - -// // Close Invitation Link Modal -// await page -// .locator(".ant-modal-wrap") -// .filter({ hasText: "Invitation Link" }) -// .locator('button[aria-label="Close"]') -// .click(); - -// // Close Invite User Modal -// await page -// .locator(".ant-modal-wrap") -// .filter({ hasText: "Invite User" }) -// .locator('button[aria-label="Close"]') -// .click(); - -// // Open invite link as new page (simulate invited user) -// const context = await page.context()?.browser()?.newContext(); -// const invitedUserPage = await context?.newPage(); -// if (!invitedUserPage) { -// throw new Error("invitedUserPage is undefined"); -// } -// await invitedUserPage?.goto(invitationUrl || ""); - -// //Insert new password -// await invitedUserPage?.fill("input#password", testPassword); - -// //Click on submit -// await invitedUserPage?.getByRole("button", { name: "Sign Up" }).click(); - -// // // --- Verify Keys Created --- -// // await invitedUserPage?.waitForSelector("table"); - -// // // Verify keyName1 (associated with user's team) IS visible in the table -// // const keyTable = invitedUserPage.locator('table'); // Locate the table element -// // await expect(keyTable).toBeVisible({ timeout: 10000 }); // Ensure table exists -// // // Use getByText within the table scope to find the key name -// // await expect(keyTable.getByText(keyName1, { exact: true })).toBeVisible({ timeout: 10000 }); -// // console.log(`Verified key ${keyName1} is visible for user ${testEmail}`); - -// // // Verify keyName2 (associated with the *other* team) IS NOT visible -// // await expect(keyTable.getByText(keyName2, { exact: true })).not.toBeVisible(); -// // console.log(`Verified key ${keyName2} is NOT visible for user ${testEmail}`); -// }); -// }); diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts deleted file mode 100644 index 832832d8ae8..00000000000 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts +++ /dev/null @@ -1,72 +0,0 @@ -/* -Test view internal user page -*/ - -import { test, expect } from "@playwright/test"; - -test("view internal user page", async ({ page }) => { - // Go to the specified URL - await page.goto("http://localhost:4000/ui"); - await page.waitForLoadState("networkidle"); - - page.screenshot({ path: "test-results/view_internal_user_before_login.png" }); - - // Enter "admin" in the username input field - await page.fill('input[placeholder="Enter your username"]', "admin"); - - // Enter "gm" in the password input field - await page.fill('input[placeholder="Enter your password"]', "gm"); - - // Click the login button - const loginButton = page.getByRole("button", { name: "Login" }); - await expect(loginButton).toBeEnabled(); - await loginButton.click(); - - // Wait for the Internal User tab and click it - const tabElement = page.locator("span.ant-menu-title-content", { - hasText: "Internal User", - }); - await tabElement.click(); - - // Wait for the table to load - await page.waitForSelector("tbody tr", { timeout: 10000 }); - await page.waitForTimeout(2000); // Additional wait for table to stabilize - await page.waitForLoadState("networkidle"); - - // Test all expected fields are present - // Verify that the API Keys column is rendered for all users - // The UI renders badges in each row - we just verify the column structure exists - const rowCount = await page.locator("tbody tr").count(); - expect(rowCount).toBeGreaterThan(0); - - const userIdHeader = await page.locator("th", { hasText: "User ID" }); - await expect(userIdHeader).toBeVisible({ timeout: 10000 }); - - // test pagination - // Wait for pagination controls to be visible - await page.waitForSelector(".flex.justify-between.items-center", { - timeout: 5000, - }); - - // Check if we're on the first page by looking at the results count - const resultsText = - (await page.locator(".text-sm.text-gray-700").textContent()) || ""; - const isFirstPage = resultsText.includes("1 -"); - - if (isFirstPage) { - // On first page, previous button should be disabled - const prevButton = page.locator("button", { hasText: "Previous" }); - await expect(prevButton).toBeDisabled(); - } - - // Next button should be enabled if there are more pages - const nextButton = page.locator("button", { hasText: "Next" }); - const totalResults = - (await page.locator(".text-sm.text-gray-700").textContent()) || ""; - const hasMorePages = - totalResults.includes("of") && !totalResults.includes("1 - 25 of 25"); - - if (hasMorePages) { - await expect(nextButton).toBeEnabled(); - } -}); diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_user_info.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/view_user_info.spec.ts deleted file mode 100644 index adda3088f12..00000000000 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_user_info.spec.ts +++ /dev/null @@ -1,124 +0,0 @@ -import { test, expect } from "@playwright/test"; -import { loginToUI } from "../utils/login"; - -test.describe("User Info View", () => { - test("should display user info when clicking on user ID", async ({ - page, - }) => { - await page.goto("http://localhost:4000/ui"); - await page.waitForLoadState("networkidle"); - - page.screenshot({ - path: "test-results/view_user_info_before_login.png", - }); - - // Enter "admin" in the username input field - await page.fill('input[placeholder="Enter your username"]', "admin"); - page.screenshot({ - path: "test-results/view_user_info_after_username_input.png", - }); - - // Enter "gm" in the password input field - await page.fill('input[placeholder="Enter your password"]', "gm"); - page.screenshot({ - path: "test-results/view_user_info_after_password_input.png", - }); - - // Click the login button - const loginButton = page.getByRole("button", { name: "Login" }); - await expect(loginButton).toBeEnabled(); - await loginButton.click(); - page.screenshot({ - path: "test-results/view_user_info_after_login_button_click.png", - }); - - // Wait for navigation to complete and dashboard to load - await page.waitForLoadState("networkidle"); - const tabElement = page.locator("span.ant-menu-title-content", { - hasText: "Internal User", - }); - await tabElement.click(); - page.screenshot({ - path: "test-results/view_user_info_after_internal_user_tab_click.png", - }); - // Wait for loading state to disappear - await page.waitForSelector('text="🚅 Loading users..."', { - state: "hidden", - timeout: 10000, - }); - page.screenshot({ path: "test-results/view_user_info_after_loading.png" }); - // Wait for users table to load - await page.waitForSelector("table"); - page.screenshot({ - path: "test-results/view_user_info_after_table_load.png", - }); - // Get the first user ID cell - const firstUserIdCell = page.locator( - "table tbody tr:first-child td:first-child" - ); - const userId = await firstUserIdCell.textContent(); - console.log("Found user ID:", userId); - - // Click on the user ID - await firstUserIdCell.click(); - await page.waitForLoadState("networkidle"); - - // Check for tabs - await expect(page.locator('button:has-text("Overview")')).toBeVisible({ - timeout: 10000, - }); - await expect(page.locator('button:has-text("Details")')).toBeVisible({ - timeout: 10000, - }); - - // Switch to details tab - await page.locator('button:has-text("Details")').click(); - - // Check details section - await expect(page.locator("text=User ID")).toBeVisible(); - await expect(page.locator("text=Email")).toBeVisible(); - - // Go back to users list - await page.locator('button:has-text("Back to Users")').click(); - - // Verify we're back on the users page - await expect(page.locator("table")).toBeVisible(); - await expect( - page.locator('input[placeholder="Search by email..."]') - ).toBeVisible(); - }); - - // test("should handle user deletion", async ({ page }) => { - // // Wait for users table to load - // await page.waitForSelector("table"); - - // // Get the first user ID cell - // const firstUserIdCell = page.locator( - // "table tbody tr:first-child td:first-child" - // ); - // const userId = await firstUserIdCell.textContent(); - - // // Click on the user ID - // await firstUserIdCell.click(); - - // // Wait for user info view to load - // await page.waitForSelector('h1:has-text("User")'); - - // // Click delete button - // await page.locator('button:has-text("Delete User")').click(); - - // // Confirm deletion in modal - // await page.locator('button:has-text("Delete")').click(); - - // // Verify success message - // await expect(page.locator("text=User deleted successfully")).toBeVisible(); - - // // Verify we're back on the users page - // await expect(page.locator('h1:has-text("Users")')).toBeVisible(); - - // // Verify user is no longer in the table - // if (userId) { - // await expect(page.locator(`text=${userId}`)).not.toBeVisible(); - // } - // }); -}); diff --git a/tests/proxy_admin_ui_tests/package-lock.json b/tests/proxy_admin_ui_tests/package-lock.json deleted file mode 100644 index 8c79edf9ad1..00000000000 --- a/tests/proxy_admin_ui_tests/package-lock.json +++ /dev/null @@ -1,97 +0,0 @@ -{ - "name": "proxy_admin_ui_tests", - "version": "1.0.0", - "lockfileVersion": 3, - "requires": true, - "packages": { - "": { - "name": "proxy_admin_ui_tests", - "version": "1.0.0", - "license": "ISC", - "devDependencies": { - "@playwright/test": "^1.47.2", - "@types/node": "^22.5.5" - } - }, - "node_modules/@playwright/test": { - "version": "1.56.1", - "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.56.1.tgz", - "integrity": "sha512-vSMYtL/zOcFpvJCW71Q/OEGQb7KYBPAdKh35WNSkaZA75JlAO8ED8UN6GUNTm3drWomcbcqRPFqQbLae8yBTdg==", - "dev": true, - "license": "Apache-2.0", - "dependencies": { - "playwright": "1.56.1" - }, - "bin": { - "playwright": "cli.js" - }, - "engines": { - "node": ">=18" - } - }, - "node_modules/@types/node": { - "version": "22.19.1", - "resolved": "https://registry.npmjs.org/@types/node/-/node-22.19.1.tgz", - "integrity": "sha512-LCCV0HdSZZZb34qifBsyWlUmok6W7ouER+oQIGBScS8EsZsQbrtFTUrDX4hOl+CS6p7cnNC4td+qrSVGSCTUfQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "undici-types": "~6.21.0" - } - }, - "node_modules/fsevents": { - "version": "2.3.2", - "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", - "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", - "dev": true, - "hasInstallScript": true, - "license": "MIT", - "optional": true, - "os": [ - "darwin" - ], - "engines": { - "node": "^8.16.0 || ^10.6.0 || >=11.0.0" - } - }, - "node_modules/playwright": { - "version": "1.56.1", - "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.56.1.tgz", - "integrity": "sha512-aFi5B0WovBHTEvpM3DzXTUaeN6eN0qWnTkKx4NQaH4Wvcmc153PdaY2UBdSYKaGYw+UyWXSVyxDUg5DoPEttjw==", - "dev": true, - "license": "Apache-2.0", - "dependencies": { - "playwright-core": "1.56.1" - }, - "bin": { - "playwright": "cli.js" - }, - "engines": { - "node": ">=18" - }, - "optionalDependencies": { - "fsevents": "2.3.2" - } - }, - "node_modules/playwright-core": { - "version": "1.56.1", - "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.56.1.tgz", - "integrity": "sha512-hutraynyn31F+Bifme+Ps9Vq59hKuUCz7H1kDOcBs+2oGguKkWTU50bBWrtz34OUWmIwpBTWDxaRPXrIXkgvmQ==", - "dev": true, - "license": "Apache-2.0", - "bin": { - "playwright-core": "cli.js" - }, - "engines": { - "node": ">=18" - } - }, - "node_modules/undici-types": { - "version": "6.21.0", - "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-6.21.0.tgz", - "integrity": "sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==", - "dev": true, - "license": "MIT" - } - } -} diff --git a/tests/proxy_admin_ui_tests/package.json b/tests/proxy_admin_ui_tests/package.json deleted file mode 100644 index 5933490fb1d..00000000000 --- a/tests/proxy_admin_ui_tests/package.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "name": "proxy_admin_ui_tests", - "version": "1.0.0", - "description": "", - "main": "index.js", - "scripts": {}, - "keywords": [], - "author": "", - "license": "ISC", - "devDependencies": { - "@playwright/test": "1.56.1", - "@types/node": "22.19.1" - } -} diff --git a/tests/proxy_admin_ui_tests/playwright.config.ts b/tests/proxy_admin_ui_tests/playwright.config.ts deleted file mode 100644 index 8b66c47394a..00000000000 --- a/tests/proxy_admin_ui_tests/playwright.config.ts +++ /dev/null @@ -1,84 +0,0 @@ -import { defineConfig, devices } from '@playwright/test'; - -/** - * Read environment variables from file. - * https://github.com/motdotla/dotenv - */ -// import dotenv from 'dotenv'; -// import path from 'path'; -// dotenv.config({ path: path.resolve(__dirname, '.env') }); - -/** - * See https://playwright.dev/docs/test-configuration. - */ -export default defineConfig({ - testDir: './e2e_ui_tests', - testIgnore: ['**/tests/pass_through_tests/**', '../pass_through_tests/**/*'], - testMatch: '**/*.spec.ts', // Only run files ending in .spec.ts - /* Run tests in files in parallel */ - fullyParallel: true, - /* Fail the build on CI if you accidentally left test.only in the source code. */ - forbidOnly: !!process.env.CI, - /* Retry on CI only */ - retries: process.env.CI ? 2 : 0, - /* Opt out of parallel tests on CI. */ - workers: process.env.CI ? 1 : undefined, - /* Reporter to use. See https://playwright.dev/docs/test-reporters */ - reporter: 'html', - /* Shared settings for all the projects below. See https://playwright.dev/docs/api/class-testoptions. */ - use: { - /* Base URL to use in actions like `await page.goto('/')`. */ - // baseURL: 'http://127.0.0.1:3000', - - /* Collect trace when retrying the failed test. See https://playwright.dev/docs/trace-viewer */ - trace: 'on-first-retry', - }, - - /* Configure projects for major browsers */ - projects: [ - { - name: 'chromium', - use: { ...devices['Desktop Chrome'] }, - }, - - { - name: 'firefox', - use: { ...devices['Desktop Firefox'] }, - }, - - { - name: 'webkit', - use: { ...devices['Desktop Safari'] }, - }, - - /* Test against mobile viewports. */ - // { - // name: 'Mobile Chrome', - // use: { ...devices['Pixel 5'] }, - // }, - // { - // name: 'Mobile Safari', - // use: { ...devices['iPhone 12'] }, - // }, - - /* Test against branded browsers. */ - // { - // name: 'Microsoft Edge', - // use: { ...devices['Desktop Edge'], channel: 'msedge' }, - // }, - // { - // name: 'Google Chrome', - // use: { ...devices['Desktop Chrome'], channel: 'chrome' }, - // }, - ], - timeout: 4*60*1000, - expect: { - timeout: 10 * 1000 - } - /* Run your local dev server before starting the tests */ - // webServer: { - // command: 'npm run start', - // url: 'http://127.0.0.1:3000', - // reuseExistingServer: !process.env.CI, - // }, -}); diff --git a/tests/proxy_admin_ui_tests/utils/login.ts b/tests/proxy_admin_ui_tests/utils/login.ts deleted file mode 100644 index 25858d9f570..00000000000 --- a/tests/proxy_admin_ui_tests/utils/login.ts +++ /dev/null @@ -1,27 +0,0 @@ -import { Page, expect } from "@playwright/test"; - -export async function loginToUI(page: Page) { - // Login first - await page.goto("http://localhost:4000/ui"); - await page.waitForLoadState("networkidle"); - console.log("Navigated to login page"); - - page.screenshot({ path: "test-results/login_utils_before.png" }); - // Wait for login form to be visible - await page.waitForSelector('input[placeholder="Enter your username"]', { - timeout: 10000, - }); - console.log("Login form is visible"); - - await page.fill('input[placeholder="Enter your username"]', "admin"); - await page.fill('input[placeholder="Enter your password"]', "gm"); - console.log("Filled login credentials"); - - const loginButton = page.getByRole("button", { name: "Login" }); - await expect(loginButton).toBeEnabled(); - await loginButton.click(); - console.log("Clicked login button"); - - // Wait for navigation to complete - await page.waitForURL("**/*"); -} diff --git a/tests/proxy_behavior/__init__.py b/tests/proxy_behavior/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_behavior/management/__init__.py b/tests/proxy_behavior/management/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_behavior/management/actors.py b/tests/proxy_behavior/management/actors.py new file mode 100644 index 00000000000..6c2f1a61ce1 --- /dev/null +++ b/tests/proxy_behavior/management/actors.py @@ -0,0 +1,279 @@ +"""Read-world seed for the authz matrix tests: 2 orgs, 3 teams, 9 actors.""" + +import enum +import uuid +from dataclasses import dataclass +from typing import Any, Dict + +from prisma import Json + +from litellm.proxy._types import LitellmUserRoles +from litellm.proxy.utils import PrismaClient, hash_token + + +class Actor(str, enum.Enum): + PROXY_ADMIN = "proxy_admin" + ORG_ADMIN = "org_admin" + TEAM_ADMIN = "team_admin" + INTERNAL_USER = "internal_user" + OWNER = "owner" + UNRELATED_SAME_ORG = "unrelated_same_org" + CROSS_ORG_USER = "cross_org_user" + SERVICE_ACCOUNT = "service_account" + ORG_B_ADMIN = "org_b_admin" + + +PREFIX = "behavior-pin-" +ORG_A = PREFIX + "org-a" +ORG_B = PREFIX + "org-b" +TEAM_ALPHA = PREFIX + "team-alpha" +TEAM_BETA = PREFIX + "team-beta" +TEAM_GAMMA = PREFIX + "team-gamma" +BUDGET_ID = PREFIX + "budget" + + +@dataclass(frozen=True) +class SeededKey: + user_id: str + cleartext: str + hashed: str + + +@dataclass(frozen=True) +class World: + org_a_id: str + org_b_id: str + team_alpha_id: str + team_beta_id: str + team_gamma_id: str + keys: Dict[Actor, SeededKey] + + +def _new_clear_key() -> str: + return "sk-" + uuid.uuid4().hex + + +def _actor_profile() -> Dict[Actor, Dict[str, Any]]: + return { + Actor.PROXY_ADMIN: { + "user_role": LitellmUserRoles.PROXY_ADMIN.value, + "team_id": None, + "organization_id": None, + }, + Actor.ORG_ADMIN: { + "user_role": LitellmUserRoles.ORG_ADMIN.value, + "team_id": None, + "organization_id": ORG_A, + }, + Actor.TEAM_ADMIN: { + "user_role": LitellmUserRoles.INTERNAL_USER.value, + "team_id": TEAM_ALPHA, + "organization_id": ORG_A, + }, + Actor.INTERNAL_USER: { + "user_role": LitellmUserRoles.INTERNAL_USER.value, + "team_id": TEAM_ALPHA, + "organization_id": ORG_A, + }, + Actor.OWNER: { + "user_role": LitellmUserRoles.INTERNAL_USER.value, + "team_id": TEAM_ALPHA, + "organization_id": ORG_A, + }, + Actor.UNRELATED_SAME_ORG: { + "user_role": LitellmUserRoles.INTERNAL_USER.value, + "team_id": TEAM_ALPHA, + "organization_id": ORG_A, + }, + Actor.CROSS_ORG_USER: { + "user_role": LitellmUserRoles.INTERNAL_USER.value, + "team_id": TEAM_BETA, + "organization_id": ORG_B, + }, + Actor.SERVICE_ACCOUNT: { + "user_role": LitellmUserRoles.INTERNAL_USER.value, + "team_id": TEAM_ALPHA, + "organization_id": ORG_A, + }, + Actor.ORG_B_ADMIN: { + "user_role": LitellmUserRoles.ORG_ADMIN.value, + "team_id": None, + "organization_id": ORG_B, + }, + } + + +async def _wipe_world(prisma: PrismaClient) -> None: + await prisma.db.litellm_verificationtoken.delete_many( + where={"user_id": {"startswith": PREFIX}} + ) + await prisma.db.litellm_organizationmembership.delete_many( + where={"user_id": {"startswith": PREFIX}} + ) + await prisma.db.litellm_teammembership.delete_many( + where={"user_id": {"startswith": PREFIX}} + ) + await prisma.db.litellm_usertable.delete_many( + where={"user_id": {"startswith": PREFIX}} + ) + await prisma.db.litellm_teamtable.delete_many( + where={"team_id": {"startswith": PREFIX}} + ) + await prisma.db.litellm_organizationtable.delete_many( + where={"organization_id": {"startswith": PREFIX}} + ) + await prisma.db.litellm_budgettable.delete_many(where={"budget_id": BUDGET_ID}) + + +async def seed_world(prisma: PrismaClient) -> World: + await _wipe_world(prisma) + + await prisma.db.litellm_budgettable.create( + data={ + "budget_id": BUDGET_ID, + "created_by": "behavior-pin-seeder", + "updated_by": "behavior-pin-seeder", + } + ) + + for org_id, alias in [(ORG_A, "alpha"), (ORG_B, "beta")]: + await prisma.db.litellm_organizationtable.create( + data={ + "organization_id": org_id, + "organization_alias": alias, + "budget_id": BUDGET_ID, + "created_by": "behavior-pin-seeder", + "updated_by": "behavior-pin-seeder", + } + ) + + profiles = _actor_profile() + user_ids: Dict[Actor, str] = {actor: PREFIX + actor.value for actor in Actor} + + for actor, profile in profiles.items(): + teams_list = [profile["team_id"]] if profile["team_id"] else [] + await prisma.db.litellm_usertable.create( + data={ + "user_id": user_ids[actor], + "user_role": profile["user_role"], + "team_id": profile["team_id"], + "organization_id": profile["organization_id"], + "teams": teams_list, + } + ) + + # _get_user_in_team in key_management_endpoints.py walks members_with_roles + # (a JSON list of {user_id, role}), not the String[] members column — + # populate both to match what /team/new produces. + await prisma.db.litellm_teamtable.create( + data={ + "team_id": TEAM_ALPHA, + "team_alias": "alpha-1", + "organization_id": ORG_A, + "admins": [user_ids[Actor.TEAM_ADMIN]], + "members": [ + user_ids[Actor.TEAM_ADMIN], + user_ids[Actor.INTERNAL_USER], + user_ids[Actor.OWNER], + user_ids[Actor.UNRELATED_SAME_ORG], + user_ids[Actor.SERVICE_ACCOUNT], + ], + "members_with_roles": Json( + [ + {"user_id": user_ids[Actor.TEAM_ADMIN], "role": "admin"}, + {"user_id": user_ids[Actor.INTERNAL_USER], "role": "user"}, + {"user_id": user_ids[Actor.OWNER], "role": "user"}, + {"user_id": user_ids[Actor.UNRELATED_SAME_ORG], "role": "user"}, + {"user_id": user_ids[Actor.SERVICE_ACCOUNT], "role": "user"}, + ] + ), + } + ) + await prisma.db.litellm_teamtable.create( + data={ + "team_id": TEAM_BETA, + "team_alias": "beta-1", + "organization_id": ORG_B, + "admins": [], + "members": [user_ids[Actor.CROSS_ORG_USER]], + "members_with_roles": Json( + [ + {"user_id": user_ids[Actor.CROSS_ORG_USER], "role": "user"}, + ] + ), + } + ) + # TEAM_GAMMA: ORG_A team with no actor members — the "same-org, + # not-my-team" read target. + await prisma.db.litellm_teamtable.create( + data={ + "team_id": TEAM_GAMMA, + "team_alias": "gamma-1", + "organization_id": ORG_A, + "admins": [], + "members": [], + "members_with_roles": Json([]), + } + ) + + for actor, org_id, role in [ + (Actor.ORG_ADMIN, ORG_A, "org_admin"), + (Actor.TEAM_ADMIN, ORG_A, "internal_user"), + (Actor.INTERNAL_USER, ORG_A, "internal_user"), + (Actor.OWNER, ORG_A, "internal_user"), + (Actor.UNRELATED_SAME_ORG, ORG_A, "internal_user"), + (Actor.SERVICE_ACCOUNT, ORG_A, "internal_user"), + (Actor.CROSS_ORG_USER, ORG_B, "internal_user"), + (Actor.ORG_B_ADMIN, ORG_B, "org_admin"), + ]: + await prisma.db.litellm_organizationmembership.create( + data={ + "user_id": user_ids[actor], + "organization_id": org_id, + "user_role": role, + } + ) + + for actor, team_id in [ + (Actor.TEAM_ADMIN, TEAM_ALPHA), + (Actor.INTERNAL_USER, TEAM_ALPHA), + (Actor.OWNER, TEAM_ALPHA), + (Actor.UNRELATED_SAME_ORG, TEAM_ALPHA), + (Actor.SERVICE_ACCOUNT, TEAM_ALPHA), + (Actor.CROSS_ORG_USER, TEAM_BETA), + ]: + await prisma.db.litellm_teammembership.create( + data={"user_id": user_ids[actor], "team_id": team_id} + ) + + keys: Dict[Actor, SeededKey] = {} + for actor, profile in profiles.items(): + cleartext = _new_clear_key() + hashed = hash_token(cleartext) + token_data: Dict[str, Any] = { + "token": hashed, + "key_name": PREFIX + actor.value + "-key", + "user_id": user_ids[actor], + # LiteLLM_VerificationTokenView's models field rejects NULL even + # though the column is nullable in Postgres. + "models": [], + } + if profile["team_id"]: + token_data["team_id"] = profile["team_id"] + if profile["organization_id"]: + token_data["organization_id"] = profile["organization_id"] + if actor == Actor.SERVICE_ACCOUNT: + token_data["metadata"] = Json({"service_account_id": user_ids[actor]}) + await prisma.db.litellm_verificationtoken.create(data=token_data) + keys[actor] = SeededKey( + user_id=user_ids[actor], cleartext=cleartext, hashed=hashed + ) + + return World( + org_a_id=ORG_A, + org_b_id=ORG_B, + team_alpha_id=TEAM_ALPHA, + team_beta_id=TEAM_BETA, + team_gamma_id=TEAM_GAMMA, + keys=keys, + ) diff --git a/tests/proxy_behavior/management/conftest.py b/tests/proxy_behavior/management/conftest.py new file mode 100644 index 00000000000..fa0bef86280 --- /dev/null +++ b/tests/proxy_behavior/management/conftest.py @@ -0,0 +1,281 @@ +"""Session-scoped async ASGI client for HTTP-boundary behavior tests.""" + +import os +import tempfile +import uuid +from dataclasses import dataclass +from typing import Any, AsyncIterator, Dict, Optional + +import httpx +import pytest_asyncio +import yaml +from prisma import Json + +from litellm.proxy.utils import hash_token + +MASTER_KEY = "sk-1234" +SCRATCH_PREFIX = "scratch-" + + +def _write_minimal_proxy_config() -> str: + config = { + "general_settings": {"master_key": MASTER_KEY}, + "litellm_settings": {}, + } + database_url = os.environ.get("DATABASE_URL") + if database_url: + config["general_settings"]["database_url"] = database_url + f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) + yaml.dump(config, f) + f.close() + return f.name + + +@pytest_asyncio.fixture(scope="session") +async def proxy_app(): + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import ( + app, + cleanup_router_config_variables, + initialize, + proxy_startup_event, + ) + + cleanup_router_config_variables() + config_path = _write_minimal_proxy_config() + + # proxy_startup_event re-reads master_key from LITELLM_MASTER_KEY and + # unconditionally overwrites the global, even when initialize() already + # set it from the config YAML. Force (not setdefault) both vars: an + # ambient LITELLM_MASTER_KEY with a different value would make the proxy + # authenticate on that key while the tests still send MASTER_KEY. + os.environ["LITELLM_MASTER_KEY"] = MASTER_KEY + os.environ["CONFIG_FILE_PATH"] = config_path + + await initialize(config=config_path) + + # /key/regenerate is gated behind premium_user; flipping it lets the matrix + # pin authz behavior instead of the licensing gate. + proxy_server.premium_user = True + + async with proxy_startup_event(app): + proxy_server.premium_user = True # lifespan re-runs _license_check + # The lifespan fires check_view_exists() as a background task; on a + # fresh DB the first auth call races it and resolves user_id=None. + if proxy_server.prisma_client is not None: + await proxy_server.prisma_client.check_view_exists() + yield app + + +@pytest_asyncio.fixture(scope="session") +async def proxy_client(proxy_app) -> AsyncIterator[httpx.AsyncClient]: + transport = httpx.ASGITransport(app=proxy_app) + async with httpx.AsyncClient( + transport=transport, base_url="http://testserver" + ) as client: + yield client + + +@pytest_asyncio.fixture(scope="session") +async def prisma(proxy_app): + from litellm.proxy import proxy_server + + assert proxy_server.prisma_client is not None + return proxy_server.prisma_client + + +@pytest_asyncio.fixture(scope="session") +async def world(prisma): + from .actors import seed_world + + return await seed_world(prisma) + + +@dataclass(frozen=True) +class Scratch: + prefix: str + + def tag(self, suffix: str = "") -> str: + return f"{self.prefix}-{suffix}" if suffix else self.prefix + + +async def create_scratch_key( + proxy_client, + seeder_cleartext: str, + scratch_prefix: str, + *, + user_id: str, + team_id: Optional[str] = None, + organization_id: Optional[str] = None, + key_alias: Optional[str] = None, +) -> str: + """Seed a scratch-tagged key via /key/generate; returns its cleartext. + + Shared by the write-scenario matrices (key update/regenerate/delete). + key_alias defaults to scratch_prefix; pass a distinct scratch-prefixed + alias when a single scenario needs more than one key (/key/generate + enforces unique aliases). + """ + body: Dict[str, Any] = { + "key_alias": key_alias or scratch_prefix, + "user_id": user_id, + } + if team_id is not None: + body["team_id"] = team_id + if organization_id is not None: + body["organization_id"] = organization_id + resp = await proxy_client.post( + "/key/generate", + headers={"Authorization": f"Bearer {seeder_cleartext}"}, + json=body, + ) + assert resp.status_code == 200, f"setup failed: {resp.text}" + return resp.json()["key"] + + +async def create_scratch_team( + prisma, + team_id: str, + *, + organization_id: Optional[str] = None, + admin_user_ids: Optional[list] = None, + member_user_ids: Optional[list] = None, + team_member_permissions: Optional[list] = None, + models: Optional[list] = None, +) -> str: + """Raw-seed a scratch-tagged team row; returns its team_id. + + The target team for the team write matrices (update / member_*). Raw + prisma (not POST /team/new) avoids creation side effects — no creator + auto-add, no membership rows written onto the world's users — so seeding + never mutates the immutable read-world. The authz gates read the team's + members_with_roles JSON, so a raw-seeded team exercises them exactly as + a /team/new-created team would. team_id must start with the scratch + prefix so the `scratch` fixture reclaims the row. + + team_member_permissions / models seed the matching raw columns — needed + by the team-key-permission and team-model matrices. + """ + admin_user_ids = list(admin_user_ids or []) + member_user_ids = list(member_user_ids or []) + members_with_roles = [ + {"user_id": uid, "role": "admin"} for uid in admin_user_ids + ] + [{"user_id": uid, "role": "user"} for uid in member_user_ids] + data: Dict[str, Any] = { + "team_id": team_id, + "team_alias": team_id, + "admins": admin_user_ids, + "members": admin_user_ids + member_user_ids, + "members_with_roles": Json(members_with_roles), + } + if organization_id is not None: + data["organization_id"] = organization_id + if team_member_permissions is not None: + data["team_member_permissions"] = team_member_permissions + if models is not None: + data["models"] = models + await prisma.db.litellm_teamtable.create(data=data) + return team_id + + +@dataclass(frozen=True) +class SeededActor: + user_id: str + cleartext: str + hashed: str + + +async def create_scratch_actor( + prisma, + scratch_prefix: str, + *, + user_role: str, + org_admin_of: tuple = (), + organization_id: Optional[str] = None, + suffix: str = "actor", +) -> SeededActor: + """Mint a scratch-prefixed user + verification token (+ org memberships). + + Reclaimed by the existing `scratch` teardown, which sweeps + litellm_usertable, litellm_verificationtoken, and + litellm_organizationmembership by scratch prefix — no bespoke cleanup + needed. Does NOT write litellm_teammembership against world teams: the + teardown reclaims that table only by team_id prefix, so a scratch actor + needing team membership must join a scratch team instead. The cleartext + is hashed with the real hash_token so the key authenticates end-to-end; + models=[] satisfies LiteLLM_VerificationTokenView. + """ + user_id = f"{scratch_prefix}-{suffix}" + cleartext = "sk-" + uuid.uuid4().hex + hashed = hash_token(cleartext) + await prisma.db.litellm_usertable.create( + data={ + "user_id": user_id, + "user_role": user_role, + "organization_id": organization_id, + } + ) + token_data: Dict[str, Any] = { + "token": hashed, + "key_name": f"{scratch_prefix}-{suffix}-key", + "key_alias": f"{scratch_prefix}-{suffix}-alias", + "user_id": user_id, + "models": [], + } + if organization_id is not None: + token_data["organization_id"] = organization_id + await prisma.db.litellm_verificationtoken.create(data=token_data) + for org_id in org_admin_of: + await prisma.db.litellm_organizationmembership.create( + data={ + "user_id": user_id, + "organization_id": org_id, + "user_role": "org_admin", + } + ) + return SeededActor(user_id=user_id, cleartext=cleartext, hashed=hashed) + + +@pytest_asyncio.fixture +async def scratch(prisma): + handle = Scratch(prefix=f"{SCRATCH_PREFIX}{uuid.uuid4().hex[:12]}") + try: + yield handle + finally: + # Children before parents to avoid FK violations. + await prisma.db.litellm_verificationtoken.delete_many( + where={ + "OR": [ + {"key_alias": {"startswith": handle.prefix}}, + {"key_name": {"startswith": handle.prefix}}, + ] + } + ) + await prisma.db.litellm_teammembership.delete_many( + where={"team_id": {"startswith": handle.prefix}} + ) + await prisma.db.litellm_organizationmembership.delete_many( + where={"user_id": {"startswith": handle.prefix}} + ) + await prisma.db.litellm_teamtable.delete_many( + where={"team_id": {"startswith": handle.prefix}} + ) + await prisma.db.litellm_usertable.delete_many( + where={"user_id": {"startswith": handle.prefix}} + ) + await prisma.db.litellm_budgettable.delete_many( + where={"budget_id": {"startswith": handle.prefix}} + ) + # /team/member_add writes LiteLLM_UserTable.teams; the available-team + # self-join writes it on a world actor whose row must survive. Strip + # dangling scratch-team refs so the read-world stays immutable. + polluted = await prisma.db.litellm_usertable.find_many( + where={"teams": {"isEmpty": False}} + ) + for user in polluted: + cleaned = [t for t in user.teams if not t.startswith(handle.prefix)] + if cleaned != list(user.teams): + await prisma.db.litellm_usertable.update( + where={"user_id": user.user_id}, + data={"teams": {"set": cleaned}}, + ) diff --git a/tests/proxy_behavior/management/test_key_aliases.py b/tests/proxy_behavior/management/test_key_aliases.py new file mode 100644 index 00000000000..38ce5cdfaf3 --- /dev/null +++ b/tests/proxy_behavior/management/test_key_aliases.py @@ -0,0 +1,119 @@ +import uuid +from typing import FrozenSet + +import pytest + +from litellm.proxy.utils import hash_token + +from .actors import TEAM_ALPHA, TEAM_BETA, Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# GET /key/aliases scopes non-admins via _apply_non_admin_alias_scope: a +# non-admin sees an alias only if it owns the key (user_id match) or the key +# belongs to one of its teams. PROXY_ADMIN sees every alias. The seeded keys: +# own — owned by INTERNAL_USER, no team -> user_id scope only +# alpha — owned by OWNER, team TEAM_ALPHA -> team scope for alpha members +# beta — owned by CROSS_ORG_USER, TEAM_BETA +async def _seed_alias_keys(prisma, prefix: str, world) -> dict: + spec = { + "own": (Actor.INTERNAL_USER, None), + "alpha": (Actor.OWNER, TEAM_ALPHA), + "beta": (Actor.CROSS_ORG_USER, TEAM_BETA), + } + out = {} + for tag, (owner, team_id) in spec.items(): + alias = f"{prefix}-{tag}" + data = { + "token": hash_token("sk-" + uuid.uuid4().hex), + "key_name": f"{prefix}-{tag}-key", + "key_alias": alias, + "user_id": world.keys[owner].user_id, + "models": [], + } + if team_id is not None: + data["team_id"] = team_id + await prisma.db.litellm_verificationtoken.create(data=data) + out[tag] = alias + return out + + +async def _fetch_aliases(proxy_client, caller_cleartext: str, query: str) -> set: + resp = await proxy_client.get( + f"/key/aliases?{query}&size=100", + headers={"Authorization": f"Bearer {caller_cleartext}"}, + ) + assert resp.status_code == 200, resp.text + return set(resp.json()["aliases"]) + + +# ORG_ADMIN-role callers are stopped 401 by the management-route gate before +# the handler runs — /key/aliases carries no org context. Every other actor +# reaches the handler and is scoped by _apply_non_admin_alias_scope. +_VISIBILITY = { + Actor.PROXY_ADMIN: (200, frozenset({"own", "alpha", "beta"})), + Actor.ORG_ADMIN: (401, None), + Actor.TEAM_ADMIN: (200, frozenset({"alpha"})), + Actor.INTERNAL_USER: (200, frozenset({"own", "alpha"})), + Actor.OWNER: (200, frozenset({"alpha"})), + Actor.UNRELATED_SAME_ORG: (200, frozenset({"alpha"})), + Actor.CROSS_ORG_USER: (200, frozenset({"beta"})), + Actor.SERVICE_ACCOUNT: (200, frozenset({"alpha"})), + Actor.ORG_B_ADMIN: (401, None), +} + + +@pytest.mark.parametrize( + "actor,expected_status,expected_tags", + [(a, s, t) for a, (s, t) in _VISIBILITY.items()], + ids=[a.value for a in _VISIBILITY], +) +async def test_key_aliases_visibility( + actor: Actor, + expected_status: int, + expected_tags: FrozenSet[str], + proxy_client, + prisma, + scratch, + world, +): + aliases = await _seed_alias_keys(prisma, scratch.prefix, world) + known = {v: k for k, v in aliases.items()} + + resp = await proxy_client.get( + f"/key/aliases?search={scratch.prefix}&size=100", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value}: {resp.status_code} {resp.text}" + if expected_status != 200: + return + + visible = {known[a] for a in resp.json()["aliases"] if a in known} + assert visible == set( + expected_tags + ), f"{actor.value}: expected {sorted(expected_tags)}, got {sorted(visible)}" + + +async def test_key_aliases_team_id_filter(proxy_client, prisma, scratch, world): + """team_id filter narrows the result to keys of that team.""" + aliases = await _seed_alias_keys(prisma, scratch.prefix, world) + returned = await _fetch_aliases( + proxy_client, + world.keys[Actor.PROXY_ADMIN].cleartext, + f"search={scratch.prefix}&team_id={TEAM_ALPHA}", + ) + assert returned & set(aliases.values()) == {aliases["alpha"]} + + +async def test_key_aliases_search_filter(proxy_client, prisma, scratch, world): + """search is a case-insensitive substring match on key_alias.""" + aliases = await _seed_alias_keys(prisma, scratch.prefix, world) + returned = await _fetch_aliases( + proxy_client, + world.keys[Actor.PROXY_ADMIN].cleartext, + f"search={aliases['beta']}", + ) + assert returned & set(aliases.values()) == {aliases["beta"]} diff --git a/tests/proxy_behavior/management/test_key_block_unblock.py b/tests/proxy_behavior/management/test_key_block_unblock.py new file mode 100644 index 00000000000..37aa0c0219a --- /dev/null +++ b/tests/proxy_behavior/management/test_key_block_unblock.py @@ -0,0 +1,159 @@ +import uuid + +import pytest + +from litellm.proxy.utils import hash_token + +from .actors import TEAM_ALPHA, TEAM_BETA, Actor +from .conftest import create_scratch_key + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /key/block + /key/unblock. PROXY_ADMIN bypasses. ORG_ADMIN-role callers +# are stopped 401 by the management-route gate BEFORE the handler runs — the +# body carries no organization_id, so the gate has no org context and falls +# back to proxy-admin-only. The handler's own _check_key_admin_access org-admin +# branch is therefore unreachable via these routes. INTERNAL_USER-role callers +# do reach _check_key_admin_access: a team admin of the key's team passes (200); +# everyone else (incl. a teamless "self" key with no team to admin) is 403. +_SCENARIOS = [ + ("self/proxy_admin", Actor.PROXY_ADMIN, "self", 200), + ("self/org_admin", Actor.ORG_ADMIN, "self", 401), + ("self/team_admin", Actor.TEAM_ADMIN, "self", 403), + ("self/internal_user", Actor.INTERNAL_USER, "self", 403), + ("self/cross_org_user", Actor.CROSS_ORG_USER, "self", 403), + ("owner/proxy_admin", Actor.PROXY_ADMIN, "owner", 200), + ("owner/org_admin", Actor.ORG_ADMIN, "owner", 401), + ("owner/team_admin", Actor.TEAM_ADMIN, "owner", 200), + ("owner/internal_user", Actor.INTERNAL_USER, "owner", 403), + ("owner/owner", Actor.OWNER, "owner", 403), + ("owner/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "owner", 403), + ("owner/cross_org_user", Actor.CROSS_ORG_USER, "owner", 403), + ("owner/service_account", Actor.SERVICE_ACCOUNT, "owner", 403), + ("owner/org_b_admin", Actor.ORG_B_ADMIN, "owner", 401), + ("cross_org/proxy_admin", Actor.PROXY_ADMIN, "cross_org", 200), + ("cross_org/org_admin", Actor.ORG_ADMIN, "cross_org", 401), + ("cross_org/team_admin", Actor.TEAM_ADMIN, "cross_org", 403), + ("cross_org/cross_org_user", Actor.CROSS_ORG_USER, "cross_org", 403), + ("cross_org/org_b_admin", Actor.ORG_B_ADMIN, "cross_org", 401), +] + + +async def _seed_target(proxy_client, seeder, scratch_prefix, world, shape, caller): + if shape == "self": + return await create_scratch_key( + proxy_client, seeder, scratch_prefix, user_id=caller.user_id + ) + if shape == "owner": + return await create_scratch_key( + proxy_client, + seeder, + scratch_prefix, + user_id=world.keys[Actor.OWNER].user_id, + team_id=TEAM_ALPHA, + ) + if shape == "cross_org": + return await create_scratch_key( + proxy_client, + seeder, + scratch_prefix, + user_id=world.keys[Actor.CROSS_ORG_USER].user_id, + team_id=TEAM_BETA, + ) + pytest.fail(f"unknown shape={shape}") # pragma: no cover + + +@pytest.mark.parametrize("route", ["block", "unblock"]) +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_key_block_unblock_authz_matrix( + route: str, + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + caller = world.keys[actor] + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + target_cleartext = await _seed_target( + proxy_client, seeder, scratch.prefix, world, shape, caller + ) + target_hashed = hash_token(target_cleartext) + + # /unblock starts from a blocked row so a 200 is observable as True->False. + if route == "unblock": + await prisma.db.litellm_verificationtoken.update( + where={"token": target_hashed}, data={"blocked": True} + ) + + resp = await proxy_client.post( + f"/key/{route}", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"key": target_cleartext}, + ) + assert ( + resp.status_code == expected_status + ), f"{route} {actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_verificationtoken.find_unique( + where={"token": target_hashed} + ) + assert row is not None + # A never-blocked key reads back blocked=None; treat that as not-blocked. + if expected_status == 200: + assert bool(row.blocked) is (route == "block") + else: + # A denial leaves the blocked column at its pre-request value. + assert bool(row.blocked) is (route == "unblock"), "denied but blocked mutated" + + +async def test_key_block_unblock_round_trip(proxy_client, prisma, scratch, world): + """PROXY_ADMIN block then unblock flips the blocked column True then False.""" + admin = world.keys[Actor.PROXY_ADMIN] + target = await create_scratch_key( + proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id + ) + hashed = hash_token(target) + headers = {"Authorization": f"Bearer {admin.cleartext}"} + + blocked = await proxy_client.post( + "/key/block", headers=headers, json={"key": target} + ) + assert blocked.status_code == 200, blocked.text + row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed}) + assert row is not None and row.blocked is True + + unblocked = await proxy_client.post( + "/key/unblock", headers=headers, json={"key": target} + ) + assert unblocked.status_code == 200, unblocked.text + row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed}) + assert row is not None and row.blocked is False + + +@pytest.mark.parametrize("route", ["block", "unblock"]) +@pytest.mark.parametrize( + "actor", [Actor.PROXY_ADMIN, Actor.TEAM_ADMIN], ids=["proxy_admin", "team_admin"] +) +async def test_key_block_unblock_missing_key_returns_404( + route: str, actor: Actor, proxy_client, world +): + """A well-formed but unseeded key is 404 — not 401/403 — for both the + PROXY_ADMIN existence check and the non-admin _check_key_admin_access path.""" + caller = world.keys[actor] + missing = "sk-" + uuid.uuid4().hex + resp = await proxy_client.post( + f"/key/{route}", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"key": missing}, + ) + assert ( + resp.status_code == 404 + ), f"{route} {actor.value}: {resp.status_code} {resp.text}" diff --git a/tests/proxy_behavior/management/test_key_bulk_update.py b/tests/proxy_behavior/management/test_key_bulk_update.py new file mode 100644 index 00000000000..1a57998cece --- /dev/null +++ b/tests/proxy_behavior/management/test_key_bulk_update.py @@ -0,0 +1,123 @@ +import uuid + +import pytest + +from litellm.proxy.utils import hash_token + +from .actors import Actor +from .conftest import create_scratch_key + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_MARKER_BUDGET = 42.0 + + +# POST /key/bulk_update is PROXY_ADMIN-only. The handler's own gate is +# user_role != PROXY_ADMIN -> 403, but ORG_ADMIN-role callers never reach it: +# the management-route gate 401s them first (the body carries no org context, +# and /key/bulk_update is an internal_user route, not an org-admin one). +# INTERNAL_USER-role callers clear the route gate and hit the handler's 403. +_MATRIX = [ + ("proxy_admin", Actor.PROXY_ADMIN, 200), + ("org_admin", Actor.ORG_ADMIN, 401), + ("team_admin", Actor.TEAM_ADMIN, 403), + ("internal_user", Actor.INTERNAL_USER, 403), + ("owner", Actor.OWNER, 403), + ("unrelated_same_org", Actor.UNRELATED_SAME_ORG, 403), + ("cross_org_user", Actor.CROSS_ORG_USER, 403), + ("service_account", Actor.SERVICE_ACCOUNT, 403), +] + + +@pytest.mark.parametrize( + "actor,expected_status", + [(a, s) for (_id, a, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_key_bulk_update_authz_matrix( + actor: Actor, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + caller = world.keys[actor] + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + target = await create_scratch_key( + proxy_client, seeder, scratch.prefix, user_id=caller.user_id + ) + hashed = hash_token(target) + + resp = await proxy_client.post( + "/key/bulk_update", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"keys": [{"key": target, "max_budget": _MARKER_BUDGET}]}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed}) + assert row is not None + if expected_status == 200: + body = resp.json() + assert len(body["successful_updates"]) == 1 + assert body["failed_updates"] == [] + assert row.max_budget == _MARKER_BUDGET + else: + assert row.max_budget != _MARKER_BUDGET, "denied but key mutated" + + +async def test_key_bulk_update_empty_keys_is_400(proxy_client, world): + """An empty batch is rejected 400 before any per-key processing.""" + resp = await proxy_client.post( + "/key/bulk_update", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"keys": []}, + ) + assert resp.status_code == 400, resp.text + + +async def test_key_bulk_update_over_max_batch_is_400(proxy_client, world): + """A batch larger than the 500-key cap is rejected 400.""" + items = [{"key": "sk-" + uuid.uuid4().hex} for _ in range(501)] + resp = await proxy_client.post( + "/key/bulk_update", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"keys": items}, + ) + assert resp.status_code == 400, resp.text + + +async def test_key_bulk_update_per_key_failure_is_isolated( + proxy_client, prisma, scratch, world +): + """One bad key in the batch does not abort the others — it lands in + failed_updates while the valid key is still updated.""" + admin = world.keys[Actor.PROXY_ADMIN] + valid = await create_scratch_key( + proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id + ) + missing = "sk-" + uuid.uuid4().hex + + resp = await proxy_client.post( + "/key/bulk_update", + headers={"Authorization": f"Bearer {admin.cleartext}"}, + json={ + "keys": [ + {"key": valid, "max_budget": _MARKER_BUDGET}, + {"key": missing, "max_budget": _MARKER_BUDGET}, + ] + }, + ) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["total_requested"] == 2 + assert len(body["successful_updates"]) == 1 + assert len(body["failed_updates"]) == 1 + + row = await prisma.db.litellm_verificationtoken.find_unique( + where={"token": hash_token(valid)} + ) + assert row is not None and row.max_budget == _MARKER_BUDGET diff --git a/tests/proxy_behavior/management/test_key_delete.py b/tests/proxy_behavior/management/test_key_delete.py new file mode 100644 index 00000000000..0b483edc056 --- /dev/null +++ b/tests/proxy_behavior/management/test_key_delete.py @@ -0,0 +1,113 @@ +import uuid + +import pytest + +from litellm.proxy.utils import hash_token + +from .actors import TEAM_ALPHA, TEAM_BETA, Actor +from .conftest import create_scratch_key + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# Same-team peers can READ each other's keys (see test_key_info) but cannot +# DELETE them — delete is stricter than read. +_SCENARIOS = [ + ("self/proxy_admin", Actor.PROXY_ADMIN, "self", 200), + ("self/org_admin", Actor.ORG_ADMIN, "self", 401), + ("self/team_admin", Actor.TEAM_ADMIN, "self", 200), + ("self/internal_user", Actor.INTERNAL_USER, "self", 200), + ("self/owner", Actor.OWNER, "self", 200), + ("self/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "self", 200), + ("self/cross_org_user", Actor.CROSS_ORG_USER, "self", 200), + ("self/service_account", Actor.SERVICE_ACCOUNT, "self", 200), + ("owner_target/proxy_admin", Actor.PROXY_ADMIN, "owner", 200), + ("owner_target/org_admin", Actor.ORG_ADMIN, "owner", 401), + ("owner_target/team_admin", Actor.TEAM_ADMIN, "owner", 200), + ("owner_target/internal_user", Actor.INTERNAL_USER, "owner", 403), + ("owner_target/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "owner", 403), + ("owner_target/cross_org_user", Actor.CROSS_ORG_USER, "owner", 403), + ("owner_target/service_account", Actor.SERVICE_ACCOUNT, "owner", 403), + ("cross_org_target/proxy_admin", Actor.PROXY_ADMIN, "cross_org", 200), + ("cross_org_target/org_admin", Actor.ORG_ADMIN, "cross_org", 401), + ("cross_org_target/team_admin", Actor.TEAM_ADMIN, "cross_org", 403), + ("cross_org_target/owner", Actor.OWNER, "cross_org", 403), + ("cross_org_target/cross_org_user", Actor.CROSS_ORG_USER, "cross_org", 200), + ("cross_org_target/service_account", Actor.SERVICE_ACCOUNT, "cross_org", 403), +] + + +@pytest.mark.parametrize( + "actor,target_shape,expected_status", + [(a, t, s) for (_id, a, t, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_key_delete_authz_matrix( + actor: Actor, + target_shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + caller = world.keys[actor] + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + + if target_shape == "self": + target_cleartext = await create_scratch_key( + proxy_client, seeder, scratch.prefix, user_id=caller.user_id + ) + elif target_shape == "owner": + target_cleartext = await create_scratch_key( + proxy_client, + seeder, + scratch.prefix, + user_id=world.keys[Actor.OWNER].user_id, + team_id=TEAM_ALPHA, + ) + elif target_shape == "cross_org": + target_cleartext = await create_scratch_key( + proxy_client, + seeder, + scratch.prefix, + user_id=world.keys[Actor.CROSS_ORG_USER].user_id, + team_id=TEAM_BETA, + ) + else: + pytest.fail(f"unknown target_shape={target_shape}") + + target_hashed = hash_token(target_cleartext) + + resp = await proxy_client.post( + "/key/delete", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"keys": [target_cleartext]}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_verificationtoken.find_unique( + where={"token": target_hashed} + ) + auth_check = await proxy_client.get( + "/key/info", headers={"Authorization": f"Bearer {target_cleartext}"} + ) + + if expected_status == 200: + # Hard- or soft-delete both produce a 401 on subsequent auth. + assert auth_check.status_code == 401 + else: + assert row is not None, f"{actor.value}: denied but row vanished" + assert auth_check.status_code == 200 + + +async def test_key_delete_missing_key_is_404(proxy_client, world): + """Deleting a key absent from the DB is a 404 — not 401/403.""" + resp = await proxy_client.post( + "/key/delete", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"keys": ["sk-" + uuid.uuid4().hex]}, + ) + assert resp.status_code == 404, resp.text diff --git a/tests/proxy_behavior/management/test_key_generate.py b/tests/proxy_behavior/management/test_key_generate.py new file mode 100644 index 00000000000..851de33d3ff --- /dev/null +++ b/tests/proxy_behavior/management/test_key_generate.py @@ -0,0 +1,70 @@ +from typing import Any, Dict + +import pytest + +from .actors import TEAM_ALPHA, TEAM_BETA, Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# (id, actor, body_extras, expected_status). Status codes pinned to observed +# handler behavior — heterogeneous (200, 400, 401) because the handler routes +# denials through three different gates (role gate, user_id mismatch, team +# member permission). +_SCENARIOS = [ + ("self/proxy_admin", Actor.PROXY_ADMIN, {}, 200), + ("self/org_admin", Actor.ORG_ADMIN, {}, 401), + ("self/team_admin", Actor.TEAM_ADMIN, {}, 200), + ("self/internal_user", Actor.INTERNAL_USER, {}, 200), + ("self/owner", Actor.OWNER, {}, 200), + ("self/unrelated_same_org", Actor.UNRELATED_SAME_ORG, {}, 200), + ("self/cross_org_user", Actor.CROSS_ORG_USER, {}, 200), + ("self/service_account", Actor.SERVICE_ACCOUNT, {}, 200), + ("team_alpha/proxy_admin", Actor.PROXY_ADMIN, {"team_id": TEAM_ALPHA}, 200), + ("team_alpha/org_admin", Actor.ORG_ADMIN, {"team_id": TEAM_ALPHA}, 401), + ("team_alpha/team_admin", Actor.TEAM_ADMIN, {"team_id": TEAM_ALPHA}, 200), + ("team_alpha/internal_user", Actor.INTERNAL_USER, {"team_id": TEAM_ALPHA}, 401), + ("team_alpha/cross_org_user", Actor.CROSS_ORG_USER, {"team_id": TEAM_ALPHA}, 400), + ("team_beta/proxy_admin", Actor.PROXY_ADMIN, {"team_id": TEAM_BETA}, 200), + ("team_beta/org_admin", Actor.ORG_ADMIN, {"team_id": TEAM_BETA}, 401), + ("team_beta/team_admin", Actor.TEAM_ADMIN, {"team_id": TEAM_BETA}, 400), + ("team_beta/internal_user", Actor.INTERNAL_USER, {"team_id": TEAM_BETA}, 400), + ("team_beta/cross_org_user", Actor.CROSS_ORG_USER, {"team_id": TEAM_BETA}, 401), +] + + +@pytest.mark.parametrize( + "actor,body_extras,expected_status", + [(actor, body, expected) for (_id, actor, body, expected) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_key_generate_authz_matrix( + actor: Actor, + body_extras: Dict[str, Any], + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + seeded = world.keys[actor] + body: Dict[str, Any] = {"key_alias": scratch.prefix, **body_extras} + + resp = await proxy_client.post( + "/key/generate", + headers={"Authorization": f"Bearer {seeded.cleartext}"}, + json=body, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {body!r} → {resp.status_code}: {resp.text}" + + rows = await prisma.db.litellm_verificationtoken.find_many( + where={"key_alias": scratch.prefix} + ) + if expected_status == 200: + cleartext = resp.json()["key"] + assert cleartext.startswith("sk-") + assert len(rows) == 1 + else: + assert rows == [], f"{actor.value}: denied but row leaked" diff --git a/tests/proxy_behavior/management/test_key_health.py b/tests/proxy_behavior/management/test_key_health.py new file mode 100644 index 00000000000..62147e7fa13 --- /dev/null +++ b/tests/proxy_behavior/management/test_key_health.py @@ -0,0 +1,24 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /key/health has no role gate — it reflects the caller's OWN key logging +# metadata. The world keys carry no "logging" metadata, so every authenticated +# actor gets 200 with key="healthy". This pins auth-required + route coverage. +@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor]) +async def test_key_health_each_actor_is_healthy(actor: Actor, proxy_client, world): + caller = world.keys[actor] + resp = await proxy_client.post( + "/key/health", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}" + assert resp.json()["key"] == "healthy" + + +async def test_key_health_requires_auth(proxy_client): + resp = await proxy_client.post("/key/health") + assert resp.status_code == 401, resp.text diff --git a/tests/proxy_behavior/management/test_key_info.py b/tests/proxy_behavior/management/test_key_info.py new file mode 100644 index 00000000000..ddcef9fd27b --- /dev/null +++ b/tests/proxy_behavior/management/test_key_info.py @@ -0,0 +1,74 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# (id, actor, target_actor, expected_status). Targets are 3 fixed seeded keys +# representing the canonical relations: own, OWNER (same org_a/team_alpha), +# and CROSS_ORG_USER (org_b/team_beta). +# +# Notable pinned behaviors (intentionally surfaced, not endorsed): +# - ORG_ADMIN 403s on individual key info even within its own org — +# visibility is "your own keys" + "your team's keys", not "your org's keys". +# - Same-team peers (internal_user, unrelated_same_org, service_account) DO +# see each other's keys. +_SCENARIOS = [ + ("own/proxy_admin", Actor.PROXY_ADMIN, Actor.PROXY_ADMIN, 200), + ("own/org_admin", Actor.ORG_ADMIN, Actor.ORG_ADMIN, 200), + ("own/team_admin", Actor.TEAM_ADMIN, Actor.TEAM_ADMIN, 200), + ("own/internal_user", Actor.INTERNAL_USER, Actor.INTERNAL_USER, 200), + ("own/owner", Actor.OWNER, Actor.OWNER, 200), + ("own/unrelated_same_org", Actor.UNRELATED_SAME_ORG, Actor.UNRELATED_SAME_ORG, 200), + ("own/cross_org_user", Actor.CROSS_ORG_USER, Actor.CROSS_ORG_USER, 200), + ("own/service_account", Actor.SERVICE_ACCOUNT, Actor.SERVICE_ACCOUNT, 200), + ("owner_key/proxy_admin", Actor.PROXY_ADMIN, Actor.OWNER, 200), + ("owner_key/org_admin", Actor.ORG_ADMIN, Actor.OWNER, 403), + ("owner_key/team_admin", Actor.TEAM_ADMIN, Actor.OWNER, 200), + ("owner_key/internal_user", Actor.INTERNAL_USER, Actor.OWNER, 200), + ("owner_key/owner", Actor.OWNER, Actor.OWNER, 200), + ("owner_key/unrelated_same_org", Actor.UNRELATED_SAME_ORG, Actor.OWNER, 200), + ("owner_key/cross_org_user", Actor.CROSS_ORG_USER, Actor.OWNER, 403), + ("owner_key/service_account", Actor.SERVICE_ACCOUNT, Actor.OWNER, 200), + ("cross_org/proxy_admin", Actor.PROXY_ADMIN, Actor.CROSS_ORG_USER, 200), + ("cross_org/org_admin", Actor.ORG_ADMIN, Actor.CROSS_ORG_USER, 403), + ("cross_org/team_admin", Actor.TEAM_ADMIN, Actor.CROSS_ORG_USER, 403), + ("cross_org/internal_user", Actor.INTERNAL_USER, Actor.CROSS_ORG_USER, 403), + ("cross_org/owner", Actor.OWNER, Actor.CROSS_ORG_USER, 403), + ( + "cross_org/unrelated_same_org", + Actor.UNRELATED_SAME_ORG, + Actor.CROSS_ORG_USER, + 403, + ), + ("cross_org/cross_org_user", Actor.CROSS_ORG_USER, Actor.CROSS_ORG_USER, 200), + ("cross_org/service_account", Actor.SERVICE_ACCOUNT, Actor.CROSS_ORG_USER, 403), +] + + +@pytest.mark.parametrize( + "actor,target_actor,expected_status", + [(a, t, s) for (_id, a, t, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_key_info_authz_matrix( + actor: Actor, target_actor: Actor, expected_status: int, proxy_client, world +): + caller = world.keys[actor] + target = world.keys[target_actor] + + resp = await proxy_client.get( + f"/key/info?key={target.cleartext}", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} → {target_actor.value}: {resp.status_code} {resp.text}" + + if expected_status == 200: + body = resp.json() + # The handler echoes back whatever ?key was passed (cleartext here), + # so accept either form — info.user_id is the canonical identity check. + assert body.get("key") in (target.cleartext, target.hashed) + assert body["info"].get("user_id") == target.user_id diff --git a/tests/proxy_behavior/management/test_key_info_v2.py b/tests/proxy_behavior/management/test_key_info_v2.py new file mode 100644 index 00000000000..b0fb27a19fa --- /dev/null +++ b/tests/proxy_behavior/management/test_key_info_v2.py @@ -0,0 +1,82 @@ +import uuid + +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /v2/key/info resolves the posted keys, then drops any key the caller +# cannot see via _can_user_query_key_info — silently, no 403. A non-admin sees +# a key it owns (user_id match) or a key whose team it belongs to. The world's +# TEAM_ALPHA members all see each other's keys; CROSS_ORG_USER and the org +# admins see only their own. The request is posted with every world key, and +# the returned info set is asserted to equal the visible subset. +_ALPHA_KEYS = frozenset( + { + Actor.TEAM_ADMIN, + Actor.INTERNAL_USER, + Actor.OWNER, + Actor.UNRELATED_SAME_ORG, + Actor.SERVICE_ACCOUNT, + } +) +_VISIBILITY = { + Actor.PROXY_ADMIN: frozenset(Actor), + Actor.ORG_ADMIN: frozenset({Actor.ORG_ADMIN}), + Actor.TEAM_ADMIN: _ALPHA_KEYS, + Actor.INTERNAL_USER: _ALPHA_KEYS, + Actor.OWNER: _ALPHA_KEYS, + Actor.UNRELATED_SAME_ORG: _ALPHA_KEYS, + Actor.SERVICE_ACCOUNT: _ALPHA_KEYS, + Actor.CROSS_ORG_USER: frozenset({Actor.CROSS_ORG_USER}), + Actor.ORG_B_ADMIN: frozenset({Actor.ORG_B_ADMIN}), +} + + +@pytest.mark.parametrize( + "actor,expected_visible", + list(_VISIBILITY.items()), + ids=[a.value for a in _VISIBILITY], +) +async def test_key_info_v2_visibility(actor, expected_visible, proxy_client, world): + caller = world.keys[actor] + user_id_to_actor = {world.keys[a].user_id: a for a in Actor} + + resp = await proxy_client.post( + "/v2/key/info", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"keys": [world.keys[a].cleartext for a in Actor]}, + ) + assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}" + + visible = { + user_id_to_actor[entry["user_id"]] + for entry in resp.json()["info"] + if entry.get("user_id") in user_id_to_actor + } + assert visible == set(expected_visible), ( + f"{actor.value}: expected {sorted(a.value for a in expected_visible)}, " + f"got {sorted(a.value for a in visible)}" + ) + + +async def test_key_info_v2_no_body_is_422(proxy_client, world): + """A request with no body is a 422 — the handler has no keys to resolve.""" + resp = await proxy_client.post( + "/v2/key/info", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 422, resp.text + + +async def test_key_info_v2_unknown_key_returns_empty_info(proxy_client, world): + """Keys that resolve to no rows yield an empty info list, not an error.""" + resp = await proxy_client.post( + "/v2/key/info", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"keys": ["sk-" + uuid.uuid4().hex]}, + ) + assert resp.status_code == 200, resp.text + assert resp.json()["info"] == [] diff --git a/tests/proxy_behavior/management/test_key_list.py b/tests/proxy_behavior/management/test_key_list.py new file mode 100644 index 00000000000..0ed101d5868 --- /dev/null +++ b/tests/proxy_behavior/management/test_key_list.py @@ -0,0 +1,171 @@ +from typing import FrozenSet + +import pytest + +from litellm.proxy.utils import hash_token + +from .actors import TEAM_ALPHA, Actor +from .conftest import create_scratch_key + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# Pinned default visibility for /key/list (no filter params): each actor's +# expected set of seeded actor keys. +_VISIBILITY = { + Actor.PROXY_ADMIN: frozenset(Actor), + Actor.ORG_ADMIN: frozenset({Actor.ORG_ADMIN}), + Actor.TEAM_ADMIN: frozenset({Actor.TEAM_ADMIN}), + Actor.INTERNAL_USER: frozenset({Actor.INTERNAL_USER}), + Actor.OWNER: frozenset({Actor.OWNER}), + Actor.UNRELATED_SAME_ORG: frozenset({Actor.UNRELATED_SAME_ORG}), + Actor.CROSS_ORG_USER: frozenset({Actor.CROSS_ORG_USER}), + Actor.SERVICE_ACCOUNT: frozenset({Actor.SERVICE_ACCOUNT}), +} + + +async def _all_visible_hashes(proxy_client, caller_cleartext) -> set: + """Walk every /key/list page — size is capped at 100 by the endpoint, so a + single request can truncate PROXY_ADMIN's view on a non-fresh DB.""" + hashes: set = set() + page = 1 + while True: + resp = await proxy_client.get( + f"/key/list?page={page}&size=100", + headers={"Authorization": f"Bearer {caller_cleartext}"}, + ) + assert resp.status_code == 200, resp.text + body = resp.json() + for entry in body.get("keys", []): + tok = entry.get("token") if isinstance(entry, dict) else entry + if tok: + hashes.add(tok) + if page >= (body.get("total_pages") or 1): + return hashes + page += 1 + + +@pytest.mark.parametrize( + "actor,expected_visible", + list(_VISIBILITY.items()), + ids=[a.value for a in _VISIBILITY], +) +async def test_key_list_visibility( + actor: Actor, expected_visible: FrozenSet[Actor], proxy_client, world +): + caller = world.keys[actor] + hashed_to_actor = {world.keys[a].hashed: a for a in Actor} + + returned_hashes = await _all_visible_hashes(proxy_client, caller.cleartext) + visible_seeded = { + hashed_to_actor[h] for h in returned_hashes if h in hashed_to_actor + } + assert visible_seeded == set(expected_visible), ( + f"{actor.value}: expected {sorted(a.value for a in expected_visible)}, " + f"got {sorted(a.value for a in visible_seeded)}" + ) + + +async def _list_hashes(proxy_client, caller_cleartext: str, query: str) -> set: + resp = await proxy_client.get( + f"/key/list?{query}&size=100", + headers={"Authorization": f"Bearer {caller_cleartext}"}, + ) + assert resp.status_code == 200, resp.text + hashes: set = set() + for entry in resp.json().get("keys", []): + tok = entry.get("token") if isinstance(entry, dict) else entry + if tok: + hashes.add(tok) + return hashes + + +async def test_key_list_admin_key_alias_substring_match(proxy_client, scratch, world): + """A PROXY_ADMIN's key_alias filter is a case-insensitive substring match; + a narrower fragment selects the subset whose alias contains it.""" + admin = world.keys[Actor.PROXY_ADMIN] + a = await create_scratch_key( + proxy_client, + admin.cleartext, + scratch.prefix, + user_id=admin.user_id, + key_alias=f"{scratch.prefix}-sub-a", + ) + b = await create_scratch_key( + proxy_client, + admin.cleartext, + scratch.prefix, + user_id=admin.user_id, + key_alias=f"{scratch.prefix}-sub-b", + ) + seeded = {hash_token(a), hash_token(b)} + + broad = await _list_hashes( + proxy_client, admin.cleartext, f"key_alias={scratch.prefix}-sub" + ) + assert broad & seeded == seeded + + narrow = await _list_hashes( + proxy_client, admin.cleartext, f"key_alias={scratch.prefix}-sub-a" + ) + assert narrow & seeded == {hash_token(a)} + + +async def test_key_list_non_admin_key_alias_is_exact_match( + proxy_client, scratch, world +): + """A non-admin's key_alias filter is exact-match only — substring filtering + is restricted to admins. The full alias matches; a fragment does not.""" + caller = world.keys[Actor.INTERNAL_USER] + alias = f"{scratch.prefix}-exact" + key = await create_scratch_key( + proxy_client, + world.keys[Actor.PROXY_ADMIN].cleartext, + scratch.prefix, + user_id=caller.user_id, + key_alias=alias, + ) + key_hash = hash_token(key) + + exact = await _list_hashes(proxy_client, caller.cleartext, f"key_alias={alias}") + assert key_hash in exact + + fragment = await _list_hashes( + proxy_client, caller.cleartext, f"key_alias={scratch.prefix}-exac" + ) + assert key_hash not in fragment + + +async def test_key_list_team_id_filter(proxy_client, scratch, world): + """A team_id filter narrows the listing to keys of that team.""" + admin = world.keys[Actor.PROXY_ADMIN] + team_key = await create_scratch_key( + proxy_client, + admin.cleartext, + scratch.prefix, + user_id=world.keys[Actor.OWNER].user_id, + team_id=TEAM_ALPHA, + key_alias=f"{scratch.prefix}-team", + ) + no_team_key = await create_scratch_key( + proxy_client, + admin.cleartext, + scratch.prefix, + user_id=admin.user_id, + key_alias=f"{scratch.prefix}-noteam", + ) + + hashes = await _list_hashes(proxy_client, admin.cleartext, f"team_id={TEAM_ALPHA}") + assert hash_token(team_key) in hashes + assert hash_token(no_team_key) not in hashes + + +async def test_key_list_non_admin_cannot_filter_other_team(proxy_client, world): + """A non-admin filtering by a team it does not belong to is rejected 403.""" + resp = await proxy_client.get( + f"/key/list?team_id={world.team_beta_id}", + headers={ + "Authorization": f"Bearer {world.keys[Actor.INTERNAL_USER].cleartext}" + }, + ) + assert resp.status_code == 403, resp.text diff --git a/tests/proxy_behavior/management/test_key_regenerate.py b/tests/proxy_behavior/management/test_key_regenerate.py new file mode 100644 index 00000000000..724b8b6d65b --- /dev/null +++ b/tests/proxy_behavior/management/test_key_regenerate.py @@ -0,0 +1,165 @@ +import litellm +import pytest + +from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, +) + +from .actors import TEAM_ALPHA, TEAM_BETA, Actor +from .conftest import create_scratch_key + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# Most denials route through team_member_permission (401), unlike /key/update +# which goes through user_id-mismatch (403). The matrix surfaces that +# divergence between the two endpoints. +_SCENARIOS = [ + ("self/proxy_admin", Actor.PROXY_ADMIN, "self", 200), + ("self/org_admin", Actor.ORG_ADMIN, "self", 401), + ("self/team_admin", Actor.TEAM_ADMIN, "self", 200), + ("self/internal_user", Actor.INTERNAL_USER, "self", 200), + ("self/owner", Actor.OWNER, "self", 200), + ("self/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "self", 200), + ("self/cross_org_user", Actor.CROSS_ORG_USER, "self", 200), + ("self/service_account", Actor.SERVICE_ACCOUNT, "self", 200), + ("owner_target/proxy_admin", Actor.PROXY_ADMIN, "owner", 200), + ("owner_target/org_admin", Actor.ORG_ADMIN, "owner", 401), + ("owner_target/team_admin", Actor.TEAM_ADMIN, "owner", 200), + ("owner_target/internal_user", Actor.INTERNAL_USER, "owner", 401), + ("owner_target/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "owner", 401), + ("owner_target/cross_org_user", Actor.CROSS_ORG_USER, "owner", 401), + ("owner_target/service_account", Actor.SERVICE_ACCOUNT, "owner", 401), + ("cross_org_target/proxy_admin", Actor.PROXY_ADMIN, "cross_org", 200), + ("cross_org_target/org_admin", Actor.ORG_ADMIN, "cross_org", 401), + ("cross_org_target/team_admin", Actor.TEAM_ADMIN, "cross_org", 401), + ("cross_org_target/owner", Actor.OWNER, "cross_org", 401), + ("cross_org_target/cross_org_user", Actor.CROSS_ORG_USER, "cross_org", 401), + ("cross_org_target/service_account", Actor.SERVICE_ACCOUNT, "cross_org", 401), +] + + +async def _info(proxy_client, cleartext: str): + return await proxy_client.get( + "/key/info", headers={"Authorization": f"Bearer {cleartext}"} + ) + + +@pytest.mark.parametrize( + "actor,target_shape,expected_status", + [(a, t, s) for (_id, a, t, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_key_regenerate_authz_matrix( + actor: Actor, + target_shape: str, + expected_status: int, + proxy_client, + scratch, + world, +): + caller = world.keys[actor] + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + + if target_shape == "self": + target_cleartext = await create_scratch_key( + proxy_client, seeder, scratch.prefix, user_id=caller.user_id + ) + elif target_shape == "owner": + target_cleartext = await create_scratch_key( + proxy_client, + seeder, + scratch.prefix, + user_id=world.keys[Actor.OWNER].user_id, + team_id=TEAM_ALPHA, + ) + elif target_shape == "cross_org": + target_cleartext = await create_scratch_key( + proxy_client, + seeder, + scratch.prefix, + user_id=world.keys[Actor.CROSS_ORG_USER].user_id, + team_id=TEAM_BETA, + ) + else: + pytest.fail(f"unknown target_shape={target_shape}") + + resp = await proxy_client.post( + "/key/regenerate", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"key": target_cleartext}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}" + + if expected_status == 200: + new_cleartext = resp.json()["key"] + assert new_cleartext.startswith("sk-") and new_cleartext != target_cleartext + assert (await _info(proxy_client, target_cleartext)).status_code == 401 + assert (await _info(proxy_client, new_cleartext)).status_code == 200 + else: + # Denied: rotation must not have leaked — old cleartext still works. + assert (await _info(proxy_client, target_cleartext)).status_code == 200 + + +async def test_key_path_regenerate_smoke(proxy_client, scratch, world): + """Pins that POST /key/{key:path}/regenerate shares the same handler.""" + caller = world.keys[Actor.PROXY_ADMIN] + target_cleartext = await create_scratch_key( + proxy_client, caller.cleartext, scratch.prefix, user_id=caller.user_id + ) + + resp = await proxy_client.post( + f"/key/{target_cleartext}/regenerate", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={}, + ) + assert resp.status_code == 200, resp.text + new_cleartext = resp.json()["key"] + assert new_cleartext.startswith("sk-") and new_cleartext != target_cleartext + assert (await _info(proxy_client, target_cleartext)).status_code == 401 + assert (await _info(proxy_client, new_cleartext)).status_code == 200 + + +async def test_key_regenerate_enforces_upperbound_key_params( + proxy_client, scratch, world, monkeypatch +): + """Regenerate runs _enforce_upperbound_key_params: a max_budget above + litellm.upperbound_key_generate_params is rejected 400, a value within the + bound is accepted. Pins #26340 (db8ef44323) — regenerate previously + bypassed the upperbound. upperbound_key_generate_params is module-level + litellm.* state, so monkeypatch save/restores it.""" + admin = world.keys[Actor.PROXY_ADMIN] + over_key = await create_scratch_key( + proxy_client, + admin.cleartext, + scratch.prefix, + user_id=admin.user_id, + key_alias=f"{scratch.prefix}-over", + ) + within_key = await create_scratch_key( + proxy_client, + admin.cleartext, + scratch.prefix, + user_id=admin.user_id, + key_alias=f"{scratch.prefix}-within", + ) + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams(max_budget=100.0), + ) + headers = {"Authorization": f"Bearer {admin.cleartext}"} + + over = await proxy_client.post( + "/key/regenerate", headers=headers, json={"key": over_key, "max_budget": 500.0} + ) + assert over.status_code == 400, over.text + + within = await proxy_client.post( + "/key/regenerate", + headers=headers, + json={"key": within_key, "max_budget": 50.0}, + ) + assert within.status_code == 200, within.text diff --git a/tests/proxy_behavior/management/test_key_reset_spend.py b/tests/proxy_behavior/management/test_key_reset_spend.py new file mode 100644 index 00000000000..fb1c266f655 --- /dev/null +++ b/tests/proxy_behavior/management/test_key_reset_spend.py @@ -0,0 +1,136 @@ +import uuid + +import pytest + +from litellm.proxy.utils import hash_token + +from .actors import TEAM_ALPHA, TEAM_BETA, Actor +from .conftest import create_scratch_key + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_SEED_SPEND = 5.0 +_RESET_TO = 2.0 + + +# POST /key/{key}/reset_spend. The target key is pre-seeded with spend=5.0 so +# reset_to=2.0 always clears _validate_reset_spend_value (which runs before +# authz). _check_proxy_or_team_admin_for_key then allows only PROXY_ADMIN or a +# team admin of the key's team — there is no org-admin branch, and a teamless +# "self" key has no team to admin. ORG_ADMIN-role callers are stopped 401 at +# the management-route gate before the handler runs. +_SCENARIOS = [ + ("self/proxy_admin", Actor.PROXY_ADMIN, "self", 200), + ("self/org_admin", Actor.ORG_ADMIN, "self", 401), + ("self/team_admin", Actor.TEAM_ADMIN, "self", 403), + ("self/internal_user", Actor.INTERNAL_USER, "self", 403), + ("self/cross_org_user", Actor.CROSS_ORG_USER, "self", 403), + ("team_alpha/proxy_admin", Actor.PROXY_ADMIN, "team_alpha", 200), + ("team_alpha/org_admin", Actor.ORG_ADMIN, "team_alpha", 401), + ("team_alpha/team_admin", Actor.TEAM_ADMIN, "team_alpha", 200), + ("team_alpha/internal_user", Actor.INTERNAL_USER, "team_alpha", 403), + ("team_alpha/owner", Actor.OWNER, "team_alpha", 403), + ("team_alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "team_alpha", 403), + ("team_alpha/cross_org_user", Actor.CROSS_ORG_USER, "team_alpha", 403), + ("team_alpha/service_account", Actor.SERVICE_ACCOUNT, "team_alpha", 403), + ("team_alpha/org_b_admin", Actor.ORG_B_ADMIN, "team_alpha", 401), + ("team_beta/proxy_admin", Actor.PROXY_ADMIN, "team_beta", 200), + ("team_beta/org_admin", Actor.ORG_ADMIN, "team_beta", 401), + ("team_beta/team_admin", Actor.TEAM_ADMIN, "team_beta", 403), + ("team_beta/cross_org_user", Actor.CROSS_ORG_USER, "team_beta", 403), + ("team_beta/org_b_admin", Actor.ORG_B_ADMIN, "team_beta", 401), +] + + +async def _seed_target(proxy_client, seeder, prefix, world, shape, caller) -> str: + if shape == "self": + return await create_scratch_key( + proxy_client, seeder, prefix, user_id=caller.user_id + ) + if shape == "team_alpha": + return await create_scratch_key( + proxy_client, + seeder, + prefix, + user_id=world.keys[Actor.OWNER].user_id, + team_id=TEAM_ALPHA, + ) + if shape == "team_beta": + return await create_scratch_key( + proxy_client, + seeder, + prefix, + user_id=world.keys[Actor.CROSS_ORG_USER].user_id, + team_id=TEAM_BETA, + ) + pytest.fail(f"unknown shape={shape}") # pragma: no cover + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_key_reset_spend_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + caller = world.keys[actor] + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + target = await _seed_target( + proxy_client, seeder, scratch.prefix, world, shape, caller + ) + hashed = hash_token(target) + await prisma.db.litellm_verificationtoken.update( + where={"token": hashed}, data={"spend": _SEED_SPEND} + ) + + resp = await proxy_client.post( + f"/key/{target}/reset_spend", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"reset_to": _RESET_TO}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed}) + assert row is not None + if expected_status == 200: + assert row.spend == _RESET_TO + else: + assert row.spend == _SEED_SPEND, "denied but spend reset" + + +@pytest.mark.parametrize( + "actor", [Actor.PROXY_ADMIN, Actor.TEAM_ADMIN], ids=["proxy_admin", "team_admin"] +) +async def test_key_reset_spend_missing_key_is_404(actor: Actor, proxy_client, world): + """A well-formed but unseeded key is 404 before any spend validation.""" + resp = await proxy_client.post( + f"/key/sk-{uuid.uuid4().hex}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 404, resp.text + + +async def test_key_reset_spend_above_current_spend_is_400( + proxy_client, prisma, scratch, world +): + """reset_to above the key's current spend is rejected 400.""" + admin = world.keys[Actor.PROXY_ADMIN] + target = await create_scratch_key( + proxy_client, admin.cleartext, scratch.prefix, user_id=admin.user_id + ) + resp = await proxy_client.post( + f"/key/{target}/reset_spend", + headers={"Authorization": f"Bearer {admin.cleartext}"}, + json={"reset_to": 1.0}, + ) + assert resp.status_code == 400, resp.text diff --git a/tests/proxy_behavior/management/test_key_service_account_generate.py b/tests/proxy_behavior/management/test_key_service_account_generate.py new file mode 100644 index 00000000000..3b5bbe39754 --- /dev/null +++ b/tests/proxy_behavior/management/test_key_service_account_generate.py @@ -0,0 +1,98 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /key/service-account/generate. PROXY_ADMIN always passes. ORG_ADMIN-role +# callers are stopped 401 by the management-route gate (the body carries a +# team_id but no organization_id, so the org-admin route branch never matches). +# INTERNAL_USER-role callers reach the handler: a team admin of the target team +# passes (200); a "user"-role member is 401 (no service-account-generate +# permission); a non-member is 400 ("not assigned to team"). A request with no +# team_id is 400 ("team_id is required") for every actor that reaches the handler. +_SCENARIOS = [ + ("own/proxy_admin", Actor.PROXY_ADMIN, "own", 200), + ("own/org_admin", Actor.ORG_ADMIN, "own", 401), + ("own/team_admin", Actor.TEAM_ADMIN, "own", 200), + ("own/internal_user", Actor.INTERNAL_USER, "own", 401), + ("own/owner", Actor.OWNER, "own", 401), + ("own/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "own", 401), + ("own/cross_org_user", Actor.CROSS_ORG_USER, "own", 400), + ("own/service_account", Actor.SERVICE_ACCOUNT, "own", 401), + ("own/org_b_admin", Actor.ORG_B_ADMIN, "own", 401), + ("cross_org/proxy_admin", Actor.PROXY_ADMIN, "cross_org", 200), + ("cross_org/org_admin", Actor.ORG_ADMIN, "cross_org", 401), + ("cross_org/team_admin", Actor.TEAM_ADMIN, "cross_org", 400), + ("cross_org/internal_user", Actor.INTERNAL_USER, "cross_org", 400), + ("cross_org/cross_org_user", Actor.CROSS_ORG_USER, "cross_org", 401), + ("cross_org/org_b_admin", Actor.ORG_B_ADMIN, "cross_org", 401), + ("none/proxy_admin", Actor.PROXY_ADMIN, "none", 400), + ("none/org_admin", Actor.ORG_ADMIN, "none", 401), + ("none/team_admin", Actor.TEAM_ADMIN, "none", 400), + ("none/internal_user", Actor.INTERNAL_USER, "none", 400), + ("none/cross_org_user", Actor.CROSS_ORG_USER, "none", 400), +] + + +@pytest.mark.parametrize( + "actor,team_target,expected_status", + [(a, t, s) for (_id, a, t, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_key_service_account_generate_authz_matrix( + actor: Actor, + team_target: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + caller = world.keys[actor] + team_id = { + "own": world.team_alpha_id, + "cross_org": world.team_beta_id, + "none": None, + }[team_target] + + body = {"key_alias": scratch.prefix} + if team_id is not None: + body["team_id"] = team_id + + resp = await proxy_client.post( + "/key/service-account/generate", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json=body, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {team_target}: {resp.status_code} {resp.text}" + + rows = await prisma.db.litellm_verificationtoken.find_many( + where={"key_alias": scratch.prefix} + ) + if expected_status == 200: + assert len(rows) == 1 + # A service-account key belongs to the team, not a user. + assert rows[0].user_id is None + assert rows[0].team_id == team_id + else: + assert rows == [], f"{actor.value}: denied but key row leaked" + + +async def test_key_service_account_generate_unknown_team_is_400( + proxy_client, prisma, scratch, world +): + """A team_id absent from the database is rejected 400.""" + resp = await proxy_client.post( + "/key/service-account/generate", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"key_alias": scratch.prefix, "team_id": scratch.tag("no-such-team")}, + ) + assert resp.status_code == 400, resp.text + rows = await prisma.db.litellm_verificationtoken.find_many( + where={"key_alias": scratch.prefix} + ) + assert rows == [] diff --git a/tests/proxy_behavior/management/test_key_update.py b/tests/proxy_behavior/management/test_key_update.py new file mode 100644 index 00000000000..7b7f6f5558b --- /dev/null +++ b/tests/proxy_behavior/management/test_key_update.py @@ -0,0 +1,184 @@ +import uuid + +import pytest + +from litellm.proxy.utils import hash_token + +from .actors import TEAM_ALPHA, TEAM_BETA, Actor +from .conftest import create_scratch_key + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# (id, actor, target_shape, expected_status). Pinned against current gating: +# proxy_admin bypasses; org_admin is blocked by an early role gate (401); +# every other (INTERNAL_USER-roled) actor hits user_id-mismatch 403, no-team- +# admin 403, or team_member_permission 401 depending on target / membership. +_SCENARIOS = [ + ("self/proxy_admin", Actor.PROXY_ADMIN, "self", 200), + ("self/org_admin", Actor.ORG_ADMIN, "self", 401), + ("self/team_admin", Actor.TEAM_ADMIN, "self", 403), + ("self/internal_user", Actor.INTERNAL_USER, "self", 403), + ("self/owner", Actor.OWNER, "self", 403), + ("self/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "self", 403), + ("self/cross_org_user", Actor.CROSS_ORG_USER, "self", 403), + ("self/service_account", Actor.SERVICE_ACCOUNT, "self", 403), + ("owner_target/proxy_admin", Actor.PROXY_ADMIN, "owner", 200), + ("owner_target/org_admin", Actor.ORG_ADMIN, "owner", 401), + ("owner_target/team_admin", Actor.TEAM_ADMIN, "owner", 403), + ("owner_target/internal_user", Actor.INTERNAL_USER, "owner", 403), + ("owner_target/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "owner", 403), + ("owner_target/cross_org_user", Actor.CROSS_ORG_USER, "owner", 403), + ("owner_target/service_account", Actor.SERVICE_ACCOUNT, "owner", 403), + ("cross_org_target/proxy_admin", Actor.PROXY_ADMIN, "cross_org", 200), + ("cross_org_target/org_admin", Actor.ORG_ADMIN, "cross_org", 401), + ("cross_org_target/team_admin", Actor.TEAM_ADMIN, "cross_org", 403), + ("cross_org_target/owner", Actor.OWNER, "cross_org", 403), + ("cross_org_target/cross_org_user", Actor.CROSS_ORG_USER, "cross_org", 401), + ("cross_org_target/service_account", Actor.SERVICE_ACCOUNT, "cross_org", 403), +] + +MARKER_MODEL = "behavior-pin-update-marker-model" + + +@pytest.mark.parametrize( + "actor,target_shape,expected_status", + [(a, t, s) for (_id, a, t, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_key_update_authz_matrix( + actor: Actor, + target_shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + caller = world.keys[actor] + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + + if target_shape == "self": + target_cleartext = await create_scratch_key( + proxy_client, seeder, scratch.prefix, user_id=caller.user_id + ) + elif target_shape == "owner": + target_cleartext = await create_scratch_key( + proxy_client, + seeder, + scratch.prefix, + user_id=world.keys[Actor.OWNER].user_id, + team_id=TEAM_ALPHA, + ) + elif target_shape == "cross_org": + target_cleartext = await create_scratch_key( + proxy_client, + seeder, + scratch.prefix, + user_id=world.keys[Actor.CROSS_ORG_USER].user_id, + team_id=TEAM_BETA, + ) + else: + pytest.fail(f"unknown target_shape={target_shape}") + + target_hashed = hash_token(target_cleartext) + + resp = await proxy_client.post( + "/key/update", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"key": target_cleartext, "models": [MARKER_MODEL]}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_verificationtoken.find_unique( + where={"token": target_hashed} + ) + assert row is not None + if expected_status == 200: + assert row.models == [MARKER_MODEL] + else: + assert row.models != [MARKER_MODEL], "denied but row mutated" + + +async def _seed_shape(proxy_client, seeder, prefix, world, shape, caller) -> str: + if shape == "self": + return await create_scratch_key( + proxy_client, seeder, prefix, user_id=caller.user_id + ) + if shape == "owner": + return await create_scratch_key( + proxy_client, + seeder, + prefix, + user_id=world.keys[Actor.OWNER].user_id, + team_id=TEAM_ALPHA, + ) + if shape == "cross_org": + return await create_scratch_key( + proxy_client, + seeder, + prefix, + user_id=world.keys[Actor.CROSS_ORG_USER].user_id, + team_id=TEAM_BETA, + ) + pytest.fail(f"unknown shape={shape}") # pragma: no cover + + +async def test_key_update_missing_key_is_404(proxy_client, world): + """An update targeting a key absent from the DB is a 404 — not 401/403.""" + resp = await proxy_client.post( + "/key/update", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"key": "sk-" + uuid.uuid4().hex, "models": [MARKER_MODEL]}, + ) + assert resp.status_code == 404, resp.text + + +# A denied /key/update must not partially apply: the budget/limit columns are +# left untouched. Each scenario is a denial cell from the matrix above. +_DENIED_BUDGET = [ + ("team_admin/self", Actor.TEAM_ADMIN, "self", 403), + ("internal_user/owner", Actor.INTERNAL_USER, "owner", 403), + ("cross_org_user/cross_org", Actor.CROSS_ORG_USER, "cross_org", 401), +] + + +@pytest.mark.parametrize( + "actor,target_shape,expected_status", + [(a, t, s) for (_id, a, t, s) in _DENIED_BUDGET], + ids=[s[0] for s in _DENIED_BUDGET], +) +async def test_key_update_denied_does_not_touch_budget_counters( + actor: Actor, + target_shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + caller = world.keys[actor] + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + target = await _seed_shape( + proxy_client, seeder, scratch.prefix, world, target_shape, caller + ) + target_hashed = hash_token(target) + + resp = await proxy_client.post( + "/key/update", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"key": target, "max_budget": 999.0, "tpm_limit": 888, "rpm_limit": 777}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {target_shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_verificationtoken.find_unique( + where={"token": target_hashed} + ) + assert row is not None + assert row.max_budget is None, "denied but max_budget applied" + assert row.tpm_limit is None, "denied but tpm_limit applied" + assert row.rpm_limit is None, "denied but rpm_limit applied" diff --git a/tests/proxy_behavior/management/test_no_management_imports.py b/tests/proxy_behavior/management/test_no_management_imports.py new file mode 100644 index 00000000000..f8c52a1c37e --- /dev/null +++ b/tests/proxy_behavior/management/test_no_management_imports.py @@ -0,0 +1,46 @@ +import pathlib +import re + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +BEHAVIOR_DIR = REPO_ROOT / "tests" / "proxy_behavior" + +FORBIDDEN_IMPORT = re.compile(r"^\s*from\s+litellm\.proxy\.management_endpoints\b") +FORBIDDEN_AUTH_MOCK = re.compile( + r"(?:mock\.[A-Za-z_]+|patch[a-z_]*)\([^)]*user_api_key_auth" +) +# This file is the only place the forbidden patterns appear as regex source; +# exclude it so it can describe what it forbids. +SELF = pathlib.Path(__file__).resolve() + + +def _iter_py_files(): + for path in BEHAVIOR_DIR.rglob("*.py"): + if path.resolve() != SELF: + yield path + + +def _scan(pattern): + violations = [] + for path in _iter_py_files(): + for lineno, line in enumerate(path.read_text().splitlines(), start=1): + if pattern.search(line): + violations.append( + f"{path.relative_to(REPO_ROOT)}:{lineno}: {line.strip()}" + ) + return violations + + +def test_no_management_endpoint_imports(): + violations = _scan(FORBIDDEN_IMPORT) + assert not violations, ( + "tests/proxy_behavior/ must not import from litellm.proxy.management_endpoints. " + "Violations:\n " + "\n ".join(violations) + ) + + +def test_no_user_api_key_auth_mocking(): + violations = _scan(FORBIDDEN_AUTH_MOCK) + assert not violations, ( + "tests/proxy_behavior/ must not mock user_api_key_auth. " + "Violations:\n " + "\n ".join(violations) + ) diff --git a/tests/proxy_behavior/management/test_route_coverage.py b/tests/proxy_behavior/management/test_route_coverage.py new file mode 100644 index 00000000000..1139e251a59 --- /dev/null +++ b/tests/proxy_behavior/management/test_route_coverage.py @@ -0,0 +1,91 @@ +"""PR3.M1 — codified route coverage. + +Every route declared in the two management-endpoint source files must be +exercised by at least one behavior-suite scenario. This is a permanent +regression guard: a future route added without a behavior test fails CI here, +the same way test_no_management_imports.py codifies the G3 import grep. +""" + +import ast +import pathlib +import re + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +SOURCE_FILES = [ + REPO_ROOT / "litellm/proxy/management_endpoints/key_management_endpoints.py", + REPO_ROOT / "litellm/proxy/management_endpoints/team_endpoints.py", +] +TEST_DIR = pathlib.Path(__file__).resolve().parent +SELF = pathlib.Path(__file__).resolve() + +# Captures the route literal from `@router.(""` — `\s*` spans +# newlines so multi-line decorators are matched too. +_ROUTE_DECORATOR = re.compile( + r"@router\.(?:get|post|put|delete|patch)\(\s*[\"']([^\"']+)[\"']" +) + + +def _source_routes() -> set: + routes: set = set() + for path in SOURCE_FILES: + routes.update(_ROUTE_DECORATOR.findall(path.read_text())) + return routes + + +def _route_to_regex(route: str) -> re.Pattern: + # A plain path param ({team_id}) matches a single path segment; a Starlette + # ':path' param ({key:path}) matches across '/'. Keeping plain params + # slash-bounded stops a loose regex from falsely reporting a future + # multi-segment route as already covered. + pattern = ["^"] + pos = 0 + for match in re.finditer(r"\{([^}]+)\}", route): + pattern.append(re.escape(route[pos : match.start()])) + pattern.append("[^?]+" if match.group(1).endswith(":path") else "[^/?]+") + pos = match.end() + pattern.append(re.escape(route[pos:]) + "$") + return re.compile("".join(pattern)) + + +def _test_urls() -> set: + """Every request-URL string literal across the behavior test suite. + + f-strings are reconstructed with each interpolation collapsed to a single + placeholder char, so f"/key/{target}/regenerate" becomes /key/X/regenerate. + Query strings are dropped — coverage is a path-level property. + """ + urls: set = set() + for path in sorted(TEST_DIR.glob("test_*.py")): + if path.resolve() == SELF: + continue + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + literal = None + if isinstance(node, ast.Constant) and isinstance(node.value, str): + literal = node.value + elif isinstance(node, ast.JoinedStr): + chunks = [] + for value in node.values: + if isinstance(value, ast.Constant) and isinstance(value.value, str): + chunks.append(value.value) + else: + chunks.append("X") # interpolated path / query segment + literal = "".join(chunks) + if literal and literal.startswith("/"): + urls.add(literal.split("?", 1)[0]) + return urls + + +def test_every_management_route_has_a_behavior_scenario(): + routes = _source_routes() + assert routes, "no @router routes parsed — the decorator regex is stale" + + urls = _test_urls() + uncovered = sorted( + route + for route in routes + if not any(_route_to_regex(route).match(url) for url in urls) + ) + assert ( + not uncovered + ), "management routes with no behavior-suite scenario:\n " + "\n ".join(uncovered) diff --git a/tests/proxy_behavior/management/test_scratch_teardown.py b/tests/proxy_behavior/management/test_scratch_teardown.py new file mode 100644 index 00000000000..bcb53935558 --- /dev/null +++ b/tests/proxy_behavior/management/test_scratch_teardown.py @@ -0,0 +1,61 @@ +import pytest + +from litellm.proxy._types import LitellmUserRoles + +from .actors import ORG_A, ORG_B +from .conftest import MASTER_KEY, SCRATCH_PREFIX, create_scratch_actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# The minting tests run in file order, then _b runs after their fixture +# teardown and asserts no scratch row survived in any reclaimed table. A leak +# in either direction fails _b on the next collection. + + +async def test_a_scratch_key_lands_in_db(proxy_client, prisma, scratch): + resp = await proxy_client.post( + "/key/generate", + headers={"Authorization": f"Bearer {MASTER_KEY}"}, + json={"key_alias": scratch.prefix}, + ) + assert resp.status_code == 200, resp.text + + rows = await prisma.db.litellm_verificationtoken.find_many( + where={"key_alias": scratch.prefix} + ) + assert len(rows) == 1 + + +async def test_a2_scratch_actor_lands_in_db(proxy_client, prisma, scratch): + actor = await create_scratch_actor( + prisma, + scratch.prefix, + user_role=LitellmUserRoles.ORG_ADMIN.value, + org_admin_of=(ORG_A, ORG_B), + ) + user_row = await prisma.db.litellm_usertable.find_unique( + where={"user_id": actor.user_id} + ) + assert user_row is not None + info = await proxy_client.get( + "/key/info", headers={"Authorization": f"Bearer {actor.cleartext}"} + ) + assert info.status_code == 200, info.text + memberships = await prisma.db.litellm_organizationmembership.find_many( + where={"user_id": actor.user_id} + ) + assert {m.organization_id for m in memberships} == {ORG_A, ORG_B} + + +async def test_b_scratch_namespace_is_clean(prisma): + tokens = await prisma.db.litellm_verificationtoken.find_many( + where={"key_alias": {"startswith": SCRATCH_PREFIX}} + ) + users = await prisma.db.litellm_usertable.find_many( + where={"user_id": {"startswith": SCRATCH_PREFIX}} + ) + memberships = await prisma.db.litellm_organizationmembership.find_many( + where={"user_id": {"startswith": SCRATCH_PREFIX}} + ) + assert tokens == [] and users == [] and memberships == [] diff --git a/tests/proxy_behavior/management/test_smoke.py b/tests/proxy_behavior/management/test_smoke.py new file mode 100644 index 00000000000..4e90986ad9f --- /dev/null +++ b/tests/proxy_behavior/management/test_smoke.py @@ -0,0 +1,28 @@ +import pytest + +from .conftest import MASTER_KEY + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +async def test_liveliness(proxy_client): + resp = await proxy_client.get("/health/liveliness") + assert resp.status_code == 200 + + +async def test_key_generate_lands_in_db(proxy_client, prisma, scratch): + from litellm.proxy.utils import hash_token + + resp = await proxy_client.post( + "/key/generate", + headers={"Authorization": f"Bearer {MASTER_KEY}"}, + json={"key_alias": scratch.prefix}, + ) + assert resp.status_code == 200, resp.text + cleartext = resp.json()["key"] + assert cleartext.startswith("sk-") + + hashed = hash_token(cleartext) + row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed}) + assert row is not None + assert row.token == hashed != cleartext diff --git a/tests/proxy_behavior/management/test_team_available.py b/tests/proxy_behavior/management/test_team_available.py new file mode 100644 index 00000000000..874c8dd4df7 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_available.py @@ -0,0 +1,21 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# GET /team/available lists teams from +# litellm.default_internal_user_params["available_teams"]. The behavior world +# configures no available_teams, so the handler returns [] for every actor +# before it even reads the caller — this is the route-coverage + default-path +# pin. /team/available is an info route, so every authenticated actor reaches +# the handler. +@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor]) +async def test_team_available_default_is_empty(actor: Actor, proxy_client, world): + resp = await proxy_client.get( + "/team/available", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + ) + assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}" + assert resp.json() == [] diff --git a/tests/proxy_behavior/management/test_team_block_unblock.py b/tests/proxy_behavior/management/test_team_block_unblock.py new file mode 100644 index 00000000000..9412e51b909 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_block_unblock.py @@ -0,0 +1,114 @@ +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /team/block + /team/unblock. The handler gate is _verify_team_access +# (proxy admin / team admin / org admin), but the management-route gate fronts +# it: the request carries the team's organization_id so an org admin of that +# org clears the gate's org-scoped branch. A team admin is an INTERNAL_USER +# and these are not internal_user routes, so a team admin can never reach the +# handler — only PROXY_ADMIN and an org admin of the team's own org pass. +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 401), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 401), + ("alpha/owner", Actor.OWNER, "alpha", 401), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 401), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 401), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 401), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 401), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 401), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_target(prisma, world, shape: str, team_id: str) -> str: + """Raw-seed the scratch target team; returns its organization_id.""" + org_id = world.org_a_id if shape == "alpha" else world.org_b_id + await create_scratch_team(prisma, team_id, organization_id=org_id) + return org_id + + +@pytest.mark.parametrize("route", ["block", "unblock"]) +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_block_unblock_authz_matrix( + route: str, + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + org_id = await _seed_target(prisma, world, shape, scratch.prefix) + caller = world.keys[actor] + + # /unblock starts from a blocked row so a 200 is observable as True->False. + if route == "unblock": + await prisma.db.litellm_teamtable.update( + where={"team_id": scratch.prefix}, data={"blocked": True} + ) + + resp = await proxy_client.post( + f"/team/{route}", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"team_id": scratch.prefix, "organization_id": org_id}, + ) + assert ( + resp.status_code == expected_status + ), f"{route} {actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert bool(row.blocked) is (route == "block") + else: + assert bool(row.blocked) is (route == "unblock"), "denied but blocked mutated" + + +async def test_team_block_unblock_round_trip(proxy_client, prisma, scratch, world): + """PROXY_ADMIN block then unblock flips the blocked column True then False.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + headers = {"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"} + + blocked = await proxy_client.post( + "/team/block", headers=headers, json={"team_id": scratch.prefix} + ) + assert blocked.status_code == 200, blocked.text + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None and row.blocked is True + + unblocked = await proxy_client.post( + "/team/unblock", headers=headers, json={"team_id": scratch.prefix} + ) + assert unblocked.status_code == 200, unblocked.text + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None and row.blocked is False + + +@pytest.mark.parametrize("route", ["block", "unblock"]) +async def test_team_block_unblock_missing_team_is_404(route: str, proxy_client, world): + """A team_id absent from the DB is 404 — the existence check precedes authz.""" + resp = await proxy_client.post( + f"/team/{route}", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"team_id": "behavior-pin-no-such-team"}, + ) + assert resp.status_code == 404, resp.text diff --git a/tests/proxy_behavior/management/test_team_bulk_member_add.py b/tests/proxy_behavior/management/test_team_bulk_member_add.py new file mode 100644 index 00000000000..fc83cd414e5 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_bulk_member_add.py @@ -0,0 +1,105 @@ +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +def _member_ids(row) -> list: + return [m["user_id"] for m in (row.members_with_roles or [])] + + +async def test_team_bulk_member_add_proxy_admin_adds_explicit_members( + proxy_client, prisma, scratch, world +): + """PROXY_ADMIN bulk-adds an explicit member list to a scratch team.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + new_member = scratch.tag("m1") + resp = await proxy_client.post( + "/team/bulk_member_add", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={ + "team_id": scratch.prefix, + "members": [{"user_id": new_member, "role": "user"}], + }, + ) + assert resp.status_code == 200, resp.text + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None and new_member in _member_ids(row) + + +async def test_team_bulk_member_add_empty_members_is_400( + proxy_client, prisma, scratch, world +): + """An empty member list (with all_users unset) is rejected 400.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + "/team/bulk_member_add", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"team_id": scratch.prefix, "members": []}, + ) + assert resp.status_code == 400, resp.text + + +async def test_team_bulk_member_add_over_max_batch_is_400( + proxy_client, prisma, scratch, world +): + """A member list larger than the 500-member cap is rejected 400.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + members = [ + {"user_id": f"{scratch.prefix}-u{i}", "role": "user"} for i in range(501) + ] + resp = await proxy_client.post( + "/team/bulk_member_add", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"team_id": scratch.prefix, "members": members}, + ) + assert resp.status_code == 400, resp.text + + +@pytest.mark.parametrize( + "actor", + [Actor.TEAM_ADMIN, Actor.INTERNAL_USER], + ids=["team_admin", "internal_user"], +) +async def test_team_bulk_member_add_non_admin_is_401( + actor: Actor, proxy_client, prisma, scratch, world +): + """/team/bulk_member_add is neither an internal_user nor a self-managed + route — a non-proxy-admin with no org context is 401 at the route gate.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + "/team/bulk_member_add", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + json={ + "team_id": scratch.prefix, + "members": [{"user_id": scratch.tag("m"), "role": "user"}], + }, + ) + assert resp.status_code == 401, f"{actor.value}: {resp.status_code} {resp.text}" + + +async def test_team_bulk_member_add_all_users_proxy_admin( + proxy_client, prisma, scratch, world +): + """all_users=True pulls every user in the DB into the team. The route is + reachable only by PROXY_ADMIN (the route gate 401s every other actor — even + an org admin with organization_id in the body), so the handler's own + all_users PROXY_ADMIN gate is never the deciding check at the boundary.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + "/team/bulk_member_add", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"team_id": scratch.prefix, "all_users": True}, + ) + assert resp.status_code == 200, resp.text + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + member_ids = _member_ids(row) + # every world actor is a user in the DB, so all are now team members + assert world.keys[Actor.INTERNAL_USER].user_id in member_ids diff --git a/tests/proxy_behavior/management/test_team_daily_activity.py b/tests/proxy_behavior/management/test_team_daily_activity.py new file mode 100644 index 00000000000..7a1e70b91fc --- /dev/null +++ b/tests/proxy_behavior/management/test_team_daily_activity.py @@ -0,0 +1,63 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# GET /team/daily/activity. A proxy admin (admin view) sees activity for any +# team. A non-admin is scoped to user_info.teams: a bare query defaults to its +# own teams (200), and an explicit team_ids filter naming a team it does not +# belong to is 404 (the VERIA-43 fix). Org admins have no team memberships, so +# they behave like a non-member for any specific team. +_MEMBERS = { + "alpha": { + Actor.TEAM_ADMIN, + Actor.INTERNAL_USER, + Actor.OWNER, + Actor.UNRELATED_SAME_ORG, + Actor.SERVICE_ACCOUNT, + }, + "beta": {Actor.CROSS_ORG_USER}, +} + + +def _expected(actor: Actor, team: str) -> int: + if team == "none" or actor == Actor.PROXY_ADMIN: + return 200 + return 200 if actor in _MEMBERS.get(team, set()) else 404 + + +_CASES = [ + (f"{team}/{actor.value}", actor, team, _expected(actor, team)) + for team in ("none", "alpha", "beta") + for actor in Actor +] + + +# start_date / end_date are required by the handler — pin only the team-scope +# authz, not the date validation. +_DATES = "start_date=2024-01-01&end_date=2024-12-31" + + +@pytest.mark.parametrize( + "actor,team,expected_status", + [(a, t, s) for (_id, a, t, s) in _CASES], + ids=[c[0] for c in _CASES], +) +async def test_team_daily_activity_matrix( + actor: Actor, team: str, expected_status: int, proxy_client, world +): + query = _DATES + if team == "alpha": + query += f"&team_ids={world.team_alpha_id}" + elif team == "beta": + query += f"&team_ids={world.team_beta_id}" + + resp = await proxy_client.get( + f"/team/daily/activity?{query}", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} -> {team}: {resp.status_code} {resp.text}" diff --git a/tests/proxy_behavior/management/test_team_delete.py b/tests/proxy_behavior/management/test_team_delete.py new file mode 100644 index 00000000000..bbf0a6563f3 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_delete.py @@ -0,0 +1,78 @@ +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /team/delete runs per-team _verify_team_access. The request carries the +# team's organization_id so an org admin of that org clears the management- +# route gate; a team admin is an INTERNAL_USER on a non-internal_user route, +# so a team admin never reaches the handler. Only PROXY_ADMIN and an org admin +# of the team's own org can delete it. +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 401), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 401), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 401), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 401), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 401), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_delete_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + org_id = world.org_a_id if shape == "alpha" else world.org_b_id + await create_scratch_team(prisma, scratch.prefix, organization_id=org_id) + caller = world.keys[actor] + + resp = await proxy_client.post( + "/team/delete", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"team_ids": [scratch.prefix], "organization_id": org_id}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + if expected_status == 200: + assert row is None, "deleted but team row survives" + else: + assert row is not None, "denied but team row vanished" + + +async def test_team_delete_batch_with_missing_id_deletes_nothing( + proxy_client, prisma, scratch, world +): + """A batch is validated whole before any deletion: one missing team_id + fails the request 404 and the accessible team in the batch survives.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + "/team/delete", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"team_ids": [scratch.prefix, "behavior-pin-no-such-team"]}, + ) + assert resp.status_code == 404, resp.text + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None, "batch aborted but the accessible team was deleted" diff --git a/tests/proxy_behavior/management/test_team_filter_ui.py b/tests/proxy_behavior/management/test_team_filter_ui.py new file mode 100644 index 00000000000..69cbabf72a1 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_filter_ui.py @@ -0,0 +1,39 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# GET /team/filter/ui (ui_view_teams) — include_in_schema=False. The handler +# body has no role/org check and never reads user_api_key_dict, but the +# endpoint is still effectively PROXY-ADMIN-only as its docstring claims: the +# management-route gate fronts it (not an internal_user / info / org-admin +# route) and 401s every non-proxy-admin before the handler runs. PROXY_ADMIN +# reaches the unscoped find_many and sees teams across every org. +@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor]) +async def test_team_filter_ui_is_proxy_admin_only(actor: Actor, proxy_client, world): + resp = await proxy_client.get( + "/team/filter/ui", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + ) + expected = 200 if actor == Actor.PROXY_ADMIN else 401 + assert ( + resp.status_code == expected + ), f"{actor.value}: {resp.status_code} {resp.text}" + + +async def test_team_filter_ui_proxy_admin_sees_cross_org_teams(proxy_client, world): + """The handler runs an unscoped query — PROXY_ADMIN sees teams from every + org, including the three seeded world teams.""" + resp = await proxy_client.get( + "/team/filter/ui", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 200, resp.text + team_ids = {t.get("team_id") for t in resp.json() if isinstance(t, dict)} + assert { + world.team_alpha_id, + world.team_beta_id, + world.team_gamma_id, + } <= team_ids diff --git a/tests/proxy_behavior/management/test_team_info.py b/tests/proxy_behavior/management/test_team_info.py new file mode 100644 index 00000000000..51809942113 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_info.py @@ -0,0 +1,70 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# GET /team/info — actor x team-target authz matrix, pinned against +# validate_membership(): a team is readable by a proxy admin, a key whose +# own team_id matches, a listed member, or an org admin of the team's org; +# everything else is 403. TEAM_GAMMA has no members, so only PROXY_ADMIN +# and ORG_A's org admin can read it. +_SCENARIOS = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 200), + ("alpha/owner", Actor.OWNER, "alpha", 200), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 200), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 200), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("gamma/proxy_admin", Actor.PROXY_ADMIN, "gamma", 200), + ("gamma/org_admin", Actor.ORG_ADMIN, "gamma", 200), + ("gamma/team_admin", Actor.TEAM_ADMIN, "gamma", 403), + ("gamma/internal_user", Actor.INTERNAL_USER, "gamma", 403), + ("gamma/owner", Actor.OWNER, "gamma", 403), + ("gamma/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "gamma", 403), + ("gamma/cross_org_user", Actor.CROSS_ORG_USER, "gamma", 403), + ("gamma/service_account", Actor.SERVICE_ACCOUNT, "gamma", 403), + ("gamma/org_b_admin", Actor.ORG_B_ADMIN, "gamma", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/internal_user", Actor.INTERNAL_USER, "beta", 403), + ("beta/owner", Actor.OWNER, "beta", 403), + ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 403), + ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 200), + ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +@pytest.mark.parametrize( + "actor,target,expected_status", + [(a, t, s) for (_id, a, t, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_team_info_authz_matrix( + actor: Actor, target: str, expected_status: int, proxy_client, world +): + caller = world.keys[actor] + target_team_id = { + "alpha": world.team_alpha_id, + "gamma": world.team_gamma_id, + "beta": world.team_beta_id, + }[target] + + resp = await proxy_client.get( + f"/team/info?team_id={target_team_id}", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} -> {target}: {resp.status_code} {resp.text}" + + if expected_status == 200: + body = resp.json() + assert body["team_id"] == target_team_id + assert body["team_info"]["team_id"] == target_team_id diff --git a/tests/proxy_behavior/management/test_team_key_bulk_update.py b/tests/proxy_behavior/management/test_team_key_bulk_update.py new file mode 100644 index 00000000000..5acf0c8185c --- /dev/null +++ b/tests/proxy_behavior/management/test_team_key_bulk_update.py @@ -0,0 +1,217 @@ +import uuid + +import pytest + +from litellm.proxy._types import KeyManagementRoutes +from litellm.proxy.utils import hash_token + +from .actors import Actor +from .conftest import create_scratch_key, create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_MARKER_BUDGET = 42.0 +_KEY_UPDATE = KeyManagementRoutes.KEY_UPDATE.value + + +# POST /team/key/bulk_update — PROXY_ADMIN bypasses; otherwise +# can_team_member_execute_key_management_endpoint runs with route=KEY_UPDATE. +# A team admin always passes; a "user"-role member passes only when the team's +# team_member_permissions grants /key/update; a non-member is 401. ORG_ADMIN is +# stopped 401 at the management-route gate before the handler (the body has a +# team_id but no organization_id, so the org-admin route branch never matches). +_MATRIX = [ + ("admin/proxy_admin", Actor.PROXY_ADMIN, "admin", 200), + ("admin/internal_user", Actor.INTERNAL_USER, "admin", 200), + ("member_allowed/internal_user", Actor.INTERNAL_USER, "member_allowed", 200), + ("member_denied/internal_user", Actor.INTERNAL_USER, "member_denied", 401), + ("nonmember/internal_user", Actor.INTERNAL_USER, "nonmember", 401), + ("nonmember/org_admin", Actor.ORG_ADMIN, "nonmember", 401), + ("nonmember/proxy_admin", Actor.PROXY_ADMIN, "nonmember", 200), +] + + +async def _seed_team_key(prisma, proxy_client, prefix: str, world, shape: str) -> str: + """Raw-seed the scratch team for `shape`, return a team key's cleartext.""" + internal = world.keys[Actor.INTERNAL_USER].user_id + owner = world.keys[Actor.OWNER].user_id + if shape == "admin": + await create_scratch_team( + prisma, prefix, organization_id=world.org_a_id, admin_user_ids=[internal] + ) + key_owner = internal + elif shape == "member_allowed": + await create_scratch_team( + prisma, + prefix, + organization_id=world.org_a_id, + admin_user_ids=[owner], + member_user_ids=[internal], + team_member_permissions=[_KEY_UPDATE], + ) + key_owner = owner + elif shape == "member_denied": + await create_scratch_team( + prisma, + prefix, + organization_id=world.org_a_id, + admin_user_ids=[owner], + member_user_ids=[internal], + team_member_permissions=[], + ) + key_owner = owner + elif shape == "nonmember": + await create_scratch_team( + prisma, prefix, organization_id=world.org_a_id, admin_user_ids=[owner] + ) + key_owner = owner + else: + pytest.fail(f"unknown shape={shape}") # pragma: no cover + return await create_scratch_key( + proxy_client, + world.keys[Actor.PROXY_ADMIN].cleartext, + prefix, + user_id=key_owner, + team_id=prefix, + ) + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_key_bulk_update_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + key = await _seed_team_key(prisma, proxy_client, scratch.prefix, world, shape) + hashed = hash_token(key) + caller = world.keys[actor] + + resp = await proxy_client.post( + "/team/key/bulk_update", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={ + "team_id": scratch.prefix, + "key_ids": [key], + "update_fields": {"max_budget": _MARKER_BUDGET}, + }, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_verificationtoken.find_unique(where={"token": hashed}) + assert row is not None + if expected_status == 200: + assert len(resp.json()["successful_updates"]) == 1 + assert row.max_budget == _MARKER_BUDGET + else: + assert row.max_budget != _MARKER_BUDGET, "denied but key mutated" + + +async def test_team_key_bulk_update_requires_team_id( + proxy_client, prisma, scratch, world +): + """An empty team_id is rejected 400.""" + resp = await proxy_client.post( + "/team/key/bulk_update", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={ + "team_id": "", + "key_ids": ["sk-" + uuid.uuid4().hex], + "update_fields": {"max_budget": _MARKER_BUDGET}, + }, + ) + assert resp.status_code == 400, resp.text + + +async def test_team_key_bulk_update_all_keys_in_team( + proxy_client, prisma, scratch, world +): + """all_keys_in_team=True broadcasts the update to every key in the team.""" + admin = world.keys[Actor.PROXY_ADMIN].cleartext + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + keys = [ + await create_scratch_key( + proxy_client, + admin, + scratch.prefix, + user_id=world.keys[Actor.OWNER].user_id, + team_id=scratch.prefix, + key_alias=f"{scratch.prefix}-k{i}", + ) + for i in range(2) + ] + + resp = await proxy_client.post( + "/team/key/bulk_update", + headers={"Authorization": f"Bearer {admin}"}, + json={ + "team_id": scratch.prefix, + "all_keys_in_team": True, + "update_fields": {"max_budget": _MARKER_BUDGET}, + }, + ) + assert resp.status_code == 200, resp.text + assert len(resp.json()["successful_updates"]) == 2 + for key in keys: + row = await prisma.db.litellm_verificationtoken.find_unique( + where={"token": hash_token(key)} + ) + assert row is not None and row.max_budget == _MARKER_BUDGET + + +async def test_team_key_bulk_update_no_keys_found_is_404( + proxy_client, prisma, scratch, world +): + """all_keys_in_team=True on a team with no keys is a top-level 404.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + "/team/key/bulk_update", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={ + "team_id": scratch.prefix, + "all_keys_in_team": True, + "update_fields": {"max_budget": _MARKER_BUDGET}, + }, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_key_bulk_update_missing_key_is_isolated( + proxy_client, prisma, scratch, world +): + """A key_id absent from the team lands in failed_updates; the batch still + returns 200 and the real key is updated.""" + admin = world.keys[Actor.PROXY_ADMIN].cleartext + real = await _seed_team_key( + prisma, proxy_client, scratch.prefix, world, "nonmember" + ) + missing = "sk-" + uuid.uuid4().hex + + resp = await proxy_client.post( + "/team/key/bulk_update", + headers={"Authorization": f"Bearer {admin}"}, + json={ + "team_id": scratch.prefix, + "key_ids": [real, missing], + "update_fields": {"max_budget": _MARKER_BUDGET}, + }, + ) + assert resp.status_code == 200, resp.text + body = resp.json() + assert body["total_requested"] == 2 + assert len(body["successful_updates"]) == 1 + assert len(body["failed_updates"]) == 1 + + row = await prisma.db.litellm_verificationtoken.find_unique( + where={"token": hash_token(real)} + ) + assert row is not None and row.max_budget == _MARKER_BUDGET diff --git a/tests/proxy_behavior/management/test_team_list.py b/tests/proxy_behavior/management/test_team_list.py new file mode 100644 index 00000000000..2bd106dd2d0 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_list.py @@ -0,0 +1,105 @@ +from typing import FrozenSet, Optional + +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# The behavior DB may hold teams beyond the three seeded ones, so every +# assertion intersects the returned team_ids with the known seeded set. +def _seeded_visible(resp_json, world) -> set: + known = { + world.team_alpha_id: "alpha", + world.team_beta_id: "beta", + world.team_gamma_id: "gamma", + } + return { + known[entry["team_id"]] + for entry in resp_json + if isinstance(entry, dict) and entry.get("team_id") in known + } + + +# Family 1 — bare GET /team/list (no query params). _authorize_and_filter_teams +# authorizes only an admin view (proxy admin) or an org admin; everyone else +# is 401. An org admin sees every team in its org(s). +_BARE = [ + ("proxy_admin", Actor.PROXY_ADMIN, 200, {"alpha", "beta", "gamma"}), + ("org_admin", Actor.ORG_ADMIN, 200, {"alpha", "gamma"}), + ("team_admin", Actor.TEAM_ADMIN, 401, None), + ("internal_user", Actor.INTERNAL_USER, 401, None), + ("owner", Actor.OWNER, 401, None), + ("unrelated_same_org", Actor.UNRELATED_SAME_ORG, 401, None), + ("cross_org_user", Actor.CROSS_ORG_USER, 401, None), + ("service_account", Actor.SERVICE_ACCOUNT, 401, None), + ("org_b_admin", Actor.ORG_B_ADMIN, 200, {"beta"}), +] + + +@pytest.mark.parametrize( + "actor,expected_status,expected_visible", + [(a, s, v) for (_id, a, s, v) in _BARE], + ids=[s[0] for s in _BARE], +) +async def test_team_list_bare_authz( + actor: Actor, + expected_status: int, + expected_visible: Optional[set], + proxy_client, + world, +): + caller = world.keys[actor] + resp = await proxy_client.get( + "/team/list", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value}: {resp.status_code} {resp.text}" + + if expected_status == 200: + visible = _seeded_visible(resp.json(), world) + assert visible == expected_visible, ( + f"{actor.value}: expected {sorted(expected_visible)}, " + f"got {sorted(visible)}" + ) + + +# Family 2 — GET /team/list?user_id= ("own query"). Every +# actor may query its own teams (200); the result is exactly the teams it +# belongs to. A user_id filter scopes proxy/org admins to their own +# membership too — the broad admin view from family 1 does not carry over. +_OWN = { + Actor.PROXY_ADMIN: frozenset(), + Actor.ORG_ADMIN: frozenset(), + Actor.TEAM_ADMIN: frozenset({"alpha"}), + Actor.INTERNAL_USER: frozenset({"alpha"}), + Actor.OWNER: frozenset({"alpha"}), + Actor.UNRELATED_SAME_ORG: frozenset({"alpha"}), + Actor.CROSS_ORG_USER: frozenset({"beta"}), + Actor.SERVICE_ACCOUNT: frozenset({"alpha"}), + Actor.ORG_B_ADMIN: frozenset(), +} + + +@pytest.mark.parametrize( + "actor,expected_visible", + list(_OWN.items()), + ids=[a.value for a in _OWN], +) +async def test_team_list_own_query( + actor: Actor, expected_visible: FrozenSet[str], proxy_client, world +): + caller = world.keys[actor] + resp = await proxy_client.get( + f"/team/list?user_id={caller.user_id}", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert resp.status_code == 200, f"{actor.value}: {resp.status_code} {resp.text}" + + visible = _seeded_visible(resp.json(), world) + assert visible == set(expected_visible), ( + f"{actor.value}: expected {sorted(expected_visible)}, " f"got {sorted(visible)}" + ) diff --git a/tests/proxy_behavior/management/test_team_list_v2.py b/tests/proxy_behavior/management/test_team_list_v2.py new file mode 100644 index 00000000000..81178ad73c0 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_list_v2.py @@ -0,0 +1,141 @@ +from typing import FrozenSet, Optional + +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +def _seeded(team_ids: set, world) -> set: + known = { + world.team_alpha_id: "alpha", + world.team_beta_id: "beta", + world.team_gamma_id: "gamma", + } + return {known[t] for t in team_ids if t in known} + + +async def _v2_team_ids(proxy_client, caller_cleartext: str, extra: str = "") -> set: + """Walk every /v2/team/list page and collect the returned team_ids.""" + ids: set = set() + page = 1 + while True: + resp = await proxy_client.get( + f"/v2/team/list?page={page}&page_size=100{extra}", + headers={"Authorization": f"Bearer {caller_cleartext}"}, + ) + assert resp.status_code == 200, resp.text + body = resp.json() + teams = body.get("teams", []) or [] + for t in teams: + tid = t.get("team_id") if isinstance(t, dict) else None + if tid: + ids.add(tid) + if page * 100 >= (body.get("total") or 0) or not teams: + return ids + page += 1 + + +# GET /v2/team/list is an info route reachable by every actor, but +# _enforce_list_team_v2_access still gates a BARE query: a proxy admin sees +# all teams, an org admin sees its orgs' teams, and a regular user — who has +# passed no user_id filter — is rejected 401 ("only admins can query all +# teams"). A regular user must scope the query to its own user_id. +_BARE = [ + ("proxy_admin", Actor.PROXY_ADMIN, 200, frozenset({"alpha", "beta", "gamma"})), + ("org_admin", Actor.ORG_ADMIN, 200, frozenset({"alpha", "gamma"})), + ("org_b_admin", Actor.ORG_B_ADMIN, 200, frozenset({"beta"})), + ("team_admin", Actor.TEAM_ADMIN, 401, None), + ("internal_user", Actor.INTERNAL_USER, 401, None), + ("owner", Actor.OWNER, 401, None), + ("unrelated_same_org", Actor.UNRELATED_SAME_ORG, 401, None), + ("cross_org_user", Actor.CROSS_ORG_USER, 401, None), + ("service_account", Actor.SERVICE_ACCOUNT, 401, None), +] + + +@pytest.mark.parametrize( + "actor,expected_status,expected_visible", + [(a, s, v) for (_id, a, s, v) in _BARE], + ids=[s[0] for s in _BARE], +) +async def test_team_list_v2_bare( + actor: Actor, + expected_status: int, + expected_visible: Optional[FrozenSet[str]], + proxy_client, + world, +): + caller = world.keys[actor] + if expected_status != 200: + resp = await proxy_client.get( + "/v2/team/list", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert resp.status_code == expected_status, resp.text + return + + visible = _seeded(await _v2_team_ids(proxy_client, caller.cleartext), world) + assert visible == set( + expected_visible + ), f"{actor.value}: expected {sorted(expected_visible)}, got {sorted(visible)}" + + +# A regular user scoping the query to its own user_id is allowed, and sees +# exactly the teams it belongs to. +_OWN = { + Actor.TEAM_ADMIN: frozenset({"alpha"}), + Actor.INTERNAL_USER: frozenset({"alpha"}), + Actor.OWNER: frozenset({"alpha"}), + Actor.UNRELATED_SAME_ORG: frozenset({"alpha"}), + Actor.CROSS_ORG_USER: frozenset({"beta"}), + Actor.SERVICE_ACCOUNT: frozenset({"alpha"}), +} + + +@pytest.mark.parametrize( + "actor,expected_visible", list(_OWN.items()), ids=[a.value for a in _OWN] +) +async def test_team_list_v2_own_user_id_query( + actor: Actor, expected_visible: FrozenSet[str], proxy_client, world +): + caller = world.keys[actor] + visible = _seeded( + await _v2_team_ids( + proxy_client, caller.cleartext, f"&user_id={caller.user_id}" + ), + world, + ) + assert visible == set( + expected_visible + ), f"{actor.value}: expected {sorted(expected_visible)}, got {sorted(visible)}" + + +async def test_team_list_v2_user_id_filter_other_user_is_401(proxy_client, world): + """A regular user filtering by another user's user_id is rejected 401.""" + resp = await proxy_client.get( + f"/v2/team/list?user_id={world.keys[Actor.OWNER].user_id}", + headers={ + "Authorization": f"Bearer {world.keys[Actor.INTERNAL_USER].cleartext}" + }, + ) + assert resp.status_code == 401, resp.text + + +async def test_team_list_v2_org_filter_foreign_org_is_403(proxy_client, world): + """An org admin filtering by an organization it does not administer is 403.""" + resp = await proxy_client.get( + f"/v2/team/list?organization_id={world.org_b_id}", + headers={"Authorization": f"Bearer {world.keys[Actor.ORG_ADMIN].cleartext}"}, + ) + assert resp.status_code == 403, resp.text + + +async def test_team_list_v2_invalid_status_is_400(proxy_client, world): + """status accepts only 'deleted' — any other value is 400.""" + resp = await proxy_client.get( + "/v2/team/list?status=bogus", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + ) + assert resp.status_code == 400, resp.text diff --git a/tests/proxy_behavior/management/test_team_member_add.py b/tests/proxy_behavior/management/test_team_member_add.py new file mode 100644 index 00000000000..a0dc4a7ecaf --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_add.py @@ -0,0 +1,149 @@ +import litellm +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /team/member_add — actor x team-shape matrix, pinned against +# _validate_team_member_add_permissions: PROXY_ADMIN, the team's team admin, +# or an org admin of the team's org may add members; everyone else is 403. +# Unlike /team/update there is no route gate in front, so the team-admin +# branch is reachable (TEAM_ADMIN, an internal_user, is allowed on its team). +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/internal_user", Actor.INTERNAL_USER, "beta", 403), + ("beta/owner", Actor.OWNER, "beta", 403), + ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 403), + ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 403), + ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_target(prisma, world, shape: str, team_id: str) -> None: + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + ) + elif shape == "beta": + await create_scratch_team(prisma, team_id, organization_id=world.org_b_id) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + + +def _member_ids(row) -> list: + return [m["user_id"] for m in (row.members_with_roles or [])] + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_add_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + await _seed_target(prisma, world, shape, scratch.prefix) + caller = world.keys[actor] + new_member_id = scratch.tag("newmember") + + resp = await proxy_client.post( + "/team/member_add", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={ + "team_id": scratch.prefix, + "member": {"user_id": new_member_id, "role": "user"}, + }, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert new_member_id in _member_ids(row) + else: + assert new_member_id not in _member_ids(row), "denied but member added" + + +# Available-team self-join: a non-admin caller may add ITSELF to a team listed +# in litellm.default_internal_user_params["available_teams"], but the bypass +# must not escalate to role=admin or inject another user. +_SELF_JOIN = [ + ("self_as_user", "self", "user", 200), + ("self_as_admin", "self", "admin", 403), + ("other_as_user", "other", "user", 403), +] + + +@pytest.mark.parametrize( + "who,role,expected_status", + [(w, r, s) for (_id, w, r, s) in _SELF_JOIN], + ids=[s[0] for s in _SELF_JOIN], +) +async def test_team_member_add_available_team_self_join( + who: str, + role: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, + monkeypatch, +): + # Org-less team with no admins: the INTERNAL_USER caller is neither team + # nor org admin, so it lands on the available-team branch. + await create_scratch_team(prisma, scratch.prefix) + monkeypatch.setattr( + litellm, "default_internal_user_params", {"available_teams": [scratch.prefix]} + ) + + caller = world.keys[Actor.INTERNAL_USER] + member_id = caller.user_id if who == "self" else world.keys[Actor.OWNER].user_id + + resp = await proxy_client.post( + "/team/member_add", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={ + "team_id": scratch.prefix, + "member": {"user_id": member_id, "role": role}, + }, + ) + assert ( + resp.status_code == expected_status + ), f"{who}/{role}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert member_id in _member_ids(row) + else: + assert member_id not in _member_ids(row), "denied but member added" diff --git a/tests/proxy_behavior/management/test_team_member_delete.py b/tests/proxy_behavior/management/test_team_member_delete.py new file mode 100644 index 00000000000..43879d9fd16 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_delete.py @@ -0,0 +1,92 @@ +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /team/member_delete — actor x team-shape matrix. The scratch team is +# raw-seeded with a victim member already in it; PROXY_ADMIN, the team's team +# admin, or an org admin of the team's org may remove members; else 403. +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/internal_user", Actor.INTERNAL_USER, "beta", 403), + ("beta/owner", Actor.OWNER, "beta", 403), + ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 403), + ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 403), + ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_target(prisma, world, shape: str, team_id: str, victim_id: str) -> None: + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + member_user_ids=[victim_id], + ) + elif shape == "beta": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_b_id, + member_user_ids=[victim_id], + ) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + + +def _member_ids(row) -> list: + return [m["user_id"] for m in (row.members_with_roles or [])] + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_delete_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + victim_id = scratch.tag("victim") + await _seed_target(prisma, world, shape, scratch.prefix, victim_id) + caller = world.keys[actor] + + resp = await proxy_client.post( + "/team/member_delete", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"team_id": scratch.prefix, "user_id": victim_id}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert victim_id not in _member_ids(row) + else: + assert victim_id in _member_ids(row), "denied but member removed" diff --git a/tests/proxy_behavior/management/test_team_member_me.py b/tests/proxy_behavior/management/test_team_member_me.py new file mode 100644 index 00000000000..bfbbe0504ae --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_me.py @@ -0,0 +1,83 @@ +import uuid + +import pytest + +from litellm.proxy.utils import hash_token + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# GET /team/{team_id}/members/me resolves the CALLER's own membership row. +# A caller that is not a member of the team is 404 — even PROXY_ADMIN, which +# is not in any seeded team. The route is self-managed, so every actor reaches +# the handler. TEAM_GAMMA has no members, so every actor is 404 there. +_MEMBERS = { + "alpha": { + Actor.TEAM_ADMIN, + Actor.INTERNAL_USER, + Actor.OWNER, + Actor.UNRELATED_SAME_ORG, + Actor.SERVICE_ACCOUNT, + }, + "beta": {Actor.CROSS_ORG_USER}, + "gamma": set(), +} + +_CASES = [ + (f"{team}/{actor.value}", actor, team, 200 if actor in members else 404) + for team, members in _MEMBERS.items() + for actor in Actor +] + + +@pytest.mark.parametrize( + "actor,team,expected_status", + [(a, t, s) for (_id, a, t, s) in _CASES], + ids=[c[0] for c in _CASES], +) +async def test_team_member_me_matrix( + actor: Actor, team: str, expected_status: int, proxy_client, world +): + team_id = { + "alpha": world.team_alpha_id, + "beta": world.team_beta_id, + "gamma": world.team_gamma_id, + }[team] + caller = world.keys[actor] + + resp = await proxy_client.get( + f"/team/{team_id}/members/me", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} -> {team}: {resp.status_code} {resp.text}" + + if expected_status == 200: + body = resp.json() + assert body["user_id"] == caller.user_id + assert body["team_id"] == team_id + + +async def test_team_member_me_team_key_without_user_id_is_400( + proxy_client, prisma, scratch, world +): + """A key with no associated user_id (a team / service-account key) cannot + resolve 'me' — the caller has no identity to look up — so it is 400.""" + cleartext = "sk-" + uuid.uuid4().hex + await prisma.db.litellm_verificationtoken.create( + data={ + "token": hash_token(cleartext), + "key_name": f"{scratch.prefix}-teamkey", + "key_alias": f"{scratch.prefix}-teamkey", + "team_id": world.team_alpha_id, + "models": [], + } + ) + resp = await proxy_client.get( + f"/team/{world.team_alpha_id}/members/me", + headers={"Authorization": f"Bearer {cleartext}"}, + ) + assert resp.status_code == 400, resp.text diff --git a/tests/proxy_behavior/management/test_team_member_update.py b/tests/proxy_behavior/management/test_team_member_update.py new file mode 100644 index 00000000000..53b245bd1e9 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_update.py @@ -0,0 +1,97 @@ +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /team/member_update — actor x team-shape matrix. The scratch team is +# raw-seeded with a "user"-role member; each scenario tries to promote it to +# "admin". PROXY_ADMIN, the team's team admin, or an org admin of the team's +# org may update members; else 403. (The harness forces premium_user, so the +# promotion does not hit the admin-role premium gate.) +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/internal_user", Actor.INTERNAL_USER, "beta", 403), + ("beta/owner", Actor.OWNER, "beta", 403), + ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 403), + ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 403), + ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None: + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + member_user_ids=[member_id], + ) + elif shape == "beta": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_b_id, + member_user_ids=[member_id], + ) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + + +def _role_of(row, user_id: str): + for m in row.members_with_roles or []: + if m["user_id"] == user_id: + return m["role"] + return None + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_update_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + member_id = scratch.tag("member") + await _seed_target(prisma, world, shape, scratch.prefix, member_id) + caller = world.keys[actor] + + resp = await proxy_client.post( + "/team/member_update", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"team_id": scratch.prefix, "user_id": member_id, "role": "admin"}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert _role_of(row, member_id) == "admin" + else: + assert _role_of(row, member_id) == "user", "denied but role changed" diff --git a/tests/proxy_behavior/management/test_team_model.py b/tests/proxy_behavior/management/test_team_model.py new file mode 100644 index 00000000000..3564e8df83a --- /dev/null +++ b/tests/proxy_behavior/management/test_team_model.py @@ -0,0 +1,78 @@ +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_MARKER_MODEL = "behavior-pin-team-model-marker" +_ROUTE_URL = {"add": "/team/model/add", "delete": "/team/model/delete"} + + +# POST /team/model/add + /team/model/delete. The handler gate is PROXY_ADMIN +# or team admin or org admin, but the management-route gate fronts it — these +# are neither internal_user nor org-admin nor info routes, so every +# non-proxy-admin is 401 before the handler runs. Only PROXY_ADMIN reaches the +# handler, making the team-admin / org-admin handler branches unreachable here. +_MATRIX = [ + ("proxy_admin", Actor.PROXY_ADMIN, 200), + ("org_admin", Actor.ORG_ADMIN, 401), + ("team_admin", Actor.TEAM_ADMIN, 401), + ("internal_user", Actor.INTERNAL_USER, 401), + ("owner", Actor.OWNER, 401), + ("unrelated_same_org", Actor.UNRELATED_SAME_ORG, 401), + ("cross_org_user", Actor.CROSS_ORG_USER, 401), + ("service_account", Actor.SERVICE_ACCOUNT, 401), + ("org_b_admin", Actor.ORG_B_ADMIN, 401), +] + + +@pytest.mark.parametrize("route", ["add", "delete"]) +@pytest.mark.parametrize( + "actor,expected_status", + [(a, s) for (_id, a, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_model_authz_matrix( + route: str, + actor: Actor, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + initial = [] if route == "add" else [_MARKER_MODEL] + await create_scratch_team( + prisma, scratch.prefix, organization_id=world.org_a_id, models=initial + ) + caller = world.keys[actor] + + resp = await proxy_client.post( + _ROUTE_URL[route], + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"team_id": scratch.prefix, "models": [_MARKER_MODEL]}, + ) + assert ( + resp.status_code == expected_status + ), f"{route} {actor.value}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert (_MARKER_MODEL in row.models) is (route == "add") + else: + assert list(row.models) == initial, "denied but models mutated" + + +@pytest.mark.parametrize("route", ["add", "delete"]) +async def test_team_model_missing_team_is_404(route: str, proxy_client, world): + """A team_id absent from the DB is 404 — the existence check precedes authz.""" + resp = await proxy_client.post( + _ROUTE_URL[route], + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"team_id": "behavior-pin-no-such-team", "models": [_MARKER_MODEL]}, + ) + assert resp.status_code == 404, resp.text diff --git a/tests/proxy_behavior/management/test_team_new.py b/tests/proxy_behavior/management/test_team_new.py new file mode 100644 index 00000000000..7b07f259641 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_new.py @@ -0,0 +1,139 @@ +from typing import Any, Dict + +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /team/new — actor x org-target matrix (org_target picks the request's +# organization_id: none / ORG_A / ORG_B). Pinned against the role gate, which +# 401s every denial: PROXY_ADMIN always passes; any other caller must name an +# organization_id AND be ORG_ADMIN of that org. +_SCENARIOS = [ + ("none/proxy_admin", Actor.PROXY_ADMIN, "none", 200), + ("none/org_admin", Actor.ORG_ADMIN, "none", 401), + ("none/team_admin", Actor.TEAM_ADMIN, "none", 401), + ("none/internal_user", Actor.INTERNAL_USER, "none", 401), + ("none/owner", Actor.OWNER, "none", 401), + ("none/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "none", 401), + ("none/cross_org_user", Actor.CROSS_ORG_USER, "none", 401), + ("none/service_account", Actor.SERVICE_ACCOUNT, "none", 401), + ("none/org_b_admin", Actor.ORG_B_ADMIN, "none", 401), + ("org_a/proxy_admin", Actor.PROXY_ADMIN, "org_a", 200), + ("org_a/org_admin", Actor.ORG_ADMIN, "org_a", 200), + ("org_a/team_admin", Actor.TEAM_ADMIN, "org_a", 401), + ("org_a/internal_user", Actor.INTERNAL_USER, "org_a", 401), + ("org_a/owner", Actor.OWNER, "org_a", 401), + ("org_a/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "org_a", 401), + ("org_a/cross_org_user", Actor.CROSS_ORG_USER, "org_a", 401), + ("org_a/service_account", Actor.SERVICE_ACCOUNT, "org_a", 401), + ("org_a/org_b_admin", Actor.ORG_B_ADMIN, "org_a", 401), + ("org_b/proxy_admin", Actor.PROXY_ADMIN, "org_b", 200), + ("org_b/org_admin", Actor.ORG_ADMIN, "org_b", 401), + ("org_b/team_admin", Actor.TEAM_ADMIN, "org_b", 401), + ("org_b/internal_user", Actor.INTERNAL_USER, "org_b", 401), + ("org_b/owner", Actor.OWNER, "org_b", 401), + ("org_b/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "org_b", 401), + ("org_b/cross_org_user", Actor.CROSS_ORG_USER, "org_b", 401), + ("org_b/service_account", Actor.SERVICE_ACCOUNT, "org_b", 401), + ("org_b/org_b_admin", Actor.ORG_B_ADMIN, "org_b", 200), +] + + +@pytest.mark.parametrize( + "actor,org_target,expected_status", + [(a, o, s) for (_id, a, o, s) in _SCENARIOS], + ids=[s[0] for s in _SCENARIOS], +) +async def test_team_new_authz_matrix( + actor: Actor, + org_target: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + caller = world.keys[actor] + org_id = { + "none": None, + "org_a": world.org_a_id, + "org_b": world.org_b_id, + }[org_target] + + body: Dict[str, Any] = {"team_id": scratch.prefix, "team_alias": scratch.prefix} + if org_id is not None: + body["organization_id"] = org_id + + resp = await proxy_client.post( + "/team/new", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json=body, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} org={org_target}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + if expected_status == 200: + assert row is not None + assert row.organization_id == org_id + else: + assert row is None, f"{actor.value}: denied but team row leaked" + + +async def test_team_new_rejects_negative_budget(proxy_client, prisma, scratch, world): + """Input-validation pin: max_budget < 0 is a 400, no row created.""" + resp = await proxy_client.post( + "/team/new", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"team_id": scratch.prefix, "max_budget": -1}, + ) + assert resp.status_code == 400, resp.text + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is None + + +async def test_team_new_rejects_duplicate_team_id(proxy_client, prisma, scratch, world): + """Input-validation pin: a colliding team_id is a 400 on the second call.""" + seeder = world.keys[Actor.PROXY_ADMIN].cleartext + first = await proxy_client.post( + "/team/new", + headers={"Authorization": f"Bearer {seeder}"}, + json={"team_id": scratch.prefix, "team_alias": scratch.prefix}, + ) + assert first.status_code == 200, first.text + + second = await proxy_client.post( + "/team/new", + headers={"Authorization": f"Bearer {seeder}"}, + json={"team_id": scratch.prefix, "team_alias": scratch.prefix}, + ) + assert second.status_code == 400, second.text + + +async def test_team_new_unknown_organization_is_500( + proxy_client, prisma, scratch, world +): + """SURFACED, NOT ENDORSED: a /team/new with an organization_id that does + not exist currently fails 500 (the role-resolution layer raises before + the handler's own 400 'Organization not found' check is reached).""" + resp = await proxy_client.post( + "/team/new", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={ + "team_id": scratch.prefix, + "organization_id": scratch.tag("no-such-org"), + }, + ) + assert resp.status_code == 500, resp.text + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is None diff --git a/tests/proxy_behavior/management/test_team_permissions.py b/tests/proxy_behavior/management/test_team_permissions.py new file mode 100644 index 00000000000..5d16702fe6c --- /dev/null +++ b/tests/proxy_behavior/management/test_team_permissions.py @@ -0,0 +1,170 @@ +import litellm +import pytest + +from litellm.proxy._types import KeyManagementRoutes + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_PERM = KeyManagementRoutes.KEY_INFO.value + + +# GET /team/permissions_list and POST /team/permissions_update are self-managed +# routes, so every actor reaches the handler. Both grant access to PROXY_ADMIN, +# the team admin, or an org admin of the team's org. The scratch team is in +# ORG_A with TEAM_ADMIN as its team admin. +_MATRIX = [ + ("proxy_admin", Actor.PROXY_ADMIN, 200), + ("org_admin", Actor.ORG_ADMIN, 200), + ("team_admin", Actor.TEAM_ADMIN, 200), + ("internal_user", Actor.INTERNAL_USER, 403), + ("owner", Actor.OWNER, 403), + ("unrelated_same_org", Actor.UNRELATED_SAME_ORG, 403), + ("cross_org_user", Actor.CROSS_ORG_USER, 403), + ("service_account", Actor.SERVICE_ACCOUNT, 403), + ("org_b_admin", Actor.ORG_B_ADMIN, 403), +] + + +async def _seed_team(prisma, scratch_prefix, world) -> None: + await create_scratch_team( + prisma, + scratch_prefix, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + member_user_ids=[ + world.keys[Actor.INTERNAL_USER].user_id, + world.keys[Actor.OWNER].user_id, + world.keys[Actor.UNRELATED_SAME_ORG].user_id, + world.keys[Actor.SERVICE_ACCOUNT].user_id, + ], + ) + + +@pytest.mark.parametrize( + "actor,expected_status", + [(a, s) for (_id, a, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_permissions_list_authz_matrix( + actor: Actor, expected_status: int, proxy_client, prisma, scratch, world +): + await _seed_team(prisma, scratch.prefix, world) + resp = await proxy_client.get( + f"/team/permissions_list?team_id={scratch.prefix}", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value}: {resp.status_code} {resp.text}" + if expected_status == 200: + assert resp.json()["team_id"] == scratch.prefix + + +@pytest.mark.parametrize( + "actor,expected_status", + [(a, s) for (_id, a, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_permissions_update_authz_matrix( + actor: Actor, expected_status: int, proxy_client, prisma, scratch, world +): + await _seed_team(prisma, scratch.prefix, world) + resp = await proxy_client.post( + "/team/permissions_update", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + json={"team_id": scratch.prefix, "team_member_permissions": [_PERM]}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert _PERM in (row.team_member_permissions or []) + else: + assert _PERM not in (row.team_member_permissions or []), "denied but mutated" + + +async def test_team_permissions_available_team_self_join_divergence( + proxy_client, prisma, scratch, world, monkeypatch +): + """permissions_list honours the available-team self-join — a non-admin can + READ an available team's permissions — but permissions_update deliberately + does not: the same caller is 403 on update. default_internal_user_params is + module-level litellm.* state, so monkeypatch save/restores it.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + monkeypatch.setattr( + litellm, "default_internal_user_params", {"available_teams": [scratch.prefix]} + ) + caller = world.keys[Actor.CROSS_ORG_USER] # non-admin, unrelated to the team + + listed = await proxy_client.get( + f"/team/permissions_list?team_id={scratch.prefix}", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + ) + assert listed.status_code == 200, listed.text + + updated = await proxy_client.post( + "/team/permissions_update", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"team_id": scratch.prefix, "team_member_permissions": [_PERM]}, + ) + assert updated.status_code == 403, updated.text + + +# POST /team/permissions_bulk_update is PROXY_ADMIN-only. ORG_ADMIN-role +# callers are stopped 401 by the management-route gate; INTERNAL_USER-role +# callers, on a route that is neither internal_user nor self-managed, are 401 +# there too — only PROXY_ADMIN reaches the handler's own admin gate. +_BULK_MATRIX = [ + ("proxy_admin", Actor.PROXY_ADMIN, 200), + ("org_admin", Actor.ORG_ADMIN, 401), + ("team_admin", Actor.TEAM_ADMIN, 401), + ("internal_user", Actor.INTERNAL_USER, 401), + ("cross_org_user", Actor.CROSS_ORG_USER, 401), + ("org_b_admin", Actor.ORG_B_ADMIN, 401), +] + + +@pytest.mark.parametrize( + "actor,expected_status", + [(a, s) for (_id, a, s) in _BULK_MATRIX], + ids=[s[0] for s in _BULK_MATRIX], +) +async def test_team_permissions_bulk_update_authz_matrix( + actor: Actor, expected_status: int, proxy_client, prisma, scratch, world +): + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + "/team/permissions_bulk_update", + headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, + json={"team_ids": [scratch.prefix], "permissions": [_PERM]}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert _PERM in (row.team_member_permissions or []) + else: + assert _PERM not in (row.team_member_permissions or []), "denied but mutated" + + +async def test_team_permissions_bulk_update_no_selector_is_400(proxy_client, world): + """Neither team_ids nor apply_to_all_teams is a 400.""" + resp = await proxy_client.post( + "/team/permissions_bulk_update", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"permissions": [_PERM]}, + ) + assert resp.status_code == 400, resp.text diff --git a/tests/proxy_behavior/management/test_team_update.py b/tests/proxy_behavior/management/test_team_update.py new file mode 100644 index 00000000000..9b21911cef2 --- /dev/null +++ b/tests/proxy_behavior/management/test_team_update.py @@ -0,0 +1,209 @@ +import pytest + +from litellm.proxy._types import LitellmUserRoles + +from .actors import Actor +from .conftest import create_scratch_actor, create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +# POST /team/update — actor x team-shape matrix (shapes built by _seed_target). +# Each request carries the team's own organization_id so a non-proxy-admin can +# reach the org-scoped branch of the route-permission gate (401 on denial), +# which fronts the handler's _verify_team_access. Only PROXY_ADMIN and an +# ORG_ADMIN of the team's org pass: an internal_user team admin is filtered by +# the route gate before _verify_team_access's team-admin branch is reached. +MARKER_ALIAS = "behavior-pin-update-marker-alias" + +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 401), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 401), + ("alpha/owner", Actor.OWNER, "alpha", 401), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 401), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 401), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 401), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 401), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 401), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 401), + ("beta/internal_user", Actor.INTERNAL_USER, "beta", 401), + ("beta/owner", Actor.OWNER, "beta", 401), + ("beta/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "beta", 401), + ("beta/cross_org_user", Actor.CROSS_ORG_USER, "beta", 401), + ("beta/service_account", Actor.SERVICE_ACCOUNT, "beta", 401), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_target(prisma, world, shape: str, team_id: str) -> str: + """Raw-seed the scratch target team; returns its organization_id.""" + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + member_user_ids=[ + world.keys[Actor.INTERNAL_USER].user_id, + world.keys[Actor.OWNER].user_id, + world.keys[Actor.UNRELATED_SAME_ORG].user_id, + world.keys[Actor.SERVICE_ACCOUNT].user_id, + ], + ) + return world.org_a_id + if shape == "beta": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_b_id, + member_user_ids=[world.keys[Actor.CROSS_ORG_USER].user_id], + ) + return world.org_b_id + pytest.fail(f"unknown shape={shape}") # pragma: no cover + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_update_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + org_id = await _seed_target(prisma, world, shape, scratch.prefix) + caller = world.keys[actor] + + resp = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={ + "team_id": scratch.prefix, + "team_alias": MARKER_ALIAS, + "organization_id": org_id, + }, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert row.team_alias == MARKER_ALIAS + else: + assert row.team_alias != MARKER_ALIAS, "denied but team mutated" + + +async def test_team_update_requires_proxy_admin_without_org_context( + proxy_client, prisma, scratch, world +): + """With no organization_id in the body the route gate has no org context + and falls back to proxy-admin-only: an org admin of the team's own org + is 401, PROXY_ADMIN is 200.""" + await _seed_target(prisma, world, "alpha", scratch.prefix) + + denied = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {world.keys[Actor.ORG_ADMIN].cleartext}"}, + json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS}, + ) + assert denied.status_code == 401, denied.text + + allowed = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"team_id": scratch.prefix, "team_alias": MARKER_ALIAS}, + ) + assert allowed.status_code == 200, allowed.text + + +# Relocation gate — moving a team to a different org. The scratch team starts +# in ORG_A; each scenario relocates it to ORG_B. PROXY_ADMIN bypasses; +# ORG_B_ADMIN clears the route gate (dest-org admin) but fails +# _verify_team_access on the source team (403); the rest fail the route gate +# (401). The relocation-*allowed* branch (caller is org admin of both orgs) is +# covered by test_team_update_org_relocation_allowed_for_dual_org_admin below. +_RELOCATION = [ + ("proxy_admin", Actor.PROXY_ADMIN, 200), + ("org_b_admin", Actor.ORG_B_ADMIN, 403), + ("org_admin", Actor.ORG_ADMIN, 401), + ("team_admin", Actor.TEAM_ADMIN, 401), + ("internal_user", Actor.INTERNAL_USER, 401), +] + + +@pytest.mark.parametrize( + "actor,expected_status", + [(a, s) for (_id, a, s) in _RELOCATION], + ids=[s[0] for s in _RELOCATION], +) +async def test_team_update_org_relocation_gate( + actor: Actor, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + await _seed_target(prisma, world, "alpha", scratch.prefix) + caller = world.keys[actor] + + resp = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"team_id": scratch.prefix, "organization_id": world.org_b_id}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teamtable.find_unique( + where={"team_id": scratch.prefix} + ) + assert row is not None + if expected_status == 200: + assert row.organization_id == world.org_b_id + else: + assert row.organization_id == world.org_a_id, "denied but team relocated" + + +async def test_team_update_org_relocation_allowed_for_dual_org_admin( + proxy_client, prisma, scratch, world +): + """Relocation-allowed branch: a caller who is org admin of BOTH the source + and destination org may relocate a team between them. Completes the + _RELOCATION matrix, whose allowed branch PR2 left open — no seeded actor is + a dual-org admin, so one is minted with create_scratch_actor.""" + actor = await create_scratch_actor( + prisma, + scratch.prefix, + user_role=LitellmUserRoles.ORG_ADMIN.value, + org_admin_of=(world.org_a_id, world.org_b_id), + ) + team_id = await create_scratch_team( + prisma, scratch.tag("team"), organization_id=world.org_a_id + ) + + resp = await proxy_client.post( + "/team/update", + headers={"Authorization": f"Bearer {actor.cleartext}"}, + json={"team_id": team_id, "organization_id": world.org_b_id}, + ) + assert resp.status_code == 200, resp.text + + row = await prisma.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + assert row is not None + assert ( + row.organization_id == world.org_b_id + ), "dual-org admin relocation not applied" diff --git a/tests/proxy_behavior/management/test_world_seed.py b/tests/proxy_behavior/management/test_world_seed.py new file mode 100644 index 00000000000..00f9540c9c3 --- /dev/null +++ b/tests/proxy_behavior/management/test_world_seed.py @@ -0,0 +1,30 @@ +import pytest + +from .actors import Actor + +pytestmark = pytest.mark.asyncio(loop_scope="session") + + +@pytest.mark.parametrize("actor", list(Actor), ids=[a.value for a in Actor]) +async def test_each_actor_can_self_info(actor, proxy_client, world): + seeded = world.keys[actor] + resp = await proxy_client.get( + "/key/info", + headers={"Authorization": f"Bearer {seeded.cleartext}"}, + ) + assert resp.status_code == 200, f"{actor.value}: {resp.text}" + body = resp.json() + assert body.get("key") == seeded.hashed + assert body["info"].get("user_id") == seeded.user_id + + +async def test_proxy_admin_actor_can_create_keys_for_others(proxy_client, world): + seeder = world.keys[Actor.PROXY_ADMIN] + target_user_id = world.keys[Actor.OWNER].user_id + + resp = await proxy_client.post( + "/key/generate", + headers={"Authorization": f"Bearer {seeder.cleartext}"}, + json={"key_alias": "smoke-proxy-admin-bypass", "user_id": target_user_id}, + ) + assert resp.status_code == 200, resp.text diff --git a/tests/proxy_unit_tests/test_reducto_ocr_route.py b/tests/proxy_unit_tests/test_reducto_ocr_route.py new file mode 100644 index 00000000000..dc658a74ee8 --- /dev/null +++ b/tests/proxy_unit_tests/test_reducto_ocr_route.py @@ -0,0 +1,137 @@ +import asyncio +import os +from unittest.mock import AsyncMock, patch + +import litellm +import pytest +from fastapi.testclient import TestClient + +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo +from litellm.proxy.proxy_server import app, initialize + + +@pytest.fixture(scope="function") +def fake_env_vars(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key") + monkeypatch.setenv("OPENAI_API_BASE", "http://fake-openai-api-base") + monkeypatch.setenv("AZURE_AI_API_BASE", "http://fake-azure-api-base") + monkeypatch.setenv("AZURE_AI_API_KEY", "fake_azure_api_key") + monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake_azure_openai_api_key") + monkeypatch.setenv("AZURE_SWEDEN_API_BASE", "http://fake-azure-sweden-api-base") + monkeypatch.setenv("AZURE_SWEDEN_API_KEY", "fake_azure_sweden_api_key") + monkeypatch.setenv("REDIS_HOST", "localhost") + + +@pytest.fixture(scope="function") +def client_no_auth(fake_env_vars): + from litellm.proxy.proxy_server import cleanup_router_config_variables + + original_disable_aiohttp = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = True + litellm.in_memory_llm_clients_cache.flush_cache() + cleanup_router_config_variables() + + filepath = os.path.dirname(os.path.abspath(__file__)) + config_fp = os.path.join(filepath, "test_configs", "test_config_no_auth.yaml") + asyncio.run(initialize(config=config_fp, debug=True)) + + # Passthrough of api_base in the JSON body is rejected by default + # (pre_db_read_auth_checks / is_request_body_safe). This test asserts + # api_base reaches aocr(). + from litellm.proxy import proxy_server as _ps + + if _ps.general_settings is None: + _ps.general_settings = {} + _ps.general_settings["allow_client_side_credentials"] = True + + try: + yield TestClient(app) + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +def test_proxy_reducto_ocr_json_rejects_reducto_id(client_no_auth): + with patch( + "litellm.proxy.proxy_server.llm_router.aocr", + new=AsyncMock(), + ) as mock_aocr: + response = client_no_auth.post( + "/v1/ocr", + json={ + "model": "reducto/parse-v3", + "document": { + "type": "document_url", + "document_url": "reducto://proxy.pdf", + }, + "api_key": "proxy-key", + "api_base": "https://platform.reducto.ai", + }, + ) + + assert response.status_code >= 400 + assert "reducto://" in response.text + assert mock_aocr.await_count == 0 + + +def test_proxy_reducto_ocr_json_rejects_reducto_id_in_image_url(client_no_auth): + with patch( + "litellm.proxy.proxy_server.llm_router.aocr", + new=AsyncMock(), + ) as mock_aocr: + response = client_no_auth.post( + "/v1/ocr", + json={ + "model": "reducto/parse-v3", + "document": { + "type": "image_url", + "image_url": "reducto://proxy.png", + }, + }, + ) + + assert response.status_code >= 400 + assert "reducto://" in response.text + assert mock_aocr.await_count == 0 + + +def test_proxy_reducto_ocr_json_passthrough_data_uri(client_no_auth): + mocked_response = OCRResponse( + pages=[OCRPage(index=0, markdown="Proxy OCR")], + model="parse-v3", + usage_info=OCRUsageInfo(pages_processed=1, credits=1), + ) + + data_uri = "data:application/pdf;base64,JVBERi0xLjQK" + + with patch( + "litellm.proxy.proxy_server.llm_router.aocr", + new=AsyncMock(return_value=mocked_response), + ) as mock_aocr: + response = client_no_auth.post( + "/v1/ocr", + json={ + "model": "reducto/parse-v3", + "document": { + "type": "document_url", + "document_url": data_uri, + }, + "api_key": "proxy-key", + "api_base": "https://platform.reducto.ai", + }, + ) + + assert response.status_code == 200 + assert mock_aocr.await_count == 1 + assert mock_aocr.await_args.kwargs["model"] == "reducto/parse-v3" + assert mock_aocr.await_args.kwargs["document"] == { + "type": "document_url", + "document_url": data_uri, + } + assert mock_aocr.await_args.kwargs["api_key"] == "proxy-key" + assert mock_aocr.await_args.kwargs["api_base"] == "https://platform.reducto.ai" + + response_body = response.json() + assert response_body["object"] == "ocr" + assert response_body["usage_info"]["credits"] == 1 + assert response_body["pages"][0]["markdown"] == "Proxy OCR" diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 210347aaf94..958b028c542 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -915,6 +915,36 @@ async def test_user_api_key_auth_websocket(): ) +@pytest.mark.asyncio +async def test_user_api_key_auth_websocket_carries_asgi_path(): + """ + The synthetic Request must carry the ASGI scope's ``path`` so + ``get_request_route`` returns the real WebSocket path, not a value + reconstructed from the (Host-poisonable) ``websocket.url``. + """ + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket + + mock_websocket = MagicMock(spec=WebSocket) + mock_websocket.query_params = {"model": "some_model"} + mock_websocket.headers = {"authorization": "Bearer some_api_key"} + mock_websocket.scope = { + "type": "websocket", + "path": "/v1/realtime", + "root_path": "", + "headers": [(b"authorization", b"Bearer some_api_key")], + } + mock_websocket.url = URL(url="/v1/realtime") + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True + ) as mock_user_api_key_auth: + await user_api_key_auth_websocket(mock_websocket) + + request_arg = mock_user_api_key_auth.call_args.kwargs["request"] + assert request_arg.scope.get("path") == "/v1/realtime" + assert request_arg.scope.get("root_path") == "" + + @pytest.mark.parametrize("enforce_rbac", [True, False]) @pytest.mark.asyncio async def test_jwt_user_api_key_auth_builder_enforce_rbac(enforce_rbac, monkeypatch): diff --git a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py new file mode 100644 index 00000000000..25bf79cd575 --- /dev/null +++ b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py @@ -0,0 +1,268 @@ +""" +Unit tests for the Responses-API streaming-fallback helpers added to Router +in PR #28215 (fix(router): wrap aresponses streaming iterator for mid-stream +fallbacks). + +Targets the four helpers introduced on Router: + - _extract_partial_responses_usage + - _combine_responses_fallback_usage + - _build_responses_continuation_input + - _aresponses_streaming_iterator +""" + +import os +import sys +from typing import Any, AsyncIterator, List +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm import Router +from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) + + +def _make_router() -> Router: + return Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + }, + }, + { + "model_name": "fallback", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-test", + }, + }, + ] + ) + + +def _make_completed_event( + input_tokens: int, output_tokens: int, total_tokens: int +) -> ResponseCompletedEvent: + response = ResponsesAPIResponse.model_construct( + usage=ResponseAPIUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + ) + ) + return ResponseCompletedEvent.model_construct( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=response, + ) + + +# -------- _extract_partial_responses_usage -------- + + +def test_extract_partial_responses_usage_native_completed(): + """Native path: completed_response carries usage → returned as-is.""" + completed = _make_completed_event(11, 7, 18) + source = MagicMock() + source.completed_response = completed + + usage = Router._extract_partial_responses_usage(source) + assert usage is not None + assert usage.input_tokens == 11 + assert usage.output_tokens == 7 + assert usage.total_tokens == 18 + + +def test_extract_partial_responses_usage_no_completed_response(): + """Native path: no completed_response → returns None.""" + source = MagicMock() + source.completed_response = None + + usage = Router._extract_partial_responses_usage(source) + assert usage is None + + +# -------- _combine_responses_fallback_usage -------- + + +def test_combine_responses_fallback_usage_sums_completed_event(): + """Partial-stream usage is summed into the fallback event's usage.""" + fallback_event = _make_completed_event(5, 3, 8) + partial = ResponseAPIUsage(input_tokens=11, output_tokens=7, total_tokens=18) + + Router._combine_responses_fallback_usage(fallback_event, partial) + + combined = fallback_event.response.usage + assert combined is not None + assert combined.input_tokens == 16 + assert combined.output_tokens == 10 + assert combined.total_tokens == 26 + + +def test_combine_responses_fallback_usage_passthrough_for_unknown_event(): + """Events that are not completed/failed/incomplete are not mutated.""" + other = MagicMock() # not a ResponseCompletedEvent etc. → isinstance false + partial = ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2) + Router._combine_responses_fallback_usage(other, partial) + # No mutation expected on the unknown event — call is a no-op. + + +# -------- _build_responses_continuation_input -------- + + +def test_build_responses_continuation_input_from_string(): + out = Router._build_responses_continuation_input( + "Hello world", "partial assistant text" + ) + assert len(out) == 3 + assert out[0]["role"] == "user" + assert out[0]["content"][0]["text"] == "Hello world" + assert out[1]["role"] == "developer" + assert out[2]["role"] == "assistant" + assert out[2]["content"][0]["text"] == "partial assistant text" + + +def test_build_responses_continuation_input_from_list_preserves_items(): + existing: List[Any] = [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "msg1"}], + } + ] + out = Router._build_responses_continuation_input(existing, "partial") + assert len(out) == 3 + assert out[0]["content"][0]["text"] == "msg1" + assert out[1]["role"] == "developer" + assert out[2]["role"] == "assistant" + + +def test_build_responses_continuation_input_from_none(): + out = Router._build_responses_continuation_input(None, "partial") + assert len(out) == 2 + assert out[0]["role"] == "developer" + assert out[1]["role"] == "assistant" + + +# -------- _aresponses_streaming_iterator (passthrough smoke test) -------- + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_passthrough(): + """ + Without MidStreamFallbackError, the wrapper yields source events + unchanged and returns a BaseResponsesAPIStreamingIterator subclass. + """ + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + events = [_make_completed_event(1, 1, 2)] + + class _FakeSource: + """Minimal source iterator. Provides every attribute the wrapper + constructor reads from source_iterator.""" + + def __init__(self) -> None: + self._i = 0 + self.completed_response = None + self.response = MagicMock() + self.model = "openai/gpt-4o-mini" + self.logging_obj = MagicMock() + self.responses_api_provider_config = MagicMock() + self.start_time = 0.0 + self.litellm_metadata = {} + self.custom_llm_provider = "openai" + self.request_data = {} + self.call_type = "aresponses" + self._hidden_params: dict = {} + + def __aiter__(self) -> AsyncIterator[Any]: + return self + + async def __anext__(self): + if self._i >= len(events): + raise StopAsyncIteration + ev = events[self._i] + self._i += 1 + return ev + + async def aclose(self): + return None + + router = _make_router() + source = _FakeSource() + + wrapper = await router._aresponses_streaming_iterator( + source, initial_kwargs={"model": "primary"} + ) + assert isinstance(wrapper, BaseResponsesAPIStreamingIterator) + + collected = [ev async for ev in wrapper] + assert len(collected) == 1 + assert collected[0].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + + +# -------- _aresponses_with_streaming_fallbacks -------- + + +@pytest.mark.asyncio +async def test_aresponses_with_streaming_fallbacks_non_streaming_passthrough(): + """Non-streaming response is returned unchanged, no wrap.""" + router = _make_router() + plain_response = MagicMock() + + async def fake_original(**_kwargs): + return plain_response + + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(return_value=plain_response), + ): + out = await router._aresponses_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=False, + ) + assert out is plain_response + + +@pytest.mark.asyncio +async def test_aresponses_with_streaming_fallbacks_wraps_streaming_iterator(): + """Streaming response is wrapped via _aresponses_streaming_iterator.""" + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + router = _make_router() + streaming_iter = MagicMock(spec=BaseResponsesAPIStreamingIterator) + wrapped = MagicMock(spec=BaseResponsesAPIStreamingIterator) + + async def fake_original(**_kwargs): + return streaming_iter + + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(return_value=streaming_iter), + ), patch.object( + router, + "_aresponses_streaming_iterator", + new=AsyncMock(return_value=wrapped), + ) as mock_wrap: + out = await router._aresponses_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=True, + ) + assert out is wrapped + mock_wrap.assert_awaited_once() diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 697a9ebc720..d335c359aa0 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -508,6 +508,308 @@ and I learn to carry this small calm home.""" print("✓ transform_response correctly handled reasoning items and output messages") +def _make_empty_responses_api_response(model: str = "gpt-5.4"): + from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse + + return ResponsesAPIResponse( + id="resp_from_stream", + created_at=1760144904, + error=None, + incomplete_details=None, + instructions=None, + metadata={}, + model=model, + object="response", + output=[], + parallel_tool_calls=True, + temperature=1.0, + tool_choice="auto", + tools=[], + top_p=1.0, + max_output_tokens=None, + previous_response_id=None, + reasoning={"effort": "low", "summary": "detailed"}, + status="completed", + text={"format": {"type": "text"}, "verbosity": "medium"}, + truncation="disabled", + usage=ResponseAPIUsage( + input_tokens=1, + input_tokens_details=None, + output_tokens=1, + output_tokens_details=None, + total_tokens=2, + cost=None, + ), + user=None, + store=True, + background=False, + billing={"payer": "developer"}, + max_tool_calls=None, + prompt_cache_key=None, + safety_identifier=None, + service_tier="default", + top_logprobs=0, + ) + + +def _make_empty_model_response(): + from litellm.types.utils import ModelResponse, Usage + + return ModelResponse( + id="chatcmpl-test-recovered", + created=1760144904, + model=None, + object="chat.completion", + system_fingerprint=None, + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + + +def test_transform_response_recovers_empty_output_from_raw_sse(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_text.done","output_index":0,"content_index":0,"item_id":"msg_from_stream","text":"Recovered from SSE"}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Recovered from SSE" + + +def test_transform_response_recovers_output_item_done_from_raw_sse(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_from_item","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Recovered from output item","annotations":[]}]}}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Recovered from output item" + + +def test_transform_response_recovers_output_item_done_from_whitespace_padded_raw_sse(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + output_item_event = { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "message", + "id": "msg_from_item", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Recovered from padded output item", + "annotations": [], + } + ], + }, + } + completed_event = { + "type": "response.completed", + "response": { + "id": "resp_from_stream", + "object": "response", + "created_at": 1760144904, + "status": "completed", + "model": "gpt-5.4", + "output": [], + }, + } + raw_sse = "\n".join( + [ + f" data: {json.dumps(output_item_event)} ", + f"\tdata: {json.dumps(completed_event)}", + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Recovered from padded output item" + + +def test_transform_response_preserves_output_item_when_text_done_arrives_later(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_from_item","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Complete output item text","annotations":[]}]}}', + 'data: {"type":"response.output_text.done","output_index":0,"content_index":0,"item_id":"msg_from_stream","text":"Late text event"}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Complete output item text" + + +def test_recover_output_items_merges_text_only_items_at_distinct_indices(): + """When OUTPUT_ITEM_DONE covers some indices and OUTPUT_TEXT_DONE covers + others, both must be preserved instead of treating them as mutually + exclusive fallbacks.""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_item_0","role":"assistant","status":"completed","content":[{"type":"output_text","text":"From OUTPUT_ITEM_DONE","annotations":[]}]}}', + 'data: {"type":"response.output_text.done","output_index":1,"content_index":0,"item_id":"msg_text_1","text":"From OUTPUT_TEXT_DONE only"}', + "data: [DONE]", + "", + ] + ) + + recovered = ( + LiteLLMResponsesTransformationHandler._recover_output_items_from_raw_sse( + raw_sse + ) + ) + + assert len(recovered) == 2 + assert recovered[0]["id"] == "msg_item_0" + assert recovered[0]["content"][0]["text"] == "From OUTPUT_ITEM_DONE" + assert recovered[1]["id"] == "msg_text_1" + assert recovered[1]["content"][0]["text"] == "From OUTPUT_TEXT_DONE only" + + +def test_transform_response_prefers_completed_output_from_raw_sse(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + handler = LiteLLMResponsesTransformationHandler() + + raw_sse = "\n".join( + [ + 'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_from_item","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Earlier stream text","annotations":[]}]}}', + 'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[{"type":"message","id":"msg_from_completed","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Authoritative completed text","annotations":[]}]}]}}', + "data: [DONE]", + "", + ] + ) + + raw_response = _make_empty_responses_api_response() + model_response = _make_empty_model_response() + logging_obj = Mock() + logging_obj.model_call_details = {"original_response": raw_sse} + + result = handler.transform_response( + model="gpt-5.4", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.4"}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Authoritative completed text" + + def test_convert_tools_to_responses_format(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, diff --git a/tests/test_litellm/containers/test_azure_container_transformation.py b/tests/test_litellm/containers/test_azure_container_transformation.py index 70181f6f03d..cdcccf7c04e 100644 --- a/tests/test_litellm/containers/test_azure_container_transformation.py +++ b/tests/test_litellm/containers/test_azure_container_transformation.py @@ -109,6 +109,31 @@ class TestAzureContainerConfig: assert "/openai/v1/containers" in url + def test_get_complete_url_strips_responses_path_and_preserves_api_version(self): + """When api_base is the responses endpoint URL, get_complete_url must: + - strip /openai/responses (no double-path) + - use the api-version from api_base query string, NOT the deployment's + older api_version (e.g. 2024-08-01-preview → containers need 2025-04-01-preview) + """ + api_base = "https://my-resource.cognitiveservices.azure.com/openai/responses?api-version=2025-04-01-preview" + + url = self.config.get_complete_url( + api_base=api_base, + litellm_params={"api_version": "2024-08-01-preview"}, + ) + + assert ( + "/openai/responses/openai/containers" not in url + ), "path must not double /openai/responses" + assert "my-resource.cognitiveservices.azure.com" in url + assert "/openai/containers" in url or "/openai/v1/containers" in url + assert ( + "2025-04-01-preview" in url + ), "must use version from api_base, not litellm_params" + assert ( + "2024-08-01-preview" not in url + ), "must not fall back to older chat api_version" + def test_get_complete_url_raises_without_api_base(self, monkeypatch): monkeypatch.delenv("AZURE_API_BASE", raising=False) monkeypatch.setattr(litellm, "api_base", None) @@ -531,6 +556,92 @@ class TestAzureContainerKnownFailureRegressions: assert qs.get("api-version") == ["v1"] assert qs.get("foo") == ["bar"] + @pytest.mark.asyncio + async def test_regression_no_container_id_does_not_use_user_supplied_model_id( + self, monkeypatch + ): + """Operations without container_id (create, list) must NOT route via + _ageneric_api_call_with_fallbacks using a caller-supplied model_id. + + Security boundary: only the path that holds a validated container_id + is trusted to fall back to the forwarded model_id. A caller setting + model_id without container_id on POST /v1/containers must not gain + access to an arbitrary deployment UUID. + """ + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "azure-model", + "litellm_params": { + "model": "azure/gpt-4", + "api_base": "https://my-resource.cognitiveservices.azure.com", + "api_key": "test-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": "deployment-uuid-123"}, + } + ] + ) + + fallback_called = {"called": False} + + async def _mock_fallback(original_function, **kwargs): + fallback_called["called"] = True + return {} + + monkeypatch.setattr(router, "_ageneric_api_call_with_fallbacks", _mock_fallback) + + original_called = {"called": False} + + async def _noop(**kwargs): + original_called["called"] = True + return {} + + # No container_id — simulates create/list; caller injects a model_id + await router._init_containers_api_endpoints( + original_function=_noop, + model_id="deployment-uuid-123", + custom_llm_provider="azure", + ) + + assert not fallback_called["called"], ( + "_ageneric_api_call_with_fallbacks must NOT be called when " + "container_id is absent, even if model_id is supplied" + ) + assert original_called["called"], "original_function must be called directly" + + def test_regression_httpx_empty_params_strips_query_string(self): + """httpx erases the URL query-string when params={} (empty dict) is passed. + + Root cause of the Azure container 404s on POST/DELETE: + _build_query_params returns {} when the endpoint has no extra params; + passing that {} as params= to httpx wiped ?api-version=2025-04-01-preview. + + Fix: every container httpx call now uses `params or None` so an empty + dict falls back to None, which tells httpx to leave the URL untouched. + """ + url = ( + "https://resource.cognitiveservices.azure.com" + "/openai/containers/cntr_123?api-version=2025-04-01-preview" + ) + client = httpx.AsyncClient() + + req_none = client.build_request("DELETE", url, params=None) + assert "api-version=2025-04-01-preview" in str(req_none.url) + + req_empty = client.build_request("DELETE", url, params={}) + assert "api-version" not in str( + req_empty.url + ), "Documents root cause: params={} strips the query string" + + effective: dict = {} + req_guarded = client.build_request("DELETE", url, params=effective or None) + assert "api-version=2025-04-01-preview" in str( + req_guarded.url + ), "`params or None` must preserve ?api-version" + def test_regression_proxy_resolves_azure_text_same_as_azure(self): """Router/proxy treat azure_text like azure for container config.""" from litellm.proxy.container_endpoints.handler_factory import ( @@ -770,3 +881,143 @@ class TestAzureContainerKnownFailureRegressions: assert captured["data"]["container_id"] == "cntr_123" assert captured["data"]["custom_llm_provider"] == "azure" assert captured["data"]["model_id"] == "model_abc123" + + @pytest.mark.asyncio + async def test_regression_get_container_forwarding_params_sets_model_id_for_managed_id( + self, + ): + """get_container_forwarding_params must extract model_id from a + LiteLLM-managed encoded container ID and include it in the forwarding + dict. This is the proxy-side half of the native-Azure-ID routing fix: + the router's _init_containers_api_endpoints reads kwargs["model_id"] + which is set here. + """ + from litellm.proxy.container_endpoints.ownership import ( + get_container_forwarding_params, + ) + + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="deployment-uuid-123", + container_id="cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df", + ) + + params = await get_container_forwarding_params( + container_id=encoded_id, + original_container_id="cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df", + custom_llm_provider="azure", + ) + + assert ( + params.get("model_id") == "deployment-uuid-123" + ), "model_id must be forwarded to the router for managed container IDs" + assert params.get("container_id") == ( + "cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df" + ) + assert params.get("custom_llm_provider") == "azure" + + @pytest.mark.asyncio + async def test_regression_get_container_forwarding_params_recovers_model_id_for_native_id( + self, monkeypatch + ): + """Native Azure IDs (``cntr_``) cannot be decoded, so model_id + must be recovered from the ownership row's ``unified_object_id`` — + the encoded form captured at create time when the router selected a + specific deployment. Without this, the router-side fallback for + native IDs in ``_init_containers_api_endpoints`` is dead code. + """ + from types import SimpleNamespace + from unittest.mock import AsyncMock + + from litellm.proxy.container_endpoints import ownership + from litellm.proxy.container_endpoints.ownership import ( + get_container_forwarding_params, + ) + + native_id = "cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df" + encoded_stored_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="deployment-uuid-123", + container_id=native_id, + ) + + ownership._CONTAINER_STORED_ID_CACHE.flush_cache() + ownership._CONTAINER_OWNER_CACHE.flush_cache() + + table = AsyncMock() + table.find_first.return_value = SimpleNamespace( + created_by="user-1", + file_purpose=ownership.CONTAINER_OBJECT_PURPOSE, + unified_object_id=encoded_stored_id, + ) + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + params = await get_container_forwarding_params( + container_id=native_id, + original_container_id=native_id, + custom_llm_provider="azure", + ) + + assert params.get("model_id") == "deployment-uuid-123", ( + "model_id must be recovered from the stored unified_object_id " + "for native upstream container IDs" + ) + assert params.get("container_id") == native_id + assert params.get("custom_llm_provider") == "azure" + + @pytest.mark.asyncio + async def test_regression_native_azure_container_id_uses_forwarded_model_id( + self, monkeypatch + ): + """Native Azure container IDs (cntr_ + hex, no LiteLLM payload) must + still route through _ageneric_api_call_with_fallbacks using the + model_id forwarded from the proxy ownership check so that deployment + credentials (api_base) are applied.""" + from litellm.router import Router + + router = Router( + model_list=[ + { + "model_name": "azure-model", + "litellm_params": { + "model": "azure/gpt-4", + "api_base": "https://my-resource.cognitiveservices.azure.com", + "api_key": "test-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": "deployment-uuid-123"}, + } + ] + ) + + called_with: dict = {} + + async def _mock_fallback(original_function, **kwargs): + called_with.update(kwargs) + return {} + + monkeypatch.setattr(router, "_ageneric_api_call_with_fallbacks", _mock_fallback) + + native_azure_id = "cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df" + + async def _noop(**kwargs): + return {} + + await router._init_containers_api_endpoints( + original_function=_noop, + container_id=native_azure_id, + model_id="deployment-uuid-123", + custom_llm_provider="azure", + ) + + assert called_with.get("model") == "deployment-uuid-123", ( + "_ageneric_api_call_with_fallbacks must be called with the forwarded " + "model_id when the container_id carries no LiteLLM routing payload" + ) diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py new file mode 100644 index 00000000000..d8669960674 --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -0,0 +1,260 @@ +"""Regression: update_batch_in_database must not persist raw provider output_file_id.""" + +import json +from types import SimpleNamespace +from typing import Optional +import pytest +from unittest.mock import AsyncMock, MagicMock + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints.common_utils import ( + ensure_batch_response_managed_file_ids, + update_batch_in_database, +) +from litellm.types.utils import LiteLLMBatch + + +def _build_batch_response( + *, + batch_id: str = "batch_managed_ids_test", + status: str = "completed", + output_file_id: Optional[str] = "file-rawoutput789", + error_file_id: Optional[str] = None, + hidden_params: Optional[dict] = None, +) -> LiteLLMBatch: + batch = LiteLLMBatch( + id=batch_id, + object="batch", + status=status, + endpoint="/v1/chat/completions", + input_file_id="file-input123", + output_file_id=output_file_id, + error_file_id=error_file_id, + completion_window="24h", + created_at=1234567890, + ) + if hidden_params is not None: + batch._hidden_params = hidden_params # type: ignore[attr-defined] + return batch + + +def _build_managed_files_mock(unified_id: str = "file-bWFuYWdlZF9vdXRwdXRfaWQ="): + mock = MagicMock() + mock.get_unified_output_file_id = MagicMock(return_value=unified_id) + mock.store_unified_file_id = AsyncMock() + return mock + + +def _build_prisma_mock(): + mock = MagicMock() + mock.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) + mock.db.litellm_managedobjecttable.update = AsyncMock() + return mock + + +@pytest.mark.asyncio +async def test_update_batch_in_database_stores_unified_output_file_id(): + raw_output_file_id = "file-rawoutput789" + unified_output_file_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" + batch_id = "batch_managed_ids_test" + unified_batch_id = ( + "litellm_proxy;model_id:my-model;llm_batch_id:batch_managed_ids_test" + ) + + response = _build_batch_response( + batch_id=batch_id, + output_file_id=raw_output_file_id, + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + + mock_managed_files = _build_managed_files_mock(unified_id=unified_output_file_id) + mock_prisma = _build_prisma_mock() + + await update_batch_in_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + response=response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + stored = json.loads( + mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs["data"][ + "file_object" + ] + ) + assert stored["output_file_id"] == unified_output_file_id + assert stored["output_file_id"] != raw_output_file_id + + +@pytest.mark.asyncio +async def test_ensure_batch_response_normalizes_error_file_id(): + """Both output_file_id and error_file_id must be normalized to managed IDs.""" + unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" + response = _build_batch_response( + output_file_id="file-raw-output", + error_file_id="file-raw-error", + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + + mock_managed_files = _build_managed_files_mock(unified_id=unified_id) + mock_prisma = _build_prisma_mock() + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert response.output_file_id == unified_id + assert response.error_file_id == unified_id + assert mock_managed_files.get_unified_output_file_id.call_count == 2 + + +@pytest.mark.asyncio +async def test_ensure_batch_response_swallows_conversion_errors(): + """When the managed-files conversion raises, the failure is logged, not propagated.""" + raw_output_file_id = "file-raw-output" + response = _build_batch_response( + output_file_id=raw_output_file_id, + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + + mock_managed_files = MagicMock() + mock_managed_files.get_unified_output_file_id = MagicMock( + side_effect=RuntimeError("boom") + ) + mock_managed_files.store_unified_file_id = AsyncMock() + + mock_logger = MagicMock() + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=mock_managed_files, + prisma_client=_build_prisma_mock(), + verbose_proxy_logger=mock_logger, + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert response.output_file_id == raw_output_file_id + mock_logger.warning.assert_called() + + +@pytest.mark.asyncio +async def test_ensure_batch_response_builds_auth_from_db_batch_object(): + """If user_api_key_dict is omitted, fall back to created_by/team_id on db_batch_object.""" + unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" + response = _build_batch_response( + output_file_id="file-raw-output", + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + + mock_managed_files = _build_managed_files_mock(unified_id=unified_id) + db_batch_object = SimpleNamespace( + created_by="user-from-db", team_id="team-from-db", status="completed" + ) + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=mock_managed_files, + prisma_client=_build_prisma_mock(), + verbose_proxy_logger=MagicMock(), + db_batch_object=db_batch_object, + ) + + forwarded_auth = mock_managed_files.store_unified_file_id.call_args.kwargs[ + "user_api_key_dict" + ] + assert forwarded_auth.user_id == "user-from-db" + assert forwarded_auth.team_id == "team-from-db" + + +@pytest.mark.asyncio +async def test_ensure_batch_response_resolves_model_name_from_unified_file_id(): + """When hidden_params lacks model_name, derive it from unified_file_id.""" + unified_id = "file-bWFuYWdlZF9vdXRwdXRfaWQ=" + response = _build_batch_response( + output_file_id="file-raw-output", + hidden_params={ + "model_id": "my-model", + "unified_file_id": "litellm_proxy:application/octet-stream;unified_id,abc;target_model_names,gpt-4o-mini,gemini-2.0-flash", + }, + ) + + mock_managed_files = _build_managed_files_mock(unified_id=unified_id) + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=mock_managed_files, + prisma_client=_build_prisma_mock(), + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert ( + mock_managed_files.get_unified_output_file_id.call_args.kwargs["model_name"] + == "gpt-4o-mini,gemini-2.0-flash" + ) + + +@pytest.mark.asyncio +async def test_ensure_batch_response_returns_early_without_managed_files_obj(): + """Without managed_files_obj, the helper is a no-op (no conversion attempted).""" + response = _build_batch_response( + output_file_id="file-raw-output", + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=None, + prisma_client=_build_prisma_mock(), + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert response.output_file_id == "file-raw-output" + + +@pytest.mark.asyncio +async def test_ensure_batch_response_returns_early_without_model_id(): + """Without model_id in hidden_params, the helper cannot create managed IDs.""" + response = _build_batch_response( + output_file_id="file-raw-output", + hidden_params={"model_name": "openai/gpt-4o"}, + ) + mock_managed_files = _build_managed_files_mock() + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=mock_managed_files, + prisma_client=_build_prisma_mock(), + verbose_proxy_logger=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-abc"), + ) + + assert response.output_file_id == "file-raw-output" + mock_managed_files.get_unified_output_file_id.assert_not_called() + + +@pytest.mark.asyncio +async def test_ensure_batch_response_returns_early_without_auth(): + """Without user_api_key_dict or db_batch_object, no conversion is attempted.""" + response = _build_batch_response( + output_file_id="file-raw-output", + hidden_params={"model_id": "my-model", "model_name": "openai/gpt-4o"}, + ) + mock_managed_files = _build_managed_files_mock() + + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=mock_managed_files, + prisma_client=_build_prisma_mock(), + verbose_proxy_logger=MagicMock(), + ) + + assert response.output_file_id == "file-raw-output" + mock_managed_files.get_unified_output_file_id.assert_not_called() diff --git a/tests/test_litellm/integrations/rubrik_test_helpers.py b/tests/test_litellm/integrations/rubrik_test_helpers.py new file mode 100644 index 00000000000..1bdb8cb247b --- /dev/null +++ b/tests/test_litellm/integrations/rubrik_test_helpers.py @@ -0,0 +1,23 @@ +"""Shared helpers for Rubrik plugin tests.""" + +from typing import Any, Dict + +from litellm.types.utils import GenericGuardrailAPIInputs + + +def make_tool_call_dict( + tc_id: str, name: str, arguments: str = "{}" +) -> Dict[str, Any]: + """Create a tool call dict matching the ChatCompletionMessageToolCall schema.""" + return { + "id": tc_id, + "type": "function", + "function": {"name": name, "arguments": arguments}, + } + + +def make_inputs_with_tools( + tool_calls: list, texts: list | None = None +) -> GenericGuardrailAPIInputs: + """Create GenericGuardrailAPIInputs with tool_calls.""" + return GenericGuardrailAPIInputs(texts=texts or [], tool_calls=tool_calls) diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 6de855262bd..b65e629c890 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -66,7 +66,7 @@ class TestOpenTelemetryGuardrails(unittest.TestCase): mock_span.set_attribute.assert_any_call("guardrail_name", "test_guardrail") mock_span.set_attribute.assert_any_call("guardrail_mode", "input") mock_span.set_attribute.assert_any_call( - "guardrail_response", "filtered_content" + "guardrail_response", safe_dumps("filtered_content") ) mock_span.set_attribute.assert_any_call( "masked_entity_count", safe_dumps({"CREDIT_CARD": 2}) @@ -87,6 +87,65 @@ class TestOpenTelemetryGuardrails(unittest.TestCase): # Verify that start_span was never called otel.tracer.start_span.assert_not_called() + @patch("litellm.integrations.opentelemetry.datetime") + def test_guardrail_response_dict_is_json_serialized(self, mock_datetime): + """Dict guardrail_response (e.g. OpenAI moderation result) must reach + the span as a JSON string so downstream pipelines can parse it for + metric extraction — this is the bug the PR fixes.""" + otel = OpenTelemetry() + otel.tracer = MagicMock() + mock_span = MagicMock() + otel.tracer.start_span.return_value = mock_span + + moderation_payload = { + "id": "modr-7740", + "model": "omni-moderation-latest", + "results": [{"categories": {"harassment": False}}], + } + guardrail_info = { + "guardrail_name": "test_guardrail", + "guardrail_mode": "input", + "guardrail_response": moderation_payload, + "start_time": 1609459200.0, + "end_time": 1609459201.0, + } + kwargs = { + "standard_logging_object": {"guardrail_information": [guardrail_info]} + } + + otel._create_guardrail_span(kwargs=kwargs, context=None) + + mock_span.set_attribute.assert_any_call( + "guardrail_response", safe_dumps(moderation_payload) + ) + + @patch("litellm.integrations.opentelemetry.datetime") + def test_guardrail_response_none_is_skipped(self, mock_datetime): + """When guardrail_response is None, the attribute must not be set — + guards against round-tripping ``"null"`` into traces.""" + otel = OpenTelemetry() + otel.tracer = MagicMock() + mock_span = MagicMock() + otel.tracer.start_span.return_value = mock_span + + guardrail_info = { + "guardrail_name": "test_guardrail", + "guardrail_mode": "input", + "guardrail_response": None, + "start_time": 1609459200.0, + "end_time": 1609459201.0, + } + kwargs = { + "standard_logging_object": {"guardrail_information": [guardrail_info]} + } + + otel._create_guardrail_span(kwargs=kwargs, context=None) + + attribute_keys = [ + call.args[0] for call in mock_span.set_attribute.call_args_list + ] + self.assertNotIn("guardrail_response", attribute_keys) + class TestOpenTelemetryTeamAttributesOnChildSpans(unittest.TestCase): """team_id / team_alias must land on every child span of a @@ -1169,7 +1228,7 @@ class TestOpenTelemetry(unittest.TestCase): mock_span.set_attribute.assert_any_call("guardrail_name", "test_guardrail") mock_span.set_attribute.assert_any_call("guardrail_mode", "input") mock_span.set_attribute.assert_any_call( - "guardrail_response", "filtered_content" + "guardrail_response", safe_dumps("filtered_content") ) mock_span.set_attribute.assert_any_call( "masked_entity_count", safe_dumps({"CREDIT_CARD": 2}) diff --git a/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py b/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py new file mode 100644 index 00000000000..ace9399cf53 --- /dev/null +++ b/tests/test_litellm/integrations/test_otel_guardrail_violation_spans.py @@ -0,0 +1,641 @@ +""" +Tests for guardrail OTEL spans on violation. + +Two distinct gaps surface together when a pre-call guardrail blocks the +request before it reaches the LLM provider: + + 1. ``async_post_call_failure_hook`` (the OTEL hook that actually runs on + the proxy failure path) only stamps attributes on the proxy parent + span. It never creates the child ``guardrail`` span, even though + ``request_data["metadata"]["standard_logging_guardrail_information"]`` + is populated by the time the hook runs. + + 2. ``_create_guardrail_span`` records ``guardrail_name`` / ``guardrail_mode`` + / ``guardrail_response`` but does not surface ``guardrail_status`` + (success / guardrail_intervened / guardrail_failed_to_respond / + not_run) or the violation categories (Bedrock topic policy names, + content filter types, etc.) as queryable span attributes — the data + is buried inside the serialised ``guardrail_response`` blob and cannot + be filtered on in the trace backend. + +The tests below use real OTEL SDK objects (TracerProvider + +InMemorySpanExporter + a real BatchSpanProcessor-equivalent) and the +real ``OpenTelemetry`` integration. No monkey patching of the integration +under test — only the OTEL exporter is in-memory. +""" + +import os +import sys +import time +import unittest +from datetime import datetime, timedelta, timezone + +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.integrations.opentelemetry import ( + LITELLM_REQUEST_SPAN_NAME, + OpenTelemetry, +) +from litellm.proxy._types import UserAPIKeyAuth + + +GUARDRAIL_SPAN_NAME = "guardrail" +PROXY_SPAN_NAME = "Received Proxy Server Request" + + +def _bedrock_block_response(): + """Realistic Bedrock ApplyGuardrail response when a topic policy fires. + + Mirrors the shape in ``litellm/types/proxy/guardrails/guardrail_hooks/ + bedrock_guardrails.py`` so the violation-category extraction can be + tested against the exact payload Bedrock returns. + """ + return { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": [ + { + "name": "Fiduciary Advice", + "type": "DENY", + "action": "BLOCKED", + } + ] + }, + "contentPolicy": { + "filters": [ + { + "type": "VIOLENCE", + "confidence": "HIGH", + "action": "BLOCKED", + } + ] + }, + "wordPolicy": { + "customWords": [{"match": "secret-codeword", "action": "BLOCKED"}], + "managedWordLists": [ + {"match": "fuck", "type": "PROFANITY", "action": "BLOCKED"} + ], + }, + } + ], + "outputs": [{"text": "Sorry, the model cannot respond to this request."}], + } + + +def _slg_entry( + guardrail_status, + guardrail_response, + *, + name="bedrock-test", + mode="pre_call", + provider="bedrock", + start=1.0, + end=2.0, + violation_categories=None, + guardrail_action=None, +): + """Build a StandardLoggingGuardrailInformation entry the way + ``add_standard_logging_guardrail_information_to_request_data`` does.""" + entry = { + "guardrail_name": name, + "guardrail_provider": provider, + "guardrail_mode": mode, + "guardrail_response": guardrail_response, + "guardrail_status": guardrail_status, + "start_time": start, + "end_time": end, + "duration": end - start, + } + if violation_categories is not None: + entry["violation_categories"] = violation_categories + if guardrail_action is not None: + entry["guardrail_action"] = guardrail_action + return entry + + +def _kwargs_with_guardrail( + *, + entries, + parent_span=None, + include_exception=False, +): + """Build the kwargs / model_call_details shape that the OTEL integration + consumes. ``litellm_params.metadata`` is the SAME dict that the proxy's + ``request_data["metadata"]`` becomes after ``update_environment_variables``, + so ``_otel_internal`` dedupe state lives there too.""" + metadata = {"standard_logging_guardrail_information": list(entries)} + if parent_span is not None: + metadata["litellm_parent_otel_span"] = parent_span + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": {}, + "litellm_params": { + "custom_llm_provider": "openai", + "metadata": metadata, + }, + "standard_logging_object": { + "id": "test-call-id", + "call_type": "completion", + "metadata": metadata, + "hidden_params": {}, + "guardrail_information": list(entries), + }, + } + if include_exception: + kwargs["exception"] = Exception("guardrail blocked the request") + return kwargs + + +def _make_otel(): + """Spin up a real OTEL pipeline backed by an in-memory exporter.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + otel = OpenTelemetry(tracer_provider=provider) + otel.tracer = provider.get_tracer(__name__) + return otel, provider, exporter + + +def _run(coro): + """Run a coroutine on a fresh event loop and close it — prevents the + "unclosed event loop" / ResourceWarning that you get from + asyncio.new_event_loop().run_until_complete() with no cleanup.""" + import asyncio + + loop = asyncio.new_event_loop() + try: + return loop.run_until_complete(coro) + finally: + loop.close() + + +def _attr(span, key): + return (span.attributes or {}).get(key) + + +class TestGuardrailSpanOnViolation(unittest.TestCase): + """Bug 1: when a pre-call guardrail blocks, the guardrail span and the + litellm_request span must both appear with the correct status.""" + + def test_handle_failure_creates_litellm_request_and_guardrail_spans(self): + """Driving ``_handle_failure`` with a populated + ``standard_logging_object['guardrail_information']`` entry must + emit both spans, parented correctly, with ERROR on the parent.""" + otel, _, exporter = _make_otel() + + kwargs = _kwargs_with_guardrail( + entries=[ + _slg_entry("guardrail_intervened", _bedrock_block_response()), + ], + include_exception=True, + ) + + start = datetime.now(timezone.utc) + end = start + timedelta(milliseconds=20) + otel._handle_failure(kwargs, response_obj=None, start_time=start, end_time=end) + + spans = exporter.get_finished_spans() + litellm_spans = [s for s in spans if s.name == LITELLM_REQUEST_SPAN_NAME] + guardrail_spans = [s for s in spans if s.name == GUARDRAIL_SPAN_NAME] + + self.assertEqual( + len(litellm_spans), + 1, + "Expected exactly one litellm_request span on guardrail block", + ) + self.assertEqual(litellm_spans[0].status.status_code, StatusCode.ERROR) + + self.assertEqual( + len(guardrail_spans), + 1, + "Expected exactly one guardrail span on guardrail block", + ) + + # Guardrail span must be a child of the litellm_request span + self.assertIsNotNone( + guardrail_spans[0].parent, + "Guardrail span must be parented (not a root span)", + ) + self.assertEqual( + guardrail_spans[0].parent.span_id, + litellm_spans[0].context.span_id, + ) + + def test_async_post_call_failure_hook_emits_guardrail_span(self): + """The production failure path on the proxy calls + ``async_post_call_failure_hook`` with the (still-populated) + ``request_data``. The hook currently only stamps attrs on the proxy + span; it must also emit the guardrail span so the violation is + visible in the trace.""" + otel, provider, exporter = _make_otel() + parent_span = provider.get_tracer(__name__).start_span(PROXY_SPAN_NAME) + + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", + parent_otel_span=parent_span, + request_route="/chat/completions", + ) + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": { + "standard_logging_guardrail_information": [ + _slg_entry("guardrail_intervened", _bedrock_block_response()) + ], + }, + } + + _run( + otel.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("guardrail blocked"), + user_api_key_dict=user_api_key_dict, + ) + ) + + spans = exporter.get_finished_spans() + guardrail_spans = [s for s in spans if s.name == GUARDRAIL_SPAN_NAME] + self.assertEqual( + len(guardrail_spans), + 1, + "async_post_call_failure_hook must emit the guardrail span when " + "request_data['metadata'] carries standard_logging_guardrail_information", + ) + + # The guardrail span must be parented to the proxy request span so + # backends correlate it with the rest of the trace. + self.assertIsNotNone(guardrail_spans[0].parent) + self.assertEqual( + guardrail_spans[0].parent.span_id, + parent_span.context.span_id, + ) + + def test_handle_failure_and_post_call_failure_hook_dedupe(self): + """When _handle_failure and async_post_call_failure_hook BOTH fire + for the same request (the production flow on a guardrail block), + exactly one guardrail span must be emitted. The dedupe relies on + request_data['metadata'] and kwargs['litellm_params']['metadata'] + referencing the SAME dict so _emit_once sees its earlier marker.""" + otel, provider, exporter = _make_otel() + parent_span = provider.get_tracer(__name__).start_span(PROXY_SPAN_NAME) + + # Shared metadata dict — same identity, mirroring how + # update_environment_variables wires them in the proxy. + shared_metadata = { + "standard_logging_guardrail_information": [ + _slg_entry( + "guardrail_intervened", + _bedrock_block_response(), + violation_categories=["Fiduciary Advice"], + ) + ], + } + + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "optional_params": {}, + "litellm_params": { + "custom_llm_provider": "openai", + "metadata": shared_metadata, + }, + "standard_logging_object": { + "id": "test-call-id", + "call_type": "completion", + "metadata": shared_metadata, + "hidden_params": {}, + "guardrail_information": shared_metadata[ + "standard_logging_guardrail_information" + ], + }, + "exception": Exception("guardrail blocked"), + } + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": shared_metadata, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", + parent_otel_span=parent_span, + request_route="/chat/completions", + ) + + start = datetime.now(timezone.utc) + end = start + timedelta(milliseconds=20) + otel._handle_failure(kwargs, response_obj=None, start_time=start, end_time=end) + _run( + otel.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("guardrail blocked"), + user_api_key_dict=user_api_key_dict, + ) + ) + + guardrail_spans = [ + s for s in exporter.get_finished_spans() if s.name == GUARDRAIL_SPAN_NAME + ] + self.assertEqual( + len(guardrail_spans), + 1, + "Dedupe must collapse the two emit calls into one span when the " + "metadata dict identity is shared between kwargs and request_data", + ) + + +class TestGuardrailSpanAttributesOnViolation(unittest.TestCase): + """Bug 2: the guardrail span must surface the violation status and + violation categories as queryable span attributes, not bury them inside + ``guardrail_response`` (which is logged as a single serialised blob).""" + + def _emit_and_get_guardrail_span(self, entry): + otel, _, exporter = _make_otel() + kwargs = _kwargs_with_guardrail(entries=[entry]) + otel._create_guardrail_span(kwargs=kwargs, context=None) + + guardrail_spans = [ + s for s in exporter.get_finished_spans() if s.name == GUARDRAIL_SPAN_NAME + ] + self.assertEqual(len(guardrail_spans), 1) + return guardrail_spans[0] + + def test_status_attribute_present_for_intervened(self): + entry = _slg_entry("guardrail_intervened", _bedrock_block_response()) + span = self._emit_and_get_guardrail_span(entry) + self.assertEqual( + _attr(span, "guardrail_status"), + "guardrail_intervened", + "guardrail_status must be exposed as a top-level span attribute", + ) + + def test_status_attribute_present_for_success(self): + entry = _slg_entry( + "success", + {"action": "NONE", "assessments": []}, + ) + span = self._emit_and_get_guardrail_span(entry) + self.assertEqual(_attr(span, "guardrail_status"), "success") + + def test_status_attribute_present_for_failed_to_respond(self): + entry = _slg_entry( + "guardrail_failed_to_respond", + {"error": "endpoint unreachable"}, + ) + span = self._emit_and_get_guardrail_span(entry) + self.assertEqual(_attr(span, "guardrail_status"), "guardrail_failed_to_respond") + + def test_violation_categories_surfaced_when_provider_populates_them(self): + """The provider hook (e.g. Bedrock) extracts violation categories + from the raw response BEFORE redaction and stamps them onto the + StandardLoggingGuardrailInformation entry. OTEL must surface that + list as a queryable span attribute so dashboards can group by + violation type without parsing the redacted guardrail_response.""" + entry = _slg_entry( + "guardrail_intervened", + _bedrock_block_response(), + violation_categories=["Fiduciary Advice", "VIOLENCE", "PROFANITY"], + ) + span = self._emit_and_get_guardrail_span(entry) + + categories = _attr(span, "guardrail_violation_categories") + self.assertIsNotNone( + categories, + "guardrail_violation_categories must be set when the entry " + "carries violation_categories", + ) + # Serialised as JSON to keep set_attribute typing simple. + as_str = categories if isinstance(categories, str) else repr(list(categories)) + self.assertIn("Fiduciary Advice", as_str) + self.assertIn("VIOLENCE", as_str) + self.assertIn("PROFANITY", as_str) + + def test_no_violation_categories_when_field_absent(self): + """When the provider didn't populate violation_categories (success + path, or provider didn't extract them), don't pollute the trace + with an empty attribute.""" + entry = _slg_entry("success", {"action": "NONE", "assessments": []}) + span = self._emit_and_get_guardrail_span(entry) + self.assertIsNone(_attr(span, "guardrail_violation_categories")) + + def test_no_violation_categories_when_field_is_empty(self): + """Empty list must not produce a span attribute either.""" + entry = _slg_entry( + "guardrail_intervened", + _bedrock_block_response(), + violation_categories=[], + ) + span = self._emit_and_get_guardrail_span(entry) + self.assertIsNone(_attr(span, "guardrail_violation_categories")) + + def test_guardrail_action_surfaced_when_provider_populates_it(self): + """The provider hook (e.g. Bedrock) writes its raw top-level + ``action`` string onto StandardLoggingGuardrailInformation as + ``guardrail_action``. OTEL must expose it as a queryable span + attribute so dashboards can pivot on the raw provider verdict + (Bedrock ``GUARDRAIL_INTERVENED`` / ``NONE``) without parsing + the redacted guardrail_response blob.""" + entry = _slg_entry( + "guardrail_intervened", + _bedrock_block_response(), + guardrail_action="GUARDRAIL_INTERVENED", + ) + span = self._emit_and_get_guardrail_span(entry) + self.assertEqual( + _attr(span, "guardrail_action"), + "GUARDRAIL_INTERVENED", + "guardrail_action must be exposed as a top-level span attribute", + ) + + def test_guardrail_action_surfaced_for_allowed_request(self): + """Even on the success path, the provider's raw action (e.g. + Bedrock ``NONE``) should be queryable so dashboards can group + allowed-vs-blocked counts off the same attribute.""" + entry = _slg_entry( + "success", + {"action": "NONE", "assessments": []}, + guardrail_action="NONE", + ) + span = self._emit_and_get_guardrail_span(entry) + self.assertEqual(_attr(span, "guardrail_action"), "NONE") + + def test_no_guardrail_action_when_field_absent(self): + """If the provider didn't populate the field (older payloads, + non-Bedrock providers without a top-level action), don't emit + an empty attribute.""" + entry = _slg_entry("success", {"action": "NONE", "assessments": []}) + span = self._emit_and_get_guardrail_span(entry) + self.assertIsNone(_attr(span, "guardrail_action")) + + +class TestMultipleGuardrailsOneBlocks(unittest.TestCase): + """When several guardrails run sequentially and only the last one + intervenes, every guardrail span must appear with its own status — + losing the early "allowed" spans would mask which checks ran.""" + + def test_all_guardrail_spans_emitted_with_per_entry_status(self): + otel, _, exporter = _make_otel() + + entries = [ + _slg_entry( + "success", + {"action": "NONE", "assessments": []}, + name="pii-mask", + start=1.0, + end=1.5, + ), + _slg_entry( + "success", + {"action": "NONE", "assessments": []}, + name="prompt-injection", + start=2.0, + end=2.2, + ), + _slg_entry( + "guardrail_intervened", + _bedrock_block_response(), + name="bedrock-policy", + start=3.0, + end=3.4, + ), + ] + kwargs = _kwargs_with_guardrail( + entries=entries, + include_exception=True, + ) + + start = datetime.now(timezone.utc) + end = start + timedelta(milliseconds=50) + otel._handle_failure(kwargs, response_obj=None, start_time=start, end_time=end) + + spans = exporter.get_finished_spans() + guardrail_spans = sorted( + (s for s in spans if s.name == GUARDRAIL_SPAN_NAME), + key=lambda s: (s.attributes or {}).get("guardrail_name", ""), + ) + self.assertEqual( + len(guardrail_spans), + 3, + "Every guardrail invocation must emit a span — even the ones " + "that allowed the request through before the blocker fired", + ) + + statuses = { + _attr(s, "guardrail_name"): _attr(s, "guardrail_status") + for s in guardrail_spans + } + self.assertEqual(statuses["pii-mask"], "success") + self.assertEqual(statuses["prompt-injection"], "success") + self.assertEqual(statuses["bedrock-policy"], "guardrail_intervened") + + +class TestCustomGuardrailEndToEnd(unittest.TestCase): + """End-to-end: a real ``CustomGuardrail`` subclass calls + ``add_standard_logging_guardrail_information_to_request_data`` and then + raises. We then drive ``_handle_failure`` with the resulting kwargs + (matching the shape ``async_failure_handler`` would build) and verify + the guardrail span carries the recorded information.""" + + def test_real_custom_guardrail_violation_path(self): + # Deliberately not importing fastapi here — the real Bedrock guardrail + # raises HTTPException, but the OTEL span flow is exception-type + # agnostic. Using a plain Exception keeps this test runnable in + # SDK-only installs that don't ship fastapi. + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.types.guardrails import GuardrailEventHooks + + class BlockingViolation(Exception): + pass + + class BlockingGuardrail(CustomGuardrail): + async def async_pre_call_hook( + self, + user_api_key_dict, + cache, + data, + call_type, + ): + start_ts = time.time() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="bedrock", + guardrail_json_response=_bedrock_block_response(), + request_data=data, + guardrail_status="guardrail_intervened", + start_time=start_ts, + end_time=start_ts + 0.01, + duration=0.01, + event_type=GuardrailEventHooks.pre_call, + tracing_detail={ + "violation_categories": ["Fiduciary Advice", "VIOLENCE"] + }, + ) + raise BlockingViolation("violation") + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hi"}], + "metadata": {}, + } + guardrail = BlockingGuardrail( + guardrail_name="blocking-test", + event_hook=GuardrailEventHooks.pre_call, + ) + + with self.assertRaises(BlockingViolation): + _run( + guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + cache=None, + data=request_data, + call_type="completion", + ) + ) + + slg_info = request_data["metadata"].get( + "standard_logging_guardrail_information" + ) + self.assertTrue( + slg_info, + "Guardrail must have recorded its information to request_data " + "BEFORE raising — otherwise the OTEL hook sees nothing", + ) + + # Now simulate the OTEL failure handler picking up this metadata + otel, _, exporter = _make_otel() + kwargs = _kwargs_with_guardrail( + entries=slg_info, + include_exception=True, + ) + start = datetime.now(timezone.utc) + end = start + timedelta(milliseconds=15) + otel._handle_failure(kwargs, response_obj=None, start_time=start, end_time=end) + + spans = exporter.get_finished_spans() + guardrail_spans = [s for s in spans if s.name == GUARDRAIL_SPAN_NAME] + self.assertEqual(len(guardrail_spans), 1) + self.assertEqual( + _attr(guardrail_spans[0], "guardrail_status"), + "guardrail_intervened", + ) + self.assertEqual( + _attr(guardrail_spans[0], "guardrail_name"), + "blocking-test", + ) + # End-to-end: the violation_categories the guardrail passed through + # tracing_detail must arrive as a queryable span attribute. + categories = _attr(guardrail_spans[0], "guardrail_violation_categories") + self.assertIsNotNone(categories) + self.assertIn("Fiduciary Advice", str(categories)) + self.assertIn("VIOLENCE", str(categories)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py new file mode 100644 index 00000000000..922d2fe8a15 --- /dev/null +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -0,0 +1,1012 @@ +""" +Tests for the Rubrik LiteLLM plugin. + +Covers initialization, apply_guardrail tool blocking (all allowed, all blocked, +partial blocking, fail-open), batch logging, and Anthropic format handling. +""" + +import os +from typing import Any, Dict +from unittest.mock import AsyncMock, Mock, patch + +import httpx +import pytest + +from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.rubrik import RubrikLogger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +from tests.test_litellm.integrations.rubrik_test_helpers import ( + make_inputs_with_tools, + make_tool_call_dict, +) + + +@pytest.fixture +def mock_env(): + """Set up environment variables for testing.""" + with patch.dict( + os.environ, + { + "RUBRIK_WEBHOOK_URL": "http://localhost:8080", + "RUBRIK_API_KEY": "test-api-key", + }, + ): + yield + + +@pytest.fixture +def handler(mock_env): + """Create a RubrikLogger instance for testing.""" + with patch("asyncio.create_task", Mock()): + return RubrikLogger() + + +# -- Initialization ----------------------------------------------------------- + + +class TestInitialization: + def test_init_success(self, mock_env): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger() + assert ( + handler.tool_blocking_endpoint + == "http://localhost:8080/v1/after_completion/openai/v1" + ) + assert handler.logging_endpoint == "http://localhost:8080/v1/litellm/batch" + assert handler.key == "test-api-key" + assert isinstance(handler.tool_blocking_client, AsyncHTTPHandler) + + def test_init_with_constructor_params(self): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(api_key="ctor-key", api_base="http://ctor-host:9090") + assert handler.key == "ctor-key" + assert ( + handler.tool_blocking_endpoint + == "http://ctor-host:9090/v1/after_completion/openai/v1" + ) + + def test_init_without_url(self): + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(ValueError, match="Rubrik webhook URL not configured"): + RubrikLogger() + + def test_init_without_api_key(self): + with patch.dict( + os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080"}, clear=True + ): + with patch("asyncio.create_task", Mock()): + assert RubrikLogger().key is None + + def test_trailing_slash_removed(self): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://localhost:8080/"}): + with patch("asyncio.create_task", Mock()): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://localhost:8080/v1/after_completion/openai/v1" + ) + + def test_v1_suffix_stripped_as_substring_not_charset(self): + with patch("asyncio.create_task", Mock()): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v1"}): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://host/v1/after_completion/openai/v1" + ) + + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host/v11"}): + assert ( + RubrikLogger().tool_blocking_endpoint + == "http://host/v11/v1/after_completion/openai/v1" + ) + + def test_sampling_rate_fractional(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "0.5"}, + ): + assert RubrikLogger().sampling_rate == 0.5 + + def test_sampling_rate_invalid_ignored(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "abc"}, + ): + assert RubrikLogger().sampling_rate == 1.0 + + def test_sampling_rate_clamped(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "2.0"}, + ): + assert RubrikLogger().sampling_rate == 1.0 + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_SAMPLING_RATE": "-0.5"}, + ): + assert RubrikLogger().sampling_rate == 0.0 + + def test_batch_size_invalid_ignored(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "abc"}, + ): + # Should use default without crashing + assert isinstance(RubrikLogger().batch_size, int) + + def test_batch_size_valid(self): + with patch("asyncio.create_task", Mock()): + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://host", "RUBRIK_BATCH_SIZE": "256"}, + ): + assert RubrikLogger().batch_size == 256 + + def test_init_outside_event_loop_does_not_raise(self): + """Instantiation without a running event loop must not raise RuntimeError.""" + with patch.dict( + os.environ, + {"RUBRIK_WEBHOOK_URL": "http://localhost:8080", "RUBRIK_API_KEY": "k"}, + ): + # Do NOT patch asyncio.create_task — the real call should be + # guarded and fall back gracefully when there is no event loop. + handler = RubrikLogger() + assert handler.tool_blocking_endpoint.startswith("http://localhost:8080") + # Without a running loop at init, the periodic flush task should be + # deferred so batches still get drained once a log event arrives. + assert handler._flush_task is None + + @pytest.mark.asyncio + async def test_periodic_flush_task_started_lazily_on_first_log(self, mock_env): + """Loggers instantiated outside an event loop must still start the + periodic flush task on first use to drain low-traffic batches.""" + # Simulate sync-init by hiding the running loop from the constructor. + with patch( + "litellm.integrations.rubrik.asyncio.get_running_loop", + side_effect=RuntimeError("no running loop"), + ): + handler = RubrikLogger() + assert handler._flush_task is None + + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "id": "litellm-id", + }, + "litellm_call_id": "litellm-id", + "litellm_params": {}, + } + with patch.object(handler, "_log_batch_to_rubrik", AsyncMock()): + await handler.async_log_success_event(kwargs, None, None, None) + + assert handler._flush_task is not None + handler._flush_task.cancel() + + def test_event_hook_defaults_to_post_call_when_none_passed(self, mock_env): + """`initialize_guardrail` always passes ``event_hook=litellm_params.mode`` + (which is ``None`` when the user omits ``mode``). The logger must coerce + a None ``event_hook`` to ``post_call`` rather than leaving it as None, + which would otherwise cause the guardrail to run on every event hook.""" + from litellm.types.guardrails import GuardrailEventHooks + + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(event_hook=None) + assert handler.event_hook == GuardrailEventHooks.post_call + + def test_explicit_event_hook_preserved(self, mock_env): + from litellm.types.guardrails import GuardrailEventHooks + + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(event_hook=GuardrailEventHooks.pre_call) + assert handler.event_hook == GuardrailEventHooks.pre_call + + def test_default_on_defaults_to_true_when_none_passed(self, mock_env): + """`initialize_guardrail` always passes ``default_on=litellm_params.default_on`` + (which is ``None`` when the user omits ``default_on``). The logger must + coerce a None ``default_on`` to True, otherwise ``should_run_guardrail`` + (which checks ``self.default_on is True``) silently skips the guardrail.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=None) + assert handler.default_on is True + + def test_explicit_default_on_false_preserved(self, mock_env): + """A user explicitly setting ``default_on: false`` in their guardrail + config must NOT be silently overridden to True.""" + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=False) + assert handler.default_on is False + + def test_explicit_default_on_true_preserved(self, mock_env): + with patch("asyncio.create_task", Mock()): + handler = RubrikLogger(default_on=True) + assert handler.default_on is True + + def test_headers_with_api_key(self, handler): + assert handler._headers["Authorization"] == "Bearer test-api-key" + assert handler._headers["Content-Type"] == "application/json" + + def test_headers_without_api_key(self): + with patch.dict(os.environ, {"RUBRIK_WEBHOOK_URL": "http://host"}, clear=True): + with patch("asyncio.create_task", Mock()): + h = RubrikLogger() + assert "Authorization" not in h._headers + + +# -- Batch Logging ------------------------------------------------------------ + + +@pytest.mark.asyncio +class TestBatchLogging: + async def test_log_success_event_appends_to_queue(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + + async def test_log_failure_event_appends_to_queue(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "error", + }, + } + await handler.async_log_failure_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + + async def test_log_success_event_sampling_skips(self, handler): + handler.sampling_rate = 0.0 + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 0 + + async def test_flush_queue_sends_batch(self, handler): + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + mock_response = Mock() + mock_response.status_code = 200 + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock(return_value=mock_response) + await handler.flush_queue() + handler.async_httpx_client.post.assert_called_once() + assert len(handler.log_queue) == 0 + + async def test_flush_queue_preserves_events_added_during_send(self, handler): + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + + async def mock_post(*_args, **_kwargs): + handler.log_queue.append({"msg": "c"}) + mock_response = Mock() + mock_response.raise_for_status = Mock() + return mock_response + + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = mock_post + + await handler.flush_queue() + + assert handler.log_queue == [{"msg": "c"}] + + async def test_async_send_batch_does_not_drain_events(self, handler): + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + + async def mock_post(*_args, **_kwargs): + handler.log_queue.append({"msg": "c"}) + mock_response = Mock() + mock_response.raise_for_status = Mock() + return mock_response + + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = mock_post + + await handler.async_send_batch() + + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}, {"msg": "c"}] + + async def test_log_batch_error_does_not_crash_and_preserves_events(self, handler): + """A failed batch send must not crash the caller AND must preserve the + original events in the queue so they can be retried on the next flush. + Previously the events were silently dropped on HTTP 5xx / network errors. + """ + handler.log_queue = [{"msg": "a"}] + mock_response = Mock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + mock_response.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + "err", request=Mock(), response=mock_response + ) + ) + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock(return_value=mock_response) + await handler.flush_queue() + assert handler.log_queue == [{"msg": "a"}] + + async def test_log_batch_network_error_preserves_events(self, handler): + """Network/timeout errors must also preserve the in-flight events.""" + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock( + side_effect=httpx.TimeoutException("timeout") + ) + await handler.flush_queue() + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}] + + async def test_enqueue_drops_oldest_when_queue_exceeds_max_size(self, handler): + """A sustained Rubrik webhook outage must not let the in-memory retry + queue grow without bound. Once max_queue_size is exceeded, the oldest + events are dropped to make room for new ones.""" + handler.max_queue_size = 3 + handler.batch_size = 10**6 # disable size-triggered flush + handler.flush_queue = AsyncMock() + for i in range(5): + await handler._enqueue_log_event( + kwargs={ + "standard_logging_object": { + "messages": [{"role": "user", "content": f"hi-{i}"}], + "response": "hello", + }, + }, + event_type="success", + ) + assert len(handler.log_queue) == 3 + retained = [item["messages"][0]["content"] for item in handler.log_queue] + assert retained == ["hi-2", "hi-3", "hi-4"] + + async def test_log_batch_failure_preserves_events_added_during_send(self, handler): + """Failure must preserve both the snapshot AND events appended mid-flush.""" + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + + async def mock_post(*_args, **_kwargs): + handler.log_queue.append({"msg": "c"}) + mock_response = Mock() + mock_response.status_code = 500 + mock_response.text = "boom" + mock_response.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + "err", request=Mock(), response=mock_response + ) + ) + return mock_response + + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = mock_post + + await handler.flush_queue() + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}, {"msg": "c"}] + + async def test_system_prompt_prepended_to_messages(self, handler): + kwargs = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "system": "You are a helpful assistant.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + msgs = handler.log_queue[0]["messages"] + assert msgs[0]["role"] == "system" + assert msgs[0]["content"] == "You are a helpful assistant." + + async def test_system_prompt_with_dict_messages(self, handler): + kwargs = { + "standard_logging_object": { + "messages": {"role": "user", "content": "hi"}, + "response": "hello", + }, + "system": "Be concise.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert len(handler.log_queue) == 1 + msgs = handler.log_queue[0]["messages"] + assert isinstance(msgs, list) + assert msgs[0]["role"] == "system" + assert msgs[1] == {"role": "user", "content": "hi"} + + async def test_anthropic_id_normalization(self, handler): + kwargs = { + "standard_logging_object": { + "id": "chatcmpl-original", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "litellm_params": { + "proxy_server_request": { + "url": "http://proxy/v1/messages", + }, + }, + "litellm_call_id": "litellm-call-123", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert handler.log_queue[0]["id"] == "litellm-call-123" + + async def test_non_anthropic_id_unchanged(self, handler): + kwargs = { + "standard_logging_object": { + "id": "chatcmpl-original", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + }, + "litellm_params": { + "proxy_server_request": { + "url": "http://proxy/v1/chat/completions", + }, + }, + "litellm_call_id": "litellm-call-123", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + assert handler.log_queue[0]["id"] == "chatcmpl-original" + + async def test_payload_deep_copied_not_mutated(self, handler): + """Verify the shared standard_logging_object is not mutated.""" + original_payload = { + "id": "original-id", + "messages": [{"role": "user", "content": "hi"}], + "response": "hello", + } + kwargs = { + "standard_logging_object": original_payload, + "system": "System prompt.", + } + await handler.async_log_success_event( + kwargs=kwargs, response_obj=None, start_time=None, end_time=None + ) + # Original payload should NOT have been mutated + assert original_payload["id"] == "original-id" + assert len(original_payload["messages"]) == 1 + + +# -- Tool Blocking (apply_guardrail) ------------------------------------------ + + +def _mock_service_response(response_json): + """Create a mock tool blocking client that returns the given JSON.""" + + async def mock_post(*_args, **kwargs): + mock_resp = Mock() + mock_resp.json.return_value = response_json + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + return mock_client + + +def _echo_service(): + """Create a mock tool blocking client that echoes the payload back.""" + + async def mock_post(*_args, **kwargs): + mock_resp = Mock() + mock_resp.json.return_value = kwargs.get("json", {}).get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + return mock_client + + +@pytest.mark.asyncio +class TestApplyGuardrail: + async def test_skips_requests(self, handler): + inputs = make_inputs_with_tools([make_tool_call_dict("call_1", "test_tool")]) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + async def test_no_tool_calls(self, handler): + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_all_allowed(self, handler): + tc1 = make_tool_call_dict("call_1", "get_weather") + tc2 = make_tool_call_dict("call_2", "get_time") + inputs = make_inputs_with_tools([tc1, tc2]) + + handler.tool_blocking_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_all_blocked(self, handler): + tc1 = make_tool_call_dict("call_1", "delete_table") + tc2 = make_tool_call_dict("call_2", "drop_database") + inputs = make_inputs_with_tools([tc1, tc2]) + + handler.tool_blocking_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "Tool blocked by policy", + "tool_calls": [], + } + } + ], + } + ) + + with pytest.raises(ModifyResponseException) as exc_info: + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert "Tool blocked by policy" in exc_info.value.message + + async def test_partial_blocking(self, handler): + tc_blocked = make_tool_call_dict("call_A", "blocked_tool") + tc_allowed = make_tool_call_dict("call_B", "allowed_tool") + inputs = make_inputs_with_tools([tc_blocked, tc_allowed]) + + async def mock_post(*_args, **kwargs): + payload = kwargs.get("json", {}).get("response", {}) + all_tcs = payload["choices"][0]["message"]["tool_calls"] + allowed = [tc for tc in all_tcs if tc.get("id") == "call_B"] + mock_resp = Mock() + mock_resp.json.return_value = { + "choices": [ + { + "message": { + "role": "assistant", + "content": "blocked", + "tool_calls": allowed, + } + } + ], + } + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + async def test_service_failure_fail_open(self, handler): + tc1 = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc1]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_service_empty_choices_fail_open(self, handler): + tc1 = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc1]) + + handler.tool_blocking_client = _mock_service_response({"choices": []}) + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_blocking_service_payload_format(self, handler): + tc1 = make_tool_call_dict("call_1", "get_weather", '{"location": "SF"}') + tc2 = make_tool_call_dict("call_2", "send_email", '{"to": "user@example.com"}') + inputs = make_inputs_with_tools([tc1, tc2]) + + captured_payload: Dict[str, Any] = {} + + async def mock_post(*_args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + mock_resp = Mock() + mock_resp.json.return_value = captured_payload.get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + # Verify envelope structure + assert "request" in captured_payload + assert "response" in captured_payload + + response_data = captured_payload["response"] + message = response_data["choices"][0]["message"] + assert message["role"] == "assistant" + assert len(message["tool_calls"]) == 2 + assert message["tool_calls"][0]["id"] == "call_1" + assert message["tool_calls"][0]["function"]["name"] == "get_weather" + assert message["tool_calls"][1]["id"] == "call_2" + assert message["tool_calls"][1]["function"]["name"] == "send_email" + + async def test_request_data_included_in_envelope(self, handler): + tc = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc]) + + captured_payload: Dict[str, Any] = {} + + async def mock_post(*_args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + mock_resp = Mock() + mock_resp.json.return_value = captured_payload.get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + logging_obj = Mock() + logging_obj.model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-4", + "litellm_params": { + "proxy_server_request": {"url": "/chat/completions"}, + }, + } + + await handler.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + + req = captured_payload["request"] + assert req["model"] == "gpt-4" + assert req["messages"] == [{"role": "user", "content": "hi"}] + + async def test_proxy_server_request_headers_stripped(self, handler): + tc = make_tool_call_dict("call_1", "test_tool") + inputs = make_inputs_with_tools([tc]) + + captured_payload: Dict[str, Any] = {} + + async def mock_post(*_args, **kwargs): + captured_payload.update(kwargs.get("json", {})) + mock_resp = Mock() + mock_resp.json.return_value = captured_payload.get("response", {}) + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + handler.tool_blocking_client = mock_client + + logging_obj = Mock() + logging_obj.model_call_details = { + "messages": [{"role": "user", "content": "hi"}], + "model": "gpt-4", + "litellm_params": { + "proxy_server_request": { + "url": "/chat/completions", + "method": "POST", + "headers": { + "authorization": "Bearer sk-litellm-secret", + "cookie": "session=abc", + "x-api-key": "leaked-key", + }, + "body": {"api_key": "sk-upstream-secret"}, + }, + }, + } + + await handler.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=logging_obj, + ) + + forwarded = captured_payload["request"]["proxy_server_request"] + assert forwarded == {"url": "/chat/completions", "method": "POST"} + + +# -- Anthropic format ---------------------------------------------------------- + + +@pytest.mark.asyncio +class TestApplyGuardrailAnthropicFormat: + """Verify blocking works correctly regardless of original provider format. + + The framework converts Anthropic tool_use blocks to OpenAI-format + tool_calls before calling apply_guardrail. + """ + + async def test_single_tool_allowed(self, handler): + tc = make_tool_call_dict( + "toolu_123", "get_weather", '{"location": "Portland, OR"}' + ) + inputs = make_inputs_with_tools([tc], texts=["I'll check the weather."]) + + handler.tool_blocking_client = _echo_service() + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + async def test_single_tool_blocked(self, handler): + tc = make_tool_call_dict("toolu_123", "dangerous_tool", '{"arg": "value"}') + inputs = make_inputs_with_tools([tc]) + + handler.tool_blocking_client = _mock_service_response( + { + "choices": [ + { + "message": { + "role": "assistant", + "content": "blocked", + "tool_calls": [], + } + } + ], + } + ) + + with pytest.raises(ModifyResponseException): + await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + async def test_text_only_response_no_blocking(self, handler): + from litellm.types.utils import GenericGuardrailAPIInputs + + inputs = GenericGuardrailAPIInputs(texts=["Hello! I'm Claude."]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock() + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + + assert result is inputs + mock_client.post.assert_not_called() + + async def test_service_failure_preserves_tools(self, handler): + tc = make_tool_call_dict("toolu_123", "get_weather", '{"location": "SF"}') + inputs = make_inputs_with_tools([tc]) + + mock_client = AsyncMock() + mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) + handler.tool_blocking_client = mock_client + + result = await handler.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + +# -- Normalize tool calls ------------------------------------------------------ + + +class TestNormalizeToolCalls: + def test_dict_input(self): + tc = make_tool_call_dict("call_1", "test", '{"a": 1}') + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_1" + assert result[0].function.name == "test" + assert result[0].function.arguments == '{"a": 1}' + + def test_typed_object_input(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_2", + type="function", + function=Function(name="fn", arguments="{}"), + ) + result = RubrikLogger._normalize_tool_calls([tc]) + assert len(result) == 1 + assert result[0].id == "call_2" + assert result[0].function.name == "fn" + + def test_unsupported_type_raises(self): + with pytest.raises(TypeError, match="Cannot normalize"): + RubrikLogger._normalize_tool_calls(["not_a_tool_call"]) + + +# -- Extract blocked tools ----------------------------------------------------- + + +class TestExtractBlockedTools: + def test_all_allowed_returns_none(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="fn", arguments="{}") + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [{"id": "call_1"}], + "content": "", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + assert result is None + + def test_some_blocked_returns_explanation(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc1 = ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function(name="fn1", arguments="{}"), + ) + tc2 = ChatCompletionMessageToolCall( + id="call_2", + type="function", + function=Function(name="fn2", arguments="{}"), + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [{"id": "call_1"}], + "content": "blocked fn2", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + assert result is not None + assert "blocked fn2" in result + + def test_empty_choices_raises(self): + with pytest.raises(Exception, match="empty response"): + RubrikLogger._extract_blocked_tools({"choices": []}, []) + + def test_null_tool_calls_treated_as_all_blocked(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc = ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="fn", arguments="{}") + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": None, + "content": "blocked everything", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc]) + assert result is not None + assert "blocked everything" in result + + def test_duplicate_ids_block_when_only_one_returned(self): + from litellm.types.utils import ChatCompletionMessageToolCall, Function + + tc1 = ChatCompletionMessageToolCall( + id="call_dup", + type="function", + function=Function(name="fn", arguments="{}"), + ) + tc2 = ChatCompletionMessageToolCall( + id="call_dup", + type="function", + function=Function(name="fn", arguments="{}"), + ) + service_resp = { + "choices": [ + { + "message": { + "tool_calls": [{"id": "call_dup"}], + "content": "blocked duplicate", + } + } + ] + } + result = RubrikLogger._extract_blocked_tools(service_resp, [tc1, tc2]) + assert result is not None + assert "blocked duplicate" in result + + +# -- Sanitize proxy server request ------------------------------------------- + + +class TestSanitizeProxyServerRequest: + def test_drops_headers_and_body(self): + proxy_request = { + "url": "/chat/completions", + "method": "POST", + "headers": { + "authorization": "Bearer sk-litellm-secret", + "cookie": "session=abc", + "content-type": "application/json", + }, + "body": {"api_key": "sk-upstream-secret", "model": "gpt-4"}, + } + result = RubrikLogger._sanitize_proxy_server_request(proxy_request) + assert result == {"url": "/chat/completions", "method": "POST"} + + def test_none_passthrough(self): + assert RubrikLogger._sanitize_proxy_server_request(None) is None + + def test_non_dict_passthrough(self): + assert RubrikLogger._sanitize_proxy_server_request("not a dict") == "not a dict" + + def test_partial_dict(self): + result = RubrikLogger._sanitize_proxy_server_request({"url": "/v1/messages"}) + assert result == {"url": "/v1/messages"} + + +# -- Resolve model ------------------------------------------------------------- + + +class TestResolveModel: + def test_model_from_response(self): + from unittest.mock import Mock + + response = Mock() + response.model = "gpt-4" + result = RubrikLogger._resolve_model({"response": response}, {}) + assert result == "gpt-4" + + def test_model_from_call_details(self): + result = RubrikLogger._resolve_model({}, {"model": "claude-3"}) + assert result == "claude-3" + + def test_fallback_to_unknown(self): + result = RubrikLogger._resolve_model({}, {}) + assert result == "unknown" + + def test_empty_model_on_response_returns_unknown(self): + from unittest.mock import Mock + + response = Mock() + response.model = "" + result = RubrikLogger._resolve_model( + {"response": response}, {"model": "fallback"} + ) + assert result == "unknown" diff --git a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py index 2e596a72158..524589abf5e 100644 --- a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py +++ b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py @@ -1,10 +1,11 @@ """ Tests for Gemini Interactions API transformation. -Covers credential leak prevention changes: -- validate_environment sets x-goog-api-key header -- get_complete_url excludes API key from URL -- get/delete/cancel interaction request URLs exclude API key +Covers: +- validate_environment: x-goog-api-key header, Api-Revision schema selection +- get_complete_url: API key excluded from URL +- get/delete/cancel interaction request URLs +- transform_request: response_mime_type coalescing, image_config migration """ import os @@ -15,6 +16,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../..")) +import litellm from litellm.interactions.litellm_responses_transformation.streaming_iterator import ( LiteLLMResponsesInteractionsStreamingIterator, ) @@ -22,7 +24,6 @@ from litellm.llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig, ) from litellm.types.llms.openai import ( - ContentPartAddedEvent, OutputTextDeltaEvent, ResponseCompletedEvent, ResponseCreatedEvent, @@ -85,6 +86,30 @@ class TestValidateEnvironment: assert headers["X-Custom"] == "value" assert headers["x-goog-api-key"] == "test-key" + def test_api_revision_new_schema_by_default(self, config): + # Default: use_legacy_interactions_schema=False → new steps schema + original = litellm.use_legacy_interactions_schema + try: + litellm.use_legacy_interactions_schema = False + headers = config.validate_environment( + headers={}, model="gemini-2.5-flash", litellm_params=None + ) + assert headers["Api-Revision"] == "2026-05-20" + finally: + litellm.use_legacy_interactions_schema = original + + def test_api_revision_legacy_schema_when_flag_set(self, config): + # Flag on → legacy outputs schema until June 8, 2026 + original = litellm.use_legacy_interactions_schema + try: + litellm.use_legacy_interactions_schema = True + headers = config.validate_environment( + headers={}, model="gemini-2.5-flash", litellm_params=None + ) + assert headers["Api-Revision"] == "2026-05-07" + finally: + litellm.use_legacy_interactions_schema = original + class TestGetCompleteUrl: def test_url_excludes_api_key(self, config): @@ -127,7 +152,12 @@ class TestTransformRequest: request_body = config.transform_request( model=None, agent="my-custom-slides-agent", - input=[{"type": "text", "text": "Create a 5-slide presentation about AI trends."}], + input=[ + { + "type": "text", + "text": "Create a 5-slide presentation about AI trends.", + } + ], optional_params={ "environment": "remote", "stream": False, @@ -172,158 +202,7 @@ class TestTransformRequest: ) assert request_body["environment"] == env_id -class TestStreamingIterator: - def _make_iterator(self) -> LiteLLMResponsesInteractionsStreamingIterator: - return LiteLLMResponsesInteractionsStreamingIterator( - model="gpt-5.4", - litellm_custom_stream_wrapper=MagicMock(), - request_input="hi", - optional_params={}, - ) - def _make_text_delta( - self, text: str, item_id: str = "item_1" - ) -> OutputTextDeltaEvent: - event = MagicMock(spec=OutputTextDeltaEvent) - event.delta = text - event.item_id = item_id - return event - - def _make_part_added(self, item_id: str = "item_1") -> ContentPartAddedEvent: - event = MagicMock(spec=ContentPartAddedEvent) - event.item_id = item_id - return event - - def _make_response_created(self) -> ResponseCreatedEvent: - event = MagicMock(spec=ResponseCreatedEvent) - event.response = MagicMock(id="resp_123") - return event - - def test_content_delta_includes_type_field(self): - """content.delta events must carry delta.type='text' so the UI can display them.""" - it = self._make_iterator() - it.sent_interaction_start = True - it.sent_content_start = True - - chunk = it._transform_responses_chunk_to_interactions_chunk( - self._make_text_delta("Hello") - ) - - assert chunk is not None - assert chunk.event_type == "content.delta" - assert chunk.delta == {"type": "text", "text": "Hello"} - - def test_response_part_added_emits_content_start(self): - """ContentPartAddedEvent (arrives before text deltas) should emit content.start - so the first OutputTextDeltaEvent immediately emits content.delta without dropping text. - """ - it = self._make_iterator() - it.sent_interaction_start = True - - chunk = it._transform_responses_chunk_to_interactions_chunk( - self._make_part_added() - ) - - assert chunk is not None - assert chunk.event_type == "content.start" - assert it.sent_content_start is True - - def test_first_text_delta_not_dropped_when_part_added_seen(self): - """After ContentPartAddedEvent, the first text delta must yield content.delta - (not content.start), preserving the token text.""" - it = self._make_iterator() - it.sent_interaction_start = True - it._transform_responses_chunk_to_interactions_chunk(self._make_part_added()) - - chunk = it._transform_responses_chunk_to_interactions_chunk( - self._make_text_delta("Hello") - ) - - assert chunk is not None - assert chunk.event_type == "content.delta" - assert chunk.delta is not None - assert chunk.delta.get("text") == "Hello" - - def test_part_added_emits_interaction_start_fallback_when_not_sent(self): - """If ContentPartAddedEvent arrives before any ResponseCreatedEvent, - the iterator must emit interaction.start before content.start to honor - the documented event ordering contract.""" - it = self._make_iterator() - - chunk = it._transform_responses_chunk_to_interactions_chunk( - self._make_part_added(item_id="item_42") - ) - - assert chunk is not None - assert chunk.event_type == "interaction.start" - assert chunk.id == "item_42" - assert chunk.status == "in_progress" - assert chunk.model == "gpt-5.4" - assert it.sent_interaction_start is True - assert it.sent_content_start is False - - def test_part_added_returns_none_when_already_started(self): - """A second ContentPartAddedEvent (after content.start was already emitted) - should be a no-op so we don't re-emit content.start.""" - it = self._make_iterator() - it.sent_interaction_start = True - it.sent_content_start = True - - chunk = it._transform_responses_chunk_to_interactions_chunk( - self._make_part_added() - ) - - assert chunk is None - - def test_part_added_without_item_id_falls_back_to_self_id(self): - """When ContentPartAddedEvent has no item_id and we emit the interaction.start - fallback, the id must default to an interaction_ string.""" - it = self._make_iterator() - event = MagicMock(spec=ContentPartAddedEvent) - event.item_id = None - - chunk = it._transform_responses_chunk_to_interactions_chunk(event) - - assert chunk is not None - assert chunk.event_type == "interaction.start" - assert chunk.id == f"interaction_{id(it)}" - - def test_first_text_delta_not_dropped_when_no_prior_start_events(self): - """When OutputTextDeltaEvent arrives before any ResponseCreatedEvent or - ContentPartAddedEvent, the iterator must emit interaction.start *and* - immediately follow with a content.start that carries this delta's text, - so the first token is never silently dropped from the stream.""" - events = [ - self._make_text_delta("Hello"), - self._make_text_delta(" World"), - ] - wrapper = MagicMock() - wrapper.__iter__ = lambda self: iter(events) - wrapper.__next__ = lambda self, _it=iter(events): next(_it) - it = LiteLLMResponsesInteractionsStreamingIterator( - model="gpt-5.4", - litellm_custom_stream_wrapper=wrapper, - request_input="hi", - optional_params={}, - ) - - first = it._transform_responses_chunk_to_interactions_chunk(events[0]) - assert first is not None - assert first.event_type == "interaction.start" - assert it.sent_interaction_start is True - assert it.sent_content_start is True - assert len(it._pending_events) == 1 - pending = it._pending_events[0] - assert pending.event_type == "content.start" - assert pending.delta == {"type": "text", "text": "Hello"} - - second = it._transform_responses_chunk_to_interactions_chunk(events[1]) - assert second is not None - assert second.event_type == "content.delta" - assert second.delta == {"type": "text", "text": " World"} - - -class TestTransformRequest: def test_stream_param_included_in_request_body(self, config): """When stream=True is in optional_params, the request body must include it so the proxy forwards the SSE streaming flag to Google's backend.""" @@ -352,6 +231,273 @@ class TestTransformRequest: assert "stream" not in body +class TestStreamingIterator: + def _make_iterator( + self, use_legacy: bool = False + ) -> LiteLLMResponsesInteractionsStreamingIterator: + original = litellm.use_legacy_interactions_schema + litellm.use_legacy_interactions_schema = use_legacy + try: + return LiteLLMResponsesInteractionsStreamingIterator( + model="gpt-5.4", + litellm_custom_stream_wrapper=MagicMock(), + request_input="hi", + optional_params={}, + ) + finally: + litellm.use_legacy_interactions_schema = original + + def _make_text_delta( + self, text: str, item_id: str = "item_1" + ) -> OutputTextDeltaEvent: + event = MagicMock(spec=OutputTextDeltaEvent) + event.delta = text + event.item_id = item_id + return event + + def _make_response_created(self) -> ResponseCreatedEvent: + event = MagicMock(spec=ResponseCreatedEvent) + event.response = MagicMock(id="resp_123") + return event + + def test_step_delta_includes_type_field(self): + """step.delta events must carry delta.type='text' so the UI can display them.""" + it = self._make_iterator(use_legacy=False) + it.sent_interaction_start = True + it.sent_content_start = True + + chunk = it._transform_responses_chunk_to_interactions_chunk( + self._make_text_delta("Hello") + ) + + assert chunk is not None + assert chunk.event_type == "step.delta" + assert chunk.delta == {"type": "text", "text": "Hello"} + + def test_content_delta_legacy_schema(self): + """Legacy schema emits content.delta with type and text fields.""" + it = self._make_iterator(use_legacy=True) + it.sent_interaction_start = True + it.sent_content_start = True + + chunk = it._transform_responses_chunk_to_interactions_chunk( + self._make_text_delta("Hello") + ) + + assert chunk is not None + assert chunk.event_type == "content.delta" + assert chunk.delta == {"type": "text", "text": "Hello"} + + def test_response_created_emits_interaction_created(self): + it = self._make_iterator(use_legacy=False) + + chunk = it._transform_responses_chunk_to_interactions_chunk( + self._make_response_created() + ) + + assert chunk is not None + assert chunk.event_type == "interaction.created" + assert chunk.id == "resp_123" + assert it.sent_interaction_start is True + + def test_response_created_emits_interaction_start_legacy(self): + it = self._make_iterator(use_legacy=True) + + chunk = it._transform_responses_chunk_to_interactions_chunk( + self._make_response_created() + ) + + assert chunk is not None + assert chunk.event_type == "interaction.start" + assert chunk.id == "resp_123" + + def test_text_delta_sequence_new_schema(self): + """First chunk yields created + step.start + step.delta; later chunks yield step.delta.""" + it = self._make_iterator(use_legacy=False) + + first_events = it._events_for_chunk(self._make_text_delta("Hello")) + assert [e.event_type for e in first_events] == [ + "interaction.created", + "step.start", + "step.delta", + ] + assert first_events[-1].delta == {"type": "text", "text": "Hello"} + assert it.sent_interaction_start is True + assert it.sent_content_start is True + + second_events = it._events_for_chunk(self._make_text_delta(" World")) + assert [e.event_type for e in second_events] == ["step.delta"] + assert second_events[0].delta == {"type": "text", "text": " World"} + + third_events = it._events_for_chunk(self._make_text_delta("!")) + assert [e.event_type for e in third_events] == ["step.delta"] + assert third_events[0].delta == {"type": "text", "text": "!"} + + def test_text_delta_sequence_legacy_schema(self): + """Legacy: first chunk yields interaction.start + content.start + content.delta.""" + it = self._make_iterator(use_legacy=True) + + first_events = it._events_for_chunk(self._make_text_delta("Hello")) + assert [e.event_type for e in first_events] == [ + "interaction.start", + "content.start", + "content.delta", + ] + assert first_events[-1].delta == {"type": "text", "text": "Hello"} + + second_events = it._events_for_chunk(self._make_text_delta(" World")) + assert [e.event_type for e in second_events] == ["content.delta"] + assert second_events[0].delta == {"type": "text", "text": " World"} + + def test_first_text_delta_without_item_id_uses_fallback_id(self): + it = self._make_iterator(use_legacy=False) + event = self._make_text_delta("Hi") + event.item_id = None + + events = it._events_for_chunk(event) + + assert events[0].event_type == "interaction.created" + assert events[0].id == f"interaction_{id(it)}" + + def test_first_text_delta_emits_text_via_compat_shim(self): + """The legacy single-chunk shim must surface the synthetic events AND the delta.""" + it = self._make_iterator(use_legacy=False) + + first = it._transform_responses_chunk_to_interactions_chunk( + self._make_text_delta("Hello") + ) + assert first is not None + assert first.event_type == "interaction.created" + + second = it.__next__() if it._pending_events else None + assert second is not None + assert second.event_type == "step.start" + + third = it.__next__() if it._pending_events else None + assert third is not None + assert third.event_type == "step.delta" + assert third.delta == {"type": "text", "text": "Hello"} + + def test_response_created_then_text_delta_emits_step_start_and_delta(self): + """Realistic flow: response.created arrives first, then text delta.""" + it = self._make_iterator(use_legacy=False) + + first = it._events_for_chunk(self._make_response_created()) + assert [e.event_type for e in first] == ["interaction.created"] + + second = it._events_for_chunk(self._make_text_delta("Hello")) + assert [e.event_type for e in second] == ["step.start", "step.delta"] + assert second[-1].delta == {"type": "text", "text": "Hello"} + + def test_no_text_token_is_dropped_during_streaming(self): + """Concatenated step.delta payloads must equal the upstream text.""" + it = self._make_iterator(use_legacy=False) + + chunks = ["Hello", " ", "world", "!"] + emitted_text = "" + for c in chunks: + for ev in it._events_for_chunk(self._make_text_delta(c)): + if ev.event_type == "step.delta": + assert ev.delta is not None + emitted_text += ev.delta["text"] + + assert emitted_text == "Hello world!" + + def test_stop_iteration_fallback_emits_completion_event(self): + """If upstream ends without ResponseCompletedEvent, terminal events still flow.""" + from unittest.mock import MagicMock + + text_event = self._make_text_delta("hi") + sync_iter = MagicMock() + sync_iter.__iter__ = lambda self: self + sync_iter.__next__ = MagicMock(side_effect=[text_event, StopIteration]) + + original = litellm.use_legacy_interactions_schema + litellm.use_legacy_interactions_schema = False + try: + it = LiteLLMResponsesInteractionsStreamingIterator( + model="gpt-5.4", + litellm_custom_stream_wrapper=sync_iter, + request_input="hi", + optional_params={}, + ) + finally: + litellm.use_legacy_interactions_schema = original + + emitted: list = [] + try: + while True: + emitted.append(next(it)) + except StopIteration: + pass + + event_types = [e.event_type for e in emitted] + assert event_types == [ + "interaction.created", + "step.start", + "step.delta", + "step.stop", + "interaction.completed", + ] + terminal = emitted[-1] + assert terminal.steps == [ + { + "type": "model_output", + "content": [{"type": "text", "text": "hi"}], + } + ] + # EOF-flushed terminal event must carry the same id as interaction.created. + assert terminal.id == emitted[0].id == "item_1" + + def test_response_completed_emits_stop_then_completion(self): + """ResponseCompletedEvent expands into step.stop + interaction.completed.""" + from unittest.mock import MagicMock + + text_event = self._make_text_delta("hi") + completed = MagicMock(spec=ResponseCompletedEvent) + completed.response = MagicMock(id="resp_999") + + sync_iter = MagicMock() + sync_iter.__iter__ = lambda self: self + sync_iter.__next__ = MagicMock(side_effect=[text_event, completed]) + + original = litellm.use_legacy_interactions_schema + litellm.use_legacy_interactions_schema = False + try: + it = LiteLLMResponsesInteractionsStreamingIterator( + model="gpt-5.4", + litellm_custom_stream_wrapper=sync_iter, + request_input="hi", + optional_params={}, + ) + finally: + litellm.use_legacy_interactions_schema = original + + emitted: list = [] + try: + while True: + emitted.append(next(it)) + except StopIteration: + pass + + event_types = [e.event_type for e in emitted] + assert event_types == [ + "interaction.created", + "step.start", + "step.delta", + "step.stop", + "interaction.completed", + ] + # StopIteration fallback path must NOT add a duplicate completion event. + assert event_types.count("interaction.completed") == 1 + # When the stream starts directly with a text delta (no preceding + # response.created), the terminal events must reuse the id derived from + # the first chunk's item_id rather than switching to response.id, so + # consumers can correlate the start and completion events by id. + assert emitted[0].id == "item_1" + assert emitted[-1].id == "item_1" + + class TestInteractionOperationUrls: """Test that get/delete/cancel interaction URLs exclude API key.""" @@ -410,3 +556,152 @@ class TestInteractionOperationUrls: litellm_params=GenericLiteLLMParams(api_key=None), headers={}, ) + + +class TestTransformRequestSchemaCoalescing: + """Test new-schema request coalescing (Api-Revision: 2026-05-20).""" + + def test_response_mime_type_folded_into_response_format(self, config): + original = litellm.use_legacy_interactions_schema + try: + litellm.use_legacy_interactions_schema = False + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="summarise", + optional_params={ + "response_mime_type": "application/json", + "response_format": {"type": "object", "properties": {}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + finally: + litellm.use_legacy_interactions_schema = original + + # response_mime_type must not appear as a top-level body key + assert "response_mime_type" not in body + rf = body["response_format"] + assert rf["type"] == "text" + assert rf["mime_type"] == "application/json" + assert "schema" in rf + + def test_image_config_moved_to_response_format(self, config): + original = litellm.use_legacy_interactions_schema + try: + litellm.use_legacy_interactions_schema = False + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw a sunset", + optional_params={ + "generation_config": { + "temperature": 0.7, + "image_config": {"aspect_ratio": "1:1", "image_size": "1K"}, + } + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + finally: + litellm.use_legacy_interactions_schema = original + + # image_config removed from generation_config + assert "image_config" not in body.get("generation_config", {}) + # moved into response_format with type=image + rf = body["response_format"] + assert rf["type"] == "image" + assert rf["aspect_ratio"] == "1:1" + + def test_response_mime_type_skipped_when_response_format_is_list(self, config): + """Lists are already polymorphic; do not wrap them into schema.""" + original = litellm.use_legacy_interactions_schema + try: + litellm.use_legacy_interactions_schema = False + rf_list = [ + {"type": "text", "mime_type": "application/json"}, + {"type": "image", "aspect_ratio": "1:1"}, + ] + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="multimodal", + optional_params={ + "response_format": rf_list, + "response_mime_type": "application/json", + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + finally: + litellm.use_legacy_interactions_schema = original + + assert body["response_format"] == rf_list + assert "response_mime_type" not in body + + def test_image_config_appended_to_response_format_list_without_mutating_input( + self, config + ): + """When response_format is already a list, image_config must not mutate optional_params.""" + original = litellm.use_legacy_interactions_schema + try: + litellm.use_legacy_interactions_schema = False + text_rf = {"type": "text", "mime_type": "application/json"} + optional_params = { + "response_format": [text_rf], + "generation_config": { + "image_config": {"aspect_ratio": "16:9", "image_size": "2K"}, + }, + } + original_rf = optional_params["response_format"] + + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw and summarise", + optional_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert optional_params["response_format"] is original_rf + assert len(optional_params["response_format"]) == 1 + assert body["response_format"] == [ + text_rf, + {"type": "image", "aspect_ratio": "16:9", "image_size": "2K"}, + ] + + # Retry must not append a second image entry into the caller's list. + body_retry = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="draw and summarise", + optional_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert len(optional_params["response_format"]) == 1 + assert body_retry["response_format"] == body["response_format"] + finally: + litellm.use_legacy_interactions_schema = original + + def test_legacy_schema_passes_fields_unchanged(self, config): + original = litellm.use_legacy_interactions_schema + try: + litellm.use_legacy_interactions_schema = True + body = config.transform_request( + model="gemini/gemini-2.5-flash", + agent=None, + input="hello", + optional_params={ + "response_mime_type": "application/json", + "generation_config": {"image_config": {"aspect_ratio": "16:9"}}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + finally: + litellm.use_legacy_interactions_schema = original + + assert body["response_mime_type"] == "application/json" + assert body["generation_config"]["image_config"]["aspect_ratio"] == "16:9" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index c6961477a58..07ab29c5231 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2078,6 +2078,146 @@ async def test_async_success_handler_preserves_response_cost_for_pass_through_en assert slo["response_cost"] > 0 +def test_process_hidden_params_recalculates_cost_after_failure_handler_zero(): + """ + Regression: PR #21844 preserved response_cost=0 set by failure_handler on failed + router retry attempts, so a later successful response with usage logged $0 spend. + """ + from datetime import datetime + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test-retry-zero-cost", + function_id="test-retry-zero-cost", + ) + logging_obj.model_call_details["litellm_params"] = {"model": "openai/gpt-4o-mini"} + logging_obj.optional_params = {} + + err = litellm.RateLimitError( + message="rate limit", + llm_provider="openai", + model="openai/gpt-4o-mini", + ) + for _ in range(2): + logging_obj._failure_handler_helper_fn( + exception=err, + traceback_exception="", + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert logging_obj.model_call_details.get("response_cost") == 0 + + result = ModelResponse( + id="success", + choices=[{"message": {"role": "assistant", "content": "ok"}}], + usage=Usage(prompt_tokens=9698, completion_tokens=30, total_tokens=9728), + ) + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) + + cost = logging_obj.model_call_details.get("response_cost") + assert cost is not None and cost > 0 + slo = logging_obj.model_call_details.get("standard_logging_object") or {} + assert slo.get("response_cost", 0) > 0 + + +def test_process_hidden_params_preserves_zero_cost_in_hidden_params(): + """Pass-through handlers often set response_cost on result._hidden_params (including 0).""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="gemini-2.5-flash-lite", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-hidden-zero-cost", + function_id="test-hidden-zero-cost", + ) + logging_obj.model_call_details["litellm_params"] = { + "model": "gemini-2.5-flash-lite" + } + logging_obj.optional_params = {} + + result = ModelResponse( + id="batch-pending", + choices=[{"message": {"role": "assistant", "content": "pending"}}], + usage=Usage(prompt_tokens=100, completion_tokens=10, total_tokens=110), + ) + result._hidden_params = {"response_cost": 0.0} + + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) + + assert logging_obj.model_call_details.get("response_cost") == 0.0 + slo = logging_obj.model_call_details.get("standard_logging_object") or {} + assert slo.get("response_cost") == 0.0 + + +def test_process_hidden_params_uses_hidden_params_cost_after_failure_handler_zero(): + """After retry failures pin model_call_details to 0, success cost on _hidden_params wins.""" + from datetime import datetime + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import ModelResponse, Usage + + logging_obj = LiteLLMLoggingObj( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test-retry-hidden-cost", + function_id="test-retry-hidden-cost", + ) + logging_obj.model_call_details["litellm_params"] = {"model": "openai/gpt-4o-mini"} + logging_obj.optional_params = {} + + err = litellm.RateLimitError( + message="rate limit", + llm_provider="openai", + model="openai/gpt-4o-mini", + ) + for _ in range(2): + logging_obj._failure_handler_helper_fn( + exception=err, + traceback_exception="", + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert logging_obj.model_call_details.get("response_cost") == 0 + + passthrough_cost = 0.00042 + result = ModelResponse( + id="success", + choices=[{"message": {"role": "assistant", "content": "ok"}}], + usage=Usage(prompt_tokens=9698, completion_tokens=30, total_tokens=9728), + ) + result._hidden_params = {"response_cost": passthrough_cost} + + logging_obj._process_hidden_params_and_response_cost( + result, datetime.now(), datetime.now() + ) + + assert logging_obj.model_call_details.get("response_cost") == passthrough_cost + slo = logging_obj.model_call_details.get("standard_logging_object") or {} + assert slo.get("response_cost") == passthrough_cost + + def test_function_setup_litellm_metadata_populates_metadata(): """ Test that function_setup() properly handles litellm_metadata (used by /v1/messages, diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index a19752dc648..7d9e4768303 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -2476,6 +2476,120 @@ def test_reasoning_effort_does_not_set_output_config_for_older_models(): ), f"output_config should not be set for {model}" +@pytest.mark.parametrize( + "reasoning_effort_value", + [ + # String shape — what callers send when using `reasoning_effort="low"` directly. + "low", + # Dict shape with `effort` only — what the Responses->Chat parser produces + # when `reasoning={"effort": "low"}` is set without `summary`. + {"effort": "low"}, + # Dict shape with `effort` AND `summary` — what the Responses->Chat parser + # produces when callers send `Reasoning(effort="low", summary="concise")`. + # PR #25359 added the dict-keeping branch for this case, but the Anthropic + # transformation must coerce the dict back to a string before mapping. + {"effort": "low", "summary": "concise"}, + {"effort": "low", "summary": "detailed"}, + ], +) +def test_reasoning_effort_accepts_dict_shape_for_adaptive_model(reasoning_effort_value): + """ + Adaptive-thinking (Claude 4.6+) branch: dict-shape reasoning_effort must + map to ``thinking.type='adaptive'`` + ``output_config.effort``. + + Regression test for the dict-shape ``reasoning_effort`` produced by the + Responses->Chat parser when ``summary`` is set on the request's + ``reasoning`` field. Before this fix, the Anthropic transformation guarded + on ``isinstance(value, str)`` and silently dropped the param — disabling + extended thinking entirely. + """ + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort_value}, + optional_params={}, + model="claude-sonnet-4-6-20260219", + drop_params=False, + ) + + # thinking must be set (adaptive for 4.6+) + assert "thinking" in result, ( + f"thinking missing for reasoning_effort={reasoning_effort_value!r}" + ) + assert result["thinking"]["type"] == "adaptive" + # output_config must carry the mapped effort + assert "output_config" in result, ( + f"output_config missing for reasoning_effort={reasoning_effort_value!r}" + ) + assert result["output_config"]["effort"] == "low" + + +@pytest.mark.parametrize( + "reasoning_effort_value", + [ + "low", + {"effort": "low"}, + {"effort": "low", "summary": "concise"}, + ], +) +def test_reasoning_effort_accepts_dict_shape_for_non_adaptive_model(reasoning_effort_value): + """ + Non-adaptive (pre-4.6) branch: dict-shape reasoning_effort must still map + to ``thinking.type='enabled'`` + ``budget_tokens``. ``output_config`` must + NOT be set on these models. + """ + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"reasoning_effort": reasoning_effort_value}, + optional_params={}, + model="claude-sonnet-4-5-20250929", + drop_params=False, + ) + + assert "thinking" in result, ( + f"thinking missing for reasoning_effort={reasoning_effort_value!r}" + ) + assert result["thinking"]["type"] == "enabled" + assert "budget_tokens" in result["thinking"] + assert result["thinking"]["budget_tokens"] > 0 + # Older models must not get adaptive-thinking output_config + assert "output_config" not in result, ( + f"output_config should not be set for non-adaptive model " + f"(reasoning_effort={reasoning_effort_value!r})" + ) + + +@pytest.mark.parametrize( + "bad_value", + [ + {"summary": "concise"}, # missing effort + {"effort": None}, # explicit None effort + {"effort": 123}, # non-string effort + ], +) +def test_reasoning_effort_unparseable_dict_is_dropped(bad_value): + """ + A dict shape that doesn't carry a usable ``effort`` key (e.g. only + ``summary`` is set, or the value is some other unexpected type) should be + silently dropped — not crash, not partially apply. + """ + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"reasoning_effort": bad_value}, + optional_params={}, + model="claude-sonnet-4-6-20260219", + drop_params=False, + ) + assert "thinking" not in result, ( + f"thinking should not be set for bad value {bad_value!r}" + ) + assert "output_config" not in result, ( + f"output_config should not be set for bad value {bad_value!r}" + ) + + @pytest.mark.parametrize( "model", [ diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 11465e6f718..44530fecebd 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1203,6 +1203,51 @@ def test_streaming_chunk_with_both_text_and_tool_calls_issue_18238(): assert content_block_start["id"] == "toolu_bdrk_013xRVejhv3ybmLEGCoZib2b" +def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta(): + """ + Some OpenAI-compatible providers emit `tool_calls: []` on regular text chunks. + + Empty tool_calls should be treated as no tool call so the Anthropic adapter + does not shadow text with an empty input_json_delta. + """ + choices = [ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + provider_specific_fields=None, + content="Hello from vLLM", + role="assistant", + function_call=None, + tool_calls=[], + audio=None, + ), + logprobs=None, + ) + ] + + adapter = LiteLLMAnthropicMessagesAdapter() + + ( + type_of_content, + content_block_delta, + ) = adapter._translate_streaming_openai_chunk_to_anthropic(choices=choices) + + assert type_of_content == "text_delta" + assert content_block_delta["type"] == "text_delta" + assert content_block_delta["text"] == "Hello from vLLM" + + ( + block_type, + content_block_start, + ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block( + choices=choices + ) + + assert block_type == "text" + assert content_block_start == {"type": "text", "text": ""} + + # ============================================================================ # Cache Control Transformation Tests # ============================================================================ diff --git a/tests/test_litellm/llms/azure/chat/test_azure_base_model_routing.py b/tests/test_litellm/llms/azure/chat/test_azure_base_model_routing.py new file mode 100644 index 00000000000..1e8e23c38ca --- /dev/null +++ b/tests/test_litellm/llms/azure/chat/test_azure_base_model_routing.py @@ -0,0 +1,274 @@ +"""Tests for decoupling Azure deployment IDs from underlying model names. + +When users name their Azure deployment something non-standard (e.g. "my-deployment-id"), +setting ``base_model`` should drive model-type detection (o-series, gpt-5, +etc.) so the correct config, supported params, and param mapping are used. +""" + +import pytest + +import litellm +from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config +from litellm.llms.azure.chat.o_series_transformation import AzureOpenAIO1Config +from litellm.utils import ProviderConfigManager, get_optional_params + + +# --------------------------------------------------------------------------- +# _get_azure_config — routes to the correct config based on base_model +# --------------------------------------------------------------------------- +class TestGetAzureConfigWithBaseModel: + """ProviderConfigManager._get_azure_config should use base_model for detection.""" + + def test_should_return_gpt5_config_when_base_model_is_gpt5(self): + config = ProviderConfigManager._get_azure_config( + model="my-deployment-id", base_model="azure/gpt-5.2" + ) + assert isinstance(config, AzureOpenAIGPT5Config) + + def test_should_return_o_series_config_when_base_model_is_o_series(self): + config = ProviderConfigManager._get_azure_config( + model="my-deployment-id", base_model="azure/o4-mini" + ) + assert isinstance(config, AzureOpenAIO1Config) + + def test_should_return_default_config_when_base_model_is_regular(self): + config = ProviderConfigManager._get_azure_config( + model="my-deployment-id", base_model="azure/gpt-4o" + ) + assert type(config).__name__ == "AzureOpenAIConfig" + + def test_should_fallback_to_model_when_base_model_is_none(self): + config = ProviderConfigManager._get_azure_config( + model="gpt-5.2", base_model=None + ) + assert isinstance(config, AzureOpenAIGPT5Config) + + def test_should_return_default_config_when_both_are_non_standard(self): + config = ProviderConfigManager._get_azure_config( + model="my-deployment-id", base_model=None + ) + assert type(config).__name__ == "AzureOpenAIConfig" + + +# --------------------------------------------------------------------------- +# get_provider_chat_config — threads base_model through for Azure +# --------------------------------------------------------------------------- +class TestGetProviderChatConfigWithBaseModel: + """get_provider_chat_config should pass base_model to Azure config selection.""" + + def test_should_return_gpt5_config_for_custom_deployment_with_base_model(self): + from litellm.types.utils import LlmProviders + + config = ProviderConfigManager.get_provider_chat_config( + model="my-deployment-id", + provider=LlmProviders.AZURE, + base_model="azure/gpt-5", + ) + assert isinstance(config, AzureOpenAIGPT5Config) + + def test_should_return_o_series_config_for_custom_deployment_with_base_model(self): + from litellm.types.utils import LlmProviders + + config = ProviderConfigManager.get_provider_chat_config( + model="my-other-deployment", + provider=LlmProviders.AZURE, + base_model="azure/o3-mini", + ) + assert isinstance(config, AzureOpenAIO1Config) + + +# --------------------------------------------------------------------------- +# get_supported_openai_params — base_model drives Azure param detection +# --------------------------------------------------------------------------- +class TestGetSupportedOpenAIParamsWithBaseModel: + """get_supported_openai_params should use base_model for Azure detection.""" + + def test_should_return_gpt5_params_for_custom_deployment_with_gpt5_base_model( + self, + ): + params = litellm.get_supported_openai_params( + model="my-deployment-id", + custom_llm_provider="azure", + base_model="azure/gpt-5", + ) + assert params is not None + assert "reasoning_effort" in params + # gpt-5 maps max_tokens -> max_completion_tokens, verifying we got GPT-5 config + assert "max_completion_tokens" in params + + def test_should_return_o_series_params_for_custom_deployment_with_o_series_base_model( + self, + ): + params = litellm.get_supported_openai_params( + model="my-other-deployment", + custom_llm_provider="azure", + base_model="azure/o4-mini", + ) + assert params is not None + assert "reasoning_effort" in params + + def test_should_return_regular_params_when_no_base_model(self): + """When base_model is not set and model is non-standard, default Azure config.""" + params = litellm.get_supported_openai_params( + model="my-deployment-id", + custom_llm_provider="azure", + ) + assert params is not None + # Default Azure config supports temperature + assert "temperature" in params + + +# --------------------------------------------------------------------------- +# get_optional_params — base_model drives Azure param mapping +# --------------------------------------------------------------------------- +class TestGetOptionalParamsWithBaseModel: + """get_optional_params should use base_model for Azure model-type detection.""" + + def test_should_map_max_tokens_for_custom_deployment_with_gpt5_base_model(self): + """A non-standard deployment name + gpt-5 base_model should map max_tokens -> max_completion_tokens.""" + params = get_optional_params( + model="my-deployment-id", + custom_llm_provider="azure", + max_tokens=100, + base_model="azure/gpt-5", + ) + assert params.get("max_completion_tokens") == 100 + assert "max_tokens" not in params + + def test_should_keep_max_tokens_for_custom_deployment_without_base_model(self): + """A non-standard deployment name without base_model should use default Azure config.""" + params = get_optional_params( + model="my-deployment-id", + custom_llm_provider="azure", + max_tokens=100, + api_version="2024-05-01-preview", + ) + # Default AzureOpenAIConfig keeps max_tokens as-is (or maps based on api_version) + assert "max_tokens" in params or "max_completion_tokens" in params + + def test_should_support_reasoning_effort_for_custom_deployment_with_o_series_base_model( + self, + ): + """A non-standard deployment name + o-series base_model should accept reasoning_effort.""" + params = get_optional_params( + model="my-other-deployment", + custom_llm_provider="azure", + reasoning_effort="low", + base_model="azure/o4-mini", + ) + assert params.get("reasoning_effort") == "low" + + def test_should_reject_temperature_for_custom_deployment_with_gpt5_base_model( + self, + ): + """A non-standard deployment + gpt-5 base_model should reject temperature.""" + with pytest.raises(litellm.UnsupportedParamsError): + get_optional_params( + model="my-deployment-id", + custom_llm_provider="azure", + temperature=0.5, + base_model="azure/gpt-5", + ) + + +# --------------------------------------------------------------------------- +# Backward compatibility — existing patterns still work +# --------------------------------------------------------------------------- +class TestBackwardCompatibility: + """Existing model-name-based and prefix-based patterns must keep working.""" + + def test_should_detect_gpt5_from_model_name(self): + config = ProviderConfigManager._get_azure_config(model="gpt-5.2") + assert isinstance(config, AzureOpenAIGPT5Config) + + def test_should_detect_gpt5_from_gpt5_series_prefix(self): + config = ProviderConfigManager._get_azure_config( + model="gpt5_series/my-deployment" + ) + assert isinstance(config, AzureOpenAIGPT5Config) + + def test_should_detect_o_series_from_model_name(self): + config = ProviderConfigManager._get_azure_config(model="o4-mini") + assert isinstance(config, AzureOpenAIO1Config) + + def test_should_detect_o_series_from_o_series_prefix(self): + config = ProviderConfigManager._get_azure_config(model="o_series/my-deployment") + assert isinstance(config, AzureOpenAIO1Config) + + def test_should_handle_gpt5_chat_model_correctly(self): + """gpt-5-chat models should NOT be routed to GPT-5 config.""" + config = ProviderConfigManager._get_azure_config(model="gpt-5-chat") + assert type(config).__name__ == "AzureOpenAIConfig" + + def test_base_model_overrides_model_detection(self): + """base_model should take priority over model for type detection.""" + # model looks like o-series, but base_model says gpt-5 + config = ProviderConfigManager._get_azure_config( + model="o3-mini", base_model="azure/gpt-5.2" + ) + assert isinstance(config, AzureOpenAIGPT5Config) + + +# --------------------------------------------------------------------------- +# Deep config method awareness — base_model flows into config internals +# --------------------------------------------------------------------------- +class TestBaseModelFlowsIntoConfigInternals: + """base_model should be used by config internal methods (e.g. is_model_gpt_5_2_model).""" + + def test_should_support_logprobs_for_prefixed_deployment_with_gpt52_base_model( + self, + ): + """Deployment 'my-gpt-5.2' with base_model='azure/gpt-5.2' should support logprobs.""" + params = litellm.get_supported_openai_params( + model="gpt5_series/my-gpt-5.2", + custom_llm_provider="azure", + base_model="azure/gpt-5.2", + ) + assert params is not None + assert "logprobs" in params + assert "top_logprobs" in params + + def test_should_support_logprobs_for_plain_deployment_with_gpt52_base_model(self): + """Deployment 'my-deployment-id' with base_model='azure/gpt-5.2' should support logprobs.""" + params = litellm.get_supported_openai_params( + model="my-deployment-id", + custom_llm_provider="azure", + base_model="azure/gpt-5.2", + ) + assert params is not None + assert "logprobs" in params + assert "top_logprobs" in params + + def test_should_not_support_logprobs_for_gpt5_base_model(self): + """Deployment with base_model='azure/gpt-5' (not 5.2) should NOT support logprobs.""" + params = litellm.get_supported_openai_params( + model="my-deployment-id", + custom_llm_provider="azure", + base_model="azure/gpt-5", + ) + assert params is not None + assert "logprobs" not in params + assert "top_logprobs" not in params + + def test_should_pass_logprobs_through_get_optional_params(self): + """logprobs should pass validation in get_optional_params when base_model is gpt-5.2.""" + params = get_optional_params( + model="gpt5_series/my-gpt-5.2", + custom_llm_provider="azure", + logprobs=True, + top_logprobs=5, + base_model="azure/gpt-5.2", + ) + assert params.get("logprobs") is True + assert params.get("top_logprobs") == 5 + + def test_should_map_max_tokens_for_prefixed_deployment_with_gpt5_base_model(self): + """my-gpt-5.2 with base_model should correctly map max_tokens -> max_completion_tokens.""" + params = get_optional_params( + model="gpt5_series/my-gpt-5.2", + custom_llm_provider="azure", + max_tokens=200, + base_model="azure/gpt-5.2", + ) + assert params.get("max_completion_tokens") == 200 + assert "max_tokens" not in params diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index 4495e3f4101..b2e254901f4 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -2,6 +2,7 @@ import asyncio import json import os import sys +from unittest.mock import patch import pytest @@ -429,6 +430,31 @@ def test_output_config_forwarded_for_bedrock_chat_invoke_request(): assert result["max_tokens"] == 100 +def test_bedrock_chat_invoke_checks_output_config_support_with_bedrock_provider(): + config = AmazonAnthropicClaudeConfig() + messages = [{"role": "user", "content": "test"}] + optional_params = {"max_tokens": 100, "output_config": {"effort": "high"}} + + with patch( + "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ) as mock_supports_factory: + result = config.transform_request( + model="us.anthropic.claude-opus-4-7", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + mock_supports_factory.assert_called_once_with( + model="us.anthropic.claude-opus-4-7", + custom_llm_provider="bedrock", + key="supports_output_config", + ) + assert result["output_config"] == {"effort": "high"} + + def test_output_format_removed_from_bedrock_invoke_request(): """ Test that output_format parameter is removed from Bedrock Invoke requests. diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 9ecdad1fcff..2e315a535f0 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -592,8 +592,15 @@ def test_remove_scope_from_cache_control(): assert request["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral" -def test_bedrock_messages_forwards_output_config(): - """Bedrock Invoke /v1/messages forwards ``output_config`` for adaptive Claude models.""" +def test_bedrock_messages_strips_output_config(): + """ + Ensure output_config is stripped from the request for models that do not + support it. + + Regression test for: https://github.com/BerriAI/litellm/issues/22797 + """ + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -605,21 +612,129 @@ def test_bedrock_messages_forwards_output_config(): }, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=False, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + "output_config" not in result + ), "output_config should be stripped for models that don't support it" + assert result.get("max_tokens") == 4096 + + +def test_bedrock_messages_preserves_output_config_for_claude_4_6(): + """ + Ensure output_config is preserved for models that support it on Bedrock Invoke. + """ + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "high", + }, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-6-v1", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + "output_config" in result + ), "output_config should be preserved for supported models" + assert result["output_config"] == {"effort": "high"} + assert result.get("max_tokens") == 4096 + + +def test_bedrock_messages_checks_output_config_support_with_bedrock_provider(): + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "high", + }, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ) as mock_supports_factory: + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + mock_supports_factory.assert_called_with( + model="us.anthropic.claude-opus-4-7", + custom_llm_provider="bedrock", + key="supports_output_config", ) + assert result["output_config"] == {"effort": "high"} + + +def test_bedrock_messages_forwards_output_config(): + """Bedrock Invoke /v1/messages forwards ``output_config`` for supported models.""" + from unittest.mock import patch + + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "high", + }, + } + + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert result.get("output_config") == {"effort": "high"} - # Other params should be preserved assert result.get("max_tokens") == 4096 def test_bedrock_messages_forwards_output_config_with_output_format(): """``output_config`` is forwarded; ``output_format`` is converted to inline schema.""" + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -636,39 +751,60 @@ def test_bedrock_messages_forwards_output_config_with_output_format(): }, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert result.get("output_config") == {"effort": "low"} assert "output_format" not in result -def test_bedrock_messages_forwards_output_config_for_non_adaptive_model(): - """``output_config`` is forwarded for non-adaptive models so the provider's error surfaces.""" +def test_bedrock_messages_strips_output_config_with_output_format(): + """ + When both output_config and output_format are present, output_format + is converted to inline schema and output_config is stripped for + unsupported models. + """ + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] optional_params = { "max_tokens": 4096, - "output_config": {"effort": "high"}, + "output_config": {"effort": "low"}, + "output_format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + }, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-3-haiku-20240307-v1:0", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=False, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) - assert result.get("output_config") == {"effort": "high"} - assert result.get("max_tokens") == 4096 + assert "output_config" not in result + assert "output_format" not in result def test_bedrock_messages_drop_params_strips_output_config_for_pre_4_5(): @@ -701,6 +837,8 @@ def test_bedrock_messages_drop_params_strips_output_config_for_pre_4_5(): def test_bedrock_messages_drop_params_keeps_output_config_for_4_7(): """``drop_params=True`` does not strip on opus-4-7 (supports effort).""" + from unittest.mock import patch + import litellm from litellm.types.router import GenericLiteLLMParams @@ -714,13 +852,17 @@ def test_bedrock_messages_drop_params_keeps_output_config_for_4_7(): original = litellm.drop_params litellm.drop_params = True try: - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) finally: litellm.drop_params = original @@ -742,6 +884,8 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( reasoning_effort, expected_effort ): """``reasoning_effort`` maps to ``thinking`` + ``output_config.effort`` on /v1/messages.""" + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -751,13 +895,17 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( "reasoning_effort": reasoning_effort, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert "reasoning_effort" not in result assert result.get("thinking") == {"type": "adaptive"} @@ -842,6 +990,8 @@ def test_bedrock_messages_invalid_reasoning_effort_raises_400(): def test_bedrock_messages_explicit_output_config_wins_over_reasoning_effort(): """Explicit ``output_config.effort`` wins over the ``reasoning_effort`` alias.""" + from unittest.mock import patch + from litellm.types.router import GenericLiteLLMParams cfg = AmazonAnthropicClaudeMessagesConfig() @@ -852,13 +1002,17 @@ def test_bedrock_messages_explicit_output_config_wins_over_reasoning_effort(): "output_config": {"effort": "max"}, } - result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-opus-4-7", - messages=messages, - anthropic_messages_optional_request_params=optional_params, - litellm_params=GenericLiteLLMParams(), - headers={}, - ) + with patch( + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation._supports_factory", + return_value=True, + ): + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-opus-4-7", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) assert "reasoning_effort" not in result assert result.get("output_config") == {"effort": "max"} @@ -994,7 +1148,7 @@ def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): } result = cfg.transform_anthropic_messages_request( - model="anthropic.claude-3-haiku-20240307-v1:0", + model="anthropic.claude-opus-4-7", messages=messages, anthropic_messages_optional_request_params=optional_params, litellm_params=GenericLiteLLMParams(), diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index a4969e5dacc..10fc358e3a5 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -2112,3 +2112,102 @@ def test_is_already_running_as_role_ssl_verify_passed(): mock_boto3_client.assert_called_once_with( "sts", verify="/path/to/ca-bundle.crt" ) + + +# --------------------------------------------------------------------------- +# LIT-3274: get_bedrock_model_id must strip "bedrock/" prefix and URL-encode +# ARNs for the invoke path (invoke-with-response-stream). Without this fix +# the Bedrock API receives a malformed URL, returns a JSON error body, and +# botocore's EventStreamBuffer raises ChecksumMismatch instead of the real +# error. 0x223a7b22 == ':{\"' — the start of a JSON object. +# --------------------------------------------------------------------------- + + +class TestGetBedrockModelIdArnHandling: + """Unit tests for get_bedrock_model_id with inference-profile ARNs.""" + + ARN = "arn:aws:bedrock:us-east-1:086734376398:inference-profile/global.anthropic.claude-sonnet-4-5-20250929-v1:0" + + def _call(self, model: str, optional_params: dict | None = None) -> str: + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + return BaseAWSLLM.get_bedrock_model_id( + model=model, + provider=provider, + optional_params=optional_params or {}, + ) + + def test_arn_with_bedrock_prefix_is_stripped_and_encoded(self): + """bedrock/arn:... must not appear verbatim in the model_id.""" + model_id = self._call(f"bedrock/{self.ARN}") + assert ( + "bedrock/arn" not in model_id + ), f"'bedrock/' prefix not stripped; got: {model_id}" + # Must be URL-encoded (colons → %3A) + assert "%3A" in model_id, f"ARN not URL-encoded; got: {model_id}" + assert "%2F" in model_id, f"ARN slashes not URL-encoded; got: {model_id}" + + def test_arn_with_compound_bedrock_invoke_prefix_is_fully_stripped_and_encoded( + self, + ): + """bedrock/invoke/arn:... — compound prefix — must be fully stripped. + + The old fix used ``break`` after the first matched prefix, so + ``bedrock/invoke/arn:...`` would only strip ``bedrock/``, leaving + ``invoke/arn:...``. The subsequent ``.replace('invoke/', '')`` call + then returned the bare unencoded ARN, reproducing the same + malformed-URL bug the fix aimed to prevent. + + strip_bedrock_routing_prefix() has no break and handles this correctly. + """ + model_id = self._call(f"bedrock/invoke/{self.ARN}") + assert ( + "invoke/" not in model_id + ), f"'invoke/' prefix not stripped; got: {model_id}" + assert ( + "bedrock/" not in model_id + ), f"'bedrock/' prefix not stripped; got: {model_id}" + assert "%3A" in model_id, f"ARN not URL-encoded; got: {model_id}" + assert "%2F" in model_id, f"ARN slashes not URL-encoded; got: {model_id}" + + def test_bare_arn_is_encoded(self): + """Direct ARN without routing prefix must also be URL-encoded.""" + model_id = self._call(self.ARN) + assert "%3A" in model_id, f"ARN not URL-encoded; got: {model_id}" + assert "%2F" in model_id, f"ARN slashes not URL-encoded; got: {model_id}" + + def test_arn_url_matches_expected(self): + """Full URL built from messages config must match expected encoded form.""" + import urllib.parse + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + config = AmazonAnthropicClaudeMessagesConfig() + url = config.get_complete_url( + api_base=None, + api_key=None, + model=f"bedrock/{self.ARN}", + optional_params={"aws_region_name": "us-east-1"}, + litellm_params={}, + stream=True, + ) + encoded_arn = urllib.parse.quote(self.ARN, safe="") + expected = ( + f"https://bedrock-runtime.us-east-1.amazonaws.com" + f"/model/{encoded_arn}/invoke-with-response-stream" + ) + assert ( + url == expected + ), f"URL mismatch:\n got: {url}\n expected: {expected}" + + def test_regular_model_id_unaffected(self): + """Non-ARN model IDs must continue to work as before.""" + model_id = self._call("anthropic.claude-3-sonnet-20240229-v1:0") + assert model_id == "anthropic.claude-3-sonnet-20240229-v1:0" + + def test_invoke_prefixed_model_unaffected(self): + """invoke/ prefix stripping still works after the fix.""" + model_id = self._call("invoke/anthropic.claude-3-sonnet-20240229-v1:0") + assert model_id == "anthropic.claude-3-sonnet-20240229-v1:0" diff --git a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index 2498946bb5c..90a1c24bada 100644 --- a/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/test_litellm/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -14,6 +14,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) +from litellm.llms.openai.common_utils import OpenAIError from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -201,3 +202,127 @@ class TestChatGPTResponsesAPITransformation: ) assert parsed.output_text == "Hello!" + + @pytest.mark.parametrize( + ("model_name", "response_model"), + [ + ("chatgpt/gpt-5.2-codex", "gpt-5.2-codex"), + ("chatgpt/gpt-5.3-codex", "gpt-5.3-codex"), + ], + ) + def test_chatgpt_non_stream_sse_response_recovers_output_items( + self, model_name: str, response_model: str + ): + config = ChatGPTResponsesAPIConfig() + response_payload = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": response_model, + "output": [], + } + streamed_output_item = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello from stream!"}], + } + sse_body = "\n".join( + [ + f"data: {json.dumps({'type': 'response.output_item.done', 'output_index': 0, 'item': streamed_output_item})}", + f"data: {json.dumps({'type': 'response.completed', 'response': response_payload})}", + "data: [DONE]", + "", + ] + ) + raw_response = httpx.Response( + 200, headers={"content-type": "text/event-stream"}, text=sse_body + ) + logging_obj = MagicMock() + + parsed = config.transform_response_api_response( + model=model_name, + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert parsed.output_text == "Hello from stream!" + + def test_chatgpt_non_stream_sse_recovers_whitespace_padded_chunks(self): + """Chunks with leading whitespace before `data:` must still parse. + + `_strip_sse_data_from_chunk` only matches the prefix at position 0, + so without an outer `.strip()` such chunks would fail JSON parsing + and silently drop the contained event. + """ + config = ChatGPTResponsesAPIConfig() + response_payload = { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.4", + "output": [], + } + streamed_output_item = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Recovered from padded"}], + } + sse_body = "\n".join( + [ + f" data: {json.dumps({'type': 'response.output_item.done', 'output_index': 0, 'item': streamed_output_item})} ", + f"\tdata: {json.dumps({'type': 'response.completed', 'response': response_payload})}", + "data: [DONE]", + "", + ] + ) + raw_response = httpx.Response( + 200, headers={"content-type": "text/event-stream"}, text=sse_body + ) + logging_obj = MagicMock() + + parsed = config.transform_response_api_response( + model="chatgpt/gpt-5.4", + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert parsed.output_text == "Recovered from padded" + + @pytest.mark.parametrize( + "error_chunk", + [ + { + "type": "response.failed", + "response": {"error": {"message": "ChatGPT upstream failed"}}, + }, + { + "type": "error", + "error": {"message": "ChatGPT upstream failed"}, + }, + ], + ) + def test_chatgpt_non_stream_sse_response_raises_openai_error(self, error_chunk): + config = ChatGPTResponsesAPIConfig() + sse_body = "\n".join( + [ + f"data: {json.dumps(error_chunk)}", + "data: [DONE]", + "", + ] + ) + raw_response = httpx.Response( + 502, headers={"content-type": "text/event-stream"}, text=sse_body + ) + logging_obj = MagicMock() + + with pytest.raises(OpenAIError) as exc_info: + config.transform_response_api_response( + model="chatgpt/gpt-5.4", + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert "ChatGPT upstream failed" in str(exc_info.value) + assert exc_info.value.status_code == 502 diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 279f16a3675..a29365544df 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -6,16 +6,29 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import litellm + sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path -from litellm import supports_reasoning +from litellm import get_model_info, supports_reasoning from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.types.llms.openai import ChatCompletionToolCallFunctionChunk from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message +@pytest.fixture(autouse=True) +def force_local_model_cost(monkeypatch): + """Force local model cost map usage for all tests in this file.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + # Refresh model_cost from local map + import litellm + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + + litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url) + + def test_handle_message_content_with_tool_calls(): config = FireworksAIConfig() message = Message( @@ -62,7 +75,6 @@ def test_handle_message_content_with_tool_calls(): def test_supports_reasoning_effort(): """Test that reasoning_effort is only supported for specific Fireworks AI models.""" - # Models that support reasoning_effort supported_models = [ "fireworks_ai/accounts/fireworks/models/qwen3-8b", "fireworks_ai/accounts/fireworks/models/qwen3-32b", @@ -72,11 +84,13 @@ def test_supports_reasoning_effort(): "fireworks_ai/accounts/fireworks/models/glm-4p5", "fireworks_ai/accounts/fireworks/models/glm-4p5-air", "fireworks_ai/accounts/fireworks/models/glm-4p6", + "fireworks_ai/accounts/fireworks/models/glm-4p7", + "fireworks_ai/accounts/fireworks/models/glm-5p1", "fireworks_ai/accounts/fireworks/models/gpt-oss-120b", "fireworks_ai/accounts/fireworks/models/gpt-oss-20b", + "fireworks_ai/glm-5p1", ] - # Models that don't support reasoning_effort unsupported_models = [ "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct", "fireworks_ai/accounts/fireworks/models/mixtral-8x7b-instruct", @@ -97,19 +111,74 @@ def test_get_supported_openai_params_reasoning_effort(): """Test that reasoning_effort is only included in supported params for models that support it.""" config = FireworksAIConfig() - # Model that supports reasoning_effort supported_params = config.get_supported_openai_params( - "fireworks_ai/accounts/fireworks/models/qwen3-8b" + "fireworks_ai/accounts/fireworks/models/glm-5p1" ) assert "reasoning_effort" in supported_params - # Model that doesn't support reasoning_effort unsupported_params = config.get_supported_openai_params( "fireworks_ai/accounts/fireworks/models/llama-v3-70b-instruct" ) assert "reasoning_effort" not in unsupported_params +def test_get_supported_openai_params_parallel_tool_calls(): + """Test that parallel_tool_calls is included for models that support function calling.""" + config = FireworksAIConfig() + + supported_params = config.get_supported_openai_params( + "fireworks_ai/accounts/fireworks/models/glm-4p6" + ) + assert "parallel_tool_calls" in supported_params + + unsupported_params = config.get_supported_openai_params( + "fireworks_ai/accounts/fireworks/models/glm-5p1" + ) + assert "parallel_tool_calls" not in unsupported_params + + +def test_get_supported_openai_params_parallel_tool_calls_without_tool_choice( + monkeypatch, +): + """Test that parallel_tool_calls is gated on tools, not tool_choice.""" + config = FireworksAIConfig() + model = "fireworks_ai/test-tools-without-tool-choice" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "supports_function_calling": True, + "supports_tool_choice": False, + }, + ) + + supported_params = config.get_supported_openai_params(model) + + assert "tools" in supported_params + assert "parallel_tool_calls" in supported_params + assert "tool_choice" not in supported_params + + +def test_get_model_info_respects_explicit_fireworks_capabilities(): + """Test that get_model_info preserves explicit capability flags from the model map.""" + model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1") + + assert model_info["supports_function_calling"] is False + assert model_info["supports_reasoning"] is True + assert model_info["supports_tool_choice"] is False + + +def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): + """Test that Fireworks only overrides supports_reasoning for supported models.""" + config = FireworksAIConfig() + model = "fireworks_ai/test-reasoning-false" + monkeypatch.setitem(litellm.model_cost, model, {"supports_reasoning": False}) + + info = config.get_provider_info(model) + + assert "supports_reasoning" not in info + + def test_add_transform_inline_image_block_skips_data_urls(): """ data: URLs must not have #transform=inline appended — doing so corrupts the @@ -234,6 +303,14 @@ def test_transform_messages_helper_removes_provider_specific_fields(): assert "provider_specific_fields" not in msg +def test_unmapped_model_fallback_function_calling(): + """Test that a model not in model_cost still defaults to supporting function calling for Fireworks.""" + config = FireworksAIConfig() + model = "fireworks_ai/unmapped-future-model" + info = config.get_provider_info(model) + assert info["supports_function_calling"] is True + + def test_transform_messages_helper_strips_thinking_blocks(): """thinking_blocks must not be forwarded to Fireworks chat completions.""" config = FireworksAIConfig() diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index acb9fa9b64c..4b2e9471fb7 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -86,6 +86,90 @@ class TestOpenAIResponsesAPIConfig: self.validate_responses_api_request_params(result, expected_fields) + def test_transform_strips_cache_control_from_input_content_blocks(self): + """`cache_control` markers (Anthropic-only) must be stripped from + Responses API input content blocks before sending to OpenAI. + + OpenAI rejects unknown params on input content blocks with HTTP 400: + "Unknown parameter: 'input[0].content[0].cache_control'" + Chat Completions strips these via + `remove_cache_control_flag_from_messages_and_tools`; the Responses + path must do the same. + """ + input_with_cache_control = [ + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": "Hello", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + + result = self.config.transform_responses_api_request( + model=self.model, + input=input_with_cache_control, + response_api_optional_request_params={}, + litellm_params={}, + headers={}, + ) + + assert "cache_control" not in result["input"][0]["content"][0] + assert result["input"][0]["content"][0]["type"] == "input_text" + assert result["input"][0]["content"][0]["text"] == "Hello" + + def test_transform_strips_cache_control_from_tools(self): + """`cache_control` markers must also be stripped from tools for + symmetry with the Chat Completions path. OpenAI currently accepts + cache_control on tools silently but stripping keeps the wire payload + clean and matches `remove_cache_control_flag_from_messages_and_tools`. + """ + tools_with_cache_control = [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + "cache_control": {"type": "ephemeral"}, + } + ] + + result = self.config.transform_responses_api_request( + model=self.model, + input="hi", + response_api_optional_request_params={"tools": tools_with_cache_control}, + litellm_params={}, + headers={}, + ) + + assert "cache_control" not in result["tools"][0] + assert result["tools"][0]["name"] == "get_weather" + + def test_transform_preserves_input_without_cache_control(self): + """Inputs without cache_control must pass through unmodified.""" + input_clean = [ + { + "role": "user", + "content": [{"type": "input_text", "text": "Hello"}], + } + ] + + result = self.config.transform_responses_api_request( + model=self.model, + input=input_clean, + response_api_optional_request_params={}, + litellm_params={}, + headers={}, + ) + + assert result["input"] == input_clean + def test_transform_streaming_response(self): """Test streaming response transformation""" # Test with a text delta event diff --git a/tests/test_litellm/llms/reducto/__init__.py b/tests/test_litellm/llms/reducto/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm/llms/reducto/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm/llms/reducto/test_cost.py b/tests/test_litellm/llms/reducto/test_cost.py new file mode 100644 index 00000000000..73340dc8729 --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_cost.py @@ -0,0 +1,122 @@ +import litellm +import pytest + +from litellm.cost_calculator import completion_cost +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo + + +def test_ocr_cost_prefers_credit_pricing_when_pages_processed_is_none(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_credit": 0.003}, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="credit priced")], + model="parse-v3", + usage_info=OCRUsageInfo(pages_processed=None, credits=10), + ) + + cost = completion_cost( + completion_response=response, + model="reducto/parse-v3", + custom_llm_provider="reducto", + call_type="ocr", + ) + + assert cost == 0.03 + + +def test_ocr_cost_prefers_zero_credit_pricing_over_page_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: { + "ocr_cost_per_credit": 0.0, + "ocr_cost_per_page": 0.5, + }, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="free credit priced")], + model="parse-v3", + usage_info=OCRUsageInfo(pages_processed=2, credits=10), + ) + + cost = completion_cost( + completion_response=response, + model="reducto/parse-v3", + custom_llm_provider="reducto", + call_type="ocr", + ) + + assert cost == 0.0 + + +def test_ocr_cost_falls_back_to_page_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5}, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="page priced")], + model="mistral-ocr-latest", + usage_info=OCRUsageInfo(pages_processed=2), + ) + + cost = completion_cost( + completion_response=response, + model="mistral/mistral-ocr-latest", + custom_llm_provider="mistral", + call_type="ocr", + ) + + assert cost == 1.0 + + +def test_ocr_cost_returns_zero_when_no_pricing_and_no_pages(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {}, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="unpriced")], + model="parse-v3", + usage_info=OCRUsageInfo(pages_processed=None, credits=5), + ) + + cost = completion_cost( + completion_response=response, + model="reducto/parse-v3", + custom_llm_provider="reducto", + call_type="ocr", + ) + + assert cost == 0.0 + + +def test_ocr_cost_raises_when_pages_processed_missing_for_page_pricing(monkeypatch): + monkeypatch.setattr( + litellm, + "get_model_info", + lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5}, + ) + + response = OCRResponse( + pages=[OCRPage(index=0, markdown="missing pages")], + model="mistral-ocr-latest", + usage_info=OCRUsageInfo(pages_processed=None), + ) + + with pytest.raises(ValueError, match="OCR response pages_processed is None"): + completion_cost( + completion_response=response, + model="mistral/mistral-ocr-latest", + custom_llm_provider="mistral", + call_type="ocr", + ) diff --git a/tests/test_litellm/llms/reducto/test_model_info.py b/tests/test_litellm/llms/reducto/test_model_info.py new file mode 100644 index 00000000000..de7a3ccba64 --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_model_info.py @@ -0,0 +1,44 @@ +import uuid + +import litellm + +from litellm.utils import _invalidate_model_cost_lowercase_map + + +def test_reducto_provider_registration(): + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model="reducto/parse-v3" + ) + + assert model == "parse-v3" + assert custom_llm_provider == "reducto" + + +def test_get_model_info_preserves_ocr_cost_per_credit(): + test_model_name = f"reducto/test-cost-propagation-{uuid.uuid4().hex[:12]}" + previous_model_entry = litellm.model_cost.get(test_model_name) + _invalidate_model_cost_lowercase_map() + + try: + litellm.register_model( + { + test_model_name: { + "litellm_provider": "reducto", + "mode": "ocr", + "ocr_cost_per_credit": 0.003, + } + } + ) + + model_info = litellm.get_model_info( + model=test_model_name, + custom_llm_provider="reducto", + ) + + assert model_info.get("ocr_cost_per_credit") == 0.003 + finally: + if previous_model_entry is None: + litellm.model_cost.pop(test_model_name, None) + else: + litellm.model_cost[test_model_name] = previous_model_entry + _invalidate_model_cost_lowercase_map() diff --git a/tests/test_litellm/llms/reducto/test_parse_legacy.py b/tests/test_litellm/llms/reducto/test_parse_legacy.py new file mode 100644 index 00000000000..db19460baa3 --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_parse_legacy.py @@ -0,0 +1,59 @@ +import json + +import litellm +import pytest + + +@pytest.fixture() +def disable_aiohttp_transport(): + original_disable_aiohttp = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = True + litellm.in_memory_llm_clients_cache.flush_cache() + try: + yield + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.asyncio +async def test_parse_legacy_wraps_enhance_under_options( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( + json={"file_id": "reducto://legacy.pdf"} + ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Legacy parse", + "blocks": [{"content": "Legacy parse", "bbox": {"page": 1}}], + } + ] + }, + } + ) + + response = await litellm.aocr( + model="reducto/parse-legacy", + document={ + "type": "file", + "file": b"%PDF-1.4 legacy", + "mime_type": "application/pdf", + }, + api_key="legacy-key", + api_base="https://platform.reducto.ai", + enhance={"agentic": [{"type": "table"}]}, + ) + + assert upload_route.called + assert parse_route.called + request_body = json.loads(parse_route.calls[0].request.read()) + assert request_body == { + "document_url": "reducto://legacy.pdf", + "options": {"enhance": {"agentic": [{"type": "table"}]}}, + } + assert response.pages[0].markdown == "Legacy parse" diff --git a/tests/test_litellm/llms/reducto/test_parse_v3.py b/tests/test_litellm/llms/reducto/test_parse_v3.py new file mode 100644 index 00000000000..140b9737dc0 --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_parse_v3.py @@ -0,0 +1,152 @@ +import json + +import litellm +import pytest + + +def _reducto_parse_response() -> dict: + return { + "job_id": "job_123", + "usage": {"num_pages": 3, "credits": 3}, + "result": { + "chunks": [ + { + "content": "Page 1 block A", + "blocks": [ + { + "content": "Page 1 block A", + "bbox": {"page": 1}, + "kind": "text", + } + ], + }, + { + "content": "Page 2 block A", + "blocks": [ + { + "content": "Page 2 block A", + "bbox": {"page": 2}, + "kind": "table", + } + ], + }, + { + "content": "Page 1 block B", + "blocks": [ + { + "content": "Page 1 block B", + "bbox": {"page": 1}, + "kind": "text", + } + ], + }, + { + "content": "Page 3 block A", + "blocks": [ + { + "content": "Page 3 block A", + "bbox": {"page": 3}, + "kind": "figure", + } + ], + }, + ] + }, + } + + +@pytest.fixture() +def disable_aiohttp_transport(): + original_disable_aiohttp = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = True + litellm.in_memory_llm_clients_cache.flush_cache() + try: + yield + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.asyncio +async def test_parse_v3_file_upload_and_response_mapping( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( + json={"file_id": "reducto://uploaded.pdf"} + ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json=_reducto_parse_response() + ) + + response = await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "file", + "file": b"%PDF-1.4 reducto", + "mime_type": "application/pdf", + }, + api_key="test-key", + api_base="https://platform.reducto.ai", + formatting={"table_output_format": "html"}, + retrieval={"chunk_mode": "section"}, + settings={"ocr_system": "standard"}, + ) + + assert upload_route.called + assert parse_route.called + assert len(upload_route.calls) == 1 + assert len(parse_route.calls) == 1 + + upload_request = upload_route.calls[0].request + assert upload_request.headers["authorization"] == "Bearer test-key" + assert "application/json" not in upload_request.headers["content-type"] + upload_body = upload_request.read() + assert b'filename="document"' in upload_body + assert b"application/pdf" in upload_body + + parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert parse_request_body["input"] == "reducto://uploaded.pdf" + assert parse_request_body["formatting"] == {"table_output_format": "html"} + assert parse_request_body["retrieval"] == {"chunk_mode": "section"} + assert parse_request_body["settings"] == {"ocr_system": "standard"} + + assert response.usage_info is not None + assert response.usage_info.credits == 3 + assert response.usage_info.pages_processed == 3 + assert len(response.pages) == 3 + assert response.pages[0].index == 0 + assert response.pages[0].markdown == "Page 1 block A\n\nPage 1 block B" + assert getattr(response.pages[0], "blocks")[0]["bbox"]["page"] == 1 + assert response.pages[1].markdown == "Page 2 block A" + assert response.pages[2].markdown == "Page 3 block A" + assert response._hidden_params["reducto_raw"]["usage"]["credits"] == 3 + + +@pytest.mark.asyncio +async def test_parse_v3_reducto_id_passthrough_skips_upload( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( + json={"file_id": "reducto://should-not-upload.pdf"} + ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json=_reducto_parse_response() + ) + + response = await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "document_url", + "document_url": "reducto://already-uploaded.pdf", + }, + api_key="test-key", + api_base="https://platform.reducto.ai", + retrieval={"chunk_mode": "section"}, + ) + + assert not upload_route.called + assert parse_route.called + parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert parse_request_body["input"] == "reducto://already-uploaded.pdf" + assert parse_request_body["retrieval"]["chunk_mode"] == "section" + assert response.pages[0].markdown.startswith("Page 1 block A") diff --git a/tests/test_litellm/llms/reducto/test_upload.py b/tests/test_litellm/llms/reducto/test_upload.py new file mode 100644 index 00000000000..4fae90436bb --- /dev/null +++ b/tests/test_litellm/llms/reducto/test_upload.py @@ -0,0 +1,213 @@ +import json +import os +from unittest.mock import AsyncMock, Mock + +import httpx +import litellm +import pytest + +from litellm.llms.reducto.common import ( + extract_file_id_or_bytes, + upload_bytes_async, + upload_bytes_sync, +) + + +@pytest.fixture() +def disable_aiohttp_transport(monkeypatch): + original_disable_aiohttp = litellm.disable_aiohttp_transport + litellm.disable_aiohttp_transport = True + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setenv("REDUCTO_API_KEY", "env-reducto-key") + try: + yield + finally: + litellm.disable_aiohttp_transport = original_disable_aiohttp + litellm.in_memory_llm_clients_cache.flush_cache() + os.environ.pop("REDUCTO_API_KEY", None) + + +@pytest.mark.asyncio +async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport): + with pytest.raises(litellm.BadRequestError, match="upload the file first"): + await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "document_url", + "document_url": "https://example.com/document.pdf", + }, + api_key="test-key", + api_base="https://platform.reducto.ai", + ) + + +@pytest.mark.asyncio +async def test_parse_v3_image_data_uri_upload_uses_image_mime( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://custom.reducto.test/upload").respond( + json={"file_id": "reducto://uploaded-image.png"} + ) + parse_route = respx_mock.post("https://custom.reducto.test/parse").respond( + json={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Image OCR", + "blocks": [{"content": "Image OCR", "bbox": {"page": 1}}], + } + ] + }, + } + ) + + response = await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "file", + "file": b"\x89PNG\r\n\x1a\npng", + "mime_type": "image/png", + }, + api_key="programmatic-key", + api_base="https://custom.reducto.test/", + ) + + assert upload_route.called + assert parse_route.called + upload_request = upload_route.calls[0].request + assert upload_request.headers["authorization"] == "Bearer programmatic-key" + assert b"image/png" in upload_request.read() + + parse_request_body = json.loads(parse_route.calls[0].request.read()) + assert parse_request_body["input"] == "reducto://uploaded-image.png" + assert response.pages[0].markdown == "Image OCR" + + +@pytest.mark.asyncio +async def test_parse_v3_uses_programmatic_api_key_over_env( + disable_aiohttp_transport, respx_mock +): + upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( + json={"file_id": "reducto://uploaded.pdf"} + ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond( + json={ + "usage": {"num_pages": 1, "credits": 1}, + "result": { + "chunks": [ + { + "content": "Programmatic auth", + "blocks": [ + {"content": "Programmatic auth", "bbox": {"page": 1}} + ], + } + ] + }, + } + ) + + await litellm.aocr( + model="reducto/parse-v3", + document={ + "type": "file", + "file": b"%PDF-1.4 auth", + "mime_type": "application/pdf", + }, + api_key="passed-key", + api_base="https://platform.reducto.ai", + ) + + assert upload_route.calls[0].request.headers["authorization"] == "Bearer passed-key" + assert parse_route.calls[0].request.headers["authorization"] == "Bearer passed-key" + + +def test_upload_bytes_sync_uses_shared_client(monkeypatch): + captured = {} + + def fake_post(*, url, headers, files, timeout): + captured["url"] = url + captured["headers"] = headers + captured["files"] = files + captured["timeout"] = timeout + return httpx.Response( + 200, + json={"file_id": "reducto://sync-upload"}, + request=httpx.Request("POST", url), + ) + + sync_post = Mock(side_effect=fake_post) + monkeypatch.setattr(litellm.module_level_client, "post", sync_post) + + class ForbiddenSyncClient: + def __init__(self, *args, **kwargs): + raise AssertionError("should not construct") + + monkeypatch.setattr(httpx, "Client", ForbiddenSyncClient) + + file_id = upload_bytes_sync( + raw_bytes=b"%PDF-1.4 sync", + mime="application/pdf", + api_key="sync-key", + api_base="https://sync.reducto.test/", + ) + + assert file_id == "reducto://sync-upload" + sync_post.assert_called_once() + assert captured["url"] == "https://sync.reducto.test/upload" + assert captured["headers"] == {"Authorization": "Bearer sync-key"} + assert captured["files"]["file"] == ( + "document", + b"%PDF-1.4 sync", + "application/pdf", + ) + + +@pytest.mark.asyncio +async def test_upload_bytes_async_uses_shared_aclient(monkeypatch): + captured = {} + + async def fake_post(*, url, headers, files, timeout): + captured["url"] = url + captured["headers"] = headers + captured["files"] = files + captured["timeout"] = timeout + return httpx.Response( + 200, + json={"file_id": "reducto://async-upload"}, + request=httpx.Request("POST", url), + ) + + async_post = AsyncMock(side_effect=fake_post) + monkeypatch.setattr(litellm.module_level_aclient, "post", async_post) + + class ForbiddenAsyncClient: + def __init__(self, *args, **kwargs): + raise AssertionError("should not construct") + + monkeypatch.setattr(httpx, "AsyncClient", ForbiddenAsyncClient) + + file_id = await upload_bytes_async( + raw_bytes=b"%PDF-1.4 async", + mime="application/pdf", + api_key="async-key", + api_base="https://async.reducto.test/", + ) + + assert file_id == "reducto://async-upload" + async_post.assert_awaited_once() + assert captured["url"] == "https://async.reducto.test/upload" + assert captured["headers"] == {"Authorization": "Bearer async-key"} + assert captured["files"]["file"] == ( + "document", + b"%PDF-1.4 async", + "application/pdf", + ) + + +def test_extract_file_id_or_bytes_raises_on_malformed_data_uri(): + with pytest.raises(litellm.BadRequestError, match="Invalid Reducto data URI"): + extract_file_id_or_bytes("data:application/pdf", model="reducto/parse-v3") + + with pytest.raises(litellm.BadRequestError, match="Invalid Reducto base64 payload"): + extract_file_id_or_bytes("data:;base64,!!!not-base64", model="reducto/parse-v3") diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py index a36aec32d13..943a3160bb7 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py @@ -17,6 +17,9 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) from litellm import embedding +from litellm.llms.sagemaker.embedding.cohere_transformation import ( + SagemakerCohereEmbeddingConfig, +) from litellm.llms.sagemaker.embedding.transformation import SagemakerEmbeddingConfig from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig from litellm.types.utils import EmbeddingResponse, Usage @@ -54,6 +57,172 @@ class TestSagemakerEmbeddingFactory: assert isinstance(config2, VoyageEmbeddingConfig) assert isinstance(config3, VoyageEmbeddingConfig) + def test_get_model_config_cohere_model(self): + """Cohere SageMaker endpoints route to SagemakerCohereEmbeddingConfig""" + for endpoint_name in ( + "cohere.embed-multilingual-v3", + "cohere-embed-english-v3-prod", + "my-cohere-marketplace-endpoint", + "COHERE-EMBED-V4", + ): + config = SagemakerEmbeddingConfig.get_model_config(endpoint_name) + assert isinstance(config, SagemakerCohereEmbeddingConfig), endpoint_name + + +class TestSagemakerCohereEmbeddingConfig: + """Cohere-specific SageMaker embedding request/response transforms""" + + def setup_method(self): + self.config = SagemakerCohereEmbeddingConfig() + + MODEL = "cohere.embed-multilingual-v3" + + def test_transform_request_uses_cohere_payload(self): + """Bug repro: request must use `texts` + `input_type`, not HF `inputs`""" + result = self.config.transform_embedding_request( + model=self.MODEL, + input=["hello"], + optional_params={"input_type": "search_query"}, + headers={}, + ) + assert "inputs" not in result + assert result["texts"] == ["hello"] + assert result["input_type"] == "search_query" + + def test_transform_request_default_input_type(self): + result = self.config.transform_embedding_request( + model=self.MODEL, + input=["hello"], + optional_params={}, + headers={}, + ) + assert result["texts"] == ["hello"] + assert result["input_type"] == "search_document" + + def test_transform_request_normalizes_string_input(self): + result = self.config.transform_embedding_request( + model=self.MODEL, + input="hello", + optional_params={}, + headers={}, + ) + assert result["texts"] == ["hello"] + + def test_map_openai_params_dimensions_to_output_dimension(self): + params = self.config.map_openai_params( + non_default_params={"dimensions": 512, "encoding_format": "float"}, + optional_params={}, + model=self.MODEL, + drop_params=False, + ) + assert params["output_dimension"] == 512 + assert params["embedding_types"] == ["float"] + + def test_map_openai_params_input_type_from_non_default_params(self): + params = self.config.map_openai_params( + non_default_params={"input_type": "search_query"}, + optional_params={}, + model=self.MODEL, + drop_params=False, + ) + assert params["input_type"] == "search_query" + + def test_get_optional_params_embeddings_preserves_input_type(self): + """Exercises get_optional_params_embeddings, not transform in isolation.""" + from litellm.utils import get_optional_params_embeddings + + optional_params = get_optional_params_embeddings( + model=self.MODEL, + custom_llm_provider="sagemaker", + input_type="search_query", + ) + assert optional_params.get("input_type") == "search_query" + + body = self.config.transform_embedding_request( + model=self.MODEL, + input=["hello"], + optional_params=optional_params, + headers={}, + ) + assert body["texts"] == ["hello"] + assert body["input_type"] == "search_query" + + def test_get_optional_params_embeddings_maps_dimensions_without_duplicate(self): + """dimensions must map to output_dimension only, not also stay as dimensions.""" + from litellm.utils import get_optional_params_embeddings + + optional_params = get_optional_params_embeddings( + model=self.MODEL, + custom_llm_provider="sagemaker", + dimensions=512, + input_type="search_query", + ) + assert optional_params.get("output_dimension") == 512 + assert "dimensions" not in optional_params + assert optional_params.get("input_type") == "search_query" + + def test_transform_response_parses_cohere_payload(self): + cohere_response = { + "embeddings": [[0.1, 0.2, 0.3]], + "meta": {"billed_units": {"input_tokens": 2}}, + } + mock_response = httpx.Response( + status_code=200, + content=json.dumps(cohere_response).encode("utf-8"), + headers={"content-type": "application/json"}, + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {"input": ["hello"]} + + result = self.config.transform_embedding_response( + model=self.MODEL, + raw_response=mock_response, + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key=None, + request_data={"texts": ["hello"], "input_type": "search_query"}, + optional_params={}, + litellm_params={}, + ) + + assert result.object == "list" + assert len(result.data) == 1 + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert result.usage.prompt_tokens == 2 + + def test_transform_response_does_not_double_call_post_call(self): + """ + Greptile review fix: SageMaker handler already calls + `logging_obj.post_call` once before invoking + `transform_embedding_response`. The transform must NOT call it again, + otherwise callbacks, cost calculators, and log handlers double-fire + for every Cohere SageMaker embedding call. + """ + cohere_response = { + "embeddings": [[0.1, 0.2, 0.3]], + "meta": {"billed_units": {"input_tokens": 2}}, + } + mock_response = httpx.Response( + status_code=200, + content=json.dumps(cohere_response).encode("utf-8"), + headers={"content-type": "application/json"}, + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {"input": ["hello"]} + + self.config.transform_embedding_response( + model=self.MODEL, + raw_response=mock_response, + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key=None, + request_data={"texts": ["hello"], "input_type": "search_query"}, + optional_params={}, + litellm_params={}, + ) + + logging_obj.post_call.assert_not_called() + class TestVoyageEmbeddingConfig: """Test Voyage-specific embedding configuration""" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index db993245426..45b9f4293fa 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2097,6 +2097,125 @@ def test_is_gemini_3_or_newer(): assert VertexGeminiConfig._is_gemini_3_or_newer("") == False +def test_forward_gemini_function_call_id_vertex_vs_google_ai_studio(): + """Vertex AI rejects `id` on function_call/function_response; Google AI Studio accepts it on Gemini 3.5+.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + model = "gemini-3.5-flash" + assert ( + VertexGeminiConfig._forward_gemini_function_call_id(model, "vertex_ai") is False + ) + assert ( + VertexGeminiConfig._forward_gemini_function_call_id(model, "vertex_ai_beta") + is False + ) + assert VertexGeminiConfig._forward_gemini_function_call_id(model, "gemini") is True + assert VertexGeminiConfig._forward_gemini_function_call_id(model, None) is False + assert ( + VertexGeminiConfig._forward_gemini_function_call_id( + "gemini-2.5-flash", "gemini" + ) + is False + ) + + +def test_vertex_ai_gemini_35_tool_calls_omit_function_call_id(): + """Regression: Vertex must not send OpenAI tool_call id inside Gemini function_call parts.""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + messages = [ + {"role": "user", "content": "Explore this directory"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_50e7e0fe0989464a89f188eda443", + "type": "function", + "function": { + "name": "read", + "arguments": '{"filePath": "/tmp"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_50e7e0fe0989464a89f188eda443", + "content": "ok", + }, + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, + model="gemini-3.5-flash", + custom_llm_provider="vertex_ai", + ) + + for content in contents: + for part in content.get("parts", []): + fc = part.get("function_call") + if fc is not None: + assert "id" not in fc, f"Vertex payload must not include id: {fc}" + fr = part.get("function_response") + if fr is not None: + assert "id" not in fr, f"Vertex payload must not include id: {fr}" + + +def test_google_ai_studio_gemini_35_tool_calls_include_function_call_id(): + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + tool_call_id = "call_50e7e0fe0989464a89f188eda443" + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": { + "name": "read", + "arguments": '{"filePath": "/tmp"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": tool_call_id, + "content": "ok", + }, + ] + + contents = _gemini_convert_messages_with_history( + messages=messages, + model="gemini-3.5-flash", + custom_llm_provider="gemini", + ) + + function_call_ids = [] + function_response_ids = [] + for content in contents: + for part in content.get("parts", []): + fc = part.get("function_call") + if fc is not None: + function_call_ids.append(fc.get("id")) + fr = part.get("function_response") + if fr is not None: + function_response_ids.append(fr.get("id")) + + assert function_call_ids == [tool_call_id] + assert function_response_ids == [tool_call_id] + + def test_reasoning_effort_maps_to_thinking_level_gemini_3(): """Test that reasoning_effort maps to thinking_level AND includeThoughts for Gemini 3+ models""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -3531,7 +3650,12 @@ def test_video_metadata_supported_for_all_gemini_models(): } ] - for model in ["gemini-1.5-pro", "gemini-2.5-flash", "gemini-2.5-pro", "gemini-3-pro-preview"]: + for model in [ + "gemini-1.5-pro", + "gemini-2.5-flash", + "gemini-2.5-pro", + "gemini-3-pro-preview", + ]: contents = _gemini_convert_messages_with_history(messages=messages, model=model) file_part = None @@ -3541,19 +3665,25 @@ def test_video_metadata_supported_for_all_gemini_models(): break assert file_part is not None, f"{model}: file part should exist" - assert "video_metadata" in file_part, f"{model}: video_metadata should be present" + assert ( + "video_metadata" in file_part + ), f"{model}: video_metadata should be present" assert file_part["video_metadata"]["fps"] == 5, f"{model}: fps should be 5" # Per-part media_resolution is Gemini 3+ only; 2.x uses generation_config global for model in ["gemini-3-pro-preview"]: contents = _gemini_convert_messages_with_history(messages=messages, model=model) file_part = next(p for p in contents[0]["parts"] if "file_data" in p) - assert "media_resolution" in file_part, f"{model}: media_resolution should be present" + assert ( + "media_resolution" in file_part + ), f"{model}: media_resolution should be present" for model in ["gemini-1.5-pro", "gemini-2.5-flash", "gemini-2.5-pro"]: contents = _gemini_convert_messages_with_history(messages=messages, model=model) file_part = next(p for p in contents[0]["parts"] if "file_data" in p) - assert "media_resolution" not in file_part, f"{model}: per-part media_resolution should not be set" + assert ( + "media_resolution" not in file_part + ), f"{model}: per-part media_resolution should not be set" def test_chunk_parser_handles_prompt_feedback_block(): @@ -4186,8 +4316,9 @@ def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt(): # DOCUMENT tokens should be included in text_tokens: 8 (TEXT) + 774 (DOCUMENT) = 782 assert result.prompt_tokens_details is not None - assert result.prompt_tokens_details.text_tokens == 782, \ - "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)" + assert ( + result.prompt_tokens_details.text_tokens == 782 + ), "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)" # Verify completion token details assert result.completion_tokens_details is not None @@ -4222,8 +4353,9 @@ def test_vertex_ai_usage_metadata_with_document_tokens_cached(): # DOCUMENT cached tokens map to cached_text_tokens, so: # text_tokens = (8 TEXT + 774 DOCUMENT) - 400 cached = 382 - assert result.prompt_tokens_details.text_tokens == 382, \ - "text_tokens should be (8 + 774) - 400 cached = 382" + assert ( + result.prompt_tokens_details.text_tokens == 382 + ), "text_tokens should be (8 + 774) - 400 cached = 382" assert result.prompt_tokens_details.cached_tokens == 400 @@ -4693,7 +4825,9 @@ def test_mid_stream_429_error_raises_during_iteration(): { "content": { "role": "model", - "parts": [{"text": "Let me think about this...", "thought": True}], + "parts": [ + {"text": "Let me think about this...", "thought": True} + ], }, "index": 0, } @@ -4713,7 +4847,9 @@ def test_mid_stream_429_error_raises_during_iteration(): { "content": { "role": "model", - "parts": [{"text": "I'll generate the image now.", "thought": True}], + "parts": [ + {"text": "I'll generate the image now.", "thought": True} + ], }, "index": 0, } diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 88aac07a0c9..2cf97081806 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -1448,3 +1449,474 @@ class TestVertexBase: aws_creds = supplier.get_aws_security_credentials(context=None, request=None) assert isinstance(aws_creds, AwsSecurityCredentials) + + @pytest.mark.asyncio + async def test_single_flight_refresh(self): + """Under high concurrency, only one coroutine should refresh expired credentials.""" + import asyncio + + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "expired-token" + mock_creds.expired = True + mock_creds.expiry = None + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + refresh_call_count = 0 + + with ( + patch.object( + vertex_base, "load_auth", return_value=(mock_creds, "project-1") + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + async def slow_refresh(creds): + nonlocal refresh_call_count + refresh_call_count += 1 + await asyncio.sleep(0.05) # simulate network latency + creds.token = "refreshed-token" + creds.expired = False + + # refresh_auth is sync, but we need to count calls. + # get_access_token_async wraps it with asyncify, so the sync side_effect works. + def sync_refresh_impl(creds): + nonlocal refresh_call_count + refresh_call_count += 1 + creds.token = "refreshed-token" + creds.expired = False + + mock_refresh.side_effect = sync_refresh_impl + + # Launch 50 concurrent requests + tasks = [ + vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + for _ in range(50) + ] + results = await asyncio.gather(*tasks) + + # All should return the refreshed token + for token, project in results: + assert token == "refreshed-token" + assert project == "project-1" + + # refresh_auth should be called exactly once (single-flight) + assert ( + refresh_call_count == 1 + ), f"Expected 1 refresh call, got {refresh_call_count}" + + @pytest.mark.asyncio + async def test_async_reauthentication_uses_async_single_flight(self): + """Concurrent async reauth should reload once without using the sync path.""" + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + stale_creds = MagicMock() + stale_creds.token = "expired-token" + stale_creds.token_state = TokenState.INVALID + stale_creds.project_id = "project-1" + stale_creds.quota_project_id = "project-1" + + refreshed_creds = MagicMock() + refreshed_creds.token = "refreshed-token" + refreshed_creds.token_state = TokenState.FRESH + refreshed_creds.project_id = "project-1" + refreshed_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + cache_key = (json.dumps(credentials), "project-1") + vertex_base._credentials_project_mapping[cache_key] = ( + stale_creds, + "project-1", + ) + + load_call_count = 0 + + def load_auth_impl(*_args, **_kwargs): + nonlocal load_call_count + load_call_count += 1 + return refreshed_creds, "project-1" + + with ( + patch.object( + vertex_base, + "refresh_auth", + side_effect=Exception("Reauthentication is needed"), + ), + patch.object(vertex_base, "load_auth", side_effect=load_auth_impl), + patch.object(vertex_base, "get_access_token") as mock_get_access_token, + ): + results = await asyncio.gather( + *[ + vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + for _ in range(10) + ] + ) + + assert results == [("refreshed-token", "project-1")] * 10 + assert load_call_count == 1 + mock_get_access_token.assert_not_called() + + @pytest.mark.asyncio + async def test_background_refresh_when_near_expiry(self): + """When token_state is STALE (within the 3:45 REFRESH_THRESHOLD window), + return the current token immediately and refresh in the background — + zero added latency.""" + import asyncio + + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + # Simulate STALE state: token is usable but near expiry. + mock_creds = MagicMock() + mock_creds.token = "near-expiry-token" + mock_creds.token_state = TokenState.STALE + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + with ( + patch.object( + vertex_base, "load_auth", return_value=(mock_creds, "project-1") + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds): + creds.token = "refreshed-token" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + token, project = await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + # Should return the current (still usable) token immediately + assert token == "near-expiry-token" + + # Let the background refresh task run + await asyncio.sleep(0.05) + + assert mock_refresh.called, "Background refresh should have been triggered" + + @pytest.mark.asyncio + async def test_stale_malformed_token_blocks_on_refresh(self): + """Malformed STALE tokens should refresh instead of failing validation.""" + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = None + mock_creds.token_state = TokenState.STALE + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + with ( + patch.object( + vertex_base, "load_auth", return_value=(mock_creds, "project-1") + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds): + creds.token = "refreshed-token" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + token, project = await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + assert mock_refresh.called + assert token == "refreshed-token" + assert project == "project-1" + + @pytest.mark.asyncio + async def test_fresh_token_skips_refresh(self): + """Credentials not marked expired by google-auth should not trigger refresh.""" + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "fresh-token" + mock_creds.expired = False + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + cache_key = (json.dumps(credentials), "project-1") + vertex_base._credentials_project_mapping[cache_key] = ( + mock_creds, + "project-1", + ) + + with patch.object(vertex_base, "refresh_auth") as mock_refresh: + token, project = await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + assert not mock_refresh.called, "Fresh token should not trigger refresh" + assert token == "fresh-token" + + @pytest.mark.asyncio + async def test_background_refresh_task_removed_after_completion(self): + """Completed background-refresh tasks must be evicted from + _background_refresh_tasks so the dict does not grow unboundedly.""" + import asyncio + + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + mock_creds = MagicMock() + mock_creds.token = "near-expiry-token" + mock_creds.token_state = TokenState.STALE + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + + with ( + patch.object( + vertex_base, "load_auth", return_value=(mock_creds, "project-1") + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds): + creds.token = "refreshed-token" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + # Allow the background task to complete. + await asyncio.sleep(0.1) + + # After completion the entry should have been removed by the done-callback. + assert len(vertex_base._background_refresh_tasks) == 0, ( + "Completed background refresh task was not removed from " + "_background_refresh_tasks" + ) + + @pytest.mark.asyncio + async def test_background_refresh_tasks_no_accumulation_across_many_keys(self): + """With many distinct credential keys the dict must not hold completed tasks.""" + import asyncio + import json as _json + + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + num_keys = 20 + + for i in range(num_keys): + mock_creds = MagicMock() + mock_creds.token = f"token-{i}" + mock_creds.token_state = TokenState.STALE + mock_creds.project_id = f"project-{i}" + mock_creds.quota_project_id = f"project-{i}" + + credentials = {"type": "service_account", "project_id": f"project-{i}"} + + with ( + patch.object( + vertex_base, + "load_auth", + return_value=(mock_creds, f"project-{i}"), + ), + patch.object(vertex_base, "refresh_auth") as mock_refresh, + ): + + def mock_refresh_impl(creds, idx=i): + creds.token = f"refreshed-{idx}" + creds.token_state = TokenState.FRESH + + mock_refresh.side_effect = mock_refresh_impl + + await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id=f"project-{i}", + custom_llm_provider="vertex_ai", + ) + + # Let all background tasks finish. + await asyncio.sleep(0.1) + + assert len(vertex_base._background_refresh_tasks) == 0, ( + f"Expected 0 tasks after all refreshes completed, " + f"found {len(vertex_base._background_refresh_tasks)}" + ) + + @pytest.mark.asyncio + async def test_async_refresh_lock_shared_while_in_use(self): + """Concurrent callers for the same key must coordinate on the same lock.""" + vertex_base = VertexBase() + key = ("creds", "project-1") + + lock_a = vertex_base._acquire_async_refresh_lock(key) + try: + async with lock_a: + lock_b = vertex_base._acquire_async_refresh_lock(key) + try: + assert lock_a is lock_b, ( + "While a coroutine still holds the lock, concurrent callers must " + "receive the same Lock instance to preserve single-flight." + ) + finally: + vertex_base._release_async_refresh_lock(key, lock_b) + finally: + vertex_base._release_async_refresh_lock(key, lock_a) + + @pytest.mark.asyncio + async def test_async_refresh_lock_pruned_after_release(self): + """get_access_token_async must drop the per-key Lock from the registry + once no coroutine is using it, so the dict stays bounded in + high-cardinality deployments. Without this, every distinct credential + leaks a Lock object for the lifetime of the process.""" + from google.auth.credentials import TokenState + + vertex_base = VertexBase() + + for i in range(10): + mock_creds = MagicMock() + mock_creds.token = f"refreshed-{i}" + mock_creds.token_state = TokenState.FRESH + mock_creds.project_id = f"project-{i}" + mock_creds.quota_project_id = f"project-{i}" + + credentials = {"type": "service_account", "project_id": f"project-{i}"} + + with ( + patch.object( + vertex_base, + "load_auth", + return_value=(mock_creds, f"project-{i}"), + ), + patch.object(vertex_base, "refresh_auth"), + ): + await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id=f"project-{i}", + custom_llm_provider="vertex_ai", + ) + + assert len(vertex_base._async_refresh_locks) == 0, ( + "expected per-key locks to be pruned once no coroutine holds or " + f"waits on them; found {len(vertex_base._async_refresh_locks)}" + ) + assert len(vertex_base._async_refresh_lock_refcounts) == 0 + + @pytest.mark.asyncio + async def test_async_refresh_lock_kept_while_waiter_pending(self): + """The prune must not run while another coroutine is still waiting on + the lock — otherwise the waiter ends up on a lock that's been replaced + in the registry and single-flight breaks.""" + vertex_base = VertexBase() + key = ("creds", "project-1") + + holder_lock = vertex_base._acquire_async_refresh_lock(key) + release_holder = asyncio.Event() + + async def hold_then_release(): + async with holder_lock: + await release_holder.wait() + vertex_base._release_async_refresh_lock(key, holder_lock) + + holder = asyncio.create_task(hold_then_release()) + await asyncio.sleep(0) # let holder grab the lock + + async def queue_for_lock(): + waiter_lock = vertex_base._acquire_async_refresh_lock(key) + try: + async with waiter_lock: + pass + finally: + vertex_base._release_async_refresh_lock(key, waiter_lock) + + waiter = asyncio.create_task(queue_for_lock()) + await asyncio.sleep(0) # let waiter queue on the lock + + assert ( + vertex_base._async_refresh_locks.get(key) is holder_lock + ), "lock with active holder/waiter must not be pruned" + + release_holder.set() + await holder + await waiter + + assert key not in vertex_base._async_refresh_locks + assert key not in vertex_base._async_refresh_lock_refcounts + + @pytest.mark.asyncio + async def test_fast_path_no_lock(self): + """Cached fresh credentials should return without acquiring the lock.""" + import datetime + + vertex_base = VertexBase() + + try: + from google.auth import _helpers as google_auth_helpers + + now = google_auth_helpers.utcnow() + except ImportError: + now = datetime.datetime.utcnow() + + mock_creds = MagicMock() + mock_creds.token = "cached-token" + mock_creds.expired = False + mock_creds.expiry = now + datetime.timedelta(minutes=30) + mock_creds.project_id = "project-1" + mock_creds.quota_project_id = "project-1" + + credentials = {"type": "service_account", "project_id": "project-1"} + cache_key = (json.dumps(credentials), "project-1") + vertex_base._credentials_project_mapping[cache_key] = ( + mock_creds, + "project-1", + ) + + # Spy on _acquire_async_refresh_lock to verify it's never called + with patch.object( + vertex_base, + "_acquire_async_refresh_lock", + wraps=vertex_base._acquire_async_refresh_lock, + ) as mock_get_lock: + token, project = await vertex_base._ensure_access_token_async( + credentials=credentials, + project_id="project-1", + custom_llm_provider="vertex_ai", + ) + + assert token == "cached-token" + assert not mock_get_lock.called, "Fast path should not acquire lock" diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py index b16fc2bc44d..f617a8db850 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gpt_oss/test_vertex_ai_gpt_oss_transformation.py @@ -118,7 +118,7 @@ async def test_vertex_ai_gpt_oss_simple_request(): "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" ) as mock_http_handler, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718"), ), patch.dict( @@ -217,7 +217,7 @@ async def test_vertex_ai_gpt_oss_reasoning_effort(): "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" ) as mock_http_handler, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token", return_value=("fake-token", "pathrise-convert-1606954137718"), ), patch.dict( diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py index bf6e0a5f2cd..5a86325b7fd 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py @@ -7,7 +7,6 @@ These tests verify that: 3. The completion() and responses() API work with Qwen models """ -import json import os import sys from unittest.mock import MagicMock, patch, AsyncMock @@ -179,7 +178,7 @@ async def test_vertex_ai_qwen_global_endpoint_url(): "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" ) as mock_http_handler, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token", return_value=("fake-token", "test-project"), ), patch.dict( diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py new file mode 100644 index 00000000000..b20442a032e --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/test_partner_models_credential_reuse.py @@ -0,0 +1,220 @@ +""" +Test that VertexBase subclasses (PartnerModels, Gemma, ModelGarden) reuse +cached credentials instead of creating a new VertexLLM instance on every request. +""" + +import sys +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.llms.vertex_ai.vertex_ai_partner_models.main import ( + VertexAIPartnerModels, +) +from litellm.llms.vertex_ai.vertex_gemma_models.main import VertexAIGemmaModels +from litellm.llms.vertex_ai.vertex_model_garden.main import VertexAIModelGardenModels + + +def _mock_vertexai(): + """Return a MagicMock that satisfies the vertexai import guards.""" + m = MagicMock() + m.preview = MagicMock() + m.preview.language_models = MagicMock() + return m + + +class TestVertexBaseSubclassInit: + """All VertexBase subclasses must call super().__init__() so that + the credential cache is initialized.""" + + @pytest.mark.parametrize( + "cls", + [VertexAIPartnerModels, VertexAIGemmaModels, VertexAIModelGardenModels], + ids=["PartnerModels", "Gemma", "ModelGarden"], + ) + def test_init_calls_super(self, cls): + instance = cls() + assert hasattr(instance, "_credentials_project_mapping") + assert isinstance(instance._credentials_project_mapping, dict) + assert hasattr(instance, "access_token") + assert hasattr(instance, "project_id") + + +class TestPartnerModelsCredentialReuse: + def test_completion_uses_self_ensure_access_token(self): + """completion() should call self._ensure_access_token, not create a + throwaway VertexLLM instance.""" + partner = VertexAIPartnerModels() + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + partner, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ) as mock_ensure, + patch( + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.base_llm_http_handler" + ) as mock_handler, + ): + mock_handler.completion.return_value = "response" + + partner.completion( + model="meta/llama-3.1-405b-instruct-maas", + messages=[{"role": "user", "content": "hello"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + mock_ensure.assert_called_once_with( + credentials='{"type": "service_account"}', + project_id="test-project", + custom_llm_provider="vertex_ai", + ) + + def test_credential_cache_shared_across_calls(self): + """Two successive completion() calls should hit load_auth only once.""" + partner = VertexAIPartnerModels() + + mock_creds = MagicMock() + mock_creds.token = "my-token" + mock_creds.expired = False + mock_creds.project_id = "proj" + mock_creds.quota_project_id = "proj" + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + partner, "load_auth", return_value=(mock_creds, "proj") + ) as mock_load, + patch( + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.base_llm_http_handler" + ) as mock_handler, + ): + mock_handler.completion.return_value = "resp" + + common_kwargs = dict( + model="meta/llama-3.1-405b-instruct-maas", + messages=[{"role": "user", "content": "hi"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="proj", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + partner.completion(**common_kwargs) + partner.completion(**common_kwargs) + + assert mock_load.call_count == 1 + + +class TestGemmaModelsCredentialReuse: + def test_completion_uses_self_ensure_access_token(self): + """completion() should call self._ensure_access_token, not create a + throwaway VertexLLM instance.""" + gemma = VertexAIGemmaModels() + + mock_gemma_config = MagicMock() + mock_gemma_config.return_value.completion.return_value = "response" + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + gemma, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ) as mock_ensure, + patch( + "litellm.llms.vertex_ai.vertex_gemma_models.transformation.VertexGemmaConfig", + mock_gemma_config, + ), + ): + gemma.completion( + model="gemma/gemma-3-12b-it-1234567890", + messages=[{"role": "user", "content": "hello"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base="https://123.us-central1-1.prediction.vertexai.goog/v1/projects/proj/locations/us-central1/endpoints/456:predict", + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + mock_ensure.assert_called_once_with( + credentials='{"type": "service_account"}', + project_id="test-project", + custom_llm_provider="vertex_ai", + ) + + +class TestModelGardenCredentialReuse: + def test_completion_uses_self_ensure_access_token(self): + """completion() should call self._ensure_access_token, not create a + throwaway VertexLLM instance.""" + garden = VertexAIModelGardenModels() + + mock_handler = MagicMock() + mock_handler.return_value.completion.return_value = "response" + + with ( + patch.dict(sys.modules, {"vertexai": _mock_vertexai()}), + patch.object( + garden, + "_ensure_access_token", + return_value=("cached-token", "test-project"), + ) as mock_ensure, + patch( + "litellm.llms.openai_like.chat.handler.OpenAILikeChatHandler", + mock_handler, + ), + ): + garden.completion( + model="openai/5464397967697903616", + messages=[{"role": "user", "content": "hello"}], + model_response=MagicMock(), + print_verbose=lambda *a, **kw: None, + encoding=MagicMock(), + logging_obj=MagicMock(), + api_base=None, + optional_params={}, + custom_prompt_dict={}, + headers=None, + timeout=30.0, + litellm_params={}, + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials='{"type": "service_account"}', + ) + + mock_ensure.assert_called_once_with( + credentials='{"type": "service_account"}', + project_id="test-project", + custom_llm_provider="vertex_ai", + ) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index 3e3e8901706..b1c8f7234ce 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -122,17 +122,19 @@ class TestVertexGemmaCompletion: # Mock the async HTTP handler and Vertex authentication with ( patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" - ) as mock_http_handler, + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_get_client, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", return_value=("fake-access-token", "PROJECT_ID"), ), ): + mock_client = Mock() mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = mock_vertex_response - mock_http_handler.return_value.post = AsyncMock(return_value=mock_response) + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client # Call litellm.acompletion() response = await litellm.acompletion( @@ -145,7 +147,7 @@ class TestVertexGemmaCompletion: ) # Verify the request sent to Vertex - call_args = mock_http_handler.return_value.post.call_args + call_args = mock_client.post.call_args assert call_args is not None, "HTTP handler was not called" request_data = call_args.kwargs["json"] @@ -210,17 +212,19 @@ class TestVertexGemmaCompletion: with ( patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler" - ) as mock_http_handler, + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_get_client, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", return_value=("fake-access-token", "test-project"), ), ): + mock_client = Mock() mock_response = Mock() mock_response.status_code = 200 mock_response.json.return_value = invalid_response - mock_http_handler.return_value.post = AsyncMock(return_value=mock_response) + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client # Should raise exception (wrapped as APIConnectionError by LiteLLM) with pytest.raises(APIConnectionError) as exc_info: @@ -286,7 +290,7 @@ class TestVertexGemmaCompletion: "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" ) as mock_get_client, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", return_value=("fake-access-token", "PROJECT_ID"), ), ): @@ -388,7 +392,7 @@ class TestVertexGemmaCompletion: "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" ) as mock_get_client, patch( - "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexLLM._ensure_access_token", + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", return_value=("fake-access-token", "PROJECT_ID"), ), ): @@ -429,3 +433,123 @@ class TestVertexGemmaCompletion: # Verify other parameters are present assert "messages" in instance assert instance["@requestFormat"] == "chatCompletions" + + @pytest.mark.asyncio + async def test_acompletion_filters_context_management(self): + """ + Test that context_management is filtered out from the request. + + Vertex AI Gemma's chatCompletions wrapper does not understand + `context_management` (an Anthropic / OpenAI Responses API concept). + It must be stripped from the request body so the upstream endpoint + does not reject the request with an unknown-field error. + """ + mock_vertex_response = { + "deployedModelId": "1207280419999999999", + "model": "projects/993702345710/locations/us-central1/models/gemma-3-12b-it-1222199011122", + "modelDisplayName": "gemma-3-12b-it-1222199011122", + "modelVersionId": "1", + "predictions": { + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "message": { + "content": "ok", + "reasoning_content": None, + "role": "assistant", + "tool_calls": [], + }, + "stop_reason": None, + } + ], + "created": 1759863903, + "id": "chatcmpl-test-ctxmgmt", + "model": "google/gemma-3-12b-it", + "object": "chat.completion", + "prompt_logprobs": None, + "usage": { + "completion_tokens": 1, + "prompt_tokens": 5, + "prompt_tokens_details": None, + "total_tokens": 6, + }, + }, + } + + with ( + patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client" + ) as mock_get_client, + patch( + "litellm.llms.vertex_ai.vertex_gemma_models.main.VertexAIGemmaModels._ensure_access_token", + return_value=("fake-access-token", "PROJECT_ID"), + ), + ): + mock_client = Mock() + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = mock_vertex_response + mock_client.post = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + # Use `allowed_openai_params` so context_management actually + # reaches the transformation layer (otherwise the upstream + # validator drops it before we can prove the transformation + # strips it). This mirrors the real-world scenario where a + # caller explicitly opts in to forwarding an arbitrary param. + await litellm.acompletion( + model="vertex_ai/gemma/gemma-3-12b-it-1222199011122", + messages=[{"role": "user", "content": "Test"}], + context_management=[ + {"type": "compaction", "compact_threshold": 200000} + ], + allowed_openai_params=["context_management"], + api_base="https://test.us-central1-project.prediction.vertexai.goog/v1/projects/PROJECT_ID/locations/us-central1/endpoints/ENDPOINT_ID:predict", + vertex_project="PROJECT_ID", + vertex_location="us-central1", + ) + + call_args = mock_client.post.call_args + assert call_args is not None, "HTTP client was not called" + + request_data = call_args.kwargs["json"] + print("request body=", json.dumps(request_data, indent=4)) + instance = request_data["instances"][0] + + assert ( + "context_management" not in instance + ), "context_management should not be forwarded to Vertex Gemma" + assert instance["@requestFormat"] == "chatCompletions" + assert "messages" in instance + + def test_transform_request_strips_context_management(self): + """ + Direct unit test for VertexGemmaConfig.transform_request: verify that + `context_management` is stripped from `optional_params` regardless of + how it was supplied to the transformation layer. + """ + from litellm.llms.vertex_ai.vertex_gemma_models.transformation import ( + VertexGemmaConfig, + ) + + config = VertexGemmaConfig() + result = config.transform_request( + model="gemma-3-12b-it", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "max_tokens": 32, + "context_management": [ + {"type": "compaction", "compact_threshold": 200000} + ], + }, + litellm_params={}, + headers={}, + ) + + assert "instances" in result + instance = result["instances"][0] + assert instance["@requestFormat"] == "chatCompletions" + assert "context_management" not in instance + assert instance.get("max_tokens") == 32 diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index 3ae8dfc3c0b..5c1f0f704d7 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -119,3 +119,19 @@ class TestXAIParallelToolCalls: assert result.get("parallel_tool_calls") is True assert len(result["messages"]) == 1 assert result["messages"][0]["role"] == "user" + + +class TestXAIUsageNormalization: + def test_preserves_reasoning_tokens_in_total_usage(self): + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=200) + + XAIChatConfig._normalize_openai_compatible_usage_totals(usage) + + assert usage.total_tokens == 200 + + def test_preserves_reasoning_tokens_in_streaming_usage(self): + usage = {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 200} + + XAIChatConfig._normalize_openai_compatible_usage_totals(usage) + + assert usage["total_tokens"] == 200 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index 66b96785f69..9f2feddb0e3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -1229,6 +1229,143 @@ def test_validate_trusted_redirect_uri_rejects_fragment_and_bad_scheme(): assert exc.value.status_code == 400, uri +def test_validate_trusted_redirect_uri_accepts_cursor_native_callback(): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + req = _make_trusted_request("http://localhost:4000/") + validate_trusted_redirect_uri(req, "cursor://anysphere.cursor-mcp/oauth/callback") + + +def test_validate_trusted_redirect_uri_rejects_unlisted_native_callback( + monkeypatch, +): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + monkeypatch.setenv("MCP_TRUSTED_NATIVE_REDIRECT_URIS", "") + # Clear defaults by patching — env-only path for this test + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.oauth_utils._DEFAULT_NATIVE_REDIRECT_URIS", + [], + ) + req = _make_trusted_request("http://localhost:4000/") + with pytest.raises(HTTPException) as exc: + validate_trusted_redirect_uri( + req, "cursor://anysphere.cursor-mcp/oauth/callback" + ) + assert exc.value.status_code == 400 + + +def test_validate_trusted_redirect_uri_accepts_env_native_redirect_uri( + monkeypatch, +): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.oauth_utils._DEFAULT_NATIVE_REDIRECT_URIS", + [], + ) + monkeypatch.setenv( + "MCP_TRUSTED_NATIVE_REDIRECT_URIS", + "vscode://my-app/oauth/callback", + ) + req = _make_trusted_request("http://localhost:4000/") + validate_trusted_redirect_uri(req, "vscode://my-app/oauth/callback") + + +def test_validate_trusted_redirect_uri_rejects_native_callback_with_fragment(): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + req = _make_trusted_request("http://localhost:4000/") + with pytest.raises(HTTPException) as exc: + validate_trusted_redirect_uri( + req, "cursor://anysphere.cursor-mcp/oauth/callback#frag" + ) + assert exc.value.status_code == 400 + + +def test_validate_trusted_redirect_uri_rejects_native_callback_with_query(): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + req = _make_trusted_request("http://localhost:4000/") + with pytest.raises(HTTPException) as exc: + validate_trusted_redirect_uri( + req, + "cursor://anysphere.cursor-mcp/oauth/callback?injected=anything", + ) + assert exc.value.status_code == 400 + + +def test_validate_trusted_redirect_uri_native_path_case_insensitive(monkeypatch): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.oauth_utils._DEFAULT_NATIVE_REDIRECT_URIS", + [], + ) + monkeypatch.setenv( + "MCP_TRUSTED_NATIVE_REDIRECT_URIS", + "myapp://host/MyPath", + ) + req = _make_trusted_request("http://localhost:4000/") + validate_trusted_redirect_uri(req, "myapp://host/MyPath") + + +def test_validate_trusted_redirect_uri_native_wildcard_respects_path_boundary( + monkeypatch, +): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.oauth_utils._DEFAULT_NATIVE_REDIRECT_URIS", + [], + ) + monkeypatch.setenv( + "MCP_TRUSTED_NATIVE_REDIRECT_URIS", + "cursor://anysphere.cursor-mcp/oauth/callback*", + ) + req = _make_trusted_request("http://localhost:4000/") + validate_trusted_redirect_uri( + req, "cursor://anysphere.cursor-mcp/oauth/callback/extra" + ) + with pytest.raises(HTTPException): + validate_trusted_redirect_uri( + req, "cursor://anysphere.cursor-mcp/oauth/callback-2" + ) + + +def test_validate_trusted_redirect_uri_native_wildcard_directory_prefix( + monkeypatch, +): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.oauth_utils._DEFAULT_NATIVE_REDIRECT_URIS", + [], + ) + monkeypatch.setenv( + "MCP_TRUSTED_NATIVE_REDIRECT_URIS", + "cursor://anysphere.cursor-mcp/oauth/*", + ) + req = _make_trusted_request("http://localhost:4000/") + validate_trusted_redirect_uri(req, "cursor://anysphere.cursor-mcp/oauth/callback") + + def test_validate_trusted_redirect_uri_rejects_scheme_mismatch_on_same_host(): """Regression: an attacker who can serve http on the proxy's own host (e.g. by MITMing an unencrypted LAN hop) must not be able to 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 b06cc7f0f12..c8789e0b0a6 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 @@ -1345,7 +1345,13 @@ def test_validate_trusted_redirect_uri_logs_diagnostic_on_rejection( "https://litellm.example.com/ui/mcp/oauth/callback", ) assert exc_info.value.status_code == 400 - assert exc_info.value.detail == "invalid_request" + detail = exc_info.value.detail + assert isinstance(detail, dict) + assert detail.get("error") == "invalid_request" + assert "error_description" in detail + assert "redirect_uri origin" in detail["error_description"] + assert "proxy origin" in detail["error_description"] + assert "hint" in detail matching = [r for r in caplog.records if "rejecting redirect_uri" in r.getMessage()] assert len(matching) == 1, ( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_header_alias_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_header_alias_utils.py new file mode 100644 index 00000000000..2627199570b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_header_alias_utils.py @@ -0,0 +1,18 @@ +"""Tests for MCP header alias sanitization and auth header lookup.""" + +from litellm.proxy._experimental.mcp_server.utils import ( + lookup_mcp_server_auth_in_headers, + sanitize_mcp_alias_for_header, +) + + +def test_sanitize_mcp_alias_for_header(): + assert sanitize_mcp_alias_for_header("My Server") == "my_server" + assert sanitize_mcp_alias_for_header("GitHub-MCP!") == "github_mcp" + assert sanitize_mcp_alias_for_header("github_mcp2") == "github_mcp2" + + +def test_lookup_mcp_server_auth_in_headers_sanitized_alias(): + headers = {"github_mcp": {"Authorization": "Bearer token"}} + result = lookup_mcp_server_auth_in_headers(headers, alias="GitHub-MCP") + assert result == {"Authorization": "Bearer token"} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index e1eddfc9c7a..f2fd73f3f22 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -774,6 +774,7 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): extra_headers=None, add_prefix=True, raw_headers=None, + user_api_key_auth=None, ): if server.name == "working_server": # Working server returns tools @@ -879,6 +880,7 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): extra_headers=None, add_prefix=True, raw_headers=None, + user_api_key_auth=None, ): # All servers fail raise Exception(f"Server {server.name} connection failed") @@ -1339,6 +1341,7 @@ async def test_list_tools_single_server_unprefixed_names(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): tool = MagicMock() tool.name = f"{server.alias}-toolA" if add_prefix else "toolA" @@ -1420,6 +1423,7 @@ async def test_list_tools_multiple_servers_prefixed_names(): extra_headers=None, add_prefix=True, raw_headers=None, + user_api_key_auth=None, ): tool = MagicMock() # When multiple servers, add_prefix should be True -> prefixed names @@ -1686,6 +1690,7 @@ async def test_list_tools_filters_by_key_team_permissions(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): # Return 4 tools, but only 2 should be allowed tool1 = MagicMock() @@ -1795,6 +1800,7 @@ async def test_list_tools_with_team_tool_permissions_inheritance(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): # Return 4 tools tool1 = MagicMock() @@ -1890,6 +1896,7 @@ async def test_list_tools_with_no_tool_permissions_shows_all(): extra_headers=None, add_prefix=False, raw_headers=None, + user_api_key_auth=None, ): # Return 3 tools tool1 = MagicMock() @@ -1988,6 +1995,7 @@ async def test_list_tools_strips_prefix_when_matching_permissions(): extra_headers=None, add_prefix=True, raw_headers=None, + user_api_key_auth=None, ): # Return tools WITH prefix (as they come from MCP server) tool1 = MagicMock() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ef1c09aa815..d7078412a44 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -322,6 +322,7 @@ class TestMCPServerManager: mcp_auth_header=None, mcp_protocol_version=None, raw_headers=None, + user_api_key_auth=None, ): if server.name == "github": tool1 = MagicMock() @@ -376,6 +377,7 @@ class TestMCPServerManager: mcp_auth_header=None, mcp_protocol_version=None, raw_headers=None, + user_api_key_auth=None, ): assert mcp_auth_header == "legacy-token" # Should use legacy header tool = MagicMock() @@ -414,6 +416,7 @@ class TestMCPServerManager: mcp_auth_header=None, mcp_protocol_version=None, raw_headers=None, + user_api_key_auth=None, ): assert ( mcp_auth_header == "server-specific-token" @@ -1004,6 +1007,7 @@ class TestMCPServerManager: mcp_auth_header=None, mcp_protocol_version=None, raw_headers=None, + user_api_key_auth=None, ): assert ( mcp_auth_header == "server-specific-token" @@ -1801,6 +1805,258 @@ class TestMCPServerManager: assert len(tools_unprefixed) == 1 assert tools_unprefixed[0].name == "send_email" + @pytest.mark.asyncio + async def test_get_tools_from_server_jwt_skipped_when_mcp_auth_header_set(self): + """When a per-user mcp_auth_header is resolved, JWT injection must be skipped. + + MCPClient._get_auth_headers() applies extra_headers AFTER writing + Authorization from auth_value, so an injected JWT would clobber the + user's per-server OAuth token. Regression test for that interaction. + """ + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="zapier", + name="zapier", + transport=MCPTransport.http, + ) + + manager._create_mcp_client = AsyncMock(return_value=object()) + manager._fetch_tools_with_timeout = AsyncMock(return_value=[]) + + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + with ( + patch( + "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.get_mcp_jwt_signer", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.inject_mcp_jwt_headers_for_upstream", + new=AsyncMock(return_value={"Authorization": "Bearer signed-jwt"}), + ) as mock_inject, + ): + # Case A: mcp_auth_header present -> JWT must NOT be injected + await manager._get_tools_from_server( + server, + mcp_auth_header="oauth-user-token", + user_api_key_auth=user_auth, + ) + mock_inject.assert_not_called() + + # Case B: no mcp_auth_header -> JWT injection runs as before + await manager._get_tools_from_server( + server, + user_api_key_auth=user_auth, + ) + mock_inject.assert_awaited_once() + + def test_resolve_mcp_server_for_tool_call_via_prefixed_name(self): + """Resolution succeeds when the prefixed tool name is in the mapping.""" + manager = MCPServerManager() + server = MCPServer( + server_id="jira", + name="jira", + transport=MCPTransport.http, + ) + manager.registry = {"jira": server} + manager.tool_name_to_mcp_server_name_mapping["jira-search_issues"] = "jira" + manager.tool_name_to_mcp_server_name_mapping["search_issues"] = "jira" + + resolved = manager._resolve_mcp_server_for_tool_call("jira", "search_issues") + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_via_alias(self): + """Resolution falls back to alias/server_name match in the registry.""" + manager = MCPServerManager() + server = MCPServer( + server_id="srv-uuid-123", + name="zapier", + alias="zapier-alias", + transport=MCPTransport.http, + ) + manager.registry = {"srv-uuid-123": server} + manager.tool_name_to_mcp_server_name_mapping["create_zap"] = "zapier" + + resolved = manager._resolve_mcp_server_for_tool_call( + "zapier-alias", "create_zap" + ) + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_unknown_tool_with_empty_mapping(self): + """Server-name match alone must not let unknown tools through when the + mapping has no entries for that server (e.g. listing has not completed + or the server is OAuth2 and the user has not yet listed tools). + """ + manager = MCPServerManager() + server = MCPServer( + server_id="srv-uuid-123", + name="zapier", + alias="zapier-alias", + transport=MCPTransport.http, + ) + manager.registry = {"srv-uuid-123": server} + + with pytest.raises(ValueError, match="Tool create_zap not found"): + manager._resolve_mcp_server_for_tool_call("zapier-alias", "create_zap") + + def test_resolve_mcp_server_for_tool_call_fallback_to_unprefixed_lookup(self): + """Fallback to unprefixed _get_mcp_server_from_tool_name when other paths fail.""" + manager = MCPServerManager() + server = MCPServer( + server_id="linear", + name="linear", + transport=MCPTransport.http, + ) + manager.registry = {"linear": server} + manager.tool_name_to_mcp_server_name_mapping["create_issue"] = "linear" + + # server_name is empty so the fallback unprefixed lookup runs and matches. + resolved = manager._resolve_mcp_server_for_tool_call("", "create_issue") + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_raises_when_not_found(self): + """ValueError is raised when no resolution path finds the tool.""" + manager = MCPServerManager() + with pytest.raises(ValueError, match="Tool .* not found"): + manager._resolve_mcp_server_for_tool_call("nonexistent", "ghost_tool") + + def test_resolve_mcp_server_for_tool_call_unknown_tool_with_known_server(self): + """Server-name match alone must not let unknown tools slip through. + + If the registry has tools for this server but neither the prefixed nor + unprefixed tool name is in the mapping, raise rather than returning the + server (would otherwise allow tool enumeration via name spoofing). + """ + manager = MCPServerManager() + server = MCPServer( + server_id="github", + name="github", + transport=MCPTransport.http, + ) + manager.registry = {"github": server} + # Mapping has *some* tools for github but not "missing_tool". + manager.tool_name_to_mcp_server_name_mapping["github-list_repos"] = "github" + manager.tool_name_to_mcp_server_name_mapping["list_repos"] = "github" + + with pytest.raises(ValueError, match="Tool missing_tool not found"): + manager._resolve_mcp_server_for_tool_call("github", "missing_tool") + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_skipped_when_not_user_oauth(self): + """Returns input headers unchanged when server does not need user OAuth.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="plain", + name="plain", + transport=MCPTransport.http, + ) + # needs_user_oauth_token defaults to False. + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="bob") + + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_returns_client_supplied_token(self): + """Returns the client's oauth2_headers as-is when already set.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + assert server.needs_user_oauth_token is True + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + supplied = {"Authorization": "Bearer client-supplied"} + + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=supplied, user_api_key_auth=user_auth + ) + assert result is supplied + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_looks_up_stored_token(self): + """Falls back to stored per-user OAuth headers when no token is supplied.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + stored = {"Authorization": "Bearer stored-user-token"} + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(return_value=stored), + ) as mock_lookup: + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + + assert result == stored + mock_lookup.assert_awaited_once() + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_swallows_lookup_exception(self): + """Returns supplied headers (None) when the stored-token lookup raises.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(side_effect=RuntimeError("redis down")), + ): + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_no_user_id(self): + """Skip lookup entirely when user_api_key_auth has no user_id.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + # user_id is None -> lookup must not happen + user_auth = UserAPIKeyAuth(api_key="sk-test") + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(return_value={"Authorization": "Bearer x"}), + ) as mock_lookup: + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + mock_lookup.assert_not_called() + def test_create_prefixed_tools_updates_mapping_for_both_forms(self): """_create_prefixed_tools should populate mapping for prefixed and original names even when not adding prefix in output.""" manager = MCPServerManager() 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 f4feac68fcc..593facd9279 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,5 +1,6 @@ import json from typing import Any, Dict, Optional +from unittest.mock import MagicMock import pytest from fastapi import HTTPException @@ -796,6 +797,25 @@ class TestCallToolRestAPI: raising=False, ) + mock_server = MagicMock() + mock_server.server_id = "server-1" + + def fake_get_mcp_server_by_id(server_id): + return mock_server if server_id == "server-1" else None + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + fake_get_mcp_server_by_id, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_name", + lambda *args, **kwargs: None, + raising=False, + ) + request_payload = { "server_id": "server-1", "name": "demo-tool", diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 26f04a4abcb..116ba83f42e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -127,6 +127,23 @@ def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_v assert expires <= now + timedelta(minutes=10, seconds=2) +def test_get_cli_jwt_auth_token_includes_team_alias(valid_sso_user_defined_values): + token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( + valid_sso_user_defined_values, + team_id="team-123", + team_alias="test-team", + ) + + decrypted_token = decrypt_value_helper( + token, key="ui_hash_key", exception_type="debug" + ) + assert decrypted_token is not None + token_data = json.loads(decrypted_token) + + assert token_data["team_id"] == "team-123" + assert token_data["team_alias"] == "test-team" + + def test_get_experimental_ui_login_jwt_auth_token_uses_10_min_expiry( valid_sso_user_defined_values, ): @@ -3016,3 +3033,340 @@ async def test_team_member_budget_check_zero_per_member_row_still_blocks(): proxy_logging_obj=proxy_logging_obj, ) assert exc_info.value.max_budget == 0.0 + + +# --- resolve_and_validate_end_user_id --------------------------------------- + + +@pytest.fixture +def _validate_flag_on(monkeypatch): + """Enable opt-in DB validation for the duration of a test.""" + import litellm + + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", True) + monkeypatch.setattr(litellm, "max_end_user_budget_id", None) + + +def _validation_cache(): + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + return cache + + +def _patch_validation_helpers(monkeypatch, *, end_user=None, user=None, fuzzy=None): + """Stub out the DB helpers resolve_and_validate_end_user_id delegates to.""" + from litellm.proxy.auth import auth_checks + + monkeypatch.setattr( + auth_checks, "get_end_user_object", AsyncMock(return_value=end_user) + ) + monkeypatch.setattr(auth_checks, "get_user_object", AsyncMock(return_value=user)) + monkeypatch.setattr( + auth_checks, "_get_fuzzy_user_object", AsyncMock(return_value=fuzzy) + ) + + +@pytest.mark.asyncio +async def test_resolve_end_user_returns_none_for_none_input( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + assert ( + await resolve_and_validate_end_user_id( + raw_end_user_id=None, + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + is None + ) + + +@pytest.mark.asyncio +async def test_resolve_end_user_passes_through_when_flag_disabled(monkeypatch): + """Default behaviour: flag is off, arbitrary ids pass through untouched.""" + import litellm + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr(litellm, "validate_end_user_id_in_db", False) + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="codex-session-abc", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "codex-session-abc" + cache.async_set_cache.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_passes_through_when_no_prisma_client( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="alice@example.com", + prisma_client=None, + user_api_key_cache=cache, + ) + assert result == "alice@example.com" + + +@pytest.mark.asyncio +async def test_resolve_end_user_matches_end_user_table(_validate_flag_on, monkeypatch): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, end_user=MagicMock()) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="customer-123", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "customer-123" + cache.async_set_cache.assert_awaited_once() + kwargs = cache.async_set_cache.await_args.kwargs + assert kwargs["key"] == "end_user_validation:customer-123" + assert kwargs["value"] == "valid" + + +@pytest.mark.asyncio +async def test_resolve_end_user_matches_user_table_by_user_id( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, user=MagicMock()) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="user-xyz", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "user-xyz" + # email fallback should not run for a non-email input + auth_checks._get_fuzzy_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_matches_user_table_by_email( + _validate_flag_on, monkeypatch +): + """Email-shaped ids route through get_user_object with user_email set. + + The fuzzy lookup must happen inside get_user_object so it shares the + _should_check_db throttle and user_api_key_cache — no direct raw + Prisma calls on the auth path. + """ + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, user=MagicMock()) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="Alice@Example.com", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "Alice@Example.com" + auth_checks.get_user_object.assert_awaited_once() + user_kwargs = auth_checks.get_user_object.await_args.kwargs + assert user_kwargs["user_id"] == "Alice@Example.com" + assert user_kwargs["user_email"] == "Alice@Example.com" + # email branch must not bypass the cached helper with a raw fuzzy call + auth_checks._get_fuzzy_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_non_email_id_does_not_pass_user_email( + _validate_flag_on, monkeypatch +): + """Non-email ids skip the email fuzzy path to avoid a pointless DB hit.""" + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, user=MagicMock()) + cache = _validation_cache() + + await resolve_and_validate_end_user_id( + raw_end_user_id="user-xyz", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + auth_checks.get_user_object.assert_awaited_once() + user_kwargs = auth_checks.get_user_object.await_args.kwargs + assert user_kwargs["user_email"] is None + + +@pytest.mark.asyncio +async def test_resolve_end_user_drops_codex_opaque_identifier( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) # all helpers return None + cache = _validation_cache() + + codex_id = ( + "user_8a4a360c36621665b341e06fb76041d9b6def732bb183eea148d4abc9d97c1de" + "_account__session_a2bce4a5-8887-44ef-b491-fbf0a55c6569" + ) + result = await resolve_and_validate_end_user_id( + raw_end_user_id=codex_id, + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result is None + cache.async_set_cache.assert_awaited_once() + kwargs = cache.async_set_cache.await_args.kwargs + assert kwargs["value"] == "invalid" + + +@pytest.mark.asyncio +async def test_resolve_end_user_preserves_id_when_default_budget_configured( + _validate_flag_on, monkeypatch +): + """Don't drop unregistered ids when litellm.max_end_user_budget_id is set. + + The default end-user budget is applied downstream when the id is present + but not found in the db — dropping the id here would bypass those limits. + """ + import litellm + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr(litellm, "max_end_user_budget_id", "default-budget") + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="new-customer", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "new-customer" + + +@pytest.mark.asyncio +async def test_resolve_end_user_drops_unknown_email(_validate_flag_on, monkeypatch): + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="stranger@example.com", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result is None + + +@pytest.mark.asyncio +async def test_resolve_end_user_uses_cached_valid_result( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch) + cache = _validation_cache() + cache.async_get_cache = AsyncMock(return_value="valid") + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="alice@example.com", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result == "alice@example.com" + auth_checks.get_end_user_object.assert_not_awaited() + auth_checks.get_user_object.assert_not_awaited() + auth_checks._get_fuzzy_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_uses_cached_invalid_result( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + _patch_validation_helpers(monkeypatch, end_user=MagicMock()) + cache = _validation_cache() + cache.async_get_cache = AsyncMock(return_value="invalid") + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="bogus", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + assert result is None + # Despite a matching row configured, helpers aren't called — cache wins. + auth_checks.get_end_user_object.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_resolve_end_user_swallows_db_errors_and_returns_none( + _validate_flag_on, monkeypatch +): + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr( + auth_checks, + "get_end_user_object", + AsyncMock(side_effect=Exception("db down")), + ) + monkeypatch.setattr( + auth_checks, + "get_user_object", + AsyncMock(side_effect=Exception("db down")), + ) + cache = _validation_cache() + + result = await resolve_and_validate_end_user_id( + raw_end_user_id="alice@example.com", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + # DB errors shouldn't raise through the auth path — treat as unknown. + assert result is None + + +@pytest.mark.asyncio +async def test_resolve_end_user_reraises_budget_exceeded( + _validate_flag_on, monkeypatch +): + """BudgetExceededError from get_end_user_object must bubble up so the + auth path enforces spend limits instead of silently dropping the id.""" + import litellm + from litellm.proxy.auth import auth_checks + from litellm.proxy.auth.auth_checks import resolve_and_validate_end_user_id + + monkeypatch.setattr( + auth_checks, + "get_end_user_object", + AsyncMock( + side_effect=litellm.BudgetExceededError(current_cost=10.0, max_budget=5.0) + ), + ) + cache = _validation_cache() + + with pytest.raises(litellm.BudgetExceededError): + await resolve_and_validate_end_user_id( + raw_end_user_id="customer-over-budget", + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 08035fb7173..68e1636d380 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -597,6 +597,315 @@ def test_get_end_user_id_falls_back_to_deprecated_user_header_name(): assert result == "user-legacy" +class TestCoerceUserIdToStr: + """Unit tests for the _coerce_user_id_to_str helper.""" + + def test_plain_string_is_returned_verbatim(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str("alice@example.com") == "alice@example.com" + + def test_string_is_stripped(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str(" bob ") == "bob" + + def test_codex_opaque_identifier_is_preserved(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + codex_id = ( + "user_8a4a360c36621665b341e06fb76041d9b6def732bb183eea148d4abc9d97c1de" + "_account__session_a2bce4a5-8887-44ef-b491-fbf0a55c6569" + ) + assert _coerce_user_id_to_str(codex_id) == codex_id + + def test_none_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str(None) is None + + def test_empty_string_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str("") is None + assert _coerce_user_id_to_str(" ") is None + + def test_dict_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + payload = { + "device_id": "abc", + "account_uuid": "", + "session_id": "c284b8cb", + } + assert _coerce_user_id_to_str(payload) is None + + def test_list_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str(["a", "b"]) is None + + def test_json_encoded_dict_string_passes_through_by_default(self): + """JSON-encoded dict strings are preserved unless opt-in flag is on. + + This preserves backwards compatibility: existing deployments that + intentionally pass JSON-encoded user identifiers keep working. + """ + import litellm + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + blob = ( + '{"device_id":"d5abe9199ee7759a0558974e9371e78c7b38d7621aae26d6609c1de61af6afb0",' + '"account_uuid":"","session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + ) + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = False + try: + assert _coerce_user_id_to_str(blob) == blob + finally: + litellm.validate_end_user_id_in_db = original + + def test_json_encoded_dict_string_returns_none_when_validation_enabled(self): + import litellm + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + # Same broken shape we saw in spend logs, but pre-stringified to JSON. + blob = ( + '{"device_id":"d5abe9199ee7759a0558974e9371e78c7b38d7621aae26d6609c1de61af6afb0",' + '"account_uuid":"","session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + ) + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = True + try: + assert _coerce_user_id_to_str(blob) is None + finally: + litellm.validate_end_user_id_in_db = original + + def test_json_encoded_list_string_passes_through_by_default(self): + import litellm + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = False + try: + assert _coerce_user_id_to_str('["a","b"]') == '["a","b"]' + finally: + litellm.validate_end_user_id_in_db = original + + def test_json_encoded_list_string_returns_none_when_validation_enabled(self): + import litellm + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = True + try: + assert _coerce_user_id_to_str('["a","b"]') is None + finally: + litellm.validate_end_user_id_in_db = original + + def test_int_returns_str(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str(12345) == "12345" + + def test_bool_returns_none(self): + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + # bool is an int subclass — reject explicitly, never produce "True"/"False". + assert _coerce_user_id_to_str(True) is None + assert _coerce_user_id_to_str(False) is None + + def test_brace_string_that_isnt_json_is_kept(self): + """A string starting with `{` but failing to parse stays as-is.""" + from litellm.proxy.auth.auth_utils import _coerce_user_id_to_str + + assert _coerce_user_id_to_str("{not json") == "{not json" + + +class TestGetEndUserIdDropsMalformedBodyValues: + """Tests that get_end_user_id_from_request_body drops dict-shaped values + rather than stringifying them into spend logs.""" + + def test_dict_user_falls_through_to_litellm_metadata(self): + request_body = { + "user": { + "device_id": "abc", + "session_id": "c284b8cb", + }, + "litellm_metadata": {"user": "alice@example.com"}, + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "alice@example.com" + + def test_dict_user_with_no_other_sources_returns_none(self): + request_body = { + "user": {"device_id": "abc", "session_id": "xyz"}, + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result is None + + def test_json_encoded_user_string_passes_through_by_default(self): + """JSON-encoded user strings pass through unless validation is opted in. + + Gating behind ``litellm.validate_end_user_id_in_db`` keeps existing + deployments that send JSON-encoded identifiers working until they + explicitly opt into the stricter extraction. + """ + import litellm + + blob = ( + '{"device_id":"d5abe9199ee7759a","account_uuid":"",' + '"session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + ) + request_body = {"user": blob} + + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = False + try: + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + finally: + litellm.validate_end_user_id_in_db = original + + assert result == blob + + def test_json_encoded_user_string_returns_none_when_validation_enabled(self): + import litellm + + request_body = { + "user": ( + '{"device_id":"d5abe9199ee7759a","account_uuid":"",' + '"session_id":"c284b8cb-a050-4278-8599-cc4e016a10ab"}' + ), + } + + original = litellm.validate_end_user_id_in_db + litellm.validate_end_user_id_in_db = True + try: + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + finally: + litellm.validate_end_user_id_in_db = original + + assert result is None + + def test_plain_string_user_is_preserved(self): + request_body = {"user": "alice@example.com"} + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "alice@example.com" + + def test_codex_opaque_user_is_preserved(self): + codex_id = ( + "user_8a4a360c36621665b341e06fb76041d9b6def732bb183eea148d4abc9d97c1de" + "_account__session_a2bce4a5-8887-44ef-b491-fbf0a55c6569" + ) + request_body = {"user": codex_id} + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == codex_id + + def test_int_user_is_coerced_to_string(self): + request_body = {"user": 12345} + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "12345" + + def test_list_user_falls_through(self): + request_body = { + "user": ["a", "b"], + "safety_identifier": "alice@example.com", + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "alice@example.com" + + def test_dict_safety_identifier_returns_none(self): + request_body = { + "safety_identifier": {"device_id": "abc"}, + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result is None + + def test_dict_metadata_user_id_returns_none(self): + request_body = { + "metadata": {"user_id": {"device_id": "abc"}}, + } + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result is None + + def test_whitespace_user_falls_through(self): + request_body = {"user": " ", "safety_identifier": "alice@example.com"} + + with patch("litellm.proxy.proxy_server.general_settings", {}): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers={} + ) + + assert result == "alice@example.com" + + def test_dict_user_header_falls_through_to_body(self): + """A dict-shaped value in a configured user-id header is dropped, not stringified.""" + general_settings = {"user_header_name": "x-custom-user-id"} + # A header value will normally be a str, but be defensive: the coercion + # must drop anything that isn't a usable identifier. + headers = {"x-custom-user-id": {"device_id": "abc"}} + request_body = {"user": "alice@example.com"} + + with ( + patch( + "litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", + return_value=None, + ), + patch("litellm.proxy.proxy_server.general_settings", general_settings), + ): + result = get_end_user_id_from_request_body( + request_body=request_body, request_headers=headers + ) + + assert result == "alice@example.com" + + def _make_deployment_dict( model_name: str, tpm: Optional[int] = None, rpm: Optional[int] = None ) -> dict: diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 77aa03032a7..f38ac5c2000 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -249,3 +249,241 @@ def test_get_complete_model_list_byok_wildcard_expansion(): assert len(result) > 0 assert all(m.startswith("openai/") for m in result) assert "openai/*" not in result + + +def test_get_complete_model_list_expands_team_scoped_wildcard_with_stored_credential( + monkeypatch, +): + """ + Team-scoped BYOK wildcard deployments are stored under an internal model_name, + with the public wildcard name in model_info.team_public_model_name. + """ + import litellm + from litellm import Router + from litellm.proxy.auth import model_checks + from litellm.proxy.auth.model_checks import get_complete_model_list + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="openai-credential", + credential_info={"provider": "openai"}, + credential_values={ + "api_key": "stored-openai-key", + "api_base": "https://example.openai.test/v1", + }, + ) + ], + ) + + captured_params = {} + + def fake_get_provider_models(provider, litellm_params=None): + captured_params["provider"] = provider + captured_params["api_key"] = litellm_params.api_key + captured_params["api_base"] = litellm_params.api_base + captured_params["credential_name"] = litellm_params.litellm_credential_name + return ["gpt-4o"] + + monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) + + router = Router( + model_list=[ + { + "model_name": "model_name_team-1_generated", + "litellm_params": { + "model": "openai/*", + "custom_llm_provider": "openai", + "litellm_credential_name": "openai-credential", + }, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "openai/*", + }, + } + ] + ) + + result = get_complete_model_list( + key_models=[], + team_models=["openai/*"], + proxy_model_list=[], + user_model=None, + infer_model_from_keys=False, + llm_router=router, + team_id="team-1", + ) + + assert "openai/gpt-4o" in result + assert captured_params == { + "provider": "openai", + "api_key": "stored-openai-key", + "api_base": "https://example.openai.test/v1", + "credential_name": None, + } + + +def test_wildcard_credential_hydration_preserves_deployment_params( + monkeypatch, +): + import litellm + from litellm.proxy.auth import model_checks + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="openai-credential", + credential_info={"provider": "openai"}, + credential_values={ + "api_key": "stored-openai-key", + "api_version": "credential-version", + "model": "openai/wrong-model", + "unexpected_field": "unexpected-value", + }, + ) + ], + ) + + captured_params = {} + + def fake_get_provider_models(provider, litellm_params=None): + captured_params["provider"] = provider + captured_params["model"] = litellm_params.model + captured_params["api_key"] = litellm_params.api_key + captured_params["api_version"] = litellm_params.api_version + captured_params["credential_name"] = litellm_params.litellm_credential_name + captured_params["has_unexpected_field"] = hasattr( + litellm_params, "unexpected_field" + ) + return ["gpt-4o"] + + monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) + + result = get_known_models_from_wildcard( + wildcard_model="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + custom_llm_provider="openai", + api_version="deployment-version", + litellm_credential_name="openai-credential", + ), + ) + + assert result == ["openai/gpt-4o"] + assert captured_params == { + "provider": "openai", + "model": "openai/*", + "api_key": "stored-openai-key", + "api_version": "deployment-version", + "credential_name": None, + "has_unexpected_field": False, + } + + +def test_wildcard_credential_hydration_preserves_missing_credential_name( + monkeypatch, +): + import litellm + from litellm.proxy.auth import model_checks + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "credential_list", []) + + captured_params = {} + + def fake_get_provider_models(provider, litellm_params=None): + captured_params["provider"] = provider + captured_params["api_key"] = litellm_params.api_key + captured_params["credential_name"] = litellm_params.litellm_credential_name + return ["gpt-4o"] + + monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) + + result = get_known_models_from_wildcard( + wildcard_model="openai/*", + litellm_params=LiteLLM_Params( + model="openai/*", + custom_llm_provider="openai", + api_key=None, + litellm_credential_name="missing-credential", + ), + ) + + assert result == ["openai/gpt-4o"] + assert captured_params == { + "provider": "openai", + "api_key": None, + "credential_name": "missing-credential", + } + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_expands_query_team_wildcard( + monkeypatch, +): + import litellm + from litellm import Router + from litellm.proxy.auth import model_checks + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import get_available_models_for_user + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="openai-credential", + credential_info={"provider": "openai"}, + credential_values={"api_key": "stored-openai-key"}, + ) + ], + ) + + def fake_get_provider_models(provider, litellm_params=None): + assert litellm_params.api_key == "stored-openai-key" + assert litellm_params.litellm_credential_name is None + return ["gpt-4o-mini"] + + monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) + + router = Router( + model_list=[ + { + "model_name": "model_name_team-1_generated", + "litellm_params": { + "model": "openai/*", + "custom_llm_provider": "openai", + "litellm_credential_name": "openai-credential", + }, + "model_info": { + "team_id": "team-1", + "team_public_model_name": "openai/*", + }, + } + ] + ) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test", + models=[], + team_id="team-1", + team_models=["openai/*"], + ), + llm_router=router, + general_settings={}, + user_model=None, + team_id="team-1", + ) + + assert "openai/gpt-4o-mini" in result diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 442625c75a7..defd3bbcdcd 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3335,3 +3335,125 @@ async def test_master_key_auth_substitutes_alias_for_api_key(): finally: for k, v in _orig.items(): setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_user_api_key_auth_sets_end_user_id_when_builder_skips_it(): + """Defense-in-depth: ``_user_api_key_auth_builder`` has multiple + early-return paths (master_key=None, /user/auth route, JWT + short-circuits) that bypass the end-user resolution block. The wrapper + must still attribute spend logs to the request-supplied end-user when + none of those paths set it. + + Krrish flagged the removal of this fallback as a regression risk; this + test pins the behaviour so future refactors don't silently drop it. + """ + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + builder_token = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + # builder did NOT set end_user_id (e.g. master_key=None early return) + assert builder_token.end_user_id is None + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = json.dumps( + {"model": "gpt-4o", "user": "alice@example.com"} + ).encode() + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + # Stub the builder so the test doesn't have to traverse the full + # auth state machine; we only care about the wrapper's safety net. + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + return_value=builder_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._run_centralized_common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + ): + result = await user_api_key_auth(request=request, api_key="Bearer sk-test") + + # Validation flag is False by default → pass-through, raw value lands + # on the auth obj instead of being silently dropped. + assert result.end_user_id == "alice@example.com" + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder(): + """When the builder already resolved the end-user id (the primary + path), the wrapper-level safety net must not run a second resolution + pass — that would re-extract from the request body and could + overwrite a value the builder explicitly chose to set.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + builder_token = UserAPIKeyAuth( + api_key="sk-test", user_id="u1", end_user_id="builder-resolved-id" + ) + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = json.dumps( + {"model": "gpt-4o", "user": "different-id-from-body"} + ).encode() + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + return_value=builder_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._run_centralized_common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + patch( + "litellm.proxy.auth.user_api_key_auth.resolve_and_validate_end_user_id", + new_callable=AsyncMock, + ) as mock_resolve, + ): + result = await user_api_key_auth(request=request, api_key="Bearer sk-test") + + assert result.end_user_id == "builder-resolved-id" + mock_resolve.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index b4343f6b2e1..3d7cb1e35f3 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -16,6 +16,7 @@ sys.path.insert( import litellm from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.http_parsing_utils import ( + _is_form_content_type, _read_request_body, _safe_get_request_headers, _safe_get_request_parsed_body, @@ -853,3 +854,145 @@ class TestGetTagsFromRequestBodyStringCoerce: tags = get_tags_from_request_body({"metadata": {"tags": ["x"]}}) assert tags == ["x"] + + +class TestIsFormContentType: + @pytest.mark.parametrize( + "content_type", + [ + "application/x-www-form-urlencoded", + "multipart/form-data", + "multipart/form-data; boundary=----WebKitFormBoundary", + "Application/X-WWW-Form-Urlencoded", + " multipart/form-data ", + "application/x-www-form-urlencoded; charset=utf-8", + ], + ) + def test_form_types_match(self, content_type): + assert _is_form_content_type(content_type) is True + + @pytest.mark.parametrize( + "content_type", + [ + "", + "application/json", + "application/json; charset=utf-8", + "application/form-json", + "multiform/anything", + "application/json; xform=1", + "application/xml-with-form-data-but-not-actually", + "text/plain", + "form", + ], + ) + def test_non_form_types_rejected(self, content_type): + assert _is_form_content_type(content_type) is False + + +class TestReadRequestBodyNonCanonicalContentType: + """A JSON body with a ``"form"``-substring Content-Type must parse as JSON.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "content_type", + [ + "application/form-json", + "application/json; xform=1", + "multiform/anything", + ], + ) + async def test_json_body_with_formlike_content_type_parses_as_json( + self, content_type + ): + payload = {"user_config": {"model_list": []}, "model": "x"} + + mock_request = MagicMock() + mock_request.body = AsyncMock(return_value=orjson.dumps(payload)) + mock_request.form = AsyncMock(return_value={}) + mock_request.headers = {"content-type": content_type} + mock_request.scope = {} + + result = await _read_request_body(mock_request) + assert result == payload + mock_request.form.assert_not_called() + + @pytest.mark.asyncio + async def test_real_form_post_still_parsed_as_form(self): + mock_request = MagicMock() + mock_request.form = AsyncMock(return_value={"k": "v"}) + mock_request.body = AsyncMock(return_value=b"") + mock_request.headers = {"content-type": "application/x-www-form-urlencoded"} + mock_request.scope = {} + + result = await _read_request_body(mock_request) + assert result == {"k": "v"} + mock_request.form.assert_awaited_once() + + +class TestReadRequestBodyFormParseFailure: + """ + A failed ``request.form()`` parse (e.g. multipart with missing boundary) + must surface as a 400, not silently return ``{}`` — otherwise the + auth-time pre-read sees an empty body while a later raw-body re-read + sees the original payload, defeating every banned-param check. + """ + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "raised_exception", + [ + ValueError("Missing boundary in multipart."), + AssertionError("malformed chunk"), + RuntimeError("form parser exploded"), + ], + ) + async def test_form_parse_failure_raises_400(self, raised_exception): + mock_request = MagicMock() + mock_request.form = AsyncMock(side_effect=raised_exception) + mock_request.headers = {"content-type": "multipart/form-data"} + mock_request.scope = {} + + with pytest.raises(ProxyException) as exc_info: + await _read_request_body(mock_request) + assert str(exc_info.value.code) == "400" + + +class TestGetRequestBody: + @pytest.mark.asyncio + async def test_json_with_charset_param_parses_as_json(self): + payload = {"k": "v"} + mock_request = MagicMock() + mock_request.method = "POST" + mock_request.body = AsyncMock(return_value=orjson.dumps(payload)) + mock_request.headers = {"content-type": "application/json; charset=utf-8"} + mock_request.scope = {} + + result = await get_request_body(mock_request) + assert result == payload + + @pytest.mark.asyncio + async def test_form_post_routes_to_form_data(self): + mock_request = MagicMock() + mock_request.method = "POST" + mock_request.headers = {"content-type": "multipart/form-data; boundary=x"} + mock_request.form = AsyncMock(return_value={"k": "v"}) + mock_request.scope = {} + + result = await get_request_body(mock_request) + assert result == {"k": "v"} + + @pytest.mark.asyncio + async def test_substring_match_no_longer_accepted(self): + mock_request = MagicMock() + mock_request.method = "POST" + mock_request.headers = {"content-type": "application/form-json"} + mock_request.scope = {} + + with pytest.raises(ValueError, match="Unsupported content type"): + await get_request_body(mock_request) + + @pytest.mark.asyncio + async def test_non_post_returns_empty(self): + mock_request = MagicMock() + mock_request.method = "GET" + assert await get_request_body(mock_request) == {} diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index a3247d2e557..71178c4826c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2073,6 +2073,226 @@ def test_get_http_exception_includes_assessments_and_identifier(): assert exc.detail["assessments"][0]["matches"][0]["match"] == "[REDACTED]" +def test_extract_violation_category_names_mixed_policies(): + """Topic names, content-filter types, PII types, and managed-word types + flatten into a single category-name list — using only the operator- + defined `name`/`type` labels.""" + g = _make_guardrail() + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": [ + {"name": "Fiduciary Advice", "action": "BLOCKED"}, + {"name": "Tax Advice", "action": "BLOCKED"}, + ] + }, + "contentPolicy": { + "filters": [{"type": "VIOLENCE", "action": "BLOCKED"}] + }, + "wordPolicy": { + "managedWordLists": [{"type": "PROFANITY", "action": "BLOCKED"}], + }, + "sensitiveInformationPolicy": { + "piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}] + }, + } + ], + } + names = g._extract_violation_category_names(response) + assert "Fiduciary Advice" in names + assert "Tax Advice" in names + assert "VIOLENCE" in names + assert "PROFANITY" in names + assert "EMAIL" in names + + +def test_extract_violation_category_names_does_not_leak_user_input(): + """SECURITY: customWords.match is the raw user-submitted word that + triggered the rule, and an unnamed regex match is the actual sensitive + value (e.g. a credit-card number). Neither must appear in + violation_categories — otherwise the content the guardrail blocked + leaks straight into telemetry backends.""" + g = _make_guardrail() + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "customWords": [ + {"match": "secret-codeword-abc-123", "action": "BLOCKED"} + ], + }, + "sensitiveInformationPolicy": { + "regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}] + }, + } + ], + } + names = g._extract_violation_category_names(response) + assert "secret-codeword-abc-123" not in names + assert "4111-1111-1111-1111" not in names + assert names == [] + + +def test_extract_violation_category_names_named_regex_uses_name(): + """A regex with a `name` field surfaces that operator-defined label + (safe to log), not the matched value.""" + g = _make_guardrail() + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "regexes": [ + { + "name": "credit-card-pattern", + "match": "4111-1111-1111-1111", + "action": "BLOCKED", + } + ] + } + } + ], + } + names = g._extract_violation_category_names(response) + assert names == ["credit-card-pattern"] + + +def test_extract_violation_category_names_skips_anonymized(): + """ANONYMIZED entries are not blocks — they must not contribute to the + violation_categories list.""" + g = _make_guardrail() + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}] + } + } + ], + } + assert g._extract_violation_category_names(response) == [] + + +def test_extract_violation_category_names_no_assessments(): + """Empty / missing assessments → empty list, not an error.""" + g = _make_guardrail() + assert g._extract_violation_category_names({"action": "NONE"}) == [] + assert g._extract_violation_category_names({"assessments": None}) == [] + + +@pytest.mark.asyncio +async def test_make_bedrock_api_request_forwards_guardrail_action(): + """Bedrock's top-level ``action`` string must be propagated through + ``tracing_detail`` so downstream loggers (OTEL, ...) can surface the + raw provider verdict as a queryable attribute without re-parsing the + redacted guardrail_response blob.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}] + } + } + ], + } + + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + } + + with ( + patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, + patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + patch.object( + guardrail, + "_get_http_exception_for_blocked_guardrail", + return_value=Exception("blocked"), + ), + ): + mock_post.return_value = mock_bedrock_response + + with pytest.raises(Exception): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data["messages"], + request_data=request_data, + ) + + tracing_detail = mock_log.call_args.kwargs["tracing_detail"] + assert tracing_detail is not None + assert tracing_detail["guardrail_action"] == "GUARDRAIL_INTERVENED" + + +@pytest.mark.asyncio +async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): + """If the Bedrock response omits ``action`` (older / partial payloads), + the field must be left off ``tracing_detail`` rather than written as + ``None`` — downstream code expects strings or absence, not nulls.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = {"assessments": []} + + with ( + patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, + patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.return_value = mock_bedrock_response + + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hi"}], + request_data={"model": "gpt-4o", "messages": []}, + ) + + tracing_detail = mock_log.call_args.kwargs["tracing_detail"] + # No violation categories and no action ⇒ tracing_detail stays None + # (the hook collapses an empty dict before forwarding). + if tracing_detail is not None: + assert "guardrail_action" not in tracing_detail + + def test_get_http_exception_no_blocked_assessments_omits_field(): """L3: when no assessments are blocked, the `assessments` key is omitted entirely.""" g = _make_guardrail() diff --git a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py index b17b3270787..cb2276ab39d 100644 --- a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py +++ b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py @@ -219,7 +219,7 @@ def test_build_claims_scope_with_tool(): def test_build_claims_scope_without_tool(): - """_build_claims() includes mcp:tools/list when no specific tool is called.""" + """_build_claims() emits only mcp:tools/list when no specific tool is called.""" signer = _make_signer() user_dict = _make_user_api_key_dict() data: Dict[str, Any] = {} @@ -227,10 +227,11 @@ def test_build_claims_scope_without_tool(): claims = signer._build_claims(user_dict, data) scopes = set(claims["scope"].split()) - assert "mcp:tools/call" in scopes assert "mcp:tools/list" in scopes + # List-only JWTs must NOT carry mcp:tools/call — least-privilege + assert "mcp:tools/call" not in scopes # No per-tool call scope when no tool name was given - assert not any(s.endswith(":call") and s != "mcp:tools/call" for s in scopes) + assert not any(s.endswith(":call") for s in scopes) def test_build_claims_act_fallback_to_litellm_proxy(): @@ -338,7 +339,7 @@ async def test_hook_skips_non_mcp_call_types(): user_dict = _make_user_api_key_dict() data = {"messages": [{"role": "user", "content": "hello"}]} - for call_type in ("completion", "acompletion", "embedding", "list_mcp_tools"): + for call_type in ("completion", "acompletion", "embedding"): original_data = {**data} result = await signer.async_pre_call_hook( user_api_key_dict=user_dict, @@ -351,6 +352,33 @@ async def test_hook_skips_non_mcp_call_types(): ), f"extra_headers should not be set for {call_type}" +@pytest.mark.asyncio +async def test_hook_signs_list_mcp_tools(): + """async_pre_call_hook() signs JWT for list_mcp_tools with list scope.""" + signer = _make_signer( + issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300 + ) + user_dict = _make_user_api_key_dict(user_id="alice", team_id="backend") + data = {"mcp_tool_name": "should_be_cleared"} + + result = await signer.async_pre_call_hook( + user_api_key_dict=user_dict, + cache=MagicMock(), + data=data, + call_type="list_mcp_tools", + ) + + assert isinstance(result, dict) + assert "extra_headers" in result + assert result["extra_headers"]["Authorization"].startswith("Bearer ") + token = result["extra_headers"]["Authorization"].removeprefix("Bearer ") + decoded = _decode_unverified(token) + scopes = set(decoded["scope"].split()) + assert "mcp:tools/list" in scopes + # List-only JWTs must NOT carry mcp:tools/call — least-privilege + assert "mcp:tools/call" not in scopes + + @pytest.mark.asyncio async def test_signed_token_is_verifiable(): """The JWT injected by the hook can be verified against the JWKS public key.""" @@ -1128,3 +1156,116 @@ async def test_hook_raises_401_when_jwt_verification_fails(): ) assert exc_info.value.status_code == 401 + + +# --- _build_scope branches: call_mcp_tool with empty tool name, list_mcp_tools --- + + +def test_build_scope_call_type_call_mcp_tool_without_tool_name(): + """call_mcp_tool with empty tool name emits a generic mcp:tools/call only.""" + signer = _make_signer() + scope = signer._build_scope("", call_type="call_mcp_tool") + scopes = set(scope.split()) + assert scopes == {"mcp:tools/call"} + + +def test_build_scope_call_type_list_mcp_tools_only_list(): + """list_mcp_tools (no tool) emits only mcp:tools/list, never tools/call.""" + signer = _make_signer() + scope = signer._build_scope("", call_type="list_mcp_tools") + scopes = set(scope.split()) + assert scopes == {"mcp:tools/list"} + + +def test_build_scope_default_is_list_only_when_no_call_type(): + """No call_type and no tool falls through to tools/list (least-privilege default).""" + signer = _make_signer() + scope = signer._build_scope("") + scopes = set(scope.split()) + assert "mcp:tools/list" in scopes + assert "mcp:tools/call" not in scopes + + +# --- inject_mcp_jwt_headers_for_upstream --- + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_returns_unchanged_when_signer_not_configured(): + """No signer configured -> return a fresh copy of extra_headers untouched.""" + import litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer as mod + from litellm.proxy._types import UserAPIKeyAuth + + mod._mcp_jwt_signer_instance = None + headers = {"X-Trace-Id": "abc"} + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await mod.inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + extra_headers=headers, + ) + assert result == headers + assert result is not headers # must be a copy + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_returns_unchanged_when_user_dict_none(): + """No user_api_key_dict -> short-circuit without invoking the signer.""" + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer() # ensure instance is created + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=None, + extra_headers={"X-Trace-Id": "abc"}, + ) + assert result == {"X-Trace-Id": "abc"} + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_signs_for_list_tools_path(): + """When for_list_tools=True, signer is invoked with list_mcp_tools call_type.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + extra_headers={"X-Trace": "1"}, + raw_headers={"Authorization": "Bearer incoming.opaque.token"}, + for_list_tools=True, + ) + assert result["X-Trace"] == "1" + assert result["Authorization"].startswith("Bearer ") + token = result["Authorization"].removeprefix("Bearer ") + decoded = _decode_unverified(token) + scopes = set(decoded["scope"].split()) + assert scopes == {"mcp:tools/list"} + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_signs_for_tool_call_path(): + """for_list_tools=False with a tool name signs a call_mcp_tool JWT.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + for_list_tools=False, + mcp_tool_name="search_web", + ) + assert result["Authorization"].startswith("Bearer ") + token = result["Authorization"].removeprefix("Bearer ") + decoded = _decode_unverified(token) + scopes = set(decoded["scope"].split()) + assert "mcp:tools/call" in scopes + assert "mcp:tools/search_web:call" in scopes diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 83317157847..c763e9c0e98 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2218,6 +2218,7 @@ class TestCLIKeyRegenerationFlow: # Mock request mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://internal-proxy.local/" # Test data session_key = "cli-session-4567890" @@ -2242,11 +2243,14 @@ class TestCLIKeyRegenerationFlow: "user_code_verified": False, "session_data": None, } - mock_request.url_for.return_value = ( - "https://test.example.com/sso/cli/complete/cli-session-4567890" - ) - with ( + patch.dict( + os.environ, + { + "PROXY_BASE_URL": "https://test.example.com", + "SERVER_ROOT_PATH": "", + }, + ), patch( "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", return_value=mock_user_info, @@ -2290,6 +2294,10 @@ class TestCLIKeyRegenerationFlow: assert result.status_code == 200 # Verify response contains success message (response is HTML) assert result.body is not None + assert ( + 'action="https://test.example.com/sso/cli/complete/cli-session-4567890"' + in result.body.decode() + ) @pytest.mark.asyncio async def test_cli_poll_key_returns_teams_for_selection(self): @@ -2430,6 +2438,7 @@ class TestCLIKeyRegenerationFlow: request=mock_request, key="cli-new-session-key-456", result=mock_result, + received_response=None, ) def test_get_redirect_url_does_not_include_existing_key_in_url(self): @@ -2488,6 +2497,11 @@ class TestCLIKeyRegenerationFlow: "user_id": "test-user-789", "user_role": "internal_user", "teams": ["team-a", "team-b", "team-c"], + "team_details": [ + {"team_id": "team-a", "team_alias": "Team A"}, + {"team_id": "team-b", "team_alias": "Team B"}, + {"team_id": "team-c", "team_alias": "Team C"}, + ], "models": ["gpt-4"], "user_email": "test@example.com", } @@ -2542,6 +2556,7 @@ class TestCLIKeyRegenerationFlow: mock_get_jwt.assert_called_once() jwt_call_args = mock_get_jwt.call_args assert jwt_call_args.kwargs["team_id"] == selected_team + assert jwt_call_args.kwargs["team_alias"] == "Team B" # Verify session was deleted after JWT generation mock_cache.delete_cache.assert_called_once() @@ -5544,6 +5559,289 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch): assert result.extra_fields["another_missing"] is None +class TestCliSsoAttributionMetadata: + """CLI SSO allowlisted OIDC claim persistence and poll exposure.""" + + def test_parse_cli_sso_claim_map(self, monkeypatch): + from litellm.proxy.management_endpoints import ui_sso + + monkeypatch.setattr( + ui_sso, + "CLI_SSO_CLAIM_MAP", + "employment_type->metadata.acme_employment_type, org_info.department -> department", + ) + assert ui_sso._parse_cli_sso_claim_map() == [ + ("employment_type", "acme_employment_type"), + ("org_info.department", "department"), + ] + + def test_build_cli_sso_attribution_metadata_filters_non_scalars(self, monkeypatch): + from litellm.proxy.management_endpoints import ui_sso + from litellm.proxy.management_endpoints.types import CustomOpenID + + monkeypatch.setattr( + ui_sso, + "CLI_SSO_CLAIM_MAP", + "employment_type->acme_employment_type,access_token->should_drop,group->groups", + ) + + result = CustomOpenID( + id="user-1", + email="user@example.com", + display_name="User", + provider="generic", + team_ids=[], + extra_fields={ + "employment_type": "full_time", + "access_token": "eyJhbGciOiJIUzI1NiJ9.payload.signature", + "group": ["team-a", "team-b"], + }, + ) + + metadata = ui_sso.build_cli_sso_attribution_metadata(result=result) + assert metadata == {"acme_employment_type": "full_time"} + + def test_build_cli_sso_attribution_metadata_from_oidc_dict(self, monkeypatch): + from litellm.proxy.management_endpoints import ui_sso + + monkeypatch.setattr( + ui_sso, + "CLI_SSO_CLAIM_MAP", + "org_info.department->department", + ) + + metadata = ui_sso.build_cli_sso_attribution_metadata( + result={ + "sub": "user-1", + "email": "user@example.com", + "org_info": {"department": "Engineering"}, + } + ) + assert metadata == {"department": "Engineering"} + + @pytest.mark.asyncio + async def test_cli_sso_callback_passes_user_defined_values_for_new_users(self): + """First CLI SSO login must supply SSOUserDefinedValues so upsert can create the user.""" + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints import ui_sso + from litellm.proxy.management_endpoints.types import CustomOpenID + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://internal-proxy.local/" + session_key = "cli-session-new-user" + mock_user_info = LiteLLM_UserTable( + user_id="cli-test-user", + user_role="internal_user", + teams=[], + models=[], + ) + mock_sso_result = CustomOpenID( + id="cli-test-user", + email="cli-test@example.com", + display_name="cli-test-user", + provider="generic", + team_ids=[], + ) + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": "poll-secret-hash", + "user_code_hash": "user-code-hash", + "sso_complete": False, + "user_code_verified": False, + "session_data": None, + } + get_user_info_mock = AsyncMock(return_value=mock_user_info) + + with ( + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + get_user_info_mock, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.user_custom_sso", None), + ): + await ui_sso.cli_sso_callback( + request=mock_request, + key=session_key, + result=mock_sso_result, + ) + + get_user_info_mock.assert_awaited_once() + assert get_user_info_mock.call_args.kwargs["user_defined_values"] is not None + assert ( + get_user_info_mock.call_args.kwargs["user_defined_values"]["user_id"] + == "cli-test-user" + ) + + @pytest.mark.asyncio + async def test_cli_sso_callback_rejects_restricted_sso_group(self): + """CLI SSO must enforce restricted_sso_group before upserting the user.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints import ui_sso + from litellm.proxy.management_endpoints.types import CustomOpenID + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://internal-proxy.local/" + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": "poll-secret-hash", + "user_code_hash": "user-code-hash", + "sso_complete": False, + "user_code_verified": False, + "session_data": None, + } + mock_sso_result = CustomOpenID( + id="cli-test-user", + email="cli-test@example.com", + display_name="cli-test-user", + provider="generic", + team_ids=["other-group"], + ) + + with ( + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + new=AsyncMock(), + ) as get_user_info_mock, + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.user_custom_sso", None), + patch( + "litellm.proxy.proxy_server.general_settings", + { + "ui_access_mode": { + "type": "restricted_sso_group", + "restricted_sso_group": "required-group", + } + }, + ), + ): + with pytest.raises(ProxyException): + await ui_sso.cli_sso_callback( + request=mock_request, + key="cli-session-restricted", + result=mock_sso_result, + received_response={"groups": ["other-group"]}, + ) + + get_user_info_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_cli_sso_callback_persists_attribution_metadata(self, monkeypatch): + from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints import ui_sso + + monkeypatch.setattr( + ui_sso, + "CLI_SSO_CLAIM_MAP", + "employment_type->acme_employment_type", + ) + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://internal-proxy.local/" + session_key = "cli-session-4567890" + mock_user_info = LiteLLM_UserTable( + user_id="test-user-123", + user_role="internal_user", + teams=["team1"], + models=["gpt-4"], + ) + mock_sso_result = { + "user_email": "test@example.com", + "user_id": "test-user-123", + "employment_type": "contractor", + } + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": "poll-secret-hash", + "user_code_hash": "user-code-hash", + "sso_complete": False, + "user_code_verified": False, + "session_data": None, + } + mock_prisma = MagicMock() + mock_prisma.db.litellm_usertable.find_unique = AsyncMock( + return_value=MagicMock(metadata={"auth_provider": "generic"}) + ) + mock_prisma.db.litellm_usertable.update_many = AsyncMock() + + with ( + patch.dict( + os.environ, + { + "PROXY_BASE_URL": "https://test.example.com", + "SERVER_ROOT_PATH": "", + }, + ), + patch( + "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", + return_value=mock_user_info, + ), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch("litellm.proxy.proxy_server.user_custom_sso", None), + patch( + "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", + return_value="Success", + ), + ): + await ui_sso.cli_sso_callback( + request=mock_request, + key=session_key, + result=mock_sso_result, + ) + + flow_data = mock_cache.set_cache.call_args.kwargs["value"] + assert flow_data["session_data"]["attribution_metadata"] == { + "acme_employment_type": "contractor" + } + mock_prisma.db.litellm_usertable.update_many.assert_awaited_once() + update_data = mock_prisma.db.litellm_usertable.update_many.call_args.kwargs[ + "data" + ] + assert update_data["metadata"]["acme_employment_type"] == "contractor" + assert update_data["metadata"]["auth_provider"] == "generic" + + @pytest.mark.asyncio + async def test_cli_poll_key_returns_attribution_metadata(self, monkeypatch): + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + session_key = "cli-session-789123" + session_data = { + "user_id": "test-user-456", + "user_role": "internal_user", + "teams": ["team-a", "team-b"], + "models": ["gpt-4"], + "attribution_metadata": { + "acme_employment_type": "full_time", + "org": {"cost_center": "CC-42"}, + }, + } + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": session_data, + } + + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + result = await cli_poll_key( + key_id=session_key, + team_id=None, + x_litellm_cli_poll_secret="poll-secret", + ) + + assert result["attribution_metadata"] == { + "acme_employment_type": "full_time", + "org.cost_center": "CC-42", + } + + class TestValidateReturnTo: """Tests for SSOAuthenticationHandler._validate_return_to""" diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 46a55fc7468..580ed95062b 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1,7 +1,11 @@ import os import sys +from pathlib import Path +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch +import click +import fastapi import pytest sys.path.insert( @@ -231,6 +235,96 @@ class TestProxyInitializationHelpers: mock_app, "localhost", 8000, "cert.pem", "key.pem", "ECDHE" ) + @patch("granian.Granian") + @patch("builtins.print") + def test_init_granian_server(self, mock_print, mock_granian_cls): + pytest.importorskip("granian") + mock_server = MagicMock() + mock_granian_cls.return_value = mock_server + fake_interfaces = SimpleNamespace(ASGI="asgi") + with patch("granian.constants.Interfaces", fake_interfaces): + ProxyInitializationHelpers._init_granian_server( + host="0.0.0.0", + port=4000, + num_workers=2, + ssl_certfile_path=None, + ssl_keyfile_path=None, + max_requests_before_restart=None, + ciphers=None, + granian_runtime_threads=None, + ) + mock_granian_cls.assert_called_once() + call_kwargs = mock_granian_cls.call_args.kwargs + assert call_kwargs["target"] == "litellm.proxy.proxy_server:app" + assert call_kwargs["address"] == "0.0.0.0" + assert call_kwargs["port"] == 4000 + assert call_kwargs["workers"] == 2 + assert call_kwargs["interface"] == "asgi" + assert call_kwargs["websockets"] is True + assert "runtime_threads" not in call_kwargs + mock_server.serve.assert_called_once() + + @patch("granian.Granian") + @patch("builtins.print") + def test_init_granian_server_runtime_threads(self, mock_print, mock_granian_cls): + pytest.importorskip("granian") + mock_server = MagicMock() + mock_granian_cls.return_value = mock_server + fake_interfaces = SimpleNamespace(ASGI="asgi") + with patch("granian.constants.Interfaces", fake_interfaces): + ProxyInitializationHelpers._init_granian_server( + host="0.0.0.0", + port=4000, + num_workers=1, + ssl_certfile_path=None, + ssl_keyfile_path=None, + max_requests_before_restart=None, + ciphers=None, + granian_runtime_threads=4, + ) + assert mock_granian_cls.call_args.kwargs["runtime_threads"] == 4 + + @patch("granian.Granian") + @patch("builtins.print") + def test_init_granian_server_ssl(self, mock_print, mock_granian_cls): + pytest.importorskip("granian") + mock_server = MagicMock() + mock_granian_cls.return_value = mock_server + fake_interfaces = SimpleNamespace(ASGI="asgi") + with patch("granian.constants.Interfaces", fake_interfaces): + ProxyInitializationHelpers._init_granian_server( + host="0.0.0.0", + port=4000, + num_workers=1, + ssl_certfile_path="/path/to/cert.pem", + ssl_keyfile_path="/path/to/key.pem", + max_requests_before_restart=None, + ciphers=None, + granian_runtime_threads=None, + ) + call_kwargs = mock_granian_cls.call_args.kwargs + assert call_kwargs["ssl_cert"] == Path("/path/to/cert.pem") + assert call_kwargs["ssl_key"] == Path("/path/to/key.pem") + mock_server.serve.assert_called_once() + + @patch("granian.Granian") + def test_init_granian_server_ssl_requires_cert_and_key(self, mock_granian_cls): + pytest.importorskip("granian") + fake_interfaces = SimpleNamespace(ASGI="asgi") + with patch("granian.constants.Interfaces", fake_interfaces): + with pytest.raises(click.ClickException, match="Both --ssl_certfile_path"): + ProxyInitializationHelpers._init_granian_server( + host="0.0.0.0", + port=4000, + num_workers=1, + ssl_certfile_path="/path/to/cert.pem", + ssl_keyfile_path=None, + max_requests_before_restart=None, + ciphers=None, + granian_runtime_threads=None, + ) + mock_granian_cls.assert_not_called() + @patch("subprocess.Popen") def test_run_ollama_serve(self, mock_popen): # Execute @@ -483,6 +577,136 @@ class TestProxyInitializationHelpers: assert appended_params["connection_limit"] == 5 assert appended_params["pool_timeout"] == expected_timeout + def test_build_db_connection_url_params_defaults(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params(connection_limit=10, pool_timeout=60) + assert params == {"connection_limit": 10, "pool_timeout": 60} + + def test_build_db_connection_url_params_omits_none_timeouts(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + connect_timeout=None, + socket_timeout=None, + ) + assert "connect_timeout" not in params + assert "socket_timeout" not in params + + def test_build_db_connection_url_params_includes_optional_timeouts(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + connect_timeout=15, + socket_timeout=120, + ) + assert params["connect_timeout"] == 15 + assert params["socket_timeout"] == 120 + + def test_build_db_connection_url_params_extras_override_defaults(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + extra_params={ + "pgbouncer": "true", + "statement_cache_size": 0, + "pool_timeout": 5, + }, + ) + assert params["pgbouncer"] == "true" + assert params["statement_cache_size"] == 0 + assert params["pool_timeout"] == 5 + + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_db_connection_extra_params_forwarded_to_url( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( + return_value={ + "general_settings": { + "database_url": "postgresql://test:test@localhost:5432/test", + "database_connect_timeout": 15, + "database_socket_timeout": 120, + "database_extra_connection_params": { + "pgbouncer": "true", + "statement_cache_size": 0, + }, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + patch( + "litellm.proxy.proxy_cli.append_query_params", + side_effect=lambda url, params: str(url), + ) as mock_append_query_params, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--config", "test-config.yaml", "--skip_server_startup"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_append_query_params.assert_called() + appended_params = mock_append_query_params.call_args.args[1] + assert appended_params["connect_timeout"] == 15 + assert appended_params["socket_timeout"] == 120 + assert appended_params["pgbouncer"] == "true" + assert appended_params["statement_cache_size"] == 0 + @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 6d10d2a6353..ae0996d16d5 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5708,6 +5708,7 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( fake_redis = AsyncMock() fake_redis.async_increment = AsyncMock(side_effect=record_increment) fake_redis.async_get_cache = AsyncMock(return_value=None) # counter missing + fake_redis.async_set_cache = AsyncMock(return_value=True) # SET NX wins counter_cache.redis_cache = fake_redis # Prisma returns spend=42.0 (authoritative) while the stale cached @@ -5744,16 +5745,131 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( where={"team_id": "team-9"} ) - # Two increments keyed on the counter: seed ($42) then request ($1.50). + # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. + # Only the per-request delta (1.5) goes through INCRBYFLOAT. + fake_redis.async_set_cache.assert_awaited_once_with( + key="spend:team:team-9", value=42.0, nx=True + ) writes = [(c["key"], c["value"]) for c in recorded_increments] - assert ("spend:team:team-9", 42.0) in writes - assert ("spend:team:team-9", 1.5) in writes + assert writes == [("spend:team:team-9", 1.5)] finally: ps.user_api_key_cache = orig_user ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma +@pytest.mark.asyncio +async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed(): + """Two pods both observing a missing Redis counter must not both + INCRBYFLOAT the full DB spend. SpendCounterReseed.coalesced uses SET NX + so the loser reads the winner's value; final Redis = db_spend, not + 2 * db_spend. + + The per-counter asyncio.Lock is per-process, so it does NOT coordinate + across pods. We simulate two pods by patching _get_lock to return a + fresh lock per call (each "pod" has its own lock registry in real life). + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed + + counter_key = "spend:team:team-concurrent-seed" + redis_store: dict = {} + db_read_count = 0 + set_results: list = [] + get_after_set_count = 0 + set_completed_count = 0 + + async def redis_set_cache(key, value, nx=False, **_): + # Yield BEFORE the membership check so two concurrent callers + # interleave the way real atomic Redis SET NX does: the first + # to resume runs check + write atomically and wins; the second + # resumes after the key exists and loses. Yielding *after* the + # check would let both callers pass the empty-store check before + # either writes, so neither would ever lose. + await asyncio.sleep(0) + if nx and key in redis_store: + set_results.append(False) + return False + redis_store[key] = float(value) + set_results.append(True) + nonlocal set_completed_count + set_completed_count += 1 + return True + + async def redis_get_cache(key): + # Track reads that happen after at least one SET NX has completed + # — those are the loser-path fallback reads we want to verify. + if set_completed_count > 0: + nonlocal get_after_set_count + get_after_set_count += 1 + return redis_store.get(key) + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) + + async def slow_find_unique(**_): + nonlocal db_read_count + db_read_count += 1 + # Both pods read DB before either's SET NX lands. + await asyncio.sleep(0) + row = MagicMock() + row.spend = 506.0 + return row + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=slow_find_unique + ) + + pod_a = DualCache() + pod_a.redis_cache = fake_redis + pod_b = DualCache() + pod_b.redis_cache = fake_redis + + # Each "pod" has its own per-process lock registry. Patch _get_lock to + # always return a fresh lock so the two coalesced calls do not serialize + # via one in-process lock (which is what would happen across pods). + async def fresh_lock(_counter_key): + return asyncio.Lock() + + with patch.object(SpendCounterReseed, "_get_lock", side_effect=fresh_lock): + results = await asyncio.gather( + SpendCounterReseed.coalesced( + prisma_client=fake_prisma, + spend_counter_cache=pod_a, + counter_key=counter_key, + ), + SpendCounterReseed.coalesced( + prisma_client=fake_prisma, + spend_counter_cache=pod_b, + counter_key=counter_key, + ), + ) + + assert all(r == 506.0 for r in results), results + assert redis_store[counter_key] == pytest.approx(506.0), redis_store + # Both pods read the DB and both attempted SET NX; exactly one wrote + # (winner) and one was rejected (loser). + assert db_read_count == 2 + assert fake_redis.async_set_cache.await_count == 2 + nx_writes = [ + call + for call in fake_redis.async_set_cache.await_args_list + if call.kwargs.get("nx") is True + ] + assert len(nx_writes) == 2 + assert sorted(set_results) == [False, True], ( + f"expected exactly one SET NX winner and one loser, got {set_results}" + ) + # Loser path executed: after the winner's SET NX returned True, the + # losing coalesced() call falls back to async_get_cache to read the + # winner's value rather than re-seeding. + assert get_after_set_count >= 1, ( + "loser branch (else: read back winner's value) was never exercised" + ) + + @pytest.mark.asyncio async def test_reseed_spend_from_db_user_and_org_prefixes(): """User and org counters reseed from their own DB tables. @@ -5877,9 +5993,16 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(return_value=None) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis db_row = MagicMock() @@ -5907,6 +6030,7 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( where={"team_id": "team-stale-local"} ) + # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. assert redis_store[counter_key] == pytest.approx(43.5) assert counter_cache.in_memory_cache.get_cache( key=counter_key @@ -6297,14 +6421,14 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): from litellm.proxy.proxy_server import get_current_spend counter_cache = DualCache() - recorded_warms: list = [] + recorded_seeds: list = [] - async def record_increment(key, value, ttl=None, **kwargs): - recorded_warms.append({"key": key, "value": value}) - return value + async def record_set_cache(key, value, nx=False, **kwargs): + recorded_seeds.append({"key": key, "value": value, "nx": nx}) + return True fake_redis = AsyncMock() - fake_redis.async_increment = AsyncMock(side_effect=record_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=record_set_cache) fake_redis.async_get_cache = AsyncMock(return_value=None) counter_cache.redis_cache = fake_redis @@ -6329,9 +6453,9 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): f"expected DB reseed to return 362.0, got {spend} " f"(fallback would have returned 30.0 and caused bypass)" ) - # Counter warmed so subsequent reads are fast - assert ("spend:team_member:user-1:team-1", 362.0) in [ - (w["key"], w["value"]) for w in recorded_warms + # Counter warmed via SET NX so subsequent reads are fast. + assert ("spend:team_member:user-1:team-1", 362.0, True) in [ + (s["key"], s["value"], s["nx"]) for s in recorded_seeds ] assert counter_cache.in_memory_cache.get_cache( key="spend:team_member:user-1:team-1" @@ -6408,8 +6532,15 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -6516,9 +6647,16 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() @@ -6621,9 +6759,16 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): redis_store[key] = (redis_store.get(key) or 0.0) + value return redis_store[key] + async def redis_set_cache(key, value, nx=False, **_): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(side_effect=redis_get) fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) counter_cache.redis_cache = fake_redis db_call_count = 0 diff --git a/tests/test_litellm/responses/test_sse_output_recovery.py b/tests/test_litellm/responses/test_sse_output_recovery.py new file mode 100644 index 00000000000..c8f3325a624 --- /dev/null +++ b/tests/test_litellm/responses/test_sse_output_recovery.py @@ -0,0 +1,57 @@ +"""Tests for litellm.responses.sse_output_recovery helpers.""" + +from litellm.responses.sse_output_recovery import ( + _MAX_CONTENT_INDEX, + record_output_text_chunk, +) + + +def test_text_chunk_with_oversized_content_index_is_dropped(): + output_items: dict = {} + text_only_items: dict = {} + record_output_text_chunk( + parsed_chunk={ + "type": "response.output_text.done", + "output_index": 0, + "content_index": _MAX_CONTENT_INDEX + 1, + "text": "ignored", + }, + output_items=output_items, + text_only_items=text_only_items, + ) + item = text_only_items[0] + assert item["content"] == [] + + +def test_text_chunk_with_negative_content_index_is_dropped(): + output_items: dict = {} + text_only_items: dict = {} + record_output_text_chunk( + parsed_chunk={ + "type": "response.output_text.done", + "output_index": 0, + "content_index": -1, + "text": "ignored", + }, + output_items=output_items, + text_only_items=text_only_items, + ) + assert text_only_items[0]["content"] == [] + + +def test_text_chunk_at_max_content_index_is_recorded(): + output_items: dict = {} + text_only_items: dict = {} + record_output_text_chunk( + parsed_chunk={ + "type": "response.output_text.done", + "output_index": 0, + "content_index": _MAX_CONTENT_INDEX, + "text": "kept", + }, + output_items=output_items, + text_only_items=text_only_items, + ) + content = text_only_items[0]["content"] + assert len(content) == _MAX_CONTENT_INDEX + 1 + assert content[_MAX_CONTENT_INDEX]["text"] == "kept" diff --git a/tests/test_litellm/test_check_licenses.py b/tests/test_litellm/test_check_licenses.py new file mode 100644 index 00000000000..4d72f185a25 --- /dev/null +++ b/tests/test_litellm/test_check_licenses.py @@ -0,0 +1,211 @@ +"""Tests for the dependency license checker at tests/code_coverage_tests/check_licenses.py. + +Focus: PEP 639 license metadata. Packages that adopt PEP 639 publish their +license as an SPDX expression in ``info.license_expression`` and often leave the +legacy ``info.license`` field null, so the checker must read the new field (and +fall back to trove classifiers) instead of reporting "Unknown license". + +PyPI HTTP responses are mocked — these tests never hit the network. +""" + +import os +import sys +from pathlib import Path + +_CODE_COVERAGE_DIR = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "code_coverage_tests" +) +sys.path.insert(0, _CODE_COVERAGE_DIR) + +import check_licenses # noqa: E402 + +_LICCHECK_INI = Path(_CODE_COVERAGE_DIR) / "liccheck.ini" + + +class _FakeResponse: + def __init__(self, payload): + self._payload = payload + + def raise_for_status(self): + return None + + def json(self): + return self._payload + + +def _make_checker(): + return check_licenses.LicenseChecker(config_file=_LICCHECK_INI) + + +def _patch_pypi(monkeypatch, info): + """Make PyPI return a JSON response with the given ``info`` block.""" + + def _fake_get(url, timeout=None): + return _FakeResponse({"info": info}) + + monkeypatch.setattr(check_licenses.requests, "get", _fake_get) + + +# -------------------------------------------------------------------------- +# get_package_license_from_pypi: license metadata resolution +# -------------------------------------------------------------------------- + + +def test_get_license_prefers_license_expression(monkeypatch): + """(a) PEP 639 packages publish the SPDX expression in license_expression.""" + _patch_pypi( + monkeypatch, + {"license_expression": "MIT", "license": None, "classifiers": []}, + ) + checker = _make_checker() + assert checker.get_package_license_from_pypi("black", "26.3.1") == "MIT" + + +def test_license_expression_wins_when_both_present(monkeypatch): + """license_expression takes precedence over the legacy license field.""" + _patch_pypi( + monkeypatch, + {"license_expression": "Apache-2.0", "license": "stale free text"}, + ) + checker = _make_checker() + assert checker.get_package_license_from_pypi("pkg", "1.0.0") == "Apache-2.0" + + +def test_get_license_falls_back_to_legacy_license(monkeypatch): + """(b) Pre-PEP-639 packages only set the legacy free-text license field.""" + _patch_pypi( + monkeypatch, + {"license_expression": None, "license": "MIT License", "classifiers": []}, + ) + checker = _make_checker() + assert checker.get_package_license_from_pypi("pkg", "1.0.0") == "MIT License" + + +def test_get_license_falls_back_to_classifiers(monkeypatch): + """(c) Some packages express the license only through trove classifiers.""" + _patch_pypi( + monkeypatch, + { + "license_expression": None, + "license": None, + "classifiers": [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", + ], + }, + ) + checker = _make_checker() + assert ( + checker.get_package_license_from_pypi("pkg", "1.0.0") + == "Apache Software License" + ) + + +def test_get_license_returns_none_when_unset(monkeypatch): + """(d) With no license metadata at all the license stays unknown.""" + _patch_pypi( + monkeypatch, + {"license_expression": None, "license": None, "classifiers": []}, + ) + checker = _make_checker() + assert checker.get_package_license_from_pypi("pkg", "1.0.0") is None + + +def test_get_license_returns_none_on_request_failure(monkeypatch): + """Network/HTTP failures are swallowed and reported as unknown.""" + + def _boom(url, timeout=None): + raise RuntimeError("network down") + + monkeypatch.setattr(check_licenses.requests, "get", _boom) + checker = _make_checker() + assert checker.get_package_license_from_pypi("pkg", "1.0.0") is None + + +# -------------------------------------------------------------------------- +# is_license_acceptable: SPDX identifiers and compound expressions +# -------------------------------------------------------------------------- + + +def test_spdx_identifiers_are_authorized(): + """Plain SPDX identifiers match the legacy-spelled authorized list as-is.""" + checker = _make_checker() + for identifier in ("MIT", "Apache-2.0", "BSD-3-Clause"): + is_ok, reason = checker.is_license_acceptable(identifier) + assert is_ok is True, f"{identifier}: {reason}" + + +def test_spdx_compound_or_expression_is_authorized(): + checker = _make_checker() + is_ok, reason = checker.is_license_acceptable("MIT OR Apache-2.0") + assert is_ok is True, reason + + +def test_spdx_with_exception_in_compound_is_authorized(): + """The 'WITH ' suffix is stripped; the base license is checked.""" + checker = _make_checker() + is_ok, reason = checker.is_license_acceptable( + "Apache-2.0 WITH LLVM-exception OR MIT" + ) + assert is_ok is True, reason + + +def test_spdx_gpl3_is_rejected(): + """GPL-3.0 spellings must fail — they match no authorized license.""" + checker = _make_checker() + for expr in ("GPL-3.0-only", "GPL-3.0-or-later"): + is_ok, reason = checker.is_license_acceptable(expr) + assert is_ok is False, f"{expr} unexpectedly accepted: {reason}" + + +def test_spdx_compound_with_copyleft_component_is_rejected(): + """A permissive-OR-copyleft expression is conservatively rejected.""" + checker = _make_checker() + is_ok, _ = checker.is_license_acceptable("MIT OR GPL-3.0-only") + assert is_ok is False + + +def test_or_later_identifier_is_not_split_as_operator(): + """The lowercase '-or-later' inside an identifier is not the SPDX OR operator.""" + assert ( + check_licenses.LicenseChecker._split_spdx_expression("GPL-2.0-or-later") is None + ) + + +def test_free_text_license_is_not_treated_as_spdx(): + """Free-text license blobs fall back to whole-string substring matching.""" + free_text = "MIT License AND additional redistribution permissions" + assert check_licenses.LicenseChecker._split_spdx_expression(free_text) is None + checker = _make_checker() + assert checker.is_license_acceptable(free_text)[0] is True + + +def test_unknown_license_is_reported(): + checker = _make_checker() + is_ok, reason = checker.is_license_acceptable(None) + assert is_ok is False + assert reason == "Unknown license" + + +# -------------------------------------------------------------------------- +# check_package: end-to-end resolution + acceptability +# -------------------------------------------------------------------------- + + +def test_check_package_accepts_pep639_package(monkeypatch): + """A PEP 639 package whose license lives only in license_expression passes.""" + _patch_pypi( + monkeypatch, + {"license_expression": "MIT", "license": None, "classifiers": []}, + ) + checker = _make_checker() + assert checker.check_package("some-pep639-pkg", "1.0.0") is True + + +def test_check_package_rejects_package_without_license(monkeypatch): + _patch_pypi( + monkeypatch, + {"license_expression": None, "license": None, "classifiers": []}, + ) + checker = _make_checker() + assert checker.check_package("mystery-pkg", "1.0.0") is False diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 1be4abbec6e..00902890da3 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -2059,6 +2059,25 @@ def test_openrouter_gemini_3_1_flash_lite_preview_pricing(): assert model_info["max_output_tokens"] == 65536 +def test_gemini_3_1_flash_lite_pricing(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + for model_name in ( + "gemini-3.1-flash-lite", + "gemini/gemini-3.1-flash-lite", + "vertex_ai/gemini-3.1-flash-lite", + ): + model_info = litellm.model_cost.get(model_name) + assert model_info is not None, f"Missing model pricing entry: {model_name}" + assert model_info["input_cost_per_token"] == 4.5e-07 + assert model_info["input_cost_per_audio_token"] == 9e-07 + assert model_info["output_cost_per_token"] == 2.7e-06 + assert model_info["output_cost_per_reasoning_token"] == 2.7e-06 + assert model_info["cache_read_input_token_cost"] == 4.5e-08 + assert model_info["max_input_tokens"] == 1048576 + + def test_custom_pricing_applies_cache_read_input_cost(): """ Bug 1 reproduction: custom_cost_per_token with cache_read_input_token_cost @@ -2371,3 +2390,34 @@ def test_custom_pricing_without_cache_keys_preserves_legacy_behavior(): expected = 1000 * 0.0000025 + 100 * 0.000015 assert cost == pytest.approx(expected) + + +def test_openrouter_gemini_3_1_flash_lite_stable_pricing(): + """ + Test that openrouter/google/gemini-3.1-flash-lite (stable, no -preview suffix) + has a pricing entry. + + Google promoted gemini-3.1-flash-lite to GA on 2026-05-07. PR #27933 added the + stable pricing for the bare, gemini/, and vertex_ai/ prefixes but missed the + openrouter/google/ variant — every other Gemini family in the file has an + openrouter/google/ sibling (2.0-flash-001, 2.5-flash, 2.5-pro, 3-flash-preview, + 3-pro-preview, 3.1-flash-lite-preview, 3.1-pro-preview), so the gap is a + consistency issue, not a design choice. Same shape as the preview-variant gap + fixed in PR #25610. + + Pricing matches the existing -preview entry one-for-one (input $0.25/M, output + $1.50/M, cache-read $0.025/M) — Google did not change costs at the GA cutover. + """ + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model_name = "openrouter/google/gemini-3.1-flash-lite" + model_info = litellm.model_cost.get(model_name) + + assert model_info is not None, f"Missing model pricing entry: {model_name}" + assert model_info["litellm_provider"] == "openrouter" + assert model_info["input_cost_per_token"] == 2.5e-07 + assert model_info["output_cost_per_token"] == 1.5e-06 + assert model_info["cache_read_input_token_cost"] == 2.5e-08 + assert model_info["max_input_tokens"] == 1048576 + assert model_info["max_output_tokens"] == 65536 diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index 1efd698fb64..719cb8eecd2 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -190,3 +190,164 @@ def test_build_custom_pricing_entry_time_based(): assert entry["litellm_provider"] == "openai" assert entry["input_cost_per_second"] == 0.01 assert entry["output_cost_per_second"] == 0.02 + + +def test_register_model_strips_none_litellm_provider(): + """``get_model_info`` returns ``litellm_provider: None`` for deployments + registered without a provider (e.g. ``Router.add_deployment`` flows). + ``register_model`` must not persist that None into ``model_cost``, + otherwise ``_check_provider_match`` will drop custom pricing on + subsequent cost lookups. + + Regression test for https://github.com/BerriAI/litellm/issues/28336. + """ + from litellm.utils import _check_provider_match + + model_key = "test-custom-pricing-no-provider-28336" + litellm.model_cost.pop(model_key, None) + + try: + litellm.register_model( + { + model_key: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + } + ) + + registered = litellm.model_cost.get(model_key) + assert registered is not None, f"{model_key} should be in model_cost" + # The key may be absent entirely, but if present it must not be None. + assert ( + "litellm_provider" not in registered + or registered["litellm_provider"] is not None + ) + # Downstream consumers must accept this entry for any provider, + # mirroring what the cost calculator does. + assert _check_provider_match(registered, "openai") is True + assert _check_provider_match(registered, "anthropic") is True + finally: + litellm.model_cost.pop(model_key, None) + + +def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeypatch): + """Directly exercise the strip in ``register_model``. + + The companion test above hits the ``except Exception`` branch where + ``existing_model`` is an empty dict, so the ``pop`` is a no-op. This + test patches ``get_model_info`` to return the failure mode the strip + was added to handle, namely a populated dict whose ``litellm_provider`` + is ``None``. Without the strip, the merged entry in + ``litellm.model_cost`` would carry ``litellm_provider: None`` and + ``_check_provider_match`` would drop custom pricing. + + Regression test for https://github.com/BerriAI/litellm/issues/28336. + """ + from litellm import utils as litellm_utils + from litellm.utils import _check_provider_match + + model_key = "test-strip-none-provider-from-get-model-info-28336" + litellm.model_cost.pop(model_key, None) + + def _fake_get_model_info(model, *args, **kwargs): + assert model == model_key + return { + "key": model_key, + "litellm_provider": None, + "mode": "chat", + "max_tokens": 4096, + } + + # ``register_model`` calls ``get_model_info.cache_clear`` via + # ``_invalidate_model_cost_lowercase_map``, so the replacement must + # expose a no-op ``cache_clear`` attribute. + _fake_get_model_info.cache_clear = lambda: None + monkeypatch.setattr(litellm_utils, "get_model_info", _fake_get_model_info) + + try: + litellm.register_model( + { + model_key: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + } + ) + + registered = litellm.model_cost.get(model_key) + assert registered is not None, f"{model_key} should be in model_cost" + # The strip must have removed the None-valued provider that + # ``get_model_info`` returned. The key may be absent entirely, but + # it must never be present with value ``None``. + assert "litellm_provider" not in registered or ( + registered["litellm_provider"] is not None + ), ( + "register_model failed to strip litellm_provider=None returned " + f"by get_model_info, got {registered.get('litellm_provider')!r}" + ) + # Metadata from the patched ``get_model_info`` must still flow + # through, so we know the strip did not nuke the rest of the entry. + assert registered.get("mode") == "chat" + assert registered.get("max_tokens") == 4096 + # And custom pricing from the registration call must be preserved. + assert registered.get("input_cost_per_token") == 0.001 + assert registered.get("output_cost_per_token") == 0.002 + # Downstream _check_provider_match must accept any provider for + # this entry, mirroring the cost calculator path. + assert _check_provider_match(registered, "openai") is True + assert _check_provider_match(registered, "anthropic") is True + finally: + litellm.model_cost.pop(model_key, None) + + +def test_register_model_router_add_deployment_custom_pricing_applies(): + """End-to-end regression for https://github.com/BerriAI/litellm/issues/28336. + + ``Router.add_deployment`` registers custom pricing without passing + ``litellm_provider``. Cost calculation must still pick up the custom + pricing instead of falling back to the default provider price. + """ + from litellm import Router + + model_key = "router-add-deployment-custom-pricing-28336" + deployment_model = f"openai/{model_key}" + litellm.model_cost.pop(model_key, None) + litellm.model_cost.pop(deployment_model, None) + + router = Router( + model_list=[ + { + "model_name": model_key, + "litellm_params": { + "model": deployment_model, + "api_key": "fake-key-for-registration", + "input_cost_per_token": 0.00042, + "output_cost_per_token": 0.00084, + }, + "model_info": {"id": "deployment-28336"}, + } + ] + ) + + try: + # ``add_deployment`` runs as part of ``Router.__init__``; the + # registered entry must not block ``_check_provider_match`` for + # the deployment's provider. + from litellm.utils import _check_provider_match + + registered_keys = [ + k for k in (deployment_model, model_key) if k in litellm.model_cost + ] + assert registered_keys, ( + "Router.add_deployment did not register custom pricing for " + f"{model_key} / {deployment_model}" + ) + for k in registered_keys: + assert _check_provider_match(litellm.model_cost[k], "openai") is True, ( + f"custom pricing for {k} was dropped by _check_provider_match" + ) + finally: + litellm.model_cost.pop(model_key, None) + litellm.model_cost.pop(deployment_model, None) + del router diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d8be527689e..5e636b86ed6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1741,6 +1741,362 @@ async def test_acompletion_streaming_iterator_pre_first_chunk_skips_continuation assert fallback_kwargs["messages"] == messages +# --------------------------------------------------------------------------- +# Shared helpers for the _aresponses_streaming_iterator test suite. +# --------------------------------------------------------------------------- +def _make_responses_iterator( + *, + chunks=(), + error=None, + bridge=False, + model="gpt-4", + hidden_params=None, + chat_chunks=None, +): + """Build a minimal mock Responses-API streaming iterator. + + Bypasses BaseResponsesAPIStreamingIterator.__init__ but mirrors every + attribute production code reads. Yields *chunks*, then raises *error* + (or StopAsyncIteration). Set bridge=True to inherit from + LiteLLMCompletionStreamingIterator so the wrapper's bridge-path + isinstance check (used by usage extraction) matches. + """ + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + base = ( + LiteLLMCompletionStreamingIterator + if bridge + else BaseResponsesAPIStreamingIterator + ) + + class _Iter(base): + def __init__(self): + self._chunks = list(chunks) + self._idx = 0 + self._hidden_params = hidden_params or {} + self.model = model + self.custom_llm_provider = "anthropic" + self.logging_obj = MagicMock() + self.litellm_metadata = None + self.responses_api_provider_config = None + self.finished = False + self.completed_response = None + self.response = None + self.start_time = None + self.request_data = {} + self.call_type = None + if chat_chunks is not None: + self.collected_chat_completion_chunks = chat_chunks + + def __aiter__(self): + return self + + async def __anext__(self): + if self._idx < len(self._chunks): + self._idx += 1 + return self._chunks[self._idx - 1] + if error is not None: + raise error + raise StopAsyncIteration + + return _Iter() + + +class _AsyncList: + """Generic async iterator over a list — used as the fallback response.""" + + def __init__(self, items=()): + self._items = list(items) + self._idx = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + if self._idx >= len(self._items): + raise StopAsyncIteration + item = self._items[self._idx] + self._idx += 1 + return item + + +def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"): + return litellm.Router( + model_list=[ + { + "model_name": primary, + "litellm_params": {"model": primary, "api_key": "k1"}, + }, + { + "model_name": secondary, + "litellm_params": {"model": secondary, "api_key": "k2"}, + }, + ], + fallbacks=[{primary: [secondary]}], + ) + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_fallback(): + """Catches MidStreamFallbackError, re-enters the fallback chain via + async_function_with_fallbacks_common_utils with the per-attempt helper + and original_generic_function preserved. Mirrors + test_acompletion_streaming_iterator for the aresponses path.""" + from litellm.exceptions import MidStreamFallbackError + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + router = _make_router_with_fallback( + "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" + ) + src = _make_responses_iterator( + chunks=[MagicMock(type="response.created")], + error=MidStreamFallbackError( + message="anthropic socket timeout", + model="anthropic/claude-sonnet-4-6", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="", + ), + model="anthropic/claude-sonnet-4-6", + hidden_params={"model_id": "src-deployment-1"}, + ) + fallback_chunks = [ + MagicMock(type="response.output_text.delta"), + MagicMock(type="response.completed"), + ] + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_AsyncList(fallback_chunks), + ) as mock_fallback_utils: + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "anthropic/claude-sonnet-4-6", + "stream": True, + "input": "Hi", + "original_generic_function": litellm.aresponses, + }, + ) + assert isinstance(wrapped, BaseResponsesAPIStreamingIterator) + assert wrapped._hidden_params.get("model_id") == "src-deployment-1" + collected = [c async for c in wrapped] + + assert len(collected) == 3 # 1 primary chunk + 2 fallback chunks + call_kwargs = mock_fallback_utils.call_args.kwargs + fbk = call_kwargs["kwargs"] + # Bound methods compare equal when they share the same instance + __func__. + assert fbk["original_function"] == router._ageneric_api_call_with_fallbacks_helper + assert fbk["original_generic_function"] is litellm.aresponses + assert call_kwargs["model_group"] == "anthropic/claude-sonnet-4-6" + assert call_kwargs["disable_fallbacks"] is False + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback(): + """Regression: model_group must land under "litellm_metadata" (the key + litellm.aresponses reads), not the default "metadata".""" + from litellm.exceptions import MidStreamFallbackError + + router = _make_router_with_fallback() + src = _make_responses_iterator( + error=MidStreamFallbackError( + message="boom", + model="gpt-4", + llm_provider="anthropic", + is_pre_first_chunk=True, + generated_content="", + ) + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_AsyncList(), + ) as mock_fallback_utils: + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + async for _ in wrapped: + pass + + fbk = mock_fallback_utils.call_args.kwargs["kwargs"] + assert "litellm_metadata" in fbk, "wrong metadata_variable_name" + assert fbk["litellm_metadata"]["model_group"] == "gpt-4" + assert "model_group" not in fbk.get( + "metadata", {} + ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_pre_first_chunk_skips_continuation(): + """Pre-first-chunk error: original input is preserved unchanged.""" + from litellm.exceptions import MidStreamFallbackError + + router = _make_router_with_fallback() + src = _make_responses_iterator( + error=MidStreamFallbackError( + message="socket timeout before first chunk", + model="gpt-4", + llm_provider="anthropic", + is_pre_first_chunk=True, + generated_content="", + ) + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_AsyncList(), + ) as mock_fallback_utils: + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "Hello", + "original_generic_function": litellm.aresponses, + }, + ) + async for _ in wrapped: + pass + + fbk = mock_fallback_utils.call_args.kwargs["kwargs"] + assert fbk["input"] == "Hello" # original input, no continuation messages + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): + """Mid-stream error: input is rewritten to include user prompt + + developer instruction + prior assistant message with partial output.""" + from litellm.exceptions import MidStreamFallbackError + + router = _make_router_with_fallback() + src = _make_responses_iterator( + chunks=[MagicMock(type="response.output_text.delta")], + error=MidStreamFallbackError( + message="socket reset mid-stream", + model="gpt-4", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="The capital of France is", + ), + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_AsyncList(), + ) as mock_fallback_utils: + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "What's the capital of France?", + "original_generic_function": litellm.aresponses, + }, + ) + async for _ in wrapped: + pass + + new_input = mock_fallback_utils.call_args.kwargs["kwargs"]["input"] + assert isinstance(new_input, list) + assert new_input[0]["role"] == "user" + assert new_input[0]["content"][0]["text"] == "What's the capital of France?" + assert new_input[1]["role"] == "developer" + assert "do not repeat" in new_input[1]["content"][0]["text"].lower() + assert new_input[2]["role"] == "assistant" + assert new_input[2]["content"][0]["type"] == "output_text" + assert new_input[2]["content"][0]["text"] == "The capital of France is" + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_combines_partial_usage(): + """Partial usage from the bridge path is normalized to ResponseAPIUsage + and summed onto the fallback's response.completed event — no token-name + split, clean ResponseAPIUsage on output.""" + from types import SimpleNamespace + + from litellm.exceptions import MidStreamFallbackError + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + router = _make_router_with_fallback() + src = _make_responses_iterator( + bridge=True, + chat_chunks=[MagicMock()], + chunks=[MagicMock(type="response.output_text.delta")], + error=MidStreamFallbackError( + message="boom", + model="gpt-4", + llm_provider="anthropic", + is_pre_first_chunk=False, + generated_content="hello", + ), + ) + + fallback_response_object = ResponsesAPIResponse( + id="resp_test", created_at=0, model="gpt-4", object="response", output=[] + ) + fallback_response_object.usage = ResponseAPIUsage( + input_tokens=20, output_tokens=15, total_tokens=35 + ) + fallback_event = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=fallback_response_object, + ) + + with ( + patch( + "litellm.main.stream_chunk_builder", + return_value=SimpleNamespace( + usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) + ), + ), + patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=_AsyncList([fallback_event]), + ), + ): + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={ + "model": "gpt-4", + "stream": True, + "input": "hi", + "original_generic_function": litellm.aresponses, + }, + ) + async for _ in wrapped: + pass + + merged = fallback_response_object.usage + assert isinstance(merged, ResponseAPIUsage) + assert merged.input_tokens == 30 # 10 (translated from prompt_tokens) + 20 + assert merged.output_tokens == 19 # 4 (translated from completion_tokens) + 15 + assert merged.total_tokens == 49 + + @pytest.mark.asyncio async def test_async_function_with_fallbacks_common_utils(): """Test the async_function_with_fallbacks_common_utils method""" @@ -3863,7 +4219,15 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): # No model_info on deployment object → treated as not blocked assert litellm.Router._is_deployment_blocked(object()) is False missing_blocked = types.SimpleNamespace() - assert litellm.Router._is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False - assert litellm.Router._is_deployment_blocked( - types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) - ) is True + assert ( + litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=missing_blocked) + ) + is False + ) + assert ( + litellm.Router._is_deployment_blocked( + types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True)) + ) + is True + ) diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index c3f93078557..9454e03e918 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -7,8 +7,10 @@ and one has explicit zero-cost pricing in model_info, the other deployment should still use the built-in pricing. """ +import copy import os import sys +from unittest.mock import patch import pytest @@ -19,6 +21,16 @@ sys.path.insert( import litellm from litellm import Router from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.utils import _invalidate_model_cost_lowercase_map + + +def _restore_model_cost_entries(original_entries): + for key, value in original_entries.items(): + if value is None: + litellm.model_cost.pop(key, None) + else: + litellm.model_cost[key] = value + _invalidate_model_cost_lowercase_map() def test_should_not_pollute_shared_key_with_zero_cost_pricing(): @@ -323,3 +335,70 @@ def test_responses_prefix_stripped_alias_registered_for_add_deployment(): ) is True ) + + +def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): + """ + ChatGPT aliases that share the same backend model should not be able to + downgrade the shared backend key from responses -> chat during router setup. + """ + from litellm.main import responses_api_bridge_check + + backend_model = "chatgpt/gpt-5.4" + model_keys = { + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + "chatgpt-shared-mode-base": copy.deepcopy( + litellm.model_cost.get("chatgpt-shared-mode-base") + ), + "chatgpt-shared-mode-alias": copy.deepcopy( + litellm.model_cost.get("chatgpt-shared-mode-alias") + ), + } + + try: + backend_entry = copy.deepcopy(model_keys[backend_model]) or {} + backend_entry["litellm_provider"] = "chatgpt" + backend_entry["mode"] = "responses" + litellm.model_cost[backend_model] = backend_entry + _invalidate_model_cost_lowercase_map() + + router = Router(model_list=[]) + with patch.object( + Router, "_add_deployment", lambda self, deployment: deployment + ): + router._create_deployment( + deployment_info={}, + _model_name="chatgpt/gpt-5.4", + _litellm_params={ + "model": "gpt-5.4", + "custom_llm_provider": "chatgpt", + }, + _model_info={ + "id": "chatgpt-shared-mode-base", + "mode": "responses", + }, + ) + router._create_deployment( + deployment_info={}, + _model_name="chatgpt/gpt-5.4-medium", + _litellm_params={ + "model": "gpt-5.4", + "custom_llm_provider": "chatgpt", + }, + _model_info={ + "id": "chatgpt-shared-mode-alias", + "mode": "chat", + }, + ) + + assert litellm.model_cost[backend_model]["mode"] == "responses" + assert "mode" in litellm.model_cost[backend_model] + + bridge_model_info, bridge_model = responses_api_bridge_check( + model="gpt-5.4", + custom_llm_provider="chatgpt", + ) + assert bridge_model == "gpt-5.4" + assert bridge_model_info["mode"] == "responses" + finally: + _restore_model_cost_entries(model_keys) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index afd75ddebb7..6f08a424151 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -754,6 +754,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_dbu_cost_per_token": {"type": "number"}, "annotation_cost_per_page": {"type": "number"}, "ocr_cost_per_page": {"type": "number"}, + "ocr_cost_per_credit": {"type": "number"}, "code_interpreter_cost_per_session": {"type": "number"}, "inference_geo": {"type": "string"}, "litellm_provider": {"type": "string"}, @@ -856,6 +857,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_adaptive_thinking": {"type": "boolean"}, "supports_service_tier": {"type": "boolean"}, "supports_preset": {"type": "boolean"}, + "supports_output_config": {"type": "boolean"}, "tool_use_system_prompt_tokens": {"type": "number"}, "tpm": {"type": "number"}, "provider_specific_entry": {"type": "object"}, @@ -1139,6 +1141,34 @@ def test_check_provider_match(): assert litellm.utils._check_provider_match(model_info, "openai") is False +def test_check_provider_match_none_value_matches_any_provider(): + """ + A ``litellm_provider`` of None must be treated the same as a missing + key: both mean "no provider constraint" and should match any + ``custom_llm_provider``. + + Regression test for https://github.com/BerriAI/litellm/issues/28336. + Before the fix, ``register_model`` persisted ``litellm_provider: None`` + via ``get_model_info`` for deployments registered without a provider + (e.g. ``Router.add_deployment``), which caused ``_check_provider_match`` + to drop custom pricing intermittently. + """ + # Missing key already returned True; None must behave identically. + assert litellm.utils._check_provider_match({}, "openai") is True + assert ( + litellm.utils._check_provider_match({"litellm_provider": None}, "openai") + is True + ) + assert ( + litellm.utils._check_provider_match({"litellm_provider": None}, "anthropic") + is True + ) + # When custom_llm_provider is also None nothing constrains the match. + assert ( + litellm.utils._check_provider_match({"litellm_provider": None}, None) is True + ) + + def test_get_provider_rerank_config(): """ Test the get_provider_rerank_config function for various providers diff --git a/tests/test_spend_logs.py b/tests/test_spend_logs.py index 8aec1d5cc60..c575fa07551 100644 --- a/tests/test_spend_logs.py +++ b/tests/test_spend_logs.py @@ -100,6 +100,9 @@ async def get_spend_logs(session, request_id=None, api_key=None): return await response.json() +@pytest.mark.skip( + reason="Flaky in CI: /spend/logs?request_id=... returns 500 even after a 20s wait for the spend log to be written. Spend-log accuracy is covered by tests/test_litellm/proxy/spend_tracking/ and the proxy_spend_accuracy_tests CircleCI job." +) @pytest.mark.asyncio async def test_spend_logs(): """ @@ -155,6 +158,9 @@ async def generate_team(session: aiohttp.ClientSession, org_id: str) -> dict: return await response.json() +@pytest.mark.skip( + reason="Flaky in CI: /spend/logs?request_id=... returns 500 even after a 20s wait for the spend log to be written. Same write-then-read race against the spend logs DB as test_spend_logs. Spend-log accuracy is covered by tests/test_litellm/proxy/spend_tracking/ and the proxy_spend_accuracy_tests CircleCI job." +) @pytest.mark.asyncio async def test_spend_logs_with_org_id(): """ diff --git a/tests/test_team_members.py b/tests/test_team_members.py index a3d64eae803..4cf85af6410 100644 --- a/tests/test_team_members.py +++ b/tests/test_team_members.py @@ -136,6 +136,9 @@ def test_add_single_member(api_client, new_team): ), f"Team size did not increase by 1 (was {initial_size}, now {updated_size})" +@pytest.mark.skip( + reason="Flaky in CI: /team/info?team_id=... intermittently returns 404/400 mid-loop after add_team_member calls. Single-member coverage in test_add_single_member is sufficient; team-member CRUD is also covered by tests/test_litellm/proxy/management_endpoints/." +) def test_add_multiple_members(api_client, new_team): """Test adding multiple members to a new team""" # Get initial team size @@ -203,6 +206,9 @@ def test_error_handling(api_client): api_client.get_team_info("invalid-team-id") +@pytest.mark.skip( + reason="Flaky in CI: /team/info?team_id=... intermittently returns 404 after add_team_member calls, same race documented for test_add_multiple_members. Duplicate-prevention is covered by test_update_team_members_list_duplicate_prevention in tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py." +) def test_duplicate_user_addition(api_client, new_team): """Test that adding the same user twice is handled appropriately""" # Add user first time diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts index ec4d3a6ddb0..6964fe52a14 100644 --- a/ui/litellm-dashboard/e2e_tests/playwright.config.ts +++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts @@ -28,6 +28,11 @@ export default defineConfig({ /* Action timeout for clicks, fills, waitForSelector, etc. */ actionTimeout: 15 * 1000, navigationTimeout: 30 * 1000, + + /* Slow down actions when SLOWMO= is set, useful for headed local debugging */ + launchOptions: { + slowMo: process.env.SLOWMO ? (parseInt(process.env.SLOWMO, 10) || 0) : 0, + }, }, /* Configure projects for major browsers */ diff --git a/ui/litellm-dashboard/e2e_tests/run_e2e.sh b/ui/litellm-dashboard/e2e_tests/run_e2e.sh index 4e3a47edfbd..36619dce9b2 100755 --- a/ui/litellm-dashboard/e2e_tests/run_e2e.sh +++ b/ui/litellm-dashboard/e2e_tests/run_e2e.sh @@ -15,7 +15,7 @@ set -euo pipefail # In CI (CI=true), expects: # - PostgreSQL already running on 127.0.0.1:5432 # - DATABASE_URL already set -# - Python/Poetry already installed +# - Python/uv already installed # - Node.js/npx already available # ================================================================ @@ -48,7 +48,7 @@ cleanup() { trap cleanup EXIT INT TERM # --- Pre-flight checks --- -for cmd in python3 npx poetry; do +for cmd in python3 npx uv; do command -v "$cmd" >/dev/null 2>&1 || { echo "Error: $cmd not found."; exit 1; } done @@ -95,6 +95,10 @@ export DISABLE_SCHEMA_UPDATE="true" export SERVER_ROOT_PATH="" # Prevent logout from redirecting to an external URL export PROXY_LOGOUT_URL="" +# Forward LITELLM_LICENSE if set in the outer env so premium-gated UI flows +# (e.g. Team-BYOK Model switch) can be exercised. Tests that depend on a +# premium proxy gate themselves on process.env.LITELLM_LICENSE. +export LITELLM_LICENSE="${LITELLM_LICENSE:-}" # --- Rebuild UI from source --- echo "=== Building UI from source ===" @@ -117,19 +121,15 @@ echo "UI build copied and restructured" # --- Python environment --- echo "=== Setting up Python environment ===" cd "$REPO_ROOT" -if ! poetry run python3 -c "import prisma" 2>/dev/null; then - echo "Installing Python dependencies (first run)..." - poetry install --with dev,proxy-dev --extras "proxy" --quiet - poetry run pip install nodejs-wheel-binaries 2>/dev/null || true - poetry run prisma generate --schema litellm/proxy/schema.prisma -fi +uv sync --group dev --group proxy-dev --extra proxy --frozen --quiet +uv run --no-sync python -m prisma generate --schema litellm/proxy/schema.prisma echo "=== Pushing Prisma schema to database ===" -poetry run prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss +uv run --no-sync python -m prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss # --- Mock LLM server --- echo "=== Starting mock LLM server ===" -poetry run python3 "$SCRIPT_DIR/fixtures/mock_llm_server/server.py" & +uv run --no-sync python "$SCRIPT_DIR/fixtures/mock_llm_server/server.py" & MOCK_PID=$! for i in $(seq 1 15); do @@ -140,7 +140,7 @@ done # --- LiteLLM proxy --- echo "=== Starting LiteLLM proxy ===" cd "$REPO_ROOT" -poetry run python3 -m litellm.proxy.proxy_cli \ +uv run --no-sync python -m litellm.proxy.proxy_cli \ --config "$SCRIPT_DIR/fixtures/config.yml" \ --port 4000 & PROXY_PID=$! diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts index 14ceb1a4a6b..1e44d9a25a0 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts @@ -126,4 +126,84 @@ test.describe("Proxy Admin - Keys", () => { await expect(page.getByText(E2E_INTERNAL_USER_KEY_ALIAS)).toBeVisible({ timeout: 10_000 }); }); + + test("Create a key with All Proxy Models (no team)", async ({ page }) => { + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + + await page.getByRole("button", { name: /Create New Key/i }).click(); + + await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + + const keyName = `e2e-admin-allproxy-${Date.now()}`; + await page.getByTestId("base-input").fill(keyName); + + // No team selection — leave team dropdown empty so the key is owned by the admin user + + // Select models — open the multi-select and pick the all-models meta-option. + // The Create Key modal labels this "All Team Models" even when no team is selected + // (see src/components/organisms/create_key_button.tsx:944), unlike the team/user + // settings screens which use "All Proxy Models". + await page.locator(".ant-select-selection-overflow").click(); + await page.locator(".ant-select-dropdown:visible").getByText("All Team Models").click(); + await page.keyboard.press("Escape"); + + await page.getByRole("button", { name: "Create Key", exact: true }).click(); + + await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); + await page.keyboard.press("Escape"); + + await expect(page.getByText(keyName)).toBeVisible({ timeout: 10_000 }); + }); + + test("Create a key with a specific proxy model (no team)", async ({ page }) => { + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + + await page.getByRole("button", { name: /Create New Key/i }).click(); + + await expect(page.getByText("Key Ownership")).toBeVisible({ timeout: 10_000 }); + + const keyName = `e2e-admin-specific-${Date.now()}`; + await page.getByTestId("base-input").fill(keyName); + + // Open the model multi-select and pick a single specific model. Use + // getByRole("option", ...) to avoid the strict-mode collision between + // the option container and its inner text node. + const modelName = "fake-openai-gpt-4"; + await page.locator(".ant-select-selection-overflow").click(); + const option = page.locator(".ant-select-dropdown:visible").getByRole("option", { name: modelName, exact: true }); + await option.waitFor({ state: "attached" }); + // Dispatch the click via the DOM — antd's dropdown can render the option + // off-viewport during the open animation, which trips Playwright's + // visibility/stability checks. The click handler fires regardless. + await option.evaluate((el: HTMLElement) => el.click()); + await page.keyboard.press("Escape"); + + await page.getByRole("button", { name: "Create Key", exact: true }).click(); + + await expect(page.getByText("Save your Key")).toBeVisible({ timeout: 10_000 }); + + // Grab the new key from the success modal (rendered inside a 

) and
+    // verify it can call /chat/completions for the model it was scoped to.
+    // The mock LLM server (fixtures/mock_llm_server/server.py) replies with
+    // a fixed "This is a mock response." body.
+    const apiKey = (await page.locator(".ant-modal:visible pre").innerText()).trim();
+    expect(apiKey).toMatch(/^sk-/);
+
+    const response = await page.request.post("/chat/completions", {
+      headers: { Authorization: `Bearer ${apiKey}` },
+      data: {
+        model: modelName,
+        messages: [{ role: "user", content: "ping" }],
+      },
+    });
+    expect(response.status()).toBe(200);
+    const body = await response.json();
+    expect(body.choices?.[0]?.message?.content).toBe("This is a mock response.");
+
+    await page.keyboard.press("Escape");
+
+    await expect(page.getByText(keyName)).toBeVisible({ timeout: 10_000 });
+  });
 });
diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts
new file mode 100644
index 00000000000..579b3cede7c
--- /dev/null
+++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/license.spec.ts
@@ -0,0 +1,37 @@
+import { test, expect } from "@playwright/test";
+import * as fs from "fs";
+import { ADMIN_STORAGE_PATH } from "../../constants";
+
+/**
+ * Sanity check that LITELLM_LICENSE is being forwarded to the proxy when set
+ * in the environment (e.g. CircleCI's `e2e_ui_testing` job). The login JWT's
+ * `premium_user` claim is the same value the dashboard reads to enable
+ * premium-gated UI surfaces (Team-BYOK switch, etc.), so asserting it here
+ * catches any future regression where the env var stops being plumbed
+ * through `run_e2e.sh` / `.circleci/config.yml`.
+ *
+ * Skips locally when no license is configured.
+ */
+test.describe("Premium license wiring", () => {
+  test("admin session JWT carries premium_user=true when LITELLM_LICENSE is set", () => {
+    test.skip(
+      !process.env.LITELLM_LICENSE,
+      "LITELLM_LICENSE not set in test env — proxy is running unlicensed",
+    );
+
+    const storage = JSON.parse(fs.readFileSync(ADMIN_STORAGE_PATH, "utf-8"));
+    const tokenCookie = storage.cookies?.find((c: { name: string }) => c.name === "token");
+    expect(tokenCookie, "token cookie missing from admin storage state").toBeDefined();
+
+    // Decode the JWT payload (no signature check — we trust globalSetup ran
+    // against our own proxy). Payload is the middle base64url segment.
+    const jwtParts = tokenCookie.value.split(".");
+    expect(jwtParts.length, "token cookie is not a 3-part JWT").toBe(3);
+    const [, payloadB64] = jwtParts;
+    const payload = JSON.parse(
+      Buffer.from(payloadB64, "base64url").toString("utf-8"),
+    );
+
+    expect(payload.premium_user).toBe(true);
+  });
+});
diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json
index b33b2a69bee..97bc797fd54 100644
--- a/ui/litellm-dashboard/package-lock.json
+++ b/ui/litellm-dashboard/package-lock.json
@@ -23,7 +23,7 @@
         "jwt-decode": "4.0.0",
         "lucide-react": "0.513.0",
         "moment": "2.30.1",
-        "next": "16.2.4",
+        "next": "16.2.6",
         "openai": "4.104.0",
         "papaparse": "5.5.3",
         "react": "18.3.1",
@@ -1883,9 +1883,9 @@
       }
     },
     "node_modules/@next/env": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/env/-/env-16.2.4.tgz",
-      "integrity": "sha512-dKkkOzOSwFYe5RX6y26fZgkSpVAlIOJKQHIiydQcrWH6y/97+RceSOAdjZ14Qa3zLduVUy0TXcn+EiM6t4rPgw==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/env/-/env-16.2.6.tgz",
+      "integrity": "sha512-gd8HoHN4ufj73WmR3JmVolrpJR47ILK6LouP5xElPglaVxir6e1a7VzvTvDWkOoPXT9rkkTzyCxBu4yeZfZwcw==",
       "license": "MIT"
     },
     "node_modules/@next/eslint-plugin-next": {
@@ -1899,9 +1899,9 @@
       }
     },
     "node_modules/@next/swc-darwin-arm64": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.2.4.tgz",
-      "integrity": "sha512-OXTFFox5EKN1Ym08vfrz+OXxmCcEjT4SFMbNRsWZE99dMqt2Kcusl5MqPXcW232RYkMLQTy0hqgAMEsfEd/l2A==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-16.2.6.tgz",
+      "integrity": "sha512-ZJGkkcNfYgrrMkqOdZ7zoLa1TOy0qpcMfk/z4Mh/FKUz40gVO+HNQWqmLxf67Z5WB64DRp0dhEbyHfel+6sJUg==",
       "cpu": [
         "arm64"
       ],
@@ -1915,9 +1915,9 @@
       }
     },
     "node_modules/@next/swc-darwin-x64": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.2.4.tgz",
-      "integrity": "sha512-XhpVnUfmYWvD3YrXu55XdcAkQtOnvaI6wtQa8fuF5fGoKoxIUZ0kWPtcOfqJEWngFF/lOS9l3+O9CcownhiQxQ==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-16.2.6.tgz",
+      "integrity": "sha512-v/YLBHIY132Ced3puBJ7YJKw1lqsCrgcNo2aRJlCEyQrrCeRJlvGlnmxhPxNQI3KE3N1DN5r9TPNPvka3nq5RQ==",
       "cpu": [
         "x64"
       ],
@@ -1931,15 +1931,12 @@
       }
     },
     "node_modules/@next/swc-linux-arm64-gnu": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.2.4.tgz",
-      "integrity": "sha512-Mx/tjlNA3G8kg14QvuGAJ4xBwPk1tUHq56JxZ8CXnZwz1Etz714soCEzGQQzVMz4bEnGPowzkV6Xrp6wAkEWOQ==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-16.2.6.tgz",
+      "integrity": "sha512-RPOvqlYBbcQjkz9VQQDZ2T2bARIjXZV1KFlt+V2Mr6SW/e4I9fcKsaA0hdyf2FHoTlsV2xnBd5Y912rP/1Ce6w==",
       "cpu": [
         "arm64"
       ],
-      "libc": [
-        "glibc"
-      ],
       "license": "MIT",
       "optional": true,
       "os": [
@@ -1950,15 +1947,12 @@
       }
     },
     "node_modules/@next/swc-linux-arm64-musl": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.2.4.tgz",
-      "integrity": "sha512-iVMMp14514u7Nup2umQS03nT/bN9HurK8ufylC3FZNykrwjtx7V1A7+4kvhbDSCeonTVqV3Txnv0Lu+m2oDXNg==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-16.2.6.tgz",
+      "integrity": "sha512-URUTu1+dMkxJsPFgm+OeEvq9wf5sujw0EvgYy80TDGHTSLTnIHeqb0Eu8A3sC95IRgjejQL+kC4mw+4yPxiAXA==",
       "cpu": [
         "arm64"
       ],
-      "libc": [
-        "musl"
-      ],
       "license": "MIT",
       "optional": true,
       "os": [
@@ -1969,15 +1963,12 @@
       }
     },
     "node_modules/@next/swc-linux-x64-gnu": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.2.4.tgz",
-      "integrity": "sha512-EZOvm1aQWgnI/N/xcWOlnS3RQBk0VtVav5Zo7n4p0A7UKyTDx047k8opDbXgBpHl4CulRqRfbw3QrX2w5UOXMQ==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-16.2.6.tgz",
+      "integrity": "sha512-DOj182mPV8G3UkrayLoREM5YEYI+Dk5wv7Ox9xl1fFibAELEsFD0lDPfHIeILlutMMfdyhlzYPELG3peuKaurw==",
       "cpu": [
         "x64"
       ],
-      "libc": [
-        "glibc"
-      ],
       "license": "MIT",
       "optional": true,
       "os": [
@@ -1988,15 +1979,12 @@
       }
     },
     "node_modules/@next/swc-linux-x64-musl": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.2.4.tgz",
-      "integrity": "sha512-h9FxsngCm9cTBf71AR4fGznDEDx1hS7+kSEiIRjq5kO1oXWm07DxVGZjCvk0SGx7TSjlUqhI8oOyz7NfwAdPoA==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-16.2.6.tgz",
+      "integrity": "sha512-HKQ5SP/V/ub73UvF7n/zeJlxk2kLmtL7Wzrg4WfmkjmNos5onJ2tKu7yZOPdL18A6Svfn3max29ym+ry7NkK4g==",
       "cpu": [
         "x64"
       ],
-      "libc": [
-        "musl"
-      ],
       "license": "MIT",
       "optional": true,
       "os": [
@@ -2007,9 +1995,9 @@
       }
     },
     "node_modules/@next/swc-win32-arm64-msvc": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.2.4.tgz",
-      "integrity": "sha512-3NdJV5OXMSOeJYijX+bjaLge3mJBlh4ybydbT4GFoB/2hAojWHtMhl3CYlYoMrjPuodp0nzFVi4Tj2+WaMg+Ow==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-16.2.6.tgz",
+      "integrity": "sha512-LZXpTlPyS5v7HhSmnvsLGP3iIYgYOBnc8r8ArlT55sGHV89bR2HlDdBjWQ+PY6SJMmk8TuVGFuxalnP3k/0Dwg==",
       "cpu": [
         "arm64"
       ],
@@ -2023,9 +2011,9 @@
       }
     },
     "node_modules/@next/swc-win32-x64-msvc": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.2.4.tgz",
-      "integrity": "sha512-kMVGgsqhO5YTYODD9IPGGhA6iprWidQckK3LmPeW08PIFENRmgfb4MjXHO+p//d+ts2rpjvK5gXWzXSMrPl9cw==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-16.2.6.tgz",
+      "integrity": "sha512-F0+4i0h9J6C4eE3EAPWsoCk7UW/dbzOjyzxY0qnDUOYFu6FFmdZ6l97/XdV3/Nz3VYyO7UWjyEJUXkGqcoXfMA==",
       "cpu": [
         "x64"
       ],
@@ -9316,12 +9304,12 @@
       "license": "MIT"
     },
     "node_modules/next": {
-      "version": "16.2.4",
-      "resolved": "https://registry.npmjs.org/next/-/next-16.2.4.tgz",
-      "integrity": "sha512-kPvz56wF5frc+FxlHI5qnklCzbq53HTwORaWBGdT0vNoKh1Aya9XC8aPauH4NJxqtzbWsS5mAbctm4cr+EkQ2Q==",
+      "version": "16.2.6",
+      "resolved": "https://registry.npmjs.org/next/-/next-16.2.6.tgz",
+      "integrity": "sha512-qOVgKJg1+At15NpeUP+eJgCHvTCgXsogweq87Ri/Ix7PkqQHg4sdaXmSFqKlgaIXE4kW0g25LE68W87UANlHtw==",
       "license": "MIT",
       "dependencies": {
-        "@next/env": "16.2.4",
+        "@next/env": "16.2.6",
         "@swc/helpers": "0.5.15",
         "baseline-browser-mapping": "^2.9.19",
         "caniuse-lite": "^1.0.30001579",
@@ -9335,14 +9323,14 @@
         "node": ">=20.9.0"
       },
       "optionalDependencies": {
-        "@next/swc-darwin-arm64": "16.2.4",
-        "@next/swc-darwin-x64": "16.2.4",
-        "@next/swc-linux-arm64-gnu": "16.2.4",
-        "@next/swc-linux-arm64-musl": "16.2.4",
-        "@next/swc-linux-x64-gnu": "16.2.4",
-        "@next/swc-linux-x64-musl": "16.2.4",
-        "@next/swc-win32-arm64-msvc": "16.2.4",
-        "@next/swc-win32-x64-msvc": "16.2.4",
+        "@next/swc-darwin-arm64": "16.2.6",
+        "@next/swc-darwin-x64": "16.2.6",
+        "@next/swc-linux-arm64-gnu": "16.2.6",
+        "@next/swc-linux-arm64-musl": "16.2.6",
+        "@next/swc-linux-x64-gnu": "16.2.6",
+        "@next/swc-linux-x64-musl": "16.2.6",
+        "@next/swc-win32-arm64-msvc": "16.2.6",
+        "@next/swc-win32-x64-msvc": "16.2.6",
         "sharp": "^0.34.5"
       },
       "peerDependencies": {
@@ -13345,16 +13333,6 @@
         "url": "https://github.com/sponsors/sindresorhus"
       }
     },
-    "node_modules/zod": {
-      "version": "3.25.76",
-      "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz",
-      "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==",
-      "extraneous": true,
-      "license": "MIT",
-      "funding": {
-        "url": "https://github.com/sponsors/colinhacks"
-      }
-    },
     "node_modules/zwitch": {
       "version": "2.0.4",
       "resolved": "https://registry.npmjs.org/zwitch/-/zwitch-2.0.4.tgz",
@@ -13364,21 +13342,6 @@
         "type": "github",
         "url": "https://github.com/sponsors/wooorm"
       }
-    },
-    "node_modules/@next/swc-win32-ia32-msvc": {
-      "version": "14.2.33",
-      "resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz",
-      "integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==",
-      "cpu": [
-        "ia32"
-      ],
-      "optional": true,
-      "os": [
-        "win32"
-      ],
-      "engines": {
-        "node": ">= 10"
-      }
     }
   }
 }
diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json
index 32c00ac62a8..72b9bc2a159 100644
--- a/ui/litellm-dashboard/package.json
+++ b/ui/litellm-dashboard/package.json
@@ -35,7 +35,7 @@
     "jwt-decode": "4.0.0",
     "lucide-react": "0.513.0",
     "moment": "2.30.1",
-    "next": "16.2.4",
+    "next": "16.2.6",
     "openai": "4.104.0",
     "papaparse": "5.5.3",
     "react": "18.3.1",
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
index 5431c196883..2626ace86d5 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
@@ -7,7 +7,7 @@ import { columns } from "@/components/molecules/models/columns";
 import { getDisplayModelName } from "@/components/view_model/model_name_display";
 import DeleteResourceModal from "@/components/common_components/DeleteResourceModal";
 import NotificationsManager from "@/components/molecules/notifications_manager";
-import { modelDeleteCall } from "@/components/networking";
+import { modelDeleteCall, modelPatchUpdateCall } from "@/components/networking";
 import { InfoCircleOutlined, SettingOutlined } from "@ant-design/icons";
 import { PaginationState, SortingState } from "@tanstack/react-table";
 import { useQueryClient } from "@tanstack/react-query";
@@ -220,6 +220,25 @@ const AllModelsTab = ({
     }
   };
 
+  const [pausingModelId, setPausingModelId] = useState(null);
+
+  const handleTogglePause = async (modelId: string, blocked: boolean) => {
+    if (!accessToken) return;
+    try {
+      setPausingModelId(modelId);
+      await modelPatchUpdateCall(accessToken, { blocked }, modelId);
+      NotificationsManager.success(blocked ? "Model paused" : "Model resumed");
+      // invalidateQueries already schedules a refetch for active observers
+      // on this key — no need to also call refetchModels() (would double-fetch).
+      queryClient.invalidateQueries({ queryKey: ["models", "list"] });
+    } catch (error) {
+      console.error("Error toggling model pause state:", error);
+      NotificationsManager.fromBackend(error);
+    } finally {
+      setPausingModelId(null);
+    }
+  };
+
   return (
     
       
@@ -536,6 +555,8 @@ const AllModelsTab = ({
                 expandedRows,
                 setExpandedRows,
                 setDeleteModalModelId,
+                handleTogglePause,
+                pausingModelId,
               )}
               data={filteredData}
               isLoading={isLoadingModelsInfo}
diff --git a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx
index 0539d6d8f19..3b3729c1ac9 100644
--- a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx
+++ b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx
@@ -4,11 +4,12 @@ import { Suspense, useEffect, useMemo } from "react";
 import { useSearchParams } from "next/navigation";
 import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
 
-// Written to sessionStorage so both the admin hook (useMcpOAuthFlow) and the
-// user hook (useUserMcpOAuthFlow) can pick up the result.  Each hook reads
-// its own namespace to avoid cross-flow collisions.
+// Written to sessionStorage so the admin hook (useMcpOAuthFlow), the user hook
+// (useUserMcpOAuthFlow), and the tools re-auth hook (useToolsOAuthFlow) can each
+// pick up the result.  Each hook reads its own namespace to avoid cross-flow collisions.
 const ADMIN_RESULT_KEY = "litellm-mcp-oauth-result";
 const USER_RESULT_KEY = "litellm-user-mcp-oauth-result";
+const TOOLS_RESULT_KEY = "litellm-tools-mcp-oauth-result";
 const RETURN_URL_STORAGE_KEY = "litellm-mcp-oauth-return-url";
 
 const resolveDefaultRedirect = () => {
@@ -50,11 +51,12 @@ const McpOAuthCallbackContent = () => {
     }
 
     try {
-      // Write to both namespace keys (admin and user) so whichever hook is
-      // active can consume the result.  sessionStorage only — no localStorage.
+      // Write to all namespace keys so whichever hook is active can consume
+      // the result.  sessionStorage only — no localStorage.
       const serialized = JSON.stringify(payload);
       setSecureItem(ADMIN_RESULT_KEY, serialized);
       setSecureItem(USER_RESULT_KEY, serialized);
+      setSecureItem(TOOLS_RESULT_KEY, serialized);
     } catch (err) {
       // Silently ignore storage errors
     }
diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx
index b9305e4723a..da00ad911b0 100644
--- a/ui/litellm-dashboard/src/components/OldTeams.tsx
+++ b/ui/litellm-dashboard/src/components/OldTeams.tsx
@@ -45,6 +45,7 @@ import OrganizationDropdown from "./common_components/OrganizationDropdown";
 import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton";
 import { teamListCall as v2TeamListCall, type TeamsResponse } from "@/app/(dashboard)/hooks/teams/useTeams";
 import AccessGroupSelector from "./common_components/AccessGroupSelector";
+import PassThroughRoutesSelector from "./common_components/PassThroughRoutesSelector";
 import AgentSelector from "./agent_management/AgentSelector";
 import ModelAliasManager from "./common_components/ModelAliasManager";
 import PremiumLoggingSettings from "./common_components/PremiumLoggingSettings";
@@ -1446,6 +1447,30 @@ const Teams: React.FC = ({
                           placeholder="Select vector stores (optional)"
                         />
                       
+                      
+                        
+                           form.setFieldValue("allowed_passthrough_routes", values)}
+                            value={form.getFieldValue("allowed_passthrough_routes")}
+                            accessToken={accessToken || ""}
+                            placeholder="Select pass through routes (optional)"
+                            disabled={!premiumUser || !isProxyAdminRole(userRole || "")}
+                          />
+                        
+                      
                     
                   
 
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx
index 7a592785a44..e98226b86bc 100644
--- a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx
@@ -182,16 +182,18 @@ export function ToolTestPanel({
 
     Object.entries(values).forEach(([key, value]) => {
       const prop = schemaToUse.properties?.[key];
-      if (prop && value !== null && value !== undefined && value !== "") {
+      // Strip leading/trailing whitespace from string inputs before submitting
+      const normalizedValue = typeof value === "string" ? value.trim() : value;
+      if (prop && normalizedValue !== null && normalizedValue !== undefined && normalizedValue !== "") {
         switch (prop.type) {
           case "boolean":
-            convertedValues[key] = value === "true" || value === true;
+            convertedValues[key] = normalizedValue === "true" || normalizedValue === true;
             break;
           case "number":
           case "integer": {
-            const numericValue = Number(value);
+            const numericValue = Number(normalizedValue);
             convertedValues[key] = Number.isNaN(numericValue)
-              ? value
+              ? normalizedValue
               : prop.type === "integer"
                 ? Math.trunc(numericValue)
                 : numericValue;
@@ -200,28 +202,28 @@ export function ToolTestPanel({
           case "object":
           case "array": {
             try {
-              const parsed = typeof value === "string" ? JSON.parse(value) : value;
+              const parsed = typeof normalizedValue === "string" ? JSON.parse(normalizedValue) : normalizedValue;
               const isValidObject =
                 prop.type === "object" && parsed !== null && typeof parsed === "object" && !Array.isArray(parsed);
               const isValidArray = prop.type === "array" && Array.isArray(parsed);
               if ((prop.type === "object" && isValidObject) || (prop.type === "array" && isValidArray)) {
                 convertedValues[key] = parsed;
               } else {
-                convertedValues[key] = value;
+                convertedValues[key] = normalizedValue;
               }
             } catch (err) {
-              convertedValues[key] = value;
+              convertedValues[key] = normalizedValue;
             }
             break;
           }
           case "string":
-            convertedValues[key] = String(value);
+            convertedValues[key] = String(normalizedValue);
             break;
           default:
-            convertedValues[key] = value;
+            convertedValues[key] = normalizedValue;
         }
-      } else if (value !== null && value !== undefined && value !== "") {
-        convertedValues[key] = value;
+      } else if (normalizedValue !== null && normalizedValue !== undefined && normalizedValue !== "") {
+        convertedValues[key] = normalizedValue;
       }
     });
 
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx
index f8b0141b25d..108911bdbf1 100644
--- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx
@@ -3,6 +3,7 @@ import { Modal, Tooltip, Form, Select, Input, Switch, Collapse } from "antd";
 import { InfoCircleOutlined } from "@ant-design/icons";
 import { Button, TextInput } from "@tremor/react";
 import { createMCPServer, registerMCPServer } from "../networking";
+import { setToken } from "@/utils/mcpTokenStore";
 import { AUTH_TYPE, DiscoverableMCPServer, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types";
 import OAuthFormFields from "./OAuthFormFields";
 import MCPServerCostConfig from "./mcp_server_cost_config";
@@ -24,6 +25,7 @@ export const mcpLogoImg = `${asset_logos_folder}mcp_logo.png`;
 
 interface CreateMCPServerProps {
   userRole: string;
+  userID?: string | null;
   accessToken: string | null;
   onCreateSuccess: (newMcpServer: MCPServer) => void;
   isModalVisible: boolean;
@@ -47,6 +49,7 @@ const reduceStaticHeaders = (list: unknown): Record => {
 };
 
 const CreateMCPServer: React.FC = ({
+  userID,
   userRole,
   accessToken,
   onCreateSuccess,
@@ -409,6 +412,21 @@ const CreateMCPServer: React.FC = ({
           ? await createMCPServer(accessToken, payload)
           : await registerMCPServer(accessToken, payload);
 
+        // Cache the OAuth token in sessionStorage so the Tools tab can use it
+        // immediately without re-authenticating.  No backend DB write.
+        if (oauthTokenResponse?.access_token && response?.server_id) {
+          setToken(
+            response.server_id,
+            {
+              access_token: oauthTokenResponse.access_token,
+              expires_in: oauthTokenResponse.expires_in,
+              refresh_token: oauthTokenResponse.refresh_token,
+              token_type: oauthTokenResponse.token_type,
+            },
+            userID,
+          );
+        }
+
         NotificationsManager.success(
           isAdmin
             ? "MCP Server created successfully"
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx
index 1f8f7f68d33..5a8035d4e0b 100644
--- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx
@@ -174,6 +174,7 @@ export const MCPServerView: React.FC = ({
               serverId={mcpServer.server_id}
               accessToken={accessToken}
               auth_type={mcpServer.auth_type}
+              tokenUrl={mcpServer.token_url}
               userRole={userRole}
               userID={userID}
               serverAlias={mcpServer.alias}
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx
index 72d5e4b5aa8..42583fdab07 100644
--- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx
@@ -287,6 +287,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID })
       
       (null);
   const [toolError, setToolError] = useState(null);
   const [toolSearchTerm, setToolSearchTerm] = useState("");
-  
+
   // State for passthrough headers
   const [passthroughHeaders, setPassthroughHeaders] = useState>({});
   const [showHeaderInput, setShowHeaderInput] = useState(false);
 
+  // OAuth session token (sessionStorage-backed, cleared on tab/browser close).
+  // Only the interactive (authorization_code/PKCE) flow needs a user-facing
+  // auth gate. M2M (client_credentials) servers are also `auth_type === "oauth2"`,
+  // but the backend fetches their token internally — gating tool listing on
+  // them would force users through a non-existent authorization endpoint.
+  // We detect M2M via the presence of `tokenUrl`, matching the heuristic in
+  // `mcp_server_edit.tsx`.
+  const isOAuth = auth_type === "oauth2" && !tokenUrl;
+  const [oauthToken, setOauthToken] = useState(() =>
+    isOAuth && isTokenValid(serverId, userID)
+      ? (getToken(serverId, userID)?.access_token ?? null)
+      : null
+  );
+
+  // Re-sync token when serverId/userID changes (useState initializer only runs on mount).
+  useEffect(() => {
+    if (!isOAuth) {
+      setOauthToken(null);
+      return;
+    }
+    setOauthToken(
+      isTokenValid(serverId, userID)
+        ? (getToken(serverId, userID)?.access_token ?? null)
+        : null
+    );
+  }, [serverId, userID, isOAuth]);
+
+  const { startOAuthFlow, status: oauthStatus, error: oauthError } = useToolsOAuthFlow({
+    accessToken: accessToken ?? "",
+    serverId,
+    serverAlias,
+    userId: userID,
+    onSuccess: setOauthToken,
+  });
+
   // Check if this server has extra headers configured
   const hasExtraHeaders = extraHeaders && extraHeaders.length > 0;
 
   // Build custom headers for MCP server requests
   const buildCustomHeaders = () => {
-    if (!serverAlias || !hasExtraHeaders) return undefined;
-    
     const customHeaders: Record = {};
-    
-    // Add passthrough headers with server-specific prefix
-    Object.entries(passthroughHeaders).forEach(([headerName, headerValue]) => {
-      if (headerValue && headerValue.trim()) {
-        // Format: x-mcp-{alias}-{header_name}
-        const mcpHeaderName = `x-mcp-${serverAlias}-${headerName.toLowerCase()}`;
-        customHeaders[mcpHeaderName] = headerValue;
+
+    // Include the session OAuth token using MCP-specific headers so it doesn't
+    // conflict with the Authorization header used by the LiteLLM proxy itself.
+    // The backend's _get_mcp_server_auth_headers_from_headers() picks up the
+    // x-mcp-{alias}-{header} pattern and forwards it to the upstream MCP server.
+    // When no alias is available, fall back to x-mcp-auth (legacy but still supported).
+    if (oauthToken) {
+      if (serverAlias) {
+        const safeAlias = sanitizeMcpAliasForHeader(serverAlias);
+        if (safeAlias) {
+          customHeaders[`x-mcp-${safeAlias}-authorization`] = `Bearer ${oauthToken}`;
+        } else {
+          customHeaders["x-mcp-auth"] = `Bearer ${oauthToken}`;
+        }
+      } else {
+        customHeaders["x-mcp-auth"] = `Bearer ${oauthToken}`;
       }
-    });
-    
+    }
+
+    // Add passthrough headers with server-specific prefix
+    if (serverAlias && hasExtraHeaders) {
+      const safeAlias = sanitizeMcpAliasForHeader(serverAlias);
+      if (safeAlias) {
+        Object.entries(passthroughHeaders).forEach(([headerName, headerValue]) => {
+          if (headerValue && headerValue.trim()) {
+            // Format: x-mcp-{alias}-{header_name}
+            const mcpHeaderName = `x-mcp-${safeAlias}-${headerName.toLowerCase()}`;
+            customHeaders[mcpHeaderName] = headerValue;
+          }
+        });
+      }
+    }
+
     return Object.keys(customHeaders).length > 0 ? customHeaders : undefined;
   };
 
@@ -54,15 +114,55 @@ const MCPToolsViewer = ({
     error: mcpToolsError,
     refetch: refetchTools,
   } = useQuery({
-    queryKey: ["mcpTools", serverId, passthroughHeaders],
-    queryFn: () => {
+    queryKey: ["mcpTools", serverId, passthroughHeaders, oauthToken],
+    queryFn: async () => {
       if (!accessToken) throw new Error("Access Token required");
-      return listMCPTools(accessToken, serverId, buildCustomHeaders());
+      const result = await listMCPTools(accessToken, serverId, buildCustomHeaders());
+      // listMCPTools never throws — surface error responses as thrown errors
+      // here so useQuery's retry/onError can react (e.g. clear the cached
+      // OAuth token on 401).
+      if (result?.error) {
+        const status = (result as { status?: number }).status;
+        if (status === 401) {
+          removeToken(serverId, userID);
+        }
+        const enhancedError = new Error(
+          result.message || result.error || "Failed to fetch MCP tools",
+        ) as Error & {
+          status?: number;
+          statusText?: string;
+          details?: any;
+        };
+        enhancedError.status = status;
+        enhancedError.statusText = (result as any).statusText;
+        enhancedError.details = (result as any).details;
+        throw enhancedError;
+      }
+      return result;
     },
-    enabled: !!accessToken,
+    // For OAuth servers, block the query until a session token is available
+    enabled: !!accessToken && (!isOAuth || oauthToken !== null),
     staleTime: 30000, // Consider data fresh for 30 seconds
+    retry: (failureCount, error: any) => {
+      // Don't retry on 401 — token is invalid, user must re-authenticate
+      if (error?.status === 401 || error?.response?.status === 401) return false;
+      return failureCount < 2;
+    },
   });
 
+  // If the tools query fails with 401, the cached OAuth token is invalid —
+  // clear it so the auth gate is shown again and the user can re-authenticate.
+  useEffect(() => {
+    const err = mcpToolsError as
+      | (Error & { status?: number; response?: { status?: number } })
+      | null;
+    const status = err?.status ?? err?.response?.status;
+    if (status === 401) {
+      removeToken(serverId, userID);
+      setOauthToken(null);
+    }
+  }, [mcpToolsError, serverId, userID]);
+
   // Mutation for calling a tool
   const { mutate: executeTool, isPending: isCallingTool } = useMutation({
     mutationFn: async (args: { tool: MCPTool; arguments: Record }) => {
@@ -85,9 +185,14 @@ const MCPToolsViewer = ({
       setToolResult(data.content);
       setToolError(null);
     },
-    onError: (error: Error) => {
+    onError: (error: Error & { status?: number; response?: { status?: number } }) => {
       setToolError(error);
       setToolResult(null);
+      // On 401, clear the cached token so the auth gate is shown again
+      if (error?.status === 401 || (error as any)?.response?.status === 401) {
+        removeToken(serverId, userID);
+        setOauthToken(null);
+      }
     },
   });
 
@@ -197,7 +302,31 @@ const MCPToolsViewer = ({
                   )}
                 
 
-                {/* Search Bar */}
+                {/* OAuth Auth Gate — shown when token is absent for OAuth servers */}
+                {isOAuth && !oauthToken && (
+                  

+                    
+                    
Authentication required

+                    

+                      Authenticate to view available tools
+                    

+                    
+                      Authorize
+                    
+                    {oauthError && (
+                      
{oauthError}

+                    )}
+                  

+                )}
+
+                {/* Search Bar — only shown when tools are loaded */}
+                {!isOAuth || oauthToken ? <>
                 {toolsData.length > 0 && (
                   

                     
-                    
Error: {mcpToolsResponse.message}

+                    

+                      Error: {mcpToolsResponse?.message || (mcpToolsError as Error)?.message}
+                    

                   

                 )}
 
                 {/* No Tools State */}
-                {!isLoadingTools && !mcpToolsResponse?.error && (!toolsData || toolsData.length === 0) && (
+                {!isLoadingTools && !mcpToolsResponse?.error && !mcpToolsError && (!toolsData || toolsData.length === 0) && (
                   

                     

                       
@@ -315,6 +446,7 @@ const MCPToolsViewer = ({
                     )}
                   
                 )}
+                 : null}
               

             

           
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx
index 7cfe08d9ee5..9a8f2e8f514 100644
--- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx
@@ -163,6 +163,13 @@ export interface MCPToolsViewerProps {
   serverId: string;
   accessToken: string | null;
   auth_type?: string | null;
+  /**
+   * When set, indicates the server uses the OAuth2 M2M (client_credentials)
+   * flow — the backend handles token acquisition internally, so the UI must
+   * not gate tool listing behind an interactive PKCE authorization. Mirrors
+   * the heuristic used in `mcp_server_edit.tsx` (`token_url` set => M2M).
+   */
+  tokenUrl?: string | null;
   userRole: string | null;
   userID: string | null;
   serverAlias?: string | null;
diff --git a/ui/litellm-dashboard/src/components/model_dashboard/types.ts b/ui/litellm-dashboard/src/components/model_dashboard/types.ts
index b1447a0634b..77a03d2c039 100644
--- a/ui/litellm-dashboard/src/components/model_dashboard/types.ts
+++ b/ui/litellm-dashboard/src/components/model_dashboard/types.ts
@@ -6,6 +6,7 @@ export interface ModelInfo {
   team_id: string;
   db_model: boolean;
   access_groups: string[] | null;
+  blocked?: boolean;
 }
 
 export interface LiteLLMParams {
diff --git a/ui/litellm-dashboard/src/components/molecules/models/columns.test.tsx b/ui/litellm-dashboard/src/components/molecules/models/columns.test.tsx
index a3dbbb2783f..c08dca1b8ce 100644
--- a/ui/litellm-dashboard/src/components/molecules/models/columns.test.tsx
+++ b/ui/litellm-dashboard/src/components/molecules/models/columns.test.tsx
@@ -944,4 +944,108 @@ describe("columns", () => {
     expect(screen.getByText("Out: $0.03")).toBeInTheDocument();
     expect(screen.queryByText(/In:/)).not.toBeInTheDocument();
   });
+
+  describe("pause/resume toggle", () => {
+    const renderWithToggle = (
+      overrides: Partial["model_info"]> = {},
+      togglePauseHandler?: ReturnType,
+      userRole: string = "Admin",
+    ) => {
+      const handler = togglePauseHandler ?? vi.fn();
+      const cols = columns(
+        userRole,
+        defaultProps.userID,
+        defaultProps.premiumUser,
+        defaultProps.setSelectedModelId,
+        defaultProps.setSelectedTeamId,
+        defaultProps.getDisplayModelName,
+        defaultProps.handleEditClick,
+        defaultProps.handleRefreshClick,
+        defaultProps.expandedRows,
+        defaultProps.setExpandedRows,
+        vi.fn(),
+        handler,
+      );
+      const model = createMockModel({
+        model_info: { ...createMockModel().model_info, ...overrides },
+      });
+      render();
+      return { handler };
+    };
+
+    it("renders the toggle ON for a db_model that is not blocked", () => {
+      renderWithToggle({ db_model: true, blocked: false });
+      const toggle = screen.getByRole("switch", { name: /pause model/i });
+      expect(toggle).toBeEnabled();
+      expect(toggle).toHaveAttribute("aria-checked", "true");
+    });
+
+    it("renders the toggle OFF for a db_model that is blocked", () => {
+      renderWithToggle({ db_model: true, blocked: true });
+      const toggle = screen.getByRole("switch", { name: /resume model/i });
+      expect(toggle).toBeEnabled();
+      expect(toggle).toHaveAttribute("aria-checked", "false");
+    });
+
+    it("calls the handler with blocked=true when an admin flips an active toggle off", async () => {
+      const handler = vi.fn();
+      renderWithToggle({ db_model: true, blocked: false }, handler);
+      await userEvent.click(screen.getByRole("switch", { name: /pause model/i }));
+      expect(handler).toHaveBeenCalledWith("test-model-id", true);
+    });
+
+    it("calls the handler with blocked=false when an admin flips a paused toggle on", async () => {
+      const handler = vi.fn();
+      renderWithToggle({ db_model: true, blocked: true }, handler);
+      await userEvent.click(screen.getByRole("switch", { name: /resume model/i }));
+      expect(handler).toHaveBeenCalledWith("test-model-id", false);
+    });
+
+    it("disables the toggle for non-admin users", () => {
+      const handler = vi.fn();
+      renderWithToggle({ db_model: true, blocked: false }, handler, "User");
+      const toggle = screen.getByRole("switch", { name: /pause model/i });
+      expect(toggle).toBeDisabled();
+    });
+
+    it("disables the toggle for config models", () => {
+      const handler = vi.fn();
+      renderWithToggle({ db_model: false, blocked: false }, handler, "Admin");
+      const toggle = screen.getByRole("switch", { name: /pause model/i });
+      expect(toggle).toBeDisabled();
+    });
+
+    it("disables the toggle while a PATCH for the same row is in-flight", () => {
+      // Regression for Greptile P1 on PR #28151 — antd's `loading` prop is
+      // visual only and does not prevent click events, so the row needs to
+      // be explicitly disabled while its PATCH is pending to avoid
+      // racing/conflicting PATCH calls on double-click.
+      const handler = vi.fn();
+      const model = createMockModel({
+        model_info: {
+          ...createMockModel().model_info,
+          db_model: true,
+          blocked: false,
+        },
+      });
+      const cols = columns(
+        "Admin",
+        defaultProps.userID,
+        defaultProps.premiumUser,
+        defaultProps.setSelectedModelId,
+        defaultProps.setSelectedTeamId,
+        defaultProps.getDisplayModelName,
+        defaultProps.handleEditClick,
+        defaultProps.handleRefreshClick,
+        defaultProps.expandedRows,
+        defaultProps.setExpandedRows,
+        vi.fn(),
+        handler,
+        model.model_info.id, // pausingModelId matches this row
+      );
+      render();
+      const toggle = screen.getByRole("switch", { name: /pause model/i });
+      expect(toggle).toBeDisabled();
+    });
+  });
 });
diff --git a/ui/litellm-dashboard/src/components/molecules/models/columns.tsx b/ui/litellm-dashboard/src/components/molecules/models/columns.tsx
index a303e1b1a44..4563e9c80dd 100644
--- a/ui/litellm-dashboard/src/components/molecules/models/columns.tsx
+++ b/ui/litellm-dashboard/src/components/molecules/models/columns.tsx
@@ -2,7 +2,7 @@ import { EditOutlined, InfoCircleOutlined, SyncOutlined } from "@ant-design/icon
 import { TrashIcon } from "@heroicons/react/outline";
 import { ColumnDef } from "@tanstack/react-table";
 import { Badge, Button, Icon } from "@tremor/react";
-import { Divider, Flex, Popover, Space, Tooltip, Typography } from "antd";
+import { Divider, Flex, Popover, Space, Switch, Tooltip, Typography } from "antd";
 import { ModelData } from "../../model_dashboard/types";
 import { ProviderLogo } from "./ProviderLogo";
 
@@ -53,6 +53,8 @@ export const columns = (
   expandedRows: Set,
   setExpandedRows: (expandedRows: Set) => void,
   onDeleteClick?: (modelId: string) => void,
+  onTogglePauseClick?: (modelId: string, blocked: boolean) => void | Promise,
+  pausingModelId?: string | null,
 ): ColumnDef[] => [
     {
       header: () => Model ID,
@@ -398,15 +400,48 @@ export const columns = (
     {
       id: "actions",
       header: () => Actions,
-      size: 60,
-      minSize: 40,
+      size: 100,
+      minSize: 80,
       enableResizing: false,
       cell: ({ row }) => {
         const model = row.original;
         const canEditModel = userRole === "Admin" || model.model_info?.created_by === userID;
         const isConfigModel = !model.model_info?.db_model;
+        const isAdmin = userRole === "Admin";
+        const isBlocked = model.model_info?.blocked === true;
+        const isPauseToggleable = !isConfigModel && isAdmin && Boolean(onTogglePauseClick);
+        const pauseTooltip = isConfigModel
+          ? "Config models cannot be paused from the dashboard. Pause is DB-backed."
+          : !isAdmin
+            ? "Only proxy admins can pause or resume a model."
+            : isBlocked
+              ? "Resume model — restore normal routing."
+              : "Pause model — stop routing requests until resumed.";
+        // antd's `loading` prop on Switch is purely cosmetic — it does not block
+        // clicks. Pair `loading` with `disabled` derived from the same condition
+        // so a double-click during a pending PATCH cannot send a second,
+        // conflicting `blocked` value.
+        const isPausing = pausingModelId === model.model_info?.id;
         return (
           

+            
+               {
+                  e.stopPropagation();
+                }}
+                onChange={(nextChecked) => {
+                  const modelId = model.model_info?.id;
+                  if (isPauseToggleable && onTogglePauseClick && modelId) {
+                    void onTogglePauseClick(modelId, !nextChecked);
+                  }
+                }}
+              />
+            
             {isConfigModel ? (
               
                 
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index 756348f4937..57e7d51123e 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -7068,46 +7068,33 @@ export const testSearchToolConnection = async (accessToken: string, litellmParam
 };
 
 export const listMCPTools = async (
-  accessToken: string, 
+  accessToken: string,
   serverId: string,
-  customHeaders?: Record
+  customHeaders?: Record,
 ) => {
+  // Construct base URL
+  let url = proxyBaseUrl
+    ? `${proxyBaseUrl}/mcp-rest/tools/list?server_id=${serverId}`
+    : `/mcp-rest/tools/list?server_id=${serverId}`;
+
+  console.log("Fetching MCP tools from:", url);
+
+  const headers: Record = {
+    [globalLitellmHeaderName]: `Bearer ${accessToken}`,
+    "Content-Type": "application/json",
+    ...customHeaders, // Merge custom headers for passthrough auth
+  };
+
+  let response: Response;
   try {
-    // Construct base URL
-    let url = proxyBaseUrl
-      ? `${proxyBaseUrl}/mcp-rest/tools/list?server_id=${serverId}`
-      : `/mcp-rest/tools/list?server_id=${serverId}`;
-
-    console.log("Fetching MCP tools from:", url);
-
-    const headers: Record = {
-      [globalLitellmHeaderName]: `Bearer ${accessToken}`,
-      "Content-Type": "application/json",
-      ...customHeaders, // Merge custom headers for passthrough auth
-    };
-
-    const response = await fetch(url, {
+    response = await fetch(url, {
       method: "GET",
       headers,
     });
-
-    const data = await response.json();
-    console.log("Fetched MCP tools response:", data);
-
-    if (!response.ok) {
-      // If the server returned an error response, use it
-      if (data.error && data.message) {
-        throw new Error(data.message);
-      }
-      // Otherwise use a generic error
-      throw new Error("Failed to fetch MCP tools");
-    }
-
-    // Return the full response object which includes tools, error, message, and stack_trace
-    return data;
   } catch (error) {
-    console.error("Failed to fetch MCP tools:", error);
-    // Return an error response in the same format as the API
+    // Network-level failure (no HTTP response). Preserve legacy shape so the
+    // caller can render a generic error message without crashing.
+    console.error("Failed to fetch MCP tools (network error):", error);
     return {
       tools: [],
       error: "network_error",
@@ -7115,6 +7102,44 @@ export const listMCPTools = async (
       stack_trace: null,
     };
   }
+
+  let data: any = null;
+  try {
+    data = await response.json();
+  } catch (parseError) {
+    console.error("Failed to parse MCP tools response:", parseError);
+    return {
+      tools: [],
+      error: "parse_error",
+      message: "Failed to parse MCP tools response",
+      status: response.status,
+      statusText: response.statusText,
+      stack_trace: null,
+    };
+  }
+  console.log("Fetched MCP tools response:", data);
+
+  if (!response.ok) {
+    // Preserve the legacy "never throws" contract so existing callers
+    // (e.g. MCPToolPermissions, MCPAppsPanel, MCPConnectPicker) can continue
+    // to inspect `result.error` / `result.message`. Attach `status` so
+    // callers that need to react to auth failures (e.g. the useQuery in
+    // mcp_tools.tsx) can still detect 401s from the returned object.
+    const errorMessage =
+      (data && (data.message || data.error)) || "Failed to fetch MCP tools";
+    return {
+      tools: [],
+      error: (data && data.error) || `http_${response.status}`,
+      message: errorMessage,
+      status: response.status,
+      statusText: response.statusText,
+      details: data,
+      stack_trace: null,
+    };
+  }
+
+  // Return the full response object which includes tools, error, message, and stack_trace
+  return data;
 };
 
 interface CallMCPToolOptions {
diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
index 8f3553a8395..ce9d15e1e19 100644
--- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx
@@ -502,6 +502,14 @@ const TeamInfoView: React.FC = ({
             (n) => !(values.guardrails || []).includes(n),
           );
 
+      // Non-proxy-admins can't set allowed_passthrough_routes; preserve the
+      // stored value so an unrelated save can't wipe it.
+      const passthroughRoutesMetadata = is_proxy_admin
+        ? { allowed_passthrough_routes: values.allowed_passthrough_routes || [] }
+        : info.metadata?.allowed_passthrough_routes
+          ? { allowed_passthrough_routes: info.metadata.allowed_passthrough_routes }
+          : {};
+
       const updateData: any = {
         team_id: teamId,
         team_alias: values.team_alias,
@@ -515,6 +523,7 @@ const TeamInfoView: React.FC = ({
         budget_duration: values.budget_duration,
         metadata: {
           ...parsedMetadata,
+          ...passthroughRoutesMetadata,
           guardrails: (values.guardrails || []).filter((n: string) => !globalGuardrailNames.has(n)),
           opted_out_global_guardrails: optedOutGlobalGuardrails,
           ...(values.logging_settings?.length > 0 ? { logging: values.logging_settings } : {}),
@@ -961,7 +970,7 @@ const TeamInfoView: React.FC = ({
                           : "",
                       metadata: info.metadata
                         ? JSON.stringify(
-                          (({ logging, secret_manager_settings, soft_budget_alerting_emails, model_tpm_limit, model_rpm_limit, ...rest }) => rest)(info.metadata),
+                          (({ logging, secret_manager_settings, soft_budget_alerting_emails, model_tpm_limit, model_rpm_limit, allowed_passthrough_routes, ...rest }) => rest)(info.metadata),
                           null,
                           2,
                         )
@@ -986,6 +995,7 @@ const TeamInfoView: React.FC = ({
                       },
                       access_group_ids: info.access_group_ids || [],
                       default_team_member_models: info.default_team_member_models || [],
+                      allowed_passthrough_routes: info.metadata?.allowed_passthrough_routes || [],
                     }}
                     layout="vertical"
                   >
@@ -1338,12 +1348,24 @@ const TeamInfoView: React.FC = ({
                     
 
                     
-                       form.setFieldValue("allowed_passthrough_routes", values)}
-                        value={form.getFieldValue("allowed_passthrough_routes")}
-                        accessToken={accessToken || ""}
-                        placeholder="Select pass through routes"
-                      />
+                      
+                         form.setFieldValue("allowed_passthrough_routes", values)}
+                          value={form.getFieldValue("allowed_passthrough_routes")}
+                          accessToken={accessToken || ""}
+                          placeholder="Select pass through routes"
+                          disabled={!premiumUser || !is_proxy_admin}
+                        />
+                      
                     
 
                     
diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx
index 03d917cd923..6c5fd03f0a0 100644
--- a/ui/litellm-dashboard/src/components/view_logs/index.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx
@@ -208,6 +208,8 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
   const deferredData = useDeferredValue(filteredData);
   const isStale = deferredData !== filteredData;
   const isButtonLoading = logsQuery.isFetching || isStale;
+  const isRefiltering = logsQuery.isPlaceholderData;
+  const isLogsLoading = logsQuery.isLoading || isRefiltering;
 
   if (!accessToken || !token || !userRole || !userID) {
     return (
@@ -277,7 +279,7 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
                     currentPage={currentPage}
                     onCurrentPageChange={setCurrentPage}
                     pageSize={pageSize}
-                    isLoading={logsQuery.isLoading}
+                    isLoading={isLogsLoading}
                     isButtonLoading={isButtonLoading}
                     onRefetch={() => logsQuery.refetch()}
                     filteredLogs={filteredLogs}
@@ -286,7 +288,7 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
                     columns={columns}
                     data={deferredData}
                     onRowClick={handleRowClick}
-                    isLoading={logsQuery.isLoading}
+                    isLoading={isLogsLoading}
                   />
                 

               
diff --git a/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts b/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts
new file mode 100644
index 00000000000..3aff8af6eef
--- /dev/null
+++ b/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts
@@ -0,0 +1,39 @@
+/**
+ * Shared utilities for MCP OAuth2 PKCE flow hooks.
+ *
+ * These helpers are used by both useToolsOAuthFlow and useUserMcpOAuthFlow
+ * to avoid divergence in URL construction and storage cleanup logic.
+ */
+
+import { getProxyBaseUrl, serverRootPath } from "@/components/networking";
+
+/**
+ * Build the OAuth callback URL for the current UI deployment.
+ *
+ * In the browser, derive the `/ui` prefix from the current pathname so the
+ * callback works regardless of how the proxy is mounted.  Outside the browser
+ * (SSR), fall back to the configured proxy base URL and server root path.
+ */
+export const buildCallbackUrl = (): string => {
+  if (typeof window !== "undefined") {
+    const path = window.location.pathname || "";
+    const idx = path.indexOf("/ui");
+    const prefix = idx >= 0 ? path.slice(0, idx + 3).replace(/\/+$/, "") : "";
+    return `${window.location.origin}${prefix}/mcp/oauth/callback`;
+  }
+  const base = (getProxyBaseUrl() || "").replace(/\/+$/, "");
+  const root = serverRootPath && serverRootPath !== "/" ? serverRootPath : "";
+  return `${base}${root}/ui/mcp/oauth/callback`;
+};
+
+/**
+ * Remove the given keys from sessionStorage, ignoring errors (e.g. storage
+ * disabled by browser privacy settings).
+ */
+export const clearStorage = (...keys: string[]): void => {
+  keys.forEach((k) => {
+    try {
+      window.sessionStorage.removeItem(k);
+    } catch (_) {}
+  });
+};
diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx
index 7edeade4cbd..11efaba53b3 100644
--- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx
+++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx
@@ -224,12 +224,21 @@ export const useMcpOAuthFlow = ({
       if (!storedPayload) {
         return;
       }
-      
+
+      // Guard: the callback page writes to the admin result key for *all* OAuth
+      // flows (including the tools re-auth flow).  Only proceed if this hook's
+      // own flow state exists, meaning startOAuthFlow() was actually called here.
+      // Without this guard, a tools re-auth redirect triggers a spurious
+      // "OAuth session state was lost" error from this hook.
+      const storedFlowState = getStorageItem(FLOW_STATE_KEY);
+      if (!storedFlowState) {
+        return;
+      }
+
       // Mark as processing
       processingRef.current = true;
       payload = JSON.parse(storedPayload);
-      const storedFlowState = getStorageItem(FLOW_STATE_KEY);
-      flowState = storedFlowState ? JSON.parse(storedFlowState) : null;
+      flowState = JSON.parse(storedFlowState);
     } catch (err) {
       clearStoredFlow();
       processingRef.current = false;
diff --git a/ui/litellm-dashboard/src/hooks/useToolsOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useToolsOAuthFlow.tsx
new file mode 100644
index 00000000000..66e59b80db4
--- /dev/null
+++ b/ui/litellm-dashboard/src/hooks/useToolsOAuthFlow.tsx
@@ -0,0 +1,232 @@
+"use client";
+
+/**
+ * OAuth2 PKCE flow for the Tools screen re-authentication path.
+ *
+ * Unlike useUserMcpOAuthFlow (used in the chat panel), this hook:
+ * - stores the resulting token in sessionStorage via mcpTokenStore only
+ * - does NOT call storeMCPOAuthUserCredential (no backend DB write)
+ * - uses "litellm-tools-mcp-oauth-result" as its result key to avoid
+ *   collisions with the admin and user flows
+ *
+ * The OAuth callback page (src/app/mcp/oauth/callback/page.tsx) writes
+ * to this key so this hook can pick up the result after the redirect.
+ */
+
+import { useCallback, useEffect, useRef, useState } from "react";
+import {
+  buildMcpOAuthAuthorizeUrl,
+  exchangeMcpOAuthToken,
+  registerMcpOAuthClient,
+} from "@/components/networking";
+import NotificationsManager from "@/components/molecules/notifications_manager";
+import { extractErrorMessage } from "@/utils/errorUtils";
+import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce";
+import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
+import { setToken } from "@/utils/mcpTokenStore";
+import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils";
+
+export type ToolsOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
+
+interface UseToolsOAuthFlowOptions {
+  accessToken: string;
+  serverId: string;
+  serverAlias?: string | null;
+  userId?: string | null;
+  scopes?: string[];
+  clientId?: string | null;
+  onSuccess: (accessToken: string) => void;
+}
+
+interface UseToolsOAuthFlowResult {
+  startOAuthFlow: () => Promise;
+  status: ToolsOAuthStatus;
+  error: string | null;
+}
+
+const FLOW_STATE_KEY = "litellm-tools-mcp-oauth-flow-state";
+const RESULT_KEY = "litellm-tools-mcp-oauth-result";
+const RETURN_URL_KEY = "litellm-mcp-oauth-return-url";
+
+type StoredFlowState = {
+  state: string;
+  codeVerifier: string;
+  serverId: string;
+  redirectUri: string;
+  clientId?: string;
+  clientSecret?: string;
+  scopes?: string[];
+};
+
+export const useToolsOAuthFlow = ({
+  accessToken,
+  serverId,
+  serverAlias,
+  userId,
+  scopes,
+  clientId: preClientId,
+  onSuccess,
+}: UseToolsOAuthFlowOptions): UseToolsOAuthFlowResult => {
+  const [status, setStatus] = useState("idle");
+  const [error, setError] = useState(null);
+  const processingRef = useRef(false);
+  const onSuccessRef = useRef(onSuccess);
+  onSuccessRef.current = onSuccess;
+
+  const startOAuthFlow = useCallback(async () => {
+    if (typeof window === "undefined") return;
+    try {
+      setStatus("authorizing");
+      setError(null);
+
+      let clientId: string | undefined = preClientId ?? undefined;
+      let clientSecret: string | undefined;
+
+      if (!clientId) {
+        try {
+          const reg = await registerMcpOAuthClient(accessToken, serverId, {
+            client_name: serverAlias || serverId,
+            grant_types: ["authorization_code", "refresh_token"],
+            response_types: ["code"],
+            token_endpoint_auth_method: "none",
+          });
+          clientId = reg?.client_id;
+          clientSecret = reg?.client_secret;
+        } catch (_) {
+          // Registration is optional; proceed without client_id
+        }
+      }
+
+      const verifier = generateCodeVerifier();
+      const challenge = await generateCodeChallenge(verifier);
+      const state = crypto.randomUUID();
+      const redirectUri = buildCallbackUrl();
+      const scopeString = scopes?.filter((s) => s.trim()).join(" ");
+
+      const authorizeUrl = buildMcpOAuthAuthorizeUrl({
+        serverId,
+        clientId,
+        redirectUri,
+        state,
+        codeChallenge: challenge,
+        scope: scopeString,
+      });
+
+      const flowState: StoredFlowState = {
+        state,
+        codeVerifier: verifier,
+        serverId,
+        redirectUri,
+        clientId,
+        clientSecret,
+        scopes,
+      };
+
+      setSecureItem(FLOW_STATE_KEY, JSON.stringify(flowState));
+      // Return to the current page (Tools tab) after the OAuth redirect
+      setSecureItem(RETURN_URL_KEY, window.location.href);
+
+      window.location.href = authorizeUrl;
+    } catch (err) {
+      const msg = extractErrorMessage(err);
+      setError(msg);
+      setStatus("error");
+      NotificationsManager.error(msg);
+    }
+  }, [accessToken, serverId, serverAlias, scopes, preClientId]);
+
+  const resumeOAuthFlow = useCallback(async () => {
+    if (typeof window === "undefined" || processingRef.current) return;
+
+    const storedResult = getSecureItem(RESULT_KEY);
+    if (!storedResult) return;
+
+    // The callback page writes to this result key for every OAuth flow (including
+    // the admin server-creation flow).  Guard: only proceed if *this* hook's flow
+    // state exists, meaning startOAuthFlow() was actually called from the Tools screen.
+    // Without this guard, a stale result written during server creation would trigger
+    // "OAuth session state was lost" when the user navigates to the Tools tab.
+    const rawFlowState = getSecureItem(FLOW_STATE_KEY);
+    if (!rawFlowState) return;
+
+    let peeked: StoredFlowState | null = null;
+    try {
+      peeked = JSON.parse(rawFlowState) as StoredFlowState;
+      if (peeked.serverId && peeked.serverId !== serverId) return;
+    } catch (_) {}
+
+    processingRef.current = true;
+    clearStorage(RESULT_KEY);
+
+    let payload: Record | null = null;
+    let flowState: StoredFlowState | null = null;
+
+    try {
+      payload = JSON.parse(storedResult);
+      flowState = peeked;
+    } catch (_) {
+      setError("Failed to resume OAuth flow. Please retry.");
+      setStatus("error");
+      processingRef.current = false;
+      clearStorage(FLOW_STATE_KEY);
+      return;
+    }
+
+    try {
+      if (!flowState?.state || !flowState.codeVerifier || !flowState.serverId) {
+        throw new Error("OAuth session state was lost. Please retry.");
+      }
+      if (!payload?.state || payload.state !== flowState.state) {
+        throw new Error("OAuth state mismatch. Please retry.");
+      }
+      if (payload.error) {
+        throw new Error((payload.error_description as string) || (payload.error as string));
+      }
+      if (!payload.code) {
+        throw new Error("Authorization code missing in callback.");
+      }
+
+      setStatus("exchanging");
+      const token = await exchangeMcpOAuthToken({
+        serverId: flowState.serverId,
+        code: payload.code as string,
+        clientId: flowState.clientId,
+        clientSecret: flowState.clientSecret,
+        codeVerifier: flowState.codeVerifier,
+        redirectUri: flowState.redirectUri,
+        accessToken,
+      });
+
+      // Store in sessionStorage only — no backend DB write
+      setToken(
+        flowState.serverId,
+        {
+          access_token: token.access_token,
+          expires_in: token.expires_in,
+          refresh_token: token.refresh_token,
+          token_type: token.token_type,
+        },
+        userId,
+      );
+
+      setStatus("success");
+      setError(null);
+      NotificationsManager.success("Connected successfully");
+      onSuccessRef.current(token.access_token);
+    } catch (err) {
+      const msg = extractErrorMessage(err);
+      setError(msg);
+      setStatus("error");
+      NotificationsManager.error(msg);
+    } finally {
+      clearStorage(FLOW_STATE_KEY);
+      setTimeout(() => { processingRef.current = false; }, 1000);
+    }
+  }, [accessToken, serverId, userId]);
+
+  useEffect(() => {
+    resumeOAuthFlow();
+  }, [resumeOAuthFlow]);
+
+  return { startOAuthFlow, status, error };
+};
diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx
index cf0a81dcadf..1dc7a5ee54a 100644
--- a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx
+++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx
@@ -16,15 +16,14 @@ import { useCallback, useEffect, useRef, useState } from "react";
 import {
   buildMcpOAuthAuthorizeUrl,
   exchangeMcpOAuthToken,
-  getProxyBaseUrl,
   registerMcpOAuthClient,
-  serverRootPath,
   storeMCPOAuthUserCredential,
 } from "@/components/networking";
 import NotificationsManager from "@/components/molecules/notifications_manager";
 import { extractErrorMessage } from "@/utils/errorUtils";
 import { generateCodeChallenge, generateCodeVerifier } from "@/utils/pkce";
 import { getSecureItem, setSecureItem } from "@/utils/secureStorage";
+import { buildCallbackUrl, clearStorage } from "./mcpOAuthUtils";
 
 export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error";
 
@@ -69,26 +68,6 @@ const getStorage = (key: string): string | null => {
   return getSecureItem(key);
 };
 
-const clearStorage = (...keys: string[]) => {
-  keys.forEach((k) => {
-    try {
-      window.sessionStorage.removeItem(k);
-    } catch (_) {}
-  });
-};
-
-const buildCallbackUrl = (): string => {
-  if (typeof window !== "undefined") {
-    const path = window.location.pathname || "";
-    const idx = path.indexOf("/ui");
-    const prefix = idx >= 0 ? path.slice(0, idx + 3).replace(/\/+$/, "") : "";
-    return `${window.location.origin}${prefix}/mcp/oauth/callback`;
-  }
-  const base = (getProxyBaseUrl() || "").replace(/\/+$/, "");
-  const root = serverRootPath && serverRootPath !== "/" ? serverRootPath : "";
-  return `${base}${root}/ui/mcp/oauth/callback`;
-};
-
 export const useUserMcpOAuthFlow = ({
   accessToken,
   serverId,
@@ -176,13 +155,17 @@ export const useUserMcpOAuthFlow = ({
     // mount and would compete for the same RESULT_KEY.  Peek at the stored
     // flow state first: only the hook instance whose serverId matches the one
     // that initiated the OAuth flow should consume the result.
+    // Guard: only proceed if this hook's flow state exists (startOAuthFlow was
+    // called from this hook).  Without the guard, a tools re-auth redirect writes
+    // to the user result key too, and every OAuth2ConnectButton instance would try
+    // to resume a flow that was never started here.
     const rawFlowState = getStorage(FLOW_STATE_KEY);
-    if (rawFlowState) {
-      try {
-        const peeked = JSON.parse(rawFlowState) as StoredFlowState;
-        if (peeked.serverId && peeked.serverId !== serverId) return;
-      } catch (_) {}
-    }
+    if (!rawFlowState) return;
+
+    try {
+      const peeked = JSON.parse(rawFlowState) as StoredFlowState;
+      if (peeked.serverId && peeked.serverId !== serverId) return;
+    } catch (_) {}
 
     processingRef.current = true;
     clearStorage(RESULT_KEY);
diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts
index c7bd27a6a85..28e2fc771c2 100644
--- a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts
+++ b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts
@@ -1,5 +1,6 @@
 import { describe, it, expect, beforeEach, vi } from "vitest";
 import { clearTokenCookies, getCookie, storeLoginToken } from "./cookieUtils";
+import { getToken, setToken } from "./mcpTokenStore";
 
 describe("cookieUtils", () => {
   beforeEach(() => {
@@ -20,6 +21,15 @@ describe("cookieUtils", () => {
       expect(getCookie("token")).toBeNull();
     });
 
+    it("should clear MCP session tokens on logout", () => {
+      setToken("server-1", { access_token: "mcp-tok" }, "user-a");
+      expect(getToken("server-1", "user-a")).not.toBeNull();
+
+      clearTokenCookies();
+
+      expect(getToken("server-1", "user-a")).toBeNull();
+    });
+
     it("should clear token cookie from /ui path", () => {
       document.cookie = "token=test-token-value; path=/ui";
       clearTokenCookies();
diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.ts b/ui/litellm-dashboard/src/utils/cookieUtils.ts
index b4493744ad4..da232e72e2a 100644
--- a/ui/litellm-dashboard/src/utils/cookieUtils.ts
+++ b/ui/litellm-dashboard/src/utils/cookieUtils.ts
@@ -2,6 +2,8 @@
  * Utility functions for managing cookies
  */
 
+import { clearAllMcpTokens } from "./mcpTokenStore";
+
 /**
  * Returns the cookie path for the UI.
  * Derives the path from window.location.pathname so it works when
@@ -67,6 +69,7 @@ export function clearTokenCookies() {
     // sessionStorage may be unavailable
   }
 
+  clearAllMcpTokens();
 }
 
 /**
diff --git a/ui/litellm-dashboard/src/utils/mcpHeaderUtils.test.ts b/ui/litellm-dashboard/src/utils/mcpHeaderUtils.test.ts
new file mode 100644
index 00000000000..b730b9c097c
--- /dev/null
+++ b/ui/litellm-dashboard/src/utils/mcpHeaderUtils.test.ts
@@ -0,0 +1,16 @@
+import { describe, expect, it } from "vitest";
+import { sanitizeMcpAliasForHeader } from "./mcpHeaderUtils";
+
+describe("sanitizeMcpAliasForHeader", () => {
+  it("lowercases and replaces spaces with underscores", () => {
+    expect(sanitizeMcpAliasForHeader("My Server")).toBe("my_server");
+  });
+
+  it("replaces invalid characters for header token segments", () => {
+    expect(sanitizeMcpAliasForHeader("GitHub-MCP!")).toBe("github_mcp");
+  });
+
+  it("preserves underscores and digits", () => {
+    expect(sanitizeMcpAliasForHeader("github_mcp2")).toBe("github_mcp2");
+  });
+});
diff --git a/ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts b/ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts
new file mode 100644
index 00000000000..76c752a4b61
--- /dev/null
+++ b/ui/litellm-dashboard/src/utils/mcpHeaderUtils.ts
@@ -0,0 +1,14 @@
+/**
+ * Sanitize an MCP server alias for use in HTTP header names (x-mcp-{alias}-...).
+ * RFC 7230 tchar allows token chars; aliases with spaces or hyphens break parsing
+ * because the backend splits on the first dash after the x-mcp- prefix.
+ * Keep in sync with litellm.proxy._experimental.mcp_server.utils.sanitize_mcp_alias_for_header.
+ */
+export function sanitizeMcpAliasForHeader(alias: string): string {
+  return alias
+    .toLowerCase()
+    .trim()
+    .replace(/[^a-z0-9_]/g, "_")
+    .replace(/_+/g, "_")
+    .replace(/^_|_$/g, "");
+}
diff --git a/ui/litellm-dashboard/src/utils/mcpTokenStore.test.ts b/ui/litellm-dashboard/src/utils/mcpTokenStore.test.ts
new file mode 100644
index 00000000000..1c61e9b1a26
--- /dev/null
+++ b/ui/litellm-dashboard/src/utils/mcpTokenStore.test.ts
@@ -0,0 +1,46 @@
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import {
+  clearAllMcpTokens,
+  getToken,
+  isTokenValid,
+  removeToken,
+  setToken,
+} from "./mcpTokenStore";
+
+describe("mcpTokenStore", () => {
+  beforeEach(() => {
+    sessionStorage.clear();
+  });
+
+  afterEach(() => {
+    sessionStorage.clear();
+  });
+
+  it("scopes tokens by user id", () => {
+    setToken("server-a", { access_token: "user1-token" }, "user-1");
+    setToken("server-a", { access_token: "user2-token" }, "user-2");
+
+    expect(getToken("server-a", "user-1")?.access_token).toBe("user1-token");
+    expect(getToken("server-a", "user-2")?.access_token).toBe("user2-token");
+    expect(getToken("server-a", "user-3")).toBeNull();
+  });
+
+  it("validates expiry per user scope", () => {
+    setToken("server-a", { access_token: "tok", expires_in: 3600 }, "user-1");
+    expect(isTokenValid("server-a", "user-1")).toBe(true);
+    removeToken("server-a", "user-1");
+    expect(isTokenValid("server-a", "user-1")).toBe(false);
+  });
+
+  it("clearAllMcpTokens removes every mcp-session-token entry", () => {
+    setToken("s1", { access_token: "a" }, "u1");
+    setToken("s2", { access_token: "b" }, "u2");
+    sessionStorage.setItem("unrelated", "keep");
+
+    clearAllMcpTokens();
+
+    expect(getToken("s1", "u1")).toBeNull();
+    expect(getToken("s2", "u2")).toBeNull();
+    expect(sessionStorage.getItem("unrelated")).toBe("keep");
+  });
+});
diff --git a/ui/litellm-dashboard/src/utils/mcpTokenStore.ts b/ui/litellm-dashboard/src/utils/mcpTokenStore.ts
new file mode 100644
index 00000000000..0279922cd07
--- /dev/null
+++ b/ui/litellm-dashboard/src/utils/mcpTokenStore.ts
@@ -0,0 +1,93 @@
+/**
+ * Session-storage-backed OAuth token store for MCP servers.
+ * Tokens are keyed by LiteLLM user id + server_id and cleared when the browser
+ * session ends (tab/window close). Never written to localStorage.
+ */
+
+const KEY_PREFIX = "mcp-session-token:";
+
+interface StoredToken {
+  access_token: string;
+  expires_at: number;
+  refresh_token?: string;
+  token_type: string;
+}
+
+interface TokenInput {
+  access_token: string;
+  expires_in?: number;
+  refresh_token?: string;
+  token_type?: string;
+}
+
+const DEFAULT_TTL_MS = 3600 * 1000; // 1 hour
+
+function storageKey(serverId: string, userId?: string | null): string {
+  const userPart = userId?.trim() || "_anonymous";
+  return `${KEY_PREFIX}${userPart}:${serverId}`;
+}
+
+export function setToken(
+  serverId: string,
+  data: TokenInput,
+  userId?: string | null,
+): void {
+  if (typeof window === "undefined") return;
+  const stored: StoredToken = {
+    access_token: data.access_token,
+    expires_at: Date.now() + (data.expires_in != null ? data.expires_in * 1000 : DEFAULT_TTL_MS),
+    token_type: data.token_type ?? "bearer",
+    ...(data.refresh_token ? { refresh_token: data.refresh_token } : {}),
+  };
+  try {
+    window.sessionStorage.setItem(storageKey(serverId, userId), JSON.stringify(stored));
+  } catch {
+    // Silently ignore storage errors (private browsing, quota exceeded, etc.)
+  }
+}
+
+export function getToken(
+  serverId: string,
+  userId?: string | null,
+): StoredToken | null {
+  if (typeof window === "undefined") return null;
+  try {
+    const raw = window.sessionStorage.getItem(storageKey(serverId, userId));
+    if (!raw) return null;
+    return JSON.parse(raw) as StoredToken;
+  } catch {
+    return null;
+  }
+}
+
+export function removeToken(serverId: string, userId?: string | null): void {
+  if (typeof window === "undefined") return;
+  try {
+    window.sessionStorage.removeItem(storageKey(serverId, userId));
+  } catch {
+    // Silently ignore
+  }
+}
+
+export function isTokenValid(serverId: string, userId?: string | null): boolean {
+  const token = getToken(serverId, userId);
+  if (!token) return false;
+  return token.expires_at > Date.now();
+}
+
+/** Remove all MCP session tokens (e.g. on logout or user switch). */
+export function clearAllMcpTokens(): void {
+  if (typeof window === "undefined") return;
+  try {
+    const keysToRemove: string[] = [];
+    for (let i = 0; i < window.sessionStorage.length; i++) {
+      const key = window.sessionStorage.key(i);
+      if (key?.startsWith(KEY_PREFIX)) {
+        keysToRemove.push(key);
+      }
+    }
+    keysToRemove.forEach((key) => window.sessionStorage.removeItem(key));
+  } catch {
+    // Silently ignore
+  }
+}
diff --git a/uv.lock b/uv.lock
index cafb6664958..fe3e0e037cf 100644
--- a/uv.lock
+++ b/uv.lock
@@ -9,7 +9,7 @@ resolution-markers = [
 ]
 
 [options]
-exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values.
+exclude-newer = "2026-05-19T01:14:41.559325863Z"
 exclude-newer-span = "P3D"
 
 [manifest]
@@ -539,7 +539,7 @@ wheels = [
 
 [[package]]
 name = "black"
-version = "24.10.0"
+version = "26.3.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "click" },
@@ -547,28 +547,33 @@ dependencies = [
     { name = "packaging" },
     { name = "pathspec" },
     { name = "platformdirs" },
+    { name = "pytokens" },
     { name = "tomli", marker = "python_full_version < '3.11'" },
     { name = "typing-extensions", marker = "python_full_version < '3.11'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/d8/0d/cc2fb42b8c50d80143221515dd7e4766995bd07c56c9a3ed30baf080b6dc/black-24.10.0.tar.gz", hash = "sha256:846ea64c97afe3bc677b761787993be4991810ecc7a4a937816dd6bddedc4875", size = 645813, upload-time = "2024-10-07T19:20:50.361Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e1/c5/61175d618685d42b005847464b8fb4743a67b1b8fdb75e50e5a96c31a27a/black-26.3.1.tar.gz", hash = "sha256:2c50f5063a9641c7eed7795014ba37b0f5fa227f3d408b968936e24bc0566b07", size = 666155, upload-time = "2026-03-12T03:36:03.593Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/a3/f3/465c0eb5cddf7dbbfe1fecd9b875d1dcf51b88923cd2c1d7e9ab95c6336b/black-24.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6668650ea4b685440857138e5fe40cde4d652633b1bdffc62933d0db4ed9812", size = 1623211, upload-time = "2024-10-07T19:26:12.43Z" },
-    { url = "https://files.pythonhosted.org/packages/df/57/b6d2da7d200773fdfcc224ffb87052cf283cec4d7102fab450b4a05996d8/black-24.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1c536fcf674217e87b8cc3657b81809d3c085d7bf3ef262ead700da345bfa6ea", size = 1457139, upload-time = "2024-10-07T19:25:06.453Z" },
-    { url = "https://files.pythonhosted.org/packages/6e/c5/9023b7673904a5188f9be81f5e129fff69f51f5515655fbd1d5a4e80a47b/black-24.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:649fff99a20bd06c6f727d2a27f401331dc0cc861fb69cde910fe95b01b5928f", size = 1753774, upload-time = "2024-10-07T19:23:58.47Z" },
-    { url = "https://files.pythonhosted.org/packages/e1/32/df7f18bd0e724e0d9748829765455d6643ec847b3f87e77456fc99d0edab/black-24.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:fe4d6476887de70546212c99ac9bd803d90b42fc4767f058a0baa895013fbb3e", size = 1414209, upload-time = "2024-10-07T19:24:42.54Z" },
-    { url = "https://files.pythonhosted.org/packages/c2/cc/7496bb63a9b06a954d3d0ac9fe7a73f3bf1cd92d7a58877c27f4ad1e9d41/black-24.10.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5a2221696a8224e335c28816a9d331a6c2ae15a2ee34ec857dcf3e45dbfa99ad", size = 1607468, upload-time = "2024-10-07T19:26:14.966Z" },
-    { url = "https://files.pythonhosted.org/packages/2b/e3/69a738fb5ba18b5422f50b4f143544c664d7da40f09c13969b2fd52900e0/black-24.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f9da3333530dbcecc1be13e69c250ed8dfa67f43c4005fb537bb426e19200d50", size = 1437270, upload-time = "2024-10-07T19:25:24.291Z" },
-    { url = "https://files.pythonhosted.org/packages/c9/9b/2db8045b45844665c720dcfe292fdaf2e49825810c0103e1191515fc101a/black-24.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4007b1393d902b48b36958a216c20c4482f601569d19ed1df294a496eb366392", size = 1737061, upload-time = "2024-10-07T19:23:52.18Z" },
-    { url = "https://files.pythonhosted.org/packages/a3/95/17d4a09a5be5f8c65aa4a361444d95edc45def0de887810f508d3f65db7a/black-24.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:394d4ddc64782e51153eadcaaca95144ac4c35e27ef9b0a42e121ae7e57a9175", size = 1423293, upload-time = "2024-10-07T19:24:41.7Z" },
-    { url = "https://files.pythonhosted.org/packages/90/04/bf74c71f592bcd761610bbf67e23e6a3cff824780761f536512437f1e655/black-24.10.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e39e0fae001df40f95bd8cc36b9165c5e2ea88900167bddf258bacef9bbdc3", size = 1644256, upload-time = "2024-10-07T19:27:53.355Z" },
-    { url = "https://files.pythonhosted.org/packages/4c/ea/a77bab4cf1887f4b2e0bce5516ea0b3ff7d04ba96af21d65024629afedb6/black-24.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d37d422772111794b26757c5b55a3eade028aa3fde43121ab7b673d050949d65", size = 1448534, upload-time = "2024-10-07T19:26:44.953Z" },
-    { url = "https://files.pythonhosted.org/packages/4e/3e/443ef8bc1fbda78e61f79157f303893f3fddf19ca3c8989b163eb3469a12/black-24.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14b3502784f09ce2443830e3133dacf2c0110d45191ed470ecb04d0f5f6fcb0f", size = 1761892, upload-time = "2024-10-07T19:24:10.264Z" },
-    { url = "https://files.pythonhosted.org/packages/52/93/eac95ff229049a6901bc84fec6908a5124b8a0b7c26ea766b3b8a5debd22/black-24.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:30d2c30dc5139211dda799758559d1b049f7f14c580c409d6ad925b74a4208a8", size = 1434796, upload-time = "2024-10-07T19:25:06.239Z" },
-    { url = "https://files.pythonhosted.org/packages/d0/a0/a993f58d4ecfba035e61fca4e9f64a2ecae838fc9f33ab798c62173ed75c/black-24.10.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1cbacacb19e922a1d75ef2b6ccaefcd6e93a2c05ede32f06a21386a04cedb981", size = 1643986, upload-time = "2024-10-07T19:28:50.684Z" },
-    { url = "https://files.pythonhosted.org/packages/37/d5/602d0ef5dfcace3fb4f79c436762f130abd9ee8d950fa2abdbf8bbc555e0/black-24.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1f93102e0c5bb3907451063e08b9876dbeac810e7da5a8bfb7aeb5a9ef89066b", size = 1448085, upload-time = "2024-10-07T19:28:12.093Z" },
-    { url = "https://files.pythonhosted.org/packages/47/6d/a3a239e938960df1a662b93d6230d4f3e9b4a22982d060fc38c42f45a56b/black-24.10.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ddacb691cdcdf77b96f549cf9591701d8db36b2f19519373d60d31746068dbf2", size = 1760928, upload-time = "2024-10-07T19:24:15.233Z" },
-    { url = "https://files.pythonhosted.org/packages/dd/cf/af018e13b0eddfb434df4d9cd1b2b7892bab119f7a20123e93f6910982e8/black-24.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:680359d932801c76d2e9c9068d05c6b107f2584b2a5b88831c83962eb9984c1b", size = 1436875, upload-time = "2024-10-07T19:24:42.762Z" },
-    { url = "https://files.pythonhosted.org/packages/8d/a7/4b27c50537ebca8bec139b872861f9d2bf501c5ec51fcf897cb924d9e264/black-24.10.0-py3-none-any.whl", hash = "sha256:3bb2b7a1f7b685f85b11fed1ef10f8a9148bceb49853e47a294a3dd963c1dd7d", size = 206898, upload-time = "2024-10-07T19:20:48.317Z" },
+    { url = "https://files.pythonhosted.org/packages/32/a8/11170031095655d36ebc6664fe0897866f6023892396900eec0e8fdc4299/black-26.3.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:86a8b5035fce64f5dcd1b794cf8ec4d31fe458cf6ce3986a30deb434df82a1d2", size = 1866562, upload-time = "2026-03-12T03:39:58.639Z" },
+    { url = "https://files.pythonhosted.org/packages/69/ce/9e7548d719c3248c6c2abfd555d11169457cbd584d98d179111338423790/black-26.3.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:5602bdb96d52d2d0672f24f6ffe5218795736dd34807fd0fd55ccd6bf206168b", size = 1703623, upload-time = "2026-03-12T03:40:00.347Z" },
+    { url = "https://files.pythonhosted.org/packages/7f/0a/8d17d1a9c06f88d3d030d0b1d4373c1551146e252afe4547ed601c0e697f/black-26.3.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6c54a4a82e291a1fee5137371ab488866b7c86a3305af4026bdd4dc78642e1ac", size = 1768388, upload-time = "2026-03-12T03:40:01.765Z" },
+    { url = "https://files.pythonhosted.org/packages/52/79/c1ee726e221c863cde5164f925bacf183dfdf0397d4e3f94889439b947b4/black-26.3.1-cp310-cp310-win_amd64.whl", hash = "sha256:6e131579c243c98f35bce64a7e08e87fb2d610544754675d4a0e73a070a5aa3a", size = 1412969, upload-time = "2026-03-12T03:40:03.252Z" },
+    { url = "https://files.pythonhosted.org/packages/73/a5/15c01d613f5756f68ed8f6d4ec0a1e24b82b18889fa71affd3d1f7fad058/black-26.3.1-cp310-cp310-win_arm64.whl", hash = "sha256:5ed0ca58586c8d9a487352a96b15272b7fa55d139fc8496b519e78023a8dab0a", size = 1220345, upload-time = "2026-03-12T03:40:04.892Z" },
+    { url = "https://files.pythonhosted.org/packages/17/57/5f11c92861f9c92eb9dddf515530bc2d06db843e44bdcf1c83c1427824bc/black-26.3.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:28ef38aee69e4b12fda8dba75e21f9b4f979b490c8ac0baa7cb505369ac9e1ff", size = 1851987, upload-time = "2026-03-12T03:40:06.248Z" },
+    { url = "https://files.pythonhosted.org/packages/54/aa/340a1463660bf6831f9e39646bf774086dbd8ca7fc3cded9d59bbdf4ad0a/black-26.3.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bf9bf162ed91a26f1adba8efda0b573bc6924ec1408a52cc6f82cb73ec2b142c", size = 1689499, upload-time = "2026-03-12T03:40:07.642Z" },
+    { url = "https://files.pythonhosted.org/packages/f3/01/b726c93d717d72733da031d2de10b92c9fa4c8d0c67e8a8a372076579279/black-26.3.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:474c27574d6d7037c1bc875a81d9be0a9a4f9ee95e62800dab3cfaadbf75acd5", size = 1754369, upload-time = "2026-03-12T03:40:09.279Z" },
+    { url = "https://files.pythonhosted.org/packages/e3/09/61e91881ca291f150cfc9eb7ba19473c2e59df28859a11a88248b5cbbc4d/black-26.3.1-cp311-cp311-win_amd64.whl", hash = "sha256:5e9d0d86df21f2e1677cc4bd090cd0e446278bcbbe49bf3659c308c3e402843e", size = 1413613, upload-time = "2026-03-12T03:40:10.943Z" },
+    { url = "https://files.pythonhosted.org/packages/16/73/544f23891b22e7efe4d8f812371ab85b57f6a01b2fc45e3ba2e52ba985b8/black-26.3.1-cp311-cp311-win_arm64.whl", hash = "sha256:9a5e9f45e5d5e1c5b5c29b3bd4265dcc90e8b92cf4534520896ed77f791f4da5", size = 1219719, upload-time = "2026-03-12T03:40:12.597Z" },
+    { url = "https://files.pythonhosted.org/packages/dc/f8/da5eae4fc75e78e6dceb60624e1b9662ab00d6b452996046dfa9b8a6025b/black-26.3.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e6f89631eb88a7302d416594a32faeee9fb8fb848290da9d0a5f2903519fc1", size = 1895920, upload-time = "2026-03-12T03:40:13.921Z" },
+    { url = "https://files.pythonhosted.org/packages/2c/9f/04e6f26534da2e1629b2b48255c264cabf5eedc5141d04516d9d68a24111/black-26.3.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:41cd2012d35b47d589cb8a16faf8a32ef7a336f56356babd9fcf70939ad1897f", size = 1718499, upload-time = "2026-03-12T03:40:15.239Z" },
+    { url = "https://files.pythonhosted.org/packages/04/91/a5935b2a63e31b331060c4a9fdb5a6c725840858c599032a6f3aac94055f/black-26.3.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f76ff19ec5297dd8e66eb64deda23631e642c9393ab592826fd4bdc97a4bce7", size = 1794994, upload-time = "2026-03-12T03:40:17.124Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/0a/86e462cdd311a3c2a8ece708d22aba17d0b2a0d5348ca34b40cdcbea512e/black-26.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:ddb113db38838eb9f043623ba274cfaf7d51d5b0c22ecb30afe58b1bb8322983", size = 1420867, upload-time = "2026-03-12T03:40:18.83Z" },
+    { url = "https://files.pythonhosted.org/packages/5b/e5/22515a19cb7eaee3440325a6b0d95d2c0e88dd180cb011b12ae488e031d1/black-26.3.1-cp312-cp312-win_arm64.whl", hash = "sha256:dfdd51fc3e64ea4f35873d1b3fb25326773d55d2329ff8449139ebaad7357efb", size = 1230124, upload-time = "2026-03-12T03:40:20.425Z" },
+    { url = "https://files.pythonhosted.org/packages/f5/77/5728052a3c0450c53d9bb3945c4c46b91baa62b2cafab6801411b6271e45/black-26.3.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:855822d90f884905362f602880ed8b5df1b7e3ee7d0db2502d4388a954cc8c54", size = 1895034, upload-time = "2026-03-12T03:40:21.813Z" },
+    { url = "https://files.pythonhosted.org/packages/52/73/7cae55fdfdfbe9d19e9a8d25d145018965fe2079fa908101c3733b0c55a0/black-26.3.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8a33d657f3276328ce00e4d37fe70361e1ec7614da5d7b6e78de5426cb56332f", size = 1718503, upload-time = "2026-03-12T03:40:23.666Z" },
+    { url = "https://files.pythonhosted.org/packages/e1/87/af89ad449e8254fdbc74654e6467e3c9381b61472cc532ee350d28cfdafb/black-26.3.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f1cd08e99d2f9317292a311dfe578fd2a24b15dbce97792f9c4d752275c1fa56", size = 1793557, upload-time = "2026-03-12T03:40:25.497Z" },
+    { url = "https://files.pythonhosted.org/packages/43/10/d6c06a791d8124b843bf325ab4ac7d2f5b98731dff84d6064eafd687ded1/black-26.3.1-cp313-cp313-win_amd64.whl", hash = "sha256:c7e72339f841b5a237ff14f7d3880ddd0fc7f98a1199e8c4327f9a4f478c1839", size = 1422766, upload-time = "2026-03-12T03:40:27.14Z" },
+    { url = "https://files.pythonhosted.org/packages/59/4f/40a582c015f2d841ac24fed6390bd68f0fc896069ff3a886317959c9daf8/black-26.3.1-cp313-cp313-win_arm64.whl", hash = "sha256:afc622538b430aa4c8c853f7f63bc582b3b8030fd8c80b70fb5fa5b834e575c2", size = 1232140, upload-time = "2026-03-12T03:40:28.882Z" },
+    { url = "https://files.pythonhosted.org/packages/8e/0d/52d98722666d6fc6c3dd4c76df339501d6efd40e0ff95e6186a7b7f0befd/black-26.3.1-py3-none-any.whl", hash = "sha256:2bd5aa94fc267d38bb21a70d7410a89f1a1d318841855f698746f8e7f51acd1b", size = 207542, upload-time = "2026-03-12T03:36:01.668Z" },
 ]
 
 [[package]]
@@ -2075,6 +2080,81 @@ grpc = [
     { name = "grpcio" },
 ]
 
+[[package]]
+name = "granian"
+version = "2.5.7"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "click" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/db/b1/100c5add0409559ddbbecca5835c17217b7a2e026eff999bfa359a630686/granian-2.5.7.tar.gz", hash = "sha256:4702a7bcc736454803426bd2c4e7a374739ae1e4b11d27bcdc49b691d316fa0c", size = 112206, upload-time = "2025-11-05T12:18:29.258Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e1/6f/7719fc97aa081915024939f0d35fdae57dfd3d7214f7ef4a7fa664abbbc3/granian-2.5.7-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:7d84a254e9c88da874ba349f7892278a871acc391ab6af21cc32f58d27cd50a9", size = 2854526, upload-time = "2025-11-05T12:15:29.721Z" },
+    { url = "https://files.pythonhosted.org/packages/9f/cd/af33b780602f962c282ba3341131f7ee3b224a6c856a9fb11a017750a48f/granian-2.5.7-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8857d5a6ed94ea64d6b92d1d5fa8f7c1676bbecd71e6ca3d71fcd7118448af1d", size = 2537151, upload-time = "2025-11-05T12:15:31.659Z" },
+    { url = "https://files.pythonhosted.org/packages/6d/58/1a0d529d3d3ddc11b2b292b8f2a7566812d8691de7b1fc8ea5c8f36fd81a/granian-2.5.7-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9914dfc93f04a53a92d8cfdb059c11d620ff83e9326a99880491a9c5bc5940ef", size = 3017277, upload-time = "2025-11-05T12:15:33.42Z" },
+    { url = "https://files.pythonhosted.org/packages/a4/78/2a3c198ee379392d9998e4ff0cfd9ffa95b2d2c683bd15a7266a09325d43/granian-2.5.7-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:24c972fe009ca3a08fd7fb182e07fcb16bffe49c87b1c3489a6986c9e9248dc1", size = 2859098, upload-time = "2025-11-05T12:15:35.15Z" },
+    { url = "https://files.pythonhosted.org/packages/6e/44/7b9fba226083170e9ba221b23ab29d7ffcb761b1ef2b6ed6dac2081bc7fe/granian-2.5.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:034df207e62f104d39db479b693e03072c7eb8e202493cdf58948ff83e753cca", size = 3119567, upload-time = "2025-11-05T12:15:36.674Z" },
+    { url = "https://files.pythonhosted.org/packages/ff/76/f1e348991c031a50d30d3ab0625fec3b7e811092cdb0d1e996885abf1605/granian-2.5.7-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0719052a27caca73bf4000ccdb0339a9d6705e7a4b6613b9fa88ba27c72ba659", size = 2901389, upload-time = "2025-11-05T12:15:39.557Z" },
+    { url = "https://files.pythonhosted.org/packages/f0/69/71b3d7d90d56fda5617fd98838ac481756ad64f76c1fc1b5e21c43a51f15/granian-2.5.7-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:be5b9224ec2583ea3b6ca90788b7f59253b6e07fcf817d14c205e6611faaf2be", size = 2989856, upload-time = "2025-11-05T12:15:41.001Z" },
+    { url = "https://files.pythonhosted.org/packages/74/42/603db3d0ede778adc979c6acc1eaafa5c670c795f5e0e14feb07772ed197/granian-2.5.7-cp310-cp310-musllinux_1_1_armv7l.whl", hash = "sha256:ff246af31840369a1d06030f4d291c6a93841f68ee1f836036bce6625ae73b30", size = 3147378, upload-time = "2025-11-05T12:15:42.432Z" },
+    { url = "https://files.pythonhosted.org/packages/35/b5/cc557e30ba23c2934c33935768dd0233ef7a10b1e8c81dbbc63d5e2562b5/granian-2.5.7-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:cf79375e37a63217f9c1dc4ad15200bc5a89860b321ca30d8a5086a6ea1202e4", size = 3210930, upload-time = "2025-11-05T12:15:45.263Z" },
+    { url = "https://files.pythonhosted.org/packages/c3/67/ba90520cafcd13b5c76d147d713556b9eef877ca001f9ccf44d5443738b6/granian-2.5.7-cp310-cp310-win_amd64.whl", hash = "sha256:b4269a390054c0f71d9ce9d7c75ce2da0c59e78cb522016eb2f5a506c3eb6573", size = 2176887, upload-time = "2025-11-05T12:15:46.615Z" },
+    { url = "https://files.pythonhosted.org/packages/61/21/da3ade91b49ae99146daac6426701cc25b2c5f1413b6c8cb1cc048877036/granian-2.5.7-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:7aa90dcda1fbf03604e229465380138954d9c000eca2947a94dcfbd765414d32", size = 2854652, upload-time = "2025-11-05T12:15:48.342Z" },
+    { url = "https://files.pythonhosted.org/packages/76/67/a6fa402ca5ebddebec5d46dacf646ce073872e5251915a725f6abf2a23bb/granian-2.5.7-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:da4f27323be1188f9e325711016ee108840e14a5971bb4b4d15b65b2d1b00a2d", size = 2537539, upload-time = "2025-11-05T12:15:50.136Z" },
+    { url = "https://files.pythonhosted.org/packages/f9/70/accb5afd83ef785bd9e32067a13547c51cb0139076a8f2857d6d436773df/granian-2.5.7-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8ca5b7028b6ebafce30419ddb6ee7fbfb236fdd0da89427811324ddd38c7d314", size = 3017554, upload-time = "2025-11-05T12:15:52.962Z" },
+    { url = "https://files.pythonhosted.org/packages/74/45/98356af5f36af2b6b47a91fef0d326c275e508bf4bcf0c08bd35ed314db8/granian-2.5.7-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b83e95b18be5dfa92296bc8acfeb353488123399c90cc5f0eccf451e88bc4caf", size = 2859127, upload-time = "2025-11-05T12:15:54.49Z" },
+    { url = "https://files.pythonhosted.org/packages/27/7a/04d3ec13b197509c40340ec80414fbbc2b0913f6e1a18c3987cc608c8571/granian-2.5.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9aad9e920441232a7b8ad33bef7f04aae986e0e386ab7f13312477c3ea2c85df", size = 3119494, upload-time = "2025-11-05T12:15:56.324Z" },
+    { url = "https://files.pythonhosted.org/packages/b9/5d/1a82a596725824f6e76b8f7b853ceb464cd0334b2b8143c278aa46f23b6d/granian-2.5.7-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:777d35961d5139d203cf54d872ad5979b171e6496a471a5bcb8032f4471bdec6", size = 2901511, upload-time = "2025-11-05T12:15:58.7Z" },
+    { url = "https://files.pythonhosted.org/packages/94/45/b53d6d7df5cd35c3b8bb329f5ee1c7b31ead7a61a6f2046f6562028d7e1b/granian-2.5.7-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:ae72c7ba1e8f35d3021dafb2ba6c4ef89f93f877218f8c6ed1cb672145cd81ad", size = 2989828, upload-time = "2025-11-05T12:16:00.341Z" },
+    { url = "https://files.pythonhosted.org/packages/7f/80/bb57b0fa24fcd518cd64442249459bd214ab1ec5f32590fd30389944261c/granian-2.5.7-cp311-cp311-musllinux_1_1_armv7l.whl", hash = "sha256:3764d87edd3fddaf557dce32be396a2a56dfc5b9ad2989b1f98952983ae4a21c", size = 3147694, upload-time = "2025-11-05T12:16:01.826Z" },
+    { url = "https://files.pythonhosted.org/packages/7f/00/f8747aaf8dcd488e4462db89f7273dd9ae702fd17a58d72193b48eff0470/granian-2.5.7-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:f5e21bbf1daebb0219253576cac4e5edc8fa8356ad85d66577c4f3ea2d5c6e3c", size = 3211169, upload-time = "2025-11-05T12:16:03.308Z" },
+    { url = "https://files.pythonhosted.org/packages/1f/69/8593d539898a870692cad447d22c2c4cc34566ad9070040ca216db6ac184/granian-2.5.7-cp311-cp311-win_amd64.whl", hash = "sha256:d210dd98852825c8a49036a6ec23cdfaa7689d1cb12ddc651c6466b412047349", size = 2176921, upload-time = "2025-11-05T12:16:04.63Z" },
+    { url = "https://files.pythonhosted.org/packages/b5/cf/f76d05e950f76924ffb6c5212561be4dd93fa569518869cc1233a0c77613/granian-2.5.7-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:41e3a293ac23c76d18628d1bd8376ce3230fb3afe3cf71126b8885e8da4e40c4", size = 2850787, upload-time = "2025-11-05T12:16:06.028Z" },
+    { url = "https://files.pythonhosted.org/packages/3f/d7/6972aa8c38d26b4cf9f35bcc9b7d3a26a3aa930e612d5913d8f4181331a1/granian-2.5.7-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8b345b539bcbe6dedf8a9323b0c960530cb1fb2cfb887139e6ae9513b6c04d8c", size = 2529552, upload-time = "2025-11-05T12:16:07.389Z" },
+    { url = "https://files.pythonhosted.org/packages/56/b4/cd5958b6af674a32296a0fef73fb499c2bf2874025062323f5dbc838f4fc/granian-2.5.7-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:12e4d7ba8e3223e2bf974860a59c29b06fa805a98ad4304be4e77180d3a28f55", size = 3009131, upload-time = "2025-11-05T12:16:08.759Z" },
+    { url = "https://files.pythonhosted.org/packages/7a/69/f3828de736c2802fd7fcac0bb1a0387b3332d432f0eeacb8116094926f06/granian-2.5.7-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e727d3518f038b64cb0352b34f43b387aafe5eb12b6c4b57ef598b811e40d4ed", size = 2852544, upload-time = "2025-11-05T12:16:10.22Z" },
+    { url = "https://files.pythonhosted.org/packages/6f/c3/b8c65cf86d473b6e99e6d985c678cb192c9b9776a966a2f4b009696bb650/granian-2.5.7-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:59fe2b352a828a2b04bcfd105e623d66786f217759d2d6245651a7b81e4ac294", size = 3131904, upload-time = "2025-11-05T12:16:13.249Z" },
+    { url = "https://files.pythonhosted.org/packages/df/7e/b60421bddf187ab2a46682423e4a94b2b22a6ddff6842bf9ca2194e62ac2/granian-2.5.7-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:ec5fb593c2d436a323e711010e79718e6d5d1491d0d660fb7c9d97f7e5900830", size = 2908851, upload-time = "2025-11-05T12:16:15.305Z" },
+    { url = "https://files.pythonhosted.org/packages/2f/cf/3f2426e19dc955a74dc94a5a47c4170e68acb060c541ac080f71a9d55d5d/granian-2.5.7-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:48fbc25f3717d01e11547afe0e9cdf9d7c41c9f316b9623a40c22ea6b2128d36", size = 2993270, upload-time = "2025-11-05T12:16:17.133Z" },
+    { url = "https://files.pythonhosted.org/packages/40/2e/67e1e05ee0d503cc6e9fe53b03f69eb2f267a589d7b40873d120c417385f/granian-2.5.7-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:770935fec3374b814d21c01508c0697842d7c3750731a8ea129738b537ac594c", size = 3134662, upload-time = "2025-11-05T12:16:18.598Z" },
+    { url = "https://files.pythonhosted.org/packages/17/d5/9d3242bbd911434c4f3d4f14c48e73774a8ddb591e0f975eaeeaef1d5081/granian-2.5.7-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:5db2600c92f74da74f624d2fdb01afe9e9365b50bd4e695a78e54961dc132f1b", size = 3220446, upload-time = "2025-11-05T12:16:20.598Z" },
+    { url = "https://files.pythonhosted.org/packages/10/27/b2baa0443a42d8eb59f3dfbe8186e8c80a090655584af4611f22f1592d7a/granian-2.5.7-cp312-cp312-win_amd64.whl", hash = "sha256:bc368bdeb21646a965adf9f43dd2f4a770647e50318ba1b7cf387d4916ed7e69", size = 2179465, upload-time = "2025-11-05T12:16:22.031Z" },
+    { url = "https://files.pythonhosted.org/packages/54/ec/bf1b7eefe824630d1d3ae9a8af397d823f2339d3adec71e9ee49d667409c/granian-2.5.7-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:fafb9c17def635bb0a5e20e145601598a6767b879bc2501663dbb45a57d1bc2e", size = 2850581, upload-time = "2025-11-05T12:16:23.516Z" },
+    { url = "https://files.pythonhosted.org/packages/28/f7/5172daf1968c3a2337c51c50f4a3013aaab564d012d3a79e8390cc66403b/granian-2.5.7-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9616a197eba637d59242661be8a46127c3f79f7c9bbfa44c0ea8c8c790a11d5e", size = 2529452, upload-time = "2025-11-05T12:16:25.088Z" },
+    { url = "https://files.pythonhosted.org/packages/92/10/4344ccacc3f8dea973d630306491de43fbd4a0248e3f7cc9ff09ed5cc524/granian-2.5.7-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:cfd7a09d5eb00a271ec79e3e0bbf069aa62ce376b64825bdeacb668d2b2a4041", size = 3008798, upload-time = "2025-11-05T12:16:26.584Z" },
+    { url = "https://files.pythonhosted.org/packages/5e/33/638cf8c7f23ab905d3f6a371b5f87d03fd611678424223a0f1d0f7766cc7/granian-2.5.7-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1438a82264690fce6e82de66a95c77f5b0a5c33b93269eb85fc69ce0112c12d5", size = 2852309, upload-time = "2025-11-05T12:16:28.064Z" },
+    { url = "https://files.pythonhosted.org/packages/18/42/6ec25d37ffc1f08679e6b325e9f9ac199ba5def948904c9205cd34fbfe6b/granian-2.5.7-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b3573121da77aac1af64cf90a88f29b2daecbf92458beec187421a382039f366", size = 3131335, upload-time = "2025-11-05T12:16:29.588Z" },
+    { url = "https://files.pythonhosted.org/packages/b0/1e/db85dac58d84d3e50e427fe5b60b4f8e8a561d9784971fa3b2879198ad88/granian-2.5.7-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:34cdb82024efbcc9de01c7505213be17e4ba5e7a3acabe74ecd93ba31de7673e", size = 2908705, upload-time = "2025-11-05T12:16:31.049Z" },
+    { url = "https://files.pythonhosted.org/packages/d9/25/a38fd12e1661bbd8535203a8b61240feac7b6b96726bff4de23b0078ab9f/granian-2.5.7-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:572451e94de69df228e4314cb91a50dee1565c4a53d33ffac5936c6ec9c5aba2", size = 2993118, upload-time = "2025-11-05T12:16:32.767Z" },
+    { url = "https://files.pythonhosted.org/packages/d3/cd/852913a0fc30efc24495453c0f973dd74ef13aa0561afb352afa4b6ecbc2/granian-2.5.7-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:6e1679a4b102511b483774397134d244108851ae7a1e8bef09a8ef927ab4d370", size = 3134260, upload-time = "2025-11-05T12:16:34.552Z" },
+    { url = "https://files.pythonhosted.org/packages/60/64/0dff100ce1e43c700918b39656cc000b1163c144eac3a12563a5f692dcd1/granian-2.5.7-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:285be70dcf3c70121afec03e691596db94bd786f9bebc229e9e0319686857d82", size = 3219987, upload-time = "2025-11-05T12:16:36.43Z" },
+    { url = "https://files.pythonhosted.org/packages/19/ab/e66cf9bf57800dd7c2a2a4b8f23124603fce561a65a176f4cf3794a85b92/granian-2.5.7-cp313-cp313-win_amd64.whl", hash = "sha256:1273c9b1d38d19bcdd550a9a846d07112e541cfa1f99be04fbb926f2a003df3d", size = 2179201, upload-time = "2025-11-05T12:16:37.869Z" },
+    { url = "https://files.pythonhosted.org/packages/da/0e/feca4a20e7b9e7de0e58103278c6581ebf3d5c1b972ed1c2dcfd25741f15/granian-2.5.7-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:75b9798bc13baa76e35165e5a778cd58a7258d5a2112ed6ef84ef84874244856", size = 2776744, upload-time = "2025-11-05T12:16:41.969Z" },
+    { url = "https://files.pythonhosted.org/packages/f7/fe/65ca38ba9b9f4805495d96ed7b774dfd300f7c944f088db39c676c16501e/granian-2.5.7-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4cb8247728680ca308b7dc41a6d27582b78e15e902377e89000711f1126524dd", size = 2465942, upload-time = "2025-11-05T12:16:43.762Z" },
+    { url = "https://files.pythonhosted.org/packages/75/d1/b9dea32fbafabe5c7b049fb0209149a37c6b8468c698d066448cbe88dc85/granian-2.5.7-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:64348b83f1ad2f7a29df7932dc518ad669cb61a08a9cde02ca8ede8e9b110506", size = 3015413, upload-time = "2025-11-05T12:16:45.265Z" },
+    { url = "https://files.pythonhosted.org/packages/cb/9e/d29485ab18896e4d911e33b006af7a9b7098316a78938d6b7455c523fea5/granian-2.5.7-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e2292d4a4661c79d471fa0ff6fe640018c923b6a6dd1bb5383b368b3d5ec2a0c", size = 2783371, upload-time = "2025-11-05T12:16:46.762Z" },
+    { url = "https://files.pythonhosted.org/packages/41/cd/58c67dc191caeecbbb15ee39d433136dd064c13778b4551661bd902b5a78/granian-2.5.7-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:45903d2f2f88a9cd4a7d0b8ec329db1fb2d9e15bf38153087a3b217b9cdb0046", size = 2979946, upload-time = "2025-11-05T12:16:48.255Z" },
+    { url = "https://files.pythonhosted.org/packages/16/0b/04e4977df3ef7607a8b6625caed7cac107a049120d2452c33392d4544875/granian-2.5.7-cp313-cp313t-musllinux_1_1_armv7l.whl", hash = "sha256:106e8988e42e527c18b763be5faae7e8f602caac6cb93657793638fc9ab41c98", size = 3123177, upload-time = "2025-11-05T12:16:49.724Z" },
+    { url = "https://files.pythonhosted.org/packages/c7/89/4e10e18fc107e5929143a06d9257646963cf5621c928b3d2774e5a85652a/granian-2.5.7-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:711632e602c4ea08b827bf6095c2c6fbe6005c7a05f142ae2b4d9e1d45cefbd9", size = 3211773, upload-time = "2025-11-05T12:16:51.438Z" },
+    { url = "https://files.pythonhosted.org/packages/57/81/94e416056d8b4b1cd09cc8065a1e240b0af99f21301c209571530cd83dd0/granian-2.5.7-cp313-cp313t-win_amd64.whl", hash = "sha256:1c571733aa0fdb6755be9ffb3cd728ef965ae565ba896e407d6019bad929d7bb", size = 2174154, upload-time = "2025-11-05T12:16:53.411Z" },
+    { url = "https://files.pythonhosted.org/packages/0e/25/2a4112983df5ce0ec8407121ad72c17d27ebfad57085749b8e4164d69e63/granian-2.5.7-pp310-pypy310_pp73-macosx_10_12_x86_64.whl", hash = "sha256:cdae1c86357bfe895ffd0065c0403913bc008f752e2f77ab363d4e3b4276009b", size = 2838744, upload-time = "2025-11-05T12:17:45.904Z" },
+    { url = "https://files.pythonhosted.org/packages/d7/0a/eb0c5b71355e8f99b89dc335f16cd5108763c554e96a2aae5e7162ef4997/granian-2.5.7-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:bc1d8aaf5bfc5fc9f8f590a42e9f88a43d19ad71f670c6969fa791b52ce1f5ec", size = 2538706, upload-time = "2025-11-05T12:17:47.471Z" },
+    { url = "https://files.pythonhosted.org/packages/f2/9c/4c592c5a813a921033a37a0f003278b1f772a6c9abd16f821bcb119151f0/granian-2.5.7-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:288b62c19aea5b162d27e229469b6307a78cb272aa8fcc296dbfca9fbbda4d8f", size = 3117369, upload-time = "2025-11-05T12:17:49.172Z" },
+    { url = "https://files.pythonhosted.org/packages/f1/35/96af9f0995a7c45f0cd31261ab6284e5d6028afa17c6fcfe757cccb0afb5/granian-2.5.7-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:66c3d2619dc5e845d658cf3ed4f7370f83d5323a85ff8338e7c7a27d9a333841", size = 2904972, upload-time = "2025-11-05T12:17:50.863Z" },
+    { url = "https://files.pythonhosted.org/packages/fc/93/45c253983c2001f534ba2c7bc1e53718fc8cecf196b1e1a0469d5874ae54/granian-2.5.7-pp310-pypy310_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:323e35d5d5054d2568fc824798471e7d33314f47aebd556c4fbf4894e539347d", size = 2991986, upload-time = "2025-11-05T12:17:52.602Z" },
+    { url = "https://files.pythonhosted.org/packages/25/77/c03e60c7bed386ab16cf15b317dea7f95dde5095af6e17cbd657cd82c21b/granian-2.5.7-pp310-pypy310_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:026ef2588a2b991b250768bf47538fd5fd864549535f885239b6908b214299c4", size = 3163649, upload-time = "2025-11-05T12:17:54.402Z" },
+    { url = "https://files.pythonhosted.org/packages/5e/c9/2bce3db4e3da8d3a697c363c8f699b71f05b7f7a0458e1ba345eaea53fcd/granian-2.5.7-pp310-pypy310_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:4717a62c0a1b79372c495b99ade18bfc3c4a365242bf75770c96a4767a9bcf66", size = 3201886, upload-time = "2025-11-05T12:17:56.553Z" },
+    { url = "https://files.pythonhosted.org/packages/78/66/997ebfd8cc4a0640befb970bc846a76437d1f0b55dff179e69f29fa4615b/granian-2.5.7-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:4b57ae0a2e1dbc7a248e3c08440b490b3f247e7e4f997faa72e82f5a89d0ea4c", size = 2175219, upload-time = "2025-11-05T12:17:58.126Z" },
+    { url = "https://files.pythonhosted.org/packages/16/0f/da2588ac78254a4d0be90a6f733d0bb7dd1edb78a10d9e59fa9837687e94/granian-2.5.7-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:bee545c9b9e38eabcdd675e3fec1a2112b8193dc864739952b9de8131433a31c", size = 2838886, upload-time = "2025-11-05T12:17:59.809Z" },
+    { url = "https://files.pythonhosted.org/packages/7d/34/75def8343534e9d48362c43c3cbd06242a2d7804fbfbc824c8aa9fb75a30/granian-2.5.7-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:73c76c0f1ee46506224e92df193b4d271ea89f0d82cd69301784ca85bc1db515", size = 2538597, upload-time = "2025-11-05T12:18:01.496Z" },
+    { url = "https://files.pythonhosted.org/packages/c3/5d/d828d97aad050cfc5b18a0163b532c289a35ad214e31f5a129695b2b4cae/granian-2.5.7-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:68879c27aed972f647a8e8ef37f9046f71d7507dc9b3ceffa97d2fbffe6a16c8", size = 3117570, upload-time = "2025-11-05T12:18:03.818Z" },
+    { url = "https://files.pythonhosted.org/packages/2d/57/b8380f3d6b6dcdcd454d720cf11dbecb0e2071a870f44eb834011f14b573/granian-2.5.7-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:ea9cbdfbd750813866dcc9c020018e5f20a57a4e3a83bd049ccc1f6da0559b75", size = 2905089, upload-time = "2025-11-05T12:18:05.567Z" },
+    { url = "https://files.pythonhosted.org/packages/0b/e9/04a7c3b83650afc4a4ad82b67e6306d99f80ac1a6aacb3a8ba182f7359d6/granian-2.5.7-pp311-pypy311_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:d142ff5ee6027515370e56f95d179ec3e81bd265d5b4958de2b19adcdf34887d", size = 2991867, upload-time = "2025-11-05T12:18:07.223Z" },
+    { url = "https://files.pythonhosted.org/packages/2b/bf/a1cdbff73cbac4fddf817d06c13ce6cdc75c22d6da1b257e3563fea4c3c5/granian-2.5.7-pp311-pypy311_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:222f0fb1688a62ca23cb3da974cefa69e7fdc40fd548d1ae87a953225e1d1cbb", size = 3164141, upload-time = "2025-11-05T12:18:09.267Z" },
+    { url = "https://files.pythonhosted.org/packages/c8/cc/35c6a55ac2c211e86a9f0c728eb81b6ad19f05a3055d79c6f11a1b71f5d5/granian-2.5.7-pp311-pypy311_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:40494c6cda1ad881ae07efbb2dc4a1ca8f12d5c6cf28d1ab8b0f2db13826617b", size = 3201599, upload-time = "2025-11-05T12:18:10.962Z" },
+    { url = "https://files.pythonhosted.org/packages/f3/0a/5a95a3889532bc5a5f652cdc78dae8ffa16d4228b4d35256a98be89e33ef/granian-2.5.7-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c3942d08af2c8b67d0ef569b6c567284433ebf09b4af3ea68388abb7caccad2b", size = 2175240, upload-time = "2025-11-05T12:18:12.956Z" },
+]
+
 [[package]]
 name = "graphene"
 version = "3.4.3"
@@ -3238,6 +3318,7 @@ proxy = [
     { name = "cryptography" },
     { name = "fastapi" },
     { name = "fastapi-sso" },
+    { name = "granian" },
     { name = "gunicorn" },
     { name = "litellm-enterprise" },
     { name = "litellm-proxy-extras" },
@@ -3401,6 +3482,7 @@ requires-dist = [
     { name = "google-cloud-iam", marker = "extra == 'extra-proxy'", specifier = "==2.19.1" },
     { name = "google-cloud-kms", marker = "extra == 'extra-proxy'", specifier = "==2.24.2" },
     { name = "google-genai", marker = "extra == 'proxy-runtime'", specifier = "==1.37.0" },
+    { name = "granian", marker = "extra == 'proxy'", specifier = "==2.5.7" },
     { name = "grpcio", marker = "extra == 'grpc'", specifier = "==1.78.0" },
     { name = "grpcio", marker = "extra == 'proxy-runtime'", specifier = "==1.78.0" },
     { name = "gunicorn", marker = "extra == 'proxy'", specifier = "==23.0.0" },
@@ -3484,7 +3566,7 @@ ci = [
     { name = "traceloop-sdk", specifier = "==0.33.12" },
 ]
 dev = [
-    { name = "black", specifier = "==24.10.0" },
+    { name = "black", specifier = "==26.3.1" },
     { name = "diff-cover", specifier = "==9.7.2" },
     { name = "fakeredis", specifier = "==2.34.1" },
     { name = "fastapi-offline", specifier = "==1.7.6" },
@@ -6162,6 +6244,35 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/6c/a0/4ed6632b70a52de845df056654162acdebaf97c20e3212c559ac43e7216e/python_ulid-3.1.0-py3-none-any.whl", hash = "sha256:e2cdc979c8c877029b4b7a38a6fba3bc4578e4f109a308419ff4d3ccf0a46619", size = 11577, upload-time = "2025-08-18T16:09:25.047Z" },
 ]
 
+[[package]]
+name = "pytokens"
+version = "0.4.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/b6/34/b4e015b99031667a7b960f888889c5bd34ef585c85e1cb56a594b92836ac/pytokens-0.4.1.tar.gz", hash = "sha256:292052fe80923aae2260c073f822ceba21f3872ced9a68bb7953b348e561179a", size = 23015, upload-time = "2026-01-30T01:03:45.924Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/42/24/f206113e05cb8ef51b3850e7ef88f20da6f4bf932190ceb48bd3da103e10/pytokens-0.4.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2a44ed93ea23415c54f3face3b65ef2b844d96aeb3455b8a69b3df6beab6acc5", size = 161522, upload-time = "2026-01-30T01:02:50.393Z" },
+    { url = "https://files.pythonhosted.org/packages/d4/e9/06a6bf1b90c2ed81a9c7d2544232fe5d2891d1cd480e8a1809ca354a8eb2/pytokens-0.4.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:add8bf86b71a5d9fb5b89f023a80b791e04fba57960aa790cc6125f7f1d39dfe", size = 246945, upload-time = "2026-01-30T01:02:52.399Z" },
+    { url = "https://files.pythonhosted.org/packages/69/66/f6fb1007a4c3d8b682d5d65b7c1fb33257587a5f782647091e3408abe0b8/pytokens-0.4.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:670d286910b531c7b7e3c0b453fd8156f250adb140146d234a82219459b9640c", size = 259525, upload-time = "2026-01-30T01:02:53.737Z" },
+    { url = "https://files.pythonhosted.org/packages/04/92/086f89b4d622a18418bac74ab5db7f68cf0c21cf7cc92de6c7b919d76c88/pytokens-0.4.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:4e691d7f5186bd2842c14813f79f8884bb03f5995f0575272009982c5ac6c0f7", size = 262693, upload-time = "2026-01-30T01:02:54.871Z" },
+    { url = "https://files.pythonhosted.org/packages/b4/7b/8b31c347cf94a3f900bdde750b2e9131575a61fdb620d3d3c75832262137/pytokens-0.4.1-cp310-cp310-win_amd64.whl", hash = "sha256:27b83ad28825978742beef057bfe406ad6ed524b2d28c252c5de7b4a6dd48fa2", size = 103567, upload-time = "2026-01-30T01:02:56.414Z" },
+    { url = "https://files.pythonhosted.org/packages/3d/92/790ebe03f07b57e53b10884c329b9a1a308648fc083a6d4a39a10a28c8fc/pytokens-0.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d70e77c55ae8380c91c0c18dea05951482e263982911fc7410b1ffd1dadd3440", size = 160864, upload-time = "2026-01-30T01:02:57.882Z" },
+    { url = "https://files.pythonhosted.org/packages/13/25/a4f555281d975bfdd1eba731450e2fe3a95870274da73fb12c40aeae7625/pytokens-0.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a58d057208cb9075c144950d789511220b07636dd2e4708d5645d24de666bdc", size = 248565, upload-time = "2026-01-30T01:02:59.912Z" },
+    { url = "https://files.pythonhosted.org/packages/17/50/bc0394b4ad5b1601be22fa43652173d47e4c9efbf0044c62e9a59b747c56/pytokens-0.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b49750419d300e2b5a3813cf229d4e5a4c728dae470bcc89867a9ad6f25a722d", size = 260824, upload-time = "2026-01-30T01:03:01.471Z" },
+    { url = "https://files.pythonhosted.org/packages/4e/54/3e04f9d92a4be4fc6c80016bc396b923d2a6933ae94b5f557c939c460ee0/pytokens-0.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d9907d61f15bf7261d7e775bd5d7ee4d2930e04424bab1972591918497623a16", size = 264075, upload-time = "2026-01-30T01:03:04.143Z" },
+    { url = "https://files.pythonhosted.org/packages/d1/1b/44b0326cb5470a4375f37988aea5d61b5cc52407143303015ebee94abfd6/pytokens-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:ee44d0f85b803321710f9239f335aafe16553b39106384cef8e6de40cb4ef2f6", size = 103323, upload-time = "2026-01-30T01:03:05.412Z" },
+    { url = "https://files.pythonhosted.org/packages/41/5d/e44573011401fb82e9d51e97f1290ceb377800fb4eed650b96f4753b499c/pytokens-0.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:140709331e846b728475786df8aeb27d24f48cbcf7bcd449f8de75cae7a45083", size = 160663, upload-time = "2026-01-30T01:03:06.473Z" },
+    { url = "https://files.pythonhosted.org/packages/f0/e6/5bbc3019f8e6f21d09c41f8b8654536117e5e211a85d89212d59cbdab381/pytokens-0.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6d6c4268598f762bc8e91f5dbf2ab2f61f7b95bdc07953b602db879b3c8c18e1", size = 255626, upload-time = "2026-01-30T01:03:08.177Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/3c/2d5297d82286f6f3d92770289fd439956b201c0a4fc7e72efb9b2293758e/pytokens-0.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:24afde1f53d95348b5a0eb19488661147285ca4dd7ed752bbc3e1c6242a304d1", size = 269779, upload-time = "2026-01-30T01:03:09.756Z" },
+    { url = "https://files.pythonhosted.org/packages/20/01/7436e9ad693cebda0551203e0bf28f7669976c60ad07d6402098208476de/pytokens-0.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5ad948d085ed6c16413eb5fec6b3e02fa00dc29a2534f088d3302c47eb59adf9", size = 268076, upload-time = "2026-01-30T01:03:10.957Z" },
+    { url = "https://files.pythonhosted.org/packages/2e/df/533c82a3c752ba13ae7ef238b7f8cdd272cf1475f03c63ac6cf3fcfb00b6/pytokens-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:3f901fe783e06e48e8cbdc82d631fca8f118333798193e026a50ce1b3757ea68", size = 103552, upload-time = "2026-01-30T01:03:12.066Z" },
+    { url = "https://files.pythonhosted.org/packages/cb/dc/08b1a080372afda3cceb4f3c0a7ba2bde9d6a5241f1edb02a22a019ee147/pytokens-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8bdb9d0ce90cbf99c525e75a2fa415144fd570a1ba987380190e8b786bc6ef9b", size = 160720, upload-time = "2026-01-30T01:03:13.843Z" },
+    { url = "https://files.pythonhosted.org/packages/64/0c/41ea22205da480837a700e395507e6a24425151dfb7ead73343d6e2d7ffe/pytokens-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5502408cab1cb18e128570f8d598981c68a50d0cbd7c61312a90507cd3a1276f", size = 254204, upload-time = "2026-01-30T01:03:14.886Z" },
+    { url = "https://files.pythonhosted.org/packages/e0/d2/afe5c7f8607018beb99971489dbb846508f1b8f351fcefc225fcf4b2adc0/pytokens-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:29d1d8fb1030af4d231789959f21821ab6325e463f0503a61d204343c9b355d1", size = 268423, upload-time = "2026-01-30T01:03:15.936Z" },
+    { url = "https://files.pythonhosted.org/packages/68/d4/00ffdbd370410c04e9591da9220a68dc1693ef7499173eb3e30d06e05ed1/pytokens-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:970b08dd6b86058b6dc07efe9e98414f5102974716232d10f32ff39701e841c4", size = 266859, upload-time = "2026-01-30T01:03:17.458Z" },
+    { url = "https://files.pythonhosted.org/packages/a7/c9/c3161313b4ca0c601eeefabd3d3b576edaa9afdefd32da97210700e47652/pytokens-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:9bd7d7f544d362576be74f9d5901a22f317efc20046efe2034dced238cbbfe78", size = 103520, upload-time = "2026-01-30T01:03:18.652Z" },
+    { url = "https://files.pythonhosted.org/packages/c6/78/397db326746f0a342855b81216ae1f0a32965deccfd7c830a2dbc66d2483/pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de", size = 13729, upload-time = "2026-01-30T01:03:45.029Z" },
+]
+
 [[package]]
 name = "pytz"
 version = "2026.2"