diff --git a/.github/workflows/ci_python.yml b/.github/workflows/ci_python.yml index 2db252e43..1b47d34a9 100644 --- a/.github/workflows/ci_python.yml +++ b/.github/workflows/ci_python.yml @@ -268,6 +268,24 @@ jobs: path: ${{ env.NEMO_RELAY_CI_WORKSPACE_TMP }}/wheels/*.whl if-no-files-found: error + - name: Package Python plugin SDK wheel + if: ${{ matrix.platform == 'linux-amd64' }} + working-directory: ${{ env.NEMO_RELAY_CI_WORKSPACE }} + run: | + set -e + just \ + --set output_dir "${{ env.NEMO_RELAY_CI_WORKSPACE_TMP }}" \ + --set ref_name "${NEMO_RELAY_PACKAGE_VERSION}" \ + package-python-plugin + + - name: Upload Python plugin SDK wheel artifact + if: ${{ matrix.platform == 'linux-amd64' }} + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: python-plugin-wheel + path: ${{ env.NEMO_RELAY_CI_WORKSPACE_TMP }}/plugin-wheels/*.whl + if-no-files-found: error + - name: Prune uv cache working-directory: ${{ env.NEMO_RELAY_CI_WORKSPACE }} run: uv cache prune --ci diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5fc0acc42..264662044 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -52,7 +52,7 @@ repos: hooks: - id: ty name: ty (type check) - entry: uv run ty check . --exclude docs/** --exclude fern/** --exclude third_party/** --exclude ./examples/** --exclude .cache/** --exclude .claude/** + entry: uv run ty check . --extra-search-path python/plugin/src --exclude docs/** --exclude fern/** --exclude third_party/** --exclude ./examples/** --exclude .cache/** --exclude .claude/** --exclude python/plugin/src/nemo_relay_plugin/_proto/** language: system types: [python] pass_filenames: false @@ -105,6 +105,13 @@ repos: files: '^(pyproject\.toml|uv\.lock)$' pass_filenames: false + - id: python-worker-proto-check + name: Python worker protobuf stubs are up to date + entry: just check-python-worker-proto + language: system + files: '^(crates/worker-proto/proto/nemo/relay/worker/v1/plugin_worker\.proto|python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2(_grpc)?\.py|justfile)$' + pass_filenames: false + - id: node-lockfile-check name: package-lock.json is up to date entry: bash -c 'npm install --package-lock-only --ignore-scripts --audit=false --fund=false' diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index 1cfac4fde..b3d06bc5e 100644 --- a/ATTRIBUTIONS-Python.md +++ b/ATTRIBUTIONS-Python.md @@ -2565,6 +2565,637 @@ Apache License limitations under the License. ``` +## grpcio (1.81.1) + +### Licenses +License: `Apache-2.0` + + - `LICENSE`: +``` +Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + +----------------------------------------------------------- + +Following applies to: +./src/objective-c/!ProtoCompiler-gRPCCppPlugin.podspec +./src/objective-c/!ProtoCompiler-gRPCPlugin.podspec +./src/objective-c/!ProtoCompiler.podspec +./src/objective-c/BoringSSL-GRPC.podspec +./templates/src/objective-c/!ProtoCompiler-gRPCCppPlugin.podspec.inja +./templates/src/objective-c/!ProtoCompiler-gRPCPlugin.podspec.inja +./templates/src/objective-c/!ProtoCompiler.podspec.inja +./templates/src/objective-c/BoringSSL-GRPC.podspec.template + +BSD 3-Clause License + +Copyright 2016, Google Inc. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, +this list of conditions and the following disclaimer. + +2. 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. + +3. Neither the name of the copyright holder nor the names of its +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 HOLDER 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. + +----------------------------------------------------------- + +Following applies to: +./etc/roots.pem + +Mozilla Public License Version 2.0 +================================== + +1. Definitions +-------------- + +1.1. "Contributor" + means each individual or legal entity that creates, contributes to + the creation of, or owns Covered Software. + +1.2. "Contributor Version" + means the combination of the Contributions of others (if any) used + by a Contributor and that particular Contributor's Contribution. + +1.3. "Contribution" + means Covered Software of a particular Contributor. + +1.4. "Covered Software" + means Source Code Form to which the initial Contributor has attached + the notice in Exhibit A, the Executable Form of such Source Code + Form, and Modifications of such Source Code Form, in each case + including portions thereof. + +1.5. "Incompatible With Secondary Licenses" + means + + (a) that the initial Contributor has attached the notice described + in Exhibit B to the Covered Software; or + + (b) that the Covered Software was made available under the terms of + version 1.1 or earlier of the License, but not also under the + terms of a Secondary License. + +1.6. "Executable Form" + means any form of the work other than Source Code Form. + +1.7. "Larger Work" + means a work that combines Covered Software with other material, in + a separate file or files, that is not Covered Software. + +1.8. "License" + means this document. + +1.9. "Licensable" + means having the right to grant, to the maximum extent possible, + whether at the time of the initial grant or subsequently, any and + all of the rights conveyed by this License. + +1.10. "Modifications" + means any of the following: + + (a) any file in Source Code Form that results from an addition to, + deletion from, or modification of the contents of Covered + Software; or + + (b) any new file in Source Code Form that contains any Covered + Software. + +1.11. "Patent Claims" of a Contributor + means any patent claim(s), including without limitation, method, + process, and apparatus claims, in any patent Licensable by such + Contributor that would be infringed, but for the grant of the + License, by the making, using, selling, offering for sale, having + made, import, or transfer of either its Contributions or its + Contributor Version. + +1.12. "Secondary License" + means either the GNU General Public License, Version 2.0, the GNU + Lesser General Public License, Version 2.1, the GNU Affero General + Public License, Version 3.0, or any later versions of those + licenses. + +1.13. "Source Code Form" + means the form of the work preferred for making modifications. + +1.14. "You" (or "Your") + means an individual or a legal entity exercising rights under this + License. For legal entities, "You" includes any entity that + controls, is controlled by, or is under common control with You. For + purposes of this definition, "control" means (a) the power, direct + or indirect, to cause the direction or management of such entity, + whether by contract or otherwise, or (b) ownership of more than + fifty percent (50%) of the outstanding shares or beneficial + ownership of such entity. + +2. License Grants and Conditions +-------------------------------- + +2.1. Grants + +Each Contributor hereby grants You a world-wide, royalty-free, +non-exclusive license: + +(a) under intellectual property rights (other than patent or trademark) + Licensable by such Contributor to use, reproduce, make available, + modify, display, perform, distribute, and otherwise exploit its + Contributions, either on an unmodified basis, with Modifications, or + as part of a Larger Work; and + +(b) under Patent Claims of such Contributor to make, use, sell, offer + for sale, have made, import, and otherwise transfer either its + Contributions or its Contributor Version. + +2.2. Effective Date + +The licenses granted in Section 2.1 with respect to any Contribution +become effective for each Contribution on the date the Contributor first +distributes such Contribution. + +2.3. Limitations on Grant Scope + +The licenses granted in this Section 2 are the only rights granted under +this License. No additional rights or licenses will be implied from the +distribution or licensing of Covered Software under this License. +Notwithstanding Section 2.1(b) above, no patent license is granted by a +Contributor: + +(a) for any code that a Contributor has removed from Covered Software; + or + +(b) for infringements caused by: (i) Your and any other third party's + modifications of Covered Software, or (ii) the combination of its + Contributions with other software (except as part of its Contributor + Version); or + +(c) under Patent Claims infringed by Covered Software in the absence of + its Contributions. + +This License does not grant any rights in the trademarks, service marks, +or logos of any Contributor (except as may be necessary to comply with +the notice requirements in Section 3.4). + +2.4. Subsequent Licenses + +No Contributor makes additional grants as a result of Your choice to +distribute the Covered Software under a subsequent version of this +License (see Section 10.2) or under the terms of a Secondary License (if +permitted under the terms of Section 3.3). + +2.5. Representation + +Each Contributor represents that the Contributor believes its +Contributions are its original creation(s) or it has sufficient rights +to grant the rights to its Contributions conveyed by this License. + +2.6. Fair Use + +This License is not intended to limit any rights You have under +applicable copyright doctrines of fair use, fair dealing, or other +equivalents. + +2.7. Conditions + +Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted +in Section 2.1. + +3. Responsibilities +------------------- + +3.1. Distribution of Source Form + +All distribution of Covered Software in Source Code Form, including any +Modifications that You create or to which You contribute, must be under +the terms of this License. You must inform recipients that the Source +Code Form of the Covered Software is governed by the terms of this +License, and how they can obtain a copy of this License. You may not +attempt to alter or restrict the recipients' rights in the Source Code +Form. + +3.2. Distribution of Executable Form + +If You distribute Covered Software in Executable Form then: + +(a) such Covered Software must also be made available in Source Code + Form, as described in Section 3.1, and You must inform recipients of + the Executable Form how they can obtain a copy of such Source Code + Form by reasonable means in a timely manner, at a charge no more + than the cost of distribution to the recipient; and + +(b) You may distribute such Executable Form under the terms of this + License, or sublicense it under different terms, provided that the + license for the Executable Form does not attempt to limit or alter + the recipients' rights in the Source Code Form under this License. + +3.3. Distribution of a Larger Work + +You may create and distribute a Larger Work under terms of Your choice, +provided that You also comply with the requirements of this License for +the Covered Software. If the Larger Work is a combination of Covered +Software with a work governed by one or more Secondary Licenses, and the +Covered Software is not Incompatible With Secondary Licenses, this +License permits You to additionally distribute such Covered Software +under the terms of such Secondary License(s), so that the recipient of +the Larger Work may, at their option, further distribute the Covered +Software under the terms of either this License or such Secondary +License(s). + +3.4. Notices + +You may not remove or alter the substance of any license notices +(including copyright notices, patent notices, disclaimers of warranty, +or limitations of liability) contained within the Source Code Form of +the Covered Software, except that You may alter any license notices to +the extent required to remedy known factual inaccuracies. + +3.5. Application of Additional Terms + +You may choose to offer, and to charge a fee for, warranty, support, +indemnity or liability obligations to one or more recipients of Covered +Software. However, You may do so only on Your own behalf, and not on +behalf of any Contributor. You must make it absolutely clear that any +such warranty, support, indemnity, or liability obligation is offered by +You alone, and You hereby agree to indemnify every Contributor for any +liability incurred by such Contributor as a result of warranty, support, +indemnity or liability terms You offer. You may include additional +disclaimers of warranty and limitations of liability specific to any +jurisdiction. + +4. Inability to Comply Due to Statute or Regulation +--------------------------------------------------- + +If it is impossible for You to comply with any of the terms of this +License with respect to some or all of the Covered Software due to +statute, judicial order, or regulation then You must: (a) comply with +the terms of this License to the maximum extent possible; and (b) +describe the limitations and the code they affect. Such description must +be placed in a text file included with all distributions of the Covered +Software under this License. Except to the extent prohibited by statute +or regulation, such description must be sufficiently detailed for a +recipient of ordinary skill to be able to understand it. + +5. Termination +-------------- + +5.1. The rights granted under this License will terminate automatically +if You fail to comply with any of its terms. However, if You become +compliant, then the rights granted under this License from a particular +Contributor are reinstated (a) provisionally, unless and until such +Contributor explicitly and finally terminates Your grants, and (b) on an +ongoing basis, if such Contributor fails to notify You of the +non-compliance by some reasonable means prior to 60 days after You have +come back into compliance. Moreover, Your grants from a particular +Contributor are reinstated on an ongoing basis if such Contributor +notifies You of the non-compliance by some reasonable means, this is the +first time You have received notice of non-compliance with this License +from such Contributor, and You become compliant prior to 30 days after +Your receipt of the notice. + +5.2. If You initiate litigation against any entity by asserting a patent +infringement claim (excluding declaratory judgment actions, +counter-claims, and cross-claims) alleging that a Contributor Version +directly or indirectly infringes any patent, then the rights granted to +You by any and all Contributors for the Covered Software under Section +2.1 of this License shall terminate. + +5.3. In the event of termination under Sections 5.1 or 5.2 above, all +end user license agreements (excluding distributors and resellers) which +have been validly granted by You or Your distributors under this License +prior to termination shall survive termination. + +************************************************************************ +* * +* 6. Disclaimer of Warranty * +* ------------------------- * +* * +* Covered Software is provided under this License on an "as is" * +* basis, without warranty of any kind, either expressed, implied, or * +* statutory, including, without limitation, warranties that the * +* Covered Software is free of defects, merchantable, fit for a * +* particular purpose or non-infringing. The entire risk as to the * +* quality and performance of the Covered Software is with You. * +* Should any Covered Software prove defective in any respect, You * +* (not any Contributor) assume the cost of any necessary servicing, * +* repair, or correction. This disclaimer of warranty constitutes an * +* essential part of this License. No use of any Covered Software is * +* authorized under this License except under this disclaimer. * +* * +************************************************************************ + +************************************************************************ +* * +* 7. Limitation of Liability * +* -------------------------- * +* * +* Under no circumstances and under no legal theory, whether tort * +* (including negligence), contract, or otherwise, shall any * +* Contributor, or anyone who distributes Covered Software as * +* permitted above, be liable to You for any direct, indirect, * +* special, incidental, or consequential damages of any character * +* including, without limitation, damages for lost profits, loss of * +* goodwill, work stoppage, computer failure or malfunction, or any * +* and all other commercial damages or losses, even if such party * +* shall have been informed of the possibility of such damages. This * +* limitation of liability shall not apply to liability for death or * +* personal injury resulting from such party's negligence to the * +* extent applicable law prohibits such limitation. Some * +* jurisdictions do not allow the exclusion or limitation of * +* incidental or consequential damages, so this exclusion and * +* limitation may not apply to You. * +* * +************************************************************************ + +8. Litigation +------------- + +Any litigation relating to this License may be brought only in the +courts of a jurisdiction where the defendant maintains its principal +place of business and such litigation shall be governed by laws of that +jurisdiction, without reference to its conflict-of-law provisions. +Nothing in this Section shall prevent a party's ability to bring +cross-claims or counter-claims. + +9. Miscellaneous +---------------- + +This License represents the complete agreement concerning the subject +matter hereof. If any provision of this License is held to be +unenforceable, such provision shall be reformed only to the extent +necessary to make it enforceable. Any law or regulation which provides +that the language of a contract shall be construed against the drafter +shall not be used to construe this License against a Contributor. + +10. Versions of the License +--------------------------- + +10.1. New Versions + +Mozilla Foundation is the license steward. Except as provided in Section +10.3, no one other than the license steward has the right to modify or +publish new versions of this License. Each version will be given a +distinguishing version number. + +10.2. Effect of New Versions + +You may distribute the Covered Software under the terms of the version +of the License under which You originally received the Covered Software, +or under the terms of any subsequent version published by the license +steward. + +10.3. Modified Versions + +If you create software not governed by this License, and you want to +create a new license for such software, you may create and use a +modified version of this License if you rename the license and remove +any references to the name of the license steward (except to note that +such modified license differs from this License). + +10.4. Distributing Source Code Form that is Incompatible With Secondary +Licenses + +If You choose to distribute Source Code Form that is Incompatible With +Secondary Licenses under the terms of this version of the License, the +notice described in Exhibit B of this License must be attached. + +Exhibit A - Source Code Form License Notice +------------------------------------------- + + This Source Code Form is subject to the terms of the Mozilla Public + License, v. 2.0. If a copy of the MPL was not distributed with this + file, You can obtain one at http://mozilla.org/MPL/2.0/. + +If it is not possible or desirable to put the notice in a particular +file, then You may include the notice in a location (such as a LICENSE +file in a relevant directory) where a recipient would be likely to look +for such a notice. + +You may add additional accurate notices of copyright ownership. + +Exhibit B - "Incompatible With Secondary Licenses" Notice +--------------------------------------------------------- + + This Source Code Form is "Incompatible With Secondary Licenses", as + defined by the Mozilla Public License, v. 2.0. +``` + ## h11 (0.16.0) ### Licenses @@ -4503,7 +5134,7 @@ Apache License limitations under the License. ``` -## protobuf (7.35.1) +## protobuf (6.33.6) ### Licenses License: `BSD-3-Clause` diff --git a/Cargo.lock b/Cargo.lock index 091a88c93..87581d293 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1348,10 +1348,12 @@ dependencies = [ "futures", "futures-util", "getrandom 0.3.4", + "hyper-util", "js-sys", "libloading", "nemo-relay-plugin", "nemo-relay-types", + "nemo-relay-worker-proto", "object_store", "openinference-semantic-conventions", "opentelemetry", @@ -1373,6 +1375,7 @@ dependencies = [ "tokio-tungstenite", "toml", "tonic", + "tower", "typed-builder", "uuid", "wasm-bindgen", diff --git a/codecov.yml b/codecov.yml index f79453093..0a927628d 100644 --- a/codecov.yml +++ b/codecov.yml @@ -54,18 +54,6 @@ component_management: threshold: 0.5% base: auto if_ci_failed: error - - component_id: plugin_sdk - name: Plugin SDK - paths: - - "crates/plugin/src" - - "crates/worker-proto/src" - - "crates/worker/src" - statuses: - - type: project - target: 95% - threshold: 0.5% - base: auto - if_ci_failed: error - component_id: shared_types name: Shared DTO Types paths: @@ -128,6 +116,20 @@ component_management: threshold: 0.5% base: auto if_ci_failed: error + - component_id: plugin_sdk + name: Dynamic Plugin SDKs + paths: + - "crates/types/src" + - "crates/plugin/src" + - "crates/worker-proto/src" + - "crates/worker/src" + - "python/plugin/src/nemo_relay_plugin" + statuses: + - type: project + target: 95% + threshold: 0.5% + base: auto + if_ci_failed: error comment: after_n_builds: 22 @@ -157,6 +159,7 @@ ignore: - "target/" - "**/*.d.ts" - "**/*.pyi" + - "python/plugin/src/nemo_relay_plugin/_proto/**" - "python/nemo_relay/lib_native*.dylib.dSYM/**" # WebAssembly Rust wrappers are covered through wasm-pack execution and # reported through generated package JavaScript coverage. diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index d3ef689d1..07e1087b9 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -25,7 +25,7 @@ default = ["atof-streaming"] atof-streaming = ["nemo-relay/atof-streaming"] [dependencies] -nemo-relay = { workspace = true, features = ["guardrails-remote", "object-store", "openinference"] } +nemo-relay = { workspace = true, features = ["guardrails-remote", "object-store", "openinference", "worker-grpc"] } nemo-relay-adaptive = { workspace = true, features = ["redis-backend"] } nemo-relay-pii-redaction.workspace = true async-stream = "0.3" diff --git a/crates/cli/src/server.rs b/crates/cli/src/server.rs index 9765cddd4..dc6977917 100644 --- a/crates/cli/src/server.rs +++ b/crates/cli/src/server.rs @@ -10,7 +10,8 @@ use axum::http::HeaderMap; use axum::routing::{get, post}; use axum::{Json, Router}; use nemo_relay::plugin::dynamic::{ - DynamicPluginKind, NativePluginActivation, NativePluginLoadSpec, load_native_plugins, + DynamicPluginKind, NativePluginActivation, NativePluginLoadSpec, WorkerPluginActivation, + WorkerPluginLoadSpec, load_native_plugins, load_worker_plugins, }; use nemo_relay::plugin::{ PluginComponentSpec, PluginConfig, clear_plugin_configuration, initialize_plugins_exact, @@ -216,6 +217,7 @@ async fn idle_shutdown_future( struct PluginActivation { active: bool, native: Option, + worker: Option, } impl PluginActivation { @@ -227,6 +229,7 @@ impl PluginActivation { return Ok(Self { active: false, native: None, + worker: None, }); }; register_adaptive_component().map_err(|error| { @@ -251,6 +254,23 @@ impl PluginActivation { }) }) .collect::, CliError>>()?; + let worker_specs = dynamic_plugins + .iter() + .filter(|plugin| plugin.kind == DynamicPluginKind::Worker) + .map(|plugin| { + let manifest_ref = plugin.manifest_ref.clone().ok_or_else(|| { + CliError::Config(format!( + "worker dynamic plugin '{}' has no manifest_ref in lifecycle state", + plugin.plugin_id + )) + })?; + Ok(WorkerPluginLoadSpec { + plugin_id: plugin.plugin_id.clone(), + manifest_ref, + config: plugin.config.clone(), + }) + }) + .collect::, CliError>>()?; let native = if native_specs.is_empty() { None @@ -259,6 +279,14 @@ impl PluginActivation { CliError::Config(format!("native plugin load failed: {error}")) })?) }; + let worker = + if worker_specs.is_empty() { + None + } else { + Some(load_worker_plugins(worker_specs).map_err(|error| { + CliError::Config(format!("worker plugin load failed: {error}")) + })?) + }; // Gateway already resolved its config; activate exactly (no re-discovery). let mut plugin_config: PluginConfig = match config { Some(config) => serde_json::from_value(config) @@ -282,6 +310,7 @@ impl PluginActivation { Ok(Self { active: true, native, + worker, }) } @@ -295,6 +324,7 @@ impl PluginActivation { Ok(()) }; self.native.take(); + self.worker.take(); result } } diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 97f668690..9082eb91c 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -73,9 +73,24 @@ openinference = [ "dep:wasm-bindgen", "dep:wasm-bindgen-futures", ] +worker-grpc = [ + "dep:nemo-relay-worker-proto", + "dep:hyper-util", + "dep:tower", + "dep:tonic", + "tokio-stream/net", + "tonic/codegen", + "tonic/router", + "tonic/transport", + "tokio/io-util", + "tokio/net", + "tokio/process", + "tokio/rt-multi-thread", +] [dependencies] nemo-relay-types.workspace = true +nemo-relay-worker-proto = { workspace = true, optional = true } uuid = { workspace = true, features = ["v7", "serde"] } serde = { version = "1", features = ["derive", "rc"] } serde_json = "1" @@ -120,6 +135,11 @@ path = "tests/integration/context_isolation_tests.rs" name = "native_plugin_integration" path = "tests/integration/native_plugin_tests.rs" +[[test]] +name = "worker_plugin_integration" +path = "tests/integration/worker_plugin_tests.rs" +required-features = ["worker-grpc"] + [[test]] name = "middleware_integration" path = "tests/integration/middleware_tests.rs" @@ -163,3 +183,5 @@ rustls = { version = "0.23", default-features = false, features = ["ring", "std" tonic = { version = "0.14.1", default-features = false, optional = true } object_store = { version = "0.13", default-features = false, features = ["aws"], optional = true } tokio-tungstenite = { version = "0.27", default-features = false, features = ["connect", "rustls-tls-native-roots"], optional = true } +tower = { version = "0.5", features = ["util"], optional = true } +hyper-util = { version = "0.1", features = ["tokio"], optional = true } diff --git a/crates/core/src/plugin/dynamic.rs b/crates/core/src/plugin/dynamic.rs index a15af4081..05182c7fb 100644 --- a/crates/core/src/plugin/dynamic.rs +++ b/crates/core/src/plugin/dynamic.rs @@ -22,11 +22,15 @@ mod manifest; #[cfg(not(target_arch = "wasm32"))] mod native; mod registry; +#[cfg(all(feature = "worker-grpc", not(target_arch = "wasm32")))] +mod worker; pub use manifest::*; #[cfg(not(target_arch = "wasm32"))] pub use native::*; pub use registry::*; +#[cfg(all(feature = "worker-grpc", not(target_arch = "wasm32")))] +pub use worker::*; /// Plugin execution lane. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash, Display)] diff --git a/crates/core/src/plugin/dynamic/worker.rs b/crates/core/src/plugin/dynamic/worker.rs new file mode 100644 index 000000000..a4e46e562 --- /dev/null +++ b/crates/core/src/plugin/dynamic/worker.rs @@ -0,0 +1,2038 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! gRPC worker dynamic plugin loader and host-side proxy adapter. + +use std::collections::HashMap; +use std::future::Future; +use std::path::{Path, PathBuf}; +use std::pin::Pin; +use std::process::{Child, Command, Stdio}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use nemo_relay_worker_proto::v1::plugin_worker_client::PluginWorkerClient; +use nemo_relay_worker_proto::v1::relay_host_runtime_server::{ + RelayHostRuntime, RelayHostRuntimeServer, +}; +use nemo_relay_worker_proto::v1::{ + CreateScopeStackRequest, CreateScopeStackResponse, DropScopeStackRequest, EmitMarkRequest, + GuardrailResult, HandshakeRequest, HealthRequest, HostAck, InvokeRequest, InvokeResponse, + JsonEnvelope, JsonResult, LlmInvocation, LlmNextRequest, LlmStreamNextRequest, PopScopeRequest, + PushScopeRequest, PushScopeResponse, RegisterRequest, RegisterResponse, Registration, + RegistrationSurface, ScopeContext, ShutdownRequest, StreamChunk, ToolInvocation, + ToolNextRequest, ValidateRequest, WorkerError, +}; +use nemo_relay_worker_proto::{WORKER_PROTOCOL_GRPC_V1, decode_json_envelope, json_envelope}; +use semver::{Version, VersionReq}; +use serde_json::{Map, Value as Json}; +use tokio::runtime::{Builder as RuntimeBuilder, Runtime}; +use tokio::sync::{mpsc, oneshot}; +use tokio_stream::StreamExt; +use tonic::transport::{Channel, Endpoint, Server}; +use tonic::{Request, Response, Status}; +use uuid::Uuid; + +#[cfg(unix)] +use hyper_util::rt::TokioIo; +#[cfg(not(unix))] +use std::net::{SocketAddr, TcpListener}; +#[cfg(unix)] +use std::os::unix::net::UnixListener as StdUnixListener; +#[cfg(not(unix))] +use tokio::net::TcpListener as TokioTcpListener; +#[cfg(unix)] +use tokio::net::{UnixListener, UnixStream}; +#[cfg(not(unix))] +use tokio_stream::wrappers::TcpListenerStream; +#[cfg(unix)] +use tokio_stream::wrappers::UnixListenerStream; +#[cfg(unix)] +use tower::service_fn; + +use crate::api::event::Event; +use crate::api::llm::LlmRequest; +use crate::api::runtime::{ + LlmExecutionNextFn, LlmJsonStream, LlmStreamExecutionNextFn, ToolExecutionNextFn, + current_scope_stack, with_scope_stack, +}; +use crate::api::scope::{ + EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeAttributes, ScopeHandle, ScopeType, + event as emit_scope_mark, pop_scope, push_scope, +}; +use crate::codec::request::AnnotatedLlmRequest; +use crate::error::{FlowError, Result as FlowResult}; +use crate::plugin::{ + ConfigDiagnostic, DiagnosticLevel, Plugin, PluginError, PluginRegistrationContext, + deregister_plugin, register_plugin, +}; + +use super::{DynamicPluginKind, DynamicPluginManifest, DynamicPluginManifestLoad, WorkerRuntime}; + +const JSON_SCHEMA: &str = "nemo.relay.Json@1"; +const EVENT_SCHEMA: &str = "nemo.relay.Event@1"; +const LLM_REQUEST_SCHEMA: &str = "nemo.relay.LlmRequest@1"; +const ANNOTATED_LLM_REQUEST_SCHEMA: &str = "nemo.relay.AnnotatedLlmRequest@1"; +const WORKER_STARTUP_TIMEOUT: Duration = Duration::from_secs(10); +const WORKER_RPC_TIMEOUT: Duration = Duration::from_secs(30); +const WORKER_CONNECT_RETRY: Duration = Duration::from_millis(25); +const PYTHON_WORKER_BOOTSTRAP: &str = r#" +import asyncio +import importlib +import inspect +import sys + +target = sys.argv[1] +module_name, separator, function_name = target.partition(":") +if not separator or not module_name or not function_name: + raise SystemExit("Python worker entrypoint must be 'module:function'") + +entrypoint = getattr(importlib.import_module(module_name), function_name) +result = entrypoint() +if inspect.isawaitable(result): + asyncio.run(result) +"#; + +/// Worker plugin load request derived from host dynamic-plugin state. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkerPluginLoadSpec { + /// Expected plugin id. + pub plugin_id: String, + /// Path to the authored `relay-plugin.toml`. + pub manifest_ref: String, + /// Resolved dynamic plugin config passed to the worker. + pub config: Map, +} + +/// Owns gRPC worker processes registered into the plugin registry. +/// +/// Dropping this value deregisters worker plugin kinds and shuts down worker +/// processes. Clear active plugin configuration before dropping it so runtime +/// callbacks cannot outlive the worker activation. +pub struct WorkerPluginActivation { + plugins: Vec>, + plugin_kinds: Vec, +} + +impl WorkerPluginActivation { + /// Returns `true` when no worker plugins were loaded. + pub fn is_empty(&self) -> bool { + self.plugins.is_empty() + } + + /// Consumes the activation; deregistration runs from `Drop`. + pub fn clear(self) {} +} + +impl Drop for WorkerPluginActivation { + fn drop(&mut self) { + for plugin_kind in self.plugin_kinds.iter().rev() { + let _ = deregister_plugin(plugin_kind); + } + } +} + +/// Loads gRPC worker plugins and registers their plugin kinds. +/// +/// The returned activation must be kept alive until after active plugin +/// configuration has been cleared. +pub fn load_worker_plugins(specs: I) -> crate::plugin::Result +where + I: IntoIterator, +{ + let mut activation = WorkerPluginActivation { + plugins: Vec::new(), + plugin_kinds: Vec::new(), + }; + for spec in specs { + let instance = load_one_worker_plugin(&spec)?; + let plugin_kind = instance.plugin_kind.clone(); + register_plugin(Arc::new(WorkerPluginAdapter { + plugin_kind: plugin_kind.clone(), + allows_multiple_components: instance.allows_multiple_components, + instance: instance.clone(), + }))?; + activation.plugins.push(instance); + activation.plugin_kinds.push(plugin_kind); + } + Ok(activation) +} + +struct WorkerPluginAdapter { + plugin_kind: String, + allows_multiple_components: bool, + instance: Arc, +} + +impl Plugin for WorkerPluginAdapter { + fn plugin_kind(&self) -> &str { + &self.plugin_kind + } + + fn allows_multiple_components(&self) -> bool { + self.allows_multiple_components + } + + fn validate(&self, plugin_config: &Map) -> Vec { + if plugin_config != &self.instance.config { + return vec![worker_error_diagnostic( + &self.plugin_kind, + "plugin.worker_config_mismatch", + "worker plugin config changed after dynamic activation; reload the worker activation", + )]; + } + self.instance.validation_diagnostics.clone() + } + + fn register<'a>( + &'a self, + plugin_config: &Map, + ctx: &'a mut PluginRegistrationContext, + ) -> Pin> + Send + 'a>> { + let config_matches = plugin_config == &self.instance.config; + Box::pin(async move { + if !config_matches { + return Err(PluginError::RegistrationFailed( + "worker plugin config changed after dynamic activation; reload the worker activation" + .into(), + )); + } + self.instance.install_registrations(ctx) + }) + } +} + +struct WorkerPluginInstance { + plugin_kind: String, + allows_multiple_components: bool, + config: Map, + validation_diagnostics: Vec, + registrations: Vec, + runtime: OwnedWorkerRuntime, + client: PluginWorkerClient, + host_state: Arc, + shutdown: Mutex>>, + process: Mutex>, + activation_dir: PathBuf, +} + +impl Drop for WorkerPluginInstance { + fn drop(&mut self) { + let mut client = self.client.clone(); + let request = ShutdownRequest { + activation_id: self.host_state.activation_id.clone(), + auth_token: self.host_state.auth_token.clone(), + reason: "plugin activation dropped".into(), + }; + let _ = block_on_runtime(self.runtime.runtime(), async move { + worker_rpc(client.shutdown(worker_rpc_request(request))).await + }); + if let Ok(mut shutdown) = self.shutdown.lock() + && let Some(sender) = shutdown.take() + { + let _ = sender.send(()); + } + if let Ok(mut process) = self.process.lock() + && let Some(mut child) = process.take() + { + let _ = child.kill(); + let _ = child.wait(); + } + let _ = std::fs::remove_dir_all(&self.activation_dir); + } +} + +fn load_one_worker_plugin( + spec: &WorkerPluginLoadSpec, +) -> crate::plugin::Result> { + let (manifest, manifest_ref) = DynamicPluginManifest::load_from_path(&spec.manifest_ref)?; + if manifest.plugin.id.trim() != spec.plugin_id { + return Err(PluginError::InvalidConfig(format!( + "dynamic plugin manifest id '{}' does not match expected id '{}'", + manifest.plugin.id, spec.plugin_id + ))); + } + if manifest.plugin.kind != DynamicPluginKind::Worker { + return Err(PluginError::InvalidConfig(format!( + "dynamic plugin '{}' is kind {}; worker loader only supports worker", + spec.plugin_id, manifest.plugin.kind + ))); + } + validate_relay_compatibility(manifest.compat.relay.as_deref())?; + let DynamicPluginManifestLoad::Worker(load) = &manifest.load else { + unreachable!("validated worker manifest must carry worker load contract"); + }; + let runtime = load + .runtime + .ok_or_else(|| PluginError::InvalidConfig("load.runtime is required".into()))?; + let entrypoint = load + .entrypoint + .as_deref() + .ok_or_else(|| PluginError::InvalidConfig("load.entrypoint is required".into()))?; + + let activation_uuid = Uuid::now_v7(); + let activation_id = activation_uuid.to_string(); + let auth_token = Uuid::now_v7().to_string(); + let activation_dir = std::env::temp_dir().join(format!("nmrw-{}", activation_uuid.simple())); + std::fs::create_dir_all(&activation_dir) + .map_err(|err| PluginError::Internal(format!("worker activation directory: {err}")))?; + let mut activation_dir_guard = ActivationDirGuard::new(activation_dir.clone()); + let runtime_handle = OwnedWorkerRuntime::new( + RuntimeBuilder::new_multi_thread() + .enable_all() + .thread_name("nemo-relay-worker-host") + .build() + .map_err(|err| PluginError::Internal(format!("worker runtime: {err}")))?, + ); + let WorkerEndpoints { + host_server, + host_advertise, + worker_advertise, + worker_connect, + worker_endpoint_file, + } = WorkerEndpoints::new(&activation_dir)?; + let host_state = Arc::new(WorkerHostRuntimeState::new( + activation_id.clone(), + auth_token.clone(), + )); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + runtime_handle.runtime().spawn(serve_host_runtime( + host_server, + host_state.clone(), + shutdown_rx, + )); + + let manifest_path = PathBuf::from(&manifest_ref); + let mut child = ChildGuard::new(spawn_worker_process(WorkerProcessLaunch { + runtime, + manifest_path: &manifest_path, + plugin_id: &spec.plugin_id, + entrypoint, + activation_id: &activation_id, + auth_token: &auth_token, + host_endpoint: &host_advertise, + worker_endpoint: &worker_advertise, + worker_endpoint_file: worker_endpoint_file.as_deref(), + })?); + let mut client = block_on_runtime( + runtime_handle.runtime(), + connect_worker_with_retry(&worker_connect), + )?; + + let health = block_on_runtime( + runtime_handle.runtime(), + worker_rpc(client.health(worker_rpc_request(HealthRequest { + activation_id: activation_id.clone(), + auth_token: auth_token.clone(), + }))), + ) + .map_err(|err| PluginError::RegistrationFailed(format!("worker health check failed: {err}")))?; + let health = health.into_inner(); + if !health.ok { + let message = format!("worker plugin health check failed: {}", health.message); + return Err(PluginError::RegistrationFailed(message)); + } + + let handshake = block_on_runtime( + runtime_handle.runtime(), + worker_rpc(client.handshake(worker_rpc_request(HandshakeRequest { + activation_id: activation_id.clone(), + plugin_id: spec.plugin_id.clone(), + relay_version: env!("CARGO_PKG_VERSION").into(), + worker_protocol: WORKER_PROTOCOL_GRPC_V1.into(), + auth_token: auth_token.clone(), + host_endpoint: host_advertise.clone(), + }))), + ) + .map_err(|err| PluginError::RegistrationFailed(format!("worker handshake failed: {err}")))?; + let handshake = handshake.into_inner(); + if handshake.plugin_id != spec.plugin_id || handshake.plugin_kind != spec.plugin_id { + return Err(PluginError::InvalidConfig(format!( + "worker plugin returned id '{}' kind '{}' but manifest id is '{}'", + handshake.plugin_id, handshake.plugin_kind, spec.plugin_id + ))); + } + if handshake.worker_protocol != WORKER_PROTOCOL_GRPC_V1 { + let message = format!( + "unsupported worker_protocol '{}'", + handshake.worker_protocol + ); + return Err(PluginError::InvalidConfig(message)); + } + + let config = Json::Object(spec.config.clone()); + let validate = block_on_runtime( + runtime_handle.runtime(), + worker_rpc(client.validate(worker_rpc_request(ValidateRequest { + activation_id: activation_id.clone(), + plugin_id: spec.plugin_id.clone(), + auth_token: auth_token.clone(), + config: Some(json_envelope(JSON_SCHEMA, &config)?), + }))), + ) + .map_err(|err| { + PluginError::RegistrationFailed(format!("worker validation RPC failed: {err}")) + })?; + let validate = validate.into_inner(); + if let Some(error) = validate.error { + return Err(worker_error_to_plugin(error, "worker validation failed")); + } + let validation_diagnostics = match validate.diagnostics { + Some(diagnostics) => decode_json_envelope::>(&diagnostics) + .map_err(PluginError::Serialization)?, + None => Vec::new(), + }; + + let registrations = if diagnostics_have_errors(&validation_diagnostics) { + Vec::new() + } else { + let register = block_on_runtime( + runtime_handle.runtime(), + worker_rpc(client.register(worker_rpc_request(RegisterRequest { + activation_id: activation_id.clone(), + plugin_id: spec.plugin_id.clone(), + auth_token: auth_token.clone(), + config: Some(json_envelope(JSON_SCHEMA, &config)?), + }))), + ) + .map_err(|err| { + PluginError::RegistrationFailed(format!("worker registration RPC failed: {err}")) + })?; + let register = register.into_inner(); + if let Some(error) = register.error { + return Err(worker_error_to_plugin(error, "worker registration failed")); + } + validate_registration_plan(&spec.plugin_id, ®ister)?; + register.registrations + }; + + Ok(Arc::new(WorkerPluginInstance { + plugin_kind: spec.plugin_id.clone(), + allows_multiple_components: handshake.allows_multiple_components, + config: spec.config.clone(), + validation_diagnostics, + registrations, + runtime: runtime_handle, + client, + host_state, + shutdown: Mutex::new(Some(shutdown_tx)), + process: Mutex::new(Some(child.take())), + activation_dir: activation_dir_guard.keep(), + })) +} + +enum HostRuntimeServer { + #[cfg(unix)] + Unix(StdUnixListener), + #[cfg(not(unix))] + Tcp(TcpListener), +} + +#[derive(Clone)] +enum WorkerConnectEndpoint { + #[cfg(unix)] + Unix(PathBuf), + #[cfg(not(unix))] + Tcp(String), + #[cfg(not(unix))] + Announced(PathBuf), +} + +struct WorkerEndpoints { + host_server: HostRuntimeServer, + host_advertise: String, + worker_advertise: String, + worker_connect: WorkerConnectEndpoint, + worker_endpoint_file: Option, +} + +impl WorkerEndpoints { + fn new(activation_dir: &Path) -> crate::plugin::Result { + #[cfg(not(unix))] + let _ = activation_dir; + + #[cfg(unix)] + { + let host_socket = activation_dir.join("host.sock"); + let worker_socket = activation_dir.join("worker.sock"); + let _ = std::fs::remove_file(&host_socket); + let host_listener = StdUnixListener::bind(&host_socket).map_err(|err| { + PluginError::RegistrationFailed(format!( + "failed to bind worker host runtime socket '{}': {err}", + host_socket.display() + )) + })?; + host_listener.set_nonblocking(true).map_err(|err| { + PluginError::RegistrationFailed(format!( + "failed to configure worker host runtime socket '{}': {err}", + host_socket.display() + )) + })?; + Ok(Self { + host_server: HostRuntimeServer::Unix(host_listener), + host_advertise: unix_endpoint_display(&host_socket), + worker_advertise: unix_endpoint_display(&worker_socket), + worker_connect: WorkerConnectEndpoint::Unix(worker_socket), + worker_endpoint_file: None, + }) + } + + #[cfg(not(unix))] + { + let (host_listener, host_addr) = bind_loopback_listener()?; + let worker_endpoint_file = activation_dir.join("worker-endpoint"); + Ok(Self { + host_server: HostRuntimeServer::Tcp(host_listener), + host_advertise: format!("http://{host_addr}"), + worker_advertise: "tcp://127.0.0.1:0".into(), + worker_connect: WorkerConnectEndpoint::Announced(worker_endpoint_file.clone()), + worker_endpoint_file: Some(worker_endpoint_file), + }) + } + } +} + +async fn serve_host_runtime( + endpoint: HostRuntimeServer, + state: Arc, + shutdown: oneshot::Receiver<()>, +) { + let service = RelayHostRuntimeServer::new(WorkerHostRuntimeService { state }); + let result = match endpoint { + #[cfg(unix)] + HostRuntimeServer::Unix(listener) => { + let listener = match UnixListener::from_std(listener) { + Ok(listener) => listener, + Err(err) => { + eprintln!("failed to attach worker host runtime socket: {err}"); + return; + } + }; + Server::builder() + .add_service(service) + .serve_with_incoming_shutdown(UnixListenerStream::new(listener), async { + let _ = shutdown.await; + }) + .await + } + #[cfg(not(unix))] + HostRuntimeServer::Tcp(listener) => { + let listener = match TokioTcpListener::from_std(listener) { + Ok(listener) => listener, + Err(err) => { + eprintln!("failed to attach worker host runtime endpoint: {err}"); + return; + } + }; + Server::builder() + .add_service(service) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown.await; + }) + .await + } + }; + if let Err(err) = result { + eprintln!("worker host runtime server failed: {err}"); + } +} + +async fn connect_worker_with_retry( + endpoint: &WorkerConnectEndpoint, +) -> crate::plugin::Result> { + let start = std::time::Instant::now(); + loop { + let connect_endpoint = match resolve_worker_connect_endpoint(endpoint) { + Ok(Some(endpoint)) => endpoint, + Ok(None) if start.elapsed() < WORKER_STARTUP_TIMEOUT => { + tokio::time::sleep(WORKER_CONNECT_RETRY).await; + continue; + } + Ok(None) => { + let message = format!( + "worker did not announce endpoint within {}s", + WORKER_STARTUP_TIMEOUT.as_secs() + ); + return Err(PluginError::RegistrationFailed(message)); + } + Err(err) => return Err(err), + }; + match connect_worker(&connect_endpoint).await { + Ok(client) => return Ok(client), + Err(err) if start.elapsed() < WORKER_STARTUP_TIMEOUT => { + let _ = err; + tokio::time::sleep(WORKER_CONNECT_RETRY).await; + } + Err(err) => { + let message = format!( + "worker did not start within {}s: {err}", + WORKER_STARTUP_TIMEOUT.as_secs() + ); + return Err(PluginError::RegistrationFailed(message)); + } + } + } +} + +#[cfg(not(unix))] +fn normalize_worker_tcp_endpoint(endpoint: &str) -> crate::plugin::Result { + let endpoint = endpoint.trim(); + if let Some(authority) = endpoint.strip_prefix("tcp://") { + if authority.is_empty() { + return Err(PluginError::RegistrationFailed( + "worker announced an empty TCP endpoint".into(), + )); + } + return Ok(format!("http://{authority}")); + } + if endpoint.starts_with("http://") { + return Ok(endpoint.to_owned()); + } + Err(PluginError::RegistrationFailed(format!( + "worker announced unsupported endpoint '{endpoint}'" + ))) +} + +fn resolve_worker_connect_endpoint( + endpoint: &WorkerConnectEndpoint, +) -> crate::plugin::Result> { + match endpoint { + #[cfg(unix)] + WorkerConnectEndpoint::Unix(path) => Ok(Some(WorkerConnectEndpoint::Unix(path.clone()))), + #[cfg(not(unix))] + WorkerConnectEndpoint::Tcp(endpoint) => Ok(Some(WorkerConnectEndpoint::Tcp( + normalize_worker_tcp_endpoint(endpoint)?, + ))), + #[cfg(not(unix))] + WorkerConnectEndpoint::Announced(path) => { + let endpoint = match std::fs::read_to_string(path) { + Ok(endpoint) => endpoint, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(err) => { + return Err(PluginError::RegistrationFailed(format!( + "failed to read worker endpoint file '{}': {err}", + path.display() + ))); + } + }; + Ok(Some(WorkerConnectEndpoint::Tcp( + normalize_worker_tcp_endpoint(endpoint.trim())?, + ))) + } + } +} + +async fn connect_worker( + endpoint: &WorkerConnectEndpoint, +) -> crate::plugin::Result> { + match endpoint { + #[cfg(unix)] + WorkerConnectEndpoint::Unix(socket) => { + let path = Arc::new(socket.to_path_buf()); + let endpoint = Endpoint::try_from("http://[::]:50051") + .map_err(|err| PluginError::Internal(format!("invalid worker endpoint: {err}")))?; + let channel = endpoint + .connect_with_connector(service_fn(move |_| { + let path = path.clone(); + async move { UnixStream::connect(&*path).await.map(TokioIo::new) } + })) + .await + .map_err(|err| { + PluginError::RegistrationFailed(format!( + "failed to connect to worker socket '{}': {err}", + socket.display() + )) + })?; + Ok(PluginWorkerClient::new(channel)) + } + #[cfg(not(unix))] + WorkerConnectEndpoint::Tcp(endpoint) => { + let channel = Endpoint::from_shared(endpoint.clone()) + .map_err(|err| PluginError::Internal(format!("invalid worker endpoint: {err}")))? + .connect() + .await + .map_err(|err| { + PluginError::RegistrationFailed(format!( + "failed to connect to worker endpoint '{endpoint}': {err}" + )) + })?; + Ok(PluginWorkerClient::new(channel)) + } + #[cfg(not(unix))] + WorkerConnectEndpoint::Announced(path) => Err(PluginError::Internal(format!( + "worker endpoint file '{}' was not resolved before connect", + path.display() + ))), + } +} + +struct WorkerProcessLaunch<'a> { + runtime: WorkerRuntime, + manifest_path: &'a Path, + plugin_id: &'a str, + entrypoint: &'a str, + activation_id: &'a str, + auth_token: &'a str, + host_endpoint: &'a str, + worker_endpoint: &'a str, + worker_endpoint_file: Option<&'a Path>, +} + +fn spawn_worker_process(spec: WorkerProcessLaunch<'_>) -> crate::plugin::Result { + let manifest_dir = spec + .manifest_path + .parent() + .unwrap_or_else(|| Path::new(".")); + let (mut command, command_display) = match spec.runtime { + WorkerRuntime::Python => { + let python = std::env::var("NEMO_RELAY_PYTHON").unwrap_or_else(|_| "python3".into()); + let mut command = Command::new(python); + command + .arg("-c") + .arg(PYTHON_WORKER_BOOTSTRAP) + .arg(spec.entrypoint); + (command, spec.entrypoint.to_string()) + } + WorkerRuntime::Rust | WorkerRuntime::Command => { + let entrypoint = resolve_manifest_relative_path(spec.manifest_path, spec.entrypoint); + let command_display = entrypoint.display().to_string(); + (Command::new(entrypoint), command_display) + } + }; + command + .current_dir(manifest_dir) + .env("NEMO_RELAY_WORKER_ID", spec.activation_id) + .env("NEMO_RELAY_PLUGIN_ID", spec.plugin_id) + .env("NEMO_RELAY_WORKER_SOCKET", spec.worker_endpoint) + .env("NEMO_RELAY_HOST_SOCKET", spec.host_endpoint) + .env("NEMO_RELAY_WORKER_TOKEN", spec.auth_token) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::inherit()); + if let Some(path) = spec.worker_endpoint_file { + command.env("NEMO_RELAY_WORKER_ENDPOINT_FILE", path); + } + command.spawn().map_err(|err| { + PluginError::RegistrationFailed(format!( + "failed to spawn {} worker '{}': {err}", + spec.runtime, command_display + )) + }) +} + +impl WorkerPluginInstance { + fn install_registrations( + &self, + ctx: &mut PluginRegistrationContext, + ) -> crate::plugin::Result<()> { + for registration in &self.registrations { + let surface = RegistrationSurface::try_from(registration.surface).map_err(|_| { + PluginError::RegistrationFailed(format!( + "worker plugin '{}' returned unsupported registration surface {}", + self.plugin_kind, registration.surface + )) + })?; + let name = registration.local_name.clone(); + let priority = registration.priority; + let break_chain = registration.break_chain; + match surface { + RegistrationSurface::Subscriber => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_subscriber( + &name, + Arc::new(move |event| { + let _ = instance.invoke_subscriber(&callback_name, event); + }), + )?; + } + RegistrationSurface::ToolSanitizeRequestGuardrail => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_tool_sanitize_request_guardrail( + &name, + priority, + Arc::new(move |tool_name, value| { + instance + .invoke_tool_json( + &callback_name, + RegistrationSurface::ToolSanitizeRequestGuardrail, + tool_name, + value.clone(), + None, + ) + .unwrap_or(value) + }), + )?; + } + RegistrationSurface::ToolSanitizeResponseGuardrail => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_tool_sanitize_response_guardrail( + &name, + priority, + Arc::new(move |tool_name, value| { + instance + .invoke_tool_json( + &callback_name, + RegistrationSurface::ToolSanitizeResponseGuardrail, + tool_name, + value.clone(), + None, + ) + .unwrap_or(value) + }), + )?; + } + RegistrationSurface::ToolConditionalExecutionGuardrail => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_tool_conditional_execution_guardrail( + &name, + priority, + Arc::new(move |tool_name, value| { + instance.invoke_tool_guardrail(&callback_name, tool_name, value.clone()) + }), + )?; + } + RegistrationSurface::ToolRequestIntercept => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_tool_request_intercept( + &name, + priority, + break_chain, + Arc::new(move |tool_name, value| { + instance.invoke_tool_json( + &callback_name, + RegistrationSurface::ToolRequestIntercept, + tool_name, + value, + None, + ) + }), + )?; + } + RegistrationSurface::ToolExecutionIntercept => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_tool_execution_intercept( + &name, + priority, + Arc::new(move |tool_name, value, next| { + let instance = instance.clone(); + let name = callback_name.clone(); + let tool_name = tool_name.to_string(); + Box::pin(async move { + instance + .invoke_tool_execution(&name, &tool_name, value, next) + .await + }) + }), + )?; + } + RegistrationSurface::LlmSanitizeRequestGuardrail => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_llm_sanitize_request_guardrail( + &name, + priority, + Arc::new(move |request| { + instance + .invoke_llm_request_json( + &callback_name, + RegistrationSurface::LlmSanitizeRequestGuardrail, + "", + request.clone(), + None, + None, + ) + .unwrap_or(request) + }), + )?; + } + RegistrationSurface::LlmSanitizeResponseGuardrail => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_llm_sanitize_response_guardrail( + &name, + priority, + Arc::new(move |value| { + instance + .invoke_llm_response_json( + &callback_name, + RegistrationSurface::LlmSanitizeResponseGuardrail, + "", + value.clone(), + ) + .unwrap_or(value) + }), + )?; + } + RegistrationSurface::LlmConditionalExecutionGuardrail => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_llm_conditional_execution_guardrail( + &name, + priority, + Arc::new(move |request| { + instance.invoke_llm_guardrail(&callback_name, request.clone()) + }), + )?; + } + RegistrationSurface::LlmRequestIntercept => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_llm_request_intercept( + &name, + priority, + break_chain, + Arc::new(move |model_name, request, annotated| { + instance.invoke_llm_request_intercept( + &callback_name, + model_name, + request, + annotated, + ) + }), + )?; + } + RegistrationSurface::LlmExecutionIntercept => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_llm_execution_intercept( + &name, + priority, + Arc::new(move |model_name, request, next| { + let instance = instance.clone(); + let name = callback_name.clone(); + let model_name = model_name.to_string(); + Box::pin(async move { + instance + .invoke_llm_execution(&name, &model_name, request, next) + .await + }) + }), + )?; + } + RegistrationSurface::LlmStreamExecutionIntercept => { + let instance = Arc::new(self.clone_for_callback()); + let callback_name = name.clone(); + ctx.register_llm_stream_execution_intercept( + &name, + priority, + Arc::new(move |model_name, request, next| { + let instance = instance.clone(); + let name = callback_name.clone(); + let model_name = model_name.to_string(); + Box::pin(async move { + instance + .invoke_llm_stream_execution(&name, &model_name, request, next) + .await + }) + }), + )?; + } + RegistrationSurface::Unspecified => { + return Err(PluginError::RegistrationFailed(format!( + "worker plugin '{}' returned unspecified registration surface", + self.plugin_kind + ))); + } + } + } + Ok(()) + } + + fn clone_for_callback(&self) -> WorkerPluginCallback { + WorkerPluginCallback { + activation_id: self.host_state.activation_id.clone(), + runtime: self.runtime.handle(), + client: self.client.clone(), + host_state: self.host_state.clone(), + } + } +} + +#[cfg(not(unix))] +fn bind_loopback_listener() -> crate::plugin::Result<(TcpListener, SocketAddr)> { + let listener = TcpListener::bind(("127.0.0.1", 0)).map_err(|err| { + PluginError::RegistrationFailed(format!( + "failed to bind worker host runtime endpoint: {err}" + )) + })?; + listener.set_nonblocking(true).map_err(|err| { + PluginError::RegistrationFailed(format!( + "failed to configure worker host runtime endpoint: {err}" + )) + })?; + let addr = listener.local_addr().map_err(|err| { + PluginError::RegistrationFailed(format!( + "failed to inspect worker host runtime endpoint: {err}" + )) + })?; + Ok((listener, addr)) +} + +#[derive(Clone)] +struct WorkerPluginCallback { + activation_id: String, + runtime: tokio::runtime::Handle, + client: PluginWorkerClient, + host_state: Arc, +} + +impl WorkerPluginCallback { + fn invoke_subscriber(&self, registration_name: &str, event: &Event) -> FlowResult<()> { + let request = self.base_request( + registration_name, + RegistrationSurface::Subscriber, + None, + Some(invoke_request_payload_event(event)), + ); + let response = self.invoke_blocking(request)?; + match response.result { + Some(invoke_response_result::Result::Empty(_)) | None => Ok(()), + Some(invoke_response_result::Result::Error(error)) => Err(worker_error_to_flow(error)), + _ => Err(FlowError::Internal( + "worker subscriber returned unexpected result".into(), + )), + } + } + + fn invoke_tool_json( + &self, + registration_name: &str, + surface: RegistrationSurface, + tool_name: &str, + value: Json, + continuation_id: Option, + ) -> FlowResult { + let request = self.base_request( + registration_name, + surface, + continuation_id, + Some(invoke_request_payload_tool(tool_name, value)), + ); + json_from_invoke_response(self.invoke_blocking(request)?) + } + + fn invoke_tool_guardrail( + &self, + registration_name: &str, + tool_name: &str, + value: Json, + ) -> FlowResult> { + let request = self.base_request( + registration_name, + RegistrationSurface::ToolConditionalExecutionGuardrail, + None, + Some(invoke_request_payload_tool(tool_name, value)), + ); + guardrail_from_invoke_response(self.invoke_blocking(request)?) + } + + async fn invoke_tool_execution( + &self, + registration_name: &str, + tool_name: &str, + value: Json, + next: ToolExecutionNextFn, + ) -> FlowResult { + let continuation_id = self + .host_state + .insert_continuation(Continuation::Tool(next))?; + let callback = self.clone(); + let registration_name = registration_name.to_string(); + let tool_name = tool_name.to_string(); + tokio::task::spawn_blocking(move || { + callback.invoke_tool_json( + ®istration_name, + RegistrationSurface::ToolExecutionIntercept, + &tool_name, + value, + Some(continuation_id.clone()), + ) + }) + .await + .map_err(|err| FlowError::Internal(format!("worker task join failed: {err}")))? + } + + fn invoke_llm_request_json( + &self, + registration_name: &str, + surface: RegistrationSurface, + model_name: &str, + request: LlmRequest, + annotated: Option, + continuation_id: Option, + ) -> FlowResult { + let invoke = self.base_request( + registration_name, + surface, + continuation_id, + Some(invoke_request_payload_llm( + model_name, + Some(request), + annotated, + None, + )), + ); + let value = json_from_invoke_response(self.invoke_blocking(invoke)?)?; + serde_json::from_value(value).map_err(|err| { + FlowError::Internal(format!("worker returned invalid LLM request: {err}")) + }) + } + + fn invoke_llm_response_json( + &self, + registration_name: &str, + surface: RegistrationSurface, + model_name: &str, + response: Json, + ) -> FlowResult { + let invoke = self.base_request( + registration_name, + surface, + None, + Some(invoke_request_payload_llm( + model_name, + None, + None, + Some(response), + )), + ); + json_from_invoke_response(self.invoke_blocking(invoke)?) + } + + fn invoke_llm_guardrail( + &self, + registration_name: &str, + request: LlmRequest, + ) -> FlowResult> { + let invoke = self.base_request( + registration_name, + RegistrationSurface::LlmConditionalExecutionGuardrail, + None, + Some(invoke_request_payload_llm("", Some(request), None, None)), + ); + guardrail_from_invoke_response(self.invoke_blocking(invoke)?) + } + + fn invoke_llm_request_intercept( + &self, + registration_name: &str, + model_name: &str, + request: LlmRequest, + annotated: Option, + ) -> FlowResult<(LlmRequest, Option)> { + let invoke = self.base_request( + registration_name, + RegistrationSurface::LlmRequestIntercept, + None, + Some(invoke_request_payload_llm( + model_name, + Some(request), + annotated, + None, + )), + ); + let response = self.invoke_blocking(invoke)?; + match response.result { + Some(invoke_response_result::Result::LlmRequest(result)) => { + let request = required_envelope(result.request, "llm request intercept request")?; + let request = decode_json_envelope::(&request).map_err(|err| { + FlowError::Internal(format!("worker returned invalid LLM request: {err}")) + })?; + let annotated = if result.has_annotated_request { + let envelope = required_envelope( + result.annotated_request, + "llm request intercept annotated request", + )?; + Some( + decode_json_envelope::(&envelope).map_err(|err| { + FlowError::Internal(format!( + "worker returned invalid annotated LLM request: {err}" + )) + })?, + ) + } else { + None + }; + Ok((request, annotated)) + } + Some(invoke_response_result::Result::Error(error)) => Err(worker_error_to_flow(error)), + _ => Err(FlowError::Internal( + "worker LLM request intercept returned unexpected result".into(), + )), + } + } + + async fn invoke_llm_execution( + &self, + registration_name: &str, + model_name: &str, + request: LlmRequest, + next: LlmExecutionNextFn, + ) -> FlowResult { + let continuation_id = self + .host_state + .insert_continuation(Continuation::Llm(next))?; + let callback = self.clone(); + let registration_name = registration_name.to_string(); + let model_name = model_name.to_string(); + tokio::task::spawn_blocking(move || { + let invoke = callback.base_request( + ®istration_name, + RegistrationSurface::LlmExecutionIntercept, + Some(continuation_id), + Some(invoke_request_payload_llm( + &model_name, + Some(request), + None, + None, + )), + ); + json_from_invoke_response(callback.invoke_blocking(invoke)?) + }) + .await + .map_err(|err| FlowError::Internal(format!("worker task join failed: {err}")))? + } + + async fn invoke_llm_stream_execution( + &self, + registration_name: &str, + model_name: &str, + request: LlmRequest, + next: LlmStreamExecutionNextFn, + ) -> FlowResult { + let continuation_id = self + .host_state + .insert_continuation(Continuation::LlmStream(next))?; + let invoke = self.base_request( + registration_name, + RegistrationSurface::LlmStreamExecutionIntercept, + Some(continuation_id.clone()), + Some(invoke_request_payload_llm( + model_name, + Some(request), + None, + None, + )), + ); + let scope_stack_id = invoke + .scope + .as_ref() + .map(|scope| scope.scope_stack_id.clone()) + .unwrap_or_default(); + let mut client = self.client.clone(); + let host_state = self.host_state.clone(); + let (tx, rx) = mpsc::channel(16); + self.runtime.spawn(async move { + let result = worker_rpc(client.invoke_stream(worker_rpc_request(invoke))).await; + match result { + Ok(response) => { + let mut stream = response.into_inner(); + while let Some(item) = stream.next().await { + let result = match item { + Ok(chunk) => json_from_stream_chunk(chunk), + Err(err) => Err(FlowError::Internal(format!( + "worker stream transport failed: {err}" + ))), + }; + if tx.send(result).await.is_err() { + break; + } + } + } + Err(err) => { + let _ = tx + .send(Err(FlowError::Internal(format!( + "worker stream invoke failed: {err}" + )))) + .await; + } + } + host_state.remove_continuation(&continuation_id); + if !scope_stack_id.is_empty() { + host_state.remove_invocation_scope_stack(&scope_stack_id); + } + }); + Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx))) + } + + fn base_request( + &self, + registration_name: &str, + surface: RegistrationSurface, + continuation_id: Option, + payload: Option, + ) -> InvokeRequest { + let scope_stack_id = self + .host_state + .insert_invocation_scope_stack(current_scope_stack()); + InvokeRequest { + activation_id: self.activation_id.clone(), + auth_token: self.host_state.auth_token.clone(), + invocation_id: Uuid::now_v7().to_string(), + registration_name: registration_name.into(), + surface: surface as i32, + continuation_id: continuation_id.unwrap_or_default(), + scope: Some(ScopeContext { + scope_stack_id, + parent_scope_id: String::new(), + }), + payload, + } + } + + fn invoke_blocking(&self, request: InvokeRequest) -> FlowResult { + let scope_stack_id = request + .scope + .as_ref() + .map(|scope| scope.scope_stack_id.clone()) + .unwrap_or_default(); + let continuation_id = request.continuation_id.clone(); + let mut client = self.client.clone(); + let result = block_on_handle(&self.runtime, async move { + worker_rpc(client.invoke(worker_rpc_request(request))).await + }) + .map(|response| response.into_inner()) + .map_err(|err| FlowError::Internal(format!("worker invoke failed: {err}"))); + if !continuation_id.is_empty() { + self.host_state.remove_continuation(&continuation_id); + } + if !scope_stack_id.is_empty() { + self.host_state + .remove_invocation_scope_stack(&scope_stack_id); + } + result + } +} + +struct OwnedWorkerRuntime { + runtime: Option, +} + +impl OwnedWorkerRuntime { + fn new(runtime: Runtime) -> Self { + Self { + runtime: Some(runtime), + } + } + + fn runtime(&self) -> &Runtime { + self.runtime + .as_ref() + .expect("worker runtime accessed after drop") + } + + fn handle(&self) -> tokio::runtime::Handle { + self.runtime().handle().clone() + } +} + +impl Drop for OwnedWorkerRuntime { + fn drop(&mut self) { + let Some(runtime) = self.runtime.take() else { + return; + }; + if tokio::runtime::Handle::try_current().is_ok() { + std::thread::scope(|scope| { + scope + .spawn(move || drop(runtime)) + .join() + .expect("worker runtime drop thread panicked"); + }); + } else { + drop(runtime); + } + } +} + +struct ChildGuard { + child: Option, +} + +impl ChildGuard { + fn new(child: Child) -> Self { + Self { child: Some(child) } + } + + fn take(&mut self) -> Child { + self.child.take().expect("worker child already taken") + } +} + +impl Drop for ChildGuard { + fn drop(&mut self) { + if let Some(mut child) = self.child.take() { + let _ = child.kill(); + let _ = child.wait(); + } + } +} + +struct ActivationDirGuard { + path: Option, +} + +impl ActivationDirGuard { + fn new(path: PathBuf) -> Self { + Self { path: Some(path) } + } + + fn keep(&mut self) -> PathBuf { + self.path + .take() + .expect("worker activation directory already taken") + } +} + +impl Drop for ActivationDirGuard { + fn drop(&mut self) { + if let Some(path) = self.path.take() { + let _ = std::fs::remove_dir_all(path); + } + } +} + +fn worker_rpc_request(message: T) -> Request { + Request::new(message) +} + +async fn worker_rpc(future: F) -> Result, Status> +where + F: Future, Status>>, +{ + match tokio::time::timeout(WORKER_RPC_TIMEOUT, future).await { + Ok(result) => result, + Err(_) => Err(Status::deadline_exceeded(format!( + "worker RPC timed out after {}s", + WORKER_RPC_TIMEOUT.as_secs() + ))), + } +} + +fn block_on_runtime(runtime: &Runtime, future: F) -> F::Output +where + F: Future + Send, + F::Output: Send, +{ + if tokio::runtime::Handle::try_current().is_ok() { + std::thread::scope(|scope| { + scope + .spawn(|| runtime.block_on(future)) + .join() + .expect("worker runtime blocking thread panicked") + }) + } else { + runtime.block_on(future) + } +} + +fn block_on_handle(handle: &tokio::runtime::Handle, future: F) -> F::Output +where + F: Future + Send, + F::Output: Send, +{ + if tokio::runtime::Handle::try_current().is_ok() { + let handle = handle.clone(); + std::thread::scope(|scope| { + scope + .spawn(move || handle.block_on(future)) + .join() + .expect("worker callback blocking thread panicked") + }) + } else { + handle.block_on(future) + } +} + +struct WorkerHostRuntimeState { + activation_id: String, + auth_token: String, + scope_stacks: Mutex>, + scope_handles: Mutex>, + continuations: Mutex>, +} + +struct StoredScopeHandle { + handle: ScopeHandle, + scope_stack_id: String, +} + +impl WorkerHostRuntimeState { + fn new(activation_id: String, auth_token: String) -> Self { + Self { + activation_id, + auth_token, + scope_stacks: Mutex::new(HashMap::new()), + scope_handles: Mutex::new(HashMap::new()), + continuations: Mutex::new(HashMap::new()), + } + } + + fn authorize(&self, activation_id: &str, token: &str) -> Result<(), Status> { + if activation_id != self.activation_id || token != self.auth_token { + return Err(Status::permission_denied("invalid worker host token")); + } + Ok(()) + } + + fn insert_invocation_scope_stack( + &self, + stack: crate::api::runtime::ScopeStackHandle, + ) -> String { + let id = format!("invoke-{}", Uuid::now_v7()); + if let Ok(mut stacks) = self.scope_stacks.lock() { + stacks.insert(id.clone(), stack); + } + id + } + + fn remove_invocation_scope_stack(&self, id: &str) { + if let Ok(mut stacks) = self.scope_stacks.lock() { + stacks.remove(id); + } + } + + fn insert_continuation(&self, continuation: Continuation) -> FlowResult { + let id = format!("next-{}", Uuid::now_v7()); + let mut continuations = self + .continuations + .lock() + .map_err(|err| FlowError::Internal(format!("continuation lock poisoned: {err}")))?; + continuations.insert(id.clone(), continuation); + Ok(id) + } + + fn remove_continuation(&self, id: &str) { + if let Ok(mut continuations) = self.continuations.lock() { + continuations.remove(id); + } + } + + fn continuation(&self, id: &str) -> Result { + self.continuations + .lock() + .map_err(|err| Status::internal(format!("continuation lock poisoned: {err}")))? + .get(id) + .cloned() + .ok_or_else(|| Status::not_found("continuation not found")) + } + + fn stack(&self, id: &str) -> Result, Status> { + if id.is_empty() { + return Ok(None); + } + self.scope_stacks + .lock() + .map_err(|err| Status::internal(format!("scope stack lock poisoned: {err}")))? + .get(id) + .cloned() + .map(Some) + .ok_or_else(|| Status::not_found("scope stack not found")) + } +} + +#[derive(Clone)] +enum Continuation { + Tool(ToolExecutionNextFn), + Llm(LlmExecutionNextFn), + LlmStream(LlmStreamExecutionNextFn), +} + +struct WorkerHostRuntimeService { + state: Arc, +} + +#[tonic::async_trait] +impl RelayHostRuntime for WorkerHostRuntimeService { + async fn emit_mark( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let result = self.with_stack(request.scope.as_ref(), || { + emit_scope_mark( + EmitMarkEventParams::builder() + .name(&request.name) + .data_opt(optional_envelope_to_json(request.data)?) + .metadata_opt(optional_envelope_to_json(request.metadata)?) + .build(), + ) + }); + Ok(Response::new(host_ack(result))) + } + + async fn push_scope( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let result = self.with_stack(request.scope.as_ref(), || { + push_scope( + PushScopeParams::builder() + .name(&request.name) + .scope_type(proto_scope_type(request.scope_type)) + .attributes(ScopeAttributes::empty()) + .data_opt(optional_envelope_to_json(request.data)?) + .metadata_opt(optional_envelope_to_json(request.metadata)?) + .input_opt(optional_envelope_to_json(request.input)?) + .build(), + ) + }); + match result { + Ok(handle) => { + let id = format!("scope-{}", handle.uuid); + let scope_stack_id = request + .scope + .as_ref() + .map(|scope| scope.scope_stack_id.clone()) + .unwrap_or_default(); + self.state + .scope_handles + .lock() + .map_err(|err| Status::internal(format!("scope handle lock poisoned: {err}")))? + .insert( + id.clone(), + StoredScopeHandle { + handle, + scope_stack_id, + }, + ); + Ok(Response::new(PushScopeResponse { + scope_handle_id: id, + error: None, + })) + } + Err(err) => Ok(Response::new(PushScopeResponse { + scope_handle_id: String::new(), + error: Some(flow_error_to_worker(err)), + })), + } + } + + async fn pop_scope( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let handle = self + .state + .scope_handles + .lock() + .map_err(|err| Status::internal(format!("scope handle lock poisoned: {err}")))? + .remove(&request.scope_handle_id) + .ok_or_else(|| Status::not_found("scope handle not found"))?; + let output = optional_envelope_to_json(request.output).map_err(status_from_flow)?; + let metadata = optional_envelope_to_json(request.metadata).map_err(status_from_flow)?; + let pop = || { + pop_scope( + PopScopeParams::builder() + .handle_uuid(&handle.handle.uuid) + .output_opt(output) + .metadata_opt(metadata) + .build(), + ) + }; + let result = if handle.scope_stack_id.is_empty() { + pop() + } else if let Some(stack) = self.state.stack(&handle.scope_stack_id)? { + with_scope_stack(stack, pop) + } else { + pop() + }; + Ok(Response::new(host_ack(result))) + } + + async fn create_scope_stack( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let id = format!("stack-{}", Uuid::now_v7()); + self.state + .scope_stacks + .lock() + .map_err(|err| Status::internal(format!("scope stack lock poisoned: {err}")))? + .insert(id.clone(), crate::api::runtime::create_scope_stack()); + Ok(Response::new(CreateScopeStackResponse { + scope_stack_id: id, + error: None, + })) + } + + async fn drop_scope_stack( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + self.state + .scope_stacks + .lock() + .map_err(|err| Status::internal(format!("scope stack lock poisoned: {err}")))? + .remove(&request.scope_stack_id); + Ok(Response::new(HostAck { + ok: true, + error: None, + })) + } + + async fn tool_next( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let continuation = self.state.continuation(&request.continuation_id)?; + let Continuation::Tool(next) = continuation else { + return Err(Status::invalid_argument( + "continuation is not a tool continuation", + )); + }; + let value = + required_envelope(request.value, "tool next value").map_err(status_from_flow)?; + let value = decode_json_envelope::(&value) + .map_err(|err| Status::invalid_argument(format!("invalid tool next JSON: {err}")))?; + let result = next(value).await; + Ok(Response::new(json_result(result))) + } + + async fn llm_next( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let continuation = self.state.continuation(&request.continuation_id)?; + let Continuation::Llm(next) = continuation else { + return Err(Status::invalid_argument( + "continuation is not an LLM continuation", + )); + }; + let request = + required_envelope(request.request, "llm next request").map_err(status_from_flow)?; + let request = decode_json_envelope::(&request) + .map_err(|err| Status::invalid_argument(format!("invalid LLM next request: {err}")))?; + let result = next(request).await; + Ok(Response::new(json_result(result))) + } + + type LlmStreamNextStream = + Pin> + Send>>; + + async fn llm_stream_next( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + self.state + .authorize(&request.activation_id, &request.auth_token)?; + let continuation = self.state.continuation(&request.continuation_id)?; + let Continuation::LlmStream(next) = continuation else { + return Err(Status::invalid_argument( + "continuation is not an LLM stream continuation", + )); + }; + let request = required_envelope(request.request, "llm stream next request") + .map_err(status_from_flow)?; + let request = decode_json_envelope::(&request).map_err(|err| { + Status::invalid_argument(format!("invalid LLM stream next request: {err}")) + })?; + let stream = next(request).await.map_err(status_from_flow)?; + let mapped = stream.map(|item| match item { + Ok(value) => Ok(StreamChunk { + item: Some(stream_chunk_item::Item::Value(json_envelope_infallible( + JSON_SCHEMA, + &value, + ))), + }), + Err(err) => Ok(StreamChunk { + item: Some(stream_chunk_item::Item::Error(flow_error_to_worker(err))), + }), + }); + Ok(Response::new(Box::pin(mapped))) + } +} + +impl WorkerHostRuntimeService { + fn with_stack( + &self, + scope: Option<&ScopeContext>, + f: impl FnOnce() -> FlowResult, + ) -> FlowResult { + let Some(stack_id) = scope.map(|scope| scope.scope_stack_id.as_str()) else { + return f(); + }; + let Some(stack) = self + .state + .stack(stack_id) + .map_err(|err| FlowError::Internal(err.to_string()))? + else { + return f(); + }; + with_scope_stack(stack, f) + } +} + +mod invoke_request_payload { + pub(crate) use nemo_relay_worker_proto::v1::invoke_request::Payload; +} + +mod invoke_response_result { + pub(crate) use nemo_relay_worker_proto::v1::invoke_response::Result; +} + +mod stream_chunk_item { + pub(crate) use nemo_relay_worker_proto::v1::stream_chunk::Item; +} + +fn invoke_request_payload_event(event: &Event) -> invoke_request_payload::Payload { + invoke_request_payload::Payload::Event(json_envelope_infallible(EVENT_SCHEMA, event)) +} + +fn invoke_request_payload_tool(tool_name: &str, value: Json) -> invoke_request_payload::Payload { + invoke_request_payload::Payload::Tool(ToolInvocation { + tool_name: tool_name.into(), + value: Some(json_envelope_infallible(JSON_SCHEMA, &value)), + }) +} + +fn invoke_request_payload_llm( + model_name: &str, + request: Option, + annotated_request: Option, + response: Option, +) -> invoke_request_payload::Payload { + invoke_request_payload::Payload::Llm(LlmInvocation { + model_name: model_name.into(), + request: request + .as_ref() + .map(|request| json_envelope_infallible(LLM_REQUEST_SCHEMA, request)), + annotated_request: annotated_request + .as_ref() + .map(|request| json_envelope_infallible(ANNOTATED_LLM_REQUEST_SCHEMA, request)), + response: response + .as_ref() + .map(|response| json_envelope_infallible(JSON_SCHEMA, response)), + }) +} + +fn json_envelope_infallible(schema: &str, value: &T) -> JsonEnvelope { + json_envelope(schema, value).expect("Relay DTO JSON serialization should be infallible") +} + +fn json_from_invoke_response(response: InvokeResponse) -> FlowResult { + match response.result { + Some(invoke_response_result::Result::Json(result)) => { + if let Some(error) = result.error { + return Err(worker_error_to_flow(error)); + } + let envelope = required_envelope(result.value, "worker JSON result")?; + decode_json_envelope::(&envelope).map_err(|err| { + FlowError::Internal(format!("worker returned invalid JSON result: {err}")) + }) + } + Some(invoke_response_result::Result::Error(error)) => Err(worker_error_to_flow(error)), + _ => Err(FlowError::Internal( + "worker returned unexpected invoke result".into(), + )), + } +} + +fn guardrail_from_invoke_response(response: InvokeResponse) -> FlowResult> { + match response.result { + Some(invoke_response_result::Result::Guardrail(GuardrailResult { block_reason })) => { + Ok((!block_reason.is_empty()).then_some(block_reason)) + } + Some(invoke_response_result::Result::Error(error)) => Err(worker_error_to_flow(error)), + _ => Err(FlowError::Internal( + "worker guardrail returned unexpected invoke result".into(), + )), + } +} + +fn json_from_stream_chunk(chunk: StreamChunk) -> FlowResult { + match chunk.item { + Some(stream_chunk_item::Item::Value(value)) => decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(format!("invalid worker stream chunk: {err}"))), + Some(stream_chunk_item::Item::Error(error)) => Err(worker_error_to_flow(error)), + None => Err(FlowError::Internal("worker stream chunk was empty".into())), + } +} + +fn required_envelope(value: Option, field: &str) -> FlowResult { + value.ok_or_else(|| FlowError::Internal(format!("{field} is missing"))) +} + +fn optional_envelope_to_json(value: Option) -> FlowResult> { + value + .map(|value| { + decode_json_envelope::(&value) + .map_err(|err| FlowError::Internal(format!("invalid JSON envelope: {err}"))) + }) + .transpose() +} + +fn host_ack(result: FlowResult<()>) -> HostAck { + match result { + Ok(()) => HostAck { + ok: true, + error: None, + }, + Err(err) => HostAck { + ok: false, + error: Some(flow_error_to_worker(err)), + }, + } +} + +fn json_result(result: FlowResult) -> JsonResult { + match result { + Ok(value) => JsonResult { + value: Some(json_envelope_infallible(JSON_SCHEMA, &value)), + error: None, + }, + Err(err) => JsonResult { + value: None, + error: Some(flow_error_to_worker(err)), + }, + } +} + +fn flow_error_to_worker(err: FlowError) -> WorkerError { + WorkerError { + code: "host.runtime_error".into(), + message: err.to_string(), + retryable: false, + } +} + +fn worker_error_to_flow(error: WorkerError) -> FlowError { + FlowError::Internal(format!("{}: {}", error.code, error.message)) +} + +fn worker_error_to_plugin(error: WorkerError, fallback: &str) -> PluginError { + let message = if error.message.is_empty() { + fallback.to_string() + } else { + format!("{}: {}", error.code, error.message) + }; + PluginError::RegistrationFailed(message) +} + +fn status_from_flow(err: FlowError) -> Status { + Status::internal(err.to_string()) +} + +fn proto_scope_type(scope_type: i32) -> ScopeType { + match nemo_relay_worker_proto::v1::ScopeType::try_from(scope_type) { + Ok(nemo_relay_worker_proto::v1::ScopeType::Agent) => ScopeType::Agent, + Ok(nemo_relay_worker_proto::v1::ScopeType::Function) => ScopeType::Function, + Ok(nemo_relay_worker_proto::v1::ScopeType::Tool) => ScopeType::Tool, + Ok(nemo_relay_worker_proto::v1::ScopeType::Llm) => ScopeType::Llm, + Ok(nemo_relay_worker_proto::v1::ScopeType::Retriever) => ScopeType::Retriever, + Ok(nemo_relay_worker_proto::v1::ScopeType::Embedder) => ScopeType::Embedder, + Ok(nemo_relay_worker_proto::v1::ScopeType::Reranker) => ScopeType::Reranker, + Ok(nemo_relay_worker_proto::v1::ScopeType::Guardrail) => ScopeType::Guardrail, + Ok(nemo_relay_worker_proto::v1::ScopeType::Evaluator) => ScopeType::Evaluator, + Ok(nemo_relay_worker_proto::v1::ScopeType::Custom) => ScopeType::Custom, + Ok(nemo_relay_worker_proto::v1::ScopeType::Unknown) => ScopeType::Unknown, + _ => ScopeType::Custom, + } +} + +fn validate_registration_plan( + plugin_id: &str, + response: &RegisterResponse, +) -> crate::plugin::Result<()> { + for registration in &response.registrations { + if registration.local_name.trim().is_empty() { + return Err(PluginError::RegistrationFailed(format!( + "worker plugin '{plugin_id}' returned a registration with empty local_name" + ))); + } + let surface = RegistrationSurface::try_from(registration.surface).map_err(|_| { + PluginError::RegistrationFailed(format!( + "worker plugin '{plugin_id}' returned unsupported registration surface {}", + registration.surface + )) + })?; + if surface == RegistrationSurface::Unspecified { + return Err(PluginError::RegistrationFailed(format!( + "worker plugin '{plugin_id}' returned unspecified registration surface" + ))); + } + } + Ok(()) +} + +fn diagnostics_have_errors(diagnostics: &[ConfigDiagnostic]) -> bool { + diagnostics + .iter() + .any(|diagnostic| diagnostic.level == DiagnosticLevel::Error) +} + +fn worker_error_diagnostic(plugin_kind: &str, code: &str, message: &str) -> ConfigDiagnostic { + ConfigDiagnostic { + level: DiagnosticLevel::Error, + code: code.into(), + component: Some(plugin_kind.into()), + field: None, + message: message.into(), + } +} + +fn validate_relay_compatibility(relay: Option<&str>) -> crate::plugin::Result<()> { + let relay = relay + .map(str::trim) + .filter(|value| !value.is_empty()) + .ok_or_else(|| PluginError::InvalidConfig("compat.relay is required".into()))?; + let req = VersionReq::parse(relay).map_err(|err| { + PluginError::InvalidConfig(format!("invalid compat.relay version requirement: {err}")) + })?; + let version = Version::parse(env!("CARGO_PKG_VERSION")) + .map_err(|err| PluginError::Internal(format!("failed to parse host version: {err}")))?; + if req.matches(&version) { + Ok(()) + } else { + Err(PluginError::InvalidConfig(format!( + "worker plugin requires relay '{relay}' but host version is {version}" + ))) + } +} + +fn resolve_manifest_relative_path(manifest_path: &Path, value: &str) -> PathBuf { + let path = PathBuf::from(value); + if path.is_absolute() { + path + } else { + manifest_path + .parent() + .map(|parent| parent.join(&path)) + .unwrap_or(path) + } +} + +#[cfg(unix)] +fn unix_endpoint_display(path: &Path) -> String { + format!("unix://{}", path.display()) +} + +#[cfg(test)] +#[path = "../../../tests/unit/dynamic_worker_tests.rs"] +mod tests; diff --git a/crates/core/tests/fixtures/worker_plugin/Cargo.lock b/crates/core/tests/fixtures/worker_plugin/Cargo.lock new file mode 100644 index 000000000..fe31ce567 --- /dev/null +++ b/crates/core/tests/fixtures/worker_plugin/Cargo.lock @@ -0,0 +1,1372 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "sync_wrapper", + "tower", + "tower-layer", + "tower-service", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", +] + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +dependencies = [ + "serde_core", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" + +[[package]] +name = "cc" +version = "1.2.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fixedbitset" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-macro", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + +[[package]] +name = "h2" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "h2", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "libc", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "mio" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + +[[package]] +name = "multimap" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" + +[[package]] +name = "nemo-relay-types" +version = "0.5.0" +dependencies = [ + "bitflags", + "chrono", + "serde", + "serde_json", + "typed-builder", + "uuid", +] + +[[package]] +name = "nemo-relay-worker" +version = "0.5.0" +dependencies = [ + "futures-util", + "hyper-util", + "nemo-relay-types", + "nemo-relay-worker-proto", + "serde", + "serde_json", + "thiserror", + "tokio", + "tokio-stream", + "tonic", + "tower", +] + +[[package]] +name = "nemo-relay-worker-plugin-fixture" +version = "0.0.0" +dependencies = [ + "nemo-relay-worker", + "serde_json", + "tokio", + "tokio-stream", +] + +[[package]] +name = "nemo-relay-worker-proto" +version = "0.5.0" +dependencies = [ + "prost", + "prost-build", + "protoc-bin-vendored", + "serde", + "serde_json", + "tonic", + "tonic-prost", + "tonic-prost-build", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "petgraph" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" +dependencies = [ + "fixedbitset", + "hashbrown 0.15.5", + "indexmap", +] + +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "prost" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-build" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" +dependencies = [ + "heck", + "itertools", + "log", + "multimap", + "petgraph", + "prettyplease", + "prost", + "prost-types", + "pulldown-cmark", + "pulldown-cmark-to-cmark", + "regex", + "syn", + "tempfile", +] + +[[package]] +name = "prost-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" +dependencies = [ + "anyhow", + "itertools", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + +[[package]] +name = "protoc-bin-vendored" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1c381df33c98266b5f08186583660090a4ffa0889e76c7e9a5e175f645a67fa" +dependencies = [ + "protoc-bin-vendored-linux-aarch_64", + "protoc-bin-vendored-linux-ppcle_64", + "protoc-bin-vendored-linux-s390_64", + "protoc-bin-vendored-linux-x86_32", + "protoc-bin-vendored-linux-x86_64", + "protoc-bin-vendored-macos-aarch_64", + "protoc-bin-vendored-macos-x86_64", + "protoc-bin-vendored-win32", +] + +[[package]] +name = "protoc-bin-vendored-linux-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c350df4d49b5b9e3ca79f7e646fde2377b199e13cfa87320308397e1f37e1a4c" + +[[package]] +name = "protoc-bin-vendored-linux-ppcle_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55a63e6c7244f19b5c6393f025017eb5d793fd5467823a099740a7a4222440c" + +[[package]] +name = "protoc-bin-vendored-linux-s390_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dba5565db4288e935d5330a07c264a4ee8e4a5b4a4e6f4e83fad824cc32f3b0" + +[[package]] +name = "protoc-bin-vendored-linux-x86_32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8854774b24ee28b7868cd71dccaae8e02a2365e67a4a87a6cd11ee6cdbdf9cf5" + +[[package]] +name = "protoc-bin-vendored-linux-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b38b07546580df720fa464ce124c4b03630a6fb83e05c336fea2a241df7e5d78" + +[[package]] +name = "protoc-bin-vendored-macos-aarch_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89278a9926ce312e51f1d999fee8825d324d603213344a9a706daa009f1d8092" + +[[package]] +name = "protoc-bin-vendored-macos-x86_64" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81745feda7ccfb9471d7a4de888f0652e806d5795b61480605d4943176299756" + +[[package]] +name = "protoc-bin-vendored-win32" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3" + +[[package]] +name = "pulldown-cmark" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f068eba8e7071c5f9511831b44f32c740d5adf574e990f946ddb53db2f314e" +dependencies = [ + "bitflags", + "memchr", + "unicase", +] + +[[package]] +name = "pulldown-cmark-to-cmark" +version = "22.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50793def1b900256624a709439404384204a5dc3a6ec580281bfaac35e882e90" +dependencies = [ + "pulldown-cmark", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "regex" +version = "1.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "async-trait", + "axum", + "base64", + "bytes", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "socket2", + "sync_wrapper", + "tokio", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-build" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322" +dependencies = [ + "prettyplease", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + +[[package]] +name = "tonic-prost-build" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27" +dependencies = [ + "prettyplease", + "proc-macro2", + "prost-build", + "prost-types", + "quote", + "syn", + "tempfile", + "tonic-build", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "indexmap", + "pin-project-lite", + "slab", + "sync_wrapper", + "tokio", + "tokio-util", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "typed-builder" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31aa81521b70f94402501d848ccc0ecaa8f93c8eb6999eb9747e72287757ffda" +dependencies = [ + "typed-builder-macro", +] + +[[package]] +name = "typed-builder-macro" +version = "0.23.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "076a02dc54dd46795c2e9c8282ed40bcfb1e22747e955de9389a1de28190fb26" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "unicase" +version = "2.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "uuid" +version = "1.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f87b8aa10b915a06587d0dec516c282ff295b475d94abf425d62b57710070a2" +dependencies = [ + "getrandom 0.3.4", + "js-sys", + "serde", + "wasm-bindgen", +] + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/crates/core/tests/fixtures/worker_plugin/Cargo.toml b/crates/core/tests/fixtures/worker_plugin/Cargo.toml new file mode 100644 index 000000000..4b9c4a908 --- /dev/null +++ b/crates/core/tests/fixtures/worker_plugin/Cargo.toml @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[workspace] + +[package] +name = "nemo-relay-worker-plugin-fixture" +version = "0.0.0" +edition = "2024" +license = "Apache-2.0" +publish = false + +[[bin]] +name = "nemo-relay-worker-plugin-fixture" +path = "src/main.rs" + +[dependencies] +nemo-relay-worker = { path = "../../../../worker" } +serde_json = "1" +tokio = { version = "1", features = ["macros", "rt-multi-thread"] } +tokio-stream = "0.1" diff --git a/crates/core/tests/fixtures/worker_plugin/src/main.rs b/crates/core/tests/fixtures/worker_plugin/src/main.rs new file mode 100644 index 000000000..c6e1e2509 --- /dev/null +++ b/crates/core/tests/fixtures/worker_plugin/src/main.rs @@ -0,0 +1,294 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use nemo_relay_worker::{ + JsonStream, LlmNext, LlmStreamNext, PluginContext, ScopeType, ToolNext, WorkerPlugin, + WorkerSdkError, serve_plugin, +}; +use nemo_relay_worker::{ConfigDiagnostic, DiagnosticLevel, Json, LlmRequest}; +use serde_json::json; + +struct FixtureWorkerPlugin; + +impl WorkerPlugin for FixtureWorkerPlugin { + fn plugin_id(&self) -> &str { + if std::env::var("FIXTURE_WORKER_PLUGIN_ID").as_deref() == Ok("other_worker") { + return "other_worker"; + } + "fixture_worker" + } + + fn validate(&self, config: &Json) -> Vec { + if config + .get("exit_in_validate") + .and_then(Json::as_bool) + .unwrap_or(false) + { + std::process::exit(42); + } + if config + .get("reject") + .and_then(Json::as_bool) + .unwrap_or(false) + { + return vec![ConfigDiagnostic { + level: DiagnosticLevel::Error, + code: "fixture.rejected".into(), + component: Some("fixture_worker".into()), + field: Some("reject".into()), + message: "fixture rejection requested".into(), + }]; + } + Vec::new() + } + + fn register(&self, ctx: &mut PluginContext, config: &Json) -> nemo_relay_worker::Result<()> { + let register_error = config + .get("register_error") + .and_then(Json::as_bool) + .unwrap_or(false); + let exit_in_register = config + .get("exit_in_register") + .and_then(Json::as_bool) + .unwrap_or(false); + if exit_in_register { + std::process::exit(43); + } + if register_error { + return Err(WorkerSdkError::Callback( + "fixture registration error requested".into(), + )); + } + + let empty_registration_name = config + .get("empty_registration_name") + .and_then(Json::as_bool) + .unwrap_or(false); + if empty_registration_name { + ctx.register_subscriber("", |_| {}); + return Ok(()); + } + + let block_tool = config + .get("block_tool") + .and_then(Json::as_bool) + .unwrap_or(false); + let tool_request_error = config + .get("tool_request_error") + .and_then(Json::as_bool) + .unwrap_or(false); + let llm_request_error = config + .get("llm_request_error") + .and_then(Json::as_bool) + .unwrap_or(false); + let llm_stream_open_error = config + .get("llm_stream_open_error") + .and_then(Json::as_bool) + .unwrap_or(false); + + let runtime = ctx + .runtime() + .ok_or_else(|| WorkerSdkError::Callback("runtime handle missing".into()))?; + + ctx.register_subscriber("fixture_subscriber", { + let runtime = runtime.clone(); + move |event| { + if event.name() == "worker-plugin-test-outer" { + let runtime = runtime.clone(); + let _ = tokio::task::block_in_place(|| { + tokio::runtime::Handle::current().block_on(async move { + runtime + .emit_mark( + "fixture.worker.subscriber.mark", + Some(json!("subscriber")), + None, + ) + .await + }) + }); + } + } + }); + + ctx.register_tool_sanitize_request_guardrail( + "fixture_tool_sanitize_request", + 0, + |_name, args| mark_json(args, "worker_plugin_tool_sanitize_request"), + ); + ctx.register_tool_sanitize_response_guardrail( + "fixture_tool_sanitize_response", + 0, + |_name, result| mark_json(result, "worker_plugin_tool_sanitize_response"), + ); + ctx.register_tool_conditional_execution_guardrail( + "fixture_tool_conditional", + 0, + move |_name, _args| { + if block_tool { + Ok(Some("fixture tool blocked".into())) + } else { + Ok(None) + } + }, + ); + ctx.register_tool_request_intercept("fixture_rewrite_args", 0, false, { + let runtime = runtime.clone(); + move |_name, args| { + if tool_request_error { + return Err(WorkerSdkError::Callback( + "fixture tool request error requested".into(), + )); + } + let runtime = runtime.clone(); + tokio::task::block_in_place(|| { + tokio::runtime::Handle::current().block_on(emit_runtime_events(runtime)) + })?; + Ok(mark_json(args, "worker_plugin")) + } + }); + ctx.register_tool_execution_intercept( + "fixture_tool_execution", + 0, + |_name, args, next: ToolNext| async move { + let result = next + .call(mark_json(args, "worker_plugin_tool_execution_request")) + .await?; + Ok(mark_json(result, "worker_plugin_tool_execution")) + }, + ); + + ctx.register_llm_sanitize_request_guardrail( + "fixture_llm_sanitize_request", + 0, + |request| mark_llm_request(request, "worker_plugin_llm_sanitize_request"), + ); + ctx.register_llm_sanitize_response_guardrail( + "fixture_llm_sanitize_response", + 0, + |response| mark_json(response, "worker_plugin_llm_sanitize_response"), + ); + ctx.register_llm_conditional_execution_guardrail( + "fixture_llm_conditional", + 0, + |_request| Ok(None), + ); + ctx.register_llm_request_intercept( + "fixture_llm_request_intercept", + 0, + false, + move |_name, request, annotated| { + if llm_request_error { + return Err(WorkerSdkError::Callback( + "fixture LLM request error requested".into(), + )); + } + let annotated = annotated.map(|mut annotated| { + annotated + .extra + .insert("worker_plugin_annotated_request".into(), json!(true)); + annotated + }); + Ok((mark_llm_request(request, "worker_plugin_llm_request_intercept"), annotated)) + }, + ); + ctx.register_llm_execution_intercept( + "fixture_llm_execution", + 0, + |_name, request, next: LlmNext| async move { + let response = next + .call(mark_llm_request( + request, + "worker_plugin_llm_execution_request", + )) + .await?; + Ok(mark_json(response, "worker_plugin_llm_execution")) + }, + ); + ctx.register_llm_stream_execution_intercept( + "fixture_llm_stream_execution", + 0, + move |_name, request, next: LlmStreamNext| async move { + if llm_stream_open_error { + return Err(WorkerSdkError::Callback( + "fixture LLM stream open error requested".into(), + )); + } + let stream = next + .call(mark_llm_request( + request, + "worker_plugin_llm_stream_execution_request", + )) + .await?; + let mapped: JsonStream = Box::pin(tokio_stream::StreamExt::map(stream, |chunk| { + chunk.map(|value| mark_json(value, "worker_plugin_llm_stream_execution")) + })); + Ok(mapped) + }, + ); + + Ok(()) + } +} + +async fn emit_runtime_events(runtime: nemo_relay_worker::PluginRuntime) -> nemo_relay_worker::Result<()> { + runtime + .emit_mark("fixture.worker.mark", Some(json!("current")), None) + .await?; + let scope = runtime + .push_scope( + None, + "fixture.worker.scope", + ScopeType::Custom, + None, + None, + Some(json!("current-scope-input")), + ) + .await?; + runtime + .pop_scope(&scope, Some(json!("current-scope-output")), None) + .await?; + + let isolated = runtime.create_scope_stack().await?; + let isolated_scope = runtime + .push_scope( + Some(&isolated), + "fixture.worker.isolated.scope", + ScopeType::Custom, + None, + None, + Some(json!("isolated-input")), + ) + .await?; + let isolated_runtime = runtime.clone(); + runtime + .with_scope_stack(&isolated, || async move { + isolated_runtime + .emit_mark("fixture.worker.isolated.mark", Some(json!("isolated")), None) + .await + }) + .await?; + runtime + .pop_scope(&isolated_scope, Some(json!("isolated-output")), None) + .await?; + runtime.drop_scope_stack(&isolated).await +} + +fn mark_llm_request(mut request: LlmRequest, key: &str) -> LlmRequest { + request.content = mark_json(request.content, key); + request +} + +fn mark_json(mut value: Json, key: &str) -> Json { + if let Json::Object(object) = &mut value { + object.insert(key.into(), json!(true)); + } + value +} + +#[tokio::main] +async fn main() { + if let Err(error) = serve_plugin(FixtureWorkerPlugin).await { + eprintln!("fixture worker failed: {error}"); + std::process::exit(1); + } +} diff --git a/crates/core/tests/integration/worker_plugin_tests.rs b/crates/core/tests/integration/worker_plugin_tests.rs new file mode 100644 index 000000000..6b0fa1735 --- /dev/null +++ b/crates/core/tests/integration/worker_plugin_tests.rs @@ -0,0 +1,1041 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Integration coverage for gRPC worker dynamic plugins. + +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::sync::{Arc, Mutex, OnceLock}; + +use futures::StreamExt; +use nemo_relay::api::event::{Event, ScopeCategory}; +use nemo_relay::api::llm::{ + LlmCallExecuteParams, LlmRequest, LlmStreamCallExecuteParams, llm_call_execute, + llm_stream_call_execute, +}; +use nemo_relay::api::runtime::{TASK_SCOPE_STACK, create_scope_stack}; +use nemo_relay::api::scope::{PopScopeParams, PushScopeParams, ScopeType, pop_scope, push_scope}; +use nemo_relay::api::subscriber::{flush_subscribers, register_subscriber}; +use nemo_relay::api::tool::{ToolCallExecuteParams, tool_call_execute, tool_request_intercepts}; +use nemo_relay::codec::request::AnnotatedLlmRequest; +use nemo_relay::codec::traits::LlmCodec; +use nemo_relay::error::Result as FlowResult; +use nemo_relay::plugin::dynamic::{ + WorkerPluginActivation, WorkerPluginLoadSpec, load_worker_plugins, +}; +use nemo_relay::plugin::{ + PluginComponentSpec, PluginConfig, clear_plugin_configuration, initialize_plugins_exact, +}; +use serde_json::{Map, Value as Json, json}; +use tempfile::TempDir; +use uuid::Uuid; + +static WORKER_PLUGIN_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +#[test] +fn worker_activation_with_no_specs_is_empty() { + let activation = load_worker_plugins(Vec::::new()) + .expect("empty worker activation should succeed"); + assert!(activation.is_empty()); + activation.clear(); +} + +#[tokio::test] +async fn rust_worker_registers_and_invokes_all_current_surfaces() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let loaded = load_and_initialize_fixture(Map::new()).await; + + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = events.clone(); + register_subscriber( + "worker_plugin_fixture_events", + Arc::new(move |event| { + captured.lock().unwrap().push(event.clone()); + }), + ) + .expect("test subscriber should register"); + + let stack = create_scope_stack(); + let (outer_uuid, rewritten, tool_result) = TASK_SCOPE_STACK + .scope(stack, async { + let outer = push_scope( + PushScopeParams::builder() + .name("worker-plugin-test-outer") + .scope_type(ScopeType::Agent) + .build(), + ) + .expect("outer scope should push"); + let outer_uuid = outer.uuid; + let rewritten = tool_request_intercepts("demo_tool", json!({ "input": "value" })) + .expect("worker request intercept should run"); + let tool_result = tool_call_execute( + ToolCallExecuteParams::builder() + .name("worker-fixture-tool") + .args(json!({ "input": "execute" })) + .func(Arc::new(|args| { + Box::pin(async move { Ok(json!({ "tool_callback": true, "args": args })) }) + })) + .build(), + ) + .await + .expect("worker tool middleware should run"); + pop_scope(PopScopeParams::builder().handle_uuid(&outer.uuid).build()) + .expect("outer scope should pop"); + (outer_uuid, rewritten, tool_result) + }) + .await; + + assert_eq!(rewritten["worker_plugin"], true); + assert_eq!(tool_result["tool_callback"], true); + assert_eq!(tool_result["worker_plugin_tool_execution"], true); + assert_eq!( + tool_result["args"]["worker_plugin_tool_execution_request"], + true + ); + + flush_subscribers().expect("worker fixture events should flush"); + let captured_events = events.lock().unwrap().clone(); + assert_parent( + &captured_events, + "fixture.worker.mark", + None, + Some(outer_uuid), + ); + assert_parent( + &captured_events, + "fixture.worker.scope", + Some(ScopeCategory::Start), + Some(outer_uuid), + ); + assert_not_parent( + &captured_events, + "fixture.worker.isolated.scope", + Some(ScopeCategory::Start), + outer_uuid, + ); + let isolated_scope = find_event( + &captured_events, + "fixture.worker.isolated.scope", + Some(ScopeCategory::Start), + ); + let isolated_mark = find_event(&captured_events, "fixture.worker.isolated.mark", None); + assert_eq!(isolated_mark.parent_uuid(), Some(isolated_scope.uuid())); + assert_ne!( + isolated_mark.parent_uuid(), + Some(outer_uuid), + "worker isolated mark should use the plugin-selected isolated stack" + ); + let tool_start = find_event( + &captured_events, + "worker-fixture-tool", + Some(ScopeCategory::Start), + ); + assert_eq!( + tool_start.input().unwrap()["worker_plugin_tool_sanitize_request"], + true + ); + let tool_end = find_event( + &captured_events, + "worker-fixture-tool", + Some(ScopeCategory::End), + ); + assert_eq!( + tool_end.output().unwrap()["worker_plugin_tool_sanitize_response"], + true + ); + + let llm_execute_response = llm_call_execute( + LlmCallExecuteParams::builder() + .name("worker-fixture-llm-execute") + .request(LlmRequest { + headers: Map::new(), + content: json!({ "prompt": "managed" }), + }) + .func(Arc::new(|request| { + Box::pin(async move { + Ok(json!({ + "id": "managed-response", + "request": request.content, + "llm_callback": true + })) + }) + })) + .build(), + ) + .await + .expect("worker LLM middleware should run"); + assert_eq!(llm_execute_response["llm_callback"], true); + assert_eq!(llm_execute_response["worker_plugin_llm_execution"], true); + assert_eq!( + llm_execute_response["request"]["worker_plugin_llm_execution_request"], + true + ); + flush_subscribers().expect("worker fixture LLM events should flush"); + let captured_events = events.lock().unwrap().clone(); + find_event(&captured_events, "fixture.worker.subscriber.mark", None); + let llm_start = find_event( + &captured_events, + "worker-fixture-llm-execute", + Some(ScopeCategory::Start), + ); + assert_eq!( + llm_start.input().unwrap()["content"]["worker_plugin_llm_sanitize_request"], + true + ); + let llm_end = find_event( + &captured_events, + "worker-fixture-llm-execute", + Some(ScopeCategory::End), + ); + assert_eq!( + llm_end.output().unwrap()["worker_plugin_llm_sanitize_response"], + true + ); + + let stream_values = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("worker-fixture-llm-stream") + .request(LlmRequest { + headers: Map::new(), + content: json!({ "prompt": "stream" }), + }) + .func(Arc::new(|request| { + Box::pin(async move { + let first = json!({ + "chunk": 1, + "request": request.content, + }); + Ok(Box::pin(tokio_stream::iter(vec![Ok(first)])) as _) + }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| json!({ "done": true }))) + .build(), + ) + .await + .expect("worker stream middleware should start") + .collect::>() + .await; + let stream_value = stream_values + .into_iter() + .next() + .expect("one stream chunk should be returned") + .expect("stream chunk should succeed"); + assert_eq!(stream_value["worker_plugin_llm_stream_execution"], true); + assert_eq!( + stream_value["request"]["worker_plugin_llm_stream_execution_request"], + true + ); + + loaded.clear(); +} + +#[tokio::test] +async fn worker_request_intercept_callback_error_surfaces_to_host() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let loaded = + load_and_initialize_fixture(Map::from_iter([("tool_request_error".into(), json!(true))])) + .await; + + let error = tool_request_intercepts("demo_tool", json!({ "input": "value" })) + .expect_err("worker callback error should surface"); + assert!( + error + .to_string() + .contains("fixture tool request error requested"), + "{error}" + ); + + loaded.clear(); +} + +#[tokio::test] +async fn worker_conditional_guardrail_blocks_tool_execution() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let loaded = + load_and_initialize_fixture(Map::from_iter([("block_tool".into(), json!(true))])).await; + + let error = tool_call_execute( + ToolCallExecuteParams::builder() + .name("worker-fixture-blocked-tool") + .args(json!({ "input": "blocked" })) + .func(Arc::new(|_| { + Box::pin(async move { Ok(json!({ "should_not_run": true })) }) + })) + .build(), + ) + .await + .expect_err("worker guardrail should block tool execution"); + assert!( + error.to_string().contains("fixture tool blocked"), + "{error}" + ); + + loaded.clear(); +} + +#[tokio::test] +async fn worker_llm_request_intercept_round_trips_annotations() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let loaded = load_and_initialize_fixture(Map::new()).await; + + let response = llm_call_execute( + LlmCallExecuteParams::builder() + .name("worker-fixture-llm-annotated") + .request(LlmRequest { + headers: Map::new(), + content: json!({ "prompt": "annotated" }), + }) + .codec(Arc::new(FixtureCodec)) + .func(Arc::new(|request| { + Box::pin(async move { + Ok(json!({ + "request": request.content, + "llm_callback": true + })) + }) + })) + .build(), + ) + .await + .expect("worker LLM request intercept should preserve annotations"); + assert_eq!(response["llm_callback"], true); + assert_eq!(response["request"]["worker_plugin_annotated_request"], true); + + loaded.clear(); +} + +#[tokio::test] +async fn worker_llm_request_intercept_callback_error_surfaces_to_host() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let loaded = + load_and_initialize_fixture(Map::from_iter([("llm_request_error".into(), json!(true))])) + .await; + + let error = llm_call_execute( + LlmCallExecuteParams::builder() + .name("worker-fixture-llm-error") + .request(LlmRequest { + headers: Map::new(), + content: json!({ "prompt": "error" }), + }) + .func(Arc::new(|request| { + Box::pin(async move { + Ok(json!({ + "request": request.content, + "should_not_complete": true + })) + }) + })) + .build(), + ) + .await + .expect_err("worker LLM request intercept error should surface"); + assert!( + error + .to_string() + .contains("fixture LLM request error requested"), + "{error}" + ); + + loaded.clear(); +} + +#[tokio::test] +async fn worker_llm_stream_open_error_surfaces_to_host() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let loaded = load_and_initialize_fixture(Map::from_iter([( + "llm_stream_open_error".into(), + json!(true), + )])) + .await; + + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("worker-fixture-llm-stream-error") + .request(LlmRequest { + headers: Map::new(), + content: json!({ "prompt": "stream-error" }), + }) + .func(Arc::new(|request| { + Box::pin(async move { + let chunk = json!({ "request": request.content }); + Ok(Box::pin(tokio_stream::iter(vec![Ok(chunk)])) as _) + }) + })) + .collector(Box::new(|_chunk| Ok(()))) + .finalizer(Box::new(|| json!({ "done": true }))) + .build(), + ) + .await + .expect("worker stream invoke should return a host stream"); + let error = stream + .next() + .await + .expect("stream should yield the worker error") + .expect_err("worker stream callback error should surface"); + assert!( + error + .to_string() + .contains("fixture LLM stream open error requested"), + "{error}" + ); + + loaded.clear(); +} + +#[tokio::test] +async fn worker_validation_diagnostics_prevent_initialization() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + let config = Map::from_iter([("reject".into(), json!(true))]); + + let activation = load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: config.clone(), + }]) + .expect("worker plugin should load with validation diagnostics"); + + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_worker".into(), + enabled: true, + config, + }); + let error = initialize_plugins_exact(plugin_config) + .await + .expect_err("validation diagnostics should prevent initialization") + .to_string(); + assert!(error.contains("fixture rejection requested"), "{error}"); + + clear_plugin_configuration().expect("worker plugin config should clear"); + activation.clear(); +} + +#[tokio::test] +async fn worker_duplicate_component_rejected_for_single_instance_plugin() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + + let activation = load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) + .expect("worker plugin should load"); + + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_worker".into(), + enabled: true, + config: Map::new(), + }); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_worker".into(), + enabled: true, + config: Map::new(), + }); + let error = initialize_plugins_exact(plugin_config) + .await + .expect_err("single-instance worker plugin should reject duplicate components") + .to_string(); + assert!(error.contains("may only appear once"), "{error}"); + + clear_plugin_configuration().expect("worker plugin config should clear"); + activation.clear(); +} + +#[tokio::test] +async fn worker_config_mismatch_prevents_initialization() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + + let activation = load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) + .expect("worker plugin should load"); + + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_worker".into(), + enabled: true, + config: Map::from_iter([("changed".into(), json!(true))]), + }); + let error = initialize_plugins_exact(plugin_config) + .await + .expect_err("config drift should prevent initialization") + .to_string(); + assert!(error.contains("config changed"), "{error}"); + + clear_plugin_configuration().expect("worker plugin config should clear"); + activation.clear(); +} + +#[tokio::test] +async fn worker_registration_error_fails_activation() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::from_iter([("register_error".into(), json!(true))]), + }]) { + Ok(activation) => { + activation.clear(); + panic!("worker registration error should fail activation"); + } + Err(error) => error.to_string(), + }; + assert!( + error.contains("fixture registration error requested"), + "{error}" + ); +} + +#[tokio::test] +async fn worker_invalid_registration_plan_fails_activation() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::from_iter([("empty_registration_name".into(), json!(true))]), + }]) { + Ok(activation) => { + activation.clear(); + panic!("empty registration name should fail activation"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("empty local_name"), "{error}"); +} + +#[tokio::test] +async fn worker_handshake_plugin_id_mismatch_reports_config_error() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let _env = EnvVarGuard::set("FIXTURE_WORKER_PLUGIN_ID", "other_worker"); + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("worker handshake id mismatch should fail activation"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("returned id 'other_worker'"), "{error}"); +} + +#[tokio::test] +async fn worker_validation_rpc_failure_reports_activation_error() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::from_iter([("exit_in_validate".into(), json!(true))]), + }]) { + Ok(activation) => { + activation.clear(); + panic!("worker validation process exit should fail activation"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("worker validation RPC failed"), "{error}"); +} + +#[tokio::test] +async fn worker_registration_rpc_failure_reports_activation_error() { + let _guard = WORKER_PLUGIN_TEST_LOCK.lock().await; + let fixture = build_fixture_worker(); + let (_manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::from_iter([("exit_in_register".into(), json!(true))]), + }]) { + Ok(activation) => { + activation.clear(); + panic!("worker registration process exit should fail activation"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("worker registration RPC failed"), "{error}"); +} + +#[test] +fn missing_worker_executable_reports_startup_error() { + let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); + let missing_binary = std::env::temp_dir().join(format!("missing-worker-{}", Uuid::now_v7())); + let (_manifest_dir, manifest_ref) = write_manifest(&missing_binary); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("missing worker executable should fail activation"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("failed to spawn"), "{error}"); +} + +#[test] +fn worker_manifest_id_mismatch_reports_config_error() { + let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); + let missing_binary = std::env::temp_dir().join(format!("unused-worker-{}", Uuid::now_v7())); + let (_manifest_dir, manifest_ref) = write_manifest(&missing_binary); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "different_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("manifest id mismatch should fail"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("does not match expected id"), "{error}"); +} + +#[test] +fn worker_manifest_kind_mismatch_reports_config_error() { + let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); + let relay = supported_relay_requirement(); + let (_manifest_dir, manifest_ref) = write_manifest_text(&format!( + r#" +manifest_version = 1 + +[plugin] +id = "fixture_worker" +kind = "rust_dynamic" + +[compat] +relay = {relay} +native_api = "1" + +[defaults] +enabled = false + +[capabilities] +items = ["plugin_native"] + +[load] +library = "missing" +symbol = "nemo_relay_plugin_entry" +"#, + relay = toml_string(&relay) + )); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("manifest kind mismatch should fail"); + } + Err(error) => error.to_string(), + }; + assert!( + error.contains("worker loader only supports worker"), + "{error}" + ); +} + +#[test] +fn unsupported_worker_relay_requirement_reports_compatibility_error() { + let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); + let missing_binary = std::env::temp_dir().join(format!("unused-worker-{}", Uuid::now_v7())); + let (_manifest_dir, manifest_ref) = + write_manifest_with_relay(&missing_binary, ">=9999.0,<10000.0"); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("unsupported relay requirement should fail"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("requires relay"), "{error}"); +} + +#[test] +fn invalid_worker_relay_requirement_reports_parse_error() { + let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); + let missing_binary = std::env::temp_dir().join(format!("unused-worker-{}", Uuid::now_v7())); + let (_manifest_dir, manifest_ref) = write_manifest_with_relay(&missing_binary, "not semver"); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("invalid relay requirement should fail"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("invalid compat.relay"), "{error}"); +} + +#[test] +fn command_worker_entrypoint_is_resolved_relative_to_manifest() { + let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); + let relay = supported_relay_requirement(); + let (manifest_dir, manifest_ref) = + write_worker_manifest("fixture_worker", &relay, "command", "missing-worker"); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("missing relative command worker should fail"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("failed to spawn command worker"), "{error}"); + assert_error_mentions_manifest_relative_entrypoint( + &error, + manifest_dir.path(), + "missing-worker", + ); +} + +#[test] +fn python_worker_uses_configured_interpreter() { + let _guard = WORKER_PLUGIN_TEST_LOCK.blocking_lock(); + let missing_python = std::env::temp_dir().join(format!("missing-python-{}", Uuid::now_v7())); + let _env = EnvVarGuard::set( + "NEMO_RELAY_PYTHON", + missing_python.to_string_lossy().as_ref(), + ); + let relay = supported_relay_requirement(); + let (_manifest_dir, manifest_ref) = write_worker_manifest( + "fixture_worker", + &relay, + "python", + "fixture_worker:create_plugin", + ); + + let error = match load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: Map::new(), + }]) { + Ok(activation) => { + activation.clear(); + panic!("missing configured Python interpreter should fail"); + } + Err(error) => error.to_string(), + }; + assert!(error.contains("failed to spawn python worker"), "{error}"); +} + +struct FixtureCodec; + +impl LlmCodec for FixtureCodec { + fn decode(&self, request: &LlmRequest) -> FlowResult { + Ok(AnnotatedLlmRequest { + messages: Vec::new(), + model: Some("fixture-model".into()), + params: None, + tools: None, + tool_choice: None, + store: None, + previous_response_id: None, + truncation: None, + reasoning: None, + include: None, + user: None, + metadata: None, + service_tier: None, + parallel_tool_calls: None, + max_output_tokens: None, + max_tool_calls: None, + top_logprobs: None, + stream: None, + extra: request.content.as_object().cloned().unwrap_or_default(), + }) + } + + fn encode( + &self, + annotated: &AnnotatedLlmRequest, + original: &LlmRequest, + ) -> FlowResult { + Ok(LlmRequest { + headers: original.headers.clone(), + content: Json::Object(annotated.extra.clone()), + }) + } +} + +struct LoadedWorker { + activation: Option, + _manifest_dir: TempDir, +} + +impl LoadedWorker { + fn clear(mut self) { + clear_plugin_configuration().expect("worker plugin config should clear"); + if let Some(activation) = self.activation.take() { + activation.clear(); + } + } +} + +impl Drop for LoadedWorker { + fn drop(&mut self) { + let _ = clear_plugin_configuration(); + if let Some(activation) = self.activation.take() { + activation.clear(); + } + } +} + +async fn load_and_initialize_fixture(config: Map) -> LoadedWorker { + let fixture = build_fixture_worker(); + let (manifest_dir, manifest_ref) = write_manifest(fixture.binary_path()); + + let activation = load_worker_plugins([WorkerPluginLoadSpec { + plugin_id: "fixture_worker".into(), + manifest_ref: manifest_ref.to_string_lossy().into_owned(), + config: config.clone(), + }]) + .expect("worker plugin should load"); + + let mut plugin_config = PluginConfig::default(); + plugin_config.components.push(PluginComponentSpec { + kind: "fixture_worker".into(), + enabled: true, + config, + }); + initialize_plugins_exact(plugin_config) + .await + .expect("worker plugin should initialize"); + + LoadedWorker { + activation: Some(activation), + _manifest_dir: manifest_dir, + } +} + +struct BuiltWorkerFixture { + binary_path: PathBuf, +} + +impl BuiltWorkerFixture { + fn binary_path(&self) -> &Path { + &self.binary_path + } +} + +fn build_fixture_worker() -> BuiltWorkerFixture { + static FIXTURE_BINARY: OnceLock = OnceLock::new(); + let binary_path = FIXTURE_BINARY.get_or_init(|| { + let fixture_dir = fixture_root(); + let target_root = + Path::new(env!("CARGO_MANIFEST_DIR")).join("../../target/worker-plugin-fixture"); + let target_dir = target_root.join("target"); + let manifest = fixture_dir.join("Cargo.toml"); + let status = Command::new("cargo") + .arg("build") + .arg("--quiet") + .arg("--locked") + .arg("--manifest-path") + .arg(&manifest) + .arg("--target-dir") + .arg(&target_dir) + .status() + .expect("fixture worker build should start"); + assert!(status.success(), "fixture worker build should succeed"); + let binary_path = target_dir.join("debug").join(format!( + "nemo-relay-worker-plugin-fixture{}", + std::env::consts::EXE_SUFFIX + )); + assert!(binary_path.exists(), "fixture worker binary should exist"); + binary_path + }); + BuiltWorkerFixture { + binary_path: binary_path.clone(), + } +} + +fn write_manifest(binary: &Path) -> (TempDir, PathBuf) { + let relay = supported_relay_requirement(); + write_manifest_with_relay(binary, &relay) +} + +fn write_manifest_with_relay(binary: &Path, relay: &str) -> (TempDir, PathBuf) { + write_worker_manifest( + "fixture_worker", + relay, + "rust", + binary.to_string_lossy().as_ref(), + ) +} + +fn write_worker_manifest( + plugin_id: &str, + relay: &str, + runtime: &str, + entrypoint: &str, +) -> (TempDir, PathBuf) { + write_manifest_text(&format!( + r#" +manifest_version = 1 + +[plugin] +id = {plugin_id} +kind = "worker" + +[compat] +relay = {relay} +worker_protocol = "grpc-v1" + +[defaults] +enabled = false + +[capabilities] +items = ["plugin_worker"] + +[load] +runtime = {runtime} +entrypoint = {entrypoint} +"#, + plugin_id = toml_string(plugin_id), + relay = toml_string(relay), + runtime = toml_string(runtime), + entrypoint = toml_string(entrypoint) + )) +} + +fn write_manifest_text(contents: &str) -> (TempDir, PathBuf) { + let temp = TempDir::new().expect("manifest tempdir should be created"); + let manifest = temp.path().join("relay-plugin.toml"); + std::fs::write(&manifest, contents).expect("manifest should be written"); + (temp, manifest) +} + +fn toml_string(value: &str) -> String { + format!("{value:?}") +} + +fn supported_relay_requirement() -> String { + format!("={}", env!("CARGO_PKG_VERSION")) +} + +fn assert_error_mentions_manifest_relative_entrypoint( + error: &str, + manifest_dir: &Path, + entrypoint: &str, +) { + let manifest_dir_name = manifest_dir + .file_name() + .expect("manifest dir should have a leaf name") + .to_string_lossy(); + let manifest_dir_pos = error.find(manifest_dir_name.as_ref()).unwrap_or_else(|| { + panic!("error did not mention manifest dir '{manifest_dir_name}': {error}") + }); + assert!( + error[manifest_dir_pos + manifest_dir_name.len()..].contains(entrypoint), + "error did not mention entrypoint '{entrypoint}' after manifest dir '{manifest_dir_name}': {error}" + ); +} + +struct EnvVarGuard { + key: &'static str, + previous: Option, +} + +impl EnvVarGuard { + fn set(key: &'static str, value: &str) -> Self { + let previous = std::env::var(key).ok(); + // SAFETY: this module serializes worker tests with WORKER_PLUGIN_TEST_LOCK. + unsafe { + std::env::set_var(key, value); + } + Self { key, previous } + } +} + +impl Drop for EnvVarGuard { + fn drop(&mut self) { + // SAFETY: this module serializes worker tests with WORKER_PLUGIN_TEST_LOCK. + unsafe { + if let Some(previous) = &self.previous { + std::env::set_var(self.key, previous); + } else { + std::env::remove_var(self.key); + } + } + } +} + +fn fixture_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/worker_plugin") +} + +fn find_event<'a>( + events: &'a [Event], + name: &str, + scope_category: Option, +) -> &'a Event { + events + .iter() + .find(|event| event.name() == name && event.scope_category() == scope_category) + .unwrap_or_else(|| panic!("event {name:?} with category {scope_category:?} not found")) +} + +fn assert_parent( + events: &[Event], + name: &str, + scope_category: Option, + expected_parent: Option, +) { + let event = find_event(events, name, scope_category); + assert_eq!(event.parent_uuid(), expected_parent); +} + +fn assert_not_parent( + events: &[Event], + name: &str, + scope_category: Option, + excluded_parent: Uuid, +) { + let event = find_event(events, name, scope_category); + assert_ne!(event.parent_uuid(), Some(excluded_parent)); +} diff --git a/crates/core/tests/unit/dynamic_worker_tests.rs b/crates/core/tests/unit/dynamic_worker_tests.rs new file mode 100644 index 000000000..aa77c0433 --- /dev/null +++ b/crates/core/tests/unit/dynamic_worker_tests.rs @@ -0,0 +1,1235 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::{Arc, Mutex}; + +use crate::api::event::{BaseEvent, MarkEvent}; +use nemo_relay_worker_proto::json_envelope; +use nemo_relay_worker_proto::v1::invoke_response::Result as InvokeResult; +use nemo_relay_worker_proto::v1::plugin_worker_server::{PluginWorker, PluginWorkerServer}; +use nemo_relay_worker_proto::v1::stream_chunk::Item as StreamItem; +use nemo_relay_worker_proto::v1::{ + CancelInvocationRequest, CreateScopeStackRequest, DropScopeStackRequest, EmitMarkRequest, + EmptyResult, GuardrailResult, HandshakeRequest, HandshakeResponse, HealthRequest, + HealthResponse, JsonEnvelope, JsonResult, LlmNextRequest, LlmRequestInterceptResult, + LlmStreamNextRequest, PopScopeRequest, PushScopeRequest, Registration, ScopeContext, + ScopeType as ProtoScopeType, ShutdownRequest, StreamChunk, ToolNextRequest, ValidateRequest, + ValidateResponse, WorkerAck, +}; +use serde_json::json; +use tokio_stream::StreamExt; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::Request; +use tonic::transport::Server; + +use super::*; + +const ACTIVATION_ID: &str = "activation-test"; +const AUTH_TOKEN: &str = "auth-test"; + +#[test] +fn response_helpers_cover_error_and_unexpected_shapes() { + let worker_error = WorkerError { + code: "worker.failed".into(), + message: "boom".into(), + retryable: false, + }; + + let error = json_from_invoke_response(InvokeResponse { + result: Some(InvokeResult::Json(JsonResult { + value: None, + error: Some(worker_error.clone()), + })), + }) + .expect_err("json result worker error should surface"); + assert!(error.to_string().contains("worker.failed: boom")); + + let error = json_from_invoke_response(InvokeResponse { + result: Some(InvokeResult::Error(worker_error.clone())), + }) + .expect_err("top-level worker error should surface"); + assert!(error.to_string().contains("worker.failed: boom")); + + let error = json_from_invoke_response(InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + }) + .expect_err("unexpected JSON result shape should fail"); + assert!(error.to_string().contains("unexpected invoke result")); + + let error = json_from_invoke_response(InvokeResponse { + result: Some(InvokeResult::Json(JsonResult { + value: Some(JsonEnvelope { + schema: JSON_SCHEMA.into(), + json: b"{".to_vec(), + }), + error: None, + })), + }) + .expect_err("invalid JSON envelope should fail"); + assert!(error.to_string().contains("invalid JSON result")); + + assert_eq!( + guardrail_from_invoke_response(InvokeResponse { + result: Some(InvokeResult::Guardrail(GuardrailResult { + block_reason: String::new(), + })), + }) + .expect("empty block reason is allowed"), + None + ); + assert_eq!( + guardrail_from_invoke_response(InvokeResponse { + result: Some(InvokeResult::Guardrail(GuardrailResult { + block_reason: "blocked".into(), + })), + }) + .expect("block reason should parse"), + Some("blocked".into()) + ); + assert!( + guardrail_from_invoke_response(InvokeResponse { + result: Some(InvokeResult::Error(worker_error.clone())), + }) + .expect_err("guardrail worker error should surface") + .to_string() + .contains("worker.failed") + ); + assert!( + guardrail_from_invoke_response(InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + }) + .expect_err("unexpected guardrail shape should fail") + .to_string() + .contains("guardrail returned unexpected") + ); + + assert!( + json_from_stream_chunk(StreamChunk { + item: Some(StreamItem::Error(worker_error.clone())), + }) + .expect_err("stream worker error should surface") + .to_string() + .contains("worker.failed") + ); + assert!( + json_from_stream_chunk(StreamChunk { + item: Some(StreamItem::Value(JsonEnvelope { + schema: JSON_SCHEMA.into(), + json: b"{".to_vec(), + })), + }) + .expect_err("invalid stream JSON envelope should fail") + .to_string() + .contains("invalid worker stream chunk") + ); + assert!( + json_from_stream_chunk(StreamChunk { item: None }) + .expect_err("empty stream chunk should fail") + .to_string() + .contains("stream chunk was empty") + ); +} + +#[test] +fn envelope_and_error_helpers_cover_failure_paths() { + assert!( + required_envelope(None, "required test") + .expect_err("missing envelope should fail") + .to_string() + .contains("required test is missing") + ); + assert!( + optional_envelope_to_json(Some(JsonEnvelope { + schema: JSON_SCHEMA.into(), + json: b"not-json".to_vec(), + })) + .expect_err("invalid optional envelope should fail") + .to_string() + .contains("invalid JSON envelope") + ); + + let ack = host_ack(Err(FlowError::Internal("host failed".into()))); + assert!(!ack.ok); + assert_eq!(ack.error.expect("host error").code, "host.runtime_error"); + + let result = json_result(Err(FlowError::Internal("json failed".into()))); + assert!(result.value.is_none()); + assert_eq!(result.error.expect("json error").code, "host.runtime_error"); + + let fallback = worker_error_to_plugin( + WorkerError { + code: "worker.empty".into(), + message: String::new(), + retryable: false, + }, + "fallback message", + ); + assert!(fallback.to_string().contains("fallback message")); + + let status = status_from_flow(FlowError::Internal("status failed".into())); + assert_eq!(status.code(), tonic::Code::Internal); + assert!(status.message().contains("status failed")); +} + +#[test] +fn registration_plan_and_scope_type_helpers_validate_edges() { + let empty_name = validate_registration_plan( + "fixture_worker", + &RegisterResponse { + registrations: vec![Registration { + local_name: " ".into(), + surface: RegistrationSurface::Subscriber as i32, + priority: 0, + break_chain: false, + }], + error: None, + }, + ) + .expect_err("empty registration names should fail"); + assert!(empty_name.to_string().contains("empty local_name")); + + let unsupported = validate_registration_plan( + "fixture_worker", + &RegisterResponse { + registrations: vec![Registration { + local_name: "bad".into(), + surface: 999, + priority: 0, + break_chain: false, + }], + error: None, + }, + ) + .expect_err("unsupported registration surfaces should fail"); + assert!( + unsupported + .to_string() + .contains("unsupported registration surface") + ); + + let unspecified = validate_registration_plan( + "fixture_worker", + &RegisterResponse { + registrations: vec![Registration { + local_name: "bad".into(), + surface: RegistrationSurface::Unspecified as i32, + priority: 0, + break_chain: false, + }], + error: None, + }, + ) + .expect_err("unspecified registration surfaces should fail"); + assert!( + unspecified + .to_string() + .contains("unspecified registration surface") + ); + + let cases = [ + (ProtoScopeType::Agent, crate::api::scope::ScopeType::Agent), + ( + ProtoScopeType::Function, + crate::api::scope::ScopeType::Function, + ), + (ProtoScopeType::Tool, crate::api::scope::ScopeType::Tool), + (ProtoScopeType::Llm, crate::api::scope::ScopeType::Llm), + ( + ProtoScopeType::Retriever, + crate::api::scope::ScopeType::Retriever, + ), + ( + ProtoScopeType::Embedder, + crate::api::scope::ScopeType::Embedder, + ), + ( + ProtoScopeType::Reranker, + crate::api::scope::ScopeType::Reranker, + ), + ( + ProtoScopeType::Guardrail, + crate::api::scope::ScopeType::Guardrail, + ), + ( + ProtoScopeType::Evaluator, + crate::api::scope::ScopeType::Evaluator, + ), + (ProtoScopeType::Custom, crate::api::scope::ScopeType::Custom), + ( + ProtoScopeType::Unknown, + crate::api::scope::ScopeType::Unknown, + ), + ]; + for (proto, expected) in cases { + assert_eq!(proto_scope_type(proto as i32), expected); + } + assert_eq!(proto_scope_type(999), crate::api::scope::ScopeType::Custom); +} + +#[test] +fn relay_compatibility_and_blocking_helpers_cover_local_edges() { + assert!( + validate_relay_compatibility(None) + .expect_err("missing relay compatibility should fail") + .to_string() + .contains("compat.relay is required") + ); + assert!( + validate_relay_compatibility(Some("not semver")) + .expect_err("invalid relay compatibility should fail") + .to_string() + .contains("invalid compat.relay") + ); + + let runtime = RuntimeBuilder::new_current_thread() + .enable_all() + .build() + .expect("runtime should build"); + assert_eq!(block_on_runtime(&runtime, async { 42 }), 42); +} + +#[test] +#[cfg(unix)] +fn worker_endpoints_fail_when_host_socket_cannot_bind() { + let activation_dir = std::env::temp_dir().join(format!("nmrw-unit-{}", Uuid::now_v7())); + let host_socket = activation_dir.join("host.sock"); + std::fs::create_dir_all(&host_socket).expect("host socket directory should be created"); + + let error = match WorkerEndpoints::new(&activation_dir) { + Ok(_) => panic!("endpoint creation should fail when host socket path is a directory"), + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("failed to bind worker host runtime socket") + ); + + let _ = std::fs::remove_dir_all(&activation_dir); +} + +#[tokio::test(flavor = "multi_thread")] +async fn callback_helpers_cover_worker_response_edges() { + let worker_error = WorkerError { + code: "worker.failed".into(), + message: "boom".into(), + retryable: false, + }; + let (callback, _shutdown) = fake_callback_service({ + let worker_error = worker_error.clone(); + move |request| match request.registration_name.as_str() { + "subscriber_error" => InvokeResponse { + result: Some(InvokeResult::Error(worker_error.clone())), + }, + "subscriber_unexpected" | "llm_intercept_unexpected" => InvokeResponse { + result: Some(InvokeResult::Json(JsonResult { + value: Some(json_envelope(JSON_SCHEMA, &json!({})).expect("json envelope")), + error: None, + })), + }, + "llm_json_invalid" => InvokeResponse { + result: Some(InvokeResult::Json(JsonResult { + value: Some(json_envelope(JSON_SCHEMA, &json!(null)).expect("json envelope")), + error: None, + })), + }, + "llm_intercept_invalid_request" => InvokeResponse { + result: Some(InvokeResult::LlmRequest(LlmRequestInterceptResult { + request: Some(JsonEnvelope { + schema: LLM_REQUEST_SCHEMA.into(), + json: b"null".to_vec(), + }), + annotated_request: None, + has_annotated_request: false, + })), + }, + "llm_intercept_missing_annotated" => InvokeResponse { + result: Some(InvokeResult::LlmRequest(LlmRequestInterceptResult { + request: Some(valid_llm_request_envelope()), + annotated_request: None, + has_annotated_request: true, + })), + }, + "llm_intercept_invalid_annotated" => InvokeResponse { + result: Some(InvokeResult::LlmRequest(LlmRequestInterceptResult { + request: Some(valid_llm_request_envelope()), + annotated_request: Some(JsonEnvelope { + schema: ANNOTATED_LLM_REQUEST_SCHEMA.into(), + json: b"null".to_vec(), + }), + has_annotated_request: true, + })), + }, + "llm_intercept_error" => InvokeResponse { + result: Some(InvokeResult::Error(worker_error.clone())), + }, + _ => InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + }, + } + }) + .await; + let event = Event::Mark(MarkEvent::new( + BaseEvent::builder().name("callback-edge").build(), + None, + None, + )); + + let error = callback + .invoke_subscriber("subscriber_error", &event) + .expect_err("subscriber worker error should surface"); + assert!(error.to_string().contains("worker.failed: boom")); + + let error = callback + .invoke_subscriber("subscriber_unexpected", &event) + .expect_err("unexpected subscriber result should fail"); + assert!(error.to_string().contains("subscriber returned unexpected")); + + let error = callback + .invoke_llm_request_json( + "llm_json_invalid", + RegistrationSurface::LlmSanitizeRequestGuardrail, + "model", + valid_llm_request(), + None, + None, + ) + .expect_err("invalid LLM JSON result should fail"); + assert!(error.to_string().contains("invalid type")); + + let error = callback + .invoke_llm_request_intercept( + "llm_intercept_invalid_request", + "model", + valid_llm_request(), + None, + ) + .expect_err("invalid LLM intercept request should fail"); + assert!(error.to_string().contains("invalid LLM request")); + + let error = callback + .invoke_llm_request_intercept( + "llm_intercept_missing_annotated", + "model", + valid_llm_request(), + None, + ) + .expect_err("missing annotated request should fail when flagged present"); + assert!( + error + .to_string() + .contains("llm request intercept annotated request is missing") + ); + + let error = callback + .invoke_llm_request_intercept( + "llm_intercept_invalid_annotated", + "model", + valid_llm_request(), + None, + ) + .expect_err("invalid annotated request should fail"); + assert!(error.to_string().contains("invalid annotated LLM request")); + + let error = callback + .invoke_llm_request_intercept("llm_intercept_error", "model", valid_llm_request(), None) + .expect_err("LLM intercept worker error should surface"); + assert!(error.to_string().contains("worker.failed: boom")); + + let error = callback + .invoke_llm_request_intercept( + "llm_intercept_unexpected", + "model", + valid_llm_request(), + None, + ) + .expect_err("unexpected LLM intercept result should fail"); + assert!( + error + .to_string() + .contains("LLM request intercept returned unexpected") + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn callback_stream_transport_error_surfaces_to_host_stream() { + let (callback, _shutdown) = fake_callback_service(|_| InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + }) + .await; + + let mut stream = callback + .invoke_llm_stream_execution( + "stream_transport_error", + "model", + valid_llm_request(), + Arc::new(|_request| { + Box::pin(async { Ok(Box::pin(tokio_stream::empty()) as LlmJsonStream) }) + }), + ) + .await + .expect("host stream should be returned"); + + let error = stream + .next() + .await + .expect("transport error should be yielded") + .expect_err("stream transport error should surface"); + assert!(error.to_string().contains("worker stream transport failed")); +} + +#[tokio::test(flavor = "multi_thread")] +async fn callback_stream_stops_when_host_receiver_is_dropped() { + let (yield_tx, yield_rx) = oneshot::channel(); + let (stream_dropped_tx, stream_dropped_rx) = oneshot::channel(); + let stream_dropped_tx = Arc::new(Mutex::new(Some(stream_dropped_tx))); + let yield_rx = Arc::new(Mutex::new(Some(yield_rx))); + let (callback, _shutdown) = fake_callback_service_with_stream( + |_| InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + }, + { + let stream_dropped_tx = stream_dropped_tx.clone(); + let yield_rx = yield_rx.clone(); + move |_| { + let dropped = stream_dropped_tx + .lock() + .expect("stream drop signal lock should not be poisoned") + .take() + .expect("test stream should be created once"); + let yield_rx = yield_rx + .lock() + .expect("stream yield signal lock should not be poisoned") + .take() + .expect("test stream should be created once"); + Box::pin(SignalChunkThenPendingStream { + yield_rx, + dropped: Some(dropped), + yielded: false, + }) as FakeInvokeStream + } + }, + ) + .await; + + let stream = callback + .invoke_llm_stream_execution( + "stream_receiver_drop", + "model", + valid_llm_request(), + Arc::new(|_request| { + Box::pin(async { Ok(Box::pin(tokio_stream::empty()) as LlmJsonStream) }) + }), + ) + .await + .expect("host stream should be returned"); + drop(stream); + yield_tx + .send(()) + .expect("worker stream yield signal should be delivered"); + tokio::time::timeout(std::time::Duration::from_secs(1), stream_dropped_rx) + .await + .expect("worker stream should be dropped after host receiver is dropped") + .expect("worker stream drop signal should be delivered"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn install_registrations_covers_registry_error_edges() { + for surface in [ + RegistrationSurface::Subscriber, + RegistrationSurface::ToolSanitizeRequestGuardrail, + RegistrationSurface::ToolSanitizeResponseGuardrail, + RegistrationSurface::ToolConditionalExecutionGuardrail, + RegistrationSurface::ToolRequestIntercept, + RegistrationSurface::ToolExecutionIntercept, + RegistrationSurface::LlmSanitizeRequestGuardrail, + RegistrationSurface::LlmSanitizeResponseGuardrail, + RegistrationSurface::LlmConditionalExecutionGuardrail, + RegistrationSurface::LlmRequestIntercept, + RegistrationSurface::LlmExecutionIntercept, + RegistrationSurface::LlmStreamExecutionIntercept, + ] { + let (instance, _shutdown) = fake_worker_instance(vec![ + registration(surface, "duplicate"), + registration(surface, "duplicate"), + ]) + .await; + let mut ctx = PluginRegistrationContext::new(); + let error = instance + .install_registrations(&mut ctx) + .expect_err("duplicate worker registration should fail"); + assert!( + error.to_string().contains("duplicate") + || error.to_string().contains("already registered"), + "{surface:?}: {error}" + ); + let mut registrations = ctx.into_registrations(); + crate::plugin::rollback_registrations(&mut registrations); + } + + let (instance, _shutdown) = fake_worker_instance(vec![Registration { + surface: 999, + ..registration(RegistrationSurface::Subscriber, "bad") + }]) + .await; + let mut ctx = PluginRegistrationContext::new(); + assert!( + instance + .install_registrations(&mut ctx) + .expect_err("unsupported registration surface should fail") + .to_string() + .contains("unsupported registration surface") + ); + + let (instance, _shutdown) = + fake_worker_instance(vec![registration(RegistrationSurface::Unspecified, "bad")]).await; + let mut ctx = PluginRegistrationContext::new(); + assert!( + instance + .install_registrations(&mut ctx) + .expect_err("unspecified registration surface should fail") + .to_string() + .contains("unspecified registration surface") + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn adapter_register_rejects_config_drift_even_without_validation_call() { + let (instance, _shutdown) = fake_worker_instance(Vec::new()).await; + let adapter = WorkerPluginAdapter { + plugin_kind: "fixture_worker".into(), + allows_multiple_components: false, + instance: Arc::new(instance), + }; + let mut ctx = PluginRegistrationContext::new(); + let changed = serde_json::Map::from_iter([("changed".into(), json!(true))]); + + let error = adapter + .register(&changed, &mut ctx) + .await + .expect_err("config drift should fail registration"); + assert!(error.to_string().contains("config changed"), "{error}"); +} + +#[tokio::test] +async fn host_runtime_service_covers_auth_scope_and_ack_errors() { + let state = Arc::new(WorkerHostRuntimeState::new( + ACTIVATION_ID.into(), + AUTH_TOKEN.into(), + )); + let service = WorkerHostRuntimeService { + state: state.clone(), + }; + + let auth_error = service + .emit_mark(Request::new(EmitMarkRequest { + activation_id: "wrong".into(), + auth_token: AUTH_TOKEN.into(), + name: "auth-failure".into(), + scope: None, + data: None, + metadata: None, + })) + .await + .expect_err("bad activation id should fail auth"); + assert_eq!(auth_error.code(), tonic::Code::PermissionDenied); + + let ack = service + .emit_mark(Request::new(EmitMarkRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + name: "missing-stack".into(), + scope: Some(ScopeContext { + scope_stack_id: "missing-stack".into(), + parent_scope_id: String::new(), + }), + data: None, + metadata: None, + })) + .await + .expect("missing stack should return host ack") + .into_inner(); + assert!(!ack.ok); + assert!( + ack.error + .expect("missing stack error") + .message + .contains("not found") + ); + + let ack = service + .emit_mark(Request::new(EmitMarkRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + name: "no-scope".into(), + scope: None, + data: None, + metadata: None, + })) + .await + .expect("no-scope mark should succeed") + .into_inner(); + assert!(ack.ok); + + let push = service + .push_scope(Request::new(PushScopeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + scope: None, + name: "invalid-json-scope".into(), + scope_type: ProtoScopeType::Custom as i32, + data: Some(JsonEnvelope { + schema: JSON_SCHEMA.into(), + json: b"not-json".to_vec(), + }), + metadata: None, + input: None, + })) + .await + .expect("invalid JSON should be structured") + .into_inner(); + assert!( + push.error + .expect("push error") + .message + .contains("invalid JSON") + ); + + let pop_error = service + .pop_scope(Request::new(PopScopeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + scope_handle_id: "missing-scope".into(), + output: None, + metadata: None, + })) + .await + .expect_err("missing scope handle should fail"); + assert_eq!(pop_error.code(), tonic::Code::NotFound); + + let created = service + .create_scope_stack(Request::new(CreateScopeStackRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + })) + .await + .expect("scope stack should be created") + .into_inner(); + let scope_stack_id = created.scope_stack_id.clone(); + assert!( + state + .stack("") + .expect("empty stack id should be valid") + .is_none() + ); + let dropped = service + .drop_scope_stack(Request::new(DropScopeStackRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + scope_stack_id: scope_stack_id.clone(), + })) + .await + .expect("scope stack should be dropped") + .into_inner(); + assert!(dropped.ok); + assert_eq!( + state + .stack(&scope_stack_id) + .expect_err("dropped stack should be removed") + .code(), + tonic::Code::NotFound + ); + + assert_eq!( + service + .with_stack( + Some(&ScopeContext { + scope_stack_id: String::new(), + parent_scope_id: String::new(), + }), + || Ok(7), + ) + .expect("empty explicit stack id should run without binding"), + 7 + ); +} + +#[tokio::test] +async fn host_runtime_service_reports_poisoned_internal_locks() { + let state = Arc::new(WorkerHostRuntimeState::new( + ACTIVATION_ID.into(), + AUTH_TOKEN.into(), + )); + poison_mutex({ + let state = state.clone(); + move || { + let _guard = state.scope_handles.lock().expect("scope handles lock"); + panic!("poison scope handles"); + } + }); + let service = WorkerHostRuntimeService { + state: state.clone(), + }; + let push_error = service + .push_scope(Request::new(PushScopeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + scope: None, + name: "poisoned".into(), + scope_type: ProtoScopeType::Custom as i32, + data: None, + metadata: None, + input: None, + })) + .await + .expect_err("poisoned scope handle lock should fail"); + assert_eq!(push_error.code(), tonic::Code::Internal); + + let pop_error = service + .pop_scope(Request::new(PopScopeRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + scope_handle_id: "missing".into(), + output: None, + metadata: None, + })) + .await + .expect_err("poisoned scope handle lock should fail"); + assert_eq!(pop_error.code(), tonic::Code::Internal); + + let state = Arc::new(WorkerHostRuntimeState::new( + ACTIVATION_ID.into(), + AUTH_TOKEN.into(), + )); + poison_mutex({ + let state = state.clone(); + move || { + let _guard = state.scope_stacks.lock().expect("scope stacks lock"); + panic!("poison scope stacks"); + } + }); + let service = WorkerHostRuntimeService { state }; + let create_error = service + .create_scope_stack(Request::new(CreateScopeStackRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + })) + .await + .expect_err("poisoned scope stack lock should fail"); + assert_eq!(create_error.code(), tonic::Code::Internal); + + let drop_error = service + .drop_scope_stack(Request::new(DropScopeStackRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + scope_stack_id: "stack".into(), + })) + .await + .expect_err("poisoned scope stack lock should fail"); + assert_eq!(drop_error.code(), tonic::Code::Internal); +} + +#[test] +fn owned_worker_runtime_drop_is_idempotent_when_runtime_already_taken() { + drop(OwnedWorkerRuntime { runtime: None }); +} + +#[tokio::test] +async fn host_runtime_service_covers_continuation_errors_and_stream_items() { + let state = Arc::new(WorkerHostRuntimeState::new( + ACTIVATION_ID.into(), + AUTH_TOKEN.into(), + )); + let service = WorkerHostRuntimeService { + state: state.clone(), + }; + + let llm_continuation = state + .insert_continuation(Continuation::Llm(Arc::new(|request| { + Box::pin(async move { Ok(request.content) }) + }))) + .expect("llm continuation should insert"); + let wrong_type = service + .tool_next(Request::new(ToolNextRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + continuation_id: llm_continuation, + value: Some(json_envelope(JSON_SCHEMA, &json!({})).expect("json envelope")), + })) + .await + .expect_err("wrong continuation type should fail"); + assert_eq!(wrong_type.code(), tonic::Code::InvalidArgument); + + let tool_continuation = state + .insert_continuation(Continuation::Tool(Arc::new(|value| { + Box::pin(async move { Ok(value) }) + }))) + .expect("tool continuation should insert"); + let invalid_tool_json = service + .tool_next(Request::new(ToolNextRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + continuation_id: tool_continuation, + value: Some(JsonEnvelope { + schema: JSON_SCHEMA.into(), + json: b"not-json".to_vec(), + }), + })) + .await + .expect_err("invalid tool next JSON should fail"); + assert_eq!(invalid_tool_json.code(), tonic::Code::InvalidArgument); + + let llm_continuation = state + .insert_continuation(Continuation::Llm(Arc::new(|request| { + Box::pin(async move { Ok(request.content) }) + }))) + .expect("llm continuation should insert"); + let invalid_llm_json = service + .llm_next(Request::new(LlmNextRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + continuation_id: llm_continuation, + request: Some(JsonEnvelope { + schema: LLM_REQUEST_SCHEMA.into(), + json: b"not-json".to_vec(), + }), + })) + .await + .expect_err("invalid LLM next request should fail"); + assert_eq!(invalid_llm_json.code(), tonic::Code::InvalidArgument); + + let stream_continuation = state + .insert_continuation(Continuation::LlmStream(Arc::new(|_request| { + Box::pin(async move { + Ok(Box::pin(tokio_stream::iter(vec![Err(FlowError::Internal( + "stream item failed".into(), + ))])) as LlmJsonStream) + }) + }))) + .expect("stream continuation should insert"); + let stream_response = service + .llm_stream_next(Request::new(LlmStreamNextRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + continuation_id: stream_continuation, + request: Some( + json_envelope( + LLM_REQUEST_SCHEMA, + &LlmRequest { + headers: serde_json::Map::new(), + content: json!({ "prompt": "stream" }), + }, + ) + .expect("llm request envelope"), + ), + })) + .await + .expect("stream next should return stream"); + let mut stream = stream_response.into_inner(); + let chunk = stream + .next() + .await + .expect("stream should yield one item") + .expect("transport should be ok"); + match chunk.item { + Some(StreamItem::Error(error)) => { + assert!(error.message.contains("stream item failed")); + } + other => panic!("expected worker stream error, got {other:?}"), + } + + let stream_continuation = state + .insert_continuation(Continuation::LlmStream(Arc::new(|_request| { + Box::pin(async move { Ok(Box::pin(tokio_stream::empty()) as LlmJsonStream) }) + }))) + .expect("stream continuation should insert"); + let invalid_stream_request = match service + .llm_stream_next(Request::new(LlmStreamNextRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + continuation_id: stream_continuation, + request: Some(JsonEnvelope { + schema: LLM_REQUEST_SCHEMA.into(), + json: b"not-json".to_vec(), + }), + })) + .await + { + Ok(_) => panic!("invalid LLM stream request should fail"), + Err(error) => error, + }; + assert_eq!(invalid_stream_request.code(), tonic::Code::InvalidArgument); +} + +fn valid_llm_request() -> LlmRequest { + LlmRequest { + headers: serde_json::Map::new(), + content: json!({ "prompt": "unit" }), + } +} + +fn valid_llm_request_envelope() -> JsonEnvelope { + json_envelope(LLM_REQUEST_SCHEMA, &valid_llm_request()).expect("llm request envelope") +} + +async fn fake_callback_service( + invoke: impl Fn(InvokeRequest) -> InvokeResponse + Send + Sync + 'static, +) -> (WorkerPluginCallback, oneshot::Sender<()>) { + let (client, shutdown_tx) = fake_worker_client(invoke).await; + callback_for_client(client, shutdown_tx) +} + +async fn fake_callback_service_with_stream( + invoke: impl Fn(InvokeRequest) -> InvokeResponse + Send + Sync + 'static, + invoke_stream: impl Fn(InvokeRequest) -> FakeInvokeStream + Send + Sync + 'static, +) -> (WorkerPluginCallback, oneshot::Sender<()>) { + let (client, shutdown_tx) = fake_worker_client_with_stream(invoke, invoke_stream).await; + callback_for_client(client, shutdown_tx) +} + +fn callback_for_client( + client: PluginWorkerClient, + shutdown_tx: oneshot::Sender<()>, +) -> (WorkerPluginCallback, oneshot::Sender<()>) { + let state = Arc::new(WorkerHostRuntimeState::new( + ACTIVATION_ID.into(), + AUTH_TOKEN.into(), + )); + ( + WorkerPluginCallback { + activation_id: ACTIVATION_ID.into(), + runtime: tokio::runtime::Handle::current(), + client, + host_state: state, + }, + shutdown_tx, + ) +} + +async fn fake_worker_instance( + registrations: Vec, +) -> (WorkerPluginInstance, oneshot::Sender<()>) { + let (client, shutdown_tx) = fake_worker_client(|_| InvokeResponse { + result: Some(InvokeResult::Empty(EmptyResult {})), + }) + .await; + let activation_dir = std::env::temp_dir().join(format!("nmrw-unit-{}", Uuid::now_v7())); + std::fs::create_dir_all(&activation_dir).expect("unit activation dir should be created"); + ( + WorkerPluginInstance { + plugin_kind: "fixture_worker".into(), + allows_multiple_components: false, + config: serde_json::Map::new(), + validation_diagnostics: Vec::new(), + registrations, + runtime: OwnedWorkerRuntime::new( + RuntimeBuilder::new_multi_thread() + .enable_all() + .build() + .expect("worker runtime should build"), + ), + client, + host_state: Arc::new(WorkerHostRuntimeState::new( + ACTIVATION_ID.into(), + AUTH_TOKEN.into(), + )), + shutdown: Mutex::new(None), + process: Mutex::new(None), + activation_dir, + }, + shutdown_tx, + ) +} + +async fn fake_worker_client( + invoke: impl Fn(InvokeRequest) -> InvokeResponse + Send + Sync + 'static, +) -> (PluginWorkerClient, oneshot::Sender<()>) { + fake_worker_client_with_stream(invoke, |_| { + Box::pin(tokio_stream::iter(vec![Err(Status::unavailable( + "stream transport down", + ))])) as FakeInvokeStream + }) + .await +} + +async fn fake_worker_client_with_stream( + invoke: impl Fn(InvokeRequest) -> InvokeResponse + Send + Sync + 'static, + invoke_stream: impl Fn(InvokeRequest) -> FakeInvokeStream + Send + Sync + 'static, +) -> (PluginWorkerClient, oneshot::Sender<()>) { + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .expect("fake worker listener should bind"); + let addr = listener + .local_addr() + .expect("fake worker listener address should be available"); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + tokio::spawn( + Server::builder() + .add_service(PluginWorkerServer::new(FakePluginWorker { + invoke: Arc::new(invoke), + invoke_stream: Arc::new(invoke_stream), + })) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }), + ); + let client = PluginWorkerClient::connect(format!("http://{addr}")) + .await + .expect("fake worker client should connect"); + (client, shutdown_tx) +} + +fn registration(surface: RegistrationSurface, local_name: &str) -> Registration { + Registration { + local_name: local_name.into(), + surface: surface as i32, + priority: 0, + break_chain: false, + } +} + +fn poison_mutex(f: impl FnOnce() + std::panic::UnwindSafe) { + let _ = std::panic::catch_unwind(f); +} + +struct FakePluginWorker { + invoke: Arc InvokeResponse + Send + Sync>, + invoke_stream: Arc FakeInvokeStream + Send + Sync>, +} + +type FakeInvokeStream = + Pin> + Send>>; + +struct SignalChunkThenPendingStream { + yield_rx: oneshot::Receiver<()>, + dropped: Option>, + yielded: bool, +} + +impl tokio_stream::Stream for SignalChunkThenPendingStream { + type Item = std::result::Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + if self.yielded { + return std::task::Poll::Pending; + } + match Pin::new(&mut self.yield_rx).poll(cx) { + std::task::Poll::Ready(_) => { + self.yielded = true; + std::task::Poll::Ready(Some(Ok(StreamChunk { + item: Some(StreamItem::Value( + json_envelope(JSON_SCHEMA, &json!({ "after_receiver_drop": true })) + .expect("test stream chunk should encode"), + )), + }))) + } + std::task::Poll::Pending => std::task::Poll::Pending, + } + } +} + +impl Drop for SignalChunkThenPendingStream { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } +} + +#[tonic::async_trait] +impl PluginWorker for FakePluginWorker { + async fn handshake( + &self, + _request: Request, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::Response::new(HandshakeResponse { + plugin_id: "fixture_worker".into(), + plugin_kind: "fixture_worker".into(), + allows_multiple_components: false, + worker_protocol: WORKER_PROTOCOL_GRPC_V1.into(), + sdk_name: "unit".into(), + sdk_version: "0".into(), + runtime_name: "unit".into(), + runtime_version: "0".into(), + supported_surfaces: Vec::new(), + })) + } + + async fn health( + &self, + _request: Request, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::Response::new(HealthResponse { + ok: true, + message: String::new(), + plugin_id: "fixture_worker".into(), + worker_protocol: WORKER_PROTOCOL_GRPC_V1.into(), + sdk_name: "unit".into(), + sdk_version: "0".into(), + runtime_name: "unit".into(), + runtime_version: "0".into(), + })) + } + + async fn validate( + &self, + _request: Request, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::Response::new(ValidateResponse { + diagnostics: None, + error: None, + })) + } + + async fn register( + &self, + _request: Request, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::Response::new(RegisterResponse { + registrations: Vec::new(), + error: None, + })) + } + + async fn invoke( + &self, + request: Request, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::Response::new((self.invoke)(request.into_inner()))) + } + + type InvokeStreamStream = + Pin> + Send>>; + + async fn invoke_stream( + &self, + request: Request, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::Response::new((self.invoke_stream)( + request.into_inner(), + ))) + } + + async fn cancel_invocation( + &self, + _request: Request, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::Response::new(WorkerAck { + accepted: false, + message: "not implemented".into(), + })) + } + + async fn shutdown( + &self, + _request: Request, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::Response::new(WorkerAck { + accepted: false, + message: "not implemented".into(), + })) + } +} diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 357f42657..4bb394c0d 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -11,7 +11,6 @@ use std::future::Future; use std::net::{SocketAddr, ToSocketAddrs}; #[cfg(unix)] use std::os::unix::fs::FileTypeExt; -#[cfg(unix)] use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::{Arc, Mutex}; @@ -38,9 +37,11 @@ use nemo_relay_worker_proto::v1::{ WorkerError, }; use nemo_relay_worker_proto::{WORKER_PROTOCOL_GRPC_V1, decode_json_envelope, json_envelope}; +use tokio::net::TcpListener; #[cfg(unix)] use tokio::net::{UnixListener, UnixStream}; use tokio::sync::OnceCell; +use tokio_stream::wrappers::TcpListenerStream; #[cfg(unix)] use tokio_stream::wrappers::UnixListenerStream; use tonic::transport::{Channel, Endpoint, Server}; @@ -501,6 +502,24 @@ impl PluginRuntime { ack_to_result(response.ok, response.error) } + /// Runs an async operation with runtime calls bound to a specific host-owned scope stack. + /// + /// This is useful for isolated stacks created with [`Self::create_scope_stack`]. The previous + /// worker invocation scope is restored after the future completes. + pub async fn with_scope_stack(&self, scope_stack_id: &str, f: F) -> T + where + F: FnOnce() -> Fut, + Fut: Future, + { + let scope = Some(scope_context(scope_stack_id)); + TASK_SCOPE_CONTEXT + .scope(scope.clone(), async move { + let future = with_thread_scope(&scope, f); + future.await + }) + .await + } + /// Pushes a scope through the host runtime. pub async fn push_scope( &self, @@ -683,7 +702,12 @@ pub async fn serve_plugin_arc(plugin: Arc) -> Result<()> { activation_id: required_env("NEMO_RELAY_WORKER_ID")?, auth_token: required_env("NEMO_RELAY_WORKER_TOKEN")?, }; - serve_plugin_arc_with_config(plugin, config).await + serve_plugin_arc_with_endpoint_file( + plugin, + config, + optional_env("NEMO_RELAY_WORKER_ENDPOINT_FILE").map(PathBuf::from), + ) + .await } /// Serves a shared worker plugin using explicit endpoint and authentication configuration. @@ -697,6 +721,14 @@ pub async fn serve_plugin_arc(plugin: Arc) -> Result<()> { pub async fn serve_plugin_arc_with_config( plugin: Arc, config: WorkerServerConfig, +) -> Result<()> { + serve_plugin_arc_with_endpoint_file(plugin, config, None).await +} + +async fn serve_plugin_arc_with_endpoint_file( + plugin: Arc, + config: WorkerServerConfig, + endpoint_file: Option, ) -> Result<()> { let runtime = PluginRuntime { activation_id: config.activation_id, @@ -709,11 +741,15 @@ pub async fn serve_plugin_arc_with_config( runtime, handlers: Arc::new(Mutex::new(WorkerHandlers::default())), }; - serve_worker_service(service, &config.worker_endpoint).await + serve_worker_service(service, &config.worker_endpoint, endpoint_file.as_deref()).await } #[cfg(unix)] -async fn serve_worker_service(service: WorkerService, endpoint: &str) -> Result<()> { +async fn serve_worker_service( + service: WorkerService, + endpoint: &str, + endpoint_file: Option<&Path>, +) -> Result<()> { if endpoint.starts_with("unix://") { let path = parse_unix_endpoint(endpoint)?; remove_stale_socket(&path)?; @@ -726,24 +762,41 @@ async fn serve_worker_service(service: WorkerService, endpoint: &str) -> Result< .await .map_err(|err| WorkerSdkError::Transport(err.to_string())); } - serve_tcp_worker_service(service, endpoint).await + serve_tcp_worker_service(service, endpoint, endpoint_file).await } #[cfg(not(unix))] -async fn serve_worker_service(service: WorkerService, endpoint: &str) -> Result<()> { +async fn serve_worker_service( + service: WorkerService, + endpoint: &str, + endpoint_file: Option<&Path>, +) -> Result<()> { if endpoint.starts_with("unix://") { return Err(WorkerSdkError::InvalidInput( "unix endpoints are not supported on this platform".into(), )); } - serve_tcp_worker_service(service, endpoint).await + serve_tcp_worker_service(service, endpoint, endpoint_file).await } -async fn serve_tcp_worker_service(service: WorkerService, endpoint: &str) -> Result<()> { +async fn serve_tcp_worker_service( + service: WorkerService, + endpoint: &str, + endpoint_file: Option<&Path>, +) -> Result<()> { let addr = parse_tcp_endpoint(endpoint)?; + let listener = TcpListener::bind(addr) + .await + .map_err(|err| WorkerSdkError::Transport(format!("failed to bind worker socket: {err}")))?; + if let Some(path) = endpoint_file { + let local_addr = listener.local_addr().map_err(|err| { + WorkerSdkError::Transport(format!("failed to inspect worker socket: {err}")) + })?; + write_endpoint_file(path, &format!("http://{local_addr}"))?; + } Server::builder() .add_service(PluginWorkerServer::new(service)) - .serve(addr) + .serve_with_incoming(TcpListenerStream::new(listener)) .await .map_err(|err| WorkerSdkError::Transport(err.to_string())) } @@ -1605,6 +1658,29 @@ fn required_env(name: &str) -> Result { }) } +fn optional_env(name: &str) -> Option { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) +} + +fn write_endpoint_file(path: &Path, endpoint: &str) -> Result<()> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent).map_err(|err| { + WorkerSdkError::Transport(format!( + "failed to create worker endpoint file directory '{}': {err}", + parent.display() + )) + })?; + } + std::fs::write(path, endpoint).map_err(|err| { + WorkerSdkError::Transport(format!( + "failed to write worker endpoint file '{}': {err}", + path.display() + )) + }) +} + fn rustc_version_runtime() -> String { option_env!("RUSTC_VERSION") .unwrap_or("unknown") diff --git a/crates/worker/tests/worker_sdk_tests.rs b/crates/worker/tests/worker_sdk_tests.rs index 6941b71fb..30ae616a4 100644 --- a/crates/worker/tests/worker_sdk_tests.rs +++ b/crates/worker/tests/worker_sdk_tests.rs @@ -5,14 +5,13 @@ use std::future::Future; use std::net::{SocketAddr, TcpListener}; +use std::path::Path; #[cfg(unix)] -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; -use std::time::Duration; -#[cfg(unix)] -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use futures_util::{Stream, StreamExt}; #[cfg(unix)] @@ -453,6 +452,8 @@ async fn worker_service_invokes_every_registration_surface() { let calls = host.calls(); assert!(calls.contains(&"mark:tool-exec:stack-1:parent-1".into())); assert!(calls.contains(&"create_scope_stack".into())); + assert!(calls.contains(&"mark:tool-exec-isolated:isolated-stack:".into())); + assert!(calls.contains(&"mark:tool-exec-restored:stack-1:parent-1".into())); assert!(calls.contains(&"push:worker-scope:stack-1:parent-1".into())); assert!(calls.contains(&"pop:scope-handle-1".into())); assert!(calls.contains(&"drop:isolated-stack".into())); @@ -689,6 +690,46 @@ async fn worker_service_validates_env_and_endpoints() { } } +#[tokio::test(flavor = "multi_thread")] +async fn worker_service_announces_ephemeral_tcp_endpoint_file() { + const ENVS: &[&str] = &[ + "NEMO_RELAY_WORKER_SOCKET", + "NEMO_RELAY_HOST_SOCKET", + "NEMO_RELAY_WORKER_ID", + "NEMO_RELAY_WORKER_TOKEN", + "NEMO_RELAY_WORKER_ENDPOINT_FILE", + ]; + let _env_guard = ENV_LOCK.lock().await; + let snapshot = EnvSnapshot::capture(ENVS); + let endpoint_file = unique_temp_file("nrw-endpoint"); + let _ = std::fs::remove_file(&endpoint_file); + + set_required_envs(); + set_env("NEMO_RELAY_WORKER_SOCKET", "tcp://127.0.0.1:0"); + set_env( + "NEMO_RELAY_WORKER_ENDPOINT_FILE", + endpoint_file.to_str().expect("endpoint path utf-8"), + ); + let handle = tokio::spawn(serve_plugin_arc(Arc::new(MinimalPlugin))); + let endpoint = wait_for_endpoint_file(&endpoint_file).await; + assert!(endpoint.starts_with("http://127.0.0.1:")); + + let mut client = connect_worker(&endpoint).await; + let health = client + .health(Request::new(HealthRequest { + activation_id: ACTIVATION_ID.into(), + auth_token: AUTH_TOKEN.into(), + })) + .await + .expect("announced endpoint should accept connections") + .into_inner(); + assert!(health.ok); + + handle.abort(); + let _ = std::fs::remove_file(endpoint_file); + snapshot.restore(); +} + #[cfg(unix)] #[tokio::test(flavor = "multi_thread")] async fn worker_service_supports_unix_socket_worker_and_host_endpoints() { @@ -1202,6 +1243,15 @@ impl WorkerPlugin for SurfacePlugin { async move { runtime.emit_mark("tool-exec", None, None).await?; let stack_id = runtime.create_scope_stack().await?; + let isolated_runtime = runtime.clone(); + runtime + .with_scope_stack(&stack_id, || async move { + isolated_runtime + .emit_mark("tool-exec-isolated", None, None) + .await + }) + .await?; + runtime.emit_mark("tool-exec-restored", None, None).await?; let handle = runtime .push_scope(None, "worker-scope", ScopeType::Function, None, None, None) .await?; @@ -1685,6 +1735,16 @@ async fn wait_for_port(endpoint: &str) { panic!("server did not start at {endpoint}"); } +async fn wait_for_endpoint_file(path: &Path) -> String { + for _ in 0..50 { + match std::fs::read_to_string(path) { + Ok(endpoint) if !endpoint.trim().is_empty() => return endpoint, + Ok(_) | Err(_) => std::thread::sleep(Duration::from_millis(20)), + } + } + panic!("endpoint file was not written at {}", path.display()); +} + #[cfg(unix)] async fn wait_for_unix_socket(path: &Path) { for _ in 0..50 { @@ -2072,6 +2132,14 @@ fn remove_env(name: &str) { } } +fn unique_temp_file(prefix: &str) -> std::path::PathBuf { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system time after epoch") + .as_nanos(); + std::env::temp_dir().join(format!("{prefix}-{}-{nanos}", std::process::id())) +} + #[cfg(unix)] fn unique_temp_path(prefix: &str) -> std::path::PathBuf { let nanos = SystemTime::now() diff --git a/examples/python-grpc-worker-plugin/.gitignore b/examples/python-grpc-worker-plugin/.gitignore new file mode 100644 index 000000000..7be5af0fa --- /dev/null +++ b/examples/python-grpc-worker-plugin/.gitignore @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +.venv/ +__pycache__/ +*.py[cod] +*.egg-info/ +nemo/ diff --git a/examples/python-grpc-worker-plugin/README.md b/examples/python-grpc-worker-plugin/README.md new file mode 100644 index 000000000..9268ddb11 --- /dev/null +++ b/examples/python-grpc-worker-plugin/README.md @@ -0,0 +1,50 @@ + + +# Python gRPC Worker Plugin + +This example shows a Python worker plugin using the `nemo-relay-plugin` SDK. It +registers a tool request intercept, emits a mark event through the host runtime, +and returns a mutated JSON tool request. + +## Set Up + +From this directory: + +```bash +python3 -m venv .venv +. .venv/bin/activate +python -m pip install -e ../../python/plugin -e . +``` + +The SDK package owns the generated protobuf stubs and gRPC server setup. + +## Register With Relay + +Point the CLI at this manifest and enable it: + +```bash +nemo-relay plugins add ./relay-plugin.toml +nemo-relay plugins enable examples.python_grpc_worker +``` + +When launching the gateway, point Relay at the Python interpreter that has +`grpcio` installed: + +```bash +NEMO_RELAY_PYTHON="$PWD/.venv/bin/python" nemo-relay gateway +``` + +You can also reference the manifest manually from `plugins.toml`: + +```toml +[[plugins.dynamic]] +manifest = "./examples/python-grpc-worker-plugin/relay-plugin.toml" +config = { tag = "demo" } +``` + +The worker process is started by Relay. Do not run `worker.py` directly unless +you also provide the worker socket, host socket, activation ID, plugin ID, and +activation token environment variables that Relay normally supplies. diff --git a/examples/python-grpc-worker-plugin/pyproject.toml b/examples/python-grpc-worker-plugin/pyproject.toml new file mode 100644 index 000000000..fd36a433a --- /dev/null +++ b/examples/python-grpc-worker-plugin/pyproject.toml @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "nemo-relay-python-grpc-worker-example" +version = "0.1.0" +description = "Example Python gRPC worker plugin for NeMo Relay" +requires-python = ">=3.11" +dependencies = [ + "nemo-relay-plugin>=0.5.0", +] + +[tool.setuptools] +py-modules = ["worker"] diff --git a/examples/python-grpc-worker-plugin/relay-plugin.toml b/examples/python-grpc-worker-plugin/relay-plugin.toml new file mode 100644 index 000000000..d260c4efc --- /dev/null +++ b/examples/python-grpc-worker-plugin/relay-plugin.toml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +manifest_version = 1 + +[plugin] +id = "examples.python_grpc_worker" +kind = "worker" + +[compat] +relay = ">=0.5,<1.0" +worker_protocol = "grpc-v1" + +[defaults] +enabled = false + +[capabilities] +items = ["plugin_worker"] + +[load] +runtime = "python" +entrypoint = "worker:main" diff --git a/examples/python-grpc-worker-plugin/worker.py b/examples/python-grpc-worker-plugin/worker.py new file mode 100644 index 000000000..fa9ed6310 --- /dev/null +++ b/examples/python-grpc-worker-plugin/worker.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example Python worker plugin using the nemo-relay-plugin SDK.""" + +from __future__ import annotations + +from nemo_relay_plugin import ConfigDiagnostic, DiagnosticLevel, Json, PluginContext, WorkerPlugin, serve_plugin + + +class ExamplePythonWorker(WorkerPlugin): + """Small worker plugin that tags tool request JSON and emits a host mark.""" + + plugin_id = "examples.python_grpc_worker" + + def validate(self, config: Json) -> list[ConfigDiagnostic]: + if isinstance(config, dict) and config.get("reject") is True: + return [ + ConfigDiagnostic( + level=DiagnosticLevel.ERROR, + code="examples.python_grpc_worker.rejected", + component=self.plugin_id, + field="reject", + message="Python gRPC worker rejection requested", + ) + ] + if isinstance(config, dict) and "tag" in config and not isinstance(config["tag"], str): + return [ + ConfigDiagnostic( + level=DiagnosticLevel.ERROR, + code="examples.python_grpc_worker.invalid_tag", + component=self.plugin_id, + field="tag", + message="tag must be a string", + ) + ] + return [] + + def register(self, ctx: PluginContext, config: Json) -> None: + tag = config.get("tag", "python_grpc_worker") if isinstance(config, dict) else "python_grpc_worker" + + async def tag_tool_request(tool_name: str, args: Json) -> Json: + await ctx.runtime.emit_mark( + "examples.python_grpc_worker.tool_request", + {"tool_name": tool_name, "source": "python-grpc-worker", "tag": tag}, + ) + return _tag_json(args, tag) + + ctx.register_tool_request_intercept("tag_tool_request", tag_tool_request) + + +def _tag_json(value: Json, tag: str) -> Json: + if not isinstance(value, dict): + raise TypeError("configured tool request tagging requires a JSON object") + if tag in value: + raise ValueError(f"tool request already contains configured tag {tag!r}") + return {**value, tag: True} + + +async def main() -> None: + """Entrypoint referenced by relay-plugin.toml.""" + await serve_plugin(ExamplePythonWorker()) + + +if __name__ == "__main__": + import asyncio + + asyncio.run(main()) diff --git a/justfile b/justfile index e8aa2d700..6c8718a22 100644 --- a/justfile +++ b/justfile @@ -64,16 +64,21 @@ uv_python_executable() { activate_project_venv() { # Ensure PATH-based tool lookups (for example, `zig`) resolve from the # synced project environment without asking uv to resync the project. - if [[ -f "$NEMO_RELAY_REPO_ROOT/.venv/bin/activate" ]]; then - # shellcheck disable=SC1091 - source "$NEMO_RELAY_REPO_ROOT/.venv/bin/activate" - elif [[ -f "$NEMO_RELAY_REPO_ROOT/.venv/Scripts/activate" ]]; then - # shellcheck disable=SC1091 - source "$NEMO_RELAY_REPO_ROOT/.venv/Scripts/activate" + local venv_dir="" + local venv_bin="" + if [[ -x "$NEMO_RELAY_REPO_ROOT/.venv/bin/python" ]]; then + venv_dir="$NEMO_RELAY_REPO_ROOT/.venv" + venv_bin="$venv_dir/bin" + elif [[ -x "$NEMO_RELAY_REPO_ROOT/.venv/Scripts/python.exe" ]]; then + venv_dir="$NEMO_RELAY_REPO_ROOT/.venv" + venv_bin="$venv_dir/Scripts" else - echo "ERROR: expected project virtualenv activation script under .venv" >&2 + echo "ERROR: expected project virtualenv Python executable under .venv" >&2 exit 1 fi + export VIRTUAL_ENV="$venv_dir" + export PATH="$venv_bin:$PATH" + unset PYTHONHOME } project_python_executable() { @@ -564,6 +569,7 @@ set_project_version() { local version="$1" set_cargo_workspace_version "$version" set_node_package_versions "$version" + set_python_plugin_package_version "$version" set_coding_agent_plugin_versions "$version" } @@ -653,6 +659,62 @@ print(f"crates/python/Cargo.toml version updated to {cargo_version}") PY } +set_python_plugin_package_version() { + local version="$1" + local python_executable="" + python_executable="$(uv_python_executable)" + + "$python_executable" - "$version" <<'PY' +from pathlib import Path +import re +import sys + +def semver_to_pep440(version: str) -> str: + pattern = re.compile( + r"^(?P\d+\.\d+\.\d+)" + r"(?:-(?Palpha|beta|rc)(?:\.(?P\d+))?)?" + r"(?:\+(?P[0-9A-Za-z.-]+))?$" + ) + match = pattern.fullmatch(version) + if not match: + raise SystemExit( + "Unsupported Python package version format. Expected SemVer with optional " + "alpha/beta/rc prerelease and optional build metadata." + ) + + pep440 = match.group("release") + pre_label = match.group("pre_label") + if pre_label: + pre_map = {"alpha": "a", "beta": "b", "rc": "rc"} + pre_num = match.group("pre_num") or "0" + pep440 += f"{pre_map[pre_label]}{pre_num}" + + local = match.group("local") + if local: + normalized_local = ".".join(part.lower() for part in re.split(r"[._-]+", local) if part) + if not normalized_local: + raise SystemExit("Python package local version metadata cannot be empty") + pep440 += f"+{normalized_local}" + + return pep440 + +version = semver_to_pep440(sys.argv[1]) +path = Path("python/plugin/pyproject.toml") +text = path.read_text() +updated, count = re.subn( + r'^version = "(.*)"$', + f'version = "{version}"', + text, + count=1, + flags=re.MULTILINE, +) +if count != 1: + raise SystemExit("Failed to update version in python/plugin/pyproject.toml") +path.write_text(updated) +print(f"python/plugin/pyproject.toml version updated to {version}") +PY +} + published_cargo_packages() { printf '%s\n' \ nemo-relay-types \ @@ -693,6 +755,71 @@ python_wheel_build_args() { esac } +python_plugin_grpc_dependencies_supported() { + local python_executable="$1" + "$python_executable" - <<'PY' +import platform +import sys + +is_windows_arm64 = sys.platform == "win32" and platform.machine().lower() in {"arm64", "aarch64"} +raise SystemExit(1 if is_windows_arm64 else 0) +PY +} + +configure_python_plugin_test_environment() { + local python_executable="$1" + if python_plugin_grpc_dependencies_supported "$python_executable"; then + unset NEMO_RELAY_SKIP_PYTHON_PLUGIN_TESTS + else + echo "Skipping nemo-relay-plugin grpc dependencies on Windows ARM64; plugin SDK tests will be skipped" + export NEMO_RELAY_SKIP_PYTHON_PLUGIN_TESTS=1 + fi +} + +python_plugin_sync_args() { + local python_executable="$1" + if ! python_plugin_grpc_dependencies_supported "$python_executable"; then + printf '%s\0' --no-install-package grpcio + fi +} + +generate_python_worker_proto_files() { + local output_dir="$1" + local python_executable="" + local proto_root="crates/worker-proto/proto/nemo/relay/worker/v1" + local proto_file="$proto_root/plugin_worker.proto" + + python_executable="$(uv_python_executable)" + mkdir -p "$output_dir" + uvx --from grpcio-tools==1.81.1 python -m grpc_tools.protoc \ + -I "$proto_root" \ + --python_out="$output_dir" \ + --grpc_python_out="$output_dir" \ + "$proto_file" + "$python_executable" - "$output_dir" <<'PY' +from pathlib import Path +import sys + +spdx = ( + "# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n" + "# SPDX-License-Identifier: Apache-2.0\n\n" +) +output_dir = Path(sys.argv[1]) + +for filename in ("plugin_worker_pb2.py", "plugin_worker_pb2_grpc.py"): + path = output_dir / filename + text = path.read_text(encoding="utf-8") + text = text.replace("# NO CHECKED-IN PROTOBUF GENCODE\n", "") + text = text.replace( + "import plugin_worker_pb2 as plugin__worker__pb2", + "from . import plugin_worker_pb2 as plugin__worker__pb2", + ) + if not text.startswith("# SPDX-FileCopyrightText:"): + text = spdx + text + path.write_text(text, encoding="utf-8") +PY +} + prepend_go_bin_to_path() { local go_bin go_bin="$(go env GOBIN)" @@ -789,6 +916,44 @@ build-python: python_executable="$(project_python_executable)" "$python_executable" -m maturin develop +build-python-plugin: + #!/usr/bin/env bash + {{ bash_helpers }} + cd "$NEMO_RELAY_REPO_ROOT" + python_executable="$(uv_python_executable)" + sync_args=(--inexact --package nemo-relay-plugin) + while IFS= read -r -d '' arg; do + sync_args+=("$arg") + done < <(python_plugin_sync_args "$python_executable") + uv sync "${sync_args[@]}" + activate_project_venv + python_executable="$(project_python_executable)" + configure_python_plugin_test_environment "$python_executable" + if python_plugin_grpc_dependencies_supported "$python_executable"; then + "$python_executable" -c 'import nemo_relay_plugin; print(f"nemo-relay-plugin import ok: {nemo_relay_plugin.__name__}")' + else + "$python_executable" -c 'import importlib.metadata; print(importlib.metadata.version("nemo-relay-plugin"))' + fi + +generate-python-worker-proto: + #!/usr/bin/env bash + {{ bash_helpers }} + cd "$NEMO_RELAY_REPO_ROOT" + generate_python_worker_proto_files python/plugin/src/nemo_relay_plugin/_proto + +check-python-worker-proto: + #!/usr/bin/env bash + {{ bash_helpers }} + cd "$NEMO_RELAY_REPO_ROOT" + tmp_dir="$(mktemp -d)" + cleanup_proto_tmp() { + rm -rf "$tmp_dir" + } + trap cleanup_proto_tmp EXIT + generate_python_worker_proto_files "$tmp_dir" + diff -u python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2.py "$tmp_dir/plugin_worker_pb2.py" + diff -u python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2_grpc.py "$tmp_dir/plugin_worker_pb2_grpc.py" + # --set [ci=true|false] build-go: @@ -828,7 +993,7 @@ build-wasm: NEMO_RELAY_WASM_RELEASE=1 npm run build:pkg --workspace=nemo-relay-wasm fi -build-all: build-rust build-python build-go build-node build-wasm +build-all: build-rust build-python build-python-plugin build-go build-node build-wasm # remove local build and test artifacts clean: @@ -861,9 +1026,21 @@ clean: python/nemo_relay/*.so \ python/nemo_relay/__pycache__ \ python/nemo_relay/_native*.pyd \ + python/plugin/build \ + python/plugin/dist \ + python/plugin/src/nemo_relay_plugin/__pycache__ \ + python/plugin/src/nemo_relay_plugin/_proto/__pycache__ \ + python/plugin/src/nemo_relay_plugin.egg-info \ + python/tests/__pycache__ \ + python/tests/plugin/__pycache__ \ + examples/python-grpc-worker-plugin/.pytest_cache \ + examples/python-grpc-worker-plugin/.venv \ + examples/python-grpc-worker-plugin/build \ + examples/python-grpc-worker-plugin/dist \ + examples/python-grpc-worker-plugin/__pycache__ \ + examples/python-grpc-worker-plugin/*.egg-info \ examples/rust-native-plugin/Cargo.lock \ examples/rust-native-plugin/target \ - python/tests/__pycache__ \ target # --set [output_dir=] [ci=true|false] @@ -937,7 +1114,7 @@ test-python: if is_true "{{ ci }}"; then coverage_out="$(prepare_artifact python-coverage.xml)" junit_out="$(prepare_artifact python-junit.xml)" - pytest_cmd+=(--cov=nemo_relay --cov-report term-missing --cov-report "xml:$coverage_out") + pytest_cmd+=(--cov=nemo_relay --cov=nemo_relay_plugin --cov-report term-missing --cov-report "xml:$coverage_out") pytest_cmd+=(--junit-xml "$junit_out") export_uv_python_runtime if rust_source_coverage_supported; then @@ -946,9 +1123,16 @@ test-python: fi cargo test -p nemo-relay-python --lib fi - uv sync --inexact --no-install-project --no-install-package nemo-relay + python_executable="$(uv_python_executable)" + sync_args=(--inexact --all-packages --no-install-project --no-install-package nemo-relay) + while IFS= read -r -d '' arg; do + sync_args+=("$arg") + done < <(python_plugin_sync_args "$python_executable") + uv sync "${sync_args[@]}" activate_project_venv + export_uv_python_runtime python_executable="$(project_python_executable)" + configure_python_plugin_test_environment "$python_executable" use_project_python_source "$python_executable" "$python_executable" -m maturin develop --skip-install "$python_executable" -m "${pytest_cmd[@]}" --ignore=python/tests/integrations @@ -960,6 +1144,28 @@ test-python: --output-path "$rust_coverage_out" fi +test-python-plugin: + #!/usr/bin/env bash + {{ bash_helpers }} + cd "$NEMO_RELAY_REPO_ROOT" + python_executable="$(uv_python_executable)" + sync_args=(--inexact --all-packages --no-install-project --no-install-package nemo-relay) + while IFS= read -r -d '' arg; do + sync_args+=("$arg") + done < <(python_plugin_sync_args "$python_executable") + uv sync "${sync_args[@]}" + activate_project_venv + python_executable="$(project_python_executable)" + configure_python_plugin_test_environment "$python_executable" + if [[ "${NEMO_RELAY_SKIP_PYTHON_PLUGIN_TESTS:-}" == "1" ]]; then + exit 0 + fi + "$python_executable" -m pytest \ + python/tests/plugin \ + --cov=nemo_relay_plugin \ + --cov-report term-missing \ + --cov-fail-under=95 + test-python-langchain: #!/usr/bin/env bash {{ bash_helpers }} @@ -967,6 +1173,7 @@ test-python-langchain: cd "$NEMO_RELAY_REPO_ROOT" uv sync --inexact --no-install-project --no-install-package nemo-relay --extra langchain --extra langgraph --extra deepagents activate_project_venv + export_uv_python_runtime python_executable="$(project_python_executable)" use_project_python_source "$python_executable" "$python_executable" -m maturin develop --skip-install @@ -1318,6 +1525,37 @@ package-python: exit 1 fi +# --set [output_dir=] [ref_name=] +package-python-plugin: + #!/usr/bin/env bash + {{ bash_helpers }} + output_dir="{{ output_dir }}" + cd "$NEMO_RELAY_REPO_ROOT" + package_dir="$(prepare_package_dir plugin-wheels)" + python_executable="$(uv_python_executable)" + sync_args=(--inexact --package nemo-relay-plugin) + while IFS= read -r -d '' arg; do + sync_args+=("$arg") + done < <(python_plugin_sync_args "$python_executable") + uv sync "${sync_args[@]}" + activate_project_venv + if [[ -z "{{ ref_name }}" ]]; then + sha="$(head_git_sha)" + version="$(read_workspace_version)" + echo "Non-release build: appending commit hash to version" + set_python_plugin_package_version "${version}+${sha}" + else + echo "Using explicit version {{ ref_name }}" + set_python_plugin_package_version "{{ ref_name }}" + fi + uv build --wheel --package nemo-relay-plugin --out-dir "$package_dir" + shopt -s nullglob + wheels=("$package_dir"/*.whl) + if ((${#wheels[@]} == 0)); then + echo "Error: No Python plugin wheels found in $package_dir" + exit 1 + fi + # --set [output_dir=] [ref_name=] package-wasm: #!/usr/bin/env bash diff --git a/pyproject.toml b/pyproject.toml index eb8afe83e..b06b69160 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,8 @@ dev = [ ] test = [ + "grpcio>=1.81.1,<2; sys_platform != 'win32' or (platform_machine != 'ARM64' and platform_machine != 'arm64' and platform_machine != 'aarch64')", + "nemo-relay-plugin; sys_platform != 'win32' or (platform_machine != 'ARM64' and platform_machine != 'arm64' and platform_machine != 'aarch64')", "opentelemetry-proto>=1.39,<2", "pydantic>=2", "pytest>=8", @@ -98,24 +100,35 @@ auditwheel = true default-groups = ["dev", "test"] package = true +[tool.uv.sources] +nemo-relay-plugin = { workspace = true } + +[tool.uv.workspace] +members = ["python/plugin"] + [tool.pytest.ini_options] testpaths = ["python/tests", "third_party/langgraph_tests"] asyncio_mode = "auto" +pythonpath = ["python", "python/plugin/src"] [tool.coverage.run] # Exclude integration tests from coverage, since we don't run these by default -omit = ["python/nemo_relay/integrations/*"] +omit = [ + "python/nemo_relay/integrations/*", + "python/plugin/src/nemo_relay_plugin/_proto/*", +] [tool.ty.analysis] # nemo_relay._native is a compiled Rust extension (built by maturin) that only # exists after `uv sync` / `pip install -e .`. Suppress unresolved-import for it. # LangChain, LangGraph, and Deep Agents are optional integration dependencies # which aren't installed by default. -allowed-unresolved-imports = ["deepagents.**", "langchain.**", "langchain_*.**", "langgraph.**", "nemo_relay._native", "pytest"] +allowed-unresolved-imports = ["deepagents.**", "grpc", "langchain.**", "langchain_*.**", "langgraph.**", "nemo_relay._native", "pytest"] [tool.ruff] line-length = 120 target-version = "py311" +extend-exclude = ["python/plugin/src/nemo_relay_plugin/_proto"] [tool.ruff.format] quote-style = "double" @@ -124,4 +137,4 @@ quote-style = "double" select = ["E", "F", "W", "I"] [tool.ruff.lint.isort] -known-first-party = ["nemo_relay"] +known-first-party = ["nemo_relay", "nemo_relay_plugin"] diff --git a/python/plugin/.gitignore b/python/plugin/.gitignore new file mode 100644 index 000000000..576001d82 --- /dev/null +++ b/python/plugin/.gitignore @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +.venv/ +.ruff_cache/ +__pycache__/ +*.egg-info/ +build/ +dist/ diff --git a/python/plugin/README.md b/python/plugin/README.md new file mode 100644 index 000000000..e85acdea5 --- /dev/null +++ b/python/plugin/README.md @@ -0,0 +1,36 @@ + + +# nemo-relay-plugin + +Python authoring SDK for NeMo Relay out-of-process dynamic worker plugins. + +Install this package in the Python environment used by a worker manifest with +`load.runtime = "python"`, then expose a `module:function` entrypoint that calls +`serve_plugin`. + +```python +from nemo_relay_plugin import Json, PluginContext, WorkerPlugin, serve_plugin + + +class PolicyPlugin(WorkerPlugin): + plugin_id = "acme.policy" + + def register(self, ctx: PluginContext, config: Json) -> None: + async def tag_tool_request(tool_name: str, args: Json) -> Json: + await ctx.runtime.emit_mark("acme.policy.tool_request", {"tool_name": tool_name}) + if isinstance(args, dict): + return {**args, "policy": "checked"} + return {"value": args, "policy": "checked"} + + ctx.register_tool_request_intercept("tag_tool_request", tag_tool_request) + + +async def main() -> None: + await serve_plugin(PolicyPlugin()) +``` + +The SDK owns gRPC serving, JSON envelope conversion, callback dispatch, +continuations, host runtime calls, and local scope-stack binding. diff --git a/python/plugin/pyproject.toml b/python/plugin/pyproject.toml new file mode 100644 index 000000000..07a5a6908 --- /dev/null +++ b/python/plugin/pyproject.toml @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + +[project] +name = "nemo-relay-plugin" +version = "0.5.0" +description = "Python SDK for NeMo Relay dynamic worker plugins." +readme = "README.md" +requires-python = ">=3.11" +license = "Apache-2.0" +authors = [ + { name = "NVIDIA Corporation & Affiliates" }, +] +keywords = [ + "agents", + "grpc", + "llm", + "middleware", + "nemo-relay", + "plugins", + "tools", +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", +] +dependencies = [ + "grpcio>=1.81.1,<2", + "protobuf>=6.33.5,<7", +] + +[project.urls] +Documentation = "https://docs.nvidia.com/nemo/relay" +Homepage = "https://github.com/NVIDIA/NeMo-Relay" +Issues = "https://github.com/NVIDIA/NeMo-Relay/issues" +Repository = "https://github.com/NVIDIA/NeMo-Relay" + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +nemo_relay_plugin = ["py.typed"] + +[tool.ruff] +line-length = 120 +target-version = "py311" +extend-exclude = ["src/nemo_relay_plugin/_proto"] + +[tool.ruff.format] +quote-style = "double" + +[tool.ruff.lint] +select = ["E", "F", "W", "I"] + +[tool.ruff.lint.isort] +known-first-party = ["nemo_relay_plugin"] diff --git a/python/plugin/src/nemo_relay_plugin/__init__.py b/python/plugin/src/nemo_relay_plugin/__init__.py new file mode 100644 index 000000000..4e0c1ac8b --- /dev/null +++ b/python/plugin/src/nemo_relay_plugin/__init__.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Python SDK for NeMo Relay dynamic worker plugins.""" + +from ._api import ( + AnnotatedLlmRequest, + ConfigDiagnostic, + DiagnosticLevel, + Event, + Json, + LlmNext, + LlmRequest, + LlmStreamNext, + PluginContext, + PluginRuntime, + ScopeType, + ToolNext, + WorkerPlugin, + WorkerSdkError, + serve_plugin, +) + +__all__ = [ + "AnnotatedLlmRequest", + "ConfigDiagnostic", + "DiagnosticLevel", + "Event", + "Json", + "LlmNext", + "LlmRequest", + "LlmStreamNext", + "PluginContext", + "PluginRuntime", + "ScopeType", + "ToolNext", + "WorkerPlugin", + "WorkerSdkError", + "serve_plugin", +] diff --git a/python/plugin/src/nemo_relay_plugin/_api.py b/python/plugin/src/nemo_relay_plugin/_api.py new file mode 100644 index 000000000..48b5e42ad --- /dev/null +++ b/python/plugin/src/nemo_relay_plugin/_api.py @@ -0,0 +1,997 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""High-level Python API for NeMo Relay grpc-v1 worker plugins.""" + +from __future__ import annotations + +import asyncio +import contextlib +import contextvars +import importlib +import inspect +import json +import os +import platform +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator, Mapping +from dataclasses import asdict, dataclass +from enum import Enum +from pathlib import Path +from typing import Any, Protocol, TypeAlias + +grpc: Any = importlib.import_module("grpc") +pb: Any = importlib.import_module("._proto.plugin_worker_pb2", __package__) +pb_grpc: Any = importlib.import_module("._proto.plugin_worker_pb2_grpc", __package__) + +Json: TypeAlias = Any +Event: TypeAlias = dict[str, Any] +LlmRequest: TypeAlias = dict[str, Any] +AnnotatedLlmRequest: TypeAlias = dict[str, Any] + +WORKER_PROTOCOL = "grpc-v1" +JSON_SCHEMA = "nemo.relay.Json@1" +EVENT_SCHEMA = "nemo.relay.Event@1" +LLM_REQUEST_SCHEMA = "nemo.relay.LlmRequest@1" +ANNOTATED_LLM_REQUEST_SCHEMA = "nemo.relay.AnnotatedLlmRequest@1" +PLUGIN_DIAGNOSTICS_SCHEMA = "nemo.relay.PluginDiagnostics@1" +_SCOPE_CONTEXT: contextvars.ContextVar[_BoundScopeContext | None] = contextvars.ContextVar( + "nemo_relay_plugin_scope_context", + default=None, +) + + +class WorkerSdkError(Exception): + """Error raised by the Python worker SDK.""" + + +class DiagnosticLevel(str, Enum): + """Plugin configuration diagnostic severity.""" + + WARNING = "warning" + ERROR = "error" + + +@dataclass(slots=True) +class ConfigDiagnostic: + """Structured plugin configuration diagnostic.""" + + level: DiagnosticLevel | str + code: str + message: str + component: str | None = None + field: str | None = None + + def to_json(self) -> dict[str, Any]: + """Return the Relay JSON representation.""" + value = asdict(self) + if isinstance(self.level, DiagnosticLevel): + value["level"] = self.level.value + return {key: item for key, item in value.items() if item is not None} + + +class ScopeType(str, Enum): + """Relay scope type accepted by host runtime scope calls.""" + + AGENT = "agent" + FUNCTION = "function" + TOOL = "tool" + LLM = "llm" + RETRIEVER = "retriever" + EMBEDDER = "embedder" + RERANKER = "reranker" + GUARDRAIL = "guardrail" + EVALUATOR = "evaluator" + CUSTOM = "custom" + UNKNOWN = "unknown" + + +@dataclass(frozen=True, slots=True) +class _BoundScopeContext: + scope_stack_id: str + parent_scope_id: str | None = None + + +class WorkerPlugin: + """Base class for Python worker plugins.""" + + plugin_id: str = "" + allows_multiple_components: bool = False + + def validate(self, config: Json) -> list[ConfigDiagnostic | dict[str, Any]]: + """Validate component config before registration.""" + del config + return [] + + def register(self, ctx: PluginContext, config: Json) -> None: + """Register callbacks into the worker plugin context.""" + del ctx, config + raise NotImplementedError("WorkerPlugin.register must be implemented") + + +class _SupportsWorkerPlugin(Protocol): + plugin_id: str + allows_multiple_components: bool + + def validate(self, config: Json) -> list[ConfigDiagnostic | dict[str, Any]]: ... + + def register(self, ctx: PluginContext, config: Json) -> None: ... + + +SubscriberCallback: TypeAlias = Callable[[Event], None | Awaitable[None]] +ToolSanitizeCallback: TypeAlias = Callable[[str, Json], Json | Awaitable[Json]] +ToolConditionalCallback: TypeAlias = Callable[[str, Json], str | None | Awaitable[str | None]] +ToolRequestCallback: TypeAlias = Callable[[str, Json], Json | Awaitable[Json]] +ToolExecutionCallback: TypeAlias = Callable[[str, Json, "ToolNext"], Json | Awaitable[Json]] +LlmSanitizeRequestCallback: TypeAlias = Callable[[LlmRequest], LlmRequest | Awaitable[LlmRequest]] +LlmSanitizeResponseCallback: TypeAlias = Callable[[Json], Json | Awaitable[Json]] +LlmConditionalCallback: TypeAlias = Callable[[LlmRequest], str | None | Awaitable[str | None]] +LlmRequestCallback: TypeAlias = Callable[ + [str, LlmRequest, AnnotatedLlmRequest | None], + LlmRequest + | tuple[LlmRequest, AnnotatedLlmRequest | None] + | Awaitable[LlmRequest | tuple[LlmRequest, AnnotatedLlmRequest | None]], +] +LlmExecutionCallback: TypeAlias = Callable[[str, LlmRequest, "LlmNext"], Json | Awaitable[Json]] +LlmStreamExecutionCallback: TypeAlias = Callable[ + [str, LlmRequest, "LlmStreamNext"], + Iterable[Json] | AsyncIterator[Json] | Awaitable[Iterable[Json] | AsyncIterator[Json]], +] + + +@dataclass(slots=True) +class _Handlers: + registrations: list[Any] + subscribers: dict[str, SubscriberCallback] + tool_sanitize_requests: dict[str, ToolSanitizeCallback] + tool_sanitize_responses: dict[str, ToolSanitizeCallback] + tool_conditionals: dict[str, ToolConditionalCallback] + tool_requests: dict[str, ToolRequestCallback] + tool_executions: dict[str, ToolExecutionCallback] + llm_sanitize_requests: dict[str, LlmSanitizeRequestCallback] + llm_sanitize_responses: dict[str, LlmSanitizeResponseCallback] + llm_conditionals: dict[str, LlmConditionalCallback] + llm_requests: dict[str, LlmRequestCallback] + llm_executions: dict[str, LlmExecutionCallback] + llm_stream_executions: dict[str, LlmStreamExecutionCallback] + + @classmethod + def empty(cls) -> _Handlers: + return cls( + registrations=[], + subscribers={}, + tool_sanitize_requests={}, + tool_sanitize_responses={}, + tool_conditionals={}, + tool_requests={}, + tool_executions={}, + llm_sanitize_requests={}, + llm_sanitize_responses={}, + llm_conditionals={}, + llm_requests={}, + llm_executions={}, + llm_stream_executions={}, + ) + + +class PluginContext: + """Registration context passed to ``WorkerPlugin.register``.""" + + def __init__(self, runtime: PluginRuntime | None = None) -> None: + self._runtime = runtime + self._handlers = _Handlers.empty() + + @property + def runtime(self) -> PluginRuntime: + """Return the host runtime handle for event and scope operations.""" + if self._runtime is None: + raise WorkerSdkError("PluginContext has no runtime handle") + return self._runtime + + def register_subscriber(self, name: str, callback: SubscriberCallback) -> None: + """Register an event subscriber.""" + self._push_registration(name, pb.SUBSCRIBER, 0, False) + self._handlers.subscribers[name] = callback + + def register_tool_sanitize_request_guardrail( + self, + name: str, + callback: ToolSanitizeCallback, + *, + priority: int = 0, + ) -> None: + """Register a tool sanitize-request guardrail.""" + self._push_registration(name, pb.TOOL_SANITIZE_REQUEST_GUARDRAIL, priority, False) + self._handlers.tool_sanitize_requests[name] = callback + + def register_tool_sanitize_response_guardrail( + self, + name: str, + callback: ToolSanitizeCallback, + *, + priority: int = 0, + ) -> None: + """Register a tool sanitize-response guardrail.""" + self._push_registration(name, pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL, priority, False) + self._handlers.tool_sanitize_responses[name] = callback + + def register_tool_conditional_execution_guardrail( + self, + name: str, + callback: ToolConditionalCallback, + *, + priority: int = 0, + ) -> None: + """Register a tool conditional-execution guardrail.""" + self._push_registration(name, pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, priority, False) + self._handlers.tool_conditionals[name] = callback + + def register_tool_request_intercept( + self, + name: str, + callback: ToolRequestCallback, + *, + priority: int = 0, + break_chain: bool = False, + ) -> None: + """Register a tool request intercept.""" + self._push_registration(name, pb.TOOL_REQUEST_INTERCEPT, priority, break_chain) + self._handlers.tool_requests[name] = callback + + def register_tool_execution_intercept( + self, + name: str, + callback: ToolExecutionCallback, + *, + priority: int = 0, + ) -> None: + """Register a tool execution intercept.""" + self._push_registration(name, pb.TOOL_EXECUTION_INTERCEPT, priority, False) + self._handlers.tool_executions[name] = callback + + def register_llm_sanitize_request_guardrail( + self, + name: str, + callback: LlmSanitizeRequestCallback, + *, + priority: int = 0, + ) -> None: + """Register an LLM sanitize-request guardrail.""" + self._push_registration(name, pb.LLM_SANITIZE_REQUEST_GUARDRAIL, priority, False) + self._handlers.llm_sanitize_requests[name] = callback + + def register_llm_sanitize_response_guardrail( + self, + name: str, + callback: LlmSanitizeResponseCallback, + *, + priority: int = 0, + ) -> None: + """Register an LLM sanitize-response guardrail.""" + self._push_registration(name, pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, priority, False) + self._handlers.llm_sanitize_responses[name] = callback + + def register_llm_conditional_execution_guardrail( + self, + name: str, + callback: LlmConditionalCallback, + *, + priority: int = 0, + ) -> None: + """Register an LLM conditional-execution guardrail.""" + self._push_registration(name, pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, priority, False) + self._handlers.llm_conditionals[name] = callback + + def register_llm_request_intercept( + self, + name: str, + callback: LlmRequestCallback, + *, + priority: int = 0, + break_chain: bool = False, + ) -> None: + """Register an LLM request intercept.""" + self._push_registration(name, pb.LLM_REQUEST_INTERCEPT, priority, break_chain) + self._handlers.llm_requests[name] = callback + + def register_llm_execution_intercept( + self, + name: str, + callback: LlmExecutionCallback, + *, + priority: int = 0, + ) -> None: + """Register an LLM execution intercept.""" + self._push_registration(name, pb.LLM_EXECUTION_INTERCEPT, priority, False) + self._handlers.llm_executions[name] = callback + + def register_llm_stream_execution_intercept( + self, + name: str, + callback: LlmStreamExecutionCallback, + *, + priority: int = 0, + ) -> None: + """Register an LLM stream execution intercept.""" + self._push_registration(name, pb.LLM_STREAM_EXECUTION_INTERCEPT, priority, False) + self._handlers.llm_stream_executions[name] = callback + + def _push_registration(self, name: str, surface: int, priority: int, break_chain: bool) -> None: + self._handlers.registrations.append( + pb.Registration( + local_name=name, + surface=surface, + priority=priority, + break_chain=break_chain, + ) + ) + + +class PluginRuntime: + """Handle for calling the Relay host runtime from worker callbacks.""" + + def __init__(self, *, activation_id: str, auth_token: str, host_stub: Any) -> None: + self._activation_id = activation_id + self._auth_token = auth_token + self._host_stub = host_stub + + async def emit_mark( + self, + name: str, + data: Json | None = None, + metadata: Json | None = None, + *, + scope_stack_id: str | None = None, + parent_scope_id: str | None = None, + ) -> None: + """Emit a mark event through the host runtime.""" + response = await self._host_stub.EmitMark( + pb.EmitMarkRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + scope=self._scope_context(scope_stack_id, parent_scope_id), + name=name, + data=_optional_json_envelope(data), + metadata=_optional_json_envelope(metadata), + ) + ) + _ack_to_result(response) + + async def create_scope_stack(self) -> str: + """Create an isolated host-owned scope stack.""" + response = await self._host_stub.CreateScopeStack( + pb.CreateScopeStackRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + ) + ) + if response.HasField("error"): + raise _worker_error_to_sdk(response.error) + return response.scope_stack_id + + async def drop_scope_stack(self, scope_stack_id: str) -> None: + """Drop an isolated host-owned scope stack.""" + response = await self._host_stub.DropScopeStack( + pb.DropScopeStackRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + scope_stack_id=scope_stack_id, + ) + ) + _ack_to_result(response) + + async def push_scope( + self, + name: str, + *, + scope_type: ScopeType = ScopeType.CUSTOM, + data: Json | None = None, + metadata: Json | None = None, + input: Json | None = None, + scope_stack_id: str | None = None, + parent_scope_id: str | None = None, + ) -> str: + """Push a scope through the host runtime and return its handle ID.""" + response = await self._host_stub.PushScope( + pb.PushScopeRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + scope=self._scope_context(scope_stack_id, parent_scope_id), + name=name, + scope_type=_proto_scope_type(scope_type), + data=_optional_json_envelope(data), + metadata=_optional_json_envelope(metadata), + input=_optional_json_envelope(input), + ) + ) + if response.HasField("error"): + raise _worker_error_to_sdk(response.error) + return response.scope_handle_id + + async def pop_scope( + self, + scope_handle_id: str, + *, + output: Json | None = None, + metadata: Json | None = None, + ) -> None: + """Pop a host scope by handle ID.""" + response = await self._host_stub.PopScope( + pb.PopScopeRequest( + activation_id=self._activation_id, + auth_token=self._auth_token, + scope_handle_id=scope_handle_id, + output=_optional_json_envelope(output), + metadata=_optional_json_envelope(metadata), + ) + ) + _ack_to_result(response) + + @contextlib.contextmanager + def bind_scope_stack(self, scope_stack_id: str | None, *, parent_scope_id: str | None = None) -> Iterator[None]: + """Temporarily bind callbacks to a worker-selected scope stack.""" + scope = _BoundScopeContext(scope_stack_id, parent_scope_id) if scope_stack_id else None + token = _SCOPE_CONTEXT.set(scope) + try: + yield + finally: + _SCOPE_CONTEXT.reset(token) + + @contextlib.contextmanager + def clear_scope_stack(self) -> Iterator[None]: + """Temporarily clear worker scope-stack correlation.""" + with self.bind_scope_stack(None): + yield + + def current_scope_stack_id(self) -> str | None: + """Return the locally bound scope stack ID, if any.""" + scope = _SCOPE_CONTEXT.get() + return scope.scope_stack_id if scope else None + + def current_parent_scope_id(self) -> str | None: + """Return the locally bound parent scope ID, if any.""" + scope = _SCOPE_CONTEXT.get() + return scope.parent_scope_id if scope else None + + def _scope_context(self, scope_stack_id: str | None = None, parent_scope_id: str | None = None) -> Any: + if scope_stack_id is not None: + effective_scope = _BoundScopeContext(scope_stack_id, parent_scope_id) + else: + bound_scope = _SCOPE_CONTEXT.get() + if bound_scope is not None: + effective_scope = _BoundScopeContext( + bound_scope.scope_stack_id, + parent_scope_id if parent_scope_id is not None else bound_scope.parent_scope_id, + ) + elif parent_scope_id is not None: + raise WorkerSdkError("parent_scope_id requires an explicit or bound scope stack") + else: + effective_scope = None + if not effective_scope: + return None + return pb.ScopeContext( + scope_stack_id=effective_scope.scope_stack_id, + parent_scope_id=effective_scope.parent_scope_id or "", + ) + + +class ToolNext: + """Continuation handle for tool execution intercepts.""" + + def __init__(self, runtime: PluginRuntime, continuation_id: str) -> None: + self._runtime = runtime + self._continuation_id = continuation_id + + async def call(self, value: Json) -> Json: + """Call the remaining tool execution chain.""" + response = await self._runtime._host_stub.ToolNext( + pb.ToolNextRequest( + activation_id=self._runtime._activation_id, + auth_token=self._runtime._auth_token, + continuation_id=self._continuation_id, + value=_json_envelope(JSON_SCHEMA, value), + ) + ) + return _json_result_to_value(response) + + +class LlmNext: + """Continuation handle for LLM execution intercepts.""" + + def __init__(self, runtime: PluginRuntime, continuation_id: str) -> None: + self._runtime = runtime + self._continuation_id = continuation_id + + async def call(self, request: LlmRequest) -> Json: + """Call the remaining LLM execution chain.""" + response = await self._runtime._host_stub.LlmNext( + pb.LlmNextRequest( + activation_id=self._runtime._activation_id, + auth_token=self._runtime._auth_token, + continuation_id=self._continuation_id, + request=_json_envelope(LLM_REQUEST_SCHEMA, request), + ) + ) + return _json_result_to_value(response) + + +class LlmStreamNext: + """Continuation handle for LLM stream execution intercepts.""" + + def __init__(self, runtime: PluginRuntime, continuation_id: str) -> None: + self._runtime = runtime + self._continuation_id = continuation_id + + def call(self, request: LlmRequest) -> AsyncIterator[Json]: + """Call the remaining LLM stream execution chain.""" + scope_context = _SCOPE_CONTEXT.get() + stream = self._runtime._host_stub.LlmStreamNext( + pb.LlmStreamNextRequest( + activation_id=self._runtime._activation_id, + auth_token=self._runtime._auth_token, + continuation_id=self._continuation_id, + request=_json_envelope(LLM_REQUEST_SCHEMA, request), + ) + ) + + async def values() -> AsyncIterator[Json]: + token = _SCOPE_CONTEXT.set(scope_context) + try: + async for chunk in stream: + yield _stream_chunk_to_value(chunk) + finally: + _SCOPE_CONTEXT.reset(token) + + return values() + + +async def serve_plugin(plugin: _SupportsWorkerPlugin) -> None: + """Serve a worker plugin using environment variables supplied by the Relay host.""" + worker_endpoint = _required_env("NEMO_RELAY_WORKER_SOCKET") + host_endpoint = _required_env("NEMO_RELAY_HOST_SOCKET") + activation_id = _required_env("NEMO_RELAY_WORKER_ID") + auth_token = _required_env("NEMO_RELAY_WORKER_TOKEN") + endpoint_file = os.environ.get("NEMO_RELAY_WORKER_ENDPOINT_FILE") + + _unlink_unix_socket(worker_endpoint) + host_channel = grpc.aio.insecure_channel(_grpc_target(host_endpoint)) + runtime = PluginRuntime( + activation_id=activation_id, + auth_token=auth_token, + host_stub=pb_grpc.RelayHostRuntimeStub(host_channel), + ) + shutdown_event = asyncio.Event() + service = _WorkerService(plugin, runtime, shutdown_event) + server = grpc.aio.server() + pb_grpc.add_PluginWorkerServicer_to_server(service, server) + bound_port = server.add_insecure_port(_grpc_target(worker_endpoint)) + if bound_port == 0: + raise WorkerSdkError(f"failed to bind worker endpoint {worker_endpoint}") + try: + await server.start() + if endpoint_file: + path = Path(endpoint_file) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(_announced_worker_endpoint(worker_endpoint, bound_port), encoding="utf-8") + await shutdown_event.wait() + finally: + await server.stop(grace=2) + await host_channel.close() + + +class _WorkerService(pb_grpc.PluginWorkerServicer): + def __init__( + self, + plugin: _SupportsWorkerPlugin, + runtime: PluginRuntime, + shutdown_event: asyncio.Event, + ) -> None: + self._plugin = plugin + self._runtime = runtime + self._shutdown_event = shutdown_event + self._handlers = _Handlers.empty() + + async def Handshake(self, request: Any, context: Any) -> Any: + await self._authorize(request, context) + plugin_id = _plugin_id(self._plugin) + return pb.HandshakeResponse( + plugin_id=plugin_id, + plugin_kind=plugin_id, + allows_multiple_components=bool(getattr(self._plugin, "allows_multiple_components", False)), + worker_protocol=WORKER_PROTOCOL, + sdk_name="nemo-relay-plugin", + sdk_version="0.5.0", + runtime_name="python", + runtime_version=platform.python_version(), + supported_surfaces=_all_surfaces(), + ) + + async def Health(self, request: Any, context: Any) -> Any: + await self._authorize(request, context) + plugin_id = _plugin_id(self._plugin) + return pb.HealthResponse( + ok=True, + message="ready", + plugin_id=plugin_id, + worker_protocol=WORKER_PROTOCOL, + sdk_name="nemo-relay-plugin", + sdk_version="0.5.0", + runtime_name="python", + runtime_version=platform.python_version(), + ) + + async def Validate(self, request: Any, context: Any) -> Any: + await self._authorize(request, context) + config = await _decode_optional_config_or_abort(request, context) + try: + diagnostics = [_diagnostic_to_json(item) for item in self._plugin.validate(config)] + return pb.ValidateResponse( + diagnostics=_json_envelope(PLUGIN_DIAGNOSTICS_SCHEMA, diagnostics), + ) + except Exception as exc: # noqa: BLE001 - callback failure is protocol data. + return pb.ValidateResponse(error=_sdk_error_to_worker(exc)) + + async def Register(self, request: Any, context: Any) -> Any: + await self._authorize(request, context) + config = await _decode_optional_config_or_abort(request, context) + try: + ctx = PluginContext(runtime=self._runtime) + self._plugin.register(ctx, config) + self._handlers = ctx._handlers + return pb.RegisterResponse(registrations=ctx._handlers.registrations) + except Exception as exc: # noqa: BLE001 - callback failure is protocol data. + return pb.RegisterResponse(error=_sdk_error_to_worker(exc)) + + async def Invoke(self, request: Any, context: Any) -> Any: + await self._authorize(request, context) + try: + return await self._invoke_result(request) + except Exception as exc: # noqa: BLE001 - callback failure is protocol data. + return pb.InvokeResponse(error=_sdk_error_to_worker(exc)) + + async def InvokeStream(self, request: Any, context: Any) -> AsyncIterator[Any]: + await self._authorize(request, context) + try: + if request.surface != pb.LLM_STREAM_EXECUTION_INTERCEPT: + raise WorkerSdkError("InvokeStream only supports LLM stream execution intercepts") + handler = self._handler(self._handlers.llm_stream_executions, request.registration_name) + payload = _require_payload(request, "llm") + llm_request = _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA) + next_call = LlmStreamNext(self._runtime, request.continuation_id) + with _bind_invocation_scope(request): + stream = await _maybe_await(handler(payload.model_name, llm_request, next_call)) + async for value in _as_async_iter(stream): + yield pb.StreamChunk(value=_json_envelope(JSON_SCHEMA, value)) + except Exception as exc: # noqa: BLE001 - callback failure is protocol data. + yield pb.StreamChunk(error=_sdk_error_to_worker(exc)) + + async def CancelInvocation(self, request: Any, context: Any) -> Any: + await self._authorize(request, context) + return pb.WorkerAck(accepted=False, message="cancel is not implemented by the Python worker SDK") + + async def Shutdown(self, request: Any, context: Any) -> Any: + await self._authorize(request, context) + asyncio.get_running_loop().call_soon(self._shutdown_event.set) + return pb.WorkerAck(accepted=True, message="shutdown accepted") + + async def _invoke_result(self, request: Any) -> Any: + with _bind_invocation_scope(request): + if request.surface == pb.SUBSCRIBER: + event = _decode_required_envelope(request.event, "event", EVENT_SCHEMA) + await _maybe_await(self._handler(self._handlers.subscribers, request.registration_name)(event)) + return pb.InvokeResponse(empty=pb.EmptyResult()) + if request.surface == pb.TOOL_SANITIZE_REQUEST_GUARDRAIL: + return _json_response( + await _maybe_await( + self._handler(self._handlers.tool_sanitize_requests, request.registration_name)( + request.tool.tool_name, + _decode_required_envelope(request.tool.value, "tool value"), + ) + ) + ) + if request.surface == pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL: + return _json_response( + await _maybe_await( + self._handler(self._handlers.tool_sanitize_responses, request.registration_name)( + request.tool.tool_name, + _decode_required_envelope(request.tool.value, "tool value"), + ) + ) + ) + if request.surface == pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL: + result = await _maybe_await( + self._handler(self._handlers.tool_conditionals, request.registration_name)( + request.tool.tool_name, + _decode_required_envelope(request.tool.value, "tool value"), + ) + ) + return pb.InvokeResponse(guardrail=pb.GuardrailResult(block_reason=result or "")) + if request.surface == pb.TOOL_REQUEST_INTERCEPT: + result = await _maybe_await( + self._handler(self._handlers.tool_requests, request.registration_name)( + request.tool.tool_name, + _decode_required_envelope(request.tool.value, "tool value"), + ) + ) + return _json_response(result) + if request.surface == pb.TOOL_EXECUTION_INTERCEPT: + result = await _maybe_await( + self._handler(self._handlers.tool_executions, request.registration_name)( + request.tool.tool_name, + _decode_required_envelope(request.tool.value, "tool value"), + ToolNext(self._runtime, request.continuation_id), + ) + ) + return _json_response(result) + if request.surface == pb.LLM_SANITIZE_REQUEST_GUARDRAIL: + return _json_response( + await _maybe_await( + self._handler(self._handlers.llm_sanitize_requests, request.registration_name)( + _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) + ) + ) + ) + if request.surface == pb.LLM_SANITIZE_RESPONSE_GUARDRAIL: + return _json_response( + await _maybe_await( + self._handler(self._handlers.llm_sanitize_responses, request.registration_name)( + _decode_required_envelope(request.llm.response, "llm response") + ) + ) + ) + if request.surface == pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL: + result = await _maybe_await( + self._handler(self._handlers.llm_conditionals, request.registration_name)( + _decode_required_envelope(request.llm.request, "llm request", LLM_REQUEST_SCHEMA) + ) + ) + return pb.InvokeResponse(guardrail=pb.GuardrailResult(block_reason=result or "")) + if request.surface == pb.LLM_REQUEST_INTERCEPT: + payload = request.llm + llm_request = _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA) + annotated = ( + _decode_required_envelope( + payload.annotated_request, + "annotated llm request", + ANNOTATED_LLM_REQUEST_SCHEMA, + ) + if payload.HasField("annotated_request") + else None + ) + result = await _maybe_await( + self._handler(self._handlers.llm_requests, request.registration_name)( + payload.model_name, + llm_request, + annotated, + ) + ) + if isinstance(result, tuple): + llm_request, annotated = result + else: + llm_request = result + return pb.InvokeResponse( + llm_request=pb.LlmRequestInterceptResult( + request=_json_envelope(LLM_REQUEST_SCHEMA, llm_request), + annotated_request=_optional_json_envelope(annotated, ANNOTATED_LLM_REQUEST_SCHEMA), + has_annotated_request=annotated is not None, + ) + ) + if request.surface == pb.LLM_EXECUTION_INTERCEPT: + payload = request.llm + result = await _maybe_await( + self._handler(self._handlers.llm_executions, request.registration_name)( + payload.model_name, + _decode_required_envelope(payload.request, "llm request", LLM_REQUEST_SCHEMA), + LlmNext(self._runtime, request.continuation_id), + ) + ) + return _json_response(result) + raise WorkerSdkError(f"unsupported registration surface {request.surface}") + + async def _authorize(self, request: Any, context: Any) -> None: + if request.activation_id != self._runtime._activation_id: + await context.abort(grpc.StatusCode.PERMISSION_DENIED, "invalid activation ID") + if request.auth_token != self._runtime._auth_token: + await context.abort(grpc.StatusCode.PERMISSION_DENIED, "invalid auth token") + + def _handler(self, handlers: dict[str, Any], name: str) -> Any: + try: + return handlers[name] + except KeyError as exc: + raise WorkerSdkError(f"handler {name!r} is not registered") from exc + + +def _plugin_id(plugin: _SupportsWorkerPlugin) -> str: + plugin_id = getattr(plugin, "plugin_id", "") + if callable(plugin_id): + plugin_id = plugin_id() + if not isinstance(plugin_id, str) or not plugin_id: + raise WorkerSdkError("plugin_id must be a non-empty string") + return plugin_id + + +def _all_surfaces() -> list[int]: + return [ + pb.SUBSCRIBER, + pb.TOOL_SANITIZE_REQUEST_GUARDRAIL, + pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL, + pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, + pb.TOOL_REQUEST_INTERCEPT, + pb.TOOL_EXECUTION_INTERCEPT, + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, + pb.LLM_REQUEST_INTERCEPT, + pb.LLM_EXECUTION_INTERCEPT, + pb.LLM_STREAM_EXECUTION_INTERCEPT, + ] + + +def _diagnostic_to_json(value: ConfigDiagnostic | dict[str, Any]) -> dict[str, Any]: + if isinstance(value, ConfigDiagnostic): + return value.to_json() + return dict(value) + + +def _json_envelope(schema: str, value: Json) -> Any: + return pb.JsonEnvelope( + schema=schema, + json=json.dumps(value, separators=(",", ":")).encode("utf-8"), + ) + + +def _optional_json_envelope(value: Json | None, schema: str = JSON_SCHEMA) -> Any: + if value is None: + return None + return _json_envelope(schema, value) + + +def _decode_required_envelope(envelope: Any, field: str, expected_schema: str = JSON_SCHEMA) -> Json: + if envelope is None or not getattr(envelope, "json", b""): + raise WorkerSdkError(f"{field} is missing") + if envelope.schema != expected_schema: + raise WorkerSdkError(f"{field} has schema {envelope.schema!r}; expected {expected_schema!r}") + return json.loads(envelope.json.decode("utf-8")) + + +def _decode_optional_json(message: Any, field: str, *, default: Json) -> Json: + if hasattr(message, "HasField") and not message.HasField(field): + return default + return _decode_required_envelope(getattr(message, field), field) + + +async def _decode_optional_config_or_abort(message: Any, context: Any) -> Json: + try: + return _decode_optional_json(message, "config", default=None) + except Exception as exc: # noqa: BLE001 - malformed config is a protocol error. + await context.abort(grpc.StatusCode.INVALID_ARGUMENT, f"invalid config: {exc}") + raise AssertionError("context.abort should not return") from exc + + +def _json_response(value: Json) -> Any: + return pb.InvokeResponse(json=pb.JsonResult(value=_json_envelope(JSON_SCHEMA, value))) + + +def _json_result_to_value(result: Any) -> Json: + if result.HasField("error"): + raise _worker_error_to_sdk(result.error) + return _decode_required_envelope(result.value, "json result") + + +def _stream_chunk_to_value(chunk: Any) -> Json: + item = chunk.WhichOneof("item") + if item == "error": + raise _worker_error_to_sdk(chunk.error) + if item != "value": + raise WorkerSdkError("stream chunk is empty") + return _decode_required_envelope(chunk.value, "stream chunk") + + +def _worker_error_to_sdk(error: Any) -> WorkerSdkError: + return WorkerSdkError(f"{error.code}: {error.message}") + + +def _sdk_error_to_worker(error: BaseException) -> Any: + code = "worker.error" + if isinstance(error, WorkerSdkError): + code = "worker.sdk_error" + return pb.WorkerError(code=code, message=str(error), retryable=False) + + +def _ack_to_result(response: Any) -> None: + if response.ok: + return + if response.HasField("error"): + raise _worker_error_to_sdk(response.error) + raise WorkerSdkError("host call failed") + + +def _require_payload(request: Any, payload: str) -> Any: + if request.WhichOneof("payload") != payload: + raise WorkerSdkError(f"expected {payload} payload") + return getattr(request, payload) + + +@contextlib.contextmanager +def _bind_invocation_scope(request: Any) -> Iterator[None]: + scope = None + if request.HasField("scope") and request.scope.scope_stack_id: + scope = _BoundScopeContext( + scope_stack_id=request.scope.scope_stack_id, + parent_scope_id=request.scope.parent_scope_id or None, + ) + token = _SCOPE_CONTEXT.set(scope) + try: + yield + finally: + _SCOPE_CONTEXT.reset(token) + + +async def _maybe_await(value: Any) -> Any: + if inspect.isawaitable(value): + return await value + return value + + +async def _as_async_iter(value: Iterable[Json] | AsyncIterator[Json]) -> AsyncIterator[Json]: + if isinstance(value, AsyncIterator): + async for item in value: + yield item + return + if isinstance(value, (str, bytes, bytearray, Mapping)): + raise WorkerSdkError("stream callback must return an iterable of JSON chunks, not a scalar or mapping") + if not isinstance(value, Iterable): + raise WorkerSdkError("stream callback must return an iterable or async iterator of JSON chunks") + for item in value: + yield item + + +def _proto_scope_type(scope_type: ScopeType | str) -> int: + value = ScopeType(scope_type) + mapping = { + ScopeType.AGENT: pb.AGENT, + ScopeType.FUNCTION: pb.FUNCTION, + ScopeType.TOOL: pb.TOOL, + ScopeType.LLM: pb.LLM, + ScopeType.RETRIEVER: pb.RETRIEVER, + ScopeType.EMBEDDER: pb.EMBEDDER, + ScopeType.RERANKER: pb.RERANKER, + ScopeType.GUARDRAIL: pb.GUARDRAIL, + ScopeType.EVALUATOR: pb.EVALUATOR, + ScopeType.CUSTOM: pb.CUSTOM, + ScopeType.UNKNOWN: pb.UNKNOWN, + } + return mapping[value] + + +def _required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise WorkerSdkError(f"environment variable {name} is required") + return value + + +def _grpc_target(endpoint: str) -> str: + if endpoint.startswith("unix://"): + return "unix:" + endpoint.removeprefix("unix://") + if endpoint.startswith("tcp://"): + return endpoint.removeprefix("tcp://") + if endpoint.startswith("http://"): + return endpoint.removeprefix("http://") + return endpoint + + +def _announced_worker_endpoint(worker_endpoint: str, bound_port: int) -> str: + target = _grpc_target(worker_endpoint) + if target.startswith("unix:"): + return worker_endpoint + host, separator, port = target.rpartition(":") + if not separator: + return f"http://{target}" + if port == "0": + return f"http://{host}:{bound_port}" + return f"http://{host}:{port}" + + +def _unlink_unix_socket(endpoint: str) -> None: + if endpoint.startswith("unix://"): + Path(endpoint.removeprefix("unix://")).unlink(missing_ok=True) diff --git a/python/plugin/src/nemo_relay_plugin/_proto/__init__.py b/python/plugin/src/nemo_relay_plugin/_proto/__init__.py new file mode 100644 index 000000000..70ab1f885 --- /dev/null +++ b/python/plugin/src/nemo_relay_plugin/_proto/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generated protobuf modules for the private worker protocol binding.""" diff --git a/python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2.py b/python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2.py new file mode 100644 index 000000000..650c6c0d9 --- /dev/null +++ b/python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: plugin_worker.proto +# Protobuf Python Version: 6.33.5 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 33, + 5, + '', + 'plugin_worker.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x13plugin_worker.proto\x12\x14nemo.relay.worker.v1\",\n\x0cJsonEnvelope\x12\x0e\n\x06schema\x18\x01 \x01(\t\x12\x0c\n\x04json\x18\x02 \x01(\x0c\"\x97\x01\n\x10HandshakeRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x11\n\tplugin_id\x18\x02 \x01(\t\x12\x15\n\rrelay_version\x18\x03 \x01(\t\x12\x17\n\x0fworker_protocol\x18\x04 \x01(\t\x12\x12\n\nauth_token\x18\x05 \x01(\t\x12\x15\n\rhost_endpoint\x18\x06 \x01(\t\"\x95\x02\n\x11HandshakeResponse\x12\x11\n\tplugin_id\x18\x01 \x01(\t\x12\x13\n\x0bplugin_kind\x18\x02 \x01(\t\x12\"\n\x1a\x61llows_multiple_components\x18\x03 \x01(\x08\x12\x17\n\x0fworker_protocol\x18\x04 \x01(\t\x12\x10\n\x08sdk_name\x18\x05 \x01(\t\x12\x13\n\x0bsdk_version\x18\x06 \x01(\t\x12\x14\n\x0cruntime_name\x18\x07 \x01(\t\x12\x17\n\x0fruntime_version\x18\x08 \x01(\t\x12\x45\n\x12supported_surfaces\x18\t \x03(\x0e\x32).nemo.relay.worker.v1.RegistrationSurface\":\n\rHealthRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\"\xaf\x01\n\x0eHealthResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x11\n\tplugin_id\x18\x03 \x01(\t\x12\x17\n\x0fworker_protocol\x18\x04 \x01(\t\x12\x10\n\x08sdk_name\x18\x05 \x01(\t\x12\x13\n\x0bsdk_version\x18\x06 \x01(\t\x12\x14\n\x0cruntime_name\x18\x07 \x01(\t\x12\x17\n\x0fruntime_version\x18\x08 \x01(\t\"\x83\x01\n\x0fValidateRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x11\n\tplugin_id\x18\x02 \x01(\t\x12\x12\n\nauth_token\x18\x03 \x01(\t\x12\x32\n\x06\x63onfig\x18\x04 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"}\n\x10ValidateResponse\x12\x37\n\x0b\x64iagnostics\x18\x01 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12\x30\n\x05\x65rror\x18\x02 \x01(\x0b\x32!.nemo.relay.worker.v1.WorkerError\"\x83\x01\n\x0fRegisterRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x11\n\tplugin_id\x18\x02 \x01(\t\x12\x12\n\nauth_token\x18\x03 \x01(\t\x12\x32\n\x06\x63onfig\x18\x04 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"\x7f\n\x10RegisterResponse\x12\x39\n\rregistrations\x18\x01 \x03(\x0b\x32\".nemo.relay.worker.v1.Registration\x12\x30\n\x05\x65rror\x18\x02 \x01(\x0b\x32!.nemo.relay.worker.v1.WorkerError\"\x85\x01\n\x0cRegistration\x12\x12\n\nlocal_name\x18\x01 \x01(\t\x12:\n\x07surface\x18\x02 \x01(\x0e\x32).nemo.relay.worker.v1.RegistrationSurface\x12\x10\n\x08priority\x18\x03 \x01(\x05\x12\x13\n\x0b\x62reak_chain\x18\x04 \x01(\x08\"\x9e\x03\n\rInvokeRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x15\n\rinvocation_id\x18\x02 \x01(\t\x12\x19\n\x11registration_name\x18\x03 \x01(\t\x12:\n\x07surface\x18\x04 \x01(\x0e\x32).nemo.relay.worker.v1.RegistrationSurface\x12\x17\n\x0f\x63ontinuation_id\x18\x05 \x01(\t\x12\x31\n\x05scope\x18\x06 \x01(\x0b\x32\".nemo.relay.worker.v1.ScopeContext\x12\x12\n\nauth_token\x18\x07 \x01(\t\x12\x33\n\x05\x65vent\x18\n \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelopeH\x00\x12\x34\n\x04tool\x18\x0b \x01(\x0b\x32$.nemo.relay.worker.v1.ToolInvocationH\x00\x12\x32\n\x03llm\x18\x0c \x01(\x0b\x32#.nemo.relay.worker.v1.LlmInvocationH\x00\x42\t\n\x07payload\"V\n\x0eToolInvocation\x12\x11\n\ttool_name\x18\x01 \x01(\t\x12\x31\n\x05value\x18\x02 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"\xcd\x01\n\rLlmInvocation\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x33\n\x07request\x18\x02 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12=\n\x11\x61nnotated_request\x18\x03 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12\x34\n\x08response\x18\x04 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"\xb8\x02\n\x0eInvokeResponse\x12\x32\n\x05\x65mpty\x18\x01 \x01(\x0b\x32!.nemo.relay.worker.v1.EmptyResultH\x00\x12\x30\n\x04json\x18\x02 \x01(\x0b\x32 .nemo.relay.worker.v1.JsonResultH\x00\x12:\n\tguardrail\x18\x03 \x01(\x0b\x32%.nemo.relay.worker.v1.GuardrailResultH\x00\x12\x46\n\x0bllm_request\x18\x04 \x01(\x0b\x32/.nemo.relay.worker.v1.LlmRequestInterceptResultH\x00\x12\x32\n\x05\x65rror\x18\x05 \x01(\x0b\x32!.nemo.relay.worker.v1.WorkerErrorH\x00\x42\x08\n\x06result\"\r\n\x0b\x45mptyResult\"q\n\nJsonResult\x12\x31\n\x05value\x18\x01 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12\x30\n\x05\x65rror\x18\x02 \x01(\x0b\x32!.nemo.relay.worker.v1.WorkerError\"\'\n\x0fGuardrailResult\x12\x14\n\x0c\x62lock_reason\x18\x01 \x01(\t\"\xae\x01\n\x19LlmRequestInterceptResult\x12\x33\n\x07request\x18\x01 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12=\n\x11\x61nnotated_request\x18\x02 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12\x1d\n\x15has_annotated_request\x18\x03 \x01(\x08\"~\n\x0bStreamChunk\x12\x33\n\x05value\x18\x01 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelopeH\x00\x12\x32\n\x05\x65rror\x18\x02 \x01(\x0b\x32!.nemo.relay.worker.v1.WorkerErrorH\x00\x42\x06\n\x04item\"k\n\x17\x43\x61ncelInvocationRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x15\n\rinvocation_id\x18\x02 \x01(\t\x12\x12\n\nauth_token\x18\x03 \x01(\t\x12\x0e\n\x06reason\x18\x04 \x01(\t\"L\n\x0fShutdownRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\x12\x0e\n\x06reason\x18\x03 \x01(\t\".\n\tWorkerAck\x12\x10\n\x08\x61\x63\x63\x65pted\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"G\n\x07HostAck\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x30\n\x05\x65rror\x18\x02 \x01(\x0b\x32!.nemo.relay.worker.v1.WorkerError\"?\n\x0bWorkerError\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x11\n\tretryable\x18\x03 \x01(\x08\"?\n\x0cScopeContext\x12\x16\n\x0escope_stack_id\x18\x01 \x01(\t\x12\x17\n\x0fparent_scope_id\x18\x02 \x01(\t\"\xe5\x01\n\x0f\x45mitMarkRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\x12\x31\n\x05scope\x18\x03 \x01(\x0b\x32\".nemo.relay.worker.v1.ScopeContext\x12\x0c\n\x04name\x18\x04 \x01(\t\x12\x30\n\x04\x64\x61ta\x18\x05 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12\x34\n\x08metadata\x18\x06 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"\xce\x02\n\x10PushScopeRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\x12\x31\n\x05scope\x18\x03 \x01(\x0b\x32\".nemo.relay.worker.v1.ScopeContext\x12\x0c\n\x04name\x18\x04 \x01(\t\x12\x33\n\nscope_type\x18\x05 \x01(\x0e\x32\x1f.nemo.relay.worker.v1.ScopeType\x12\x30\n\x04\x64\x61ta\x18\x06 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12\x34\n\x08metadata\x18\x07 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12\x31\n\x05input\x18\x08 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"^\n\x11PushScopeResponse\x12\x17\n\x0fscope_handle_id\x18\x01 \x01(\t\x12\x30\n\x05\x65rror\x18\x02 \x01(\x0b\x32!.nemo.relay.worker.v1.WorkerError\"\xbf\x01\n\x0fPopScopeRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\x12\x17\n\x0fscope_handle_id\x18\x03 \x01(\t\x12\x32\n\x06output\x18\x04 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\x12\x34\n\x08metadata\x18\x05 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"D\n\x17\x43reateScopeStackRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\"d\n\x18\x43reateScopeStackResponse\x12\x16\n\x0escope_stack_id\x18\x01 \x01(\t\x12\x30\n\x05\x65rror\x18\x02 \x01(\x0b\x32!.nemo.relay.worker.v1.WorkerError\"Z\n\x15\x44ropScopeStackRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\x12\x16\n\x0escope_stack_id\x18\x03 \x01(\t\"\x88\x01\n\x0fToolNextRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\x12\x17\n\x0f\x63ontinuation_id\x18\x03 \x01(\t\x12\x31\n\x05value\x18\x04 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"\x89\x01\n\x0eLlmNextRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\x12\x17\n\x0f\x63ontinuation_id\x18\x03 \x01(\t\x12\x33\n\x07request\x18\x04 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope\"\x8f\x01\n\x14LlmStreamNextRequest\x12\x15\n\ractivation_id\x18\x01 \x01(\t\x12\x12\n\nauth_token\x18\x02 \x01(\t\x12\x17\n\x0f\x63ontinuation_id\x18\x03 \x01(\t\x12\x33\n\x07request\x18\x04 \x01(\x0b\x32\".nemo.relay.worker.v1.JsonEnvelope*\xc8\x03\n\x13RegistrationSurface\x12$\n REGISTRATION_SURFACE_UNSPECIFIED\x10\x00\x12\x0e\n\nSUBSCRIBER\x10\x01\x12#\n\x1fTOOL_SANITIZE_REQUEST_GUARDRAIL\x10\n\x12$\n TOOL_SANITIZE_RESPONSE_GUARDRAIL\x10\x0b\x12(\n$TOOL_CONDITIONAL_EXECUTION_GUARDRAIL\x10\x0c\x12\x1a\n\x16TOOL_REQUEST_INTERCEPT\x10\r\x12\x1c\n\x18TOOL_EXECUTION_INTERCEPT\x10\x0e\x12\"\n\x1eLLM_SANITIZE_REQUEST_GUARDRAIL\x10\x14\x12#\n\x1fLLM_SANITIZE_RESPONSE_GUARDRAIL\x10\x15\x12\'\n#LLM_CONDITIONAL_EXECUTION_GUARDRAIL\x10\x16\x12\x19\n\x15LLM_REQUEST_INTERCEPT\x10\x17\x12\x1b\n\x17LLM_EXECUTION_INTERCEPT\x10\x18\x12\"\n\x1eLLM_STREAM_EXECUTION_INTERCEPT\x10\x19*\xb5\x01\n\tScopeType\x12\x1a\n\x16SCOPE_TYPE_UNSPECIFIED\x10\x00\x12\t\n\x05\x41GENT\x10\x01\x12\x0c\n\x08\x46UNCTION\x10\x02\x12\x08\n\x04TOOL\x10\x03\x12\x07\n\x03LLM\x10\x04\x12\r\n\tRETRIEVER\x10\x05\x12\x0c\n\x08\x45MBEDDER\x10\x06\x12\x0c\n\x08RERANKER\x10\x07\x12\r\n\tGUARDRAIL\x10\x08\x12\r\n\tEVALUATOR\x10\t\x12\n\n\x06\x43USTOM\x10\n\x12\x0b\n\x07UNKNOWN\x10\x0b\x32\xde\x05\n\x0cPluginWorker\x12\\\n\tHandshake\x12&.nemo.relay.worker.v1.HandshakeRequest\x1a\'.nemo.relay.worker.v1.HandshakeResponse\x12S\n\x06Health\x12#.nemo.relay.worker.v1.HealthRequest\x1a$.nemo.relay.worker.v1.HealthResponse\x12Y\n\x08Validate\x12%.nemo.relay.worker.v1.ValidateRequest\x1a&.nemo.relay.worker.v1.ValidateResponse\x12Y\n\x08Register\x12%.nemo.relay.worker.v1.RegisterRequest\x1a&.nemo.relay.worker.v1.RegisterResponse\x12S\n\x06Invoke\x12#.nemo.relay.worker.v1.InvokeRequest\x1a$.nemo.relay.worker.v1.InvokeResponse\x12X\n\x0cInvokeStream\x12#.nemo.relay.worker.v1.InvokeRequest\x1a!.nemo.relay.worker.v1.StreamChunk0\x01\x12\x62\n\x10\x43\x61ncelInvocation\x12-.nemo.relay.worker.v1.CancelInvocationRequest\x1a\x1f.nemo.relay.worker.v1.WorkerAck\x12R\n\x08Shutdown\x12%.nemo.relay.worker.v1.ShutdownRequest\x1a\x1f.nemo.relay.worker.v1.WorkerAck2\xef\x05\n\x10RelayHostRuntime\x12P\n\x08\x45mitMark\x12%.nemo.relay.worker.v1.EmitMarkRequest\x1a\x1d.nemo.relay.worker.v1.HostAck\x12\\\n\tPushScope\x12&.nemo.relay.worker.v1.PushScopeRequest\x1a\'.nemo.relay.worker.v1.PushScopeResponse\x12P\n\x08PopScope\x12%.nemo.relay.worker.v1.PopScopeRequest\x1a\x1d.nemo.relay.worker.v1.HostAck\x12q\n\x10\x43reateScopeStack\x12-.nemo.relay.worker.v1.CreateScopeStackRequest\x1a..nemo.relay.worker.v1.CreateScopeStackResponse\x12\\\n\x0e\x44ropScopeStack\x12+.nemo.relay.worker.v1.DropScopeStackRequest\x1a\x1d.nemo.relay.worker.v1.HostAck\x12S\n\x08ToolNext\x12%.nemo.relay.worker.v1.ToolNextRequest\x1a .nemo.relay.worker.v1.JsonResult\x12Q\n\x07LlmNext\x12$.nemo.relay.worker.v1.LlmNextRequest\x1a .nemo.relay.worker.v1.JsonResult\x12`\n\rLlmStreamNext\x12*.nemo.relay.worker.v1.LlmStreamNextRequest\x1a!.nemo.relay.worker.v1.StreamChunk0\x01\x62\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'plugin_worker_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + DESCRIPTOR._loaded_options = None + _globals['_REGISTRATIONSURFACE']._serialized_start=4914 + _globals['_REGISTRATIONSURFACE']._serialized_end=5370 + _globals['_SCOPETYPE']._serialized_start=5373 + _globals['_SCOPETYPE']._serialized_end=5554 + _globals['_JSONENVELOPE']._serialized_start=45 + _globals['_JSONENVELOPE']._serialized_end=89 + _globals['_HANDSHAKEREQUEST']._serialized_start=92 + _globals['_HANDSHAKEREQUEST']._serialized_end=243 + _globals['_HANDSHAKERESPONSE']._serialized_start=246 + _globals['_HANDSHAKERESPONSE']._serialized_end=523 + _globals['_HEALTHREQUEST']._serialized_start=525 + _globals['_HEALTHREQUEST']._serialized_end=583 + _globals['_HEALTHRESPONSE']._serialized_start=586 + _globals['_HEALTHRESPONSE']._serialized_end=761 + _globals['_VALIDATEREQUEST']._serialized_start=764 + _globals['_VALIDATEREQUEST']._serialized_end=895 + _globals['_VALIDATERESPONSE']._serialized_start=897 + _globals['_VALIDATERESPONSE']._serialized_end=1022 + _globals['_REGISTERREQUEST']._serialized_start=1025 + _globals['_REGISTERREQUEST']._serialized_end=1156 + _globals['_REGISTERRESPONSE']._serialized_start=1158 + _globals['_REGISTERRESPONSE']._serialized_end=1285 + _globals['_REGISTRATION']._serialized_start=1288 + _globals['_REGISTRATION']._serialized_end=1421 + _globals['_INVOKEREQUEST']._serialized_start=1424 + _globals['_INVOKEREQUEST']._serialized_end=1838 + _globals['_TOOLINVOCATION']._serialized_start=1840 + _globals['_TOOLINVOCATION']._serialized_end=1926 + _globals['_LLMINVOCATION']._serialized_start=1929 + _globals['_LLMINVOCATION']._serialized_end=2134 + _globals['_INVOKERESPONSE']._serialized_start=2137 + _globals['_INVOKERESPONSE']._serialized_end=2449 + _globals['_EMPTYRESULT']._serialized_start=2451 + _globals['_EMPTYRESULT']._serialized_end=2464 + _globals['_JSONRESULT']._serialized_start=2466 + _globals['_JSONRESULT']._serialized_end=2579 + _globals['_GUARDRAILRESULT']._serialized_start=2581 + _globals['_GUARDRAILRESULT']._serialized_end=2620 + _globals['_LLMREQUESTINTERCEPTRESULT']._serialized_start=2623 + _globals['_LLMREQUESTINTERCEPTRESULT']._serialized_end=2797 + _globals['_STREAMCHUNK']._serialized_start=2799 + _globals['_STREAMCHUNK']._serialized_end=2925 + _globals['_CANCELINVOCATIONREQUEST']._serialized_start=2927 + _globals['_CANCELINVOCATIONREQUEST']._serialized_end=3034 + _globals['_SHUTDOWNREQUEST']._serialized_start=3036 + _globals['_SHUTDOWNREQUEST']._serialized_end=3112 + _globals['_WORKERACK']._serialized_start=3114 + _globals['_WORKERACK']._serialized_end=3160 + _globals['_HOSTACK']._serialized_start=3162 + _globals['_HOSTACK']._serialized_end=3233 + _globals['_WORKERERROR']._serialized_start=3235 + _globals['_WORKERERROR']._serialized_end=3298 + _globals['_SCOPECONTEXT']._serialized_start=3300 + _globals['_SCOPECONTEXT']._serialized_end=3363 + _globals['_EMITMARKREQUEST']._serialized_start=3366 + _globals['_EMITMARKREQUEST']._serialized_end=3595 + _globals['_PUSHSCOPEREQUEST']._serialized_start=3598 + _globals['_PUSHSCOPEREQUEST']._serialized_end=3932 + _globals['_PUSHSCOPERESPONSE']._serialized_start=3934 + _globals['_PUSHSCOPERESPONSE']._serialized_end=4028 + _globals['_POPSCOPEREQUEST']._serialized_start=4031 + _globals['_POPSCOPEREQUEST']._serialized_end=4222 + _globals['_CREATESCOPESTACKREQUEST']._serialized_start=4224 + _globals['_CREATESCOPESTACKREQUEST']._serialized_end=4292 + _globals['_CREATESCOPESTACKRESPONSE']._serialized_start=4294 + _globals['_CREATESCOPESTACKRESPONSE']._serialized_end=4394 + _globals['_DROPSCOPESTACKREQUEST']._serialized_start=4396 + _globals['_DROPSCOPESTACKREQUEST']._serialized_end=4486 + _globals['_TOOLNEXTREQUEST']._serialized_start=4489 + _globals['_TOOLNEXTREQUEST']._serialized_end=4625 + _globals['_LLMNEXTREQUEST']._serialized_start=4628 + _globals['_LLMNEXTREQUEST']._serialized_end=4765 + _globals['_LLMSTREAMNEXTREQUEST']._serialized_start=4768 + _globals['_LLMSTREAMNEXTREQUEST']._serialized_end=4911 + _globals['_PLUGINWORKER']._serialized_start=5557 + _globals['_PLUGINWORKER']._serialized_end=6291 + _globals['_RELAYHOSTRUNTIME']._serialized_start=6294 + _globals['_RELAYHOSTRUNTIME']._serialized_end=7045 +# @@protoc_insertion_point(module_scope) diff --git a/python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2_grpc.py b/python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2_grpc.py new file mode 100644 index 000000000..0b942b145 --- /dev/null +++ b/python/plugin/src/nemo_relay_plugin/_proto/plugin_worker_pb2_grpc.py @@ -0,0 +1,774 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! +"""Client and server classes corresponding to protobuf-defined services.""" +import grpc +import warnings + +from . import plugin_worker_pb2 as plugin__worker__pb2 + +GRPC_GENERATED_VERSION = '1.81.1' +GRPC_VERSION = grpc.__version__ +_version_not_supported = False + +try: + from grpc._utilities import first_version_is_lower + _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) +except ImportError: + _version_not_supported = True + +if _version_not_supported: + raise RuntimeError( + f'The grpc package installed is at version {GRPC_VERSION},' + + ' but the generated code in plugin_worker_pb2_grpc.py depends on' + + f' grpcio>={GRPC_GENERATED_VERSION}.' + + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' + ) + + +class PluginWorkerStub: + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.Handshake = channel.unary_unary( + '/nemo.relay.worker.v1.PluginWorker/Handshake', + request_serializer=plugin__worker__pb2.HandshakeRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.HandshakeResponse.FromString, + _registered_method=True) + self.Health = channel.unary_unary( + '/nemo.relay.worker.v1.PluginWorker/Health', + request_serializer=plugin__worker__pb2.HealthRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.HealthResponse.FromString, + _registered_method=True) + self.Validate = channel.unary_unary( + '/nemo.relay.worker.v1.PluginWorker/Validate', + request_serializer=plugin__worker__pb2.ValidateRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.ValidateResponse.FromString, + _registered_method=True) + self.Register = channel.unary_unary( + '/nemo.relay.worker.v1.PluginWorker/Register', + request_serializer=plugin__worker__pb2.RegisterRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.RegisterResponse.FromString, + _registered_method=True) + self.Invoke = channel.unary_unary( + '/nemo.relay.worker.v1.PluginWorker/Invoke', + request_serializer=plugin__worker__pb2.InvokeRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.InvokeResponse.FromString, + _registered_method=True) + self.InvokeStream = channel.unary_stream( + '/nemo.relay.worker.v1.PluginWorker/InvokeStream', + request_serializer=plugin__worker__pb2.InvokeRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.StreamChunk.FromString, + _registered_method=True) + self.CancelInvocation = channel.unary_unary( + '/nemo.relay.worker.v1.PluginWorker/CancelInvocation', + request_serializer=plugin__worker__pb2.CancelInvocationRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.WorkerAck.FromString, + _registered_method=True) + self.Shutdown = channel.unary_unary( + '/nemo.relay.worker.v1.PluginWorker/Shutdown', + request_serializer=plugin__worker__pb2.ShutdownRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.WorkerAck.FromString, + _registered_method=True) + + +class PluginWorkerServicer: + """Missing associated documentation comment in .proto file.""" + + def Handshake(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Health(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Validate(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Register(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Invoke(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def InvokeStream(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def CancelInvocation(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def Shutdown(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_PluginWorkerServicer_to_server(servicer, server): + rpc_method_handlers = { + 'Handshake': grpc.unary_unary_rpc_method_handler( + servicer.Handshake, + request_deserializer=plugin__worker__pb2.HandshakeRequest.FromString, + response_serializer=plugin__worker__pb2.HandshakeResponse.SerializeToString, + ), + 'Health': grpc.unary_unary_rpc_method_handler( + servicer.Health, + request_deserializer=plugin__worker__pb2.HealthRequest.FromString, + response_serializer=plugin__worker__pb2.HealthResponse.SerializeToString, + ), + 'Validate': grpc.unary_unary_rpc_method_handler( + servicer.Validate, + request_deserializer=plugin__worker__pb2.ValidateRequest.FromString, + response_serializer=plugin__worker__pb2.ValidateResponse.SerializeToString, + ), + 'Register': grpc.unary_unary_rpc_method_handler( + servicer.Register, + request_deserializer=plugin__worker__pb2.RegisterRequest.FromString, + response_serializer=plugin__worker__pb2.RegisterResponse.SerializeToString, + ), + 'Invoke': grpc.unary_unary_rpc_method_handler( + servicer.Invoke, + request_deserializer=plugin__worker__pb2.InvokeRequest.FromString, + response_serializer=plugin__worker__pb2.InvokeResponse.SerializeToString, + ), + 'InvokeStream': grpc.unary_stream_rpc_method_handler( + servicer.InvokeStream, + request_deserializer=plugin__worker__pb2.InvokeRequest.FromString, + response_serializer=plugin__worker__pb2.StreamChunk.SerializeToString, + ), + 'CancelInvocation': grpc.unary_unary_rpc_method_handler( + servicer.CancelInvocation, + request_deserializer=plugin__worker__pb2.CancelInvocationRequest.FromString, + response_serializer=plugin__worker__pb2.WorkerAck.SerializeToString, + ), + 'Shutdown': grpc.unary_unary_rpc_method_handler( + servicer.Shutdown, + request_deserializer=plugin__worker__pb2.ShutdownRequest.FromString, + response_serializer=plugin__worker__pb2.WorkerAck.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'nemo.relay.worker.v1.PluginWorker', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('nemo.relay.worker.v1.PluginWorker', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class PluginWorker: + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def Handshake(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.PluginWorker/Handshake', + plugin__worker__pb2.HandshakeRequest.SerializeToString, + plugin__worker__pb2.HandshakeResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def Health(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.PluginWorker/Health', + plugin__worker__pb2.HealthRequest.SerializeToString, + plugin__worker__pb2.HealthResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def Validate(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.PluginWorker/Validate', + plugin__worker__pb2.ValidateRequest.SerializeToString, + plugin__worker__pb2.ValidateResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def Register(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.PluginWorker/Register', + plugin__worker__pb2.RegisterRequest.SerializeToString, + plugin__worker__pb2.RegisterResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def Invoke(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.PluginWorker/Invoke', + plugin__worker__pb2.InvokeRequest.SerializeToString, + plugin__worker__pb2.InvokeResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def InvokeStream(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/nemo.relay.worker.v1.PluginWorker/InvokeStream', + plugin__worker__pb2.InvokeRequest.SerializeToString, + plugin__worker__pb2.StreamChunk.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def CancelInvocation(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.PluginWorker/CancelInvocation', + plugin__worker__pb2.CancelInvocationRequest.SerializeToString, + plugin__worker__pb2.WorkerAck.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def Shutdown(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.PluginWorker/Shutdown', + plugin__worker__pb2.ShutdownRequest.SerializeToString, + plugin__worker__pb2.WorkerAck.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + +class RelayHostRuntimeStub: + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.EmitMark = channel.unary_unary( + '/nemo.relay.worker.v1.RelayHostRuntime/EmitMark', + request_serializer=plugin__worker__pb2.EmitMarkRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.HostAck.FromString, + _registered_method=True) + self.PushScope = channel.unary_unary( + '/nemo.relay.worker.v1.RelayHostRuntime/PushScope', + request_serializer=plugin__worker__pb2.PushScopeRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.PushScopeResponse.FromString, + _registered_method=True) + self.PopScope = channel.unary_unary( + '/nemo.relay.worker.v1.RelayHostRuntime/PopScope', + request_serializer=plugin__worker__pb2.PopScopeRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.HostAck.FromString, + _registered_method=True) + self.CreateScopeStack = channel.unary_unary( + '/nemo.relay.worker.v1.RelayHostRuntime/CreateScopeStack', + request_serializer=plugin__worker__pb2.CreateScopeStackRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.CreateScopeStackResponse.FromString, + _registered_method=True) + self.DropScopeStack = channel.unary_unary( + '/nemo.relay.worker.v1.RelayHostRuntime/DropScopeStack', + request_serializer=plugin__worker__pb2.DropScopeStackRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.HostAck.FromString, + _registered_method=True) + self.ToolNext = channel.unary_unary( + '/nemo.relay.worker.v1.RelayHostRuntime/ToolNext', + request_serializer=plugin__worker__pb2.ToolNextRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.JsonResult.FromString, + _registered_method=True) + self.LlmNext = channel.unary_unary( + '/nemo.relay.worker.v1.RelayHostRuntime/LlmNext', + request_serializer=plugin__worker__pb2.LlmNextRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.JsonResult.FromString, + _registered_method=True) + self.LlmStreamNext = channel.unary_stream( + '/nemo.relay.worker.v1.RelayHostRuntime/LlmStreamNext', + request_serializer=plugin__worker__pb2.LlmStreamNextRequest.SerializeToString, + response_deserializer=plugin__worker__pb2.StreamChunk.FromString, + _registered_method=True) + + +class RelayHostRuntimeServicer: + """Missing associated documentation comment in .proto file.""" + + def EmitMark(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def PushScope(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def PopScope(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def CreateScopeStack(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def DropScopeStack(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def ToolNext(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def LlmNext(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + def LlmStreamNext(self, request, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_RelayHostRuntimeServicer_to_server(servicer, server): + rpc_method_handlers = { + 'EmitMark': grpc.unary_unary_rpc_method_handler( + servicer.EmitMark, + request_deserializer=plugin__worker__pb2.EmitMarkRequest.FromString, + response_serializer=plugin__worker__pb2.HostAck.SerializeToString, + ), + 'PushScope': grpc.unary_unary_rpc_method_handler( + servicer.PushScope, + request_deserializer=plugin__worker__pb2.PushScopeRequest.FromString, + response_serializer=plugin__worker__pb2.PushScopeResponse.SerializeToString, + ), + 'PopScope': grpc.unary_unary_rpc_method_handler( + servicer.PopScope, + request_deserializer=plugin__worker__pb2.PopScopeRequest.FromString, + response_serializer=plugin__worker__pb2.HostAck.SerializeToString, + ), + 'CreateScopeStack': grpc.unary_unary_rpc_method_handler( + servicer.CreateScopeStack, + request_deserializer=plugin__worker__pb2.CreateScopeStackRequest.FromString, + response_serializer=plugin__worker__pb2.CreateScopeStackResponse.SerializeToString, + ), + 'DropScopeStack': grpc.unary_unary_rpc_method_handler( + servicer.DropScopeStack, + request_deserializer=plugin__worker__pb2.DropScopeStackRequest.FromString, + response_serializer=plugin__worker__pb2.HostAck.SerializeToString, + ), + 'ToolNext': grpc.unary_unary_rpc_method_handler( + servicer.ToolNext, + request_deserializer=plugin__worker__pb2.ToolNextRequest.FromString, + response_serializer=plugin__worker__pb2.JsonResult.SerializeToString, + ), + 'LlmNext': grpc.unary_unary_rpc_method_handler( + servicer.LlmNext, + request_deserializer=plugin__worker__pb2.LlmNextRequest.FromString, + response_serializer=plugin__worker__pb2.JsonResult.SerializeToString, + ), + 'LlmStreamNext': grpc.unary_stream_rpc_method_handler( + servicer.LlmStreamNext, + request_deserializer=plugin__worker__pb2.LlmStreamNextRequest.FromString, + response_serializer=plugin__worker__pb2.StreamChunk.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'nemo.relay.worker.v1.RelayHostRuntime', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('nemo.relay.worker.v1.RelayHostRuntime', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class RelayHostRuntime: + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def EmitMark(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.RelayHostRuntime/EmitMark', + plugin__worker__pb2.EmitMarkRequest.SerializeToString, + plugin__worker__pb2.HostAck.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def PushScope(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.RelayHostRuntime/PushScope', + plugin__worker__pb2.PushScopeRequest.SerializeToString, + plugin__worker__pb2.PushScopeResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def PopScope(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.RelayHostRuntime/PopScope', + plugin__worker__pb2.PopScopeRequest.SerializeToString, + plugin__worker__pb2.HostAck.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def CreateScopeStack(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.RelayHostRuntime/CreateScopeStack', + plugin__worker__pb2.CreateScopeStackRequest.SerializeToString, + plugin__worker__pb2.CreateScopeStackResponse.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def DropScopeStack(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.RelayHostRuntime/DropScopeStack', + plugin__worker__pb2.DropScopeStackRequest.SerializeToString, + plugin__worker__pb2.HostAck.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def ToolNext(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.RelayHostRuntime/ToolNext', + plugin__worker__pb2.ToolNextRequest.SerializeToString, + plugin__worker__pb2.JsonResult.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def LlmNext(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_unary( + request, + target, + '/nemo.relay.worker.v1.RelayHostRuntime/LlmNext', + plugin__worker__pb2.LlmNextRequest.SerializeToString, + plugin__worker__pb2.JsonResult.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + @staticmethod + def LlmStreamNext(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/nemo.relay.worker.v1.RelayHostRuntime/LlmStreamNext', + plugin__worker__pb2.LlmStreamNextRequest.SerializeToString, + plugin__worker__pb2.StreamChunk.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) diff --git a/python/plugin/src/nemo_relay_plugin/py.typed b/python/plugin/src/nemo_relay_plugin/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/python/tests/plugin/test_python_worker_example.py b/python/tests/plugin/test_python_worker_example.py new file mode 100644 index 000000000..1056882bb --- /dev/null +++ b/python/tests/plugin/test_python_worker_example.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the Python gRPC worker plugin example.""" + +from __future__ import annotations + +import runpy +from pathlib import Path +from typing import Any + +import pytest + + +@pytest.fixture(scope="module") +def example() -> dict[str, Any]: + path = Path(__file__).parents[3] / "examples/python-grpc-worker-plugin/worker.py" + return runpy.run_path(str(path)) + + +def test_tag_json_rejects_non_objects_and_collisions(example: dict[str, Any]) -> None: + tag_json = example["_tag_json"] + + assert tag_json({"query": "relay"}, "demo") == {"query": "relay", "demo": True} + with pytest.raises(TypeError, match="requires a JSON object"): + tag_json(["relay"], "demo") + with pytest.raises(TypeError, match="requires a JSON object"): + tag_json(None, "demo") + with pytest.raises(ValueError, match="already contains configured tag"): + tag_json({"demo": False}, "demo") + + +def test_example_validates_tag_configuration(example: dict[str, Any]) -> None: + plugin = example["ExamplePythonWorker"]() + + assert plugin.validate({"tag": "demo"}) == [] + diagnostics = plugin.validate({"tag": 42}) + assert len(diagnostics) == 1 + assert diagnostics[0].code == "examples.python_grpc_worker.invalid_tag" + + +async def test_example_register_propagates_configured_tag(example: dict[str, Any]) -> None: + marks: list[tuple[str, Any]] = [] + + class Runtime: + async def emit_mark(self, name: str, data: Any) -> None: + marks.append((name, data)) + + class Context: + runtime = Runtime() + callback: Any = None + + def register_tool_request_intercept(self, name: str, callback: Any) -> None: + assert name == "tag_tool_request" + self.callback = callback + + context = Context() + plugin = example["ExamplePythonWorker"]() + plugin.register(context, {"tag": "demo"}) + + assert await context.callback("lookup", {"query": "relay"}) == {"query": "relay", "demo": True} + assert marks == [ + ( + "examples.python_grpc_worker.tool_request", + {"tool_name": "lookup", "source": "python-grpc-worker", "tag": "demo"}, + ) + ] diff --git a/python/tests/plugin/test_worker_sdk.py b/python/tests/plugin/test_worker_sdk.py new file mode 100644 index 000000000..9176d5fcb --- /dev/null +++ b/python/tests/plugin/test_worker_sdk.py @@ -0,0 +1,1382 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the Python dynamic worker plugin SDK.""" + +from __future__ import annotations + +import asyncio +import contextlib +import json +import os +from collections.abc import AsyncIterator +from typing import Any + +import pytest + +if os.environ.get("NEMO_RELAY_SKIP_PYTHON_PLUGIN_TESTS") == "1": + pytest.skip("grpcio is unavailable for Python plugin SDK tests on this runner", allow_module_level=True) + +grpc = pytest.importorskip("grpc") + +from nemo_relay_plugin import ( # noqa: E402 + ConfigDiagnostic, + DiagnosticLevel, + Json, + PluginContext, + PluginRuntime, + ScopeType, + ToolNext, + WorkerPlugin, + WorkerSdkError, + serve_plugin, +) +from nemo_relay_plugin import _api as plugin_api # noqa: E402 +from nemo_relay_plugin._api import ( # noqa: E402 + ANNOTATED_LLM_REQUEST_SCHEMA, + EVENT_SCHEMA, + JSON_SCHEMA, + LLM_REQUEST_SCHEMA, + WORKER_PROTOCOL, + _announced_worker_endpoint, + _grpc_target, + _json_envelope, + _required_env, + _unlink_unix_socket, + _WorkerService, + pb, + pb_grpc, +) + +ACTIVATION_ID = "act" +AUTH_TOKEN = "token" + + +class GrpcAbort(Exception): + def __init__(self, code: object, details: str) -> None: + super().__init__(f"{code}: {details}") + self.code = code + self.details = details + + +class AbortContext: + async def abort(self, code: object, details: str) -> None: + raise GrpcAbort(code, details) + + +class RecordingHostStub: + def __init__(self) -> None: + self.requests: list[Any] = [] + self.failures: dict[str, str] = {} + + async def EmitMark(self, request: Any) -> Any: + self.requests.append(request) + return self._host_ack("EmitMark") + + async def CreateScopeStack(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("CreateScopeStack") == "error": + return pb.CreateScopeStackResponse(error=_worker_error("CreateScopeStack failed")) + return pb.CreateScopeStackResponse(scope_stack_id="stack-1") + + async def DropScopeStack(self, request: Any) -> Any: + self.requests.append(request) + return self._host_ack("DropScopeStack") + + async def PushScope(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("PushScope") == "error": + return pb.PushScopeResponse(error=_worker_error("PushScope failed")) + return pb.PushScopeResponse(scope_handle_id="scope-1") + + async def PopScope(self, request: Any) -> Any: + self.requests.append(request) + return self._host_ack("PopScope") + + async def ToolNext(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("ToolNext") == "error": + return pb.JsonResult(error=_worker_error("ToolNext failed")) + value = json.loads(request.value.json.decode("utf-8")) + return pb.JsonResult(value=_json_envelope(JSON_SCHEMA, {"next_tool": value})) + + async def LlmNext(self, request: Any) -> Any: + self.requests.append(request) + if self.failures.get("LlmNext") == "error": + return pb.JsonResult(error=_worker_error("LlmNext failed")) + value = json.loads(request.request.json.decode("utf-8")) + return pb.JsonResult(value=_json_envelope(JSON_SCHEMA, {"next_llm": value})) + + def LlmStreamNext(self, request: Any) -> AsyncIterator[Any]: + self.requests.append(request) + + async def stream() -> AsyncIterator[Any]: + failure = self.failures.get("LlmStreamNext") + if failure == "error": + yield pb.StreamChunk(error=_worker_error("LlmStreamNext failed")) + return + if failure == "empty": + yield pb.StreamChunk() + return + value = json.loads(request.request.json.decode("utf-8")) + yield pb.StreamChunk(value=_json_envelope(JSON_SCHEMA, {"next_stream": value})) + + return stream() + + def _host_ack(self, method: str) -> Any: + failure = self.failures.get(method) + if failure == "empty": + return pb.HostAck(ok=False) + if failure == "error": + return pb.HostAck(ok=False, error=_worker_error(f"{method} failed")) + return pb.HostAck(ok=True) + + +class AllSurfacesPlugin(WorkerPlugin): + plugin_id = "tests.python_worker" + + def validate(self, config: Json) -> list[ConfigDiagnostic | dict[str, Any]]: + if isinstance(config, dict) and config.get("warn"): + return [ + ConfigDiagnostic( + level=DiagnosticLevel.WARNING, + code="tests.warn", + message="warning requested", + ) + ] + return [] + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + async def subscriber(event: Json) -> None: + await ctx.runtime.emit_mark("tests.subscriber", event) + + def tool_sanitize(name: str, value: Json) -> Json: + return _tag(value, f"sanitize_{name}") + + def tool_block(name: str, value: Json) -> str | None: + del name, value + return "tool blocked" + + async def tool_request(name: str, value: Json) -> Json: + return _tag(value, f"request_{name}") + + async def tool_execution(name: str, value: Json, next_call: ToolNext) -> Json: + result = await next_call.call(_tag(value, f"execute_{name}")) + return _tag(result, "tool_execution") + + def llm_sanitize_request(request: Json) -> Json: + return _tag_llm_request(request, "llm_sanitize_request") + + async def llm_sanitize_response(response: Json) -> Json: + return _tag(response, "llm_sanitize_response") + + def llm_block(request: Json) -> str | None: + del request + return "llm blocked" + + def llm_request(name: str, request: Json, annotated: Json | None) -> tuple[Json, Json]: + del name + return _tag_llm_request(request, "llm_request"), _tag(annotated or {}, "annotated") + + async def llm_execution(name: str, request: Json, next_call: Any) -> Json: + result = await next_call.call(_tag_llm_request(request, f"llm_execute_{name}")) + return _tag(result, "llm_execution") + + async def llm_stream_execution(name: str, request: Json, next_call: Any) -> AsyncIterator[Json]: + stream = next_call.call(_tag_llm_request(request, f"llm_stream_{name}")) + async for chunk in stream: + yield _tag(chunk, "llm_stream_execution") + + ctx.register_subscriber("subscriber", subscriber) + ctx.register_tool_sanitize_request_guardrail("tool_sanitize", tool_sanitize, priority=1) + ctx.register_tool_sanitize_response_guardrail("tool_sanitize", tool_sanitize, priority=2) + ctx.register_tool_conditional_execution_guardrail("tool_conditional", tool_block, priority=3) + ctx.register_tool_request_intercept("tool_request", tool_request, priority=4, break_chain=True) + ctx.register_tool_execution_intercept("tool_execution", tool_execution, priority=5) + ctx.register_llm_sanitize_request_guardrail("llm_sanitize_request", llm_sanitize_request, priority=6) + ctx.register_llm_sanitize_response_guardrail("llm_sanitize_response", llm_sanitize_response, priority=7) + ctx.register_llm_conditional_execution_guardrail("llm_conditional", llm_block, priority=8) + ctx.register_llm_request_intercept("llm_request", llm_request, priority=9, break_chain=True) + ctx.register_llm_execution_intercept("llm_execution", llm_execution, priority=10) + ctx.register_llm_stream_execution_intercept("llm_stream_execution", llm_stream_execution, priority=11) + + +@pytest.fixture(name="host_stub") +def host_stub_fixture() -> RecordingHostStub: + return RecordingHostStub() + + +@pytest.fixture(name="service") +def service_fixture(host_stub: RecordingHostStub) -> _WorkerService: + return _service(AllSurfacesPlugin(), host_stub) + + +def test_generated_proto_matches_worker_contract() -> None: + methods = {method.name for method in pb.DESCRIPTOR.services_by_name["PluginWorker"].methods} + assert methods == { + "Handshake", + "Health", + "Validate", + "Register", + "Invoke", + "InvokeStream", + "CancelInvocation", + "Shutdown", + } + assert pb.InvokeRequest.DESCRIPTOR.fields_by_name["auth_token"].number == 7 + assert pb.HealthRequest.DESCRIPTOR.fields_by_name["activation_id"].number == 1 + assert pb.HealthRequest.DESCRIPTOR.fields_by_name["auth_token"].number == 2 + assert pb.SUBSCRIBER == 1 + assert pb.TOOL_SANITIZE_REQUEST_GUARDRAIL == 10 + assert pb.LLM_STREAM_EXECUTION_INTERCEPT == 25 + assert pb.CUSTOM == 10 + + +async def test_health_handshake_validate_register_and_all_surfaces(service: _WorkerService) -> None: + health = await service.Health(pb.HealthRequest(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN), AbortContext()) + assert health.ok + assert health.plugin_id == "tests.python_worker" + assert health.worker_protocol == WORKER_PROTOCOL + assert health.sdk_name == "nemo-relay-plugin" + assert health.runtime_name == "python" + + handshake = await service.Handshake(_handshake_request(), AbortContext()) + assert handshake.plugin_id == "tests.python_worker" + assert handshake.plugin_kind == "tests.python_worker" + assert handshake.worker_protocol == WORKER_PROTOCOL + assert set(handshake.supported_surfaces) == set(_all_expected_surfaces()) + + validate = await service.Validate( + pb.ValidateRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + auth_token=AUTH_TOKEN, + config=_json_envelope(JSON_SCHEMA, {"warn": True}), + ), + AbortContext(), + ) + diagnostics = _envelope_value(validate.diagnostics) + assert diagnostics == [{"level": "warning", "code": "tests.warn", "message": "warning requested"}] + + register = await _register(service) + registrations = [ + (registration.local_name, registration.surface, registration.priority, registration.break_chain) + for registration in register.registrations + ] + assert registrations == [ + ("subscriber", pb.SUBSCRIBER, 0, False), + ("tool_sanitize", pb.TOOL_SANITIZE_REQUEST_GUARDRAIL, 1, False), + ("tool_sanitize", pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL, 2, False), + ("tool_conditional", pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, 3, False), + ("tool_request", pb.TOOL_REQUEST_INTERCEPT, 4, True), + ("tool_execution", pb.TOOL_EXECUTION_INTERCEPT, 5, False), + ("llm_sanitize_request", pb.LLM_SANITIZE_REQUEST_GUARDRAIL, 6, False), + ("llm_sanitize_response", pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, 7, False), + ("llm_conditional", pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, 8, False), + ("llm_request", pb.LLM_REQUEST_INTERCEPT, 9, True), + ("llm_execution", pb.LLM_EXECUTION_INTERCEPT, 10, False), + ("llm_stream_execution", pb.LLM_STREAM_EXECUTION_INTERCEPT, 11, False), + ] + + +@pytest.mark.parametrize( + ("rpc_name", "request_factory", "streaming"), + [ + ("Handshake", lambda: _handshake_request(), False), + ("Health", lambda: pb.HealthRequest(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN), False), + ( + "Validate", + lambda: pb.ValidateRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + auth_token=AUTH_TOKEN, + config=_json_envelope(JSON_SCHEMA, {}), + ), + False, + ), + ( + "Register", + lambda: pb.RegisterRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + auth_token=AUTH_TOKEN, + config=_json_envelope(JSON_SCHEMA, {}), + ), + False, + ), + ("Invoke", lambda: _tool_request("missing", pb.TOOL_REQUEST_INTERCEPT, {}), False), + ( + "InvokeStream", + lambda: _invoke_request( + "missing", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + True, + ), + ( + "CancelInvocation", + lambda: pb.CancelInvocationRequest( + activation_id=ACTIVATION_ID, + invocation_id="invoke-1", + auth_token=AUTH_TOKEN, + reason="test", + ), + False, + ), + ( + "Shutdown", + lambda: pb.ShutdownRequest(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN, reason="test"), + False, + ), + ], +) +async def test_auth_and_activation_failures_for_every_rpc( + service: _WorkerService, + rpc_name: str, + request_factory: Any, + streaming: bool, +) -> None: + for field in ("activation_id", "auth_token"): + request = request_factory() + setattr(request, field, "wrong") + with pytest.raises(GrpcAbort) as exc_info: + result = getattr(service, rpc_name)(request, AbortContext()) + if streaming: + async for _chunk in result: + pass + else: + await result + assert exc_info.value.code == grpc.StatusCode.PERMISSION_DENIED + assert field.split("_")[0] in exc_info.value.details + + +async def test_validate_and_register_decode_errors_are_grpc_protocol_errors(service: _WorkerService) -> None: + bad_config = _json_envelope(JSON_SCHEMA, {}) + bad_config.json = b"{" + + with pytest.raises(GrpcAbort) as validate_error: + await service.Validate( + pb.ValidateRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + auth_token=AUTH_TOKEN, + config=bad_config, + ), + AbortContext(), + ) + assert validate_error.value.code == grpc.StatusCode.INVALID_ARGUMENT + + with pytest.raises(GrpcAbort) as register_error: + await service.Register( + pb.RegisterRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + auth_token=AUTH_TOKEN, + config=bad_config, + ), + AbortContext(), + ) + assert register_error.value.code == grpc.StatusCode.INVALID_ARGUMENT + + wrong_schema = _json_envelope(EVENT_SCHEMA, {}) + with pytest.raises(GrpcAbort) as schema_error: + await service.Validate( + pb.ValidateRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + auth_token=AUTH_TOKEN, + config=wrong_schema, + ), + AbortContext(), + ) + assert schema_error.value.code == grpc.StatusCode.INVALID_ARGUMENT + assert "expected 'nemo.relay.Json@1'" in schema_error.value.details + + +async def test_base_plugin_defaults_context_errors_and_plugin_id_validation() -> None: + base = WorkerPlugin() + assert base.validate({"unused": True}) == [] + with pytest.raises(NotImplementedError): + base.register(PluginContext(), {}) + with pytest.raises(WorkerSdkError, match="no runtime handle"): + _ = PluginContext().runtime + + class CallablePluginId(WorkerPlugin): + allows_multiple_components = True + + def plugin_id(self) -> str: + return "tests.callable_id" + + def register(self, ctx: PluginContext, config: Json) -> None: + del ctx, config + + class InvalidPluginId(WorkerPlugin): + plugin_id = "" + + def register(self, ctx: PluginContext, config: Json) -> None: + del ctx, config + + callable_service = _service(CallablePluginId(), RecordingHostStub()) + health = await callable_service.Health( + pb.HealthRequest(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN), + AbortContext(), + ) + assert health.plugin_id == "tests.callable_id" + + invalid_service = _service(InvalidPluginId(), RecordingHostStub()) + with pytest.raises(WorkerSdkError, match="plugin_id"): + await invalid_service.Health( + pb.HealthRequest(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN), + AbortContext(), + ) + + +async def test_validate_accepts_missing_config_and_dict_diagnostics() -> None: + class DictDiagnosticPlugin(WorkerPlugin): + plugin_id = "tests.dict_diagnostic" + + def validate(self, config: Json) -> list[ConfigDiagnostic | dict[str, Any]]: + assert config is None + return [{"level": "error", "code": "dict.diag", "message": "dict diagnostic"}] + + def register(self, ctx: PluginContext, config: Json) -> None: + del ctx, config + + service = _service(DictDiagnosticPlugin(), RecordingHostStub()) + response = await service.Validate( + pb.ValidateRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.dict_diagnostic", + auth_token=AUTH_TOKEN, + ), + AbortContext(), + ) + assert _envelope_value(response.diagnostics) == [ + {"level": "error", "code": "dict.diag", "message": "dict diagnostic"} + ] + + +async def test_validate_register_and_invoke_callback_errors_are_structured() -> None: + class FailingValidatePlugin(WorkerPlugin): + plugin_id = "tests.failing_validate" + + def validate(self, config: Json) -> list[ConfigDiagnostic | dict[str, Any]]: + del config + raise RuntimeError("validate boom") + + def register(self, ctx: PluginContext, config: Json) -> None: + del ctx, config + + class FailingRegisterPlugin(WorkerPlugin): + plugin_id = "tests.failing_register" + + def register(self, ctx: PluginContext, config: Json) -> None: + del ctx, config + raise RuntimeError("register boom") + + class FailingInvokePlugin(WorkerPlugin): + plugin_id = "tests.failing_invoke" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + def fail(tool_name: str, args: Json) -> Json: + del tool_name, args + raise RuntimeError("invoke boom") + + ctx.register_tool_request_intercept("fail", fail) + + validate_service = _service(FailingValidatePlugin(), RecordingHostStub()) + validate = await validate_service.Validate(_validate_request(), AbortContext()) + assert validate.HasField("error") + assert "validate boom" in validate.error.message + + register_service = _service(FailingRegisterPlugin(), RecordingHostStub()) + register = await register_service.Register(_register_request(), AbortContext()) + assert register.HasField("error") + assert "register boom" in register.error.message + + invoke_service = _service(FailingInvokePlugin(), RecordingHostStub()) + await _register(invoke_service) + response = await invoke_service.Invoke(_tool_request("fail", pb.TOOL_REQUEST_INTERCEPT, {}), AbortContext()) + assert response.WhichOneof("result") == "error" + assert "invoke boom" in response.error.message + + +async def test_unary_invoke_success_paths(service: _WorkerService, host_stub: RecordingHostStub) -> None: + await _register(service) + + subscriber = await service.Invoke( + _invoke_request( + "subscriber", + pb.SUBSCRIBER, + event=_json_envelope(EVENT_SCHEMA, {"name": "event"}), + scope=pb.ScopeContext(scope_stack_id="invoke-stack", parent_scope_id="parent-scope"), + ), + AbortContext(), + ) + assert subscriber.WhichOneof("result") == "empty" + mark_request = _last_request(host_stub, pb.EmitMarkRequest) + assert mark_request.name == "tests.subscriber" + assert mark_request.scope.scope_stack_id == "invoke-stack" + assert mark_request.scope.parent_scope_id == "parent-scope" + + tool_sanitize_request = await _invoke_json_async(service, "tool_sanitize", pb.TOOL_SANITIZE_REQUEST_GUARDRAIL) + assert tool_sanitize_request["tag"] == "sanitize_lookup" + tool_sanitize_response = await _invoke_json_async(service, "tool_sanitize", pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL) + assert tool_sanitize_response["tag"] == "sanitize_lookup" + + tool_conditional = await service.Invoke( + _tool_request("tool_conditional", pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, {"query": "relay"}), + AbortContext(), + ) + assert tool_conditional.guardrail.block_reason == "tool blocked" + + tool_request = await _invoke_json_async(service, "tool_request", pb.TOOL_REQUEST_INTERCEPT) + assert tool_request["tag"] == "request_lookup" + tool_execution = await _invoke_json_async(service, "tool_execution", pb.TOOL_EXECUTION_INTERCEPT) + assert tool_execution["tag"] == "tool_execution" + assert tool_execution["next_tool"]["tag"] == "execute_lookup" + + llm_sanitize_request = await _invoke_json_async( + service, + "llm_sanitize_request", + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + payload=_llm_payload(request={"content": {"prompt": "hello"}}), + ) + assert llm_sanitize_request["content"]["llm_sanitize_request"] + + llm_sanitize_response = await _invoke_json_async( + service, + "llm_sanitize_response", + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + payload=_llm_payload(response={"answer": "hello"}), + ) + assert llm_sanitize_response["tag"] == "llm_sanitize_response" + + llm_conditional = await service.Invoke( + _invoke_request( + "llm_conditional", + pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + assert llm_conditional.guardrail.block_reason == "llm blocked" + + llm_request = await service.Invoke( + _invoke_request( + "llm_request", + pb.LLM_REQUEST_INTERCEPT, + llm=_llm_payload( + request={"content": {"prompt": "hello"}}, + annotated={"messages": [], "extra": {"before": True}}, + ), + ), + AbortContext(), + ) + assert _envelope_value(llm_request.llm_request.request)["content"]["llm_request"] + assert _envelope_value(llm_request.llm_request.annotated_request)["tag"] == "annotated" + assert llm_request.llm_request.has_annotated_request + + llm_execution = await _invoke_json_async( + service, + "llm_execution", + pb.LLM_EXECUTION_INTERCEPT, + payload=_llm_payload(model_name="gpt-test", request={"content": {"prompt": "hello"}}), + ) + assert llm_execution["tag"] == "llm_execution" + assert llm_execution["next_llm"]["content"]["llm_execute_gpt-test"] + + +async def test_unary_invoke_failure_paths(service: _WorkerService) -> None: + await _register(service) + + invalid = _tool_request("tool_request", pb.TOOL_REQUEST_INTERCEPT, {}) + invalid.tool.value.json = b"{" + invalid_payload = await service.Invoke(invalid, AbortContext()) + assert invalid_payload.WhichOneof("result") == "error" + assert "Expecting" in invalid_payload.error.message + + missing_handler = await service.Invoke(_tool_request("missing", pb.TOOL_REQUEST_INTERCEPT, {}), AbortContext()) + assert missing_handler.WhichOneof("result") == "error" + assert "not registered" in missing_handler.error.message + + unsupported = await service.Invoke( + _tool_request("tool_request", pb.REGISTRATION_SURFACE_UNSPECIFIED, {}), + AbortContext(), + ) + assert unsupported.WhichOneof("result") == "error" + assert "unsupported registration surface" in unsupported.error.message + + missing_event = await service.Invoke(_invoke_request("subscriber", pb.SUBSCRIBER), AbortContext()) + assert missing_event.WhichOneof("result") == "error" + assert "event is missing" in missing_event.error.message + + +@pytest.mark.parametrize( + ("request_factory", "expected_message"), + [ + ( + lambda: _invoke_request( + "subscriber", + pb.SUBSCRIBER, + event=_json_envelope(JSON_SCHEMA, {"name": "event"}), + ), + "expected 'nemo.relay.Event@1'", + ), + ( + lambda: _tool_request("tool_request", pb.TOOL_REQUEST_INTERCEPT, {}), + "expected 'nemo.relay.Json@1'", + ), + ( + lambda: _invoke_request( + "llm_sanitize_request", + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + llm=_llm_payload(request={"content": {}}), + ), + "expected 'nemo.relay.LlmRequest@1'", + ), + ( + lambda: _invoke_request( + "llm_conditional", + pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, + llm=_llm_payload(request={"content": {}}), + ), + "expected 'nemo.relay.LlmRequest@1'", + ), + ( + lambda: _invoke_request( + "llm_execution", + pb.LLM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {}}), + ), + "expected 'nemo.relay.LlmRequest@1'", + ), + ( + lambda: _invoke_request( + "llm_request", + pb.LLM_REQUEST_INTERCEPT, + llm=_llm_payload(request={"content": {}}, annotated={"messages": []}), + ), + "expected 'nemo.relay.AnnotatedLlmRequest@1'", + ), + ], +) +async def test_invoke_rejects_mismatched_envelope_schemas( + service: _WorkerService, + request_factory: Any, + expected_message: str, +) -> None: + await _register(service) + request = request_factory() + if request.surface == pb.TOOL_REQUEST_INTERCEPT: + request.tool.value.schema = EVENT_SCHEMA + elif request.surface in { + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, + pb.LLM_EXECUTION_INTERCEPT, + }: + request.llm.request.schema = JSON_SCHEMA + elif request.surface == pb.LLM_REQUEST_INTERCEPT: + request.llm.annotated_request.schema = JSON_SCHEMA + + response = await service.Invoke(request, AbortContext()) + assert response.WhichOneof("result") == "error" + assert expected_message in response.error.message + + +async def test_invoke_stream_rejects_mismatched_llm_request_schema(service: _WorkerService) -> None: + await _register(service) + request = _invoke_request( + "llm_stream_execution", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ) + request.llm.request.schema = JSON_SCHEMA + + chunks = [chunk async for chunk in service.InvokeStream(request, AbortContext())] + assert len(chunks) == 1 + assert "expected 'nemo.relay.LlmRequest@1'" in chunks[0].error.message + + +async def test_llm_request_intercept_can_return_request_without_annotation() -> None: + class RequestOnlyPlugin(WorkerPlugin): + plugin_id = "tests.request_only" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + def llm_request(name: str, request: Json, annotated: Json | None) -> Json: + del name, annotated + return _tag_llm_request(request, "request_only") + + ctx.register_llm_request_intercept("request_only", llm_request) + + service = _service(RequestOnlyPlugin(), RecordingHostStub()) + await _register(service) + response = await service.Invoke( + _invoke_request( + "request_only", + pb.LLM_REQUEST_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + assert _envelope_value(response.llm_request.request)["content"]["request_only"] + assert not response.llm_request.has_annotated_request + + +async def test_stream_invoke_success_and_failures(service: _WorkerService, host_stub: RecordingHostStub) -> None: + await _register(service) + + chunks = [ + chunk + async for chunk in service.InvokeStream( + _invoke_request( + "llm_stream_execution", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(model_name="gpt-test", request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + ] + assert [_stream_value(chunk)["tag"] for chunk in chunks] == ["llm_stream_execution"] + assert _stream_value(chunks[0])["next_stream"]["content"]["llm_stream_gpt-test"] + + wrong_surface = [ + chunk + async for chunk in service.InvokeStream( + _tool_request("tool_request", pb.TOOL_REQUEST_INTERCEPT, {}), + AbortContext(), + ) + ] + assert "only supports LLM stream" in wrong_surface[0].error.message + + missing_handler = [ + chunk + async for chunk in service.InvokeStream( + _invoke_request( + "missing", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + ] + assert "not registered" in missing_handler[0].error.message + + missing_payload = [ + chunk + async for chunk in service.InvokeStream( + _invoke_request("llm_stream_execution", pb.LLM_STREAM_EXECUTION_INTERCEPT), + AbortContext(), + ) + ] + assert "expected llm payload" in missing_payload[0].error.message + + host_stub.failures["LlmStreamNext"] = "error" + host_error = [ + chunk + async for chunk in service.InvokeStream( + _invoke_request( + "llm_stream_execution", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + ] + assert "LlmStreamNext failed" in host_error[0].error.message + + host_stub.failures["LlmStreamNext"] = "empty" + empty_chunk = [ + chunk + async for chunk in service.InvokeStream( + _invoke_request( + "llm_stream_execution", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + ] + assert "stream chunk is empty" in empty_chunk[0].error.message + + +async def test_stream_callback_exception_is_structured() -> None: + class FailingStreamPlugin(WorkerPlugin): + plugin_id = "tests.stream_fail" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + async def fail(name: str, request: Json, next_call: Any) -> AsyncIterator[Json]: + del name, request, next_call + raise RuntimeError("stream boom") + yield {} + + ctx.register_llm_stream_execution_intercept("fail_stream", fail) + + service = _service(FailingStreamPlugin(), RecordingHostStub()) + await _register(service) + chunks = [ + chunk + async for chunk in service.InvokeStream( + _invoke_request( + "fail_stream", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + ] + assert "stream boom" in chunks[0].error.message + + +async def test_stream_callback_can_return_sync_iterable() -> None: + class SyncStreamPlugin(WorkerPlugin): + plugin_id = "tests.sync_stream" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + def stream(name: str, request: Json, next_call: Any) -> list[Json]: + del name, request, next_call + return [{"sync": True}] + + ctx.register_llm_stream_execution_intercept("sync_stream", stream) + + service = _service(SyncStreamPlugin(), RecordingHostStub()) + await _register(service) + chunks = [ + chunk + async for chunk in service.InvokeStream( + _invoke_request( + "sync_stream", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + ] + assert _stream_value(chunks[0]) == {"sync": True} + + +@pytest.mark.parametrize("invalid_stream", [{"chunk": True}, "chunk", b"chunk", 42]) +async def test_stream_callback_rejects_scalar_and_mapping_results(invalid_stream: Any) -> None: + class InvalidStreamPlugin(WorkerPlugin): + plugin_id = "tests.invalid_stream" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + def stream(name: str, request: Json, next_call: Any) -> Any: + del name, request, next_call + return invalid_stream + + ctx.register_llm_stream_execution_intercept("invalid_stream", stream) + + service = _service(InvalidStreamPlugin(), RecordingHostStub()) + await _register(service) + chunks = [ + chunk + async for chunk in service.InvokeStream( + _invoke_request( + "invalid_stream", + pb.LLM_STREAM_EXECUTION_INTERCEPT, + llm=_llm_payload(request={"content": {"prompt": "hello"}}), + ), + AbortContext(), + ) + ] + assert len(chunks) == 1 + assert "stream callback must return" in chunks[0].error.message + + +async def test_runtime_host_calls_and_scope_context(host_stub: RecordingHostStub) -> None: + runtime = PluginRuntime(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN, host_stub=host_stub) + assert runtime.current_scope_stack_id() is None + assert runtime.current_parent_scope_id() is None + + stack_id = await runtime.create_scope_stack() + assert stack_id == "stack-1" + with runtime.bind_scope_stack(stack_id, parent_scope_id="parent-1"): + assert runtime.current_scope_stack_id() == stack_id + assert runtime.current_parent_scope_id() == "parent-1" + await runtime.emit_mark("mark", {"ok": True}) + await runtime.emit_mark("override-parent", parent_scope_id="parent-2") + scope_id = await runtime.push_scope("scope", scope_type=ScopeType.TOOL, input={"in": True}) + await runtime.pop_scope(scope_id, output={"out": True}) + tool_next = await ToolNext(runtime, "tool-next").call({"value": 1}) + llm_next = await _llm_next(runtime, {"content": {"prompt": "hello"}}) + stream_next = [chunk async for chunk in _llm_stream_next(runtime, {"content": {"prompt": "hello"}})] + with runtime.clear_scope_stack(): + assert runtime.current_scope_stack_id() is None + assert runtime.current_parent_scope_id() is None + assert runtime.current_scope_stack_id() == stack_id + assert runtime.current_scope_stack_id() is None + assert tool_next["next_tool"]["value"] == 1 + assert llm_next["next_llm"]["content"]["prompt"] == "hello" + assert stream_next[0]["next_stream"]["content"]["prompt"] == "hello" + + await runtime.emit_mark("explicit", scope_stack_id="explicit-stack", parent_scope_id="explicit-parent") + await runtime.drop_scope_stack(stack_id) + mark_request = _last_request(host_stub, pb.EmitMarkRequest) + assert mark_request.scope.scope_stack_id == "explicit-stack" + assert mark_request.scope.parent_scope_id == "explicit-parent" + push_request = _last_request(host_stub, pb.PushScopeRequest) + assert push_request.scope.scope_stack_id == "stack-1" + assert push_request.scope.parent_scope_id == "parent-1" + + override_mark = next( + request + for request in host_stub.requests + if isinstance(request, pb.EmitMarkRequest) and request.name == "override-parent" + ) + assert override_mark.scope.scope_stack_id == "stack-1" + assert override_mark.scope.parent_scope_id == "parent-2" + + with pytest.raises(WorkerSdkError, match="parent_scope_id requires"): + await runtime.emit_mark("missing-stack", parent_scope_id="parent") + + +async def test_invocation_scope_context_is_isolated_across_concurrent_requests(host_stub: RecordingHostStub) -> None: + started = 0 + both_started = asyncio.Event() + release = asyncio.Event() + + class ConcurrentScopePlugin(WorkerPlugin): + plugin_id = "tests.concurrent_scope" + + def register(self, ctx: PluginContext, config: Json) -> None: + del config + + async def subscriber(event: Json) -> None: + nonlocal started + started += 1 + if started == 2: + both_started.set() + await release.wait() + await ctx.runtime.emit_mark(event["name"]) + + ctx.register_subscriber("subscriber", subscriber) + + service = _service(ConcurrentScopePlugin(), host_stub) + await _register(service) + + async def invoke(name: str, stack: str, parent: str) -> None: + response = await service.Invoke( + _invoke_request( + "subscriber", + pb.SUBSCRIBER, + event=_json_envelope(EVENT_SCHEMA, {"name": name}), + scope=pb.ScopeContext(scope_stack_id=stack, parent_scope_id=parent), + ), + AbortContext(), + ) + assert response.WhichOneof("result") == "empty" + + tasks = [ + asyncio.create_task(invoke("first", "stack-1", "parent-1")), + asyncio.create_task(invoke("second", "stack-2", "parent-2")), + ] + try: + await asyncio.wait_for(both_started.wait(), timeout=1) + release.set() + await asyncio.gather(*tasks) + finally: + release.set() + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + marks = [ + (request.name, request.scope.scope_stack_id, request.scope.parent_scope_id) + for request in host_stub.requests + if isinstance(request, pb.EmitMarkRequest) + ] + assert len(marks) == 2 + assert sorted(marks) == [ + ("first", "stack-1", "parent-1"), + ("second", "stack-2", "parent-2"), + ] + + +async def test_runtime_host_call_error_paths(host_stub: RecordingHostStub) -> None: + runtime = PluginRuntime(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN, host_stub=host_stub) + + host_stub.failures["EmitMark"] = "error" + with pytest.raises(WorkerSdkError, match="EmitMark failed"): + await runtime.emit_mark("mark") + + host_stub.failures["EmitMark"] = "empty" + with pytest.raises(WorkerSdkError, match="host call failed"): + await runtime.emit_mark("mark") + + host_stub.failures["CreateScopeStack"] = "error" + with pytest.raises(WorkerSdkError, match="CreateScopeStack failed"): + await runtime.create_scope_stack() + + host_stub.failures["PushScope"] = "error" + with pytest.raises(WorkerSdkError, match="PushScope failed"): + await runtime.push_scope("scope") + + host_stub.failures["PopScope"] = "error" + with pytest.raises(WorkerSdkError, match="PopScope failed"): + await runtime.pop_scope("scope") + + host_stub.failures["DropScopeStack"] = "error" + with pytest.raises(WorkerSdkError, match="DropScopeStack failed"): + await runtime.drop_scope_stack("stack") + + host_stub.failures["ToolNext"] = "error" + with pytest.raises(WorkerSdkError, match="ToolNext failed"): + await ToolNext(runtime, "tool-next").call({"value": 1}) + + host_stub.failures["LlmNext"] = "error" + with pytest.raises(WorkerSdkError, match="LlmNext failed"): + await _llm_next(runtime, {"content": {}}) + + host_stub.failures["LlmStreamNext"] = "error" + with pytest.raises(WorkerSdkError, match="LlmStreamNext failed"): + async for _chunk in _llm_stream_next(runtime, {"content": {}}): + pass + + host_stub.failures["LlmStreamNext"] = "empty" + with pytest.raises(WorkerSdkError, match="stream chunk is empty"): + async for _chunk in _llm_stream_next(runtime, {"content": {}}): + pass + + +async def test_lifecycle_acks(service: _WorkerService) -> None: + cancel = await service.CancelInvocation( + pb.CancelInvocationRequest( + activation_id=ACTIVATION_ID, + invocation_id="invoke-1", + auth_token=AUTH_TOKEN, + reason="test", + ), + AbortContext(), + ) + assert not cancel.accepted + assert "not implemented" in cancel.message + + shutdown = await service.Shutdown( + pb.ShutdownRequest(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN, reason="test"), + AbortContext(), + ) + assert shutdown.accepted + assert "shutdown accepted" in shutdown.message + + +def test_required_environment_reports_missing_value(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("NEMO_RELAY_WORKER_SOCKET", raising=False) + with pytest.raises(Exception, match="NEMO_RELAY_WORKER_SOCKET"): + _required_env("NEMO_RELAY_WORKER_SOCKET") + + +def test_endpoint_helpers_normalize_announce_and_unlink_tcp_and_unix_targets(tmp_path: Any) -> None: + assert _grpc_target("tcp://127.0.0.1:50051") == "127.0.0.1:50051" + assert _grpc_target("http://127.0.0.1:50051") == "127.0.0.1:50051" + assert _grpc_target("unix:///tmp/worker.sock") == "unix:/tmp/worker.sock" + assert _announced_worker_endpoint("tcp://127.0.0.1:0", 43123) == "http://127.0.0.1:43123" + assert _announced_worker_endpoint("http://127.0.0.1:50051", 43123) == "http://127.0.0.1:50051" + assert _announced_worker_endpoint("127.0.0.1:50051", 43123) == "http://127.0.0.1:50051" + assert _announced_worker_endpoint("unix:///tmp/worker.sock", 43123) == "unix:///tmp/worker.sock" + assert _announced_worker_endpoint("worker-endpoint", 43123) == "http://worker-endpoint" + + socket_path = tmp_path / "worker.sock" + socket_path.write_text("", encoding="utf-8") + _unlink_unix_socket(f"unix://{socket_path}") + assert not socket_path.exists() + + +async def test_serve_plugin_announces_endpoint_only_after_server_start( + tmp_path: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + endpoint_file = tmp_path / "endpoint.txt" + start_entered = asyncio.Event() + allow_start = asyncio.Event() + + class FakeChannel: + async def close(self) -> None: + pass + + class FakeServer: + def add_insecure_port(self, target: str) -> int: + assert target == "127.0.0.1:0" + return 43123 + + async def start(self) -> None: + start_entered.set() + await allow_start.wait() + + async def stop(self, grace: int) -> None: + assert grace == 2 + + fake_server = FakeServer() + monkeypatch.setattr(plugin_api.grpc.aio, "insecure_channel", lambda target: FakeChannel()) + monkeypatch.setattr(plugin_api.grpc.aio, "server", lambda: fake_server) + monkeypatch.setattr(plugin_api.pb_grpc, "RelayHostRuntimeStub", lambda channel: object()) + monkeypatch.setattr(plugin_api.pb_grpc, "add_PluginWorkerServicer_to_server", lambda service, server: None) + monkeypatch.setenv("NEMO_RELAY_WORKER_SOCKET", "tcp://127.0.0.1:0") + monkeypatch.setenv("NEMO_RELAY_HOST_SOCKET", "http://127.0.0.1:9") + monkeypatch.setenv("NEMO_RELAY_WORKER_ID", ACTIVATION_ID) + monkeypatch.setenv("NEMO_RELAY_WORKER_TOKEN", AUTH_TOKEN) + monkeypatch.setenv("NEMO_RELAY_WORKER_ENDPOINT_FILE", str(endpoint_file)) + + task = asyncio.create_task(serve_plugin(AllSurfacesPlugin())) + try: + await asyncio.wait_for(start_entered.wait(), timeout=1) + assert not endpoint_file.exists() + allow_start.set() + for _ in range(100): + if endpoint_file.exists(): + break + await asyncio.sleep(0.01) + assert endpoint_file.read_text(encoding="utf-8") == "http://127.0.0.1:43123" + finally: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +async def test_serve_plugin_closes_resources_when_server_start_fails(monkeypatch: pytest.MonkeyPatch) -> None: + class FakeChannel: + closed = False + + async def close(self) -> None: + self.closed = True + + class FakeServer: + stopped = False + + def add_insecure_port(self, target: str) -> int: + assert target == "127.0.0.1:0" + return 43123 + + async def start(self) -> None: + raise RuntimeError("start failed") + + async def stop(self, grace: int) -> None: + assert grace == 2 + self.stopped = True + + fake_channel = FakeChannel() + fake_server = FakeServer() + monkeypatch.setattr(plugin_api.grpc.aio, "insecure_channel", lambda target: fake_channel) + monkeypatch.setattr(plugin_api.grpc.aio, "server", lambda: fake_server) + monkeypatch.setattr(plugin_api.pb_grpc, "RelayHostRuntimeStub", lambda channel: object()) + monkeypatch.setattr(plugin_api.pb_grpc, "add_PluginWorkerServicer_to_server", lambda service, server: None) + monkeypatch.setenv("NEMO_RELAY_WORKER_SOCKET", "tcp://127.0.0.1:0") + monkeypatch.setenv("NEMO_RELAY_HOST_SOCKET", "http://127.0.0.1:9") + monkeypatch.setenv("NEMO_RELAY_WORKER_ID", ACTIVATION_ID) + monkeypatch.setenv("NEMO_RELAY_WORKER_TOKEN", AUTH_TOKEN) + monkeypatch.delenv("NEMO_RELAY_WORKER_ENDPOINT_FILE", raising=False) + + with pytest.raises(RuntimeError, match="start failed"): + await serve_plugin(AllSurfacesPlugin()) + assert fake_server.stopped + assert fake_channel.closed + + +async def test_serve_plugin_announces_tcp_endpoint_and_accepts_health_shutdown( + tmp_path: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + endpoint_file = tmp_path / "endpoint.txt" + monkeypatch.setenv("NEMO_RELAY_WORKER_SOCKET", "tcp://127.0.0.1:0") + monkeypatch.setenv("NEMO_RELAY_HOST_SOCKET", "http://127.0.0.1:9") + monkeypatch.setenv("NEMO_RELAY_WORKER_ID", ACTIVATION_ID) + monkeypatch.setenv("NEMO_RELAY_WORKER_TOKEN", AUTH_TOKEN) + monkeypatch.setenv("NEMO_RELAY_WORKER_ENDPOINT_FILE", str(endpoint_file)) + + task = asyncio.create_task(serve_plugin(AllSurfacesPlugin())) + channel = None + try: + for _ in range(100): + if endpoint_file.exists(): + break + await asyncio.sleep(0.05) + assert endpoint_file.exists() + endpoint = endpoint_file.read_text(encoding="utf-8") + assert endpoint.startswith("http://127.0.0.1:") + + channel = grpc.aio.insecure_channel(_grpc_target(endpoint)) + stub = pb_grpc.PluginWorkerStub(channel) + health = await stub.Health(pb.HealthRequest(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN), timeout=5) + assert health.ok + shutdown = await stub.Shutdown( + pb.ShutdownRequest(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN, reason="test"), + timeout=5, + ) + assert shutdown.accepted + await asyncio.wait_for(task, timeout=5) + finally: + if channel is not None: + await channel.close() + if not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + +def _service(plugin: WorkerPlugin, host_stub: RecordingHostStub) -> _WorkerService: + runtime = PluginRuntime(activation_id=ACTIVATION_ID, auth_token=AUTH_TOKEN, host_stub=host_stub) + return _WorkerService(plugin, runtime, asyncio.Event()) + + +def _worker_error(message: str) -> Any: + return pb.WorkerError(code="test.error", message=message, retryable=False) + + +def _tag(value: Json, tag: str) -> Json: + if isinstance(value, dict): + return {**value, "tag": tag} + return {"value": value, "tag": tag} + + +def _tag_llm_request(request: Json, tag: str) -> Json: + request = dict(request) + content = request.get("content") + if isinstance(content, dict): + request["content"] = {**content, tag: True} + else: + request["content"] = {"value": content, tag: True} + return request + + +def _handshake_request() -> Any: + return pb.HandshakeRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + relay_version="0.5.0", + worker_protocol=WORKER_PROTOCOL, + auth_token=AUTH_TOKEN, + host_endpoint="http://127.0.0.1:9", + ) + + +def _validate_request(config: Json | None = None) -> Any: + return pb.ValidateRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + auth_token=AUTH_TOKEN, + config=_json_envelope(JSON_SCHEMA, {} if config is None else config), + ) + + +def _register_request(config: Json | None = None) -> Any: + return pb.RegisterRequest( + activation_id=ACTIVATION_ID, + plugin_id="tests.python_worker", + auth_token=AUTH_TOKEN, + config=_json_envelope(JSON_SCHEMA, {} if config is None else config), + ) + + +async def _register(service: _WorkerService) -> Any: + response = await service.Register(_register_request(), AbortContext()) + assert not response.HasField("error"), response.error + return response + + +def _invoke_request(registration_name: str, surface: int, **kwargs: Any) -> Any: + return pb.InvokeRequest( + activation_id=ACTIVATION_ID, + invocation_id="invoke-1", + registration_name=registration_name, + surface=surface, + continuation_id="next-1", + auth_token=AUTH_TOKEN, + **kwargs, + ) + + +def _tool_request(registration_name: str, surface: int, value: Json) -> Any: + return _invoke_request( + registration_name, + surface, + tool=pb.ToolInvocation(tool_name="lookup", value=_json_envelope(JSON_SCHEMA, value)), + ) + + +def _llm_payload( + *, + model_name: str = "model", + request: Json | None = None, + response: Json | None = None, + annotated: Json | None = None, +) -> Any: + kwargs: dict[str, Any] = { + "model_name": model_name, + "request": _json_envelope(LLM_REQUEST_SCHEMA, request or {"content": {}}), + "response": _json_envelope(JSON_SCHEMA, response or {}), + } + if annotated is not None: + kwargs["annotated_request"] = _json_envelope(ANNOTATED_LLM_REQUEST_SCHEMA, annotated) + return pb.LlmInvocation(**kwargs) + + +async def _invoke_json_async( + service: _WorkerService, + registration_name: str, + surface: int, + *, + payload: Any | None = None, +) -> Json: + if payload is None: + request = _tool_request(registration_name, surface, {"query": "relay"}) + else: + request = _invoke_request(registration_name, surface, llm=payload) + response = await service.Invoke(request, AbortContext()) + assert response.WhichOneof("result") == "json", response + return _envelope_value(response.json.value) + + +def _envelope_value(envelope: Any) -> Json: + return json.loads(envelope.json.decode("utf-8")) + + +def _stream_value(chunk: Any) -> Json: + assert chunk.WhichOneof("item") == "value", chunk + return _envelope_value(chunk.value) + + +def _last_request(host_stub: RecordingHostStub, request_type: Any) -> Any: + return next(request for request in reversed(host_stub.requests) if isinstance(request, request_type)) + + +async def _llm_next(runtime: PluginRuntime, request: Json) -> Json: + from nemo_relay_plugin import LlmNext + + return await LlmNext(runtime, "llm-next").call(request) + + +def _llm_stream_next(runtime: PluginRuntime, request: Json) -> AsyncIterator[Json]: + from nemo_relay_plugin import LlmStreamNext + + return LlmStreamNext(runtime, "llm-stream-next").call(request) + + +def _all_expected_surfaces() -> list[int]: + return [ + pb.SUBSCRIBER, + pb.TOOL_SANITIZE_REQUEST_GUARDRAIL, + pb.TOOL_SANITIZE_RESPONSE_GUARDRAIL, + pb.TOOL_CONDITIONAL_EXECUTION_GUARDRAIL, + pb.TOOL_REQUEST_INTERCEPT, + pb.TOOL_EXECUTION_INTERCEPT, + pb.LLM_SANITIZE_REQUEST_GUARDRAIL, + pb.LLM_SANITIZE_RESPONSE_GUARDRAIL, + pb.LLM_CONDITIONAL_EXECUTION_GUARDRAIL, + pb.LLM_REQUEST_INTERCEPT, + pb.LLM_EXECUTION_INTERCEPT, + pb.LLM_STREAM_EXECUTION_INTERCEPT, + ] diff --git a/uv.lock b/uv.lock index 075306d2d..82ef45391 100644 --- a/uv.lock +++ b/uv.lock @@ -6,6 +6,12 @@ resolution-markers = [ "sys_platform != 'linux'", ] +[manifest] +members = [ + "nemo-relay", + "nemo-relay-plugin", +] + [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -784,6 +790,57 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2d/b6/552d40e96da22921eb1fead7c14b00b5b5473a20e45959488660fab35ee2/google_genai-1.75.0-py3-none-any.whl", hash = "sha256:8dc4c096e7d6288c3087f6893f582fe52468932464781edb8193bd92b9fefb2c", size = 793726, upload-time = "2026-05-04T22:48:53.033Z" }, ] +[[package]] +name = "grpcio" +version = "1.81.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b0/b5/1ff353970a87eda4c98251e34d2dfd214abd4982dc89119c9252a2a482d2/grpcio-1.81.1.tar.gz", hash = "sha256:6fa10a767143a5e82e8eaab53918af0cd8909a57a27f8cb2288b80a613ac671b", size = 13026582, upload-time = "2026-06-11T12:46:51.673Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/52/ea/1c2fa386b718ff493225e61cfc052ef400b4d6ffc54cbe261026432624b5/grpcio-1.81.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:d71d30f2d92f67d944631c523713934fee37292469e182ebcd2c1dd8a64ce53f", size = 6093112, upload-time = "2026-06-11T12:44:52.131Z" }, + { url = "https://files.pythonhosted.org/packages/2b/18/acf45fa8bd1bc5d7b0c2fd3dc4c209379fbd5bb396b440b68a83342226b7/grpcio-1.81.1-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:b137f4bf3ada9dc44d411478decc6ff09a79ed30b306cd2abaa98408c3588137", size = 12074277, upload-time = "2026-06-11T12:44:55.354Z" }, + { url = "https://files.pythonhosted.org/packages/48/d7/ee86a60699b7db039f772a2c4a7e4facc7138984ff42c0130933a0063884/grpcio-1.81.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:a3acb384427816dd5d470f47e62137b87f74da694faa8a50147012cf40df276a", size = 6640348, upload-time = "2026-06-11T12:44:59.223Z" }, + { url = "https://files.pythonhosted.org/packages/26/ee/d2de5e47378ffc207d476c230fea3be4d2601edbce9995f4fe45535d4896/grpcio-1.81.1-cp311-cp311-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:f9a0ebbe45c29b5e5866593c12b78bd9035f0f0f0d4bc8361680cd580d99db49", size = 7331842, upload-time = "2026-06-11T12:45:02.001Z" }, + { url = "https://files.pythonhosted.org/packages/23/d6/abeda5c2b896a0b341584fe5ac411bbf72e197a9a374c355fb90965e08d2/grpcio-1.81.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0a37165cc80b1a368384b383e63a4c38116a10467ae44c904d2d7468c4470ec2", size = 6842229, upload-time = "2026-06-11T12:45:04.76Z" }, + { url = "https://files.pythonhosted.org/packages/10/1c/1f0da7d590b4aeee006826ba568d0e419ca14b23e18f901a3da3e9fba613/grpcio-1.81.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:6282caffb41ec326d4cb67ca9cf53b739d1b2f975a2acb498c7418e9f7d9a416", size = 7446096, upload-time = "2026-06-11T12:45:07.499Z" }, + { url = "https://files.pythonhosted.org/packages/6a/81/5c505d508f7c887aa7982d21443a4126597c80d34b0bcf40f9cec576d7f3/grpcio-1.81.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:a35009284d0d3d5c2c9601c164a911b8b4331608d98a9a66d47d97bb2f522b70", size = 8445238, upload-time = "2026-06-11T12:45:10.243Z" }, + { url = "https://files.pythonhosted.org/packages/f7/b2/524847365122ee509ca17bcc4e092198b700e94af7bfd5bb5e6dd9f3ee66/grpcio-1.81.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1b22c80559854b789a01fd89e8929b3798a156c0829b5282a8939f33ad4115ad", size = 7873989, upload-time = "2026-06-11T12:45:13.102Z" }, + { url = "https://files.pythonhosted.org/packages/18/fa/07c037c50b006909d1d13a5848774f8aa7b242f70dc03a035c64eea0e6db/grpcio-1.81.1-cp311-cp311-win32.whl", hash = "sha256:428bec0161b48d8cf583c068591bc0016d0d9cfff52462b72b3884861ea768c5", size = 4202223, upload-time = "2026-06-11T12:45:16.166Z" }, + { url = "https://files.pythonhosted.org/packages/41/ed/6bff15376920942fac6b95b9802752b837437172c9e8fc2d3170546b89cc/grpcio-1.81.1-cp311-cp311-win_amd64.whl", hash = "sha256:30e825f6848d9f18bba350ed6c75c1b02a0b5184474a31db9a32b1fa66fd8c79", size = 4941303, upload-time = "2026-06-11T12:45:18.724Z" }, + { url = "https://files.pythonhosted.org/packages/85/07/9a979c81738863a738dc23d65177056e71fbb2db817740ed870b33434e7a/grpcio-1.81.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:8b39472beafc0bdcafc4c8c73ad082ebfdb449d566897a61e7acb4fa88089115", size = 6053264, upload-time = "2026-06-11T12:45:21.017Z" }, + { url = "https://files.pythonhosted.org/packages/75/95/539706ca0d3bd40dbad583dc56fd883da941f37556b629132da5762781b9/grpcio-1.81.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:12b7524c88d4026d3dcb7b0ebe16b6714f3b4af402ddd0f0639ab064a00c87c3", size = 12052560, upload-time = "2026-06-11T12:45:23.652Z" }, + { url = "https://files.pythonhosted.org/packages/e0/44/f257b7e0bd69c93b06c6cb8ac8d1b901ccb42bedabd83c1a4c77a71f8810/grpcio-1.81.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1e123f9b37edb8375fd74130d1f69c944bbf0a7b06761ae7211154b8759e94d2", size = 6595983, upload-time = "2026-06-11T12:45:26.963Z" }, + { url = "https://files.pythonhosted.org/packages/b9/f3/19782aa04c960968bef8c5539329d8e3bbc3364e2e46d19eb5e5cc5e43b7/grpcio-1.81.1-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:2c2e2ae6867c2966b8daccc836d54a13218e0007e9a490aeb81dd05be64d22d7", size = 7303455, upload-time = "2026-06-11T12:45:29.707Z" }, + { url = "https://files.pythonhosted.org/packages/eb/8c/dea020b6d91508cd84463917a63149ec196ee7db505d032ae43fcb3303b9/grpcio-1.81.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:766bc7c9a9c340342f4c864ccbda8e78111e4751f13b895812b9c148fb79e9d0", size = 6809167, upload-time = "2026-06-11T12:45:32.52Z" }, + { url = "https://files.pythonhosted.org/packages/1c/c7/3030dd940408083bd32cd95d634777a71605ade4887154d93e8a89244946/grpcio-1.81.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b259a04a737cb3496be0901328eb8b7552ed8df4865d8c8f1cf1bffcfc0776a3", size = 7412536, upload-time = "2026-06-11T12:45:35.403Z" }, + { url = "https://files.pythonhosted.org/packages/e0/dd/1172a9e42b168edcafefad6115346ef619a3fc02158bb170e66ced24bcdd/grpcio-1.81.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:85b10a45b8993d195c4f3ff57025b8d1e11834909ee475c403bfa60cb4caefaf", size = 8408276, upload-time = "2026-06-11T12:45:37.78Z" }, + { url = "https://files.pythonhosted.org/packages/25/7a/71437c7f3596e5246155c515852795a85a1a8d228190212432b13b97a95d/grpcio-1.81.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8ea1936c26b99999b27479853039a7f34713f56c49375ad52b38535ec93a796c", size = 7849660, upload-time = "2026-06-11T12:45:40.627Z" }, + { url = "https://files.pythonhosted.org/packages/65/40/7debc0da45d2efebafb82da75644be347497fe4ee250514b8cd3b86ae8bf/grpcio-1.81.1-cp312-cp312-win32.whl", hash = "sha256:a185a04039df6cae8648bc8ab6d6fde7bf94f7188ecf7828e76ac52eef1e41d6", size = 4185819, upload-time = "2026-06-11T12:45:43.027Z" }, + { url = "https://files.pythonhosted.org/packages/2e/b9/8fe3ba5ed462067774ebc1f9c7f26aa7ebcc280ddd476be107153de1339e/grpcio-1.81.1-cp312-cp312-win_amd64.whl", hash = "sha256:3ad74f8bb1a18963914c5452d289422830b39459e8776ebbcd207be1fbfb1d94", size = 4930461, upload-time = "2026-06-11T12:45:45.775Z" }, + { url = "https://files.pythonhosted.org/packages/7a/42/dcc2e4b600538ef18327c0839d56b7d3c3812337c5d710df5877dbb39b1e/grpcio-1.81.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:b10e1ff4756ed27d5a29d7fc79cfce7ef1ff56ad20025b89bac7cf79e09abbbe", size = 6054466, upload-time = "2026-06-11T12:45:48.43Z" }, + { url = "https://files.pythonhosted.org/packages/7b/4a/a36e03210183a8a7d4c80c3936acee679f4bd77d5861f369db47b2cc5f05/grpcio-1.81.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:819edbdcb42ab8598b494bcf0222684bbb7a3c772bd1b1f0be7e029a6063c28e", size = 12048795, upload-time = "2026-06-11T12:45:54.011Z" }, + { url = "https://files.pythonhosted.org/packages/b0/d5/d68e30b29098f63beab6fe501100fe82674ff142b32c672532da86a99b3a/grpcio-1.81.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c5bf2dc311127d91230cc79b92188c082634a06cf66c5234db49a43b910183b0", size = 6599094, upload-time = "2026-06-11T12:45:57.799Z" }, + { url = "https://files.pythonhosted.org/packages/3d/b3/e837954d279754f638a11cca5dcf6b24a005efb398984cefaf7735945a54/grpcio-1.81.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:e8ca6a1fcdb2943c9cbc1804a1baf3acb6071d72a471591678ded84218006e14", size = 7307182, upload-time = "2026-06-11T12:46:00.568Z" }, + { url = "https://files.pythonhosted.org/packages/0d/1e/b47957057e729adc6cdf519a47f8be2562b7140e280f1418443eb4022192/grpcio-1.81.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e64dd101d380a115cc5a0c7856788adb535f1a4e21fc543775602f8be95180ae", size = 6810962, upload-time = "2026-06-11T12:46:03.312Z" }, + { url = "https://files.pythonhosted.org/packages/40/26/569868e364e05b19ec8f969da53d230bcd89c962cd198f7c29943155c4d3/grpcio-1.81.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:98a07f9bf591e3a8919797bee1c53f026ba4acd587e5a4404c8e57c9ec36b2a5", size = 7415698, upload-time = "2026-06-11T12:46:06.005Z" }, + { url = "https://files.pythonhosted.org/packages/36/0c/5440a0582cb5653fc42a6e262eeb22700943313f8076f9dc927491b20a59/grpcio-1.81.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:c261d74b1a945cf895a9d6eccd1685a8e837531beaab782da4d630a8d12deffb", size = 8407779, upload-time = "2026-06-11T12:46:08.84Z" }, + { url = "https://files.pythonhosted.org/packages/ff/aa/66fe9f39871d766987d869a03ee0842a026f499c7b1e62decb9e78a8088e/grpcio-1.81.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:58ad1131c300d3c9b933802b3cc4dc69d380822935ba50b28703156ea826fbf7", size = 7844521, upload-time = "2026-06-11T12:46:12.171Z" }, + { url = "https://files.pythonhosted.org/packages/f0/9e/69bb7194861bcd28fb3193261d4f9c3831b4446993f002cf59068943e7ab/grpcio-1.81.1-cp313-cp313-win32.whl", hash = "sha256:78e29211f26da2fdd0e9c6d2b79f489476140cf7029b6a64808ade7ca4156a42", size = 4182786, upload-time = "2026-06-11T12:46:15.192Z" }, + { url = "https://files.pythonhosted.org/packages/0d/20/3da8bb0d637feccdc3e1e419bb511ce93651ce7d54164f95de22cc0b8b34/grpcio-1.81.1-cp313-cp313-win_amd64.whl", hash = "sha256:edb59506291b647a30884b1d51a599d605f40b20af4a7dc3d33786a47a31de60", size = 4928648, upload-time = "2026-06-11T12:46:17.823Z" }, + { url = "https://files.pythonhosted.org/packages/b6/58/19414622b1bf6981bc9c05a365bd548e71876c89000083b3af489251e9c0/grpcio-1.81.1-cp314-cp314-linux_armv7l.whl", hash = "sha256:506f48f2f9c29b143fca3dad7b0d518c188b6c9648c75a2ae6e2d9f2c13a060b", size = 6055336, upload-time = "2026-06-11T12:46:20.557Z" }, + { url = "https://files.pythonhosted.org/packages/32/f1/2ec88adb92b0eba970dd0e0e7dd086341daa3c75eba4f735f9e44bf684b0/grpcio-1.81.1-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:d865db4a6318e1c1bea83292e0ed231090538fc4ca45425b0f0480eb338bbc6e", size = 12056279, upload-time = "2026-06-11T12:46:24.255Z" }, + { url = "https://files.pythonhosted.org/packages/41/36/e8c5f8c6ec71de73733695ebc809e98b178b534ec6d8eaa31a7ebab4ad4c/grpcio-1.81.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e2aa72e3ce1770317ef534f63d397b55e130725f5149bd36077c3b539019db27", size = 6608225, upload-time = "2026-06-11T12:46:27.601Z" }, + { url = "https://files.pythonhosted.org/packages/30/22/96fc577a845ab093326d9ab1adb874bd4936c8cf98ac8ed2f3db13a0a2fb/grpcio-1.81.1-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:0490c30c261eded63f3f354979f9dc4502a9fb944cccb60cd9dc85f5a7349854", size = 7306576, upload-time = "2026-06-11T12:46:30.514Z" }, + { url = "https://files.pythonhosted.org/packages/76/7b/61dab5d5969f28d97fb1009cead1df0a5cd987d3315e1b37f18a4449f8bc/grpcio-1.81.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:410482da976329fe5f4067270401b12cf2bd552ff8020f054ecfaddb5475f9d6", size = 6812165, upload-time = "2026-06-11T12:46:33.699Z" }, + { url = "https://files.pythonhosted.org/packages/82/78/6e501929d4f5f96462fd82fd9f0f06e5f9612207582b862868d68757b27d/grpcio-1.81.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:e3657301562ac3cb8018d30d0d3ebfa39932239f7b5703422057ef14b69949f5", size = 7422962, upload-time = "2026-06-11T12:46:36.511Z" }, + { url = "https://files.pythonhosted.org/packages/2a/7e/f2157589e66daa78ebb3165942d05a08bdea93b9d11c2bc1e172aef89685/grpcio-1.81.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:24c8e57504c8f45b237e40b99262d181071e5099a07053695b75d97bb53053a0", size = 8408176, upload-time = "2026-06-11T12:46:39.803Z" }, + { url = "https://files.pythonhosted.org/packages/da/df/c6717fef716e00d235ffb96123baf6dce76d6004f6233fa767c502861460/grpcio-1.81.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b427c19380991a4eaab2f6144b64b99b412043314c6bf4ab544f97bb31ee4190", size = 7846681, upload-time = "2026-06-11T12:46:43.013Z" }, + { url = "https://files.pythonhosted.org/packages/36/84/3502e9f210a6a5c4438c8aca3f88edd2e04f6a27f3d41b26cf0a0024b096/grpcio-1.81.1-cp314-cp314-win32.whl", hash = "sha256:61233fe8951e5c85dff81c2458b6528624760166946b5b47ea150a589168411f", size = 4264615, upload-time = "2026-06-11T12:46:45.741Z" }, + { url = "https://files.pythonhosted.org/packages/ff/b0/4af731ff7492c68a96e4c71bfd0f4590acde92b31c6fe4894e6465c10ff6/grpcio-1.81.1-cp314-cp314-win_amd64.whl", hash = "sha256:3768a5ff1b2125e6f552e561b6b2dca0e64982d8949689b4df145cf8b98d7821", size = 5070275, upload-time = "2026-06-11T12:46:48.486Z" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -1357,6 +1414,8 @@ dev = [ { name = "uv" }, ] test = [ + { name = "grpcio", marker = "(platform_machine != 'ARM64' and platform_machine != 'aarch64' and platform_machine != 'arm64') or sys_platform != 'win32'" }, + { name = "nemo-relay-plugin", marker = "(platform_machine != 'ARM64' and platform_machine != 'aarch64' and platform_machine != 'arm64') or sys_platform != 'win32'" }, { name = "opentelemetry-proto" }, { name = "pydantic" }, { name = "pytest" }, @@ -1391,6 +1450,8 @@ dev = [ { name = "uv", specifier = "~=0.10.0" }, ] test = [ + { name = "grpcio", marker = "(platform_machine != 'ARM64' and platform_machine != 'aarch64' and platform_machine != 'arm64') or sys_platform != 'win32'", specifier = ">=1.81.1,<2" }, + { name = "nemo-relay-plugin", marker = "(platform_machine != 'ARM64' and platform_machine != 'aarch64' and platform_machine != 'arm64') or sys_platform != 'win32'", editable = "python/plugin" }, { name = "opentelemetry-proto", specifier = ">=1.39,<2" }, { name = "pydantic", specifier = ">=2" }, { name = "pytest", specifier = ">=8" }, @@ -1398,6 +1459,21 @@ test = [ { name = "pytest-cov", specifier = "~=7.0" }, ] +[[package]] +name = "nemo-relay-plugin" +version = "0.5.0" +source = { editable = "python/plugin" } +dependencies = [ + { name = "grpcio" }, + { name = "protobuf" }, +] + +[package.metadata] +requires-dist = [ + { name = "grpcio", specifier = ">=1.81.1,<2" }, + { name = "protobuf", specifier = ">=6.33.5,<7" }, +] + [[package]] name = "nodeenv" version = "1.10.0" @@ -1736,17 +1812,17 @@ wheels = [ [[package]] name = "protobuf" -version = "7.35.1" +version = "6.33.6" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/da/01/9ef0afd7999eb9badb3a768b4aedd78c86d4c65cfaf1958ab276199e76b4/protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a", size = 458717, upload-time = "2026-06-11T21:55:40.257Z" } +sdist = { url = "https://files.pythonhosted.org/packages/66/70/e908e9c5e52ef7c3a6c7902c9dfbb34c7e29c25d2f81ade3856445fd5c94/protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135", size = 444531, upload-time = "2026-03-18T19:05:00.988Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/10/03/8aeeb7458d22546bf64b5250ca1daeb5ff757d900e8e4a7476c6f0db843e/protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6", size = 433226, upload-time = "2026-06-11T21:55:31.719Z" }, - { url = "https://files.pythonhosted.org/packages/37/4b/dfb89eb0e652a1ff073c39a59fb5e3a83cfe9b57a2c83fa6d78270101767/protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799", size = 328847, upload-time = "2026-06-11T21:55:34.035Z" }, - { url = "https://files.pythonhosted.org/packages/0f/58/dc12f2cd484951524af6e3382c785869b9b3fb5e52ee95ae23add53ee8f9/protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4", size = 344030, upload-time = "2026-06-11T21:55:34.941Z" }, - { url = "https://files.pythonhosted.org/packages/e4/be/5b3cfe508bfab6761414ff944e3366eb13be4fd71efcd69450f89ba39f43/protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4", size = 327130, upload-time = "2026-06-11T21:55:35.921Z" }, - { url = "https://files.pythonhosted.org/packages/d8/bc/6d6c7ba8709c85f8f2c390b2b118d6fb08a783676a572271851bf45a7d22/protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30", size = 428945, upload-time = "2026-06-11T21:55:37.034Z" }, - { url = "https://files.pythonhosted.org/packages/0a/19/8d0cb6f20a1ef7b18f1c8986ad5783f22f84cce39c6ce9a6e645ea55192e/protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87", size = 439996, upload-time = "2026-06-11T21:55:38.123Z" }, - { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" }, + { url = "https://files.pythonhosted.org/packages/fc/9f/2f509339e89cfa6f6a4c4ff50438db9ca488dec341f7e454adad60150b00/protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3", size = 425739, upload-time = "2026-03-18T19:04:48.373Z" }, + { url = "https://files.pythonhosted.org/packages/76/5d/683efcd4798e0030c1bab27374fd13a89f7c2515fb1f3123efdfaa5eab57/protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326", size = 437089, upload-time = "2026-03-18T19:04:50.381Z" }, + { url = "https://files.pythonhosted.org/packages/5c/01/a3c3ed5cd186f39e7880f8303cc51385a198a81469d53d0fdecf1f64d929/protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a", size = 427737, upload-time = "2026-03-18T19:04:51.866Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/b3c01fdec7d2f627b3a6884243ba328c1217ed2d978def5c12dc50d328a3/protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2", size = 324610, upload-time = "2026-03-18T19:04:53.096Z" }, + { url = "https://files.pythonhosted.org/packages/9b/ca/25afc144934014700c52e05103c2421997482d561f3101ff352e1292fb81/protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3", size = 339381, upload-time = "2026-03-18T19:04:54.616Z" }, + { url = "https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593", size = 323436, upload-time = "2026-03-18T19:04:55.768Z" }, + { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, ] [[package]]